From 6f6f7c3efc55c19afdd0021c87c9bf3b3da09707 Mon Sep 17 00:00:00 2001 From: tbphp Date: Mon, 10 Aug 2026 08:22:15 +0800 Subject: [PATCH] =?UTF-8?q?refactor(channel):=20=E5=88=87=E6=8D=A2?= =?UTF-8?q?=E4=B8=BA=20Bifrost=20=E6=B8=A0=E9=81=93=E6=89=A7=E8=A1=8C?= =?UTF-8?q?=E6=9E=B6=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 2 +- .github/scripts/release-docker-smoke.sh | 43 +- .github/workflows/ci.yml | 13 +- .github/workflows/release.yml | 135 +- Dockerfile | 2 + LICENSES/Apache-2.0.txt | 201 + README.md | 50 +- README_CN.md | 52 +- README_JP.md | 52 +- SECURITY.md | 6 +- THIRD_PARTY_NOTICES.md | 19 + go.mod | 55 +- go.sum | 145 +- internal/app/app.go | 14 + internal/app/app_test.go | 25 +- internal/app/runtime_checkpoint.go | 12 +- internal/app/runtime_checkpoint_test.go | 32 +- internal/catalog/curated.go | 149 - internal/catalog/curated_test.go | 251 - internal/catalog/official.go | 247 - internal/catalog/official_test.go | 174 - internal/catalog/parse.go | 12 +- internal/catalog/parse_test.go | 3 +- internal/catalog/runtime.go | 34 - internal/catalog/runtime_test.go | 42 +- internal/catalog/types.go | 15 +- internal/channel/channel.go | 657 +++ internal/channel/presets.go | 433 ++ internal/channel/registry_test.go | 506 ++ internal/container/container.go | 27 +- internal/container/container_test.go | 31 +- internal/container/route_contract_test.go | 4 +- internal/control/access_key_idempotency.go | 2 +- internal/control/access_keys_test.go | 15 +- internal/control/autoweight.go | 2 +- internal/control/autoweight_test.go | 14 +- .../control/catalog_discovery_merge_test.go | 65 +- internal/control/catalog_sync.go | 21 +- internal/control/catalog_sync_test.go | 148 +- .../control/channel_group_foundation_test.go | 451 ++ internal/control/channels.go | 69 + .../control/credential_api_foundation_test.go | 403 ++ .../control/credential_collection_helpers.go | 232 + internal/control/credential_mutations.go | 600 +++ internal/control/credentials.go | 373 ++ internal/control/discover.go | 55 +- internal/control/discover_executor.go | 268 +- internal/control/discover_executor_test.go | 527 ++- internal/control/discover_group.go | 83 +- internal/control/discover_group_test.go | 147 +- internal/control/discover_test.go | 145 +- .../control/execution_test_helpers_test.go | 144 + .../external_database_integration_test.go | 11 +- internal/control/group_collection.go | 251 +- internal/control/group_collection_http.go | 25 +- .../control/group_collection_http_test.go | 101 +- internal/control/group_collection_query.go | 48 +- .../control/group_collection_query_test.go | 39 +- internal/control/group_collection_test.go | 716 --- .../group_collection_test_helpers_test.go | 126 + internal/control/group_create.go | 174 +- internal/control/group_create_test.go | 809 ---- internal/control/group_delete.go | 22 +- internal/control/group_delete_test.go | 32 +- internal/control/group_detail.go | 35 +- .../group_detail_ledger_routes_test.go | 32 +- internal/control/group_detail_test.go | 23 +- internal/control/group_idempotency.go | 175 +- .../control/group_idempotency_phase1_test.go | 83 +- internal/control/group_key_collection.go | 357 -- internal/control/group_key_collection_test.go | 300 -- internal/control/group_key_mutations.go | 371 -- internal/control/group_key_mutations_test.go | 871 ---- internal/control/group_keys.go | 63 - internal/control/group_keys_test.go | 316 -- internal/control/group_models.go | 14 +- internal/control/group_models_test.go | 75 +- internal/control/group_options.go | 55 +- internal/control/group_options_test.go | 58 +- internal/control/group_settings.go | 204 +- internal/control/group_settings_test.go | 129 +- internal/control/group_update.go | 39 +- internal/control/group_update_test.go | 62 +- internal/control/group_write.go | 199 +- internal/control/health.go | 101 +- internal/control/health_test.go | 199 +- internal/control/home.go | 116 +- internal/control/home_test.go | 87 +- internal/control/http_routes.go | 91 +- internal/control/idempotency_digest.go | 11 +- internal/control/idempotency_digest_test.go | 17 +- internal/control/idempotency_operation.go | 6 +- .../control/idempotency_operation_test.go | 2 +- internal/control/model_candidate.go | 41 + internal/control/model_price.go | 39 +- internal/control/model_price_sync_http.go | 35 + .../control/model_price_sync_http_test.go | 67 + internal/control/model_price_test.go | 111 +- internal/control/model_prices_audit_test.go | 2 + internal/control/model_prices_http_test.go | 33 +- internal/control/mutation_audit.go | 8 +- .../control/mutation_audit_routes_test.go | 81 +- internal/control/mutation_audit_test.go | 16 +- internal/control/operation_error.go | 15 +- internal/control/operation_error_test.go | 20 +- internal/control/operation_recovery.go | 5 +- internal/control/operation_recovery_test.go | 34 +- internal/control/price_identity.go | 23 +- internal/control/price_reconcile.go | 26 +- internal/control/price_reconcile_test.go | 45 +- internal/control/price_reference.go | 95 +- internal/control/price_reference_test.go | 87 +- internal/control/pricing_runtime.go | 6 +- internal/control/pricing_status.go | 14 +- internal/control/pricing_status_test.go | 13 +- internal/control/project_model_collection.go | 188 +- .../control/project_model_collection_test.go | 105 +- internal/control/provider_catalog.go | 225 - internal/control/provider_catalog_http.go | 174 - .../control/provider_catalog_http_test.go | 340 -- internal/control/request_logs.go | 228 +- internal/control/request_logs_test.go | 224 +- internal/control/route_inspect.go | 61 +- internal/control/route_inspect_test.go | 171 +- internal/control/runtime.go | 27 +- internal/control/runtime_observation.go | 4 +- internal/control/runtime_observation_test.go | 41 +- internal/control/runtime_test.go | 23 +- internal/control/server.go | 132 +- internal/control/server_phase1_test.go | 23 +- internal/control/server_test.go | 310 +- internal/control/service.go | 51 +- internal/control/service_integration_test.go | 32 +- internal/control/service_test.go | 311 +- internal/control/settings.go | 4 +- internal/control/settings_occ_test.go | 13 +- internal/control/test_helpers_test.go | 57 +- .../idempotency/group-key-import-k-twice.json | 11 - .../idempotency/group-key-import-k.json | 11 - internal/control/upstream_keys.go | 404 -- internal/control/upstream_keys_test.go | 1565 ------- internal/control/usage.go | 41 +- internal/control/usage_test.go | 55 +- internal/control/validation.go | 152 +- internal/control/validation_test.go | 389 +- internal/control/wire_v3_contract_test.go | 60 +- internal/dialect/anthropic.go | 190 +- internal/dialect/anthropic_test.go | 547 +-- internal/dialect/anthropic_usage_test.go | 49 +- internal/dialect/classify.go | 72 - internal/dialect/classify_test.go | 90 - internal/dialect/dialect.go | 24 +- internal/dialect/gemini.go | 229 +- internal/dialect/gemini_test.go | 524 +-- internal/dialect/gemini_usage_test.go | 13 +- internal/dialect/model_list.go | 117 - internal/dialect/model_list_test.go | 179 - internal/dialect/model_rewrite_test.go | 40 +- internal/dialect/openai.go | 142 +- internal/dialect/openai_responses.go | 120 +- .../dialect/openai_responses_stream_test.go | 5 +- internal/dialect/openai_responses_test.go | 362 +- .../dialect/openai_responses_usage_test.go | 7 +- internal/dialect/openai_test.go | 493 +- internal/dialect/openai_usage_options_test.go | 16 +- internal/dialect/openai_usage_test.go | 19 +- internal/dialect/pipeline.go | 31 - internal/dialect/pipeline_test.go | 70 - internal/dialect/probe.go | 69 - internal/dialect/probe_test.go | 255 - internal/dialect/request_execution.go | 243 + internal/dialect/request_execution_test.go | 188 + internal/dialect/request_reasoning_test.go | 18 +- internal/dialect/response_model_test.go | 29 +- internal/dialect/set_test.go | 3 +- internal/dialect/upstream_url.go | 31 - internal/dialect/upstream_url_test.go | 91 - internal/dialect/usage_test.go | 26 +- .../execution/bifrost/cloud_providers_test.go | 86 + .../bifrost/compatible_matrix_test.go | 330 ++ internal/execution/bifrost/converted.go | 523 +++ .../bifrost/converted_protocols_test.go | 601 +++ internal/execution/bifrost/executor.go | 1135 +++++ internal/execution/bifrost/executor_test.go | 1065 +++++ internal/execution/bifrost/list_models.go | 112 + internal/execution/bifrost/mapping.go | 694 +++ internal/execution/bifrost/model_alias.go | 279 ++ .../execution/bifrost/model_alias_test.go | 92 + .../bifrost/native_protocols_test.go | 245 + internal/execution/bifrost/passthrough.go | 676 +++ .../bifrost/responses_passthrough_test.go | 129 + internal/execution/bifrost/responses_test.go | 298 ++ internal/execution/bifrost/runtime.go | 302 ++ .../bifrost/utility_operations_test.go | 150 + internal/execution/contracts.go | 562 +++ internal/execution/contracts_test.go | 548 +++ internal/execution/validation.go | 372 ++ internal/gateway/channel_test_helpers_test.go | 79 + internal/gateway/dialects_integration_test.go | 115 +- internal/gateway/execution_boundary.go | 154 + internal/gateway/execution_boundary_test.go | 99 + internal/gateway/execution_forward.go | 413 ++ internal/gateway/execution_forward_test.go | 278 ++ internal/gateway/forward.go | 773 +--- internal/gateway/forward_test.go | 4115 ----------------- internal/gateway/handler.go | 161 +- internal/gateway/handler_test.go | 567 ++- internal/gateway/models.go | 19 +- internal/gateway/models_test.go | 46 +- .../gateway/openai_responses_stream_test.go | 545 --- .../representation_test_helpers_test.go | 55 + internal/gateway/request_id_test.go | 34 - internal/gateway/request_log.go | 67 +- internal/gateway/request_log_test.go | 322 +- .../response_model_observation_test.go | 178 - internal/gateway/response_representation.go | 16 +- .../gateway/response_representation_test.go | 1163 ----- internal/gateway/security_integration_test.go | 17 +- internal/gateway/stream_integration_test.go | 254 +- internal/gateway/transport_replay_test.go | 243 - internal/gateway/usage_capture_test.go | 1082 ----- internal/health/decision.go | 29 + internal/health/decision_test.go | 22 + internal/health/execution_judge.go | 159 + internal/health/execution_judge_test.go | 208 + internal/health/judge.go | 145 - internal/health/judge_test.go | 252 - internal/health/mutation.go | 14 +- internal/health/runtime_checkpoint.go | 24 +- internal/health/stats.go | 54 +- internal/health/stats_test.go | 34 +- internal/platform/errors/errors.go | 7 +- internal/platform/errors/errors_test.go | 16 +- internal/platform/i18n/i18n_test.go | 10 +- internal/platform/i18n/locales/en-US.go | 59 +- internal/platform/i18n/locales/ja-JP.go | 59 +- internal/platform/i18n/locales/zh-CN.go | 59 +- internal/platform/response/response_test.go | 6 +- internal/pricing/compile.go | 73 +- internal/pricing/compile_test.go | 61 +- internal/pricing/quote.go | 13 +- internal/pricing/quote_test.go | 34 +- internal/pricing/receipt.go | 4 +- internal/pricing/receipt_test.go | 88 + internal/pricing/types.go | 17 +- internal/reasoning/config.go | 13 + internal/requestlog/channel_identity_test.go | 328 ++ .../external_database_integration_test.go | 33 +- internal/requestlog/home_statistics_test.go | 10 +- internal/requestlog/mapper.go | 89 +- internal/requestlog/mapper_test.go | 8 +- internal/requestlog/phase4_v3_test.go | 8 +- internal/requestlog/process_log.go | 19 +- internal/requestlog/process_log_test.go | 14 +- internal/requestlog/query.go | 57 +- internal/requestlog/query_test.go | 70 +- internal/requestlog/service_test.go | 4 +- internal/requestlog/test_helpers_test.go | 11 +- internal/requestlog/types.go | 48 +- internal/requestlog/usage_query.go | 37 +- internal/requestlog/worker.go | 28 +- internal/requestlog/worker_test.go | 3 +- internal/scheduler/channel_scheduler_test.go | 354 ++ internal/scheduler/doc.go | 2 +- internal/scheduler/inspect.go | 294 +- internal/scheduler/inspect_test.go | 153 +- internal/scheduler/scheduler.go | 270 +- internal/scheduler/scheduler_test.go | 342 +- internal/state/channel_snapshot_test.go | 262 ++ internal/state/credential_registry_test.go | 79 + internal/state/loader/channel_loader_test.go | 140 + internal/state/loader/loader.go | 177 +- internal/state/loader/loader_test.go | 236 +- internal/state/loader/provider_id_test.go | 32 - internal/state/loader/weights_test.go | 46 +- internal/state/manager_test.go | 22 +- internal/state/registry.go | 418 +- internal/state/registry_test.go | 710 +-- internal/state/runtime_checkpoint.go | 14 +- internal/state/runtime_settings_test.go | 6 +- internal/state/runtime_view.go | 70 +- internal/state/snapshot.go | 293 +- internal/state/snapshot_test.go | 1029 +---- internal/state/weights_test.go | 238 +- internal/storage/database.go | 8 +- internal/storage/database_integration_test.go | 33 +- internal/storage/db_test.go | 235 +- internal/storage/migration_0001_initial_v2.go | 2 +- .../storage/migration_0001_initial_v2_test.go | 52 +- .../migration_0002_request_log_reasoning.go | 26 - .../migration_0003_global_model_prices.go | 166 - ...ation_0003_global_model_prices_sql_test.go | 58 - ...migration_0003_global_model_prices_test.go | 108 - ...gration_0004_mysql_model_price_identity.go | 38 - ...on_0004_mysql_model_price_identity_test.go | 39 - ...tion_0005_request_log_model_consistency.go | 51 - .../migration_0006_channel_execution.go | 731 +++ .../migration_0006_channel_execution_test.go | 263 ++ .../storage/migration_compatibility_test.go | 99 - internal/storage/migrations.go | 21 +- internal/storage/models/group.go | 70 +- internal/storage/models/model_price.go | 6 +- internal/storage/models/request.go | 58 +- internal/telemetry/requestlog.go | 47 +- internal/telemetry/requestlog_test.go | 11 +- internal/webui/workflow_test.go | 12 +- web/src/api/control/types.ts | 99 +- web/src/app/mutation-outcome.ts | 4 +- web/src/app/query-keys.ts | 38 +- web/src/app/resources/channels.ts | 163 + .../{upstream-keys.ts => credentials.ts} | 339 +- web/src/app/resources/groups.ts | 183 +- web/src/app/resources/health.ts | 39 +- web/src/app/resources/home.ts | 14 +- web/src/app/resources/invalidation.ts | 12 +- web/src/app/resources/model-prices.ts | 25 +- web/src/app/resources/models.ts | 30 +- web/src/app/resources/providers.ts | 244 +- web/src/app/resources/request-log-filters.ts | 3 +- web/src/app/resources/request-logs.ts | 101 +- web/src/app/resources/route-inspection.ts | 44 +- web/src/app/resources/usage.ts | 31 +- web/src/components/ui/CredentialHealthBar.vue | 187 + web/src/components/ui/KeyHealthBar.vue | 164 - .../features/access-keys/AccessKeyDrawer.vue | 4 +- .../features/access-keys/AccessKeysView.vue | 20 +- .../access-keys/access-key-options.ts | 5 +- web/src/features/groups/GroupDetailView.vue | 16 +- web/src/features/groups/GroupHeader.vue | 35 +- web/src/features/groups/GroupTabs.vue | 4 +- web/src/features/groups/GroupsView.vue | 124 +- .../GroupCredentialBatchBar.vue} | 28 +- .../GroupCredentialRecord.vue} | 248 +- .../GroupCredentialsTab.vue} | 344 +- .../credential-failure-presenter.ts | 18 + .../features/groups/group-collection-route.ts | 14 +- web/src/features/groups/group-route.ts | 86 +- .../groups/keys/key-failure-presenter.ts | 18 - web/src/features/groups/keys/key-patch.ts | 28 - .../features/groups/models/GroupModelsTab.vue | 4 +- .../groups/settings/GroupSettingsBaseForm.vue | 287 +- .../groups/settings/GroupSettingsTab.vue | 123 +- .../groups/settings/group-settings-patch.ts | 25 +- web/src/features/groups/upstream-host.ts | 7 - web/src/features/home/HomeSummary.vue | 6 +- web/src/features/home/HomeView.vue | 4 +- web/src/features/home/HomeWelcome.vue | 2 +- .../features/import/ChannelCatalogDrawer.vue | 255 + ...esetPicker.vue => ChannelPresetPicker.vue} | 131 +- .../features/import/CredentialTextarea.vue | 261 ++ .../features/import/ExistingGroupImport.vue | 68 +- .../import/ImportConnectionSection.vue | 195 +- web/src/features/import/KeyTextarea.vue | 207 - web/src/features/import/NewGroupImport.vue | 352 +- .../features/import/ProviderCatalogDrawer.vue | 304 -- .../features/import/credential-analysis.ts | 43 + .../features/import/import-operation-owner.ts | 29 +- web/src/features/import/import-recovery.ts | 57 +- web/src/features/import/import-route.ts | 12 +- web/src/features/import/key-analysis.ts | 32 - web/src/features/import/model-draft.ts | 11 +- web/src/features/import/recent-providers.ts | 29 - web/src/features/models/ModelTree.vue | 12 +- .../features/models/ModelUpstreamDrawer.vue | 10 + .../monitor/GroupHealthCollection.vue | 29 +- .../monitor/HealthProblemCollection.vue | 71 +- .../features/monitor/HealthSummaryStrip.vue | 24 +- web/src/features/monitor/HealthTab.vue | 49 +- web/src/features/monitor/InspectorTab.vue | 251 +- web/src/features/monitor/LogDetailDrawer.vue | 81 +- .../monitor/LogsAdvancedFilterDrawer.vue | 59 +- web/src/features/monitor/LogsFilterForm.vue | 5 + web/src/features/monitor/LogsTab.vue | 56 +- web/src/features/monitor/MonitorView.vue | 4 +- web/src/features/monitor/UsageFilterForm.vue | 78 +- web/src/features/monitor/UsageTab.vue | 88 +- web/src/features/monitor/log-filters.ts | 22 +- web/src/features/monitor/monitor-route.ts | 80 +- web/src/features/monitor/usage-filters.ts | 43 +- web/src/i18n/locales/en-US/core.ts | 26 +- web/src/i18n/locales/en-US/group.ts | 122 +- web/src/i18n/locales/en-US/import.ts | 97 +- web/src/i18n/locales/en-US/models.ts | 1 + web/src/i18n/locales/en-US/monitor.ts | 145 +- web/src/i18n/locales/ja-JP/core.ts | 27 +- web/src/i18n/locales/ja-JP/group.ts | 124 +- web/src/i18n/locales/ja-JP/import.ts | 99 +- web/src/i18n/locales/ja-JP/models.ts | 1 + web/src/i18n/locales/ja-JP/monitor.ts | 145 +- web/src/i18n/locales/zh-CN/core.ts | 25 +- web/src/i18n/locales/zh-CN/group.ts | 122 +- web/src/i18n/locales/zh-CN/import.ts | 70 +- web/src/i18n/locales/zh-CN/models.ts | 1 + web/src/i18n/locales/zh-CN/monitor.ts | 137 +- 394 files changed, 28959 insertions(+), 31035 deletions(-) create mode 100644 LICENSES/Apache-2.0.txt create mode 100644 THIRD_PARTY_NOTICES.md delete mode 100644 internal/catalog/curated.go delete mode 100644 internal/catalog/curated_test.go delete mode 100644 internal/catalog/official.go delete mode 100644 internal/catalog/official_test.go create mode 100644 internal/channel/channel.go create mode 100644 internal/channel/presets.go create mode 100644 internal/channel/registry_test.go create mode 100644 internal/control/channel_group_foundation_test.go create mode 100644 internal/control/channels.go create mode 100644 internal/control/credential_api_foundation_test.go create mode 100644 internal/control/credential_collection_helpers.go create mode 100644 internal/control/credential_mutations.go create mode 100644 internal/control/credentials.go create mode 100644 internal/control/execution_test_helpers_test.go delete mode 100644 internal/control/group_collection_test.go create mode 100644 internal/control/group_collection_test_helpers_test.go delete mode 100644 internal/control/group_create_test.go delete mode 100644 internal/control/group_key_collection.go delete mode 100644 internal/control/group_key_collection_test.go delete mode 100644 internal/control/group_key_mutations.go delete mode 100644 internal/control/group_key_mutations_test.go delete mode 100644 internal/control/group_keys.go delete mode 100644 internal/control/group_keys_test.go create mode 100644 internal/control/model_candidate.go create mode 100644 internal/control/model_price_sync_http.go create mode 100644 internal/control/model_price_sync_http_test.go delete mode 100644 internal/control/provider_catalog.go delete mode 100644 internal/control/provider_catalog_http.go delete mode 100644 internal/control/provider_catalog_http_test.go delete mode 100644 internal/control/testdata/idempotency/group-key-import-k-twice.json delete mode 100644 internal/control/testdata/idempotency/group-key-import-k.json delete mode 100644 internal/control/upstream_keys.go delete mode 100644 internal/control/upstream_keys_test.go delete mode 100644 internal/dialect/classify.go delete mode 100644 internal/dialect/classify_test.go delete mode 100644 internal/dialect/model_list.go delete mode 100644 internal/dialect/model_list_test.go delete mode 100644 internal/dialect/pipeline.go delete mode 100644 internal/dialect/pipeline_test.go delete mode 100644 internal/dialect/probe.go delete mode 100644 internal/dialect/probe_test.go create mode 100644 internal/dialect/request_execution.go create mode 100644 internal/dialect/request_execution_test.go delete mode 100644 internal/dialect/upstream_url.go delete mode 100644 internal/dialect/upstream_url_test.go create mode 100644 internal/execution/bifrost/cloud_providers_test.go create mode 100644 internal/execution/bifrost/compatible_matrix_test.go create mode 100644 internal/execution/bifrost/converted.go create mode 100644 internal/execution/bifrost/converted_protocols_test.go create mode 100644 internal/execution/bifrost/executor.go create mode 100644 internal/execution/bifrost/executor_test.go create mode 100644 internal/execution/bifrost/list_models.go create mode 100644 internal/execution/bifrost/mapping.go create mode 100644 internal/execution/bifrost/model_alias.go create mode 100644 internal/execution/bifrost/model_alias_test.go create mode 100644 internal/execution/bifrost/native_protocols_test.go create mode 100644 internal/execution/bifrost/passthrough.go create mode 100644 internal/execution/bifrost/responses_passthrough_test.go create mode 100644 internal/execution/bifrost/responses_test.go create mode 100644 internal/execution/bifrost/runtime.go create mode 100644 internal/execution/bifrost/utility_operations_test.go create mode 100644 internal/execution/contracts.go create mode 100644 internal/execution/contracts_test.go create mode 100644 internal/execution/validation.go create mode 100644 internal/gateway/channel_test_helpers_test.go create mode 100644 internal/gateway/execution_boundary.go create mode 100644 internal/gateway/execution_boundary_test.go create mode 100644 internal/gateway/execution_forward.go create mode 100644 internal/gateway/execution_forward_test.go delete mode 100644 internal/gateway/forward_test.go delete mode 100644 internal/gateway/openai_responses_stream_test.go create mode 100644 internal/gateway/representation_test_helpers_test.go delete mode 100644 internal/gateway/response_model_observation_test.go delete mode 100644 internal/gateway/response_representation_test.go delete mode 100644 internal/gateway/transport_replay_test.go delete mode 100644 internal/gateway/usage_capture_test.go create mode 100644 internal/health/decision.go create mode 100644 internal/health/decision_test.go create mode 100644 internal/health/execution_judge.go create mode 100644 internal/health/execution_judge_test.go delete mode 100644 internal/health/judge.go delete mode 100644 internal/health/judge_test.go create mode 100644 internal/pricing/receipt_test.go create mode 100644 internal/requestlog/channel_identity_test.go create mode 100644 internal/scheduler/channel_scheduler_test.go create mode 100644 internal/state/channel_snapshot_test.go create mode 100644 internal/state/credential_registry_test.go create mode 100644 internal/state/loader/channel_loader_test.go delete mode 100644 internal/state/loader/provider_id_test.go delete mode 100644 internal/storage/migration_0002_request_log_reasoning.go delete mode 100644 internal/storage/migration_0003_global_model_prices.go delete mode 100644 internal/storage/migration_0003_global_model_prices_sql_test.go delete mode 100644 internal/storage/migration_0003_global_model_prices_test.go delete mode 100644 internal/storage/migration_0004_mysql_model_price_identity.go delete mode 100644 internal/storage/migration_0004_mysql_model_price_identity_test.go delete mode 100644 internal/storage/migration_0005_request_log_model_consistency.go create mode 100644 internal/storage/migration_0006_channel_execution.go create mode 100644 internal/storage/migration_0006_channel_execution_test.go delete mode 100644 internal/storage/migration_compatibility_test.go create mode 100644 web/src/app/resources/channels.ts rename web/src/app/resources/{upstream-keys.ts => credentials.ts} (59%) create mode 100644 web/src/components/ui/CredentialHealthBar.vue delete mode 100644 web/src/components/ui/KeyHealthBar.vue rename web/src/features/groups/{keys/GroupKeyBatchBar.vue => credentials/GroupCredentialBatchBar.vue} (73%) rename web/src/features/groups/{keys/GroupKeyRecord.vue => credentials/GroupCredentialRecord.vue} (59%) rename web/src/features/groups/{keys/GroupKeysTab.vue => credentials/GroupCredentialsTab.vue} (64%) create mode 100644 web/src/features/groups/credentials/credential-failure-presenter.ts delete mode 100644 web/src/features/groups/keys/key-failure-presenter.ts delete mode 100644 web/src/features/groups/keys/key-patch.ts delete mode 100644 web/src/features/groups/upstream-host.ts create mode 100644 web/src/features/import/ChannelCatalogDrawer.vue rename web/src/features/import/{ProviderPresetPicker.vue => ChannelPresetPicker.vue} (56%) create mode 100644 web/src/features/import/CredentialTextarea.vue delete mode 100644 web/src/features/import/KeyTextarea.vue delete mode 100644 web/src/features/import/ProviderCatalogDrawer.vue create mode 100644 web/src/features/import/credential-analysis.ts delete mode 100644 web/src/features/import/key-analysis.ts delete mode 100644 web/src/features/import/recent-providers.ts diff --git a/.env.example b/.env.example index d135f8e5..c68c94d2 100644 --- a/.env.example +++ b/.env.example @@ -48,7 +48,7 @@ AUTH_KEY= # Optional explicit encryption master key. When empty, GPT-Load creates # ${DATA_DIR}/encryption.key with 0600 permissions on POSIX or a current-user # ACL on Windows. Back it up together with the selected database; losing or -# replacing it makes encrypted upstream keys unreadable. 2.0.0 does not +# replacing it makes encrypted channel credentials unreadable. 2.0.0 does not # support master-key rotation. ENCRYPTION_KEY= diff --git a/.github/scripts/release-docker-smoke.sh b/.github/scripts/release-docker-smoke.sh index 21f3cdcb..ecbc3174 100755 --- a/.github/scripts/release-docker-smoke.sh +++ b/.github/scripts/release-docker-smoke.sh @@ -148,6 +148,9 @@ docker run --name "${probe}" \ --entrypoint /bin/sh \ "${image}" -ceu ' test "$(id -u):$(id -g)" = "10001:10001" + test -r /app/licenses/LICENSE + test -r /app/licenses/THIRD_PARTY_NOTICES.md + test -r /app/licenses/Apache-2.0.txt printf canary >/app/data/release-write-canary test "$(cat /app/data/release-write-canary)" = canary rm /app/data/release-write-canary @@ -233,21 +236,20 @@ test "$( api_get "/api/usage?range=24h" >"${task_tmp}/usage-empty.json" api_get "/api/model-prices" >"${task_tmp}/prices-empty.json" -upstream_key="task13-upstream-${suffix}-$(openssl rand -hex 12)" +credential_secret="task13-credential-${suffix}-$(openssl rand -hex 12)" group_idempotency_key="$(python3 -c 'import uuid; print(uuid.uuid4())')" group_response="$( api_write POST "/api/groups" "$( node -e ' process.stdout.write(JSON.stringify({ name:"Task13 Release Smoke Group", - upstream_url:process.argv[1], - protocols:["openai-completions"], - models:[{id:"task13-release-model",alias:""}], - config:{}, - keys:process.argv[2], - confirm_same_upstream_url:false, + channel_id:"openai_compatible", + params:{base_url:process.argv[1]}, + models:[{id:"task13-release-model",alias:"",alias_enabled:false}], + credentials:process.argv[2], + confirm_same_target:false, })); - ' "http://host.docker.internal:${fake_port}/v1" "${upstream_key}" + ' "http://host.docker.internal:${fake_port}/v1" "${credential_secret}" )" "${group_idempotency_key}" )" group_id="$( @@ -256,7 +258,8 @@ group_id="$( const value=JSON.parse(fs.readFileSync(0,"utf8")); const group=value.data; if(value.code!==0||!group||group.group_name!=="Task13 Release Smoke Group"|| - !Number.isSafeInteger(group.group_id)||group.group_id<=0) process.exit(1); + !Number.isSafeInteger(group.group_id)||group.group_id<=0|| + group.credentials_added!==1||group.credentials_duplicated!==0) process.exit(1); process.stdout.write(String(group.group_id)); ' )" @@ -269,19 +272,18 @@ model_price_id="$( printf '%s' "${model_price_list}" | node -e ' const fs=require("fs"); const value=JSON.parse(fs.readFileSync(0,"utf8")); - const groupID=process.argv[1]; + const channelID=process.argv[1]; const modelID=process.argv[2]; const items=value.data&&value.data.items; if(value.code!==0||!Array.isArray(items)) process.exit(1); const matches=items.filter(item=> - item&&item.scope&&item.scope.kind==="group"&&item.scope.id===groupID&& - item.model_id===modelID + item&&item.channel_id===channelID&&item.model_id===modelID ); if(matches.length!==1||!Number.isSafeInteger(matches[0].id)||matches[0].id<=0) { process.exit(1); } process.stdout.write(String(matches[0].id)); - ' "${group_id}" "task13-release-model" + ' "openai_compatible" "task13-release-model" )" unset model_price_list test -n "${model_price_id}" @@ -384,15 +386,18 @@ node -e ' const groups=JSON.parse(fs.readFileSync(process.argv[1],"utf8")); const prices=JSON.parse(fs.readFileSync(process.argv[2],"utf8")); const usage=JSON.parse(fs.readFileSync(process.argv[3],"utf8")); - const groupID=process.argv[4]; + const groupID=Number(process.argv[4]); const priceID=Number(process.argv[5]); const modelID=process.argv[6]; - if(!groups.data.some(item=>item.name==="Task13 Release Smoke Group")) process.exit(1); + const groupItems=groups.data&&groups.data.items; + if(!Array.isArray(groupItems)||!groupItems.some(item=> + item.id===groupID&&item.name==="Task13 Release Smoke Group"&& + item.channel_id==="openai_compatible"&&item.credential_counts.total===1 + )) process.exit(1); const items=prices.data&&prices.data.items; if(!Array.isArray(items)||!Number.isSafeInteger(priceID)||priceID<=0) process.exit(1); const matches=items.filter(item=> - item&&item.scope&&item.scope.kind==="group"&&item.scope.id===groupID&& - item.model_id===modelID + item&&item.channel_id==="openai_compatible"&&item.model_id===modelID ); if(matches.length!==1||matches[0].id!==priceID|| matches[0].method!=="user_set"||matches[0].pricing_status!=="configured") { @@ -435,8 +440,8 @@ summary_file="${task_tmp}/summary.txt" printf 'graceful_stop_exit=0\n' } >"${summary_file}" -secret_labels=(auth_key encryption_key access_key upstream_key) -secret_values=("${auth_key}" "${encryption_key}" "${access_key}" "${upstream_key}") +secret_labels=(auth_key encryption_key access_key credential_secret) +secret_values=("${auth_key}" "${encryption_key}" "${access_key}" "${credential_secret}") for index in "${!secret_labels[@]}"; do secret_label="${secret_labels[${index}]}" secret_value="${secret_values[${index}]}" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a4c43e48..efb7ad9f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -47,8 +47,8 @@ jobs: - name: Build backend run: go build -o gpt-load . - - name: Run race-enabled tests - run: go test -race -count=1 . ./internal/... + - name: Run backend tests + run: go test -count=1 . ./internal/... - name: Check repository invariants run: git diff --check @@ -62,9 +62,9 @@ jobs: matrix: include: - driver: mysql - dsn: mysql://root:root@mysql:3306/gpt_load?charset=utf8mb4&transaction_isolation=%27READ-COMMITTED%27 + dsn: mysql://root:root@127.0.0.1:3306/gpt_load?charset=utf8mb4&transaction_isolation=%27READ-COMMITTED%27 - driver: postgres - dsn: postgres://postgres:postgres@postgres:5432/gpt_load?sslmode=disable + dsn: postgres://postgres:postgres@127.0.0.1:5432/gpt_load?sslmode=disable services: mysql: image: mysql:8.4 @@ -103,7 +103,10 @@ jobs: - name: Run external database contract env: GPT_LOAD_DATABASE_TEST_DSN: ${{ matrix.dsn }} - run: go test -v -count=1 ./internal/storage ./internal/control ./internal/requestlog -run '^TestExternalDatabase(ConcurrentMigrations|Lifecycle|ModelPriceIdentityUsesExactComparison|ReadSnapshotUsesRepeatableRead|GroupPriceReconciliation|RequestLogLifecycle)$' + GPT_LOAD_DATABASE_MIGRATION_TEST_DSN: ${{ matrix.dsn }} + run: | + go test -v -count=1 ./internal/storage -run '^TestExternalDatabaseChannelExecutionMigration$' + go test -v -count=1 ./internal/storage ./internal/control ./internal/requestlog -run '^TestExternalDatabase(ConcurrentMigrations|Lifecycle|ModelPriceIdentityUsesExactComparison|ReadSnapshotUsesRepeatableRead|GroupPriceReconciliation|RequestLogLifecycle)$' windows-encryption-acl: runs-on: windows-2025 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index f4c1824f..57132f2f 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -108,8 +108,8 @@ jobs: - name: Build backend run: go build -trimpath -o gpt-load . - - name: Run race-enabled tests - run: go test -race -count=1 . ./internal/... + - name: Run backend tests + run: go test -count=1 . ./internal/... - name: Check repository invariants run: git diff --check @@ -122,6 +122,62 @@ jobs: if-no-files-found: error retention-days: 14 + database-contract: + name: Database contract (${{ matrix.driver }}) + needs: validate-tag + runs-on: ubuntu-24.04 + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + include: + - driver: mysql + dsn: mysql://root:root@127.0.0.1:3306/gpt_load?charset=utf8mb4&transaction_isolation=%27READ-COMMITTED%27 + - driver: postgres + dsn: postgres://postgres:postgres@127.0.0.1:5432/gpt_load?sslmode=disable + services: + mysql: + image: mysql:8.4 + env: + MYSQL_ROOT_PASSWORD: root + MYSQL_DATABASE: gpt_load + ports: + - 3306:3306 + options: >- + --health-cmd="mysqladmin ping -h 127.0.0.1 -proot" + --health-interval=5s + --health-timeout=10s + --health-retries=20 + postgres: + image: postgres:18 + env: + POSTGRES_PASSWORD: postgres + POSTGRES_DB: gpt_load + ports: + - 5432:5432 + options: >- + --health-cmd="pg_isready -U postgres -d gpt_load" + --health-interval=5s + --health-timeout=10s + --health-retries=20 + steps: + - name: Checkout tag commit + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Set up Go + uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 + with: + go-version-file: go.mod + cache: true + + - name: Run external database contract + env: + GPT_LOAD_DATABASE_TEST_DSN: ${{ matrix.dsn }} + GPT_LOAD_DATABASE_MIGRATION_TEST_DSN: ${{ matrix.dsn }} + run: | + go test -v -count=1 ./internal/storage -run '^TestExternalDatabaseChannelExecutionMigration$' + go test -v -count=1 ./internal/storage ./internal/control ./internal/requestlog -run '^TestExternalDatabase(ConcurrentMigrations|Lifecycle|ModelPriceIdentityUsesExactComparison|ReadSnapshotUsesRepeatableRead|GroupPriceReconciliation|RequestLogLifecycle)$' + build-binaries: needs: verify-and-build-web runs-on: ubuntu-24.04 @@ -183,8 +239,17 @@ jobs: package-checksums: needs: build-binaries runs-on: ubuntu-24.04 - timeout-minutes: 5 + timeout-minutes: 10 steps: + - name: Checkout tag commit + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Set up Go + uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 + with: + go-version-file: go.mod + cache: true + - name: Download all release binaries uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: @@ -192,10 +257,22 @@ jobs: path: release merge-multiple: true + - name: Generate release SBOM and license assets + run: | + cp LICENSE release/LICENSE + cp THIRD_PARTY_NOTICES.md release/THIRD_PARTY_NOTICES.md + cp LICENSES/Apache-2.0.txt release/Apache-2.0.txt + go run github.com/CycloneDX/cyclonedx-gomod/cmd/cyclonedx-gomod@v1.10.0 \ + mod -json -noserial -notimestamp -output release/bom.cdx.json . + - name: Generate SHA256SUMS working-directory: release run: | sha256sum \ + Apache-2.0.txt \ + LICENSE \ + THIRD_PARTY_NOTICES.md \ + bom.cdx.json \ gpt-load-linux-amd64 \ gpt-load-linux-arm64 \ gpt-load-macos-amd64 \ @@ -287,6 +364,7 @@ jobs: - package-checksums - native-artifact-smoke - docker-smoke + - database-contract runs-on: ubuntu-24.04 timeout-minutes: 20 permissions: @@ -328,7 +406,11 @@ jobs: git merge-base --is-ancestor "${GITHUB_SHA}" origin/v2 cat >/tmp/expected-release-assets <<'ASSETS' + Apache-2.0.txt + LICENSE SHA256SUMS + THIRD_PARTY_NOTICES.md + bom.cdx.json gpt-load-linux-amd64 gpt-load-linux-arm64 gpt-load-macos-amd64 @@ -339,7 +421,7 @@ jobs: sort >/tmp/current-run-release-assets cmp /tmp/expected-release-assets /tmp/current-run-release-assets cd release - test "$(wc -l /tmp/expected-release-assets <<'ASSETS' + Apache-2.0.txt + LICENSE SHA256SUMS + THIRD_PARTY_NOTICES.md + bom.cdx.json gpt-load-linux-amd64 gpt-load-linux-arm64 gpt-load-macos-amd64 @@ -855,13 +957,14 @@ jobs: find public-release -mindepth 1 -maxdepth 1 -type f -exec basename {} \; | sort >/tmp/actual-release-assets asset_count="$(wc -l /tmp/checksummed-release-assets - tail -n +2 /tmp/expected-release-assets >/tmp/expected-checksummed-release-assets + grep -v '^SHA256SUMS$' /tmp/expected-release-assets \ + >/tmp/expected-checksummed-release-assets cmp /tmp/expected-checksummed-release-assets /tmp/checksummed-release-assets sha256sum --check SHA256SUMS diff --git a/Dockerfile b/Dockerfile index 6af06cda..8d3d515a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -47,6 +47,8 @@ RUN apk add --no-cache ca-certificates tzdata \ ENV HOST=0.0.0.0 ENV DATA_DIR=/app/data COPY --from=go-builder /build/gpt-load . +COPY LICENSE THIRD_PARTY_NOTICES.md /app/licenses/ +COPY LICENSES/Apache-2.0.txt /app/licenses/Apache-2.0.txt EXPOSE 3001 USER 10001:10001 ENTRYPOINT ["/app/gpt-load"] diff --git a/LICENSES/Apache-2.0.txt b/LICENSES/Apache-2.0.txt new file mode 100644 index 00000000..72eb37cd --- /dev/null +++ b/LICENSES/Apache-2.0.txt @@ -0,0 +1,201 @@ +Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2025 H3 Labs Inc. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md index f4b4ff30..9b517160 100644 --- a/README.md +++ b/README.md @@ -3,10 +3,10 @@ English | [中文](README_CN.md) | [日本語](README_JP.md) [![Release](https://img.shields.io/github/v/release/tbphp/gpt-load)](https://github.com/tbphp/gpt-load/releases) -![Go Version](https://img.shields.io/badge/Go-1.25-blue.svg) +![Go Version](https://img.shields.io/badge/Go-1.26-blue.svg) [![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE) -GPT-Load is a self-hosted AI API key aggregator and native-protocol gateway written in Go. A single binary with an embedded admin UI manages keys for OpenAI, Anthropic, Gemini, and compatible upstreams while exposing each provider's native data-plane endpoints. +GPT-Load is a self-hosted multi-provider AI gateway written in Go. A single binary with an embedded admin UI manages channel presets and encrypted credentials while exposing OpenAI, Anthropic, and Gemini-compatible client endpoints. For the maintained 1.4.x release documentation, visit the [official documentation](https://www.gpt-load.com/docs?lang=en). @@ -41,13 +41,14 @@ For the maintained 1.4.x release documentation, visit the [official documentatio > [!WARNING] > 2.0 is a **pre-release local candidate**. M3/M4 candidate code and retained local verification evidence exist, but release exit and publication are not complete. No `v2.0.0` tag, GitHub Release, public binary, or public container image has been verified as available. A checkout or branch is not release evidence. -2.0 is a greenfield rewrite whose data is incompatible with 1.x. `main` remains the 1.4.x maintenance line. The release contract reserves explicit `2`, `2.0`, and `v2.0.0` container tags and does not move `latest` automatically; these names do not imply that images have been published. +2.0 is a greenfield rewrite whose data is incompatible with 1.x. Existing 2.x databases migrate forward in place; stop the service and back up the database with its matching encryption key before upgrading. `main` remains the 1.4.x maintenance line. The release contract reserves explicit `2`, `2.0`, and `v2.0.0` container tags and does not move `latest` automatically; these names do not imply that images have been published. ## 2.0 capabilities - **Two planes:** provider-native paths on the data plane; management APIs under `/api`, with the admin UI embedded in the same Go binary. -- **Four selectable native protocols:** OpenAI Completions, OpenAI Responses, Anthropic Messages, and Gemini requests are forwarded in their respective protocols. A Group may enable any combination. GPT-Load does not translate between protocols. -- **Key and traffic management:** Groups, encrypted upstream keys, AccessKeys, model discovery, filtering and rate limits, scheduling, health state, cooldown, blacklist, and automatic weights. +- **Four client protocols:** OpenAI Completions, OpenAI Responses, Anthropic Messages, and Gemini requests keep their public endpoints. Each channel declares its supported native and converted operations; unsupported feature combinations fail before dispatch. +- **Channel and traffic management:** Users select a searchable code-defined channel, then enter models and encrypted credentials. GPT-Load owns AccessKey filtering, cross-Group credential scheduling, retry decisions, health state, cooldown, blacklist, and automatic weights. +- **Provider execution:** The official Bifrost Core Go SDK maintains provider-specific authentication, request/response conversion, streaming, normalized usage, and provider errors. GPT-Load pins one selected credential per logical attempt and disables Bifrost's configured retry and fallback policy. - **Control and observability:** runtime settings, route inspection, health views, RequestLog, and a Chinese, English, and Japanese admin UI. - **Usage and estimated cost:** usage extraction for the four protocols where the endpoint returns generation usage, 24-hour/30-day reports, per-request quality states, exact four-slot model prices synchronized from Models.dev where available, and user-managed prices. @@ -58,12 +59,11 @@ The M3 control-plane UI and M4 usage/pricing scope are present in the local cand - Correctness is guaranteed for a **single application instance** only; multi-instance coordination is not supported. - Supports SQLite, MySQL, and PostgreSQL through one unified `DATABASE_DSN` configuration. The current release targets the latest stable versions of these three engines; other database products are not supported. - The AccessKey and runtime configuration select the Group. A Group never appears in the data-plane URL. -- Protocol configuration is a clean break: use `openai-completions`, `openai-responses`, `anthropic`, or `gemini`. The old `openai`, `openai-response`, and `openai-chat-completions` values are invalid and have no compatibility path. -- A stored old protocol value causes the complete `ConfigSnapshot` compilation, and therefore startup/publication, to fail. The error identifies the Group or AccessKey and invalid value. Rebuild the pre-release 2.0 data before starting; there is no in-place protocol-value migration. +- Client protocol filters use `openai-completions`, `openai-responses`, `anthropic`, or `gemini`. A Group stores one `channel_id`; the code-owned channel preset determines provider behavior, accepted client protocols, credential schema, discovery, and Models.dev mapping. - OpenAI Responses resource routing has no Key affinity. Stateful turns using `previous_response_id` or `conversation`, and later retrieve/delete/cancel/input-item calls, are reliable only with one upstream Key or an upstream that shares resource storage across Keys. Otherwise the selected upstream may return a resource-not-found error. -- Upstream keys must be encrypted at rest with no plaintext fallback. 2.0.0 has no master-key rotation; `migrate-keys` remains an explicitly failing deferred command. -- There is no automatic 1.x migration, in-place upgrade, or reverse synchronization. -- There is no protocol conversion, online billing reconciliation, online backup API, or backup CLI. Models.dev synchronization supplies estimation metadata only; it is not a provider bill or invoice. +- Channel credentials must be encrypted at rest with no plaintext fallback. 2.0.0 has no master-key rotation; `migrate-keys` remains an explicitly failing deferred command. +- There is no automatic 1.x migration or reverse synchronization. Existing 2.x schemas migrate forward through the migration ledger. +- There is no online billing reconciliation, online backup API, or backup CLI. Models.dev synchronization supplies estimation metadata only; it is not a provider bill or invoice. ## Quick start @@ -128,9 +128,9 @@ Data-plane requests use an AccessKey. Provider-compatible credentials are accept | Gemini | `POST /v1beta/models/{model}:generateContent` | Gemini non-streaming generation | | Gemini | `POST /v1beta/models/{model}:streamGenerateContent` | Gemini streaming generation | -GPT-Load does not translate one dialect into another. The AccessKey and runtime configuration select the Group; it is not passed as a URL path segment. +The AccessKey and runtime configuration select the Group; it is not passed as a URL path segment. The selected channel determines whether an operation is native or converted. Conversion is capability-gated and never treated as arbitrary lossless JSON transformation. -The canonical protocol configuration values and display names are: +The canonical client protocol filter values and display names are: | Configuration value | Display name | |---|---| @@ -139,16 +139,16 @@ The canonical protocol configuration values and display names are: | `anthropic` | Anthropic | | `gemini` | Gemini | -The built-in OpenAI provider preset keeps the `openai` preset ID, uses `https://api.openai.com/v1` as its URL, and enables both OpenAI protocols by default. They remain ordinary independent checkboxes: either one or both may be selected. +The built-in `openai` channel supports both OpenAI client protocols. Other official, curated, compatible, and cloud channels expose only the operations and features declared by their code-owned presets; users do not configure protocol checkboxes per Group. -Responses routing uses the namespace boundary, not a per-resource allowlist. After AccessKey authentication, `/v1/responses` and its ordinary subpaths are sent through the same scheduler and forwarding pipeline. Decoded `.` or `..` path segments are rejected locally so normalization or redirects cannot escape the authorized namespace. `OPTIONS`, `CONNECT`, and `TRACE` are also rejected locally; other methods, including `GET`, `POST`, `DELETE`, and `HEAD`, are forwarded. Paths and queries are preserved within Go URL normalization: decoded `URL.Path` is re-encoded and `RawPath` is not retained. GPT-Load does not search other Keys for a resource ID; the selected upstream's response, including a resource-not-found error, is returned through the normal response-safety boundary. +Responses routing uses the namespace boundary, not a per-resource allowlist. After AccessKey authentication, `/v1/responses` and its ordinary subpaths are sent through the same scheduler and execution pipeline. Decoded `.` or `..` path segments are rejected locally so normalization or redirects cannot escape the authorized namespace. `OPTIONS`, `CONNECT`, and `TRACE` are also rejected locally; other methods, including `GET`, `POST`, `DELETE`, and `HEAD`, are forwarded. Paths and queries are preserved within Go URL normalization: decoded `URL.Path` is re-encoded and `RawPath` is not retained. GPT-Load does not search other Credentials for a resource ID; the selected upstream's response, including a resource-not-found error, is returned through the normal response-safety boundary. -A Group that enables Responses may keep an empty model list and still serve model-free Responses resource endpoints. Requests that include a model, including ordinary create requests, still require a configured model route. +A Group whose channel supports Responses may keep an empty model list and still serve model-free Responses resource endpoints. Requests that include a model, including ordinary create requests, still require a configured model route. > [!WARNING] -> 2.0.0 does not implement Responses affinity. Stateful multi-turn requests using `previous_response_id` or `conversation`, and resource operations on an earlier response ID, may reach a different Group/Key and receive an upstream 404. Use a single Key, stateless item replay with `store: false`, or an upstream with shared resource storage until affinity is implemented. +> 2.0.0 does not implement Responses affinity. Stateful multi-turn requests using `previous_response_id` or `conversation`, and resource operations on an earlier response ID, may reach a different Group/Credential and receive an upstream 404. Use a single Credential, stateless item replay with `store: false`, or an upstream with shared resource storage until affinity is implemented. -Responses create and compact requests participate in usage extraction. Retrieve, delete, cancel, input-items, input-token-count, and unknown extension subpaths are recorded with usage `not_applicable`. `InjectUsageOptions` remains capability-based: the Responses dialect does not support OpenAI Completions' `stream_options.include_usage`, so that Group setting is ignored for Responses. A Responses-only Group probe sends `input: "ping"`, `max_output_tokens: 16`, and `store: false`; when both OpenAI protocols are selected, OpenAI Completions is the representative Group/Key probe. Health is not tracked per protocol. +Responses create and compact requests participate in usage extraction. Retrieve, delete, cancel, input-items, input-token-count, and unknown extension subpaths are recorded with usage `not_applicable`. Discovery and validation are selected by the channel preset rather than by user-configured protocol checkboxes. OpenAI Completions example: @@ -188,7 +188,7 @@ print(response.output_text) ## Management, usage, and cost -The admin UI is served at `/`, and management APIs are under `/api`; both use `AUTH_KEY`. The UI covers Groups, upstream keys, AccessKeys, runtime settings, health, logs, route inspection, Usage, and model-price management. Current code and UI are the management API reference; this README intentionally avoids copying a route list that can drift. +The admin UI is served at `/`, and management APIs are under `/api`; both use `AUTH_KEY`. The UI covers the searchable channel directory, Groups, encrypted credentials, AccessKeys, runtime settings, health, logs, route inspection, Usage, and model-price management. Current code and UI are the management API reference; this README intentionally avoids copying a route list that can drift. Automatic catalog synchronization is enabled by default and uses the control plane to fetch the fixed endpoint `https://models.dev/api.json`; startup remains asynchronous and can use the durable last-known-good catalog. Manual synchronization remains available. Data-plane requests never contact Models.dev. @@ -197,7 +197,7 @@ Usage/Cost quality boundaries: - `complete` and `partial` usage contribute their known token dimensions; `missing` usage contributes only request and quality counts. - `priced` requests contribute their known estimated cost. `pricing_partial` retains the calculable portion while reporting incomplete price coverage; `unpriced` requests are never assigned guessed prices. - A clean EOF on a stream does not guarantee complete usage, and compatible relays may omit the provider's official terminal usage. -- Prices match the exact upstream model within the Group's Provider or custom-Group scope. The four flat slots are input, output, cache read, and cache write; an explicit zero means free, while an unset slot remains unavailable. +- Prices match the exact `(channel_id, upstream model)` identity. Models.dev is the sole automatic source; explicit user overrides and the existing fallback rules remain available. The four flat slots are input, output, cache read, and cache write; an explicit zero means free, while an unset slot remains unavailable. - Price changes affect future writes only. Historical RequestLog and UsageStat rows are not recalculated. - Current-process dropped/write-failure counters and durable database-window aggregates have different scopes. @@ -211,7 +211,7 @@ Usage/Cost quality boundaries: | `DATA_DIR` | `./data` | Native persistent directory; the container overrides it to `/app/data` | | `DATABASE_DSN` | empty → `${DATA_DIR}/gpt-load.db` | Empty selects managed SQLite. Non-empty values use one of `sqlite`, `mysql`, or `postgres` URL forms and are operator-managed | | `AUTH_KEY` | generated keyfile | Management bearer credential; an explicit value cannot contain whitespace; otherwise reads or creates `${DATA_DIR}/auth.key` | -| `ENCRYPTION_KEY` | generated keyfile | Master key for encrypted upstream keys; otherwise reads or creates `${DATA_DIR}/encryption.key` | +| `ENCRYPTION_KEY` | generated keyfile | Master key for encrypted channel credentials; otherwise reads or creates `${DATA_DIR}/encryption.key` | | `MODELS_DEV_AUTO_SYNC_ENABLED` | unset | Optional strict boolean override for Models.dev automatic synchronization; unset uses the runtime setting, which defaults to enabled | | `GRACEFUL_SHUTDOWN_TIMEOUT` | `10` | Graceful shutdown timeout in seconds | | `READ_TIMEOUT` | `60` | Maximum time to read a complete request, in seconds | @@ -237,9 +237,9 @@ The empty value is the only application-managed database mode. A non-empty SQLit - Database ownership follows only the raw `DATABASE_DSN`: empty means the managed SQLite DB/WAL/SHM under `${DATA_DIR}`; every non-empty value means an external, operator-owned database that GPT-Load does not mkdir, chmod, or create database users for and that must be backed up separately. - Secret ownership is independent of database ownership. For each secret, `/api/system/info` reports `key_file` or `environment`: archive a reported `key_file` (`auth.key` or `encryption.key`) from `DATA_DIR` regardless of the database source, or restore an `environment` secret separately from the protected external secret system. - On POSIX, managed `${DATA_DIR}` is restricted to `0700` and managed DB/WAL/SHM plus application-created key files to `0600`. Windows uses current-user-only ACLs, but the Windows runtime stop/ACL gate has not been executed for this candidate. -- Losing the matching `encryption.key`, from either source, makes encrypted upstream keys unrecoverable. 2.0.0 has no automatic repair or master-key rotation. +- Losing the matching `encryption.key`, from either source, makes encrypted channel credentials unrecoverable. 2.0.0 has no automatic repair or master-key rotation. - Managed SQLite uses WAL. Before backing it up, stop incoming traffic and wait for a clean exit: use `SIGTERM` on POSIX, or Ctrl+C, Ctrl+Break, or the service manager's stop action on Windows. External MySQL/PostgreSQL backups must follow the operator's engine-native backup procedure. -- Never paste AUTH_KEY, ENCRYPTION_KEY, AccessKeys, or upstream keys into logs, public issues, screenshots, or ordinary backup manifests. +- Never paste AUTH_KEY, ENCRYPTION_KEY, AccessKeys, or channel credentials into logs, public issues, screenshots, or ordinary backup manifests. ### Public operations baseline @@ -259,7 +259,7 @@ This checklist is self-contained and does not require access to the project's pr 1. Keep 1.x running and verify that its backup can be restored. 2. Give 2.0 a separate port, `DATA_DIR`, database, and Compose project/named volume. Do not share any of these with 1.x. -3. Manually rebuild the minimum Groups, upstream keys, AccessKeys, and rules; validate all four protocol variants, logs, and usage/cost in isolation. +3. Manually rebuild the minimum Groups, channel credentials, AccessKeys, and rules; validate the required client protocols, logs, and usage/cost in isolation. 4. Move entry traffic during a maintenance window or small rollout. On failure, stop 2.0 and switch back to the original 1.x deployment; do not reverse-import new 2.0 data. `latest` is not a safe 1.x-to-2.0 upgrade channel. Use the public operations baseline above for backup and restore, and keep the original 1.x deployment and data intact until the rollback window closes. @@ -282,7 +282,7 @@ make check Frontend unit tests and browser E2E tests are not part of the project workflow. Frontend verification consists of dependency installation, linting, formatting, type-checking, and building. -2.0.0 is expected to provide five native raw binaries plus `SHA256SUMS`: +2.0.0 is expected to provide five native raw binaries plus `SHA256SUMS`, a CycloneDX SBOM, and project/third-party license notices: - `gpt-load-linux-amd64` - `gpt-load-linux-arm64` @@ -294,4 +294,4 @@ These are the expected names in the release contract, not a claim that a downloa ## License and security -GPT-Load is released under the [MIT License](LICENSE). Report vulnerabilities through the process in [SECURITY.md](SECURITY.md). +GPT-Load is released under the [MIT License](LICENSE). Distributed third-party notices are listed in [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md), with license texts under [`LICENSES/`](LICENSES/). Report vulnerabilities through the process in [SECURITY.md](SECURITY.md). diff --git a/README_CN.md b/README_CN.md index e1a23674..7408cb37 100644 --- a/README_CN.md +++ b/README_CN.md @@ -3,10 +3,10 @@ [English](README.md) | 中文 | [日本語](README_JP.md) [![Release](https://img.shields.io/github/v/release/tbphp/gpt-load)](https://github.com/tbphp/gpt-load/releases) -![Go Version](https://img.shields.io/badge/Go-1.25-blue.svg) +![Go Version](https://img.shields.io/badge/Go-1.26-blue.svg) [![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE) -GPT-Load 是一个用 Go 构建的自托管 AI API Key 聚合与原生协议网关。它通过单个内嵌管理 UI 的二进制，管理 OpenAI、Anthropic、Gemini 及其兼容上游的密钥，并分别暴露三家的原生数据面端点。 +GPT-Load 是一个用 Go 构建的自托管多渠道 AI 网关。它通过单个内嵌管理 UI 的二进制管理渠道预设和加密凭据，并对外提供 OpenAI、Anthropic 与 Gemini 兼容端点。 已发布的 1.4.x 维护线文档请访问[官方文档](https://www.gpt-load.com/docs?lang=zh)。 @@ -41,13 +41,14 @@ GPT-Load 是一个用 Go 构建的自托管 AI API Key 聚合与原生协议网 > [!WARNING] > 2.0 当前是**发布前本地候选**。M3/M4 候选代码及本地验证留存证据已经存在，但正式出口和发布仍未完成；目前没有证据表明 `v2.0.0` tag、GitHub Release、公开二进制或公开容器镜像已经可用。checkout 或分支状态不能作为发布证据。 -2.0 是与 1.x 数据不兼容的 greenfield rewrite。`main` 继续承载 1.4.x 维护线。发布契约预留显式的 `2`、`2.0`、`v2.0.0` 容器 tag，且不会自动移动 `latest`；这些名称本身不代表镜像已经发布。 +2.0 是与 1.x 数据不兼容的 greenfield rewrite。已有 2.x 数据库支持原地向前迁移；升级前必须停机，并把数据库与匹配的 encryption key 一起备份。`main` 继续承载 1.4.x 维护线。发布契约预留显式的 `2`、`2.0`、`v2.0.0` 容器 tag，且不会自动移动 `latest`；这些名称本身不代表镜像已经发布。 ## 2.0 能力 - **双平面**：数据面保留服务商原生路径；管理 API 统一位于 `/api`，管理 UI 内嵌在同一个 Go 二进制中。 -- **四种可选原生协议**：OpenAI Completions、OpenAI Responses、Anthropic Messages 与 Gemini 请求分别按对应协议转发；Group 可以任意多选，不做协议互转。 -- **密钥与流量管理**：Group、加密上游 Key、AccessKey、模型发现、筛选与限流、调度、健康状态、cooldown、blacklist 和自动权重。 +- **四种客户端协议**：继续提供 OpenAI Completions、OpenAI Responses、Anthropic Messages 与 Gemini 入口。每个渠道明确声明可原生执行或转换的 operation；不支持的能力组合会在发送前拒绝。 +- **渠道与流量管理**：用户只需搜索并选择代码内置渠道，再填写模型和加密凭据。GPT-Load 继续负责 AccessKey 筛选、跨 Group 的凭据调度、重试决策、健康状态、cooldown、blacklist 与自动权重。 +- **Provider 执行**：官方 Bifrost Core Go SDK 负责不同 Provider 的认证、请求/响应转换、流式、usage 归一化和错误解析。每个逻辑 attempt 由 GPT-Load 固定一个凭据，并关闭 Bifrost 可配置的重试和 fallback。 - **控制与可观测性**：运行设置、路由检查、健康视图、RequestLog，以及中文、英文、日文管理 UI。 - **用量与估算成本**：对四种协议中会返回生成 usage 的接口进行采集，提供 24 小时/30 天汇总、明细质量状态、可用时从 Models.dev 同步的精确四槽模型价格，以及用户管理的价格。 @@ -58,12 +59,11 @@ M3 控制面 UI 与 M4 用量/定价范围已经进入本地候选，但正式 - 只保证**单应用实例**正确性，不支持多实例协调。 - 通过统一的 `DATABASE_DSN` 同时支持 SQLite、MySQL 和 PostgreSQL；当前版本面向这三种数据库的最新稳定版，不支持其他数据库产品。 - Group 由 AccessKey 和运行时配置选择，不出现在数据面 URL 中。 -- 协议配置采用 clean break：只允许 `openai-completions`、`openai-responses`、`anthropic`、`gemini`。旧值 `openai`、`openai-response` 与 `openai-chat-completions` 均无效，不提供兼容。 -- 数据库中只要保留一个旧协议值，整个 `ConfigSnapshot` 就会编译失败，进而阻止启动或配置发布；错误会包含 Group 或 AccessKey ID 及非法值。启动前需要重建发布前 2.0 数据，不提供协议值原地迁移。 -- OpenAI Responses 资源路由暂时没有 Key 亲和。使用 `previous_response_id` 或 `conversation` 的有状态多轮，以及后续 retrieve/delete/cancel/input-items 请求，只有在单上游 Key 或上游跨 Key 共享资源存储时才可靠；否则可能由被选中的上游返回资源不存在。 -- 上游密钥必须静态加密，不允许明文回退；2.0.0 不支持主密钥轮换，`migrate-keys` 仍是明确失败的延后命令。 -- 不支持 1.x 数据自动迁移、原地升级或反向同步。 -- 不提供协议转换、在线账单对账、在线备份 API 或备份 CLI。Models.dev 同步只提供估算元数据，不是服务商账单或发票。 +- AccessKey 的客户端协议筛选使用 `openai-completions`、`openai-responses`、`anthropic`、`gemini`。Group 只保存一个 `channel_id`；代码中的渠道预设统一决定 Provider 行为、客户端协议、凭据结构、模型发现与 Models.dev 映射。 +- OpenAI Responses 资源路由暂时没有 Credential 亲和。使用 `previous_response_id` 或 `conversation` 的有状态多轮，以及后续 retrieve/delete/cancel/input-items 请求，只有在单 Credential 或上游跨 Credential 共享资源存储时才可靠；否则可能由被选中的上游返回资源不存在。 +- 渠道凭据必须静态加密，不允许明文回退；2.0.0 不支持主密钥轮换，`migrate-keys` 仍是明确失败的延后命令。 +- 不支持 1.x 数据自动迁移或反向同步；已有 2.x schema 按 migration ledger 向前迁移。 +- 不提供在线账单对账、在线备份 API 或备份 CLI。Models.dev 同步只提供估算元数据，不是服务商账单或发票。 ## 快速开始 @@ -127,9 +127,9 @@ curl --fail http://localhost:3001/health | Gemini | `POST /v1beta/models/{model}:generateContent` | Gemini 非流式生成 | | Gemini | `POST /v1beta/models/{model}:streamGenerateContent` | Gemini 流式生成 | -GPT-Load 不把一种方言转换为另一种方言。Group 由 AccessKey 与运行时配置选择，不作为 URL 路径段传入。 +Group 由 AccessKey 与运行时配置选择，不作为 URL 路径段传入。被选中的渠道决定一次 operation 使用原生执行还是协议转换；转换必须通过能力门禁，不承诺任意 JSON 的无损互转。 -规范协议配置值与展示名如下： +规范的客户端协议筛选值与展示名如下： | 配置值 | 展示名 | |---|---| @@ -138,16 +138,16 @@ GPT-Load 不把一种方言转换为另一种方言。Group 由 AccessKey 与运 | `anthropic` | Anthropic | | `gemini` | Gemini | -内置 OpenAI 供应商预设仍使用 `openai` 作为 preset ID，URL 为 `https://api.openai.com/v1`，并默认同时勾选两种 OpenAI 协议。它们仍是普通、相互独立的多选项，用户可以只选任意一种或两种都选。 +内置 `openai` 渠道支持两种 OpenAI 客户端协议。其他官方、常用中转、兼容和云渠道只暴露代码预设中明确声明的 operation 与能力；用户不再为每个 Group 配置协议复选框。 -Responses 路由按命名空间边界匹配，不维护资源接口白名单。AccessKey 认证完成后，`/v1/responses` 及其普通子路径都会进入同一套调度和转发管线；已解码的 `.` 或 `..` 路径段会在本地拒绝，避免路径规范化或重定向逃逸已授权命名空间。`OPTIONS`、`CONNECT`、`TRACE` 也在本地拒绝，其他 method（包括 `GET`、`POST`、`DELETE`、`HEAD`）直接转发。路径与 query 在 Go URL 规范化边界内保留：已解码的 `URL.Path` 会重新编码，`RawPath` 不保留。GPT-Load 不会根据资源 ID 跨 Key 查找；被选中上游的响应（包括资源不存在错误）会经过统一响应安全边界后返回。 +Responses 路由按命名空间边界匹配，不维护资源接口白名单。AccessKey 认证完成后，`/v1/responses` 及其普通子路径都会进入同一套调度和执行管线；已解码的 `.` 或 `..` 路径段会在本地拒绝，避免路径规范化或重定向逃逸已授权命名空间。`OPTIONS`、`CONNECT`、`TRACE` 也在本地拒绝，其他 method（包括 `GET`、`POST`、`DELETE`、`HEAD`）直接转发。路径与 query 在 Go URL 规范化边界内保留：已解码的 `URL.Path` 会重新编码，`RawPath` 不保留。GPT-Load 不会根据资源 ID 跨 Credential 查找；被选中上游的响应（包括资源不存在错误）会经过统一响应安全边界后返回。 -选择 Responses 的 Group 可以不配置模型，并继续服务不含 model 的 Responses 资源接口；含 model 的请求（包括常规 create）仍需存在可匹配的模型路由。 +渠道支持 Responses 的 Group 可以不配置模型，并继续服务不含 model 的 Responses 资源接口；含 model 的请求（包括常规 create）仍需存在可匹配的模型路由。 > [!WARNING] -> 2.0.0 尚未实现 Responses 亲和。携带 `previous_response_id` 或 `conversation` 的有状态多轮，以及对既有 response ID 的资源操作，可能命中不同 Group/Key 并收到上游 404。在亲和完成前，请使用单 Key、`store: false` 的无状态 item replay，或使用跨 Key 共享资源存储的上游。 +> 2.0.0 尚未实现 Responses 亲和。携带 `previous_response_id` 或 `conversation` 的有状态多轮，以及对既有 response ID 的资源操作，可能命中不同 Group/Credential 并收到上游 404。在亲和完成前，请使用单 Credential、`store: false` 的无状态 item replay，或使用跨 Credential 共享资源存储的上游。 -Responses create 与 compact 请求参与 usage 抽取；retrieve、delete、cancel、input-items、input-token-count 及未知扩展子路径在 RequestLog 中记录为 usage `not_applicable`。`InjectUsageOptions` 继续按能力接口生效：Responses dialect 不支持 OpenAI Completions 的 `stream_options.include_usage`，因此该 Group 设置对 Responses 会被忽略。仅选择 Responses 的 Group 会用 `input: "ping"`、`max_output_tokens: 16`、`store: false` 进行探测；同时选择两种 OpenAI 协议时，以 OpenAI Completions 作为 Group/Key 的代表性探测。健康状态不细分到协议级别。 +Responses create 与 compact 请求参与 usage 抽取；retrieve、delete、cancel、input-items、input-token-count 及未知扩展子路径在 RequestLog 中记录为 usage `not_applicable`。模型发现与验活由渠道预设决定，不再依赖用户选择协议。 OpenAI Completions 示例： @@ -187,7 +187,7 @@ print(response.output_text) ## 管理面与用量成本 -管理 UI 位于 `/`，管理 API 位于 `/api`；两者都使用 `AUTH_KEY`。UI 包含 Group、上游 Key、AccessKey、运行设置、健康、日志、路由检查、Usage 与模型价格管理。完整管理 API 以当前代码和 UI 为准，本 README 不复制容易漂移的路由清单。 +管理 UI 位于 `/`，管理 API 位于 `/api`；两者都使用 `AUTH_KEY`。UI 包含可搜索渠道目录、Group、加密凭据、AccessKey、运行设置、健康、日志、路由检查、Usage 与模型价格管理。完整管理 API 以当前代码和 UI 为准，本 README 不复制容易漂移的路由清单。 模型目录自动同步默认启用，只在控制面访问固定端点 `https://models.dev/api.json`；启动过程保持异步，并可使用持久化的 last-known-good 目录。手动同步始终可用，数据面请求不会访问 Models.dev。 @@ -196,7 +196,7 @@ Usage/Cost 质量边界： - `complete` 与 `partial` usage 的已知 token 维度进入汇总；`missing` usage 只进入请求数和质量计数。 - `priced` 请求的已知估算成本进入汇总；`pricing_partial` 会保留可计算部分并标记价格覆盖不完整，`unpriced` 请求不会被猜价。 - 流式连接 clean EOF 不代表一定获得完整 usage；兼容中转站也可能不返回官方终态 usage。 -- 价格按 Group 的 Provider 或自定义 Group 作用域精确匹配上游模型。四个平面价格槽为输入、输出、缓存读取和缓存写入；显式 `0` 表示免费，未设置表示不可估算。 +- 价格按 `(channel_id, 上游模型)` 精确匹配。Models.dev 是唯一自动价格源，同时保留用户显式覆盖和现有 fallback 规则。四个平面价格槽为输入、输出、缓存读取和缓存写入；显式 `0` 表示免费，未设置表示不可估算。 - 修改模型价格只影响后续写入，不重算历史 RequestLog 或 UsageStat。 - 当前进程的 dropped/write-failure 计数与数据库窗口内的耐久汇总是不同口径。 @@ -210,7 +210,7 @@ Usage/Cost 质量边界： | `DATA_DIR` | `./data` | 原生持久目录；容器内覆盖为 `/app/data` | | `DATABASE_DSN` | 空 → `${DATA_DIR}/gpt-load.db` | 空值选择 managed SQLite；非空值必须使用统一的 `sqlite`、`mysql` 或 `postgres` URL，并由 operator 管理 | | `AUTH_KEY` | 自动生成 keyfile | 管理 bearer 凭据；显式值不能包含空白，留空时读取或创建 `${DATA_DIR}/auth.key` | -| `ENCRYPTION_KEY` | 自动生成 keyfile | 加密上游 Key 的主密钥；留空时读取或创建 `${DATA_DIR}/encryption.key` | +| `ENCRYPTION_KEY` | 自动生成 keyfile | 加密渠道凭据的主密钥；留空时读取或创建 `${DATA_DIR}/encryption.key` | | `MODELS_DEV_AUTO_SYNC_ENABLED` | 未设置 | Models.dev 自动同步的可选严格布尔覆盖；未设置时使用运行设置，其默认启用 | | `GRACEFUL_SHUTDOWN_TIMEOUT` | `10` | 优雅停机超时，单位为秒 | | `READ_TIMEOUT` | `60` | 读取完整请求的最长时间，单位为秒 | @@ -236,9 +236,9 @@ postgres://user:password@db.example:5432/gpt_load?sslmode=require - 数据库归属只由 raw `DATABASE_DSN` 决定：空值表示 `${DATA_DIR}` 下 managed SQLite DB/WAL/SHM；任何非空值都表示 operator 管理的 external 数据库，GPT-Load 不为其 mkdir、chmod、创建数据库或用户，必须单独备份。 - secret 归属与数据库归属相互独立。`/api/system/info` 会分别报告 secret source：无论数据库是哪种 source，`key_file` 都必须归档 `DATA_DIR` 中对应的 `auth.key` / `encryption.key`；`environment` 则必须从受保护的外部 secret system 单独恢复。 - POSIX 下 managed `${DATA_DIR}` 权限收紧为 `0700`，managed DB/WAL/SHM 及应用创建的 key 文件为 `0600`。Windows 使用当前用户专属 ACL，但当前候选尚未执行 Windows runtime 停机/ACL 门禁。 -- 无论来自哪种 source，丢失匹配的 `encryption.key` 都会使已加密上游 Key 无法恢复。2.0.0 没有自动修复或主密钥轮换。 +- 无论来自哪种 source，丢失匹配的 `encryption.key` 都会使已加密渠道凭据无法恢复。2.0.0 没有自动修复或主密钥轮换。 - Managed SQLite 使用 WAL。备份前先停止入口流量并等待 clean exit：POSIX 使用 `SIGTERM`，Windows 使用 Ctrl+C、Ctrl+Break 或 service manager stop。MySQL/PostgreSQL 必须按 operator 的数据库原生备份流程执行。 -- 不要把 AUTH_KEY、ENCRYPTION_KEY、AccessKey 或上游 Key 粘贴到日志、公开 issue、截图或普通备份清单中。 +- 不要把 AUTH_KEY、ENCRYPTION_KEY、AccessKey 或渠道凭据粘贴到日志、公开 issue、截图或普通备份清单中。 ### 公开运维基线 @@ -258,7 +258,7 @@ postgres://user:password@db.example:5432/gpt_load?sslmode=require 1. 保持 1.x 运行并先验证其备份可恢复。 2. 为 2.0 使用独立端口、`DATA_DIR`、数据库、Compose project 与 named volume，不与 1.x 共享任何一项。 -3. 手工重建最小 Group、上游 Key、AccessKey 与规则，在隔离环境验证四种协议、日志及 usage/cost。 +3. 手工重建最小 Group、渠道凭据、AccessKey 与规则，在隔离环境验证所需客户端协议、日志及 usage/cost。 4. 在维护窗或小流量阶段切换入口；失败时停止 2.0 并切回原 1.x，不把 2.0 新数据反向导入 1.x。 `latest` 不是 1.x → 2.0 的安全升级通道。备份和恢复按上面的公开运维基线执行，并在回滚窗口关闭前保留原 1.x 部署及其数据。 @@ -281,7 +281,7 @@ make check 项目工作流不包含前端单元测试或浏览器 E2E；前端验证范围为依赖安装、lint、format、type-check 与 build。 -2.0.0 预期提供五个原生 raw binary 和一个 `SHA256SUMS`： +2.0.0 预期提供五个原生 raw binary，以及 `SHA256SUMS`、CycloneDX SBOM、项目和第三方许可证声明： - `gpt-load-linux-amd64` - `gpt-load-linux-arm64` @@ -293,4 +293,4 @@ make check ## 许可证与安全 -GPT-Load 使用 [MIT License](LICENSE)。安全漏洞请按 [SECURITY.md](SECURITY.md) 中的流程报告。 +GPT-Load 使用 [MIT License](LICENSE)。发行物中的第三方声明见 [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md)，完整许可文本位于 [`LICENSES/`](LICENSES/)。安全漏洞请按 [SECURITY.md](SECURITY.md) 中的流程报告。 diff --git a/README_JP.md b/README_JP.md index 6747f468..c2823d75 100644 --- a/README_JP.md +++ b/README_JP.md @@ -3,10 +3,10 @@ [English](README.md) | [中文](README_CN.md) | 日本語 [![Release](https://img.shields.io/github/v/release/tbphp/gpt-load)](https://github.com/tbphp/gpt-load/releases) -![Go Version](https://img.shields.io/badge/Go-1.25-blue.svg) +![Go Version](https://img.shields.io/badge/Go-1.26-blue.svg) [![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE) -GPT-Loadは、Goで構築されたセルフホスト型のAI APIキー集約・ネイティブプロトコルゲートウェイです。管理UIを内蔵した単一バイナリでOpenAI、Anthropic、Geminiおよび互換上流のキーを管理し、各プロバイダーのネイティブなデータプレーンエンドポイントを公開します。 +GPT-Loadは、Goで構築されたセルフホスト型のマルチチャネルAIゲートウェイです。管理UIを内蔵した単一バイナリでチャネルプリセットと暗号化された認証情報を管理し、OpenAI、Anthropic、Gemini互換のクライアントエンドポイントを公開します。 公開済みの1.4.xメンテナンスラインについては、[公式ドキュメント](https://www.gpt-load.com/docs?lang=ja)をご覧ください。 @@ -41,13 +41,14 @@ GPT-Loadは、Goで構築されたセルフホスト型のAI APIキー集約・ > [!WARNING] > 2.0は現在、**プレリリースのローカル候補**です。M3/M4の候補コードとローカル検証の保存済み証跡はありますが、正式なリリース判断・承認と公開は完了していません。`v2.0.0`タグ、GitHub Release、公開バイナリ、公開コンテナイメージが利用可能だと確認できる証拠はありません。checkoutやブランチの状態はリリースの証拠にはなりません。 -2.0は1.xとデータ互換性のないgreenfield rewriteです。`main`は引き続き1.4.xメンテナンスラインを提供します。リリース契約では明示的な`2`、`2.0`、`v2.0.0`コンテナタグを予約し、`latest`を自動更新しませんが、これらの名前自体はイメージの公開を意味しません。 +2.0は1.xとデータ互換性のないgreenfield rewriteです。既存の2.xデータベースはインプレースで前方移行できます。アップグレード前にサービスを停止し、対応するencryption keyとともにデータベースをバックアップしてください。`main`は引き続き1.4.xメンテナンスラインを提供します。リリース契約では明示的な`2`、`2.0`、`v2.0.0`コンテナタグを予約し、`latest`を自動更新しませんが、これらの名前自体はイメージの公開を意味しません。 ## 2.0の機能 - **2つのプレーン**：データプレーンはプロバイダーのネイティブパスを維持し、管理APIは`/api`に統一されます。管理UIは同じGoバイナリに内蔵されます。 -- **4つの選択可能なネイティブプロトコル**：OpenAI Completions、OpenAI Responses、Anthropic Messages、Geminiのリクエストをそれぞれのプロトコルで転送します。Groupでは任意に複数選択でき、プロトコル間の変換は行いません。 -- **キーとトラフィックの管理**：Group、暗号化された上流Key、AccessKey、モデル検出、フィルターとレート制限、スケジューリング、ヘルス状態、cooldown、blacklist、自動重み付け。 +- **4つのクライアントプロトコル**：OpenAI Completions、OpenAI Responses、Anthropic Messages、Geminiの公開エンドポイントを維持します。各チャネルはネイティブまたは変換可能なoperationを明示し、未対応の機能組み合わせは送信前に拒否します。 +- **チャネルとトラフィックの管理**：ユーザーは検索可能なコード定義チャネルを選び、モデルと暗号化認証情報を入力します。GPT-LoadがAccessKeyフィルター、Groupをまたぐ認証情報スケジューリング、リトライ判断、ヘルス、cooldown、blacklist、自動重み付けを所有します。 +- **Provider実行**：公式Bifrost Core Go SDKがProvider固有の認証、request/response変換、streaming、usage正規化、エラー解析を担当します。GPT-Loadは論理attemptごとに1つの認証情報を固定し、Bifrostの設定可能なretryとfallbackを無効にします。 - **制御と可観測性**：ランタイム設定、ルート検査、ヘルス表示、RequestLog、中国語・英語・日本語の管理UI。 - **使用量と推定コスト**：4プロトコルのうち生成usageを返すエンドポイントからusageを取得し、24時間/30日レポート、リクエスト単位の品質状態、利用可能な場合にModels.devから同期する完全一致の4価格スロット、ユーザー管理価格を提供します。 @@ -58,12 +59,11 @@ M3のコントロールプレーンUIとM4のusage/pricing範囲はローカル - 正しさを保証するのは**単一アプリケーションインスタンス**のみで、複数インスタンスの協調には対応しません。 - 統一された `DATABASE_DSN` で SQLite、MySQL、PostgreSQL をサポートします。現在のリリースはこの3種類の最新安定版を対象とし、その他のデータベース製品には対応しません。 - GroupはAccessKeyとランタイム設定で選択され、データプレーンURLには含まれません。 -- プロトコル設定はclean breakです。使用できる値は`openai-completions`、`openai-responses`、`anthropic`、`gemini`だけです。旧値`openai`、`openai-response`、`openai-chat-completions`に互換処理はありません。 -- データベースに旧プロトコル値が1件でも残ると、`ConfigSnapshot`全体のコンパイルが失敗し、起動または設定公開を停止します。エラーにはGroupまたはAccessKeyのIDと不正な値が含まれます。起動前にプレリリース2.0のデータを再構築してください。プロトコル値のインプレース移行はありません。 -- OpenAI ResponsesのリソースルーティングにはKey affinityがありません。`previous_response_id`または`conversation`を使うステートフルな複数ターンと、後続のretrieve/delete/cancel/input-itemsは、上流Keyが1つの場合、または上流がKey間でリソースストレージを共有する場合だけ確実です。それ以外では、選択された上流からresource-not-foundが返る可能性があります。 -- 上流キーは必ず保存時に暗号化され、平文へのフォールバックはありません。2.0.0はマスターキーのローテーションに対応せず、`migrate-keys`は明示的に失敗する延期コマンドのままです。 -- 1.xデータの自動移行、インプレースアップグレード、逆同期には対応しません。 -- プロトコル変換、オンライン請求照合、オンラインバックアップAPI、バックアップCLIは提供しません。Models.dev同期が提供するのは推定用メタデータだけで、プロバイダー請求書やインボイスではありません。 +- AccessKeyのクライアントプロトコルフィルターは`openai-completions`、`openai-responses`、`anthropic`、`gemini`を使用します。Groupは1つの`channel_id`だけを保存し、コード定義プリセットがProvider動作、クライアントプロトコル、認証情報schema、モデル検出、Models.dev対応を決定します。 +- OpenAI ResponsesのリソースルーティングにはCredential affinityがありません。`previous_response_id`または`conversation`を使うステートフルな複数ターンと、後続のretrieve/delete/cancel/input-itemsは、Credentialが1つの場合、または上流がCredential間でリソースストレージを共有する場合だけ確実です。それ以外では、選択された上流からresource-not-foundが返る可能性があります。 +- チャネル認証情報は必ず保存時に暗号化され、平文へのフォールバックはありません。2.0.0はマスターキーのローテーションに対応せず、`migrate-keys`は明示的に失敗する延期コマンドのままです。 +- 1.xデータの自動移行や逆同期には対応しません。既存2.x schemaはmigration ledgerで前方移行します。 +- オンライン請求照合、オンラインバックアップAPI、バックアップCLIは提供しません。Models.dev同期が提供するのは推定用メタデータだけで、プロバイダー請求書やインボイスではありません。 ## クイックスタート @@ -128,9 +128,9 @@ curl --fail http://localhost:3001/health | Gemini | `POST /v1beta/models/{model}:generateContent` | Gemini非ストリーミング生成 | | Gemini | `POST /v1beta/models/{model}:streamGenerateContent` | Geminiストリーミング生成 | -GPT-Loadは方言間の変換を行いません。GroupはAccessKeyとランタイム設定で選択され、URLパスセグメントとして渡しません。 +GroupはAccessKeyとランタイム設定で選択され、URLパスセグメントとして渡しません。選択されたチャネルがoperationをネイティブ実行するか変換するかを決定します。変換はcapability gateを通り、任意JSONの無損失変換を保証しません。 -正規のプロトコル設定値と表示名は次のとおりです。 +正規のクライアントプロトコルフィルター値と表示名は次のとおりです。 | 設定値 | 表示名 | |---|---| @@ -139,16 +139,16 @@ GPT-Loadは方言間の変換を行いません。GroupはAccessKeyとランタ | `anthropic` | Anthropic | | `gemini` | Gemini | -組み込みOpenAIプロバイダープリセットは、preset IDとして`openai`を維持し、URLに`https://api.openai.com/v1`を使用して、デフォルトで2つのOpenAIプロトコルを有効にします。両方とも通常の独立した選択肢で、どちらか一方または両方を選択できます。 +組み込み`openai`チャネルは2つのOpenAIクライアントプロトコルをサポートします。その他の公式、主要中継、互換、クラウドチャネルはコードプリセットで宣言されたoperationと機能だけを公開し、Groupごとのプロトコルチェックボックスはありません。 -Responsesルーティングはリソースごとのallowlistではなく、名前空間境界で一致します。AccessKey認証後、`/v1/responses`と通常のサブパスは同じスケジューラーおよび転送パイプラインに入ります。デコード済みの`.`または`..`パスセグメントは、正規化やリダイレクトによる認可済み名前空間からの逸脱を防ぐためローカルで拒否します。`OPTIONS`、`CONNECT`、`TRACE`もローカルで拒否し、`GET`、`POST`、`DELETE`、`HEAD`を含むその他のmethodは転送します。パスとqueryはGo URL正規化の範囲内で保持され、デコード済み`URL.Path`は再エンコードされ、`RawPath`は保持されません。GPT-LoadはリソースIDを使って別のKeyを検索しません。選択された上流の応答（resource-not-foundを含む）は、共通の応答安全境界を通して返されます。 +Responsesルーティングはリソースごとのallowlistではなく、名前空間境界で一致します。AccessKey認証後、`/v1/responses`と通常のサブパスは同じスケジューラーおよび実行パイプラインに入ります。デコード済みの`.`または`..`パスセグメントは、正規化やリダイレクトによる認可済み名前空間からの逸脱を防ぐためローカルで拒否します。`OPTIONS`、`CONNECT`、`TRACE`もローカルで拒否し、`GET`、`POST`、`DELETE`、`HEAD`を含むその他のmethodは転送します。パスとqueryはGo URL正規化の範囲内で保持され、デコード済み`URL.Path`は再エンコードされ、`RawPath`は保持されません。GPT-LoadはリソースIDを使って別のCredentialを検索しません。選択された上流の応答（resource-not-foundを含む）は、共通の応答安全境界を通して返されます。 -Responsesを有効にしたGroupはモデル一覧が空でも、modelを含まないResponsesリソースAPIを処理できます。通常のcreateを含むmodel付き要求には、引き続き一致するモデルルートが必要です。 +ResponsesをサポートするチャネルのGroupはモデル一覧が空でも、modelを含まないResponsesリソースAPIを処理できます。通常のcreateを含むmodel付き要求には、引き続き一致するモデルルートが必要です。 > [!WARNING] -> 2.0.0はResponses affinityを実装していません。`previous_response_id`または`conversation`を使うステートフルな複数ターン、および既存response IDへのリソース操作は、別のGroup/Keyに到達して上流404を受ける場合があります。affinity実装までは、単一Key、`store: false`によるステートレスなitem replay、またはKey間でリソースストレージを共有する上流を使用してください。 +> 2.0.0はResponses affinityを実装していません。`previous_response_id`または`conversation`を使うステートフルな複数ターン、および既存response IDへのリソース操作は、別のGroup/Credentialに到達して上流404を受ける場合があります。affinity実装までは、単一Credential、`store: false`によるステートレスなitem replay、またはCredential間でリソースストレージを共有する上流を使用してください。 -Responsesのcreateとcompactはusage抽出の対象です。retrieve、delete、cancel、input-items、input-token-count、不明な拡張サブパスはRequestLogでusage `not_applicable`として記録されます。`InjectUsageOptions`は引き続きcapabilityベースです。Responses dialectはOpenAI Completionsの`stream_options.include_usage`を実装しないため、このGroup設定はResponsesでは無視されます。Responsesだけを選択したGroupは`input: "ping"`、`max_output_tokens: 16`、`store: false`でprobeします。2つのOpenAIプロトコルを選択した場合は、OpenAI CompletionsをGroup/Keyの代表probeとして使用します。プロトコル単位のヘルス状態はありません。 +Responsesのcreateとcompactはusage抽出の対象です。retrieve、delete、cancel、input-items、input-token-count、不明な拡張サブパスはRequestLogでusage `not_applicable`として記録されます。モデル検出とvalidationはユーザーのプロトコル選択ではなく、チャネルプリセットが決定します。 OpenAI Completionsの例： @@ -188,7 +188,7 @@ print(response.output_text) ## 管理、使用量、コスト -管理UIは`/`、管理APIは`/api`で提供され、どちらも`AUTH_KEY`を使用します。UIにはGroup、上流キー、AccessKey、ランタイム設定、ヘルス、ログ、ルート検査、Usage、モデル価格管理があります。管理APIの事実源は現在のコードとUIであり、このREADMEでは変化しやすいルート一覧を複製しません。 +管理UIは`/`、管理APIは`/api`で提供され、どちらも`AUTH_KEY`を使用します。UIには検索可能なチャネル一覧、Group、暗号化認証情報、AccessKey、ランタイム設定、ヘルス、ログ、ルート検査、Usage、モデル価格管理があります。管理APIの事実源は現在のコードとUIであり、このREADMEでは変化しやすいルート一覧を複製しません。 モデルカタログの自動同期はデフォルトで有効で、コントロールプレーンから固定エンドポイント`https://models.dev/api.json`へアクセスします。起動処理は非同期のままで、永続化したlast-known-goodカタログを利用できます。手動同期も常に利用でき、データプレーンのリクエストがModels.devへアクセスすることはありません。 @@ -197,7 +197,7 @@ Usage/Costの品質境界： - `complete`と`partial`のusageは既知のtoken次元を集計し、`missing` usageはリクエスト数と品質カウントだけに入ります。 - `priced`リクエストは既知の推定コストを集計します。`pricing_partial`は計算可能な部分を保持しながら価格カバレッジ不足を示し、`unpriced`に推測価格を割り当てることはありません。 - ストリームのclean EOFは完全なusageを保証せず、互換中継サービスがプロバイダー公式の終端usageを返さない場合もあります。 -- 価格はGroupのProviderまたはカスタムGroupスコープ内で上流モデルに完全一致します。4つの平面価格スロットは入力、出力、キャッシュ読み取り、キャッシュ書き込みで、明示的な`0`は無料、未設定は推定不可を表します。 +- 価格は`(channel_id, 上流モデル)`に完全一致します。Models.devが唯一の自動価格ソースで、ユーザーの明示的overrideと既存fallbackルールは維持されます。4つの平面価格スロットは入力、出力、キャッシュ読み取り、キャッシュ書き込みで、明示的な`0`は無料、未設定は推定不可を表します。 - 価格変更は今後の書き込みにだけ影響し、過去のRequestLogやUsageStatは再計算されません。 - 現在のプロセスにおけるdropped/write-failureカウンターと、データベース期間内の永続集計は異なる範囲です。 @@ -211,7 +211,7 @@ Usage/Costの品質境界： | `DATA_DIR` | `./data` | ネイティブの永続ディレクトリ。コンテナ内では`/app/data`に上書き | | `DATABASE_DSN` | 空 → `${DATA_DIR}/gpt-load.db` | 空ならmanaged SQLiteを選択。非空値は統一された`sqlite`、`mysql`、`postgres` URLで指定し、operatorが管理 | | `AUTH_KEY` | keyfileを自動生成 | 管理bearer認証情報。明示値に空白は使用不可。空の場合`${DATA_DIR}/auth.key`を読み取りまたは作成 | -| `ENCRYPTION_KEY` | keyfileを自動生成 | 上流キー暗号化用マスターキー。空の場合`${DATA_DIR}/encryption.key`を読み取りまたは作成 | +| `ENCRYPTION_KEY` | keyfileを自動生成 | チャネル認証情報の暗号化用マスターキー。空の場合`${DATA_DIR}/encryption.key`を読み取りまたは作成 | | `MODELS_DEV_AUTO_SYNC_ENABLED` | 未設定 | Models.dev自動同期の任意の厳密なboolean override。未設定ではデフォルト有効のランタイム設定を使用 | | `GRACEFUL_SHUTDOWN_TIMEOUT` | `10` | グレースフルシャットダウンの秒数 | | `READ_TIMEOUT` | `60` | リクエスト全体を読み取る最大秒数 | @@ -237,9 +237,9 @@ postgres://user:password@db.example:5432/gpt_load?sslmode=require - データベースの所有区分はraw `DATABASE_DSN`だけで決まります。空なら`${DATA_DIR}`配下のmanaged SQLite DB/WAL/SHM、非空ならoperator所有のexternalデータベースです。GPT-Loadはexternalに対してmkdir、chmod、データベースやユーザー作成を行わず、operatorが別途バックアップします。 - secretの所有区分はデータベースとは独立しています。`/api/system/info`がsecretごとにsourceを返します。データベースのsourceにかかわらず、`key_file`なら`DATA_DIR`内の対応する`auth.key` / `encryption.key`をアーカイブし、`environment`なら保護された外部secret systemから別途復元します。 - POSIXではmanaged `${DATA_DIR}`を`0700`、managed DB/WAL/SHMとアプリケーションが作成したkeyファイルを`0600`に制限します。Windowsでは実行ユーザー専用ACLを使用しますが、この候補についてWindows runtimeの停止/ACLゲートは未実行です。 -- sourceにかかわらず、対応する`encryption.key`を失うと、暗号化済みの上流キーは復旧できません。2.0.0には自動修復やマスターキーのローテーションがありません。 +- sourceにかかわらず、対応する`encryption.key`を失うと、暗号化済みのチャネル認証情報は復旧できません。2.0.0には自動修復やマスターキーのローテーションがありません。 - Managed SQLiteはWALを使用します。バックアップ前に新規トラフィックを止め、clean exitを待ちます。POSIXでは`SIGTERM`、WindowsではCtrl+C、Ctrl+Break、またはservice managerの停止操作を使用します。MySQL/PostgreSQLはoperatorのデータベースネイティブなバックアップ手順に従います。 -- AUTH_KEY、ENCRYPTION_KEY、AccessKey、上流キーをログ、公開issue、スクリーンショット、通常のバックアップ一覧に貼り付けないでください。 +- AUTH_KEY、ENCRYPTION_KEY、AccessKey、チャネル認証情報をログ、公開issue、スクリーンショット、通常のバックアップ一覧に貼り付けないでください。 ### 公開運用ベースライン @@ -259,7 +259,7 @@ postgres://user:password@db.example:5432/gpt_load?sslmode=require 1. 1.xを稼働させたまま、バックアップから復元できることを確認します。 2. 2.0には別のポート、`DATA_DIR`、データベース、Compose project、named volumeを用意し、1.xと共有しません。 -3. 最小限のGroup、上流キー、AccessKey、ルールを手動で再構築し、4つのプロトコル、ログ、usage/costを隔離環境で検証します。 +3. 最小限のGroup、チャネル認証情報、AccessKey、ルールを手動で再構築し、必要なクライアントプロトコル、ログ、usage/costを隔離環境で検証します。 4. メンテナンス時間または小規模なロールアウトで入口トラフィックを切り替えます。失敗した場合は2.0を停止して元の1.xへ戻し、2.0で新たに生成されたデータを1.xへ逆インポートしません。 `latest`は1.xから2.0への安全なアップグレードチャネルではありません。バックアップと復元は上記の公開運用ベースラインに従い、ロールバック期間が終了するまで元の1.xデプロイとデータを保持してください。 @@ -282,7 +282,7 @@ make check プロジェクトのワークフローには、フロントエンドのユニットテストとブラウザE2Eテストを含めません。フロントエンドの検証範囲は、依存関係のインストール、lint、format、type-check、buildです。 -2.0.0では5つのネイティブraw binaryと`SHA256SUMS`を提供する予定です。 +2.0.0では5つのネイティブraw binaryに加え、`SHA256SUMS`、CycloneDX SBOM、プロジェクトと第三者ライセンス通知を提供する予定です。 - `gpt-load-linux-amd64` - `gpt-load-linux-arm64` @@ -294,4 +294,4 @@ make check ## ライセンスとセキュリティ -GPT-Loadは[MIT License](LICENSE)で公開されています。脆弱性は[SECURITY.md](SECURITY.md)の手順に従って報告してください。 +GPT-Loadは[MIT License](LICENSE)で公開されています。第三者通知は[THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md)、完全なライセンス本文は[`LICENSES/`](LICENSES/)にあります。脆弱性は[SECURITY.md](SECURITY.md)の手順に従って報告してください。 diff --git a/SECURITY.md b/SECURITY.md index 0345affe..6bd471d2 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -15,8 +15,8 @@ Report suspected vulnerabilities through [GitHub Private Vulnerability Reporting Please do not open a public issue for an undisclosed vulnerability. Use the private report so the maintainers can investigate and coordinate a fix before public disclosure. -## Provider credential exposure boundary +## Channel credential exposure boundary -Provider keys and any derived fragments, including prefix/suffix masks, must not appear in data-plane response headers or bodies, request logs, or the request-log management API. Request-log attempts identify an upstream key only by its internal `key_id`. +Channel credentials and any derived secret fragments must not appear in data-plane response headers or bodies, request logs, or the request-log management API. Request-log attempts identify the selected credential only by its internal `credential_id`, together with the non-secret `channel_id`. -The authenticated Group upstream-key management list is the sole approved exception: it may return the existing mask so an operator can identify a key in the provider console. It remains protected by management authentication and never exposes plaintext credentials. +The authenticated Group credential list is the sole approved mask exception: it may return a safe identifier mask so an operator can identify a credential in the provider console. Plaintext reveal is a separate authenticated, explicitly invoked endpoint and must never be logged or returned by collection, health, usage, or request-log APIs. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 00000000..619d957d --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,19 @@ +# Third-Party Notices + +GPT-Load includes third-party open-source software. The release SBOM contains +the complete resolved dependency inventory. + +## Bifrost Core + +- Module: `github.com/maximhq/bifrost/core` +- Version: `v1.7.7` +- Copyright: 2025 H3 Labs Inc. +- License: Apache License 2.0 + +GPT-Load uses Bifrost Core as an infrastructure adapter for provider execution +and protocol conversion. GPT-Load's domain models, persisted channel IDs, +scheduling, retry policy, health state, usage accounting, and pricing remain +owned by GPT-Load. + +The complete Apache License 2.0 text is distributed in +`LICENSES/Apache-2.0.txt`. diff --git a/go.mod b/go.mod index f09ea5a5..23c42c5d 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( github.com/glebarez/sqlite v1.11.0 github.com/joho/godotenv v1.5.1 github.com/klauspost/compress v1.19.1 + github.com/maximhq/bifrost/core v1.7.7 github.com/nicksnyder/go-i18n/v2 v2.6.1 github.com/sirupsen/logrus v1.9.4 go.uber.org/dig v1.19.0 @@ -20,12 +21,41 @@ require ( ) require ( + cloud.google.com/go v0.123.0 // indirect + cloud.google.com/go/compute/metadata v0.9.0 // indirect filippo.io/edwards25519 v1.1.0 // indirect + github.com/Azure/azure-sdk-for-go/sdk/azcore v1.20.0 // indirect + github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1 // indirect + github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2 // indirect + github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0 // indirect + github.com/aws/aws-sdk-go-v2 v1.42.0 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.10 // indirect + github.com/aws/aws-sdk-go-v2/config v1.32.14 // indirect + github.com/aws/aws-sdk-go-v2/credentials v1.19.14 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.29 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.29 // indirect + github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.22 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.13 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.21 // indirect + github.com/aws/aws-sdk-go-v2/service/s3 v1.99.0 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.30.15 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.19 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 // indirect + github.com/aws/smithy-go v1.27.1 // indirect + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/buger/jsonparser v1.2.0 // indirect github.com/bytedance/gopkg v0.1.4 // indirect github.com/bytedance/sonic v1.15.2 // indirect github.com/bytedance/sonic/loader v0.5.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cloudwego/base64x v0.1.7 // indirect github.com/dustin/go-humanize v1.0.1 // indirect + github.com/fasthttp/websocket v1.5.12 // indirect github.com/gabriel-vasile/mimetype v1.4.15 // indirect github.com/gin-contrib/sse v1.1.1 // indirect github.com/glebarez/go-sqlite v1.22.0 // indirect @@ -35,7 +65,9 @@ require ( github.com/go-sql-driver/mysql v1.8.1 // indirect github.com/goccy/go-json v0.10.6 // indirect github.com/goccy/go-yaml v1.19.2 // indirect + github.com/golang-jwt/jwt/v5 v5.3.1 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/invopop/jsonschema v0.13.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/pgx/v5 v5.10.0 // indirect @@ -44,22 +76,41 @@ require ( github.com/jinzhu/now v1.1.5 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect + github.com/kylelemons/godebug v1.1.0 // indirect github.com/leodido/go-urn v1.4.0 // indirect + github.com/mailru/easyjson v0.9.1 // indirect + github.com/mark3labs/mcp-go v0.43.2 // indirect + github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.24 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect github.com/pelletier/go-toml/v2 v2.2.4 // indirect + github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect github.com/quic-go/qpack v0.6.0 // indirect github.com/quic-go/quic-go v0.59.1 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + github.com/rs/zerolog v1.34.0 // indirect + github.com/savsgio/gotils v0.0.0-20250408102913-196191ec6287 // indirect + github.com/spf13/cast v1.10.0 // indirect + github.com/tidwall/gjson v1.18.0 // indirect + github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/pretty v1.2.1 // indirect + github.com/tidwall/sjson v1.2.5 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.3.1 // indirect + github.com/valyala/bytebufferpool v1.0.0 // indirect + github.com/valyala/fasthttp v1.71.0 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect go.mongodb.org/mongo-driver/v2 v2.5.1 // indirect - golang.org/x/arch v0.22.0 // indirect + go.starlark.net v0.0.0-20260102030733-3fee463870c9 // indirect + golang.org/x/arch v0.23.0 // indirect golang.org/x/net v0.56.0 // indirect + golang.org/x/oauth2 v0.36.0 // indirect golang.org/x/sync v0.22.0 // indirect - google.golang.org/protobuf v1.36.11 // indirect + google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect modernc.org/libc v1.74.4 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.11.0 // indirect diff --git a/go.sum b/go.sum index 66adbde8..6de806dc 100644 --- a/go.sum +++ b/go.sum @@ -1,22 +1,88 @@ +cloud.google.com/go v0.123.0 h1:2NAUJwPR47q+E35uaJeYoNhuNEM9kM8SjgRgdeOJUSE= +cloud.google.com/go v0.123.0/go.mod h1:xBoMV08QcqUGuPW65Qfm1o9Y4zKZBpGS+7bImXLTAZU= +cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs= +cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10= filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA= filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4= +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.20.0 h1:JXg2dwJUmPB9JmtVmdEB16APJ7jurfbY5jnfXpJoRMc= +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.20.0/go.mod h1:YD5h/ldMsG0XiIw7PdyNhLxaM317eFh5yNLccNfGdyw= +github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1 h1:Hk5QBxZQC1jb2Fwj6mpzme37xbCDdNTxU7O9eb5+LB4= +github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1/go.mod h1:IYus9qsFobWIc2YVwe/WPjcnyCkPKtnHAqUYeebc8z0= +github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache v0.3.2 h1:yz1bePFlP5Vws5+8ez6T3HWXPmwOK7Yvq8QxDBD3SKY= +github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache v0.3.2/go.mod h1:Pa9ZNPuoNu/GztvBSKk9J1cDJW6vk/n0zLtV4mgd8N8= +github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2 h1:9iefClla7iYpfYWdzPCRDozdmndjTm8DXdpCzPajMgA= +github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2/go.mod h1:XtLgD3ZD34DAaVIIAyG3objl5DynM3CQ/vMcbBNJZGI= +github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1 h1:WJTmL004Abzc5wDB5VtZG2PJk5ndYDgVacGqfirKxjM= +github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1/go.mod h1:tCcJZ0uHAmvjsVYzEFivsRTN00oz5BEsRgQHu5JZ9WE= +github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0 h1:XRzhVemXdgvJqCH0sFfrBUTnUJSBrBf7++ypk+twtRs= +github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0/go.mod h1:HKpQxkWaGLJ+D/5H8QRpyQXA1eKjxkFlOMwck5+33Jk= github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk= github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= github.com/andybalholm/brotli v1.2.2 h1:HzTuoo2ErYQqf5qvcJInB8uvqSVxRttzkFexPWtnceM= github.com/andybalholm/brotli v1.2.2/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= +github.com/aws/aws-sdk-go-v2 v1.42.0 h1:XvXMJTkFQtpBKIWZnmr9ZEOc2InWM2yldjXEJ/bymhA= +github.com/aws/aws-sdk-go-v2 v1.42.0/go.mod h1:27+ACypSLljLAEKsCYOmrjKh83vuTRkuAe9Uv/3A4bg= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.10 h1:gx1AwW1Iyk9Z9dD9F4akX5gnN3QZwUB20GGKH/I+Rho= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.10/go.mod h1:qqY157uZoqm5OXq/amuaBJyC9hgBCBQnsaWnPe905GY= +github.com/aws/aws-sdk-go-v2/config v1.32.14 h1:opVIRo/ZbbI8OIqSOKmpFaY7IwfFUOCCXBsUpJOwDdI= +github.com/aws/aws-sdk-go-v2/config v1.32.14/go.mod h1:U4/V0uKxh0Tl5sxmCBZ3AecYny4UNlVmObYjKuuaiOo= +github.com/aws/aws-sdk-go-v2/credentials v1.19.14 h1:n+UcGWAIZHkXzYt87uMFBv/l8THYELoX6gVcUvgl6fI= +github.com/aws/aws-sdk-go-v2/credentials v1.19.14/go.mod h1:cJKuyWB59Mqi0jM3nFYQRmnHVQIcgoxjEMAbLkpr62w= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 h1:NUS3K4BTDArQqNu2ih7yeDLaS3bmHD0YndtA6UP884g= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21/go.mod h1:YWNWJQNjKigKY1RHVJCuupeWDrrHjRqHm0N9rdrWzYI= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.29 h1:f3vKqSo13fhTYb+JEcXwXefZQE26I1FB5eTSniU67ko= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.29/go.mod h1:MzoLFUArKGpGD+ukmPiTPG1X5x4o6M2kq4v2dr1FiEc= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.29 h1:RdwIf/CuUsvJX3RgJagbOyotl/cxoLY4xviKuE7p2GY= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.29/go.mod h1:71wt8W2EgswdZy9Mf9KNnzxZ3TiZlv4caKghPktDOkA= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 h1:qYQ4pzQ2Oz6WpQ8T3HvGHnZydA72MnLuFK9tJwmrbHw= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6/go.mod h1:O3h0IK87yXci+kg6flUKzJnWeziQUKciKrLjcatSNcY= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.22 h1:rWyie/PxDRIdhNf4DzRk0lvjVOqFJuNnO8WwaIRVxzQ= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.22/go.mod h1:zd/JsJ4P7oGfUhXn1VyLqaRZwPmZwg44Jf2dS84Dm3Y= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 h1:5EniKhLZe4xzL7a+fU3C2tfUN4nWIqlLesfrjkuPFTY= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7/go.mod h1:x0nZssQ3qZSnIcePWLvcoFisRXJzcTVvYpAAdYX8+GI= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.13 h1:JRaIgADQS/U6uXDqlPiefP32yXTda7Kqfx+LgspooZM= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.13/go.mod h1:CEuVn5WqOMilYl+tbccq8+N2ieCy0gVn3OtRb0vBNNM= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 h1:c31//R3xgIJMSC8S6hEVq+38DcvUlgFY0FM6mSI5oto= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21/go.mod h1:r6+pf23ouCB718FUxaqzZdbpYFyDtehyZcmP5KL9FkA= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.21 h1:ZlvrNcHSFFWURB8avufQq9gFsheUgjVD9536obIknfM= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.21/go.mod h1:cv3TNhVrssKR0O/xxLJVRfd2oazSnZnkUeTf6ctUwfQ= +github.com/aws/aws-sdk-go-v2/service/s3 v1.99.0 h1:hlSuz394kV0vhv9drL5lhuEFbEOEP1VyQpy15qWh1Pk= +github.com/aws/aws-sdk-go-v2/service/s3 v1.99.0/go.mod h1:uoA43SdFwacedBfSgfFSjjCvYe8aYBS7EnU5GZ/YKMM= +github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 h1:QKZH0S178gCmFEgst8hN0mCX1KxLgHBKKY/CLqwP8lg= +github.com/aws/aws-sdk-go-v2/service/signin v1.0.9/go.mod h1:7yuQJoT+OoH8aqIxw9vwF+8KpvLZ8AWmvmUWHsGQZvI= +github.com/aws/aws-sdk-go-v2/service/sso v1.30.15 h1:lFd1+ZSEYJZYvv9d6kXzhkZu07si3f+GQ1AaYwa2LUM= +github.com/aws/aws-sdk-go-v2/service/sso v1.30.15/go.mod h1:WSvS1NLr7JaPunCXqpJnWk1Bjo7IxzZXrZi1QQCkuqM= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.19 h1:dzztQ1YmfPrxdrOiuZRMF6fuOwWlWpD2StNLTceKpys= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.19/go.mod h1:YO8TrYtFdl5w/4vmjL8zaBSsiNp3w0L1FfKVKenZT7w= +github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 h1:p8ogvvLugcR/zLBXTXrTkj0RYBUdErbMnAFFp12Lm/U= +github.com/aws/aws-sdk-go-v2/service/sts v1.41.10/go.mod h1:60dv0eZJfeVXfbT1tFJinbHrDfSJ2GZl4Q//OSSNAVw= +github.com/aws/smithy-go v1.27.1 h1:4T340VFndXtADGF52gYa1POyL7s9E4Z1OeZ1hCscIw8= +github.com/aws/smithy-go v1.27.1/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/buger/jsonparser v1.2.0 h1:4EFcvK1kD4jyj6YqNK6skK6w+y7FHHBR+XBCtxwu/6g= +github.com/buger/jsonparser v1.2.0/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/bytedance/gopkg v0.1.4 h1:oZnQwnX82KAIWb7033bEwtxvTqXcYMxDBaQxo5JJHWM= github.com/bytedance/gopkg v0.1.4/go.mod h1:v1zWfPm21Fb+OsyXN2VAHdL6TBb2L88anLQgdyje6R4= github.com/bytedance/sonic v1.15.2 h1:90H+rcF/FwLXwfB1cudOLq/je83n683Utf4Cbp0xHCo= github.com/bytedance/sonic v1.15.2/go.mod h1:mT2NbXunuaEbnZ+mRIX/vYqKISmgEuHFDI4UzmKx2SA= github.com/bytedance/sonic/loader v0.5.1 h1:Ygpfa9zwRCCKSlrp5bBP/b/Xzc3VxsAW+5NIYXrOOpI= github.com/bytedance/sonic/loader v0.5.1/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cloudwego/base64x v0.1.7 h1:NppS+Fgzg5ovhn4NkUXaDT3x9jldgH5ToMCqzBSi2zI= github.com/cloudwego/base64x v0.1.7/go.mod h1:Cu1PV9zfrSf7ET2tIbWbbEy7jO7HHJ13q4X2SQ8aWYg= +github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/fasthttp/websocket v1.5.12 h1:e4RGPpWW2HTbL3zV0Y/t7g0ub294LkiuXXUuTOUInlE= +github.com/fasthttp/websocket v1.5.12/go.mod h1:I+liyL7/4moHojiOgUOIKEWm9EIxHqxZChS+aMFltyg= +github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= +github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/gabriel-vasile/mimetype v1.4.15 h1:05iP/CYtZ/w455R/KZM6rZ5ieAdh99UPtd+d3YzLmaI= github.com/gabriel-vasile/mimetype v1.4.15/go.mod h1:azpTcoLcDZRNgFou5j+APrqQx9HqVPWa6ijYQIIVswQ= github.com/gin-contrib/sse v1.1.1 h1:uGYpNwTacv5R68bSGMapo62iLTRa9l5zxGCps4hK6ko= @@ -41,6 +107,9 @@ github.com/goccy/go-json v0.10.6 h1:p8HrPJzOakx/mn/bQtjgNjdTcN+/S6FcG2CTtQOrHVU= github.com/goccy/go-json v0.10.6/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM= github.com/goccy/go-yaml v1.19.2/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA= +github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= +github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= @@ -48,8 +117,12 @@ github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17k github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hajimehoshi/go-mp3 v0.3.4 h1:NUP7pBYH8OguP4diaTZ9wJbUbk3tC0KlfzsEpWmYj68= +github.com/hajimehoshi/go-mp3 v0.3.4/go.mod h1:fRtZraRFcWb0pu7ok0LqyFhCUrPeMsGRSVop0eemFmo= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= +github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= @@ -66,16 +139,35 @@ github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/keybase/go-keychain v0.0.1 h1:way+bWYa6lDppZoZcgMbYsvC7GxljxrskdNInRtuthU= +github.com/keybase/go-keychain v0.0.1/go.mod h1:PdEILRW3i9D8JcdM+FmY6RwkHGnhHxXwkPPMeUgOK1k= github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk= github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= +github.com/mailru/easyjson v0.9.1 h1:LbtsOm5WAswyWbvTEOqhypdPeZzHavpZx96/n553mR8= +github.com/mailru/easyjson v0.9.1/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU= +github.com/mark3labs/mcp-go v0.43.2 h1:21PUSlWWiSbUPQwXIJ5WKlETixpFpq+WBpbMGDSVy/I= +github.com/mark3labs/mcp-go v0.43.2/go.mod h1:YnJfOL382MIWDx1kMY+2zsRHU/q78dBg9aFb8W6Thdw= +github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= +github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE= +github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8= +github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= +github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI= github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A= github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU= github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= +github.com/maximhq/bifrost/core v1.7.7 h1:5mMspsjvP/yJ4XBvuuAZ99R0BF+Dyy2l6FFO6UZxLfk= +github.com/maximhq/bifrost/core v1.7.7/go.mod h1:XQGQ99V6iW2yRbq4zg5+LfDcJMR06x/Ie4HuC8FWEDI= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= @@ -87,16 +179,29 @@ github.com/nicksnyder/go-i18n/v2 v2.6.1 h1:JDEJraFsQE17Dut9HFDHzCoAWGEQJom5s0TRd github.com/nicksnyder/go-i18n/v2 v2.6.1/go.mod h1:Vee0/9RD3Quc/NmwEjzzD7VTZ+Ir7QbXocrkhOzmUKA= github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4= github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= -github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ= +github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII= github.com/quic-go/quic-go v0.59.1 h1:0Gmua0HW1Tv7ANR7hUYwRyD0MG5OJfgvYSZasGZzBic= github.com/quic-go/quic-go v0.59.1/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= +github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= +github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0= +github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY= +github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ= +github.com/savsgio/gotils v0.0.0-20250408102913-196191ec6287 h1:qIQ0tWF9vxGtkJa24bR+2i53WBCz1nW/Pc47oVYauC4= +github.com/savsgio/gotils v0.0.0-20250408102913-196191ec6287/go.mod h1:sM7Mt7uEoCeFSCBM+qBrqvEo+/9vdmj19wzp3yzUhmg= github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w= github.com/sirupsen/logrus v1.9.4/go.mod h1:ftWc9WdOfJ0a92nsE2jF5u5ZwH8Bv2zdeOC42RjbV2g= +github.com/spf13/cast v1.10.0 h1:h2x0u2shc1QuLHfxi+cTJvs30+ZAHOGRic8uyGTDWxY= +github.com/spf13/cast v1.10.0/go.mod h1:jNfB8QC9IA6ZuY2ZjDp0KtFO2LZZlg4S/7bzP6qqeHo= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= @@ -109,39 +214,67 @@ github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXl github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= +github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= +github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= +github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= +github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= github.com/ugorji/go/codec v1.3.1 h1:waO7eEiFDwidsBN6agj1vJQ4AG7lh2yqXyOXqhgQuyY= github.com/ugorji/go/codec v1.3.1/go.mod h1:pRBVtBSKl77K30Bv8R2P+cLSGaTtex6fsA2Wjqmfxj4= +github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw= +github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= +github.com/valyala/fasthttp v1.71.0 h1:tepR7H+Guh9VUqxxcPggYi8R3lGUu2Rsdh+z7/FCY3k= +github.com/valyala/fasthttp v1.71.0/go.mod h1:z1sDUvOShhXq/C9mwH/fSm1Vb71tUJwmQdgkBrBNwnA= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= go.mongodb.org/mongo-driver/v2 v2.5.1 h1:j2U/Qp+wvueSpqitLCSZPT/+ZpVc1xzuwdHWwl7d8ro= go.mongodb.org/mongo-driver/v2 v2.5.1/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0= +go.starlark.net v0.0.0-20260102030733-3fee463870c9 h1:nV1OyvU+0CYrp5eKfQ3rD03TpFYYhH08z31NK1HmtTk= +go.starlark.net v0.0.0-20260102030733-3fee463870c9/go.mod h1:YKMCv9b1WrfWmeqdV5MAuEHWsu5iC+fe6kYl2sQjdI8= go.uber.org/dig v1.19.0 h1:BACLhebsYdpQ7IROQ1AGPjrXcP5dF80U3gKoFzbaq/4= go.uber.org/dig v1.19.0/go.mod h1:Us0rSJiThwCv2GteUN0Q7OKvU7n5J4dxZ9JKUXozFdE= go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU= go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= -golang.org/x/arch v0.22.0 h1:c/Zle32i5ttqRXjdLyyHZESLD/bB90DCU1g9l/0YBDI= -golang.org/x/arch v0.22.0/go.mod h1:dNHoOeKiyja7GTvF9NJS1l3Z2yntpQNzgrjh1cU103A= +golang.org/x/arch v0.23.0 h1:lKF64A2jF6Zd8L0knGltUnegD62JMFBiCPBmQpToHhg= +golang.org/x/arch v0.23.0/go.mod h1:dNHoOeKiyja7GTvF9NJS1l3Z2yntpQNzgrjh1cU103A= golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw= golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk= golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ= golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0= golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= +golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= +golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q= golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA= -google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= -google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/app/app.go b/internal/app/app.go index 208fbef4..0a9b077c 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -35,6 +35,7 @@ type App struct { startupBootstrap StartupBootstrap startupRecovery StartupRecovery requestLogs RequestLogRuntime + executionRuntime ExecutionRuntime listen func(network, address string) (net.Listener, error) mu sync.Mutex @@ -71,6 +72,12 @@ type RequestLogRuntime interface { Stop(context.Context) error } +// ExecutionRuntime owns provider-execution resources that must outlive all +// in-flight data-plane handlers and be released before process exit. +type ExecutionRuntime interface { + Shutdown() +} + // AppParams defines dependencies injected into App. type AppParams struct { dig.In @@ -85,6 +92,7 @@ type AppParams struct { Lifecycle *httplifecycle.Coordinator `optional:"true"` ControlRuntime ControlRuntime RequestLogs RequestLogRuntime + ExecutionRuntime ExecutionRuntime `optional:"true"` } // NewEngine creates the process HTTP engine and global middleware. @@ -125,6 +133,7 @@ func NewApp(params AppParams) *App { startupBootstrap: params.StartupBootstrap, startupRecovery: params.StartupRecovery, requestLogs: params.RequestLogs, + executionRuntime: params.ExecutionRuntime, listen: net.Listen, serveErrors: make(chan error, 1), } @@ -256,6 +265,7 @@ func (a *App) Stop(ctx context.Context) error { cancelRuntime := a.runtimeCancel runtimeDone := a.runtimeDone requestLogs := a.requestLogs + executionRuntime := a.executionRuntime runtimeCheckpoint := a.runtimeCheckpoint lifecycle := a.lifecycle a.mu.Unlock() @@ -335,6 +345,10 @@ func (a *App) Stop(ctx context.Context) error { logrus.WithField("event", "shutdown.request_log_drain").Info("request log runtime drained") } } + if executionRuntime != nil { + executionRuntime.Shutdown() + logrus.WithField("event", "shutdown.execution_runtime").Info("execution runtime stopped") + } if a.db != nil { sqlDB, err := a.db.DB() if err != nil { diff --git a/internal/app/app_test.go b/internal/app/app_test.go index b3981867..5a2c3082 100644 --- a/internal/app/app_test.go +++ b/internal/app/app_test.go @@ -59,6 +59,27 @@ type requestLogRuntimeFake struct { stopCalls atomic.Int32 } +type executionRuntimeFake struct { + shutdownCalls atomic.Int32 +} + +func (fake *executionRuntimeFake) Shutdown() { + fake.shutdownCalls.Add(1) +} + +func TestAppStopShutsDownExecutionRuntime(t *testing.T) { + t.Parallel() + + runtime := &executionRuntimeFake{} + application := NewApp(AppParams{ExecutionRuntime: runtime}) + if err := application.Stop(context.Background()); err != nil { + t.Fatalf("Stop() error = %v", err) + } + if got := runtime.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } +} + func newRequestLogRuntimeFake( startFunc func() error, stopFunc func(context.Context) error, @@ -309,7 +330,7 @@ func TestAppStartMigratesDatabaseAndServesHTTP(t *testing.T) { t.Fatalf("WriteTimeout = %s, want 0 for streaming responses", application.httpServer.WriteTimeout) } for _, table := range []string{ - "groups", "upstream_keys", "access_keys", "request_logs", "usage_stats", + "groups", "credentials", "access_keys", "request_logs", "usage_stats", "model_prices", "system_settings", "jobs", "control_operations", "schema_migrations", } { if !db.Migrator().HasTable(table) { @@ -1081,7 +1102,7 @@ func TestAppStopDeadlineJoinsRequestLogErrorAndClosesDatabase(t *testing.T) { func newTestRuntimeState(db *gorm.DB) (*state.Manager, *loader.Loader) { manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() return manager, loader.New(db, manager, registry) } diff --git a/internal/app/runtime_checkpoint.go b/internal/app/runtime_checkpoint.go index 8f52b28e..dd5e2dc7 100644 --- a/internal/app/runtime_checkpoint.go +++ b/internal/app/runtime_checkpoint.go @@ -22,8 +22,8 @@ type RuntimeStateCheckpoint interface { } type runtimeStateCheckpointDocument struct { - Keys []state.KeyRuntimeCheckpoint `json:"keys,omitempty"` - Stats []health.StatsRuntimeCheckpoint `json:"stats,omitempty"` + Credentials []state.CredentialRuntimeCheckpoint `json:"credentials,omitempty"` + Stats []health.StatsRuntimeCheckpoint `json:"stats,omitempty"` } // FileRuntimeStateCheckpoint stores the small, disposable runtime checkpoint @@ -31,14 +31,14 @@ type runtimeStateCheckpointDocument struct { // malformed or partially written file cannot be retried on every restart. type FileRuntimeStateCheckpoint struct { path string - registry *state.KeyRegistry + registry *state.CredentialRegistry stats *health.StatsStore removeFile func(string) error } func NewFileRuntimeStateCheckpoint( dataDir string, - registry *state.KeyRegistry, + registry *state.CredentialRegistry, stats *health.StatsStore, ) *FileRuntimeStateCheckpoint { return &FileRuntimeStateCheckpoint{ @@ -71,7 +71,7 @@ func (checkpoint *FileRuntimeStateCheckpoint) Restore(ctx context.Context) error return fmt.Errorf("decode runtime state checkpoint: %w", err) } if checkpoint.registry != nil { - checkpoint.registry.RestoreRuntimeCheckpoint(document.Keys) + checkpoint.registry.RestoreRuntimeCheckpoint(document.Credentials) } if checkpoint.stats != nil { checkpoint.stats.RestoreRuntimeCheckpoint(document.Stats) @@ -88,7 +88,7 @@ func (checkpoint *FileRuntimeStateCheckpoint) Save(ctx context.Context) error { } document := runtimeStateCheckpointDocument{} if checkpoint.registry != nil { - document.Keys = checkpoint.registry.CaptureRuntimeCheckpoint() + document.Credentials = checkpoint.registry.CaptureRuntimeCheckpoint() } if checkpoint.stats != nil { document.Stats = checkpoint.stats.CaptureRuntimeCheckpoint() diff --git a/internal/app/runtime_checkpoint_test.go b/internal/app/runtime_checkpoint_test.go index d88b49ac..53423eb6 100644 --- a/internal/app/runtime_checkpoint_test.go +++ b/internal/app/runtime_checkpoint_test.go @@ -96,7 +96,7 @@ func TestAppRestoresCheckpointAfterRuntimeRecovery(t *testing.T) { func TestAppLogsCheckpointRestoreFailureOnce(t *testing.T) { dataDir := t.TempDir() path := filepath.Join(dataDir, runtimeStateCheckpointFileName) - if err := os.WriteFile(path, []byte(`{"keys":[`), 0o600); err != nil { + if err := os.WriteFile(path, []byte(`{"credentials":[`), 0o600); err != nil { t.Fatalf("write malformed checkpoint fixture: %v", err) } @@ -184,9 +184,9 @@ func TestFileRuntimeStateCheckpointRestoresAndConsumesFile(t *testing.T) { dataDir := t.TempDir() path := filepath.Join(dataDir, runtimeStateCheckpointFileName) - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 10, Status: state.KeyStatusActive, WeightAuto: 37, + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, WeightAuto: 37, EncryptedValue: "cipher-one", }}); err != nil { t.Fatalf("replace registry: %v", err) @@ -205,12 +205,12 @@ func TestFileRuntimeStateCheckpointRestoresAndConsumesFile(t *testing.T) { t.Fatalf("checkpoint file stat error = %v", err) } - loadedRegistry := state.NewKeyRegistry() - if err := loadedRegistry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 10, Status: state.KeyStatusActive, WeightAuto: 1, + loadedRegistry := state.NewCredentialRegistry() + if err := loadedRegistry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, WeightAuto: 1, EncryptedValue: "cipher-one", }, { - ID: 2, GroupID: 20, Status: state.KeyStatusActive, WeightAuto: 2, + ID: 2, GroupID: 20, Version: 1, IdentityGeneration: 2, Fingerprint: "test-2", Status: state.CredentialStatusActive, WeightAuto: 2, EncryptedValue: "cipher-two", }}); err != nil { t.Fatalf("replace loaded registry: %v", err) @@ -239,7 +239,7 @@ func TestFileRuntimeStateCheckpointReturnsErrorWhenDeleteFails(t *testing.T) { dataDir := t.TempDir() path := filepath.Join(dataDir, runtimeStateCheckpointFileName) raw, err := json.Marshal(runtimeStateCheckpointDocument{ - Keys: []state.KeyRuntimeCheckpoint{{ID: 1, GroupID: 10, WeightAuto: 37}}, + Credentials: []state.CredentialRuntimeCheckpoint{{ID: 1, GroupID: 10, WeightAuto: 37}}, }) if err != nil { t.Fatalf("marshal checkpoint fixture: %v", err) @@ -248,9 +248,9 @@ func TestFileRuntimeStateCheckpointReturnsErrorWhenDeleteFails(t *testing.T) { t.Fatalf("write checkpoint fixture: %v", err) } - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 10, Status: state.KeyStatusActive, WeightAuto: 1, + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, WeightAuto: 1, EncryptedValue: "cipher-one", }}); err != nil { t.Fatalf("replace registry: %v", err) @@ -271,12 +271,12 @@ func TestFileRuntimeStateCheckpointReturnsErrorWhenDeleteFails(t *testing.T) { func TestFileRuntimeStateCheckpointConsumesMalformedFileAndReturnsError(t *testing.T) { dataDir := t.TempDir() path := filepath.Join(dataDir, runtimeStateCheckpointFileName) - if err := os.WriteFile(path, []byte(`{"keys":[`), 0o600); err != nil { + if err := os.WriteFile(path, []byte(`{"credentials":[`), 0o600); err != nil { t.Fatalf("write malformed checkpoint fixture: %v", err) } - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 10, Status: state.KeyStatusActive, WeightAuto: 1, + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, WeightAuto: 1, EncryptedValue: "cipher-one", }}); err != nil { t.Fatalf("replace registry: %v", err) diff --git a/internal/catalog/curated.go b/internal/catalog/curated.go deleted file mode 100644 index 5eba23e6..00000000 --- a/internal/catalog/curated.go +++ /dev/null @@ -1,149 +0,0 @@ -package catalog - -import ( - "strings" - - "gpt-load/internal/protocol" -) - -// curatedProvider is the maintainer-curated provider directory: services we -// choose to surface in the import picker beyond the three official presets -// and whatever models.dev happens to index. Entries here rank above -// models.dev catalog matches (see ListProviderSuggestions) and remain -// browsable with an empty search query, unlike the models.dev catalog. -// -// An ID may intentionally coincide with a models.dev provider ID (e.g. -// "deepseek"): the curated entry only overrides display metadata and -// priority, while pricing continues to key off the same provider scope and -// can still inherit models.dev cost data. -type curatedProvider struct { - ID string - Name string - APIURL string - Mark string - Protocols []protocol.Protocol -} - -// curatedProviders is the static seed table. IDs must be lowercase slugs -// accepted by pricing.ProviderScopeKey, APIURL must already be normalized -// per normalizeSuggestionURL, and Protocols must be non-empty data-plane -// protocols; TestCuratedProvidersSelfCheck enforces this at test time. -var curatedProviders = []curatedProvider{ - { - ID: "deepseek", - Name: "DeepSeek", - APIURL: "https://api.deepseek.com/v1", - Mark: "DS", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - }, - { - ID: "moonshotai", - Name: "Moonshot AI", - APIURL: "https://api.moonshot.cn/v1", - Mark: "MS", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - }, - { - ID: "siliconflow", - Name: "SiliconFlow", - APIURL: "https://api.siliconflow.cn/v1", - Mark: "SF", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "zhipuai", - Name: "Zhipu AI", - APIURL: "https://open.bigmodel.cn/api/paas/v4", - Mark: "ZP", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - }, - { - ID: "alibaba", - Name: "Alibaba Cloud Bailian", - APIURL: "https://dashscope.aliyuncs.com/compatible-mode/v1", - Mark: "BL", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "volcengine", - Name: "Volcengine Ark", - APIURL: "https://ark.cn-beijing.volces.com/api/v3", - Mark: "VE", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "openrouter", - Name: "OpenRouter", - APIURL: "https://openrouter.ai/api/v1", - Mark: "OR", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "groq", - Name: "Groq", - APIURL: "https://api.groq.com/openai/v1", - Mark: "GQ", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "xai", - Name: "xAI", - APIURL: "https://api.x.ai/v1", - Mark: "XA", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, -} - -var curatedProviderIDSet = buildCuratedProviderIDSet() - -func buildCuratedProviderIDSet() map[string]struct{} { - set := make(map[string]struct{}, len(curatedProviders)) - for _, provider := range curatedProviders { - set[provider.ID] = struct{}{} - } - return set -} - -// IsCuratedProviderID reports whether providerID is one of the maintainer -// curated presets, distinct from the three official provider IDs. -func IsCuratedProviderID(providerID string) bool { - _, ok := curatedProviderIDSet[providerID] - return ok -} - -// CuratedProviders returns cloned curated provider entries in fixed -// declaration order. Each call returns independent slice storage so callers -// may freely mutate the result. -func CuratedProviders() []Provider { - providers := make([]Provider, len(curatedProviders)) - for index, provider := range curatedProviders { - providers[index] = Provider{ - ID: provider.ID, - Name: provider.Name, - APIURL: provider.APIURL, - Mark: provider.Mark, - Protocols: append([]protocol.Protocol(nil), provider.Protocols...), - } - } - return providers -} - -// SearchCurated returns curated providers matching query by case-insensitive -// substring match against ID or Name, in fixed declaration order. An empty -// query returns the full table so the curated directory stays browsable -// without typing, unlike the unbounded models.dev catalog. -func SearchCurated(query string) []Provider { - normalizedQuery := strings.ToLower(strings.TrimSpace(query)) - all := CuratedProviders() - if normalizedQuery == "" { - return all - } - matched := make([]Provider, 0, len(all)) - for _, provider := range all { - if strings.Contains(strings.ToLower(provider.ID), normalizedQuery) || - strings.Contains(strings.ToLower(provider.Name), normalizedQuery) { - matched = append(matched, provider) - } - } - return matched -} diff --git a/internal/catalog/curated_test.go b/internal/catalog/curated_test.go deleted file mode 100644 index d5adb8c5..00000000 --- a/internal/catalog/curated_test.go +++ /dev/null @@ -1,251 +0,0 @@ -package catalog - -import ( - "slices" - "strings" - "testing" - - "gpt-load/internal/pricing" - "gpt-load/internal/protocol" -) - -func TestCuratedProvidersSelfCheck(t *testing.T) { - providers := CuratedProviders() - if len(providers) == 0 { - t.Fatal("CuratedProviders() is empty, want at least one curated entry") - } - - seenIDs := make(map[string]struct{}, len(providers)) - for _, provider := range providers { - if _, err := pricing.ProviderScopeKey(provider.ID); err != nil { - t.Errorf("curated provider %q: ID is not a valid price scope key: %v", provider.ID, err) - } - if IsOfficialProviderID(provider.ID) { - t.Errorf("curated provider %q collides with an official provider ID", provider.ID) - } - if _, duplicate := seenIDs[provider.ID]; duplicate { - t.Errorf("curated provider ID %q is duplicated", provider.ID) - } - seenIDs[provider.ID] = struct{}{} - - if strings.TrimSpace(provider.Name) == "" { - t.Errorf("curated provider %q: Name must not be blank", provider.ID) - } - if strings.TrimSpace(provider.Mark) == "" { - t.Errorf("curated provider %q: Mark must not be blank", provider.ID) - } - - if normalized, ok := normalizeSuggestionURL(provider.APIURL); !ok { - t.Errorf("curated provider %q: APIURL %q is not a valid suggestion URL", provider.ID, provider.APIURL) - } else if normalized != provider.APIURL { - t.Errorf("curated provider %q: APIURL %q is not already normalized (want %q)", provider.ID, provider.APIURL, normalized) - } - - if len(provider.Protocols) == 0 { - t.Errorf("curated provider %q: Protocols must not be empty", provider.ID) - } - seenProtocols := make(map[string]struct{}, len(provider.Protocols)) - for _, proto := range provider.Protocols { - if !proto.DataPlaneEnabled() { - t.Errorf("curated provider %q: protocol %q is not data-plane enabled", provider.ID, proto) - } - if _, duplicate := seenProtocols[string(proto)]; duplicate { - t.Errorf("curated provider %q: protocol %q is duplicated", provider.ID, proto) - } - seenProtocols[string(proto)] = struct{}{} - } - } -} - -func TestCuratedProvidersDeclaration(t *testing.T) { - want := []Provider{ - { - ID: "deepseek", - Name: "DeepSeek", - APIURL: "https://api.deepseek.com/v1", - Mark: "DS", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - }, - { - ID: "moonshotai", - Name: "Moonshot AI", - APIURL: "https://api.moonshot.cn/v1", - Mark: "MS", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - }, - { - ID: "siliconflow", - Name: "SiliconFlow", - APIURL: "https://api.siliconflow.cn/v1", - Mark: "SF", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "zhipuai", - Name: "Zhipu AI", - APIURL: "https://open.bigmodel.cn/api/paas/v4", - Mark: "ZP", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - }, - { - ID: "alibaba", - Name: "Alibaba Cloud Bailian", - APIURL: "https://dashscope.aliyuncs.com/compatible-mode/v1", - Mark: "BL", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "volcengine", - Name: "Volcengine Ark", - APIURL: "https://ark.cn-beijing.volces.com/api/v3", - Mark: "VE", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "openrouter", - Name: "OpenRouter", - APIURL: "https://openrouter.ai/api/v1", - Mark: "OR", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "groq", - Name: "Groq", - APIURL: "https://api.groq.com/openai/v1", - Mark: "GQ", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - { - ID: "xai", - Name: "xAI", - APIURL: "https://api.x.ai/v1", - Mark: "XA", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - }, - } - - got := CuratedProviders() - if len(got) != len(want) { - t.Fatalf("CuratedProviders() returned %d providers, want %d", len(got), len(want)) - } - for index, wantProvider := range want { - gotProvider := got[index] - if gotProvider.ID != wantProvider.ID { - t.Errorf("provider[%d].ID = %q, want %q", index, gotProvider.ID, wantProvider.ID) - } - if gotProvider.Name != wantProvider.Name { - t.Errorf("provider[%d].Name = %q, want %q", index, gotProvider.Name, wantProvider.Name) - } - if gotProvider.APIURL != wantProvider.APIURL { - t.Errorf("provider[%d].APIURL = %q, want %q", index, gotProvider.APIURL, wantProvider.APIURL) - } - if gotProvider.Mark != wantProvider.Mark { - t.Errorf("provider[%d].Mark = %q, want %q", index, gotProvider.Mark, wantProvider.Mark) - } - if !slices.Equal(gotProvider.Protocols, wantProvider.Protocols) { - t.Errorf("provider[%d].Protocols = %q, want %q", index, gotProvider.Protocols, wantProvider.Protocols) - } - } -} - -func TestCuratedProvidersUseModelsDevProviderIDs(t *testing.T) { - wantIDs := map[string]string{ - "Moonshot AI": "moonshotai", - "Zhipu AI": "zhipuai", - "Alibaba Cloud Bailian": "alibaba", - } - found := make(map[string]struct{}, len(wantIDs)) - for _, provider := range CuratedProviders() { - wantID, tracked := wantIDs[provider.Name] - if !tracked { - continue - } - found[provider.Name] = struct{}{} - if provider.ID != wantID { - t.Errorf("curated provider %q ID = %q, want %q", provider.Name, provider.ID, wantID) - } - if _, err := pricing.ProviderScopeKey(wantID); err != nil { - t.Errorf("models.dev provider ID %q has an invalid price scope: %v", wantID, err) - } - } - for name := range wantIDs { - if _, ok := found[name]; !ok { - t.Errorf("curated provider %q is missing", name) - } - } -} - -func TestVolcengineRemainsCuratedOutsideAutomaticPricePriority(t *testing.T) { - if !IsCuratedProviderID("volcengine") { - t.Fatal("IsCuratedProviderID(\"volcengine\") = false, want true") - } - if slices.Contains(AutomaticPriceProviderPriority(), "volcengine") { - t.Error("AutomaticPriceProviderPriority() contains volcengine, want false") - } -} - -func TestCuratedProvidersReturnsIndependentClones(t *testing.T) { - first := CuratedProviders() - if len(first) == 0 { - t.Fatal("CuratedProviders() is empty") - } - first[0].Protocols[0] = "mutated" - second := CuratedProviders() - if second[0].Protocols[0] == "mutated" { - t.Fatal("CuratedProviders() shares mutable protocol storage across calls") - } -} - -func TestSearchCuratedEmptyQueryReturnsFullTable(t *testing.T) { - all := CuratedProviders() - got := SearchCurated("") - if len(got) != len(all) { - t.Fatalf("SearchCurated(\"\") returned %d providers, want full table of %d", len(got), len(all)) - } -} - -func TestSearchCuratedMatchesIDOrNameCaseInsensitively(t *testing.T) { - all := CuratedProviders() - if len(all) == 0 { - t.Fatal("CuratedProviders() is empty") - } - target := all[0] - - byID := SearchCurated(strings.ToUpper(target.ID)) - if !containsProviderID(byID, target.ID) { - t.Fatalf("SearchCurated(%q) = %#v, want it to contain %q", strings.ToUpper(target.ID), byID, target.ID) - } - - byName := SearchCurated(strings.ToUpper(target.Name)) - if !containsProviderID(byName, target.ID) { - t.Fatalf("SearchCurated(%q) = %#v, want it to contain %q", strings.ToUpper(target.Name), byName, target.ID) - } - - none := SearchCurated("no-such-curated-provider-xyz") - if len(none) != 0 { - t.Fatalf("SearchCurated(no-match) = %#v, want empty", none) - } -} - -func TestIsCuratedProviderID(t *testing.T) { - for _, provider := range CuratedProviders() { - if !IsCuratedProviderID(provider.ID) { - t.Errorf("IsCuratedProviderID(%q) = false, want true", provider.ID) - } - } - if IsCuratedProviderID("openai") { - t.Error("IsCuratedProviderID(\"openai\") = true, want false (official, not curated)") - } - if IsCuratedProviderID("no-such-id") { - t.Error("IsCuratedProviderID(\"no-such-id\") = true, want false") - } -} - -func containsProviderID(providers []Provider, id string) bool { - for _, provider := range providers { - if provider.ID == id { - return true - } - } - return false -} diff --git a/internal/catalog/official.go b/internal/catalog/official.go deleted file mode 100644 index addd5f97..00000000 --- a/internal/catalog/official.go +++ /dev/null @@ -1,247 +0,0 @@ -package catalog - -import ( - "net/url" - "sort" - "strings" - "unicode" - - "gpt-load/internal/protocol" -) - -var officialProviderIDs = map[string]struct{}{ - "openai": {}, - "anthropic": {}, - "google": {}, -} - -// IsOfficialProviderID reports whether providerID is one of the built-in -// provider presets that remain selectable without a loaded catalog snapshot. -func IsOfficialProviderID(providerID string) bool { - _, ok := officialProviderIDs[providerID] - return ok -} - -// OfficialSuggestions returns the three local provider presets with safe -// field-level name and API URL updates from the current catalog. -func OfficialSuggestions(snapshot *Snapshot) []Provider { - providers := []Provider{ - {ID: "openai", Name: "OpenAI", APIURL: "https://api.openai.com/v1", Mark: "OA", Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.OpenAIResponses}}, - {ID: "anthropic", Name: "Anthropic", APIURL: "https://api.anthropic.com/v1", Mark: "AN", Protocols: []protocol.Protocol{protocol.Anthropic}}, - {ID: "google", Name: "Google", APIURL: "https://generativelanguage.googleapis.com/v1beta", Mark: "GE", Protocols: []protocol.Protocol{protocol.Gemini}}, - } - if snapshot == nil { - return providers - } - for index := range providers { - catalogProvider, ok := snapshot.Providers[providers[index].ID] - if !ok { - continue - } - if name := strings.TrimSpace(catalogProvider.Name); name != "" { - providers[index].Name = name - } - if apiURL, ok := normalizeSuggestionURL(catalogProvider.APIURL); ok { - providers[index].APIURL = apiURL - } - } - return providers -} - -// SearchProviders returns cloned non-official catalog suggestions in stable -// display order with only conservative protocol inference. -func SearchProviders(snapshot *Snapshot, query string) []Provider { - if snapshot == nil || len(snapshot.Providers) == 0 { - return nil - } - normalizedQuery := strings.ToLower(strings.TrimSpace(query)) - providers := make([]Provider, 0, len(snapshot.Providers)) - for id, provider := range snapshot.Providers { - if _, official := officialProviderIDs[id]; official { - continue - } - if normalizedQuery != "" && - !strings.Contains(strings.ToLower(provider.ID), normalizedQuery) && - !strings.Contains(strings.ToLower(provider.Name), normalizedQuery) { - continue - } - provider = cloneProvider(provider) - provider.Mark = providerMark(provider.Name, provider.ID) - provider.Protocols = inferProtocols(provider.NPM) - if apiURL, ok := normalizeSuggestionURL(provider.APIURL); ok { - provider.APIURL = apiURL - } else { - provider.APIURL = "" - } - providers = append(providers, provider) - } - sort.Slice(providers, func(left, right int) bool { - leftName := strings.ToLower(strings.TrimSpace(providers[left].Name)) - rightName := strings.ToLower(strings.TrimSpace(providers[right].Name)) - if leftName == rightName { - return providers[left].ID < providers[right].ID - } - return leftName < rightName - }) - return providers -} - -// SearchProviderMetadataBounded returns a deterministic top-N metadata result -// without cloning provider model maps. Empty queries deliberately return no -// catalog-only providers; callers should present local official suggestions -// until the user searches. -func SearchProviderMetadataBounded(snapshot *Snapshot, query string, limit int) []Provider { - if snapshot == nil || limit <= 0 { - return nil - } - normalizedQuery := strings.ToLower(strings.TrimSpace(query)) - if normalizedQuery == "" { - return nil - } - type candidate struct { - id string - nameSort string - provider Provider - } - selected := make([]candidate, 0, limit) - less := func(left, right candidate) bool { - if left.nameSort == right.nameSort { - return left.id < right.id - } - return left.nameSort < right.nameSort - } - for id, provider := range snapshot.Providers { - if _, official := officialProviderIDs[id]; official { - continue - } - if !strings.Contains(strings.ToLower(provider.ID), normalizedQuery) && - !strings.Contains(strings.ToLower(provider.Name), normalizedQuery) { - continue - } - entry := candidate{ - id: id, nameSort: strings.ToLower(strings.TrimSpace(provider.Name)), provider: provider, - } - position := sort.Search(len(selected), func(index int) bool { - return !less(selected[index], entry) - }) - if position >= limit { - continue - } - selected = append(selected, candidate{}) - copy(selected[position+1:], selected[position:]) - selected[position] = entry - if len(selected) > limit { - selected = selected[:limit] - } - } - providers := make([]Provider, 0, len(selected)) - for _, entry := range selected { - providers = append(providers, providerMetadata(entry.provider)) - } - return providers -} - -// SearchProviderMetadataByIDs returns the requested Models.dev providers in -// caller order. Missing IDs are omitted, and each result contains only the -// metadata needed by the provider directory. -func SearchProviderMetadataByIDs(snapshot *Snapshot, ids []string) []Provider { - if snapshot == nil || len(ids) == 0 { - return nil - } - providers := make([]Provider, 0, len(ids)) - seen := make(map[string]struct{}, len(ids)) - for _, id := range ids { - if _, duplicate := seen[id]; duplicate { - continue - } - seen[id] = struct{}{} - provider, ok := snapshot.Providers[id] - if !ok { - continue - } - providers = append(providers, providerMetadata(provider)) - } - return providers -} - -func providerMetadata(provider Provider) Provider { - provider.Models = nil - provider.Protocols = inferProtocols(provider.NPM) - provider.Mark = providerMark(provider.Name, provider.ID) - if apiURL, ok := normalizeSuggestionURL(provider.APIURL); ok { - provider.APIURL = apiURL - } else { - provider.APIURL = "" - } - return provider -} - -// SearchProvidersBounded is retained as the metadata-only bounded search -// boundary used by provider suggestion callers. -func SearchProvidersBounded(snapshot *Snapshot, query string, limit int) []Provider { - return SearchProviderMetadataBounded(snapshot, query, limit) -} - -func normalizeSuggestionURL(value string) (string, bool) { - value = strings.TrimSpace(value) - parsed, err := url.Parse(value) - if err != nil || parsed == nil || parsed.Host == "" || parsed.User != nil || - (parsed.Scheme != "http" && parsed.Scheme != "https") || - parsed.RawQuery != "" || parsed.Fragment != "" || parsed.Opaque != "" { - return "", false - } - parsed.Path = strings.TrimRight(parsed.Path, "/") - parsed.RawPath = strings.TrimRight(parsed.RawPath, "/") - return parsed.String(), true -} - -func inferProtocols(npm string) []protocol.Protocol { - switch strings.ToLower(strings.TrimSpace(npm)) { - case "@ai-sdk/anthropic": - return []protocol.Protocol{protocol.Anthropic} - case "@ai-sdk/google", "@ai-sdk/google-vertex": - return []protocol.Protocol{protocol.Gemini} - case "@ai-sdk/openai", "@ai-sdk/openai-compatible", "@openrouter/ai-sdk-provider": - return []protocol.Protocol{protocol.OpenAICompletions} - default: - return nil - } -} - -func providerMark(name, id string) string { - value := strings.TrimSpace(name) - if value == "" { - value = id - } - mark := markFromText(value) - if mark == "" && value != id { - mark = markFromText(id) - } - return mark -} - -func markFromText(value string) string { - words := strings.FieldsFunc(value, func(character rune) bool { - return !unicode.IsLetter(character) && !unicode.IsDigit(character) - }) - mark := make([]rune, 0, 2) - if len(words) > 1 { - for _, word := range words { - for _, character := range word { - mark = append(mark, unicode.ToUpper(character)) - break - } - if len(mark) == 2 { - break - } - } - } else if len(words) == 1 { - for _, character := range words[0] { - mark = append(mark, unicode.ToUpper(character)) - if len(mark) == 2 { - break - } - } - } - return string(mark) -} diff --git a/internal/catalog/official_test.go b/internal/catalog/official_test.go deleted file mode 100644 index b93d42b4..00000000 --- a/internal/catalog/official_test.go +++ /dev/null @@ -1,174 +0,0 @@ -package catalog - -import ( - "reflect" - "testing" - - "gpt-load/internal/protocol" -) - -func TestOfficialSuggestionsUseOrderedFallbacksAndFieldLevelOverrides(t *testing.T) { - snapshot := &Snapshot{Providers: map[string]Provider{ - "openai": { - ID: "openai", - Name: "OpenAI Catalog", - APIURL: "https://catalog.openai.example/v1///", - NPM: "catalog-must-not-replace-local-preset", - Mark: "XX", - Protocols: []protocol.Protocol{protocol.Gemini}, - Models: map[string]Model{"gpt-x": {ID: "gpt-x", Name: "GPT X"}}, - }, - "anthropic": { - ID: "anthropic", - Name: " ", - APIURL: "https://user:secret@api.example.com/v1?unsafe=true", - }, - "google": { - ID: "google", - Name: "Google Catalog", - APIURL: "http://google.example.test/v1beta/", - }, - }} - - got := OfficialSuggestions(snapshot) - want := []Provider{ - {ID: "openai", Name: "OpenAI Catalog", APIURL: "https://catalog.openai.example/v1", Mark: "OA", Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.OpenAIResponses}}, - {ID: "anthropic", Name: "Anthropic", APIURL: "https://api.anthropic.com/v1", Mark: "AN", Protocols: []protocol.Protocol{protocol.Anthropic}}, - {ID: "google", Name: "Google Catalog", APIURL: "http://google.example.test/v1beta", Mark: "GE", Protocols: []protocol.Protocol{protocol.Gemini}}, - } - if !reflect.DeepEqual(got, want) { - t.Fatalf("OfficialSuggestions() = %#v, want %#v", got, want) - } - - got[0].Protocols[0] = protocol.Gemini - again := OfficialSuggestions(snapshot) - if again[0].Protocols[0] != protocol.OpenAICompletions { - t.Fatalf("official suggestions share mutable protocol storage: %#v", again[0].Protocols) - } -} - -func TestOfficialSuggestionsReturnFallbacksForNilOrEmptySnapshot(t *testing.T) { - for _, snapshot := range []*Snapshot{nil, {}, {Providers: map[string]Provider{}}} { - got := OfficialSuggestions(snapshot) - if len(got) != 3 || got[0].ID != "openai" || got[1].ID != "anthropic" || got[2].ID != "google" { - t.Fatalf("OfficialSuggestions(%#v) = %#v", snapshot, got) - } - } -} - -func TestIsOfficialProviderID(t *testing.T) { - for _, providerID := range []string{"openai", "anthropic", "google"} { - if !IsOfficialProviderID(providerID) { - t.Fatalf("IsOfficialProviderID(%q) = false, want true", providerID) - } - } - for _, providerID := range []string{"", "OpenAI", "private-company"} { - if IsOfficialProviderID(providerID) { - t.Fatalf("IsOfficialProviderID(%q) = true, want false", providerID) - } - } -} - -func TestSearchProvidersExcludesOfficialSortsAndInfersConservatively(t *testing.T) { - snapshot := &Snapshot{Providers: map[string]Provider{ - "openai": {ID: "openai", Name: "OpenAI"}, - "anthropic-cloud": {ID: "anthropic-cloud", Name: "Zeta Anthropic", NPM: "@ai-sdk/anthropic", Models: map[string]Model{"claude": {ID: "claude", Name: "Claude"}}}, - "gemini-cloud": {ID: "gemini-cloud", Name: "alpha Google", NPM: "@ai-sdk/google"}, - "compatible": {ID: "compatible", Name: "Alpha OpenAI Compatible", NPM: "@ai-sdk/openai-compatible"}, - "openrouter": {ID: "openrouter", Name: "Beta OpenRouter", NPM: "@openrouter/ai-sdk-provider"}, - "openai-sdk": {ID: "openai-sdk", Name: "Delta OpenAI SDK", NPM: "@ai-sdk/openai"}, - "unknown": {ID: "unknown", Name: "Gamma Unknown", NPM: "@vendor/unknown", Protocols: []protocol.Protocol{protocol.OpenAIResponses}}, - }} - - got := SearchProviders(snapshot, "") - wantIDs := []string{"gemini-cloud", "compatible", "openrouter", "openai-sdk", "unknown", "anthropic-cloud"} - if len(got) != len(wantIDs) { - t.Fatalf("SearchProviders() count = %d, want %d: %#v", len(got), len(wantIDs), got) - } - for index, wantID := range wantIDs { - if got[index].ID != wantID { - t.Fatalf("SearchProviders()[%d].ID = %q, want %q", index, got[index].ID, wantID) - } - if got[index].Mark == "" { - t.Fatalf("SearchProviders()[%d].Mark is empty", index) - } - } - assertProtocols(t, got[0].Protocols, protocol.Gemini) - assertProtocols(t, got[1].Protocols, protocol.OpenAICompletions) - assertProtocols(t, got[2].Protocols, protocol.OpenAICompletions) - assertProtocols(t, got[3].Protocols, protocol.OpenAICompletions) - assertProtocols(t, got[4].Protocols) - assertProtocols(t, got[5].Protocols, protocol.Anthropic) - for _, provider := range got { - for _, inferred := range provider.Protocols { - if inferred == protocol.OpenAIResponses { - t.Fatalf("non-official provider %q inferred responses", provider.ID) - } - } - } - - filtered := SearchProviders(snapshot, "gOoGlE") - if len(filtered) != 1 || filtered[0].ID != "gemini-cloud" { - t.Fatalf("case-insensitive name search = %#v", filtered) - } - filtered = SearchProviders(snapshot, "OPENROUTER") - if len(filtered) != 1 || filtered[0].ID != "openrouter" { - t.Fatalf("case-insensitive ID search = %#v", filtered) - } - - got[len(got)-1].Models["claude"] = Model{ID: "mutated"} - again := SearchProviders(snapshot, "anthropic") - if again[0].Models["claude"].ID != "claude" { - t.Fatalf("search result shares model map: %#v", again[0].Models) - } -} - -func TestSearchProvidersReturnsNoMoreProvidersForNilSnapshot(t *testing.T) { - if got := SearchProviders(nil, ""); len(got) != 0 { - t.Fatalf("SearchProviders(nil) = %#v, want empty", got) - } -} - -func TestSearchProvidersAlwaysGeneratesLocalTextMark(t *testing.T) { - snapshot := &Snapshot{Providers: map[string]Provider{ - "symbols": {ID: "symbols", Name: "🔥", Models: map[string]Model{}}, - }} - got := SearchProviders(snapshot, "") - if len(got) != 1 || got[0].Mark != "SY" { - t.Fatalf("SearchProviders() mark = %#v, want local ID fallback mark SY", got) - } -} - -func TestSearchProviderMetadataByIDsPreservesRequestedOrderAndOnlyMatches(t *testing.T) { - snapshot := &Snapshot{Providers: map[string]Provider{ - "alpha": { - ID: "alpha", Name: "Alpha", APIURL: "https://alpha.example/v1///", - NPM: "@ai-sdk/openai-compatible", Models: map[string]Model{"alpha-model": {ID: "alpha-model"}}, - }, - "beta": { - ID: "beta", Name: "Beta", APIURL: "https://beta.example/v1", - NPM: "@ai-sdk/anthropic", Models: map[string]Model{"beta-model": {ID: "beta-model"}}, - }, - }} - - got := SearchProviderMetadataByIDs(snapshot, []string{"beta", "missing", "beta", "alpha"}) - if len(got) != 2 || got[0].ID != "beta" || got[1].ID != "alpha" { - t.Fatalf("metadata by IDs = %#v, want beta then alpha", got) - } - if got[0].APIURL != "https://beta.example/v1" || got[0].Protocols[0] != protocol.Anthropic || got[0].Mark != "BE" { - t.Fatalf("beta metadata = %#v", got[0]) - } - if got[1].APIURL != "https://alpha.example/v1" || got[1].Protocols[0] != protocol.OpenAICompletions || got[1].Mark != "AL" { - t.Fatalf("alpha metadata = %#v", got[1]) - } - if got[0].Models != nil || got[1].Models != nil { - t.Fatalf("metadata by IDs exposes model maps: %#v", got) - } -} - -func assertProtocols(t *testing.T, got []protocol.Protocol, want ...protocol.Protocol) { - t.Helper() - if !reflect.DeepEqual(got, want) { - t.Fatalf("protocols = %#v, want %#v", got, want) - } -} diff --git a/internal/catalog/parse.go b/internal/catalog/parse.go index 2251092a..7f167447 100644 --- a/internal/catalog/parse.go +++ b/internal/catalog/parse.go @@ -97,20 +97,14 @@ func parseProvider(key string, raw json.RawMessage) (Provider, error) { return Provider{}, err } - apiURL, err := decodeStringField(fields, "api") - if err != nil { + if _, err := decodeStringField(fields, "api"); err != nil { return Provider{}, err } - npm, err := decodeStringField(fields, "npm") - if err != nil { + if _, err := decodeStringField(fields, "npm"); err != nil { return Provider{}, err } provider := Provider{ - ID: id, - Name: name, - APIURL: strings.TrimSpace(apiURL), - NPM: strings.TrimSpace(npm), - Models: make(map[string]Model, len(modelValues)), + ID: id, Name: name, Models: make(map[string]Model, len(modelValues)), } for modelKey, rawModel := range modelValues { model, err := parseModel(modelKey, rawModel) diff --git a/internal/catalog/parse_test.go b/internal/catalog/parse_test.go index b1c39b27..0bfc18e3 100644 --- a/internal/catalog/parse_test.go +++ b/internal/catalog/parse_test.go @@ -54,8 +54,7 @@ func TestParseModelsDevUsesExactDecimalsAndCanonicalTiers(t *testing.T) { snapshot := mustParse(t, raw) provider := snapshot.Providers["openai"] - if provider.ID != "openai" || provider.Name != "OpenAI" || - provider.APIURL != "https://api.openai.com/v1/" || provider.NPM != "@ai-sdk/openai" { + if provider.ID != "openai" || provider.Name != "OpenAI" { t.Fatalf("provider = %#v", provider) } diff --git a/internal/catalog/runtime.go b/internal/catalog/runtime.go index c7289810..fec51b35 100644 --- a/internal/catalog/runtime.go +++ b/internal/catalog/runtime.go @@ -24,40 +24,6 @@ func (runtime *Runtime) HasGeneration() bool { return runtime != nil && runtime.snapshot.Load() != nil } -// OfficialSuggestions reads only provider metadata from the current immutable -// generation. -func (runtime *Runtime) OfficialSuggestions() []Provider { - if runtime == nil { - return OfficialSuggestions(nil) - } - return OfficialSuggestions(runtime.snapshot.Load()) -} - -// SearchProviderMetadata returns bounded provider metadata without copying any -// provider model maps. -func (runtime *Runtime) SearchProviderMetadata(query string, limit int) []Provider { - if runtime == nil { - return nil - } - return SearchProviderMetadataBounded(runtime.snapshot.Load(), query, limit) -} - -// LoadProvider returns a caller-owned clone of one provider only. -func (runtime *Runtime) LoadProvider(providerID string) (Provider, bool) { - if runtime == nil { - return Provider{}, false - } - snapshot := runtime.snapshot.Load() - if snapshot == nil { - return Provider{}, false - } - provider, ok := snapshot.Providers[providerID] - if !ok { - return Provider{}, false - } - return cloneProvider(provider), true -} - // Publish deep-clones and atomically replaces the current snapshot. func (runtime *Runtime) Publish(snapshot *Snapshot) { if runtime == nil { diff --git a/internal/catalog/runtime_test.go b/internal/catalog/runtime_test.go index af654908..f26061a7 100644 --- a/internal/catalog/runtime_test.go +++ b/internal/catalog/runtime_test.go @@ -7,7 +7,6 @@ import ( "testing" "gpt-load/internal/pricing" - "gpt-load/internal/protocol" ) func TestRuntimePublishAndLoadDeepCloneEveryMutableLayer(t *testing.T) { @@ -17,9 +16,8 @@ func TestRuntimePublishAndLoadDeepCloneEveryMutableLayer(t *testing.T) { contextLimit := int64(1_000_000) source := &Snapshot{Providers: map[string]Provider{ "openai": { - ID: "openai", - Name: "OpenAI", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: "openai", + Name: "OpenAI", Models: map[string]Model{ "gpt-x": { ID: "gpt-x", @@ -50,7 +48,6 @@ func TestRuntimePublishAndLoadDeepCloneEveryMutableLayer(t *testing.T) { runtime.Publish(source) provider := source.Providers["openai"] - provider.Protocols[0] = protocol.Gemini model := provider.Models["gpt-x"] *model.Metadata.Capabilities.Attachment = false model.Metadata.Modalities.Input[0] = "audio" @@ -66,7 +63,6 @@ func TestRuntimePublishAndLoadDeepCloneEveryMutableLayer(t *testing.T) { first := runtime.Load() assertRuntimeSnapshot(t, first) loadedProvider := first.Providers["openai"] - loadedProvider.Protocols[0] = protocol.Gemini loadedModel := loadedProvider.Models["gpt-x"] *loadedModel.Metadata.Capabilities.Attachment = false loadedModel.Metadata.Modalities.Input[0] = "audio" @@ -121,43 +117,11 @@ func TestRuntimeConcurrentLoadPublishReturnsWholeSnapshots(t *testing.T) { } } -func TestRuntimeProviderReadBoundariesCopyOnlyRequestedData(t *testing.T) { - runtime := &Runtime{} - runtime.Publish(&Snapshot{Providers: map[string]Provider{ - "alpha": { - ID: "alpha", Name: "Alpha", NPM: "@ai-sdk/openai-compatible", - Models: map[string]Model{"alpha-model": {ID: "alpha-model", Name: "Alpha Model"}}, - }, - "beta": { - ID: "beta", Name: "Beta", - Models: map[string]Model{"beta-model": {ID: "beta-model", Name: "Beta Model"}}, - }, - }}) - - metadata := runtime.SearchProviderMetadata("a", 1) - if len(metadata) != 1 || metadata[0].ID != "alpha" || metadata[0].Models != nil { - t.Fatalf("metadata search = %#v, want one provider without Models", metadata) - } - metadata[0].Name = "mutated" - provider, ok := runtime.LoadProvider("alpha") - if !ok || provider.Name != "Alpha" || len(provider.Models) != 1 { - t.Fatalf("single provider = %#v/%t", provider, ok) - } - provider.Models["other"] = Model{ID: "other"} - again, ok := runtime.LoadProvider("alpha") - if !ok || len(again.Models) != 1 { - t.Fatalf("single provider read exposed runtime state: %#v/%t", again, ok) - } - if _, ok := runtime.LoadProvider("missing"); ok { - t.Fatal("missing provider reported present") - } -} - func assertRuntimeSnapshot(t *testing.T, snapshot *Snapshot) { t.Helper() provider := snapshot.Providers["openai"] model := provider.Models["gpt-x"] - if provider.Name != "OpenAI" || provider.Protocols[0] != protocol.OpenAICompletions || + if provider.Name != "OpenAI" || model.Metadata.Capabilities.Attachment == nil || !*model.Metadata.Capabilities.Attachment || !reflect.DeepEqual(model.Metadata.Modalities.Input, []string{"text", "image"}) || !reflect.DeepEqual(model.Metadata.Modalities.Output, []string{"text"}) || diff --git a/internal/catalog/types.go b/internal/catalog/types.go index 57d6aaaa..6d4f230d 100644 --- a/internal/catalog/types.go +++ b/internal/catalog/types.go @@ -5,7 +5,6 @@ import ( "encoding/json" "gpt-load/internal/pricing" - "gpt-load/internal/protocol" ) // ModelCost is the exact provider price data retained for one model. @@ -62,16 +61,11 @@ type Model struct { Cost *ModelCost } -// Provider is the retained Models.dev provider identity and model map. Mark and -// Protocols are local suggestion metadata and are never read from Models.dev. +// Provider is the retained Models.dev provider identity and model map. type Provider struct { - ID string - Name string - APIURL string - NPM string - Mark string - Protocols []protocol.Protocol - Models map[string]Model + ID string + Name string + Models map[string]Model } // Snapshot is one complete validated catalog generation. @@ -114,7 +108,6 @@ func cloneSnapshot(snapshot *Snapshot) *Snapshot { } func cloneProvider(provider Provider) Provider { - provider.Protocols = append([]protocol.Protocol(nil), provider.Protocols...) if provider.Models != nil { models := make(map[string]Model, len(provider.Models)) for id, model := range provider.Models { diff --git a/internal/channel/channel.go b/internal/channel/channel.go new file mode 100644 index 00000000..44a98633 --- /dev/null +++ b/internal/channel/channel.go @@ -0,0 +1,657 @@ +// Package channel defines the code-owned directory of supported upstream channels. +package channel + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "net/url" + "strings" + + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +// ID is a stable product channel identifier. +type ID string + +const ( + OpenAI ID = "openai" + Anthropic ID = "anthropic" + Gemini ID = "gemini" + AzureOpenAI ID = "azure_openai" + AWSBedrock ID = "aws_bedrock" + GoogleVertex ID = "google_vertex" + OpenAICompatible ID = "openai_compatible" + AnthropicCompatible ID = "anthropic_compatible" + GeminiCompatible ID = "gemini_compatible" + DeepSeek ID = "deepseek" + MoonshotAI ID = "moonshotai" + SiliconFlow ID = "siliconflow" + ZhipuAI ID = "zhipuai" + Alibaba ID = "alibaba" + Volcengine ID = "volcengine" + OpenRouter ID = "openrouter" + Groq ID = "groq" + XAI ID = "xai" +) + +// InputKind describes how a field is collected without exposing its value. +type InputKind string + +const ( + InputText InputKind = "text" + InputURL InputKind = "url" + InputSecret InputKind = "secret" +) + +// FieldDescriptor is the public, value-free schema for one channel field. +type FieldDescriptor struct { + Key string `json:"key"` + Label string `json:"label"` + InputKind InputKind `json:"input_kind"` + Required bool `json:"required"` + Sensitive bool `json:"sensitive"` +} + +// Descriptor contains only information safe to expose to the management UI. +type Descriptor struct { + ID ID `json:"channel_id"` + Name string `json:"name"` + Mark string `json:"mark"` + Description string `json:"description"` + ParamFields []FieldDescriptor `json:"param_fields"` + CredentialFields []FieldDescriptor `json:"credential_fields"` + ClientProtocols []protocol.Protocol `json:"client_protocols"` +} + +// Params is one validated, canonical channel parameter object. +type Params struct { + values map[string]string +} + +// Value returns one normalized parameter without exposing the backing map. +func (p Params) Value(key string) (string, bool) { + value, ok := p.values[key] + return value, ok +} + +// CanonicalJSON returns an independent canonical JSON object. +func (p Params) CanonicalJSON() json.RawMessage { return canonicalJSON(p.values) } + +// Credential is one validated credential object. Its fields are never marshaled implicitly. +type Credential struct { + values map[string]string +} + +// Value returns one normalized credential field without exposing the backing map. +func (c Credential) Value(key string) (string, bool) { + value, ok := c.values[key] + return value, ok +} + +// CanonicalJSON explicitly returns the credential object for encryption or execution. +func (c Credential) CanonicalJSON() json.RawMessage { return canonicalJSON(c.values) } + +// MarshalJSON prevents accidental serialization of credential names or values. +func (c Credential) MarshalJSON() ([]byte, error) { return []byte(`{}`), nil } + +// RouteMode identifies whether an operation uses the channel's native wire +// contract or requires protocol conversion. +type RouteMode string + +const ( + RouteNative RouteMode = "native" + RouteConverted RouteMode = "converted" +) + +// Valid reports whether mode is recognized. +func (m RouteMode) Valid() bool { return m == RouteNative || m == RouteConverted } + +// ProviderKind identifies the upstream API family selected by a code-owned +// channel preset. It is runtime metadata, not a persisted SDK identifier. +type ProviderKind string + +const ( + ProviderOpenAI ProviderKind = "openai" + ProviderAnthropic ProviderKind = "anthropic" + ProviderGemini ProviderKind = "gemini" + ProviderOpenAICompatible ProviderKind = "openai_compatible" + ProviderAnthropicCompatible ProviderKind = "anthropic_compatible" + ProviderGeminiCompatible ProviderKind = "gemini_compatible" + ProviderAzureOpenAI ProviderKind = "azure_openai" + ProviderAWSBedrock ProviderKind = "aws_bedrock" + ProviderGoogleVertex ProviderKind = "google_vertex" +) + +// Valid reports whether the provider kind is recognized. +func (k ProviderKind) Valid() bool { + switch k { + case ProviderOpenAI, + ProviderAnthropic, + ProviderGemini, + ProviderOpenAICompatible, + ProviderAnthropicCompatible, + ProviderGeminiCompatible, + ProviderAzureOpenAI, + ProviderAWSBedrock, + ProviderGoogleVertex: + return true + default: + return false + } +} + +// ResolvedTarget is the provider-neutral execution target derived from a channel preset. +type ResolvedTarget struct { + ChannelID ID `json:"channel_id"` + ProviderKind ProviderKind `json:"-"` + TargetConfig json.RawMessage `json:"-"` + CatalogProviderID string `json:"-"` + + capabilities map[protocol.Protocol]execution.CapabilitySet + modes map[protocol.Protocol]map[execution.Operation]RouteMode +} + +// Supports reports whether one client protocol can execute the operation with +// every required feature through this target. +func (t ResolvedTarget) Supports( + clientProtocol protocol.Protocol, + operation execution.Operation, + required execution.FeatureSet, +) bool { + capabilities, ok := t.capabilities[clientProtocol] + return ok && capabilities.Supports(operation, required) +} + +// Capabilities returns an independent capability set for one client protocol. +func (t ResolvedTarget) Capabilities(clientProtocol protocol.Protocol) execution.CapabilitySet { + return t.capabilities[clientProtocol].Clone() +} + +// Mode returns how an operation reaches this channel. +func (t ResolvedTarget) Mode( + clientProtocol protocol.Protocol, + operation execution.Operation, +) (RouteMode, bool) { + byOperation, ok := t.modes[clientProtocol] + if !ok { + return "", false + } + mode, ok := byOperation[operation] + return mode, ok +} + +// Registry is an immutable code-owned channel directory. +type Registry struct { + byID map[ID]definition + order []ID +} + +// NewRegistry constructs the built-in, read-only channel registry. +func NewRegistry() *Registry { + registry, err := newRegistry(builtinDefinitions()) + if err != nil { + panic(fmt.Sprintf("build channel registry: %v", err)) + } + return registry +} + +// Get returns an independent public descriptor. +func (r *Registry) Get(id ID) (Descriptor, bool) { + definition, ok := r.lookup(id) + if !ok { + return Descriptor{}, false + } + return cloneDescriptor(definition.descriptor), true +} + +// List returns every built-in preset in stable product order. +func (r *Registry) List() []Descriptor { + if r == nil { + return nil + } + result := make([]Descriptor, 0, len(r.order)) + for _, id := range r.order { + result = append(result, cloneDescriptor(r.byID[id].descriptor)) + } + return result +} + +// Search returns case-insensitive ID/name/keyword matches in stable product order. +func (r *Registry) Search(query string) []Descriptor { + if r == nil { + return nil + } + normalized := strings.ToLower(strings.TrimSpace(query)) + result := make([]Descriptor, 0, len(r.order)) + for _, id := range r.order { + definition := r.byID[id] + if normalized != "" && !definition.matches(normalized) { + continue + } + result = append(result, cloneDescriptor(definition.descriptor)) + } + return result +} + +// CatalogProviderID returns the exact Models.dev provider mapping for a known +// channel. A known compatible channel intentionally returns an empty mapping. +func (r *Registry) CatalogProviderID(id ID) (string, bool) { + definition, ok := r.lookup(id) + if !ok { + return "", false + } + return definition.catalogProviderID, true +} + +// ValidateParams validates and normalizes the channel parameter object. +func (r *Registry) ValidateParams(id ID, raw json.RawMessage) (Params, error) { + definition, ok := r.lookup(id) + if !ok { + return Params{}, &ValidationError{Field: "channel_id", Reason: "unknown channel"} + } + values, err := definition.params.validate("params", raw) + if err != nil { + return Params{}, err + } + return Params{values: values}, nil +} + +// ValidateCredential validates and normalizes one credential object. +func (r *Registry) ValidateCredential(id ID, raw json.RawMessage) (Credential, error) { + definition, ok := r.lookup(id) + if !ok { + return Credential{}, &ValidationError{Field: "channel_id", Reason: "unknown channel"} + } + values, err := definition.credentials.validate("credential", raw) + if err != nil { + return Credential{}, err + } + if definition.validateCredential != nil { + if err := definition.validateCredential(values); err != nil { + return Credential{}, err + } + } + return Credential{values: values}, nil +} + +// Resolve validates params and produces a provider-neutral execution target. +func (r *Registry) Resolve(id ID, raw json.RawMessage) (ResolvedTarget, error) { + definition, ok := r.lookup(id) + if !ok { + return ResolvedTarget{}, &ValidationError{Field: "channel_id", Reason: "unknown channel"} + } + params, err := r.ValidateParams(id, raw) + if err != nil { + return ResolvedTarget{}, err + } + targetConfig := json.RawMessage(`{}`) + if len(definition.fixedTargetConfig) > 0 { + targetConfig = append(json.RawMessage(nil), definition.fixedTargetConfig...) + } else if len(definition.params) > 0 { + targetConfig = params.CanonicalJSON() + } + return ResolvedTarget{ + ChannelID: id, + ProviderKind: definition.providerKind, + TargetConfig: append(json.RawMessage(nil), targetConfig...), + CatalogProviderID: definition.catalogProviderID, + capabilities: cloneCapabilityMatrix(definition.capabilities), + modes: cloneRouteModes(definition.modes), + }, nil +} + +// ResolveExecutionTarget validates a frozen runtime target produced by Resolve. +// Fixed presets must match their code-owned target exactly; configurable +// channels re-validate the target through their parameter schema. +func (r *Registry) ResolveExecutionTarget(id ID, raw json.RawMessage) (ResolvedTarget, error) { + definition, ok := r.lookup(id) + if !ok { + return ResolvedTarget{}, &ValidationError{Field: "channel_id", Reason: "unknown channel"} + } + if len(definition.fixedTargetConfig) > 0 { + if !bytes.Equal(bytes.TrimSpace(raw), definition.fixedTargetConfig) { + return ResolvedTarget{}, &ValidationError{Field: "target_config", Reason: "does not match channel preset"} + } + return r.Resolve(id, nil) + } + return r.Resolve(id, raw) +} + +// ValidationError describes one invalid field without embedding its value. +type ValidationError struct { + Field string + Reason string +} + +func (e *ValidationError) Error() string { return fmt.Sprintf("%s: %s", e.Field, e.Reason) } + +type normalizer func(string) (string, error) + +type fieldSpec struct { + descriptor FieldDescriptor + normalize normalizer +} + +type objectSchema []fieldSpec + +func (s objectSchema) descriptors() []FieldDescriptor { + descriptors := make([]FieldDescriptor, len(s)) + for index, field := range s { + descriptors[index] = field.descriptor + } + return descriptors +} + +func (s objectSchema) validate(prefix string, raw json.RawMessage) (map[string]string, error) { + object, err := decodeStrictObject(raw) + if err != nil { + return nil, &ValidationError{Field: prefix, Reason: err.Error()} + } + byKey := make(map[string]fieldSpec, len(s)) + for _, field := range s { + byKey[field.descriptor.Key] = field + } + for key := range object { + if _, ok := byKey[key]; !ok { + return nil, &ValidationError{Field: prefix + "." + key, Reason: "unknown field"} + } + } + values := make(map[string]string, len(object)) + for _, field := range s { + rawValue, exists := object[field.descriptor.Key] + if !exists { + if field.descriptor.Required { + return nil, &ValidationError{Field: prefix + "." + field.descriptor.Key, Reason: "required"} + } + continue + } + var value string + if err := json.Unmarshal(rawValue, &value); err != nil { + return nil, &ValidationError{Field: prefix + "." + field.descriptor.Key, Reason: "must be a string"} + } + normalized, err := field.normalize(value) + if err != nil { + return nil, &ValidationError{Field: prefix + "." + field.descriptor.Key, Reason: err.Error()} + } + values[field.descriptor.Key] = normalized + } + return values, nil +} + +type definition struct { + descriptor Descriptor + searchTerms []string + params objectSchema + credentials objectSchema + validateCredential func(map[string]string) error + catalogProviderID string + providerKind ProviderKind + compatible bool + fixedTargetConfig json.RawMessage + capabilities map[protocol.Protocol]execution.CapabilitySet + modes map[protocol.Protocol]map[execution.Operation]RouteMode +} + +func (d definition) matches(query string) bool { + values := []string{string(d.descriptor.ID), d.descriptor.Name, d.descriptor.Description} + values = append(values, d.searchTerms...) + for _, value := range values { + if strings.Contains(strings.ToLower(value), query) { + return true + } + } + return false +} + +func newRegistry(definitions []definition) (*Registry, error) { + registry := &Registry{byID: make(map[ID]definition, len(definitions)), order: make([]ID, 0, len(definitions))} + for _, definition := range definitions { + if err := validateDefinition(definition); err != nil { + return nil, err + } + id := definition.descriptor.ID + if _, duplicate := registry.byID[id]; duplicate { + return nil, fmt.Errorf("duplicate channel ID %q", id) + } + definition.descriptor = cloneDescriptor(definition.descriptor) + definition.searchTerms = append([]string(nil), definition.searchTerms...) + definition.capabilities = cloneCapabilityMatrix(definition.capabilities) + definition.modes = cloneRouteModes(definition.modes) + definition.fixedTargetConfig = append(json.RawMessage(nil), definition.fixedTargetConfig...) + registry.byID[id] = definition + registry.order = append(registry.order, id) + } + return registry, nil +} + +func validateDefinition(definition definition) error { + id := string(definition.descriptor.ID) + if id == "" || strings.Trim(id, "abcdefghijklmnopqrstuvwxyz0123456789_") != "" || strings.HasPrefix(id, "_") || strings.HasSuffix(id, "_") { + return fmt.Errorf("invalid channel ID %q", id) + } + if strings.TrimSpace(definition.descriptor.Name) == "" { + return fmt.Errorf("channel %q has no name", id) + } + if !definition.providerKind.Valid() { + return fmt.Errorf("channel %q has invalid provider kind", id) + } + if len(definition.fixedTargetConfig) > 0 && !json.Valid(definition.fixedTargetConfig) { + return fmt.Errorf("channel %q has invalid fixed target config", id) + } + for name, schema := range map[string]objectSchema{"params": definition.params, "credential": definition.credentials} { + seen := make(map[string]struct{}, len(schema)) + for _, field := range schema { + key := field.descriptor.Key + if key == "" || field.normalize == nil || (field.descriptor.InputKind != InputText && field.descriptor.InputKind != InputURL && field.descriptor.InputKind != InputSecret) { + return fmt.Errorf("channel %q has invalid %s field", id, name) + } + if _, duplicate := seen[key]; duplicate { + return fmt.Errorf("channel %q has duplicate %s field %q", id, name, key) + } + seen[key] = struct{}{} + } + } + if len(definition.capabilities) == 0 { + return fmt.Errorf("channel %q has no client protocols", id) + } + for clientProtocol, capabilities := range definition.capabilities { + if !clientProtocol.Valid() { + return fmt.Errorf("channel %q has invalid client protocol %q", id, clientProtocol) + } + if err := capabilities.Validate(); err != nil { + return fmt.Errorf("channel %q capabilities: %w", id, err) + } + modes := definition.modes[clientProtocol] + for _, operation := range capabilities.Operations() { + mode, exists := modes[operation] + if !exists || !mode.Valid() { + return fmt.Errorf("channel %q has no valid route mode for %q/%q", id, clientProtocol, operation) + } + } + for operation := range modes { + if !capabilities.Has(operation) { + return fmt.Errorf("channel %q has route mode for unsupported %q/%q", id, clientProtocol, operation) + } + } + } + return nil +} + +func (r *Registry) lookup(id ID) (definition, bool) { + if r == nil { + return definition{}, false + } + definition, ok := r.byID[id] + return definition, ok +} + +func cloneDescriptor(source Descriptor) Descriptor { + source.ParamFields = append([]FieldDescriptor{}, source.ParamFields...) + source.CredentialFields = append([]FieldDescriptor{}, source.CredentialFields...) + source.ClientProtocols = append([]protocol.Protocol{}, source.ClientProtocols...) + return source +} + +func cloneCapabilityMatrix(source map[protocol.Protocol]execution.CapabilitySet) map[protocol.Protocol]execution.CapabilitySet { + clone := make(map[protocol.Protocol]execution.CapabilitySet, len(source)) + for clientProtocol, capabilities := range source { + clone[clientProtocol] = capabilities.Clone() + } + return clone +} + +func cloneRouteModes(source map[protocol.Protocol]map[execution.Operation]RouteMode) map[protocol.Protocol]map[execution.Operation]RouteMode { + clone := make(map[protocol.Protocol]map[execution.Operation]RouteMode, len(source)) + for clientProtocol, modes := range source { + clonedModes := make(map[execution.Operation]RouteMode, len(modes)) + for operation, mode := range modes { + clonedModes[operation] = mode + } + clone[clientProtocol] = clonedModes + } + return clone +} + +func canonicalJSON(values map[string]string) json.RawMessage { + if len(values) == 0 { + return json.RawMessage(`{}`) + } + encoded, err := json.Marshal(values) + if err != nil { + panic(fmt.Sprintf("encode validated channel values: %v", err)) + } + return append(json.RawMessage(nil), encoded...) +} + +func decodeStrictObject(raw json.RawMessage) (map[string]json.RawMessage, error) { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 { + trimmed = []byte(`{}`) + } + if trimmed[0] != '{' { + return nil, errors.New("must be a JSON object") + } + if err := rejectDuplicateFields(trimmed); err != nil { + return nil, err + } + decoder := json.NewDecoder(bytes.NewReader(trimmed)) + decoder.UseNumber() + object := make(map[string]json.RawMessage) + if err := decoder.Decode(&object); err != nil { + return nil, errors.New("invalid JSON object") + } + var extra json.RawMessage + if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) { + return nil, errors.New("must contain one JSON object") + } + return object, nil +} + +func rejectDuplicateFields(raw []byte) error { + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var walk func() error + walk = func() error { + token, err := decoder.Token() + if err != nil { + return errors.New("invalid JSON object") + } + delimiter, ok := token.(json.Delim) + if !ok { + return nil + } + switch delimiter { + case '{': + seen := make(map[string]struct{}) + for decoder.More() { + keyToken, err := decoder.Token() + if err != nil { + return errors.New("invalid JSON object") + } + key, ok := keyToken.(string) + if !ok { + return errors.New("invalid JSON object") + } + if _, duplicate := seen[key]; duplicate { + return fmt.Errorf("duplicate field %q", key) + } + seen[key] = struct{}{} + if err := walk(); err != nil { + return err + } + } + _, err = decoder.Token() + return err + case '[': + for decoder.More() { + if err := walk(); err != nil { + return err + } + } + _, err = decoder.Token() + return err + default: + return errors.New("invalid JSON object") + } + } + if err := walk(); err != nil { + return err + } + var extra json.RawMessage + if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) { + return errors.New("must contain one JSON object") + } + return nil +} + +func normalizeNonEmpty(value string) (string, error) { + value = strings.TrimSpace(value) + if value == "" { + return "", errors.New("must not be empty") + } + return value, nil +} + +func normalizeBaseURL(value string) (string, error) { + value = strings.TrimSpace(value) + parsed, err := url.Parse(value) + if err != nil || parsed == nil || parsed.Opaque != "" || parsed.User != nil || parsed.Host == "" || parsed.Fragment != "" { + return "", errors.New("must be an absolute HTTP(S) URL without credentials or fragment") + } + parsed.Scheme = strings.ToLower(parsed.Scheme) + if parsed.Scheme != "http" && parsed.Scheme != "https" { + return "", errors.New("must be an absolute HTTP(S) URL without credentials or fragment") + } + if containsCredentialQuery(parsed.RawQuery) { + return "", errors.New("must not contain credential query parameters") + } + parsed.Host = strings.ToLower(parsed.Host) + parsed.Path = strings.TrimRight(parsed.Path, "/") + parsed.RawPath = strings.TrimRight(parsed.RawPath, "/") + parsed.ForceQuery = false + return parsed.String(), nil +} + +func containsCredentialQuery(rawQuery string) bool { + for _, segment := range strings.Split(rawQuery, "&") { + key := segment + if separator := strings.IndexByte(key, '='); separator >= 0 { + key = key[:separator] + } + decoded, err := url.QueryUnescape(key) + if err != nil { + return true + } + normalized := strings.ReplaceAll(strings.ToLower(decoded), "-", "_") + switch normalized { + case "authorization", "proxy_authorization", "api_key", "apikey", "x_api_key", + "x_goog_api_key", "key", "access_token", "provider", "fallback", "fallbacks": + return true + } + } + return false +} diff --git a/internal/channel/presets.go b/internal/channel/presets.go new file mode 100644 index 00000000..b17f1562 --- /dev/null +++ b/internal/channel/presets.go @@ -0,0 +1,433 @@ +package channel + +import ( + "encoding/json" + "fmt" + "strings" + + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func builtinDefinitions() []definition { + apiKeySchema := objectSchema{{ + descriptor: FieldDescriptor{ + Key: "api_key", Label: "API Key", InputKind: InputSecret, + Required: true, Sensitive: true, + }, + normalize: normalizeNonEmpty, + }} + compatibleParams := objectSchema{{ + descriptor: FieldDescriptor{ + Key: "base_url", Label: "Base URL", InputKind: InputURL, + Required: true, + }, + normalize: normalizeBaseURL, + }} + azureParams := objectSchema{{ + descriptor: FieldDescriptor{Key: "endpoint", Label: "Azure endpoint", InputKind: InputURL, Required: true}, + normalize: normalizeBaseURL, + }} + azureCredentials := objectSchema{ + secretField("api_key", "API Key"), + secretField("client_id", "Entra client ID"), + secretField("client_secret", "Entra client secret"), + secretField("tenant_id", "Entra tenant ID"), + } + bedrockParams := objectSchema{{ + descriptor: FieldDescriptor{Key: "region", Label: "AWS region", InputKind: InputText, Required: true}, + normalize: normalizeCloudIdentifier, + }} + bedrockCredentials := objectSchema{ + secretField("api_key", "Bedrock API Key"), + secretField("access_key", "AWS access key"), + secretField("secret_key", "AWS secret key"), + secretField("session_token", "AWS session token"), + secretField("role_arn", "AWS role ARN"), + secretField("external_id", "AWS external ID"), + secretField("session_name", "AWS role session name"), + } + vertexParams := objectSchema{ + { + descriptor: FieldDescriptor{Key: "project_id", Label: "Google Cloud project ID", InputKind: InputText, Required: true}, + normalize: normalizeCloudIdentifier, + }, + { + descriptor: FieldDescriptor{Key: "location", Label: "Vertex location", InputKind: InputText, Required: true}, + normalize: normalizeCloudIdentifier, + }, + { + descriptor: FieldDescriptor{Key: "project_number", Label: "Google Cloud project number", InputKind: InputText}, + normalize: normalizeCloudIdentifier, + }, + } + vertexCredentials := objectSchema{{ + descriptor: FieldDescriptor{ + Key: "service_account_json", Label: "Service account JSON", InputKind: InputSecret, + Required: true, Sensitive: true, + }, + normalize: normalizeServiceAccountJSON, + }} + + azureDefinition := newDefinition( + AzureOpenAI, "Azure OpenAI", "AZ", "Azure OpenAI and Microsoft Foundry", []string{"azure", "foundry", "entra"}, + azureParams, azureCredentials, "azure", ProviderAzureOpenAI, false, nil, false, + ) + azureDefinition.validateCredential = validateAzureCredential + bedrockDefinition := newDefinition( + AWSBedrock, "AWS Bedrock", "BR", "Amazon Bedrock", []string{"aws", "amazon", "sigv4", "iam"}, + bedrockParams, bedrockCredentials, "amazon-bedrock", ProviderAWSBedrock, false, nil, false, + ) + bedrockDefinition.validateCredential = validateBedrockCredential + vertexDefinition := newDefinition( + GoogleVertex, "Google Vertex AI", "VX", "Google Cloud Vertex AI", []string{"google", "gcp", "vertex"}, + vertexParams, vertexCredentials, "google-vertex", ProviderGoogleVertex, false, nil, false, + ) + + return []definition{ + newDefinition( + OpenAI, "OpenAI", "OA", "OpenAI official API", []string{"gpt"}, + nil, apiKeySchema, "openai", ProviderOpenAI, false, + map[protocol.Protocol]bool{ + protocol.OpenAICompletions: true, + protocol.OpenAIResponses: true, + }, + true, + ), + newDefinition( + Anthropic, "Anthropic", "AN", "Anthropic official API", []string{"claude"}, + nil, apiKeySchema, "anthropic", ProviderAnthropic, false, + map[protocol.Protocol]bool{protocol.Anthropic: true}, false, + ), + newDefinition( + Gemini, "Google Gemini", "GE", "Google Gemini official API", []string{"google"}, + nil, apiKeySchema, "google", ProviderGemini, false, + map[protocol.Protocol]bool{protocol.Gemini: true}, false, + ), + azureDefinition, + bedrockDefinition, + vertexDefinition, + newFixedCompatibleDefinition( + DeepSeek, "DeepSeek", "DS", []string{"deep seek"}, + "https://api.deepseek.com/v1", "deepseek", apiKeySchema, + ), + newFixedCompatibleDefinition( + MoonshotAI, "Moonshot AI", "MS", []string{"kimi", "moonshot"}, + "https://api.moonshot.cn/v1", "moonshotai", apiKeySchema, + ), + newFixedCompatibleDefinition( + SiliconFlow, "SiliconFlow", "SF", []string{"silicon flow"}, + "https://api.siliconflow.cn/v1", "siliconflow", apiKeySchema, + ), + newFixedCompatibleDefinition( + ZhipuAI, "Zhipu AI", "ZP", []string{"glm", "bigmodel"}, + "https://open.bigmodel.cn/api/paas/v4", "zhipuai", apiKeySchema, + ), + newFixedCompatibleDefinition( + Alibaba, "Alibaba Cloud Bailian", "BL", []string{"dashscope", "qwen", "bailian"}, + "https://dashscope.aliyuncs.com/compatible-mode/v1", "alibaba", apiKeySchema, + ), + newFixedCompatibleDefinition( + Volcengine, "Volcengine Ark", "VE", []string{"doubao", "ark"}, + "https://ark.cn-beijing.volces.com/api/v3", "volcengine", apiKeySchema, + ), + newFixedCompatibleDefinition( + OpenRouter, "OpenRouter", "OR", []string{"router"}, + "https://openrouter.ai/api/v1", "openrouter", apiKeySchema, + ), + newFixedCompatibleDefinition( + Groq, "Groq", "GQ", []string{"groqcloud"}, + "https://api.groq.com/openai/v1", "groq", apiKeySchema, + ), + newFixedCompatibleDefinition( + XAI, "xAI", "XA", []string{"grok"}, + "https://api.x.ai/v1", "xai", apiKeySchema, + ), + newDefinition( + OpenAICompatible, "OpenAI Compatible", "OC", "Custom OpenAI-compatible API", []string{"custom", "proxy", "gateway"}, + compatibleParams, apiKeySchema, "", ProviderOpenAICompatible, true, + map[protocol.Protocol]bool{ + protocol.OpenAICompletions: true, + protocol.OpenAIResponses: true, + }, + true, + ), + newDefinition( + AnthropicCompatible, "Anthropic Compatible", "AC", "Custom Anthropic-compatible API", []string{"custom", "proxy", "gateway", "claude"}, + compatibleParams, apiKeySchema, "", ProviderAnthropicCompatible, true, + map[protocol.Protocol]bool{protocol.Anthropic: true}, false, + ), + newDefinition( + GeminiCompatible, "Gemini Compatible", "GC", "Custom Gemini-compatible API", []string{"custom", "proxy", "gateway", "google"}, + compatibleParams, apiKeySchema, "", ProviderGeminiCompatible, true, + map[protocol.Protocol]bool{protocol.Gemini: true}, false, + ), + } +} + +func secretField(key, label string) fieldSpec { + return fieldSpec{ + descriptor: FieldDescriptor{Key: key, Label: label, InputKind: InputSecret, Sensitive: true}, + normalize: normalizeNonEmpty, + } +} + +func validateAzureCredential(values map[string]string) error { + hasAPIKey := values["api_key"] != "" + entraFields := []string{"client_id", "client_secret", "tenant_id"} + entraCount := 0 + for _, field := range entraFields { + if values[field] != "" { + entraCount++ + } + } + if hasAPIKey && entraCount > 0 { + return &ValidationError{Field: "credential", Reason: "must use either API key or Entra credentials"} + } + if !hasAPIKey && entraCount != len(entraFields) { + return &ValidationError{Field: "credential", Reason: "requires an API key or complete Entra credentials"} + } + return nil +} + +func validateBedrockCredential(values map[string]string) error { + hasAPIKey := values["api_key"] != "" + hasAccessKey := values["access_key"] != "" + hasSecretKey := values["secret_key"] != "" + hasRole := values["role_arn"] != "" + hasSigV4Fields := hasAccessKey || hasSecretKey || hasRole || values["session_token"] != "" || + values["external_id"] != "" || values["session_name"] != "" + if hasAPIKey && hasSigV4Fields { + return &ValidationError{Field: "credential", Reason: "must use either API key or SigV4 credentials"} + } + if hasAPIKey { + return nil + } + if hasAccessKey != hasSecretKey { + return &ValidationError{Field: "credential", Reason: "access_key and secret_key must be provided together"} + } + if !hasAccessKey && !hasRole { + return &ValidationError{Field: "credential", Reason: "requires an API key, access key pair, or role ARN"} + } + if values["session_token"] != "" && !hasAccessKey { + return &ValidationError{Field: "credential.session_token", Reason: "requires access_key and secret_key"} + } + if (values["external_id"] != "" || values["session_name"] != "") && !hasRole { + return &ValidationError{Field: "credential.role_arn", Reason: "required for role options"} + } + return nil +} + +func normalizeCloudIdentifier(value string) (string, error) { + value = strings.TrimSpace(value) + if value == "" || len(value) > 255 { + return "", fmt.Errorf("must be between 1 and 255 bytes") + } + for _, character := range value { + if character <= ' ' || character == 0x7f { + return "", fmt.Errorf("must not contain whitespace or control characters") + } + } + return value, nil +} + +func normalizeServiceAccountJSON(value string) (string, error) { + object, err := decodeStrictObject(json.RawMessage(strings.TrimSpace(value))) + if err != nil { + return "", fmt.Errorf("must be a valid JSON object") + } + for _, key := range []string{"type", "project_id", "client_email", "private_key"} { + raw, exists := object[key] + if !exists { + return "", fmt.Errorf("must contain %s", key) + } + var field string + if json.Unmarshal(raw, &field) != nil || strings.TrimSpace(field) == "" { + return "", fmt.Errorf("must contain non-empty %s", key) + } + } + encoded, err := json.Marshal(object) + if err != nil { + return "", fmt.Errorf("must be a valid JSON object") + } + return string(encoded), nil +} + +func newFixedCompatibleDefinition( + id ID, + name string, + mark string, + searchTerms []string, + baseURL string, + catalogProviderID string, + credentials objectSchema, +) definition { + definition := newDefinition( + id, + name, + mark, + "Managed API preset", + searchTerms, + nil, + credentials, + catalogProviderID, + ProviderOpenAICompatible, + false, + map[protocol.Protocol]bool{protocol.OpenAICompletions: true}, + false, + ) + definition.fixedTargetConfig = canonicalJSON(map[string]string{"base_url": baseURL}) + return definition +} + +func newDefinition( + id ID, + name string, + mark string, + description string, + searchTerms []string, + params objectSchema, + credentials objectSchema, + catalogProviderID string, + providerKind ProviderKind, + compatible bool, + nativeProtocols map[protocol.Protocol]bool, + responsesLifecycle bool, +) definition { + capabilities, modes := buildRouteMatrix(nativeProtocols, responsesLifecycle) + return definition{ + descriptor: Descriptor{ + ID: id, + Name: name, + Mark: mark, + Description: description, + ParamFields: params.descriptors(), + CredentialFields: credentials.descriptors(), + ClientProtocols: orderedProtocols(capabilities), + }, + searchTerms: append([]string(nil), searchTerms...), + params: append(objectSchema(nil), params...), + credentials: append(objectSchema(nil), credentials...), + catalogProviderID: catalogProviderID, + providerKind: providerKind, + compatible: compatible, + capabilities: capabilities, + modes: modes, + } +} + +func buildRouteMatrix( + nativeProtocols map[protocol.Protocol]bool, + responsesLifecycle bool, +) ( + map[protocol.Protocol]execution.CapabilitySet, + map[protocol.Protocol]map[execution.Operation]RouteMode, +) { + capabilities := make(map[protocol.Protocol]execution.CapabilitySet) + modes := make(map[protocol.Protocol]map[execution.Operation]RouteMode) + for _, clientProtocol := range protocol.DataPlaneProtocols() { + native := nativeProtocols[clientProtocol] + mode := RouteConverted + features := mustFeatureSet(execution.FeatureStreaming) + if native { + mode = RouteNative + features = mustFeatureSet( + execution.FeatureStreaming, + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + ) + } + operation := execution.OperationChatCompletion + if clientProtocol == protocol.OpenAIResponses { + operation = execution.OperationResponsesCreate + if native && responsesLifecycle { + features = mustFeatureSet( + execution.FeatureStreaming, + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + execution.FeatureNativeResourceSemantics, + ) + } + } + + items := []execution.Capability{ + {Operation: operation, Features: features}, + {Operation: execution.OperationListModels}, + {Operation: execution.OperationProbe}, + } + operationModes := map[execution.Operation]RouteMode{ + operation: mode, + execution.OperationListModels: RouteNative, + execution.OperationProbe: RouteNative, + } + if clientProtocol == protocol.OpenAIResponses && native && responsesLifecycle { + nativeFeatures := mustFeatureSet( + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + ) + for _, nativeOperation := range []execution.Operation{ + execution.OperationResponsesCompact, + execution.OperationResponsesInputTokens, + } { + items = append(items, execution.Capability{ + Operation: nativeOperation, + Features: nativeFeatures, + }) + operationModes[nativeOperation] = RouteNative + } + resourceFeatures := mustFeatureSet(execution.FeatureNativeResourceSemantics) + for _, resourceOperation := range []execution.Operation{ + execution.OperationResponsesRetrieve, + execution.OperationResponsesDelete, + execution.OperationResponsesCancel, + execution.OperationResponsesInputItems, + } { + items = append(items, execution.Capability{Operation: resourceOperation, Features: resourceFeatures}) + operationModes[resourceOperation] = RouteNative + } + items = append(items, execution.Capability{ + Operation: execution.OperationResponsesPassthrough, + Features: mustFeatureSet( + execution.FeatureStreaming, + execution.FeatureNativeResourceSemantics, + ), + }) + operationModes[execution.OperationResponsesPassthrough] = RouteNative + } + capabilities[clientProtocol] = mustCapabilitySet(items...) + modes[clientProtocol] = operationModes + } + return capabilities, modes +} + +func mustFeatureSet(features ...execution.Feature) execution.FeatureSet { + set, err := execution.NewFeatureSet(features...) + if err != nil { + panic(fmt.Sprintf("build channel feature set: %v", err)) + } + return set +} + +func mustCapabilitySet(capabilities ...execution.Capability) execution.CapabilitySet { + set, err := execution.NewCapabilitySet(capabilities...) + if err != nil { + panic(fmt.Sprintf("build channel capability set: %v", err)) + } + return set +} + +func orderedProtocols(capabilities map[protocol.Protocol]execution.CapabilitySet) []protocol.Protocol { + ordered := make([]protocol.Protocol, 0, len(capabilities)) + for _, clientProtocol := range protocol.DataPlaneProtocols() { + if _, ok := capabilities[clientProtocol]; ok { + ordered = append(ordered, clientProtocol) + } + } + return ordered +} diff --git a/internal/channel/registry_test.go b/internal/channel/registry_test.go new file mode 100644 index 00000000..68802000 --- /dev/null +++ b/internal/channel/registry_test.go @@ -0,0 +1,506 @@ +package channel + +import ( + "encoding/json" + "reflect" + "strings" + "testing" + + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestRegistryHasStableBuiltInOrderAndSearch(t *testing.T) { + registry := NewRegistry() + wantIDs := []ID{ + OpenAI, + Anthropic, + Gemini, + ID("azure_openai"), + ID("aws_bedrock"), + ID("google_vertex"), + DeepSeek, + MoonshotAI, + SiliconFlow, + ZhipuAI, + Alibaba, + Volcengine, + OpenRouter, + Groq, + XAI, + OpenAICompatible, + AnthropicCompatible, + GeminiCompatible, + } + if got := descriptorIDs(registry.List()); !reflect.DeepEqual(got, wantIDs) { + t.Fatalf("List() IDs = %v, want %v", got, wantIDs) + } + if got := descriptorIDs(registry.Search("GeMiNi")); !reflect.DeepEqual(got, []ID{Gemini, GeminiCompatible}) { + t.Fatalf("Search(gemini) IDs = %v", got) + } + if got := descriptorIDs(registry.Search("compatible")); !reflect.DeepEqual(got, []ID{OpenAICompatible, AnthropicCompatible, GeminiCompatible}) { + t.Fatalf("Search(compatible) IDs = %v", got) + } + if got := registry.Search("does-not-exist"); len(got) != 0 { + t.Fatalf("Search(unknown) = %#v, want empty", got) + } + if got := descriptorIDs(registry.Search("deep")); !reflect.DeepEqual(got, []ID{DeepSeek}) { + t.Fatalf("Search(deep) IDs = %v, want [%s]", got, DeepSeek) + } + if got := descriptorIDs(registry.Search("azure")); !reflect.DeepEqual(got, []ID{ID("azure_openai")}) { + t.Fatalf("Search(azure) IDs = %v", got) + } + if got := descriptorIDs(registry.Search("aws")); !reflect.DeepEqual(got, []ID{ID("aws_bedrock")}) { + t.Fatalf("Search(aws) IDs = %v", got) + } + if got := descriptorIDs(registry.Search("vertex")); !reflect.DeepEqual(got, []ID{ID("google_vertex")}) { + t.Fatalf("Search(vertex) IDs = %v", got) + } + + first := registry.List() + first[0].ClientProtocols[0] = protocol.Protocol("mutated") + first[0].ParamFields = append(first[0].ParamFields, FieldDescriptor{Key: "mutated"}) + second := registry.List() + if second[0].ClientProtocols[0] == protocol.Protocol("mutated") || len(second[0].ParamFields) != 0 { + t.Fatal("mutating List() output changed registry state") + } +} + +func TestRegistryPublicDescriptorsContainSchemasButNoInternalOrSecretValues(t *testing.T) { + registry := NewRegistry() + official, ok := registry.Get(OpenAI) + if !ok { + t.Fatal("Get(openai) missing") + } + if len(official.ParamFields) != 0 || len(official.CredentialFields) != 1 { + t.Fatalf("openai schemas = params %#v credentials %#v", official.ParamFields, official.CredentialFields) + } + credentialField := official.CredentialFields[0] + if credentialField.Key != "api_key" || credentialField.InputKind != InputSecret || !credentialField.Required || !credentialField.Sensitive { + t.Fatalf("openai credential field = %#v", credentialField) + } + compatible, ok := registry.Get(OpenAICompatible) + if !ok || len(compatible.ParamFields) != 1 { + t.Fatalf("Get(openai_compatible) = %#v, %t", compatible, ok) + } + if field := compatible.ParamFields[0]; field.Key != "base_url" || field.InputKind != InputURL || !field.Required || field.Sensitive { + t.Fatalf("openai compatible param field = %#v", field) + } + + encoded, err := json.Marshal(compatible) + if err != nil { + t.Fatalf("json.Marshal(descriptor) error = %v", err) + } + for _, forbidden := range []string{"target_config", "catalog_provider", "secret-value", "default_value"} { + if strings.Contains(strings.ToLower(string(encoded)), forbidden) { + t.Fatalf("public descriptor leaked %q: %s", forbidden, encoded) + } + } +} + +func TestRegistryReturnsExactCatalogProviderMappingWithoutResolvingParams(t *testing.T) { + registry := NewRegistry() + for id, want := range map[ID]string{ + OpenAI: "openai", Anthropic: "anthropic", Gemini: "google", + ID("azure_openai"): "azure", ID("aws_bedrock"): "amazon-bedrock", ID("google_vertex"): "google-vertex", + OpenAICompatible: "", AnthropicCompatible: "", GeminiCompatible: "", + } { + got, ok := registry.CatalogProviderID(id) + if !ok || got != want { + t.Errorf("CatalogProviderID(%q) = %q, %t, want %q, true", id, got, ok, want) + } + } + if got, ok := registry.CatalogProviderID(ID("missing")); ok || got != "" { + t.Fatalf("CatalogProviderID(missing) = %q, %t, want empty, false", got, ok) + } +} + +func TestRegistryValidatesCloudChannelParamsAndStructuredCredentials(t *testing.T) { + registry := NewRegistry() + tests := []struct { + name string + channelID ID + params string + credential string + wantParams string + wantCredential string + invalidCredentials []string + }{ + { + name: "Azure API key", channelID: ID("azure_openai"), + params: `{"endpoint":" HTTPS://Example.OPENAI.Azure.COM/ "}`, + credential: `{"api_key":" azure-secret "}`, + wantParams: `{"endpoint":"https://example.openai.azure.com"}`, + wantCredential: `{"api_key":"azure-secret"}`, + invalidCredentials: []string{ + `{}`, + `{"client_id":"client","tenant_id":"tenant"}`, + `{"api_key":"key","client_id":"client","client_secret":"secret","tenant_id":"tenant"}`, + }, + }, + { + name: "Azure Entra", channelID: ID("azure_openai"), + params: `{"endpoint":"https://resource.services.ai.azure.com"}`, + credential: `{"client_id":" client ","client_secret":" secret ","tenant_id":" tenant "}`, + wantParams: `{"endpoint":"https://resource.services.ai.azure.com"}`, + wantCredential: `{"client_id":"client","client_secret":"secret","tenant_id":"tenant"}`, + }, + { + name: "Bedrock API key", channelID: ID("aws_bedrock"), + params: `{"region":" us-east-1 "}`, + credential: `{"api_key":" bedrock-secret "}`, + wantParams: `{"region":"us-east-1"}`, + wantCredential: `{"api_key":"bedrock-secret"}`, + invalidCredentials: []string{ + `{}`, + `{"access_key":"access"}`, + `{"secret_key":"secret"}`, + `{"api_key":"key","access_key":"access","secret_key":"secret"}`, + }, + }, + { + name: "Bedrock SigV4", channelID: ID("aws_bedrock"), + params: `{"region":"eu-west-1"}`, + credential: `{"access_key":" access ","secret_key":" secret ","session_token":" token ","role_arn":" arn:aws:iam::123456789012:role/test "}`, + wantParams: `{"region":"eu-west-1"}`, + wantCredential: `{"access_key":"access","role_arn":"arn:aws:iam::123456789012:role/test","secret_key":"secret","session_token":"token"}`, + }, + { + name: "Bedrock ambient role", channelID: ID("aws_bedrock"), + params: `{"region":"ap-southeast-1"}`, + credential: `{"role_arn":" arn:aws:iam::123456789012:role/runtime "}`, + wantParams: `{"region":"ap-southeast-1"}`, + wantCredential: `{"role_arn":"arn:aws:iam::123456789012:role/runtime"}`, + }, + { + name: "Vertex service account", channelID: ID("google_vertex"), + params: `{"project_id":" project-one ","location":" us-central1 "}`, + credential: `{"service_account_json":" {\"type\":\"service_account\",\"project_id\":\"project-one\",\"client_email\":\"svc@example.iam.gserviceaccount.com\",\"private_key\":\"secret\"} "}`, + wantParams: `{"location":"us-central1","project_id":"project-one"}`, + wantCredential: `{"service_account_json":"{\"client_email\":\"svc@example.iam.gserviceaccount.com\",\"private_key\":\"secret\",\"project_id\":\"project-one\",\"type\":\"service_account\"}"}`, + invalidCredentials: []string{ + `{}`, + `{"service_account_json":"not-json"}`, + `{"service_account_json":"{}"}`, + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + params, err := registry.ValidateParams(test.channelID, json.RawMessage(test.params)) + if err != nil { + t.Fatalf("ValidateParams() error = %v", err) + } + if got := string(params.CanonicalJSON()); got != test.wantParams { + t.Fatalf("params = %s, want %s", got, test.wantParams) + } + credential, err := registry.ValidateCredential(test.channelID, json.RawMessage(test.credential)) + if err != nil { + t.Fatalf("ValidateCredential() error = %v", err) + } + if got := string(credential.CanonicalJSON()); got != test.wantCredential { + t.Fatalf("credential = %s, want %s", got, test.wantCredential) + } + for _, raw := range test.invalidCredentials { + if _, err := registry.ValidateCredential(test.channelID, json.RawMessage(raw)); err == nil { + t.Fatalf("ValidateCredential(%s) error = nil", raw) + } + } + }) + } +} + +func TestCloudTargetsCarryOnlyNeutralPresetMetadata(t *testing.T) { + registry := NewRegistry() + tests := []struct { + channelID ID + params string + providerKind ProviderKind + catalogProvider string + }{ + {ID("azure_openai"), `{"endpoint":"https://resource.openai.azure.com"}`, ProviderKind("azure_openai"), "azure"}, + {ID("aws_bedrock"), `{"region":"us-east-1"}`, ProviderKind("aws_bedrock"), "amazon-bedrock"}, + {ID("google_vertex"), `{"location":"us-central1","project_id":"project-one"}`, ProviderKind("google_vertex"), "google-vertex"}, + } + for _, test := range tests { + t.Run(string(test.channelID), func(t *testing.T) { + target, err := registry.Resolve(test.channelID, json.RawMessage(test.params)) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if target.ChannelID != test.channelID || target.ProviderKind != test.providerKind || target.CatalogProviderID != test.catalogProvider || string(target.TargetConfig) != test.params { + t.Fatalf("target = %#v, config=%s", target, target.TargetConfig) + } + if mode, ok := target.Mode(protocol.OpenAICompletions, execution.OperationChatCompletion); !ok || mode != RouteConverted { + t.Fatalf("OpenAI chat mode = %q, %t", mode, ok) + } + if target.Supports(protocol.OpenAIResponses, execution.OperationResponsesRetrieve, execution.FeatureSet{}) { + t.Fatal("cloud converted target unexpectedly exposes native Responses lifecycle") + } + }) + } +} + +func TestRegistryStrictlyValidatesAndCanonicalizesParams(t *testing.T) { + registry := NewRegistry() + params, err := registry.ValidateParams(OpenAI, nil) + if err != nil { + t.Fatalf("ValidateParams(openai, nil) error = %v", err) + } + if got := string(params.CanonicalJSON()); got != `{}` { + t.Fatalf("openai params = %s, want {}", got) + } + + valid, err := registry.ValidateParams(OpenAICompatible, json.RawMessage(`{"base_url":" HTTPS://Example.COM/v1/ "}`)) + if err != nil { + t.Fatalf("ValidateParams(compatible) error = %v", err) + } + if got, ok := valid.Value("base_url"); !ok || got != "https://example.com/v1" { + t.Fatalf("base_url = %q, %t", got, ok) + } + if got := string(valid.CanonicalJSON()); got != `{"base_url":"https://example.com/v1"}` { + t.Fatalf("canonical params = %s", got) + } + + withQuery, err := registry.ValidateParams(OpenAICompatible, json.RawMessage(`{"base_url":" HTTPS://Example.COM/v1/?api-version=2025-01-01&tenant=one "}`)) + if err != nil { + t.Fatalf("ValidateParams(compatible query) error = %v", err) + } + if got, ok := withQuery.Value("base_url"); !ok || got != "https://example.com/v1?api-version=2025-01-01&tenant=one" { + t.Fatalf("query base_url = %q, %t", got, ok) + } + + tests := []struct { + name string + id ID + raw string + }{ + {name: "missing required", id: OpenAICompatible, raw: `{}`}, + {name: "unknown", id: OpenAICompatible, raw: `{"base_url":"https://example.com","extra":"value"}`}, + {name: "wrong type", id: OpenAICompatible, raw: `{"base_url":42}`}, + {name: "duplicate", id: OpenAICompatible, raw: `{"base_url":"https://one.example","base_url":"https://two.example"}`}, + {name: "not object", id: OpenAICompatible, raw: `[]`}, + {name: "unsafe URL", id: OpenAICompatible, raw: `{"base_url":"https://user:secret@example.com"}`}, + {name: "query credential", id: OpenAICompatible, raw: `{"base_url":"https://example.com/v1?api_key=secret"}`}, + {name: "encoded query credential", id: OpenAICompatible, raw: `{"base_url":"https://example.com/v1?api%5Fkey=secret"}`}, + {name: "official unknown", id: OpenAI, raw: `{"base_url":"https://example.com"}`}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if _, err := registry.ValidateParams(test.id, json.RawMessage(test.raw)); err == nil { + t.Fatal("ValidateParams() error = nil") + } + }) + } + if _, err := registry.ValidateParams(ID("missing"), nil); err == nil { + t.Fatal("ValidateParams(unknown channel) error = nil") + } +} + +func TestRegistryStrictlyValidatesCredentialsWithoutLeakingSecrets(t *testing.T) { + registry := NewRegistry() + const secret = "sk-secret-value" + credential, err := registry.ValidateCredential(OpenAI, json.RawMessage(`{"api_key":" `+secret+` "}`)) + if err != nil { + t.Fatalf("ValidateCredential() error = %v", err) + } + if got, ok := credential.Value("api_key"); !ok || got != secret { + t.Fatalf("api_key = %q, %t", got, ok) + } + if got := string(credential.CanonicalJSON()); got != `{"api_key":"`+secret+`"}` { + t.Fatalf("canonical credential = %s", got) + } + encoded, err := json.Marshal(credential) + if err != nil { + t.Fatalf("json.Marshal(credential) error = %v", err) + } + if strings.Contains(string(encoded), secret) || strings.Contains(string(encoded), "api_key") { + t.Fatalf("credential JSON leaked secret fields: %s", encoded) + } + + for _, raw := range []string{ + `{}`, + `{"api_key":""}`, + `{"api_key":42}`, + `{"api_key":"one","extra":"value"}`, + `{"api_key":"one","api_key":"two"}`, + } { + _, err := registry.ValidateCredential(OpenAI, json.RawMessage(raw)) + if err == nil { + t.Fatalf("ValidateCredential(%s) error = nil", raw) + } + if strings.Contains(err.Error(), secret) || strings.Contains(err.Error(), "one") || strings.Contains(err.Error(), "two") { + t.Fatalf("validation error leaked credential value: %v", err) + } + } +} + +func TestResolvedTargetsUseExactCatalogMappingAndProtocolSpecificCapabilities(t *testing.T) { + registry := NewRegistry() + openAI, err := registry.Resolve(OpenAI, nil) + if err != nil { + t.Fatalf("Resolve(openai) error = %v", err) + } + if openAI.ChannelID != OpenAI || openAI.ProviderKind != ProviderOpenAI || + openAI.CatalogProviderID != "openai" || string(openAI.TargetConfig) != `{}` { + t.Fatalf("openai target = %#v", openAI) + } + if !openAI.Supports(protocol.OpenAIResponses, execution.OperationResponsesRetrieve, featureSet(t, execution.FeatureNativeResourceSemantics)) { + t.Fatal("openai target should support native Responses retrieval") + } + for _, operation := range []execution.Operation{ + execution.OperationResponsesCompact, + execution.OperationResponsesInputTokens, + execution.OperationResponsesPassthrough, + } { + if !openAI.Supports(protocol.OpenAIResponses, operation, execution.FeatureSet{}) { + t.Fatalf("openai target should support native %q", operation) + } + if mode, ok := openAI.Mode(protocol.OpenAIResponses, operation); !ok || mode != RouteNative { + t.Fatalf("openai %q mode = %q, %t", operation, mode, ok) + } + } + if !openAI.Supports(protocol.OpenAICompletions, execution.OperationChatCompletion, featureSet(t, execution.FeatureReasoning)) { + t.Fatal("openai native completions should support reasoning") + } + if mode, ok := openAI.Mode(protocol.OpenAICompletions, execution.OperationChatCompletion); !ok || mode != RouteNative { + t.Fatalf("openai completions mode = %q, %t", mode, ok) + } + + anthropic, err := registry.Resolve(Anthropic, nil) + if err != nil { + t.Fatalf("Resolve(anthropic) error = %v", err) + } + if anthropic.ProviderKind != ProviderAnthropic || anthropic.CatalogProviderID != "anthropic" { + t.Fatalf("anthropic catalog provider = %q", anthropic.CatalogProviderID) + } + if anthropic.Supports(protocol.OpenAIResponses, execution.OperationResponsesRetrieve, execution.FeatureSet{}) { + t.Fatal("anthropic target unexpectedly supports native Responses retrieval") + } + if anthropic.Supports(protocol.OpenAICompletions, execution.OperationChatCompletion, featureSet(t, execution.FeatureReasoning)) { + t.Fatal("converted OpenAI-to-Anthropic route unexpectedly advertises reasoning") + } + if !anthropic.Supports(protocol.OpenAICompletions, execution.OperationChatCompletion, featureSet(t, execution.FeatureStreaming)) { + t.Fatal("converted OpenAI-to-Anthropic route should support streaming") + } + if mode, ok := anthropic.Mode(protocol.OpenAICompletions, execution.OperationChatCompletion); !ok || mode != RouteConverted { + t.Fatalf("converted OpenAI-to-Anthropic mode = %q, %t", mode, ok) + } + if !anthropic.Supports(protocol.Anthropic, execution.OperationChatCompletion, featureSet(t, execution.FeatureReasoning)) { + t.Fatal("native Anthropic route should support reasoning") + } + + gemini, err := registry.Resolve(Gemini, nil) + if err != nil || gemini.ProviderKind != ProviderGemini || gemini.CatalogProviderID != "google" { + t.Fatalf("Resolve(gemini) = %#v, %v", gemini, err) + } + compatible, err := registry.Resolve(GeminiCompatible, json.RawMessage(`{"base_url":"https://proxy.example/v1/"}`)) + if err != nil { + t.Fatalf("Resolve(gemini compatible) error = %v", err) + } + if compatible.ProviderKind != ProviderGeminiCompatible || compatible.CatalogProviderID != "" || + string(compatible.TargetConfig) != `{"base_url":"https://proxy.example/v1"}` { + t.Fatalf("compatible target = %#v", compatible) + } + encodedTarget, err := json.Marshal(compatible) + if err != nil { + t.Fatalf("json.Marshal(compatible target) error = %v", err) + } + if strings.Contains(string(encodedTarget), "base_url") || strings.Contains(string(encodedTarget), "proxy.example") { + t.Fatalf("resolved target JSON leaked internal config: %s", encodedTarget) + } + if mode, ok := compatible.Mode(protocol.OpenAIResponses, execution.OperationResponsesCreate); !ok || mode != RouteConverted { + t.Fatalf("Gemini-compatible Responses create mode = %q, %t", mode, ok) + } + if compatible.Supports(protocol.OpenAIResponses, execution.OperationResponsesRetrieve, execution.FeatureSet{}) { + t.Fatal("Gemini-compatible target unexpectedly supports Responses lifecycle") + } + + openAICompatible, err := registry.Resolve(OpenAICompatible, json.RawMessage(`{"base_url":"https://proxy.example/v1"}`)) + if err != nil { + t.Fatalf("Resolve(openai compatible) error = %v", err) + } + if mode, ok := openAICompatible.Mode(protocol.OpenAIResponses, execution.OperationResponsesCreate); !ok || mode != RouteNative { + t.Fatalf("OpenAI-compatible Responses create mode = %q, %t", mode, ok) + } + if !openAICompatible.Supports(protocol.OpenAIResponses, execution.OperationResponsesRetrieve, featureSet(t, execution.FeatureNativeResourceSemantics)) { + t.Fatal("OpenAI-compatible target should preserve native Responses lifecycle") + } + for _, operation := range []execution.Operation{ + execution.OperationResponsesCompact, + execution.OperationResponsesInputTokens, + } { + if !openAICompatible.Supports(protocol.OpenAIResponses, operation, execution.FeatureSet{}) { + t.Fatalf("OpenAI-compatible target should preserve native %q", operation) + } + } + + capabilities := openAI.Capabilities(protocol.OpenAIResponses) + if !capabilities.Has(execution.OperationResponsesRetrieve) { + t.Fatal("Capabilities() omitted native operation") + } + capabilities = execution.CapabilitySet{} + if !openAI.Supports(protocol.OpenAIResponses, execution.OperationResponsesRetrieve, execution.FeatureSet{}) { + t.Fatal("mutating returned capability value changed target") + } +} + +func TestCuratedChannelResolvesCodeOwnedTargetAndCatalogMapping(t *testing.T) { + registry := NewRegistry() + target, err := registry.Resolve(DeepSeek, nil) + if err != nil { + t.Fatalf("Resolve(deepseek) error = %v", err) + } + if target.ChannelID != DeepSeek || target.ProviderKind != ProviderOpenAICompatible || + target.CatalogProviderID != "deepseek" || + string(target.TargetConfig) != `{"base_url":"https://api.deepseek.com/v1"}` { + t.Fatalf("Resolve(deepseek) = %#v", target) + } + if mode, ok := target.Mode(protocol.OpenAICompletions, execution.OperationChatCompletion); !ok || mode != RouteNative { + t.Fatalf("DeepSeek native completions mode = %q, %t", mode, ok) + } + if mode, ok := target.Mode(protocol.Anthropic, execution.OperationChatCompletion); !ok || mode != RouteConverted { + t.Fatalf("DeepSeek converted Anthropic mode = %q, %t", mode, ok) + } + descriptor, ok := registry.Get(DeepSeek) + if !ok || len(descriptor.ParamFields) != 0 || len(descriptor.CredentialFields) != 1 { + t.Fatalf("Get(deepseek) = %#v, %t", descriptor, ok) + } + encoded, err := json.Marshal(descriptor) + if err != nil { + t.Fatalf("json.Marshal(deepseek descriptor) error = %v", err) + } + if strings.Contains(string(encoded), "api.deepseek.com") { + t.Fatalf("public descriptor leaked fixed target: %s", encoded) + } + validated, err := registry.ResolveExecutionTarget(DeepSeek, target.TargetConfig) + if err != nil || validated.ProviderKind != ProviderOpenAICompatible { + t.Fatalf("ResolveExecutionTarget(deepseek) = %#v, %v", validated, err) + } + for _, tampered := range []json.RawMessage{ + nil, + json.RawMessage(`{}`), + json.RawMessage(`{"base_url":"https://attacker.example/v1"}`), + } { + if _, err := registry.ResolveExecutionTarget(DeepSeek, tampered); err == nil { + t.Fatalf("ResolveExecutionTarget(deepseek, %s) error = nil", tampered) + } + } +} + +func descriptorIDs(descriptors []Descriptor) []ID { + ids := make([]ID, len(descriptors)) + for index, descriptor := range descriptors { + ids[index] = descriptor.ID + } + return ids +} + +func featureSet(t *testing.T, features ...execution.Feature) execution.FeatureSet { + t.Helper() + set, err := execution.NewFeatureSet(features...) + if err != nil { + t.Fatalf("NewFeatureSet() error = %v", err) + } + return set +} diff --git a/internal/container/container.go b/internal/container/container.go index aadb2561..e58f7b43 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -2,6 +2,7 @@ package container import ( + "context" "fmt" "net/http" "time" @@ -12,8 +13,11 @@ import ( "gpt-load/internal/app" "gpt-load/internal/catalog" + "gpt-load/internal/channel" "gpt-load/internal/control" "gpt-load/internal/dialect" + "gpt-load/internal/execution" + bifrostexecutor "gpt-load/internal/execution/bifrost" "gpt-load/internal/gateway" "gpt-load/internal/health" "gpt-load/internal/httplifecycle" @@ -54,7 +58,8 @@ func BuildContainer() (*dig.Container, error) { app.NewEngineWithLifecycle, webui.NewServer, state.NewManager, - state.NewKeyRegistry, + state.NewCredentialRegistry, + channel.NewRegistry, control.NewPriceRuntime, control.NewCatalogBootstrap, func(bootstrap *control.CatalogBootstrap) *catalog.Runtime { return bootstrap.Runtime }, @@ -94,7 +99,7 @@ func BuildContainer() (*dig.Container, error) { }, func( cfg *config.Config, - registry *state.KeyRegistry, + registry *state.CredentialRegistry, stats *health.StatsStore, ) app.RuntimeStateCheckpoint { return app.NewFileRuntimeStateCheckpoint(cfg.DataDir, registry, stats) @@ -133,8 +138,13 @@ func BuildContainer() (*dig.Container, error) { ) dialect.Set { return dialect.NewSet(openAI, openAIResponses, anthropic, gemini) }, - gateway.NewForwarder, - func(forwarder *gateway.Forwarder) gateway.AttemptForwarder { return forwarder }, + func() (*bifrostexecutor.Runtime, error) { + return bifrostexecutor.NewRuntime(context.Background()) + }, + func(runtime *bifrostexecutor.Runtime) execution.Executor { return runtime }, + func(runtime *bifrostexecutor.Runtime) app.ExecutionRuntime { return runtime }, + gateway.NewExecutionForwarder, + func(forwarder *gateway.ExecutionForwarder) gateway.AttemptForwarder { return forwarder }, gateway.NewHandlerWithLifecycle, control.NewService, control.NewCatalogSyncCoordinator, @@ -142,8 +152,13 @@ func BuildContainer() (*dig.Container, error) { func(service *control.Service) app.StartupRecovery { return service }, control.NewServer, newHTTPRegistry, - func(db *gorm.DB, manager *state.Manager, registry *state.KeyRegistry) app.RuntimeStateLoader { - return stateloader.New(db, manager, registry) + func( + db *gorm.DB, + manager *state.Manager, + registry *state.CredentialRegistry, + channelRegistry *channel.Registry, + ) app.RuntimeStateLoader { + return stateloader.New(db, manager, registry, channelRegistry) }, app.NewApp, } diff --git a/internal/container/container_test.go b/internal/container/container_test.go index 7597c798..1baf3695 100644 --- a/internal/container/container_test.go +++ b/internal/container/container_test.go @@ -16,6 +16,7 @@ import ( "gorm.io/gorm" "gpt-load/internal/app" + "gpt-load/internal/channel" "gpt-load/internal/control" "gpt-load/internal/dialect" "gpt-load/internal/gateway" @@ -266,7 +267,7 @@ func TestBuildContainerResolvesRuntimeDependencies(t *testing.T) { db *gorm.DB, _ *gin.Engine, manager *state.Manager, - registry *state.KeyRegistry, + registry *state.CredentialRegistry, startupBootstrap app.StartupBootstrap, runtimeState app.RuntimeStateLoader, gatewayHandler *gateway.Handler, @@ -316,7 +317,7 @@ func TestBuildContainerResolvesRuntimeDependencies(t *testing.T) { if _, ok := snapshot.AccessKeysByHash[keyService.Hash(plaintext)]; !ok { t.Fatal("first snapshot cannot authenticate default AccessKey") } - if got := registry.CollectCandidates(nil, nil, time.Time{}); len(got) != 0 { + if got := registry.CollectCredentialCandidates(nil, nil, time.Time{}); len(got) != 0 { t.Fatalf("empty registry candidates = %#v", got) } if attemptForwarder == nil { @@ -802,8 +803,9 @@ func TestContainerHealthEndpointReadsSharedStatsStore(t *testing.T) { } err = dependencyContainer.Invoke(func( engine *gin.Engine, + channels *channel.Registry, manager *state.Manager, - registry *state.KeyRegistry, + registry *state.CredentialRegistry, stats *health.StatsStore, keyService encryption.Service, db *gorm.DB, @@ -814,8 +816,8 @@ func TestContainerHealthEndpointReadsSharedStatsStore(t *testing.T) { _ = sqlDB.Close() } }) - if _, publishErr := manager.Publish(state.CompileInput{Groups: []state.GroupConfig{{ - ID: 1, Name: "shared", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + if _, publishErr := manager.Publish(state.CompileInput{ChannelRegistry: channels, Groups: []state.GroupConfig{{ + ID: 1, Name: "shared", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }}}); publishErr != nil { t.Fatalf("Publish() error = %v", publishErr) @@ -824,8 +826,9 @@ func TestContainerHealthEndpointReadsSharedStatsStore(t *testing.T) { if encryptErr != nil { t.Fatalf("Encrypt() error = %v", encryptErr) } - if replaceErr := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + if replaceErr := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Status: state.CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "container-credential", Blacklisted: true, EncryptedValue: ciphertext, }}); replaceErr != nil { t.Fatalf("Replace() error = %v", replaceErr) @@ -841,19 +844,19 @@ func TestContainerHealthEndpointReadsSharedStatsStore(t *testing.T) { } var envelope struct { Data struct { - BlacklistedKeys []struct { - KeyID uint `json:"key_id"` + BlacklistedCredentials []struct { + CredentialID uint `json:"credential_id"` RecentProblemCount uint64 `json:"recent_problem_count"` - } `json:"blacklisted_keys"` + } `json:"blacklisted_credentials"` } `json:"data"` } if decodeErr := json.Unmarshal(recorder.Body.Bytes(), &envelope); decodeErr != nil { t.Fatalf("decode response: %v", decodeErr) } - if len(envelope.Data.BlacklistedKeys) != 1 || - envelope.Data.BlacklistedKeys[0].KeyID != 1 || - envelope.Data.BlacklistedKeys[0].RecentProblemCount != 1 { - t.Fatalf("blacklisted stats = %#v", envelope.Data.BlacklistedKeys) + if len(envelope.Data.BlacklistedCredentials) != 1 || + envelope.Data.BlacklistedCredentials[0].CredentialID != 1 || + envelope.Data.BlacklistedCredentials[0].RecentProblemCount != 1 { + t.Fatalf("blacklisted stats = %#v", envelope.Data.BlacklistedCredentials) } }) if err != nil { diff --git a/internal/container/route_contract_test.go b/internal/container/route_contract_test.go index c6107d80..a0e6306a 100644 --- a/internal/container/route_contract_test.go +++ b/internal/container/route_contract_test.go @@ -101,11 +101,11 @@ func TestBuildContainerExposesUnifiedRouteCatalog(t *testing.T) { }, { ModuleName: "control", - RouteName: "control.group-keys.reveal", + RouteName: "control.group-credentials.reveal", Owner: httproute.OwnerControl, Auth: httproute.AuthControl, Methods: []string{http.MethodPost}, - Path: "/api/groups/:group_id/keys/:key_id/reveal", + Path: "/api/groups/:group_id/credentials/:credential_id/reveal", }, { ModuleName: "control", diff --git a/internal/control/access_key_idempotency.go b/internal/control/access_key_idempotency.go index 06788aa8..4d0079cc 100644 --- a/internal/control/access_key_idempotency.go +++ b/internal/control/access_key_idempotency.go @@ -95,7 +95,7 @@ func (s *Service) CreateAccessKeyIdempotent( if err != nil { return idempotentMutationResult{}, err } - input, err := stateloader.BuildCompileInput(ctx, tx) + input, err := stateloader.BuildCompileInput(ctx, tx, s.channelRegistry) if err != nil { return idempotentMutationResult{}, err } diff --git a/internal/control/access_keys_test.go b/internal/control/access_keys_test.go index 2c8c265c..d664523e 100644 --- a/internal/control/access_keys_test.go +++ b/internal/control/access_keys_test.go @@ -30,8 +30,10 @@ func TestCreateAccessKeyGeneratesEncryptedSKGLToken(t *testing.T) { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, }) - if err := fixture.registry.ApplyImport(77, []state.KeyEntry{{ - ID: 88, GroupID: 77, Status: state.KeyStatusActive, EncryptedValue: "existing-upstream-cipher", + if err := fixture.registry.ApplyCredentialImport(77, []state.CredentialEntry{{ + ID: 88, GroupID: 77, Version: 1, IdentityGeneration: 1, + Fingerprint: "existing-upstream", Status: state.CredentialStatusActive, + EncryptedValue: "existing-upstream-cipher", }}); err != nil { t.Fatalf("seed Registry: %v", err) } @@ -75,7 +77,7 @@ func TestCreateAccessKeyGeneratesEncryptedSKGLToken(t *testing.T) { if _, ok := fixture.manager.Current().AccessKeysByHash[row.KeyHash]; !ok { t.Fatalf("published snapshot lacks hash %q", row.KeyHash) } - if got, ok := fixture.registry.EncryptedValue(88); !ok || got != "existing-upstream-cipher" { + if got, ok := fixture.registry.EncryptedCredentialData(88); !ok || got != "existing-upstream-cipher" { t.Fatalf("Registry value = %q, %t, want unchanged", got, ok) } if count, ok := fixture.registry.IncrFailure(88); !ok || count != 2 { @@ -383,8 +385,9 @@ func TestUpdateAccessKeyStatusAndDeletePublishWithoutMutatingRegistry(t *testing t.Fatalf("CreateAccessKey() error = %v", err) } row := loadAccessKeyRow(t, fixture.db, created.ID) - if err := fixture.registry.ApplyImport(77, []state.KeyEntry{{ - ID: 88, GroupID: 77, Status: state.KeyStatusActive, EncryptedValue: "registry-cipher", + if err := fixture.registry.ApplyCredentialImport(77, []state.CredentialEntry{{ + ID: 88, GroupID: 77, Version: 1, IdentityGeneration: 1, + Fingerprint: "registry", Status: state.CredentialStatusActive, EncryptedValue: "registry-cipher", }}); err != nil { t.Fatalf("seed Registry: %v", err) } @@ -445,7 +448,7 @@ func TestUpdateAccessKeyStatusAndDeletePublishWithoutMutatingRegistry(t *testing if got := fixture.manager.Current().Revision; got != before { t.Fatalf("revision after missing mutations = %d, want %d", got, before) } - if value, ok := fixture.registry.EncryptedValue(88); !ok || value != "registry-cipher" { + if value, ok := fixture.registry.EncryptedCredentialData(88); !ok || value != "registry-cipher" { t.Fatalf("Registry value = %q, %t, want unchanged", value, ok) } if failures, ok := fixture.registry.IncrFailure(88); !ok || failures != 2 { diff --git a/internal/control/autoweight.go b/internal/control/autoweight.go index 8ce44367..01cb2ba1 100644 --- a/internal/control/autoweight.go +++ b/internal/control/autoweight.go @@ -7,7 +7,7 @@ import ( "gpt-load/internal/state" ) -func calculateAutoWeight(stats health.KeyStats) int { +func calculateAutoWeight(stats health.CredentialStats) int { total := stats.Success + stats.Failure if total < 10 { return state.DefaultWeight diff --git a/internal/control/autoweight_test.go b/internal/control/autoweight_test.go index ead31881..ecbd0cb3 100644 --- a/internal/control/autoweight_test.go +++ b/internal/control/autoweight_test.go @@ -9,15 +9,15 @@ import ( func TestCalculateAutoWeight(t *testing.T) { tests := []struct { name string - stats health.KeyStats + stats health.CredentialStats want int }{ - {name: "empty window", stats: health.KeyStats{}, want: 50}, - {name: "insufficient samples", stats: health.KeyStats{Success: 9}, want: 50}, - {name: "all successes", stats: health.KeyStats{Success: 10}, want: 92}, - {name: "mixed results", stats: health.KeyStats{Success: 8, Failure: 2}, want: 75}, - {name: "failure streak penalty", stats: health.KeyStats{Success: 8, Failure: 2, ConsecutiveFailure: 1}, want: 38}, - {name: "minimum weight", stats: health.KeyStats{Failure: 10, ConsecutiveFailure: 10}, want: 1}, + {name: "empty window", stats: health.CredentialStats{}, want: 50}, + {name: "insufficient samples", stats: health.CredentialStats{Success: 9}, want: 50}, + {name: "all successes", stats: health.CredentialStats{Success: 10}, want: 92}, + {name: "mixed results", stats: health.CredentialStats{Success: 8, Failure: 2}, want: 75}, + {name: "failure streak penalty", stats: health.CredentialStats{Success: 8, Failure: 2, ConsecutiveFailure: 1}, want: 38}, + {name: "minimum weight", stats: health.CredentialStats{Failure: 10, ConsecutiveFailure: 10}, want: 1}, } for _, tt := range tests { diff --git a/internal/control/catalog_discovery_merge_test.go b/internal/control/catalog_discovery_merge_test.go index 9da7575f..e5d8416f 100644 --- a/internal/control/catalog_discovery_merge_test.go +++ b/internal/control/catalog_discovery_merge_test.go @@ -2,11 +2,12 @@ package control import ( "context" + "encoding/json" "reflect" "testing" "gpt-load/internal/catalog" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" "gpt-load/internal/pricing" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -29,23 +30,20 @@ func TestDraftDiscoveryMergesLiveAndLocalCatalogByExactIDWithoutURLInference(t * }}) zero := int64(0) if err := fixture.db.Create(&models.ModelPrice{ + ChannelID: string(channel.OpenAI), ModelID: "shared", InputPriceNanoUSDPerMillionTokens: &zero, }).Error; err != nil { t.Fatal(err) } - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { return []string{"shared", "live-only", "shared"}, nil }, }) - providerID := "openai" got, err := fixture.service.DiscoverModels(t.Context(), ModelDiscoveryRequest{ - ProviderID: &providerID, - UpstreamURL: "https://proxy.example/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "secret", + ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Credentials: "secret", }) if err != nil { t.Fatal(err) @@ -65,15 +63,14 @@ func TestDraftDiscoveryMergesLiveAndLocalCatalogByExactIDWithoutURLInference(t * } withoutProvider, err := fixture.service.DiscoverModels(t.Context(), ModelDiscoveryRequest{ - UpstreamURL: "https://api.openai.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "secret", + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://api.openai.com/v1"}`), Credentials: "secret", }) if err != nil { t.Fatal(err) } if !reflect.DeepEqual(withoutProvider.Models, []ModelCandidate{ - {ID: "shared", Name: "shared", Sources: []string{"live"}, PricingStatus: PricingStatusConfigured}, + {ID: "shared", Name: "shared", Sources: []string{"live"}, PricingStatus: PricingStatusPending}, {ID: "live-only", Name: "live-only", Sources: []string{"live"}, PricingStatus: PricingStatusPending}, }) { t.Fatalf("URL-inferred candidates = %#v", withoutProvider.Models) @@ -98,17 +95,24 @@ func TestDiscoveryPricingStatusUsesAutomaticPriceReferenceForDraftScopes(t *test }, }, }}) - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ - value: protocol.OpenAICompletions, - listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { - return []string{"shared-model"}, nil + fixture.service.executor = newRecordingDiscoveryExecutor( + &recordingDiscoveryExecutorTarget{ + value: protocol.OpenAICompletions, + listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { + return []string{"shared-model"}, nil + }, }, - }) + &recordingDiscoveryExecutorTarget{ + value: protocol.Anthropic, + listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { + return []string{"shared-model"}, nil + }, + }, + ) request := ModelDiscoveryRequest{ - UpstreamURL: "https://proxy.example/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "secret", + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://proxy.example/v1"}`), Credentials: "secret", } custom, err := fixture.service.DiscoverModels(t.Context(), request) if err != nil { @@ -124,10 +128,10 @@ func TestDiscoveryPricingStatusUsesAutomaticPriceReferenceForDraftScopes(t *test } savedCustom, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://saved-custom.example/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://saved-custom.example/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "saved-secret", + Credentials: "saved-secret", }) if err != nil { t.Fatal(err) @@ -145,8 +149,8 @@ func TestDiscoveryPricingStatusUsesAutomaticPriceReferenceForDraftScopes(t *test t.Fatalf("custom Group candidates = %#v", customGroup.Models) } - providerID := "anthropic" - request.ProviderID = &providerID + request.ChannelID = channel.Anthropic + request.Params = json.RawMessage(`{}`) provider, err := fixture.service.DiscoverModels(t.Context(), request) if err != nil { t.Fatal(err) @@ -154,7 +158,7 @@ func TestDiscoveryPricingStatusUsesAutomaticPriceReferenceForDraftScopes(t *test if !reflect.DeepEqual(provider.Models, []ModelCandidate{ { ID: "shared-model", Name: "Anthropic model", Sources: []string{"live", "catalog"}, - PricingStatus: PricingStatusConfigured, PricingSource: &openAI, + PricingStatus: PricingStatusPending, }, }) { t.Fatalf("provider draft candidates = %#v", provider.Models) @@ -174,7 +178,7 @@ func TestSavedGroupDiscoveryUsesPersistedProviderAndSharedPricingStatus(t *testi }, }, }}) - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { return []string{"shared", "live-only"}, nil @@ -182,18 +186,17 @@ func TestSavedGroupDiscoveryUsesPersistedProviderAndSharedPricingStatus(t *testi }) providerID := "openai" created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - ProviderID: &providerID, - UpstreamURL: "https://saved-proxy.example/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{{ID: "shared"}}}, - Keys: "saved-secret", + Credentials: "saved-secret", }) if err != nil { t.Fatal(err) } zero := int64(0) if err := fixture.db.Model(&models.ModelPrice{}). - Where("model_id = ?", "shared"). + Where("channel_id = ? AND model_id = ?", channel.OpenAI, "shared"). Update("input_price_nano_usd_per_million_tokens", &zero).Error; err != nil { t.Fatal(err) } diff --git a/internal/control/catalog_sync.go b/internal/control/catalog_sync.go index c78892c6..66f96835 100644 --- a/internal/control/catalog_sync.go +++ b/internal/control/catalog_sync.go @@ -666,16 +666,26 @@ func logMissingAutomaticPricePriorityProviders(snapshot *catalog.Snapshot) { } func reconcileCatalogAutomaticPrices(tx *gorm.DB, snapshot *catalog.Snapshot) error { + references, err := loadReferencedPrices(tx) + if err != nil { + return err + } var rows []models.ModelPrice if err := tx.Where("is_manual = ?", false).Order("id ASC").Find(&rows).Error; err != nil { return fmt.Errorf("load automatic model prices: %w", app_errors.ParseDBError(err)) } for _, row := range rows { - cost, _, ok := resolveAutomaticPrice(snapshot, row.ModelID) + identity, err := PriceIdentityForChannelModel(row.ChannelID, row.ModelID) + if err != nil { + return fmt.Errorf("validate automatic model price identity: %w", app_errors.ErrInternalServer) + } + if _, referenced := references[identity]; !referenced { + continue + } + match, ok := resolveAutomaticPriceForIdentity(snapshot, identity) desired := models.ModelPrice{} if ok { - var err error - desired, err = automaticCatalogValues(cost) + desired, err = automaticCatalogValues(match.cost) if err != nil { return fmt.Errorf("normalize catalog price: %w", app_errors.ErrInternalServer) } @@ -754,7 +764,10 @@ func cleanupUnreferencedAutomaticPrices(tx *gorm.DB) error { } ids := make([]uint, 0) for _, row := range rows { - identity := pricing.Identity{ModelID: row.ModelID} + identity, err := PriceIdentityForChannelModel(row.ChannelID, row.ModelID) + if err != nil { + return fmt.Errorf("validate automatic price cleanup identity: %w", app_errors.ErrInternalServer) + } if _, referenced := references[identity]; !referenced { ids = append(ids, row.ID) } diff --git a/internal/control/catalog_sync_test.go b/internal/control/catalog_sync_test.go index c934f918..3d0afd80 100644 --- a/internal/control/catalog_sync_test.go +++ b/internal/control/catalog_sync_test.go @@ -17,6 +17,7 @@ import ( "github.com/sirupsen/logrus" "gpt-load/internal/catalog" + "gpt-load/internal/channel" "gpt-load/internal/pricing" "gpt-load/internal/storage/models" ) @@ -226,8 +227,10 @@ func TestApplyCatalogSnapshotDoesNotLogCompletePriorityAndLogsEachSuccess(t *tes func TestApplyCatalogSnapshotFailureDoesNotLogPriorityWarningOrPublishRuntime(t *testing.T) { fixture := newServiceFixture(t) + seedCatalogPriceGroup(t, fixture, "failure", nil, []string{"gpt"}) oldPrice := int64(1) if err := fixture.db.Create(&models.ModelPrice{ + ChannelID: string(channel.OpenAICompatible), ModelID: "gpt", InputPriceNanoUSDPerMillionTokens: &oldPrice, }).Error; err != nil { @@ -319,6 +322,7 @@ func TestCatalogStartupReconcilesDurableLKGBeforeAnyNetworkSync(t *testing.T) { seedCatalogPriceGroup(t, fixture, "startup-lkg", &providerID, []string{"gpt"}) old := int64(1) if err := fixture.db.Create(&models.ModelPrice{ + ChannelID: string(channel.OpenAICompatible), ModelID: "gpt", InputPriceNanoUSDPerMillionTokens: &old, }).Error; err != nil { @@ -338,7 +342,7 @@ func TestCatalogStartupReconcilesDurableLKGBeforeAnyNetworkSync(t *testing.T) { if err := fixture.service.EnsureInitialState(t.Context()); err != nil { t.Fatal(err) } - assertCatalogPriceRow(t, fixture, "provider:openai", "gpt", func(row models.ModelPrice) { + assertCatalogPriceRow(t, fixture, "gpt", func(row models.ModelPrice) { if row.InputPriceNanoUSDPerMillionTokens == nil || *row.InputPriceNanoUSDPerMillionTokens != 9 { t.Fatalf("startup reconciled row = %#v", row) @@ -466,7 +470,7 @@ func TestCatalogSync304PreservesPublishedGenerationsAndDoesNotStoreCache(t *test }} fixture.catalogRuntime.Publish(oldSnapshot) oldTable, err := pricing.NewTable([]pricing.Rule{{ - Identity: pricing.Identity{ModelID: "old"}, + Identity: pricing.Identity{ChannelID: string(channel.OpenAICompatible), ModelID: "old"}, Prices: pricing.Prices{Input: pricing.Price{Set: true, NanoUSDPerMillion: 1}}, }}) if err != nil { @@ -1009,13 +1013,13 @@ func TestCatalogSyncReconcilesSlotsLKGManualCleanupAndStableTimestamps(t *testin oldInput, oldOutput, oldRead, oldWrite := int64(1), int64(2), int64(3), int64(4) oldTiers := models.JSON(`[{"threshold_tokens":128000,"input_price_nano_usd_per_million_tokens":5}]`) rows := []models.ModelPrice{ - {ModelID: "changed", InputPriceNanoUSDPerMillionTokens: &oldInput, OutputPriceNanoUSDPerMillionTokens: &oldOutput, CacheReadPriceNanoUSDPerMillionTokens: &oldRead, CacheWritePriceNanoUSDPerMillionTokens: &oldWrite, UpdatedAtMS: 111}, - {ModelID: "missing-model", InputPriceNanoUSDPerMillionTokens: &oldInput, OutputPriceNanoUSDPerMillionTokens: &oldOutput, CacheReadPriceNanoUSDPerMillionTokens: &oldRead, CacheWritePriceNanoUSDPerMillionTokens: &oldWrite, ContextPriceTiers: oldTiers, UpdatedAtMS: 112}, - {ModelID: "missing-provider-model", InputPriceNanoUSDPerMillionTokens: &oldInput, OutputPriceNanoUSDPerMillionTokens: &oldOutput, CacheReadPriceNanoUSDPerMillionTokens: &oldRead, CacheWritePriceNanoUSDPerMillionTokens: &oldWrite, ContextPriceTiers: oldTiers, UpdatedAtMS: 113}, - {ModelID: "manual", IsManual: true, UpdatedAtMS: 114}, - {ModelID: "stale", InputPriceNanoUSDPerMillionTokens: &oldInput, UpdatedAtMS: 115}, - {ModelID: "manual-stale", IsManual: true, UpdatedAtMS: 116}, - {ModelID: "unchanged", InputPriceNanoUSDPerMillionTokens: &oldInput, UpdatedAtMS: 117}, + {ChannelID: string(channel.OpenAI), ModelID: "changed", InputPriceNanoUSDPerMillionTokens: &oldInput, OutputPriceNanoUSDPerMillionTokens: &oldOutput, CacheReadPriceNanoUSDPerMillionTokens: &oldRead, CacheWritePriceNanoUSDPerMillionTokens: &oldWrite, UpdatedAtMS: 111}, + {ChannelID: string(channel.OpenAI), ModelID: "missing-model", InputPriceNanoUSDPerMillionTokens: &oldInput, OutputPriceNanoUSDPerMillionTokens: &oldOutput, CacheReadPriceNanoUSDPerMillionTokens: &oldRead, CacheWritePriceNanoUSDPerMillionTokens: &oldWrite, ContextPriceTiers: oldTiers, UpdatedAtMS: 112}, + {ChannelID: string(channel.OpenAICompatible), ModelID: "missing-provider-model", InputPriceNanoUSDPerMillionTokens: &oldInput, OutputPriceNanoUSDPerMillionTokens: &oldOutput, CacheReadPriceNanoUSDPerMillionTokens: &oldRead, CacheWritePriceNanoUSDPerMillionTokens: &oldWrite, ContextPriceTiers: oldTiers, UpdatedAtMS: 113}, + {ChannelID: string(channel.OpenAI), ModelID: "manual", IsManual: true, UpdatedAtMS: 114}, + {ChannelID: string(channel.OpenAICompatible), ModelID: "stale", InputPriceNanoUSDPerMillionTokens: &oldInput, UpdatedAtMS: 115}, + {ChannelID: string(channel.OpenAICompatible), ModelID: "manual-stale", IsManual: true, UpdatedAtMS: 116}, + {ChannelID: string(channel.OpenAI), ModelID: "unchanged", InputPriceNanoUSDPerMillionTokens: &oldInput, UpdatedAtMS: 117}, } for index := range rows { if err := fixture.db.Create(&rows[index]).Error; err != nil { @@ -1040,34 +1044,34 @@ func TestCatalogSyncReconcilesSlotsLKGManualCleanupAndStableTimestamps(t *testin t.Fatal(err) } - assertCatalogPriceRow(t, fixture, "provider:openai", "changed", func(row models.ModelPrice) { + assertCatalogPriceRow(t, fixture, "changed", func(row models.ModelPrice) { if row.InputPriceNanoUSDPerMillionTokens == nil || *row.InputPriceNanoUSDPerMillionTokens != 9 || row.OutputPriceNanoUSDPerMillionTokens != nil || row.CacheReadPriceNanoUSDPerMillionTokens != nil || row.CacheWritePriceNanoUSDPerMillionTokens != nil { t.Fatalf("changed row = %#v, want exact replacement and cleared missing slots", row) } }) - assertCatalogPriceRow(t, fixture, "provider:openai", "missing-model", func(row models.ModelPrice) { + assertCatalogPriceRow(t, fixture, "missing-model", func(row models.ModelPrice) { if priceTestRowHasValue(row) { t.Fatalf("missing model retained automatic catalog values: %#v", row) } }) - assertCatalogPriceRow(t, fixture, "provider:missing-provider", "missing-provider-model", func(row models.ModelPrice) { + assertCatalogPriceRow(t, fixture, "missing-provider-model", func(row models.ModelPrice) { if priceTestRowHasValue(row) { t.Fatalf("missing provider retained automatic catalog values: %#v", row) } }) - assertCatalogPriceRow(t, fixture, "provider:openai", "manual", func(row models.ModelPrice) { + assertCatalogPriceRow(t, fixture, "manual", func(row models.ModelPrice) { if !row.IsManual || row.InputPriceNanoUSDPerMillionTokens != nil || row.UpdatedAtMS != 114 { t.Fatalf("manual all-null row changed: %#v", row) } }) - assertCatalogPriceRow(t, fixture, "provider:openai", "unchanged", func(row models.ModelPrice) { + assertCatalogPriceRow(t, fixture, "unchanged", func(row models.ModelPrice) { if row.UpdatedAtMS != 117 { t.Fatalf("unchanged row timestamp = %d, want 117", row.UpdatedAtMS) } }) - assertCatalogPriceMissing(t, fixture, "provider:openai", "stale") - assertCatalogPriceRow(t, fixture, "provider:openai", "manual-stale", func(row models.ModelPrice) { + assertCatalogPriceMissing(t, fixture, "stale") + assertCatalogPriceRow(t, fixture, "manual-stale", func(row models.ModelPrice) { if !row.IsManual { t.Fatal("unreferenced manual row lost ownership") } @@ -1077,17 +1081,17 @@ func TestCatalogSyncReconcilesSlotsLKGManualCleanupAndStableTimestamps(t *testin func TestCatalogSyncReconcilesGroupAutomaticPricesAndPreservesManualRows(t *testing.T) { fixture := newServiceFixture(t) group := createPriceTestGroup(t, fixture.db, models.Group{ - Name: "custom", UpstreamURL: "https://custom.example/v1", - Protocols: models.JSON(`["openai-completions"]`), + Name: "custom", ChannelID: string(channel.OpenAICompatible), + Params: models.JSON(`{"base_url":"https://custom.example/v1"}`), Models: models.JSON(`[{"id":"group-model","alias":""},{"id":"manual-group-model","alias":""}]`), - Config: models.JSON(`{}`), Enabled: true, + Overrides: models.JSON(`{}`), Enabled: true, }) - scopeKey, _ := pricing.GroupScopeKey(group.ID) oldInput, oldOutput, oldRead, oldWrite := int64(1), int64(2), int64(3), int64(4) oldTiers := models.JSON(`[{"threshold_tokens":128000,"input_price_nano_usd_per_million_tokens":5}]`) manualInput := int64(99) for _, row := range []models.ModelPrice{ { + ChannelID: group.ChannelID, ModelID: "group-model", InputPriceNanoUSDPerMillionTokens: &oldInput, OutputPriceNanoUSDPerMillionTokens: &oldOutput, @@ -1096,6 +1100,7 @@ func TestCatalogSyncReconcilesGroupAutomaticPricesAndPreservesManualRows(t *test ContextPriceTiers: oldTiers, UpdatedAtMS: 111, }, { + ChannelID: group.ChannelID, ModelID: "manual-group-model", InputPriceNanoUSDPerMillionTokens: &manualInput, IsManual: true, UpdatedAtMS: 112, @@ -1136,7 +1141,7 @@ func TestCatalogSyncReconcilesGroupAutomaticPricesAndPreservesManualRows(t *test if err := fixture.service.applyCatalogSnapshot(t.Context(), snapshot); err != nil { t.Fatal(err) } - assertCatalogPriceRow(t, fixture, scopeKey, "group-model", func(row models.ModelPrice) { + assertCatalogPriceRow(t, fixture, "group-model", func(row models.ModelPrice) { if row.IsManual || row.InputPriceNanoUSDPerMillionTokens == nil || *row.InputPriceNanoUSDPerMillionTokens != 8 || row.OutputPriceNanoUSDPerMillionTokens != nil || @@ -1155,7 +1160,7 @@ func TestCatalogSyncReconcilesGroupAutomaticPricesAndPreservesManualRows(t *test t.Fatalf("automatic Group tiers = %#v", tiers) } }) - assertCatalogPriceRow(t, fixture, scopeKey, "manual-group-model", func(row models.ModelPrice) { + assertCatalogPriceRow(t, fixture, "manual-group-model", func(row models.ModelPrice) { if !row.IsManual || row.InputPriceNanoUSDPerMillionTokens == nil || *row.InputPriceNanoUSDPerMillionTokens != 99 || row.UpdatedAtMS != 112 { t.Fatalf("manual Group row changed: %#v", row) @@ -1163,12 +1168,88 @@ func TestCatalogSyncReconcilesGroupAutomaticPricesAndPreservesManualRows(t *test }) } +func TestCatalogSyncScopesSameModelByChannelAndPreservesManualOverride(t *testing.T) { + fixture := newServiceFixture(t) + for _, group := range []models.Group{ + { + Name: "official-openai", ChannelID: string(channel.OpenAI), Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"shared"}]`), Enabled: true, + }, + { + Name: "official-anthropic", ChannelID: string(channel.Anthropic), Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"shared"}]`), Enabled: true, + }, + } { + createPriceTestGroup(t, fixture.db, group) + } + manual := int64(99) + if err := fixture.db.Create(&models.ModelPrice{ + ChannelID: string(channel.OpenAI), ModelID: "shared", + InputPriceNanoUSDPerMillionTokens: &manual, IsManual: true, + }).Error; err != nil { + t.Fatal(err) + } + snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ + "openai": catalogProviderFixture("openai", "OpenAI", "shared", 10), + "anthropic": catalogProviderFixture("anthropic", "Anthropic", "shared", 20), + }} + + if err := fixture.service.applyCatalogSnapshot(t.Context(), snapshot); err != nil { + t.Fatal(err) + } + var rows []models.ModelPrice + if err := fixture.db.Where("model_id = ?", "shared").Order("channel_id ASC").Find(&rows).Error; err != nil { + t.Fatal(err) + } + if len(rows) != 2 { + t.Fatalf("channel prices = %#v", rows) + } + byChannel := make(map[string]models.ModelPrice, len(rows)) + for _, row := range rows { + byChannel[row.ChannelID] = row + } + openAI := byChannel[string(channel.OpenAI)] + if !openAI.IsManual || openAI.InputPriceNanoUSDPerMillionTokens == nil || + *openAI.InputPriceNanoUSDPerMillionTokens != 99 { + t.Fatalf("OpenAI manual price = %#v", openAI) + } + anthropic := byChannel[string(channel.Anthropic)] + if anthropic.IsManual || anthropic.InputPriceNanoUSDPerMillionTokens == nil || + *anthropic.InputPriceNanoUSDPerMillionTokens != 20 { + t.Fatalf("Anthropic automatic price = %#v", anthropic) + } + + listed, err := fixture.service.ListModelPrices(t.Context(), ModelPriceListQuery{ + Usage: ModelPriceUsageAll, Status: ModelPriceStatusAll, Page: 1, PageSize: 20, + }) + if err != nil { + t.Fatal(err) + } + if len(listed.Items) != 2 { + t.Fatalf("listed channel prices = %#v", listed.Items) + } + for _, item := range listed.Items { + if item.ChannelID == string(channel.OpenAI) { + if item.MatchSource != nil || item.MatchedProviderID != nil { + t.Fatalf("manual price exposes automatic match = %#v", item) + } + continue + } + if item.ChannelID != string(channel.Anthropic) || item.MatchSource == nil || + *item.MatchSource != ModelPriceMatchSourceChannelCatalogProvider || + item.MatchedProviderID == nil || *item.MatchedProviderID != "anthropic" { + t.Fatalf("automatic price match = %#v", item) + } + } +} + func TestCatalogSyncReconcileFailurePublishesNeitherRuntimeAndKeepsPendingLKG(t *testing.T) { fixture := newServiceFixture(t) providerID := "openai" seedCatalogPriceGroup(t, fixture, "openai", &providerID, []string{"gpt"}) oldValue := int64(1) if err := fixture.db.Create(&models.ModelPrice{ + ChannelID: string(channel.OpenAI), ModelID: "gpt", InputPriceNanoUSDPerMillionTokens: &oldValue, }).Error; err != nil { @@ -1250,7 +1331,7 @@ func catalogResultFixture( } rawProviders[providerID] = map[string]any{ "id": provider.ID, "name": provider.Name, - "api": provider.APIURL, "npm": provider.NPM, "models": rawModels, + "api": "", "npm": "", "models": rawModels, } } raw, err := json.Marshal(rawProviders) @@ -1271,7 +1352,7 @@ func catalogResultFixture( } func catalogProviderFixture(id, name, modelID string, input pricing.NanoUSD) catalog.Provider { - provider := catalog.Provider{ID: id, Name: name, APIURL: "https://example.com/v1", Models: map[string]catalog.Model{}} + provider := catalog.Provider{ID: id, Name: name, Models: map[string]catalog.Model{}} if modelID != "" { provider.Models[modelID] = catalog.Model{ ID: modelID, Name: modelID, @@ -1299,9 +1380,21 @@ func seedCatalogPriceGroup( if err != nil { t.Fatal(err) } + channelID := channel.OpenAICompatible + params := models.JSON(`{"base_url":"https://` + name + `.example.com/v1"}`) + if providerID != nil { + switch *providerID { + case "openai": + channelID, params = channel.OpenAI, models.JSON(`{}`) + case "anthropic": + channelID, params = channel.Anthropic, models.JSON(`{}`) + case "google": + channelID, params = channel.Gemini, models.JSON(`{}`) + } + } if err := fixture.db.Create(&models.Group{ - Name: name, ProviderID: providerID, UpstreamURL: "https://" + name + ".example.com/v1", - Protocols: models.JSON(`["openai-completions"]`), Models: models.JSON(encoded), Enabled: true, + Name: name, ChannelID: string(channelID), Params: params, + Models: models.JSON(encoded), Enabled: true, }).Error; err != nil { t.Fatal(err) } @@ -1310,7 +1403,6 @@ func seedCatalogPriceGroup( func assertCatalogPriceRow( t *testing.T, fixture serviceFixture, - scopeKey string, modelID string, assert func(models.ModelPrice), ) { @@ -1322,7 +1414,7 @@ func assertCatalogPriceRow( assert(row) } -func assertCatalogPriceMissing(t *testing.T, fixture serviceFixture, scopeKey, modelID string) { +func assertCatalogPriceMissing(t *testing.T, fixture serviceFixture, modelID string) { t.Helper() var count int64 if err := fixture.db.Model(&models.ModelPrice{}). diff --git a/internal/control/channel_group_foundation_test.go b/internal/control/channel_group_foundation_test.go new file mode 100644 index 00000000..c170ad6c --- /dev/null +++ b/internal/control/channel_group_foundation_test.go @@ -0,0 +1,451 @@ +package control + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "reflect" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" + + "gpt-load/internal/channel" + "gpt-load/internal/platform/config" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/storage/models" +) + +func TestServiceUsesInjectedChannelRegistryAndListsSafeDescriptors(t *testing.T) { + t.Parallel() + + fixture := newServiceFixture(t) + if fixture.service.channelRegistry != fixture.channelRegistry { + t.Fatal("Service did not retain the injected ChannelRegistry") + } + + result, err := fixture.service.ListChannels(context.Background(), "openai") + if err != nil { + t.Fatalf("ListChannels() error = %v", err) + } + if result.Total != 3 || len(result.Items) != 3 || + result.Items[0].ID != channel.OpenAI || result.Items[1].ID != channel.AzureOpenAI || + result.Items[2].ID != channel.OpenAICompatible { + t.Fatalf("ListChannels(openai) = %#v", result) + } + encoded, err := json.Marshal(result) + if err != nil { + t.Fatalf("json.Marshal() error = %v", err) + } + text := strings.ToLower(string(encoded)) + implementationName := "bif" + "rost" + if strings.Contains(text, "target_config") || strings.Contains(text, "catalog_provider_id") || + strings.Contains(text, implementationName) || strings.Contains(text, "credential_value") { + t.Fatalf("channel descriptors leaked internal execution data: %s", encoded) + } + common, err := fixture.service.ListChannels(context.Background(), "deep seek") + if err != nil || common.Total != 1 || len(common.Items) != 1 || + common.Items[0].ID != channel.DeepSeek || len(common.Items[0].ParamFields) != 0 { + t.Fatalf("ListChannels(deep seek) = %#v, %v", common, err) + } + commonJSON, err := json.Marshal(common) + if err != nil || strings.Contains(strings.ToLower(string(commonJSON)), "deepseek.com") { + t.Fatalf("common preset leaked fixed target: %s, %v", commonJSON, err) + } +} + +func TestParseChannelQueryIsStrict(t *testing.T) { + t.Parallel() + + if got, apiErr := parseChannelQuery("q=%20claude%20", false); apiErr != nil || got != "claude" { + t.Fatalf("parseChannelQuery() = %q, %v", got, apiErr) + } + for _, raw := range []string{"unknown=1", "q=a&q=b"} { + if _, apiErr := parseChannelQuery(raw, false); apiErr != app_errors.ErrBadRequest { + t.Fatalf("parseChannelQuery(%q) error = %v, want bad request", raw, apiErr) + } + } + if _, apiErr := parseChannelQuery("", true); apiErr != app_errors.ErrBadRequest { + t.Fatalf("parseChannelQuery(force query) error = %v, want bad request", apiErr) + } +} + +func TestChannelsHTTPIsAuthenticatedAndStrict(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + engine := gin.New() + NewServer(&config.Config{AuthKey: "channel-auth"}, fixture.service).RegisterRoutes(engine) + + request := httptest.NewRequest(http.MethodGet, "/api/channels?q=deep%20seek", nil) + request.Header.Set("Authorization", "Bearer channel-auth") + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + body := recorder.Body.String() + if recorder.Code != http.StatusOK || !strings.Contains(body, `"channel_id":"deepseek"`) { + t.Fatalf("authenticated channels response = %d %s", recorder.Code, recorder.Body.String()) + } + if strings.Contains(body, `"param_fields":null`) || + strings.Contains(body, `"credential_fields":null`) || + strings.Contains(body, `"client_protocols":null`) { + t.Fatalf("channels response must encode descriptor collections as arrays: %s", body) + } + + for name, path := range map[string]string{ + "unauthenticated": "/api/channels", + "unknown query": "/api/channels?unknown=1", + } { + t.Run(name, func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, path, nil) + if name != "unauthenticated" { + request.Header.Set("Authorization", "Bearer channel-auth") + } + engine.ServeHTTP(recorder, request) + if name == "unauthenticated" && recorder.Code != http.StatusUnauthorized { + t.Fatalf("response = %d %s, want 401", recorder.Code, recorder.Body.String()) + } + if name == "unknown query" && recorder.Code != http.StatusBadRequest { + t.Fatalf("response = %d %s, want 400", recorder.Code, recorder.Body.String()) + } + }) + } + + for _, path := range []string{ + "/api/provider-suggestions", + "/api/providers/openai/models", + } { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, path, nil) + request.Header.Set("Authorization", "Bearer channel-auth") + engine.ServeHTTP(recorder, request) + if recorder.Code != http.StatusNotFound { + t.Fatalf("retired provider catalog route %q = %d %s, want 404", path, recorder.Code, recorder.Body.String()) + } + } +} + +func TestGroupCreateWireAcceptsOnlyChannelContract(t *testing.T) { + t.Parallel() + + var request GroupCreateRequest + valid := []byte(`{ + "channel_id":"openai_compatible", + "params":{"base_url":"https://proxy.example/v1"}, + "models":[], + "credentials":"sk-one", + "confirm_same_target":true + }`) + if err := decodeStrictControlJSONObject(valid, &request); err != nil { + t.Fatalf("decode channel GroupCreateRequest: %v", err) + } + if request.ChannelID != channel.OpenAICompatible || request.Credentials != "sk-one" || !request.ConfirmSameTarget { + t.Fatalf("GroupCreateRequest = %#v", request) + } + for _, legacy := range []string{"keys", "provider_id", "upstream_url", "protocols", "confirm_same_upstream_url"} { + body := []byte(`{"channel_id":"openai","models":[],"credentials":"sk-one","` + legacy + `":null}`) + if err := decodeStrictControlJSONObject(body, &GroupCreateRequest{}); err == nil { + t.Fatalf("legacy field %q was accepted", legacy) + } + } +} + +func TestCreateChannelGroupPersistsCanonicalCredentialsAndPublishes(t *testing.T) { + fixture := newServiceFixture(t) + fixture.service.now = func() time.Time { return time.UnixMilli(1_786_000_000_000) } + beforeRevision := fixture.manager.Current().Revision + + result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer(" channel group "), + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":" HTTPS://Proxy.Example/v1/ "}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{ + {ID: " provider-model ", Alias: " public ", AliasEnabled: true}, + }}, + Credentials: " sk-one \n sk-one\nsk-two\n", + }) + if err != nil { + t.Fatalf("CreateGroup() error = %v", err) + } + if result.CredentialsAdded != 2 || result.CredentialsDuplicated != 1 { + t.Fatalf("CreateGroup() result = %#v", result) + } + encodedResult, err := json.Marshal(result) + if err != nil { + t.Fatalf("json.Marshal(result) error = %v", err) + } + var resultFields map[string]json.RawMessage + if err := json.Unmarshal(encodedResult, &resultFields); err != nil { + t.Fatalf("json.Unmarshal(result) error = %v", err) + } + if len(resultFields) != 4 || resultFields["credentials_added"] == nil || + resultFields["credentials_duplicated"] == nil || resultFields["keys_added"] != nil { + t.Fatalf("result fields = %s", encodedResult) + } + + var group models.Group + if err := fixture.db.First(&group, result.GroupID).Error; err != nil { + t.Fatalf("load group: %v", err) + } + if group.ChannelID != string(channel.OpenAICompatible) || string(group.Params) != `{"base_url":"https://proxy.example/v1"}` { + t.Fatalf("stored group = %#v", group) + } + var credentials []models.Credential + if err := fixture.db.Where("group_id = ?", group.ID).Order("id ASC").Find(&credentials).Error; err != nil { + t.Fatalf("load credentials: %v", err) + } + if len(credentials) != 2 { + t.Fatalf("credentials = %#v, want two", credentials) + } + wantCanonical := []string{`{"api_key":"sk-one"}`, `{"api_key":"sk-two"}`} + for index, row := range credentials { + plaintext, decryptErr := fixture.encryption.Decrypt(row.Data) + if decryptErr != nil { + t.Fatalf("decrypt credential %d: %v", row.ID, decryptErr) + } + if plaintext != wantCanonical[index] || row.Fingerprint != fixture.encryption.Hash(wantCanonical[index]) || + row.UpdatedAtMS < 1 { + t.Fatalf("credential %d = plaintext %q row %#v", index, plaintext, row) + } + } + snapshot := fixture.manager.Current() + view, ok := snapshot.Groups[group.ID] + if snapshot.Revision != beforeRevision+1 || !ok || view.ChannelID != channel.OpenAICompatible || + string(view.Params) != `{"base_url":"https://proxy.example/v1"}` { + t.Fatalf("published snapshot = revision %d group %#v exists=%t", snapshot.Revision, view, ok) + } + refs := fixture.registry.CaptureActiveCredentialRefs([]uint{group.ID}) + if len(refs) != 2 || refs[0].Version == 0 || refs[0].IdentityGeneration == 0 || refs[0].Fingerprint == "" { + t.Fatalf("credential refs = %#v", refs) + } + if got := loadCreatedGroupModels(t, fixture, group.ID); !reflect.DeepEqual(got, []GroupModel{{ID: "provider-model", Alias: "public"}}) { + t.Fatalf("stored models = %#v", got) + } +} + +func TestCreateChannelGroupUsesChannelAndCanonicalParamsForSimilarity(t *testing.T) { + fixture := newServiceFixture(t) + first, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"HTTPS://Same.Example/v1/"}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, + Credentials: "first", + }) + if err != nil { + t.Fatalf("first CreateGroup() error = %v", err) + } + _, err = fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://same.example/v1"}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, + Credentials: "second", + }) + var apiErr *app_errors.APIError + if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrChannelTargetConflict.Code { + t.Fatalf("same target error = %v", err) + } + data, ok := apiErr.Data.(SameTargetConflictData) + if !ok || len(data.Groups) != 1 || data.Groups[0].ID != first.GroupID { + t.Fatalf("same target conflict data = %#v", apiErr.Data) + } + + confirmed, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://same.example/v1"}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, + Credentials: "second", + ConfirmSameTarget: true, + }) + if err != nil || confirmed.GroupID == 0 { + t.Fatalf("confirmed CreateGroup() = %#v, %v", confirmed, err) + } + + otherChannel, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("same URL other channel"), + ChannelID: channel.AnthropicCompatible, + Params: json.RawMessage(`{"base_url":"https://same.example/v1"}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, + Credentials: "third", + }) + if err != nil || otherChannel.GroupID == 0 { + t.Fatalf("other-channel CreateGroup() = %#v, %v", otherChannel, err) + } +} + +func TestCreateChannelGroupIdempotencyReplaysCredentialCounts(t *testing.T) { + fixture := newServiceFixture(t) + request := GroupCreateRequest{ + Name: stringPointer("idempotent channel"), + ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, + Credentials: " repeated \nrepeated\n", + } + const idempotencyKey = "328f47a2-9c35-4d6e-8b1a-1234567890ab" + first, err := fixture.service.CreateGroupIdempotent(t.Context(), idempotencyKey, request) + if err != nil { + t.Fatalf("first CreateGroupIdempotent() error = %v", err) + } + replayed, err := fixture.service.CreateGroupIdempotent(t.Context(), idempotencyKey, request) + if err != nil { + t.Fatalf("replay CreateGroupIdempotent() error = %v", err) + } + if !reflect.DeepEqual(replayed, first) || first.CredentialsAdded != 1 || first.CredentialsDuplicated != 1 { + t.Fatalf("first/replayed = %#v / %#v", first, replayed) + } + var credentials int64 + if err := fixture.db.Model(&models.Credential{}).Where("group_id = ?", first.GroupID).Count(&credentials).Error; err != nil { + t.Fatalf("count credentials: %v", err) + } + if credentials != 1 || len(fixture.registry.CaptureActiveCredentialRefs([]uint{first.GroupID})) != 1 { + t.Fatalf("credential state = db %d registry %#v", credentials, fixture.registry.Snapshot()) + } +} + +func TestChannelGroupSettingsExposeImmutableChannelAndEditableParams(t *testing.T) { + fixture := newServiceFixture(t) + created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("settings channel"), + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://first.example/v1"}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, + Credentials: "settings-key", + }) + if err != nil { + t.Fatalf("CreateGroup() error = %v", err) + } + settings, err := fixture.service.GetGroupSettings(t.Context(), created.GroupID) + if err != nil { + t.Fatalf("GetGroupSettings() error = %v", err) + } + if settings.ChannelID != channel.OpenAICompatible || string(settings.Params) != `{"base_url":"https://first.example/v1"}` { + t.Fatalf("settings = %#v", settings) + } + encoded, err := json.Marshal(settings) + if err != nil { + t.Fatalf("json.Marshal(settings) error = %v", err) + } + text := string(encoded) + for _, legacy := range []string{"provider_id", "upstream_url", "protocols"} { + if strings.Contains(text, `"`+legacy+`"`) { + t.Fatalf("settings leaked legacy field %q: %s", legacy, text) + } + } + + updated, err := fixture.service.UpdateGroupSettings(t.Context(), created.GroupID, GroupSettingsUpdateRequest{ + Params: optionalField[json.RawMessage]{ + Set: true, Value: json.RawMessage(`{"base_url":" HTTPS://Second.Example/v1/ "}`), + }, + }) + if err != nil { + t.Fatalf("UpdateGroupSettings(params) error = %v", err) + } + if updated.ChannelID != channel.OpenAICompatible || string(updated.Params) != `{"base_url":"https://second.example/v1"}` { + t.Fatalf("updated settings = %#v", updated) + } + + for _, body := range []string{ + `{"channel_id":"anthropic"}`, + `{"upstream_url":"https://legacy.example/v1"}`, + `{"protocols":["anthropic"]}`, + `{"provider_id":"anthropic"}`, + } { + if err := decodeStrictControlJSONObject([]byte(body), &GroupSettingsUpdateRequest{}); err == nil { + t.Fatalf("legacy or immutable settings body accepted: %s", body) + } + } +} + +func TestChannelGroupCollectionDetailAndOptionsUseChannelCredentialContract(t *testing.T) { + fixture := newServiceFixture(t) + created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("collection channel"), + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://collection.example/v1"}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{ + {ID: "provider-model", Alias: "public-model", AliasEnabled: true}, + }}, + Credentials: "collection-key", + }) + if err != nil { + t.Fatalf("CreateGroup() error = %v", err) + } + legacyCredentialQueries := 0 + var credentialSelects []string + const callbackName = "test:channel_collection_no_legacy_credentials" + if err := fixture.db.Callback().Query().Before("gorm:query").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "upstream_keys" { + legacyCredentialQueries++ + } + if tx.Statement.Table == "credentials" { + credentialSelects = append([]string(nil), tx.Statement.Selects...) + } + }); err != nil { + t.Fatalf("register query observer: %v", err) + } + t.Cleanup(func() { _ = fixture.db.Callback().Query().Remove(callbackName) }) + + collection, err := fixture.service.ListGroupCollection(t.Context(), GroupCollectionQuery{ + Page: 1, PageSize: 20, + }) + if err != nil || len(collection.Items) != 1 { + t.Fatalf("ListGroupCollection() = %#v, %v", collection, err) + } + if legacyCredentialQueries != 0 { + t.Fatalf("legacy upstream key queries = %d, want 0", legacyCredentialQueries) + } + if want := []string{"id", "group_id", "fingerprint", "status", "weight_manual", "updated_at_ms"}; !reflect.DeepEqual(credentialSelects, want) { + t.Fatalf("credential SELECT columns = %#v, want %#v", credentialSelects, want) + } + item := collection.Items[0] + if item.ChannelID != channel.OpenAICompatible || string(item.Params) != `{"base_url":"https://collection.example/v1"}` || + item.CredentialCounts.Total != 1 || item.CredentialCounts.Available != 1 { + t.Fatalf("collection item = %#v", item) + } + assertNoLegacyGroupFields(t, item) + + summary, err := fixture.service.GetGroupSummary(t.Context(), created.GroupID) + if err != nil || summary.ChannelID != channel.OpenAICompatible || + string(summary.Params) != `{"base_url":"https://collection.example/v1"}` || summary.CredentialCount != 1 { + t.Fatalf("GetGroupSummary() = %#v, %v", summary, err) + } + assertNoLegacyGroupFields(t, summary) + + options, err := fixture.service.ListGroupOptions(t.Context()) + if err != nil || len(options) != 1 || options[0].ChannelID != channel.OpenAICompatible || + string(options[0].Params) != `{"base_url":"https://collection.example/v1"}` || + !reflect.DeepEqual(options[0].Models, []string{"public-model"}) { + t.Fatalf("ListGroupOptions() = %#v, %v", options, err) + } + assertNoLegacyGroupFields(t, options[0]) + + if err := fixture.service.DeleteGroup(t.Context(), created.GroupID); err != nil { + t.Fatalf("DeleteGroup() error = %v", err) + } + var credentials int64 + if err := fixture.db.Model(&models.Credential{}).Where("group_id = ?", created.GroupID).Count(&credentials).Error; err != nil { + t.Fatalf("count credentials after delete: %v", err) + } + if credentials != 0 || len(fixture.registry.CaptureActiveCredentialRefs([]uint{created.GroupID})) != 0 { + t.Fatalf("deleted credential state = db %d registry %#v", credentials, fixture.registry.Snapshot()) + } + if _, exists := fixture.manager.Current().Groups[created.GroupID]; exists { + t.Fatalf("deleted group %d remains in snapshot", created.GroupID) + } +} + +func assertNoLegacyGroupFields(t *testing.T, value any) { + t.Helper() + encoded, err := json.Marshal(value) + if err != nil { + t.Fatalf("json.Marshal() error = %v", err) + } + for _, legacy := range []string{"provider_id", "upstream_url", "protocols", "key_count", "key_counts"} { + if strings.Contains(string(encoded), `"`+legacy+`"`) { + t.Fatalf("response leaked legacy field %q: %s", legacy, encoded) + } + } +} diff --git a/internal/control/channels.go b/internal/control/channels.go new file mode 100644 index 00000000..e0e8e8c6 --- /dev/null +++ b/internal/control/channels.go @@ -0,0 +1,69 @@ +package control + +import ( + "context" + "net/url" + "strings" + "unicode/utf8" + + "github.com/gin-gonic/gin" + + "gpt-load/internal/channel" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/platform/response" +) + +const channelMaxQueryRunes = 200 + +type ChannelListResponse struct { + Items []channel.Descriptor `json:"items"` + Total int `json:"total"` +} + +func (s *Service) ListChannels(ctx context.Context, query string) (ChannelListResponse, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return ChannelListResponse{}, err + } + if s == nil || s.channelRegistry == nil { + return ChannelListResponse{}, app_errors.ErrInternalServer + } + items := s.channelRegistry.Search(query) + return ChannelListResponse{Items: items, Total: len(items)}, nil +} + +func parseChannelQuery(rawQuery string, forceQuery bool) (string, *app_errors.APIError) { + if forceQuery && rawQuery == "" { + return "", app_errors.ErrBadRequest + } + values, err := url.ParseQuery(rawQuery) + if err != nil { + return "", app_errors.ErrBadRequest + } + for key, entries := range values { + if key != "q" || len(entries) != 1 { + return "", app_errors.ErrBadRequest + } + } + query := strings.TrimSpace(values.Get("q")) + if utf8.RuneCountInString(query) > channelMaxQueryRunes { + return "", app_errors.ErrBadRequest + } + return query, nil +} + +func (s *Server) handleListChannels(c *gin.Context) { + query, apiErr := parseChannelQuery(c.Request.URL.RawQuery, c.Request.URL.ForceQuery) + if apiErr != nil { + writeServiceError(c, "list_channels", apiErr) + return + } + result, err := s.service.ListChannels(c.Request.Context(), query) + if err != nil { + writeServiceError(c, "list_channels", err) + return + } + response.SuccessI18n(c, "common.success", result) +} diff --git a/internal/control/credential_api_foundation_test.go b/internal/control/credential_api_foundation_test.go new file mode 100644 index 00000000..edeee95d --- /dev/null +++ b/internal/control/credential_api_foundation_test.go @@ -0,0 +1,403 @@ +package control + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + + "gpt-load/internal/channel" + "gpt-load/internal/platform/config" + "gpt-load/internal/state" + "gpt-load/internal/storage/models" +) + +func TestCredentialRoutesReplaceLegacyGroupKeyRoutes(t *testing.T) { + fixture := newServiceFixture(t) + module := NewServer(&config.Config{AuthKey: "credential-auth"}, fixture.service).HTTPModule() + want := map[string]string{ + "control.group-credentials.list": "/groups/:group_id/credentials", + "control.group-credentials.reveal": "/groups/:group_id/credentials/:credential_id/reveal", + "control.group-credentials.update": "/groups/:group_id/credentials/:credential_id", + "control.group-credentials.restore": "/groups/:group_id/credentials/:credential_id/restore", + "control.group-credentials.batch": "/groups/:group_id/credentials/batch", + "control.group-credentials.delete": "/groups/:group_id/credentials/:credential_id", + "control.group-credentials.import": "/groups/:group_id/credentials/import", + } + seen := make(map[string]string) + for _, route := range module.Routes { + if strings.Contains(route.Path, "/keys") || strings.HasPrefix(route.Name, "control.group-keys") { + t.Fatalf("legacy key route remains: %s %s", route.Name, route.Path) + } + if path, ok := want[route.Name]; ok { + if len(route.Methods) != 1 || route.Path != path { + t.Fatalf("route %s = %#v", route.Name, route) + } + seen[route.Name] = route.Methods[0] + } + } + if len(seen) != len(want) || seen["control.group-credentials.list"] != http.MethodGet { + t.Fatalf("credential routes = %#v, want %#v", seen, want) + } +} + +func TestImportAndListGroupCredentialsUseCanonicalCredentialStorage(t *testing.T) { + fixture := newServiceFixture(t) + created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("credential api"), ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: optionalGroupModels{Set: true}, Credentials: "first-secret", + }) + if err != nil { + t.Fatalf("CreateGroup() error = %v", err) + } + result, err := fixture.service.ImportGroupCredentials(t.Context(), created.GroupID, CredentialImportRequest{ + Credentials: " second-secret \nfirst-secret\nsecond-secret\n", + }) + if err != nil { + t.Fatalf("ImportGroupCredentials() error = %v", err) + } + if result.CredentialsAdded != 1 || result.CredentialsDuplicated != 2 { + t.Fatalf("import result = %#v", result) + } + var stored []models.Credential + if err := fixture.db.Where("group_id = ?", created.GroupID).Order("id ASC").Find(&stored).Error; err != nil { + t.Fatalf("load credentials: %v", err) + } + if len(stored) != 2 { + t.Fatalf("stored credentials = %#v", stored) + } + response, err := fixture.service.ListGroupCredentials(t.Context(), created.GroupID, CredentialCollectionQuery{Page: 1, PageSize: 20}) + if err != nil || response.Summary.Total != 2 || len(response.Items) != 2 { + t.Fatalf("ListGroupCredentials() = %#v, %v", response, err) + } + encoded, err := json.Marshal(response) + if err != nil { + t.Fatal(err) + } + text := string(encoded) + if strings.Contains(text, "first-secret") || strings.Contains(text, "second-secret") || + strings.Contains(text, `"id"`) || strings.Contains(text, "key_id") { + t.Fatalf("credential collection leaked legacy or secret data: %s", encoded) + } +} + +func TestCloudCredentialImportAcceptsOneStrictJSONObjectPerLine(t *testing.T) { + fixture := newServiceFixture(t) + tests := []struct { + name string + channelID channel.ID + params string + credentials string + want string + }{ + { + name: "Azure Entra", channelID: channel.AzureOpenAI, + params: `{"endpoint":"https://resource.openai.azure.com"}`, + credentials: `{"tenant_id":" tenant ","client_secret":" secret ","client_id":" client "}`, + want: `{"client_id":"client","client_secret":"secret","tenant_id":"tenant"}`, + }, + { + name: "Bedrock SigV4", channelID: channel.AWSBedrock, + params: `{"region":"us-east-1"}`, + credentials: `{"secret_key":" secret ","access_key":" access "}`, + want: `{"access_key":"access","secret_key":"secret"}`, + }, + { + name: "Vertex service account", channelID: channel.GoogleVertex, + params: `{"location":"us-central1","project_id":"project-one"}`, + credentials: `{"service_account_json":"{\"type\":\"service_account\",\"project_id\":\"project-one\",\"client_email\":\"svc@example.iam.gserviceaccount.com\",\"private_key\":\"private-secret\"}"}`, + want: `{"service_account_json":"{\"client_email\":\"svc@example.iam.gserviceaccount.com\",\"private_key\":\"private-secret\",\"project_id\":\"project-one\",\"type\":\"service_account\"}"}`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("cloud " + test.name), ChannelID: test.channelID, + Params: json.RawMessage(test.params), Models: optionalGroupModels{Set: true}, + Credentials: test.credentials + "\n" + test.credentials, + }) + if err != nil { + t.Fatalf("CreateGroup() error = %v", err) + } + if created.CredentialsAdded != 1 || created.CredentialsDuplicated != 1 { + t.Fatalf("create result = %#v", created) + } + var row models.Credential + if err := fixture.db.Where("group_id = ?", created.GroupID).Take(&row).Error; err != nil { + t.Fatal(err) + } + plaintext, err := fixture.encryption.Decrypt(row.Data) + if err != nil { + t.Fatal(err) + } + if plaintext != test.want { + t.Fatalf("stored credential = %s, want %s", plaintext, test.want) + } + }) + } +} + +func TestVertexCredentialImportAcceptsPastedServiceAccountJSON(t *testing.T) { + fixture := newServiceFixture(t) + raw := `{ + "type": "service_account", + "project_id": "project-one", + "client_email": "svc@example.iam.gserviceaccount.com", + "private_key": "private-secret" +}` + created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("vertex pasted credential"), ChannelID: channel.GoogleVertex, + Params: json.RawMessage(`{"location":"us-central1","project_id":"project-one"}`), + Models: optionalGroupModels{Set: true}, Credentials: raw, + }) + if err != nil { + t.Fatalf("CreateGroup() error = %v", err) + } + if created.CredentialsAdded != 1 || created.CredentialsDuplicated != 0 { + t.Fatalf("create result = %#v", created) + } + var row models.Credential + if err := fixture.db.Where("group_id = ?", created.GroupID).Take(&row).Error; err != nil { + t.Fatal(err) + } + plaintext, err := fixture.encryption.Decrypt(row.Data) + if err != nil { + t.Fatal(err) + } + want := `{"service_account_json":"{\"client_email\":\"svc@example.iam.gserviceaccount.com\",\"private_key\":\"private-secret\",\"project_id\":\"project-one\",\"type\":\"service_account\"}"}` + if plaintext != want { + t.Fatalf("stored credential = %s, want %s", plaintext, want) + } +} + +func TestGroupCredentialMutationsPreserveRuntimeIdentityAndHealthContracts(t *testing.T) { + fixture := newServiceFixture(t) + created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("credential mutations"), ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: optionalGroupModels{Set: true}, + Credentials: "first-secret\nsecond-secret", + }) + if err != nil { + t.Fatalf("CreateGroup() error = %v", err) + } + var rows []models.Credential + if err := fixture.db.Where("group_id = ?", created.GroupID).Order("id ASC").Find(&rows).Error; err != nil { + t.Fatal(err) + } + if len(rows) != 2 { + t.Fatalf("credentials = %#v", rows) + } + before, ok := findRuntimeCredential(fixture.registry.Snapshot(), rows[0].ID) + if !ok { + t.Fatal("first credential missing from Registry") + } + beforeEntries, err := fixture.registry.SnapshotGroupCredentialEntriesExact(created.GroupID, []uint{rows[0].ID}) + if err != nil { + t.Fatal(err) + } + fixture.service.now = func() time.Time { return time.UnixMilli(rows[0].UpdatedAtMS) } + weight := 17 + updated, err := fixture.service.UpdateGroupCredential(t.Context(), created.GroupID, rows[0].ID, CredentialUpdateRequest{ + Status: optionalField[state.CredentialStatus]{Set: true, Value: state.CredentialStatusDisabled}, + WeightManual: optionalField[int]{Set: true, Value: weight}, + }) + if err != nil { + t.Fatalf("UpdateGroupCredential() error = %v", err) + } + if updated.CredentialID != rows[0].ID || updated.ConfiguredStatus != "disabled" || + updated.WeightMode != "manual" { + t.Fatalf("updated credential = %#v", updated) + } + var committed models.Credential + if err := fixture.db.Take(&committed, rows[0].ID).Error; err != nil { + t.Fatal(err) + } + after, ok := findRuntimeCredential(fixture.registry.Snapshot(), rows[0].ID) + afterEntries, entryErr := fixture.registry.SnapshotGroupCredentialEntriesExact(created.GroupID, []uint{rows[0].ID}) + if !ok || committed.UpdatedAtMS <= rows[0].UpdatedAtMS || + committed.WeightManual == nil || *committed.WeightManual != weight || + after.WeightManual == nil || *after.WeightManual != weight || + after.Version != uint64(committed.UpdatedAtMS) || + after.IdentityGeneration != before.IdentityGeneration || entryErr != nil || + afterEntries[0].Fingerprint != beforeEntries[0].Fingerprint { + t.Fatalf("identity/version after update: row=%#v before=%#v after=%#v", committed, before, after) + } + + batch, err := fixture.service.BatchGroupCredentials(t.Context(), created.GroupID, CredentialBatchRequest{ + Action: CredentialBatchEnable, CredentialIDs: []uint{rows[1].ID, rows[0].ID}, + }) + if err != nil { + t.Fatalf("BatchGroupCredentials() error = %v", err) + } + if len(batch.AffectedCredentialIDs) != 2 || batch.AffectedCredentialIDs[0] != rows[0].ID || + batch.AffectedCredentialIDs[1] != rows[1].ID || batch.Summary.Available != 2 { + t.Fatalf("batch response = %#v", batch) + } + + if !fixture.registry.SetBlacklisted(rows[0].ID) { + t.Fatal("blacklist first credential") + } + restored, err := fixture.service.RestoreGroupCredential(t.Context(), created.GroupID, rows[0].ID) + if err != nil { + t.Fatalf("RestoreGroupCredential() error = %v", err) + } + if restored.CredentialID != rows[0].ID || restored.EffectiveStatus != "available" { + t.Fatalf("restored credential = %#v", restored) + } + + revealed, err := fixture.service.RevealGroupCredential(t.Context(), created.GroupID, rows[0].ID) + if err != nil { + t.Fatalf("RevealGroupCredential() error = %v", err) + } + if revealed.CredentialID != rows[0].ID || string(revealed.Credential) != `{"api_key":"first-secret"}` { + t.Fatalf("revealed credential = %#v", revealed) + } + + fixture.stats.RecordSuccess(rows[1].ID, fixture.service.now()) + if err := fixture.service.DeleteGroupCredential(t.Context(), created.GroupID, rows[1].ID); err != nil { + t.Fatalf("DeleteGroupCredential() error = %v", err) + } + if _, exists := findRuntimeCredential(fixture.registry.Snapshot(), rows[1].ID); exists { + t.Fatal("deleted credential remains in Registry") + } + var count int64 + if err := fixture.db.Model(&models.Credential{}).Where("id = ?", rows[1].ID).Count(&count).Error; err != nil || count != 0 { + t.Fatalf("deleted credential count = %d, err=%v", count, err) + } + if got := fixture.stats.Snapshot(rows[1].ID, fixture.service.now()); got.Success != 0 { + t.Fatalf("deleted credential stats = %#v", got) + } +} + +func TestCredentialHTTPUsesCanonicalWireAndRejectsLegacyFields(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("credential http"), ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: optionalGroupModels{Set: true}, Credentials: "first-secret", + }) + if err != nil { + t.Fatal(err) + } + var row models.Credential + if err := fixture.db.Where("group_id = ?", created.GroupID).Take(&row).Error; err != nil { + t.Fatal(err) + } + engine := gin.New() + const auth = "credential-http-auth" + NewServer(&config.Config{AuthKey: auth}, fixture.service).RegisterRoutes(engine) + + legacy := serveCredentialRequest(t, engine, http.MethodGet, + fmt.Sprintf("/api/groups/%d/keys", created.GroupID), "", auth, "") + if legacy.Code != http.StatusNotFound { + t.Fatalf("legacy key route = %d %s", legacy.Code, legacy.Body.String()) + } + + const idempotencyKey = "00000000-0000-4000-8000-0000000000c1" + imported := serveCredentialRequest(t, engine, http.MethodPost, + fmt.Sprintf("/api/groups/%d/credentials/import", created.GroupID), + `{"credentials":"second-secret"}`, auth, idempotencyKey) + if imported.Code != http.StatusOK || !strings.Contains(imported.Body.String(), `"credentials_added":1`) || + strings.Contains(imported.Body.String(), "keys_added") { + t.Fatalf("credential import = %d %s", imported.Code, imported.Body.String()) + } + replayed := serveCredentialRequest(t, engine, http.MethodPost, + fmt.Sprintf("/api/groups/%d/credentials/import", created.GroupID), + `{"credentials":"second-secret"}`, auth, idempotencyKey) + if replayed.Code != http.StatusOK || replayed.Body.String() != imported.Body.String() { + t.Fatalf("credential import replay = %d %s, want %s", replayed.Code, replayed.Body.String(), imported.Body.String()) + } + legacyImportField := serveCredentialRequest(t, engine, http.MethodPost, + fmt.Sprintf("/api/groups/%d/credentials/import", created.GroupID), + `{"keys":"third-secret"}`, auth, "00000000-0000-4000-8000-0000000000c2") + if legacyImportField.Code != http.StatusBadRequest { + t.Fatalf("legacy import field = %d %s", legacyImportField.Code, legacyImportField.Body.String()) + } + + list := serveCredentialRequest(t, engine, http.MethodGet, + fmt.Sprintf("/api/groups/%d/credentials", created.GroupID), "", auth, "") + if list.Code != http.StatusOK || !strings.Contains(list.Body.String(), `"credential_id":`) || + strings.Contains(list.Body.String(), "key_id") || strings.Contains(list.Body.String(), "first-secret") || + strings.Contains(list.Body.String(), "second-secret") { + t.Fatalf("credential list = %d %s", list.Code, list.Body.String()) + } + + reveal := serveCredentialRequest(t, engine, http.MethodPost, + fmt.Sprintf("/api/groups/%d/credentials/%d/reveal", created.GroupID, row.ID), "{}", auth, "") + if reveal.Code != http.StatusOK || reveal.Header().Get("Cache-Control") != "no-store" || + reveal.Header().Get("Pragma") != "no-cache" || + !strings.Contains(reveal.Body.String(), `"credential":{"api_key":"first-secret"}`) || + strings.Contains(reveal.Body.String(), `"key":`) { + t.Fatalf("credential reveal = %d headers=%v body=%s", reveal.Code, reveal.Header(), reveal.Body.String()) + } + legacyRevealField := serveCredentialRequest(t, engine, http.MethodPost, + fmt.Sprintf("/api/groups/%d/credentials/%d/reveal", created.GroupID, row.ID), + `{"key_id":1}`, auth, "") + if legacyRevealField.Code != http.StatusBadRequest { + t.Fatalf("legacy reveal field = %d %s", legacyRevealField.Code, legacyRevealField.Body.String()) + } + legacyBatchField := serveCredentialRequest(t, engine, http.MethodPost, + fmt.Sprintf("/api/groups/%d/credentials/batch", created.GroupID), + fmt.Sprintf(`{"action":"disable","key_ids":[%d]}`, row.ID), auth, "") + if legacyBatchField.Code != http.StatusBadRequest { + t.Fatalf("legacy batch field = %d %s", legacyBatchField.Code, legacyBatchField.Body.String()) + } +} + +func TestNormalizeStoredCredentialAcceptsLegacyAPIKeyAndCanonicalObjects(t *testing.T) { + t.Parallel() + + registry := channel.NewRegistry() + legacy, err := normalizeStoredCredential(registry, channel.OpenAI, " sk-legacy ") + if err != nil { + t.Fatalf("normalizeStoredCredential(legacy) error = %v", err) + } + if got := string(legacy.CanonicalJSON()); got != `{"api_key":"sk-legacy"}` { + t.Fatalf("legacy canonical credential = %s", got) + } + + typed, err := normalizeStoredCredential( + registry, + channel.AWSBedrock, + `{"access_key":"AKIAEXAMPLE","secret_key":"secret"}`, + ) + if err != nil { + t.Fatalf("normalizeStoredCredential(typed) error = %v", err) + } + if got := string(typed.CanonicalJSON()); got != `{"access_key":"AKIAEXAMPLE","secret_key":"secret"}` { + t.Fatalf("typed canonical credential = %s", got) + } + + if _, err := normalizeStoredCredential(registry, channel.OpenAI, `{"api_key":`); err == nil { + t.Fatal("normalizeStoredCredential(invalid object) error = nil") + } +} + +func serveCredentialRequest( + t *testing.T, + engine *gin.Engine, + method string, + path string, + body string, + auth string, + idempotencyKey string, +) *httptest.ResponseRecorder { + t.Helper() + request := httptest.NewRequest(method, path, bytes.NewBufferString(body)) + request.Header.Set("Authorization", "Bearer "+auth) + if body != "" { + request.Header.Set("Content-Type", "application/json") + } + if idempotencyKey != "" { + request.Header.Set("Idempotency-Key", idempotencyKey) + } + response := httptest.NewRecorder() + engine.ServeHTTP(response, request) + return response +} diff --git a/internal/control/credential_collection_helpers.go b/internal/control/credential_collection_helpers.go new file mode 100644 index 00000000..b6b7a7d9 --- /dev/null +++ b/internal/control/credential_collection_helpers.go @@ -0,0 +1,232 @@ +package control + +import ( + "encoding/json" + "fmt" + "net/url" + "strconv" + "strings" + "time" + + "gpt-load/internal/channel" + "gpt-load/internal/health" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/platform/utils" + "gpt-load/internal/state" +) + +// normalizeStoredCredential accepts both the canonical object format and the +// legacy encrypted API-key string used before channel credentials existed. +// Invalid object-shaped data fails closed instead of being reinterpreted as a +// raw key. +func normalizeStoredCredential( + registry *channel.Registry, + channelID channel.ID, + plaintext string, +) (channel.Credential, error) { + if registry == nil || channelID == "" { + return channel.Credential{}, fmt.Errorf("credential registry is unavailable") + } + trimmed := strings.TrimSpace(plaintext) + if trimmed == "" { + return channel.Credential{}, fmt.Errorf("credential is empty") + } + if strings.HasPrefix(trimmed, "{") { + return registry.ValidateCredential(channelID, json.RawMessage(trimmed)) + } + encoded, err := json.Marshal(map[string]string{"api_key": trimmed}) + if err != nil { + return channel.Credential{}, fmt.Errorf("encode legacy credential: %w", err) + } + return registry.ValidateCredential(channelID, encoded) +} + +const ( + credentialCollectionDefaultPage = 1 + credentialCollectionDefaultPageSize = 50 + credentialCollectionStatsWindow = int64(health.StatsWindow / time.Second) +) + +func parseCredentialCollectionQuery(rawQuery string) (CredentialCollectionQuery, *app_errors.APIError) { + query := CredentialCollectionQuery{ + Page: credentialCollectionDefaultPage, PageSize: credentialCollectionDefaultPageSize, + } + values, err := url.ParseQuery(rawQuery) + if err != nil { + return CredentialCollectionQuery{}, app_errors.ErrBadRequest + } + for key, entries := range values { + switch key { + case "q", "status", "page", "page_size": + default: + return CredentialCollectionQuery{}, app_errors.ErrBadRequest + } + if len(entries) != 1 { + return CredentialCollectionQuery{}, app_errors.ErrBadRequest + } + } + if entries, exists := values["q"]; exists { + query.Query = strings.TrimSpace(entries[0]) + } + if entries, exists := values["status"]; exists { + status := entries[0] + switch status { + case string(healthBucketAvailable), string(healthBucketCooldown), + string(healthBucketBlacklisted), string(healthBucketDisabled): + query.Status = &status + default: + return CredentialCollectionQuery{}, app_errors.ErrBadRequest + } + } + if entries, exists := values["page"]; exists { + page, ok := parseCredentialCollectionPositiveInt(entries[0]) + if !ok { + return CredentialCollectionQuery{}, app_errors.ErrBadRequest + } + query.Page = page + } + if entries, exists := values["page_size"]; exists { + pageSize, ok := parseCredentialCollectionPositiveInt(entries[0]) + if !ok || pageSize != 20 && pageSize != 50 && pageSize != 100 { + return CredentialCollectionQuery{}, app_errors.ErrBadRequest + } + query.PageSize = pageSize + } + return query, nil +} + +func parseCredentialCollectionPositiveInt(value string) (int, bool) { + if value == "" { + return 0, false + } + for index := range len(value) { + if value[index] < '0' || value[index] > '9' { + return 0, false + } + } + parsed, err := strconv.Atoi(value) + return parsed, err == nil && parsed > 0 +} + +func maskCredential(plaintext string) (string, error) { + mask := utils.MaskAPIKey(plaintext) + if mask == "" { + return "", fmt.Errorf("credential is empty: %w", app_errors.ErrInternalServer) + } + return mask, nil +} + +func maskCanonicalCredential(canonical json.RawMessage) (string, error) { + var values map[string]string + if err := json.Unmarshal(canonical, &values); err != nil || values == nil { + return "", fmt.Errorf("credential is invalid: %w", app_errors.ErrInternalServer) + } + for _, field := range []string{"api_key", "access_key", "client_id", "role_arn"} { + if value := strings.TrimSpace(values[field]); value != "" { + return maskCredential(value) + } + } + if raw := strings.TrimSpace(values["service_account_json"]); raw != "" { + var serviceAccount struct { + ClientEmail string `json:"client_email"` + } + if json.Unmarshal([]byte(raw), &serviceAccount) == nil && strings.TrimSpace(serviceAccount.ClientEmail) != "" { + return maskCredential(serviceAccount.ClientEmail) + } + } + return "", fmt.Errorf("credential identity is unavailable: %w", app_errors.ErrInternalServer) +} + +func mapCredentialRuntimeItem( + mask string, + credentialID uint, + view state.CredentialRuntimeView, + bucket healthBucket, + stats health.CredentialStats, + observedAt time.Time, +) (CredentialItemResponse, error) { + weightMode := "auto" + if view.WeightManual != nil { + weightMode = "manual" + } + item := CredentialItemResponse{ + CredentialID: credentialID, + Mask: mask, + ConfiguredStatus: string(view.Status), + EffectiveStatus: string(bucket), + WeightMode: weightMode, + RecentSuccessCount: stats.Success, + RecentFailureCount: stats.Failure, + ConsecutiveFailureCount: stats.ConsecutiveFailure, + LastFailureCategory: normalizeCredentialFailureCategory(stats.LastFailureCategory).String(), + LastStatusCode: optionalHealthStatusCode(stats.LastStatusCode), + } + switch bucket { + case healthBucketAvailable: + weight := view.WeightAuto + if view.WeightManual != nil { + weight = *view.WeightManual + } + if weight < 1 || weight > state.MaxWeight { + return CredentialItemResponse{}, fmt.Errorf( + "map credential %d weight: %w", credentialID, app_errors.ErrInternalServer, + ) + } + item.Weight = &weight + item.Recovery = CredentialRecoveryResponse{Mode: "none"} + case healthBucketCooldown: + cooldownUntilMS, err := optionalSafeEpochMilliseconds(view.CooldownUntil) + if err != nil { + return CredentialItemResponse{}, fmt.Errorf( + "map credential %d cooldown_until_ms: %w", credentialID, err, + ) + } + if cooldownUntilMS == nil || !view.CooldownUntil.After(observedAt) { + return CredentialItemResponse{}, fmt.Errorf( + "map credential %d invalid cooldown: %w", credentialID, app_errors.ErrInternalServer, + ) + } + item.CooldownUntilMS = cooldownUntilMS + item.Recovery = CredentialRecoveryResponse{ + Mode: "cooldown", Automatic: true, AtMS: cooldownUntilMS, + } + case healthBucketBlacklisted: + item.Recovery = CredentialRecoveryResponse{Mode: "probe", Automatic: true} + case healthBucketDisabled: + item.Recovery = CredentialRecoveryResponse{Mode: "manual"} + default: + return CredentialItemResponse{}, fmt.Errorf( + "map credential %d unknown status: %w", credentialID, app_errors.ErrInternalServer, + ) + } + return item, nil +} + +func normalizeCredentialFailureCategory(category health.FailureCategory) health.FailureCategory { + if category == health.FailureCategoryOK { + return health.FailureCategoryAmbiguous + } + return category +} + +func credentialCollectionBucketOrder(bucket healthBucket) int { + switch bucket { + case healthBucketBlacklisted: + return 0 + case healthBucketCooldown: + return 1 + case healthBucketAvailable: + return 2 + case healthBucketDisabled: + return 3 + default: + return 4 + } +} + +func credentialCollectionTotalPages(totalItems, pageSize int) int { + if totalItems == 0 { + return 0 + } + return (totalItems + pageSize - 1) / pageSize +} diff --git a/internal/control/credential_mutations.go b/internal/control/credential_mutations.go new file mode 100644 index 00000000..8c4bd209 --- /dev/null +++ b/internal/control/credential_mutations.go @@ -0,0 +1,600 @@ +package control + +import ( + "context" + "errors" + "fmt" + "math" + "sort" + "time" + + "gorm.io/gorm" + + "gpt-load/internal/health" + "gpt-load/internal/platform/epochms" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/state" + "gpt-load/internal/storage/models" +) + +func normalizeCredentialUpdate( + request CredentialUpdateRequest, +) (status *state.CredentialStatus, weight *int, weightSet bool, err error) { + if !request.Status.Set && !request.WeightManual.Set { + return nil, nil, false, app_errors.ErrBadRequest + } + if request.Status.Set { + if request.Status.Null || + (request.Status.Value != state.CredentialStatusActive && request.Status.Value != state.CredentialStatusDisabled) { + return nil, nil, false, app_errors.ErrValidation + } + value := request.Status.Value + status = &value + } + if request.WeightManual.Set { + weightSet = true + if !request.WeightManual.Null { + if request.WeightManual.Value < 1 || request.WeightManual.Value > state.MaxWeight { + return nil, nil, false, app_errors.ErrValidation + } + value := request.WeightManual.Value + weight = &value + } + } + return status, weight, weightSet, nil +} + +func nextCredentialUpdatedAtMS(now time.Time, previous int64) (int64, error) { + nowMS, err := epochms.FromTime(now) + if err != nil { + return 0, err + } + if nowMS < 1 { + nowMS = 1 + } + if nowMS <= previous { + if previous == math.MaxInt64 { + return 0, fmt.Errorf("credential version exhausted") + } + nowMS = previous + 1 + } + return nowMS, nil +} + +func equalOptionalWeight(left, right *int) bool { + if left == nil || right == nil { + return left == nil && right == nil + } + return *left == *right +} + +func findRuntimeCredential( + views []state.CredentialRuntimeView, + credentialID uint, +) (state.CredentialRuntimeView, bool) { + for _, view := range views { + if view.ID == credentialID { + return view, true + } + } + return state.CredentialRuntimeView{}, false +} + +func (s *Service) RevealGroupCredential( + ctx context.Context, + groupID uint, + credentialID uint, +) (CredentialRevealResult, error) { + if groupID == 0 || credentialID == 0 { + return CredentialRevealResult{}, app_errors.ErrBadRequest + } + s.writeMu.RLock() + defer s.writeMu.RUnlock() + group, err := loadGroupRow(s.db.WithContext(ctx), groupID) + if err != nil { + return CredentialRevealResult{}, err + } + if group.ChannelID == "" { + return CredentialRevealResult{}, app_errors.ErrValidation + } + var row models.Credential + if err := s.db.WithContext(ctx).Select("id", "group_id", "data", "fingerprint", "status", "weight_manual", "updated_at_ms"). + Where("id = ? AND group_id = ?", credentialID, groupID).Take(&row).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return CredentialRevealResult{}, credentialNotFoundError() + } + return CredentialRevealResult{}, app_errors.ParseDBError(err) + } + credential, _, err := s.decodeCredential(group, row) + if err != nil { + return CredentialRevealResult{}, err + } + revealedAtMS, err := safeEpochMilliseconds(s.now()) + if err != nil { + return CredentialRevealResult{}, app_errors.ErrInternalServer + } + return CredentialRevealResult{ + CredentialID: row.ID, Credential: append([]byte(nil), credential...), RevealedAtMS: revealedAtMS, + }, nil +} + +func (s *Service) UpdateGroupCredential( + ctx context.Context, + groupID uint, + credentialID uint, + request CredentialUpdateRequest, +) (CredentialItemResponse, error) { + if groupID == 0 || credentialID == 0 { + return CredentialItemResponse{}, app_errors.ErrBadRequest + } + status, weight, weightSet, err := normalizeCredentialUpdate(request) + if err != nil { + return CredentialItemResponse{}, err + } + var committed models.Credential + err = s.writeCredentialConfig(ctx, groupID, credentialID, func(tx *gorm.DB) error { + group, err := loadGroupRow(tx, groupID) + if err != nil { + return err + } + if group.ChannelID == "" { + return app_errors.ErrValidation + } + if err := tx.Where("id = ? AND group_id = ?", credentialID, groupID).Take(&committed).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return credentialNotFoundError() + } + return app_errors.ParseDBError(err) + } + view, exists := findRuntimeCredential(s.registry.Snapshot(), credentialID) + if err := validateCredentialRuntimeRow(groupID, committed, view, exists); err != nil { + return err + } + updatedAtMS, err := nextCredentialUpdatedAtMS(s.now(), committed.UpdatedAtMS) + if err != nil { + return app_errors.ErrInternalServer + } + updates := map[string]any{"updated_at_ms": updatedAtMS} + if status != nil { + committed.Status = models.CredentialStatus(*status) + updates["status"] = committed.Status + } + if weightSet { + committed.WeightManual = cloneInt(weight) + updates["weight_manual"] = committed.WeightManual + } + committed.UpdatedAtMS = updatedAtMS + if err := tx.Model(&models.Credential{}).Where("id = ? AND group_id = ?", credentialID, groupID). + Updates(updates).Error; err != nil { + return app_errors.ParseDBError(err) + } + return nil + }, func() error { + var applyErr error + apply := func() { + entries, snapshotErr := s.registry.SnapshotGroupCredentialEntriesExact(groupID, []uint{credentialID}) + if snapshotErr != nil { + applyErr = dbRegistryMismatch(mismatchMissingRegistry, groupID, credentialID) + return + } + entry := entries[0] + entry.Status = state.CredentialStatus(committed.Status) + entry.WeightManual = cloneInt(committed.WeightManual) + entry.Version = groupCollectionCredentialVersion(committed.UpdatedAtMS) + entry.IdentityGeneration = groupCollectionCredentialIdentity(committed.Fingerprint) + entry.Fingerprint = committed.Fingerprint + entry.EncryptedValue = committed.Data + applyErr = s.registry.RestoreGroupCredentialEntriesExact(groupID, []state.CredentialEntry{entry}) + } + if s.mutations == nil { + apply() + } else { + s.mutations.Do(credentialID, apply) + } + return applyErr + }) + if err != nil { + return CredentialItemResponse{}, err + } + return s.loadCredentialItem(ctx, groupID, credentialID) +} + +func (s *Service) DeleteGroupCredential(ctx context.Context, groupID, credentialID uint) error { + if groupID == 0 || credentialID == 0 { + return app_errors.ErrBadRequest + } + return s.writeCredentialConfig(ctx, groupID, credentialID, func(tx *gorm.DB) error { + group, err := loadGroupRow(tx, groupID) + if err != nil { + return err + } + if group.ChannelID == "" { + return app_errors.ErrValidation + } + var row models.Credential + if err := tx.Where("id = ? AND group_id = ?", credentialID, groupID).Take(&row).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return credentialNotFoundError() + } + return app_errors.ParseDBError(err) + } + view, exists := findRuntimeCredential(s.registry.Snapshot(), credentialID) + if err := validateCredentialRuntimeRow(groupID, row, view, exists); err != nil { + return err + } + if err := tx.Delete(&row).Error; err != nil { + return app_errors.ParseDBError(err) + } + return nil + }, func() error { + var applyErr error + apply := func() { + if !s.registry.RemoveCredential(credentialID) { + applyErr = dbRegistryMismatch(mismatchMissingRegistry, groupID, credentialID) + return + } + s.stats.Reset(credentialID) + } + if s.mutations == nil { + apply() + } else { + s.mutations.Do(credentialID, apply) + } + return applyErr + }) +} + +func (s *Service) RestoreGroupCredential( + ctx context.Context, + groupID uint, + credentialID uint, +) (CredentialItemResponse, error) { + if groupID == 0 || credentialID == 0 { + return CredentialItemResponse{}, app_errors.ErrBadRequest + } + s.writeMu.Lock() + defer s.writeMu.Unlock() + group, err := loadGroupRow(s.db.WithContext(ctx), groupID) + if err != nil { + return CredentialItemResponse{}, err + } + if group.ChannelID == "" { + return CredentialItemResponse{}, app_errors.ErrValidation + } + var row models.Credential + if err := s.db.WithContext(ctx).Where("id = ? AND group_id = ?", credentialID, groupID).Take(&row).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return CredentialItemResponse{}, credentialNotFoundError() + } + return CredentialItemResponse{}, app_errors.ParseDBError(err) + } + view, exists := findRuntimeCredential(s.registry.Snapshot(), credentialID) + if err := validateCredentialRuntimeRow(groupID, row, view, exists); err != nil { + return CredentialItemResponse{}, err + } + groupView := state.GroupCatalogView{ID: group.ID, Name: group.Name, Enabled: group.Enabled, + WeightManual: cloneInt(group.WeightManual)} + var observedAt time.Time + var restoreErr error + restore := func() { + observedAt = s.now().UTC() + current, exists := findRuntimeCredential(s.registry.Snapshot(), credentialID) + if !exists { + restoreErr = dbRegistryMismatch(mismatchMissingRegistry, groupID, credentialID) + return + } + bucket := classifyHealthKey(groupView, current, observedAt) + if bucket != healthBucketCooldown && bucket != healthBucketBlacklisted { + restoreErr = app_errors.ErrInvalidCredentialState + return + } + stats := s.stats.Snapshot(credentialID, observedAt) + stats.ConsecutiveFailure = 0 + stats.ConsecutiveProblem = 0 + stats.LastFailureCategory = 0 + stats.LastStatusCode = 0 + if !s.registry.RestoreRuntimeState(credentialID, calculateAutoWeight(stats)) { + restoreErr = dbRegistryMismatch(mismatchMissingRegistry, groupID, credentialID) + return + } + s.stats.ClearProblemState(credentialID) + } + if s.mutations == nil { + restore() + } else { + s.mutations.Do(credentialID, restore) + } + if restoreErr != nil { + return CredentialItemResponse{}, restoreErr + } + view, exists = findRuntimeCredential(s.registry.Snapshot(), credentialID) + if !exists { + return CredentialItemResponse{}, dbRegistryMismatch(mismatchMissingRegistry, groupID, credentialID) + } + return s.mapCredentialItem(row, view, group, s.stats.Snapshot(credentialID, observedAt), observedAt) +} + +func validateCredentialRuntimeRow( + groupID uint, + row models.Credential, + view state.CredentialRuntimeView, + exists bool, +) error { + if !exists { + return dbRegistryMismatch(mismatchMissingRegistry, groupID, row.ID) + } + if view.GroupID != groupID { + return dbRegistryMismatch(mismatchGroupID, groupID, row.ID) + } + if view.Status != state.CredentialStatus(row.Status) { + return dbRegistryMismatch(mismatchStatus, groupID, row.ID) + } + if !equalOptionalWeight(view.WeightManual, row.WeightManual) { + return dbRegistryMismatch(mismatchWeightManual, groupID, row.ID) + } + if view.Version != groupCollectionCredentialVersion(row.UpdatedAtMS) || + view.IdentityGeneration != groupCollectionCredentialIdentity(row.Fingerprint) { + return dbRegistryMismatch(mismatchIdentity, groupID, row.ID) + } + return nil +} + +func (s *Service) loadCredentialItem(ctx context.Context, groupID, credentialID uint) (CredentialItemResponse, error) { + capture, err := s.captureCredentials(ctx, groupID) + if err != nil { + return CredentialItemResponse{}, err + } + observation, err := validateCredentialCapture(capture) + if err != nil { + return CredentialItemResponse{}, err + } + for _, row := range observation.rows { + if row.ID == credentialID { + return s.mapCredentialItem(row, observation.runtime[credentialID], observation.group, + s.stats.Snapshot(credentialID, observation.observedAt), observation.observedAt) + } + } + return CredentialItemResponse{}, credentialNotFoundError() +} + +func (s *Service) mapCredentialItem( + row models.Credential, + view state.CredentialRuntimeView, + group models.Group, + stats health.CredentialStats, + observedAt time.Time, +) (CredentialItemResponse, error) { + canonical, _, err := s.decodeCredential(group, row) + if err != nil { + return CredentialItemResponse{}, err + } + mask, err := maskCanonicalCredential(canonical) + if err != nil { + return CredentialItemResponse{}, err + } + bucket := classifyHealthKey(state.GroupCatalogView{ID: group.ID, Name: group.Name, Enabled: group.Enabled, + WeightManual: cloneInt(group.WeightManual)}, view, observedAt) + return mapCredentialRuntimeItem(mask, row.ID, view, bucket, stats, observedAt) +} + +func normalizeCredentialBatchRequest(request CredentialBatchRequest) ([]uint, error) { + if request.Action != CredentialBatchEnable && request.Action != CredentialBatchDisable && request.Action != CredentialBatchDelete { + return nil, app_errors.ErrValidation + } + if len(request.CredentialIDs) < 1 || len(request.CredentialIDs) > 100 { + return nil, app_errors.ErrValidation + } + ids := append([]uint(nil), request.CredentialIDs...) + sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) + for index, id := range ids { + if id == 0 || index > 0 && id == ids[index-1] { + return nil, app_errors.ErrValidation + } + } + return ids, nil +} + +func (s *Service) BatchGroupCredentials( + ctx context.Context, + groupID uint, + request CredentialBatchRequest, +) (CredentialBatchResponse, error) { + if groupID == 0 { + return CredentialBatchResponse{}, app_errors.ErrBadRequest + } + ids, err := normalizeCredentialBatchRequest(request) + if err != nil { + return CredentialBatchResponse{}, err + } + s.writeMu.Lock() + defer s.writeMu.Unlock() + if err := s.enforceOperationRecoveryBarrierLocked(ctx, 0); err != nil { + return CredentialBatchResponse{}, err + } + group, err := loadGroupRow(s.db.WithContext(ctx), groupID) + if err != nil { + return CredentialBatchResponse{}, err + } + if group.ChannelID == "" { + return CredentialBatchResponse{}, app_errors.ErrValidation + } + var rows []models.Credential + if err := s.db.WithContext(ctx).Where("id IN ?", ids).Find(&rows).Error; err != nil { + return CredentialBatchResponse{}, app_errors.ParseDBError(err) + } + if len(rows) != len(ids) { + return CredentialBatchResponse{}, credentialNotFoundError() + } + rowByID := make(map[uint]models.Credential, len(rows)) + viewByID := make(map[uint]state.CredentialRuntimeView) + for _, view := range s.registry.Snapshot() { + viewByID[view.ID] = view + } + for _, row := range rows { + if row.GroupID != groupID { + return CredentialBatchResponse{}, credentialNotFoundError() + } + view, exists := viewByID[row.ID] + if err := validateCredentialRuntimeRow(groupID, row, view, exists); err != nil { + return CredentialBatchResponse{}, err + } + rowByID[row.ID] = row + } + coordinator, ok := s.mutations.(interface{ DoMany([]uint, func()) }) + if !ok { + return CredentialBatchResponse{}, fmt.Errorf("batch mutation coordinator unavailable: %w", app_errors.ErrInternalServer) + } + var mutationErr error + coordinator.DoMany(ids, func() { + before, snapshotErr := s.registry.SnapshotGroupCredentialEntriesExact(groupID, ids) + if snapshotErr != nil { + mutationErr = fmt.Errorf("snapshot credential registry entries: %w", app_errors.ErrInternalServer) + return + } + nowMS := int64(0) + for _, id := range ids { + candidate, versionErr := nextCredentialUpdatedAtMS(s.now(), rowByID[id].UpdatedAtMS) + if versionErr != nil { + mutationErr = app_errors.ErrInternalServer + return + } + if candidate > nowMS { + nowMS = candidate + } + } + desired := make([]state.CredentialEntry, len(before)) + for index, entry := range before { + desired[index] = entry + if request.Action == CredentialBatchEnable { + desired[index].Status = state.CredentialStatusActive + } else if request.Action == CredentialBatchDisable { + desired[index].Status = state.CredentialStatusDisabled + } + desired[index].Version = uint64(nowMS) + } + if s.applyBatchRegistryMutation == nil { + mutationErr = app_errors.ErrInternalServer + return + } + if applyErr := s.applyBatchRegistryMutation(groupID, ids, request.Action); applyErr != nil { + mutationErr = withControlOperationContext(newControlOperationError(stageApplyCommittedRegistryMutation), groupID, 0) + return + } + if request.Action != CredentialBatchDelete { + if restoreErr := s.registry.RestoreGroupCredentialEntriesExact(groupID, desired); restoreErr != nil { + mutationErr = compensateCredentialBatchRegistry(s, groupID, before, restoreErr) + return + } + } + mutationErr = s.withControlTransaction(ctx, func(tx *gorm.DB) error { + query := tx.Where("group_id = ? AND id IN ?", groupID, ids) + var result *gorm.DB + switch request.Action { + case CredentialBatchEnable: + result = query.Model(&models.Credential{}).Updates(map[string]any{ + "status": models.CredentialStatusActive, "updated_at_ms": nowMS, + }) + case CredentialBatchDisable: + result = query.Model(&models.Credential{}).Updates(map[string]any{ + "status": models.CredentialStatusDisabled, "updated_at_ms": nowMS, + }) + case CredentialBatchDelete: + result = query.Delete(&models.Credential{}) + } + if result.Error != nil { + return app_errors.ParseDBError(result.Error) + } + if result.RowsAffected != int64(len(ids)) { + return fmt.Errorf("batch credential rows affected = %d, want %d: %w", result.RowsAffected, len(ids), app_errors.ErrDatabase) + } + return nil + }) + if mutationErr != nil { + mutationErr = compensateCredentialBatchRegistry(s, groupID, before, mutationErr) + return + } + if request.Action == CredentialBatchDelete { + for _, id := range ids { + s.stats.Reset(id) + } + } + }) + if mutationErr != nil { + return CredentialBatchResponse{}, mutationErr + } + return CredentialBatchResponse{ + AffectedCredentialIDs: ids, + Summary: summarizeGroupRuntimeCredentials(group, s.registry.Snapshot(), s.now().UTC()), + }, nil +} + +func (s *Service) applyCredentialBatchRegistryMutation( + groupID uint, + credentialIDs []uint, + action CredentialBatchAction, +) error { + switch action { + case CredentialBatchEnable: + return s.registry.UpdateGroupCredentialStatuses( + groupID, credentialIDs, state.CredentialStatusActive, + ) + case CredentialBatchDisable: + return s.registry.UpdateGroupCredentialStatuses( + groupID, credentialIDs, state.CredentialStatusDisabled, + ) + case CredentialBatchDelete: + return s.registry.RemoveGroupCredentials(groupID, credentialIDs) + default: + return fmt.Errorf("unsupported batch credential action %q", action) + } +} + +func compensateCredentialBatchRegistry( + s *Service, + groupID uint, + before []state.CredentialEntry, + cause error, +) error { + if s.restoreBatchRegistryEntries == nil { + return errors.Join( + cause, + fmt.Errorf("compensate batch credential Registry mutation: %w", app_errors.ErrInternalServer), + ) + } + if err := s.restoreBatchRegistryEntries(groupID, before); err != nil { + return errors.Join( + cause, + fmt.Errorf("compensate batch credential Registry mutation: %w", err), + ) + } + return cause +} + +func summarizeGroupRuntimeCredentials( + group models.Group, + views []state.CredentialRuntimeView, + observedAt time.Time, +) CredentialSummaryResponse { + summary := CredentialSummaryResponse{} + groupView := state.GroupCatalogView{ + ID: group.ID, Name: group.Name, Enabled: group.Enabled, + WeightManual: cloneInt(group.WeightManual), + } + for _, view := range views { + if view.GroupID != group.ID { + continue + } + summary.Total++ + switch classifyHealthKey(groupView, view, observedAt) { + case healthBucketAvailable: + summary.Available++ + case healthBucketCooldown: + summary.Cooldown++ + case healthBucketBlacklisted: + summary.Blacklisted++ + case healthBucketDisabled: + summary.Disabled++ + } + } + return summary +} diff --git a/internal/control/credentials.go b/internal/control/credentials.go new file mode 100644 index 00000000..66783b7f --- /dev/null +++ b/internal/control/credentials.go @@ -0,0 +1,373 @@ +package control + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sort" + "strings" + "time" + + "gorm.io/gorm" + + "gpt-load/internal/channel" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/state" + stateloader "gpt-load/internal/state/loader" + "gpt-load/internal/storage/models" +) + +type CredentialImportRequest struct { + Credentials string `json:"credentials"` +} + +type CredentialImportResult struct { + GroupID uint `json:"group_id"` + CredentialsAdded int `json:"credentials_added"` + CredentialsDuplicated int `json:"credentials_duplicated"` +} + +type CredentialUpdateRequest struct { + Status optionalField[state.CredentialStatus] `json:"status"` + WeightManual optionalField[int] `json:"weight_manual"` +} + +type CredentialRevealResult struct { + CredentialID uint `json:"credential_id"` + Credential json.RawMessage `json:"credential"` + RevealedAtMS int64 `json:"revealed_at_ms"` +} + +type CredentialCollectionQuery struct { + Query string + Status *string + Page int + PageSize int +} + +type CredentialCollectionResponse struct { + ObservedAtMS int64 `json:"observed_at_ms"` + StatsWindowSeconds int64 `json:"stats_window_seconds"` + Summary CredentialSummaryResponse `json:"summary"` + Items []CredentialItemResponse `json:"items"` + Pagination CredentialPaginationResponse `json:"pagination"` +} + +type CredentialSummaryResponse struct { + Total int `json:"total"` + Available int `json:"available"` + Cooldown int `json:"cooldown"` + Blacklisted int `json:"blacklisted"` + Disabled int `json:"disabled"` +} + +type CredentialItemResponse struct { + CredentialID uint `json:"credential_id"` + Mask string `json:"mask"` + ConfiguredStatus string `json:"configured_status"` + EffectiveStatus string `json:"effective_status"` + WeightMode string `json:"weight_mode"` + Weight *int `json:"weight"` + RecentSuccessCount uint64 `json:"recent_success_count"` + RecentFailureCount uint64 `json:"recent_failure_count"` + ConsecutiveFailureCount uint64 `json:"consecutive_failure_count"` + LastFailureCategory string `json:"last_failure_category"` + LastStatusCode *int `json:"last_status_code"` + CooldownUntilMS *int64 `json:"cooldown_until_ms"` + Recovery CredentialRecoveryResponse `json:"recovery"` +} + +type CredentialRecoveryResponse struct { + Mode string `json:"mode"` + Automatic bool `json:"automatic"` + AtMS *int64 `json:"at_ms"` +} + +type CredentialPaginationResponse struct { + Page int `json:"page"` + PageSize int `json:"page_size"` + TotalItems int `json:"total_items"` + TotalPages int `json:"total_pages"` +} + +type CredentialBatchAction string + +const ( + CredentialBatchEnable CredentialBatchAction = "enable" + CredentialBatchDisable CredentialBatchAction = "disable" + CredentialBatchDelete CredentialBatchAction = "delete" +) + +type CredentialBatchRequest struct { + Action CredentialBatchAction `json:"action"` + CredentialIDs []uint `json:"credential_ids"` +} + +type CredentialBatchResponse struct { + AffectedCredentialIDs []uint `json:"affected_credential_ids"` + Summary CredentialSummaryResponse `json:"summary"` +} + +type credentialCapture struct { + group models.Group + rows []models.Credential + views []state.CredentialRuntimeView + observedAt time.Time +} + +type credentialObservation struct { + group models.Group + rows []models.Credential + runtime map[uint]state.CredentialRuntimeView + observedAt time.Time +} + +type credentialCollectionRecord struct { + item CredentialItemResponse + bucket healthBucket +} + +func (s *Service) ImportGroupCredentials( + ctx context.Context, + groupID uint, + request CredentialImportRequest, +) (CredentialImportResult, error) { + if groupID == 0 { + return CredentialImportResult{}, app_errors.ErrValidation + } + result := CredentialImportResult{GroupID: groupID} + var entries []state.CredentialEntry + err := s.writeCredentialConfig(ctx, groupID, 0, func(tx *gorm.DB) error { + group, err := loadGroupRow(tx, groupID) + if err != nil { + return err + } + if group.ChannelID == "" { + return app_errors.ErrValidation + } + normalized, err := s.normalizeCredentials(channel.ID(group.ChannelID), request.Credentials) + if err != nil { + return err + } + result.CredentialsAdded, result.CredentialsDuplicated, err = + s.persistCredentials(tx, groupID, normalized) + if err != nil { + return err + } + entries, err = stateloader.BuildGroupCredentialEntries(ctx, tx, groupID) + if err != nil { + return err + } + return state.ValidateCredentialEntries(entries) + }, func() error { + _, err := s.reconcileRegistryGroup(groupID, entries) + return err + }) + if err != nil { + return CredentialImportResult{}, err + } + return result, nil +} + +func (s *Service) ListGroupCredentials( + ctx context.Context, + groupID uint, + query CredentialCollectionQuery, +) (CredentialCollectionResponse, error) { + if groupID == 0 { + return CredentialCollectionResponse{}, app_errors.ErrBadRequest + } + if query.Page < 1 || (query.PageSize != 20 && query.PageSize != 50 && query.PageSize != 100) { + return CredentialCollectionResponse{}, app_errors.ErrBadRequest + } + capture, err := s.captureCredentials(ctx, groupID) + if err != nil { + return CredentialCollectionResponse{}, err + } + observation, err := validateCredentialCapture(capture) + if err != nil { + return CredentialCollectionResponse{}, err + } + return s.mapCredentialCollection(observation, query) +} + +func (s *Service) captureCredentials(ctx context.Context, groupID uint) (credentialCapture, error) { + s.writeMu.RLock() + var group models.Group + groupErr := s.db.WithContext(ctx).Where("id = ?", groupID).Take(&group).Error + rows := make([]models.Credential, 0) + var rowsErr error + if groupErr == nil { + rowsErr = s.db.WithContext(ctx).Where("group_id = ?", groupID).Order("id ASC").Find(&rows).Error + } + var views []state.CredentialRuntimeView + var observedAt time.Time + if groupErr == nil && rowsErr == nil { + views = s.registry.Snapshot() + observedAt = s.now().UTC() + } + s.writeMu.RUnlock() + if groupErr != nil { + if errors.Is(groupErr, gorm.ErrRecordNotFound) { + return credentialCapture{}, groupNotFoundError() + } + return credentialCapture{}, app_errors.ParseDBError(groupErr) + } + if group.ChannelID == "" { + return credentialCapture{}, app_errors.ErrValidation + } + if rowsErr != nil { + return credentialCapture{}, app_errors.ParseDBError(rowsErr) + } + return credentialCapture{group: group, rows: rows, views: views, observedAt: observedAt}, nil +} + +func validateCredentialCapture(capture credentialCapture) (credentialObservation, error) { + groupID := capture.group.ID + byID := make(map[uint]state.CredentialRuntimeView, len(capture.views)) + for _, view := range capture.views { + byID[view.ID] = view + } + persisted := make(map[uint]struct{}, len(capture.rows)) + for _, row := range capture.rows { + view, exists := byID[row.ID] + if !exists { + return credentialObservation{}, dbRegistryMismatch(mismatchMissingRegistry, groupID, row.ID) + } + if view.GroupID != groupID { + return credentialObservation{}, dbRegistryMismatch(mismatchGroupID, groupID, row.ID) + } + if view.Status != state.CredentialStatus(row.Status) { + return credentialObservation{}, dbRegistryMismatch(mismatchStatus, groupID, row.ID) + } + if !equalOptionalWeight(view.WeightManual, row.WeightManual) { + return credentialObservation{}, dbRegistryMismatch(mismatchWeightManual, groupID, row.ID) + } + if view.Version != groupCollectionCredentialVersion(row.UpdatedAtMS) || + view.IdentityGeneration != groupCollectionCredentialIdentity(row.Fingerprint) { + return credentialObservation{}, dbRegistryMismatch(mismatchIdentity, groupID, row.ID) + } + persisted[row.ID] = struct{}{} + } + for _, view := range capture.views { + if view.GroupID == groupID { + if _, exists := persisted[view.ID]; !exists { + return credentialObservation{}, dbRegistryMismatch(mismatchExtraRegistry, groupID, view.ID) + } + } + } + return credentialObservation{group: capture.group, rows: capture.rows, runtime: byID, observedAt: capture.observedAt}, nil +} + +func (s *Service) decodeCredential(group models.Group, row models.Credential) (json.RawMessage, string, error) { + plaintext, err := s.encryption.Decrypt(row.Data) + if err != nil { + return nil, "", fmt.Errorf("decrypt credential %d: %w", row.ID, app_errors.ErrInternalServer) + } + validated, err := normalizeStoredCredential(s.channelRegistry, channel.ID(group.ChannelID), plaintext) + if err != nil { + return nil, "", fmt.Errorf("validate credential %d: %w", row.ID, app_errors.ErrInternalServer) + } + apiKey, _ := validated.Value("api_key") + return validated.CanonicalJSON(), apiKey, nil +} + +func (s *Service) mapCredentialCollection( + observation credentialObservation, + query CredentialCollectionQuery, +) (CredentialCollectionResponse, error) { + if s == nil || s.encryption == nil || s.stats == nil || s.channelRegistry == nil { + return CredentialCollectionResponse{}, fmt.Errorf("credential collection dependencies unavailable: %w", app_errors.ErrInternalServer) + } + observedAtMS, err := safeEpochMilliseconds(observation.observedAt) + if err != nil { + return CredentialCollectionResponse{}, err + } + group := state.GroupCatalogView{ID: observation.group.ID, Name: observation.group.Name, + Enabled: observation.group.Enabled, WeightManual: cloneInt(observation.group.WeightManual)} + records := make([]credentialCollectionRecord, 0, len(observation.rows)) + for _, row := range observation.rows { + canonical, _, err := s.decodeCredential(observation.group, row) + if err != nil { + return CredentialCollectionResponse{}, err + } + mask, err := maskCanonicalCredential(canonical) + if err != nil { + return CredentialCollectionResponse{}, err + } + view := observation.runtime[row.ID] + bucket := classifyHealthKey(group, view, observation.observedAt) + item, err := mapCredentialRuntimeItem( + mask, row.ID, view, bucket, s.stats.Snapshot(row.ID, observation.observedAt), observation.observedAt, + ) + if err != nil { + return CredentialCollectionResponse{}, err + } + records = append(records, credentialCollectionRecord{item: item, bucket: bucket}) + } + summary := summarizeCredentialCollection(records) + filtered := make([]credentialCollectionRecord, 0, len(records)) + for _, record := range records { + if credentialCollectionMatches(record, query) { + filtered = append(filtered, record) + } + } + sort.Slice(filtered, func(i, j int) bool { + left, right := filtered[i], filtered[j] + if credentialCollectionBucketOrder(left.bucket) != credentialCollectionBucketOrder(right.bucket) { + return credentialCollectionBucketOrder(left.bucket) < credentialCollectionBucketOrder(right.bucket) + } + return left.item.CredentialID < right.item.CredentialID + }) + total := len(filtered) + return CredentialCollectionResponse{ + ObservedAtMS: observedAtMS, StatsWindowSeconds: credentialCollectionStatsWindow, + Summary: summary, Items: credentialCollectionPage(filtered, query.Page, query.PageSize), + Pagination: CredentialPaginationResponse{Page: query.Page, PageSize: query.PageSize, + TotalItems: total, TotalPages: credentialCollectionTotalPages(total, query.PageSize)}, + }, nil +} + +func summarizeCredentialCollection(records []credentialCollectionRecord) CredentialSummaryResponse { + summary := CredentialSummaryResponse{Total: len(records)} + for _, record := range records { + switch record.bucket { + case healthBucketAvailable: + summary.Available++ + case healthBucketCooldown: + summary.Cooldown++ + case healthBucketBlacklisted: + summary.Blacklisted++ + case healthBucketDisabled: + summary.Disabled++ + } + } + return summary +} + +func credentialCollectionMatches(record credentialCollectionRecord, query CredentialCollectionQuery) bool { + if query.Status != nil && record.item.EffectiveStatus != *query.Status { + return false + } + return query.Query == "" || strings.Contains(strings.ToLower(record.item.Mask), strings.ToLower(query.Query)) +} + +func credentialCollectionPage(records []credentialCollectionRecord, page, pageSize int) []CredentialItemResponse { + if page < 1 || pageSize < 1 || page-1 > len(records)/pageSize { + return []CredentialItemResponse{} + } + offset := (page - 1) * pageSize + if offset < 0 || offset >= len(records) { + return []CredentialItemResponse{} + } + end := offset + pageSize + if end < offset || end > len(records) { + end = len(records) + } + items := make([]CredentialItemResponse, end-offset) + for index, record := range records[offset:end] { + items[index] = record.item + } + return items +} diff --git a/internal/control/discover.go b/internal/control/discover.go index ecbc3a7d..07bcb267 100644 --- a/internal/control/discover.go +++ b/internal/control/discover.go @@ -2,20 +2,20 @@ package control import ( "context" + "encoding/json" "fmt" - "gpt-load/internal/platform/config" + "gpt-load/internal/channel" + "gpt-load/internal/execution" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" ) type ModelDiscoveryRequest struct { - ProviderID *string `json:"provider_id,omitempty"` - UpstreamURL string `json:"upstream_url"` - Protocols []protocol.Protocol `json:"protocols"` - Keys string `json:"keys"` + ChannelID channel.ID `json:"channel_id"` + Params json.RawMessage `json:"params"` + Credentials string `json:"credentials"` } type ModelDiscoveryResult struct { @@ -30,19 +30,14 @@ func (s *Service) DiscoverModels( return ModelDiscoveryResult{}, err } - baseURL, _, err := normalizeUpstreamBaseURL(request.UpstreamURL) - if err != nil { - return ModelDiscoveryResult{}, err + if s == nil || s.channelRegistry == nil || request.ChannelID == "" { + return ModelDiscoveryResult{}, app_errors.ErrValidation } - protocols, err := normalizeGroupProtocols(request.Protocols) - if err != nil { - return ModelDiscoveryResult{}, err - } - providerID, err := normalizeProviderID(request.ProviderID) + resolvedTarget, err := s.channelRegistry.Resolve(request.ChannelID, request.Params) if err != nil { return ModelDiscoveryResult{}, app_errors.ErrValidation } - keys, err := s.normalizeUpstreamKeys(request.Keys) + credentials, err := s.normalizeCredentials(request.ChannelID, request.Credentials) if err != nil { return ModelDiscoveryResult{}, err } @@ -54,10 +49,10 @@ func (s *Service) DiscoverModels( return ModelDiscoveryResult{}, fmt.Errorf("load model discovery settings: %w", err) } snapshot, err := state.Compile(state.CompileInput{ - SystemSettings: systemSettings, + SystemSettings: systemSettings, ChannelRegistry: s.channelRegistry, Groups: []state.GroupConfig{{ - ID: 1, Name: "draft", UpstreamURL: baseURL, - Protocols: protocols, Settings: config.Settings{}, Enabled: true, + ID: 1, Name: "draft", ChannelID: request.ChannelID, + Params: append(json.RawMessage(nil), request.Params...), Enabled: true, }}, }) if err != nil { @@ -68,15 +63,23 @@ func (s *Service) DiscoverModels( return ModelDiscoveryResult{}, fmt.Errorf("compiled model discovery draft is missing") } - plaintextKeys := make([]string, 0, len(keys.candidates)) - for _, candidate := range keys.candidates { - plaintextKeys = append(plaintextKeys, candidate.plaintext) + discoveryCredentials := make([]discoveryCredential, 0, len(credentials.candidates)) + for index, candidate := range credentials.candidates { + credential, err := s.channelRegistry.ValidateCredential(request.ChannelID, candidate.canonical) + if err != nil { + return ModelDiscoveryResult{}, app_errors.ErrValidation + } + apiKey, _ := credential.Value("api_key") + discoveryCredentials = append(discoveryCredentials, discoveryCredential{ + snapshot: execution.NewCredentialSnapshot( + uint(index+1), 1, uint64(index+1), credential.CanonicalJSON(), + ), + apiKey: apiKey, + }) } return s.executeModelDiscovery(ctx, discoveryTarget{ - baseURL: baseURL, - protocols: protocols, - keys: plaintextKeys, - headerRules: group.HeaderRules, - providerID: providerID, + channelID: request.ChannelID, resolvedTarget: resolvedTarget, + credentials: discoveryCredentials, headerRules: group.HeaderRules, + timeouts: group.Timeouts, catalogProviderID: resolvedTarget.CatalogProviderID, }) } diff --git a/internal/control/discover_executor.go b/internal/control/discover_executor.go index c427e762..02cc09df 100644 --- a/internal/control/discover_executor.go +++ b/internal/control/discover_executor.go @@ -2,25 +2,36 @@ package control import ( "context" + cryptorand "crypto/rand" + "encoding/json" "errors" "fmt" + "io" + "net/http" "sort" "strings" "gpt-load/internal/catalog" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" + "gpt-load/internal/execution" app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/pricing" "gpt-load/internal/protocol" "gpt-load/internal/state" - "gpt-load/internal/storage/models" ) +type discoveryCredential struct { + snapshot execution.CredentialSnapshot + apiKey string +} + type discoveryTarget struct { - baseURL string - protocols []protocol.Protocol - keys []string - headerRules state.HeaderRules - providerID *string + channelID channel.ID + resolvedTarget channel.ResolvedTarget + credentials []discoveryCredential + headerRules state.HeaderRules + timeouts state.TimeoutConfig + catalogProviderID string } func (s *Service) executeModelDiscovery( @@ -30,57 +41,183 @@ func (s *Service) executeModelDiscovery( if err := ctx.Err(); err != nil { return ModelDiscoveryResult{}, err } - if strings.TrimSpace(target.baseURL) == "" || len(target.protocols) == 0 || len(target.keys) == 0 { + if s == nil || s.executor == nil { + return ModelDiscoveryResult{}, app_errors.ErrInternalServer + } + if target.channelID == "" || + target.resolvedTarget.ChannelID != target.channelID || len(target.credentials) == 0 { return ModelDiscoveryResult{}, app_errors.ErrValidation } - - orderedProtocols := canonicalProtocolOrder(target.protocols) - selectedDialects := make([]dialect.Dialect, len(orderedProtocols)) - for index, value := range orderedProtocols { - selected, ok := s.dialects[value] - if !ok || selected == nil { - return ModelDiscoveryResult{}, fmt.Errorf( - "dialect for protocol %q is not configured", - value, - ) - } - selectedDialects[index] = selected + clientProtocol, method, path, body, err := utilityRequestShape( + target.resolvedTarget.ProviderKind, + execution.OperationListModels, + "", + ) + if err != nil { + return ModelDiscoveryResult{}, err } discoveryCtx, cancel := context.WithTimeout(ctx, s.modelDiscoveryTimeout) defer cancel() - for _, selected := range selectedDialects { - for _, apiKey := range target.keys { - models, err := selected.ListModels( - discoveryCtx, - target.baseURL, - apiKey, - target.headerRules, - ) - if parentErr := ctx.Err(); parentErr != nil { - return ModelDiscoveryResult{}, parentErr - } - if errors.Is(discoveryCtx.Err(), context.DeadlineExceeded) { - if parentErr := ctx.Err(); parentErr != nil { - return ModelDiscoveryResult{}, parentErr - } - return ModelDiscoveryResult{}, fmt.Errorf( - "discover upstream models: %w", - app_errors.ErrBadGateway, - ) - } - if err == nil { - return s.mergeDiscoveredModels(discoveryCtx, normalizeDiscoveredModels(models), target) - } + requestID, err := s.newExecutionID() + if err != nil { + return ModelDiscoveryResult{}, fmt.Errorf("create discovery request identity: %w", app_errors.ErrInternalServer) + } + for index, credential := range target.credentials { + attemptID, identityErr := s.newExecutionID() + if identityErr != nil { + return ModelDiscoveryResult{}, fmt.Errorf("create discovery attempt identity: %w", app_errors.ErrInternalServer) } + spec := execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: requestID, AttemptID: attemptID, Sequence: uint32(index + 1), + ChannelID: string(target.channelID), ClientProtocol: clientProtocol, + Operation: execution.OperationListModels, Method: method, Path: path, + Header: applyControlHeaderRules(target.headerRules, credential.apiKey), Body: body, + TargetConfig: target.resolvedTarget.TargetConfig, + Timeouts: executionTimeouts(target.timeouts), + Credential: credential.snapshot, + }) + if validationErr := spec.Validate(); validationErr != nil { + return ModelDiscoveryResult{}, fmt.Errorf("build discovery attempt: %w", app_errors.ErrInternalServer) + } + result := s.executor.Execute(discoveryCtx, spec) + if parentErr := ctx.Err(); parentErr != nil { + return ModelDiscoveryResult{}, parentErr + } + if errors.Is(discoveryCtx.Err(), context.DeadlineExceeded) { + return ModelDiscoveryResult{}, fmt.Errorf("discover upstream models: %w", app_errors.ErrBadGateway) + } + if result.Validate() != nil || result.Error != nil || + result.StatusCode < http.StatusOK || result.StatusCode >= http.StatusMultipleChoices { + continue + } + models, parseErr := parseDiscoveredModels(clientProtocol, result.Body) + if parseErr != nil { + continue + } + return s.mergeDiscoveredModels(discoveryCtx, normalizeDiscoveredModels(models), target) } if parentErr := ctx.Err(); parentErr != nil { return ModelDiscoveryResult{}, parentErr } - return ModelDiscoveryResult{}, fmt.Errorf( - "discover upstream models: %w", - app_errors.ErrBadGateway, - ) + return ModelDiscoveryResult{}, fmt.Errorf("discover upstream models: %w", app_errors.ErrBadGateway) +} + +func (s *Service) newExecutionID() (string, error) { + random := io.Reader(cryptorand.Reader) + if s != nil && s.random != nil { + random = s.random + } + return newOperationID(random) +} + +func executionTimeouts(value state.TimeoutConfig) execution.AttemptTimeouts { + return execution.AttemptTimeouts{ + Connect: value.Connect, FirstByte: value.FirstByte, + Request: value.Request, StreamIdle: value.StreamIdle, + } +} + +func applyControlHeaderRules(rules state.HeaderRules, apiKey string) http.Header { + headers := make(http.Header, len(rules.Set)) + for name, value := range rules.Set { + headers.Set(name, strings.ReplaceAll(value, "${API_KEY}", apiKey)) + } + for _, name := range rules.Remove { + headers.Del(name) + } + headers.Set("Accept-Encoding", "identity") + return headers +} + +func utilityRequestShape( + providerKind channel.ProviderKind, + operation execution.Operation, + model string, +) (protocol.Protocol, string, string, []byte, error) { + switch operation { + case execution.OperationListModels: + switch providerKind { + case channel.ProviderOpenAI, channel.ProviderOpenAICompatible: + return protocol.OpenAICompletions, http.MethodGet, "/v1/models", nil, nil + case channel.ProviderAzureOpenAI, channel.ProviderAWSBedrock, channel.ProviderGoogleVertex: + return protocol.OpenAICompletions, http.MethodGet, "/v1/models", nil, nil + case channel.ProviderAnthropic, channel.ProviderAnthropicCompatible: + return protocol.Anthropic, http.MethodGet, "/v1/models", nil, nil + case channel.ProviderGemini, channel.ProviderGeminiCompatible: + return protocol.Gemini, http.MethodGet, "/v1beta/models", nil, nil + } + case execution.OperationProbe: + model = strings.TrimSpace(model) + if model == "" { + return "", "", "", nil, app_errors.ErrValidation + } + switch providerKind { + case channel.ProviderOpenAI, channel.ProviderOpenAICompatible: + return protocol.OpenAICompletions, http.MethodPost, "/v1/chat/completions", + []byte(`{"model":"probe","messages":[{"role":"user","content":"ping"}],"max_tokens":1}`), nil + case channel.ProviderAzureOpenAI, channel.ProviderAWSBedrock, channel.ProviderGoogleVertex: + return protocol.OpenAICompletions, http.MethodPost, "/v1/chat/completions", + []byte(`{"model":"probe","messages":[{"role":"user","content":"ping"}],"max_tokens":1}`), nil + case channel.ProviderAnthropic, channel.ProviderAnthropicCompatible: + return protocol.Anthropic, http.MethodPost, "/v1/messages", + []byte(`{"model":"probe","max_tokens":1,"messages":[{"role":"user","content":"ping"}]}`), nil + case channel.ProviderGemini, channel.ProviderGeminiCompatible: + return protocol.Gemini, http.MethodPost, "/v1beta/models/probe:generateContent", + []byte(`{"contents":[{"parts":[{"text":"ping"}]}],"generationConfig":{"maxOutputTokens":1}}`), nil + } + } + return "", "", "", nil, app_errors.ErrValidation +} + +func parseDiscoveredModels(clientProtocol protocol.Protocol, body []byte) ([]string, error) { + switch clientProtocol { + case protocol.OpenAICompletions, protocol.OpenAIResponses, protocol.Anthropic: + var payload struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + if err := decodeSingleJSON(body, &payload); err != nil { + return nil, err + } + models := make([]string, 0, len(payload.Data)) + for _, item := range payload.Data { + models = append(models, item.ID) + } + return models, nil + case protocol.Gemini: + var payload struct { + Models []struct { + Name string `json:"name"` + } `json:"models"` + } + if err := decodeSingleJSON(body, &payload); err != nil { + return nil, err + } + models := make([]string, 0, len(payload.Models)) + for _, item := range payload.Models { + models = append(models, strings.TrimPrefix(item.Name, "models/")) + } + return models, nil + default: + return nil, app_errors.ErrValidation + } +} + +func decodeSingleJSON(body []byte, target any) error { + decoder := json.NewDecoder(strings.NewReader(string(body))) + if err := decoder.Decode(target); err != nil { + return err + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + if err == nil { + return fmt.Errorf("multiple JSON values") + } + return err + } + return nil } func (s *Service) mergeDiscoveredModels( @@ -93,12 +230,12 @@ func (s *Service) mergeDiscoveredModels( catalogSnapshot = s.catalogRuntime.Load() } var providerModels map[string]catalog.Model - if target.providerID != nil && catalogSnapshot != nil { - if provider, exists := catalogSnapshot.Providers[*target.providerID]; exists { + if target.catalogProviderID != "" && catalogSnapshot != nil { + if provider, exists := catalogSnapshot.Providers[target.catalogProviderID]; exists { providerModels = provider.Models } } - rows := map[string]*models.ModelPrice{} + rows := modelPriceRows{} if s.db != nil { loaded, err := loadModelPriceRows(ctx, s.db) if err != nil { @@ -111,7 +248,8 @@ func (s *Service) mergeDiscoveredModels( seen := make(map[string]int, len(live)+len(providerModels)) for _, id := range live { model, catalogMatch := providerModels[id] - pricingStatus, pricingSource := resolveCandidatePricing(rows[id], catalogSnapshot, id) + identity := pricing.Identity{ChannelID: string(target.channelID), ModelID: id} + pricingStatus, pricingSource := resolveCandidatePricing(rows[identity], catalogSnapshot, identity) candidate := ModelCandidate{ ID: id, Name: id, Sources: []string{"live"}, PricingStatus: pricingStatus, PricingSource: pricingSource, @@ -134,7 +272,8 @@ func (s *Service) mergeDiscoveredModels( if name == "" { name = id } - pricingStatus, pricingSource := resolveCandidatePricing(rows[id], catalogSnapshot, id) + identity := pricing.Identity{ChannelID: string(target.channelID), ModelID: id} + pricingStatus, pricingSource := resolveCandidatePricing(rows[identity], catalogSnapshot, identity) catalogOnly = append(catalogOnly, ModelCandidate{ ID: id, Name: name, Sources: []string{"catalog"}, PricingStatus: pricingStatus, PricingSource: pricingSource, @@ -168,28 +307,3 @@ func normalizeDiscoveredModels(values []string) []string { } return result } - -func canonicalProtocolOrder( - values []protocol.Protocol, -) []protocol.Protocol { - present := make(map[protocol.Protocol]struct{}, len(values)) - for _, value := range values { - present[value] = struct{}{} - } - result := make([]protocol.Protocol, 0, len(present)) - for _, value := range protocol.DataPlaneProtocols() { - if _, exists := present[value]; !exists { - continue - } - result = append(result, value) - delete(present, value) - } - for _, value := range values { - if _, exists := present[value]; !exists { - continue - } - result = append(result, value) - delete(present, value) - } - return result -} diff --git a/internal/control/discover_executor_test.go b/internal/control/discover_executor_test.go index 246015a4..2e4dacdb 100644 --- a/internal/control/discover_executor_test.go +++ b/internal/control/discover_executor_test.go @@ -2,335 +2,328 @@ package control import ( "context" + "encoding/json" "errors" - "fmt" "net/http" "reflect" "strings" + "sync" "testing" "time" - "gpt-load/internal/dialect" - "gpt-load/internal/health" + "gpt-load/internal/channel" + "gpt-load/internal/execution" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/protocol" "gpt-load/internal/state" ) -type recordingDiscoveryDialect struct { +// recordingDiscoveryExecutorTarget acts as a provider-neutral execution target. +// This lets the discovery service tests focus +// on target freezing and credential failover rather than SDK wire details. +type recordingDiscoveryExecutorTarget struct { value protocol.Protocol listFn func(context.Context, string, string, state.HeaderRules) ([]string, error) } -func (d *recordingDiscoveryDialect) Protocol() protocol.Protocol { - return d.value -} - -func (*recordingDiscoveryDialect) InspectRequest( - *dialect.ParsedRequest, -) (dialect.RequestMetadata, error) { - return dialect.RequestMetadata{}, nil -} - -func (*recordingDiscoveryDialect) BuildUpstreamURL(string, *dialect.ParsedRequest) (string, error) { - return "", nil -} - -func (*recordingDiscoveryDialect) InjectCredential(http.Header, string) {} - -func (d *recordingDiscoveryDialect) ListModels( +func (recorder *recordingDiscoveryExecutorTarget) ListModels( ctx context.Context, - baseURL, apiKey string, + baseURL string, + apiKey string, rules state.HeaderRules, ) ([]string, error) { - return d.listFn(ctx, baseURL, apiKey, rules) + return recorder.listFn(ctx, baseURL, apiKey, rules) } -func (*recordingDiscoveryDialect) Probe(context.Context, string, string, state.HeaderRules, string) error { - return nil +type recordingDiscoveryExecutor struct { + byProtocol map[protocol.Protocol]*recordingDiscoveryExecutorTarget } -func (*recordingDiscoveryDialect) ClassifyStatus(int, []byte) health.FailureCategory { - return health.FailureCategoryClientError +func newRecordingDiscoveryExecutor(values ...*recordingDiscoveryExecutorTarget) execution.Executor { + byProtocol := make(map[protocol.Protocol]*recordingDiscoveryExecutorTarget, len(values)) + for _, value := range values { + if value != nil { + byProtocol[value.value] = value + } + } + return &recordingDiscoveryExecutor{byProtocol: byProtocol} } -func TestExecuteModelDiscoveryUsesProtocolOuterKeyInnerFallback(t *testing.T) { - var calls []string - newRecorder := func(value protocol.Protocol) *recordingDiscoveryDialect { - return &recordingDiscoveryDialect{ - value: value, - listFn: func( - _ context.Context, - baseURL, apiKey string, - rules state.HeaderRules, - ) ([]string, error) { - calls = append(calls, string(value)+":"+apiKey) - if baseURL != "https://api.example.com/v1" || rules.Set["X-Test"] != "draft" { - t.Fatalf("ListModels target = %q, %#v", baseURL, rules) - } - if value == protocol.Anthropic && apiKey == "key-a" { - return make([]string, 0), nil - } - return nil, errors.New("try next combination") +func (*recordingDiscoveryExecutor) Capabilities() execution.CapabilitySet { + return execution.CapabilitySet{} +} + +func (executor *recordingDiscoveryExecutor) Execute( + ctx context.Context, + spec execution.AttemptSpec, +) execution.AttemptResult { + recorder := executor.byProtocol[spec.ClientProtocol] + if recorder == nil || recorder.listFn == nil { + return execution.AttemptResult{ + DispatchState: execution.DispatchNotSent, + Error: &execution.ErrorEvidence{ + Kind: execution.ErrorKindInvalidRequest, Summary: "missing test executor", }, } } - service := &Service{ - dialects: dialect.NewSet(newRecorder(protocol.OpenAICompletions), newRecorder(protocol.Anthropic)), - modelDiscoveryTimeout: time.Second, + var target struct { + BaseURL string `json:"base_url"` } - result, err := service.executeModelDiscovery(context.Background(), discoveryTarget{ - baseURL: "https://api.example.com/v1", - protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - keys: []string{"key-a", "key-b"}, - headerRules: state.HeaderRules{Set: map[string]string{"X-Test": "draft"}}, - }) + _ = json.Unmarshal(spec.TargetConfig, &target) + var credential struct { + APIKey string `json:"api_key"` + } + _ = json.Unmarshal(spec.Credential.Data(), &credential) + rules := state.HeaderRules{Set: make(map[string]string)} + for name, values := range spec.Header { + if strings.EqualFold(name, "Accept-Encoding") || len(values) == 0 { + continue + } + rules.Set[name] = values[len(values)-1] + } + models, err := recorder.listFn(ctx, target.BaseURL, credential.APIKey, rules) if err != nil { - t.Fatalf("executeModelDiscovery() error = %v", err) - } - if result.Models == nil || len(result.Models) != 0 { - t.Fatalf("models = %#v, want non-nil empty success", result.Models) - } - wantCalls := []string{ - "openai-completions:key-a", - "openai-completions:key-b", - "anthropic:key-a", - } - if !reflect.DeepEqual(calls, wantCalls) { - t.Fatalf("calls = %#v, want %#v", calls, wantCalls) - } -} - -func TestExecuteModelDiscoveryUsesCanonicalOpenAIRepresentativeFirst(t *testing.T) { - t.Parallel() - - var calls []protocol.Protocol - newRecorder := func(value protocol.Protocol) *recordingDiscoveryDialect { - return &recordingDiscoveryDialect{ - value: value, - listFn: func( - context.Context, - string, - string, - state.HeaderRules, - ) ([]string, error) { - calls = append(calls, value) - return []string{"gpt-5"}, nil + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + Error: &execution.ErrorEvidence{ + Kind: execution.ErrorKindTransport, Summary: "test discovery failed", }, } } - service := &Service{ - dialects: dialect.NewSet( - newRecorder(protocol.OpenAIResponses), - newRecorder(protocol.OpenAICompletions), - ), - modelDiscoveryTimeout: time.Second, - } - result, err := service.executeModelDiscovery( - context.Background(), - discoveryTarget{ - baseURL: "https://api.example.com", - protocols: []protocol.Protocol{ - protocol.OpenAIResponses, - protocol.OpenAICompletions, - }, - keys: []string{"key-a"}, - }, - ) - if err != nil { - t.Fatalf("executeModelDiscovery() error = %v", err) - } - if !reflect.DeepEqual(result.Models, []ModelCandidate{ - {ID: "gpt-5", Name: "gpt-5", Sources: []string{"live"}, PricingStatus: PricingStatusPending}, - }) { - t.Fatalf("models = %#v", result.Models) - } - if !reflect.DeepEqual(calls, []protocol.Protocol{ - protocol.OpenAICompletions, - }) { - t.Fatalf("ListModels protocols = %#v, want Chat only", calls) + body := encodeDiscoveryModelsForTest(spec.ClientProtocol, models) + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusOK, Header: http.Header{}, Body: body, } } -func TestExecuteModelDiscoveryRejectsMissingDialectBeforeHTTP(t *testing.T) { - calls := 0 - openAI := &recordingDiscoveryDialect{ - value: protocol.OpenAICompletions, - listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { - calls++ - return nil, nil - }, - } - service := &Service{ - dialects: dialect.NewSet(openAI), - modelDiscoveryTimeout: time.Second, - } - target := discoveryTarget{ - baseURL: "https://api.example.com", - protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - keys: []string{"secret-key"}, - } - _, err := service.executeModelDiscovery(context.Background(), target) - if err == nil || errors.Is(err, app_errors.ErrBadGateway) { - t.Fatalf("missing Dialect error = %v, want internal invariant error", err) - } - if calls != 0 { - t.Fatalf("ListModels calls = %d, want preflight rejection", calls) - } +func (*recordingDiscoveryExecutor) ExecuteStream( + context.Context, + execution.AttemptSpec, + execution.StreamSink, +) execution.StreamResult { + panic("unexpected stream execution") +} - service.dialects = dialect.Set{protocol.OpenAICompletions: openAI, protocol.Anthropic: nil} - _, err = service.executeModelDiscovery(context.Background(), target) - if err == nil || errors.Is(err, app_errors.ErrBadGateway) || calls != 0 { - t.Fatalf("nil Dialect result = error %v, calls %d", err, calls) +func encodeDiscoveryModelsForTest(value protocol.Protocol, models []string) []byte { + items := make([]map[string]string, 0, len(models)) + for _, model := range models { + if value == protocol.Gemini { + items = append(items, map[string]string{"name": "models/" + model}) + } else { + items = append(items, map[string]string{"id": model}) + } } + payload := map[string]any{"data": items} + if value == protocol.Gemini { + payload = map[string]any{"models": items} + } + body, _ := json.Marshal(payload) + return body +} - for name, invalid := range map[string]discoveryTarget{ - "base URL": {protocols: []protocol.Protocol{protocol.OpenAICompletions}, keys: []string{"key"}}, - "protocols": {baseURL: "https://api.example.com", keys: []string{"key"}}, - "keys": {baseURL: "https://api.example.com", protocols: []protocol.Protocol{protocol.OpenAICompletions}}, - } { - t.Run(name, func(t *testing.T) { - if _, err := service.executeModelDiscovery(context.Background(), invalid); err == nil { - t.Fatal("empty target error = nil") - } +func discoveryTargetForTest( + t *testing.T, + channelID channel.ID, + baseURL string, + keys []string, + rules state.HeaderRules, +) discoveryTarget { + t.Helper() + registry := channel.NewRegistry() + params := json.RawMessage(`{}`) + if baseURL != "" { + encoded, err := json.Marshal(map[string]string{"base_url": baseURL}) + if err != nil { + t.Fatal(err) + } + params = encoded + } + resolved, err := registry.Resolve(channelID, params) + if err != nil { + t.Fatalf("Resolve(%q) error = %v", channelID, err) + } + credentials := make([]discoveryCredential, 0, len(keys)) + for index, key := range keys { + data, err := json.Marshal(map[string]string{"api_key": key}) + if err != nil { + t.Fatal(err) + } + credentials = append(credentials, discoveryCredential{ + snapshot: execution.NewCredentialSnapshot(uint(index+1), 1, uint64(index+1), data), + apiKey: key, }) } - if calls != 0 { - t.Fatalf("ListModels calls after empty targets = %d, want 0", calls) + return discoveryTarget{ + channelID: channelID, resolvedTarget: resolved, credentials: credentials, + headerRules: rules, catalogProviderID: resolved.CatalogProviderID, + } +} + +func TestExecuteModelDiscoveryUsesNeutralExecutor(t *testing.T) { + t.Parallel() + + var observed execution.AttemptSpec + executor := scriptedDiscoveryExecutor{execute: func( + _ context.Context, + spec execution.AttemptSpec, + ) execution.AttemptResult { + observed = spec.Clone() + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusOK, Header: http.Header{}, + Body: []byte(`{"object":"list","data":[{"id":"gpt-test"}]}`), + } + }} + service := &Service{executor: executor, modelDiscoveryTimeout: time.Second} + target := discoveryTargetForTest( + t, channel.OpenAICompatible, "https://api.example.com/v1", []string{"secret-key"}, state.HeaderRules{}, + ) + result, err := service.executeModelDiscovery(context.Background(), target) + if err != nil { + t.Fatalf("executeModelDiscovery() error = %v", err) + } + if observed.Operation != execution.OperationListModels || + observed.ChannelID != string(channel.OpenAICompatible) || + observed.Credential.ID == 0 || string(observed.TargetConfig) != `{"base_url":"https://api.example.com/v1"}` { + t.Fatalf("observed spec = %#v", observed) + } + if got := result.Models; len(got) != 1 || got[0].ID != "gpt-test" { + t.Fatalf("models = %#v", got) + } +} + +type scriptedDiscoveryExecutor struct { + execute func(context.Context, execution.AttemptSpec) execution.AttemptResult +} + +func (scriptedDiscoveryExecutor) Capabilities() execution.CapabilitySet { + return execution.CapabilitySet{} +} + +func (executor scriptedDiscoveryExecutor) Execute(ctx context.Context, spec execution.AttemptSpec) execution.AttemptResult { + return executor.execute(ctx, spec) +} + +func (scriptedDiscoveryExecutor) ExecuteStream(context.Context, execution.AttemptSpec, execution.StreamSink) execution.StreamResult { + panic("unexpected stream execution") +} + +func TestExecuteModelDiscoveryFallsBackAcrossCredentialsInStableOrder(t *testing.T) { + var calls []string + recorder := &recordingDiscoveryExecutorTarget{ + value: protocol.OpenAICompletions, + listFn: func(_ context.Context, baseURL, apiKey string, rules state.HeaderRules) ([]string, error) { + calls = append(calls, apiKey) + if baseURL != "https://api.example.com/v1" || rules.Set["X-Test"] != "draft" { + t.Fatalf("target = %q/%#v", baseURL, rules) + } + if apiKey == "key-b" { + return []string{" z-model ", "a-model", "z-model"}, nil + } + return nil, errors.New("try next credential") + }, + } + service := &Service{executor: newRecordingDiscoveryExecutor(recorder), modelDiscoveryTimeout: time.Second} + target := discoveryTargetForTest(t, channel.OpenAICompatible, "https://api.example.com/v1", + []string{"key-a", "key-b"}, state.HeaderRules{Set: map[string]string{"X-Test": "draft"}}) + result, err := service.executeModelDiscovery(context.Background(), target) + if err != nil { + t.Fatalf("executeModelDiscovery() error = %v", err) + } + if !reflect.DeepEqual(calls, []string{"key-a", "key-b"}) { + t.Fatalf("calls = %#v", calls) + } + if got := []string{result.Models[0].ID, result.Models[1].ID}; !reflect.DeepEqual(got, []string{"z-model", "a-model"}) { + t.Fatalf("models = %#v", result.Models) } } func TestExecuteModelDiscoverySharesOneTotalTimeout(t *testing.T) { + var mu sync.Mutex var deadlines []time.Time - newRecorder := func(value protocol.Protocol) *recordingDiscoveryDialect { - return &recordingDiscoveryDialect{ - value: value, - listFn: func(ctx context.Context, _, _ string, _ state.HeaderRules) ([]string, error) { - deadline, ok := ctx.Deadline() - if !ok { - t.Fatal("ListModels context has no deadline") - } - deadlines = append(deadlines, deadline) - time.Sleep(2 * time.Millisecond) - return nil, errors.New("retry") - }, - } + recorder := &recordingDiscoveryExecutorTarget{ + value: protocol.Anthropic, + listFn: func(ctx context.Context, _, _ string, _ state.HeaderRules) ([]string, error) { + deadline, ok := ctx.Deadline() + if !ok { + t.Fatal("discovery context has no deadline") + } + mu.Lock() + deadlines = append(deadlines, deadline) + mu.Unlock() + return nil, errors.New("retry") + }, } - service := &Service{ - dialects: dialect.NewSet(newRecorder(protocol.OpenAICompletions), newRecorder(protocol.Anthropic)), - modelDiscoveryTimeout: 200 * time.Millisecond, - } - _, err := service.executeModelDiscovery(context.Background(), discoveryTarget{ - baseURL: "https://api.example.com", - protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - keys: []string{"key-a", "key-b"}, - }) + service := &Service{executor: newRecordingDiscoveryExecutor(recorder), modelDiscoveryTimeout: time.Second} + target := discoveryTargetForTest(t, channel.AnthropicCompatible, "https://api.example.com", + []string{"key-a", "key-b"}, state.HeaderRules{}) + _, err := service.executeModelDiscovery(context.Background(), target) if !errors.Is(err, app_errors.ErrBadGateway) { - t.Fatalf("executeModelDiscovery() error = %v, want ErrBadGateway", err) + t.Fatalf("executeModelDiscovery() error = %v", err) } - if len(deadlines) != 4 { - t.Fatalf("observed deadlines = %d, want all four combinations", len(deadlines)) - } - for index := 1; index < len(deadlines); index++ { - if !deadlines[index].Equal(deadlines[0]) { - t.Fatalf("deadline[%d] = %v, want shared %v", index, deadlines[index], deadlines[0]) - } - } -} - -func TestExecuteModelDiscoveryRejectsSuccessAfterInternalTimeout(t *testing.T) { - service := &Service{ - dialects: dialect.NewSet(&recordingDiscoveryDialect{ - value: protocol.OpenAICompletions, - listFn: func(ctx context.Context, _, _ string, _ state.HeaderRules) ([]string, error) { - <-ctx.Done() - return []string{"late-model"}, nil - }, - }), - modelDiscoveryTimeout: 10 * time.Millisecond, - } - result, err := service.executeModelDiscovery(context.Background(), discoveryTarget{ - baseURL: "https://api.example.com", - protocols: []protocol.Protocol{protocol.OpenAICompletions}, - keys: []string{"key-a"}, - }) - if !errors.Is(err, app_errors.ErrBadGateway) { - t.Fatalf("executeModelDiscovery() = %#v, %v, want zero result and ErrBadGateway", result, err) - } - if result.Models != nil { - t.Fatalf("models = %#v, want nil after internal timeout", result.Models) + if len(deadlines) != 2 || !deadlines[0].Equal(deadlines[1]) { + t.Fatalf("deadlines = %#v", deadlines) } } func TestExecuteModelDiscoveryReturnsParentCancellation(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) calls := 0 - service := &Service{ - dialects: dialect.NewSet(&recordingDiscoveryDialect{ - value: protocol.OpenAICompletions, - listFn: func(discoveryCtx context.Context, _, _ string, _ state.HeaderRules) ([]string, error) { - calls++ - cancel() - <-discoveryCtx.Done() - return nil, discoveryCtx.Err() - }, - }), - modelDiscoveryTimeout: time.Second, + recorder := &recordingDiscoveryExecutorTarget{ + value: protocol.Gemini, + listFn: func(discoveryCtx context.Context, _, _ string, _ state.HeaderRules) ([]string, error) { + calls++ + cancel() + <-discoveryCtx.Done() + return nil, discoveryCtx.Err() + }, } - _, err := service.executeModelDiscovery(ctx, discoveryTarget{ - baseURL: "https://api.example.com", - protocols: []protocol.Protocol{protocol.OpenAICompletions}, - keys: []string{"key-a", "key-b"}, - }) - if err != context.Canceled { - t.Fatalf("executeModelDiscovery() error = %v, want exact context.Canceled", err) - } - if calls != 1 { - t.Fatalf("ListModels calls = %d, want stop after parent cancellation", calls) + service := &Service{executor: newRecordingDiscoveryExecutor(recorder), modelDiscoveryTimeout: time.Second} + target := discoveryTargetForTest(t, channel.GeminiCompatible, "https://api.example.com", + []string{"key-a", "key-b"}, state.HeaderRules{}) + _, err := service.executeModelDiscovery(ctx, target) + if err != context.Canceled || calls != 1 { + t.Fatalf("error/calls = %v/%d", err, calls) } } -func TestExecuteModelDiscoverySanitizesAllCombinationFailures(t *testing.T) { - const ( - baseURLSecret = "https://secret.example.com/v1?token=query-secret" - bodySecret = "distinctive-upstream-body" - partialSecret = "partial-secret-model" - ) - keys := []string{"key-secret-a", "key-secret-b"} - newRecorder := func(value protocol.Protocol) *recordingDiscoveryDialect { - return &recordingDiscoveryDialect{ - value: value, - listFn: func(_ context.Context, baseURL, apiKey string, _ state.HeaderRules) ([]string, error) { - return []string{partialSecret}, fmt.Errorf( - "provider failure url=%s key=%s body=%s", - baseURL, - apiKey, - bodySecret, - ) - }, - } - } +func TestExecuteModelDiscoveryRejectsInvalidTargetBeforeDispatch(t *testing.T) { + calls := 0 service := &Service{ - dialects: dialect.NewSet(newRecorder(protocol.OpenAICompletions), newRecorder(protocol.Anthropic)), + executor: scriptedDiscoveryExecutor{execute: func(context.Context, execution.AttemptSpec) execution.AttemptResult { + calls++ + return execution.AttemptResult{} + }}, modelDiscoveryTimeout: time.Second, } - result, err := service.executeModelDiscovery(context.Background(), discoveryTarget{ - baseURL: baseURLSecret, - protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - keys: keys, - }) - if !errors.Is(err, app_errors.ErrBadGateway) { - t.Fatalf("executeModelDiscovery() error = %v, want ErrBadGateway", err) + for name, target := range map[string]discoveryTarget{ + "empty channel": {}, + "empty credentials": func() discoveryTarget { + value := discoveryTargetForTest(t, channel.OpenAICompatible, "https://api.example.com", []string{"key"}, state.HeaderRules{}) + value.credentials = nil + return value + }(), + "mismatched target": func() discoveryTarget { + value := discoveryTargetForTest(t, channel.OpenAICompatible, "https://api.example.com", []string{"key"}, state.HeaderRules{}) + value.channelID = channel.AnthropicCompatible + return value + }(), + } { + t.Run(name, func(t *testing.T) { + if _, err := service.executeModelDiscovery(context.Background(), target); err == nil { + t.Fatal("invalid target error = nil") + } + }) } - if result.Models != nil { - t.Fatalf("partial result = %#v, want zero result", result) - } - for _, forbidden := range append( - []string{baseURLSecret, "secret.example.com", "query-secret", bodySecret, partialSecret}, - keys..., - ) { - if strings.Contains(err.Error(), forbidden) { - t.Fatalf("error exposes %q: %s", forbidden, err) - } + if calls != 0 { + t.Fatalf("dispatch calls = %d", calls) + } +} + +func TestNormalizeDiscoveredModels(t *testing.T) { + got := normalizeDiscoveredModels([]string{" model-b ", "", "model-a", "model-b", "\t"}) + if !reflect.DeepEqual(got, []string{"model-b", "model-a"}) { + t.Fatalf("normalizeDiscoveredModels() = %#v", got) } } diff --git a/internal/control/discover_group.go b/internal/control/discover_group.go index 3baf0c7c..f40b7427 100644 --- a/internal/control/discover_group.go +++ b/internal/control/discover_group.go @@ -11,9 +11,10 @@ import ( "gorm.io/gorm" "gorm.io/gorm/clause" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/config" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" "gpt-load/internal/storage/models" @@ -31,10 +32,10 @@ func (s *Service) DiscoverGroupModels( } type groupDiscoverySnapshotRows struct { - found bool - group models.Group - keys []models.UpstreamKey - settings []models.SystemSetting + found bool + group models.Group + credentials []models.Credential + settings []models.SystemSetting } func (s *Service) buildGroupDiscoveryTarget( @@ -78,14 +79,14 @@ func (s *Service) readGroupDiscoverySnapshot( rows.found = true rows.group = cloneGroupRows(groups)[0] - var keys []models.UpstreamKey + var credentials []models.Credential if err := tx. - Where("group_id = ? AND status = ?", groupID, models.UpstreamKeyStatusActive). + Where("group_id = ? AND status = ?", groupID, models.CredentialStatusActive). Order("id ASC"). - Find(&keys).Error; err != nil { + Find(&credentials).Error; err != nil { return err } - rows.keys = cloneUpstreamKeyRows(keys) + rows.credentials = cloneDiscoveryCredentialRows(credentials) var settings []models.SystemSetting if err := tx. @@ -108,10 +109,11 @@ func (s *Service) readGroupDiscoverySnapshot( return rows, nil } -func cloneUpstreamKeyRows(rows []models.UpstreamKey) []models.UpstreamKey { - cloned := make([]models.UpstreamKey, len(rows)) +func cloneDiscoveryCredentialRows(rows []models.Credential) []models.Credential { + cloned := make([]models.Credential, len(rows)) for index := range rows { cloned[index] = rows[index] + cloned[index].Group = nil if rows[index].WeightManual != nil { value := *rows[index].WeightManual cloned[index].WeightManual = &value @@ -126,21 +128,27 @@ func (s *Service) mapGroupDiscoveryTarget( if !rows.found { return discoveryTarget{}, app_errors.ErrResourceNotFound } - if len(rows.keys) == 0 { - return discoveryTarget{}, app_errors.ErrNoActiveUpstreamKey + if len(rows.credentials) == 0 { + return discoveryTarget{}, app_errors.ErrNoActiveCredential } - var protocols []protocol.Protocol - if err := decodeGroupDiscoveryJSON(rows.group.Protocols, &protocols); err != nil { - return discoveryTarget{}, fmt.Errorf("decode persisted discovery protocols: %w", app_errors.ErrInternalServer) + if s == nil || s.channelRegistry == nil { + return discoveryTarget{}, app_errors.ErrInternalServer } - groupConfig := rows.group.Config - if len(bytes.TrimSpace(groupConfig)) == 0 { - groupConfig = models.JSON(`{}`) + resolvedTarget, err := s.channelRegistry.Resolve( + channel.ID(rows.group.ChannelID), + json.RawMessage(rows.group.Params), + ) + if err != nil { + return discoveryTarget{}, fmt.Errorf("resolve persisted discovery channel: %w", app_errors.ErrInternalServer) + } + overrides := rows.group.Overrides + if len(bytes.TrimSpace(overrides)) == 0 { + overrides = models.JSON(`{}`) } settings := make(config.Settings) - if err := decodeGroupDiscoveryJSON(groupConfig, &settings); err != nil { - return discoveryTarget{}, fmt.Errorf("decode persisted discovery config: %w", app_errors.ErrInternalServer) + if err := decodeGroupDiscoveryJSON(overrides, &settings); err != nil { + return discoveryTarget{}, fmt.Errorf("decode persisted discovery overrides: %w", app_errors.ErrInternalServer) } systemSettings, err := stateloader.MapSystemSettings(rows.settings) @@ -148,11 +156,11 @@ func (s *Service) mapGroupDiscoveryTarget( return discoveryTarget{}, fmt.Errorf("load persisted discovery settings: %w", app_errors.ErrInternalServer) } snapshot, err := state.Compile(state.CompileInput{ - SystemSettings: systemSettings, + SystemSettings: systemSettings, ChannelRegistry: s.channelRegistry, Groups: []state.GroupConfig{{ - ID: rows.group.ID, Name: rows.group.Name, UpstreamURL: rows.group.UpstreamURL, - ProviderID: cloneString(rows.group.ProviderID), - Protocols: protocols, Settings: settings, Enabled: true, + ID: rows.group.ID, Name: rows.group.Name, + ChannelID: channel.ID(rows.group.ChannelID), Params: append(json.RawMessage(nil), rows.group.Params...), + Settings: settings, Enabled: true, }}, }) if err != nil { @@ -163,18 +171,27 @@ func (s *Service) mapGroupDiscoveryTarget( return discoveryTarget{}, fmt.Errorf("compiled persisted discovery Group is missing: %w", app_errors.ErrInternalServer) } - plaintextKeys := make([]string, 0, len(rows.keys)) - for _, keyRow := range rows.keys { - plaintext, err := s.encryption.Decrypt(keyRow.KeyValue) + discoveryCredentials := make([]discoveryCredential, 0, len(rows.credentials)) + for _, credentialRow := range rows.credentials { + canonical, apiKey, err := s.decodeCredential(rows.group, credentialRow) if err != nil { - return discoveryTarget{}, fmt.Errorf("decrypt persisted discovery key: %w", app_errors.ErrInternalServer) + return discoveryTarget{}, err } - plaintextKeys = append(plaintextKeys, plaintext) + discoveryCredentials = append(discoveryCredentials, discoveryCredential{ + snapshot: execution.NewCredentialSnapshot( + credentialRow.ID, + groupCollectionCredentialVersion(credentialRow.UpdatedAtMS), + groupCollectionCredentialIdentity(credentialRow.Fingerprint), + canonical, + ), + apiKey: apiKey, + }) } return discoveryTarget{ - baseURL: rows.group.UpstreamURL, protocols: protocols, - keys: plaintextKeys, headerRules: compiledGroup.HeaderRules, - providerID: cloneString(rows.group.ProviderID), + channelID: channel.ID(rows.group.ChannelID), + resolvedTarget: resolvedTarget, credentials: discoveryCredentials, + headerRules: compiledGroup.HeaderRules, timeouts: compiledGroup.Timeouts, + catalogProviderID: resolvedTarget.CatalogProviderID, }, nil } diff --git a/internal/control/discover_group_test.go b/internal/control/discover_group_test.go index 220597ae..c671f66c 100644 --- a/internal/control/discover_group_test.go +++ b/internal/control/discover_group_test.go @@ -12,7 +12,7 @@ import ( "gorm.io/gorm" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -24,18 +24,17 @@ func TestDiscoveryUsesSingleReadSnapshot(t *testing.T) { fixture, dsn := newFileServiceFixture(t) group := seedPersistedDiscoveryGroup(t, fixture, true, models.JSON(`{}`)) group.Name = "discovery-snapshot-old" - group.UpstreamURL = "https://discovery-old.example/v1" - group.Protocols = models.JSON(`["openai-completions"]`) + group.Params = models.JSON(`{"base_url":"https://discovery-old.example/v1"}`) if err := fixture.db.Save(&group).Error; err != nil { t.Fatalf("seed old Group version: %v", err) } - seedPersistedDiscoveryKey( + seedPersistedDiscoveryCredential( t, fixture, group.ID, 1, "discovery-key-old", - models.UpstreamKeyStatusActive, + models.CredentialStatusActive, ) if err := fixture.db.Create(&models.SystemSetting{ Key: "header_rules", Value: `{"set":{"X-Version":"old"}}`, @@ -56,7 +55,7 @@ func TestDiscoveryUsesSingleReadSnapshot(t *testing.T) { t.Errorf("close writer database: %v", err) } }) - newCiphertext, err := fixture.encryption.Encrypt("discovery-key-new") + newCiphertext, err := fixture.encryption.Encrypt(`{"api_key":"discovery-key-new"}`) if err != nil { t.Fatalf("encrypt new key version: %v", err) } @@ -86,7 +85,7 @@ func TestDiscoveryUsesSingleReadSnapshot(t *testing.T) { version string } observed := make(chan observedTarget, 1) - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func( _ context.Context, @@ -117,13 +116,13 @@ func TestDiscoveryUsesSingleReadSnapshot(t *testing.T) { writeDone <- writer.Transaction(func(tx *gorm.DB) error { if err := tx.Model(&models.Group{}).Where("id = ?", group.ID). Updates(map[string]any{ - "name": "discovery-snapshot-new", - "upstream_url": "https://discovery-new.example/v1", + "name": "discovery-snapshot-new", + "params": models.JSON(`{"base_url":"https://discovery-new.example/v1"}`), }).Error; err != nil { return err } - if err := tx.Model(&models.UpstreamKey{}).Where("group_id = ?", group.ID). - Update("key_value", newCiphertext).Error; err != nil { + if err := tx.Model(&models.Credential{}).Where("group_id = ?", group.ID). + Update("data", newCiphertext).Error; err != nil { return err } return tx.Model(&models.SystemSetting{}).Where("key = ?", "header_rules"). @@ -172,17 +171,13 @@ func TestDiscoveryUsesSingleReadSnapshot(t *testing.T) { func TestDiscoveryReleasesReadSnapshotBeforeDecrypt(t *testing.T) { fixture, _ := newFileServiceFixture(t) group := seedPersistedDiscoveryGroup(t, fixture, true, models.JSON(`{}`)) - group.Protocols = models.JSON(`["openai-completions"]`) - if err := fixture.db.Save(&group).Error; err != nil { - t.Fatal(err) - } - seedPersistedDiscoveryKey( + seedPersistedDiscoveryCredential( t, fixture, group.ID, 1, "release-snapshot-key", - models.UpstreamKeyStatusActive, + models.CredentialStatusActive, ) decryptStarted := make(chan struct{}) releaseDecrypt := make(chan struct{}) @@ -191,7 +186,7 @@ func TestDiscoveryReleasesReadSnapshotBeforeDecrypt(t *testing.T) { started: decryptStarted, release: releaseDecrypt, } - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { return []string{"model"}, nil @@ -234,14 +229,14 @@ func TestDiscoveryReleasesReadSnapshotBeforeDecrypt(t *testing.T) { } } -func TestDiscoverGroupModelsUsesDisabledGroupAndActiveKeysInIDOrder(t *testing.T) { +func TestDiscoverGroupModelsUsesDisabledGroupAndActiveCredentialsInIDOrder(t *testing.T) { fixture := newServiceFixture(t) group := seedPersistedDiscoveryGroup(t, fixture, false, models.JSON( `{"header_rules":{"set":{"X-Group":"group"},"remove":["X-Remove"]}}`, )) - seedPersistedDiscoveryKey(t, fixture, group.ID, 1, "key-1", models.UpstreamKeyStatusActive) - seedPersistedDiscoveryKey(t, fixture, group.ID, 2, "key-2", models.UpstreamKeyStatusDisabled) - seedPersistedDiscoveryKey(t, fixture, group.ID, 3, "key-3", models.UpstreamKeyStatusActive) + seedPersistedDiscoveryCredential(t, fixture, group.ID, 1, "key-1", models.CredentialStatusActive) + seedPersistedDiscoveryCredential(t, fixture, group.ID, 2, "key-2", models.CredentialStatusDisabled) + seedPersistedDiscoveryCredential(t, fixture, group.ID, 3, "key-3", models.CredentialStatusActive) if err := fixture.db.Create(&models.SystemSetting{ Key: "header_rules", Value: `{"set":{"X-System":"system"}}`, }).Error; err != nil { @@ -249,8 +244,8 @@ func TestDiscoverGroupModelsUsesDisabledGroupAndActiveKeysInIDOrder(t *testing.T } var calls []string - newRecorder := func(value protocol.Protocol) *recordingDiscoveryDialect { - return &recordingDiscoveryDialect{ + newRecorder := func(value protocol.Protocol) *recordingDiscoveryExecutorTarget { + return &recordingDiscoveryExecutorTarget{ value: value, listFn: func( _ context.Context, @@ -258,24 +253,23 @@ func TestDiscoverGroupModelsUsesDisabledGroupAndActiveKeysInIDOrder(t *testing.T rules state.HeaderRules, ) ([]string, error) { calls = append(calls, string(value)+":"+apiKey) - if baseURL != group.UpstreamURL { - t.Fatalf("base URL = %q, want %q", baseURL, group.UpstreamURL) + if baseURL != "https://persisted.example.com/v1" { + t.Fatalf("base URL = %q, want persisted channel target", baseURL) } wantRules := state.HeaderRules{ - Set: map[string]string{"X-Group": "group"}, Remove: []string{"X-Remove"}, + Set: map[string]string{"X-Group": "group"}, } if !reflect.DeepEqual(rules, wantRules) { t.Fatalf("HeaderRules = %#v, want persisted Group override %#v", rules, wantRules) } - if value == protocol.Anthropic && apiKey == "key-1" { + if value == protocol.OpenAICompletions && apiKey == "key-3" { return []string{"z-model", "a-model"}, nil } return nil, errors.New("try next candidate") }, } } - fixture.service.dialects = dialect.NewSet( - newRecorder(protocol.Anthropic), + fixture.service.executor = newRecordingDiscoveryExecutor( newRecorder(protocol.OpenAICompletions), ) @@ -292,7 +286,6 @@ func TestDiscoverGroupModelsUsesDisabledGroupAndActiveKeysInIDOrder(t *testing.T wantCalls := []string{ "openai-completions:key-1", "openai-completions:key-3", - "anthropic:key-1", } if !reflect.DeepEqual(calls, wantCalls) { t.Fatalf("calls = %#v, want protocol-outer active-key-ID-inner order %#v", calls, wantCalls) @@ -311,10 +304,10 @@ func TestDiscoverGroupModelsReturnsNotFoundAndNoActiveUpstreamKey(t *testing.T) t.Run("no active upstream key", func(t *testing.T) { fixture := newServiceFixture(t) group := seedPersistedDiscoveryGroup(t, fixture, true, models.JSON(`{}`)) - seedPersistedDiscoveryKey(t, fixture, group.ID, 1, "disabled", models.UpstreamKeyStatusDisabled) + seedPersistedDiscoveryCredential(t, fixture, group.ID, 1, "disabled", models.CredentialStatusDisabled) _, err := fixture.service.DiscoverGroupModels(t.Context(), group.ID) - if !errors.Is(err, app_errors.ErrNoActiveUpstreamKey) { - t.Fatalf("DiscoverGroupModels() error = %v, want ErrNoActiveUpstreamKey", err) + if !errors.Is(err, app_errors.ErrNoActiveCredential) { + t.Fatalf("DiscoverGroupModels() error = %v, want ErrNoActiveCredential", err) } }) } @@ -322,16 +315,16 @@ func TestDiscoverGroupModelsReturnsNotFoundAndNoActiveUpstreamKey(t *testing.T) func TestDiscoverGroupModelsDecryptsEveryKeyBeforeHTTP(t *testing.T) { fixture := newServiceFixture(t) group := seedPersistedDiscoveryGroup(t, fixture, true, models.JSON(`{}`)) - seedPersistedDiscoveryKey(t, fixture, group.ID, 1, "key-1", models.UpstreamKeyStatusActive) - if err := fixture.db.Create(&models.UpstreamKey{ - ID: 3, GroupID: group.ID, KeyValue: "corrupt-second-active-ciphertext", - KeyHash: "corrupt-hash", Status: models.UpstreamKeyStatusActive, + seedPersistedDiscoveryCredential(t, fixture, group.ID, 1, "key-1", models.CredentialStatusActive) + if err := fixture.db.Create(&models.Credential{ + ID: 3, GroupID: group.ID, Data: "corrupt-second-active-ciphertext", + Fingerprint: "corrupt-hash", Status: models.CredentialStatusActive, }).Error; err != nil { t.Fatalf("seed corrupt active key: %v", err) } calls := 0 - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.Anthropic, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { calls++ @@ -385,9 +378,9 @@ func TestDiscoverGroupModelsDoesNotMutateDatabaseSnapshotOrRegistry(t *testing.T t.Run(test.name, func(t *testing.T) { fixture := newServiceFixture(t) created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://state.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "state-key", + ChannelID: channel.OpenAICompatible, + Params: []byte(`{"base_url":"https://state.example.com/v1"}`), + Credentials: "state-key", Models: optionalGroupModels{ Set: true, Values: []GroupModel{{ID: "persisted-model"}}, }, @@ -398,7 +391,7 @@ func TestDiscoverGroupModelsDoesNotMutateDatabaseSnapshotOrRegistry(t *testing.T before := captureGroupDiscoveryState(t, fixture, created.GroupID) ctx, cancel := context.WithCancel(t.Context()) defer cancel() - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(ctx context.Context, _, _ string, _ state.HeaderRules) ([]string, error) { return test.listFn(ctx, cancel) @@ -423,16 +416,16 @@ func TestDiscoverGroupModelsDoesNotMutateDatabaseSnapshotOrRegistry(t *testing.T func TestDiscoverGroupModelsDoesNotAcquireWriteMuOrBlockWrites(t *testing.T) { fixture := newServiceFixture(t) created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://discovery-lock.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: []byte(`{"base_url":"https://discovery-lock.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "key-1", + Credentials: "key-1", }) if err != nil { t.Fatalf("seed CreateGroup() error = %v", err) } fixture.service.modelDiscoveryTimeout = 3 * time.Second - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { return []string{"model"}, nil @@ -458,7 +451,7 @@ func TestDiscoverGroupModelsDoesNotAcquireWriteMuOrBlockWrites(t *testing.T) { entered := make(chan struct{}) release := make(chan struct{}) - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(ctx context.Context, _, _ string, _ state.HeaderRules) ([]string, error) { close(entered) @@ -484,10 +477,10 @@ func TestDiscoverGroupModelsDoesNotAcquireWriteMuOrBlockWrites(t *testing.T) { groupWriteDone := make(chan error, 1) go func() { _, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://concurrent-group-write.example.com/v1", - Protocols: []protocol.Protocol{protocol.Anthropic}, + ChannelID: channel.AnthropicCompatible, + Params: []byte(`{"base_url":"https://concurrent-group-write.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "concurrent-group-key", + Credentials: "concurrent-group-key", }) groupWriteDone <- err }() @@ -502,8 +495,8 @@ func TestDiscoverGroupModelsDoesNotAcquireWriteMuOrBlockWrites(t *testing.T) { keyWriteDone := make(chan error, 1) go func() { - _, err := fixture.service.ImportGroupKeys(t.Context(), created.GroupID, GroupKeyImportRequest{ - Keys: "concurrent-imported-key", + _, err := fixture.service.ImportGroupCredentials(t.Context(), created.GroupID, CredentialImportRequest{ + Credentials: "concurrent-imported-key", }) keyWriteDone <- err }() @@ -533,7 +526,7 @@ type groupDiscoveryState struct { config string snapshot *state.ConfigSnapshot snapshotRevision uint64 - registryKeys []state.KeyMeta + registryKeys []state.CredentialMeta registryValues map[uint]string } @@ -543,24 +536,24 @@ func captureGroupDiscoveryState(t *testing.T, fixture serviceFixture, groupID ui if err := fixture.db.First(&group, groupID).Error; err != nil { t.Fatalf("load Group state: %v", err) } - var keyRows []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&keyRows).Error; err != nil { - t.Fatalf("load key state: %v", err) + var credentialRows []models.Credential + if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&credentialRows).Error; err != nil { + t.Fatalf("load credential state: %v", err) } - registryValues := make(map[uint]string, len(keyRows)) - for _, keyRow := range keyRows { - if ciphertext, ok := fixture.registry.EncryptedValue(keyRow.ID); ok { - registryValues[keyRow.ID] = ciphertext + registryValues := make(map[uint]string, len(credentialRows)) + for _, credentialRow := range credentialRows { + if ciphertext, ok := fixture.registry.EncryptedCredentialData(credentialRow.ID); ok { + registryValues[credentialRow.ID] = ciphertext } } snapshot := fixture.manager.Current() return groupDiscoveryState{ rowCounts: discoveryRowCounts(t, fixture.db), models: string(group.Models), - config: string(group.Config), + config: string(group.Overrides), snapshot: snapshot, snapshotRevision: snapshot.Revision, - registryKeys: fixture.registry.CollectCandidates([]uint{groupID}, nil, time.Time{}), + registryKeys: fixture.registry.CollectCredentialCandidates([]uint{groupID}, nil, time.Time{}), registryValues: registryValues, } } @@ -569,13 +562,14 @@ func seedPersistedDiscoveryGroup( t *testing.T, fixture serviceFixture, enabled bool, - groupConfig models.JSON, + overrides models.JSON, ) models.Group { t.Helper() group := models.Group{ - Name: "persisted-discovery", UpstreamURL: "https://persisted.example.com/v1", - Protocols: models.JSON(`["anthropic","openai-completions"]`), - Models: models.JSON(`[{"id":"persisted-only"}]`), Config: groupConfig, Enabled: enabled, + Name: "persisted-discovery", ChannelID: string(channel.OpenAICompatible), + Params: models.JSON(`{"base_url":"https://persisted.example.com/v1"}`), + Models: models.JSON(`[{"id":"persisted-only"}]`), + Overrides: overrides, Enabled: enabled, } if err := fixture.db.Create(&group).Error; err != nil { t.Fatalf("seed persisted discovery Group: %v", err) @@ -583,22 +577,23 @@ func seedPersistedDiscoveryGroup( return group } -func seedPersistedDiscoveryKey( +func seedPersistedDiscoveryCredential( t *testing.T, fixture serviceFixture, - groupID, keyID uint, + groupID, credentialID uint, plaintext string, - status models.UpstreamKeyStatus, + status models.CredentialStatus, ) { t.Helper() - ciphertext, err := fixture.encryption.Encrypt(plaintext) + canonical := `{"api_key":` + fmt.Sprintf("%q", plaintext) + `}` + ciphertext, err := fixture.encryption.Encrypt(canonical) if err != nil { - t.Fatalf("encrypt seeded discovery key: %v", err) + t.Fatalf("encrypt seeded discovery credential: %v", err) } - if err := fixture.db.Create(&models.UpstreamKey{ - ID: keyID, GroupID: groupID, KeyValue: ciphertext, - KeyHash: fixture.encryption.Hash(plaintext), Status: status, + if err := fixture.db.Create(&models.Credential{ + ID: credentialID, GroupID: groupID, Data: ciphertext, + Fingerprint: fixture.encryption.Hash(canonical), Status: status, }).Error; err != nil { - t.Fatalf("seed persisted discovery key: %v", err) + t.Fatalf("seed persisted discovery credential: %v", err) } } diff --git a/internal/control/discover_test.go b/internal/control/discover_test.go index e38f7af9..603b4eda 100644 --- a/internal/control/discover_test.go +++ b/internal/control/discover_test.go @@ -2,7 +2,9 @@ package control import ( "context" + "encoding/json" "errors" + "net/http" "reflect" "sync/atomic" "testing" @@ -10,7 +12,8 @@ import ( "gorm.io/gorm" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" + "gpt-load/internal/execution" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -28,8 +31,8 @@ func TestDiscoverModelsUsesSystemDefaultsAndNormalizesSuccessfulResult(t *testin } var calls []string - newRecorder := func(value protocol.Protocol) *recordingDiscoveryDialect { - return &recordingDiscoveryDialect{ + newRecorder := func(value protocol.Protocol) *recordingDiscoveryExecutorTarget { + return &recordingDiscoveryExecutorTarget{ value: value, listFn: func( _ context.Context, @@ -45,31 +48,24 @@ func TestDiscoverModelsUsesSystemDefaultsAndNormalizesSuccessfulResult(t *testin "X-System": "system", "X-Override": "system", }, - Remove: []string{"X-System-Remove"}, } if !reflect.DeepEqual(rules, wantRules) { t.Fatalf("HeaderRules = %#v, want system defaults %#v", rules, wantRules) } - if value == protocol.Anthropic && apiKey == "key-a" { + if value == protocol.OpenAICompletions && apiKey == "key-b" { return []string{" claude-z ", "claude-a", "claude-z", "", "claude-a"}, nil } return nil, errors.New("try next") }, } } - fixture.service.dialects = dialect.NewSet( + fixture.service.executor = newRecordingDiscoveryExecutor( newRecorder(protocol.OpenAICompletions), - newRecorder(protocol.Anthropic), ) result, err := fixture.service.DiscoverModels(context.Background(), ModelDiscoveryRequest{ - UpstreamURL: " HTTPS://API.Example.COM/v1/?fixed=1 ", - Protocols: []protocol.Protocol{ - protocol.OpenAICompletions, - protocol.OpenAICompletions, - protocol.Anthropic, - protocol.OpenAICompletions, - }, - Keys: " key-a \nkey-a\n\n key-b \nkey-b", + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":" HTTPS://API.Example.COM/v1/?fixed=1 "}`), + Credentials: " key-a \nkey-a\n\n key-b \nkey-b", }) if err != nil { t.Fatalf("DiscoverModels() error = %v", err) @@ -83,7 +79,6 @@ func TestDiscoverModelsUsesSystemDefaultsAndNormalizesSuccessfulResult(t *testin wantCalls := []string{ "openai-completions:key-a", "openai-completions:key-b", - "anthropic:key-a", } if !reflect.DeepEqual(calls, wantCalls) { t.Fatalf("calls = %#v, want stable normalized order %#v", calls, wantCalls) @@ -93,7 +88,7 @@ func TestDiscoverModelsUsesSystemDefaultsAndNormalizesSuccessfulResult(t *testin func TestDiscoverModelsRejectsInvalidDraftBeforeHTTP(t *testing.T) { fixture := newServiceFixture(t) var calls atomic.Int64 - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { calls.Add(1) @@ -101,21 +96,23 @@ func TestDiscoverModelsRejectsInvalidDraftBeforeHTTP(t *testing.T) { }, }) valid := ModelDiscoveryRequest{ - UpstreamURL: "https://api.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "key-a", + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://api.example.com"}`), + Credentials: "key-a", } tests := []struct { name string mutate func(*ModelDiscoveryRequest) }{ - {name: "empty URL", mutate: func(value *ModelDiscoveryRequest) { value.UpstreamURL = "" }}, - {name: "relative URL", mutate: func(value *ModelDiscoveryRequest) { value.UpstreamURL = "/v1" }}, - {name: "empty protocols", mutate: func(value *ModelDiscoveryRequest) { value.Protocols = nil }}, - {name: "unknown protocol", mutate: func(value *ModelDiscoveryRequest) { - value.Protocols = []protocol.Protocol{"unknown"} + {name: "empty channel", mutate: func(value *ModelDiscoveryRequest) { value.ChannelID = "" }}, + {name: "relative URL", mutate: func(value *ModelDiscoveryRequest) { + value.Params = json.RawMessage(`{"base_url":"/v1"}`) }}, - {name: "empty keys", mutate: func(value *ModelDiscoveryRequest) { value.Keys = " \n\t" }}, + {name: "missing base URL", mutate: func(value *ModelDiscoveryRequest) { value.Params = json.RawMessage(`{}`) }}, + {name: "unknown channel", mutate: func(value *ModelDiscoveryRequest) { + value.ChannelID = channel.ID("unknown") + }}, + {name: "empty credentials", mutate: func(value *ModelDiscoveryRequest) { value.Credentials = " \n\t" }}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { @@ -132,11 +129,11 @@ func TestDiscoverModelsRejectsInvalidDraftBeforeHTTP(t *testing.T) { } } -func TestDiscoverModelsSupportsResponsesOnly(t *testing.T) { +func TestDiscoverModelsUsesChannelNativeDiscoveryProtocol(t *testing.T) { fixture := newServiceFixture(t) calls := 0 - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ - value: protocol.OpenAIResponses, + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ + value: protocol.OpenAICompletions, listFn: func( _ context.Context, baseURL, apiKey string, @@ -154,9 +151,9 @@ func TestDiscoverModelsSupportsResponsesOnly(t *testing.T) { result, err := fixture.service.DiscoverModels( context.Background(), ModelDiscoveryRequest{ - UpstreamURL: "https://api.example.com", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Keys: "key-a", + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://api.example.com"}`), + Credentials: "key-a", }, ) if err != nil { @@ -170,12 +167,51 @@ func TestDiscoverModelsSupportsResponsesOnly(t *testing.T) { } } +func TestDiscoverModelsPassesStructuredCloudCredentialToExecutor(t *testing.T) { + fixture := newServiceFixture(t) + var observed execution.AttemptSpec + fixture.service.executor = scriptedDiscoveryExecutor{execute: func( + _ context.Context, + spec execution.AttemptSpec, + ) execution.AttemptResult { + observed = spec.Clone() + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: []byte(`{"object":"list","data":[{"id":"anthropic.claude-test"}]}`), + } + }} + + result, err := fixture.service.DiscoverModels(context.Background(), ModelDiscoveryRequest{ + ChannelID: channel.AWSBedrock, + Params: json.RawMessage(`{"region":"us-east-1"}`), + Credentials: `{"access_key":"AKIA_TEST","secret_key":"bedrock-secret",` + + `"session_token":"bedrock-session"}`, + }) + if err != nil { + t.Fatalf("DiscoverModels() error = %v", err) + } + if got, want := string(observed.Credential.Data()), + `{"access_key":"AKIA_TEST","secret_key":"bedrock-secret","session_token":"bedrock-session"}`; got != want { + t.Fatalf("credential = %s, want %s", got, want) + } + if observed.ChannelID != string(channel.AWSBedrock) || + observed.Operation != execution.OperationListModels || + observed.ClientProtocol != protocol.OpenAICompletions { + t.Fatalf("attempt = %#v", observed) + } + if got := result.Models; len(got) != 1 || got[0].ID != "anthropic.claude-test" { + t.Fatalf("models = %#v", got) + } +} + func TestDiscoverModelsDoesNotReadOrMutateRuntimeState(t *testing.T) { fixture := newServiceFixture(t) created, err := fixture.service.CreateGroup(context.Background(), GroupCreateRequest{ - UpstreamURL: "https://state.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "sk-state", + Name: stringPointer("discover-runtime-state"), ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://state.example.com"}`), Credentials: "sk-state", Models: optionalGroupModels{ Set: true, Values: []GroupModel{{ID: "gpt-4o"}}, }, @@ -185,11 +221,11 @@ func TestDiscoverModelsDoesNotReadOrMutateRuntimeState(t *testing.T) { } beforeRows := discoveryRowCounts(t, fixture.db) beforeSnapshot := fixture.manager.Current() - var keyRow models.UpstreamKey + var keyRow models.Credential if err := fixture.db.First(&keyRow).Error; err != nil { t.Fatalf("query seeded key: %v", err) } - beforeCipher, ok := fixture.registry.EncryptedValue(keyRow.ID) + beforeCipher, ok := fixture.registry.EncryptedCredentialData(keyRow.ID) if !ok { t.Fatal("seeded Registry key missing") } @@ -226,16 +262,16 @@ func TestDiscoverModelsDoesNotReadOrMutateRuntimeState(t *testing.T) { } fixture.service.manager = nil fixture.service.registry = nil - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { return []string{"remote-only"}, nil }, }) result, err := fixture.service.DiscoverModels(context.Background(), ModelDiscoveryRequest{ - UpstreamURL: "https://discover.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "sk-discovery", + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://discover.example.com"}`), + Credentials: "sk-discovery", }) if err != nil || !reflect.DeepEqual(result.Models, []ModelCandidate{ {ID: "remote-only", Name: "remote-only", Sources: []string{"live"}, PricingStatus: PricingStatusPending}, @@ -252,13 +288,13 @@ func TestDiscoverModelsDoesNotReadOrMutateRuntimeState(t *testing.T) { fixture.manager.Current().Revision != beforeSnapshot.Revision { t.Fatal("discovery replaced or revised ConfigSnapshot") } - if got, ok := fixture.registry.EncryptedValue(keyRow.ID); !ok || got != beforeCipher { + if got, ok := fixture.registry.EncryptedCredentialData(keyRow.ID); !ok || got != beforeCipher { t.Fatalf("Registry value = %q, %t, want unchanged", got, ok) } if afterRows := discoveryRowCounts(t, fixture.db); afterRows != beforeRows { t.Fatalf("row counts = %#v, want %#v", afterRows, beforeRows) } - if _, exists := fixture.manager.Current().Candidates[protocol.OpenAICompletions]["remote-only"]; exists { + if _, exists := fixture.manager.Current().ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion]["remote-only"]; exists { t.Fatal("discovered model leaked into ConfigSnapshot") } if created.GroupID == 0 { @@ -268,7 +304,7 @@ func TestDiscoverModelsDoesNotReadOrMutateRuntimeState(t *testing.T) { func TestDiscoverModelsDoesNotAcquireWriteMu(t *testing.T) { fixture := newServiceFixture(t) - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { return []string{"gpt-4o"}, nil @@ -282,9 +318,9 @@ func TestDiscoverModelsDoesNotAcquireWriteMu(t *testing.T) { done := make(chan error, 1) go func() { _, err := fixture.service.DiscoverModels(ctx, ModelDiscoveryRequest{ - UpstreamURL: "https://discover.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "sk-discovery", + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://discover.example.com"}`), + Credentials: "sk-discovery", }) done <- err }() @@ -302,7 +338,7 @@ func TestDiscoverModelsDoesNotBlockMutation(t *testing.T) { entered := make(chan struct{}) release := make(chan struct{}) fixture := newServiceFixture(t) - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.Anthropic, listFn: func( ctx context.Context, @@ -322,9 +358,9 @@ func TestDiscoverModelsDoesNotBlockMutation(t *testing.T) { discoveryDone := make(chan error, 1) go func() { _, err := fixture.service.DiscoverModels(context.Background(), ModelDiscoveryRequest{ - UpstreamURL: "https://discover.example.com", - Protocols: []protocol.Protocol{protocol.Anthropic}, - Keys: "sk-discovery", + ChannelID: channel.AnthropicCompatible, + Params: json.RawMessage(`{"base_url":"https://discover.example.com"}`), + Credentials: "sk-discovery", }) discoveryDone <- err }() @@ -337,9 +373,8 @@ func TestDiscoverModelsDoesNotBlockMutation(t *testing.T) { mutationDone := make(chan error, 1) go func() { _, err := fixture.service.CreateGroup(context.Background(), GroupCreateRequest{ - UpstreamURL: "https://mutation.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "sk-mutation", + Name: stringPointer("discover-concurrent-mutation"), ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://mutation.example.com"}`), Credentials: "sk-mutation", Models: optionalGroupModels{ Set: true, Values: []GroupModel{{ID: "gpt-4o"}}, }, @@ -368,7 +403,7 @@ func TestDiscoverModelsDoesNotBlockMutation(t *testing.T) { func discoveryRowCounts(t *testing.T, db *gorm.DB) [3]int64 { t.Helper() var result [3]int64 - for index, model := range []any{&models.Group{}, &models.UpstreamKey{}, &models.AccessKey{}} { + for index, model := range []any{&models.Group{}, &models.Credential{}, &models.AccessKey{}} { if err := db.Model(model).Count(&result[index]).Error; err != nil { t.Fatalf("count %T rows: %v", model, err) } diff --git a/internal/control/execution_test_helpers_test.go b/internal/control/execution_test_helpers_test.go new file mode 100644 index 00000000..0aac372e --- /dev/null +++ b/internal/control/execution_test_helpers_test.go @@ -0,0 +1,144 @@ +package control + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "strings" + + "gpt-load/internal/execution" +) + +// controlHTTPExecutor is a narrow test double for control/data-plane publication +// tests. Provider protocol behavior is covered by the production Bifrost adapter +// tests; this double only proves that a frozen channel target and credential are +// visible atomically to the gateway. +type controlHTTPExecutor struct { + client *http.Client +} + +func (controlHTTPExecutor) Capabilities() execution.CapabilitySet { + return execution.CapabilitySet{} +} + +func (executor controlHTTPExecutor) Execute( + ctx context.Context, + spec execution.AttemptSpec, +) execution.AttemptResult { + response, err := executor.do(ctx, spec) + if err != nil { + return controlExecutionFailure(execution.ErrorKindTransport) + } + defer response.Body.Close() + body, err := io.ReadAll(response.Body) + if err != nil { + return controlExecutionFailure(execution.ErrorKindTransport) + } + result := execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: response.StatusCode, + Header: response.Header.Clone(), + Body: body, + Model: spec.UpstreamModel, + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + result.Error = &execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, StatusCode: response.StatusCode, + Summary: "test upstream rejected request", + } + } + return result +} + +func (executor controlHTTPExecutor) ExecuteStream( + ctx context.Context, + spec execution.AttemptSpec, + sink execution.StreamSink, +) execution.StreamResult { + response, err := executor.do(ctx, spec) + if err != nil { + return controlStreamExecutionFailure(execution.ErrorKindTransport) + } + defer response.Body.Close() + body, err := io.ReadAll(response.Body) + if err != nil { + return controlStreamExecutionFailure(execution.ErrorKindTransport) + } + if err := sink(execution.StreamEvent{ + Sequence: 1, Kind: execution.StreamEventReady, + StatusCode: response.StatusCode, Header: response.Header.Clone(), + }); err != nil { + return controlStreamExecutionFailure(execution.ErrorKindCanceled) + } + if len(body) > 0 { + if err := sink(execution.StreamEvent{ + Sequence: 2, Kind: execution.StreamEventData, Data: body, + }); err != nil { + return controlStreamExecutionFailure(execution.ErrorKindCanceled) + } + } + result := execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: response.StatusCode, + Header: response.Header.Clone(), + Model: spec.UpstreamModel, + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + result.Error = &execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, StatusCode: response.StatusCode, + Summary: "test upstream rejected request", + } + } + return result +} + +func (executor controlHTTPExecutor) do( + ctx context.Context, + spec execution.AttemptSpec, +) (*http.Response, error) { + var target struct { + BaseURL string `json:"base_url"` + } + if err := json.Unmarshal(spec.TargetConfig, &target); err != nil { + return nil, err + } + var credential struct { + APIKey string `json:"api_key"` + } + if err := json.Unmarshal(spec.Credential.Data(), &credential); err != nil { + return nil, err + } + requestURL := strings.TrimSuffix(target.BaseURL, "/") + "/chat/completions" + request, err := http.NewRequestWithContext(ctx, spec.Method, requestURL, bytes.NewReader(spec.Body)) + if err != nil { + return nil, err + } + request.Header = spec.Header.Clone() + if request.Header == nil { + request.Header = make(http.Header) + } + request.Header.Set("Authorization", "Bearer "+credential.APIKey) + client := executor.client + if client == nil { + client = http.DefaultClient + } + return client.Do(request) +} + +func controlExecutionFailure(kind execution.ErrorKind) execution.AttemptResult { + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + Error: &execution.ErrorEvidence{Kind: kind, Summary: "test upstream execution failed"}, + } +} + +func controlStreamExecutionFailure(kind execution.ErrorKind) execution.StreamResult { + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + Error: &execution.ErrorEvidence{Kind: kind, Summary: "test upstream execution failed"}, + } +} diff --git a/internal/control/external_database_integration_test.go b/internal/control/external_database_integration_test.go index af95d837..8359f8f2 100644 --- a/internal/control/external_database_integration_test.go +++ b/internal/control/external_database_integration_test.go @@ -1,13 +1,14 @@ package control import ( + "encoding/json" "fmt" "os" "strings" "testing" "time" - "gpt-load/internal/protocol" + "gpt-load/internal/channel" "gpt-load/internal/storage/models" ) @@ -26,13 +27,13 @@ func TestExternalDatabaseGroupPriceReconciliation(t *testing.T) { modelID := fmt.Sprintf("external-control-model-%d", suffix) create := func(name, upstreamURL string) GroupCreateResult { result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - Name: &name, - UpstreamURL: upstreamURL, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + Name: &name, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"` + upstreamURL + `"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{{ ID: modelID, }}}, - Keys: fmt.Sprintf("sk-external-control-%d", suffix), + Credentials: fmt.Sprintf("sk-external-control-%d", suffix), }) if err != nil { t.Fatalf("CreateGroup(%q) error = %v", name, err) diff --git a/internal/control/group_collection.go b/internal/control/group_collection.go index 805c1b86..66264c5c 100644 --- a/internal/control/group_collection.go +++ b/internal/control/group_collection.go @@ -4,11 +4,14 @@ import ( "context" "encoding/json" "fmt" + "hash/fnv" "strings" "time" "gorm.io/gorm" + "gpt-load/internal/channel" + "gpt-load/internal/execution" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -23,7 +26,7 @@ const ( GroupCollectionStatusDisabled GroupCollectionStatus = "disabled" ) -type GroupCollectionKeyCounts struct { +type GroupCollectionCredentialCounts struct { Total int64 `json:"total"` Available int64 `json:"available"` Cooldown int64 `json:"cooldown"` @@ -32,14 +35,13 @@ type GroupCollectionKeyCounts struct { } type GroupCollectionItem struct { - ID uint `json:"id"` - Name string `json:"name"` - ProviderID *string `json:"provider_id"` - Status GroupCollectionStatus `json:"status"` - UpstreamURL string `json:"upstream_url"` - Protocols []protocol.Protocol `json:"protocols"` - ModelCount int64 `json:"model_count"` - KeyCounts GroupCollectionKeyCounts `json:"key_counts"` + ID uint `json:"id"` + Name string `json:"name"` + ChannelID channel.ID `json:"channel_id"` + Params json.RawMessage `json:"params"` + Status GroupCollectionStatus `json:"status"` + ModelCount int64 `json:"model_count"` + CredentialCounts GroupCollectionCredentialCounts `json:"credential_counts"` } type groupCollectionRecord struct { @@ -48,30 +50,26 @@ type groupCollectionRecord struct { } type groupCollectionRows struct { - groups []models.Group - keys []models.UpstreamKey + groups []models.Group + credentials []models.Credential } func cloneGroupRows(rows []models.Group) []models.Group { cloned := make([]models.Group, len(rows)) for index := range rows { cloned[index] = rows[index] - cloned[index].Protocols = append(models.JSON(nil), rows[index].Protocols...) + cloned[index].Params = append(models.JSON(nil), rows[index].Params...) cloned[index].Models = append(models.JSON(nil), rows[index].Models...) - cloned[index].Config = append(models.JSON(nil), rows[index].Config...) + cloned[index].Overrides = append(models.JSON(nil), rows[index].Overrides...) if rows[index].WeightManual != nil { value := *rows[index].WeightManual cloned[index].WeightManual = &value } - if rows[index].ProviderID != nil { - value := *rows[index].ProviderID - cloned[index].ProviderID = &value - } if rows[index].ValidationModel != nil { value := *rows[index].ValidationModel cloned[index].ValidationModel = &value } - cloned[index].UpstreamKeys = nil + cloned[index].Credentials = nil } return cloned } @@ -109,7 +107,7 @@ func (s *Service) captureGroupCollectionRecords( if err != nil { return 0, nil, app_errors.ParseDBError(err) } - records, err := mapGroupCollectionRecords(snapshot, runtimeKeys, rows, observedAt) + records, err := mapGroupCollectionRecords(snapshot, runtimeKeys, rows, observedAt, s.channelRegistry) if parentErr := ctx.Err(); parentErr != nil { return 0, nil, parentErr } @@ -128,14 +126,14 @@ func (s *Service) readGroupCollectionRows( if err := tx.Order("id ASC").Find(&groups).Error; err != nil { return err } - var keys []models.UpstreamKey - if err := tx.Model(&models.UpstreamKey{}). - Select("id", "group_id", "status", "weight_manual"). - Order("group_id ASC, id ASC").Find(&keys).Error; err != nil { + rows.groups = cloneGroupRows(groups) + var credentials []models.Credential + if err := tx.Model(&models.Credential{}). + Select("id", "group_id", "fingerprint", "status", "weight_manual", "updated_at_ms"). + Order("group_id ASC, id ASC").Find(&credentials).Error; err != nil { return err } - rows.groups = cloneGroupRows(groups) - rows.keys = cloneUpstreamKeyRows(keys) + rows.credentials = cloneCredentialRows(credentials) return nil }) return rows, err @@ -143,9 +141,10 @@ func (s *Service) readGroupCollectionRows( func mapGroupCollectionRecords( snapshot *state.ConfigSnapshot, - runtimeKeys []state.KeyRuntimeView, + runtimeKeys []state.CredentialRuntimeView, rows groupCollectionRows, observedAt time.Time, + registries ...*channel.Registry, ) ([]groupCollectionRecord, error) { if snapshot == nil { return nil, groupCollectionDataError("runtime snapshot is nil") @@ -195,7 +194,7 @@ func mapGroupCollectionRecords( } } - runtimeByID := make(map[uint]state.KeyRuntimeView, len(runtimeKeys)) + runtimeByID := make(map[uint]state.CredentialRuntimeView, len(runtimeKeys)) for _, key := range runtimeKeys { if key.ID == 0 { return nil, groupCollectionDataError("runtime key has zero id") @@ -205,68 +204,81 @@ func mapGroupCollectionRecords( } runtimeByID[key.ID] = key } - persistedByID := make(map[uint]models.UpstreamKey, len(rows.keys)) - keysByGroup := make(map[uint][]models.UpstreamKey) - for _, key := range rows.keys { - if key.ID == 0 { - return nil, groupCollectionDataError("persisted key has zero id") + var registry *channel.Registry + for _, candidate := range registries { + if candidate != nil { + registry = candidate + break } - if _, duplicate := persistedByID[key.ID]; duplicate { - return nil, groupCollectionDataError("duplicate persisted key %d", key.ID) - } - if _, exists := persistedGroups[key.GroupID]; !exists { - return nil, groupCollectionDataError( - "persisted key %d references missing group %d", - key.ID, - key.GroupID, - ) - } - persistedByID[key.ID] = key - keysByGroup[key.GroupID] = append(keysByGroup[key.GroupID], key) } - if len(persistedByID) != len(runtimeByID) { - return nil, groupCollectionDataError("persisted and runtime key sets differ") - } - for keyID, persistedKey := range persistedByID { - runtimeKey, exists := runtimeByID[keyID] + credentialsByGroup := make(map[uint][]models.Credential) + persistedCredentialByID := make(map[uint]models.Credential, len(rows.credentials)) + for _, credential := range rows.credentials { + if credential.ID == 0 { + return nil, groupCollectionDataError("persisted credential has zero id") + } + if _, duplicate := persistedCredentialByID[credential.ID]; duplicate { + return nil, groupCollectionDataError("duplicate persisted credential %d", credential.ID) + } + _, exists := persistedGroups[credential.GroupID] if !exists { return nil, groupCollectionDataError( - "persisted key %d is missing from runtime registry", - keyID, + "persisted credential %d references missing group %d", + credential.ID, + credential.GroupID, ) } - status, err := groupCollectionRuntimeKeyStatus(persistedKey.Status) + persistedCredentialByID[credential.ID] = credential + credentialsByGroup[credential.GroupID] = append(credentialsByGroup[credential.GroupID], credential) + } + if len(persistedCredentialByID) != len(runtimeByID) { + return nil, groupCollectionDataError("persisted and runtime credential sets differ") + } + for credentialID, persistedCredential := range persistedCredentialByID { + runtimeCredential, exists := runtimeByID[credentialID] + if !exists { + return nil, groupCollectionDataError( + "persisted credential %d is missing from runtime registry", + credentialID, + ) + } + status, err := groupCollectionRuntimeCredentialStatus(persistedCredential.Status) if err != nil { return nil, err } - if runtimeKey.ID != persistedKey.ID || - runtimeKey.GroupID != persistedKey.GroupID || - runtimeKey.Status != status || - !equalGroupCollectionWeight(runtimeKey.WeightManual, persistedKey.WeightManual) { + if runtimeCredential.ID != persistedCredential.ID || + runtimeCredential.GroupID != persistedCredential.GroupID || + runtimeCredential.Status != status || + runtimeCredential.Version != groupCollectionCredentialVersion(persistedCredential.UpdatedAtMS) || + runtimeCredential.IdentityGeneration != groupCollectionCredentialIdentity(persistedCredential.Fingerprint) || + !equalGroupCollectionWeight(runtimeCredential.WeightManual, persistedCredential.WeightManual) { return nil, groupCollectionDataError( - "persisted key %d differs from runtime registry", - keyID, + "persisted credential %d differs from runtime registry", + credentialID, ) } } - records := make([]groupCollectionRecord, 0, len(rows.groups)) for _, group := range rows.groups { - var protocols []protocol.Protocol - if err := json.Unmarshal(group.Protocols, &protocols); err != nil { - return nil, groupCollectionDataError( - "decode group %d protocols: %v", - group.ID, - err, - ) + supportsModelOptionalRequests := false + channelID := channel.ID(group.ChannelID) + if registry == nil { + return nil, groupCollectionDataError("channel registry is nil") } - if err := validateGroupCollectionProtocols(protocols); err != nil { - return nil, groupCollectionDataError( - "validate group %d protocols: %v", - group.ID, - err, - ) + validated, err := registry.ValidateParams(channelID, json.RawMessage(group.Params)) + if err != nil { + return nil, groupCollectionDataError("validate group %d params: %v", group.ID, err) } + params := validated.CanonicalJSON() + resolvedTarget, err := registry.Resolve(channelID, params) + if err != nil { + return nil, groupCollectionDataError("resolve group %d channel: %v", group.ID, err) + } + supportsModelOptionalRequests = resolvedTarget.Supports( + protocol.OpenAIResponses, + execution.OperationResponsesRetrieve, + execution.FeatureSet{}, + ) var groupModels []GroupModel if err := json.Unmarshal(group.Models, &groupModels); err != nil { return nil, groupCollectionDataError( @@ -286,21 +298,20 @@ func mapGroupCollectionRecords( catalog := snapshot.GroupCatalog[group.ID] record := groupCollectionRecord{ GroupCollectionItem: GroupCollectionItem{ - ID: group.ID, Name: group.Name, UpstreamURL: group.UpstreamURL, - ProviderID: cloneString(group.ProviderID), - Protocols: append([]protocol.Protocol(nil), protocols...), + ID: group.ID, Name: group.Name, ChannelID: channelID, + Params: append(json.RawMessage(nil), params...), ModelCount: int64(len(groupModels)), }, CreatedAtMS: group.CreatedAtMS, } - for _, persistedKey := range keysByGroup[group.ID] { - bucket := classifyHealthKey(catalog, runtimeByID[persistedKey.ID], observedAt) - addGroupCollectionKeyCount(&record.KeyCounts, bucket) + for _, persistedCredential := range credentialsByGroup[group.ID] { + bucket := classifyHealthKey(catalog, runtimeByID[persistedCredential.ID], observedAt) + addGroupCollectionCredentialCount(&record.CredentialCounts, bucket) } record.Status = groupCollectionStatus( catalog, - record.KeyCounts, - record.Protocols, + record.CredentialCounts, + supportsModelOptionalRequests, record.ModelCount, ) records = append(records, record) @@ -308,6 +319,44 @@ func mapGroupCollectionRecords( return records, nil } +func cloneCredentialRows(rows []models.Credential) []models.Credential { + cloned := make([]models.Credential, len(rows)) + for index := range rows { + cloned[index] = rows[index] + cloned[index].Group = nil + cloned[index].Data = "" + cloned[index].WeightManual = cloneInt(rows[index].WeightManual) + } + return cloned +} + +func groupCollectionRuntimeCredentialStatus(status models.CredentialStatus) (state.CredentialStatus, error) { + switch status { + case models.CredentialStatusActive: + return state.CredentialStatusActive, nil + case models.CredentialStatusDisabled: + return state.CredentialStatusDisabled, nil + default: + return "", groupCollectionDataError("invalid persisted credential status %q", status) + } +} + +func groupCollectionCredentialVersion(updatedAtMS int64) uint64 { + if updatedAtMS < 1 { + return 1 + } + return uint64(updatedAtMS) +} + +func groupCollectionCredentialIdentity(fingerprint string) uint64 { + hasher := fnv.New64a() + _, _ = hasher.Write([]byte(fingerprint)) + if value := hasher.Sum64(); value != 0 { + return value + } + return 1 +} + func groupCollectionDataError(format string, args ...any) error { return fmt.Errorf("%s: %w", fmt.Sprintf(format, args...), app_errors.ErrInternalServer) } @@ -319,36 +368,6 @@ func equalGroupCollectionWeight(left, right *int) bool { return *left == *right } -func groupCollectionRuntimeKeyStatus( - status models.UpstreamKeyStatus, -) (state.KeyStatus, error) { - switch status { - case models.UpstreamKeyStatusActive: - return state.KeyStatusActive, nil - case models.UpstreamKeyStatusDisabled: - return state.KeyStatusDisabled, nil - default: - return "", groupCollectionDataError("invalid persisted key status %q", status) - } -} - -func validateGroupCollectionProtocols(values []protocol.Protocol) error { - if len(values) == 0 { - return fmt.Errorf("protocols are required") - } - seen := make(map[protocol.Protocol]struct{}, len(values)) - for _, value := range values { - if !value.Valid() { - return fmt.Errorf("invalid protocol %q", value) - } - if _, duplicate := seen[value]; duplicate { - return fmt.Errorf("duplicate protocol %q", value) - } - seen[value] = struct{}{} - } - return nil -} - func validateGroupCollectionModels(values []GroupModel) error { seen := make(map[string]struct{}, len(values)) for _, value := range values { @@ -368,7 +387,7 @@ func validateGroupCollectionModels(values []GroupModel) error { return nil } -func addGroupCollectionKeyCount(counts *GroupCollectionKeyCounts, bucket healthBucket) { +func addGroupCollectionCredentialCount(counts *GroupCollectionCredentialCounts, bucket healthBucket) { counts.Total++ switch bucket { case healthBucketAvailable: @@ -384,8 +403,8 @@ func addGroupCollectionKeyCount(counts *GroupCollectionKeyCounts, bucket healthB func groupCollectionStatus( group state.GroupCatalogView, - counts GroupCollectionKeyCounts, - protocols []protocol.Protocol, + counts GroupCollectionCredentialCounts, + supportsModelOptionalRequests bool, modelCount int64, ) GroupCollectionStatus { if !group.Enabled || (group.WeightManual != nil && *group.WeightManual == 0) { @@ -397,10 +416,8 @@ func groupCollectionStatus( if modelCount > 0 { return GroupCollectionStatusAvailable } - for _, value := range protocols { - if value.SupportsModelOptionalRequests() { - return GroupCollectionStatusAvailable - } + if supportsModelOptionalRequests { + return GroupCollectionStatusAvailable } return GroupCollectionStatusUnavailable } diff --git a/internal/control/group_collection_http.go b/internal/control/group_collection_http.go index 31546c1d..983349e2 100644 --- a/internal/control/group_collection_http.go +++ b/internal/control/group_collection_http.go @@ -10,7 +10,6 @@ import ( app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/platform/response" - "gpt-load/internal/protocol" ) const ( @@ -68,7 +67,7 @@ func parseGroupCollectionQuery( } for key, entries := range values { switch key { - case "q", "status", "protocol", "sort", "page", "page_size": + case "q", "status", "sort", "page", "page_size": default: return GroupCollectionQuery{}, app_errors.ErrBadRequest } @@ -90,19 +89,12 @@ func parseGroupCollectionQuery( } query.Status = status } - if entries, exists := values["protocol"]; exists { - protocolValue, ok := parseGroupCollectionProtocol(entries[0]) - if !ok { - return GroupCollectionQuery{}, app_errors.ErrBadRequest - } - query.Protocol = &protocolValue - } if entries, exists := values["sort"]; exists { sortValue := GroupCollectionSort(entries[0]) switch sortValue { case GroupCollectionSortStatus, GroupCollectionSortName, - GroupCollectionSortKeys, + GroupCollectionSortCredentials, GroupCollectionSortCreated: query.Sort = sortValue default: @@ -140,19 +132,6 @@ func parseGroupCollectionStatus( } } -func parseGroupCollectionProtocol(value string) (protocol.Protocol, bool) { - protocolValue := protocol.Protocol(value) - switch protocolValue { - case protocol.OpenAICompletions, - protocol.OpenAIResponses, - protocol.Anthropic, - protocol.Gemini: - return protocolValue, true - default: - return "", false - } -} - func parseGroupCollectionPositiveInt(value string) (int64, bool) { if value == "" || value[0] == '0' { return 0, false diff --git a/internal/control/group_collection_http_test.go b/internal/control/group_collection_http_test.go index 40a04d94..5c6417fc 100644 --- a/internal/control/group_collection_http_test.go +++ b/internal/control/group_collection_http_test.go @@ -4,15 +4,14 @@ import ( "encoding/json" "net/http" "net/http/httptest" - "reflect" "strings" "testing" "time" "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" - "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/storage/models" ) @@ -21,10 +20,6 @@ func TestParseGroupCollectionQueryAcceptsStrictContract(t *testing.T) { available := GroupCollectionStatusAvailable unavailable := GroupCollectionStatusUnavailable disabled := GroupCollectionStatusDisabled - openAICompletions := protocol.OpenAICompletions - openAIResponses := protocol.OpenAIResponses - anthropic := protocol.Anthropic - gemini := protocol.Gemini q200 := strings.Repeat("猫", 200) tests := []struct { @@ -81,34 +76,6 @@ func TestParseGroupCollectionQueryAcceptsStrictContract(t *testing.T) { Status: &disabled, Sort: GroupCollectionSortStatus, Page: 1, PageSize: 20, }, }, - { - name: "protocol openai completions", - rawQuery: "protocol=openai-completions", - want: GroupCollectionQuery{ - Protocol: &openAICompletions, Sort: GroupCollectionSortStatus, Page: 1, PageSize: 20, - }, - }, - { - name: "protocol openai responses", - rawQuery: "protocol=openai-responses", - want: GroupCollectionQuery{ - Protocol: &openAIResponses, Sort: GroupCollectionSortStatus, Page: 1, PageSize: 20, - }, - }, - { - name: "protocol anthropic", - rawQuery: "protocol=anthropic", - want: GroupCollectionQuery{ - Protocol: &anthropic, Sort: GroupCollectionSortStatus, Page: 1, PageSize: 20, - }, - }, - { - name: "protocol gemini", - rawQuery: "protocol=gemini", - want: GroupCollectionQuery{ - Protocol: &gemini, Sort: GroupCollectionSortStatus, Page: 1, PageSize: 20, - }, - }, { name: "sort status", rawQuery: "sort=status", @@ -124,10 +91,10 @@ func TestParseGroupCollectionQueryAcceptsStrictContract(t *testing.T) { }, }, { - name: "sort keys", - rawQuery: "sort=keys", + name: "sort credentials", + rawQuery: "sort=credentials", want: GroupCollectionQuery{ - Sort: GroupCollectionSortKeys, Page: 1, PageSize: 20, + Sort: GroupCollectionSortCredentials, Page: 1, PageSize: 20, }, }, { @@ -160,9 +127,9 @@ func TestParseGroupCollectionQueryAcceptsStrictContract(t *testing.T) { }, { name: "all fields combine", - rawQuery: "q=+alpha+&status=available&protocol=anthropic&sort=created&page=2&page_size=100", + rawQuery: "q=+alpha+&status=available&sort=created&page=2&page_size=100", want: GroupCollectionQuery{ - Query: "alpha", Status: &available, Protocol: &anthropic, + Query: "alpha", Status: &available, Sort: GroupCollectionSortCreated, Page: 2, PageSize: 100, }, }, @@ -227,6 +194,21 @@ func TestParseGroupCollectionQueryRejectsEveryInvalidForm(t *testing.T) { } } +func TestGroupCollectionQueryUsesCredentialSortAndRejectsProtocolFilter(t *testing.T) { + got, apiErr := parseGroupCollectionQuery("sort=credentials", false) + if apiErr != nil { + t.Fatalf("parseGroupCollectionQuery(sort=credentials) error = %v", apiErr) + } + if got.Sort != GroupCollectionSort("credentials") { + t.Fatalf("sort = %q, want credentials", got.Sort) + } + for _, rawQuery := range []string{"sort=keys", "protocol=anthropic"} { + if _, apiErr := parseGroupCollectionQuery(rawQuery, false); apiErr == nil || apiErr.Code != "BAD_REQUEST" { + t.Fatalf("parseGroupCollectionQuery(%q) error = %v, want BAD_REQUEST", rawQuery, apiErr) + } + } +} + func assertGroupCollectionQueryEqual( t *testing.T, got GroupCollectionQuery, @@ -241,10 +223,6 @@ func assertGroupCollectionQueryEqual( got.Status != nil && *got.Status != *want.Status { t.Fatalf("query status = %#v, want %#v", got.Status, want.Status) } - if (got.Protocol == nil) != (want.Protocol == nil) || - got.Protocol != nil && *got.Protocol != *want.Protocol { - t.Fatalf("query protocol = %#v, want %#v", got.Protocol, want.Protocol) - } } func TestGroupCollectionHTTPReturnsExactCollectionAndOptionsContracts(t *testing.T) { @@ -256,7 +234,8 @@ func TestGroupCollectionHTTPReturnsExactCollectionAndOptionsContracts(t *testing 10, "alpha", true, - `["openai-completions","gemini"]`, + channel.OpenAICompatible, + `{"base_url":"https://alpha.example/v1"}`, `[{"id":"private-alpha","alias":"public-alpha"},{"id":"second-alpha","alias":""}]`, ) createGroupOptionGroup( @@ -265,17 +244,18 @@ func TestGroupCollectionHTTPReturnsExactCollectionAndOptionsContracts(t *testing 20, "zulu", false, - `["anthropic"]`, + channel.Anthropic, + `{}`, `[{"id":"private-zulu","alias":"public-zulu"}]`, ) entry := createGroupCollectionKey( t, fixture, 10, - models.UpstreamKeyStatusActive, + models.CredentialStatusActive, nil, ) - publishGroupCollectionRuntime(t, fixture, []state.KeyEntry{entry}) + publishGroupCollectionRuntime(t, fixture, []state.CredentialEntry{entry}) observedAt := time.Date(2026, time.August, 1, 9, 10, 11, 0, time.UTC) fixture.service.now = func() time.Time { return observedAt } @@ -309,13 +289,10 @@ func TestGroupCollectionHTTPReturnsExactCollectionAndOptionsContracts(t *testing item := collectionData.Items[0] if item.ID != 10 || item.Name != "alpha" || item.Status != GroupCollectionStatusAvailable || - item.UpstreamURL != "https://alpha.example/v1" || - !reflect.DeepEqual(item.Protocols, []protocol.Protocol{ - protocol.OpenAICompletions, - protocol.Gemini, - }) || + item.ChannelID != channel.OpenAICompatible || + string(item.Params) != `{"base_url":"https://alpha.example/v1"}` || item.ModelCount != 2 || - item.KeyCounts != (GroupCollectionKeyCounts{Total: 1, Available: 1}) { + item.CredentialCounts != (GroupCollectionCredentialCounts{Total: 1, Available: 1}) { t.Fatalf("collection item = %#v, want exact public fields", item) } @@ -331,13 +308,14 @@ func TestGroupCollectionHTTPReturnsExactCollectionAndOptionsContracts(t *testing if len(optionData) != 2 || optionData[0].ID != 10 || optionData[0].Name != "alpha" || !optionData[0].Enabled || - !reflect.DeepEqual(optionData[0].Protocols, []protocol.Protocol{protocol.OpenAICompletions, protocol.Gemini}) || + optionData[0].ChannelID != channel.OpenAICompatible || + string(optionData[0].Params) != `{"base_url":"https://alpha.example/v1"}` || len(optionData[0].Models) != 2 || optionData[0].Models[0] != "public-alpha" || optionData[0].Models[1] != "second-alpha" || optionData[1].ID != 20 || optionData[1].Name != "zulu" || optionData[1].Enabled || - !reflect.DeepEqual(optionData[1].Protocols, []protocol.Protocol{protocol.Anthropic}) || + optionData[1].ChannelID != channel.Anthropic || string(optionData[1].Params) != `{}` || len(optionData[1].Models) != 1 || optionData[1].Models[0] != "public-zulu" { t.Fatalf("options data = %#v, want exact ID-ordered directory", optionData) } @@ -356,9 +334,10 @@ func TestGroupCollectionHTTPAllowsAvailableKeysInAnUnavailableStatus(t *testing. initControlI18n(t) fixture := newServiceFixture(t) group := createGroupCollectionGroup(t, fixture, "http-zero-model-completions", true, nil) + setGroupCollectionChannel(t, fixture, group, channel.Anthropic, models.JSON(`{}`)) setGroupCollectionRoute(t, fixture, group, `["openai-completions"]`, `[]`) - publishGroupCollectionRuntime(t, fixture, []state.KeyEntry{ - createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil), + publishGroupCollectionRuntime(t, fixture, []state.CredentialEntry{ + createGroupCollectionKey(t, fixture, group.ID, models.CredentialStatusActive, nil), }) engine := gin.New() @@ -379,7 +358,7 @@ func TestGroupCollectionHTTPAllowsAvailableKeysInAnUnavailableStatus(t *testing. len(data.Items) != 1 || data.Items[0].Status != GroupCollectionStatusUnavailable || data.Items[0].ModelCount != 0 || - data.Items[0].KeyCounts != (GroupCollectionKeyCounts{Total: 1, Available: 1}) { + data.Items[0].CredentialCounts != (GroupCollectionCredentialCounts{Total: 1, Available: 1}) { t.Fatalf("collection data = %#v, want unavailable route with available key bucket", data) } } @@ -512,10 +491,10 @@ func decodeGroupOptionsSuccess( t.Fatalf("decode raw options data: %v", err) } for _, rawOption := range rawOptions { - if len(rawOption) != 5 { - t.Fatalf("option fields = %#v, want exactly id/name/enabled/protocols/models", rawOption) + if len(rawOption) != 6 { + t.Fatalf("option fields = %#v, want exactly id/name/channel_id/params/enabled/models", rawOption) } - for _, field := range []string{"id", "name", "enabled", "protocols", "models"} { + for _, field := range []string{"id", "name", "channel_id", "params", "enabled", "models"} { if _, ok := rawOption[field]; !ok { t.Fatalf("option fields = %#v, missing %q", rawOption, field) } diff --git a/internal/control/group_collection_query.go b/internal/control/group_collection_query.go index 883cbb79..0fa2e860 100644 --- a/internal/control/group_collection_query.go +++ b/internal/control/group_collection_query.go @@ -5,23 +5,20 @@ import ( "sort" "strings" "unicode" - - "gpt-load/internal/protocol" ) type GroupCollectionSort string const ( - GroupCollectionSortStatus GroupCollectionSort = "status" - GroupCollectionSortName GroupCollectionSort = "name" - GroupCollectionSortKeys GroupCollectionSort = "keys" - GroupCollectionSortCreated GroupCollectionSort = "created" + GroupCollectionSortStatus GroupCollectionSort = "status" + GroupCollectionSortName GroupCollectionSort = "name" + GroupCollectionSortCredentials GroupCollectionSort = "credentials" + GroupCollectionSortCreated GroupCollectionSort = "created" ) type GroupCollectionQuery struct { Query string Status *GroupCollectionStatus - Protocol *protocol.Protocol Sort GroupCollectionSort Page int64 PageSize int64 @@ -118,33 +115,14 @@ func matchesGroupCollectionQuery( if query.Status != nil && record.Status != *query.Status { return false } - if query.Protocol != nil && !groupCollectionHasProtocol(record.Protocols, *query.Protocol) { - return false - } if query.Query == "" { return true } if groupCollectionContainsFold(record.Name, query.Query) || - groupCollectionContainsFold(record.UpstreamURL, query.Query) { + groupCollectionContainsFold(string(record.ChannelID), query.Query) || + groupCollectionContainsFold(string(record.Params), query.Query) { return true } - for _, value := range record.Protocols { - if groupCollectionContainsFold(string(value), query.Query) { - return true - } - } - return false -} - -func groupCollectionHasProtocol( - values []protocol.Protocol, - want protocol.Protocol, -) bool { - for _, value := range values { - if value == want { - return true - } - } return false } @@ -180,9 +158,10 @@ func sortGroupCollectionRecords( switch sortBy { case GroupCollectionSortName: return compareGroupCollectionNames(left, right) - case GroupCollectionSortKeys: - if left.KeyCounts.Total != right.KeyCounts.Total { - return left.KeyCounts.Total > right.KeyCounts.Total + case GroupCollectionSortCredentials: + leftTotal, rightTotal := groupCollectionCredentialTotal(left), groupCollectionCredentialTotal(right) + if leftTotal != rightTotal { + return leftTotal > rightTotal } return compareGroupCollectionNames(left, right) case GroupCollectionSortCreated: @@ -201,6 +180,10 @@ func sortGroupCollectionRecords( }) } +func groupCollectionCredentialTotal(record groupCollectionRecord) int64 { + return record.CredentialCounts.Total +} + func groupCollectionStatusOrder(value GroupCollectionStatus) int { switch value { case GroupCollectionStatusUnavailable: @@ -264,7 +247,6 @@ func groupCollectionPageItems( func cloneGroupCollectionItem(value GroupCollectionItem) GroupCollectionItem { cloned := value - cloned.ProviderID = cloneString(value.ProviderID) - cloned.Protocols = append([]protocol.Protocol(nil), value.Protocols...) + cloned.Params = append([]byte(nil), value.Params...) return cloned } diff --git a/internal/control/group_collection_query_test.go b/internal/control/group_collection_query_test.go index c4a3acab..94d82307 100644 --- a/internal/control/group_collection_query_test.go +++ b/internal/control/group_collection_query_test.go @@ -8,6 +8,7 @@ import ( "testing" "time" + "gpt-load/internal/channel" "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/storage/models" @@ -19,9 +20,9 @@ func TestListGroupCollectionCapturesThenQueries(t *testing.T) { fixture.service.now = func() time.Time { return time.UnixMilli(observedAt) } available := createGroupCollectionGroup(t, fixture, "available", true, nil) disabled := createGroupCollectionGroup(t, fixture, "disabled", false, nil) - entries := []state.KeyEntry{ - createGroupCollectionKey(t, fixture, available.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, disabled.ID, models.UpstreamKeyStatusActive, nil), + entries := []state.CredentialEntry{ + createGroupCollectionKey(t, fixture, available.ID, models.CredentialStatusActive, nil), + createGroupCollectionKey(t, fixture, disabled.ID, models.CredentialStatusActive, nil), } publishGroupCollectionRuntime(t, fixture, entries) @@ -51,8 +52,8 @@ func TestListGroupCollectionQueryDoesNotSearchPersistedModelIDOrAlias(t *testing if err := fixture.db.Model(group).Update("models", group.Models).Error; err != nil { t.Fatalf("update persisted models: %v", err) } - entry := createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil) - publishGroupCollectionRuntime(t, fixture, []state.KeyEntry{entry}) + entry := createGroupCollectionKey(t, fixture, group.ID, models.CredentialStatusActive, nil) + publishGroupCollectionRuntime(t, fixture, []state.CredentialEntry{entry}) for _, query := range []string{"private-upstream-model", "public-model-alias"} { t.Run(query, func(t *testing.T) { @@ -114,7 +115,6 @@ func TestGroupCollectionQueryMatchesUnicodeSimpleFoldAndSortsWithIt(t *testing.T func TestGroupCollectionQueryCombinesFiltersAndSummarizesCompleteCollection(t *testing.T) { available := GroupCollectionStatusAvailable - anthropic := protocol.Anthropic records := []groupCollectionRecord{ groupCollectionQueryRecord(1, "Alpha", GroupCollectionStatusAvailable, "https://alpha.example/v1", []protocol.Protocol{protocol.OpenAICompletions}, 1, 100), groupCollectionQueryRecord(2, "Alpha anthropic", GroupCollectionStatusAvailable, "https://alpha.example/v1", []protocol.Protocol{protocol.Anthropic}, 2, 200), @@ -123,10 +123,10 @@ func TestGroupCollectionQueryCombinesFiltersAndSummarizesCompleteCollection(t *t } result := queryGroupCollectionRecords(1_700, records, GroupCollectionQuery{ - Query: "ALPHA", Status: &available, Protocol: &anthropic, + Query: "ALPHA", Status: &available, Sort: GroupCollectionSortName, Page: 1, PageSize: 20, }) - if got, want := groupCollectionItemIDs(result.Items), []uint{2}; !reflect.DeepEqual(got, want) { + if got, want := groupCollectionItemIDs(result.Items), []uint{1, 2}; !reflect.DeepEqual(got, want) { t.Fatalf("combined filter item IDs = %#v, want %#v", got, want) } if got, want := result.Summary, (GroupCollectionSummary{ @@ -154,7 +154,7 @@ func TestGroupCollectionQueryUsesFixedSortsWithIDTieBreak(t *testing.T) { }{ {sort: GroupCollectionSortStatus, want: []uint{2, 3, 5, 4, 1}}, {sort: GroupCollectionSortName, want: []uint{2, 5, 4, 1, 3}}, - {sort: GroupCollectionSortKeys, want: []uint{2, 3, 5, 4, 1}}, + {sort: GroupCollectionSortCredentials, want: []uint{2, 3, 5, 4, 1}}, {sort: GroupCollectionSortCreated, want: []uint{5, 4, 3, 2, 1}}, } for _, test := range tests { @@ -196,7 +196,7 @@ func TestGroupCollectionQueryUsesIDTieBreakAfterEachSortPrimaryAndName(t *testin }, { name: "keys ascending ID after same keys and name", - sort: GroupCollectionSortKeys, + sort: GroupCollectionSortCredentials, records: []groupCollectionRecord{ groupCollectionQueryRecord(10, "same", GroupCollectionStatusAvailable, "https://ten.example/v1", nil, 5, 100), groupCollectionQueryRecord(9, "same", GroupCollectionStatusDisabled, "https://nine.example/v1", nil, 5, 200), @@ -300,11 +300,24 @@ func groupCollectionQueryRecord( keys int64, createdAtMS int64, ) groupCollectionRecord { + channelID := channel.OpenAICompatible + for _, selectedProtocol := range protocols { + switch selectedProtocol { + case protocol.Anthropic: + channelID = channel.AnthropicCompatible + case protocol.Gemini: + channelID = channel.GeminiCompatible + } + } + params, err := json.Marshal(map[string]string{"base_url": upstreamURL}) + if err != nil { + panic(err) + } return groupCollectionRecord{ GroupCollectionItem: GroupCollectionItem{ - ID: id, Name: name, Status: status, UpstreamURL: upstreamURL, - Protocols: protocols, ModelCount: 7, - KeyCounts: GroupCollectionKeyCounts{Total: keys}, + ID: id, Name: name, Status: status, ChannelID: channelID, + Params: params, ModelCount: 7, + CredentialCounts: GroupCollectionCredentialCounts{Total: keys}, }, CreatedAtMS: createdAtMS, } diff --git a/internal/control/group_collection_test.go b/internal/control/group_collection_test.go deleted file mode 100644 index a9acccde..00000000 --- a/internal/control/group_collection_test.go +++ /dev/null @@ -1,716 +0,0 @@ -package control - -import ( - "context" - "errors" - "reflect" - "strings" - "sync" - "testing" - "time" - - "gorm.io/gorm" - - "gpt-load/internal/platform/encryption" - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" - "gpt-load/internal/state" - stateloader "gpt-load/internal/state/loader" - "gpt-load/internal/storage" - "gpt-load/internal/storage/models" -) - -func TestGroupCollectionStatusReflectsRouteCapabilityAndHealth(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.July, 31, 12, 0, 0, 0, time.FixedZone("test", 8*60*60)) - fixture.service.now = func() time.Time { return now } - - zero := 0 - available := createGroupCollectionGroup(t, fixture, "available", true, nil) - zeroModelCompletions := createGroupCollectionGroup(t, fixture, "zero-model-completions", true, nil) - zeroModelResponses := createGroupCollectionGroup(t, fixture, "zero-model-responses", true, nil) - mixedZeroModels := createGroupCollectionGroup(t, fixture, "mixed-zero-models", true, nil) - responsesWithoutHealthyKey := createGroupCollectionGroup(t, fixture, "responses-without-healthy-key", true, nil) - zeroKeys := createGroupCollectionGroup(t, fixture, "zero-keys", true, nil) - cooldown := createGroupCollectionGroup(t, fixture, "cooldown", true, nil) - blacklisted := createGroupCollectionGroup(t, fixture, "blacklisted", true, nil) - disabled := createGroupCollectionGroup(t, fixture, "disabled", false, nil) - groupWeightZero := createGroupCollectionGroup(t, fixture, "group-weight-zero", true, &zero) - keyWeightZero := createGroupCollectionGroup(t, fixture, "key-weight-zero", true, nil) - setGroupCollectionRoute(t, fixture, zeroModelCompletions, `["openai-completions"]`, `[]`) - setGroupCollectionRoute(t, fixture, zeroModelResponses, `["openai-responses"]`, `[]`) - setGroupCollectionRoute( - t, - fixture, - mixedZeroModels, - `["openai-completions","openai-responses"]`, - `[]`, - ) - setGroupCollectionRoute(t, fixture, responsesWithoutHealthyKey, `["openai-responses"]`, `[]`) - setGroupCollectionRoute(t, fixture, disabled, `["openai-responses"]`, `[]`) - setGroupCollectionRoute(t, fixture, groupWeightZero, `["openai-responses"]`, `[]`) - - entries := []state.KeyEntry{ - createGroupCollectionKey(t, fixture, available.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, zeroModelCompletions.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, zeroModelResponses.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, mixedZeroModels.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, cooldown.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, blacklisted.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, disabled.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, groupWeightZero.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, keyWeightZero.ID, models.UpstreamKeyStatusActive, &zero), - } - entries[4].CooldownUntil = now.Add(time.Minute) - entries[5].Blacklisted = true - publishGroupCollectionRuntime(t, fixture, entries) - - observedAtMS, records, err := fixture.service.captureGroupCollectionRecords(context.Background()) - if err != nil { - t.Fatalf("captureGroupCollectionRecords() error = %v", err) - } - if observedAtMS != now.UTC().UnixMilli() { - t.Fatalf("observedAtMS = %d, want %d", observedAtMS, now.UTC().UnixMilli()) - } - if len(records) != 11 { - t.Fatalf("record count = %d, want 11", len(records)) - } - - want := map[uint]GroupCollectionStatus{ - available.ID: GroupCollectionStatusAvailable, - zeroModelCompletions.ID: GroupCollectionStatusUnavailable, - zeroModelResponses.ID: GroupCollectionStatusAvailable, - mixedZeroModels.ID: GroupCollectionStatusAvailable, - responsesWithoutHealthyKey.ID: GroupCollectionStatusUnavailable, - zeroKeys.ID: GroupCollectionStatusUnavailable, - cooldown.ID: GroupCollectionStatusUnavailable, - blacklisted.ID: GroupCollectionStatusUnavailable, - disabled.ID: GroupCollectionStatusDisabled, - groupWeightZero.ID: GroupCollectionStatusDisabled, - keyWeightZero.ID: GroupCollectionStatusUnavailable, - } - for _, record := range records { - if got := record.Status; got != want[record.ID] { - t.Errorf("group %d status = %q, want %q", record.ID, got, want[record.ID]) - } - if record.ID == groupWeightZero.ID { - if got, want := record.KeyCounts, (GroupCollectionKeyCounts{ - Total: 1, Disabled: 1, - }); got != want { - t.Errorf("group weight zero key counts = %#v, want %#v", got, want) - } - } - if record.ID == zeroModelCompletions.ID { - if got, want := record.KeyCounts, (GroupCollectionKeyCounts{ - Total: 1, Available: 1, - }); got != want { - t.Errorf("zero-model Completions key counts = %#v, want %#v", got, want) - } - } - assertGroupCollectionKeyCountInvariant(t, record) - } -} - -func TestCaptureGroupCollectionRecordsKeyBuckets(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.July, 31, 4, 0, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - - group := createGroupCollectionGroup(t, fixture, "all-buckets", true, nil) - disabledGroup := createGroupCollectionGroup(t, fixture, "disabled-all-keys", false, nil) - zero := 0 - entries := []state.KeyEntry{ - createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusDisabled, nil), - createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, &zero), - createGroupCollectionKey(t, fixture, disabledGroup.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, disabledGroup.ID, models.UpstreamKeyStatusActive, nil), - } - entries[1].CooldownUntil = now.Add(time.Minute) - entries[2].Blacklisted = true - entries[5].CooldownUntil = now.Add(time.Minute) - entries[6].Blacklisted = true - publishGroupCollectionRuntime(t, fixture, entries) - - _, records, err := fixture.service.captureGroupCollectionRecords(context.Background()) - if err != nil { - t.Fatalf("captureGroupCollectionRecords() error = %v", err) - } - if len(records) != 2 { - t.Fatalf("record count = %d, want 2", len(records)) - } - if got, want := records[0].KeyCounts, (GroupCollectionKeyCounts{ - Total: 5, Available: 1, Cooldown: 1, Blacklisted: 1, Disabled: 2, - }); got != want { - t.Fatalf("all-buckets key counts = %#v, want %#v", got, want) - } - if got, want := records[1].KeyCounts, (GroupCollectionKeyCounts{ - Total: 2, Disabled: 2, - }); got != want { - t.Fatalf("disabled group key counts = %#v, want %#v", got, want) - } - for _, record := range records { - assertGroupCollectionKeyCountInvariant(t, record) - } -} - -func TestCaptureGroupCollectionRecordsMapsOwnedMetadata(t *testing.T) { - fixture := newServiceFixture(t) - group := createGroupCollectionGroup(t, fixture, "metadata", true, nil) - group.Protocols = models.JSON(`["openai-completions","anthropic"]`) - group.Models = models.JSON(`[{"id":"one"},{"id":"two","alias":"public-two"}]`) - if err := fixture.db.Model(group).Updates(map[string]any{ - "protocols": group.Protocols, - "models": group.Models, - }).Error; err != nil { - t.Fatalf("update group metadata: %v", err) - } - entry := createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil) - publishGroupCollectionRuntime(t, fixture, []state.KeyEntry{entry}) - - _, records, err := fixture.service.captureGroupCollectionRecords(t.Context()) - if err != nil { - t.Fatalf("captureGroupCollectionRecords() error = %v", err) - } - if len(records) != 1 { - t.Fatalf("records = %#v, want one", records) - } - record := records[0] - if record.ID != group.ID || record.Name != "metadata" || - record.UpstreamURL != group.UpstreamURL || record.CreatedAtMS != group.CreatedAtMS || - record.ModelCount != 2 || - !reflect.DeepEqual(record.Protocols, []protocol.Protocol{ - protocol.OpenAICompletions, - protocol.Anthropic, - }) { - t.Fatalf("record metadata = %#v, want persisted metadata", record) - } -} - -func TestCaptureGroupCollectionRecordsFailsClosedOnInconsistentObservations(t *testing.T) { - tests := []struct { - name string - mutate func(*state.ConfigSnapshot, *[]state.KeyRuntimeView, *groupCollectionRows) - wantErrorContains string - }{ - { - name: "persisted group zero id", - mutate: func(_ *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.groups[0].ID = 0 - }, - }, - { - name: "persisted group duplicate id", - mutate: func(_ *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.groups = append(rows.groups, rows.groups[0]) - }, - }, - { - name: "catalog group missing at equal cardinality", - mutate: func(snapshot *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, _ *groupCollectionRows) { - delete(snapshot.GroupCatalog, 1) - snapshot.GroupCatalog[2] = state.GroupCatalogView{ID: 2, Name: "replacement", Enabled: true} - }, - wantErrorContains: "missing from runtime catalog", - }, - { - name: "catalog group extra", - mutate: func(snapshot *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, _ *groupCollectionRows) { - snapshot.GroupCatalog[2] = state.GroupCatalogView{ID: 2, Name: "extra", Enabled: true} - }, - }, - { - name: "catalog group zero id", - mutate: func(snapshot *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, _ *groupCollectionRows) { - view := snapshot.GroupCatalog[1] - view.ID = 0 - snapshot.GroupCatalog[1] = view - }, - wantErrorContains: "zero id", - }, - { - name: "catalog id mismatch", - mutate: func(snapshot *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, _ *groupCollectionRows) { - view := snapshot.GroupCatalog[1] - view.ID = 2 - snapshot.GroupCatalog[1] = view - }, - }, - { - name: "catalog name mismatch", - mutate: func(snapshot *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, _ *groupCollectionRows) { - view := snapshot.GroupCatalog[1] - view.Name = "other" - snapshot.GroupCatalog[1] = view - }, - }, - { - name: "catalog enabled mismatch", - mutate: func(snapshot *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, _ *groupCollectionRows) { - view := snapshot.GroupCatalog[1] - view.Enabled = false - snapshot.GroupCatalog[1] = view - }, - }, - { - name: "catalog weight mismatch", - mutate: func(snapshot *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, _ *groupCollectionRows) { - other := 26 - view := snapshot.GroupCatalog[1] - view.WeightManual = &other - snapshot.GroupCatalog[1] = view - }, - }, - { - name: "persisted key zero id", - mutate: func(_ *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.keys[0].ID = 0 - }, - }, - { - name: "persisted key duplicate id", - mutate: func(_ *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.keys = append(rows.keys, rows.keys[0]) - }, - }, - { - name: "runtime key zero id", - mutate: func(_ *state.ConfigSnapshot, runtimeKeys *[]state.KeyRuntimeView, _ *groupCollectionRows) { - (*runtimeKeys)[0].ID = 0 - }, - }, - { - name: "runtime key duplicate id", - mutate: func(_ *state.ConfigSnapshot, runtimeKeys *[]state.KeyRuntimeView, _ *groupCollectionRows) { - *runtimeKeys = append(*runtimeKeys, (*runtimeKeys)[0]) - }, - }, - { - name: "runtime key missing at equal cardinality", - mutate: func(_ *state.ConfigSnapshot, runtimeKeys *[]state.KeyRuntimeView, _ *groupCollectionRows) { - *runtimeKeys = []state.KeyRuntimeView{{ - ID: 11, GroupID: 1, Status: state.KeyStatusActive, - }} - }, - wantErrorContains: "missing from runtime registry", - }, - { - name: "runtime key extra", - mutate: func(_ *state.ConfigSnapshot, runtimeKeys *[]state.KeyRuntimeView, _ *groupCollectionRows) { - *runtimeKeys = append(*runtimeKeys, state.KeyRuntimeView{ - ID: 11, GroupID: 1, Status: state.KeyStatusActive, - }) - }, - }, - { - name: "persisted key references missing group", - mutate: func(_ *state.ConfigSnapshot, runtimeKeys *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.keys[0].GroupID = 2 - (*runtimeKeys)[0].GroupID = 2 - }, - wantErrorContains: "references missing group", - }, - { - name: "runtime group mismatch", - mutate: func(_ *state.ConfigSnapshot, runtimeKeys *[]state.KeyRuntimeView, _ *groupCollectionRows) { - (*runtimeKeys)[0].GroupID = 2 - }, - }, - { - name: "runtime status mismatch", - mutate: func(_ *state.ConfigSnapshot, runtimeKeys *[]state.KeyRuntimeView, _ *groupCollectionRows) { - (*runtimeKeys)[0].Status = state.KeyStatusDisabled - }, - }, - { - name: "runtime weight mismatch", - mutate: func(_ *state.ConfigSnapshot, runtimeKeys *[]state.KeyRuntimeView, _ *groupCollectionRows) { - other := 51 - (*runtimeKeys)[0].WeightManual = &other - }, - }, - { - name: "invalid persisted key status", - mutate: func(_ *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.keys[0].Status = models.UpstreamKeyStatus("broken") - }, - }, - { - name: "invalid protocols json", - mutate: func(_ *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.groups[0].Protocols = models.JSON(`{"not":"an array"}`) - }, - }, - { - name: "invalid protocol value", - mutate: func(_ *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.groups[0].Protocols = models.JSON(`["legacy-openai"]`) - }, - }, - { - name: "invalid models json", - mutate: func(_ *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.groups[0].Models = models.JSON(`{"not":"an array"}`) - }, - }, - { - name: "invalid model value", - mutate: func(_ *state.ConfigSnapshot, _ *[]state.KeyRuntimeView, rows *groupCollectionRows) { - rows.groups[0].Models = models.JSON(`[{"id":" "}]`) - }, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - snapshot, runtimeKeys, rows := validGroupCollectionObservation() - test.mutate(snapshot, &runtimeKeys, &rows) - _, err := mapGroupCollectionRecords(snapshot, runtimeKeys, rows, time.Unix(1, 0)) - if !errors.Is(err, app_errors.ErrInternalServer) { - t.Fatalf("mapGroupCollectionRecords() error = %v, want ErrInternalServer", err) - } - if test.wantErrorContains != "" && !strings.Contains(err.Error(), test.wantErrorContains) { - t.Fatalf( - "mapGroupCollectionRecords() error = %v, want diagnostic containing %q", - err, - test.wantErrorContains, - ) - } - }) - } -} - -func TestCaptureGroupCollectionRecordsRejectsNilSnapshotWithoutPanic(t *testing.T) { - _, runtimeKeys, rows := validGroupCollectionObservation() - var panicValue any - var err error - func() { - defer func() { panicValue = recover() }() - _, err = mapGroupCollectionRecords(nil, runtimeKeys, rows, time.Unix(1, 0)) - }() - if panicValue != nil { - t.Fatalf("mapGroupCollectionRecords() panic = %v", panicValue) - } - if !errors.Is(err, app_errors.ErrInternalServer) { - t.Fatalf("mapGroupCollectionRecords() error = %v, want ErrInternalServer", err) - } -} - -func TestCaptureGroupCollectionRecordsWaitsForControlWrite(t *testing.T) { - fixture := newServiceFixture(t) - group := createGroupCollectionGroup(t, fixture, "wait-for-write", true, nil) - entry := createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil) - publishGroupCollectionRuntime(t, fixture, []state.KeyEntry{entry}) - - fixture.service.writeMu.Lock() - locked := true - t.Cleanup(func() { - if locked { - fixture.service.writeMu.Unlock() - } - }) - done := make(chan error, 1) - go func() { - _, _, err := fixture.service.captureGroupCollectionRecords(t.Context()) - done <- err - }() - select { - case err := <-done: - t.Fatalf("capture returned during control write: %v", err) - case <-time.After(50 * time.Millisecond): - } - fixture.service.writeMu.Unlock() - locked = false - select { - case err := <-done: - if err != nil { - t.Fatalf("capture after control write error = %v", err) - } - case <-time.After(time.Second): - t.Fatal("capture did not finish after control write") - } -} - -func TestCaptureGroupCollectionRecordsUsesSingleReadSnapshot(t *testing.T) { - fixture, dsn := newFileServiceFixture(t) - group := createGroupCollectionGroup(t, fixture, "collection-snapshot-old", true, nil) - entry := createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil) - publishGroupCollectionRuntime(t, fixture, []state.KeyEntry{entry}) - - writer, err := storage.Open(dsn) - if err != nil { - t.Fatalf("storage.Open(writer) error = %v", err) - } - writerSQL, err := writer.DB() - if err != nil { - t.Fatalf("writer.DB() error = %v", err) - } - t.Cleanup(func() { - if err := writerSQL.Close(); err != nil { - t.Errorf("close writer database: %v", err) - } - }) - - firstSelect := make(chan struct{}) - releaseReader := make(chan struct{}) - var barrierOnce sync.Once - const callbackName = "test:group_collection_single_read_snapshot" - if err := fixture.db.Callback().Query().After("gorm:query").Register( - callbackName, - func(tx *gorm.DB) { - if tx.Statement.Table != "groups" { - return - } - barrierOnce.Do(func() { - close(firstSelect) - <-releaseReader - }) - }, - ); err != nil { - t.Fatalf("register collection query barrier: %v", err) - } - - type captureResult struct { - records []groupCollectionRecord - err error - } - done := make(chan captureResult, 1) - go func() { - _, records, err := fixture.service.captureGroupCollectionRecords(t.Context()) - done <- captureResult{records: records, err: err} - }() - select { - case <-firstSelect: - case <-time.After(time.Second): - t.Fatal("capture did not pause after Group SELECT") - } - - writeDone := make(chan error, 1) - go func() { - writeDone <- writer.Transaction(func(tx *gorm.DB) error { - if err := tx.Model(&models.Group{}).Where("id = ?", group.ID). - Update("name", "collection-snapshot-new").Error; err != nil { - return err - } - return tx.Create(&models.UpstreamKey{ - GroupID: group.ID, KeyValue: "cipher-new", KeyHash: "hash-new", - Status: models.UpstreamKeyStatusActive, - }).Error - }) - }() - select { - case err := <-writeDone: - if err != nil { - t.Fatalf("concurrent collection update error = %v", err) - } - case <-time.After(time.Second): - t.Fatal("collection read blocked WAL writer") - } - close(releaseReader) - - select { - case result := <-done: - if result.err != nil { - t.Fatalf("captureGroupCollectionRecords() error = %v", result.err) - } - if len(result.records) != 1 || result.records[0].Name != "collection-snapshot-old" || - result.records[0].KeyCounts.Total != 1 { - t.Fatalf("capture mixed database versions: %#v", result.records) - } - case <-time.After(time.Second): - t.Fatal("capture did not finish after query barrier release") - } -} - -func TestCaptureGroupCollectionRecordsContextCancellationTakesPrecedence(t *testing.T) { - fixture := newServiceFixture(t) - ctx, cancel := context.WithCancel(t.Context()) - cancel() - fixture.service.db = nil - _, _, err := fixture.service.captureGroupCollectionRecords(ctx) - if !errors.Is(err, context.Canceled) { - t.Fatalf("captureGroupCollectionRecords() error = %v, want context.Canceled", err) - } -} - -func TestCaptureGroupCollectionRecordsMapsDatabaseErrors(t *testing.T) { - fixture := newServiceFixture(t) - if err := fixture.db.Exec("DROP TABLE groups").Error; err != nil { - t.Fatalf("drop groups table: %v", err) - } - _, _, err := fixture.service.captureGroupCollectionRecords(t.Context()) - if !errors.Is(err, app_errors.ErrDatabase) { - t.Fatalf("captureGroupCollectionRecords() error = %v, want ErrDatabase", err) - } -} - -type groupCollectionEncryptionProbe struct { - encryption.Service - decryptCalls int -} - -func (probe *groupCollectionEncryptionProbe) Decrypt(ciphertext string) (string, error) { - probe.decryptCalls++ - return probe.Service.Decrypt(ciphertext) -} - -func TestCaptureGroupCollectionRecordsNeverReadsOrDecryptsSecrets(t *testing.T) { - fixture := newServiceFixture(t) - group := createGroupCollectionGroup(t, fixture, "no-secrets", true, nil) - entry := createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil) - publishGroupCollectionRuntime(t, fixture, []state.KeyEntry{entry}) - probe := &groupCollectionEncryptionProbe{Service: fixture.encryption} - fixture.service.encryption = probe - fixture.service.registry = nil - var upstreamKeySelects []string - const callbackName = "test:group_collection_selected_key_columns" - if err := fixture.db.Callback().Query().Before("gorm:query").Register( - callbackName, - func(tx *gorm.DB) { - if tx.Statement.Table == "upstream_keys" { - upstreamKeySelects = append([]string(nil), tx.Statement.Selects...) - } - }, - ); err != nil { - t.Fatalf("register key query observer: %v", err) - } - - _, _, err := fixture.service.captureGroupCollectionRecords(t.Context()) - if err != nil { - t.Fatalf("captureGroupCollectionRecords() error = %v", err) - } - if probe.decryptCalls != 0 { - t.Fatalf("Decrypt calls = %d, want 0", probe.decryptCalls) - } - if want := []string{"id", "group_id", "status", "weight_manual"}; !reflect.DeepEqual( - upstreamKeySelects, - want, - ) { - t.Fatalf("upstream key SELECT columns = %#v, want %#v", upstreamKeySelects, want) - } -} - -func validGroupCollectionObservation() ( - *state.ConfigSnapshot, - []state.KeyRuntimeView, - groupCollectionRows, -) { - groupWeight := 25 - keyWeight := 50 - return &state.ConfigSnapshot{ - GroupCatalog: map[uint]state.GroupCatalogView{ - 1: {ID: 1, Name: "group", Enabled: true, WeightManual: &groupWeight}, - }, - }, []state.KeyRuntimeView{{ - ID: 10, GroupID: 1, Status: state.KeyStatusActive, WeightManual: &keyWeight, - }}, groupCollectionRows{ - groups: []models.Group{{ - ID: 1, Name: "group", Enabled: true, WeightManual: &groupWeight, - UpstreamURL: "https://group.example/v1", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[{"id":"gpt-4o"}]`), - CreatedAtMS: 123, - }}, - keys: []models.UpstreamKey{{ - ID: 10, GroupID: 1, Status: models.UpstreamKeyStatusActive, - WeightManual: &keyWeight, - }}, - } -} - -func createGroupCollectionGroup( - t *testing.T, - fixture serviceFixture, - name string, - enabled bool, - weight *int, -) *models.Group { - t.Helper() - group := validControlGroup(name) - group.Enabled = enabled - group.WeightManual = weight - if err := fixture.db.Create(group).Error; err != nil { - t.Fatalf("create group %q: %v", name, err) - } - if !enabled { - if err := fixture.db.Model(group).Update("enabled", false).Error; err != nil { - t.Fatalf("disable group %q: %v", name, err) - } - group.Enabled = false - } - return group -} - -func setGroupCollectionRoute( - t *testing.T, - fixture serviceFixture, - group *models.Group, - rawProtocols string, - rawModels string, -) { - t.Helper() - group.Protocols = models.JSON(rawProtocols) - group.Models = models.JSON(rawModels) - if err := fixture.db.Model(group).Updates(map[string]any{ - "protocols": group.Protocols, - "models": group.Models, - }).Error; err != nil { - t.Fatalf("update group %q route: %v", group.Name, err) - } -} - -func createGroupCollectionKey( - t *testing.T, - fixture serviceFixture, - groupID uint, - status models.UpstreamKeyStatus, - weight *int, -) state.KeyEntry { - t.Helper() - row := models.UpstreamKey{ - GroupID: groupID, - KeyValue: "ciphertext-not-observed", - KeyHash: "hash-not-observed-" + time.Now().String(), - Status: status, - WeightManual: weight, - } - if err := fixture.db.Create(&row).Error; err != nil { - t.Fatalf("create key for group %d: %v", groupID, err) - } - runtimeStatus := state.KeyStatusActive - if status == models.UpstreamKeyStatusDisabled { - runtimeStatus = state.KeyStatusDisabled - } - return state.KeyEntry{ - ID: row.ID, GroupID: groupID, Status: runtimeStatus, - WeightManual: weight, EncryptedValue: row.KeyValue, - } -} - -func publishGroupCollectionRuntime( - t *testing.T, - fixture serviceFixture, - entries []state.KeyEntry, -) { - t.Helper() - input, err := stateloader.BuildCompileInput(t.Context(), fixture.db) - if err != nil { - t.Fatalf("BuildCompileInput() error = %v", err) - } - if _, err := fixture.manager.Publish(input); err != nil { - t.Fatalf("manager.Publish() error = %v", err) - } - if err := fixture.registry.Replace(entries); err != nil { - t.Fatalf("registry.Replace() error = %v", err) - } -} - -func assertGroupCollectionKeyCountInvariant(t *testing.T, record groupCollectionRecord) { - t.Helper() - counts := record.KeyCounts - if got := counts.Available + counts.Cooldown + counts.Blacklisted + counts.Disabled; got != counts.Total { - t.Errorf("group %d bucket sum = %d, want total %d", record.ID, got, counts.Total) - } -} diff --git a/internal/control/group_collection_test_helpers_test.go b/internal/control/group_collection_test_helpers_test.go new file mode 100644 index 00000000..7a90ffbb --- /dev/null +++ b/internal/control/group_collection_test_helpers_test.go @@ -0,0 +1,126 @@ +package control + +import ( + "encoding/json" + "fmt" + "testing" + + "gpt-load/internal/channel" + "gpt-load/internal/state" + stateloader "gpt-load/internal/state/loader" + "gpt-load/internal/storage/models" +) + +func createGroupCollectionGroup( + t *testing.T, + fixture serviceFixture, + name string, + enabled bool, + weight *int, +) *models.Group { + t.Helper() + group := validControlGroup(name) + group.Params = models.JSON(fmt.Sprintf( + `{"base_url":"https://group-%d.example/v1"}`, + testIdempotencySequence.Add(1), + )) + group.Enabled = enabled + group.WeightManual = weight + if err := fixture.db.Create(group).Error; err != nil { + t.Fatalf("create group %q: %v", name, err) + } + if !enabled { + if err := fixture.db.Model(group).Update("enabled", false).Error; err != nil { + t.Fatalf("disable group %q: %v", name, err) + } + group.Enabled = false + } + return group +} + +func setGroupCollectionChannel( + t *testing.T, + fixture serviceFixture, + group *models.Group, + channelID channel.ID, + params models.JSON, +) { + t.Helper() + group.ChannelID = string(channelID) + group.Params = append(models.JSON(nil), params...) + if err := fixture.db.Model(group).Updates(map[string]any{ + "channel_id": group.ChannelID, + "params": group.Params, + }).Error; err != nil { + t.Fatalf("update group %q channel: %v", group.Name, err) + } +} + +func setGroupCollectionRoute( + t *testing.T, + fixture serviceFixture, + group *models.Group, + _ string, + rawModels string, +) { + t.Helper() + group.Models = models.JSON(rawModels) + if err := fixture.db.Model(group).Update("models", group.Models).Error; err != nil { + t.Fatalf("update group %q route: %v", group.Name, err) + } +} + +func createGroupCollectionKey( + t *testing.T, + fixture serviceFixture, + groupID uint, + status models.CredentialStatus, + weight *int, +) state.CredentialEntry { + t.Helper() + sequence := testIdempotencySequence.Add(1) + plain, err := json.Marshal(map[string]any{"api_key": fmt.Sprintf("sk-%d", sequence)}) + if err != nil { + t.Fatalf("encode credential: %v", err) + } + encrypted, err := fixture.encryption.Encrypt(string(plain)) + if err != nil { + t.Fatalf("encrypt credential: %v", err) + } + row := models.Credential{ + GroupID: groupID, Data: encrypted, Fingerprint: fixture.encryption.Hash(string(plain)), + Status: status, WeightManual: weight, + } + if err := fixture.db.Create(&row).Error; err != nil { + t.Fatalf("create credential for group %d: %v", groupID, err) + } + runtimeStatus := state.CredentialStatusActive + if status == models.CredentialStatusDisabled { + runtimeStatus = state.CredentialStatusDisabled + } + return state.CredentialEntry{ + ID: row.ID, GroupID: groupID, Status: runtimeStatus, WeightManual: weight, + Version: groupCollectionCredentialVersion(row.UpdatedAtMS), + IdentityGeneration: groupCollectionCredentialIdentity(row.Fingerprint), + Fingerprint: row.Fingerprint, + EncryptedValue: row.Data, + } +} + +func publishGroupCollectionRuntime( + t *testing.T, + fixture serviceFixture, + entries []state.CredentialEntry, +) { + t.Helper() + input, err := stateloader.BuildCompileInput(t.Context(), fixture.db, fixture.service.channelRegistry) + if err != nil { + t.Fatalf("BuildCompileInput() error = %v", err) + } + if _, err := fixture.manager.Publish(input); err != nil { + t.Fatalf("manager.Publish() error = %v", err) + } + if err := fixture.registry.ReplaceCredentials(entries); err != nil { + t.Fatalf("registry.ReplaceCredentials() error = %v", err) + } +} diff --git a/internal/control/group_create.go b/internal/control/group_create.go index ad446c49..4cff2a00 100644 --- a/internal/control/group_create.go +++ b/internal/control/group_create.go @@ -10,30 +10,29 @@ import ( "github.com/sirupsen/logrus" "gorm.io/gorm" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/platform/utils" - "gpt-load/internal/protocol" "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" "gpt-load/internal/storage/models" ) type GroupCreateRequest struct { - Name *string `json:"name"` - ProviderID *string `json:"provider_id"` - UpstreamURL string `json:"upstream_url"` - Protocols []protocol.Protocol `json:"protocols"` - Models optionalGroupModels `json:"models"` - Keys string `json:"keys"` - ConfirmSameUpstreamURL bool `json:"confirm_same_upstream_url"` + Name *string `json:"name"` + ChannelID channel.ID `json:"channel_id"` + Params json.RawMessage `json:"params"` + Models optionalGroupModels `json:"models"` + Credentials string `json:"credentials"` + ConfirmSameTarget bool `json:"confirm_same_target"` } type GroupCreateResult struct { - GroupID uint `json:"group_id"` - GroupName string `json:"group_name"` - KeysAdded int `json:"keys_added"` - KeysDuplicated int `json:"keys_duplicated"` + GroupID uint `json:"group_id"` + GroupName string `json:"group_name"` + CredentialsAdded int `json:"credentials_added"` + CredentialsDuplicated int `json:"credentials_duplicated"` } type ExistingGroupSummary struct { @@ -41,20 +40,20 @@ type ExistingGroupSummary struct { Name string `json:"name"` } -type UpstreamURLConflictData struct { +type SameTargetConflictData struct { Groups []ExistingGroupSummary `json:"groups"` } type normalizedGroupCreate struct { - providerID *string - upstreamURL string - hostname string - protocols []protocol.Protocol - explicitName *string - models []GroupModel - encodedConfig models.JSON - keys normalizedUpstreamKeys - confirmSameUpstreamURL bool + channelID channel.ID + params models.JSON + defaultName string + hostname string + explicitName *string + models []GroupModel + encodedOverrides models.JSON + credentials normalizedCredentials + confirmSameTarget bool } func (s *Service) CreateGroup(ctx context.Context, request GroupCreateRequest) (GroupCreateResult, error) { @@ -68,46 +67,41 @@ func (s *Service) CreateGroup(ctx context.Context, request GroupCreateRequest) ( logrus.WarnLevel, utils.LogPlaneControl, logrus.Fields{"host": normalized.hostname}, - "Creating upstream group with a private or local host", + "Creating channel group with a private or local host", ) } result := GroupCreateResult{} - requestedEntries := make([]state.KeyEntry, 0, len(normalized.keys.candidates)) + requestedEntries := make([]state.CredentialEntry, 0, len(normalized.credentials.candidates)) _, err = s.writeGroupConfig(ctx, func(tx *gorm.DB) error { - if !normalized.confirmSameUpstreamURL { - conflicts, err := findGroupsByUpstreamURL(tx, normalized.upstreamURL) + if !normalized.confirmSameTarget { + conflicts, err := findGroupsByTarget(tx, normalized.channelID, normalized.params) if err != nil { return err } if len(conflicts) > 0 { return app_errors.NewAPIErrorWithData( - app_errors.ErrUpstreamURLConflict, - UpstreamURLConflictData{Groups: conflicts}, + app_errors.ErrChannelTargetConflict, + SameTargetConflictData{Groups: conflicts}, ) } } - name, err := resolveGroupCreateName(tx, normalized.explicitName, normalized.hostname) + name, err := resolveGroupCreateName(tx, normalized.explicitName, normalized.defaultName) if err != nil { return err } - encodedProtocols, err := json.Marshal(normalized.protocols) - if err != nil { - return fmt.Errorf("encode group protocols: %w", err) - } encodedModels, err := json.Marshal(normalized.models) if err != nil { return fmt.Errorf("encode group models: %w", err) } group := models.Group{ - Name: name, - ProviderID: cloneString(normalized.providerID), - UpstreamURL: normalized.upstreamURL, - Protocols: models.JSON(encodedProtocols), - Models: models.JSON(encodedModels), - Config: normalized.encodedConfig, - Enabled: true, + Name: name, + ChannelID: string(normalized.channelID), + Params: append(models.JSON(nil), normalized.params...), + Models: models.JSON(encodedModels), + Overrides: normalized.encodedOverrides, + Enabled: true, } if err := tx.Create(&group).Error; err != nil { return app_errors.ParseDBError(err) @@ -115,22 +109,21 @@ func (s *Service) CreateGroup(ctx context.Context, request GroupCreateRequest) ( result.GroupID = group.ID result.GroupName = group.Name - requestedEntries, result.KeysAdded, result.KeysDuplicated, err = - s.persistUpstreamKeys(tx, group.ID, normalized.keys) + result.CredentialsAdded, result.CredentialsDuplicated, err = + s.persistCredentials(tx, group.ID, normalized.credentials) if err != nil { return err } - if err := state.ValidateKeyEntries(requestedEntries); err != nil { - return fmt.Errorf("validate created group keys: %w", err) - } - return nil + requestedEntries, err = stateloader.BuildGroupCredentialEntries(ctx, tx, group.ID) + return err }, func() error { - return s.applyMissingRegistryKeys(result.GroupID, requestedEntries) + _, reconcileErr := s.reconcileRegistryGroup(result.GroupID, requestedEntries) + return reconcileErr }) if err != nil { return GroupCreateResult{}, err } - if normalized.providerID != nil && len(normalized.models) > 0 && s.catalogSync != nil { + if len(normalized.models) > 0 && s.catalogSync != nil { s.catalogSync.RequestGroupSync() } return result, nil @@ -140,23 +133,19 @@ func (s *Service) normalizeGroupCreate( ctx context.Context, request GroupCreateRequest, ) (normalizedGroupCreate, error) { - upstreamURL, hostname, err := normalizeUpstreamBaseURL(request.UpstreamURL) - if err != nil { - return normalizedGroupCreate{}, err + if s == nil || s.channelRegistry == nil || request.ChannelID == "" { + return normalizedGroupCreate{}, app_errors.ErrValidation } - protocols, err := normalizeGroupProtocols(request.Protocols) - if err != nil { - return normalizedGroupCreate{}, err - } - explicitName, err := normalizeGroupName(request.Name) - if err != nil { - return normalizedGroupCreate{}, err - } - providerID, err := normalizeProviderID(request.ProviderID) + params, err := s.channelRegistry.ValidateParams(request.ChannelID, request.Params) if err != nil { return normalizedGroupCreate{}, app_errors.ErrValidation } - if err := s.validateSelectableProviderID(providerID); err != nil { + descriptor, ok := s.channelRegistry.Get(request.ChannelID) + if !ok { + return normalizedGroupCreate{}, app_errors.ErrValidation + } + explicitName, err := normalizeGroupName(request.Name) + if err != nil { return normalizedGroupCreate{}, err } if !request.Models.Set { @@ -166,7 +155,7 @@ func (s *Service) normalizeGroupCreate( if err != nil { return normalizedGroupCreate{}, err } - keys, err := s.normalizeUpstreamKeys(request.Keys) + credentials, err := s.normalizeCredentials(request.ChannelID, request.Credentials) if err != nil { return normalizedGroupCreate{}, err } @@ -182,23 +171,38 @@ func (s *Service) normalizeGroupCreate( if err != nil { return normalizedGroupCreate{}, app_errors.ParseDBError(err) } + canonicalParams := params.CanonicalJSON() _, err = state.Compile(state.CompileInput{ - SystemSettings: systemSettings, + SystemSettings: systemSettings, + ChannelRegistry: s.channelRegistry, Groups: []state.GroupConfig{{ - ID: 1, Name: "candidate", UpstreamURL: upstreamURL, - ProviderID: providerID, - Protocols: protocols, Models: runtimeModels, Settings: config.Settings{}, Enabled: true, - }}}) + ID: 1, Name: "candidate", ChannelID: request.ChannelID, + Params: canonicalParams, Models: runtimeModels, Settings: config.Settings{}, Enabled: true, + }}, + }) if err != nil { return normalizedGroupCreate{}, app_errors.ErrValidation } + defaultName := descriptor.Name + hostname := "" + if baseURL, exists := params.Value("base_url"); exists { + _, hostname, err = normalizeUpstreamBaseURL(baseURL) + if err != nil { + return normalizedGroupCreate{}, app_errors.ErrValidation + } + defaultName = hostname + } return normalizedGroupCreate{ - providerID: providerID, - upstreamURL: upstreamURL, hostname: hostname, protocols: protocols, - explicitName: explicitName, models: groupModels, - encodedConfig: models.JSON(`{}`), keys: keys, - confirmSameUpstreamURL: request.ConfirmSameUpstreamURL, + channelID: request.ChannelID, + params: models.JSON(canonicalParams), + defaultName: defaultName, + hostname: hostname, + explicitName: explicitName, + models: groupModels, + encodedOverrides: models.JSON(`{}`), + credentials: credentials, + confirmSameTarget: request.ConfirmSameTarget, }, nil } @@ -249,15 +253,31 @@ func canonicalizeGroupSettingNumbers(value any) any { } } -func findGroupsByUpstreamURL(tx *gorm.DB, upstreamURL string) ([]ExistingGroupSummary, error) { - groups := make([]ExistingGroupSummary, 0) +func findGroupsByTarget( + tx *gorm.DB, + channelID channel.ID, + params models.JSON, +) ([]ExistingGroupSummary, error) { + type targetRow struct { + ID uint + Name string + Params models.JSON + } + var rows []targetRow if err := tx.Model(&models.Group{}). - Select("id", "name"). - Where("upstream_url = ?", upstreamURL). + Select("id", "name", "params"). + Where("channel_id = ?", string(channelID)). Order("id ASC"). - Scan(&groups).Error; err != nil { + Find(&rows).Error; err != nil { return nil, app_errors.ParseDBError(err) } + groups := make([]ExistingGroupSummary, 0, len(rows)) + for _, row := range rows { + if !bytes.Equal(bytes.TrimSpace(row.Params), bytes.TrimSpace(params)) { + continue + } + groups = append(groups, ExistingGroupSummary{ID: row.ID, Name: row.Name}) + } return groups, nil } diff --git a/internal/control/group_create_test.go b/internal/control/group_create_test.go deleted file mode 100644 index 3fe7e282..00000000 --- a/internal/control/group_create_test.go +++ /dev/null @@ -1,809 +0,0 @@ -package control - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "net/http" - "reflect" - "strconv" - "strings" - "testing" - "time" - - "github.com/sirupsen/logrus" - - "gpt-load/internal/catalog" - "gpt-load/internal/dialect" - "gpt-load/internal/platform/encryption" - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" - "gpt-load/internal/state" - "gpt-load/internal/storage/models" -) - -func TestCreateGroupNormalizesAndPublishesOnce(t *testing.T) { - fixture := newServiceFixture(t) - if err := fixture.db.Create(&models.SystemSetting{ - Key: state.SettingInjectUsageOptions, - Value: "false", - }).Error; err != nil { - t.Fatalf("seed system setting: %v", err) - } - name := " primary upstream " - beforeRevision := fixture.manager.Current().Revision - - result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - Name: &name, - UpstreamURL: " HTTPS://API.Example.COM/v1/ ", - Protocols: []protocol.Protocol{ - protocol.OpenAIResponses, - protocol.Anthropic, - protocol.OpenAICompletions, - protocol.OpenAIResponses, - }, - Models: optionalGroupModels{ - Set: true, - Values: []GroupModel{ - {ID: " gpt-4o ", Alias: " primary ", AliasEnabled: true}, - {ID: "claude-3-5", Alias: "must-be-ignored"}, - }, - }, - Keys: " sk-a \n\n sk-a\nsk-b\n", - }) - if err != nil { - t.Fatalf("CreateGroup() error = %v", err) - } - if result.GroupID == 0 || result.GroupName != "primary upstream" || - result.KeysAdded != 2 || result.KeysDuplicated != 1 { - t.Fatalf("CreateGroup() result = %#v", result) - } - wantModels := []GroupModel{{ID: "gpt-4o", Alias: "primary"}, {ID: "claude-3-5", Alias: ""}} - encodedResult, err := json.Marshal(result) - if err != nil { - t.Fatalf("encode create result: %v", err) - } - var resultFields map[string]json.RawMessage - if err := json.Unmarshal(encodedResult, &resultFields); err != nil { - t.Fatalf("decode create result fields: %v", err) - } - if len(resultFields) != 4 || resultFields["models"] != nil { - t.Fatalf("create result fields = %#v, want only group and key counts", resultFields) - } - - var group models.Group - if err := fixture.db.First(&group, result.GroupID).Error; err != nil { - t.Fatalf("query created group: %v", err) - } - if group.Name != "primary upstream" || group.UpstreamURL != "https://api.example.com/v1" || - string(group.Config) != `{}` || !group.Enabled { - t.Fatalf("stored group = %#v", group) - } - var storedProtocols []protocol.Protocol - if err := json.Unmarshal(group.Protocols, &storedProtocols); err != nil { - t.Fatalf("decode stored protocols: %v", err) - } - wantProtocols := []protocol.Protocol{ - protocol.OpenAICompletions, - protocol.OpenAIResponses, - protocol.Anthropic, - } - if !reflect.DeepEqual(storedProtocols, wantProtocols) { - t.Fatalf("stored protocols = %#v", storedProtocols) - } - if storedModels := loadCreatedGroupModels(t, fixture, result.GroupID); !reflect.DeepEqual(storedModels, wantModels) { - t.Fatalf("stored models = %#v, want %#v", storedModels, wantModels) - } - - snapshot := fixture.manager.Current() - if snapshot.Revision != beforeRevision+1 { - t.Fatalf("snapshot revision = %d, want %d", snapshot.Revision, beforeRevision+1) - } - view, ok := snapshot.Groups[result.GroupID] - if !ok || view.UpstreamURL != group.UpstreamURL || - view.InjectUsageOptions || - !reflect.DeepEqual(view.Protocols, wantProtocols) { - t.Fatalf("snapshot group = %#v, exists=%t", view, ok) - } - settings, err := fixture.service.GetGroupSettings(t.Context(), result.GroupID) - if err != nil { - t.Fatal(err) - } - if len(settings.Overrides) != 0 || - settings.Effective.InjectUsageOptions { - t.Fatalf("created group config/effective = %#v/%#v", settings.Overrides, settings.Effective) - } - if candidates := fixture.registry.CollectCandidates([]uint{result.GroupID}, nil, time.Time{}); len(candidates) != 2 { - t.Fatalf("Registry candidates = %#v, want two", candidates) - } -} - -func TestCreateGroupRejectsDuplicateExternalModelWithoutMutation(t *testing.T) { - fixture := newServiceFixture(t) - beforeRevision := fixture.manager.Current().Revision - beforeRegistry := fixture.registry.Snapshot() - - _, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://duplicate-external-model.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{ - Set: true, - Values: []GroupModel{ - {ID: "provider-a", Alias: "public", AliasEnabled: true}, - {ID: "provider-b", Alias: "public", AliasEnabled: true}, - }, - }, - Keys: "sk-duplicate-external-model", - }) - var apiErr *app_errors.APIError - if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrModelNameConflict.Code { - t.Fatalf("CreateGroup() error = %#v, want MODEL_NAME_CONFLICT", err) - } - assertGroupCount(t, fixture.db, 0) - var keyCount int64 - if err := fixture.db.Model(&models.UpstreamKey{}).Count(&keyCount).Error; err != nil { - t.Fatal(err) - } - if keyCount != 0 { - t.Fatalf("upstream key count = %d, want 0", keyCount) - } - if fixture.manager.Current().Revision != beforeRevision { - t.Fatal("invalid models published a Snapshot") - } - if !reflect.DeepEqual(fixture.registry.Snapshot(), beforeRegistry) { - t.Fatal("invalid models changed Registry") - } - - result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://case-sensitive-models.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{ - Set: true, - Values: []GroupModel{ - {ID: "provider-a", Alias: "Public", AliasEnabled: true}, - {ID: "provider-b", Alias: "public", AliasEnabled: true}, - }, - }, - Keys: "sk-case-sensitive-models", - }) - if err != nil { - t.Fatalf("CreateGroup(case-distinct client names) error = %v", err) - } - if stored := loadCreatedGroupModels(t, fixture, result.GroupID); !reflect.DeepEqual(stored, []GroupModel{ - {ID: "provider-a", Alias: "Public"}, - {ID: "provider-b", Alias: "public"}, - }) { - t.Fatalf("case-distinct stored models = %#v", stored) - } -} - -func TestCreateGroupReturnsUpstreamURLConflictWithoutMutation(t *testing.T) { - fixture := newServiceFixture(t) - first, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://api.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-first", - }) - if err != nil { - t.Fatalf("initial CreateGroup() error = %v", err) - } - - var existingKey models.UpstreamKey - if err := fixture.db.Where("group_id = ?", first.GroupID).First(&existingKey).Error; err != nil { - t.Fatalf("query existing key: %v", err) - } - beforeCiphertext := existingKey.KeyValue - if count, ok := fixture.registry.IncrFailure(existingKey.ID); !ok || count != 1 { - t.Fatalf("first failure count = %d, %t", count, ok) - } - beforeRevision := fixture.manager.Current().Revision - beforeCandidates := fixture.registry.CollectCandidates([]uint{first.GroupID}, nil, time.Time{}) - - _, err = fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: " https://API.example.com/v1/ ", - Protocols: []protocol.Protocol{protocol.Anthropic}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-second", - }) - - var apiErr *app_errors.APIError - if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrUpstreamURLConflict.Code { - t.Fatalf("CreateGroup() error = %#v", err) - } - data, ok := apiErr.Data.(UpstreamURLConflictData) - if !ok || len(data.Groups) != 1 || data.Groups[0].ID != first.GroupID || - data.Groups[0].Name != first.GroupName { - t.Fatalf("conflict data = %#v", apiErr.Data) - } - assertGroupCount(t, fixture.db, 1) - if fixture.manager.Current().Revision != beforeRevision { - t.Fatal("conflict published a snapshot") - } - if !reflect.DeepEqual(fixture.registry.CollectCandidates([]uint{first.GroupID}, nil, time.Time{}), beforeCandidates) { - t.Fatal("conflict changed Registry") - } - if err := fixture.db.First(&existingKey, existingKey.ID).Error; err != nil { - t.Fatalf("reload existing key: %v", err) - } - if existingKey.KeyValue != beforeCiphertext { - t.Fatalf("ciphertext changed from %q to %q", beforeCiphertext, existingKey.KeyValue) - } - if count, ok := fixture.registry.IncrFailure(existingKey.ID); !ok || count != 2 { - t.Fatalf("second failure count = %d, %t, want 2", count, ok) - } - var keyCount int64 - if err := fixture.db.Model(&models.UpstreamKey{}).Count(&keyCount).Error; err != nil { - t.Fatalf("count upstream keys: %v", err) - } - if keyCount != 1 { - t.Fatalf("upstream key count = %d, want 1", keyCount) - } -} - -func TestCreateGroupAllowsConfirmedDuplicateURLAsIndependentGroup(t *testing.T) { - fixture := newServiceFixture(t) - first, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://same.example.com/v1/", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{ - Set: true, Values: []GroupModel{{ID: "gpt-4o"}}, - }, - Keys: "sk-first", - }) - if err != nil { - t.Fatalf("first CreateGroup() error = %v", err) - } - second, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: " HTTPS://SAME.example.com/v1 ", - Protocols: []protocol.Protocol{protocol.Anthropic}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{{ID: "claude-sonnet"}}}, - Keys: "sk-second", - ConfirmSameUpstreamURL: true, - }) - if err != nil { - t.Fatalf("confirmed CreateGroup() error = %v", err) - } - if first.GroupID == second.GroupID { - t.Fatalf("confirmed groups share ID %d", first.GroupID) - } - if first.GroupName == second.GroupName { - t.Fatalf("automatic names are not independent: %q", first.GroupName) - } - - var groups []models.Group - if err := fixture.db.Order("id ASC").Find(&groups).Error; err != nil { - t.Fatalf("query groups: %v", err) - } - if len(groups) != 2 || groups[0].UpstreamURL != "https://same.example.com/v1" || - groups[1].UpstreamURL != groups[0].UpstreamURL { - t.Fatalf("stored groups = %#v", groups) - } - assertStoredGroupProtocols(t, groups[0], []protocol.Protocol{protocol.OpenAICompletions}) - assertStoredGroupProtocols(t, groups[1], []protocol.Protocol{protocol.Anthropic}) - if !reflect.DeepEqual(loadCreatedGroupModels(t, fixture, first.GroupID), []GroupModel{{ID: "gpt-4o"}}) || - !reflect.DeepEqual(loadCreatedGroupModels(t, fixture, second.GroupID), []GroupModel{{ID: "claude-sonnet"}}) { - t.Fatal("confirmed groups do not retain independent models") - } - - _, err = fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://same.example.com/v1", - Protocols: []protocol.Protocol{protocol.Gemini}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-conflicting-third", - }) - var apiErr *app_errors.APIError - if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrUpstreamURLConflict.Code { - t.Fatalf("third CreateGroup() error = %#v, want URL conflict", err) - } - data, ok := apiErr.Data.(UpstreamURLConflictData) - if !ok || !reflect.DeepEqual(data.Groups, []ExistingGroupSummary{ - {ID: first.GroupID, Name: first.GroupName}, - {ID: second.GroupID, Name: second.GroupName}, - }) { - t.Fatalf("multiple conflict groups = %#v", apiErr.Data) - } -} - -func TestCreateGroupPersistsEncryptedKeysAndDoesNotDiscoverModels(t *testing.T) { - fixture := newServiceFixture(t) - discoveryCalls := 0 - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ - value: protocol.OpenAICompletions, - listFn: func(_ context.Context, _, _ string, _ state.HeaderRules) ([]string, error) { - discoveryCalls++ - return []string{"upstream-only"}, nil - }, - }) - - result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://secure.example.com/v1/", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{ - Set: true, Values: []GroupModel{{ID: "user-configured"}}, - }, - Keys: "sk-secret-a\nsk-secret-b", - }) - if err != nil { - t.Fatalf("CreateGroup() error = %v", err) - } - if discoveryCalls != 0 { - t.Fatalf("Dialect.ListModels calls = %d, want 0", discoveryCalls) - } - var rows []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", result.GroupID).Order("id ASC").Find(&rows).Error; err != nil { - t.Fatalf("query upstream keys: %v", err) - } - if len(rows) != 2 { - t.Fatalf("stored keys = %#v, want two", rows) - } - for index, plaintext := range []string{"sk-secret-a", "sk-secret-b"} { - row := rows[index] - if row.KeyValue == "" || row.KeyHash == "" || strings.Contains(row.KeyValue, plaintext) || - strings.Contains(row.KeyHash, plaintext) { - t.Fatalf("stored key %d exposes plaintext: %#v", row.ID, row) - } - decrypted, err := fixture.encryption.Decrypt(row.KeyValue) - if err != nil || decrypted != plaintext { - t.Fatalf("Decrypt(key %d) = %q, %v, want %q", row.ID, decrypted, err, plaintext) - } - if got, ok := fixture.registry.EncryptedValue(row.ID); !ok || got != row.KeyValue { - t.Fatalf("Registry key %d = %q, %t, want persisted cipher", row.ID, got, ok) - } - } -} - -func TestCreateGroupPreservesUserGeminiModelPrefix(t *testing.T) { - fixture := newServiceFixture(t) - result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://generativelanguage.googleapis.com/v1beta", - Protocols: []protocol.Protocol{protocol.Gemini}, - Models: optionalGroupModels{ - Set: true, Values: []GroupModel{{ID: " models/foo ", Alias: " user alias ", AliasEnabled: true}}, - }, - Keys: "gemini-secret", - }) - if err != nil { - t.Fatalf("CreateGroup() error = %v", err) - } - want := []GroupModel{{ID: "models/foo", Alias: "user alias"}} - if stored := loadCreatedGroupModels(t, fixture, result.GroupID); !reflect.DeepEqual(stored, want) { - t.Fatalf("stored models = %#v, want %#v", stored, want) - } - snapshot := fixture.manager.Current() - targets := snapshot.Candidates[protocol.Gemini]["user alias"] - if len(targets) != 1 || targets[0].UpstreamModelID != "models/foo" { - t.Fatalf("Gemini snapshot targets = %#v", targets) - } - if _, exists := snapshot.Candidates[protocol.Gemini]["models/foo"]; exists { - t.Fatal("aliased upstream model ID must not appear in the external candidate index") - } -} - -func TestCreateGroupPrivateHostLogContainsOnlyHostname(t *testing.T) { - fixture := newServiceFixture(t) - const ( - rawURL = "http://127.0.0.1/base?tenant_secret=known-query-secret" - plaintext = "known-upstream-key" - ) - var logs bytes.Buffer - previousOutput := logrus.StandardLogger().Out - logrus.SetOutput(&logs) - t.Cleanup(func() { logrus.SetOutput(previousOutput) }) - - result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: rawURL, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: plaintext, - }) - if err != nil { - t.Fatalf("CreateGroup() error = %v", err) - } - var row models.UpstreamKey - if err := fixture.db.Where("group_id = ?", result.GroupID).First(&row).Error; err != nil { - t.Fatalf("query created key: %v", err) - } - logText := logs.String() - if !strings.Contains(logText, "127.0.0.1") { - t.Fatalf("private warning missing hostname: %s", logText) - } - for _, required := range []string{"[CONTROL]"} { - if !strings.Contains(logText, required) { - t.Fatalf("private warning missing %q: %s", required, logText) - } - } - if strings.Contains(logText, "plane=control") { - t.Fatalf("private warning contains redundant plane field: %s", logText) - } - for _, forbidden := range []string{ - rawURL, - "known-query-secret", - plaintext, - row.KeyValue, - row.KeyHash, - "signature", - } { - if strings.Contains(logText, forbidden) { - t.Fatalf("private warning exposes %q: %s", forbidden, logText) - } - } -} - -func TestCreateGroupRejectsNameCollisionWithoutMutation(t *testing.T) { - fixture := newServiceFixture(t) - name := "shared-name" - first, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - Name: &name, - UpstreamURL: "https://first.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-first", - }) - if err != nil { - t.Fatalf("initial CreateGroup() error = %v", err) - } - beforeRevision := fixture.manager.Current().Revision - beforeCandidates := fixture.registry.CollectCandidates([]uint{first.GroupID}, nil, time.Time{}) - - _, err = fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - Name: &name, - UpstreamURL: "https://second.example.com", - Protocols: []protocol.Protocol{protocol.Anthropic}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-second", - }) - var apiErr *app_errors.APIError - if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrDuplicateResource.Code || - apiErr.HTTPStatus != http.StatusConflict { - t.Fatalf("name collision error = %#v", err) - } - assertGroupCount(t, fixture.db, 1) - if fixture.manager.Current().Revision != beforeRevision { - t.Fatal("name collision published a snapshot") - } - if !reflect.DeepEqual(fixture.registry.CollectCandidates([]uint{first.GroupID}, nil, time.Time{}), beforeCandidates) { - t.Fatal("name collision changed Registry") - } -} - -func TestCreateGroupModelsPresenceRejectsNull(t *testing.T) { - tests := []struct { - name string - body string - wantSet bool - wantCount int - wantError bool - }{ - { - name: "omitted", - body: `{"upstream_url":"https://example.com","protocols":["openai-completions"],"keys":"sk-a"}`, - }, - { - name: "empty", - body: `{"upstream_url":"https://example.com","protocols":["openai-completions"],"keys":"sk-a","models":[]}`, - wantSet: true, - }, - { - name: "values", - body: `{"upstream_url":"https://example.com","protocols":["openai-completions"],"keys":"sk-a","models":[{"id":" gpt-4o ","alias":" primary ","alias_enabled":true}]}`, - wantSet: true, - wantCount: 1, - }, - { - name: "missing alias enabled", - body: `{"upstream_url":"https://example.com","protocols":["openai-completions"],"keys":"sk-a","models":[{"id":"gpt-4o","alias":"primary"}]}`, - wantError: true, - }, - { - name: "null", - body: `{"upstream_url":"https://example.com","protocols":["openai-completions"],"keys":"sk-a","models":null}`, - wantError: true, - }, - { - name: "unknown model field", - body: `{"upstream_url":"https://example.com","protocols":["openai-completions"],"keys":"sk-a","models":[{"id":"gpt-4o","extra":true}]}`, - wantError: true, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - decoder := json.NewDecoder(strings.NewReader(test.body)) - decoder.DisallowUnknownFields() - var request GroupCreateRequest - err := decoder.Decode(&request) - if test.wantError { - if err == nil { - t.Fatal("Decode() error = nil") - } - return - } - if err != nil { - t.Fatalf("Decode() error = %v", err) - } - if request.Models.Set != test.wantSet || len(request.Models.Values) != test.wantCount { - t.Fatalf("Models = %#v, want set=%t count=%d", request.Models, test.wantSet, test.wantCount) - } - }) - } - - var modelValue optionalGroupModels - if err := modelValue.UnmarshalJSON([]byte(`[{"id":"gpt-4o","alias":""}] []`)); err == nil { - t.Fatal("UnmarshalJSON() accepted a trailing JSON value") - } -} - -func TestCreateGroupRequiresModelsAndAllowsEmptyArrayAtomically(t *testing.T) { - fixture := newServiceFixture(t) - beforeRevision := fixture.manager.Current().Revision - _, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://empty-models.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Keys: "sk-empty-models", - }) - if !errors.Is(err, app_errors.ErrValidation) { - t.Fatalf("CreateGroup(omitted models) error = %#v, want validation", err) - } - assertGroupCount(t, fixture.db, 0) - if fixture.manager.Current().Revision != beforeRevision { - t.Fatal("omitted models published a Snapshot") - } - - result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://empty-models.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-empty-models", - }) - if err != nil { - t.Fatalf("CreateGroup(empty models) error = %v", err) - } - if result.GroupID == 0 || result.KeysAdded != 1 || result.KeysDuplicated != 0 { - t.Fatalf("CreateGroup(empty models) result = %#v", result) - } - if stored := loadCreatedGroupModels(t, fixture, result.GroupID); len(stored) != 0 { - t.Fatalf("stored models = %#v, want empty", stored) - } - var group models.Group - if err := fixture.db.First(&group, result.GroupID).Error; err != nil { - t.Fatalf("query created Group: %v", err) - } - if string(group.Config) != `{}` { - t.Fatalf("stored config = %s, want empty override", group.Config) - } - if fixture.manager.Current().Revision != beforeRevision+1 { - t.Fatalf("snapshot revision = %d, want %d", fixture.manager.Current().Revision, beforeRevision+1) - } -} - -func TestCreateGroupRequiresSelectableProviderID(t *testing.T) { - fixture := newServiceFixture(t) - providerID := "private-company" - request := GroupCreateRequest{ - ProviderID: &providerID, - UpstreamURL: "https://private-company.example/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-private-company", - } - - if _, err := fixture.service.CreateGroup(t.Context(), request); !errors.Is(err, app_errors.ErrValidation) { - t.Fatalf("CreateGroup(unknown provider) error = %v, want ErrValidation", err) - } - assertGroupCount(t, fixture.db, 0) - - fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ - providerID: {ID: providerID}, - }}) - result, err := fixture.service.CreateGroup(t.Context(), request) - if err != nil { - t.Fatalf("CreateGroup(catalog provider) error = %v", err) - } - var stored models.Group - if err := fixture.db.First(&stored, result.GroupID).Error; err != nil { - t.Fatal(err) - } - if stored.ProviderID == nil || *stored.ProviderID != providerID { - t.Fatalf("stored provider_id = %v, want %q", stored.ProviderID, providerID) - } -} - -func TestCreateGroupRejectsInvalidNormalizedInputs(t *testing.T) { - fixture := newServiceFixture(t) - blank := " " - tooLong := strings.Repeat("名", 86) - controlName := "bad\nname" - emptyModels := optionalGroupModels{Set: true, Values: []GroupModel{}} - tests := []struct { - name string - request GroupCreateRequest - }{ - {name: "relative URL", request: GroupCreateRequest{UpstreamURL: "/v1", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: "sk"}}, - {name: "FTP URL", request: GroupCreateRequest{UpstreamURL: "ftp://example.com", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: "sk"}}, - {name: "userinfo URL", request: GroupCreateRequest{UpstreamURL: "https://user:pass@example.com", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: "sk"}}, - {name: "fragment URL", request: GroupCreateRequest{UpstreamURL: "https://example.com/#secret", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: "sk"}}, - {name: "missing host URL", request: GroupCreateRequest{UpstreamURL: "https:///v1", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: "sk"}}, - {name: "empty protocols", request: GroupCreateRequest{UpstreamURL: "https://example.com", Models: emptyModels, Keys: "sk"}}, - {name: "replaced completions protocol", request: GroupCreateRequest{UpstreamURL: "https://example.com", Protocols: []protocol.Protocol{"openai-chat-completions"}, Models: emptyModels, Keys: "sk"}}, - {name: "unknown protocol", request: GroupCreateRequest{UpstreamURL: "https://example.com", Protocols: []protocol.Protocol{"unknown"}, Models: emptyModels, Keys: "sk"}}, - {name: "empty keys", request: GroupCreateRequest{UpstreamURL: "https://example.com", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: " \n\t"}}, - {name: "too many keys", request: GroupCreateRequest{UpstreamURL: "https://example.com", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: groupCreateLines(1001)}}, - {name: "blank name", request: GroupCreateRequest{Name: &blank, UpstreamURL: "https://example.com", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: "sk"}}, - {name: "long name", request: GroupCreateRequest{Name: &tooLong, UpstreamURL: "https://example.com", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: "sk"}}, - {name: "control name", request: GroupCreateRequest{Name: &controlName, UpstreamURL: "https://example.com", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: emptyModels, Keys: "sk"}}, - {name: "blank model", request: GroupCreateRequest{UpstreamURL: "https://example.com", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: optionalGroupModels{Set: true, Values: []GroupModel{{ID: " "}}}, Keys: "sk"}}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - beforeRevision := fixture.manager.Current().Revision - _, err := fixture.service.CreateGroup(t.Context(), test.request) - var apiErr *app_errors.APIError - if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrValidation.Code { - t.Fatalf("CreateGroup() error = %#v", err) - } - assertGroupCount(t, fixture.db, 0) - if fixture.manager.Current().Revision != beforeRevision { - t.Fatal("invalid request published a snapshot") - } - }) - } -} - -func TestCreateGroupPreservesApprovedURLNormalizationRules(t *testing.T) { - tests := []struct { - name string - raw string - want string - }{ - {name: "trim lower host and slash", raw: " HTTPS://API.Example.COM/v1/ ", want: "https://api.example.com/v1"}, - {name: "preserve fixed query", raw: "https://API.Example.com/base/?api-version=2026-01", want: "https://api.example.com/base?api-version=2026-01"}, - {name: "preserve internal double slash", raw: "https://api.example.com/a//b/", want: "https://api.example.com/a//b"}, - {name: "IPv6 without port", raw: "https://[::1]/v1/", want: "https://[::1]/v1"}, - {name: "IPv6 with port", raw: "http://[2001:db8::1]:8080/v1/", want: "http://[2001:db8::1]:8080/v1"}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - got, _, err := normalizeUpstreamBaseURL(test.raw) - if err != nil { - t.Fatalf("normalizeUpstreamBaseURL(%q) error = %v", test.raw, err) - } - if got != test.want { - t.Fatalf("normalizeUpstreamBaseURL(%q) = %q, want %q", test.raw, got, test.want) - } - }) - } -} - -func TestCreateGroupPrivateHostClassificationUsesOnlyLiterals(t *testing.T) { - tests := []struct { - host string - want bool - }{ - {host: "localhost", want: true}, - {host: "LOCALHOST.", want: true}, - {host: "127.0.0.1", want: true}, - {host: "10.1.2.3", want: true}, - {host: "172.16.0.1", want: true}, - {host: "192.168.1.1", want: true}, - {host: "::1", want: true}, - {host: "fe80::1", want: true}, - {host: "0.0.0.0", want: true}, - {host: "8.8.8.8", want: false}, - {host: "api.example.com", want: false}, - } - for _, test := range tests { - t.Run(test.host, func(t *testing.T) { - if got := isLiteralPrivateHost(test.host); got != test.want { - t.Fatalf("isLiteralPrivateHost(%q) = %t, want %t", test.host, got, test.want) - } - }) - } -} - -func TestCreateGroupAllowsExactlyOneThousandNonEmptyKeyLines(t *testing.T) { - fixture := newServiceFixture(t) - result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://thousand.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: groupCreateLines(1000), - }) - if err != nil { - t.Fatalf("CreateGroup() error = %v", err) - } - if result.KeysAdded+result.KeysDuplicated != 1000 { - t.Fatalf("key line count = %d, want 1000", result.KeysAdded+result.KeysDuplicated) - } -} - -func TestCreateGroupRollsBackWhenKeyEncryptionFails(t *testing.T) { - fixture := newServiceFixture(t) - seed, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://existing-rollback.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-existing-rollback", - }) - if err != nil { - t.Fatalf("seed CreateGroup() error = %v", err) - } - var existingKey models.UpstreamKey - if err := fixture.db.Where("group_id = ?", seed.GroupID).First(&existingKey).Error; err != nil { - t.Fatalf("query seeded key: %v", err) - } - if count, ok := fixture.registry.IncrFailure(existingKey.ID); !ok || count != 1 { - t.Fatalf("first failure count = %d, %t", count, ok) - } - beforeCandidates := fixture.registry.CollectCandidates([]uint{seed.GroupID}, nil, time.Time{}) - fixture.service.encryption = groupCreateFailingEncryptService{Service: fixture.encryption} - beforeRevision := fixture.manager.Current().Revision - - _, err = fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://rollback.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{ - Set: true, Values: []GroupModel{{ID: "must-not-persist"}}, - }, - Keys: "sk-must-not-persist", - }) - if err == nil { - t.Fatal("CreateGroup() error = nil, want encryption failure") - } - assertGroupCount(t, fixture.db, 1) - if fixture.manager.Current().Revision != beforeRevision { - t.Fatal("failed creation published a snapshot") - } - if candidates := fixture.registry.CollectCandidates([]uint{seed.GroupID}, nil, time.Time{}); !reflect.DeepEqual(candidates, beforeCandidates) { - t.Fatalf("failed creation changed Registry: %#v", candidates) - } - if count, ok := fixture.registry.IncrFailure(existingKey.ID); !ok || count != 2 { - t.Fatalf("second failure count = %d, %t, want 2", count, ok) - } -} - -type groupCreateFailingEncryptService struct { - encryption.Service -} - -func (groupCreateFailingEncryptService) Encrypt(string) (string, error) { - return "", errors.New("forced encryption failure") -} - -func loadCreatedGroupModels(t *testing.T, fixture serviceFixture, groupID uint) []GroupModel { - t.Helper() - var group models.Group - if err := fixture.db.First(&group, groupID).Error; err != nil { - t.Fatalf("query group %d: %v", groupID, err) - } - var result []GroupModel - if err := json.Unmarshal(group.Models, &result); err != nil { - t.Fatalf("decode group %d models: %v", groupID, err) - } - if result == nil { - result = make([]GroupModel, 0) - } - return result -} - -func assertStoredGroupProtocols(t *testing.T, group models.Group, want []protocol.Protocol) { - t.Helper() - var got []protocol.Protocol - if err := json.Unmarshal(group.Protocols, &got); err != nil { - t.Fatalf("decode group %d protocols: %v", group.ID, err) - } - if !reflect.DeepEqual(got, want) { - t.Fatalf("group %d protocols = %#v, want %#v", group.ID, got, want) - } -} - -func groupCreateLines(count int) string { - var builder strings.Builder - for index := range count { - builder.WriteString("sk-") - builder.WriteString(strconv.Itoa(index)) - builder.WriteByte('\n') - } - return builder.String() -} diff --git a/internal/control/group_delete.go b/internal/control/group_delete.go index e26ff163..a604477d 100644 --- a/internal/control/group_delete.go +++ b/internal/control/group_delete.go @@ -62,20 +62,18 @@ func (s *Service) DeleteGroup(ctx context.Context, groupID uint) error { if groupID == 0 { return app_errors.ErrBadRequest } - var deletedKeyIDs []uint + var deletedCredentialIDs []uint providerReferencesChanged := false _, err := s.writeGroupConfig(ctx, func(tx *gorm.DB) error { var group models.Group if err := tx.Where("id = ?", groupID).Take(&group).Error; err != nil { return app_errors.ParseDBError(err) } - if group.ProviderID != nil { - var groupModels []GroupModel - if err := decodeGroupDiscoveryJSON(group.Models, &groupModels); err != nil { - return fmt.Errorf("decode group %d models: %w", groupID, app_errors.ErrInternalServer) - } - providerReferencesChanged = len(groupModels) > 0 + var groupModels []GroupModel + if err := decodeGroupDiscoveryJSON(group.Models, &groupModels); err != nil { + return fmt.Errorf("decode group %d models: %w", groupID, app_errors.ErrInternalServer) } + providerReferencesChanged = len(groupModels) > 0 references, err := explicitGroupReferences(tx, groupID) if err != nil { return err @@ -86,10 +84,10 @@ func (s *Service) DeleteGroup(ctx context.Context, groupID uint) error { GroupInUseData{AccessKeys: references}, ) } - if err := tx.Model(&models.UpstreamKey{}). + if err := tx.Model(&models.Credential{}). Where("group_id = ?", groupID). Order("id ASC"). - Pluck("id", &deletedKeyIDs).Error; err != nil { + Pluck("id", &deletedCredentialIDs).Error; err != nil { return app_errors.ParseDBError(err) } if err := tx.Delete(&group).Error; err != nil { @@ -98,9 +96,9 @@ func (s *Service) DeleteGroup(ctx context.Context, groupID uint) error { return nil }, func() error { s.registry.RemoveGroup(groupID) - for _, keyID := range deletedKeyIDs { - if _, exists := s.registry.EncryptedValue(keyID); exists { - return fmt.Errorf("deleted Registry key %d remains", keyID) + for _, credentialID := range deletedCredentialIDs { + if _, exists := s.registry.EncryptedCredentialData(credentialID); exists { + return fmt.Errorf("deleted Registry credential %d remains", credentialID) } } return nil diff --git a/internal/control/group_delete_test.go b/internal/control/group_delete_test.go index aff6d9c2..d377332b 100644 --- a/internal/control/group_delete_test.go +++ b/internal/control/group_delete_test.go @@ -23,7 +23,7 @@ import ( func TestDeleteGroupRejectsActiveAndDisabledExplicitAccessKeyReferences(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-in-use") + groupID := createGroupWithCredentials(t, fixture, "sk-in-use") other := validControlGroup("other-reference") if err := fixture.db.Create(other).Error; err != nil { t.Fatal(err) @@ -88,9 +88,9 @@ func TestDeleteGroupRejectsActiveAndDisabledExplicitAccessKeyReferences(t *testi func TestDeleteGroupCommitsCascadeThenRemovesRegistryThenPublishes(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-delete-a\nsk-delete-b") + groupID := createGroupWithCredentials(t, fixture, "sk-delete-a\nsk-delete-b") beforeRevision := fixture.manager.Current().Revision - var keyRows []models.UpstreamKey + var keyRows []models.Credential if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&keyRows).Error; err != nil { t.Fatal(err) } @@ -99,7 +99,7 @@ func TestDeleteGroupCommitsCascadeThenRemovesRegistryThenPublishes(t *testing.T) t.Fatalf("DeleteGroup() error = %v", err) } for _, row := range keyRows { - if value, ok := fixture.registry.EncryptedValue(row.ID); ok || value != "" { + if value, ok := fixture.registry.EncryptedCredentialData(row.ID); ok || value != "" { t.Fatalf("Registry key %d remains = %q, %t", row.ID, value, ok) } } @@ -112,7 +112,7 @@ func TestDeleteGroupCommitsCascadeThenRemovesRegistryThenPublishes(t *testing.T) if err := fixture.db.Model(&models.Group{}).Where("id = ?", groupID).Count(&groupCount).Error; err != nil { t.Fatal(err) } - if err := fixture.db.Model(&models.UpstreamKey{}).Where("group_id = ?", groupID).Count(&keyCount).Error; err != nil { + if err := fixture.db.Model(&models.Credential{}).Where("group_id = ?", groupID).Count(&keyCount).Error; err != nil { t.Fatal(err) } if groupCount != 0 || keyCount != 0 { @@ -122,14 +122,14 @@ func TestDeleteGroupCommitsCascadeThenRemovesRegistryThenPublishes(t *testing.T) func TestDeleteGroupCompileFailurePreservesDatabaseRegistryAndSnapshot(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-delete-rollback") + groupID := createGroupWithCredentials(t, fixture, "sk-delete-rollback") corrupt := validControlGroup("delete-corrupt-other") if err := fixture.db.Create(corrupt).Error; err != nil { t.Fatal(err) } if err := fixture.db.Exec( - "UPDATE groups SET protocols = ? WHERE id = ?", - `[]`, + "UPDATE groups SET channel_id = ? WHERE id = ?", + "unknown-channel", corrupt.ID, ).Error; err != nil { t.Fatal(err) @@ -156,8 +156,8 @@ func TestDeleteGroupCompileFailurePreservesDatabaseRegistryAndSnapshot(t *testin func TestDeleteGroupCommitFailurePreservesDatabaseRegistryAndSnapshot(t *testing.T) { fixture, dsn := newFileServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-delete-commit-busy") - var row models.UpstreamKey + groupID := createGroupWithCredentials(t, fixture, "sk-delete-commit-busy") + var row models.Credential if err := fixture.db.Where("group_id = ?", groupID).Take(&row).Error; err != nil { t.Fatal(err) } @@ -181,21 +181,21 @@ func TestDeleteGroupCommitFailurePreservesDatabaseRegistryAndSnapshot(t *testing Where("id = ?", groupID).Count(&groupCount).Error; err != nil { t.Fatal(err) } - if err := fixture.db.Model(&models.UpstreamKey{}). + if err := fixture.db.Model(&models.Credential{}). Where("id = ?", row.ID).Count(&keyCount).Error; err != nil { t.Fatal(err) } if groupCount != 1 || keyCount != 1 { t.Fatalf("rollback counts = group:%d key:%d", groupCount, keyCount) } - if value, exists := fixture.registry.EncryptedValue(row.ID); !exists || value != row.KeyValue { + if value, exists := fixture.registry.EncryptedCredentialData(row.ID); !exists || value != row.Data { t.Fatalf("Registry credential = %q, %t", value, exists) } } func TestDeleteGroupCorruptAccessKeyFiltersPreservesDatabaseRegistryAndSnapshot(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-delete-corrupt-filter") + groupID := createGroupWithCredentials(t, fixture, "sk-delete-corrupt-filter") corrupt := models.AccessKey{ Name: "Corrupt filters", KeyValue: "cipher-corrupt-filter", KeyHash: "hash-corrupt-filter", KeySuffix: "0004", Status: string(state.AccessKeyStatusDisabled), Filters: models.JSON(`{}`), @@ -229,8 +229,8 @@ func TestDeleteGroupCorruptAccessKeyFiltersPreservesDatabaseRegistryAndSnapshot( func TestDeleteGroupAcceptsRegistryAlreadyAtTargetStateAndRetainsHistory(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-delete-target-state") - var key models.UpstreamKey + groupID := createGroupWithCredentials(t, fixture, "sk-delete-target-state") + var key models.Credential if err := fixture.db.Where("group_id = ?", groupID).Take(&key).Error; err != nil { t.Fatal(err) } @@ -275,7 +275,7 @@ func TestDeleteGroupAcceptsRegistryAlreadyAtTargetStateAndRetainsHistory(t *test func TestDeleteGroupEndpointAuthenticationValidationNotFoundConflictAndSuccess(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-delete-http") + groupID := createGroupWithCredentials(t, fixture, "sk-delete-http") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) path := "/api/groups/" + strconv.FormatUint(uint64(groupID), 10) diff --git a/internal/control/group_detail.go b/internal/control/group_detail.go index 259253c3..16e3cb51 100644 --- a/internal/control/group_detail.go +++ b/internal/control/group_detail.go @@ -2,12 +2,14 @@ package control import ( "context" + "encoding/json" + "errors" "gorm.io/gorm" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/storage/models" ) @@ -25,14 +27,13 @@ type GroupEffectiveConfigResponse struct { // It deliberately excludes models and configuration that are loaded through focused // resources. type GroupSummaryResponse struct { - ID uint `json:"id"` - Name string `json:"name"` - ProviderID *string `json:"provider_id"` - ServiceStatus GroupCollectionStatus `json:"service_status"` - UpstreamURL string `json:"upstream_url"` - Protocols []protocol.Protocol `json:"protocols"` - KeyCount int64 `json:"key_count"` - ModelCount int `json:"model_count"` + ID uint `json:"id"` + Name string `json:"name"` + ChannelID channel.ID `json:"channel_id"` + Params json.RawMessage `json:"params"` + ServiceStatus GroupCollectionStatus `json:"service_status"` + CredentialCount int64 `json:"credential_count"` + ModelCount int `json:"model_count"` } func (s *Service) GetGroupSummary(ctx context.Context, groupID uint) (GroupSummaryResponse, error) { @@ -48,14 +49,11 @@ func (s *Service) GetGroupSummary(ctx context.Context, groupID uint) (GroupSumma continue } return GroupSummaryResponse{ - ID: record.ID, - Name: record.Name, - ProviderID: cloneString(record.ProviderID), - ServiceStatus: record.Status, - UpstreamURL: record.UpstreamURL, - Protocols: append([]protocol.Protocol(nil), record.Protocols...), - KeyCount: record.KeyCounts.Total, - ModelCount: int(record.ModelCount), + ID: record.ID, Name: record.Name, + ChannelID: record.ChannelID, Params: append(json.RawMessage(nil), record.Params...), + ServiceStatus: record.Status, + CredentialCount: record.CredentialCounts.Total, + ModelCount: int(record.ModelCount), }, nil } return GroupSummaryResponse{}, app_errors.ErrResourceNotFound @@ -89,6 +87,9 @@ func effectiveGroupConfig( func loadGroupRow(db *gorm.DB, groupID uint) (models.Group, error) { var group models.Group if err := db.Where("id = ?", groupID).Take(&group).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return models.Group{}, groupNotFoundError() + } return models.Group{}, app_errors.ParseDBError(err) } return group, nil diff --git a/internal/control/group_detail_ledger_routes_test.go b/internal/control/group_detail_ledger_routes_test.go index 580fbcd1..5ffa3cfc 100644 --- a/internal/control/group_detail_ledger_routes_test.go +++ b/internal/control/group_detail_ledger_routes_test.go @@ -10,6 +10,7 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" "gpt-load/internal/storage/models" ) @@ -18,9 +19,16 @@ func TestGroupDetailLedgerRoutesReplaceLegacyContracts(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) mustEnsureInitialPrices(t, fixture) - groupID := createGroupForKeyImport(t, fixture, "sk-ledger-route") - var key models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Take(&key).Error; err != nil { + created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("credential-ledger-route"), ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: optionalGroupModels{Set: true}, Credentials: "sk-ledger-route", + }) + if err != nil { + t.Fatal(err) + } + groupID := created.GroupID + var credential models.Credential + if err := fixture.db.Where("group_id = ?", groupID).Take(&credential).Error; err != nil { t.Fatal(err) } engine := gin.New() @@ -47,10 +55,10 @@ func TestGroupDetailLedgerRoutesReplaceLegacyContracts(t *testing.T) { {name: "models put", method: http.MethodPut, path: fmt.Sprintf("/api/groups/%d/models", groupID), body: `{}`}, {name: "settings get", method: http.MethodGet, path: fmt.Sprintf("/api/groups/%d/settings", groupID)}, {name: "settings put", method: http.MethodPut, path: fmt.Sprintf("/api/groups/%d/settings", groupID), body: `{}`}, - {name: "keys collection", method: http.MethodGet, path: fmt.Sprintf("/api/groups/%d/keys", groupID)}, - {name: "key update", method: http.MethodPut, path: fmt.Sprintf("/api/groups/%d/keys/%d", groupID, key.ID), body: `{}`}, - {name: "key restore", method: http.MethodPost, path: fmt.Sprintf("/api/groups/%d/keys/%d/restore", groupID, key.ID), body: `{}`}, - {name: "key batch", method: http.MethodPost, path: fmt.Sprintf("/api/groups/%d/keys/batch", groupID), body: `{}`}, + {name: "credentials collection", method: http.MethodGet, path: fmt.Sprintf("/api/groups/%d/credentials", groupID)}, + {name: "credential update", method: http.MethodPut, path: fmt.Sprintf("/api/groups/%d/credentials/%d", groupID, credential.ID), body: `{}`}, + {name: "credential restore", method: http.MethodPost, path: fmt.Sprintf("/api/groups/%d/credentials/%d/restore", groupID, credential.ID), body: `{}`}, + {name: "credential batch", method: http.MethodPost, path: fmt.Sprintf("/api/groups/%d/credentials/batch", groupID), body: `{}`}, } for _, route := range canonical { t.Run(route.name+" requires authentication", func(t *testing.T) { @@ -77,7 +85,7 @@ func TestGroupDetailLedgerRoutesReplaceLegacyContracts(t *testing.T) { t, engine, http.MethodPut, - fmt.Sprintf("/api/groups/%d/keys/%d", groupID, key.ID), + fmt.Sprintf("/api/groups/%d/credentials/%d", groupID, credential.ID), `{"weight_manual":25}`, "Bearer test-auth-key", ) @@ -89,19 +97,19 @@ func TestGroupDetailLedgerRoutesReplaceLegacyContracts(t *testing.T) { t.Fatal(err) } for _, field := range []string{ - "id", "mask", "configured_status", "effective_status", "weight_mode", "weight", + "credential_id", "mask", "configured_status", "effective_status", "weight_mode", "weight", "recent_success_count", "recent_failure_count", "consecutive_failure_count", "last_failure_category", "last_status_code", "cooldown_until_ms", "recovery", } { if _, exists := envelope.Data[field]; !exists { - t.Fatalf("updated key missing ledger field %q: %s", field, updated.Body.String()) + t.Fatalf("updated credential missing ledger field %q: %s", field, updated.Body.String()) } } for _, field := range []string{ "group_id", "status", "weight_manual", "weight_auto", "blacklisted", "failure_count", } { if _, exists := envelope.Data[field]; exists { - t.Fatalf("updated key exposes legacy field %q: %s", field, updated.Body.String()) + t.Fatalf("updated credential exposes legacy field %q: %s", field, updated.Body.String()) } } @@ -109,7 +117,7 @@ func TestGroupDetailLedgerRoutesReplaceLegacyContracts(t *testing.T) { t, engine, http.MethodPut, - fmt.Sprintf("/api/groups/%d/keys/%d", groupID, key.ID), + fmt.Sprintf("/api/groups/%d/credentials/%d", groupID, credential.ID), `{"weight_manual":0}`, "Bearer test-auth-key", ) diff --git a/internal/control/group_detail_test.go b/internal/control/group_detail_test.go index 80514a46..97fc19d1 100644 --- a/internal/control/group_detail_test.go +++ b/internal/control/group_detail_test.go @@ -10,6 +10,7 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" "gpt-load/internal/state" "gpt-load/internal/storage/models" @@ -20,13 +21,14 @@ func TestGetGroupSummaryUsesCollectionServiceStatusAndOnlyReturnsHeaderCounts(t available := createGroupCollectionGroup(t, fixture, "summary-available", true, nil) unavailable := createGroupCollectionGroup(t, fixture, "summary-unavailable", true, nil) disabled := createGroupCollectionGroup(t, fixture, "summary-disabled", false, nil) + setGroupCollectionChannel(t, fixture, unavailable, channel.Anthropic, models.JSON(`{}`)) setGroupCollectionRoute(t, fixture, available, `["openai-responses"]`, `[]`) setGroupCollectionRoute(t, fixture, unavailable, `["openai-completions"]`, `[]`) setGroupCollectionRoute(t, fixture, disabled, `["openai-responses"]`, `[]`) - publishGroupCollectionRuntime(t, fixture, []state.KeyEntry{ - createGroupCollectionKey(t, fixture, available.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, unavailable.ID, models.UpstreamKeyStatusActive, nil), - createGroupCollectionKey(t, fixture, disabled.ID, models.UpstreamKeyStatusActive, nil), + publishGroupCollectionRuntime(t, fixture, []state.CredentialEntry{ + createGroupCollectionKey(t, fixture, available.ID, models.CredentialStatusActive, nil), + createGroupCollectionKey(t, fixture, unavailable.ID, models.CredentialStatusActive, nil), + createGroupCollectionKey(t, fixture, disabled.ID, models.CredentialStatusActive, nil), }) for _, test := range []struct { @@ -44,7 +46,7 @@ func TestGetGroupSummaryUsesCollectionServiceStatusAndOnlyReturnsHeaderCounts(t if err != nil { t.Fatalf("GetGroupSummary() error = %v", err) } - if got.ID != test.groupID || got.ServiceStatus != test.wantStatus || got.KeyCount != test.wantKeys { + if got.ID != test.groupID || got.ServiceStatus != test.wantStatus || got.CredentialCount != test.wantKeys { t.Fatalf("GetGroupSummary() = %#v", got) } @@ -57,8 +59,8 @@ func TestGetGroupSummaryUsesCollectionServiceStatusAndOnlyReturnsHeaderCounts(t t.Fatal(err) } wantFields := map[string]struct{}{ - "id": {}, "name": {}, "provider_id": {}, "service_status": {}, "upstream_url": {}, - "protocols": {}, "key_count": {}, "model_count": {}, + "id": {}, "name": {}, "channel_id": {}, "params": {}, "service_status": {}, + "credential_count": {}, "model_count": {}, } for name := range fields { if _, exists := wantFields[name]; !exists { @@ -79,16 +81,17 @@ func TestGetGroupSummaryUsesCollectionServiceStatusAndOnlyReturnsHeaderCounts(t func TestGroupDetailStatusRequiresAHealthyKeyAndRouteCapability(t *testing.T) { fixture := newServiceFixture(t) group := createGroupCollectionGroup(t, fixture, "detail-zero-model-completions", true, nil) + setGroupCollectionChannel(t, fixture, group, channel.Anthropic, models.JSON(`{}`)) setGroupCollectionRoute(t, fixture, group, `["openai-completions"]`, `[]`) - publishGroupCollectionRuntime(t, fixture, []state.KeyEntry{ - createGroupCollectionKey(t, fixture, group.ID, models.UpstreamKeyStatusActive, nil), + publishGroupCollectionRuntime(t, fixture, []state.CredentialEntry{ + createGroupCollectionKey(t, fixture, group.ID, models.CredentialStatusActive, nil), }) got, err := fixture.service.GetGroupSummary(t.Context(), group.ID) if err != nil { t.Fatalf("GetGroupSummary() error = %v", err) } - if got.ServiceStatus != GroupCollectionStatusUnavailable || got.KeyCount != 1 || got.ModelCount != 0 { + if got.ServiceStatus != GroupCollectionStatusUnavailable || got.CredentialCount != 1 || got.ModelCount != 0 { t.Fatalf("GetGroupSummary() = %#v, want unavailable with one key and zero models", got) } } diff --git a/internal/control/group_idempotency.go b/internal/control/group_idempotency.go index 1f441d6c..4f8dbd0b 100644 --- a/internal/control/group_idempotency.go +++ b/internal/control/group_idempotency.go @@ -3,35 +3,33 @@ package control import ( "context" "encoding/json" - "fmt" - "sort" + "errors" "strconv" "github.com/sirupsen/logrus" "gorm.io/gorm" "gpt-load/internal/catalog" + "gpt-load/internal/channel" "gpt-load/internal/platform/canonicaljson" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/platform/utils" - "gpt-load/internal/protocol" "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" "gpt-load/internal/storage/models" ) type groupCreateDigestBody struct { - Name *string `json:"name"` - ProviderID *string `json:"provider_id"` - UpstreamURL string `json:"upstream_url"` - Protocols []protocol.Protocol `json:"protocols"` - Models []GroupModel `json:"models"` - Keys []string `json:"keys"` - ConfirmSameUpstreamURL bool `json:"confirm_same_upstream_url"` + Name *string `json:"name"` + ChannelID channel.ID `json:"channel_id"` + Params json.RawMessage `json:"params"` + Models []GroupModel `json:"models"` + Credentials []string `json:"credentials"` + ConfirmSameTarget bool `json:"confirm_same_target,omitempty"` } -type groupKeyImportDigestBody struct { - Keys []string `json:"keys"` +type credentialImportDigestBody struct { + Credentials []string `json:"credentials"` } func (s *Service) CreateGroupIdempotent( @@ -43,23 +41,19 @@ func (s *Service) CreateGroupIdempotent( if err != nil { return GroupCreateResult{}, err } - keyLines, err := normalizeIdempotencyKeyLines(request.Keys) + credentialLines, err := normalizeIdempotencyKeyLines(request.Credentials) if err != nil { return GroupCreateResult{}, err } - protocols := append([]protocol.Protocol(nil), normalized.protocols...) - sort.Slice(protocols, func(left, right int) bool { - return string(protocols[left]) < string(protocols[right]) - }) - canonicalBody, err := canonicalIdempotencyBody(groupCreateDigestBody{ - Name: normalized.explicitName, - ProviderID: normalized.providerID, - UpstreamURL: normalized.upstreamURL, - Protocols: protocols, - Models: append([]GroupModel(nil), normalized.models...), - Keys: keyLines, - ConfirmSameUpstreamURL: normalized.confirmSameUpstreamURL, - }) + digestBody := groupCreateDigestBody{ + Name: normalized.explicitName, + ChannelID: normalized.channelID, + Params: append(json.RawMessage(nil), normalized.params...), + Models: append([]GroupModel(nil), normalized.models...), + Credentials: credentialLines, + ConfirmSameTarget: normalized.confirmSameTarget, + } + canonicalBody, err := canonicalIdempotencyBody(digestBody) if err != nil { return GroupCreateResult{}, app_errors.ErrInternalServer } @@ -81,7 +75,7 @@ func (s *Service) CreateGroupIdempotent( logrus.WarnLevel, utils.LogPlaneControl, logrus.Fields{"host": normalized.hostname}, - "Creating upstream group with a private or local host", + "Creating channel group with a private or local host", ) } var catalogSnapshot *catalog.Snapshot @@ -97,64 +91,52 @@ func (s *Service) CreateGroupIdempotent( } }, Mutate: func(tx *gorm.DB) (idempotentMutationResult, error) { - if !normalized.confirmSameUpstreamURL { - conflicts, err := findGroupsByUpstreamURL(tx, normalized.upstreamURL) + if !normalized.confirmSameTarget { + conflicts, err := findGroupsByTarget(tx, normalized.channelID, normalized.params) if err != nil { return idempotentMutationResult{}, err } if len(conflicts) > 0 { return idempotentMutationResult{}, app_errors.NewAPIErrorWithData( - app_errors.ErrUpstreamURLConflict, - UpstreamURLConflictData{Groups: conflicts}, + app_errors.ErrChannelTargetConflict, + SameTargetConflictData{Groups: conflicts}, ) } } - name, err := resolveGroupCreateName( - tx, - normalized.explicitName, - normalized.hostname, - ) + name, err := resolveGroupCreateName(tx, normalized.explicitName, normalized.defaultName) if err != nil { return idempotentMutationResult{}, err } - encodedProtocols, err := json.Marshal(normalized.protocols) - if err != nil { - return idempotentMutationResult{}, app_errors.ErrInternalServer - } encodedModels, err := json.Marshal(normalized.models) if err != nil { return idempotentMutationResult{}, app_errors.ErrInternalServer } group := models.Group{ - Name: name, - ProviderID: cloneString(normalized.providerID), - UpstreamURL: normalized.upstreamURL, - Protocols: models.JSON(encodedProtocols), - Models: models.JSON(encodedModels), - Config: normalized.encodedConfig, - Enabled: true, + Name: name, + ChannelID: string(normalized.channelID), + Params: append(models.JSON(nil), normalized.params...), + Models: models.JSON(encodedModels), + Overrides: normalized.encodedOverrides, + Enabled: true, } if err := tx.Create(&group).Error; err != nil { return idempotentMutationResult{}, app_errors.ParseDBError(err) } - entries, added, duplicated, err := s.persistUpstreamKeys( - tx, - group.ID, - normalized.keys, - ) + added, duplicated, err := s.persistCredentials(tx, group.ID, normalized.credentials) if err != nil { return idempotentMutationResult{}, err } - if err := state.ValidateKeyEntries(entries); err != nil { - return idempotentMutationResult{}, fmt.Errorf( - "validate created group keys: %w", - err, - ) + entries, err := stateloader.BuildGroupCredentialEntries(ctx, tx, group.ID) + if err != nil { + return idempotentMutationResult{}, err + } + if err := state.ValidateCredentialEntries(entries); err != nil { + return idempotentMutationResult{}, err } if err := reconcileReferencedPrices(tx, catalogSnapshot); err != nil { return idempotentMutationResult{}, err } - input, err := stateloader.BuildCompileInput(ctx, tx) + input, err := stateloader.BuildCompileInput(ctx, tx, s.channelRegistry) if err != nil { return idempotentMutationResult{}, err } @@ -165,8 +147,10 @@ func (s *Service) CreateGroupIdempotent( return idempotentMutationResult{}, err } result := GroupCreateResult{ - GroupID: group.ID, GroupName: group.Name, - KeysAdded: added, KeysDuplicated: duplicated, + GroupID: group.ID, + GroupName: group.Name, + CredentialsAdded: added, + CredentialsDuplicated: duplicated, } canonicalResult, err := canonicaljson.Marshal(result) if err != nil { @@ -185,76 +169,75 @@ func (s *Service) CreateGroupIdempotent( if err := json.Unmarshal(operationResult.CanonicalResult, &result); err != nil { return GroupCreateResult{}, app_errors.ErrInternalServer } - if !operationResult.Replayed && normalized.providerID != nil && - len(normalized.models) > 0 && s.catalogSync != nil { + if !operationResult.Replayed && len(normalized.models) > 0 && s.catalogSync != nil { s.catalogSync.RequestGroupSync() } return result, nil } -func (s *Service) ImportGroupKeysIdempotent( +func (s *Service) ImportGroupCredentialsIdempotent( ctx context.Context, idempotencyKey string, groupID uint, - request GroupKeyImportRequest, -) (GroupKeyImportResult, error) { + request CredentialImportRequest, +) (CredentialImportResult, error) { if groupID == 0 { - return GroupKeyImportResult{}, app_errors.ErrValidation + return CredentialImportResult{}, app_errors.ErrValidation } - keys, err := s.normalizeUpstreamKeys(request.Keys) + credentialLines, err := normalizeIdempotencyKeyLines(request.Credentials) if err != nil { - return GroupKeyImportResult{}, err + return CredentialImportResult{}, err } - keyLines, err := normalizeIdempotencyKeyLines(request.Keys) + canonicalBody, err := canonicalIdempotencyBody(credentialImportDigestBody{Credentials: credentialLines}) if err != nil { - return GroupKeyImportResult{}, err - } - canonicalBody, err := canonicalIdempotencyBody( - groupKeyImportDigestBody{Keys: keyLines}, - ) - if err != nil { - return GroupKeyImportResult{}, app_errors.ErrInternalServer + return CredentialImportResult{}, app_errors.ErrInternalServer } resourceIdentity := "group:" + strconv.FormatUint(uint64(groupID), 10) digest, err := buildIdempotencyDigest(idempotencyDigestInput{ Version: 1, Method: "POST", - OperationKind: operationKindGroupKeyImport, - PathTemplate: "/api/groups/:group_id/keys/import", + OperationKind: operationKindCredentialImport, + PathTemplate: "/api/groups/:group_id/credentials/import", ResourceLocator: resourceIdentity, AuthScopeID: idempotencyAuthScopeID, CanonicalBody: canonicalBody, }) if err != nil { - return GroupKeyImportResult{}, app_errors.ErrInternalServer + return CredentialImportResult{}, app_errors.ErrInternalServer } operationResult, err := s.executeIdempotentOperation(ctx, idempotentOperationInput{ IdempotencyKey: idempotencyKey, DigestVersion: 1, RequestDigest: digest.Digest, - Kind: operationKindGroupKeyImport, + Kind: operationKindCredentialImport, Mutate: func(tx *gorm.DB) (idempotentMutationResult, error) { - var group struct{ ID uint } - if err := tx.Model(&models.Group{}). - Select("id"). - Where("id = ?", groupID). - Take(&group).Error; err != nil { + var group models.Group + if err := tx.Select("id", "channel_id").Where("id = ?", groupID).Take(&group).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return idempotentMutationResult{}, groupNotFoundError() + } return idempotentMutationResult{}, app_errors.ParseDBError(err) } - entries, added, duplicated, err := s.persistUpstreamKeys( - tx, - groupID, - keys, - ) + normalized, err := s.normalizeCredentials(channel.ID(group.ChannelID), request.Credentials) if err != nil { return idempotentMutationResult{}, err } - if err := state.ValidateKeyEntries(entries); err != nil { + added, duplicated, err := s.persistCredentials(tx, groupID, normalized) + if err != nil { return idempotentMutationResult{}, err } - result := GroupKeyImportResult{ - GroupID: groupID, KeysAdded: added, KeysDuplicated: duplicated, + entries, err := stateloader.BuildGroupCredentialEntries(ctx, tx, groupID) + if err != nil { + return idempotentMutationResult{}, err + } + if err := state.ValidateCredentialEntries(entries); err != nil { + return idempotentMutationResult{}, err + } + result := CredentialImportResult{ + GroupID: groupID, + CredentialsAdded: added, + CredentialsDuplicated: duplicated, } canonicalResult, err := canonicaljson.Marshal(result) if err != nil { @@ -267,11 +250,11 @@ func (s *Service) ImportGroupKeysIdempotent( }, }) if err != nil { - return GroupKeyImportResult{}, err + return CredentialImportResult{}, err } - var result GroupKeyImportResult + var result CredentialImportResult if err := json.Unmarshal(operationResult.CanonicalResult, &result); err != nil { - return GroupKeyImportResult{}, app_errors.ErrInternalServer + return CredentialImportResult{}, app_errors.ErrInternalServer } return result, nil } diff --git a/internal/control/group_idempotency_phase1_test.go b/internal/control/group_idempotency_phase1_test.go index 99906137..4bc27dd4 100644 --- a/internal/control/group_idempotency_phase1_test.go +++ b/internal/control/group_idempotency_phase1_test.go @@ -3,24 +3,23 @@ package control import ( "bytes" "encoding/json" - "errors" "reflect" "testing" + "gpt-load/internal/channel" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/storage/models" ) -func TestCreateGroupIdempotentReplaysOriginalCountsAndPreservesKeyMultiplicity(t *testing.T) { +func TestCreateGroupIdempotentReplaysOriginalCountsAndPreservesCredentialMultiplicity(t *testing.T) { fixture := newServiceFixture(t) fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x71}, 16)) request := GroupCreateRequest{ Name: stringPointer("idempotent-group"), - UpstreamURL: "https://idempotent.example.com/v1/", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://idempotent.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: " K \r\nK\n", + Credentials: " K \r\nK\n", } const key = "218f47a2-9c35-4d6e-8b1a-1234567890ab" @@ -28,7 +27,7 @@ func TestCreateGroupIdempotentReplaysOriginalCountsAndPreservesKeyMultiplicity(t if err != nil { t.Fatalf("first CreateGroupIdempotent() error = %v", err) } - if first.KeysAdded != 1 || first.KeysDuplicated != 1 { + if first.CredentialsAdded != 1 || first.CredentialsDuplicated != 1 { t.Fatalf("first result = %#v", first) } replayed, err := fixture.service.CreateGroupIdempotent(t.Context(), key, request) @@ -42,15 +41,15 @@ func TestCreateGroupIdempotentReplaysOriginalCountsAndPreservesKeyMultiplicity(t if err := fixture.db.Model(&models.Group{}).Count(&groupCount).Error; err != nil { t.Fatalf("count groups: %v", err) } - if err := fixture.db.Model(&models.UpstreamKey{}).Count(&keyCount).Error; err != nil { - t.Fatalf("count keys: %v", err) + if err := fixture.db.Model(&models.Credential{}).Count(&keyCount).Error; err != nil { + t.Fatalf("count credentials: %v", err) } if groupCount != 1 || keyCount != 1 { t.Fatalf("resource counts = group:%d key:%d, want 1/1", groupCount, keyCount) } different := request - different.Keys = "K" + different.Credentials = "K" _, err = fixture.service.CreateGroupIdempotent(t.Context(), key, different) assertAPIErrorCode(t, err, app_errors.ErrIdempotencyKeyReused.Code) } @@ -58,16 +57,16 @@ func TestCreateGroupIdempotentReplaysOriginalCountsAndPreservesKeyMultiplicity(t func TestCreateGroupIdempotentCanonicalizesDisabledAliasesAndReplaysNarrowResult(t *testing.T) { fixture := newServiceFixture(t) request := GroupCreateRequest{ - Name: stringPointer("canonical-models"), - UpstreamURL: "https://canonical-models.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + Name: stringPointer("canonical-models"), + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://canonical-models.example.com"}`), Models: optionalGroupModels{ Set: true, Values: []GroupModel{{ ID: "provider-model", Alias: "discarded-a", AliasEnabled: false, }}, }, - Keys: "key-one", + Credentials: "key-one", } const key = "238f47a2-9c35-4d6e-8b1a-1234567890ab" @@ -105,64 +104,10 @@ func TestCreateGroupIdempotentCanonicalizesDisabledAliasesAndReplaysNarrowResult if err := fixture.db.First(&group, first.GroupID).Error; err != nil { t.Fatalf("read created group: %v", err) } - if got := string(group.Config); got != `{}` { + if got := string(group.Overrides); got != `{}` { t.Fatalf("stored config = %s, want empty override", got) } if stored := loadCreatedGroupModels(t, fixture, first.GroupID); !reflect.DeepEqual(stored, []GroupModel{{ID: "provider-model"}}) { t.Fatalf("stored models = %#v, want disabled alias omitted", stored) } } - -func TestImportGroupKeysIdempotentReplaysOriginalCountsWithoutSecondInsert(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "seed-existing") - fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x72}, 16)) - request := GroupKeyImportRequest{Keys: " imported \n imported \n"} - const key = "228f47a2-9c35-4d6e-8b1a-1234567890ab" - - first, err := fixture.service.ImportGroupKeysIdempotent( - t.Context(), - key, - groupID, - request, - ) - if err != nil { - t.Fatalf("first ImportGroupKeysIdempotent() error = %v", err) - } - if first.KeysAdded != 1 || first.KeysDuplicated != 1 { - t.Fatalf("first result = %#v", first) - } - replayed, err := fixture.service.ImportGroupKeysIdempotent( - t.Context(), - key, - groupID, - request, - ) - if err != nil { - t.Fatalf("replay ImportGroupKeysIdempotent() error = %v", err) - } - if !reflect.DeepEqual(replayed, first) { - t.Fatalf("replay = %#v, first = %#v", replayed, first) - } - var count int64 - if err := fixture.db.Model(&models.UpstreamKey{}). - Where("group_id = ?", groupID). - Count(&count).Error; err != nil { - t.Fatalf("count group keys: %v", err) - } - if count != 2 { - t.Fatalf("group key count = %d, want seed + imported = 2", count) - } - - _, err = fixture.service.ImportGroupKeysIdempotent( - t.Context(), - key, - groupID, - GroupKeyImportRequest{Keys: "imported"}, - ) - var apiErr *app_errors.APIError - if !errors.As(err, &apiErr) || - apiErr.Code != app_errors.ErrIdempotencyKeyReused.Code { - t.Fatalf("different multiplicity error = %v", err) - } -} diff --git a/internal/control/group_key_collection.go b/internal/control/group_key_collection.go deleted file mode 100644 index 472f55f1..00000000 --- a/internal/control/group_key_collection.go +++ /dev/null @@ -1,357 +0,0 @@ -package control - -import ( - "context" - "fmt" - "net/url" - "sort" - "strconv" - "strings" - "time" - - "gpt-load/internal/health" - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/platform/utils" - "gpt-load/internal/state" -) - -const ( - groupKeyCollectionDefaultPage = 1 - groupKeyCollectionDefaultPageSize = 50 - groupKeyCollectionStatsWindow = int64(health.StatsWindow / time.Second) -) - -type GroupKeyCollectionQuery struct { - Query string - Status *string - Page int - PageSize int -} - -type GroupKeyCollectionResponse struct { - ObservedAtMS int64 `json:"observed_at_ms"` - StatsWindowSeconds int64 `json:"stats_window_seconds"` - Summary GroupKeySummaryResponse `json:"summary"` - Items []GroupKeyItemResponse `json:"items"` - Pagination GroupKeyPaginationResponse `json:"pagination"` -} - -type GroupKeySummaryResponse struct { - Total int `json:"total"` - Available int `json:"available"` - Cooldown int `json:"cooldown"` - Blacklisted int `json:"blacklisted"` - Disabled int `json:"disabled"` -} - -type GroupKeyItemResponse struct { - ID uint `json:"id"` - Mask string `json:"mask"` - ConfiguredStatus string `json:"configured_status"` - EffectiveStatus string `json:"effective_status"` - WeightMode string `json:"weight_mode"` - Weight *int `json:"weight"` - RecentSuccessCount uint64 `json:"recent_success_count"` - RecentFailureCount uint64 `json:"recent_failure_count"` - ConsecutiveFailureCount uint64 `json:"consecutive_failure_count"` - LastFailureCategory string `json:"last_failure_category"` - LastStatusCode *int `json:"last_status_code"` - CooldownUntilMS *int64 `json:"cooldown_until_ms"` - Recovery GroupKeyRecoveryResponse `json:"recovery"` -} - -type GroupKeyRecoveryResponse struct { - Mode string `json:"mode"` - Automatic bool `json:"automatic"` - AtMS *int64 `json:"at_ms"` -} - -type GroupKeyPaginationResponse struct { - Page int `json:"page"` - PageSize int `json:"page_size"` - TotalItems int `json:"total_items"` - TotalPages int `json:"total_pages"` -} - -type groupKeyCollectionRecord struct { - item GroupKeyItemResponse - bucket healthBucket -} - -func parseGroupKeyCollectionQuery(rawQuery string) (GroupKeyCollectionQuery, *app_errors.APIError) { - query := GroupKeyCollectionQuery{ - Page: groupKeyCollectionDefaultPage, PageSize: groupKeyCollectionDefaultPageSize, - } - values, err := url.ParseQuery(rawQuery) - if err != nil { - return GroupKeyCollectionQuery{}, app_errors.ErrBadRequest - } - for key, entries := range values { - switch key { - case "q", "status", "page", "page_size": - default: - return GroupKeyCollectionQuery{}, app_errors.ErrBadRequest - } - if len(entries) != 1 { - return GroupKeyCollectionQuery{}, app_errors.ErrBadRequest - } - } - if entries, exists := values["q"]; exists { - query.Query = strings.TrimSpace(entries[0]) - } - if entries, exists := values["status"]; exists { - status := entries[0] - switch status { - case string(healthBucketAvailable), string(healthBucketCooldown), string(healthBucketBlacklisted), string(healthBucketDisabled): - query.Status = &status - default: - return GroupKeyCollectionQuery{}, app_errors.ErrBadRequest - } - } - if entries, exists := values["page"]; exists { - page, ok := parseGroupKeyCollectionPositiveInt(entries[0]) - if !ok { - return GroupKeyCollectionQuery{}, app_errors.ErrBadRequest - } - query.Page = page - } - if entries, exists := values["page_size"]; exists { - pageSize, ok := parseGroupKeyCollectionPositiveInt(entries[0]) - if !ok || (pageSize != 20 && pageSize != 50 && pageSize != 100) { - return GroupKeyCollectionQuery{}, app_errors.ErrBadRequest - } - query.PageSize = pageSize - } - return query, nil -} - -func parseGroupKeyCollectionPositiveInt(value string) (int, bool) { - if value == "" { - return 0, false - } - for index := range len(value) { - if value[index] < '0' || value[index] > '9' { - return 0, false - } - } - parsed, err := strconv.Atoi(value) - if err != nil || parsed < 1 { - return 0, false - } - return parsed, true -} - -func (s *Service) ListGroupKeys( - ctx context.Context, - groupID uint, - query GroupKeyCollectionQuery, -) (GroupKeyCollectionResponse, error) { - if groupID == 0 { - return GroupKeyCollectionResponse{}, app_errors.ErrBadRequest - } - if query.Page < 1 || (query.PageSize != 20 && query.PageSize != 50 && query.PageSize != 100) { - return GroupKeyCollectionResponse{}, app_errors.ErrBadRequest - } - capture, err := s.captureGroupKeys(ctx, groupID) - if err != nil { - return GroupKeyCollectionResponse{}, err - } - observation, err := validateGroupKeysCapture(capture) - if err != nil { - return GroupKeyCollectionResponse{}, err - } - return s.mapGroupKeyCollection(observation, query) -} - -func (s *Service) mapGroupKeyCollection( - observation groupKeysObservation, - query GroupKeyCollectionQuery, -) (GroupKeyCollectionResponse, error) { - if s == nil || s.encryption == nil || s.stats == nil { - return GroupKeyCollectionResponse{}, fmt.Errorf("group key collection dependencies unavailable: %w", app_errors.ErrInternalServer) - } - observedAtMS, err := safeEpochMilliseconds(observation.observedAt) - if err != nil { - return GroupKeyCollectionResponse{}, fmt.Errorf("map group key collection observed_at_ms: %w", err) - } - group := state.GroupCatalogView{ - ID: observation.group.ID, - Name: observation.group.Name, - Enabled: observation.group.Enabled, - WeightManual: cloneInt(observation.group.WeightManual), - } - records := make([]groupKeyCollectionRecord, 0, len(observation.rows)) - for _, row := range observation.rows { - view := observation.runtime[row.ID] - plaintext, err := s.encryption.Decrypt(row.KeyValue) - if err != nil { - return GroupKeyCollectionResponse{}, fmt.Errorf("decrypt group key %d: %w", row.ID, app_errors.ErrInternalServer) - } - mask, err := maskGroupKeyCollection(plaintext) - if err != nil { - return GroupKeyCollectionResponse{}, fmt.Errorf("mask group key %d: %w", row.ID, err) - } - bucket := classifyHealthKey(group, view, observation.observedAt) - item, err := mapGroupKeyCollectionItem(mask, row.ID, view, bucket, s.stats.Snapshot(row.ID, observation.observedAt), observation.observedAt) - if err != nil { - return GroupKeyCollectionResponse{}, err - } - records = append(records, groupKeyCollectionRecord{item: item, bucket: bucket}) - } - - summary := summarizeGroupKeyCollection(records) - filtered := make([]groupKeyCollectionRecord, 0, len(records)) - for _, record := range records { - if groupKeyCollectionMatches(record, query) { - filtered = append(filtered, record) - } - } - sort.Slice(filtered, func(leftIndex, rightIndex int) bool { - left, right := filtered[leftIndex], filtered[rightIndex] - if groupKeyCollectionBucketOrder(left.bucket) != groupKeyCollectionBucketOrder(right.bucket) { - return groupKeyCollectionBucketOrder(left.bucket) < groupKeyCollectionBucketOrder(right.bucket) - } - return left.item.ID < right.item.ID - }) - - totalItems := len(filtered) - pagination := GroupKeyPaginationResponse{ - Page: query.Page, PageSize: query.PageSize, TotalItems: totalItems, - TotalPages: groupKeyCollectionTotalPages(totalItems, query.PageSize), - } - items := groupKeyCollectionPage(filtered, query.Page, query.PageSize) - return GroupKeyCollectionResponse{ - ObservedAtMS: observedAtMS, StatsWindowSeconds: groupKeyCollectionStatsWindow, - Summary: summary, Items: items, Pagination: pagination, - }, nil -} - -func maskGroupKeyCollection(plaintext string) (string, error) { - mask := utils.MaskAPIKey(plaintext) - if mask == "" { - return "", fmt.Errorf("credential is empty: %w", app_errors.ErrInternalServer) - } - return mask, nil -} - -func mapGroupKeyCollectionItem( - mask string, - keyID uint, - view state.KeyRuntimeView, - bucket healthBucket, - stats health.KeyStats, - observedAt time.Time, -) (GroupKeyItemResponse, error) { - weightMode := "auto" - if view.WeightManual != nil { - weightMode = "manual" - } - item := GroupKeyItemResponse{ - ID: keyID, Mask: mask, ConfiguredStatus: string(view.Status), EffectiveStatus: string(bucket), - WeightMode: weightMode, RecentSuccessCount: stats.Success, RecentFailureCount: stats.Failure, - ConsecutiveFailureCount: stats.ConsecutiveFailure, - LastFailureCategory: normalizeGroupKeyFailureCategory(stats.LastFailureCategory).String(), - LastStatusCode: optionalHealthStatusCode(stats.LastStatusCode), - } - switch bucket { - case healthBucketAvailable: - weight := view.WeightAuto - if view.WeightManual != nil { - weight = *view.WeightManual - } - if weight < 1 || weight > state.MaxWeight { - return GroupKeyItemResponse{}, fmt.Errorf("map group key %d weight: %w", keyID, app_errors.ErrInternalServer) - } - item.Weight = &weight - item.Recovery = GroupKeyRecoveryResponse{Mode: "none"} - case healthBucketCooldown: - cooldownUntilMS, err := optionalSafeEpochMilliseconds(view.CooldownUntil) - if err != nil { - return GroupKeyItemResponse{}, fmt.Errorf("map group key %d cooldown_until_ms: %w", keyID, err) - } - if cooldownUntilMS == nil || !view.CooldownUntil.After(observedAt) { - return GroupKeyItemResponse{}, fmt.Errorf("map group key %d invalid cooldown: %w", keyID, app_errors.ErrInternalServer) - } - item.CooldownUntilMS = cooldownUntilMS - item.Recovery = GroupKeyRecoveryResponse{Mode: "cooldown", Automatic: true, AtMS: cooldownUntilMS} - case healthBucketBlacklisted: - item.Recovery = GroupKeyRecoveryResponse{Mode: "probe", Automatic: true} - case healthBucketDisabled: - item.Recovery = GroupKeyRecoveryResponse{Mode: "manual"} - default: - return GroupKeyItemResponse{}, fmt.Errorf("map group key %d unknown status: %w", keyID, app_errors.ErrInternalServer) - } - return item, nil -} - -func normalizeGroupKeyFailureCategory(category health.FailureCategory) health.FailureCategory { - if category == health.FailureCategoryOK { - return health.FailureCategoryAmbiguous - } - return category -} - -func summarizeGroupKeyCollection(records []groupKeyCollectionRecord) GroupKeySummaryResponse { - summary := GroupKeySummaryResponse{Total: len(records)} - for _, record := range records { - switch record.bucket { - case healthBucketAvailable: - summary.Available++ - case healthBucketCooldown: - summary.Cooldown++ - case healthBucketBlacklisted: - summary.Blacklisted++ - case healthBucketDisabled: - summary.Disabled++ - } - } - return summary -} - -func groupKeyCollectionMatches(record groupKeyCollectionRecord, query GroupKeyCollectionQuery) bool { - if query.Status != nil && record.item.EffectiveStatus != *query.Status { - return false - } - return query.Query == "" || strings.Contains(strings.ToLower(record.item.Mask), strings.ToLower(query.Query)) -} - -func groupKeyCollectionBucketOrder(bucket healthBucket) int { - switch bucket { - case healthBucketBlacklisted: - return 0 - case healthBucketCooldown: - return 1 - case healthBucketAvailable: - return 2 - case healthBucketDisabled: - return 3 - default: - return 4 - } -} - -func groupKeyCollectionTotalPages(totalItems, pageSize int) int { - if totalItems == 0 { - return 0 - } - return (totalItems + pageSize - 1) / pageSize -} - -func groupKeyCollectionPage(records []groupKeyCollectionRecord, page, pageSize int) []GroupKeyItemResponse { - if page-1 > len(records)/pageSize { - return []GroupKeyItemResponse{} - } - offset := (page - 1) * pageSize - if offset < 0 || offset >= len(records) { - return []GroupKeyItemResponse{} - } - end := offset + pageSize - if end < offset || end > len(records) { - end = len(records) - } - items := make([]GroupKeyItemResponse, end-offset) - for index, record := range records[offset:end] { - items[index] = record.item - } - return items -} diff --git a/internal/control/group_key_collection_test.go b/internal/control/group_key_collection_test.go deleted file mode 100644 index 322f9e8a..00000000 --- a/internal/control/group_key_collection_test.go +++ /dev/null @@ -1,300 +0,0 @@ -package control - -import ( - "bytes" - "encoding/json" - "errors" - "fmt" - "net/http" - "reflect" - "strconv" - "testing" - "time" - - "github.com/gin-gonic/gin" - - "gpt-load/internal/health" - "gpt-load/internal/platform/config" - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/storage/models" -) - -func TestGroupKeyCollectionQuery(t *testing.T) { - initControlI18n(t) - fixture := newServiceFixture(t) - group := validControlGroup("key-collection-query") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - engine := gin.New() - NewServer(&config.Config{AuthKey: groupKeyHTTPAuth}, fixture.service).RegisterRoutes(engine) - path := fmt.Sprintf("/api/groups/%d/keys", group.ID) - - for _, test := range []struct { - name string - query string - wantCode int - wantObject bool - }{ - {name: "empty", wantCode: http.StatusOK, wantObject: true}, - {name: "mask search", query: "?q=sk-gl-%2A%2A%2A%2Aabcd", wantCode: http.StatusOK, wantObject: true}, - {name: "available", query: "?status=available", wantCode: http.StatusOK, wantObject: true}, - {name: "cooldown", query: "?status=cooldown", wantCode: http.StatusOK, wantObject: true}, - {name: "blacklisted", query: "?status=blacklisted", wantCode: http.StatusOK, wantObject: true}, - {name: "disabled", query: "?status=disabled", wantCode: http.StatusOK, wantObject: true}, - {name: "minimum page", query: "?page=1", wantCode: http.StatusOK, wantObject: true}, - {name: "page size 20", query: "?page_size=20", wantCode: http.StatusOK, wantObject: true}, - {name: "page size 50", query: "?page_size=50", wantCode: http.StatusOK, wantObject: true}, - {name: "page size 100", query: "?page_size=100", wantCode: http.StatusOK, wantObject: true}, - {name: "duplicate parameter", query: "?page=1&page=2", wantCode: http.StatusBadRequest}, - {name: "unknown parameter", query: "?unexpected=value", wantCode: http.StatusBadRequest}, - {name: "unknown status", query: "?status=unknown", wantCode: http.StatusBadRequest}, - {name: "zero page", query: "?page=0", wantCode: http.StatusBadRequest}, - {name: "unsupported page size", query: "?page_size=21", wantCode: http.StatusBadRequest}, - } { - t.Run(test.name, func(t *testing.T) { - recorder := serveGroupKeyHTTPRequest( - t, engine, http.MethodGet, path+test.query, "", groupKeyHTTPAuth, "en-US", - ) - if recorder.Code != test.wantCode { - t.Fatalf("GET %s = %d %s, want %d", test.query, recorder.Code, recorder.Body.String(), test.wantCode) - } - if !test.wantObject { - return - } - var envelope struct { - Code int `json:"code"` - Data json.RawMessage `json:"data"` - } - if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil { - t.Fatal(err) - } - if envelope.Code != 0 || len(envelope.Data) == 0 || envelope.Data[0] != '{' { - t.Fatalf("collection envelope = %s, want object data", recorder.Body.String()) - } - }) - } - - unauthenticated := serveGroupKeyHTTPRequest( - t, engine, http.MethodGet, path+"?unexpected=value", "", "wrong-auth-key", "en-US", - ) - if unauthenticated.Code != http.StatusUnauthorized { - t.Fatalf("unauthenticated request = %d %s, want %d before query handling", unauthenticated.Code, unauthenticated.Body.String(), http.StatusUnauthorized) - } -} - -func TestListGroupKeysCollection(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.August, 1, 10, 0, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - group := validControlGroup("key-collection") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - - manualWeight := 73 - blacklistedOne := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-0001", models.UpstreamKeyStatusActive, nil) - blacklistedTwo := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-0002", models.UpstreamKeyStatusActive, nil) - cooldownOne := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-0003", models.UpstreamKeyStatusActive, nil) - cooldownTwo := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-0004", models.UpstreamKeyStatusActive, nil) - availableAuto := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-0005", models.UpstreamKeyStatusActive, nil) - availableManual := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-0006", models.UpstreamKeyStatusActive, &manualWeight) - disabledOne := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-0007", models.UpstreamKeyStatusDisabled, nil) - disabledTwo := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-0008", models.UpstreamKeyStatusDisabled, nil) - - if !fixture.registry.SetBlacklisted(blacklistedOne.ID) || !fixture.registry.SetBlacklisted(blacklistedTwo.ID) || - !fixture.registry.SetCooldown(cooldownOne.ID, now.Add(2*time.Minute)) || - !fixture.registry.SetCooldown(cooldownTwo.ID, now.Add(time.Minute)) || - !fixture.registry.SetAutoWeight(availableAuto.ID, 61) { - t.Fatal("seed runtime state") - } - fixture.stats.RecordSuccess(blacklistedOne.ID, now.Add(-2*time.Minute)) - fixture.stats.RecordFailure(blacklistedOne.ID, health.FailureCategoryRateLimited, http.StatusTooManyRequests, now.Add(-time.Minute)) - fixture.stats.RecordFailure(blacklistedOne.ID, health.FailureCategoryInvalidKey, http.StatusUnauthorized, now) - - response, err := fixture.service.ListGroupKeys(t.Context(), group.ID, GroupKeyCollectionQuery{Page: 1, PageSize: 20}) - if err != nil { - t.Fatalf("ListGroupKeys() error = %v", err) - } - if response.ObservedAtMS != now.UnixMilli() || response.StatsWindowSeconds != 300 { - t.Fatalf("observation metadata = %#v", response) - } - if response.Summary != (GroupKeySummaryResponse{Total: 8, Available: 2, Cooldown: 2, Blacklisted: 2, Disabled: 2}) { - t.Fatalf("summary = %#v", response.Summary) - } - wantIDs := []uint{ - blacklistedOne.ID, blacklistedTwo.ID, - cooldownOne.ID, cooldownTwo.ID, - availableAuto.ID, availableManual.ID, - disabledOne.ID, disabledTwo.ID, - } - gotIDs := make([]uint, len(response.Items)) - for index, item := range response.Items { - gotIDs[index] = item.ID - } - if !reflect.DeepEqual(gotIDs, wantIDs) { - t.Fatalf("sorted ids = %v, want %v", gotIDs, wantIDs) - } - - blacklisted := response.Items[0] - if blacklisted.Mask != "prov****0001" || blacklisted.ConfiguredStatus != "active" || - blacklisted.EffectiveStatus != "blacklisted" || blacklisted.Weight != nil || - blacklisted.RecentSuccessCount != 1 || blacklisted.RecentFailureCount != 2 || - blacklisted.ConsecutiveFailureCount != 2 || blacklisted.LastFailureCategory != "invalid_key" || - blacklisted.LastStatusCode == nil || *blacklisted.LastStatusCode != http.StatusUnauthorized || - blacklisted.CooldownUntilMS != nil || blacklisted.Recovery.Mode != "probe" || - !blacklisted.Recovery.Automatic || blacklisted.Recovery.AtMS != nil { - t.Fatalf("blacklisted item = %#v", blacklisted) - } - cooldown := response.Items[2] - if cooldown.EffectiveStatus != "cooldown" || cooldown.Weight != nil || cooldown.CooldownUntilMS == nil || - *cooldown.CooldownUntilMS != now.Add(2*time.Minute).UnixMilli() || cooldown.Recovery.Mode != "cooldown" || - !cooldown.Recovery.Automatic || cooldown.Recovery.AtMS == nil || *cooldown.Recovery.AtMS != *cooldown.CooldownUntilMS { - t.Fatalf("cooldown item = %#v", cooldown) - } - auto := response.Items[4] - if auto.EffectiveStatus != "available" || auto.WeightMode != "auto" || auto.Weight == nil || *auto.Weight != 61 || - auto.Recovery.Mode != "none" || auto.Recovery.Automatic || auto.Recovery.AtMS != nil { - t.Fatalf("automatic available item = %#v", auto) - } - manual := response.Items[5] - if manual.EffectiveStatus != "available" || manual.WeightMode != "manual" || manual.Weight == nil || *manual.Weight != manualWeight { - t.Fatalf("manual available item = %#v", manual) - } - disabled := response.Items[6] - if disabled.ConfiguredStatus != "disabled" || disabled.EffectiveStatus != "disabled" || disabled.Weight != nil || - disabled.Recovery.Mode != "manual" || disabled.Recovery.Automatic || disabled.Recovery.AtMS != nil { - t.Fatalf("disabled item = %#v", disabled) - } - - blacklistedStatus := "blacklisted" - filtered, err := fixture.service.ListGroupKeys(t.Context(), group.ID, GroupKeyCollectionQuery{ - Status: &blacklistedStatus, Page: 1, PageSize: 20, - }) - if err != nil { - t.Fatalf("ListGroupKeys(status) error = %v", err) - } - if filtered.Summary != response.Summary || filtered.Pagination.TotalItems != 2 || filtered.Pagination.TotalPages != 1 || - len(filtered.Items) != 2 || filtered.Items[0].ID != blacklistedOne.ID || filtered.Items[1].ID != blacklistedTwo.ID { - t.Fatalf("filtered response = %#v", filtered) - } - - matched, err := fixture.service.ListGroupKeys(t.Context(), group.ID, GroupKeyCollectionQuery{ - Query: "PROV****0006", Page: 1, PageSize: 20, - }) - if err != nil { - t.Fatalf("ListGroupKeys(q) error = %v", err) - } - if len(matched.Items) != 1 || matched.Items[0].ID != availableManual.ID { - t.Fatalf("mask query response = %#v", matched) - } - - outOfRange, err := fixture.service.ListGroupKeys(t.Context(), group.ID, GroupKeyCollectionQuery{Page: 2, PageSize: 20}) - if err != nil { - t.Fatalf("ListGroupKeys(out of range) error = %v", err) - } - if outOfRange.Pagination.TotalItems != 8 || outOfRange.Pagination.TotalPages != 1 || len(outOfRange.Items) != 0 || outOfRange.Items == nil { - t.Fatalf("out of range response = %#v", outOfRange) - } - - encoded, err := json.Marshal(response) - if err != nil { - t.Fatal(err) - } - for _, forbidden := range []string{"provider-secret-", "key_value", "key_hash", "ciphertext", blacklistedOne.KeyValue, blacklistedOne.KeyHash} { - if bytes.Contains(encoded, []byte(forbidden)) { - t.Fatalf("collection response exposes %q: %s", forbidden, encoded) - } - } -} - -func TestListGroupKeysCollectionMasksImportedShortKeysWithoutLeakingPlaintext(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-collection-short-mask") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - for _, plaintext := range []string{"x", "xy", "xyz"} { - if _, err := fixture.service.ImportGroupKeys( - t.Context(), - group.ID, - GroupKeyImportRequest{Keys: plaintext}, - ); err != nil { - t.Fatalf("ImportGroupKeys(%q) error = %v", plaintext, err) - } - } - - response, err := fixture.service.ListGroupKeys( - t.Context(), - group.ID, - GroupKeyCollectionQuery{Page: 1, PageSize: 20}, - ) - if err != nil { - t.Fatalf("ListGroupKeys() error = %v", err) - } - if len(response.Items) != 3 { - t.Fatalf("ListGroupKeys() items = %d, want 3", len(response.Items)) - } - for _, item := range response.Items { - if item.Mask != "****" { - t.Fatalf("key %d mask = %q, want full mask", item.ID, item.Mask) - } - } - encoded, err := json.Marshal(response) - if err != nil { - t.Fatal(err) - } - for _, plaintext := range []string{"x", "xy", "xyz"} { - if bytes.Contains(encoded, []byte(`"mask":"`+plaintext+`"`)) { - t.Fatalf("collection response exposes plaintext %q: %s", plaintext, encoded) - } - } -} - -func TestListGroupKeysCollectionFailsAtomicallyWhenMaskingCannotDecrypt(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-collection-decrypt") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - first := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-first", models.UpstreamKeyStatusActive, nil) - second := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-second", models.UpstreamKeyStatusActive, nil) - if err := fixture.db.Model(&models.UpstreamKey{}).Where("id = ?", second.ID).Update("key_value", "not-valid-ciphertext").Error; err != nil { - t.Fatal(err) - } - - response, err := fixture.service.ListGroupKeys(t.Context(), group.ID, GroupKeyCollectionQuery{Page: 1, PageSize: 20}) - if !errors.Is(err, app_errors.ErrInternalServer) { - t.Fatalf("ListGroupKeys() error = %v, want internal error", err) - } - if !reflect.DeepEqual(response, GroupKeyCollectionResponse{}) { - t.Fatalf("ListGroupKeys() response = %#v, want zero response", response) - } - if first.ID == 0 { - t.Fatal("first key was not persisted") - } -} - -func TestListGroupKeysCollectionReturnsEmptyItemsForOverflowingOffset(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-collection-overflow") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-overflow", models.UpstreamKeyStatusActive, nil) - - overflowPage := 1 + (1 << (strconv.IntSize - 2)) - query, apiErr := parseGroupKeyCollectionQuery(fmt.Sprintf("page=%d&page_size=20", overflowPage)) - if apiErr != nil { - t.Fatalf("parseGroupKeyCollectionQuery() error = %v", apiErr) - } - response, err := fixture.service.ListGroupKeys(t.Context(), group.ID, query) - if err != nil { - t.Fatalf("ListGroupKeys() error = %v", err) - } - if response.Pagination != (GroupKeyPaginationResponse{ - Page: overflowPage, PageSize: 20, TotalItems: 1, TotalPages: 1, - }) || response.Items == nil || len(response.Items) != 0 { - t.Fatalf("overflow page response = %#v, want empty out-of-range page", response) - } -} diff --git a/internal/control/group_key_mutations.go b/internal/control/group_key_mutations.go deleted file mode 100644 index 726c66f6..00000000 --- a/internal/control/group_key_mutations.go +++ /dev/null @@ -1,371 +0,0 @@ -package control - -import ( - "context" - "errors" - "fmt" - "sort" - "time" - - "gorm.io/gorm" - - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/state" - "gpt-load/internal/storage/models" -) - -type GroupKeyBatchAction string - -const ( - GroupKeyBatchEnable GroupKeyBatchAction = "enable" - GroupKeyBatchDisable GroupKeyBatchAction = "disable" - GroupKeyBatchDelete GroupKeyBatchAction = "delete" -) - -type GroupKeyBatchRequest struct { - Action GroupKeyBatchAction `json:"action"` - KeyIDs []uint `json:"key_ids"` -} - -type GroupKeyBatchResponse struct { - AffectedIDs []uint `json:"affected_ids"` - Summary GroupKeySummaryResponse `json:"summary"` -} - -func (s *Service) RestoreGroupKey( - ctx context.Context, - groupID uint, - keyID uint, -) (GroupKeyItemResponse, error) { - if groupID == 0 || keyID == 0 { - return GroupKeyItemResponse{}, app_errors.ErrBadRequest - } - - s.writeMu.Lock() - defer s.writeMu.Unlock() - - var group models.Group - if err := s.db.WithContext(ctx).Where("id = ?", groupID).Take(&group).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return GroupKeyItemResponse{}, groupNotFoundError() - } - return GroupKeyItemResponse{}, app_errors.ParseDBError(err) - } - var row models.UpstreamKey - if err := s.db.WithContext(ctx). - Where("id = ? AND group_id = ?", keyID, groupID). - Take(&row).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return GroupKeyItemResponse{}, keyNotFoundError() - } - return GroupKeyItemResponse{}, app_errors.ParseDBError(err) - } - view, exists := findRuntimeKey(s.registry.Snapshot(), keyID) - if !exists { - return GroupKeyItemResponse{}, dbRegistryMismatch(mismatchMissingRegistry, groupID, keyID) - } - if view.GroupID != groupID { - return GroupKeyItemResponse{}, dbRegistryMismatch(mismatchGroupID, groupID, keyID) - } - if view.Status != state.KeyStatus(row.Status) { - return GroupKeyItemResponse{}, dbRegistryMismatch(mismatchStatus, groupID, keyID) - } - if !equalOptionalWeight(view.WeightManual, row.WeightManual) { - return GroupKeyItemResponse{}, dbRegistryMismatch(mismatchWeightManual, groupID, keyID) - } - - var observedAt time.Time - groupView := state.GroupCatalogView{ - ID: group.ID, Name: group.Name, Enabled: group.Enabled, - WeightManual: cloneInt(group.WeightManual), - } - var restoreErr error - restore := func() { - observedAt = s.now().UTC() - current, exists := findRuntimeKey(s.registry.Snapshot(), keyID) - if !exists { - restoreErr = dbRegistryMismatch(mismatchMissingRegistry, groupID, keyID) - return - } - bucket := classifyHealthKey(groupView, current, observedAt) - if bucket != healthBucketCooldown && bucket != healthBucketBlacklisted { - restoreErr = app_errors.ErrInvalidKeyState - return - } - stats := s.stats.Snapshot(keyID, observedAt) - stats.ConsecutiveFailure = 0 - stats.ConsecutiveProblem = 0 - stats.LastFailureCategory = 0 - stats.LastStatusCode = 0 - weight := calculateAutoWeight(stats) - if !s.registry.RestoreRuntimeState(keyID, weight) { - restoreErr = dbRegistryMismatch(mismatchMissingRegistry, groupID, keyID) - return - } - s.stats.ClearProblemState(keyID) - } - if s.mutations == nil { - restore() - } else { - s.mutations.Do(keyID, restore) - } - if restoreErr != nil { - return GroupKeyItemResponse{}, restoreErr - } - - view, exists = findRuntimeKey(s.registry.Snapshot(), keyID) - if !exists { - return GroupKeyItemResponse{}, dbRegistryMismatch(mismatchMissingRegistry, groupID, keyID) - } - return s.mapGroupKeyItem(row, view, group, s.stats.Snapshot(keyID, observedAt), observedAt) -} - -func (s *Service) BatchGroupKeys( - ctx context.Context, - groupID uint, - request GroupKeyBatchRequest, -) (GroupKeyBatchResponse, error) { - if groupID == 0 { - return GroupKeyBatchResponse{}, app_errors.ErrBadRequest - } - keyIDs, err := normalizeGroupKeyBatchRequest(request) - if err != nil { - return GroupKeyBatchResponse{}, err - } - - s.writeMu.Lock() - defer s.writeMu.Unlock() - if err := s.enforceOperationRecoveryBarrierLocked(ctx, 0); err != nil { - return GroupKeyBatchResponse{}, err - } - - var group models.Group - if err := s.db.WithContext(ctx).Where("id = ?", groupID).Take(&group).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return GroupKeyBatchResponse{}, groupNotFoundError() - } - return GroupKeyBatchResponse{}, app_errors.ParseDBError(err) - } - var rows []models.UpstreamKey - if err := s.db.WithContext(ctx).Where("id IN ?", keyIDs).Find(&rows).Error; err != nil { - return GroupKeyBatchResponse{}, app_errors.ParseDBError(err) - } - if len(rows) != len(keyIDs) { - return GroupKeyBatchResponse{}, keyNotFoundError() - } - rowByID := make(map[uint]models.UpstreamKey, len(rows)) - for _, row := range rows { - if row.GroupID != groupID { - return GroupKeyBatchResponse{}, keyNotFoundError() - } - rowByID[row.ID] = row - } - viewByID := make(map[uint]state.KeyRuntimeView) - for _, view := range s.registry.Snapshot() { - viewByID[view.ID] = view - } - if err := validateBatchRegistryRows(groupID, keyIDs, rowByID, viewByID); err != nil { - return GroupKeyBatchResponse{}, err - } - - persist := func() error { - return s.withControlTransaction(ctx, func(tx *gorm.DB) error { - query := tx.Where("group_id = ? AND id IN ?", groupID, keyIDs) - var result *gorm.DB - switch request.Action { - case GroupKeyBatchEnable: - result = query.Model(&models.UpstreamKey{}).Update("status", models.UpstreamKeyStatusActive) - case GroupKeyBatchDisable: - result = query.Model(&models.UpstreamKey{}).Update("status", models.UpstreamKeyStatusDisabled) - case GroupKeyBatchDelete: - result = query.Delete(&models.UpstreamKey{}) - } - if result.Error != nil { - return app_errors.ParseDBError(result.Error) - } - if result.RowsAffected != int64(len(keyIDs)) { - return fmt.Errorf("batch group key rows affected = %d, want %d: %w", result.RowsAffected, len(keyIDs), app_errors.ErrDatabase) - } - return nil - }) - } - - coordinator, ok := s.mutations.(interface { - DoMany([]uint, func()) - }) - if !ok { - return GroupKeyBatchResponse{}, fmt.Errorf( - "batch mutation coordinator unavailable: %w", - app_errors.ErrInternalServer, - ) - } - apply := func() { - linearizedViews := make(map[uint]state.KeyRuntimeView) - for _, view := range s.registry.Snapshot() { - linearizedViews[view.ID] = view - } - if err = validateBatchRegistryRows(groupID, keyIDs, rowByID, linearizedViews); err != nil { - return - } - beforeEntries, snapshotErr := s.registry.SnapshotGroupKeyEntriesExact(groupID, keyIDs) - if snapshotErr != nil { - err = fmt.Errorf( - "snapshot batch Registry entries: %v: %w", - snapshotErr, - app_errors.ErrInternalServer, - ) - return - } - if s.applyBatchRegistryMutation == nil { - err = fmt.Errorf("batch Registry mutation unavailable: %w", app_errors.ErrInternalServer) - return - } - applyErr := s.applyBatchRegistryMutation(groupID, keyIDs, request.Action) - if applyErr != nil { - err = compensateBatchRegistry( - s, - groupID, - beforeEntries, - fmt.Errorf( - "apply batch Registry mutation: %v: %w", - applyErr, - withControlOperationContext( - newControlOperationError(stageApplyCommittedRegistryMutation), - groupID, - 0, - ), - ), - ) - return - } - if err = persist(); err != nil { - err = compensateBatchRegistry(s, groupID, beforeEntries, err) - return - } - if request.Action == GroupKeyBatchDelete { - for _, keyID := range keyIDs { - s.stats.Reset(keyID) - } - } - } - coordinator.DoMany(keyIDs, apply) - if err != nil { - return GroupKeyBatchResponse{}, err - } - - return GroupKeyBatchResponse{ - AffectedIDs: keyIDs, - Summary: summarizeGroupRuntimeKeys(group, s.registry.Snapshot(), s.now().UTC()), - }, nil -} - -func validateBatchRegistryRows( - groupID uint, - keyIDs []uint, - rows map[uint]models.UpstreamKey, - views map[uint]state.KeyRuntimeView, -) error { - for _, keyID := range keyIDs { - row := rows[keyID] - view, exists := views[keyID] - if !exists { - return dbRegistryMismatch(mismatchMissingRegistry, groupID, keyID) - } - if view.GroupID != groupID { - return dbRegistryMismatch(mismatchGroupID, groupID, keyID) - } - if view.Status != state.KeyStatus(row.Status) { - return dbRegistryMismatch(mismatchStatus, groupID, keyID) - } - if !equalOptionalWeight(view.WeightManual, row.WeightManual) { - return dbRegistryMismatch(mismatchWeightManual, groupID, keyID) - } - } - return nil -} - -func (s *Service) applyGroupKeyBatchRegistryMutation( - groupID uint, - keyIDs []uint, - action GroupKeyBatchAction, -) error { - switch action { - case GroupKeyBatchEnable: - return s.registry.UpdateGroupKeyStatuses(groupID, keyIDs, state.KeyStatusActive) - case GroupKeyBatchDisable: - return s.registry.UpdateGroupKeyStatuses(groupID, keyIDs, state.KeyStatusDisabled) - case GroupKeyBatchDelete: - return s.registry.RemoveGroupKeys(groupID, keyIDs) - default: - return fmt.Errorf("unsupported batch Registry action %q", action) - } -} - -func compensateBatchRegistry( - s *Service, - groupID uint, - before []state.KeyEntry, - cause error, -) error { - if s.restoreBatchRegistryEntries == nil { - return errors.Join( - cause, - fmt.Errorf("compensate batch Registry mutation: %w", app_errors.ErrInternalServer), - ) - } - if err := s.restoreBatchRegistryEntries(groupID, before); err != nil { - return errors.Join( - cause, - fmt.Errorf("compensate batch Registry mutation: %w", err), - ) - } - return cause -} - -func normalizeGroupKeyBatchRequest(request GroupKeyBatchRequest) ([]uint, error) { - if request.Action != GroupKeyBatchEnable && - request.Action != GroupKeyBatchDisable && - request.Action != GroupKeyBatchDelete { - return nil, app_errors.ErrValidation - } - if len(request.KeyIDs) < 1 || len(request.KeyIDs) > 100 { - return nil, app_errors.ErrValidation - } - keyIDs := append([]uint(nil), request.KeyIDs...) - sort.Slice(keyIDs, func(left, right int) bool { return keyIDs[left] < keyIDs[right] }) - for index, keyID := range keyIDs { - if keyID == 0 || (index > 0 && keyID == keyIDs[index-1]) { - return nil, app_errors.ErrValidation - } - } - return keyIDs, nil -} - -func summarizeGroupRuntimeKeys( - group models.Group, - views []state.KeyRuntimeView, - observedAt time.Time, -) GroupKeySummaryResponse { - summary := GroupKeySummaryResponse{} - groupView := state.GroupCatalogView{ - ID: group.ID, Name: group.Name, Enabled: group.Enabled, - WeightManual: cloneInt(group.WeightManual), - } - for _, view := range views { - if view.GroupID != group.ID { - continue - } - summary.Total++ - switch classifyHealthKey(groupView, view, observedAt) { - case healthBucketAvailable: - summary.Available++ - case healthBucketCooldown: - summary.Cooldown++ - case healthBucketBlacklisted: - summary.Blacklisted++ - case healthBucketDisabled: - summary.Disabled++ - } - } - return summary -} diff --git a/internal/control/group_key_mutations_test.go b/internal/control/group_key_mutations_test.go deleted file mode 100644 index 9e642f94..00000000 --- a/internal/control/group_key_mutations_test.go +++ /dev/null @@ -1,871 +0,0 @@ -package control - -import ( - "bytes" - "encoding/json" - "errors" - "fmt" - "net/http" - "reflect" - "sync" - "testing" - "time" - - "github.com/gin-gonic/gin" - - "gpt-load/internal/health" - "gpt-load/internal/platform/config" - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/state" - "gpt-load/internal/storage/models" -) - -func TestRestoreGroupKeyClearsProblemStateAndReturnsItem(t *testing.T) { - for _, test := range []struct { - name string - seed func(*state.KeyRegistry, uint, time.Time) bool - }{ - { - name: "cooldown", - seed: func(registry *state.KeyRegistry, keyID uint, now time.Time) bool { - return registry.SetCooldown(keyID, now.Add(time.Minute)) - }, - }, - { - name: "blacklisted", - seed: func(registry *state.KeyRegistry, keyID uint, _ time.Time) bool { - return registry.SetBlacklisted(keyID) - }, - }, - } { - t.Run(test.name, func(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.August, 1, 12, 0, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - group := validControlGroup("restore-" + test.name) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-restore", models.UpstreamKeyStatusActive, nil) - if !test.seed(fixture.registry, row.ID, now) { - t.Fatal("seed runtime problem state") - } - for range 4 { - if _, ok := fixture.registry.IncrFailure(row.ID); !ok { - t.Fatal("seed Registry failure") - } - } - for index := 0; index < 8; index++ { - fixture.stats.RecordSuccess(row.ID, now.Add(-2*time.Minute+time.Duration(index)*time.Second)) - } - for index := 0; index < 4; index++ { - fixture.stats.RecordFailure(row.ID, health.FailureCategoryInvalidKey, http.StatusUnauthorized, now.Add(-time.Minute+time.Duration(index)*time.Second)) - } - - got, err := fixture.service.RestoreGroupKey(t.Context(), group.ID, row.ID) - if err != nil { - t.Fatalf("RestoreGroupKey() error = %v", err) - } - if got.ID != row.ID || got.Mask != "prov****tore" || got.ConfiguredStatus != "active" || - got.EffectiveStatus != "available" || got.WeightMode != "auto" || got.Weight == nil || *got.Weight != 64 || - got.RecentSuccessCount != 8 || got.RecentFailureCount != 4 || got.ConsecutiveFailureCount != 0 || - got.LastFailureCategory != "ambiguous" || got.LastStatusCode != nil || got.CooldownUntilMS != nil || - got.Recovery != (GroupKeyRecoveryResponse{Mode: "none"}) { - t.Fatalf("restored item = %#v", got) - } - view, exists := findRuntimeKey(fixture.registry.Snapshot(), row.ID) - if !exists || view.Blacklisted || !view.CooldownUntil.IsZero() || view.FailureCount != 0 || view.WeightAuto != 64 { - t.Fatalf("restored Registry view = %#v, exists=%t", view, exists) - } - if stats := fixture.stats.Snapshot(row.ID, now); stats != (health.KeyStats{Success: 8, Failure: 4, Problem: 4}) { - t.Fatalf("restored stats = %#v", stats) - } - }) - } -} - -func TestRestoreGroupKeyRejectsNormalAndDisabledWithoutMutation(t *testing.T) { - for _, status := range []models.UpstreamKeyStatus{ - models.UpstreamKeyStatusActive, - models.UpstreamKeyStatusDisabled, - } { - t.Run(string(status), func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("restore-invalid-" + string(status)) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-invalid", status, nil) - before := fixture.registry.Snapshot() - - _, err := fixture.service.RestoreGroupKey(t.Context(), group.ID, row.ID) - if !errors.Is(err, app_errors.ErrInvalidKeyState) { - t.Fatalf("RestoreGroupKey() error = %#v, want INVALID_KEY_STATE", err) - } - if after := fixture.registry.Snapshot(); len(after) != len(before) || after[0] != before[0] { - t.Fatalf("Registry after rejection = %#v, want %#v", after, before) - } - }) - } -} - -func TestRestoreGroupKeyReturnsPreciseNotFound(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("restore-not-found") - other := validControlGroup("restore-not-found-other") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - if err := fixture.db.Create(other).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, other.ID, "provider-secret-other", models.UpstreamKeyStatusActive, nil) - - for _, test := range []struct { - name string - groupID uint - keyID uint - want string - }{ - {name: "group", groupID: group.ID + other.ID + 1000, keyID: row.ID, want: "group"}, - {name: "key", groupID: group.ID, keyID: row.ID, want: "key"}, - } { - t.Run(test.name, func(t *testing.T) { - _, err := fixture.service.RestoreGroupKey(t.Context(), test.groupID, test.keyID) - var notFound *controlResourceNotFoundError - if !errors.As(err, ¬Found) || notFound.resource != test.want { - t.Fatalf("RestoreGroupKey() error = %#v, want %s not found", err, test.want) - } - }) - } -} - -func TestRestoreGroupKeyAllowsFailureAfterRestorePoint(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.August, 1, 12, 0, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - group := validControlGroup("restore-new-failure") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-new-failure", models.UpstreamKeyStatusActive, nil) - if !fixture.registry.SetBlacklisted(row.ID) { - t.Fatal("seed blacklist") - } - if _, err := fixture.service.RestoreGroupKey(t.Context(), group.ID, row.ID); err != nil { - t.Fatal(err) - } - - if _, ok := fixture.registry.IncrFailure(row.ID); !ok || !fixture.registry.SetBlacklisted(row.ID) { - t.Fatal("record post-restore Registry failure") - } - fixture.stats.RecordFailure(row.ID, health.FailureCategoryInvalidKey, http.StatusUnauthorized, now.Add(time.Second)) - - view, exists := findRuntimeKey(fixture.registry.Snapshot(), row.ID) - if !exists || !view.Blacklisted || view.FailureCount != 1 { - t.Fatalf("post-restore Registry view = %#v, exists=%t", view, exists) - } - stats := fixture.stats.Snapshot(row.ID, now.Add(time.Second)) - if stats.Failure != 1 || stats.ConsecutiveFailure != 1 || stats.LastFailureCategory != health.FailureCategoryInvalidKey { - t.Fatalf("post-restore stats = %#v", stats) - } -} - -type postRestoreFailureCoordinator struct { - after func() - calls int -} - -type recoverBeforeRestoreCoordinator struct { - registry *state.KeyRegistry - stats *health.StatsStore - ref state.KeyRef - recovered bool -} - -func (coordinator *recoverBeforeRestoreCoordinator) Do(_ uint, restore func()) { - coordinator.recovered = coordinator.registry.RecoverIfMatch( - coordinator.ref, - state.DefaultWeight, - ) - if coordinator.recovered { - coordinator.stats.Reset(coordinator.ref.ID) - } - restore() -} - -func TestRestoreGroupKeyRejectsWhenValidationRecoveryWinsCoordinator(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.August, 1, 12, 15, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - group := validControlGroup("restore-validation-wins") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-validation-wins", models.UpstreamKeyStatusActive, nil) - if !fixture.registry.SetBlacklisted(row.ID) { - t.Fatal("seed blacklist") - } - refs := fixture.registry.BlacklistedKeys() - if len(refs) != 1 { - t.Fatalf("blacklisted refs = %#v, want one", refs) - } - coordinator := &recoverBeforeRestoreCoordinator{ - registry: fixture.registry, - stats: fixture.stats, - ref: refs[0], - } - fixture.service.mutations = coordinator - - _, err := fixture.service.RestoreGroupKey(t.Context(), group.ID, row.ID) - if !coordinator.recovered { - t.Fatal("validation recovery did not complete before restore callback") - } - if !errors.Is(err, app_errors.ErrInvalidKeyState) { - t.Fatalf("RestoreGroupKey() error = %#v, want INVALID_KEY_STATE", err) - } - view, exists := findRuntimeKey(fixture.registry.Snapshot(), row.ID) - if !exists || view.Blacklisted || view.FailureCount != 0 || view.WeightAuto != state.DefaultWeight { - t.Fatalf("validation-won Registry view = %#v, exists=%t", view, exists) - } -} - -type beforeRestoreCoordinator struct { - before func() -} - -func (coordinator *beforeRestoreCoordinator) Do(_ uint, restore func()) { - coordinator.before() - restore() -} - -func TestRestoreGroupKeyEvaluatesCooldownAtCoordinatorLinearizationPoint(t *testing.T) { - fixture := newServiceFixture(t) - before := time.Date(2026, time.August, 1, 12, 20, 0, 0, time.UTC) - current := before - fixture.service.now = func() time.Time { return current } - group := validControlGroup("restore-expired-at-linearization") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-expired-linearization", models.UpstreamKeyStatusActive, nil) - if !fixture.registry.SetCooldown(row.ID, before.Add(time.Second)) { - t.Fatal("seed cooldown") - } - fixture.service.mutations = &beforeRestoreCoordinator{before: func() { - current = before.Add(2 * time.Second) - }} - - _, err := fixture.service.RestoreGroupKey(t.Context(), group.ID, row.ID) - if !errors.Is(err, app_errors.ErrInvalidKeyState) { - t.Fatalf("RestoreGroupKey() error = %#v, want INVALID_KEY_STATE", err) - } - view, exists := findRuntimeKey(fixture.registry.Snapshot(), row.ID) - if !exists || !view.CooldownUntil.Equal(before.Add(time.Second)) { - t.Fatalf("expired cooldown view = %#v, exists=%t", view, exists) - } -} - -func (coordinator *postRestoreFailureCoordinator) Do(_ uint, mutate func()) { - coordinator.calls++ - mutate() - coordinator.after() -} - -func TestRestoreGroupKeySerializesNewFailureAfterCompleteRestore(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.August, 1, 12, 30, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - group := validControlGroup("restore-linearized-failure") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-linearized", models.UpstreamKeyStatusActive, nil) - if !fixture.registry.SetBlacklisted(row.ID) { - t.Fatal("seed blacklist") - } - coordinator := &postRestoreFailureCoordinator{after: func() { - if _, ok := fixture.registry.IncrFailure(row.ID); !ok || !fixture.registry.SetBlacklisted(row.ID) { - t.Fatal("record serialized post-restore Registry failure") - } - fixture.stats.RecordFailure(row.ID, health.FailureCategoryInvalidKey, http.StatusUnauthorized, now.Add(time.Second)) - }} - fixture.service.mutations = coordinator - - if _, err := fixture.service.RestoreGroupKey(t.Context(), group.ID, row.ID); err != nil { - t.Fatal(err) - } - if coordinator.calls != 1 { - t.Fatalf("mutation coordinator calls = %d, want 1", coordinator.calls) - } - view, exists := findRuntimeKey(fixture.registry.Snapshot(), row.ID) - if !exists || !view.Blacklisted || view.FailureCount != 1 { - t.Fatalf("serialized post-restore Registry view = %#v, exists=%t", view, exists) - } - stats := fixture.stats.Snapshot(row.ID, now.Add(time.Second)) - if stats.Failure != 1 || stats.ConsecutiveFailure != 1 || stats.LastFailureCategory != health.FailureCategoryInvalidKey { - t.Fatalf("serialized post-restore stats = %#v", stats) - } -} - -func TestRestoreGroupKeyHTTPStrictBodyAndStateError(t *testing.T) { - initControlI18n(t) - fixture := newServiceFixture(t) - now := time.Date(2026, time.August, 1, 12, 45, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - group := validControlGroup("restore-http") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - cooldown := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-restore-http", models.UpstreamKeyStatusActive, nil) - normal := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-normal-http", models.UpstreamKeyStatusActive, nil) - if !fixture.registry.SetCooldown(cooldown.ID, now.Add(time.Minute)) { - t.Fatal("seed cooldown") - } - engine := gin.New() - NewServer(&config.Config{AuthKey: groupKeyHTTPAuth}, fixture.service).RegisterRoutes(engine) - - success := serveGroupKeyHTTPRequest(t, engine, http.MethodPost, fmt.Sprintf("/api/groups/%d/keys/%d/restore", group.ID, cooldown.ID), "", groupKeyHTTPAuth, "en-US") - if success.Code != http.StatusOK || !bytes.Contains(success.Body.Bytes(), []byte(`"effective_status":"available"`)) { - t.Fatalf("restore success = %d %s, want available item", success.Code, success.Body.String()) - } - invalidBody := serveGroupKeyHTTPRequest(t, engine, http.MethodPost, fmt.Sprintf("/api/groups/%d/keys/%d/restore", group.ID, normal.ID), `{"unexpected":true}`, groupKeyHTTPAuth, "en-US") - if invalidBody.Code != http.StatusBadRequest { - t.Fatalf("restore invalid body = %d %s, want 400", invalidBody.Code, invalidBody.Body.String()) - } - invalidState := serveGroupKeyHTTPRequest(t, engine, http.MethodPost, fmt.Sprintf("/api/groups/%d/keys/%d/restore", group.ID, normal.ID), "", groupKeyHTTPAuth, "en-US") - if invalidState.Code != http.StatusConflict || !bytes.Contains(invalidState.Body.Bytes(), []byte(`"code":"INVALID_KEY_STATE"`)) { - t.Fatalf("restore invalid state = %d %s, want INVALID_KEY_STATE", invalidState.Code, invalidState.Body.String()) - } -} - -func TestBatchGroupKeysEnableDisableAndDelete(t *testing.T) { - for _, action := range []GroupKeyBatchAction{ - GroupKeyBatchEnable, - GroupKeyBatchDisable, - GroupKeyBatchDelete, - } { - t.Run(string(action), func(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.August, 1, 13, 0, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - group := validControlGroup("batch-" + string(action)) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - initialStatus := models.UpstreamKeyStatusActive - if action == GroupKeyBatchEnable { - initialStatus = models.UpstreamKeyStatusDisabled - } - first := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-batch-first", initialStatus, nil) - second := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-batch-second", initialStatus, nil) - untouched := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-batch-third", models.UpstreamKeyStatusActive, nil) - fixture.stats.RecordFailure(first.ID, health.FailureCategoryInvalidKey, http.StatusUnauthorized, now) - fixture.stats.RecordSuccess(second.ID, now) - - got, err := fixture.service.BatchGroupKeys(t.Context(), group.ID, GroupKeyBatchRequest{ - Action: action, KeyIDs: []uint{second.ID, first.ID}, - }) - if err != nil { - t.Fatalf("BatchGroupKeys() error = %v", err) - } - if len(got.AffectedIDs) != 2 || got.AffectedIDs[0] != first.ID || got.AffectedIDs[1] != second.ID { - t.Fatalf("affected ids = %v, want sorted [%d %d]", got.AffectedIDs, first.ID, second.ID) - } - - var rows []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", group.ID).Order("id ASC").Find(&rows).Error; err != nil { - t.Fatal(err) - } - views := fixture.registry.Snapshot() - switch action { - case GroupKeyBatchEnable: - if len(rows) != 3 || rows[0].Status != models.UpstreamKeyStatusActive || rows[1].Status != models.UpstreamKeyStatusActive { - t.Fatalf("enabled DB rows = %#v", rows) - } - if got.Summary != (GroupKeySummaryResponse{Total: 3, Available: 3}) { - t.Fatalf("enabled summary = %#v", got.Summary) - } - case GroupKeyBatchDisable: - if len(rows) != 3 || rows[0].Status != models.UpstreamKeyStatusDisabled || rows[1].Status != models.UpstreamKeyStatusDisabled { - t.Fatalf("disabled DB rows = %#v", rows) - } - if got.Summary != (GroupKeySummaryResponse{Total: 3, Available: 1, Disabled: 2}) { - t.Fatalf("disabled summary = %#v", got.Summary) - } - case GroupKeyBatchDelete: - if len(rows) != 1 || rows[0].ID != untouched.ID || len(views) != 1 || views[0].ID != untouched.ID { - t.Fatalf("delete DB rows / Registry views = %#v / %#v", rows, views) - } - if got.Summary != (GroupKeySummaryResponse{Total: 1, Available: 1}) { - t.Fatalf("delete summary = %#v", got.Summary) - } - if stats := fixture.stats.Snapshot(first.ID, now); stats != (health.KeyStats{}) { - t.Fatalf("deleted first stats = %#v", stats) - } - if stats := fixture.stats.Snapshot(second.ID, now); stats != (health.KeyStats{}) { - t.Fatalf("deleted second stats = %#v", stats) - } - } - }) - } -} - -func TestBatchGroupKeysRejectsInvalidRequestsWithoutMutation(t *testing.T) { - for _, test := range []struct { - name string - action GroupKeyBatchAction - ids func(first, second, other uint) []uint - wantErr *app_errors.APIError - }{ - {name: "unknown action", action: "archive", ids: func(first, _, _ uint) []uint { return []uint{first} }, wantErr: app_errors.ErrValidation}, - {name: "empty", action: GroupKeyBatchDisable, ids: func(_, _, _ uint) []uint { return nil }, wantErr: app_errors.ErrValidation}, - {name: "duplicate", action: GroupKeyBatchDisable, ids: func(first, _, _ uint) []uint { return []uint{first, first} }, wantErr: app_errors.ErrValidation}, - {name: "cross group", action: GroupKeyBatchDisable, ids: func(first, _, other uint) []uint { return []uint{first, other} }, wantErr: app_errors.ErrResourceNotFound}, - {name: "missing", action: GroupKeyBatchDisable, ids: func(first, _, _ uint) []uint { return []uint{first, 999999} }, wantErr: app_errors.ErrResourceNotFound}, - {name: "over limit", action: GroupKeyBatchDisable, ids: func(_, _, _ uint) []uint { - ids := make([]uint, 101) - for index := range ids { - ids[index] = uint(index + 1) - } - return ids - }, wantErr: app_errors.ErrValidation}, - } { - t.Run(test.name, func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("batch-reject-" + test.name) - otherGroup := validControlGroup("batch-reject-other-" + test.name) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - if err := fixture.db.Create(otherGroup).Error; err != nil { - t.Fatal(err) - } - first := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-reject-first", models.UpstreamKeyStatusActive, nil) - second := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-reject-second", models.UpstreamKeyStatusActive, nil) - other := seedManagedUpstreamKey(t, fixture, otherGroup.ID, "provider-secret-reject-other", models.UpstreamKeyStatusActive, nil) - before := fixture.registry.Snapshot() - - _, err := fixture.service.BatchGroupKeys(t.Context(), group.ID, GroupKeyBatchRequest{Action: test.action, KeyIDs: test.ids(first.ID, second.ID, other.ID)}) - if !errors.Is(err, test.wantErr) { - t.Fatalf("BatchGroupKeys() error = %#v, want %s", err, test.wantErr.Code) - } - if after := fixture.registry.Snapshot(); len(after) != len(before) { - t.Fatalf("Registry length after rejection = %d, want %d", len(after), len(before)) - } else { - for index := range after { - if after[index] != before[index] { - t.Fatalf("Registry after rejection = %#v, want %#v", after, before) - } - } - } - var persisted []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", group.ID).Order("id ASC").Find(&persisted).Error; err != nil { - t.Fatal(err) - } - if len(persisted) != 2 || persisted[0].Status != models.UpstreamKeyStatusActive || persisted[1].Status != models.UpstreamKeyStatusActive { - t.Fatalf("DB rows after rejection = %#v", persisted) - } - }) - } -} - -func TestBatchGroupKeysDBAndRegistryFailuresAreAtomic(t *testing.T) { - t.Run("Registry validation before DB", func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("batch-registry-failure") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - first := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-registry-first", models.UpstreamKeyStatusActive, nil) - second := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-registry-second", models.UpstreamKeyStatusActive, nil) - fixture.registry.RemoveKey(second.ID) - - _, err := fixture.service.BatchGroupKeys(t.Context(), group.ID, GroupKeyBatchRequest{Action: GroupKeyBatchDisable, KeyIDs: []uint{first.ID, second.ID}}) - if err == nil { - t.Fatal("BatchGroupKeys() error = nil, want Registry mismatch") - } - var rows []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", group.ID).Order("id ASC").Find(&rows).Error; err != nil { - t.Fatal(err) - } - if len(rows) != 2 || rows[0].Status != models.UpstreamKeyStatusActive || rows[1].Status != models.UpstreamKeyStatusActive { - t.Fatalf("DB rows after Registry failure = %#v", rows) - } - }) - - t.Run("DB failure before Registry", func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("batch-db-failure") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - first := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-db-first", models.UpstreamKeyStatusActive, nil) - second := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-db-second", models.UpstreamKeyStatusActive, nil) - if err := fixture.db.Exec(`CREATE TRIGGER reject_batch_status BEFORE UPDATE OF status ON upstream_keys BEGIN SELECT RAISE(ABORT, 'blocked'); END`).Error; err != nil { - t.Fatal(err) - } - - _, err := fixture.service.BatchGroupKeys(t.Context(), group.ID, GroupKeyBatchRequest{Action: GroupKeyBatchDisable, KeyIDs: []uint{first.ID, second.ID}}) - if !errors.Is(err, app_errors.ErrDatabase) { - t.Fatalf("BatchGroupKeys() error = %#v, want DATABASE_ERROR", err) - } - for _, view := range fixture.registry.Snapshot() { - if view.GroupID == group.ID && view.Status != state.KeyStatusActive { - t.Fatalf("Registry view after DB failure = %#v", view) - } - } - var rows []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", group.ID).Order("id ASC").Find(&rows).Error; err != nil { - t.Fatal(err) - } - if len(rows) != 2 || rows[0].Status != models.UpstreamKeyStatusActive || rows[1].Status != models.UpstreamKeyStatusActive { - t.Fatalf("DB rows after DB failure = %#v", rows) - } - }) -} - -func TestBatchGroupKeysRegistryApplyFailureRestoresDBAndRegistry(t *testing.T) { - for _, action := range []GroupKeyBatchAction{ - GroupKeyBatchDisable, - GroupKeyBatchDelete, - } { - t.Run(string(action), func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("batch-apply-failure-" + string(action)) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - first := seedManagedUpstreamKey( - t, fixture, group.ID, "provider-secret-apply-first", models.UpstreamKeyStatusActive, nil, - ) - second := seedManagedUpstreamKey( - t, fixture, group.ID, "provider-secret-apply-second", models.UpstreamKeyStatusActive, nil, - ) - before := fixture.registry.Snapshot() - apply := fixture.service.applyBatchRegistryMutation - fixture.service.applyBatchRegistryMutation = func( - groupID uint, - keyIDs []uint, - action GroupKeyBatchAction, - ) error { - if !fixture.registry.RemoveKey(second.ID) { - t.Fatal("remove Registry key before batch apply") - } - return apply(groupID, keyIDs, action) - } - - _, err := fixture.service.BatchGroupKeys(t.Context(), group.ID, GroupKeyBatchRequest{ - Action: action, - KeyIDs: []uint{first.ID, second.ID}, - }) - if !errors.Is(err, app_errors.ErrInternalServer) { - t.Fatalf("BatchGroupKeys() error = %#v, want INTERNAL_SERVER_ERROR", err) - } - - var rows []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", group.ID).Order("id ASC").Find(&rows).Error; err != nil { - t.Fatal(err) - } - if len(rows) != 2 || rows[0].Status != models.UpstreamKeyStatusActive || - rows[1].Status != models.UpstreamKeyStatusActive { - t.Fatalf("DB rows after Registry apply failure = %#v, want old state", rows) - } - after := fixture.registry.Snapshot() - if !reflect.DeepEqual(after, before) { - t.Fatalf("Registry after apply failure = %#v, want %#v", after, before) - } - }) - } -} - -func TestBatchGroupKeysCompensationRestoresSelectedRuntimeGenerationExactly(t *testing.T) { - for _, test := range []struct { - name string - action GroupKeyBatchAction - failApply bool - triggerSQL string - }{ - { - name: "disable apply failure", action: GroupKeyBatchDisable, failApply: true, - }, - { - name: "delete apply failure", action: GroupKeyBatchDelete, failApply: true, - }, - { - name: "disable DB failure", - action: GroupKeyBatchDisable, - triggerSQL: `CREATE TRIGGER reject_exact_status BEFORE UPDATE OF status ON upstream_keys BEGIN SELECT RAISE(ABORT, 'blocked'); END`, - }, - { - name: "delete DB failure", - action: GroupKeyBatchDelete, - triggerSQL: `CREATE TRIGGER reject_exact_delete BEFORE DELETE ON upstream_keys BEGIN SELECT RAISE(ABORT, 'blocked'); END`, - }, - } { - t.Run(test.name, func(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.August, 1, 15, 0, 0, 0, time.UTC) - group := validControlGroup("batch-exact-compensation-" + test.name) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - first := seedManagedUpstreamKey( - t, fixture, group.ID, "provider-secret-exact-first", models.UpstreamKeyStatusActive, nil, - ) - second := seedManagedUpstreamKey( - t, fixture, group.ID, "provider-secret-exact-second", models.UpstreamKeyStatusActive, nil, - ) - unselected := seedManagedUpstreamKey( - t, fixture, group.ID, "provider-secret-exact-unselected", models.UpstreamKeyStatusActive, nil, - ) - for _, keyID := range []uint{first.ID, second.ID} { - if !fixture.registry.SetCooldown(keyID, now.Add(time.Hour)) || - !fixture.registry.SetAutoWeight(keyID, 37) { - t.Fatalf("seed selected runtime key %d", keyID) - } - if _, ok := fixture.registry.IncrFailure(keyID); !ok { - t.Fatalf("seed selected failure key %d", keyID) - } - if !fixture.registry.SetBlacklisted(keyID) { - t.Fatalf("seed selected blacklist key %d", keyID) - } - } - beforeViews := map[uint]state.KeyRuntimeView{} - beforeGenerations := map[uint]uint64{} - for _, keyID := range []uint{first.ID, second.ID} { - view, exists := findRuntimeKey(fixture.registry.Snapshot(), keyID) - if !exists { - t.Fatalf("selected key %d missing before batch", keyID) - } - beforeViews[keyID] = view - beforeGenerations[keyID] = blacklistedGeneration(t, fixture.registry, keyID) - if beforeGenerations[keyID] == 0 { - t.Fatalf("selected key %d generation = 0 before batch", keyID) - } - } - - if test.failApply { - apply := fixture.service.applyBatchRegistryMutation - fixture.service.applyBatchRegistryMutation = func( - groupID uint, - keyIDs []uint, - action GroupKeyBatchAction, - ) error { - if err := apply(groupID, keyIDs, action); err != nil { - return err - } - if _, ok := fixture.registry.IncrFailure(unselected.ID); !ok || - !fixture.registry.SetBlacklisted(unselected.ID) { - t.Fatal("mutate unselected runtime key during apply") - } - return errors.New("injected batch Registry apply failure") - } - } else if err := fixture.db.Exec(test.triggerSQL).Error; err != nil { - t.Fatal(err) - } - - _, err := fixture.service.BatchGroupKeys(t.Context(), group.ID, GroupKeyBatchRequest{ - Action: test.action, - KeyIDs: []uint{first.ID, second.ID}, - }) - if test.failApply { - if !errors.Is(err, app_errors.ErrInternalServer) { - t.Fatalf("BatchGroupKeys() error = %#v, want INTERNAL_SERVER_ERROR", err) - } - } else if !errors.Is(err, app_errors.ErrDatabase) { - t.Fatalf("BatchGroupKeys() error = %#v, want DATABASE_ERROR", err) - } - - for _, keyID := range []uint{first.ID, second.ID} { - afterView, exists := findRuntimeKey(fixture.registry.Snapshot(), keyID) - if !exists || !reflect.DeepEqual(afterView, beforeViews[keyID]) { - t.Fatalf("selected key %d after compensation = %#v, %t; want %#v", keyID, afterView, exists, beforeViews[keyID]) - } - if generation := blacklistedGeneration(t, fixture.registry, keyID); generation != beforeGenerations[keyID] { - t.Fatalf("selected key %d generation = %d, want %d", keyID, generation, beforeGenerations[keyID]) - } - } - if test.failApply { - unselectedView, exists := findRuntimeKey(fixture.registry.Snapshot(), unselected.ID) - if !exists || !unselectedView.Blacklisted || unselectedView.FailureCount != 1 || - blacklistedGeneration(t, fixture.registry, unselected.ID) == 0 { - t.Fatalf("unselected runtime update was overwritten: %#v, exists=%t", unselectedView, exists) - } - } - }) - } -} - -func blacklistedGeneration(t *testing.T, registry *state.KeyRegistry, keyID uint) uint64 { - t.Helper() - for _, ref := range registry.BlacklistedKeys() { - if ref.ID == keyID { - return ref.FailureGeneration - } - } - t.Fatalf("blacklisted key %d not found", keyID) - return 0 -} - -type heldGroupKeyMutationCoordinator struct { - mu sync.Mutex - manyAttempted chan struct{} - manyOnce sync.Once -} - -func (coordinator *heldGroupKeyMutationCoordinator) Do(_ uint, mutate func()) { - coordinator.mu.Lock() - defer coordinator.mu.Unlock() - mutate() -} - -func (coordinator *heldGroupKeyMutationCoordinator) DoMany(_ []uint, mutate func()) { - coordinator.manyOnce.Do(func() { close(coordinator.manyAttempted) }) - coordinator.mu.Lock() - defer coordinator.mu.Unlock() - mutate() -} - -func TestBatchGroupKeysDeleteWaitsForInflightMutationThenResetsStats(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.August, 1, 14, 0, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - group := validControlGroup("batch-delete-inflight") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-delete-inflight", models.UpstreamKeyStatusActive, nil) - coordinator := &heldGroupKeyMutationCoordinator{manyAttempted: make(chan struct{})} - fixture.service.mutations = coordinator - holderEntered := make(chan struct{}) - releaseHolder := make(chan struct{}) - holderDone := make(chan struct{}) - go func() { - coordinator.Do(row.ID, func() { - if _, ok := fixture.registry.IncrFailure(row.ID); !ok { - t.Error("in-flight Registry failure did not start") - } - close(holderEntered) - <-releaseHolder - fixture.stats.RecordFailure(row.ID, health.FailureCategoryInvalidKey, http.StatusUnauthorized, now) - }) - close(holderDone) - }() - <-holderEntered - - type batchResult struct { - response GroupKeyBatchResponse - err error - } - batchDone := make(chan batchResult, 1) - go func() { - response, err := fixture.service.BatchGroupKeys(t.Context(), group.ID, GroupKeyBatchRequest{ - Action: GroupKeyBatchDelete, - KeyIDs: []uint{row.ID}, - }) - batchDone <- batchResult{response: response, err: err} - }() - - select { - case <-coordinator.manyAttempted: - select { - case result := <-batchDone: - t.Fatalf("batch delete completed while mutation held key: %#v", result) - default: - } - case result := <-batchDone: - t.Fatalf("batch delete bypassed per-key coordinator: %#v", result) - case <-time.After(time.Second): - t.Fatal("timed out waiting for batch delete to reach per-key coordinator") - } - - close(releaseHolder) - <-holderDone - result := <-batchDone - if result.err != nil { - t.Fatalf("BatchGroupKeys() error = %v", result.err) - } - if len(result.response.AffectedIDs) != 1 || result.response.AffectedIDs[0] != row.ID { - t.Fatalf("batch response = %#v", result.response) - } - if _, exists := fixture.registry.EncryptedValue(row.ID); exists { - t.Fatal("deleted key remains in Registry") - } - if stats := fixture.stats.Snapshot(row.ID, now); stats != (health.KeyStats{}) { - t.Fatalf("deleted key stats = %#v, want zero", stats) - } -} - -func TestBatchGroupKeysDeleteFailsClosedWithoutMultiKeyCoordinator(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("batch-delete-no-multi-coordinator") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-no-multi", models.UpstreamKeyStatusActive, nil) - fixture.service.mutations = &postRestoreFailureCoordinator{after: func() {}} - - _, err := fixture.service.BatchGroupKeys(t.Context(), group.ID, GroupKeyBatchRequest{ - Action: GroupKeyBatchDelete, - KeyIDs: []uint{row.ID}, - }) - if !errors.Is(err, app_errors.ErrInternalServer) { - t.Fatalf("BatchGroupKeys() error = %#v, want INTERNAL_SERVER_ERROR", err) - } - var count int64 - if err := fixture.db.Model(&models.UpstreamKey{}).Where("id = ?", row.ID).Count(&count).Error; err != nil { - t.Fatal(err) - } - if count != 1 { - t.Fatalf("persisted key count = %d, want 1", count) - } - if _, exists := fixture.registry.EncryptedValue(row.ID); !exists { - t.Fatal("Registry key removed without multi-key coordinator") - } -} - -func TestBatchGroupKeysHTTPStrictJSONAndSecretFreeResponse(t *testing.T) { - initControlI18n(t) - fixture := newServiceFixture(t) - group := validControlGroup("batch-http") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey(t, fixture, group.ID, "provider-secret-http", models.UpstreamKeyStatusActive, nil) - engine := gin.New() - NewServer(&config.Config{AuthKey: groupKeyHTTPAuth}, fixture.service).RegisterRoutes(engine) - path := fmt.Sprintf("/api/groups/%d/keys/batch", group.ID) - - for _, body := range []string{ - `{"action":"disable","key_ids":[],"unexpected":true}`, - `{"action":"disable","key_ids":[]} {}`, - `{"action":"disable","key_ids":null}`, - } { - recorder := serveGroupKeyHTTPRequest(t, engine, http.MethodPost, path, body, groupKeyHTTPAuth, "en-US") - if recorder.Code != http.StatusBadRequest { - t.Fatalf("POST body %s = %d %s, want 400", body, recorder.Code, recorder.Body.String()) - } - } - - body := fmt.Sprintf(`{"action":"disable","key_ids":[%d]}`, row.ID) - recorder := serveGroupKeyHTTPRequest(t, engine, http.MethodPost, path, body, groupKeyHTTPAuth, "en-US") - if recorder.Code != http.StatusOK { - t.Fatalf("POST = %d %s, want 200", recorder.Code, recorder.Body.String()) - } - var envelope struct { - Code int `json:"code"` - Data GroupKeyBatchResponse `json:"data"` - } - if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil { - t.Fatal(err) - } - if envelope.Code != 0 || len(envelope.Data.AffectedIDs) != 1 || envelope.Data.AffectedIDs[0] != row.ID || envelope.Data.Summary.Disabled != 1 { - t.Fatalf("batch envelope = %#v", envelope) - } - for _, forbidden := range [][]byte{[]byte("provider-secret-http"), []byte(row.KeyValue), []byte("key_value"), []byte("key_hash"), []byte("mask")} { - if bytes.Contains(recorder.Body.Bytes(), forbidden) { - t.Fatalf("response exposes %q: %s", forbidden, recorder.Body.String()) - } - } -} diff --git a/internal/control/group_keys.go b/internal/control/group_keys.go deleted file mode 100644 index 3df00f45..00000000 --- a/internal/control/group_keys.go +++ /dev/null @@ -1,63 +0,0 @@ -package control - -import ( - "context" - - "gorm.io/gorm" - - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/state" - "gpt-load/internal/storage/models" -) - -type GroupKeyImportRequest struct { - Keys string `json:"keys"` -} - -type GroupKeyImportResult struct { - GroupID uint `json:"group_id"` - KeysAdded int `json:"keys_added"` - KeysDuplicated int `json:"keys_duplicated"` -} - -func (s *Service) ImportGroupKeys( - ctx context.Context, - groupID uint, - request GroupKeyImportRequest, -) (GroupKeyImportResult, error) { - if groupID == 0 { - return GroupKeyImportResult{}, app_errors.ErrValidation - } - keys, err := s.normalizeUpstreamKeys(request.Keys) - if err != nil { - return GroupKeyImportResult{}, err - } - - result := GroupKeyImportResult{GroupID: groupID} - var entries []state.KeyEntry - err = s.writeKeyConfig( - ctx, - groupID, - 0, - func(tx *gorm.DB) error { - var group struct{ ID uint } - if err := tx.Model(&models.Group{}). - Select("id").Where("id = ?", groupID).Take(&group).Error; err != nil { - return app_errors.ParseDBError(err) - } - entries, result.KeysAdded, result.KeysDuplicated, err = - s.persistUpstreamKeys(tx, groupID, keys) - if err != nil { - return err - } - return state.ValidateKeyEntries(entries) - }, - func() error { - return s.applyMissingRegistryKeys(groupID, entries) - }, - ) - if err != nil { - return GroupKeyImportResult{}, err - } - return result, nil -} diff --git a/internal/control/group_keys_test.go b/internal/control/group_keys_test.go deleted file mode 100644 index 15bc60ec..00000000 --- a/internal/control/group_keys_test.go +++ /dev/null @@ -1,316 +0,0 @@ -package control - -import ( - "errors" - "reflect" - "sync" - "testing" - "time" - - "gpt-load/internal/platform/encryption" - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" - "gpt-load/internal/storage/models" -) - -func TestImportGroupKeysAddsAndCountsDuplicatesWithoutPublishing(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-existing") - beforeSnapshot := fixture.manager.Current() - var beforeGroup models.Group - if err := fixture.db.First(&beforeGroup, groupID).Error; err != nil { - t.Fatal(err) - } - - result, err := fixture.service.ImportGroupKeys(t.Context(), groupID, GroupKeyImportRequest{ - Keys: "sk-existing\nsk-new\nsk-new", - }) - if err != nil { - t.Fatal(err) - } - if result.KeysAdded != 1 || result.KeysDuplicated != 2 { - t.Fatalf("result = %#v", result) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("key-only import published Snapshot") - } - - var afterGroup models.Group - if err := fixture.db.First(&afterGroup, groupID).Error; err != nil { - t.Fatal(err) - } - if !reflect.DeepEqual(beforeGroup, afterGroup) { - t.Fatalf("Group changed:\nbefore=%#v\nafter=%#v", beforeGroup, afterGroup) - } - assertImportedKeyState(t, fixture, groupID, 2) -} - -func TestImportGroupKeysRepairsMissingRegistryWithoutResettingRuntime(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-existing\nsk-registry-missing") - var keys []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&keys).Error; err != nil { - t.Fatal(err) - } - if len(keys) != 2 { - t.Fatalf("seeded keys = %#v", keys) - } - if count, ok := fixture.registry.IncrFailure(keys[0].ID); !ok || count != 1 { - t.Fatalf("seed failure count = %d, %t", count, ok) - } - if !fixture.registry.RemoveKey(keys[1].ID) { - t.Fatalf("remove Registry key %d = false", keys[1].ID) - } - beforeSnapshot := fixture.manager.Current() - - result, err := fixture.service.ImportGroupKeys(t.Context(), groupID, GroupKeyImportRequest{ - Keys: "sk-existing\nsk-registry-missing", - }) - if err != nil { - t.Fatal(err) - } - if result.KeysAdded != 0 || result.KeysDuplicated != 2 { - t.Fatalf("result = %#v", result) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("Registry repair published Snapshot") - } - if count, ok := fixture.registry.IncrFailure(keys[0].ID); !ok || count != 2 { - t.Fatalf("existing Registry failure count = %d, %t, want 2", count, ok) - } - if value, ok := fixture.registry.EncryptedValue(keys[1].ID); !ok || value != keys[1].KeyValue { - t.Fatalf("repaired Registry value = %q, %t, want %q", value, ok, keys[1].KeyValue) - } -} - -func TestImportGroupKeysReturnsNotFoundWithoutMutation(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-existing") - var existingKey models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).First(&existingKey).Error; err != nil { - t.Fatal(err) - } - if count, ok := fixture.registry.IncrFailure(existingKey.ID); !ok || count != 1 { - t.Fatalf("seed failure count = %d, %t", count, ok) - } - beforeSnapshot := fixture.manager.Current() - beforeCandidates := fixture.registry.CollectCandidates([]uint{groupID}, nil, time.Time{}) - - _, err := fixture.service.ImportGroupKeys(t.Context(), groupID+1000, GroupKeyImportRequest{Keys: "sk-new"}) - var apiErr *app_errors.APIError - if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrResourceNotFound.Code { - t.Fatalf("ImportGroupKeys() error = %#v, want NOT_FOUND", err) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("not-found import published Snapshot") - } - if got := fixture.registry.CollectCandidates([]uint{groupID}, nil, time.Time{}); !reflect.DeepEqual(got, beforeCandidates) { - t.Fatalf("not-found import changed Registry: %#v", got) - } - if count, ok := fixture.registry.IncrFailure(existingKey.ID); !ok || count != 2 { - t.Fatalf("Registry failure count = %d, %t, want 2", count, ok) - } - assertImportedKeyState(t, fixture, groupID, 1) - - _, err = fixture.service.ImportGroupKeys(t.Context(), 0, GroupKeyImportRequest{Keys: "sk-zero"}) - if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrValidation.Code { - t.Fatalf("ImportGroupKeys(group 0) error = %#v, want VALIDATION_FAILED", err) - } -} - -func TestImportGroupKeysRollsBackOnEncryptionFailure(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-existing") - var existingKey models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).First(&existingKey).Error; err != nil { - t.Fatal(err) - } - if count, ok := fixture.registry.IncrFailure(existingKey.ID); !ok || count != 1 { - t.Fatalf("seed failure count = %d, %t", count, ok) - } - fixture.service.encryption = groupCreateFailingEncryptService{Service: fixture.encryption} - beforeSnapshot := fixture.manager.Current() - beforeCandidates := fixture.registry.CollectCandidates([]uint{groupID}, nil, time.Time{}) - - _, err := fixture.service.ImportGroupKeys(t.Context(), groupID, GroupKeyImportRequest{Keys: "sk-new"}) - if err == nil { - t.Fatal("ImportGroupKeys() error = nil, want encryption failure") - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("failed import published Snapshot") - } - if got := fixture.registry.CollectCandidates([]uint{groupID}, nil, time.Time{}); !reflect.DeepEqual(got, beforeCandidates) { - t.Fatalf("failed import changed Registry: %#v", got) - } - if count, ok := fixture.registry.IncrFailure(existingKey.ID); !ok || count != 2 { - t.Fatalf("Registry failure count = %d, %t, want 2", count, ok) - } - assertImportedKeyState(t, fixture, groupID, 1) -} - -func TestImportGroupKeysCleansPoolAfterCommitBusy(t *testing.T) { - fixture, dsn := newFileServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-existing") - beforeSnapshot := fixture.manager.Current() - releaseReader := holdRollbackJournalReadLock(t, fixture.db, dsn) - - _, err := fixture.service.ImportGroupKeys( - t.Context(), groupID, GroupKeyImportRequest{Keys: "sk-failed-busy"}, - ) - if err == nil { - t.Fatal("ImportGroupKeys() error = nil, want COMMIT BUSY") - } - var apiErr *app_errors.APIError - if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrDatabase.Code { - t.Fatalf("ImportGroupKeys() error = %#v, want DATABASE_ERROR", err) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("key-only import changed Snapshot") - } - - releaseReader() - assertImportedKeyState(t, fixture, groupID, 1) - result, err := fixture.service.ImportGroupKeys( - t.Context(), groupID, GroupKeyImportRequest{Keys: "sk-after-busy"}, - ) - if err != nil || result.KeysAdded != 1 || result.KeysDuplicated != 0 { - t.Fatalf("ImportGroupKeys() = %#v, %v", result, err) - } - assertImportedKeyState(t, fixture, groupID, 2) -} - -func TestImportGroupKeysLeavesGroupFieldsUnchanged(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-existing") - weight := 37 - validationModel := "gpt-field-check" - if err := fixture.db.Model(&models.Group{}).Where("id = ?", groupID).Updates(map[string]any{ - "convert_enabled": true, - "weight_manual": weight, - "validation_model": validationModel, - "config": models.JSON(`{"max_retries":3}`), - "enabled": false, - }).Error; err != nil { - t.Fatal(err) - } - beforeSnapshot := fixture.manager.Current() - var beforeGroup models.Group - if err := fixture.db.First(&beforeGroup, groupID).Error; err != nil { - t.Fatal(err) - } - - result, err := fixture.service.ImportGroupKeys(t.Context(), groupID, GroupKeyImportRequest{ - Keys: "sk-existing\nsk-new\nsk-new", - }) - if err != nil { - t.Fatal(err) - } - if result.KeysAdded != 1 || result.KeysDuplicated != 2 { - t.Fatalf("result = %#v", result) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("key-only import published Snapshot") - } - - var afterGroup models.Group - if err := fixture.db.First(&afterGroup, groupID).Error; err != nil { - t.Fatal(err) - } - if !reflect.DeepEqual(beforeGroup, afterGroup) { - t.Fatalf("Group changed:\nbefore=%#v\nafter=%#v", beforeGroup, afterGroup) - } -} - -func TestConcurrentGroupKeyImportsRemainSerialized(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-existing") - blocking := &groupKeyBlockingEncryptService{ - Service: fixture.encryption, - entered: make(chan struct{}), - release: make(chan struct{}), - } - fixture.service.encryption = blocking - - results := make(chan GroupKeyImportResult, 2) - errors := make(chan error, 2) - startImport := func() { - result, err := fixture.service.ImportGroupKeys(t.Context(), groupID, GroupKeyImportRequest{Keys: "sk-concurrent"}) - results <- result - errors <- err - } - go startImport() - select { - case <-blocking.entered: - case <-time.After(2 * time.Second): - t.Fatal("first import did not reach encryption") - } - if fixture.service.writeMu.TryLock() { - fixture.service.writeMu.Unlock() - t.Fatal("writeMu was not held during key persistence") - } - go startImport() - close(blocking.release) - - added := 0 - duplicated := 0 - for range 2 { - if err := <-errors; err != nil { - t.Fatalf("concurrent ImportGroupKeys() error = %v", err) - } - result := <-results - added += result.KeysAdded - duplicated += result.KeysDuplicated - } - if added != 1 || duplicated != 1 { - t.Fatalf("concurrent totals = added %d, duplicated %d", added, duplicated) - } - assertImportedKeyState(t, fixture, groupID, 2) -} - -type groupKeyBlockingEncryptService struct { - encryption.Service - entered chan struct{} - release chan struct{} - once sync.Once -} - -func (s *groupKeyBlockingEncryptService) Encrypt(plaintext string) (string, error) { - s.once.Do(func() { - close(s.entered) - <-s.release - }) - return s.Service.Encrypt(plaintext) -} - -func createGroupForKeyImport(t *testing.T, fixture serviceFixture, keys string) uint { - t.Helper() - result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://key-import.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: keys, - }) - if err != nil { - t.Fatalf("CreateGroup() error = %v", err) - } - return result.GroupID -} - -func assertImportedKeyState(t *testing.T, fixture serviceFixture, groupID uint, want int) { - t.Helper() - var rows []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&rows).Error; err != nil { - t.Fatal(err) - } - if len(rows) != want { - t.Fatalf("persisted keys = %d, want %d", len(rows), want) - } - if candidates := fixture.registry.CollectCandidates([]uint{groupID}, nil, time.Time{}); len(candidates) != want { - t.Fatalf("Registry candidates = %#v, want %d", candidates, want) - } - for _, row := range rows { - if value, ok := fixture.registry.EncryptedValue(row.ID); !ok || value != row.KeyValue { - t.Fatalf("Registry key %d = %q, %t, want %q", row.ID, value, ok, row.KeyValue) - } - } -} diff --git a/internal/control/group_models.go b/internal/control/group_models.go index d0e62801..3665f1f4 100644 --- a/internal/control/group_models.go +++ b/internal/control/group_models.go @@ -8,6 +8,7 @@ import ( "gorm.io/gorm" app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/pricing" "gpt-load/internal/storage/models" ) @@ -59,12 +60,13 @@ func (s *Service) GetGroupModels(ctx context.Context, groupID uint) (GroupModels return GroupModelsResponse{}, err } - return mapGroupModelsResponse(groupModels, rows) + return mapGroupModelsResponse(group.ChannelID, groupModels, rows) } func mapGroupModelsResponse( + channelID string, groupModels []GroupModel, - rows map[string]*models.ModelPrice, + rows modelPriceRows, ) (GroupModelsResponse, error) { result := GroupModelsResponse{Items: make([]GroupModelResponse, 0, len(groupModels))} for _, model := range groupModels { @@ -78,7 +80,7 @@ func mapGroupModelsResponse( if item.AliasEnabled { item.ClientModel = model.Alias } - item.PricingStatus = resolvePricingStatus(rows[model.ID]) + item.PricingStatus = resolvePricingStatus(rows[pricing.Identity{ChannelID: channelID, ModelID: model.ID}]) if item.PricingStatus == PricingStatusPending { result.Pending++ } @@ -114,17 +116,17 @@ func (s *Service) UpdateGroupModels( if err != nil { return err } - if err := validateGroupRowCandidate(ctx, tx, group); err != nil { + if err := validateGroupRowCandidate(ctx, tx, group, s.channelRegistry); err != nil { return fmt.Errorf("validate existing group %d: %w", groupID, app_errors.ErrInternalServer) } var previous []GroupModel if err := decodeGroupDiscoveryJSON(group.Models, &previous); err != nil { return fmt.Errorf("decode group %d models: %w", groupID, app_errors.ErrInternalServer) } - modelIDsChanged = group.ProviderID != nil && !sameGroupModelIDs(previous, normalized) + modelIDsChanged = !sameGroupModelIDs(previous, normalized) group.Models = models.JSON(encoded) - if err := validateGroupRowCandidate(ctx, tx, group); err != nil { + if err := validateGroupRowCandidate(ctx, tx, group, s.channelRegistry); err != nil { return app_errors.ErrValidation } if err := tx.Model(&models.Group{}). diff --git a/internal/control/group_models_test.go b/internal/control/group_models_test.go index 7f3c389e..765a09bb 100644 --- a/internal/control/group_models_test.go +++ b/internal/control/group_models_test.go @@ -9,7 +9,8 @@ import ( "time" "gpt-load/internal/catalog" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" + "gpt-load/internal/execution" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/pricing" "gpt-load/internal/protocol" @@ -19,7 +20,6 @@ import ( func TestGetGroupModelsReturnsClientNamesAndPricingStatus(t *testing.T) { fixture := newServiceFixture(t) - providerID := "openai" fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ "openai": { ID: "openai", @@ -34,14 +34,13 @@ func TestGetGroupModelsReturnsClientNamesAndPricingStatus(t *testing.T) { }, }}) created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - ProviderID: &providerID, - UpstreamURL: "https://model-read.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{ {ID: "gpt-4o", Alias: "default", AliasEnabled: true}, {ID: "missing-price", Alias: ""}, }}, - Keys: "sk-model-read", + Credentials: "sk-model-read", }) if err != nil { t.Fatal(err) @@ -66,9 +65,11 @@ func TestGetGroupModelsReturnsClientNamesAndPricingStatus(t *testing.T) { func TestMapGroupModelsResponseTreatsContextTierOnlyPriceAsConfigured(t *testing.T) { result, err := mapGroupModelsResponse( + string(channel.OpenAI), []GroupModel{{ID: "tiered-model"}}, - map[string]*models.ModelPrice{ - "tiered-model": { + modelPriceRows{ + {ChannelID: string(channel.OpenAI), ModelID: "tiered-model"}: { + ChannelID: string(channel.OpenAI), ModelID: "tiered-model", ContextPriceTiers: models.JSON(`[{"threshold_tokens":1000,"input_price_nano_usd_per_million_tokens":1,"output_price_nano_usd_per_million_tokens":null,"cache_read_price_nano_usd_per_million_tokens":null,"cache_write_price_nano_usd_per_million_tokens":null}]`), }, @@ -198,7 +199,7 @@ func TestNormalizeGroupModelsRejectsDuplicateExternalNames(t *testing.T) { func TestUpdateGroupModelsRequiresNonNullModelsField(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-required-models") + groupID := createGroupForCredentialImport(t, fixture, "sk-required-models") for _, request := range []GroupModelsUpdateRequest{ {}, {Models: optionalGroupModels{Set: false}}, @@ -219,13 +220,13 @@ func TestUpdateGroupModelsReplacesAuthoritativeListAndPublishesOnce(t *testing.T fixture := newServiceFixture(t) mustEnsureInitialPrices(t, fixture) created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://model-save.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://model-save.example.com/v1"}`), Models: optionalGroupModels{ Set: true, Values: []GroupModel{{ID: "provider-old", Alias: "old-public", AliasEnabled: true}}, }, - Keys: "sk-model-save-a\nsk-model-save-b", + Credentials: "sk-model-save-a\nsk-model-save-b", }) if err != nil { t.Fatal(err) @@ -238,7 +239,7 @@ func TestUpdateGroupModelsReplacesAuthoritativeListAndPublishesOnce(t *testing.T } if err := fixture.db.Model(&models.Group{}). Where("id = ?", created.GroupID). - Update("config", models.JSON(`{ + Update("overrides", models.JSON(`{ "stream_idle_timeout":45, "inject_usage_options":false, "header_rules":{"remove":["X-Trace"]} @@ -252,8 +253,8 @@ func TestUpdateGroupModelsReplacesAuthoritativeListAndPublishesOnce(t *testing.T } beforeRevision := fixture.manager.Current().Revision beforeRegistry := fixture.registry.Snapshot() - var beforeKeys []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", created.GroupID).Order("id ASC").Find(&beforeKeys).Error; err != nil { + var beforeCredentials []models.Credential + if err := fixture.db.Where("group_id = ?", created.GroupID).Order("id ASC").Find(&beforeCredentials).Error; err != nil { t.Fatal(err) } @@ -292,7 +293,7 @@ func TestUpdateGroupModelsReplacesAuthoritativeListAndPublishesOnce(t *testing.T if err != nil { t.Fatalf("GetGroupSummary() error = %v", err) } - if settings.ValidationModel == nil || *settings.ValidationModel != validation || summary.KeyCount != 2 { + if settings.ValidationModel == nil || *settings.ValidationModel != validation || summary.CredentialCount != 2 { t.Fatalf("settings/summary = %#v/%#v", settings, summary) } streamIdle, ok := settings.Overrides[state.SettingStreamIdleTimeout].(json.Number) @@ -322,12 +323,12 @@ func TestUpdateGroupModelsReplacesAuthoritativeListAndPublishesOnce(t *testing.T if !reflect.DeepEqual(fixture.registry.Snapshot(), beforeRegistry) { t.Fatal("models save changed Registry") } - var afterKeys []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", created.GroupID).Order("id ASC").Find(&afterKeys).Error; err != nil { + var afterCredentials []models.Credential + if err := fixture.db.Where("group_id = ?", created.GroupID).Order("id ASC").Find(&afterCredentials).Error; err != nil { t.Fatal(err) } - if !reflect.DeepEqual(afterKeys, beforeKeys) { - t.Fatalf("upstream keys changed: got=%#v want=%#v", afterKeys, beforeKeys) + if !reflect.DeepEqual(afterCredentials, beforeCredentials) { + t.Fatalf("credentials changed: got=%#v want=%#v", afterCredentials, beforeCredentials) } snapshot := fixture.manager.Current() view := snapshot.Groups[created.GroupID] @@ -337,7 +338,7 @@ func TestUpdateGroupModelsReplacesAuthoritativeListAndPublishesOnce(t *testing.T !reflect.DeepEqual(settings.Effective.HeaderRules.Remove, view.HeaderRules.Remove) { t.Fatalf("effective/snapshot = %#v/%#v", settings.Effective, view) } - targets := snapshot.Candidates[protocol.OpenAICompletions] + targets := snapshot.ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion] if len(targets) != 2 || targets["public-a"][0].UpstreamModelID != "provider-a" || targets["public-b"][0].UpstreamModelID != "provider-b" { @@ -346,7 +347,7 @@ func TestUpdateGroupModelsReplacesAuthoritativeListAndPublishesOnce(t *testing.T if _, exists := targets["old-public"]; exists { t.Fatalf("authoritative replacement retained old model: %#v", targets) } - routes := snapshot.RouteCatalog[protocol.OpenAICompletions] + routes := snapshot.ExecutionRouteCatalog[protocol.OpenAICompletions][execution.OperationChatCompletion] if len(routes) != 2 || routes["public-a"][0].UpstreamModelID != "provider-a" || routes["public-b"][0].UpstreamModelID != "provider-b" { @@ -358,13 +359,13 @@ func TestUpdateGroupModelsAllowsEmptyList(t *testing.T) { fixture := newServiceFixture(t) mustEnsureInitialPrices(t, fixture) created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://empty-models.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://empty-models.example.com/v1"}`), Models: optionalGroupModels{ Set: true, Values: []GroupModel{{ID: "provider-old", Alias: "old-public", AliasEnabled: true}}, }, - Keys: "sk-empty-models", + Credentials: "sk-empty-models", }) if err != nil { t.Fatal(err) @@ -383,17 +384,17 @@ func TestUpdateGroupModelsAllowsEmptyList(t *testing.T) { if fixture.manager.Current().Revision != before+1 { t.Fatalf("revision = %d, want %d", fixture.manager.Current().Revision, before+1) } - if len(fixture.manager.Current().Candidates[protocol.OpenAICompletions]) != 0 || - len(fixture.manager.Current().RouteCatalog[protocol.OpenAICompletions]) != 0 { + if len(fixture.manager.Current().ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion]) != 0 || + len(fixture.manager.Current().ExecutionRouteCatalog[protocol.OpenAICompletions][execution.OperationChatCompletion]) != 0 { t.Fatalf("model indexes = candidates:%#v routes:%#v", - fixture.manager.Current().Candidates, fixture.manager.Current().RouteCatalog) + fixture.manager.Current().ExecutionCandidates, fixture.manager.Current().ExecutionRouteCatalog) } } func TestUpdateGroupModelsNeverCallsDiscoveryOrChangesAccessKeyFilters(t *testing.T) { fixture := newServiceFixture(t) mustEnsureInitialPrices(t, fixture) - groupID := createGroupForKeyImport(t, fixture, "sk-no-discovery") + groupID := createGroupForCredentialImport(t, fixture, "sk-no-discovery") access, err := fixture.service.CreateAccessKey(t.Context(), AccessKeyCreateRequest{ Name: "filtered", Filters: &AccessKeyFilters{ @@ -408,7 +409,7 @@ func TestUpdateGroupModelsNeverCallsDiscoveryOrChangesAccessKeyFilters(t *testin if err := fixture.db.First(&beforeAccess, access.ID).Error; err != nil { t.Fatal(err) } - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { t.Fatal("UpdateGroupModels must not call model discovery") @@ -444,7 +445,7 @@ func TestUpdateGroupModelsNeverCallsDiscoveryOrChangesAccessKeyFilters(t *testin func TestUpdateGroupModelsFailuresDoNotPublish(t *testing.T) { t.Run("external collision", func(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-invalid-models") + groupID := createGroupForCredentialImport(t, fixture, "sk-invalid-models") beforeRevision := fixture.manager.Current().Revision beforeRegistry := fixture.registry.Snapshot() beforeModels := loadCreatedGroupModels(t, fixture, groupID) @@ -467,12 +468,12 @@ func TestUpdateGroupModelsFailuresDoNotPublish(t *testing.T) { t.Run("full compile failure", func(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-compile-models") + groupID := createGroupForCredentialImport(t, fixture, "sk-compile-models") corrupt := validControlGroup("model-save-corrupt-other") if err := fixture.db.Create(corrupt).Error; err != nil { t.Fatal(err) } - if err := fixture.db.Exec("UPDATE groups SET protocols = ? WHERE id = ?", `[]`, corrupt.ID).Error; err != nil { + if err := fixture.db.Exec("UPDATE groups SET channel_id = ? WHERE id = ?", "unknown", corrupt.ID).Error; err != nil { t.Fatal(err) } beforeRevision := fixture.manager.Current().Revision @@ -494,13 +495,13 @@ func TestUpdateGroupModelsFailuresDoNotPublish(t *testing.T) { t.Run("commit failure", func(t *testing.T) { fixture, dsn := newFileServiceFixture(t) created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://commit-failure-models.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://commit-failure-models.example.com/v1"}`), Models: optionalGroupModels{ Set: true, Values: []GroupModel{{ID: "provider-old", Alias: "old-public", AliasEnabled: true}}, }, - Keys: "sk-commit-models", + Credentials: "sk-commit-models", }) if err != nil { t.Fatal(err) @@ -532,7 +533,7 @@ func assertModelsUpdateStateUnchanged( fixture serviceFixture, groupID uint, wantRevision uint64, - wantRegistry []state.KeyRuntimeView, + wantRegistry []state.CredentialRuntimeView, wantModels []GroupModel, ) { t.Helper() diff --git a/internal/control/group_options.go b/internal/control/group_options.go index e4ff3047..2628171f 100644 --- a/internal/control/group_options.go +++ b/internal/control/group_options.go @@ -8,24 +8,26 @@ import ( "gorm.io/gorm" + "gpt-load/internal/channel" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/storage/models" ) type GroupOption struct { - ID uint `json:"id"` - Name string `json:"name"` - Enabled bool `json:"enabled"` - Protocols []protocol.Protocol `json:"protocols"` - Models []string `json:"models"` + ID uint `json:"id"` + Name string `json:"name"` + ChannelID channel.ID `json:"channel_id"` + Params json.RawMessage `json:"params"` + Enabled bool `json:"enabled"` + Models []string `json:"models"` } type groupOptionRow struct { ID uint Name string + ChannelID string + Params models.JSON Enabled bool - Protocols models.JSON Models models.JSON } @@ -49,7 +51,7 @@ func (s *Service) ListGroupOptions(ctx context.Context) ([]GroupOption, error) { if err != nil { return nil, app_errors.ParseDBError(err) } - options, err := mapGroupOptions(rows) + options, err := mapGroupOptions(rows, s.channelRegistry) if parentErr := ctx.Err(); parentErr != nil { return nil, parentErr } @@ -63,7 +65,7 @@ func (s *Service) readGroupOptionRows(ctx context.Context) ([]groupOptionRow, er var rows []groupOptionRow err := s.withReadSnapshot(ctx, func(tx *gorm.DB) error { return tx.Model(&models.Group{}). - Select("id", "name", "enabled", "protocols", "models"). + Select("id", "name", "channel_id", "params", "enabled", "models"). Order("id ASC"). Find(&rows).Error }) @@ -71,26 +73,31 @@ func (s *Service) readGroupOptionRows(ctx context.Context) ([]groupOptionRow, er return nil, err } for index := range rows { - rows[index].Protocols = append(models.JSON(nil), rows[index].Protocols...) + rows[index].Params = append(models.JSON(nil), rows[index].Params...) rows[index].Models = append(models.JSON(nil), rows[index].Models...) } return rows, nil } -func mapGroupOptions(rows []groupOptionRow) ([]GroupOption, error) { +func mapGroupOptions(rows []groupOptionRow, registries ...*channel.Registry) ([]GroupOption, error) { + var registry *channel.Registry + for _, candidate := range registries { + if candidate != nil { + registry = candidate + break + } + } options := make([]GroupOption, 0, len(rows)) for _, row := range rows { - var protocols []protocol.Protocol - if err := json.Unmarshal(row.Protocols, &protocols); err != nil { - return nil, groupCollectionDataError( - "decode group %d protocols: %v", row.ID, err, - ) + channelID := channel.ID(row.ChannelID) + if registry == nil { + return nil, groupCollectionDataError("channel registry is nil") } - if err := validateGroupCollectionProtocols(protocols); err != nil { - return nil, groupCollectionDataError( - "validate group %d protocols: %v", row.ID, err, - ) + validated, err := registry.ValidateParams(channelID, json.RawMessage(row.Params)) + if err != nil { + return nil, groupCollectionDataError("validate group %d params: %v", row.ID, err) } + params := validated.CanonicalJSON() var models []GroupModel if err := json.Unmarshal(row.Models, &models); err != nil { return nil, groupCollectionDataError( @@ -103,11 +110,9 @@ func mapGroupOptions(rows []groupOptionRow) ([]GroupOption, error) { ) } option := GroupOption{ - ID: row.ID, - Name: row.Name, - Enabled: row.Enabled, - Protocols: append([]protocol.Protocol(nil), protocols...), - Models: make([]string, 0, len(models)), + ID: row.ID, Name: row.Name, ChannelID: channelID, + Params: append(json.RawMessage(nil), params...), Enabled: row.Enabled, + Models: make([]string, 0, len(models)), } for _, model := range models { alias := strings.TrimSpace(model.Alias) diff --git a/internal/control/group_options_test.go b/internal/control/group_options_test.go index bbd527e0..5ae3293d 100644 --- a/internal/control/group_options_test.go +++ b/internal/control/group_options_test.go @@ -8,24 +8,24 @@ import ( "strings" "testing" + "gpt-load/internal/channel" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/storage/models" ) func TestListGroupOptionsReturnsAllGroupsByIDWithExternalModels(t *testing.T) { fixture := newServiceFixture(t) createGroupOptionGroup(t, fixture, 20, "later enabled", true, - `["openai-completions","gemini"]`, + channel.OpenAICompatible, `{"base_url":"https://later-enabled.example/v1"}`, `[{"id":" upstream-first ","alias":" public-first "},{"id":" upstream-second ","alias":" "}]`, ) createGroupOptionGroup(t, fixture, 10, "earlier disabled", false, - `["anthropic"]`, + channel.Anthropic, `{}`, `[{"id":" upstream-third ","alias":""},{"id":" upstream-fourth ","alias":" public-fourth "}]`, ) - if err := fixture.db.Create(&models.UpstreamKey{ - GroupID: 10, KeyValue: "ciphertext-secret", KeyHash: "hash-secret", - Status: models.UpstreamKeyStatusActive, + if err := fixture.db.Create(&models.Credential{ + GroupID: 10, Data: "ciphertext-secret", Fingerprint: "hash-secret", + Status: models.CredentialStatusActive, }).Error; err != nil { t.Fatalf("create unrelated upstream key: %v", err) } @@ -34,20 +34,14 @@ func TestListGroupOptionsReturnsAllGroupsByIDWithExternalModels(t *testing.T) { if err != nil { t.Fatalf("ListGroupOptions() error = %v", err) } - want := []GroupOption{ - { - ID: 10, Name: "earlier disabled", Enabled: false, - Protocols: []protocol.Protocol{protocol.Anthropic}, - Models: []string{"upstream-third", "public-fourth"}, - }, - { - ID: 20, Name: "later enabled", Enabled: true, - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Gemini}, - Models: []string{"public-first", "upstream-second"}, - }, - } - if !reflect.DeepEqual(options, want) { - t.Fatalf("ListGroupOptions() = %#v, want %#v", options, want) + if len(options) != 2 || options[0].ID != 10 || options[0].Name != "earlier disabled" || + options[0].ChannelID != channel.Anthropic || string(options[0].Params) != `{}` || + options[0].Enabled || !reflect.DeepEqual(options[0].Models, []string{"upstream-third", "public-fourth"}) || + options[1].ID != 20 || options[1].Name != "later enabled" || + options[1].ChannelID != channel.OpenAICompatible || + string(options[1].Params) != `{"base_url":"https://later-enabled.example/v1"}` || + !options[1].Enabled || !reflect.DeepEqual(options[1].Models, []string{"public-first", "upstream-second"}) { + t.Fatalf("ListGroupOptions() = %#v, want terminal channel options", options) } encoded, err := json.Marshal(options) @@ -55,7 +49,7 @@ func TestListGroupOptionsReturnsAllGroupsByIDWithExternalModels(t *testing.T) { t.Fatalf("json.Marshal(options) error = %v", err) } for _, forbiddenField := range []string{ - "upstream_url", "key_value", "key_hash", "keyvalue", "keyhash", + "upstream_url", "protocols", "key_value", "key_hash", "keyvalue", "keyhash", } { if containsJSONToken(encoded, forbiddenField) { t.Fatalf("options JSON exposes field %q: %s", forbiddenField, encoded) @@ -71,7 +65,7 @@ func TestListGroupOptionsReturnsAllGroupsByIDWithExternalModels(t *testing.T) { func TestListGroupOptionsFailsClosedForInvalidDataDatabaseAndCancellation(t *testing.T) { t.Run("invalid models JSON", func(t *testing.T) { fixture := newServiceFixture(t) - createGroupOptionGroup(t, fixture, 1, "invalid", true, `["openai-completions"]`, `{"not":"an array"}`) + createGroupOptionGroup(t, fixture, 1, "invalid", true, channel.OpenAI, `{}`, `{"not":"an array"}`) options, err := fixture.service.ListGroupOptions(t.Context()) if options != nil || !errors.Is(err, app_errors.ErrInternalServer) { @@ -79,14 +73,17 @@ func TestListGroupOptionsFailsClosedForInvalidDataDatabaseAndCancellation(t *tes } }) - for name, rawProtocols := range map[string]string{ - "invalid protocols JSON": `{"not":"an array"}`, - "invalid protocol": `["openai"]`, - "duplicate protocol": `["anthropic","anthropic"]`, + for name, test := range map[string]struct { + channelID channel.ID + params string + }{ + "invalid params shape": {channelID: channel.OpenAICompatible, params: `[]`}, + "invalid params": {channelID: channel.OpenAICompatible, params: `{"base_url":"relative"}`}, + "unknown channel": {channelID: channel.ID("unknown"), params: `{}`}, } { t.Run(name, func(t *testing.T) { fixture := newServiceFixture(t) - createGroupOptionGroup(t, fixture, 1, "invalid", true, rawProtocols, `[]`) + createGroupOptionGroup(t, fixture, 1, "invalid", true, test.channelID, test.params, `[]`) options, err := fixture.service.ListGroupOptions(t.Context()) if options != nil || !errors.Is(err, app_errors.ErrInternalServer) { @@ -129,15 +126,16 @@ func createGroupOptionGroup( id uint, name string, enabled bool, - rawProtocols string, + channelID channel.ID, + rawParams string, rawModels string, ) { t.Helper() group := validControlGroup(name) group.ID = id - group.Protocols = models.JSON(rawProtocols) + group.ChannelID = string(channelID) + group.Params = models.JSON(rawParams) group.Models = models.JSON(rawModels) - group.UpstreamURL = "https://" + strings.ReplaceAll(name, " ", "-") + ".example/v1" if err := fixture.db.Create(group).Error; err != nil { t.Fatalf("create group %q: %v", name, err) } diff --git a/internal/control/group_settings.go b/internal/control/group_settings.go index acd5e346..2362a0e3 100644 --- a/internal/control/group_settings.go +++ b/internal/control/group_settings.go @@ -4,9 +4,11 @@ import ( "context" "encoding/json" "fmt" + "slices" "gorm.io/gorm" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/protocol" @@ -15,10 +17,9 @@ import ( ) type GroupSettingsResponse struct { + ChannelID channel.ID `json:"channel_id"` + Params json.RawMessage `json:"params"` Name string `json:"name"` - ProviderID *string `json:"provider_id"` - UpstreamURL string `json:"upstream_url"` - Protocols []protocol.Protocol `json:"protocols"` ValidationModel *string `json:"validation_model"` Enabled bool `json:"enabled"` WeightManual *int `json:"weight_manual"` @@ -27,32 +28,24 @@ type GroupSettingsResponse struct { } type GroupSettingsUpdateRequest struct { - Name optionalField[string] `json:"name"` - ProviderID optionalField[string] `json:"provider_id"` - UpstreamURL optionalField[string] `json:"upstream_url"` - Protocols optionalField[[]protocol.Protocol] `json:"protocols"` - ValidationModel optionalField[string] `json:"validation_model"` - Enabled optionalField[bool] `json:"enabled"` - WeightManual optionalField[int] `json:"weight_manual"` - Overrides optionalField[config.Settings] `json:"overrides"` - // Deprecated: retained for compatibility; URL changes no longer require confirmation. - ConfirmUpstreamChange bool `json:"confirm_upstream_change"` + Name optionalField[string] `json:"name"` + Params optionalField[json.RawMessage] `json:"params"` + ValidationModel optionalField[string] `json:"validation_model"` + Enabled optionalField[bool] `json:"enabled"` + WeightManual optionalField[int] `json:"weight_manual"` + Overrides optionalField[config.Settings] `json:"overrides"` } type normalizedGroupSettingsUpdate struct { name *string - providerID *string - providerIDSet bool - upstreamURL *string - protocols []protocol.Protocol - protocolsSet bool + params json.RawMessage + paramsSet bool validationModel *string validationModelSet bool enabled *bool weightManual *int weightManualSet bool encodedOverrides models.JSON - overrides config.Settings overridesSet bool } @@ -71,24 +64,32 @@ func (s *Service) GetGroupSettings(ctx context.Context, groupID uint) (GroupSett snapshot := s.manager.Current() if snapshot == nil { return GroupSettingsResponse{}, fmt.Errorf( - "runtime snapshot unavailable: %w", - app_errors.ErrInternalServer, + "runtime snapshot unavailable: %w", app_errors.ErrInternalServer, ) } - return groupSettingsResponse(group, snapshot.Settings) + return groupSettingsResponse(group, snapshot.Settings, s.channelRegistry) } func groupSettingsResponse( group models.Group, system state.RuntimeSettings, + registry *channel.Registry, ) (GroupSettingsResponse, error) { - protocols := make([]protocol.Protocol, 0) - if err := decodeGroupDiscoveryJSON(group.Protocols, &protocols); err != nil { - return GroupSettingsResponse{}, fmt.Errorf("decode group %d protocols: %w", group.ID, err) + if registry == nil || group.ChannelID == "" { + return GroupSettingsResponse{}, fmt.Errorf( + "resolve group %d channel: %w", group.ID, app_errors.ErrInternalServer, + ) + } + channelID := channel.ID(group.ChannelID) + validated, err := registry.ValidateParams(channelID, json.RawMessage(group.Params)) + if err != nil { + return GroupSettingsResponse{}, fmt.Errorf( + "validate group %d params: %w", group.ID, app_errors.ErrInternalServer, + ) } overrides := make(config.Settings) - if len(group.Config) > 0 { - if err := decodeGroupDiscoveryJSON(group.Config, &overrides); err != nil { + if len(group.Overrides) > 0 { + if err := decodeGroupDiscoveryJSON(group.Overrides, &overrides); err != nil { return GroupSettingsResponse{}, fmt.Errorf("decode group %d config: %w", group.ID, err) } } @@ -98,16 +99,13 @@ func groupSettingsResponse( effective, err := effectiveGroupConfig(system, overrides) if err != nil { return GroupSettingsResponse{}, fmt.Errorf( - "resolve group %d effective config: %w", - group.ID, - app_errors.ErrInternalServer, + "resolve group %d effective config: %w", group.ID, app_errors.ErrInternalServer, ) } return GroupSettingsResponse{ + ChannelID: channelID, + Params: validated.CanonicalJSON(), Name: group.Name, - ProviderID: cloneString(group.ProviderID), - UpstreamURL: group.UpstreamURL, - Protocols: protocols, ValidationModel: cloneString(group.ValidationModel), Enabled: group.Enabled, WeightManual: cloneInt(group.WeightManual), @@ -129,8 +127,7 @@ func normalizeGroupSettingsUpdate( ) (normalizedGroupSettingsUpdate, error) { for _, nullable := range []bool{ request.Name.Set && request.Name.Null, - request.UpstreamURL.Set && request.UpstreamURL.Null, - request.Protocols.Set && request.Protocols.Null, + request.Params.Set && request.Params.Null, request.Enabled.Set && request.Enabled.Null, request.Overrides.Set && request.Overrides.Null, } { @@ -138,9 +135,8 @@ func normalizeGroupSettingsUpdate( return normalizedGroupSettingsUpdate{}, app_errors.ErrValidation } } - if !request.Name.Set && !request.ProviderID.Set && !request.UpstreamURL.Set && !request.Protocols.Set && - !request.ValidationModel.Set && !request.Enabled.Set && !request.WeightManual.Set && - !request.Overrides.Set { + if !request.Name.Set && !request.Params.Set && !request.ValidationModel.Set && + !request.Enabled.Set && !request.WeightManual.Set && !request.Overrides.Set { return normalizedGroupSettingsUpdate{}, app_errors.ErrBadRequest } @@ -152,30 +148,9 @@ func normalizeGroupSettingsUpdate( } result.name = value } - if request.ProviderID.Set { - result.providerIDSet = true - if !request.ProviderID.Null { - value, err := normalizeProviderID(&request.ProviderID.Value) - if err != nil { - return normalizedGroupSettingsUpdate{}, app_errors.ErrValidation - } - result.providerID = value - } - } - if request.UpstreamURL.Set { - value, _, err := normalizeUpstreamBaseURL(request.UpstreamURL.Value) - if err != nil { - return normalizedGroupSettingsUpdate{}, err - } - result.upstreamURL = &value - } - if request.Protocols.Set { - values, err := normalizeGroupProtocols(request.Protocols.Value) - if err != nil { - return normalizedGroupSettingsUpdate{}, err - } - result.protocols = values - result.protocolsSet = true + if request.Params.Set { + result.paramsSet = true + result.params = append(json.RawMessage(nil), request.Params.Value...) } if request.ValidationModel.Set { result.validationModelSet = true @@ -202,12 +177,11 @@ func normalizeGroupSettingsUpdate( } } if request.Overrides.Set { - settings, encoded, err := normalizeGroupSettings(request.Overrides.Value) + _, encoded, err := normalizeGroupSettings(request.Overrides.Value) if err != nil { return normalizedGroupSettingsUpdate{}, err } result.encodedOverrides = encoded - result.overrides = settings result.overridesSet = true } return result, nil @@ -227,52 +201,29 @@ func (s *Service) UpdateGroupSettings( } var committed models.Group - providerReferencesChanged := false snapshot, err := s.writeGroupConfig(ctx, func(tx *gorm.DB) error { group, err := loadGroupRow(tx, groupID) if err != nil { return err } - if err := validateGroupRowCandidate(ctx, tx, group); err != nil { + if err := validateGroupRowCandidate(ctx, tx, group, s.channelRegistry); err != nil { return fmt.Errorf("validate existing group %d: %w", groupID, app_errors.ErrInternalServer) } - updates := make(map[string]any, 8) + updates := make(map[string]any, 6) if normalized.name != nil { group.Name = *normalized.name updates["name"] = group.Name } - if normalized.providerIDSet { - if !optionalStringsEqual(group.ProviderID, normalized.providerID) { - if err := s.validateSelectableProviderID(normalized.providerID); err != nil { - return err - } - var groupModels []GroupModel - if err := decodeGroupDiscoveryJSON(group.Models, &groupModels); err != nil { - return fmt.Errorf("decode group %d models: %w", groupID, app_errors.ErrInternalServer) - } - providerReferencesChanged = len(groupModels) > 0 + if normalized.paramsSet { + params, validateErr := s.channelRegistry.ValidateParams( + channel.ID(group.ChannelID), normalized.params, + ) + if validateErr != nil { + return app_errors.ErrValidation } - group.ProviderID = cloneString(normalized.providerID) - updates["provider_id"] = cloneString(normalized.providerID) - } - if normalized.upstreamURL != nil { - currentURL, _, normalizeErr := normalizeUpstreamBaseURL(group.UpstreamURL) - if normalizeErr != nil { - return fmt.Errorf("normalize persisted group URL: %w", app_errors.ErrInternalServer) - } - if currentURL != *normalized.upstreamURL { - group.UpstreamURL = *normalized.upstreamURL - updates["upstream_url"] = group.UpstreamURL - } - } - if normalized.protocolsSet { - encoded, encodeErr := json.Marshal(normalized.protocols) - if encodeErr != nil { - return encodeErr - } - group.Protocols = models.JSON(encoded) - updates["protocols"] = group.Protocols + group.Params = models.JSON(params.CanonicalJSON()) + updates["params"] = append(models.JSON(nil), group.Params...) } if normalized.validationModelSet { group.ValidationModel = normalized.validationModel @@ -287,17 +238,19 @@ func (s *Service) UpdateGroupSettings( updates["weight_manual"] = normalized.weightManual } if normalized.overridesSet { - group.Config = normalized.encodedOverrides - updates["config"] = group.Config + group.Overrides = normalized.encodedOverrides + updates["overrides"] = group.Overrides } - if err := validateGroupInjectUsageOptionsConstraint(group); err != nil { + if err := validateGroupInjectUsageOptionsConstraint(group, s.channelRegistry); err != nil { return err } - if err := validateGroupRowCandidate(ctx, tx, group); err != nil { + if err := validateGroupRowCandidate(ctx, tx, group, s.channelRegistry); err != nil { return app_errors.ErrValidation } - if err := tx.Model(&models.Group{}).Where("id = ?", groupID).Updates(updates).Error; err != nil { - return app_errors.ParseDBError(err) + if len(updates) > 0 { + if err := tx.Model(&models.Group{}).Where("id = ?", groupID).Updates(updates).Error; err != nil { + return app_errors.ParseDBError(err) + } } committed = group return nil @@ -305,46 +258,25 @@ func (s *Service) UpdateGroupSettings( if err != nil { return GroupSettingsResponse{}, withControlOperationContext(err, groupID, 0) } - if providerReferencesChanged && s.catalogSync != nil { - s.catalogSync.RequestGroupSync() - } - result, err := groupSettingsResponse(committed, snapshot.Settings) - if err != nil { - return GroupSettingsResponse{}, err - } - return result, nil + return groupSettingsResponse(committed, snapshot.Settings, s.channelRegistry) } -func optionalStringsEqual(left, right *string) bool { - if left == nil || right == nil { - return left == nil && right == nil - } - return *left == *right -} - -func validateGroupInjectUsageOptionsConstraint(group models.Group) error { +func validateGroupInjectUsageOptionsConstraint(group models.Group, registry *channel.Registry) error { settings := make(config.Settings) - if len(group.Config) > 0 { - if err := decodeGroupDiscoveryJSON(group.Config, &settings); err != nil { + if len(group.Overrides) > 0 { + if err := decodeGroupDiscoveryJSON(group.Overrides, &settings); err != nil { return app_errors.ErrValidation } } - if _, injectUsageOptionsSet := settings[state.SettingInjectUsageOptions]; injectUsageOptionsSet && - !groupSupportsOpenAICompletions(group.Protocols) { + if _, set := settings[state.SettingInjectUsageOptions]; !set { + return nil + } + if registry == nil { + return app_errors.ErrValidation + } + descriptor, ok := registry.Get(channel.ID(group.ChannelID)) + if !ok || !slices.Contains(descriptor.ClientProtocols, protocol.OpenAICompletions) { return app_errors.ErrValidation } return nil } - -func groupSupportsOpenAICompletions(raw models.JSON) bool { - var protocols []protocol.Protocol - if err := decodeGroupDiscoveryJSON(raw, &protocols); err != nil { - return false - } - for _, value := range protocols { - if value == protocol.OpenAICompletions { - return true - } - } - return false -} diff --git a/internal/control/group_settings_test.go b/internal/control/group_settings_test.go index 3612bf32..dfd2e56a 100644 --- a/internal/control/group_settings_test.go +++ b/internal/control/group_settings_test.go @@ -12,9 +12,9 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/storage/models" ) @@ -22,7 +22,7 @@ import ( func TestGetGroupSettingsReturnsPersistedDraftOverridesAndEffectiveConfig(t *testing.T) { fixture := newServiceFixture(t) group := validControlGroup("settings-read") - group.Config = models.JSON(`{ + group.Overrides = models.JSON(`{ "connect_timeout":20, "first_byte_timeout":180, "request_timeout":480, @@ -41,8 +41,8 @@ func TestGetGroupSettingsReturnsPersistedDraftOverridesAndEffectiveConfig(t *tes if err != nil { t.Fatalf("GetGroupSettings() error = %v", err) } - if got.Name != group.Name || got.UpstreamURL != group.UpstreamURL || - !reflect.DeepEqual(got.Protocols, []protocol.Protocol{protocol.OpenAICompletions}) || + if got.Name != group.Name || got.ChannelID != channel.OpenAICompatible || + string(got.Params) != `{"base_url":"https://settings-read.example/v1"}` || !got.Enabled || got.WeightManual != nil { t.Fatalf("persisted settings = %#v", got) } @@ -69,23 +69,31 @@ func TestGetGroupSettingsReturnsPersistedDraftOverridesAndEffectiveConfig(t *tes func TestUpdateGroupSettingsPublishesOnceAndReturnsNewSettings(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-settings-update") + created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://settings-before.example.com/v1"}`), + Models: optionalGroupModels{Set: true}, Credentials: "sk-settings-update", + }) + if err != nil { + t.Fatal(err) + } + groupID := created.GroupID beforeRevision := fixture.manager.Current().Revision weight := 75 result, err := fixture.service.UpdateGroupSettings(t.Context(), groupID, GroupSettingsUpdateRequest{ - Name: optionalField[string]{Set: true, Value: " updated settings "}, - UpstreamURL: optionalField[string]{Set: true, Value: " HTTPS://SETTINGS-UPDATED.EXAMPLE.COM/v1/ "}, - Protocols: optionalField[[]protocol.Protocol]{Set: true, Value: []protocol.Protocol{protocol.OpenAICompletions}}, - ValidationModel: optionalField[string]{Set: true, Value: " gpt-4.1 "}, - Enabled: optionalField[bool]{Set: true, Value: false}, - WeightManual: optionalField[int]{Set: true, Value: weight}, - Overrides: optionalField[config.Settings]{Set: true, Value: config.Settings{"request_timeout": json.Number("720")}}, - ConfirmUpstreamChange: true, + Name: optionalField[string]{Set: true, Value: " updated settings "}, + Params: optionalField[json.RawMessage]{Set: true, + Value: json.RawMessage(`{"base_url":" HTTPS://SETTINGS-UPDATED.EXAMPLE.COM/v1/ "}`)}, + ValidationModel: optionalField[string]{Set: true, Value: " gpt-4.1 "}, + Enabled: optionalField[bool]{Set: true, Value: false}, + WeightManual: optionalField[int]{Set: true, Value: weight}, + Overrides: optionalField[config.Settings]{Set: true, Value: config.Settings{"request_timeout": json.Number("720")}}, }) if err != nil { t.Fatalf("UpdateGroupSettings() error = %v", err) } - if result.Name != "updated settings" || result.UpstreamURL != "https://settings-updated.example.com/v1" || + if result.Name != "updated settings" || result.ChannelID != channel.OpenAICompatible || + string(result.Params) != `{"base_url":"https://settings-updated.example.com/v1"}` || result.ValidationModel == nil || *result.ValidationModel != "gpt-4.1" || result.Enabled || result.WeightManual == nil || *result.WeightManual != weight || result.Effective.RequestTimeout != 720 { t.Fatalf("UpdateGroupSettings() = %#v", result) @@ -105,20 +113,20 @@ func TestUpdateGroupSettingsPublishesOnceAndReturnsNewSettings(t *testing.T) { func TestUpdateGroupSettingsAllowsDuplicateUpstreamURLWithoutConfirmation(t *testing.T) { fixture := newServiceFixture(t) first, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://settings-first.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://settings-first.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-settings-first", + Credentials: "sk-settings-first", }) if err != nil { t.Fatalf("first CreateGroup() error = %v", err) } sharedURL := "https://settings-shared.example.com/v1" _, err = fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: sharedURL, - Protocols: []protocol.Protocol{protocol.Anthropic}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://settings-shared.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-settings-second", + Credentials: "sk-settings-second", }) if err != nil { t.Fatalf("second CreateGroup() error = %v", err) @@ -126,25 +134,25 @@ func TestUpdateGroupSettingsAllowsDuplicateUpstreamURLWithoutConfirmation(t *tes beforeRevision := fixture.manager.Current().Revision result, err := fixture.service.UpdateGroupSettings(t.Context(), first.GroupID, GroupSettingsUpdateRequest{ - UpstreamURL: optionalField[string]{ + Params: optionalField[json.RawMessage]{ Set: true, - Value: " HTTPS://SETTINGS-SHARED.EXAMPLE.COM/v1/ ", + Value: json.RawMessage(`{"base_url":" HTTPS://SETTINGS-SHARED.EXAMPLE.COM/v1/ "}`), }, }) if err != nil { t.Fatalf("UpdateGroupSettings() error = %v", err) } - if result.UpstreamURL != sharedURL { - t.Fatalf("updated upstream URL = %q, want %q", result.UpstreamURL, sharedURL) + if string(result.Params) != `{"base_url":"`+sharedURL+`"}` { + t.Fatalf("updated params = %s, want shared target", result.Params) } if got := fixture.manager.Current().Revision; got != beforeRevision+1 { t.Fatalf("snapshot revision = %d, want %d", got, beforeRevision+1) } } -func TestUpdateGroupSettingsValidatesWeightAndInjectUsage(t *testing.T) { +func TestUpdateGroupSettingsValidatesWeightAndAllowsUsageObservationAcrossChannels(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-settings-validation") + groupID := createGroupForCredentialImport(t, fixture, "sk-settings-validation") beforeRevision := fixture.manager.Current().Revision for _, weight := range []optionalField[int]{ @@ -176,75 +184,28 @@ func TestUpdateGroupSettingsValidatesWeightAndInjectUsage(t *testing.T) { } }) } - anthropic := validControlGroup("settings-anthropic") - anthropic.Protocols = models.JSON(`["anthropic"]`) - if err := fixture.db.Create(anthropic).Error; err != nil { - t.Fatal(err) - } - if _, err := fixture.manager.Publish(mustBuildCompileInput(t, fixture.db)); err != nil { - t.Fatal(err) - } - if _, err := fixture.service.UpdateGroupSettings(t.Context(), anthropic.ID, GroupSettingsUpdateRequest{ - Overrides: optionalField[config.Settings]{Set: true, Value: config.Settings{state.SettingInjectUsageOptions: true}}, - }); !errors.Is(err, app_errors.ErrValidation) { - t.Fatalf("inject_usage_options without OpenAI completions error = %v", err) - } - if got := fixture.manager.Current().Revision; got != beforeRevision+4 { - t.Fatalf("invalid settings mutation revision = %d, want %d", got, beforeRevision+4) - } -} - -func TestUpdateGroupSettingsRejectsProtocolChangeThatLeavesInjectUsageOptionsOnNonOpenAIGroup(t *testing.T) { - fixture := newServiceFixture(t) - created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://settings-protocol-constraint.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-settings-protocol-constraint", + anthropic, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("settings-anthropic"), ChannelID: channel.Anthropic, + Params: json.RawMessage(`{}`), Models: optionalGroupModels{Set: true}, Credentials: "sk-anthropic", }) if err != nil { - t.Fatalf("CreateGroup() error = %v", err) - } - if _, err := fixture.service.UpdateGroupSettings(t.Context(), created.GroupID, GroupSettingsUpdateRequest{ - Overrides: optionalField[config.Settings]{ - Set: true, - Value: config.Settings{ - state.SettingInjectUsageOptions: true, - }, - }, - }); err != nil { - t.Fatalf("enable inject_usage_options override: %v", err) - } - var before models.Group - if err := fixture.db.First(&before, created.GroupID).Error; err != nil { t.Fatal(err) } - beforeRevision := fixture.manager.Current().Revision - - _, err = fixture.service.UpdateGroupSettings(t.Context(), created.GroupID, GroupSettingsUpdateRequest{ - Protocols: optionalField[[]protocol.Protocol]{ - Set: true, Value: []protocol.Protocol{protocol.Anthropic}, - }, + updated, err := fixture.service.UpdateGroupSettings(t.Context(), anthropic.GroupID, GroupSettingsUpdateRequest{ + Overrides: optionalField[config.Settings]{Set: true, Value: config.Settings{state.SettingInjectUsageOptions: true}}, }) - if !errors.Is(err, app_errors.ErrValidation) { - t.Fatalf("UpdateGroupSettings() error = %v, want validation", err) + if err != nil || !updated.Effective.InjectUsageOptions { + t.Fatalf("cross-channel usage observation update = %#v, %v", updated, err) } - var after models.Group - if err := fixture.db.First(&after, created.GroupID).Error; err != nil { - t.Fatal(err) - } - if !reflect.DeepEqual(after, before) { - t.Fatalf("protocol constraint changed persisted group: got=%#v want=%#v", after, before) - } - if got := fixture.manager.Current().Revision; got != beforeRevision { - t.Fatalf("protocol constraint published revision %d, want %d", got, beforeRevision) + if got := fixture.manager.Current().Revision; got != beforeRevision+5 { + t.Fatalf("settings mutation revision = %d, want %d", got, beforeRevision+5) } } func TestGroupSettingsHTTPRejectsStrictJSONAndUnauthorizedWithoutMutation(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-settings-http") + groupID := createGroupForCredentialImport(t, fixture, "sk-settings-http") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) path := "/api/groups/" + stringGroupID(groupID) + "/settings" @@ -272,7 +233,7 @@ func TestGroupSettingsHTTPRejectsStrictJSONAndUnauthorizedWithoutMutation(t *tes func TestGroupSettingsHTTPNotFoundAndDatabaseFailureDoNotMutate(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-settings-errors") + groupID := createGroupForCredentialImport(t, fixture, "sk-settings-errors") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) beforeRevision := fixture.manager.Current().Revision diff --git a/internal/control/group_update.go b/internal/control/group_update.go index 973fa8d6..7ef54ac9 100644 --- a/internal/control/group_update.go +++ b/internal/control/group_update.go @@ -2,14 +2,15 @@ package control import ( "context" + "encoding/json" "strings" "unicode" "gorm.io/gorm" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" "gpt-load/internal/storage/models" @@ -29,17 +30,13 @@ func normalizeValidationModel(raw string) (string, error) { } func mapGroupRowToState(group models.Group) (state.GroupConfig, error) { - var protocols []protocol.Protocol - if err := decodeGroupDiscoveryJSON(group.Protocols, &protocols); err != nil { - return state.GroupConfig{}, err - } var storedModels []GroupModel if err := decodeGroupDiscoveryJSON(group.Models, &storedModels); err != nil { return state.GroupConfig{}, err } settings := make(config.Settings) - if len(group.Config) > 0 { - if err := decodeGroupDiscoveryJSON(group.Config, &settings); err != nil { + if len(group.Overrides) > 0 { + if err := decodeGroupDiscoveryJSON(group.Overrides, &settings); err != nil { return state.GroupConfig{}, err } } @@ -51,16 +48,25 @@ func mapGroupRowToState(group models.Group) (state.GroupConfig, error) { if group.ValidationModel != nil { validationModel = *group.ValidationModel } - return state.GroupConfig{ - ID: group.ID, Name: group.Name, UpstreamURL: group.UpstreamURL, - ProviderID: cloneString(group.ProviderID), + result := state.GroupConfig{ + ID: group.ID, + Name: group.Name, + ChannelID: channel.ID(group.ChannelID), + Params: append(json.RawMessage(nil), group.Params...), ValidationModel: validationModel, - Protocols: protocols, Models: runtimeModels, Settings: settings, - WeightManual: cloneInt(group.WeightManual), Enabled: group.Enabled, - }, nil + Models: runtimeModels, + Settings: settings, + WeightManual: cloneInt(group.WeightManual), Enabled: group.Enabled, + } + return result, nil } -func validateGroupRowCandidate(ctx context.Context, tx *gorm.DB, group models.Group) error { +func validateGroupRowCandidate( + ctx context.Context, + tx *gorm.DB, + group models.Group, + registry *channel.Registry, +) error { candidate, err := mapGroupRowToState(group) if err != nil { return err @@ -70,8 +76,9 @@ func validateGroupRowCandidate(ctx context.Context, tx *gorm.DB, group models.Gr return err } _, err = state.Compile(state.CompileInput{ - SystemSettings: systemSettings, - Groups: []state.GroupConfig{candidate}, + SystemSettings: systemSettings, + ChannelRegistry: registry, + Groups: []state.GroupConfig{candidate}, }) return err } diff --git a/internal/control/group_update_test.go b/internal/control/group_update_test.go index 8d86efbf..64a55c30 100644 --- a/internal/control/group_update_test.go +++ b/internal/control/group_update_test.go @@ -1,49 +1,41 @@ package control import ( + "encoding/json" "path/filepath" "testing" "gpt-load/internal/catalog" - "gpt-load/internal/protocol" + "gpt-load/internal/channel" "gpt-load/internal/storage/models" ) -func TestMapGroupRowToStateCarriesDeepClonedProviderID(t *testing.T) { - providerID := "openai" +func TestMapGroupRowToStateCarriesDeepClonedParams(t *testing.T) { group := models.Group{ - ID: 1, - Name: "provider-group", - ProviderID: &providerID, - UpstreamURL: "https://api.openai.com/v1", - Protocols: models.JSON(`["openai-responses"]`), - Models: models.JSON(`[]`), - Config: models.JSON(`{}`), - Enabled: false, + ID: 1, Name: "channel-group", ChannelID: string(channel.OpenAICompatible), + Params: models.JSON(`{"base_url":"https://api.example.com/v1"}`), + Models: models.JSON(`[]`), Overrides: models.JSON(`{}`), Enabled: false, } got, err := mapGroupRowToState(group) if err != nil { t.Fatalf("mapGroupRowToState() error = %v", err) } - providerID = "mutated-source" - *group.ProviderID = "mutated-row" - if got.ProviderID == nil || *got.ProviderID != "openai" { - t.Fatalf("GroupConfig.ProviderID = %v, want independent openai", got.ProviderID) + group.Params[0] = '[' + if string(got.Params) != `{"base_url":"https://api.example.com/v1"}` { + t.Fatalf("GroupConfig.Params = %s, want independent canonical params", got.Params) } } func TestGroupCatalogSyncTriggerOnlyTracksProviderAndModelIDChanges(t *testing.T) { fixture := newServiceFixture(t) - initialProviderID := "openai" created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - ProviderID: &initialProviderID, - UpstreamURL: "https://catalog-trigger.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{ {ID: "model-a", Alias: "public-a", AliasEnabled: true}, }}, - Keys: "sk-catalog-trigger", + Credentials: "sk-catalog-trigger", }) if err != nil { t.Fatal(err) @@ -55,13 +47,11 @@ func TestGroupCatalogSyncTriggerOnlyTracksProviderAndModelIDChanges(t *testing.T catalog.Metadata{}, false, ) - createdProviderID := "anthropic" if _, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - ProviderID: &createdProviderID, - UpstreamURL: "https://catalog-trigger-created.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.Anthropic, + Params: json.RawMessage(`{}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{{ID: "created-model"}}}, - Keys: "sk-catalog-trigger-created", + Credentials: "sk-catalog-trigger-created", }); err != nil { t.Fatal(err) } @@ -77,7 +67,7 @@ func TestGroupCatalogSyncTriggerOnlyTracksProviderAndModelIDChanges(t *testing.T assertNoCatalogGroupWake(t, coordinator) if _, err := fixture.service.UpdateGroupSettings(t.Context(), created.GroupID, GroupSettingsUpdateRequest{ - UpstreamURL: optionalField[string]{Set: true, Value: "https://catalog-trigger-new-url.example.com/v1"}, + Params: optionalField[json.RawMessage]{Set: true, Value: json.RawMessage(`{}`)}, }); err != nil { t.Fatal(err) } @@ -97,34 +87,26 @@ func TestGroupCatalogSyncTriggerOnlyTracksProviderAndModelIDChanges(t *testing.T } assertCatalogGroupWake(t, coordinator) - providerID := "anthropic" - if _, err := fixture.service.UpdateGroupSettings(t.Context(), created.GroupID, GroupSettingsUpdateRequest{ - ProviderID: optionalField[string]{Set: true, Value: providerID}, - }); err != nil { - t.Fatal(err) - } - assertCatalogGroupWake(t, coordinator) - custom, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://catalog-trigger-custom.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://catalog-trigger-custom.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{{ID: "custom-a"}}}, - Keys: "sk-catalog-trigger-custom", + Credentials: "sk-catalog-trigger-custom", }) if err != nil { t.Fatal(err) } - assertNoCatalogGroupWake(t, coordinator) + assertCatalogGroupWake(t, coordinator) if _, err := fixture.service.UpdateGroupModels(t.Context(), custom.GroupID, GroupModelsUpdateRequest{ Models: optionalGroupModels{Set: true, Values: []GroupModel{{ID: "custom-b"}}}, }); err != nil { t.Fatal(err) } - assertNoCatalogGroupWake(t, coordinator) + assertCatalogGroupWake(t, coordinator) if err := fixture.service.DeleteGroup(t.Context(), custom.GroupID); err != nil { t.Fatal(err) } - assertNoCatalogGroupWake(t, coordinator) + assertCatalogGroupWake(t, coordinator) if err := fixture.service.DeleteGroup(t.Context(), created.GroupID); err != nil { t.Fatal(err) } diff --git a/internal/control/group_write.go b/internal/control/group_write.go index c7ee4484..7789a4cc 100644 --- a/internal/control/group_write.go +++ b/internal/control/group_write.go @@ -12,13 +12,13 @@ import ( "gorm.io/gorm" + "gpt-load/internal/channel" + "gpt-load/internal/platform/epochms" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" - "gpt-load/internal/state" "gpt-load/internal/storage/models" ) -const maxUpstreamKeyLines = 1000 +const maxCredentialLines = 1000 type GroupModel struct { ID string `json:"id"` @@ -50,13 +50,13 @@ func (model *GroupModel) UnmarshalJSON(data []byte) error { return nil } -type upstreamKeyCandidate struct { - plaintext string - hash string +type credentialCandidate struct { + canonical json.RawMessage + fingerprint string } -type normalizedUpstreamKeys struct { - candidates []upstreamKeyCandidate +type normalizedCredentials struct { + candidates []credentialCandidate duplicateLines int } @@ -180,26 +180,6 @@ func normalizeUpstreamBaseURL(raw string) (normalized, hostname string, err erro return parsed.String(), hostname, nil } -func normalizeGroupProtocols(values []protocol.Protocol) ([]protocol.Protocol, error) { - seen := make(map[protocol.Protocol]struct{}, len(values)) - for _, value := range values { - if !value.DataPlaneEnabled() { - return nil, app_errors.ErrValidation - } - seen[value] = struct{}{} - } - if len(seen) == 0 { - return nil, app_errors.ErrValidation - } - result := make([]protocol.Protocol, 0, len(seen)) - for _, value := range protocol.DataPlaneProtocols() { - if _, exists := seen[value]; exists { - result = append(result, value) - } - } - return result, nil -} - func normalizeGroupModels(values []GroupModel) ([]GroupModel, error) { result := make([]GroupModel, 0, len(values)) indexesByClientModel := make(map[string][]int, len(values)) @@ -265,102 +245,131 @@ func normalizeGroupName(value *string) (*string, error) { return &normalized, nil } -func (s *Service) normalizeUpstreamKeys(raw string) (normalizedUpstreamKeys, error) { - result := normalizedUpstreamKeys{ - candidates: make([]upstreamKeyCandidate, 0), +func (s *Service) normalizeCredentials( + channelID channel.ID, + raw string, +) (normalizedCredentials, error) { + if s == nil || s.channelRegistry == nil || s.encryption == nil { + return normalizedCredentials{}, app_errors.ErrInternalServer } + result := normalizedCredentials{candidates: make([]credentialCandidate, 0)} seen := make(map[string]struct{}) nonEmptyLines := 0 - for _, line := range strings.Split(raw, "\n") { + for _, line := range credentialInputEntries(channelID, raw) { plaintext := strings.TrimSpace(line) if plaintext == "" { continue } nonEmptyLines++ - if nonEmptyLines > maxUpstreamKeyLines { - return normalizedUpstreamKeys{}, app_errors.ErrValidation + if nonEmptyLines > maxCredentialLines { + return normalizedCredentials{}, app_errors.ErrValidation } - hash := s.encryption.Hash(plaintext) - if _, duplicate := seen[hash]; duplicate { + var encoded json.RawMessage + if strings.HasPrefix(plaintext, "{") { + encoded = json.RawMessage(plaintext) + } else { + marshaled, err := json.Marshal(map[string]string{"api_key": plaintext}) + if err != nil { + return normalizedCredentials{}, app_errors.ErrInternalServer + } + encoded = marshaled + } + credential, err := s.channelRegistry.ValidateCredential(channelID, encoded) + if err != nil { + return normalizedCredentials{}, app_errors.ErrValidation + } + canonical := credential.CanonicalJSON() + fingerprint := s.encryption.Hash(string(canonical)) + if _, duplicate := seen[fingerprint]; duplicate { result.duplicateLines++ continue } - seen[hash] = struct{}{} - result.candidates = append(result.candidates, upstreamKeyCandidate{ - plaintext: plaintext, - hash: hash, + seen[fingerprint] = struct{}{} + result.candidates = append(result.candidates, credentialCandidate{ + canonical: append(json.RawMessage(nil), canonical...), fingerprint: fingerprint, }) } if len(result.candidates) == 0 { - return normalizedUpstreamKeys{}, app_errors.ErrValidation + return normalizedCredentials{}, app_errors.ErrValidation } return result, nil } -func (s *Service) persistUpstreamKeys( - tx *gorm.DB, - groupID uint, - normalized normalizedUpstreamKeys, -) ([]state.KeyEntry, int, int, error) { - hashes := make([]string, 0, len(normalized.candidates)) - for _, candidate := range normalized.candidates { - hashes = append(hashes, candidate.hash) +func credentialInputEntries(channelID channel.ID, raw string) []string { + trimmed := strings.TrimSpace(raw) + if !strings.HasPrefix(trimmed, "{") || !json.Valid([]byte(trimmed)) { + return strings.Split(raw, "\n") } - var existingRows []models.UpstreamKey - if err := tx.Where("group_id = ? AND key_hash IN ?", groupID, hashes).Find(&existingRows).Error; err != nil { - return nil, 0, 0, app_errors.ParseDBError(err) - } - existingByHash := make(map[string]models.UpstreamKey, len(existingRows)) - for _, row := range existingRows { - existingByHash[row.KeyHash] = row + if channelID != channel.GoogleVertex { + return []string{trimmed} } - entries := make([]state.KeyEntry, 0, len(normalized.candidates)) + var object map[string]json.RawMessage + if json.Unmarshal([]byte(trimmed), &object) != nil { + return []string{trimmed} + } + if _, wrapped := object["service_account_json"]; wrapped { + return []string{trimmed} + } + var credentialType string + if json.Unmarshal(object["type"], &credentialType) != nil || credentialType != "service_account" { + return []string{trimmed} + } + encoded, err := json.Marshal(map[string]string{"service_account_json": trimmed}) + if err != nil { + return []string{trimmed} + } + return []string{string(encoded)} +} + +func (s *Service) persistCredentials( + tx *gorm.DB, + groupID uint, + normalized normalizedCredentials, +) (int, int, error) { + fingerprints := make([]string, 0, len(normalized.candidates)) + for _, candidate := range normalized.candidates { + fingerprints = append(fingerprints, candidate.fingerprint) + } + var existingRows []models.Credential + if err := tx.Where("group_id = ? AND fingerprint IN ?", groupID, fingerprints). + Find(&existingRows).Error; err != nil { + return 0, 0, app_errors.ParseDBError(err) + } + existingByFingerprint := make(map[string]struct{}, len(existingRows)) + for _, row := range existingRows { + existingByFingerprint[row.Fingerprint] = struct{}{} + } + + nowMS, err := epochms.FromTime(s.now()) + if err != nil { + return 0, 0, app_errors.ErrInternalServer + } + if nowMS < 1 { + nowMS = 1 + } added := 0 duplicated := normalized.duplicateLines for _, candidate := range normalized.candidates { - row, exists := existingByHash[candidate.hash] - if exists { + if _, exists := existingByFingerprint[candidate.fingerprint]; exists { duplicated++ - } else { - ciphertext, err := s.encryption.Encrypt(candidate.plaintext) - if err != nil { - return nil, 0, 0, fmt.Errorf("encrypt upstream key: %w", err) - } - row = models.UpstreamKey{ - GroupID: groupID, - KeyValue: ciphertext, - KeyHash: candidate.hash, - Status: models.UpstreamKeyStatusActive, - } - if err := tx.Create(&row).Error; err != nil { - return nil, 0, 0, app_errors.ParseDBError(err) - } - added++ + continue } - entries = append(entries, state.KeyEntry{ - ID: row.ID, - GroupID: row.GroupID, - WeightManual: row.WeightManual, - WeightAuto: state.DefaultWeight, - Status: state.KeyStatus(row.Status), - EncryptedValue: row.KeyValue, - }) - } - return entries, added, duplicated, nil -} - -func (s *Service) applyMissingRegistryKeys(groupID uint, entries []state.KeyEntry) error { - missing := make([]state.KeyEntry, 0, len(entries)) - for _, entry := range entries { - if _, exists := s.registry.EncryptedValue(entry.ID); !exists { - missing = append(missing, entry) + ciphertext, err := s.encryption.Encrypt(string(candidate.canonical)) + if err != nil { + return 0, 0, fmt.Errorf("encrypt credential: %w", err) } + row := models.Credential{ + GroupID: groupID, Data: ciphertext, Fingerprint: candidate.fingerprint, + Status: models.CredentialStatusActive, CreatedAtMS: nowMS, UpdatedAtMS: nowMS, + } + if err := tx.Create(&row).Error; err != nil { + return 0, 0, app_errors.ParseDBError(err) + } + existingByFingerprint[candidate.fingerprint] = struct{}{} + added++ } - if len(missing) == 0 { - return nil - } - return s.registry.ApplyImport(groupID, missing) + return added, duplicated, nil } func isLiteralPrivateHost(hostname string) bool { diff --git a/internal/control/health.go b/internal/control/health.go index 4cf83eb0..d170d91e 100644 --- a/internal/control/health.go +++ b/internal/control/health.go @@ -7,10 +7,10 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/health" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/platform/response" - "gpt-load/internal/platform/utils" "gpt-load/internal/platform/version" "gpt-load/internal/requestlog" "gpt-load/internal/state" @@ -21,11 +21,10 @@ type RequestLogStatsReader interface { } type healthCountsResponse struct { - Total int `json:"total"` + Credentials int `json:"credentials"` Available int `json:"available"` Cooldown int `json:"cooldown"` Blacklisted int `json:"blacklisted"` - Disabled int `json:"disabled"` } type healthGroupResponse struct { @@ -41,8 +40,8 @@ type healthRecoveryResponse struct { AtMS *int64 `json:"at_ms"` } -type healthProblemKeyResponse struct { - KeyID uint `json:"key_id"` +type healthProblemCredentialResponse struct { + CredentialID uint `json:"credential_id"` GroupID uint `json:"group_id"` GroupName string `json:"group_name"` Mask string `json:"mask"` @@ -76,16 +75,16 @@ type requestLogHealthResponse struct { } type runtimeHealthResponse struct { - ObservedAtMS int64 `json:"observed_at_ms"` - Version string `json:"version"` - UptimeSeconds int64 `json:"uptime_seconds"` - SnapshotRevision uint64 `json:"snapshot_revision"` - StatsWindowSeconds int64 `json:"stats_window_seconds"` - Counts healthCountsResponse `json:"counts"` - Groups []healthGroupResponse `json:"groups"` - CooldownKeys []healthProblemKeyResponse `json:"cooldown_keys"` - BlacklistedKeys []healthProblemKeyResponse `json:"blacklisted_keys"` - RequestLog requestLogHealthResponse `json:"request_log"` + ObservedAtMS int64 `json:"observed_at_ms"` + Version string `json:"version"` + UptimeSeconds int64 `json:"uptime_seconds"` + SnapshotRevision uint64 `json:"snapshot_revision"` + StatsWindowSeconds int64 `json:"stats_window_seconds"` + Counts healthCountsResponse `json:"counts"` + Groups []healthGroupResponse `json:"groups"` + CooldownCredentials []healthProblemCredentialResponse `json:"cooldown_credentials"` + BlacklistedCredentials []healthProblemCredentialResponse `json:"blacklisted_credentials"` + RequestLog requestLogHealthResponse `json:"request_log"` } type healthBucket string @@ -99,19 +98,19 @@ const ( func classifyHealthKey( group state.GroupCatalogView, - key state.KeyRuntimeView, + key state.CredentialRuntimeView, now time.Time, ) healthBucket { if !group.Enabled || (group.WeightManual != nil && *group.WeightManual == 0) || - key.Status != state.KeyStatusActive || + key.Status != state.CredentialStatusActive || (key.WeightManual != nil && *key.WeightManual == 0) { return healthBucketDisabled } switch key.RuntimeState(now) { - case state.KeyRuntimeBlacklisted: + case state.CredentialRuntimeBlacklisted: return healthBucketBlacklisted - case state.KeyRuntimeCooldown: + case state.CredentialRuntimeCooldown: return healthBucketCooldown default: return healthBucketAvailable @@ -119,16 +118,16 @@ func classifyHealthKey( } func addHealthCount(counts *healthCountsResponse, bucket healthBucket) { - counts.Total++ switch bucket { case healthBucketAvailable: + counts.Credentials++ counts.Available++ case healthBucketCooldown: + counts.Credentials++ counts.Cooldown++ case healthBucketBlacklisted: + counts.Credentials++ counts.Blacklisted++ - case healthBucketDisabled: - counts.Disabled++ } } @@ -150,19 +149,20 @@ func optionalHealthStatusCode(value int) *int { func (service *Service) healthProblemMask( ciphertexts map[uint]string, - keyID uint, + credentialID uint, + channelID channel.ID, ) (string, error) { - if service == nil || service.encryption == nil || keyID == 0 { + if service == nil || service.encryption == nil || credentialID == 0 { return "", fmt.Errorf( "map runtime health problem key: %w", app_errors.ErrInternalServer, ) } - ciphertext, exists := ciphertexts[keyID] + ciphertext, exists := ciphertexts[credentialID] if !exists || ciphertext == "" { return "", fmt.Errorf( "map runtime health problem key %d: ciphertext unavailable: %w", - keyID, + credentialID, app_errors.ErrInternalServer, ) } @@ -170,20 +170,21 @@ func (service *Service) healthProblemMask( if err != nil { return "", fmt.Errorf( "map runtime health problem key %d: decrypt credential: %v: %w", - keyID, + credentialID, err, app_errors.ErrInternalServer, ) } - mask := utils.MaskAPIKey(plaintext) - if mask == "" { + validated, err := normalizeStoredCredential(service.channelRegistry, channelID, plaintext) + if err != nil { return "", fmt.Errorf( - "map runtime health problem key %d: empty credential: %w", - keyID, + "map runtime health problem credential %d: validate credential: %v: %w", + credentialID, + err, app_errors.ErrInternalServer, ) } - return mask, nil + return maskCanonicalCredential(validated.CanonicalJSON()) } func (service *Service) RuntimeHealth() (runtimeHealthResponse, error) { @@ -202,12 +203,12 @@ func (service *Service) RuntimeHealth() (runtimeHealthResponse, error) { return runtimeHealthResponse{}, fmt.Errorf("map runtime health observed_at_ms: %w", err) } result := runtimeHealthResponse{ - ObservedAtMS: observedAtMS, - SnapshotRevision: observation.snapshot.Revision, - StatsWindowSeconds: int64(health.StatsWindow / time.Second), - Groups: []healthGroupResponse{}, - CooldownKeys: []healthProblemKeyResponse{}, - BlacklistedKeys: []healthProblemKeyResponse{}, + ObservedAtMS: observedAtMS, + SnapshotRevision: observation.snapshot.Revision, + StatsWindowSeconds: int64(health.StatsWindow / time.Second), + Groups: []healthGroupResponse{}, + CooldownCredentials: []healthProblemCredentialResponse{}, + BlacklistedCredentials: []healthProblemCredentialResponse{}, } groupIDs := make([]uint, 0, len(observation.snapshot.GroupCatalog)) for groupID := range observation.snapshot.GroupCatalog { @@ -231,7 +232,12 @@ func (service *Service) RuntimeHealth() (runtimeHealthResponse, error) { if bucket != healthBucketCooldown && bucket != healthBucketBlacklisted { continue } - mask, err := service.healthProblemMask(observation.problemCiphertexts, key.ID) + groupView := observation.snapshot.Groups[key.GroupID] + mask, err := service.healthProblemMask( + observation.problemCiphertexts, + key.ID, + groupView.ChannelID, + ) if err != nil { return runtimeHealthResponse{}, err } @@ -242,8 +248,8 @@ func (service *Service) RuntimeHealth() (runtimeHealthResponse, error) { lastFailureCategory = health.FailureCategoryAmbiguous lastStatusCode = 0 } - detail := healthProblemKeyResponse{ - KeyID: key.ID, + detail := healthProblemCredentialResponse{ + CredentialID: key.ID, GroupID: key.GroupID, GroupName: group.Name, Mask: mask, @@ -270,18 +276,17 @@ func (service *Service) RuntimeHealth() (runtimeHealthResponse, error) { Mode: "cooldown_expiry", AtMS: cooldownUntilMS, } - result.CooldownKeys = append(result.CooldownKeys, detail) + result.CooldownCredentials = append(result.CooldownCredentials, detail) } else { detail.Recovery = healthRecoveryResponse{Mode: "configuration_required"} - if validationGroup, exists := observation.snapshot.Groups[key.GroupID]; exists { - if _, valid := buildGroupValidationTarget(validationGroup); valid { - detail.Recovery = healthRecoveryResponse{ - Automatic: true, - Mode: "validation_probe", + if service.executor != nil && service.channelRegistry != nil { + if validationGroup, exists := observation.snapshot.Groups[key.GroupID]; exists { + if _, valid := buildGroupValidationTarget(validationGroup); valid { + detail.Recovery = healthRecoveryResponse{Automatic: true, Mode: "validation_probe"} } } } - result.BlacklistedKeys = append(result.BlacklistedKeys, detail) + result.BlacklistedCredentials = append(result.BlacklistedCredentials, detail) } } requestLog, err := mapRequestLogHealth(service.requestLogStats.Stats()) diff --git a/internal/control/health_test.go b/internal/control/health_test.go index 3d53e8fe..7dd93968 100644 --- a/internal/control/health_test.go +++ b/internal/control/health_test.go @@ -12,10 +12,10 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/health" "gpt-load/internal/platform/config" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/requestlog" "gpt-load/internal/state" ) @@ -41,22 +41,23 @@ func TestRuntimeHealthReturnsMutuallyExclusiveCurrentState(t *testing.T) { blacklistedPlaintext := "invalid-key-secret-lock" zero := 0 if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{ { - ID: 1, Name: "active", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 1, Name: "active", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }, { - ID: 2, Name: "disabled", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 2, Name: "disabled", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: false, }, { - ID: 3, Name: "zero-weight", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 3, Name: "zero-weight", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, WeightManual: &zero, Enabled: true, }, { - ID: 4, Name: "empty", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 4, Name: "empty", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }, }, @@ -64,30 +65,30 @@ func TestRuntimeHealthReturnsMutuallyExclusiveCurrentState(t *testing.T) { t.Fatalf("Publish() error = %v", err) } keyWeightZero := 0 - if err := fixture.registry.Replace([]state.KeyEntry{ - {ID: 11, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "available"}, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{ + {ID: 11, GroupID: 1, Version: 1, IdentityGeneration: 11, Fingerprint: "test-11", Status: state.CredentialStatusActive, EncryptedValue: "available"}, { - ID: 12, GroupID: 1, Status: state.KeyStatusActive, + ID: 12, GroupID: 1, Version: 1, IdentityGeneration: 12, Fingerprint: "test-12", Status: state.CredentialStatusActive, CooldownUntil: now.Add(time.Minute), FailureCount: 1, EncryptedValue: encryptHealthKey(t, fixture, cooldownPlaintext), }, { - ID: 13, GroupID: 1, Status: state.KeyStatusActive, + ID: 13, GroupID: 1, Version: 1, IdentityGeneration: 13, Fingerprint: "test-13", Status: state.CredentialStatusActive, Blacklisted: true, CooldownUntil: now.Add(time.Hour), FailureCount: 3, EncryptedValue: encryptHealthKey(t, fixture, blacklistedPlaintext), }, { - ID: 14, GroupID: 1, Status: state.KeyStatusDisabled, + ID: 14, GroupID: 1, Version: 1, IdentityGeneration: 14, Fingerprint: "test-14", Status: state.CredentialStatusDisabled, Blacklisted: true, EncryptedValue: "disabled", }, { - ID: 15, GroupID: 1, Status: state.KeyStatusActive, + ID: 15, GroupID: 1, Version: 1, IdentityGeneration: 15, Fingerprint: "test-15", Status: state.CredentialStatusActive, WeightManual: &keyWeightZero, Blacklisted: true, EncryptedValue: "weight-zero", }, - {ID: 21, GroupID: 2, Status: state.KeyStatusActive, EncryptedValue: "disabled-group"}, - {ID: 31, GroupID: 3, Status: state.KeyStatusActive, EncryptedValue: "zero-group"}, + {ID: 21, GroupID: 2, Version: 1, IdentityGeneration: 21, Fingerprint: "test-21", Status: state.CredentialStatusActive, EncryptedValue: "disabled-group"}, + {ID: 31, GroupID: 3, Version: 1, IdentityGeneration: 31, Fingerprint: "test-31", Status: state.CredentialStatusActive, EncryptedValue: "zero-group"}, }); err != nil { t.Fatalf("Replace() error = %v", err) } @@ -111,7 +112,7 @@ func TestRuntimeHealthReturnsMutuallyExclusiveCurrentState(t *testing.T) { t.Fatalf("observation metadata = %#v", got) } wantCounts := healthCountsResponse{ - Total: 7, Available: 1, Cooldown: 1, Blacklisted: 1, Disabled: 4, + Credentials: 3, Available: 1, Cooldown: 1, Blacklisted: 1, } if got.Counts != wantCounts { t.Fatalf("global counts = %#v, want %#v", got.Counts, wantCounts) @@ -121,34 +122,34 @@ func TestRuntimeHealthReturnsMutuallyExclusiveCurrentState(t *testing.T) { t.Fatalf("group order = %#v", got.Groups) } if got.Groups[0].Counts != (healthCountsResponse{ - Total: 5, Available: 1, Cooldown: 1, Blacklisted: 1, Disabled: 2, + Credentials: 3, Available: 1, Cooldown: 1, Blacklisted: 1, }) { t.Fatalf("active group counts = %#v", got.Groups[0].Counts) } - if got.Groups[1].Counts.Disabled != 1 || got.Groups[2].Counts.Disabled != 1 || - got.Groups[3].Counts.Total != 0 { + if got.Groups[1].Counts.Credentials != 0 || got.Groups[2].Counts.Credentials != 0 || + got.Groups[3].Counts.Credentials != 0 { t.Fatalf("disabled/zero/empty group counts = %#v", got.Groups) } - if len(got.CooldownKeys) != 1 || got.CooldownKeys[0].KeyID != 12 || - got.CooldownKeys[0].Mask != "rate****safe" || - got.CooldownKeys[0].LastFailureCategory != "rate_limited" || - got.CooldownKeys[0].LastStatusCode == nil || - *got.CooldownKeys[0].LastStatusCode != 429 || - got.CooldownKeys[0].RecentSuccessCount != 1 || - got.CooldownKeys[0].RecentProblemCount != 1 || - got.CooldownKeys[0].ConsecutiveProblemCount != 1 || - got.CooldownKeys[0].Recovery.Mode != "cooldown_expiry" { - t.Fatalf("cooldown details = %#v", got.CooldownKeys) + if len(got.CooldownCredentials) != 1 || got.CooldownCredentials[0].CredentialID != 12 || + got.CooldownCredentials[0].Mask != "rate****safe" || + got.CooldownCredentials[0].LastFailureCategory != "rate_limited" || + got.CooldownCredentials[0].LastStatusCode == nil || + *got.CooldownCredentials[0].LastStatusCode != 429 || + got.CooldownCredentials[0].RecentSuccessCount != 1 || + got.CooldownCredentials[0].RecentProblemCount != 1 || + got.CooldownCredentials[0].ConsecutiveProblemCount != 1 || + got.CooldownCredentials[0].Recovery.Mode != "cooldown_expiry" { + t.Fatalf("cooldown details = %#v", got.CooldownCredentials) } - if len(got.BlacklistedKeys) != 1 || got.BlacklistedKeys[0].KeyID != 13 || - got.BlacklistedKeys[0].Mask != "inva****lock" || - got.BlacklistedKeys[0].LastFailureCategory != "invalid_key" || - got.BlacklistedKeys[0].LastStatusCode == nil || - *got.BlacklistedKeys[0].LastStatusCode != 401 || - got.BlacklistedKeys[0].ConsecutiveProblemCount != 1 || - got.BlacklistedKeys[0].Recovery.Mode != "validation_probe" || - got.BlacklistedKeys[0].Recovery.AtMS != nil { - t.Fatalf("blacklisted details = %#v", got.BlacklistedKeys) + if len(got.BlacklistedCredentials) != 1 || got.BlacklistedCredentials[0].CredentialID != 13 || + got.BlacklistedCredentials[0].Mask != "inva****lock" || + got.BlacklistedCredentials[0].LastFailureCategory != "invalid_key" || + got.BlacklistedCredentials[0].LastStatusCode == nil || + *got.BlacklistedCredentials[0].LastStatusCode != 401 || + got.BlacklistedCredentials[0].ConsecutiveProblemCount != 1 || + got.BlacklistedCredentials[0].Recovery.Mode != "validation_probe" || + got.BlacklistedCredentials[0].Recovery.AtMS != nil { + t.Fatalf("blacklisted details = %#v", got.BlacklistedCredentials) } if got.RequestLog.DroppedTotal != 2 || got.RequestLog.LastWriteFailureAtMS == nil || @@ -158,19 +159,22 @@ func TestRuntimeHealthReturnsMutuallyExclusiveCurrentState(t *testing.T) { } } -func TestRuntimeHealthRequiresValidationConfigurationForBlacklistedKeyWithoutTarget(t *testing.T) { +func TestRuntimeHealthAdvertisesExecutorValidationForChannelCredential(t *testing.T) { fixture := newServiceFixture(t) now := healthNow() fixture.service.now = func() time.Time { return now } - if _, err := fixture.manager.Publish(state.CompileInput{Groups: []state.GroupConfig{{ - ID: 1, Name: "resources-only", Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Enabled: true, - }}}); err != nil { + if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, + Groups: []state.GroupConfig{{ + ID: 1, Name: "channel", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + }}, + }); err != nil { t.Fatalf("Publish() error = %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{{ - ID: 11, GroupID: 1, Status: state.KeyStatusActive, Blacklisted: true, - EncryptedValue: encryptHealthKey(t, fixture, "blacklisted-without-target"), + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 11, GroupID: 1, Version: 1, IdentityGeneration: 11, Fingerprint: "test-11", Status: state.CredentialStatusActive, Blacklisted: true, + EncryptedValue: encryptHealthKey(t, fixture, `{"api_key":"blacklisted-channel-credential"}`), }}); err != nil { t.Fatalf("Replace() error = %v", err) } @@ -180,12 +184,12 @@ func TestRuntimeHealthRequiresValidationConfigurationForBlacklistedKeyWithoutTar if err != nil { t.Fatalf("RuntimeHealth() error = %v", err) } - if len(got.BlacklistedKeys) != 1 { - t.Fatalf("blacklisted keys = %#v", got.BlacklistedKeys) + if len(got.BlacklistedCredentials) != 1 { + t.Fatalf("blacklisted keys = %#v", got.BlacklistedCredentials) } - recovery := got.BlacklistedKeys[0].Recovery - if recovery.Automatic || recovery.Mode != "configuration_required" || recovery.AtMS != nil { - t.Fatalf("recovery = %#v, want manual configuration requirement", recovery) + recovery := got.BlacklistedCredentials[0].Recovery + if !recovery.Automatic || recovery.Mode != "validation_probe" || recovery.AtMS != nil { + t.Fatalf("recovery = %#v, want automatic validation probe", recovery) } } @@ -193,14 +197,14 @@ func TestRuntimeHealthExposesProblemCountsInsteadOfFailureAliases(t *testing.T) fixture := newServiceFixture(t) now := healthNow() fixture.service.now = func() time.Time { return now } - if _, err := fixture.manager.Publish(state.CompileInput{Groups: []state.GroupConfig{{ - ID: 1, Name: "active", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + if _, err := fixture.manager.Publish(state.CompileInput{ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ + ID: 1, Name: "active", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }}}); err != nil { t.Fatalf("Publish() error = %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{{ - ID: 11, GroupID: 1, Status: state.KeyStatusActive, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 11, GroupID: 1, Version: 1, IdentityGeneration: 11, Fingerprint: "test-11", Status: state.CredentialStatusActive, CooldownUntil: now.Add(time.Minute), EncryptedValue: encryptHealthKey(t, fixture, "rate-limit-secret-safe"), }}); err != nil { @@ -219,15 +223,15 @@ func TestRuntimeHealthExposesProblemCountsInsteadOfFailureAliases(t *testing.T) t.Fatalf("Marshal() error = %v", err) } var document struct { - CooldownKeys []map[string]json.RawMessage `json:"cooldown_keys"` + CooldownCredentials []map[string]json.RawMessage `json:"cooldown_credentials"` } if err := json.Unmarshal(body, &document); err != nil { t.Fatalf("Unmarshal() error = %v", err) } - if len(document.CooldownKeys) != 1 { - t.Fatalf("cooldown key count = %d, want 1", len(document.CooldownKeys)) + if len(document.CooldownCredentials) != 1 { + t.Fatalf("cooldown key count = %d, want 1", len(document.CooldownCredentials)) } - key := document.CooldownKeys[0] + key := document.CooldownCredentials[0] for field, want := range map[string]uint64{ "recent_success_count": 1, "recent_problem_count": 2, @@ -253,42 +257,43 @@ func TestRuntimeHealthSortsProblemKeysByGroupAndKey(t *testing.T) { now := healthNow() fixture.service.now = func() time.Time { return now } if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{ { - ID: 2, Name: "two", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 2, Name: "two", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }, { - ID: 1, Name: "one", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 1, Name: "one", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }, }, }); err != nil { t.Fatalf("Publish() error = %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{ + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{ { - ID: 22, GroupID: 2, Status: state.KeyStatusActive, + ID: 22, GroupID: 2, Version: 1, IdentityGeneration: 22, Fingerprint: "test-22", Status: state.CredentialStatusActive, CooldownUntil: now.Add(time.Minute), EncryptedValue: encryptHealthKey(t, fixture, "cooldown-secret-0022"), }, { - ID: 13, GroupID: 1, Status: state.KeyStatusActive, + ID: 13, GroupID: 1, Version: 1, IdentityGeneration: 13, Fingerprint: "test-13", Status: state.CredentialStatusActive, Blacklisted: true, EncryptedValue: encryptHealthKey(t, fixture, "blacklisted-secret-0013"), }, { - ID: 12, GroupID: 1, Status: state.KeyStatusActive, + ID: 12, GroupID: 1, Version: 1, IdentityGeneration: 12, Fingerprint: "test-12", Status: state.CredentialStatusActive, CooldownUntil: now.Add(time.Minute), EncryptedValue: encryptHealthKey(t, fixture, "cooldown-secret-0012"), }, { - ID: 21, GroupID: 2, Status: state.KeyStatusActive, + ID: 21, GroupID: 2, Version: 1, IdentityGeneration: 21, Fingerprint: "test-21", Status: state.CredentialStatusActive, Blacklisted: true, EncryptedValue: encryptHealthKey(t, fixture, "blacklisted-secret-0021"), }, { - ID: 11, GroupID: 1, Status: state.KeyStatusActive, + ID: 11, GroupID: 1, Version: 1, IdentityGeneration: 11, Fingerprint: "test-11", Status: state.CredentialStatusActive, CooldownUntil: now.Add(time.Minute), EncryptedValue: encryptHealthKey(t, fixture, "cooldown-secret-0011"), }, @@ -299,19 +304,19 @@ func TestRuntimeHealthSortsProblemKeysByGroupAndKey(t *testing.T) { if err != nil { t.Fatalf("RuntimeHealth() error = %v", err) } - pairs := func(items []healthProblemKeyResponse) [][2]uint { + pairs := func(items []healthProblemCredentialResponse) [][2]uint { result := make([][2]uint, 0, len(items)) for _, item := range items { - result = append(result, [2]uint{item.GroupID, item.KeyID}) + result = append(result, [2]uint{item.GroupID, item.CredentialID}) } return result } - if gotPairs, want := pairs(got.CooldownKeys), [][2]uint{ + if gotPairs, want := pairs(got.CooldownCredentials), [][2]uint{ {1, 11}, {1, 12}, {2, 22}, }; !reflect.DeepEqual(gotPairs, want) { t.Fatalf("cooldown order = %v, want %v", gotPairs, want) } - if gotPairs, want := pairs(got.BlacklistedKeys), [][2]uint{ + if gotPairs, want := pairs(got.BlacklistedCredentials), [][2]uint{ {1, 13}, {2, 21}, }; !reflect.DeepEqual(gotPairs, want) { t.Fatalf("blacklisted order = %v, want %v", gotPairs, want) @@ -323,14 +328,14 @@ func TestRuntimeHealthJSONOmitsScoresCredentialsAndZeroTimes(t *testing.T) { fixture.service.now = healthNow plaintext := "provider-secret-credential-tail" ciphertext := encryptHealthKey(t, fixture, plaintext) - if _, err := fixture.manager.Publish(state.CompileInput{Groups: []state.GroupConfig{{ - ID: 1, Name: "safe", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + if _, err := fixture.manager.Publish(state.CompileInput{ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ + ID: 1, Name: "safe", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }}}); err != nil { t.Fatalf("Publish() error = %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, Blacklisted: true, EncryptedValue: ciphertext, }}); err != nil { t.Fatalf("Replace() error = %v", err) @@ -343,11 +348,11 @@ func TestRuntimeHealthJSONOmitsScoresCredentialsAndZeroTimes(t *testing.T) { if err != nil { t.Fatalf("RuntimeHealth() error = %v", err) } - if len(result.BlacklistedKeys) != 1 || - result.BlacklistedKeys[0].Mask != "prov****tail" || - result.BlacklistedKeys[0].LastFailureCategory != "ambiguous" || - result.BlacklistedKeys[0].LastStatusCode != nil { - t.Fatalf("blacklisted details = %#v", result.BlacklistedKeys) + if len(result.BlacklistedCredentials) != 1 || + result.BlacklistedCredentials[0].Mask != "prov****tail" || + result.BlacklistedCredentials[0].LastFailureCategory != "ambiguous" || + result.BlacklistedCredentials[0].LastStatusCode != nil { + t.Fatalf("blacklisted details = %#v", result.BlacklistedCredentials) } body, err := json.Marshal(result) if err != nil { @@ -404,14 +409,14 @@ func TestRuntimeHealthJSONOmitsScoresCredentialsAndZeroTimes(t *testing.T) { func TestRuntimeHealthFailsClosedWhenProblemKeyCannotBeDecrypted(t *testing.T) { fixture := newServiceFixture(t) fixture.service.now = healthNow - if _, err := fixture.manager.Publish(state.CompileInput{Groups: []state.GroupConfig{{ - ID: 1, Name: "safe", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + if _, err := fixture.manager.Publish(state.CompileInput{ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ + ID: 1, Name: "safe", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }}}); err != nil { t.Fatalf("Publish() error = %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, Blacklisted: true, EncryptedValue: "not-a-valid-ciphertext", }}); err != nil { t.Fatalf("Replace() error = %v", err) @@ -427,7 +432,7 @@ func TestRuntimeHealthFailsClosedWhenProblemKeyCannotBeDecrypted(t *testing.T) { func TestHealthProblemMaskFailsClosedWhenCiphertextIsMissing(t *testing.T) { fixture := newServiceFixture(t) - if _, err := fixture.service.healthProblemMask(nil, 1); !errors.Is( + if _, err := fixture.service.healthProblemMask(nil, 1, ""); !errors.Is( err, app_errors.ErrInternalServer, ) { @@ -435,6 +440,26 @@ func TestHealthProblemMaskFailsClosedWhenCiphertextIsMissing(t *testing.T) { } } +func TestHealthProblemMaskExtractsTypedCredentialSecret(t *testing.T) { + fixture := newServiceFixture(t) + ciphertext := encryptHealthKey( + t, + fixture, + `{"api_key":"provider-secret-credential-tail"}`, + ) + mask, err := fixture.service.healthProblemMask( + map[uint]string{1: ciphertext}, + 1, + channel.OpenAI, + ) + if err != nil { + t.Fatalf("healthProblemMask() error = %v", err) + } + if mask != "prov****tail" || strings.Contains(mask, "api_key") || strings.Contains(mask, "{") { + t.Fatalf("healthProblemMask() = %q, want api_key-only mask", mask) + } +} + func TestRuntimeHealthDTOHasNoCredentialOrScoreFields(t *testing.T) { forbidden := map[string]struct{}{ "EncryptedValue": {}, "KeyHash": {}, "AccessKey": {}, @@ -444,7 +469,7 @@ func TestRuntimeHealthDTOHasNoCredentialOrScoreFields(t *testing.T) { reflect.TypeOf(runtimeHealthResponse{}), reflect.TypeOf(healthCountsResponse{}), reflect.TypeOf(healthGroupResponse{}), - reflect.TypeOf(healthProblemKeyResponse{}), + reflect.TypeOf(healthProblemCredentialResponse{}), reflect.TypeOf(healthRecoveryResponse{}), reflect.TypeOf(requestLogHealthResponse{}), } @@ -460,8 +485,8 @@ func TestRuntimeHealthDTOHasNoCredentialOrScoreFields(t *testing.T) { func TestRuntimeHealthFailsLoudForRegistryCatalogMismatch(t *testing.T) { fixture := newServiceFixture(t) - if err := fixture.registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 999, Status: state.KeyStatusActive, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 999, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, EncryptedValue: "cipher", }}); err != nil { t.Fatalf("Replace() error = %v", err) @@ -519,7 +544,7 @@ func TestRuntimeHealthEndpointRequiresManagementAuthentication(t *testing.T) { } body := success.Body.String() for _, emptyArray := range []string{ - `"groups":[]`, `"cooldown_keys":[]`, `"blacklisted_keys":[]`, + `"groups":[]`, `"cooldown_credentials":[]`, `"blacklisted_credentials":[]`, } { if !strings.Contains(body, emptyArray) { t.Fatalf("response must contain %s: %s", emptyArray, body) diff --git a/internal/control/home.go b/internal/control/home.go index fa4db328..5ec1e6a0 100644 --- a/internal/control/home.go +++ b/internal/control/home.go @@ -19,10 +19,10 @@ import ( ) type HomeInventory struct { - GroupCount int64 `json:"group_count"` - UpstreamKeyCount int64 `json:"upstream_key_count"` - AvailableUpstreamKeyCount int64 `json:"available_upstream_key_count"` - ModelCount int64 `json:"model_count"` + GroupCount int64 `json:"group_count"` + CredentialCount int64 `json:"credential_count"` + AvailableCredentialCount int64 `json:"available_credential_count"` + ModelCount int64 `json:"model_count"` } type HomeAccessKey struct { @@ -45,10 +45,12 @@ type homeResponse struct { HomeBase } -type homeUpstreamKeyRow struct { - ID uint - GroupID uint - Status models.UpstreamKeyStatus +type homeCredentialRow struct { + ID uint + GroupID uint + Fingerprint string + Status models.CredentialStatus + UpdatedAtMS int64 } type homeAccessKeyRow struct { @@ -64,9 +66,9 @@ type homeAccessKeyRow struct { } type homeReadRows struct { - groupCount int64 - upstreamKeys []homeUpstreamKeyRow - accessKeys []homeAccessKeyRow + groupCount int64 + credentials []homeCredentialRow + accessKeys []homeAccessKeyRow } func (s *Service) ReadHomeBase( @@ -143,16 +145,16 @@ func (s *Service) readHomeBase( inventory := HomeInventory{} if accessKeyID == nil { inventory.GroupCount = rows.groupCount - inventory.UpstreamKeyCount = int64(len(rows.upstreamKeys)) + inventory.CredentialCount = int64(len(rows.credentials)) inventory.ModelCount = countHomeModels(snapshot) } else { inventory.GroupCount = int64(len(allowedGroups)) - inventory.UpstreamKeyCount = countHomeUpstreamKeys(rows.upstreamKeys, allowedGroups) + inventory.CredentialCount = countHomeCredentials(rows.credentials, allowedGroups) inventory.ModelCount = countScopedHomeModels(snapshot, allowedGroups, snapshot.AccessKeysByID[*accessKeyID]) } - inventory.AvailableUpstreamKeyCount, err = countAvailableHomeKeysInGroups( + inventory.AvailableCredentialCount, err = countAvailableHomeCredentialsInGroups( snapshot, - rows.upstreamKeys, + rows.credentials, runtimeKeys, now, allowedGroups, @@ -233,11 +235,11 @@ func (s *Service) readHomeRows( if err := tx.Model(&models.Group{}).Count(&result.groupCount).Error; err != nil { return fmt.Errorf("count home groups: %w", err) } - if err := tx.Model(&models.UpstreamKey{}). - Select("id", "group_id", "status"). + if err := tx.Model(&models.Credential{}). + Select("id", "group_id", "fingerprint", "status", "updated_at_ms"). Order("id ASC"). - Find(&result.upstreamKeys).Error; err != nil { - return fmt.Errorf("query home upstream keys: %w", err) + Find(&result.credentials).Error; err != nil { + return fmt.Errorf("query home credentials: %w", err) } if err := tx.Model(&models.AccessKey{}). Select( @@ -286,7 +288,7 @@ func accessibleHomeGroups( } } protocolAllowed := len(accessKey.Filters.Protocols) == 0 - for _, value := range group.Protocols { + for _, value := range group.ClientProtocols { if _, allowed := accessKey.Filters.Protocols[value]; allowed { protocolAllowed = true break @@ -346,8 +348,8 @@ func homeExternalModelName(model state.ModelConfig) string { return strings.TrimSpace(model.ID) } -func countHomeUpstreamKeys( - rows []homeUpstreamKeyRow, +func countHomeCredentials( + rows []homeCredentialRow, allowedGroups map[uint]struct{}, ) int64 { var count int64 @@ -359,38 +361,38 @@ func countHomeUpstreamKeys( return count } -func countAvailableHomeKeys( +func countAvailableHomeCredentials( snapshot *state.ConfigSnapshot, - persisted []homeUpstreamKeyRow, - runtime []state.KeyRuntimeView, + persisted []homeCredentialRow, + runtime []state.CredentialRuntimeView, now time.Time, ) (int64, error) { - return countAvailableHomeKeysInGroups(snapshot, persisted, runtime, now, nil) + return countAvailableHomeCredentialsInGroups(snapshot, persisted, runtime, now, nil) } -func countAvailableHomeKeysInGroups( +func countAvailableHomeCredentialsInGroups( snapshot *state.ConfigSnapshot, - persisted []homeUpstreamKeyRow, - runtime []state.KeyRuntimeView, + persisted []homeCredentialRow, + runtime []state.CredentialRuntimeView, now time.Time, allowedGroups map[uint]struct{}, ) (int64, error) { - runtimeByID := make(map[uint]state.KeyRuntimeView, len(runtime)) - for _, key := range runtime { - if key.ID == 0 { + runtimeByID := make(map[uint]state.CredentialRuntimeView, len(runtime)) + for _, credential := range runtime { + if credential.ID == 0 { return 0, fmt.Errorf( - "count available home keys: runtime key has zero id: %w", + "count available home credentials: runtime credential has zero id: %w", app_errors.ErrInternalServer, ) } - if _, duplicate := runtimeByID[key.ID]; duplicate { + if _, duplicate := runtimeByID[credential.ID]; duplicate { return 0, fmt.Errorf( - "count available home keys: duplicate runtime key %d: %w", - key.ID, + "count available home credentials: duplicate runtime credential %d: %w", + credential.ID, app_errors.ErrInternalServer, ) } - runtimeByID[key.ID] = key + runtimeByID[credential.ID] = credential } seen := make(map[uint]struct{}, len(persisted)) @@ -398,58 +400,60 @@ func countAvailableHomeKeysInGroups( for _, row := range persisted { if row.ID == 0 { return 0, fmt.Errorf( - "count available home keys: persisted key has zero id: %w", + "count available home credentials: persisted credential has zero id: %w", app_errors.ErrInternalServer, ) } if _, duplicate := seen[row.ID]; duplicate { return 0, fmt.Errorf( - "count available home keys: duplicate persisted key %d: %w", + "count available home credentials: duplicate persisted credential %d: %w", row.ID, app_errors.ErrInternalServer, ) } seen[row.ID] = struct{}{} group, groupExists := snapshot.GroupCatalog[row.GroupID] - key, keyExists := runtimeByID[row.ID] - if !groupExists || !keyExists || key.GroupID != row.GroupID { + credential, credentialExists := runtimeByID[row.ID] + if !groupExists || !credentialExists || credential.GroupID != row.GroupID { return 0, fmt.Errorf( - "count available home key %d: runtime configuration mismatch: %w", + "count available home credential %d: runtime configuration mismatch: %w", row.ID, app_errors.ErrInternalServer, ) } - var status state.KeyStatus + var status state.CredentialStatus switch row.Status { - case models.UpstreamKeyStatusActive: - status = state.KeyStatusActive - case models.UpstreamKeyStatusDisabled: - status = state.KeyStatusDisabled + case models.CredentialStatusActive: + status = state.CredentialStatusActive + case models.CredentialStatusDisabled: + status = state.CredentialStatusDisabled default: return 0, fmt.Errorf( - "count available home key %d: invalid persisted status: %w", + "count available home credential %d: invalid persisted status: %w", row.ID, app_errors.ErrInternalServer, ) } - if key.Status != status { + if credential.Status != status || + credential.Version != groupCollectionCredentialVersion(row.UpdatedAtMS) || + credential.IdentityGeneration != groupCollectionCredentialIdentity(row.Fingerprint) { return 0, fmt.Errorf( - "count available home key %d: runtime status mismatch: %w", + "count available home credential %d: runtime identity mismatch: %w", row.ID, app_errors.ErrInternalServer, ) } _, groupAllowed := allowedGroups[row.GroupID] if (allowedGroups == nil || groupAllowed) && group.Enabled && - status == state.KeyStatusActive && - key.RuntimeState(now) == state.KeyRuntimeAvailable { + status == state.CredentialStatusActive && + credential.RuntimeState(now) == state.CredentialRuntimeAvailable { available++ } } if len(runtimeByID) != len(seen) { return 0, fmt.Errorf( - "count available home keys: runtime key set mismatch: %w", + "count available home credentials: runtime credential set mismatch: %w", app_errors.ErrInternalServer, ) } @@ -535,10 +539,10 @@ func validateHomeInventory(value HomeInventory) error { value int64 }{ {name: "group count", value: value.GroupCount}, - {name: "upstream key count", value: value.UpstreamKeyCount}, + {name: "credential count", value: value.CredentialCount}, { - name: "available upstream key count", - value: value.AvailableUpstreamKeyCount, + name: "available credential count", + value: value.AvailableCredentialCount, }, {name: "model count", value: value.ModelCount}, } { diff --git a/internal/control/home_test.go b/internal/control/home_test.go index fded3738..ff1f1851 100644 --- a/internal/control/home_test.go +++ b/internal/control/home_test.go @@ -2,13 +2,16 @@ package control import ( "context" + "encoding/json" "reflect" + "strings" "sync/atomic" "testing" "time" "gorm.io/gorm" + "gpt-load/internal/channel" "gpt-load/internal/protocol" "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" @@ -45,46 +48,54 @@ func TestReadHomeBaseUsesPersistedAndRuntimeSnapshots(t *testing.T) { t.Fatalf("disable group: %v", err) } - keys := []models.UpstreamKey{ + credentials := []models.Credential{ { - ID: 1, GroupID: enabled.ID, KeyValue: "cipher-1", KeyHash: "hash-1", - Status: models.UpstreamKeyStatusActive, + ID: 1, GroupID: enabled.ID, Data: "cipher-1", Fingerprint: "hash-1", + Status: models.CredentialStatusActive, }, { - ID: 2, GroupID: enabled.ID, KeyValue: "cipher-2", KeyHash: "hash-2", - Status: models.UpstreamKeyStatusActive, + ID: 2, GroupID: enabled.ID, Data: "cipher-2", Fingerprint: "hash-2", + Status: models.CredentialStatusActive, }, { - ID: 3, GroupID: enabledTwo.ID, KeyValue: "cipher-3", KeyHash: "hash-3", - Status: models.UpstreamKeyStatusDisabled, + ID: 3, GroupID: enabledTwo.ID, Data: "cipher-3", Fingerprint: "hash-3", + Status: models.CredentialStatusDisabled, }, { - ID: 4, GroupID: disabled.ID, KeyValue: "cipher-4", KeyHash: "hash-4", - Status: models.UpstreamKeyStatusActive, + ID: 4, GroupID: disabled.ID, Data: "cipher-4", Fingerprint: "hash-4", + Status: models.CredentialStatusActive, }, } - if err := fixture.db.Create(&keys).Error; err != nil { - t.Fatalf("create upstream keys: %v", err) + if err := fixture.db.Create(&credentials).Error; err != nil { + t.Fatalf("create credentials: %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{ + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{ { - ID: 1, GroupID: enabled.ID, Status: state.KeyStatusActive, + ID: 1, GroupID: enabled.ID, Status: state.CredentialStatusActive, + Version: groupCollectionCredentialVersion(credentials[0].UpdatedAtMS), + IdentityGeneration: groupCollectionCredentialIdentity(credentials[0].Fingerprint), Fingerprint: credentials[0].Fingerprint, EncryptedValue: "cipher-1", }, { - ID: 2, GroupID: enabled.ID, Status: state.KeyStatusActive, + ID: 2, GroupID: enabled.ID, Status: state.CredentialStatusActive, + Version: groupCollectionCredentialVersion(credentials[1].UpdatedAtMS), + IdentityGeneration: groupCollectionCredentialIdentity(credentials[1].Fingerprint), Fingerprint: credentials[1].Fingerprint, CooldownUntil: now.Add(time.Hour), EncryptedValue: "cipher-2", }, { - ID: 3, GroupID: enabledTwo.ID, Status: state.KeyStatusDisabled, + ID: 3, GroupID: enabledTwo.ID, Status: state.CredentialStatusDisabled, + Version: groupCollectionCredentialVersion(credentials[2].UpdatedAtMS), + IdentityGeneration: groupCollectionCredentialIdentity(credentials[2].Fingerprint), Fingerprint: credentials[2].Fingerprint, EncryptedValue: "cipher-3", }, { - ID: 4, GroupID: disabled.ID, Status: state.KeyStatusActive, + ID: 4, GroupID: disabled.ID, Status: state.CredentialStatusActive, + Version: groupCollectionCredentialVersion(credentials[3].UpdatedAtMS), + IdentityGeneration: groupCollectionCredentialIdentity(credentials[3].Fingerprint), Fingerprint: credentials[3].Fingerprint, EncryptedValue: "cipher-4", }, }); err != nil { - t.Fatalf("registry.Replace() error = %v", err) + t.Fatalf("registry.ReplaceCredentials() error = %v", err) } accessKeys := []models.AccessKey{ @@ -116,7 +127,7 @@ func TestReadHomeBaseUsesPersistedAndRuntimeSnapshots(t *testing.T) { if err := fixture.db.Create(&accessKeys).Error; err != nil { t.Fatalf("create access keys: %v", err) } - input, err := stateloader.BuildCompileInput(context.Background(), fixture.db) + input, err := stateloader.BuildCompileInput(context.Background(), fixture.db, fixture.channelRegistry) if err != nil { t.Fatalf("BuildCompileInput() error = %v", err) } @@ -125,7 +136,7 @@ func TestReadHomeBaseUsesPersistedAndRuntimeSnapshots(t *testing.T) { } registrySnapshotCalls := 0 - fixture.service.registrySnapshot = func() []state.KeyRuntimeView { + fixture.service.registrySnapshot = func() []state.CredentialRuntimeView { registrySnapshotCalls++ return fixture.registry.Snapshot() } @@ -136,10 +147,10 @@ func TestReadHomeBaseUsesPersistedAndRuntimeSnapshots(t *testing.T) { } want := HomeBase{ Inventory: HomeInventory{ - GroupCount: 3, - UpstreamKeyCount: 4, - AvailableUpstreamKeyCount: 1, - ModelCount: 4, + GroupCount: 3, + CredentialCount: 4, + AvailableCredentialCount: 1, + ModelCount: 4, }, AccessKeys: []HomeAccessKey{ { @@ -165,22 +176,40 @@ func TestReadHomeBaseUsesPersistedAndRuntimeSnapshots(t *testing.T) { } } +func TestHomeInventoryWireUsesCredentialTerms(t *testing.T) { + encoded, err := json.Marshal(HomeInventory{ + GroupCount: 1, CredentialCount: 2, AvailableCredentialCount: 1, ModelCount: 3, + }) + if err != nil { + t.Fatalf("json.Marshal() error = %v", err) + } + text := string(encoded) + for _, field := range []string{`"credential_count":2`, `"available_credential_count":1`} { + if !strings.Contains(text, field) { + t.Fatalf("home inventory missing %s: %s", field, text) + } + } + for _, legacy := range []string{"upstream_key_count", "available_upstream_key_count"} { + if strings.Contains(text, legacy) { + t.Fatalf("home inventory exposes %s: %s", legacy, text) + } + } +} + func TestReadAccessKeyHomeBaseScopesInventoryToRoutableModels(t *testing.T) { fixture := newServiceFixture(t) allowed := createPriceTestGroup(t, fixture.db, models.Group{ - Name: "allowed", UpstreamURL: "https://allowed.example/v1", - Protocols: models.JSON(`["openai-completions"]`), + Name: "allowed", ChannelID: string(channel.OpenAI), Params: models.JSON(`{}`), Models: models.JSON(`[ {"id":"upstream-allowed","alias":"client-allowed"}, {"id":"upstream-hidden","alias":"client-hidden"} ]`), - Config: models.JSON(`{}`), Enabled: true, + Overrides: models.JSON(`{}`), Enabled: true, }) createPriceTestGroup(t, fixture.db, models.Group{ - Name: "private", UpstreamURL: "https://private.example", - Protocols: models.JSON(`["gemini"]`), + Name: "private", ChannelID: string(channel.Gemini), Params: models.JSON(`{}`), Models: models.JSON(`[{"id":"private-upstream","alias":"private-client"}]`), - Config: models.JSON(`{}`), Enabled: true, + Overrides: models.JSON(`{}`), Enabled: true, }) created, err := fixture.service.CreateAccessKey(t.Context(), AccessKeyCreateRequest{ Name: "scoped home", diff --git a/internal/control/http_routes.go b/internal/control/http_routes.go index 3bad57f7..c27df14e 100644 --- a/internal/control/http_routes.go +++ b/internal/control/http_routes.go @@ -43,6 +43,7 @@ func (s *Server) HTTPModule() httproute.Module { MethodNotAllowed: controlMethodNotAllowed, Routes: []httproute.Route{ controlRoute("control.auth.session", http.MethodGet, "/auth/session", s.handleAuthSession), + controlRoute("control.channels.list", http.MethodGet, "/channels", s.handleListChannels), controlRoute("control.models.list", http.MethodGet, "/models", s.handleListProjectModels), controlRoute( "control.model-prices.detail", @@ -56,18 +57,6 @@ func (s *Server) HTTPModule() httproute.Module { "/model-prices", s.handleListModelPrices, ), - controlRoute( - "control.provider-suggestions.list", - http.MethodGet, - "/provider-suggestions", - s.handleListProviderSuggestions, - ), - controlRoute( - "control.provider-models.list", - http.MethodGet, - "/providers/:provider_id/models", - s.handleListProviderModels, - ), controlRoute( "control.model-prices.sync", http.MethodPost, @@ -213,76 +202,76 @@ func (s *Server) HTTPModule() httproute.Module { s.handleDeleteGroup, ), controlRoute( - "control.group-keys.list", + "control.group-credentials.list", http.MethodGet, - "/groups/:group_id/keys", - s.handleListGroupKeys, + "/groups/:group_id/credentials", + s.handleListGroupCredentials, ), controlRoute( - "control.group-keys.reveal", + "control.group-credentials.reveal", http.MethodPost, - "/groups/:group_id/keys/:key_id/reveal", + "/groups/:group_id/credentials/:credential_id/reveal", s.auditMutation(newMutationDescriptor( - "group_key_reveal", - "group_key", - groupKeyMutationLocator, + "group_credential_reveal", + "group_credential", + groupCredentialMutationLocator, )), - s.handleRevealGroupKey, + s.handleRevealGroupCredential, ), controlRoute( - "control.group-keys.update", + "control.group-credentials.update", http.MethodPut, - "/groups/:group_id/keys/:key_id", + "/groups/:group_id/credentials/:credential_id", s.auditMutation(newMutationDescriptor( - "group_key_update", - "group_key", - groupKeyMutationLocator, + "group_credential_update", + "group_credential", + groupCredentialMutationLocator, )), - s.handleUpdateGroupKey, + s.handleUpdateGroupCredential, ), controlRoute( - "control.group-keys.restore", + "control.group-credentials.restore", http.MethodPost, - "/groups/:group_id/keys/:key_id/restore", + "/groups/:group_id/credentials/:credential_id/restore", s.auditMutation(newMutationDescriptor( - "group_key_restore", - "group_key", - groupKeyMutationLocator, + "group_credential_restore", + "group_credential", + groupCredentialMutationLocator, )), - s.handleRestoreGroupKey, + s.handleRestoreGroupCredential, ), controlRoute( - "control.group-keys.batch", + "control.group-credentials.batch", http.MethodPost, - "/groups/:group_id/keys/batch", + "/groups/:group_id/credentials/batch", s.auditMutation(newMutationDescriptor( - "group_keys_batch", - "group_key", - groupKeysMutationLocator, + "group_credentials_batch", + "group_credential", + groupCredentialsMutationLocator, )), - s.handleBatchGroupKeys, + s.handleBatchGroupCredentials, ), controlRoute( - "control.group-keys.delete", + "control.group-credentials.delete", http.MethodDelete, - "/groups/:group_id/keys/:key_id", + "/groups/:group_id/credentials/:credential_id", s.auditMutation(newMutationDescriptor( - "group_key_delete", - "group_key", - groupKeyMutationLocator, + "group_credential_delete", + "group_credential", + groupCredentialMutationLocator, )), - s.handleDeleteGroupKey, + s.handleDeleteGroupCredential, ), controlRoute( - "control.group-keys.import", + "control.group-credentials.import", http.MethodPost, - "/groups/:group_id/keys/import", + "/groups/:group_id/credentials/import", s.auditMutation(newMutationDescriptor( - "group_key_import", - "group_key", - groupKeysMutationLocator, + "group_credential_import", + "group_credential", + groupCredentialsMutationLocator, )), - s.handleImportGroupKeys, + s.handleImportGroupCredentials, ), controlRoute( "control.group-models.discover", diff --git a/internal/control/idempotency_digest.go b/internal/control/idempotency_digest.go index ca456792..b87c1fbc 100644 --- a/internal/control/idempotency_digest.go +++ b/internal/control/idempotency_digest.go @@ -23,9 +23,9 @@ const ( type operationKind string const ( - operationKindAccessKeyCreate operationKind = "access_key_create" - operationKindGroupCreate operationKind = "group_create" - operationKindGroupKeyImport operationKind = "group_key_import" + operationKindAccessKeyCreate operationKind = "access_key_create" + operationKindGroupCreate operationKind = "group_create" + operationKindCredentialImport operationKind = "credential_import" ) type idempotencyDigestInput struct { @@ -104,7 +104,8 @@ func buildIdempotencyDigest( func (kind operationKind) valid() bool { switch kind { - case operationKindAccessKeyCreate, operationKindGroupCreate, operationKindGroupKeyImport: + case operationKindAccessKeyCreate, operationKindGroupCreate, + operationKindCredentialImport: return true default: return false @@ -129,7 +130,7 @@ func normalizeIdempotencyKeyLines(raw string) ([]string, error) { continue } lines = append(lines, normalized) - if len(lines) > maxUpstreamKeyLines { + if len(lines) > maxCredentialLines { return nil, app_errors.ErrValidation } } diff --git a/internal/control/idempotency_digest_test.go b/internal/control/idempotency_digest_test.go index 0fb8b0c9..0f04d3a3 100644 --- a/internal/control/idempotency_digest_test.go +++ b/internal/control/idempotency_digest_test.go @@ -24,8 +24,8 @@ func TestIdempotencyDigestFixtures(t *testing.T) { if err != nil { t.Fatalf("glob fixtures: %v", err) } - if len(paths) != 5 { - t.Fatalf("fixture count = %d, want 5", len(paths)) + if len(paths) != 3 { + t.Fatalf("fixture count = %d, want 3", len(paths)) } digests := make(map[string]string, len(paths)) for _, path := range paths { @@ -70,7 +70,6 @@ func TestIdempotencyDigestFixtures(t *testing.T) { } for _, pair := range [][2]string{ {"group-create-k", "group-create-k-twice"}, - {"group-key-import-k", "group-key-import-k-twice"}, } { if digests[pair[0]] == digests[pair[1]] { t.Fatalf("%s and %s digests are equal", pair[0], pair[1]) @@ -132,25 +131,25 @@ func TestNormalizeIdempotencyKeyLinesPreservesSortedMultiplicity(t *testing.T) { t.Fatalf("two = %#v, want [K K]", two) } - oneBody, err := canonicalIdempotencyBody(map[string]any{"keys": one}) + oneBody, err := canonicalIdempotencyBody(map[string]any{"credentials": one}) if err != nil { t.Fatalf("canonicalIdempotencyBody(one) error = %v", err) } - twoBody, err := canonicalIdempotencyBody(map[string]any{"keys": two}) + twoBody, err := canonicalIdempotencyBody(map[string]any{"credentials": two}) if err != nil { t.Fatalf("canonicalIdempotencyBody(two) error = %v", err) } oneDigest, err := buildIdempotencyDigest(idempotencyDigestInput{ - Version: 1, Method: "POST", OperationKind: operationKindGroupKeyImport, - PathTemplate: "/api/groups/:group_id/keys/import", ResourceLocator: "group:7", + Version: 1, Method: "POST", OperationKind: operationKindCredentialImport, + PathTemplate: "/api/groups/:group_id/credentials/import", ResourceLocator: "group:7", AuthScopeID: "control-admin-v1", CanonicalBody: oneBody, }) if err != nil { t.Fatalf("buildIdempotencyDigest(one) error = %v", err) } twoDigest, err := buildIdempotencyDigest(idempotencyDigestInput{ - Version: 1, Method: "POST", OperationKind: operationKindGroupKeyImport, - PathTemplate: "/api/groups/:group_id/keys/import", ResourceLocator: "group:7", + Version: 1, Method: "POST", OperationKind: operationKindCredentialImport, + PathTemplate: "/api/groups/:group_id/credentials/import", ResourceLocator: "group:7", AuthScopeID: "control-admin-v1", CanonicalBody: twoBody, }) if err != nil { diff --git a/internal/control/idempotency_operation.go b/internal/control/idempotency_operation.go index ad91ff77..4db48f99 100644 --- a/internal/control/idempotency_operation.go +++ b/internal/control/idempotency_operation.go @@ -84,7 +84,7 @@ func operationRequiredStages(kind operationKind) ([]operationStage, error) { operationStageSnapshotPublished, operationStageCompleted, } - case operationKindGroupKeyImport: + case operationKindCredentialImport: stages = []operationStage{ operationStageDBCommitted, operationStageRegistryApplied, @@ -355,7 +355,7 @@ func (s *Service) recoverOperationLocked( case operationStageRegistryApplied: stageErr = s.recoverRegistryOperation(ctx, operation) case operationStageSnapshotPublished: - compileInput, inputErr := stateloader.BuildCompileInput(ctx, s.db) + compileInput, inputErr := stateloader.BuildCompileInput(ctx, s.db, s.channelRegistry) if inputErr != nil { stageErr = inputErr } else { @@ -389,7 +389,7 @@ func (s *Service) recoverRegistryOperation( if err != nil { return err } - entries, err := stateloader.BuildGroupKeyEntries(ctx, s.db, groupID) + entries, err := stateloader.BuildGroupCredentialEntries(ctx, s.db, groupID) if err != nil { return err } diff --git a/internal/control/idempotency_operation_test.go b/internal/control/idempotency_operation_test.go index 03add227..edb984ee 100644 --- a/internal/control/idempotency_operation_test.go +++ b/internal/control/idempotency_operation_test.go @@ -36,7 +36,7 @@ func TestOperationRequiredStagesAreKindSpecificAndDetached(t *testing.T) { }, }, { - kind: operationKindGroupKeyImport, + kind: operationKindCredentialImport, want: []operationStage{ operationStageDBCommitted, operationStageRegistryApplied, diff --git a/internal/control/model_candidate.go b/internal/control/model_candidate.go new file mode 100644 index 00000000..b4f67fa9 --- /dev/null +++ b/internal/control/model_candidate.go @@ -0,0 +1,41 @@ +package control + +import ( + "context" + "fmt" + + "gorm.io/gorm" + + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/pricing" + "gpt-load/internal/storage/models" +) + +type ModelCandidate struct { + ID string `json:"id"` + Name string `json:"name"` + Sources []string `json:"sources"` + PricingStatus PricingStatus `json:"pricing_status"` + PricingSource *string `json:"pricing_source"` +} + +func loadModelPriceRows( + ctx context.Context, + db *gorm.DB, +) (modelPriceRows, error) { + var rows []models.ModelPrice + if err := db.WithContext(ctx).Order("id ASC").Find(&rows).Error; err != nil { + return nil, fmt.Errorf("load model price status: %w", app_errors.ParseDBError(err)) + } + result := make(modelPriceRows, len(rows)) + for index := range rows { + identity := pricing.Identity{ChannelID: rows[index].ChannelID, ModelID: rows[index].ModelID} + if _, duplicate := result[identity]; duplicate { + return nil, fmt.Errorf("duplicate persisted price identity: %w", app_errors.ErrInternalServer) + } + result[identity] = &rows[index] + } + return result, nil +} + +type modelPriceRows map[pricing.Identity]*models.ModelPrice diff --git a/internal/control/model_price.go b/internal/control/model_price.go index dd803581..570e42af 100644 --- a/internal/control/model_price.go +++ b/internal/control/model_price.go @@ -24,11 +24,13 @@ type PriceSlotsDTO struct { type ModelPriceDTO struct { ID uint `json:"id"` + ChannelID string `json:"channel_id"` ModelID string `json:"model_id"` Prices PriceSlotsDTO `json:"prices"` PricingStatus PricingStatus `json:"pricing_status"` Method *string `json:"method"` MatchedProviderID *string `json:"matched_provider_id"` + MatchSource *ModelPriceMatchSource `json:"match_source"` Referenced bool `json:"referenced"` ReferenceCount int `json:"reference_count"` ReferenceGroupCount int `json:"reference_group_count"` @@ -198,7 +200,11 @@ func (s *Service) ResetModelPrice( if err != nil { return err } - desired, err := resetModelPriceValues(row.ModelID, catalogSnapshot) + identity, err := PriceIdentityForChannelModel(row.ChannelID, row.ModelID) + if err != nil { + return fmt.Errorf("validate model price identity: %w", app_errors.ErrInternalServer) + } + desired, err := resetModelPriceValues(identity, catalogSnapshot) if err != nil { return fmt.Errorf("normalize catalog model price: %w", app_errors.ErrInternalServer) } @@ -247,12 +253,12 @@ func (s *Service) ResetModelPrice( return result, nil } -func resetModelPriceValues(modelID string, snapshot *catalog.Snapshot) (models.ModelPrice, error) { - cost, _, ok := resolveAutomaticPrice(snapshot, modelID) +func resetModelPriceValues(identity pricing.Identity, snapshot *catalog.Snapshot) (models.ModelPrice, error) { + match, ok := resolveAutomaticPriceForIdentity(snapshot, identity) if !ok { return models.ModelPrice{}, nil } - return automaticCatalogValues(cost) + return automaticCatalogValues(match.cost) } func (s *Service) DeleteModelPrice(ctx context.Context, id uint) error { @@ -275,8 +281,8 @@ func (s *Service) DeleteModelPrice(ctx context.Context, id uint) error { if err := tx.First(&row, id).Error; err != nil { return fmt.Errorf("load model price: %w", app_errors.ParseDBError(err)) } - identity := pricing.Identity{ModelID: row.ModelID} - if _, err := pricing.NewTable([]pricing.Rule{{Identity: identity}}); err != nil { + identity, err := PriceIdentityForChannelModel(row.ChannelID, row.ModelID) + if err != nil { return fmt.Errorf("validate model price identity: %w", app_errors.ErrInternalServer) } references, err := loadPriceReferenceSnapshot(tx) @@ -472,7 +478,7 @@ func projectModelPriceRow( if _, err := pricing.NewTable([]pricing.Rule{rule}); err != nil { return modelPriceListRecord{}, fmt.Errorf("validate persisted model price: %w", app_errors.ErrInternalServer) } - identity := pricing.Identity{ModelID: row.ModelID} + identity := rule.Identity reference := references.references[identity] status := resolvePricingStatus(&row) prices := PriceSlotsDTO{ @@ -483,22 +489,27 @@ func projectModelPriceRow( } configured := modelPriceHasConfiguredValue(row) var matchedProviderID *string + var matchSource *ModelPriceMatchSource matchedAutomaticPrice := false - matchedProvider := "" if !row.IsManual { - _, matchedProvider, matchedAutomaticPrice = resolveAutomaticPrice(catalogSnapshot, row.ModelID) + match, matched := resolveAutomaticPriceForIdentity(catalogSnapshot, identity) + matchedAutomaticPrice = matched if configured && matchedAutomaticPrice { - matchedProviderID = &matchedProvider + providerID := match.providerID + matchedProviderID = &providerID + source := match.source + matchSource = &source } else { matchedAutomaticPrice = false } } dto := ModelPriceDTO{ - ID: row.ID, ModelID: row.ModelID, + ID: row.ID, ChannelID: row.ChannelID, ModelID: row.ModelID, Prices: prices, PricingStatus: status, Method: modelPriceMethod(row, configured, matchedAutomaticPrice), MatchedProviderID: matchedProviderID, + MatchSource: matchSource, Referenced: reference.referenceCount > 0, ReferenceCount: reference.referenceCount, ReferenceGroupCount: reference.referenceGroupCount(), @@ -624,7 +635,8 @@ func modelPriceRecordMatches(record modelPriceListRecord, query ModelPriceListQu return false } return query.Search == "" || - accessKeyCollectionContainsFold(record.dto.ModelID, query.Search) + accessKeyCollectionContainsFold(record.dto.ModelID, query.Search) || + accessKeyCollectionContainsFold(record.dto.ChannelID, query.Search) } func sortModelPriceRecords(records []modelPriceListRecord) { @@ -639,6 +651,9 @@ func sortModelPriceRecords(records []modelPriceListRecord) { if left.dto.ModelID != right.dto.ModelID { return left.dto.ModelID < right.dto.ModelID } + if left.dto.ChannelID != right.dto.ChannelID { + return left.dto.ChannelID < right.dto.ChannelID + } return left.dto.ID < right.dto.ID }) } diff --git a/internal/control/model_price_sync_http.go b/internal/control/model_price_sync_http.go new file mode 100644 index 00000000..d5c75ccb --- /dev/null +++ b/internal/control/model_price_sync_http.go @@ -0,0 +1,35 @@ +package control + +import ( + "context" + "fmt" + + "github.com/gin-gonic/gin" + + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/platform/response" +) + +func (s *Server) handleSyncModelPrices(c *gin.Context) { + if c.Request.URL.ForceQuery || c.Request.URL.RawQuery != "" { + writeServiceError(c, "sync_model_prices", app_errors.ErrBadRequest) + return + } + result, err := s.service.SyncModelPrices(c.Request.Context()) + if err != nil { + writeServiceError(c, "sync_model_prices", err) + return + } + response.SuccessI18n(c, "common.success", result) +} + +func (s *Service) SyncModelPrices(ctx context.Context) (CatalogSyncStatus, error) { + if s.catalogSync == nil { + return CatalogSyncStatus{}, fmt.Errorf("catalog sync is unavailable: %w", app_errors.ErrInternalServer) + } + status, err := s.catalogSync.Sync(ctx, CatalogSyncManual) + if err != nil { + return status, app_errors.NewAPIErrorWithData(app_errors.ErrBadGateway, status) + } + return status, nil +} diff --git a/internal/control/model_price_sync_http_test.go b/internal/control/model_price_sync_http_test.go new file mode 100644 index 00000000..f01cc7ed --- /dev/null +++ b/internal/control/model_price_sync_http_test.go @@ -0,0 +1,67 @@ +package control + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/sirupsen/logrus" + + "gpt-load/internal/catalog" + "gpt-load/internal/platform/config" +) + +func TestModelPriceSyncRouteSanitizesFailure(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + const rawFailure = "secret upstream response body" + client := catalogSyncClientFunc(func(context.Context, catalog.Metadata) (catalog.SyncResult, error) { + return catalog.SyncResult{}, errors.New(rawFailure) + }) + coordinator := newCatalogSyncCoordinator(fixture.service, client, "unused", catalog.Metadata{ + CheckedAtMillis: 100, SuccessfulFetchAtMillis: 90, + }, true) + coordinator.now = func() time.Time { return time.UnixMilli(250) } + var serviceLogs bytes.Buffer + previousOutput := logrus.StandardLogger().Out + logrus.SetOutput(&serviceLogs) + t.Cleanup(func() { logrus.SetOutput(previousOutput) }) + engine := gin.New() + NewServer(&config.Config{AuthKey: authTestKey}, fixture.service).RegisterRoutes(engine) + + syncResponse := serveModelPriceSyncRequest(engine, authTestKey) + if syncResponse.Code != http.StatusBadGateway || + strings.Contains(syncResponse.Body.String(), rawFailure) { + t.Fatalf("sync failure = %d %s", syncResponse.Code, syncResponse.Body.String()) + } + var envelope struct { + Code string `json:"code"` + Data CatalogSyncStatus `json:"data"` + } + if err := json.Unmarshal(syncResponse.Body.Bytes(), &envelope); err != nil { + t.Fatal(err) + } + if envelope.Code != "BAD_GATEWAY" || envelope.Data.Trigger != CatalogSyncManual || + envelope.Data.CheckedAtMS != 250 || envelope.Data.SuccessfulFetchAtMS != 90 || + envelope.Data.ErrorCode != "catalog_sync_failed" || envelope.Data.NotModified || + envelope.Data.Skipped { + t.Fatalf("sync failure envelope = %#v", envelope) + } + assertControlLogExcludes(t, serviceLogs.String(), rawFailure) +} + +func serveModelPriceSyncRequest(engine *gin.Engine, authKey string) *httptest.ResponseRecorder { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/api/model-prices/sync", nil) + request.Header.Set("Authorization", "Bearer "+authKey) + request.Header.Set("Accept-Language", "en-US") + engine.ServeHTTP(recorder, request) + return recorder +} diff --git a/internal/control/model_price_test.go b/internal/control/model_price_test.go index 63ccd30e..2ef49584 100644 --- a/internal/control/model_price_test.go +++ b/internal/control/model_price_test.go @@ -4,24 +4,25 @@ import ( "encoding/json" "testing" + "gpt-load/internal/catalog" + "gpt-load/internal/channel" "gpt-load/internal/pricing" "gpt-load/internal/storage/models" ) -func TestGlobalModelPriceIsSharedAcrossGroupsAndAliases(t *testing.T) { +func TestChannelModelPriceIsSharedAcrossGroupsAndAliases(t *testing.T) { fixture := newServiceFixture(t) - providerID := "openai" createPriceTestGroup(t, fixture.db, models.Group{ - Name: "first", ProviderID: &providerID, UpstreamURL: "https://first.example/v1", - Protocols: models.JSON(`["openai-completions"]`), + Name: "first", ChannelID: string(channel.OpenAICompatible), + Params: models.JSON(`{"base_url":"https://first.example/v1"}`), Models: models.JSON(`[{"id":"shared-upstream","alias":"client-a"},{"id":"shared-upstream","alias":"client-b"}]`), - Config: models.JSON(`{}`), Enabled: true, + Overrides: models.JSON(`{}`), Enabled: true, }) createPriceTestGroup(t, fixture.db, models.Group{ - Name: "second", ProviderID: &providerID, UpstreamURL: "https://second.example/v1", - Protocols: models.JSON(`["openai-completions"]`), + Name: "second", ChannelID: string(channel.OpenAICompatible), + Params: models.JSON(`{"base_url":"https://second.example/v1"}`), Models: models.JSON(`[{"id":"shared-upstream","alias":"client-a"}]`), - Config: models.JSON(`{}`), Enabled: true, + Overrides: models.JSON(`{}`), Enabled: true, }) mustEnsureInitialPrices(t, fixture) @@ -29,7 +30,7 @@ func TestGlobalModelPriceIsSharedAcrossGroupsAndAliases(t *testing.T) { if err := fixture.db.Order("id ASC").Find(&rows).Error; err != nil { t.Fatal(err) } - if len(rows) != 1 || rows[0].ModelID != "shared-upstream" { + if len(rows) != 1 || rows[0].ChannelID != string(channel.OpenAICompatible) || rows[0].ModelID != "shared-upstream" { t.Fatalf("persisted prices = %#v, want one shared upstream price", rows) } @@ -51,6 +52,9 @@ func TestGlobalModelPriceIsSharedAcrossGroupsAndAliases(t *testing.T) { t.Fatal(err) } item := wire["items"].([]any)[0].(map[string]any) + if item["channel_id"] != string(channel.OpenAICompatible) { + t.Fatalf("price response channel_id = %#v", item["channel_id"]) + } if _, exists := item["scope"]; exists { t.Fatalf("global price response leaked scope: %s", encoded) } @@ -59,7 +63,7 @@ func TestGlobalModelPriceIsSharedAcrossGroupsAndAliases(t *testing.T) { func TestProjectModelPriceRowMarksIncompleteTierPartial(t *testing.T) { value := int64(1) record, err := projectModelPriceRow(models.ModelPrice{ - ID: 1, ModelID: "tiered", InputPriceNanoUSDPerMillionTokens: &value, + ID: 1, ChannelID: string(channel.OpenAICompatible), ModelID: "tiered", InputPriceNanoUSDPerMillionTokens: &value, ContextPriceTiers: models.JSON(`[{"threshold_tokens":1000,"input_price_nano_usd_per_million_tokens":1,"output_price_nano_usd_per_million_tokens":null,"cache_read_price_nano_usd_per_million_tokens":null,"cache_write_price_nano_usd_per_million_tokens":null}]`), }, priceReferenceSnapshot{references: map[pricing.Identity]referencedPrice{}}, nil) if err != nil { @@ -70,6 +74,93 @@ func TestProjectModelPriceRowMarksIncompleteTierPartial(t *testing.T) { } } +func TestResetModelPriceChangesOnlySelectedChannelIdentity(t *testing.T) { + fixture := newServiceFixture(t) + manual := int64(99) + rows := []models.ModelPrice{ + { + ChannelID: string(channel.OpenAI), ModelID: "shared", + InputPriceNanoUSDPerMillionTokens: &manual, IsManual: true, + }, + { + ChannelID: string(channel.Anthropic), ModelID: "shared", + InputPriceNanoUSDPerMillionTokens: &manual, IsManual: true, + }, + } + for index := range rows { + if err := fixture.db.Create(&rows[index]).Error; err != nil { + t.Fatal(err) + } + } + fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ + "openai": {ID: "openai", Models: map[string]catalog.Model{ + "shared": {ID: "shared", Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(10)}}}, + }}, + "anthropic": {ID: "anthropic", Models: map[string]catalog.Model{ + "shared": {ID: "shared", Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(20)}}}, + }}, + }}) + + reset, err := fixture.service.ResetModelPrice(t.Context(), rows[1].ID) + if err != nil { + t.Fatal(err) + } + if reset.ChannelID != string(channel.Anthropic) || reset.Prices.Input == nil || + *reset.Prices.Input != "0.00000002" || reset.MatchSource == nil || + *reset.MatchSource != ModelPriceMatchSourceChannelCatalogProvider || + reset.MatchedProviderID == nil || *reset.MatchedProviderID != "anthropic" { + t.Fatalf("reset Anthropic price = %#v", reset) + } + var openAI models.ModelPrice + if err := fixture.db.First(&openAI, rows[0].ID).Error; err != nil { + t.Fatal(err) + } + if !openAI.IsManual || openAI.InputPriceNanoUSDPerMillionTokens == nil || + *openAI.InputPriceNanoUSDPerMillionTokens != 99 { + t.Fatalf("OpenAI price changed with Anthropic reset: %#v", openAI) + } + table := fixture.priceRuntime.Load() + for _, expected := range []struct { + identity pricing.Identity + value pricing.NanoUSD + }{ + {identity: pricing.Identity{ChannelID: string(channel.OpenAI), ModelID: "shared"}, value: 99}, + {identity: pricing.Identity{ChannelID: string(channel.Anthropic), ModelID: "shared"}, value: 20}, + } { + rule, ok := table.Lookup(expected.identity) + if !ok || !rule.Prices.Input.Set || rule.Prices.Input.NanoUSDPerMillion != expected.value { + t.Fatalf("runtime rule %#v = %#v, %t", expected.identity, rule, ok) + } + } +} + +func TestDeleteModelPriceIgnoresSameModelReferenceFromDifferentChannel(t *testing.T) { + fixture := newServiceFixture(t) + manual := int64(1) + row := models.ModelPrice{ + ChannelID: string(channel.OpenAI), ModelID: "shared", + InputPriceNanoUSDPerMillionTokens: &manual, IsManual: true, + } + if err := fixture.db.Create(&row).Error; err != nil { + t.Fatal(err) + } + createPriceTestGroup(t, fixture.db, models.Group{ + Name: "anthropic", ChannelID: string(channel.Anthropic), Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"shared"}]`), Enabled: true, + }) + + if err := fixture.service.DeleteModelPrice(t.Context(), row.ID); err != nil { + t.Fatalf("DeleteModelPrice() error = %v", err) + } + var count int64 + if err := fixture.db.Model(&models.ModelPrice{}).Where("id = ?", row.ID).Count(&count).Error; err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatalf("deleted OpenAI price count = %d", count) + } +} + func mustModelPriceUpdateRequest(t *testing.T, body string) ModelPriceUpdateRequest { t.Helper() var request ModelPriceUpdateRequest diff --git a/internal/control/model_prices_audit_test.go b/internal/control/model_prices_audit_test.go index 216e2f98..2eed24e4 100644 --- a/internal/control/model_prices_audit_test.go +++ b/internal/control/model_prices_audit_test.go @@ -8,6 +8,7 @@ import ( "strings" "testing" + "gpt-load/internal/channel" "gpt-load/internal/storage/models" ) @@ -239,6 +240,7 @@ func newModelPriceAuditFixture(t *testing.T) (serviceFixture, models.ModelPrice) fixture := newServiceFixture(t) input := int64(1_000_000_000) row := models.ModelPrice{ + ChannelID: string(channel.OpenAICompatible), ModelID: "model-audit-secret", InputPriceNanoUSDPerMillionTokens: &input, } diff --git a/internal/control/model_prices_http_test.go b/internal/control/model_prices_http_test.go index 0d030777..f835f5eb 100644 --- a/internal/control/model_prices_http_test.go +++ b/internal/control/model_prices_http_test.go @@ -11,6 +11,7 @@ import ( "github.com/gin-gonic/gin" "gpt-load/internal/catalog" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" "gpt-load/internal/pricing" "gpt-load/internal/storage/models" @@ -66,7 +67,8 @@ func TestModelPriceHTTPListAndUpdateUseFinalWireContract(t *testing.T) { t.Fatalf("list items = %#v", listData["items"]) } item, ok := items[0].(map[string]any) - if !ok || item["id"] != float64(row.ID) || item["model_id"] != "gpt-wire" { + if !ok || item["id"] != float64(row.ID) || + item["channel_id"] != string(channel.OpenAICompatible) || item["model_id"] != "gpt-wire" { t.Fatalf("list item identity = %#v", items[0]) } prices, ok := item["prices"].(map[string]any) @@ -76,6 +78,9 @@ func TestModelPriceHTTPListAndUpdateUseFinalWireContract(t *testing.T) { if matchedProviderID, exists := item["matched_provider_id"]; !exists || matchedProviderID != "openai" { t.Fatalf("list matched provider wire = %#v", item) } + if matchSource, exists := item["match_source"]; !exists || matchSource != string(ModelPriceMatchSourceProviderPriorityFallback) { + t.Fatalf("list match source wire = %#v", item) + } if _, leaked := item["price_scope_key"]; leaked { t.Fatalf("list leaked internal scope: %#v", item) } @@ -108,6 +113,9 @@ func TestModelPriceHTTPListAndUpdateUseFinalWireContract(t *testing.T) { if matchedProviderID, exists := updated["matched_provider_id"]; !exists || matchedProviderID != nil { t.Fatalf("updated matched provider wire = %#v", updated) } + if matchSource, exists := updated["match_source"]; !exists || matchSource != nil { + t.Fatalf("updated match source wire = %#v", updated) + } conflictResponse := serveModelPriceHTTPRequest( engine, @@ -265,6 +273,9 @@ func TestModelPriceHTTPResetDeleteAndStableErrorData(t *testing.T) { if matchedProviderID, exists := data["matched_provider_id"]; !exists || matchedProviderID != nil { t.Fatalf("reset matched provider wire = %#v", data) } + if matchSource, exists := data["match_source"]; !exists || matchSource != nil { + t.Fatalf("reset match source wire = %#v", data) + } }) @@ -283,20 +294,17 @@ func TestModelPriceHTTPResetDeleteAndStableErrorData(t *testing.T) { t.Run("referenced delete conflict includes both counts", func(t *testing.T) { fixture, engine, row := newModelPriceHTTPFixture(t, true) - providerID := "openai" createPriceTestGroup(t, fixture.db, models.Group{ - Name: "reference-one", ProviderID: &providerID, - UpstreamURL: "https://one.example/v1", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[{"id":"gpt-wire","alias":"one"},{"id":"gpt-wire","alias":"two"}]`), - Config: models.JSON(`{}`), Enabled: true, + Name: "reference-one", ChannelID: string(channel.OpenAICompatible), + Params: models.JSON(`{"base_url":"https://one.example/v1"}`), + Models: models.JSON(`[{"id":"gpt-wire","alias":"one"},{"id":"gpt-wire","alias":"two"}]`), + Overrides: models.JSON(`{}`), Enabled: true, }) createPriceTestGroup(t, fixture.db, models.Group{ - Name: "reference-two", ProviderID: &providerID, - UpstreamURL: "https://two.example/v1", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[{"id":"gpt-wire","alias":"three"}]`), - Config: models.JSON(`{}`), Enabled: false, + Name: "reference-two", ChannelID: string(channel.OpenAICompatible), + Params: models.JSON(`{"base_url":"https://two.example/v1"}`), + Models: models.JSON(`[{"id":"gpt-wire","alias":"three"}]`), + Overrides: models.JSON(`{}`), Enabled: false, }) response := serveModelPriceHTTPRequest( engine, @@ -435,6 +443,7 @@ func newModelPriceHTTPFixture( if seed { input := int64(2_500_000_000) row = models.ModelPrice{ + ChannelID: string(channel.OpenAICompatible), ModelID: "gpt-wire", InputPriceNanoUSDPerMillionTokens: &input, } diff --git a/internal/control/mutation_audit.go b/internal/control/mutation_audit.go index eae285a2..de01baf5 100644 --- a/internal/control/mutation_audit.go +++ b/internal/control/mutation_audit.go @@ -75,13 +75,13 @@ func groupMutationLocator(c *gin.Context) string { return "group:" + mutationID(c.Param("group_id")) } -func groupKeyMutationLocator(c *gin.Context) string { +func groupCredentialMutationLocator(c *gin.Context) string { return "group:" + mutationID(c.Param("group_id")) + - "/key:" + mutationID(c.Param("key_id")) + "/credential:" + mutationID(c.Param("credential_id")) } -func groupKeysMutationLocator(c *gin.Context) string { - return "group:" + mutationID(c.Param("group_id")) + "/keys" +func groupCredentialsMutationLocator(c *gin.Context) string { + return "group:" + mutationID(c.Param("group_id")) + "/credentials" } func accessKeyMutationLocator(c *gin.Context) string { diff --git a/internal/control/mutation_audit_routes_test.go b/internal/control/mutation_audit_routes_test.go index 2f505731..89de778b 100644 --- a/internal/control/mutation_audit_routes_test.go +++ b/internal/control/mutation_audit_routes_test.go @@ -189,7 +189,7 @@ func TestGroupMutationAuditIncompleteAndBlocked(t *testing.T) { fixture := newServiceFixture(t) fixture.service.reconcileRegistryGroup = func( uint, - []state.KeyEntry, + []state.CredentialEntry, ) (bool, error) { return false, errors.New("injected registry failure") } @@ -251,7 +251,7 @@ func TestGroupMutationAuditIncompleteAndBlocked(t *testing.T) { func TestGroupMutationAuditExcludesInternalCalls(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport( + groupID := createGroupWithCredentials( t, fixture, "sk-direct-service", @@ -1222,8 +1222,8 @@ func groupMutationAuditCases() []groupMutationAuditCase { database: groupAuditSeedRequest(http.MethodDelete, "", ""), }, { - operation: "group_key_update", - success: groupKeyAuditSeedRequest( + operation: "group_credential_update", + success: groupCredentialAuditSeedRequest( http.MethodPut, `{"status":"disabled"}`, ), @@ -1233,20 +1233,20 @@ func groupMutationAuditCases() []groupMutationAuditCase { ) (mutationAuditRequest, string, string) { return mutationAuditRequest{ method: http.MethodPut, - path: "/api/groups/raw-secret/keys/1", + path: "/api/groups/raw-secret/credentials/1", body: `{"status":"disabled"}`, }, - "group:unknown/key:1", + "group:unknown/credential:1", app_errors.ErrBadRequest.Code }, - database: groupKeyAuditSeedRequest( + database: groupCredentialAuditSeedRequest( http.MethodPut, `{"status":"disabled"}`, ), }, { - operation: "group_key_delete", - success: groupKeyAuditSeedRequest( + operation: "group_credential_delete", + success: groupCredentialAuditSeedRequest( http.MethodDelete, "", ), @@ -1256,26 +1256,26 @@ func groupMutationAuditCases() []groupMutationAuditCase { ) (mutationAuditRequest, string, string) { return mutationAuditRequest{ method: http.MethodDelete, - path: "/api/groups/1/keys/raw-secret", + path: "/api/groups/1/credentials/raw-secret", }, - "group:1/key:unknown", + "group:1/credential:unknown", app_errors.ErrBadRequest.Code }, - database: groupKeyAuditSeedRequest( + database: groupCredentialAuditSeedRequest( http.MethodDelete, "", ), }, { - operation: "group_key_import", - success: groupKeyImportAuditSeedRequest( + operation: "group_credential_import", + success: groupCredentialImportAuditSeedRequest( "00000000-0000-4000-8000-000000001003", ), rejected: func( t *testing.T, fixture serviceFixture, ) (mutationAuditRequest, string, string) { - groupID := createGroupForKeyImport( + groupID := createGroupForCredentialImport( t, fixture, "sk-import-existing", @@ -1284,13 +1284,13 @@ func groupMutationAuditCases() []groupMutationAuditCase { method: http.MethodPost, path: "/api/groups/" + strconv.FormatUint(uint64(groupID), 10) + - "/keys/import", - body: `{"keys":"sk-new-audit-key"}`, + "/credentials/import", + body: `{"credentials":"sk-new-audit-key"}`, }, - fmt.Sprintf("group:%d/keys", groupID), + fmt.Sprintf("group:%d/credentials", groupID), app_errors.ErrIdempotencyKeyRequired.Code }, - database: groupKeyImportAuditSeedRequest( + database: groupCredentialImportAuditSeedRequest( "00000000-0000-4000-8000-000000001004", ), }, @@ -1299,10 +1299,9 @@ func groupMutationAuditCases() []groupMutationAuditCase { func groupCreateAuditBody(name string, upstreamURL string) string { return fmt.Sprintf( - `{"name":%q,"upstream_url":%q,`+ - `"protocols":["openai-completions"],`+ + `{"name":%q,"channel_id":"openai_compatible","params":{"base_url":%q},`+ `"models":[{"id":"gpt-4o","alias_enabled":false}],`+ - `"keys":"sk-audit-upstream"}`, + `"credentials":"sk-audit-upstream"}`, name, upstreamURL, ) @@ -1317,7 +1316,7 @@ func groupAuditSeedRequest( t *testing.T, fixture serviceFixture, ) (mutationAuditRequest, string) { - groupID := createGroupForKeyImport( + groupID := createGroupForCredentialImport( t, fixture, "sk-group-audit", @@ -1333,7 +1332,7 @@ func groupAuditSeedRequest( } } -func groupKeyAuditSeedRequest( +func groupCredentialAuditSeedRequest( method string, body string, ) func(*testing.T, serviceFixture) (mutationAuditRequest, string) { @@ -1341,41 +1340,41 @@ func groupKeyAuditSeedRequest( t *testing.T, fixture serviceFixture, ) (mutationAuditRequest, string) { - groupID := createGroupForKeyImport( + groupID := createGroupForCredentialImport( t, fixture, "sk-group-key-audit", ) - var key models.UpstreamKey + var credential models.Credential if err := fixture.db.Where("group_id = ?", groupID). - First(&key).Error; err != nil { - t.Fatalf("query seeded key: %v", err) + First(&credential).Error; err != nil { + t.Fatalf("query seeded credential: %v", err) } locator := fmt.Sprintf( - "group:%d/key:%d", + "group:%d/credential:%d", groupID, - key.ID, + credential.ID, ) return mutationAuditRequest{ method: method, path: fmt.Sprintf( - "/api/groups/%d/keys/%d", + "/api/groups/%d/credentials/%d", groupID, - key.ID, + credential.ID, ), body: body, }, locator } } -func groupKeyImportAuditSeedRequest( +func groupCredentialImportAuditSeedRequest( idempotencyKey string, ) func(*testing.T, serviceFixture) (mutationAuditRequest, string) { return func( t *testing.T, fixture serviceFixture, ) (mutationAuditRequest, string) { - groupID := createGroupForKeyImport( + groupID := createGroupForCredentialImport( t, fixture, "sk-import-audit", @@ -1383,20 +1382,20 @@ func groupKeyImportAuditSeedRequest( return mutationAuditRequest{ method: http.MethodPost, path: fmt.Sprintf( - "/api/groups/%d/keys/import", + "/api/groups/%d/credentials/import", groupID, ), - body: `{"keys":"sk-new-audit-key"}`, + body: `{"credentials":"sk-new-audit-key"}`, idempotencyKey: idempotencyKey, - }, fmt.Sprintf("group:%d/keys", groupID) + }, fmt.Sprintf("group:%d/credentials", groupID) } } func groupMutationResourceType(operation string) string { - if operation == "group_key_update" || - operation == "group_key_delete" || - operation == "group_key_import" { - return "group_key" + if operation == "group_credential_update" || + operation == "group_credential_delete" || + operation == "group_credential_import" { + return "group_credential" } return "group" } diff --git a/internal/control/mutation_audit_test.go b/internal/control/mutation_audit_test.go index 3e3c0a97..e62d88bb 100644 --- a/internal/control/mutation_audit_test.go +++ b/internal/control/mutation_audit_test.go @@ -93,25 +93,25 @@ func TestMutationLocatorsOnlyExposeValidatedIDs(t *testing.T) { locate: groupMutationLocator, want: "group:unknown", }, { - name: "group key", + name: "group credential", params: gin.Params{ {Key: "group_id", Value: "7"}, - {Key: "key_id", Value: "9"}, + {Key: "credential_id", Value: "9"}, }, - locate: groupKeyMutationLocator, want: "group:7/key:9", + locate: groupCredentialMutationLocator, want: "group:7/credential:9", }, { - name: "group key invalid key", + name: "group credential invalid credential", params: gin.Params{ {Key: "group_id", Value: "7"}, - {Key: "key_id", Value: "-7"}, + {Key: "credential_id", Value: "-7"}, }, - locate: groupKeyMutationLocator, want: "group:7/key:unknown", + locate: groupCredentialMutationLocator, want: "group:7/credential:unknown", }, { - name: "group keys", + name: "group credentials", params: gin.Params{{Key: "group_id", Value: "12"}}, - locate: groupKeysMutationLocator, want: "group:12/keys", + locate: groupCredentialsMutationLocator, want: "group:12/credentials", }, { name: "access key", diff --git a/internal/control/operation_error.go b/internal/control/operation_error.go index 28e5ab93..d55ad4c5 100644 --- a/internal/control/operation_error.go +++ b/internal/control/operation_error.go @@ -20,13 +20,14 @@ const ( mismatchGroupID = "group_id" mismatchStatus = "status" mismatchWeightManual = "weight_manual" + mismatchIdentity = "identity" ) type controlOperationError struct { stage string mismatchKind string groupID uint - keyID uint + credentialID uint } func (e *controlOperationError) Error() string { @@ -57,30 +58,30 @@ func groupNotFoundError() error { return &controlResourceNotFoundError{resource: "group"} } -func keyNotFoundError() error { - return &controlResourceNotFoundError{resource: "key"} +func credentialNotFoundError() error { + return &controlResourceNotFoundError{resource: "credential"} } func dbRegistryMismatch( kind dbRegistryMismatchKind, groupID uint, - keyID uint, + credentialID uint, ) error { return &controlOperationError{ stage: stageValidateDBRegistryPair, mismatchKind: string(kind), groupID: groupID, - keyID: keyID, + credentialID: credentialID, } } -func withControlOperationContext(err error, groupID, keyID uint) error { +func withControlOperationContext(err error, groupID, credentialID uint) error { var operationErr *controlOperationError if !errors.As(err, &operationErr) { return err } cloned := *operationErr cloned.groupID = groupID - cloned.keyID = keyID + cloned.credentialID = credentialID return &cloned } diff --git a/internal/control/operation_error_test.go b/internal/control/operation_error_test.go index d7a474b3..8e3bd7aa 100644 --- a/internal/control/operation_error_test.go +++ b/internal/control/operation_error_test.go @@ -26,7 +26,7 @@ func TestControlOperationErrorCarriesOnlyFixedContext(t *testing.T) { t.Fatalf("error = %T, want *controlOperationError", err) } if operationErr.stage != stageApplyCommittedRegistryMutation || - operationErr.groupID != 12 || operationErr.keyID != 34 || + operationErr.groupID != 12 || operationErr.credentialID != 34 || operationErr.mismatchKind != "" { t.Fatalf("operation error = %#v", operationErr) } @@ -40,11 +40,11 @@ func TestServiceErrorMessageIDUsesTypedResourceNotFound(t *testing.T) { want string }{ { - name: "list Group keys missing Group", operation: "list_group_keys", + name: "list Group credentials missing Group", operation: "list_group_credentials", err: groupNotFoundError(), want: "group.not_found", }, { - name: "update Group key missing Group", operation: "update_group_key", + name: "update Group credential missing Group", operation: "update_group_credential", err: groupNotFoundError(), want: "group.not_found", }, { @@ -52,12 +52,12 @@ func TestServiceErrorMessageIDUsesTypedResourceNotFound(t *testing.T) { err: groupNotFoundError(), want: "group.not_found", }, { - name: "update Group key missing Key", operation: "update_group_key", - err: keyNotFoundError(), want: "key.not_found", + name: "update Group credential missing Credential", operation: "update_group_credential", + err: credentialNotFoundError(), want: "credential.not_found", }, { - name: "delete wrong-owner Key", operation: "delete_group_key", - err: keyNotFoundError(), want: "key.not_found", + name: "delete wrong-owner Credential", operation: "delete_group_credential", + err: credentialNotFoundError(), want: "credential.not_found", }, } { t.Run(test.name, func(t *testing.T) { @@ -99,15 +99,15 @@ func TestLogServiceErrorUsesOnlyFixedOperationContext(t *testing.T) { logrus.SetFormatter(previousFormatter) }) - logServiceError("list_group_keys", err, app_errors.ErrInternalServer.Code) + logServiceError("list_group_credentials", err, app_errors.ErrInternalServer.Code) logText := logs.String() for _, required := range []string{ `"msg":"[CONTROL] Operation failed"`, - `"operation":"list_group_keys"`, + `"operation":"list_group_credentials"`, `"stage":"validate_db_registry_pair"`, `"mismatch_kind":"weight_manual"`, `"group_id":12`, - `"key_id":34`, + `"credential_id":34`, } { if !strings.Contains(logText, required) { t.Fatalf("log output missing %q: %s", required, logText) diff --git a/internal/control/operation_recovery.go b/internal/control/operation_recovery.go index 385ea03b..572e7add 100644 --- a/internal/control/operation_recovery.go +++ b/internal/control/operation_recovery.go @@ -259,7 +259,8 @@ func validateRecoverableOperation(operation *models.ControlOperation) error { func operationGroupID(operation *models.ControlOperation) (uint, error) { kind := operationKind(operation.OperationKind) - if kind != operationKindGroupCreate && kind != operationKindGroupKeyImport { + if kind != operationKindGroupCreate && + kind != operationKindCredentialImport { return 0, fmt.Errorf("operation %q has no group registry stage", kind) } return parseResourceIdentity(operation.ResourceIdentity, "group") @@ -270,7 +271,7 @@ func validateOperationResourceIdentity(kind operationKind, identity string) erro case operationKindAccessKeyCreate: _, err := parseResourceIdentity(identity, "access-key") return err - case operationKindGroupCreate, operationKindGroupKeyImport: + case operationKindGroupCreate, operationKindCredentialImport: _, err := parseResourceIdentity(identity, "group") return err default: diff --git a/internal/control/operation_recovery_test.go b/internal/control/operation_recovery_test.go index 393d1fd9..185e06b5 100644 --- a/internal/control/operation_recovery_test.go +++ b/internal/control/operation_recovery_test.go @@ -12,8 +12,8 @@ import ( "time" "gpt-load/internal/catalog" + "gpt-load/internal/channel" app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/storage/models" ) @@ -22,7 +22,7 @@ func TestOperationReplayRepairsPostCommitRegistryFailureWithoutRerunningMutation fixture := newServiceFixture(t) fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x31}, 16)) reconcileCalls := 0 - fixture.service.reconcileRegistryGroup = func(uint, []state.KeyEntry) (bool, error) { + fixture.service.reconcileRegistryGroup = func(uint, []state.CredentialEntry) (bool, error) { reconcileCalls++ return false, errors.New("injected registry failure") } @@ -65,7 +65,7 @@ func TestOperationReplayRepairsPostCommitRegistryFailureWithoutRerunningMutation t.Fatalf("replay/mutation = %t/%d, want true/1", replayed.Replayed, mutations) } groupID := mustResourceGroupID(t, replayed.ResourceIdentity) - if value, ok := fixture.registry.EncryptedValue(1); !ok || value != "cipher-one" { + if value, ok := fixture.registry.EncryptedCredentialData(1); !ok || value != "cipher-one" { t.Fatalf("recovered registry key = %q, %t", value, ok) } if current := fixture.manager.Current(); current == nil || current.Groups[groupID].ID != groupID { @@ -79,7 +79,7 @@ func TestOperationReplayOnlyRecordsStageWhenRegistrySideEffectAlreadySatisfied(t reconcileCalls := 0 fixture.service.reconcileRegistryGroup = func( groupID uint, - entries []state.KeyEntry, + entries []state.CredentialEntry, ) (bool, error) { reconcileCalls++ return fixture.registry.ReconcileGroup(groupID, entries) @@ -109,7 +109,7 @@ func TestOperationReplayOnlyRecordsStageWhenRegistrySideEffectAlreadySatisfied(t if reconcileCalls != 1 { t.Fatalf("initial reconcile calls = %d, want 1", reconcileCalls) } - if value, ok := fixture.registry.EncryptedValue(1); !ok || value != "cipher-one" { + if value, ok := fixture.registry.EncryptedCredentialData(1); !ok || value != "cipher-one" { t.Fatalf("side effect did not complete before stage failure: %q, %t", value, ok) } @@ -129,7 +129,7 @@ func TestOperationReplayOnlyRecordsStageWhenRegistrySideEffectAlreadySatisfied(t func TestOperationBarrierBlocksNewMutationBehindOlderRecovery(t *testing.T) { fixture := newServiceFixture(t) fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x33}, 32)) - fixture.service.reconcileRegistryGroup = func(uint, []state.KeyEntry) (bool, error) { + fixture.service.reconcileRegistryGroup = func(uint, []state.CredentialEntry) (bool, error) { return false, errors.New("registry remains unavailable") } firstMutations := 0 @@ -189,7 +189,7 @@ func TestOperationBarrierBlocksNewMutationBehindOlderRecovery(t *testing.T) { t.Fatalf("catalog apply published behind recovery barrier: before=%#v after=%#v", beforeCatalog, afterCatalog) } - fixture.service.reconcileRegistryGroup = func(uint, []state.KeyEntry) (bool, error) { + fixture.service.reconcileRegistryGroup = func(uint, []state.CredentialEntry) (bool, error) { return true, nil } recoveredCatalog := &catalog.Snapshot{Providers: map[string]catalog.Provider{ @@ -370,21 +370,21 @@ func newDurableGroupOperationInput( Mutate: func(tx operationTransaction) (idempotentMutationResult, error) { *mutations++ group := models.Group{ - Name: "group-" + idempotencyKey[:4], - UpstreamURL: "https://upstream.example.com", - Protocols: models.JSON(`["` + string(protocol.OpenAICompletions) + `"]`), - Models: models.JSON(`[]`), - Config: models.JSON(`{}`), - Enabled: true, + Name: "group-" + idempotencyKey[:4], + ChannelID: string(channel.OpenAI), + Params: models.JSON(`{}`), + Models: models.JSON(`[]`), + Overrides: models.JSON(`{}`), + Enabled: true, } if err := tx.Create(&group).Error; err != nil { return idempotentMutationResult{}, err } - key := models.UpstreamKey{ - GroupID: group.ID, KeyValue: "cipher-one", KeyHash: "hash-one", - Status: models.UpstreamKeyStatusActive, + credential := models.Credential{ + GroupID: group.ID, Data: "cipher-one", Fingerprint: "fingerprint-one", + Status: models.CredentialStatusActive, UpdatedAtMS: 1, } - if err := tx.Create(&key).Error; err != nil { + if err := tx.Create(&credential).Error; err != nil { return idempotentMutationResult{}, err } result := []byte(fmt.Sprintf(`{"group_id":%d}`, group.ID)) diff --git a/internal/control/price_identity.go b/internal/control/price_identity.go index b84a6976..e0ec59d8 100644 --- a/internal/control/price_identity.go +++ b/internal/control/price_identity.go @@ -3,25 +3,20 @@ package control import ( "fmt" + "gpt-load/internal/channel" "gpt-load/internal/pricing" ) -// PriceIdentityForModel returns the persisted global pricing identity. -func PriceIdentityForModel(modelID string) (pricing.Identity, error) { - identity := pricing.Identity{ModelID: modelID} +var modelPriceChannelRegistry = channel.NewRegistry() + +// PriceIdentityForChannelModel returns one exact persisted pricing identity. +func PriceIdentityForChannelModel(channelID, modelID string) (pricing.Identity, error) { + identity := pricing.Identity{ChannelID: channelID, ModelID: modelID} if _, err := pricing.NewTable([]pricing.Rule{{Identity: identity}}); err != nil { return pricing.Identity{}, fmt.Errorf("validate price identity: %w", err) } + if _, ok := modelPriceChannelRegistry.CatalogProviderID(channel.ID(channelID)); !ok { + return pricing.Identity{}, fmt.Errorf("validate price identity: unknown channel %q", channelID) + } return identity, nil } - -func normalizeProviderID(value *string) (*string, error) { - if value == nil { - return nil, nil - } - if _, err := pricing.ProviderScopeKey(*value); err != nil { - return nil, err - } - cloned := *value - return &cloned, nil -} diff --git a/internal/control/price_reconcile.go b/internal/control/price_reconcile.go index c84b002d..b183242d 100644 --- a/internal/control/price_reconcile.go +++ b/internal/control/price_reconcile.go @@ -26,7 +26,7 @@ type priceReferenceSnapshot struct { references map[pricing.Identity]referencedPrice } -// reconcileReferencedPrices materializes every exact Group/model identity that +// reconcileReferencedPrices materializes every exact channel/model identity that // is currently referenced. Existing rows are immutable under ordinary Group // writes; catalog refresh and broad cleanup are owned by the sync workflow. func reconcileReferencedPrices(tx *gorm.DB, snapshot *catalog.Snapshot) error { @@ -41,7 +41,10 @@ func reconcileReferencedPrices(tx *gorm.DB, snapshot *catalog.Snapshot) error { } existingIdentities := make(map[pricing.Identity]struct{}, len(existing)) for _, row := range existing { - identity := pricing.Identity{ModelID: row.ModelID} + identity, err := PriceIdentityForChannelModel(row.ChannelID, row.ModelID) + if err != nil { + return fmt.Errorf("validate persisted price identity: %w", app_errors.ErrInternalServer) + } if _, duplicate := existingIdentities[identity]; duplicate { return fmt.Errorf("duplicate persisted price identity: %w", app_errors.ErrInternalServer) } @@ -55,6 +58,9 @@ func reconcileReferencedPrices(tx *gorm.DB, snapshot *catalog.Snapshot) error { } } sort.Slice(ordered, func(left, right int) bool { + if ordered[left].ChannelID != ordered[right].ChannelID { + return ordered[left].ChannelID < ordered[right].ChannelID + } return ordered[left].ModelID < ordered[right].ModelID }) for _, identity := range ordered { @@ -89,17 +95,12 @@ func loadPriceReferenceSnapshot(tx *gorm.DB) (priceReferenceSnapshot, error) { func buildPriceReferenceSnapshot(groups []models.Group) (priceReferenceSnapshot, error) { references := make(map[pricing.Identity]referencedPrice) for _, group := range groups { - if group.ProviderID != nil { - if _, err := pricing.ProviderScopeKey(*group.ProviderID); err != nil { - return priceReferenceSnapshot{}, fmt.Errorf("validate persisted Group provider identity: %w", app_errors.ErrInternalServer) - } - } var groupModels []GroupModel if err := decodeGroupDiscoveryJSON(group.Models, &groupModels); err != nil { return priceReferenceSnapshot{}, fmt.Errorf("decode group price references: %w", app_errors.ErrInternalServer) } for _, model := range groupModels { - identity, err := PriceIdentityForModel(model.ID) + identity, err := PriceIdentityForChannelModel(group.ChannelID, model.ID) if err != nil { return priceReferenceSnapshot{}, fmt.Errorf("validate group price reference: %w", app_errors.ErrInternalServer) } @@ -120,14 +121,15 @@ func newReconciledModelPrice( snapshot *catalog.Snapshot, ) (models.ModelPrice, error) { row := models.ModelPrice{ - ModelID: reference.identity.ModelID, - IsManual: false, + ChannelID: reference.identity.ChannelID, + ModelID: reference.identity.ModelID, + IsManual: false, } - cost, _, ok := resolveAutomaticPrice(snapshot, reference.identity.ModelID) + match, ok := resolveAutomaticPriceForIdentity(snapshot, reference.identity) if !ok { return row, nil } - desired, err := automaticCatalogValues(cost) + desired, err := automaticCatalogValues(match.cost) if err != nil { return models.ModelPrice{}, err } diff --git a/internal/control/price_reconcile_test.go b/internal/control/price_reconcile_test.go index 5e0b8230..128c789b 100644 --- a/internal/control/price_reconcile_test.go +++ b/internal/control/price_reconcile_test.go @@ -6,25 +6,28 @@ import ( "gorm.io/gorm" "gpt-load/internal/catalog" + "gpt-load/internal/channel" "gpt-load/internal/pricing" "gpt-load/internal/storage/models" ) -func TestReconcileReferencedPricesUsesOneGlobalIdentity(t *testing.T) { +func TestReconcileReferencedPricesUsesChannelModelIdentity(t *testing.T) { fixture := newServiceFixture(t) - providerID := "openai" createPriceTestGroup(t, fixture.db, models.Group{ - Name: "provider", ProviderID: &providerID, UpstreamURL: "https://provider.example/v1", - Protocols: models.JSON(`["openai-completions"]`), Models: models.JSON(`[{"id":"shared","alias":"a"}]`), Config: models.JSON(`{}`), Enabled: true, + Name: "openai", ChannelID: string(channel.OpenAI), Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"shared","alias":"a"}]`), Overrides: models.JSON(`{}`), Enabled: true, }) createPriceTestGroup(t, fixture.db, models.Group{ - Name: "custom", UpstreamURL: "https://custom.example/v1", - Protocols: models.JSON(`["openai-completions"]`), Models: models.JSON(`[{"id":"shared","alias":"b"}]`), Config: models.JSON(`{}`), Enabled: true, + Name: "anthropic", ChannelID: string(channel.Anthropic), Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"shared","alias":"b"}]`), Overrides: models.JSON(`{}`), Enabled: true, }) snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ "openai": {ID: "openai", Models: map[string]catalog.Model{ "shared": {ID: "shared", Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(2_000_000_000)}}}, }}, + "anthropic": {ID: "anthropic", Models: map[string]catalog.Model{ + "shared": {ID: "shared", Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(3_000_000_000)}}}, + }}, }} if err := fixture.service.withControlTransaction(t.Context(), func(tx *gorm.DB) error { return reconcileReferencedPrices(tx, snapshot) @@ -35,16 +38,40 @@ func TestReconcileReferencedPricesUsesOneGlobalIdentity(t *testing.T) { if err := fixture.db.Find(&rows).Error; err != nil { t.Fatal(err) } - if len(rows) != 1 || rows[0].ModelID != "shared" || rows[0].InputPriceNanoUSDPerMillionTokens == nil || *rows[0].InputPriceNanoUSDPerMillionTokens != 2_000_000_000 { + if len(rows) != 2 { t.Fatalf("reconciled rows = %#v", rows) } - if _, ok := mustLoadPriceTable(t, fixture.db).Lookup(pricing.Identity{ModelID: "shared"}); !ok { - t.Fatal("global runtime table did not contain shared model") + byChannel := make(map[string]models.ModelPrice, len(rows)) + for _, row := range rows { + byChannel[row.ChannelID] = row + } + if row := byChannel[string(channel.OpenAI)]; row.InputPriceNanoUSDPerMillionTokens == nil || + *row.InputPriceNanoUSDPerMillionTokens != 2_000_000_000 { + t.Fatalf("OpenAI price = %#v", row) + } + if row := byChannel[string(channel.Anthropic)]; row.InputPriceNanoUSDPerMillionTokens == nil || + *row.InputPriceNanoUSDPerMillionTokens != 3_000_000_000 { + t.Fatalf("Anthropic price = %#v", row) + } + table := mustLoadPriceTable(t, fixture.db) + for _, identity := range []pricing.Identity{ + {ChannelID: string(channel.OpenAI), ModelID: "shared"}, + {ChannelID: string(channel.Anthropic), ModelID: "shared"}, + } { + if _, ok := table.Lookup(identity); !ok { + t.Fatalf("runtime table does not contain %#v", identity) + } } } func createPriceTestGroup(t *testing.T, db *gorm.DB, group models.Group) models.Group { t.Helper() + if group.ChannelID == "" { + group.ChannelID = string(channel.OpenAICompatible) + } + if len(group.Params) == 0 { + group.Params = models.JSON(`{"base_url":"https://` + group.Name + `.example/v1"}`) + } if err := db.Create(&group).Error; err != nil { t.Fatalf("create price test group: %v", err) } diff --git a/internal/control/price_reference.go b/internal/control/price_reference.go index 2ca53055..ee17b23c 100644 --- a/internal/control/price_reference.go +++ b/internal/control/price_reference.go @@ -4,38 +4,82 @@ import ( "sort" "gpt-load/internal/catalog" + "gpt-load/internal/channel" + "gpt-load/internal/pricing" ) -// resolveAutomaticPrice selects a global exact catalog model price without -// changing the supplied snapshot. Routing providers never influence the -// result: the fixed catalog priority wins, followed by provider ID. -func resolveAutomaticPrice( - snapshot *catalog.Snapshot, - modelID string, -) (cost *catalog.ModelCost, matchedProviderID string, ok bool) { - if snapshot == nil || len(snapshot.Providers) == 0 || modelID == "" { - return nil, "", false - } - for _, providerID := range catalogProviderLookupOrder(snapshot, "") { - provider, exists := snapshot.Providers[providerID] - if !exists { - continue - } - model, exists := provider.Models[modelID] - if !exists || model.Cost == nil { - continue - } - return model.Cost, providerID, true - } - return nil, "", false +type ModelPriceMatchSource string + +const ( + ModelPriceMatchSourceChannelCatalogProvider ModelPriceMatchSource = "channel_catalog_provider" + ModelPriceMatchSourceProviderPriorityFallback ModelPriceMatchSource = "provider_priority_fallback" +) + +type automaticPriceMatch struct { + cost *catalog.ModelCost + providerID string + source ModelPriceMatchSource } -func catalogProviderLookupOrder(snapshot *catalog.Snapshot, preferredProviderID string) []string { +func resolveAutomaticPriceForIdentity( + snapshot *catalog.Snapshot, + identity pricing.Identity, +) (automaticPriceMatch, bool) { + model, providerID, source, ok := resolveCatalogModelForIdentity(snapshot, identity, true) + if !ok { + return automaticPriceMatch{}, false + } + return automaticPriceMatch{ + cost: model.Cost, providerID: providerID, source: source, + }, true +} + +func resolveCatalogModelForIdentity( + snapshot *catalog.Snapshot, + identity pricing.Identity, + requirePrice bool, +) (catalog.Model, string, ModelPriceMatchSource, bool) { + if snapshot == nil || len(snapshot.Providers) == 0 || identity.ModelID == "" { + return catalog.Model{}, "", "", false + } + exactProviderID, known := modelPriceChannelRegistry.CatalogProviderID(channel.ID(identity.ChannelID)) + if !known { + return catalog.Model{}, "", "", false + } + lookup := func(providerID string) (catalog.Model, bool) { + provider, exists := snapshot.Providers[providerID] + if !exists { + return catalog.Model{}, false + } + model, exists := provider.Models[identity.ModelID] + if !exists || (requirePrice && model.Cost == nil) { + return catalog.Model{}, false + } + return model, true + } + if exactProviderID != "" { + model, ok := lookup(exactProviderID) + if !ok { + return catalog.Model{}, "", "", false + } + return model, exactProviderID, ModelPriceMatchSourceChannelCatalogProvider, true + } + for _, providerID := range catalogProviderLookupOrder(snapshot) { + model, ok := lookup(providerID) + if !ok { + continue + } + return model, providerID, ModelPriceMatchSourceProviderPriorityFallback, true + } + return catalog.Model{}, "", "", false +} + +func catalogProviderLookupOrder(snapshot *catalog.Snapshot) []string { if snapshot == nil || len(snapshot.Providers) == 0 { return nil } - result := make([]string, 0, len(snapshot.Providers)+1) - seen := make(map[string]struct{}, len(snapshot.Providers)+1) + result := make([]string, 0, len(snapshot.Providers)) + seen := make(map[string]struct{}, len(snapshot.Providers)) appendProvider := func(providerID string) { if providerID == "" { return @@ -46,7 +90,6 @@ func catalogProviderLookupOrder(snapshot *catalog.Snapshot, preferredProviderID seen[providerID] = struct{}{} result = append(result, providerID) } - appendProvider(preferredProviderID) priority := catalog.AutomaticPriceProviderPriority() for _, providerID := range priority { appendProvider(providerID) diff --git a/internal/control/price_reference_test.go b/internal/control/price_reference_test.go index ef856923..36be757d 100644 --- a/internal/control/price_reference_test.go +++ b/internal/control/price_reference_test.go @@ -5,10 +5,75 @@ import ( "testing" "gpt-load/internal/catalog" + "gpt-load/internal/channel" "gpt-load/internal/pricing" ) -func TestResolveAutomaticPriceUsesGlobalPriorityNotRouteProvider(t *testing.T) { +func TestPriceIdentityForChannelModelRequiresRegisteredChannel(t *testing.T) { + identity, err := PriceIdentityForChannelModel(string(channel.OpenAICompatible), "shared") + if err != nil || identity != (pricing.Identity{ + ChannelID: string(channel.OpenAICompatible), ModelID: "shared", + }) { + t.Fatalf("PriceIdentityForChannelModel() = %#v, %v", identity, err) + } + if _, err := PriceIdentityForChannelModel("unknown", "shared"); err == nil { + t.Fatal("PriceIdentityForChannelModel() accepted an unregistered channel") + } +} + +func TestResolveAutomaticPriceForIdentityUsesOfficialChannelCatalogProvider(t *testing.T) { + openAICost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}} + anthropicCost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(2)}} + snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ + "openai": {ID: "openai", Models: map[string]catalog.Model{ + "shared": {ID: "shared", Cost: openAICost}, + }}, + "anthropic": {ID: "anthropic", Models: map[string]catalog.Model{ + "shared": {ID: "shared", Cost: anthropicCost}, + }}, + }} + + match, ok := resolveAutomaticPriceForIdentity(snapshot, pricing.Identity{ + ChannelID: string(channel.Anthropic), ModelID: "shared", + }) + if !ok || match.cost != anthropicCost || match.providerID != "anthropic" || + match.source != ModelPriceMatchSourceChannelCatalogProvider { + t.Fatalf("official channel match = %#v, %t", match, ok) + } +} + +func TestResolveAutomaticPriceForIdentityDoesNotFallbackForOfficialChannel(t *testing.T) { + snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ + "openai": {ID: "openai", Models: map[string]catalog.Model{ + "shared": {ID: "shared", Cost: &catalog.ModelCost{}}, + }}, + }} + + if match, ok := resolveAutomaticPriceForIdentity(snapshot, pricing.Identity{ + ChannelID: string(channel.Anthropic), ModelID: "shared", + }); ok { + t.Fatalf("official channel unexpectedly fell back: %#v", match) + } +} + +func TestResolveAutomaticPriceForIdentityFallsBackForCompatibleChannel(t *testing.T) { + openAICost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}} + snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ + "openai": {ID: "openai", Models: map[string]catalog.Model{ + "shared": {ID: "shared", Cost: openAICost}, + }}, + }} + + match, ok := resolveAutomaticPriceForIdentity(snapshot, pricing.Identity{ + ChannelID: string(channel.OpenAICompatible), ModelID: "shared", + }) + if !ok || match.cost != openAICost || match.providerID != "openai" || + match.source != ModelPriceMatchSourceProviderPriorityFallback { + t.Fatalf("compatible channel match = %#v, %t", match, ok) + } +} + +func TestCompatibleAutomaticPriceUsesGlobalPriority(t *testing.T) { openAICost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}} alphaCost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(2)}} snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ @@ -20,23 +85,29 @@ func TestResolveAutomaticPriceUsesGlobalPriorityNotRouteProvider(t *testing.T) { for id, provider := range snapshot.Providers { want.Providers[id] = provider } - cost, providerID, ok := resolveAutomaticPrice(snapshot, "shared") - if !ok || cost != openAICost || providerID != "openai" { - t.Fatalf("global price = (%p, %q, %t), want openai", cost, providerID, ok) + match, ok := resolveAutomaticPriceForIdentity(snapshot, pricing.Identity{ + ChannelID: string(channel.OpenAICompatible), ModelID: "shared", + }) + if !ok || match.cost != openAICost || match.providerID != "openai" || + match.source != ModelPriceMatchSourceProviderPriorityFallback { + t.Fatalf("compatible price = (%p, %q, %q, %t), want openai fallback", match.cost, match.providerID, match.source, ok) } if !reflect.DeepEqual(snapshot, &want) { t.Fatal("global price resolution mutated catalog snapshot") } } -func TestResolveAutomaticPriceFallsBackToStableProviderID(t *testing.T) { +func TestCompatibleAutomaticPriceFallsBackToStableProviderID(t *testing.T) { alphaCost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}} snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ "zeta": {ID: "zeta", Models: map[string]catalog.Model{"shared": {ID: "shared", Cost: &catalog.ModelCost{}}}}, "alpha": {ID: "alpha", Models: map[string]catalog.Model{"shared": {ID: "shared", Cost: alphaCost}}}, }} - cost, providerID, ok := resolveAutomaticPrice(snapshot, "shared") - if !ok || cost != alphaCost || providerID != "alpha" { - t.Fatalf("global fallback = (%p, %q, %t)", cost, providerID, ok) + match, ok := resolveAutomaticPriceForIdentity(snapshot, pricing.Identity{ + ChannelID: string(channel.OpenAICompatible), ModelID: "shared", + }) + if !ok || match.cost != alphaCost || match.providerID != "alpha" || + match.source != ModelPriceMatchSourceProviderPriorityFallback { + t.Fatalf("compatible fallback = (%p, %q, %q, %t)", match.cost, match.providerID, match.source, ok) } } diff --git a/internal/control/pricing_runtime.go b/internal/control/pricing_runtime.go index fcb61cec..70c64545 100644 --- a/internal/control/pricing_runtime.go +++ b/internal/control/pricing_runtime.go @@ -57,8 +57,12 @@ func loadPriceTable(ctx context.Context, tx *gorm.DB) (*pricing.Table, error) { } func persistedPriceRule(row models.ModelPrice) (pricing.Rule, error) { + identity, err := PriceIdentityForChannelModel(row.ChannelID, row.ModelID) + if err != nil { + return pricing.Rule{}, err + } rule := pricing.Rule{ - Identity: pricing.Identity{ModelID: row.ModelID}, + Identity: identity, Prices: pricing.Prices{ Input: priceFromStoragePointer(row.InputPriceNanoUSDPerMillionTokens), Output: priceFromStoragePointer(row.OutputPriceNanoUSDPerMillionTokens), diff --git a/internal/control/pricing_status.go b/internal/control/pricing_status.go index 7c6ffdee..edbf1119 100644 --- a/internal/control/pricing_status.go +++ b/internal/control/pricing_status.go @@ -4,6 +4,7 @@ import ( "strings" "gpt-load/internal/catalog" + "gpt-load/internal/pricing" "gpt-load/internal/storage/models" ) @@ -35,21 +36,24 @@ func modelPriceHasConfiguredValue(row models.ModelPrice) bool { func resolveCandidatePricing( row *models.ModelPrice, snapshot *catalog.Snapshot, - modelID string, + identity pricing.Identity, ) (PricingStatus, *string) { if row != nil { status := resolvePricingStatus(row) if status != PricingStatusConfigured || !modelPriceHasConfiguredValue(*row) || row.IsManual { return status, nil } - if _, providerID, ok := resolveAutomaticPrice(snapshot, modelID); ok { - source := pricingSourceName(snapshot, providerID, providerID) + if row.ChannelID != identity.ChannelID || row.ModelID != identity.ModelID { + return status, nil + } + if match, ok := resolveAutomaticPriceForIdentity(snapshot, identity); ok { + source := pricingSourceName(snapshot, match.providerID, match.providerID) return status, &source } return status, nil } - if _, providerID, ok := resolveAutomaticPrice(snapshot, modelID); ok { - source := pricingSourceName(snapshot, providerID, providerID) + if match, ok := resolveAutomaticPriceForIdentity(snapshot, identity); ok { + source := pricingSourceName(snapshot, match.providerID, match.providerID) return PricingStatusConfigured, &source } return PricingStatusPending, nil diff --git a/internal/control/pricing_status_test.go b/internal/control/pricing_status_test.go index fcd5e587..3c87cd22 100644 --- a/internal/control/pricing_status_test.go +++ b/internal/control/pricing_status_test.go @@ -4,6 +4,7 @@ import ( "testing" "gpt-load/internal/catalog" + "gpt-load/internal/channel" "gpt-load/internal/pricing" "gpt-load/internal/storage/models" ) @@ -24,9 +25,13 @@ func TestResolveCandidatePricingUsesCurrentAutomaticPriceMatch(t *testing.T) { }, }} status, source := resolveCandidatePricing( - &models.ModelPrice{InputPriceNanoUSDPerMillionTokens: &price}, + &models.ModelPrice{ + ChannelID: string(channel.OpenAI), + ModelID: "gpt-4o", + InputPriceNanoUSDPerMillionTokens: &price, + }, snapshot, - "gpt-4o", + pricing.Identity{ChannelID: string(channel.OpenAI), ModelID: "gpt-4o"}, ) if status != PricingStatusConfigured || source == nil || *source != "OpenAI" { t.Fatalf("pricing = %q, %v, want configured and automatic Provider match", status, source) @@ -35,9 +40,9 @@ func TestResolveCandidatePricingUsesCurrentAutomaticPriceMatch(t *testing.T) { func TestResolveCandidatePricingHidesManuallyUnpricedSource(t *testing.T) { status, source := resolveCandidatePricing( - &models.ModelPrice{IsManual: true}, + &models.ModelPrice{ChannelID: string(channel.OpenAI), ModelID: "gpt-4o", IsManual: true}, nil, - "gpt-4o", + pricing.Identity{ChannelID: string(channel.OpenAI), ModelID: "gpt-4o"}, ) if status != PricingStatusConfigured || source != nil { t.Fatalf("pricing = %q, %v, want configured without source", status, source) diff --git a/internal/control/project_model_collection.go b/internal/control/project_model_collection.go index 7bfcc565..648e2758 100644 --- a/internal/control/project_model_collection.go +++ b/internal/control/project_model_collection.go @@ -10,6 +10,7 @@ import ( "gorm.io/gorm" "gpt-load/internal/catalog" + "gpt-load/internal/channel" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/pricing" "gpt-load/internal/protocol" @@ -18,11 +19,12 @@ import ( ) type ProjectModelGroupDTO struct { - ID uint `json:"id"` - Name string `json:"name"` - ProviderID *string `json:"provider_id"` - Enabled bool `json:"enabled"` - Protocols []protocol.Protocol `json:"protocols"` + ID uint `json:"id"` + Name string `json:"name"` + ChannelID channel.ID `json:"channel_id"` + Params json.RawMessage `json:"params"` + Enabled bool `json:"enabled"` + ClientProtocols []protocol.Protocol `json:"client_protocols"` } type ProjectModelCapabilitiesDTO struct { @@ -147,7 +149,7 @@ type projectModelUpstreamRecord struct { type projectModelRecord struct { clientModel string protocols []protocol.Protocol - upstreams map[string]*projectModelUpstreamRecord + upstreams map[pricing.Identity]*projectModelUpstreamRecord } func (s *Service) ListProjectModels(ctx context.Context, query ProjectModelListQuery) (ProjectModelListResponse, error) { @@ -178,14 +180,16 @@ func (s *Service) ListProjectModels(ctx context.Context, query ProjectModelListQ groups = cloneGroupRows(groups) return nil }) + protocolOverrides := map[uint][]protocol.Protocol(nil) if err == nil && query.AccessKeyID != nil { if s.manager == nil { err = app_errors.ErrInternalServer } else { - groups, err = scopeProjectModelGroups( + groups, protocolOverrides, err = scopeProjectModelGroups( groups, s.manager.Current(), *query.AccessKeyID, + s.channelRegistry, ) } } @@ -201,13 +205,16 @@ func (s *Service) ListProjectModels(ctx context.Context, query ProjectModelListQ if err != nil { return ProjectModelListResponse{}, err } - groupRecords, groupDTOs, err := projectModelGroups(groups) + groupRecords, groupDTOs, err := projectModelGroups(groups, s.channelRegistry, protocolOverrides) if err != nil { return ProjectModelListResponse{}, err } priceRecords := make(map[pricing.Identity]modelPriceListRecord, len(prices)) for _, row := range prices { - identity := pricing.Identity{ModelID: row.ModelID} + identity, err := PriceIdentityForChannelModel(row.ChannelID, row.ModelID) + if err != nil { + return ProjectModelListResponse{}, fmt.Errorf("validate persisted price identity: %w", app_errors.ErrInternalServer) + } if query.AccessKeyID != nil { if _, referenced := references.references[identity]; !referenced { continue @@ -232,7 +239,7 @@ func (s *Service) ListProjectModels(ctx context.Context, query ProjectModelListQ if err := ctx.Err(); err != nil { return ProjectModelListResponse{}, err } - identity, err := PriceIdentityForModel(model.ID) + identity, err := PriceIdentityForChannelModel(group.row.ChannelID, model.ID) if err != nil { return ProjectModelListResponse{}, fmt.Errorf("validate group %d model %q: %w", group.row.ID, model.ID, app_errors.ErrInternalServer) } @@ -249,12 +256,12 @@ func (s *Service) ListProjectModels(ctx context.Context, query ProjectModelListQ root = &projectModelRecord{ clientModel: clientModel, protocols: []protocol.Protocol{}, - upstreams: make(map[string]*projectModelUpstreamRecord), + upstreams: make(map[pricing.Identity]*projectModelUpstreamRecord), } records[clientModel] = root } - root.protocols = mergeProjectModelProtocols(root.protocols, group.dto.Protocols) - upstream := root.upstreams[model.ID] + root.protocols = mergeProjectModelProtocols(root.protocols, group.dto.ClientProtocols) + upstream := root.upstreams[identity] if upstream == nil { affectedGroups, err := projectModelAffectedGroups(identity, references, groupDTOs) if err != nil { @@ -263,10 +270,10 @@ func (s *Service) ListProjectModels(ctx context.Context, query ProjectModelListQ upstream = &projectModelUpstreamRecord{ modelID: model.ID, aliasApplied: strings.TrimSpace(model.Alias) != "", price: priceRecord.dto, affectedGroups: affectedGroups, - catalogReference: projectModelCatalogReference(priceRecord.dto, model.ID, catalogSnapshot), + catalogReference: projectModelCatalogReference(priceRecord.dto, identity, catalogSnapshot), groupSeen: make(map[uint]struct{}), } - root.upstreams[model.ID] = upstream + root.upstreams[identity] = upstream } if _, duplicate := upstream.groupSeen[group.row.ID]; !duplicate { upstream.groupSeen[group.row.ID] = struct{}{} @@ -297,7 +304,10 @@ func (s *Service) ListProjectModels(ctx context.Context, query ProjectModelListQ dto.UpstreamModels = append(dto.UpstreamModels, upstreamDTO) } sort.SliceStable(dto.UpstreamModels, func(left, right int) bool { - return dto.UpstreamModels[left].ModelID < dto.UpstreamModels[right].ModelID + if dto.UpstreamModels[left].ModelID != dto.UpstreamModels[right].ModelID { + return dto.UpstreamModels[left].ModelID < dto.UpstreamModels[right].ModelID + } + return dto.UpstreamModels[left].Price.ChannelID < dto.UpstreamModels[right].Price.ChannelID }) if len(dto.UpstreamModels) == 0 || !projectModelSearchMatch(dto, query.Search) { continue @@ -337,15 +347,17 @@ func scopeProjectModelGroups( groups []models.Group, snapshot *state.ConfigSnapshot, accessKeyID uint, -) ([]models.Group, error) { + registry *channel.Registry, +) ([]models.Group, map[uint][]protocol.Protocol, error) { if snapshot == nil || accessKeyID == 0 { - return nil, app_errors.ErrUnauthorized + return nil, nil, app_errors.ErrUnauthorized } accessKey, exists := snapshot.AccessKeysByID[accessKeyID] if !exists || accessKey.Status != state.AccessKeyStatusActive { - return nil, app_errors.ErrUnauthorized + return nil, nil, app_errors.ErrUnauthorized } result := make([]models.Group, 0, len(groups)) + protocolsByGroup := make(map[uint][]protocol.Protocol, len(groups)) for _, group := range groups { if !group.Enabled { continue @@ -355,9 +367,9 @@ func scopeProjectModelGroups( continue } } - protocols, err := projectModelGroupProtocols(group) + protocols, err := projectModelGroupProtocols(registry, group) if err != nil { - return nil, err + return nil, nil, err } filteredProtocols := make([]protocol.Protocol, 0, len(protocols)) for _, value := range protocols { @@ -373,7 +385,7 @@ func scopeProjectModelGroups( } var groupModels []GroupModel if err := decodeGroupDiscoveryJSON(group.Models, &groupModels); err != nil { - return nil, fmt.Errorf( + return nil, nil, fmt.Errorf( "decode group %d models for access scope: %w", group.ID, app_errors.ErrInternalServer, @@ -395,19 +407,15 @@ func scopeProjectModelGroups( if len(filteredModels) == 0 { continue } - protocolJSON, err := json.Marshal(filteredProtocols) - if err != nil { - return nil, fmt.Errorf("encode scoped group protocols: %w", err) - } modelJSON, err := json.Marshal(filteredModels) if err != nil { - return nil, fmt.Errorf("encode scoped group models: %w", err) + return nil, nil, fmt.Errorf("encode scoped group models: %w", err) } - group.Protocols = models.JSON(protocolJSON) group.Models = models.JSON(modelJSON) + protocolsByGroup[group.ID] = append([]protocol.Protocol(nil), filteredProtocols...) result = append(result, group) } - return result, nil + return result, protocolsByGroup, nil } func (s *Service) GetUpstreamModelDetail(ctx context.Context, priceID uint) (UpstreamModelDetailDTO, error) { @@ -446,14 +454,21 @@ func (s *Service) GetUpstreamModelDetail(ctx context.Context, priceID uint) (Ups if err != nil { return UpstreamModelDetailDTO{}, err } - groupRecords, _, err := projectModelGroups(groups) + groupRecords, _, err := projectModelGroups(groups, s.channelRegistry, nil) if err != nil { return UpstreamModelDetailDTO{}, err } associations := make([]UpstreamModelAssociationDTO, 0) clientModels := make(map[string]struct{}) groupIDs := make(map[uint]struct{}) + identity, err := PriceIdentityForChannelModel(row.ChannelID, row.ModelID) + if err != nil { + return UpstreamModelDetailDTO{}, fmt.Errorf("validate model detail price identity: %w", app_errors.ErrInternalServer) + } for _, group := range groupRecords { + if group.row.ChannelID != row.ChannelID { + continue + } for _, model := range group.models { if model.ID != row.ModelID { continue @@ -481,26 +496,45 @@ func (s *Service) GetUpstreamModelDetail(ctx context.Context, priceID uint) (Ups }) return UpstreamModelDetailDTO{ ModelID: row.ModelID, Price: price.dto, - CatalogReference: projectModelCatalogReference(price.dto, row.ModelID, catalogSnapshot), + CatalogReference: projectModelCatalogReference(price.dto, identity, catalogSnapshot), Associations: associations, ClientModelCount: len(clientModels), GroupCount: len(groupIDs), }, nil } -func projectModelGroups(groups []models.Group) ([]projectModelGroupRecord, map[uint]ProjectModelGroupDTO, error) { +func projectModelGroups( + groups []models.Group, + registry *channel.Registry, + protocolOverrides map[uint][]protocol.Protocol, +) ([]projectModelGroupRecord, map[uint]ProjectModelGroupDTO, error) { records := make([]projectModelGroupRecord, 0, len(groups)) dtos := make(map[uint]ProjectModelGroupDTO, len(groups)) for _, group := range groups { - protocols, err := projectModelGroupProtocols(group) + protocols, overridden := protocolOverrides[group.ID] + var err error + if !overridden { + protocols, err = projectModelGroupProtocols(registry, group) + } if err != nil { return nil, nil, err } + if len(protocols) == 0 { + return nil, nil, fmt.Errorf("group %d has no client protocols: %w", group.ID, app_errors.ErrInternalServer) + } + if registry == nil { + return nil, nil, app_errors.ErrInternalServer + } + params, err := registry.ValidateParams(channel.ID(group.ChannelID), json.RawMessage(group.Params)) + if err != nil { + return nil, nil, fmt.Errorf("validate group %d channel params: %w", group.ID, app_errors.ErrInternalServer) + } var groupModels []GroupModel if err := decodeGroupDiscoveryJSON(group.Models, &groupModels); err != nil { return nil, nil, fmt.Errorf("decode group %d models: %w", group.ID, app_errors.ErrInternalServer) } dto := ProjectModelGroupDTO{ - ID: group.ID, Name: group.Name, ProviderID: cloneString(group.ProviderID), Enabled: group.Enabled, - Protocols: append([]protocol.Protocol(nil), protocols...), + ID: group.ID, Name: group.Name, ChannelID: channel.ID(group.ChannelID), + Params: params.CanonicalJSON(), Enabled: group.Enabled, + ClientProtocols: append([]protocol.Protocol(nil), protocols...), } dtos[group.ID] = dto records = append(records, projectModelGroupRecord{row: group, dto: dto, models: groupModels}) @@ -515,7 +549,12 @@ func projectModelAffectedGroups( ) ([]ProjectModelGroupDTO, error) { reference, exists := references.references[identity] if !exists || len(reference.groupIDs) == 0 { - return nil, fmt.Errorf("missing model price references for %s: %w", identity.ModelID, app_errors.ErrInternalServer) + return nil, fmt.Errorf( + "missing model price references for %s/%s: %w", + identity.ChannelID, + identity.ModelID, + app_errors.ErrInternalServer, + ) } result := make([]ProjectModelGroupDTO, 0, len(reference.groupIDs)) for groupID := range reference.groupIDs { @@ -563,12 +602,18 @@ func projectModelPricingVisible(status PricingStatus, filter ProjectModelPricing return filter == ProjectModelPricingStatusAll || (filter == ProjectModelPricingStatusPending && status == PricingStatusPending) || (filter == ProjectModelPricingStatusConfigured && status == PricingStatusConfigured) } -func projectModelGroupProtocols(group models.Group) ([]protocol.Protocol, error) { - var values []protocol.Protocol - if err := decodeGroupDiscoveryJSON(group.Protocols, &values); err != nil { - return nil, fmt.Errorf("decode group %d protocols: %w", group.ID, app_errors.ErrInternalServer) +func projectModelGroupProtocols(registry *channel.Registry, group models.Group) ([]protocol.Protocol, error) { + if registry == nil || group.ChannelID == "" { + return nil, fmt.Errorf("group %d has no channel: %w", group.ID, app_errors.ErrInternalServer) } - result := mergeProjectModelProtocols(nil, values) + descriptor, ok := registry.Get(channel.ID(group.ChannelID)) + if !ok { + return nil, fmt.Errorf("group %d has unknown channel: %w", group.ID, app_errors.ErrInternalServer) + } + if _, err := registry.ValidateParams(channel.ID(group.ChannelID), json.RawMessage(group.Params)); err != nil { + return nil, fmt.Errorf("group %d has invalid channel params: %w", group.ID, app_errors.ErrInternalServer) + } + result := mergeProjectModelProtocols(nil, descriptor.ClientProtocols) if len(result) == 0 { return nil, fmt.Errorf("group %d has no valid protocols: %w", group.ID, app_errors.ErrInternalServer) } @@ -594,48 +639,42 @@ func mergeProjectModelProtocols(left, right []protocol.Protocol) []protocol.Prot func projectModelCatalogReference( price ModelPriceDTO, - modelID string, + identity pricing.Identity, snapshot *catalog.Snapshot, ) *ProjectModelCatalogReferenceDTO { if snapshot == nil { return nil } - lookup := func(providerID, source string) *ProjectModelCatalogReferenceDTO { - provider, exists := snapshot.Providers[providerID] - if !exists { - return nil - } - model, exists := provider.Models[modelID] - if !exists { - return nil - } - providerName := strings.TrimSpace(provider.Name) - if providerName == "" { - providerName = providerID - } - return &ProjectModelCatalogReferenceDTO{ - Source: source, ProviderID: providerID, ProviderName: providerName, - Model: projectModelCatalogModel(model), - } - } - - tested := make(map[string]struct{}) + var model catalog.Model + providerID := "" if price.MatchedProviderID != nil { - providerID := *price.MatchedProviderID - tested[providerID] = struct{}{} - if reference := lookup(providerID, "reference_provider"); reference != nil { - return reference + providerID = *price.MatchedProviderID + provider, exists := snapshot.Providers[providerID] + if exists { + model, exists = provider.Models[identity.ModelID] + if !exists { + providerID = "" + } + } else { + providerID = "" } } - for _, providerID := range catalogProviderLookupOrder(snapshot, "") { - if _, alreadyTested := tested[providerID]; alreadyTested { - continue - } - if reference := lookup(providerID, "reference_provider"); reference != nil { - return reference + if providerID == "" { + var ok bool + model, providerID, _, ok = resolveCatalogModelForIdentity(snapshot, identity, false) + if !ok { + return nil } } - return nil + provider := snapshot.Providers[providerID] + providerName := strings.TrimSpace(provider.Name) + if providerName == "" { + providerName = providerID + } + return &ProjectModelCatalogReferenceDTO{ + Source: "reference_provider", ProviderID: providerID, ProviderName: providerName, + Model: projectModelCatalogModel(model), + } } func projectModelCatalogModel(model catalog.Model) ProjectModelCatalogDTO { @@ -697,6 +736,9 @@ func projectModelSearchMatch(record ProjectModelDTO, search string) bool { if strings.Contains(strings.ToLower(upstream.ModelID), search) { return true } + if strings.Contains(strings.ToLower(upstream.Price.ChannelID), search) { + return true + } if reference := upstream.CatalogReference; reference != nil { if strings.Contains(strings.ToLower(reference.Model.Name), search) || strings.Contains(strings.ToLower(reference.ProviderID), search) || diff --git a/internal/control/project_model_collection_test.go b/internal/control/project_model_collection_test.go index e1d0dd6f..96ea4b83 100644 --- a/internal/control/project_model_collection_test.go +++ b/internal/control/project_model_collection_test.go @@ -8,20 +8,23 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/catalog" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" + "gpt-load/internal/pricing" "gpt-load/internal/protocol" "gpt-load/internal/storage/models" ) -func TestProjectModelsUseOnePricePerUpstreamAndDetailShowsAllRelationships(t *testing.T) { +func TestProjectModelsSeparateSameUpstreamModelByChannelAndDetail(t *testing.T) { fixture := newServiceFixture(t) createPriceTestGroup(t, fixture.db, models.Group{ - Name: "enabled", UpstreamURL: "https://enabled.example/v1", Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[{"id":"upstream-x","alias":"client-a"},{"id":"upstream-x","alias":"client-b"}]`), Config: models.JSON(`{}`), Enabled: true, + Name: "enabled", ChannelID: string(channel.OpenAICompatible), Params: models.JSON(`{"base_url":"https://enabled.example/v1"}`), + Models: models.JSON(`[{"id":"upstream-x","alias":"client-a"},{"id":"upstream-x","alias":"client-b"}]`), Overrides: models.JSON(`{}`), Enabled: true, }) createPriceTestGroup(t, fixture.db, models.Group{ - Name: "disabled", UpstreamURL: "https://disabled.example/v1", Protocols: models.JSON(`["anthropic"]`), - Models: models.JSON(`[{"id":"upstream-x","alias":"client-a"}]`), Config: models.JSON(`{}`), Enabled: false, + Name: "disabled", ChannelID: string(channel.AnthropicCompatible), Params: models.JSON(`{"base_url":"https://disabled.example/v1"}`), + Models: models.JSON(`[{"id":"upstream-x","alias":"client-a"}]`), Overrides: models.JSON(`{}`), Enabled: false, }) mustEnsureInitialPrices(t, fixture) all, err := fixture.service.ListProjectModels(t.Context(), ProjectModelListQuery{ @@ -33,55 +36,75 @@ func TestProjectModelsUseOnePricePerUpstreamAndDetailShowsAllRelationships(t *te if len(all.Items) != 2 || all.Items[0].ClientModel != "client-a" || all.Items[1].ClientModel != "client-b" { t.Fatalf("client model overview = %#v", all.Items) } - upstream := all.Items[0].UpstreamModels[0] - if upstream.ModelID != "upstream-x" || upstream.Price.ModelID != "upstream-x" || len(upstream.RouteGroups) != 2 || len(upstream.AffectedGroups) != 2 { - t.Fatalf("upstream overview = %#v", upstream) + if len(all.Items[0].UpstreamModels) != 2 || len(all.Items[1].UpstreamModels) != 1 { + t.Fatalf("channel-scoped upstream overview = %#v", all.Items) + } + upstream := all.Items[1].UpstreamModels[0] + if upstream.ModelID != "upstream-x" || upstream.Price.ModelID != "upstream-x" || + upstream.Price.ChannelID != "openai_compatible" || + len(upstream.RouteGroups) != 1 || len(upstream.AffectedGroups) != 1 { + t.Fatalf("OpenAI-compatible upstream overview = %#v", upstream) } detail, err := fixture.service.GetUpstreamModelDetail(t.Context(), upstream.Price.ID) if err != nil { t.Fatal(err) } - if detail.ModelID != "upstream-x" || detail.ClientModelCount != 2 || detail.GroupCount != 2 || len(detail.Associations) != 3 { + if detail.ModelID != "upstream-x" || detail.Price.ChannelID != "openai_compatible" || + detail.ClientModelCount != 2 || detail.GroupCount != 1 || len(detail.Associations) != 2 { t.Fatalf("upstream detail = %#v", detail) } - if detail.Associations[0].ClientModel != "client-a" || detail.Associations[0].Group.Name != "disabled" { + if detail.Associations[0].ClientModel != "client-a" || detail.Associations[0].Group.Name != "enabled" { t.Fatalf("detail associations are not stable/unfiltered: %#v", detail.Associations) } + encoded, err := json.Marshal(upstream.RouteGroups[0]) + if err != nil { + t.Fatalf("encode route Group: %v", err) + } + var wire map[string]json.RawMessage + if err := json.Unmarshal(encoded, &wire); err != nil { + t.Fatalf("decode route Group wire: %v", err) + } + for _, field := range []string{"channel_id", "params", "client_protocols"} { + if _, ok := wire[field]; !ok { + t.Fatalf("route Group wire misses %q: %s", field, encoded) + } + } + for _, field := range []string{"provider_id", "upstream_url", "protocols"} { + if _, ok := wire[field]; ok { + t.Fatalf("route Group wire exposes legacy %q: %s", field, encoded) + } + } } func TestProjectModelsHTTPScopesAccessKeyFiltersAndRelationships(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) allowedOpenAI := createPriceTestGroup(t, fixture.db, models.Group{ - Name: "allowed-openai", UpstreamURL: "https://openai.example/v1", - Protocols: models.JSON(`["openai-completions"]`), + Name: "allowed-openai", ChannelID: string(channel.OpenAICompatible), Params: models.JSON(`{"base_url":"https://openai.example/v1"}`), Models: models.JSON(`[ {"id":"upstream-a","alias":"client-a"}, {"id":"private-hidden","alias":"client-hidden"}, {"id":"upstream-shared","alias":"shared"} ]`), - Config: models.JSON(`{}`), Enabled: true, + Overrides: models.JSON(`{}`), Enabled: true, }) allowedAnthropic := createPriceTestGroup(t, fixture.db, models.Group{ - Name: "allowed-anthropic", UpstreamURL: "https://anthropic.example", - Protocols: models.JSON(`["anthropic"]`), + Name: "allowed-anthropic", ChannelID: string(channel.AnthropicCompatible), Params: models.JSON(`{"base_url":"https://anthropic.example"}`), Models: models.JSON(`[ {"id":"private-client-b","alias":"client-b"}, {"id":"upstream-shared","alias":"shared"} ]`), - Config: models.JSON(`{}`), Enabled: true, + Overrides: models.JSON(`{}`), Enabled: true, }) createPriceTestGroup(t, fixture.db, models.Group{ - Name: "private-gemini", UpstreamURL: "https://gemini.example", - Protocols: models.JSON(`["gemini"]`), + Name: "private-gemini", ChannelID: string(channel.GeminiCompatible), Params: models.JSON(`{"base_url":"https://gemini.example"}`), Models: models.JSON(`[{"id":"private-gemini-model","alias":"gemini-private"}]`), - Config: models.JSON(`{}`), Enabled: true, + Overrides: models.JSON(`{}`), Enabled: true, }) disabled := createPriceTestGroup(t, fixture.db, models.Group{ - Name: "private-disabled", UpstreamURL: "https://disabled.example/v1", - Protocols: models.JSON(`["openai-completions"]`), + Name: "private-disabled", ChannelID: string(channel.OpenAICompatible), Params: models.JSON(`{"base_url":"https://disabled.example/v1"}`), Models: models.JSON(`[{"id":"private-disabled-model","alias":"client-a"}]`), - Config: models.JSON(`{}`), Enabled: false, + Overrides: models.JSON(`{}`), Enabled: false, }) if err := fixture.db.Model(&disabled).Update("enabled", false).Error; err != nil { t.Fatalf("disable private Group: %v", err) @@ -121,19 +144,22 @@ func TestProjectModelsHTTPScopesAccessKeyFiltersAndRelationships(t *testing.T) { t.Fatalf("decode AccessKey models: %v", err) } if envelope.Data.Summary.ClientModelCount != 2 || - envelope.Data.Summary.UpstreamModelCount != 2 || - envelope.Data.Summary.PriceCount != 2 || + envelope.Data.Summary.UpstreamModelCount != 3 || + envelope.Data.Summary.PriceCount != 3 || len(envelope.Data.Items) != 2 || envelope.Data.Items[0].ClientModel != "client-a" || envelope.Data.Items[1].ClientModel != "shared" { t.Fatalf("AccessKey model response = %#v", envelope.Data) } shared := envelope.Data.Items[1] - if len(shared.Protocols) != 2 || len(shared.UpstreamModels) != 1 || - len(shared.UpstreamModels[0].RouteGroups) != 2 || - len(shared.UpstreamModels[0].AffectedGroups) != 2 { + if len(shared.Protocols) != 2 || len(shared.UpstreamModels) != 2 { t.Fatalf("AccessKey shared model = %#v", shared) } + for _, upstream := range shared.UpstreamModels { + if len(upstream.RouteGroups) != 1 || len(upstream.AffectedGroups) != 1 { + t.Fatalf("channel-scoped shared upstream = %#v", upstream) + } + } for _, privateValue := range []string{ "client-hidden", "client-b", "gemini-private", "private-gemini", "private-disabled", "private-hidden", "private-client-b", @@ -144,3 +170,28 @@ func TestProjectModelsHTTPScopesAccessKeyFiltersAndRelationships(t *testing.T) { } } } + +func TestProjectModelCatalogReferenceUsesTheRecordedPriceProvider(t *testing.T) { + matchedProviderID := "anthropic" + snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ + "openai": {ID: "openai", Models: map[string]catalog.Model{ + "shared": {ID: "shared", Name: "OpenAI metadata without a price"}, + }}, + "anthropic": {ID: "anthropic", Models: map[string]catalog.Model{ + "shared": { + ID: "shared", Name: "Anthropic priced model", + Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}}, + }, + }}, + }} + + reference := projectModelCatalogReference( + ModelPriceDTO{MatchedProviderID: &matchedProviderID}, + pricing.Identity{ChannelID: string(channel.OpenAICompatible), ModelID: "shared"}, + snapshot, + ) + if reference == nil || reference.ProviderID != "anthropic" || + reference.Model.Name != "Anthropic priced model" { + t.Fatalf("catalog reference = %#v", reference) + } +} diff --git a/internal/control/provider_catalog.go b/internal/control/provider_catalog.go deleted file mode 100644 index b3e0448a..00000000 --- a/internal/control/provider_catalog.go +++ /dev/null @@ -1,225 +0,0 @@ -package control - -import ( - "context" - "fmt" - "sort" - "strings" - - "gorm.io/gorm" - - "gpt-load/internal/catalog" - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/pricing" - "gpt-load/internal/protocol" - "gpt-load/internal/storage/models" -) - -const ( - providerSuggestionCatalogLimit = 25 - providerModelResultLimit = 100 -) - -// ProviderSuggestionSource identifies which directory a provider suggestion -// came from. Suggestions are merged official-first, then curated, then -// catalog, and deduplicated by provider ID so a higher-priority source's -// metadata always wins for a shared ID. -type ProviderSuggestionSource string - -const ( - ProviderSuggestionSourceOfficial ProviderSuggestionSource = "official" - ProviderSuggestionSourceCurated ProviderSuggestionSource = "curated" - ProviderSuggestionSourceCatalog ProviderSuggestionSource = "catalog" -) - -type ProviderSuggestion struct { - ProviderID string `json:"provider_id"` - Name string `json:"name"` - APIURL string `json:"api_url,omitempty"` - Protocols []protocol.Protocol `json:"protocols"` - Mark string `json:"mark"` - Source ProviderSuggestionSource `json:"source"` -} - -type ProviderSuggestionListResponse struct { - Items []ProviderSuggestion `json:"items"` - Total int `json:"total"` -} - -type ModelCandidate struct { - ID string `json:"id"` - Name string `json:"name"` - Sources []string `json:"sources"` - PricingStatus PricingStatus `json:"pricing_status"` - PricingSource *string `json:"pricing_source"` -} - -type ProviderModelQuery struct { - Query string - Status PricingStatus -} - -type ProviderModelListResponse struct { - Items []ModelCandidate `json:"items"` - Total int `json:"total"` -} - -func (s *Service) ListProviderSuggestions( - ctx context.Context, - query string, -) (ProviderSuggestionListResponse, error) { - if err := ctx.Err(); err != nil { - return ProviderSuggestionListResponse{}, err - } - query = strings.TrimSpace(query) - if len([]rune(query)) > 200 { - return ProviderSuggestionListResponse{}, app_errors.ErrValidation - } - curated := catalog.SearchCurated(query) - catalogOnly := s.catalogRuntime.SearchProviderMetadata(query, providerSuggestionCatalogLimit) - items := make([]ProviderSuggestion, 0, 3+len(curated)+len(catalogOnly)) - seen := make(map[string]struct{}, 3+len(curated)+len(catalogOnly)) - - appendSuggestions := func(providers []catalog.Provider, source ProviderSuggestionSource) { - for _, provider := range providers { - if _, duplicate := seen[provider.ID]; duplicate { - continue - } - seen[provider.ID] = struct{}{} - items = append(items, mapProviderSuggestion(provider, source)) - } - } - - // Official first, then curated, then catalog: a shared provider ID keeps - // the higher-priority source's display metadata (mapProviderSuggestion), - // while pricing continues to key off the same provider ID either way. - appendSuggestions(s.catalogRuntime.OfficialSuggestions(), ProviderSuggestionSourceOfficial) - appendSuggestions(curated, ProviderSuggestionSourceCurated) - appendSuggestions(catalogOnly, ProviderSuggestionSourceCatalog) - - return ProviderSuggestionListResponse{Items: items, Total: len(items)}, nil -} - -func (s *Service) ListProviderSuggestionsByIDs( - ctx context.Context, - providerIDs []string, -) (ProviderSuggestionListResponse, error) { - if err := ctx.Err(); err != nil { - return ProviderSuggestionListResponse{}, err - } - providers := catalog.SearchProviderMetadataByIDs(s.catalogRuntime.Load(), providerIDs) - items := make([]ProviderSuggestion, 0, len(providers)) - for _, provider := range providers { - items = append(items, mapProviderSuggestion(provider, ProviderSuggestionSourceCatalog)) - } - return ProviderSuggestionListResponse{Items: items, Total: len(items)}, nil -} - -func mapProviderSuggestion(provider catalog.Provider, source ProviderSuggestionSource) ProviderSuggestion { - protocols := make([]protocol.Protocol, len(provider.Protocols)) - copy(protocols, provider.Protocols) - return ProviderSuggestion{ - ProviderID: provider.ID, - Name: provider.Name, - APIURL: provider.APIURL, - Protocols: protocols, - Mark: provider.Mark, - Source: source, - } -} - -func (s *Service) validateSelectableProviderID(providerID *string) error { - if providerID == nil || catalog.IsOfficialProviderID(*providerID) || catalog.IsCuratedProviderID(*providerID) { - return nil - } - if s.catalogRuntime != nil { - if _, exists := s.catalogRuntime.LoadProvider(*providerID); exists { - return nil - } - } - return app_errors.ErrValidation -} - -func (s *Service) ListProviderModels( - ctx context.Context, - providerID string, - query ProviderModelQuery, -) (ProviderModelListResponse, error) { - if _, err := pricing.ProviderScopeKey(providerID); err != nil { - return ProviderModelListResponse{}, app_errors.ErrValidation - } - query.Query = strings.TrimSpace(query.Query) - if len([]rune(query.Query)) > 200 || - (query.Status != "" && query.Status != PricingStatusPending && query.Status != PricingStatusConfigured) { - return ProviderModelListResponse{}, app_errors.ErrValidation - } - - s.writeMu.RLock() - defer s.writeMu.RUnlock() - snapshot := s.catalogRuntime.Load() - if snapshot == nil { - return ProviderModelListResponse{}, app_errors.ErrResourceNotFound - } - provider, exists := snapshot.Providers[providerID] - if !exists { - return ProviderModelListResponse{}, app_errors.ErrResourceNotFound - } - rows, err := loadModelPriceRows(ctx, s.db) - if err != nil { - return ProviderModelListResponse{}, err - } - - normalizedQuery := strings.ToLower(query.Query) - items := make([]ModelCandidate, 0, len(provider.Models)) - for _, model := range provider.Models { - if normalizedQuery != "" && - !strings.Contains(strings.ToLower(model.ID), normalizedQuery) && - !strings.Contains(strings.ToLower(model.Name), normalizedQuery) { - continue - } - row := rows[model.ID] - status, pricingSource := resolveCandidatePricing(row, snapshot, model.ID) - if query.Status != "" && query.Status != status { - continue - } - name := strings.TrimSpace(model.Name) - if name == "" { - name = model.ID - } - items = append(items, ModelCandidate{ - ID: model.ID, Name: name, Sources: []string{"catalog"}, - PricingStatus: status, PricingSource: pricingSource, - }) - } - sort.Slice(items, func(left, right int) bool { - leftName := strings.ToLower(items[left].Name) - rightName := strings.ToLower(items[right].Name) - if leftName == rightName { - return items[left].ID < items[right].ID - } - return leftName < rightName - }) - total := len(items) - if len(items) > providerModelResultLimit { - items = items[:providerModelResultLimit] - } - return ProviderModelListResponse{Items: items, Total: total}, nil -} - -func loadModelPriceRows( - ctx context.Context, - db *gorm.DB, -) (map[string]*models.ModelPrice, error) { - var rows []models.ModelPrice - if err := db.WithContext(ctx).Order("id ASC").Find(&rows).Error; err != nil { - return nil, fmt.Errorf("load model price status: %w", app_errors.ParseDBError(err)) - } - result := make(map[string]*models.ModelPrice, len(rows)) - for index := range rows { - if _, duplicate := result[rows[index].ModelID]; duplicate { - return nil, fmt.Errorf("duplicate persisted price identity: %w", app_errors.ErrInternalServer) - } - result[rows[index].ModelID] = &rows[index] - } - return result, nil -} diff --git a/internal/control/provider_catalog_http.go b/internal/control/provider_catalog_http.go deleted file mode 100644 index dc44bdd2..00000000 --- a/internal/control/provider_catalog_http.go +++ /dev/null @@ -1,174 +0,0 @@ -package control - -import ( - "context" - "fmt" - "net/url" - "strings" - "unicode/utf8" - - "github.com/gin-gonic/gin" - - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/platform/response" - "gpt-load/internal/pricing" -) - -const ( - providerCatalogMaxQueryRunes = 200 - providerSuggestionMaxIDs = 20 -) - -func (s *Server) handleListProviderSuggestions(c *gin.Context) { - query, providerIDs, apiErr := parseProviderSuggestionQuery( - c.Request.URL.RawQuery, - c.Request.URL.ForceQuery, - ) - if apiErr != nil { - writeServiceError(c, "list_provider_suggestions", apiErr) - return - } - var result ProviderSuggestionListResponse - var err error - if len(providerIDs) > 0 { - result, err = s.service.ListProviderSuggestionsByIDs(c.Request.Context(), providerIDs) - } else { - result, err = s.service.ListProviderSuggestions(c.Request.Context(), query) - } - if err != nil { - writeServiceError(c, "list_provider_suggestions", err) - return - } - response.SuccessI18n(c, "common.success", result) -} - -func (s *Server) handleListProviderModels(c *gin.Context) { - query, apiErr := parseProviderModelQuery(c.Request.URL.RawQuery, c.Request.URL.ForceQuery) - if apiErr != nil { - writeServiceError(c, "list_provider_models", apiErr) - return - } - result, err := s.service.ListProviderModels(c.Request.Context(), c.Param("provider_id"), query) - if err != nil { - writeServiceError(c, "list_provider_models", err) - return - } - response.SuccessI18n(c, "common.success", result) -} - -func (s *Server) handleSyncModelPrices(c *gin.Context) { - if c.Request.URL.ForceQuery || c.Request.URL.RawQuery != "" { - writeServiceError(c, "sync_model_prices", app_errors.ErrBadRequest) - return - } - result, err := s.service.SyncModelPrices(c.Request.Context()) - if err != nil { - writeServiceError(c, "sync_model_prices", err) - return - } - response.SuccessI18n(c, "common.success", result) -} - -func parseProviderSuggestionQuery(rawQuery string, forceQuery bool) (string, []string, *app_errors.APIError) { - values, apiErr := parseProviderCatalogValues( - rawQuery, - forceQuery, - map[string]struct{}{"q": {}, "provider_ids": {}}, - ) - if apiErr != nil { - return "", nil, apiErr - } - query := strings.TrimSpace(values.Get("q")) - if utf8.RuneCountInString(query) > providerCatalogMaxQueryRunes { - return "", nil, app_errors.ErrBadRequest - } - providerIDs := []string(nil) - if values.Has("provider_ids") { - if query != "" { - return "", nil, app_errors.ErrBadRequest - } - var parseErr *app_errors.APIError - providerIDs, parseErr = parseProviderSuggestionIDs(values.Get("provider_ids")) - if parseErr != nil { - return "", nil, parseErr - } - } - return query, providerIDs, nil -} - -func parseProviderSuggestionIDs(raw string) ([]string, *app_errors.APIError) { - parts := strings.Split(raw, ",") - if len(parts) == 0 || len(parts) > providerSuggestionMaxIDs { - return nil, app_errors.ErrBadRequest - } - ids := make([]string, 0, len(parts)) - seen := make(map[string]struct{}, len(parts)) - for _, part := range parts { - id := strings.TrimSpace(part) - if _, err := pricing.ProviderScopeKey(id); err != nil { - return nil, app_errors.ErrBadRequest - } - if _, duplicate := seen[id]; duplicate { - continue - } - seen[id] = struct{}{} - ids = append(ids, id) - } - if len(ids) == 0 { - return nil, app_errors.ErrBadRequest - } - return ids, nil -} - -func parseProviderModelQuery(rawQuery string, forceQuery bool) (ProviderModelQuery, *app_errors.APIError) { - values, apiErr := parseProviderCatalogValues( - rawQuery, - forceQuery, - map[string]struct{}{"q": {}, "status": {}}, - ) - if apiErr != nil { - return ProviderModelQuery{}, apiErr - } - query := ProviderModelQuery{Query: strings.TrimSpace(values.Get("q"))} - if utf8.RuneCountInString(query.Query) > providerCatalogMaxQueryRunes { - return ProviderModelQuery{}, app_errors.ErrBadRequest - } - if values.Has("status") { - query.Status = PricingStatus(values.Get("status")) - if query.Status != PricingStatusPending && query.Status != PricingStatusConfigured { - return ProviderModelQuery{}, app_errors.ErrBadRequest - } - } - return query, nil -} - -func parseProviderCatalogValues( - rawQuery string, - forceQuery bool, - allowed map[string]struct{}, -) (url.Values, *app_errors.APIError) { - if forceQuery && rawQuery == "" { - return nil, app_errors.ErrBadRequest - } - values, err := url.ParseQuery(rawQuery) - if err != nil { - return nil, app_errors.ErrBadRequest - } - for key, entries := range values { - if _, ok := allowed[key]; !ok || len(entries) != 1 { - return nil, app_errors.ErrBadRequest - } - } - return values, nil -} - -func (s *Service) SyncModelPrices(ctx context.Context) (CatalogSyncStatus, error) { - if s.catalogSync == nil { - return CatalogSyncStatus{}, fmt.Errorf("catalog sync is unavailable: %w", app_errors.ErrInternalServer) - } - status, err := s.catalogSync.Sync(ctx, CatalogSyncManual) - if err != nil { - return status, app_errors.NewAPIErrorWithData(app_errors.ErrBadGateway, status) - } - return status, nil -} diff --git a/internal/control/provider_catalog_http_test.go b/internal/control/provider_catalog_http_test.go deleted file mode 100644 index a7eddb57..00000000 --- a/internal/control/provider_catalog_http_test.go +++ /dev/null @@ -1,340 +0,0 @@ -package control - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "net/http" - "net/http/httptest" - "strings" - "testing" - "time" - - "github.com/gin-gonic/gin" - "github.com/sirupsen/logrus" - - "gpt-load/internal/catalog" - "gpt-load/internal/platform/config" - "gpt-load/internal/pricing" - "gpt-load/internal/protocol" - "gpt-load/internal/storage/models" -) - -func TestProviderSuggestionsUseOfficialFallbacksMergeCatalogAndStayBounded(t *testing.T) { - fixture := newServiceFixture(t) - withoutCatalog, err := fixture.service.ListProviderSuggestions(t.Context(), "") - if err != nil { - t.Fatal(err) - } - wantWithoutCatalogCount := 3 + len(catalog.CuratedProviders()) - if len(withoutCatalog.Items) != wantWithoutCatalogCount || withoutCatalog.Items[0].ProviderID != "openai" || - withoutCatalog.Items[1].ProviderID != "anthropic" || withoutCatalog.Items[2].ProviderID != "google" { - t.Fatalf("fallback suggestions = %#v, want %d items (3 official + full curated table)", - withoutCatalog.Items, wantWithoutCatalogCount) - } - - providers := map[string]catalog.Provider{ - "openai": { - ID: "openai", Name: "OpenAI Catalog", APIURL: "https://catalog.openai.example/v1", - Models: map[string]catalog.Model{}, - }, - } - for index := 0; index < providerSuggestionCatalogLimit+5; index++ { - id := "provider-" + strings.Repeat("a", index/26) + string(rune('a'+index%26)) - providers[id] = catalog.Provider{ - ID: id, Name: "Provider " + id, APIURL: "https://" + id + ".example/v1", - NPM: "@ai-sdk/openai-compatible", Models: map[string]catalog.Model{}, - } - } - fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: providers}) - empty, err := fixture.service.ListProviderSuggestions(t.Context(), "") - if err != nil { - t.Fatal(err) - } - wantEmptyCount := 3 + len(catalog.CuratedProviders()) - if len(empty.Items) != wantEmptyCount { - t.Fatalf("empty search suggestions = %d, want three official plus full curated table (%d)", - len(empty.Items), wantEmptyCount) - } - for _, item := range empty.Items[3:] { - if item.Source != ProviderSuggestionSourceCurated { - t.Fatalf("empty search suggestion %#v, want source=curated (models.dev catalog stays hidden until searched)", item) - } - } - got, err := fixture.service.ListProviderSuggestions(t.Context(), "provider") - if err != nil { - t.Fatal(err) - } - if len(got.Items) != 3+providerSuggestionCatalogLimit || got.Items[0].Name != "OpenAI Catalog" || - got.Items[0].APIURL != "https://catalog.openai.example/v1" || - got.Items[0].Source != ProviderSuggestionSourceOfficial || - len(got.Items[0].Protocols) != 2 || got.Items[0].Protocols[0] != protocol.OpenAICompletions { - t.Fatalf("merged suggestions = %#v", got.Items) - } - encoded, err := json.Marshal(got) - if err != nil { - t.Fatal(err) - } - for _, forbidden := range []string{"svg", `"models"`, `"npm"`} { - if strings.Contains(strings.ToLower(string(encoded)), forbidden) { - t.Fatalf("suggestions expose %q: %s", forbidden, encoded) - } - } -} - -func TestProviderSuggestionsCuratedOutranksCatalogForSharedID(t *testing.T) { - fixture := newServiceFixture(t) - curated := catalog.CuratedProviders() - if len(curated) == 0 { - t.Fatal("catalog.CuratedProviders() is empty, cannot exercise dedup priority") - } - sharedID := curated[0].ID - - // A models.dev snapshot entry sharing a curated provider's ID must lose - // the display metadata race: the curated entry's Name/APIURL/Mark/Source - // wins, and the catalog duplicate never appears as a second item. - fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ - sharedID: { - ID: sharedID, - Name: "Should Not Win", - APIURL: "https://should-not-win.example/v1", - Models: map[string]catalog.Model{}, - }, - }}) - - got, err := fixture.service.ListProviderSuggestions(t.Context(), sharedID) - if err != nil { - t.Fatal(err) - } - matches := 0 - for _, item := range got.Items { - if item.ProviderID != sharedID { - continue - } - matches++ - if item.Source != ProviderSuggestionSourceCurated || item.Name == "Should Not Win" { - t.Fatalf("shared-ID suggestion = %#v, want curated metadata to win", item) - } - } - if matches != 1 { - t.Fatalf("shared-ID provider %q appeared %d times, want exactly 1 (deduplicated)", sharedID, matches) - } -} - -func TestProviderSuggestionsByIDsUseModelsDevMetadataInRequestedOrder(t *testing.T) { - fixture := newServiceFixture(t) - fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ - "deepseek": { - ID: "deepseek", Name: "DeepSeek Models.dev", APIURL: "https://models.example/deepseek/v1", - NPM: "@ai-sdk/openai-compatible", Models: map[string]catalog.Model{}, - }, - "anthropic-cloud": { - ID: "anthropic-cloud", Name: "Anthropic Cloud", APIURL: "https://models.example/anthropic/v1", - NPM: "@ai-sdk/anthropic", Models: map[string]catalog.Model{}, - }, - }}) - - got, err := fixture.service.ListProviderSuggestionsByIDs( - t.Context(), - []string{"anthropic-cloud", "deepseek", "missing", "deepseek"}, - ) - if err != nil { - t.Fatal(err) - } - if len(got.Items) != 2 || got.Items[0].ProviderID != "anthropic-cloud" || got.Items[1].ProviderID != "deepseek" { - t.Fatalf("provider suggestions by IDs = %#v, want requested matched order", got.Items) - } - if got.Items[0].Source != ProviderSuggestionSourceCatalog || got.Items[0].Name != "Anthropic Cloud" || - got.Items[0].APIURL != "https://models.example/anthropic/v1" || - len(got.Items[0].Protocols) != 1 || got.Items[0].Protocols[0] != protocol.Anthropic { - t.Fatalf("anthropic catalog metadata = %#v", got.Items[0]) - } - if got.Items[1].Source != ProviderSuggestionSourceCatalog || got.Items[1].Name != "DeepSeek Models.dev" || - got.Items[1].APIURL != "https://models.example/deepseek/v1" || - len(got.Items[1].Protocols) != 1 || got.Items[1].Protocols[0] != protocol.OpenAICompletions { - t.Fatalf("deepseek catalog metadata = %#v", got.Items[1]) - } -} - -func TestValidateSelectableProviderIDAllowsCuratedProviderIDs(t *testing.T) { - fixture := newServiceFixture(t) - for _, provider := range catalog.CuratedProviders() { - id := provider.ID - if err := fixture.service.validateSelectableProviderID(&id); err != nil { - t.Fatalf("validateSelectableProviderID(%q) = %v, want nil for curated provider", id, err) - } - } - unknown := "no-such-provider-anywhere" - if err := fixture.service.validateSelectableProviderID(&unknown); err == nil { - t.Fatal("validateSelectableProviderID(unknown) = nil, want an error") - } -} - -func TestProviderSuggestionsEncodeUnknownProtocolsAsEmptyArray(t *testing.T) { - fixture := newServiceFixture(t) - fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ - "custom-provider": { - ID: "custom-provider", - Name: "Custom Provider", - APIURL: "https://custom.example/v1", - NPM: "@vendor/unknown-sdk", - Models: map[string]catalog.Model{}, - }, - }}) - - got, err := fixture.service.ListProviderSuggestions(t.Context(), "custom") - if err != nil { - t.Fatal(err) - } - encoded, err := json.Marshal(got) - if err != nil { - t.Fatal(err) - } - if strings.Contains(string(encoded), `"protocols":null`) || - !strings.Contains(string(encoded), `"provider_id":"custom-provider"`) || - !strings.Contains(string(encoded), `"protocols":[]`) { - t.Fatalf("unknown provider protocols = %s, want an empty JSON array", encoded) - } -} - -func TestProviderModelsAreStrictBoundedAndUsePersistedOrCatalogCandidatePricingStatus(t *testing.T) { - fixture := newServiceFixture(t) - providerModels := make(map[string]catalog.Model) - for index := 0; index < providerModelResultLimit+5; index++ { - id := "model-" + strings.Repeat("a", index/26) + string(rune('a'+index%26)) - providerModels[id] = catalog.Model{ID: id, Name: "Model " + id} - } - providerModels["configured"] = catalog.Model{ID: "configured", Name: "Configured"} - providerModels["catalog-priced"] = catalog.Model{ - ID: "catalog-priced", Name: "Catalog Priced", - Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}}, - } - providerModels["pending"] = catalog.Model{ID: "pending", Name: "Pending"} - fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ - "openai": {ID: "openai", Name: "OpenAI", Models: providerModels}, - }}) - configuredValue := int64(0) - for _, row := range []models.ModelPrice{ - {ModelID: "configured", InputPriceNanoUSDPerMillionTokens: &configuredValue}, - {ModelID: "pending"}, - } { - if err := fixture.db.Create(&row).Error; err != nil { - t.Fatal(err) - } - } - - got, err := fixture.service.ListProviderModels(t.Context(), "openai", ProviderModelQuery{}) - if err != nil { - t.Fatal(err) - } - if len(got.Items) != providerModelResultLimit || got.Total != len(providerModels) { - t.Fatalf("provider model bounds = %d/%d, want %d/%d", len(got.Items), got.Total, providerModelResultLimit, len(providerModels)) - } - configured, err := fixture.service.ListProviderModels(t.Context(), "openai", ProviderModelQuery{ - Query: "configured", Status: PricingStatusConfigured, - }) - if err != nil || len(configured.Items) != 1 || configured.Items[0].PricingStatus != PricingStatusConfigured { - t.Fatalf("configured filter = %#v, %v", configured, err) - } - catalogConfigured, err := fixture.service.ListProviderModels(t.Context(), "openai", ProviderModelQuery{ - Query: "catalog-priced", Status: PricingStatusConfigured, - }) - if err != nil || len(catalogConfigured.Items) != 1 || - catalogConfigured.Items[0].PricingStatus != PricingStatusConfigured { - t.Fatalf("catalog configured filter = %#v, %v", catalogConfigured, err) - } - pending, err := fixture.service.ListProviderModels(t.Context(), "openai", ProviderModelQuery{ - Query: "pending", Status: PricingStatusPending, - }) - if err != nil || len(pending.Items) != 1 || pending.Items[0].PricingStatus != PricingStatusPending { - t.Fatalf("pending filter = %#v, %v", pending, err) - } - for _, providerID := range []string{"OpenAI", "openai:alias", "missing"} { - if _, err := fixture.service.ListProviderModels(t.Context(), providerID, ProviderModelQuery{}); err == nil { - t.Fatalf("ListProviderModels(%q) error = nil", providerID) - } - } -} - -func TestProviderCatalogRoutesRequireAuthRejectInvalidQueriesAndSanitizeSyncFailure(t *testing.T) { - initControlI18n(t) - fixture := newServiceFixture(t) - fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ - "openai": {ID: "openai", Name: "OpenAI", Models: map[string]catalog.Model{}}, - }}) - const rawFailure = "secret upstream response body" - client := catalogSyncClientFunc(func(context.Context, catalog.Metadata) (catalog.SyncResult, error) { - return catalog.SyncResult{}, errors.New(rawFailure) - }) - coordinator := newCatalogSyncCoordinator(fixture.service, client, "unused", catalog.Metadata{ - CheckedAtMillis: 100, SuccessfulFetchAtMillis: 90, - }, true) - coordinator.now = func() time.Time { return time.UnixMilli(250) } - var serviceLogs bytes.Buffer - previousOutput := logrus.StandardLogger().Out - logrus.SetOutput(&serviceLogs) - t.Cleanup(func() { logrus.SetOutput(previousOutput) }) - engine := gin.New() - NewServer(&config.Config{AuthKey: authTestKey}, fixture.service).RegisterRoutes(engine) - - unauthorized := serveProviderCatalogRequest(engine, http.MethodGet, "/api/provider-suggestions", "") - if unauthorized.Code != http.StatusUnauthorized { - t.Fatalf("unauthorized status = %d", unauthorized.Code) - } - for _, path := range []string{ - "/api/provider-suggestions?unknown=1", - "/api/provider-suggestions?q=a&q=b", - "/api/provider-suggestions?q=a&provider_ids=openai", - "/api/provider-suggestions?provider_ids=", - "/api/providers/OpenAI/models", - "/api/providers/openai/models?status=priced", - } { - response := serveProviderCatalogRequest(engine, http.MethodGet, path, authTestKey) - if response.Code != http.StatusBadRequest { - t.Fatalf("GET %s = %d %s, want 400", path, response.Code, response.Body.String()) - } - } - catalogByID := serveProviderCatalogRequest( - engine, - http.MethodGet, - "/api/provider-suggestions?provider_ids=openai", - authTestKey, - ) - if catalogByID.Code != http.StatusOK || - !strings.Contains(catalogByID.Body.String(), `"provider_id":"openai"`) || - !strings.Contains(catalogByID.Body.String(), `"source":"catalog"`) { - t.Fatalf("provider suggestions by ID = %d %s", catalogByID.Code, catalogByID.Body.String()) - } - - syncResponse := serveProviderCatalogRequest(engine, http.MethodPost, "/api/model-prices/sync", authTestKey) - if syncResponse.Code != http.StatusBadGateway || - strings.Contains(syncResponse.Body.String(), rawFailure) { - t.Fatalf("sync failure = %d %s", syncResponse.Code, syncResponse.Body.String()) - } - var envelope struct { - Code string `json:"code"` - Data CatalogSyncStatus `json:"data"` - } - if err := json.Unmarshal(syncResponse.Body.Bytes(), &envelope); err != nil { - t.Fatal(err) - } - if envelope.Code != "BAD_GATEWAY" || envelope.Data.Trigger != CatalogSyncManual || - envelope.Data.CheckedAtMS != 250 || envelope.Data.SuccessfulFetchAtMS != 90 || - envelope.Data.ErrorCode != "catalog_sync_failed" || envelope.Data.NotModified || - envelope.Data.Skipped { - t.Fatalf("sync failure envelope = %#v", envelope) - } - assertControlLogExcludes(t, serviceLogs.String(), rawFailure) -} - -func serveProviderCatalogRequest(engine *gin.Engine, method, path, authKey string) *httptest.ResponseRecorder { - recorder := httptest.NewRecorder() - request := httptest.NewRequest(method, path, nil) - if authKey != "" { - request.Header.Set("Authorization", "Bearer "+authKey) - } - request.Header.Set("Accept-Language", "en-US") - engine.ServeHTTP(recorder, request) - return recorder -} diff --git a/internal/control/request_logs.go b/internal/control/request_logs.go index e1fa597d..6982e209 100644 --- a/internal/control/request_logs.go +++ b/internal/control/request_logs.go @@ -14,6 +14,8 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" + "gpt-load/internal/execution" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/platform/response" "gpt-load/internal/pricing" @@ -33,6 +35,8 @@ var canonicalLowercaseUUIDv4 = regexp.MustCompile( `^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$`, ) +var requestLogChannelRegistry = channel.NewRegistry() + type RequestLogReader interface { List(context.Context, requestlog.ListQuery) (requestlog.Page, error) Get(context.Context, string) (requestlog.Record, error) @@ -57,24 +61,32 @@ type requestLogReasoningResponse struct { } type requestLogAttemptResponse struct { - Sequence int `json:"sequence"` - GroupID uint `json:"group_id"` - GroupName string `json:"group_name"` - KeyID uint `json:"key_id"` - UpstreamModel *string `json:"upstream_model"` - StatusCode int `json:"status_code"` - DurationMs int64 `json:"duration_ms"` - FailureCategory telemetry.FailureCategory `json:"failure_category"` - Action telemetry.Action `json:"action"` - WillRetry bool `json:"will_retry"` - ErrorCode string `json:"error_code"` - ErrorSummary string `json:"error_summary"` - PricingReceipt *requestLogPricingReceiptResponse `json:"pricing_receipt"` + Sequence int `json:"sequence"` + GroupID uint `json:"group_id"` + GroupName string `json:"group_name"` + ChannelID *channel.ID `json:"channel_id"` + CredentialID *uint `json:"credential_id"` + Operation *execution.Operation `json:"operation"` + RouteMode *channel.RouteMode `json:"route_mode"` + UpstreamModel *string `json:"upstream_model"` + UpstreamRequestID *string `json:"upstream_request_id"` + DispatchState *execution.DispatchState `json:"dispatch_state"` + ResponseStarted bool `json:"response_started"` + StatusCode int `json:"status_code"` + DurationMs int64 `json:"duration_ms"` + FailureCategory telemetry.FailureCategory `json:"failure_category"` + Action string `json:"action"` + WillRetry bool `json:"will_retry"` + ErrorCode string `json:"error_code"` + ErrorSummary string `json:"error_summary"` + Committed bool `json:"committed"` + PricingReceipt *requestLogPricingReceiptResponse `json:"pricing_receipt"` } type requestLogPricingIdentityResponse struct { - ScopeKey *string `json:"scope_key,omitempty"` - ModelID string `json:"model_id"` + ScopeKey *string `json:"scope_key,omitempty"` + ChannelID *string `json:"channel_id,omitempty"` + ModelID string `json:"model_id"` } type requestLogPricingMultiplierResponse struct { @@ -122,6 +134,8 @@ type requestLogItemResponse struct { ErrorSummary string `json:"error_summary"` AffinityHit bool `json:"affinity_hit"` GroupID *uint `json:"group_id"` + ChannelID *channel.ID `json:"channel_id"` + CredentialID *uint `json:"credential_id"` UsageState usage.State `json:"usage_state"` CostState pricing.CostState `json:"cost_state"` PricingCompleteness pricing.Completeness `json:"pricing_completeness"` @@ -259,9 +273,10 @@ func requestLogQueryUsesInternalFields(rawQuery string) bool { } for _, key := range []string{ "group_id", + "channel_id", + "credential_id", "upstream_model", "access_key_id", - "upstream_key_id", "attempt_status_code", "failure_category", "error_code", @@ -283,6 +298,8 @@ func sanitizeAccessKeyRequestLog(record requestlog.Record) requestlog.Record { record.AttemptCount = 0 record.AffinityHit = false record.GroupID = 0 + record.ChannelID = "" + record.CredentialID = 0 record.Attempts = []requestlog.Attempt{} return record } @@ -293,10 +310,11 @@ func parseRequestLogQuery(rawQuery string) (requestlog.ListQuery, *app_errors.AP return requestlog.ListQuery{}, app_errors.ErrBadRequest } allowed := map[string]struct{}{ - "from_ms": {}, "to_ms": {}, "group_id": {}, "client_model": {}, "upstream_model": {}, "access_key_id": {}, + "from_ms": {}, "to_ms": {}, "group_id": {}, "channel_id": {}, "credential_id": {}, + "client_model": {}, "upstream_model": {}, "access_key_id": {}, "status": {}, "request_id": {}, "protocol": {}, "stream": {}, "final_status_code": {}, "usage_state": {}, "cost_state": {}, "pricing_completeness": {}, "cache_present": {}, - "upstream_key_id": {}, "attempt_status_code": {}, "failure_category": {}, "error_code": {}, + "attempt_status_code": {}, "failure_category": {}, "error_code": {}, "retry_state": {}, "retry_count_min": {}, "retry_count_max": {}, "first_response_min_ms": {}, "first_response_max_ms": {}, "duration_min_ms": {}, "duration_max_ms": {}, @@ -336,6 +354,20 @@ func parseRequestLogQuery(rawQuery string) (requestlog.ListQuery, *app_errors.AP } query.GroupID = &parsed } + if value, ok := singleQueryValue(values, "channel_id"); ok { + parsed, apiErr := parseRequestLogChannelID(value) + if apiErr != nil { + return requestlog.ListQuery{}, apiErr + } + query.ChannelID = parsed + } + if value, ok := singleQueryValue(values, "credential_id"); ok { + parsed, apiErr := parseRequestLogID(value) + if apiErr != nil { + return requestlog.ListQuery{}, apiErr + } + query.CredentialID = &parsed + } if value, ok := singleQueryValue(values, "client_model"); ok { if !validUsageModel(value) { return requestlog.ListQuery{}, app_errors.ErrValidation @@ -425,13 +457,6 @@ func parseRequestLogQuery(rawQuery string) (requestlog.ListQuery, *app_errors.AP } query.CachePresent = &parsed } - if value, ok := singleQueryValue(values, "upstream_key_id"); ok { - parsed, apiErr := parseRequestLogID(value) - if apiErr != nil { - return requestlog.ListQuery{}, apiErr - } - query.UpstreamKeyID = &parsed - } if value, ok := singleQueryValue(values, "attempt_status_code"); ok { parsed, apiErr := parseRequestLogStatusCode(value) if apiErr != nil { @@ -552,6 +577,14 @@ func parseRequestLogID(value string) (uint, *app_errors.APIError) { return uint(parsed), nil } +func parseRequestLogChannelID(value string) (channel.ID, *app_errors.APIError) { + channelID := channel.ID(value) + if _, ok := requestLogChannelRegistry.Get(channelID); !ok { + return "", app_errors.ErrValidation + } + return channelID, nil +} + func parseRequestLogBool(value string) (bool, bool) { switch value { case "true": @@ -811,6 +844,8 @@ func mapRequestLogItemResponse( ErrorSummary: record.ErrorSummary, AffinityHit: record.AffinityHit, GroupID: usageCost.groupID, + ChannelID: usageCost.channelID, + CredentialID: usageCost.credentialID, UsageState: record.UsageState, CostState: record.CostState, PricingCompleteness: record.PricingCompleteness, @@ -875,29 +910,68 @@ func mapRequestLogDetailResponse(record requestlog.Record) (requestLogDetailResp } attempts := make([]requestLogAttemptResponse, 0, len(record.Attempts)) for _, attempt := range record.Attempts { - receipt, err := mapRequestLogPricingReceipt(attempt.PricingReceipt) + mapped, err := mapRequestLogAttempt(attempt) if err != nil { return requestLogDetailResponse{}, err } - attempts = append(attempts, requestLogAttemptResponse{ - Sequence: attempt.Sequence, - GroupID: attempt.GroupID, - GroupName: attempt.GroupName, - KeyID: attempt.KeyID, - UpstreamModel: nullableRequestLogModel(attempt.UpstreamModel), - StatusCode: attempt.StatusCode, - DurationMs: attempt.DurationMs, - FailureCategory: attempt.FailureCategory, - Action: attempt.Action, - WillRetry: attempt.WillRetry, - ErrorCode: attempt.ErrorCode, - ErrorSummary: attempt.ErrorSummary, - PricingReceipt: receipt, - }) + attempts = append(attempts, mapped) } return requestLogDetailResponse{requestLogItemResponse: item, Attempts: attempts}, nil } +func mapRequestLogAttempt(attempt requestlog.Attempt) (requestLogAttemptResponse, error) { + channelID, err := nullableRequestLogChannelID(attempt.ChannelID) + if err != nil { + return requestLogAttemptResponse{}, fmt.Errorf("map request log attempt: %w", err) + } + credentialID, err := nullableRequestLogID(attempt.CredentialID, "credential") + if err != nil { + return requestLogAttemptResponse{}, fmt.Errorf("map request log attempt: %w", err) + } + operation, err := nullableRequestLogOperation(attempt.Operation) + if err != nil { + return requestLogAttemptResponse{}, fmt.Errorf("map request log attempt: %w", err) + } + routeMode, err := nullableRequestLogRouteMode(attempt.RouteMode) + if err != nil { + return requestLogAttemptResponse{}, fmt.Errorf("map request log attempt: %w", err) + } + dispatchState, err := nullableRequestLogDispatchState(attempt.DispatchState) + if err != nil { + return requestLogAttemptResponse{}, fmt.Errorf("map request log attempt: %w", err) + } + receipt, err := mapRequestLogPricingReceipt(attempt.PricingReceipt) + if err != nil { + return requestLogAttemptResponse{}, err + } + return requestLogAttemptResponse{ + Sequence: attempt.Sequence, + GroupID: attempt.GroupID, + GroupName: attempt.GroupName, + ChannelID: channelID, + CredentialID: credentialID, + Operation: operation, + RouteMode: routeMode, + UpstreamModel: nullableRequestLogModel(attempt.UpstreamModel), + UpstreamRequestID: nullableRequestLogModel(attempt.UpstreamRequestID), + DispatchState: dispatchState, + ResponseStarted: attempt.ResponseStarted, + StatusCode: attempt.StatusCode, + DurationMs: attempt.DurationMs, + FailureCategory: attempt.FailureCategory, + Action: requestLogAttemptAction(attempt.Action), + WillRetry: attempt.WillRetry, + ErrorCode: attempt.ErrorCode, + ErrorSummary: attempt.ErrorSummary, + Committed: attempt.Committed, + PricingReceipt: receipt, + }, nil +} + +func requestLogAttemptAction(action telemetry.Action) string { + return string(action) +} + func mapRequestLogPricingReceipt( receipt *pricing.Receipt, ) (*requestLogPricingReceiptResponse, error) { @@ -920,6 +994,14 @@ func mapRequestLogPricingReceipt( scopeKey := receipt.Rule.ScopeKey result.Rule.ScopeKey = &scopeKey } + if receipt.SchemaVersion == 3 { + channelID := channel.ID(receipt.Rule.ChannelID) + if _, ok := requestLogChannelRegistry.Get(channelID); !ok { + return nil, fmt.Errorf("map request log pricing receipt: unknown channel ID") + } + value := string(channelID) + result.Rule.ChannelID = &value + } if receipt.ContextThresholdTokens != nil { value := strconv.FormatInt(*receipt.ContextThresholdTokens, 10) result.ContextThresholdTokens = &value @@ -954,8 +1036,62 @@ func nullableRequestLogModel(value string) *string { return &value } +func nullableRequestLogChannelID(value channel.ID) (*channel.ID, error) { + if value == "" { + return nil, nil + } + if _, ok := requestLogChannelRegistry.Get(value); !ok { + return nil, fmt.Errorf("unknown channel ID") + } + return &value, nil +} + +func nullableRequestLogID(value uint, field string) (*uint, error) { + if value == 0 { + return nil, nil + } + if uint64(value) > uint64(maxSafeInteger) { + return nil, fmt.Errorf("unsafe %s ID", field) + } + return &value, nil +} + +func nullableRequestLogOperation(value execution.Operation) (*execution.Operation, error) { + if value == "" { + return nil, nil + } + if !value.Valid() { + return nil, fmt.Errorf("invalid operation") + } + return &value, nil +} + +func nullableRequestLogRouteMode(value channel.RouteMode) (*channel.RouteMode, error) { + if value == "" { + return nil, nil + } + if !value.Valid() { + return nil, fmt.Errorf("invalid route mode") + } + return &value, nil +} + +func nullableRequestLogDispatchState( + value execution.DispatchState, +) (*execution.DispatchState, error) { + if value == "" { + return nil, nil + } + if !value.Valid() { + return nil, fmt.Errorf("invalid dispatch state") + } + return &value, nil +} + type requestLogUsageCostResponse struct { groupID *uint + channelID *channel.ID + credentialID *uint estimatedCostNanoUSD string } @@ -993,12 +1129,22 @@ func mapRequestLogUsageCost(record requestlog.Record) (requestLogUsageCostRespon if uint64(record.GroupID) > uint64(maxSafeInteger) { return requestLogUsageCostResponse{}, fmt.Errorf("map request log final Group ID: unsafe value") } + channelID, err := nullableRequestLogChannelID(record.ChannelID) + if err != nil { + return requestLogUsageCostResponse{}, fmt.Errorf("map request log final channel: %w", err) + } + credentialID, err := nullableRequestLogID(record.CredentialID, "credential") + if err != nil { + return requestLogUsageCostResponse{}, fmt.Errorf("map request log final credential: %w", err) + } if record.EstimatedCostNanoUSD < 0 { return requestLogUsageCostResponse{}, fmt.Errorf( "map request log usage/cost: negative estimated cost", ) } result := requestLogUsageCostResponse{ + channelID: channelID, + credentialID: credentialID, estimatedCostNanoUSD: strconv.FormatInt(record.EstimatedCostNanoUSD, 10), } if record.GroupID != 0 { diff --git a/internal/control/request_logs_test.go b/internal/control/request_logs_test.go index 38465f52..74102979 100644 --- a/internal/control/request_logs_test.go +++ b/internal/control/request_logs_test.go @@ -16,6 +16,8 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/config" "gpt-load/internal/pricing" "gpt-load/internal/protocol" @@ -53,6 +55,8 @@ func TestRequestLogEndpointRejectsUnknownDuplicateAndMalformedQueries(t *testing {name: "access key plus", query: "access_key_id=%2B1"}, {name: "access key whitespace", query: "access_key_id=%201"}, {name: "access key duplicate", query: "access_key_id=1&access_key_id=2"}, + {name: "legacy upstream key", query: "upstream_key_id=9"}, + {name: "legacy key", query: "key_id=9"}, {name: "limit leading zero", query: "limit=01"}, {name: "limit plus", query: "limit=%2B1"}, {name: "limit whitespace", query: "limit=%201"}, @@ -100,6 +104,9 @@ func TestRequestLogEndpointRejectsInvalidDomainValues(t *testing.T) { {name: "group unsafe", query: "group_id=9007199254740992"}, {name: "access key zero", query: "access_key_id=0"}, {name: "access key unsafe", query: "access_key_id=9007199254740992"}, + {name: "unknown channel", query: "channel_id=unknown"}, + {name: "credential zero", query: "credential_id=0"}, + {name: "credential unsafe", query: "credential_id=9007199254740992"}, {name: "empty client model", query: "client_model="}, {name: "empty upstream model", query: "upstream_model="}, {name: "equal range", query: "from_ms=" + equalTime + "&to_ms=" + equalTime}, @@ -154,7 +161,8 @@ func TestRequestLogEndpointParsesAdvancedFilters(t *testing.T) { "cost_state=priced", "pricing_completeness=partial", "cache_present=true", - "upstream_key_id=9", + "channel_id=openai", + "credential_id=9", "attempt_status_code=429", "failure_category=rate_limited", "error_code=provider_rate_limit", @@ -185,7 +193,8 @@ func TestRequestLogEndpointParsesAdvancedFilters(t *testing.T) { got.UsageState != usage.StatePartial || got.CostState != pricing.CostStatePriced || got.PricingCompleteness != pricing.CompletenessPartial || got.CachePresent == nil || !*got.CachePresent || - got.UpstreamKeyID == nil || *got.UpstreamKeyID != 9 || + got.ChannelID != channel.OpenAI || + got.CredentialID == nil || *got.CredentialID != 9 || got.AttemptStatusCode == nil || *got.AttemptStatusCode != 429 || got.FailureCategory != telemetry.FailureCategoryRateLimited || got.AttemptErrorCode != "provider_rate_limit" || got.RetryState != requestlog.RetryStateRetried || @@ -214,7 +223,8 @@ func TestRequestLogEndpointRejectsInvalidAdvancedFilters(t *testing.T) { "cost_state=unknown", "pricing_completeness=unknown", "cache_present=yes", - "upstream_key_id=0", + "channel_id=unknown", + "credential_id=0", "attempt_status_code=-1", "failure_category=unknown", "error_code=", @@ -386,11 +396,14 @@ func TestRequestLogDetailEndpointReturnsAttemptsAndFrozenPricingReceipt(t *testi rate := int64(1_000_000_000) amount := int64(1_000_000) receipt := &pricing.Receipt{ - SchemaVersion: 1, + SchemaVersion: 3, Method: pricing.ReceiptMethodUnitRateSum, MethodVersion: 1, Currency: "USD", - Rule: pricing.ReceiptRule{ScopeKey: "group:12", ModelID: "gpt-4.1"}, + Rule: pricing.ReceiptRule{ + ChannelID: string(channel.OpenAI), + ModelID: "gpt-4.1", + }, LineItems: []pricing.ReceiptLine{{ Code: "input", Quantity: 1000, @@ -405,6 +418,9 @@ func TestRequestLogDetailEndpointReturnsAttemptsAndFrozenPricingReceipt(t *testi requestID: { RequestID: requestID, CompletedAtMS: 1_784_894_400_000, + GroupID: 12, + ChannelID: channel.OpenAI, + CredentialID: 99, Protocol: protocol.OpenAICompletions, Status: telemetry.RequestStatusSuccess, StatusCode: http.StatusOK, @@ -415,16 +431,23 @@ func TestRequestLogDetailEndpointReturnsAttemptsAndFrozenPricingReceipt(t *testi UncachedInputTokens: 1000, EstimatedCostNanoUSD: amount, Attempts: []requestlog.Attempt{{ - Sequence: 1, - GroupID: 12, - GroupName: "Primary", - KeyID: 99, - UpstreamModel: "gpt-4.1", - StatusCode: http.StatusOK, - DurationMs: 1200, - FailureCategory: telemetry.FailureCategoryOK, - Action: telemetry.ActionTerminate, - PricingReceipt: receipt, + Sequence: 1, + GroupID: 12, + GroupName: "Primary", + ChannelID: channel.OpenAI, + CredentialID: 99, + Operation: execution.OperationChatCompletion, + RouteMode: channel.RouteNative, + UpstreamModel: "gpt-4.1", + UpstreamRequestID: "upstream-request-1", + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + DurationMs: 1200, + FailureCategory: telemetry.FailureCategoryOK, + Action: telemetry.ActionTerminate, + Committed: true, + PricingReceipt: receipt, }}, }, }} @@ -438,13 +461,26 @@ func TestRequestLogDetailEndpointReturnsAttemptsAndFrozenPricingReceipt(t *testi } var envelope struct { Data struct { - RequestID string `json:"request_id"` - Attempts []struct { - KeyID uint `json:"key_id"` - PricingReceipt struct { + RequestID string `json:"request_id"` + ChannelID string `json:"channel_id"` + CredentialID uint `json:"credential_id"` + Attempts []struct { + ChannelID string `json:"channel_id"` + CredentialID uint `json:"credential_id"` + Operation string `json:"operation"` + RouteMode string `json:"route_mode"` + UpstreamRequestID string `json:"upstream_request_id"` + DispatchState string `json:"dispatch_state"` + ResponseStarted bool `json:"response_started"` + Committed bool `json:"committed"` + PricingReceipt struct { Method string `json:"method"` TotalNanoUSD string `json:"total_nano_usd"` - LineItems []struct { + Rule struct { + ChannelID string `json:"channel_id"` + ModelID string `json:"model_id"` + } `json:"rule"` + LineItems []struct { Quantity string `json:"quantity"` AmountNanoUSD string `json:"amount_nano_usd"` } `json:"line_items"` @@ -455,17 +491,97 @@ func TestRequestLogDetailEndpointReturnsAttemptsAndFrozenPricingReceipt(t *testi if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil { t.Fatalf("decode response: %v", err) } - if envelope.Data.RequestID != requestID || len(envelope.Data.Attempts) != 1 || - envelope.Data.Attempts[0].KeyID != 99 || + if envelope.Data.RequestID != requestID || envelope.Data.ChannelID != string(channel.OpenAI) || + envelope.Data.CredentialID != 99 || len(envelope.Data.Attempts) != 1 || + envelope.Data.Attempts[0].ChannelID != string(channel.OpenAI) || + envelope.Data.Attempts[0].CredentialID != 99 || + envelope.Data.Attempts[0].Operation != string(execution.OperationChatCompletion) || + envelope.Data.Attempts[0].RouteMode != string(channel.RouteNative) || + envelope.Data.Attempts[0].UpstreamRequestID != "upstream-request-1" || + envelope.Data.Attempts[0].DispatchState != string(execution.DispatchMaybeSent) || + !envelope.Data.Attempts[0].ResponseStarted || !envelope.Data.Attempts[0].Committed || envelope.Data.Attempts[0].PricingReceipt.Method != pricing.ReceiptMethodUnitRateSum || + envelope.Data.Attempts[0].PricingReceipt.Rule.ChannelID != string(channel.OpenAI) || + envelope.Data.Attempts[0].PricingReceipt.Rule.ModelID != "gpt-4.1" || envelope.Data.Attempts[0].PricingReceipt.TotalNanoUSD != "1000000" || len(envelope.Data.Attempts[0].PricingReceipt.LineItems) != 1 || envelope.Data.Attempts[0].PricingReceipt.LineItems[0].Quantity != "1000" || envelope.Data.Attempts[0].PricingReceipt.LineItems[0].AmountNanoUSD != "1000000" { t.Fatalf("detail projection = %#v", envelope.Data) } - if strings.Contains(strings.ToLower(recorder.Body.String()), "committed") { - t.Fatalf("detail exposes internal committed flag: %s", recorder.Body.String()) + for _, legacyField := range []string{`"key_id"`, `"upstream_key_id"`} { + if strings.Contains(recorder.Body.String(), legacyField) { + t.Fatalf("detail exposes legacy field %s: %s", legacyField, recorder.Body.String()) + } + } +} + +func TestRequestLogPricingReceiptKeepsHistoricalSchemasReadable(t *testing.T) { + tests := []struct { + name string + receipt pricing.Receipt + wantScopeKey string + wantChannelID string + }{ + { + name: "v1 scoped identity", + receipt: pricing.Receipt{ + SchemaVersion: 1, + Method: pricing.ReceiptMethodUnitRateSum, + MethodVersion: 1, + Currency: "USD", + Rule: pricing.ReceiptRule{ + ScopeKey: "provider:openai", + ModelID: "gpt-4.1", + }, + }, + wantScopeKey: "provider:openai", + }, + { + name: "v2 global identity", + receipt: pricing.Receipt{ + SchemaVersion: 2, + Method: pricing.ReceiptMethodUnitRateSum, + MethodVersion: 1, + Currency: "USD", + Rule: pricing.ReceiptRule{ModelID: "gpt-4.1"}, + }, + }, + { + name: "v3 channel identity", + receipt: pricing.Receipt{ + SchemaVersion: 3, + Method: pricing.ReceiptMethodUnitRateSum, + MethodVersion: 1, + Currency: "USD", + Rule: pricing.ReceiptRule{ + ChannelID: string(channel.OpenAI), + ModelID: "gpt-4.1", + }, + }, + wantChannelID: string(channel.OpenAI), + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + mapped, err := mapRequestLogPricingReceipt(&test.receipt) + if err != nil { + t.Fatalf("mapRequestLogPricingReceipt() error = %v", err) + } + gotScopeKey := "" + if mapped.Rule.ScopeKey != nil { + gotScopeKey = *mapped.Rule.ScopeKey + } + gotChannelID := "" + if mapped.Rule.ChannelID != nil { + gotChannelID = *mapped.Rule.ChannelID + } + if mapped.SchemaVersion != test.receipt.SchemaVersion || + mapped.Rule.ModelID != test.receipt.Rule.ModelID || + gotScopeKey != test.wantScopeKey || gotChannelID != test.wantChannelID { + t.Fatalf("mapped receipt = %#v", mapped) + } + }) } } @@ -502,17 +618,19 @@ func TestRequestLogEndpointProjectsUsageCostAndNullGroupZero(t *testing.T) { var envelope struct { Data struct { Items []struct { - GroupID *uint `json:"group_id"` - UsageState string `json:"usage_state"` - CostState string `json:"cost_state"` - PricingCompleteness string `json:"pricing_completeness"` - InputTokens string `json:"input_tokens"` - CacheReadTokens string `json:"cache_read_tokens"` - CacheWrite5MTokens string `json:"cache_write_5m_tokens"` - CacheWrite1HTokens string `json:"cache_write_1h_tokens"` - CacheWriteUnknownTokens string `json:"cache_write_unknown_tokens"` - OutputTokens string `json:"output_tokens"` - EstimatedCostNanoUSD string `json:"estimated_cost_nano_usd"` + GroupID *uint `json:"group_id"` + ChannelID *string `json:"channel_id"` + CredentialID *uint `json:"credential_id"` + UsageState string `json:"usage_state"` + CostState string `json:"cost_state"` + PricingCompleteness string `json:"pricing_completeness"` + InputTokens string `json:"input_tokens"` + CacheReadTokens string `json:"cache_read_tokens"` + CacheWrite5MTokens string `json:"cache_write_5m_tokens"` + CacheWrite1HTokens string `json:"cache_write_1h_tokens"` + CacheWriteUnknownTokens string `json:"cache_write_unknown_tokens"` + OutputTokens string `json:"output_tokens"` + EstimatedCostNanoUSD string `json:"estimated_cost_nano_usd"` Reasoning *struct { Mode *string `json:"mode"` Effort *string `json:"effort"` @@ -528,7 +646,8 @@ func TestRequestLogEndpointProjectsUsageCostAndNullGroupZero(t *testing.T) { t.Fatalf("items = %#v", envelope.Data.Items) } item := envelope.Data.Items[0] - if item.GroupID != nil || item.UsageState != "complete" || item.CostState != "priced" || + if item.GroupID != nil || item.ChannelID != nil || item.CredentialID != nil || + item.UsageState != "complete" || item.CostState != "priced" || item.PricingCompleteness != "complete" || item.InputTokens != "16" || item.CacheReadTokens != "2" || item.CacheWrite5MTokens != "3" || item.CacheWrite1HTokens != "4" || @@ -752,6 +871,8 @@ func TestRequestLogEndpointsBindAccessKeyScopeAndRedactRoutingInternals(t *testi AttemptCount: 2, AffinityHit: true, GroupID: 99, + ChannelID: channel.OpenAI, + CredentialID: 101, UsageState: usage.StateComplete, CostState: pricing.CostStatePriced, PricingCompleteness: pricing.CompletenessComplete, @@ -759,7 +880,8 @@ func TestRequestLogEndpointsBindAccessKeyScopeAndRedactRoutingInternals(t *testi OutputTokens: 2, EstimatedCostNanoUSD: 50, Attempts: []requestlog.Attempt{{ - Sequence: 1, GroupID: 99, GroupName: "private group", KeyID: 101, + Sequence: 1, GroupID: 99, GroupName: "private group", + ChannelID: channel.OpenAI, CredentialID: 101, UpstreamModel: "private-upstream-model", }}, } @@ -788,6 +910,8 @@ func TestRequestLogEndpointsBindAccessKeyScopeAndRedactRoutingInternals(t *testi for _, query := range []string{ "access_key_id=" + strconv.FormatUint(uint64(other.ID), 10), "group_id=99", + "channel_id=openai", + "credential_id=101", "upstream_model=private-upstream-model", "retry_state=retried", } { @@ -853,6 +977,8 @@ func assertAccessKeyLogRedaction(t *testing.T, body []byte, detail bool) { "model_consistency": `"not_applicable"`, "affinity_hit": "false", "group_id": "null", + "channel_id": "null", + "credential_id": "null", "attempt_count": "0", } { if string(item[field]) != want { @@ -953,6 +1079,30 @@ func encodeTestCursorPayload(payload string) string { return base64.RawURLEncoding.EncodeToString([]byte(payload)) } +func TestRequestLogAttemptActionUsesCredentialTerms(t *testing.T) { + for input, want := range map[telemetry.Action]string{ + telemetry.Action("cooldown_credential"): "cooldown_credential", + telemetry.Action("fail_credential"): "fail_credential", + telemetry.ActionTerminate: string(telemetry.ActionTerminate), + } { + mapped, err := mapRequestLogAttempt(requestlog.Attempt{Action: input}) + if err != nil { + t.Fatalf("mapRequestLogAttempt(%q) error = %v", input, err) + } + if mapped.Action != want { + t.Fatalf("mapRequestLogAttempt(%q).Action = %q, want %q", input, mapped.Action, want) + } + encoded, err := json.Marshal(mapped) + if err != nil { + t.Fatalf("json.Marshal(%q) error = %v", input, err) + } + if strings.Contains(string(encoded), `"action":"cooldown_key"`) || + strings.Contains(string(encoded), `"action":"fail_key"`) { + t.Fatalf("attempt wire exposes legacy key action: %s", encoded) + } + } +} + func Example_requestLogOpaqueCursor() { payload := `{"v":2,"completed_at_ms":1784894400000,"request_id":"00000000-0000-4000-8000-000000000001"}` fmt.Println(encodeTestCursorPayload(payload)) diff --git a/internal/control/route_inspect.go b/internal/control/route_inspect.go index 4f6734bc..81dc4f9e 100644 --- a/internal/control/route_inspect.go +++ b/internal/control/route_inspect.go @@ -6,6 +6,7 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/execution" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/platform/response" "gpt-load/internal/protocol" @@ -25,8 +26,8 @@ type routeInspectAccessKeyResponse struct { Status state.AccessKeyStatus `json:"status"` } -type routeInspectKeyResponse struct { - KeyID uint `json:"key_id"` +type routeInspectCredentialResponse struct { + CredentialID uint `json:"credential_id"` Available bool `json:"available"` ReasonCode *scheduler.ReasonCode `json:"reason_code"` WeightManual *int `json:"weight_manual"` @@ -36,14 +37,14 @@ type routeInspectKeyResponse struct { } type routeInspectGroupResponse struct { - GroupID uint `json:"group_id"` - GroupName string `json:"group_name"` - UpstreamModel *string `json:"upstream_model"` - WeightManual *int `json:"weight_manual"` - Included bool `json:"included"` - Routable bool `json:"routable"` - ReasonCode *scheduler.ReasonCode `json:"reason_code"` - Keys []routeInspectKeyResponse `json:"keys"` + GroupID uint `json:"group_id"` + GroupName string `json:"group_name"` + UpstreamModel *string `json:"upstream_model"` + WeightManual *int `json:"weight_manual"` + Included bool `json:"included"` + Routable bool `json:"routable"` + ReasonCode *scheduler.ReasonCode `json:"reason_code"` + Credentials []routeInspectCredentialResponse `json:"credentials"` } type routeInspectResponse struct { @@ -94,9 +95,10 @@ func (service *Service) InspectRoute( observation.snapshot, observation.keys, scheduler.Query{ - Protocol: request.Protocol, - ExternalModel: cloneRouteModel(request.ExternalModel), - AccessKey: accessKey, + ClientProtocol: request.Protocol, + Operation: routeInspectOperation(request.Protocol, request.ExternalModel), + ExternalModel: cloneRouteModel(request.ExternalModel), + AccessKey: accessKey, }, observation.observedAt, ) @@ -117,6 +119,19 @@ func (service *Service) InspectRoute( ) } +func routeInspectOperation( + clientProtocol protocol.Protocol, + externalModel *string, +) execution.Operation { + if clientProtocol == protocol.OpenAIResponses { + if externalModel == nil { + return execution.OperationResponsesRetrieve + } + return execution.OperationResponsesCreate + } + return execution.OperationChatCompletion +} + func mapRouteInspectResponse( observation runtimeObservation, request routeInspectRequest, @@ -148,23 +163,23 @@ func mapRouteInspectResponse( Included: group.Included, Routable: group.Routable, ReasonCode: optionalReason(group.Reason), - Keys: []routeInspectKeyResponse{}, + Credentials: []routeInspectCredentialResponse{}, } - for _, key := range group.Keys { - cooldownUntilMS, err := optionalSafeEpochMilliseconds(key.CooldownUntil) + for _, credential := range group.Credentials { + cooldownUntilMS, err := optionalSafeEpochMilliseconds(credential.CooldownUntil) if err != nil { return routeInspectResponse{}, fmt.Errorf( "map route inspection cooldown_until_ms: %w", err, ) } - groupResponse.Keys = append(groupResponse.Keys, routeInspectKeyResponse{ - KeyID: key.KeyID, - Available: key.Available, - ReasonCode: optionalReason(key.Reason), - WeightManual: cloneInt(key.WeightManual), - WeightAuto: key.WeightAuto, - EffectiveWeight: key.EffectiveWeight, + groupResponse.Credentials = append(groupResponse.Credentials, routeInspectCredentialResponse{ + CredentialID: credential.CredentialID, + Available: credential.Available, + ReasonCode: optionalReason(credential.Reason), + WeightManual: cloneInt(credential.WeightManual), + WeightAuto: credential.WeightAuto, + EffectiveWeight: credential.EffectiveWeight, CooldownUntilMS: cooldownUntilMS, }) } diff --git a/internal/control/route_inspect_test.go b/internal/control/route_inspect_test.go index 72cb8f2f..b6dd139d 100644 --- a/internal/control/route_inspect_test.go +++ b/internal/control/route_inspect_test.go @@ -15,7 +15,8 @@ import ( "github.com/gin-gonic/gin" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/config" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/protocol" @@ -197,12 +198,10 @@ func TestRouteInspectSupportsResponsesWithoutModel(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ - ID: 1, - Name: "responses", - UpstreamURL: "https://responses.invalid", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Enabled: true, + ID: 1, Name: "responses", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Enabled: true, }}, AccessKeys: []state.AccessKeyConfig{{ ID: 10, @@ -213,10 +212,10 @@ func TestRouteInspectSupportsResponsesWithoutModel(t *testing.T) { }); err != nil { t.Fatalf("Publish() error = %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{{ - ID: 21, - GroupID: 1, - Status: state.KeyStatusActive, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 21, GroupID: 1, Version: 1, IdentityGeneration: 21, + Fingerprint: "test-21", + Status: state.CredentialStatusActive, EncryptedValue: "cipher", }}); err != nil { t.Fatalf("Replace() error = %v", err) @@ -255,12 +254,10 @@ func TestRouteInspectSupportsResponsesWithoutModel(t *testing.T) { func TestRouteInspectWithoutModelReturnsSharedFilterReason(t *testing.T) { fixture := newServiceFixture(t) if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ - ID: 1, - Name: "responses", - UpstreamURL: "https://responses.invalid", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Enabled: true, + ID: 1, Name: "responses", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Enabled: true, }}, AccessKeys: []state.AccessKeyConfig{{ ID: 10, @@ -297,18 +294,19 @@ func TestRouteInspectEndpointReturnsCurrentSafeExplanation(t *testing.T) { fixture.service.now = func() time.Time { return now } groupWeight := 20 if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{ { - ID: 2, Name: "backup", UpstreamURL: "https://backup.invalid", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 2, Name: "backup", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "provider-backup", Alias: "public-model"}}, WeightManual: &groupWeight, Enabled: true, }, { - ID: 1, Name: "primary", UpstreamURL: "https://not-called.invalid", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "provider-model", Alias: "public-model"}}, - Enabled: true, + ID: 1, Name: "primary", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "provider-model", Alias: "public-model"}}, + Enabled: true, }, }, AccessKeys: []state.AccessKeyConfig{{ @@ -319,18 +317,18 @@ func TestRouteInspectEndpointReturnsCurrentSafeExplanation(t *testing.T) { t.Fatalf("Publish() error = %v", err) } keyWeight := 25 - if err := fixture.registry.Replace([]state.KeyEntry{ + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{ { - ID: 31, GroupID: 2, Status: state.KeyStatusActive, + ID: 31, GroupID: 2, Version: 1, IdentityGeneration: 31, Fingerprint: "test-31", Status: state.CredentialStatusActive, WeightAuto: 30, EncryptedValue: "cipher-three", }, { - ID: 22, GroupID: 1, Status: state.KeyStatusActive, + ID: 22, GroupID: 1, Version: 1, IdentityGeneration: 22, Fingerprint: "test-22", Status: state.CredentialStatusActive, CooldownUntil: now.Add(time.Minute), WeightAuto: 40, EncryptedValue: "cipher-two", }, { - ID: 21, GroupID: 1, Status: state.KeyStatusActive, + ID: 21, GroupID: 1, Version: 1, IdentityGeneration: 21, Fingerprint: "test-21", Status: state.CredentialStatusActive, WeightManual: &keyWeight, WeightAuto: 90, EncryptedValue: "cipher-one", }, @@ -364,36 +362,36 @@ func TestRouteInspectEndpointReturnsCurrentSafeExplanation(t *testing.T) { routeModelValue(primary.UpstreamModel) != "provider-model" || primary.WeightManual != nil || !primary.Included || !primary.Routable || primary.ReasonCode != nil || - len(primary.Keys) != 2 || - primary.Keys[0].KeyID != 21 || primary.Keys[1].KeyID != 22 { + len(primary.Credentials) != 2 || + primary.Credentials[0].CredentialID != 21 || primary.Credentials[1].CredentialID != 22 { t.Fatalf("primary group = %#v", primary) } - available := primary.Keys[0] + available := primary.Credentials[0] if !available.Available || available.ReasonCode != nil || available.WeightManual == nil || *available.WeightManual != 25 || available.WeightAuto != 90 || available.EffectiveWeight != 50*25 || available.CooldownUntilMS != nil { t.Fatalf("available key = %#v", available) } - cooldown := primary.Keys[1] + cooldown := primary.Credentials[1] if cooldown.Available || cooldown.WeightManual != nil || cooldown.WeightAuto != 40 || cooldown.EffectiveWeight != 0 || cooldown.CooldownUntilMS == nil || *cooldown.CooldownUntilMS != now.Add(time.Minute).UnixMilli() { t.Fatalf("cooldown key = %#v", cooldown) } - assertRouteReason(t, cooldown.ReasonCode, scheduler.ReasonKeyCooldown) + assertRouteReason(t, cooldown.ReasonCode, scheduler.ReasonCredentialCooldown) backup := got.Groups[1] if backup.GroupName != "backup" || routeModelValue(backup.UpstreamModel) != "provider-backup" || backup.WeightManual == nil || *backup.WeightManual != 20 || !backup.Included || !backup.Routable || backup.ReasonCode != nil || - len(backup.Keys) != 1 || backup.Keys[0].KeyID != 31 || - !backup.Keys[0].Available || backup.Keys[0].ReasonCode != nil || - backup.Keys[0].WeightManual != nil || - backup.Keys[0].WeightAuto != 30 || - backup.Keys[0].EffectiveWeight != 20*30 || - backup.Keys[0].CooldownUntilMS != nil { + len(backup.Credentials) != 1 || backup.Credentials[0].CredentialID != 31 || + !backup.Credentials[0].Available || backup.Credentials[0].ReasonCode != nil || + backup.Credentials[0].WeightManual != nil || + backup.Credentials[0].WeightAuto != 30 || + backup.Credentials[0].EffectiveWeight != 20*30 || + backup.Credentials[0].CooldownUntilMS != nil { t.Fatalf("backup group = %#v", backup) } body := recorder.Body.String() @@ -449,14 +447,15 @@ func TestRouteInspectEndpointReturnsFilterExplanations(t *testing.T) { fixture.service.now = func() time.Time { return now } manual := 15 if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{ { - ID: 2, Name: "second", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 2, Name: "second", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "provider-two", Alias: "public-model"}}, WeightManual: &manual, Enabled: true, }, { - ID: 1, Name: "first", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 1, Name: "first", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "provider-one", Alias: "public-model"}}, Enabled: true, }, @@ -468,9 +467,9 @@ func TestRouteInspectEndpointReturnsFilterExplanations(t *testing.T) { }); err != nil { t.Fatalf("Publish() error = %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{ - {ID: 22, GroupID: 2, Status: state.KeyStatusActive, EncryptedValue: "two"}, - {ID: 11, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "one"}, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{ + {ID: 22, GroupID: 2, Version: 1, IdentityGeneration: 22, Fingerprint: "test-22", Status: state.CredentialStatusActive, EncryptedValue: "two"}, + {ID: 11, GroupID: 1, Version: 1, IdentityGeneration: 11, Fingerprint: "test-11", Status: state.CredentialStatusActive, EncryptedValue: "one"}, }); err != nil { t.Fatalf("Replace() error = %v", err) } @@ -504,7 +503,7 @@ func TestRouteInspectEndpointReturnsFilterExplanations(t *testing.T) { } for index, group := range got.Groups { if group.Included || group.Routable || - group.Keys == nil || len(group.Keys) != 0 { + group.Credentials == nil || len(group.Credentials) != 0 { t.Fatalf("filtered group %d = %#v", index, group) } assertRouteReason(t, group.ReasonCode, scheduler.ReasonGroupFiltered) @@ -526,8 +525,9 @@ func TestRouteInspectEndpointReturnsNoRouteTargetExplanation(t *testing.T) { now := healthNow() fixture.service.now = func() time.Time { return now } if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ - ID: 1, Name: "primary", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 1, Name: "primary", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "configured-model"}}, Enabled: true, }}, AccessKeys: []state.AccessKeyConfig{{ @@ -563,8 +563,9 @@ func TestRouteInspectEndpointReturnsNoAvailableKeyExplanation(t *testing.T) { fixture.service.now = func() time.Time { return now } groupWeight := 25 if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ - ID: 1, Name: "primary", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 1, Name: "primary", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "provider-model", Alias: "public-model"}}, WeightManual: &groupWeight, Enabled: true, }}, @@ -582,21 +583,21 @@ func TestRouteInspectEndpointReturnsNoAvailableKeyExplanation(t *testing.T) { if cooldownAt.Location() == time.UTC { t.Fatal("cooldown fixture must use a non-UTC location") } - if err := fixture.registry.Replace([]state.KeyEntry{ + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{ { - ID: 14, GroupID: 1, Status: state.KeyStatusActive, + ID: 14, GroupID: 1, Version: 1, IdentityGeneration: 14, Fingerprint: "test-14", Status: state.CredentialStatusActive, WeightAuto: 70, CooldownUntil: cooldownAt, EncryptedValue: "cooldown", }, { - ID: 12, GroupID: 1, Status: state.KeyStatusActive, + ID: 12, GroupID: 1, Version: 1, IdentityGeneration: 12, Fingerprint: "test-12", Status: state.CredentialStatusActive, WeightManual: &zero, WeightAuto: 45, EncryptedValue: "zero", }, { - ID: 13, GroupID: 1, Status: state.KeyStatusActive, + ID: 13, GroupID: 1, Version: 1, IdentityGeneration: 13, Fingerprint: "test-13", Status: state.CredentialStatusActive, WeightAuto: 60, Blacklisted: true, EncryptedValue: "blacklisted", }, { - ID: 11, GroupID: 1, Status: state.KeyStatusDisabled, + ID: 11, GroupID: 1, Version: 1, IdentityGeneration: 11, Fingerprint: "test-11", Status: state.CredentialStatusDisabled, WeightManual: &disabledManual, WeightAuto: 30, EncryptedValue: "disabled", }, }); err != nil { @@ -615,7 +616,7 @@ func TestRouteInspectEndpointReturnsNoAvailableKeyExplanation(t *testing.T) { got.Routable { t.Fatalf("unavailable response = %#v", got) } - assertRouteReason(t, got.ReasonCode, scheduler.ReasonNoAvailableKey) + assertRouteReason(t, got.ReasonCode, scheduler.ReasonNoAvailableCredential) if len(got.Groups) != 1 { t.Fatalf("groups = %#v", got.Groups) } @@ -623,40 +624,40 @@ func TestRouteInspectEndpointReturnsNoAvailableKeyExplanation(t *testing.T) { if group.GroupID != 1 || group.GroupName != "primary" || routeModelValue(group.UpstreamModel) != "provider-model" || group.WeightManual == nil || *group.WeightManual != 25 || - !group.Included || group.Routable || len(group.Keys) != 4 { + !group.Included || group.Routable || len(group.Credentials) != 4 { t.Fatalf("unavailable group = %#v", group) } - assertRouteReason(t, group.ReasonCode, scheduler.ReasonNoAvailableKey) + assertRouteReason(t, group.ReasonCode, scheduler.ReasonNoAvailableCredential) wantReasons := []scheduler.ReasonCode{ - scheduler.ReasonKeyDisabled, - scheduler.ReasonKeyWeightZero, - scheduler.ReasonKeyBlacklisted, - scheduler.ReasonKeyCooldown, + scheduler.ReasonCredentialDisabled, + scheduler.ReasonCredentialWeightZero, + scheduler.ReasonCredentialBlacklisted, + scheduler.ReasonCredentialCooldown, } wantManual := []*int{&disabledManual, &zero, nil, nil} wantAuto := []int{30, 45, 60, 70} - for index, key := range group.Keys { - if key.KeyID != uint(11+index) || key.Available || - key.EffectiveWeight != 0 || - key.WeightAuto != wantAuto[index] { - t.Fatalf("unavailable key %d = %#v", index, key) + for index, credential := range group.Credentials { + if credential.CredentialID != uint(11+index) || credential.Available || + credential.EffectiveWeight != 0 || + credential.WeightAuto != wantAuto[index] { + t.Fatalf("unavailable credential %d = %#v", index, credential) } - assertRouteReason(t, key.ReasonCode, wantReasons[index]) + assertRouteReason(t, credential.ReasonCode, wantReasons[index]) if wantManual[index] == nil { - if key.WeightManual != nil { - t.Fatalf("key %d manual weight = %v, want nil", index, key.WeightManual) + if credential.WeightManual != nil { + t.Fatalf("key %d manual weight = %v, want nil", index, credential.WeightManual) } - } else if key.WeightManual == nil || - *key.WeightManual != *wantManual[index] { - t.Fatalf("key %d manual weight = %v, want %d", index, key.WeightManual, *wantManual[index]) + } else if credential.WeightManual == nil || + *credential.WeightManual != *wantManual[index] { + t.Fatalf("key %d manual weight = %v, want %d", index, credential.WeightManual, *wantManual[index]) } if index == 3 { - if key.CooldownUntilMS == nil || - *key.CooldownUntilMS != cooldownAt.UnixMilli() { - t.Fatalf("cooldown = %v, want %v", key.CooldownUntilMS, cooldownAt.UnixMilli()) + if credential.CooldownUntilMS == nil || + *credential.CooldownUntilMS != cooldownAt.UnixMilli() { + t.Fatalf("cooldown = %v, want %v", credential.CooldownUntilMS, cooldownAt.UnixMilli()) } - } else if key.CooldownUntilMS != nil { - t.Fatalf("key %d cooldown = %v, want nil", index, key.CooldownUntilMS) + } else if credential.CooldownUntilMS != nil { + t.Fatalf("key %d cooldown = %v, want nil", index, credential.CooldownUntilMS) } } if strings.Count(recorder.Body.String(), `"cooldown_until_ms":null`) != 3 { @@ -758,20 +759,11 @@ func (spy *routeInspectEncryptionSpy) Hash(string) string { func TestRouteInspectNeverCallsUpstreamOrMutatesRuntime(t *testing.T) { var upstreamCalls atomic.Int64 - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - request *http.Request, - ) { - upstreamCalls.Add(1) - writer.WriteHeader(http.StatusNoContent) - })) - defer upstream.Close() - fixture := newServiceFixture(t) encryptionSpy := &routeInspectEncryptionSpy{} fixture.service.encryption = encryptionSpy dialectCalls := 0 - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func( context.Context, @@ -784,10 +776,10 @@ func TestRouteInspectNeverCallsUpstreamOrMutatesRuntime(t *testing.T) { }, }) if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ - ID: 1, Name: "upstream", UpstreamURL: upstream.URL, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, + ID: 1, Name: "upstream", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }}, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: "hash", @@ -796,8 +788,8 @@ func TestRouteInspectNeverCallsUpstreamOrMutatesRuntime(t *testing.T) { }); err != nil { t.Fatalf("Publish() error = %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, EncryptedValue: "cipher", }}); err != nil { t.Fatalf("Replace() error = %v", err) @@ -859,8 +851,9 @@ func TestRouteInspectCatalogMismatchReturnsInternalServerError(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ - ID: 1, Name: "primary", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 1, Name: "primary", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }}, AccessKeys: []state.AccessKeyConfig{{ @@ -869,7 +862,7 @@ func TestRouteInspectCatalogMismatchReturnsInternalServerError(t *testing.T) { }); err != nil { t.Fatalf("Publish() error = %v", err) } - fixture.manager.Current().RouteCatalog[protocol.OpenAICompletions]["model"] = []state.RouteTarget{{ + fixture.manager.Current().ExecutionRouteCatalog[protocol.OpenAICompletions][execution.OperationChatCompletion]["model"] = []state.RouteTarget{{ GroupID: 999, UpstreamModelID: "model", }} engine := gin.New() diff --git a/internal/control/runtime.go b/internal/control/runtime.go index 693bc070..d866a755 100644 --- a/internal/control/runtime.go +++ b/internal/control/runtime.go @@ -6,7 +6,8 @@ import ( "sync" "time" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/health" "gpt-load/internal/platform/encryption" "gpt-load/internal/state" @@ -20,11 +21,11 @@ const ( ) type autoWeightRegistry interface { - ActiveKeyIDs() []uint - SetAutoWeight(keyID uint, weight int) bool + ActiveCredentialIDs() []uint + SetAutoWeight(credentialID uint, weight int) bool } -type keyMutationCoordinator interface { +type credentialMutationCoordinator interface { Do(uint, func()) } @@ -62,7 +63,7 @@ func (ticker standardRuntimeTicker) Stop() { type Runtime struct { registry autoWeightRegistry stats *health.StatsStore - mutations keyMutationCoordinator + mutations credentialMutationCoordinator validator validationSweep requestLogCleaner RequestLogCleaner operationRecovery operationRecoveryRuntime @@ -75,12 +76,13 @@ type Runtime struct { } func NewRuntime( - registry *state.KeyRegistry, + registry *state.CredentialRegistry, stats *health.StatsStore, mutations *health.MutationCoordinator, manager *state.Manager, encryptionService encryption.Service, - dialects dialect.Set, + channelRegistry *channel.Registry, + executor execution.Executor, requestLogCleaner RequestLogCleaner, operationRecovery *Service, catalogSync *CatalogSyncCoordinator, @@ -108,7 +110,8 @@ func NewRuntime( stats, mutations, encryptionService, - dialects, + channelRegistry, + executor, ) return runtime } @@ -204,10 +207,10 @@ func (runtime *Runtime) runRetention(ctx context.Context, ticker runtimeTicker) } func (runtime *Runtime) recompute(now time.Time) { - for _, keyID := range runtime.registry.ActiveKeyIDs() { - runtime.mutations.Do(keyID, func() { - stats := runtime.stats.Snapshot(keyID, now) - runtime.registry.SetAutoWeight(keyID, calculateAutoWeight(stats)) + for _, credentialID := range runtime.registry.ActiveCredentialIDs() { + runtime.mutations.Do(credentialID, func() { + stats := runtime.stats.Snapshot(credentialID, now) + runtime.registry.SetAutoWeight(credentialID, calculateAutoWeight(stats)) }) } } diff --git a/internal/control/runtime_observation.go b/internal/control/runtime_observation.go index f8631975..af5d6dab 100644 --- a/internal/control/runtime_observation.go +++ b/internal/control/runtime_observation.go @@ -11,7 +11,7 @@ import ( type runtimeObservation struct { observedAt time.Time snapshot *state.ConfigSnapshot - keys []state.KeyRuntimeView + keys []state.CredentialRuntimeView } type runtimeHealthObservation struct { @@ -96,7 +96,7 @@ func (service *Service) captureRuntimeHealthObservation() ( if bucket != healthBucketCooldown && bucket != healthBucketBlacklisted { continue } - ciphertext, exists := service.registry.EncryptedValue(key.ID) + ciphertext, exists := service.registry.EncryptedCredentialData(key.ID) if !exists || ciphertext == "" { return runtimeHealthObservation{}, fmt.Errorf( "capture runtime health observation: key %d ciphertext unavailable: %w", diff --git a/internal/control/runtime_observation_test.go b/internal/control/runtime_observation_test.go index 67d0cca9..30538542 100644 --- a/internal/control/runtime_observation_test.go +++ b/internal/control/runtime_observation_test.go @@ -1,14 +1,15 @@ package control import ( + "encoding/json" "sync" "testing" "time" "gorm.io/gorm" + "gpt-load/internal/channel" "gpt-load/internal/platform/encryption" - "gpt-load/internal/protocol" "gpt-load/internal/requestlog" "gpt-load/internal/state" "gpt-load/internal/storage/models" @@ -27,7 +28,7 @@ func (probe healthDecryptLockProbe) Decrypt(ciphertext string) (string, error) { func TestCaptureRuntimeObservationWaitsForPublishedConfigPair(t *testing.T) { fixture := newServiceFixture(t) group := validControlGroup("runtime-observation") - var key models.UpstreamKey + var key models.Credential runtimeApplied := make(chan struct{}) allowPublish := make(chan struct{}) var releaseOnce sync.Once @@ -41,16 +42,15 @@ func TestCaptureRuntimeObservationWaitsForPublishedConfigPair(t *testing.T) { if err := tx.Create(group).Error; err != nil { return err } - key = models.UpstreamKey{ - GroupID: group.ID, KeyValue: "cipher-runtime-observation", - KeyHash: "hash-runtime-observation", - Status: models.UpstreamKeyStatusActive, + key = models.Credential{ + GroupID: group.ID, Data: "cipher-runtime-observation", + Fingerprint: "hash-runtime-observation", Status: models.CredentialStatusActive, } return tx.Create(&key).Error }, func() error { - if err := fixture.registry.ApplyImport(group.ID, []state.KeyEntry{{ - ID: key.ID, GroupID: group.ID, Status: state.KeyStatusActive, - EncryptedValue: key.KeyValue, + if err := fixture.registry.ApplyCredentialImport(group.ID, []state.CredentialEntry{{ + ID: key.ID, GroupID: group.ID, Status: state.CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: key.Fingerprint, EncryptedValue: key.Data, }}); err != nil { return err } @@ -114,7 +114,7 @@ func TestCaptureRuntimeObservationWaitsForPublishedConfigPair(t *testing.T) { func TestCaptureRuntimeHealthObservationWaitsForPublishedConfigPair(t *testing.T) { fixture := newServiceFixture(t) group := validControlGroup("runtime-health-observation") - var key models.UpstreamKey + var key models.Credential runtimeApplied := make(chan struct{}) allowPublish := make(chan struct{}) var releaseOnce sync.Once @@ -128,16 +128,15 @@ func TestCaptureRuntimeHealthObservationWaitsForPublishedConfigPair(t *testing.T if err := tx.Create(group).Error; err != nil { return err } - key = models.UpstreamKey{ - GroupID: group.ID, KeyValue: "cipher-runtime-health-observation", - KeyHash: "hash-runtime-health-observation", - Status: models.UpstreamKeyStatusActive, + key = models.Credential{ + GroupID: group.ID, Data: "cipher-runtime-health-observation", + Fingerprint: "hash-runtime-health-observation", Status: models.CredentialStatusActive, } return tx.Create(&key).Error }, func() error { - if err := fixture.registry.ApplyImport(group.ID, []state.KeyEntry{{ - ID: key.ID, GroupID: group.ID, Status: state.KeyStatusActive, - EncryptedValue: key.KeyValue, + if err := fixture.registry.ApplyCredentialImport(group.ID, []state.CredentialEntry{{ + ID: key.ID, GroupID: group.ID, Status: state.CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: key.Fingerprint, EncryptedValue: key.Data, }}); err != nil { return err } @@ -222,14 +221,14 @@ func TestRuntimeHealthReleasesReadLockBeforeDecryptingProblemKeys(t *testing.T) fixture := newServiceFixture(t) now := healthNow() fixture.service.now = func() time.Time { return now } - if _, err := fixture.manager.Publish(state.CompileInput{Groups: []state.GroupConfig{{ - ID: 1, Name: "health-lock", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + if _, err := fixture.manager.Publish(state.CompileInput{ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ + ID: 1, Name: "health-lock", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "model"}}, Enabled: true, }}}); err != nil { t.Fatalf("Publish() error = %v", err) } - if err := fixture.registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, CooldownUntil: now.Add(time.Minute), EncryptedValue: encryptHealthKey(t, fixture, "health-lock-secret-safe"), }}); err != nil { diff --git a/internal/control/runtime_test.go b/internal/control/runtime_test.go index a8916266..12422b4e 100644 --- a/internal/control/runtime_test.go +++ b/internal/control/runtime_test.go @@ -8,7 +8,7 @@ import ( "testing" "time" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" "gpt-load/internal/health" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -35,7 +35,7 @@ func newFakeAutoWeightRegistry(activeIDs ...uint) *fakeAutoWeightRegistry { } } -func (registry *fakeAutoWeightRegistry) ActiveKeyIDs() []uint { +func (registry *fakeAutoWeightRegistry) ActiveCredentialIDs() []uint { registry.mu.Lock() defer registry.mu.Unlock() return append([]uint(nil), registry.activeIDs...) @@ -425,9 +425,9 @@ func TestRuntimeContinuesWhenKeyDisappears(t *testing.T) { func TestRuntimeCooldownProblemDoesNotAffectAutoWeightSeenByCandidateCollection(t *testing.T) { base := time.Date(2026, time.July, 22, 12, 0, 0, 0, time.UTC) - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 10, Status: state.KeyStatusActive, EncryptedValue: "cipher-one", + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, EncryptedValue: "cipher-one", }}); err != nil { t.Fatalf("Replace() error = %v", err) } @@ -441,7 +441,7 @@ func TestRuntimeCooldownProblemDoesNotAffectAutoWeightSeenByCandidateCollection( registry: registry, stats: stats, mutations: health.NewMutationCoordinator(), } runtime.recompute(base) - candidates := registry.CollectCandidates([]uint{10}, nil, base) + candidates := registry.CollectCredentialCandidates([]uint{10}, nil, base) if len(candidates) != 1 || candidates[0].WeightAuto != 92 { t.Fatalf("CollectCandidates() = %#v, want one candidate with WeightAuto 92", candidates) } @@ -511,7 +511,8 @@ func TestRuntimeCoordinatesAutoWeightWithValidationRecovery(t *testing.T) { stats: stats, mutations: mutations, decryptor: validationDecryptor{}, - dialects: dialect.Set{protocol.OpenAICompletions: &validationTestDialect{protocol: protocol.OpenAICompletions, probes: probes}}, + channels: channel.NewRegistry(), + executor: &validationTestExecutor{probes: probes}, } autoDone := make(chan struct{}) @@ -729,7 +730,7 @@ func newInterleavingRegistry() *interleavingRegistry { } } -func (*interleavingRegistry) ActiveKeyIDs() []uint { +func (*interleavingRegistry) ActiveCredentialIDs() []uint { return []uint{1} } @@ -744,11 +745,11 @@ func (registry *interleavingRegistry) SetAutoWeight(_ uint, weight int) bool { return true } -func (*interleavingRegistry) BlacklistedKeys() []state.KeyRef { - return []state.KeyRef{{ID: 1, GroupID: 1, EncryptedValue: "cipher-one"}} +func (*interleavingRegistry) BlacklistedCredentials() []state.CredentialRef { + return []state.CredentialRef{{ID: 1, GroupID: 1, EncryptedValue: "cipher-one"}} } -func (registry *interleavingRegistry) RecoverIfMatch(_ state.KeyRef, weight int) bool { +func (registry *interleavingRegistry) RecoverIfMatch(_ state.CredentialRef, weight int) bool { registry.mu.Lock() registry.currentWeight = weight registry.recoveries++ diff --git a/internal/control/server.go b/internal/control/server.go index e9a0fe13..698d557e 100644 --- a/internal/control/server.go +++ b/internal/control/server.go @@ -206,150 +206,146 @@ func (s *Server) handleDeleteGroup(c *gin.Context) { response.SuccessI18n(c, "common.success", nil) } -func (s *Server) handleListGroupKeys(c *gin.Context) { - id, ok := groupID(c, "list_group_keys") +func (s *Server) handleListGroupCredentials(c *gin.Context) { + id, ok := groupID(c, "list_group_credentials") if !ok { return } - query, apiErr := parseGroupKeyCollectionQuery(c.Request.URL.RawQuery) + query, apiErr := parseCredentialCollectionQuery(c.Request.URL.RawQuery) if apiErr != nil { - writeServiceError(c, "list_group_keys", apiErr) + writeServiceError(c, "list_group_credentials", apiErr) return } - result, err := s.service.ListGroupKeys(c.Request.Context(), id, query) + result, err := s.service.ListGroupCredentials(c.Request.Context(), id, query) if err != nil { - writeServiceError(c, "list_group_keys", err) + writeServiceError(c, "list_group_credentials", err) return } response.SuccessI18n(c, "common.success", result) } -func (s *Server) handleRevealGroupKey(c *gin.Context) { - groupID, ok := groupID(c, "reveal_group_key") +func (s *Server) handleRevealGroupCredential(c *gin.Context) { + groupID, ok := groupID(c, "reveal_group_credential") if !ok { return } - keyID, ok := keyID(c, "reveal_group_key") + credentialID, ok := credentialID(c, "reveal_group_credential") if !ok { return } - result, err := s.service.RevealGroupKey(c.Request.Context(), groupID, keyID) + if err := bindOptionalEmptyJSONObject(c); err != nil { + writeServiceError(c, "reveal_group_credential", mapControlJSONError(err)) + return + } + result, err := s.service.RevealGroupCredential(c.Request.Context(), groupID, credentialID) if err != nil { - writeServiceError(c, "reveal_group_key", err) + writeServiceError(c, "reveal_group_credential", err) return } setSecretResponseHeaders(c) response.SuccessI18n(c, "common.success", result) } -func (s *Server) handleUpdateGroupKey(c *gin.Context) { - groupID, ok := groupID(c, "update_group_key") +func (s *Server) handleUpdateGroupCredential(c *gin.Context) { + groupID, ok := groupID(c, "update_group_credential") if !ok { return } - keyID, ok := keyID(c, "update_group_key") + credentialID, ok := credentialID(c, "update_group_credential") if !ok { return } - var request UpstreamKeyUpdateRequest + var request CredentialUpdateRequest if err := bindStrictJSON(c, &request); err != nil { - writeServiceError(c, "update_group_key", mapControlJSONError(err)) + writeServiceError(c, "update_group_credential", mapControlJSONError(err)) return } - result, err := s.service.UpdateGroupKey( - c.Request.Context(), - groupID, - keyID, - request, - ) + result, err := s.service.UpdateGroupCredential(c.Request.Context(), groupID, credentialID, request) if err != nil { - writeServiceError(c, "update_group_key", err) + writeServiceError(c, "update_group_credential", err) return } response.SuccessI18n(c, "common.success", result) } -func (s *Server) handleDeleteGroupKey(c *gin.Context) { - groupID, ok := groupID(c, "delete_group_key") +func (s *Server) handleDeleteGroupCredential(c *gin.Context) { + groupID, ok := groupID(c, "delete_group_credential") if !ok { return } - keyID, ok := keyID(c, "delete_group_key") + credentialID, ok := credentialID(c, "delete_group_credential") if !ok { return } - if err := s.service.DeleteGroupKey( - c.Request.Context(), - groupID, - keyID, - ); err != nil { - writeServiceError(c, "delete_group_key", err) + if err := bindOptionalEmptyJSONObject(c); err != nil { + writeServiceError(c, "delete_group_credential", mapControlJSONError(err)) + return + } + if err := s.service.DeleteGroupCredential(c.Request.Context(), groupID, credentialID); err != nil { + writeServiceError(c, "delete_group_credential", err) return } response.SuccessI18n(c, "common.success", nil) } -func (s *Server) handleRestoreGroupKey(c *gin.Context) { - groupID, ok := groupID(c, "restore_group_key") +func (s *Server) handleRestoreGroupCredential(c *gin.Context) { + groupID, ok := groupID(c, "restore_group_credential") if !ok { return } - keyID, ok := keyID(c, "restore_group_key") + credentialID, ok := credentialID(c, "restore_group_credential") if !ok { return } if err := bindOptionalEmptyJSONObject(c); err != nil { - writeServiceError(c, "restore_group_key", mapControlJSONError(err)) + writeServiceError(c, "restore_group_credential", mapControlJSONError(err)) return } - result, err := s.service.RestoreGroupKey(c.Request.Context(), groupID, keyID) + result, err := s.service.RestoreGroupCredential(c.Request.Context(), groupID, credentialID) if err != nil { - writeServiceError(c, "restore_group_key", err) + writeServiceError(c, "restore_group_credential", err) return } response.SuccessI18n(c, "common.success", result) } -func (s *Server) handleBatchGroupKeys(c *gin.Context) { - groupID, ok := groupID(c, "batch_group_keys") +func (s *Server) handleBatchGroupCredentials(c *gin.Context) { + groupID, ok := groupID(c, "batch_group_credentials") if !ok { return } - var request GroupKeyBatchRequest + var request CredentialBatchRequest if err := bindStrictJSON(c, &request); err != nil { - writeServiceError(c, "batch_group_keys", mapControlJSONError(err)) + writeServiceError(c, "batch_group_credentials", mapControlJSONError(err)) return } - result, err := s.service.BatchGroupKeys(c.Request.Context(), groupID, request) + result, err := s.service.BatchGroupCredentials(c.Request.Context(), groupID, request) if err != nil { - writeServiceError(c, "batch_group_keys", err) + writeServiceError(c, "batch_group_credentials", err) return } response.SuccessI18n(c, "common.success", result) } -func (s *Server) handleImportGroupKeys(c *gin.Context) { - id, ok := groupID(c, "import_group_keys") +func (s *Server) handleImportGroupCredentials(c *gin.Context) { + id, ok := groupID(c, "import_group_credentials") if !ok { return } - idempotencyKey, ok := requiredIdempotencyKey(c, "import_group_keys") + idempotencyKey, ok := requiredIdempotencyKey(c, "import_group_credentials") if !ok { return } - var request GroupKeyImportRequest + var request CredentialImportRequest if err := bindStrictJSON(c, &request); err != nil { - writeServiceError(c, "import_group_keys", mapControlJSONError(err)) + writeServiceError(c, "import_group_credentials", mapControlJSONError(err)) return } - result, err := s.service.ImportGroupKeysIdempotent( - c.Request.Context(), - idempotencyKey, - id, - request, + result, err := s.service.ImportGroupCredentialsIdempotent( + c.Request.Context(), idempotencyKey, id, request, ) if err != nil { - writeServiceError(c, "import_group_keys", err) + writeServiceError(c, "import_group_credentials", err) return } response.SuccessI18n(c, "common.success", result) @@ -589,8 +585,8 @@ func groupID(c *gin.Context, operation string) (uint, bool) { return uint(parsed), true } -func keyID(c *gin.Context, operation string) (uint, bool) { - parsed, err := strconv.ParseUint(c.Param("key_id"), 10, strconv.IntSize) +func credentialID(c *gin.Context, operation string) (uint, bool) { + parsed, err := strconv.ParseUint(c.Param("credential_id"), 10, strconv.IntSize) if err != nil || parsed == 0 { writeServiceError(c, operation, app_errors.ErrBadRequest) return 0, false @@ -672,30 +668,28 @@ func serviceErrorMessageID( if resourceErr.resource == "group" { return "group.not_found" } - return "key.not_found" + return "credential.not_found" } switch operation { case "list_groups", "get_group_summary", "get_group_settings", "get_group_models", "update_group_settings", "delete_group", - "update_group_models", "import_group_keys", - "discover_group_models", "list_group_keys": + "update_group_models", "import_group_credentials", + "discover_group_models", "list_group_credentials": return "group.not_found" default: - return "key.not_found" + return "credential.not_found" } - case app_errors.ErrNoActiveUpstreamKey.Code: - return "group.no_active_upstream_key" + case app_errors.ErrNoActiveCredential.Code: + return "group.no_active_credential" case app_errors.ErrDuplicateResource.Code: if operation == "create_group" || operation == "update_group_settings" { return "group.name_exists" } return "bad_request" - case app_errors.ErrUpstreamURLChangeConfirmationRequired.Code: - return "group.upstream_url_change_confirmation_required" case app_errors.ErrGroupInUse.Code: return "group.in_use" - case app_errors.ErrUpstreamURLConflict.Code: - return "group.upstream_url_conflict" + case app_errors.ErrChannelTargetConflict.Code: + return "group.channel_target_conflict" case app_errors.ErrModelNameConflict.Code: return "group.model_name_conflict" case app_errors.ErrBadGateway.Code: @@ -722,8 +716,8 @@ func logServiceError(operation string, err error, code string) { if operationErr.groupID != 0 { fields["group_id"] = operationErr.groupID } - if operationErr.keyID != 0 { - fields["key_id"] = operationErr.keyID + if operationErr.credentialID != 0 { + fields["credential_id"] = operationErr.credentialID } } utils.LogPlaneBestEffort( diff --git a/internal/control/server_phase1_test.go b/internal/control/server_phase1_test.go index b92dc8e5..0ad0cf4c 100644 --- a/internal/control/server_phase1_test.go +++ b/internal/control/server_phase1_test.go @@ -13,7 +13,6 @@ import ( "github.com/gin-gonic/gin" - "gpt-load/internal/dialect" "gpt-load/internal/platform/config" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -23,7 +22,7 @@ import ( func TestControlRoutesRequireAuthenticationForGroupCreateAndModelDiscovery(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { t.Fatal("ListModels called without valid management authentication") @@ -82,7 +81,7 @@ func TestControlRoutesRequireAuthenticationForGroupCreateAndModelDiscovery(t *te func TestCreateAndImportEndpointsRequireCanonicalIdempotencyKeyBeforeMutation(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "seed-idempotency-header") + groupID := createGroupWithCredentials(t, fixture, "seed-idempotency-header") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) tests := []struct { @@ -98,12 +97,12 @@ func TestCreateAndImportEndpointsRequireCanonicalIdempotencyKeyBeforeMutation(t { name: "Group create", path: "/api/groups", - body: `{"upstream_url":"https://header.example.com","protocols":["openai-completions"],"keys":"K"}`, + body: `{"channel_id":"openai","params":{},"models":[],"credentials":"K"}`, }, { - name: "Group key import", - path: "/api/groups/" + strconv.FormatUint(uint64(groupID), 10) + "/keys/import", - body: `{"keys":"K"}`, + name: "Group credential import", + path: "/api/groups/" + strconv.FormatUint(uint64(groupID), 10) + "/credentials/import", + body: `{"credentials":"K"}`, }, } for _, test := range tests { @@ -231,10 +230,10 @@ func TestAccessKeyCreateReplayOptionsAndRevealWireContracts(t *testing.T) { } type createImportRowCounts struct { - groups int64 - upstream int64 - access int64 - operations int64 + groups int64 + credentials int64 + access int64 + operations int64 } func countCreateImportRows( @@ -245,7 +244,7 @@ func countCreateImportRows( var result createImportRowCounts for model, target := range map[any]*int64{ &models.Group{}: &result.groups, - &models.UpstreamKey{}: &result.upstream, + &models.Credential{}: &result.credentials, &models.AccessKey{}: &result.access, &models.ControlOperation{}: &result.operations, } { diff --git a/internal/control/server_test.go b/internal/control/server_test.go index 0d67c906..339f459f 100644 --- a/internal/control/server_test.go +++ b/internal/control/server_test.go @@ -20,7 +20,8 @@ import ( "github.com/gin-gonic/gin" "github.com/sirupsen/logrus" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/config" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/protocol" @@ -93,7 +94,7 @@ func TestGroupCollectionHTTPRoutesDeclareStaticOptionsBeforeDynamicDetail(t *tes {name: "control.groups.get", path: "/groups/:group_id"}, {name: "control.groups.settings.get", path: "/groups/:group_id/settings"}, {name: "control.groups.models.get", path: "/groups/:group_id/models"}, - {name: "control.group-keys.list", path: "/groups/:group_id/keys"}, + {name: "control.group-credentials.list", path: "/groups/:group_id/credentials"}, } if !reflect.DeepEqual(got, want) { t.Fatalf("GET group routes = %#v, want %#v", got, want) @@ -326,9 +327,9 @@ func TestControlMutationRejectsDuplicateJSONWithoutSideEffects(t *testing.T) { } type duplicateJSONControlState struct { - rowCounts [5]int64 + rowCounts [4]int64 snapshot *state.ConfigSnapshot - registry []state.KeyRuntimeView + registry []state.CredentialRuntimeView } func captureDuplicateJSONControlState( @@ -336,10 +337,9 @@ func captureDuplicateJSONControlState( fixture serviceFixture, ) duplicateJSONControlState { t.Helper() - var counts [5]int64 + var counts [4]int64 for index, model := range []any{ &models.Group{}, - &models.UpstreamKey{}, &models.AccessKey{}, &models.SystemSetting{}, &models.ModelPrice{}, @@ -400,8 +400,8 @@ type controlJSONBodyLimitState struct { rowCounts [3]int64 snapshot *state.ConfigSnapshot group models.Group - upstreamKey models.UpstreamKey - registryRuntime []state.KeyRuntimeView + credential models.Credential + registryRuntime []state.CredentialRuntimeView accessKey models.AccessKey } @@ -416,16 +416,16 @@ func captureControlJSONBodyLimitState( if err := fixture.db.First(&group, groupID).Error; err != nil { t.Fatalf("load body-limit Group: %v", err) } - var upstreamKey models.UpstreamKey + var credential models.Credential if err := fixture.db.Where("group_id = ?", groupID). - Order("id ASC").Take(&upstreamKey).Error; err != nil { - t.Fatalf("load body-limit UpstreamKey: %v", err) + Order("id ASC").Take(&credential).Error; err != nil { + t.Fatalf("load body-limit Credential: %v", err) } return controlJSONBodyLimitState{ rowCounts: discoveryRowCounts(t, fixture.db), snapshot: fixture.manager.Current(), group: group, - upstreamKey: upstreamKey, + credential: credential, registryRuntime: fixture.registry.Snapshot(), accessKey: loadAccessKeyRow(t, fixture.db, accessKeyID), } @@ -455,8 +455,8 @@ func assertControlJSONBodyLimitStateUnchanged( if !reflect.DeepEqual(got.group, want.group) { t.Errorf("persisted Group changed") } - if !reflect.DeepEqual(got.upstreamKey, want.upstreamKey) { - t.Errorf("persisted UpstreamKey changed") + if !reflect.DeepEqual(got.credential, want.credential) { + t.Errorf("persisted Credential changed") } if !reflect.DeepEqual(got.registryRuntime, want.registryRuntime) { t.Errorf( @@ -486,11 +486,11 @@ func TestControlJSONBodyLimitAppliesToEveryJSONEndpoint(t *testing.T) { `"protocols":["openai-completions"],"models":[],"keys":"sk-body-limit-create"}`, }, { - name: "import group keys", method: http.MethodPost, + name: "import group credentials", method: http.MethodPost, path: func(groupID, _, _ uint) string { - return "/api/groups/" + strconv.FormatUint(uint64(groupID), 10) + "/keys/import" + return "/api/groups/" + strconv.FormatUint(uint64(groupID), 10) + "/credentials/import" }, - jsonPrefix: `{"keys":"sk-body-limit-import"}`, + jsonPrefix: `{"credentials":"sk-body-limit-import"}`, }, { name: "update group settings", method: http.MethodPut, @@ -514,17 +514,17 @@ func TestControlJSONBodyLimitAppliesToEveryJSONEndpoint(t *testing.T) { jsonPrefix: `{"models":[]}`, }, { - name: "update upstream key", method: http.MethodPut, + name: "update credential", method: http.MethodPut, path: func(groupID, upstreamKeyID, _ uint) string { - return fmt.Sprintf("/api/groups/%d/keys/%d", groupID, upstreamKeyID) + return fmt.Sprintf("/api/groups/%d/credentials/%d", groupID, upstreamKeyID) }, jsonPrefix: `{"status":"disabled"}`, }, { name: "discover draft models", method: http.MethodPost, path: func(uint, uint, uint) string { return "/api/models/discover" }, - jsonPrefix: `{"upstream_url":"https://body-limit-discover.example.com/v1",` + - `"protocols":["openai-completions"],"keys":"sk-body-limit-discover"}`, + jsonPrefix: `{"channel_id":"openai_compatible","params":{"base_url":"https://body-limit-discover.example.com/v1"},` + + `"credentials":"sk-body-limit-discover"}`, }, { name: "create access key", method: http.MethodPost, @@ -551,7 +551,7 @@ func TestControlJSONBodyLimitAppliesToEveryJSONEndpoint(t *testing.T) { } { t.Run(endpoint.method+" "+endpoint.name, func(t *testing.T) { fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-body-limit-seed") + groupID := createGroupForCredentialImport(t, fixture, "sk-body-limit-seed") accessKey, err := fixture.service.CreateAccessKey(t.Context(), AccessKeyCreateRequest{ Name: "body-limit-seed-access-key", }) @@ -559,7 +559,7 @@ func TestControlJSONBodyLimitAppliesToEveryJSONEndpoint(t *testing.T) { t.Fatalf("seed CreateAccessKey() error = %v", err) } discoveryCalls := 0 - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { discoveryCalls++ @@ -570,7 +570,7 @@ func TestControlJSONBodyLimitAppliesToEveryJSONEndpoint(t *testing.T) { NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) before := captureControlJSONBodyLimitState(t, fixture, groupID, accessKey.ID) - path := endpoint.path(groupID, before.upstreamKey.ID, accessKey.ID) + path := endpoint.path(groupID, before.credential.ID, accessKey.ID) request := httptest.NewRequest(endpoint.method, path, oversizedControlJSONBody(endpoint.jsonPrefix)) request.ContentLength = -1 request.Header.Set("Authorization", "Bearer test-auth-key") @@ -795,7 +795,7 @@ func TestGroupCreateHTTPReturnsNarrowSuccessAndConflictEnvelopes(t *testing.T) { NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) request := httptest.NewRequest(http.MethodPost, "/api/groups", strings.NewReader( - `{"name":"primary","upstream_url":"https://api.example.com/v1/","protocols":["openai-completions"],"models":[{"id":"gpt-4o","alias":"public-gpt","alias_enabled":true}],"keys":"sk-first"}`, + `{"name":"primary","channel_id":"openai_compatible","params":{"base_url":"https://api.example.com/v1/"},"models":[{"id":"gpt-4o","alias":"public-gpt","alias_enabled":true}],"credentials":"sk-first"}`, )) request.Header.Set("Authorization", "Bearer test-auth-key") request.Header.Set("Content-Type", "application/json") @@ -815,12 +815,12 @@ func TestGroupCreateHTTPReturnsNarrowSuccessAndConflictEnvelopes(t *testing.T) { if success.Code != 0 || len(success.Data) != 4 { t.Fatalf("success response = %#v", success) } - for _, field := range []string{"group_id", "group_name", "keys_added", "keys_duplicated"} { + for _, field := range []string{"group_id", "group_name", "credentials_added", "credentials_duplicated"} { if success.Data[field] == nil { t.Fatalf("success data lacks %q: %#v", field, success.Data) } } - for _, forbidden := range []string{`"models"`, `"config"`, `"keys"`, "sk-first"} { + for _, forbidden := range []string{`"models"`, `"config"`, `"credentials"`, `"keys"`, "sk-first"} { if strings.Contains(recorder.Body.String(), forbidden) { t.Fatalf("success response exposes %q: %s", forbidden, recorder.Body.String()) } @@ -835,7 +835,7 @@ func TestGroupCreateHTTPReturnsNarrowSuccessAndConflictEnvelopes(t *testing.T) { } request = httptest.NewRequest(http.MethodPost, "/api/groups", strings.NewReader( - `{"upstream_url":" HTTPS://API.example.com/v1 ","protocols":["anthropic"],"models":[],"keys":"sk-second"}`, + `{"channel_id":"openai_compatible","params":{"base_url":" HTTPS://API.example.com/v1 "},"models":[],"credentials":"sk-second"}`, )) request.Header.Set("Authorization", "Bearer test-auth-key") request.Header.Set("Content-Type", "application/json") @@ -852,7 +852,7 @@ func TestGroupCreateHTTPReturnsNarrowSuccessAndConflictEnvelopes(t *testing.T) { if err := json.Unmarshal(recorder.Body.Bytes(), &conflict); err != nil { t.Fatalf("decode conflict response: %v", err) } - if conflict.Code != "UPSTREAM_URL_CONFLICT" || len(conflict.Data) != 1 { + if conflict.Code != "CHANNEL_TARGET_CONFLICT" || len(conflict.Data) != 1 { t.Fatalf("conflict response = %#v", conflict) } groups := conflict.Data["groups"] @@ -877,26 +877,26 @@ func TestGroupCreateHTTPRejectsLegacyMissingAndMalformedContractsWithoutMutation }, { name: "omitted models", - body: `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],` + - `"keys":"sk-secret"}`, + body: `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},` + + `"credentials":"sk-secret"}`, code: app_errors.ErrValidation.Code, }, { name: "null models", - body: `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],` + - `"models":null,"keys":"sk-secret"}`, + body: `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},` + + `"models":null,"credentials":"sk-secret"}`, code: app_errors.ErrValidation.Code, }, { name: "unknown model field", - body: `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],` + - `"models":[{"id":"gpt-4o","unknown":true}],"keys":"sk-secret"}`, + body: `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},` + + `"models":[{"id":"gpt-4o","unknown":true}],"credentials":"sk-secret"}`, code: app_errors.ErrInvalidJSON.Code, }, { name: "malformed JSON", - body: `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],` + - `"models":[],"keys":"sk-secret"`, + body: `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},` + + `"models":[],"credentials":"sk-secret"`, code: app_errors.ErrInvalidJSON.Code, }, } @@ -926,17 +926,17 @@ func TestGroupCreateHTTPRejectsLegacyMissingAndMalformedContractsWithoutMutation } } -func TestImportGroupKeysEndpointReturnsSuccessEnvelope(t *testing.T) { +func TestImportGroupCredentialsEndpointReturnsSuccessEnvelope(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-existing") + groupID := createGroupForCredentialImport(t, fixture, "sk-existing") beforeSnapshot := fixture.manager.Current() engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) recorder := httptest.NewRecorder() - request := httptest.NewRequest(http.MethodPost, "/api/groups/"+strconv.FormatUint(uint64(groupID), 10)+"/keys/import", strings.NewReader( - `{"keys":"sk-existing\nsk-new\nsk-new"}`, + request := httptest.NewRequest(http.MethodPost, "/api/groups/"+strconv.FormatUint(uint64(groupID), 10)+"/credentials/import", strings.NewReader( + `{"credentials":"sk-existing\nsk-new\nsk-new"}`, )) request.Header.Set("Authorization", "Bearer test-auth-key") request.Header.Set("Content-Type", "application/json") @@ -956,12 +956,12 @@ func TestImportGroupKeysEndpointReturnsSuccessEnvelope(t *testing.T) { if envelope.Code != 0 || len(envelope.Data) != 3 { t.Fatalf("success envelope = %#v", envelope) } - for _, field := range []string{"group_id", "keys_added", "keys_duplicated"} { + for _, field := range []string{"group_id", "credentials_added", "credentials_duplicated"} { if _, ok := envelope.Data[field]; !ok { t.Fatalf("success data lacks %q: %#v", field, envelope.Data) } } - var result GroupKeyImportResult + var result CredentialImportResult data, err := json.Marshal(envelope.Data) if err != nil { t.Fatal(err) @@ -969,7 +969,7 @@ func TestImportGroupKeysEndpointReturnsSuccessEnvelope(t *testing.T) { if err := json.Unmarshal(data, &result); err != nil { t.Fatal(err) } - if result.GroupID != groupID || result.KeysAdded != 1 || result.KeysDuplicated != 2 { + if result.GroupID != groupID || result.CredentialsAdded != 1 || result.CredentialsDuplicated != 2 { t.Fatalf("result = %#v", result) } if fixture.manager.Current() != beforeSnapshot { @@ -977,10 +977,10 @@ func TestImportGroupKeysEndpointReturnsSuccessEnvelope(t *testing.T) { } } -func TestImportGroupKeysEndpointRejectsUnknownFieldsAndInvalidGroupID(t *testing.T) { +func TestImportGroupCredentialsEndpointRejectsUnknownFieldsAndInvalidGroupID(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-existing") + groupID := createGroupForCredentialImport(t, fixture, "sk-existing") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) @@ -989,17 +989,17 @@ func TestImportGroupKeysEndpointRejectsUnknownFieldsAndInvalidGroupID(t *testing groupID string body string }{ - {name: "unknown field", groupID: strconv.FormatUint(uint64(groupID), 10), body: `{"keys":"sk-new","name":"must-not-change"}`}, - {name: "multiple JSON values", groupID: strconv.FormatUint(uint64(groupID), 10), body: `{"keys":"sk-new"} {"keys":"sk-other"}`}, - {name: "zero group ID", groupID: "0", body: `{"keys":"sk-new"}`}, - {name: "non-numeric group ID", groupID: "not-a-number", body: `{"keys":"sk-new"}`}, - {name: "overflowing group ID", groupID: "18446744073709551616", body: `{"keys":"sk-new"}`}, + {name: "unknown field", groupID: strconv.FormatUint(uint64(groupID), 10), body: `{"credentials":"sk-new","name":"must-not-change"}`}, + {name: "multiple JSON values", groupID: strconv.FormatUint(uint64(groupID), 10), body: `{"credentials":"sk-new"} {"credentials":"sk-other"}`}, + {name: "zero group ID", groupID: "0", body: `{"credentials":"sk-new"}`}, + {name: "non-numeric group ID", groupID: "not-a-number", body: `{"credentials":"sk-new"}`}, + {name: "overflowing group ID", groupID: "18446744073709551616", body: `{"credentials":"sk-new"}`}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { beforeSnapshot := fixture.manager.Current() recorder := httptest.NewRecorder() - request := httptest.NewRequest(http.MethodPost, "/api/groups/"+test.groupID+"/keys/import", strings.NewReader(test.body)) + request := httptest.NewRequest(http.MethodPost, "/api/groups/"+test.groupID+"/credentials/import", strings.NewReader(test.body)) request.Header.Set("Authorization", "Bearer test-auth-key") request.Header.Set("Content-Type", "application/json") setRequiredTestIdempotencyHeader(request) @@ -1011,19 +1011,19 @@ func TestImportGroupKeysEndpointRejectsUnknownFieldsAndInvalidGroupID(t *testing if fixture.manager.Current() != beforeSnapshot { t.Fatal("invalid endpoint request published Snapshot") } - assertImportedKeyState(t, fixture, groupID, 1) + assertImportedCredentialState(t, fixture, groupID, 1) }) } } -func TestImportGroupKeysEndpointReturnsGroupNotFound(t *testing.T) { +func TestImportGroupCredentialsEndpointReturnsGroupNotFound(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) recorder := httptest.NewRecorder() - request := httptest.NewRequest(http.MethodPost, "/api/groups/999/keys/import", strings.NewReader(`{"keys":"sk-new"}`)) + request := httptest.NewRequest(http.MethodPost, "/api/groups/999/credentials/import", strings.NewReader(`{"credentials":"sk-new"}`)) request.Header.Set("Authorization", "Bearer test-auth-key") request.Header.Set("Content-Type", "application/json") request.Header.Set("Accept-Language", "zh-CN") @@ -1045,6 +1045,35 @@ func TestImportGroupKeysEndpointReturnsGroupNotFound(t *testing.T) { } } +func createGroupForCredentialImport(t *testing.T, fixture serviceFixture, credentials string) uint { + t.Helper() + result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("credential-import-" + strconv.FormatUint(testIdempotencySequence.Add(1), 10)), + ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: optionalGroupModels{Set: true}, Credentials: credentials, + }) + if err != nil { + t.Fatalf("CreateGroup() error = %v", err) + } + return result.GroupID +} + +func assertImportedCredentialState(t *testing.T, fixture serviceFixture, groupID uint, want int) { + t.Helper() + var rows []models.Credential + if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&rows).Error; err != nil { + t.Fatal(err) + } + if len(rows) != want { + t.Fatalf("persisted credentials = %d, want %d", len(rows), want) + } + for _, row := range rows { + if value, ok := fixture.registry.EncryptedCredentialData(row.ID); !ok || value != row.Data { + t.Fatalf("Registry credential %d = %q, %t, want %q", row.ID, value, ok, row.Data) + } + } +} + func TestLegacyImportRouteIsNotRegistered(t *testing.T) { fixture := newServiceFixture(t) engine := gin.New() @@ -1151,7 +1180,7 @@ func TestManagementWritesRejectUnknownFieldsAndMultipleJSONValues(t *testing.T) func TestUpdateGroupSettingsEndpointRejectsStrictInvalidBodies(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-update-http") + groupID := createGroupForCredentialImport(t, fixture, "sk-update-http") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) @@ -1161,12 +1190,12 @@ func TestUpdateGroupSettingsEndpointRejectsStrictInvalidBodies(t *testing.T) { code string }{ {name: "empty object", body: `{}`, code: app_errors.ErrBadRequest.Code}, - {name: "confirmation only", body: `{"confirm_upstream_change":true}`, code: app_errors.ErrBadRequest.Code}, + {name: "retired confirmation", body: `{"confirm_upstream_change":true}`, code: app_errors.ErrInvalidJSON.Code}, {name: "unknown field", body: `{"name":"changed","unknown":true}`, code: app_errors.ErrInvalidJSON.Code}, {name: "multiple JSON values", body: `{"name":"changed"} {"enabled":false}`, code: app_errors.ErrInvalidJSON.Code}, {name: "null name", body: `{"name":null}`, code: app_errors.ErrValidation.Code}, {name: "negative weight", body: `{"weight_manual":-1}`, code: app_errors.ErrValidation.Code}, - {name: "empty protocols", body: `{"protocols":[]}`, code: app_errors.ErrValidation.Code}, + {name: "retired protocols", body: `{"protocols":[]}`, code: app_errors.ErrInvalidJSON.Code}, {name: "invalid overrides", body: `{"overrides":{"first_byte_timeout":-1}}`, code: app_errors.ErrValidation.Code}, } { t.Run(test.name, func(t *testing.T) { @@ -1200,7 +1229,7 @@ func TestUpdateGroupSettingsEndpointRejectsStrictInvalidBodies(t *testing.T) { func TestUpdateGroupSettingsEndpointRejectsTopLevelNullWithoutMutation(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-update-null") + groupID := createGroupForCredentialImport(t, fixture, "sk-update-null") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) @@ -1208,12 +1237,12 @@ func TestUpdateGroupSettingsEndpointRejectsTopLevelNullWithoutMutation(t *testin if err := fixture.db.First(&beforeGroup, groupID).Error; err != nil { t.Fatal(err) } - var beforeKeys []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&beforeKeys).Error; err != nil { + var beforeCredentials []models.Credential + if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&beforeCredentials).Error; err != nil { t.Fatal(err) } beforeRevision := fixture.manager.Current().Revision - beforeRegistry := fixture.registry.CollectCandidates([]uint{groupID}, nil, time.Time{}) + beforeRegistry := fixture.registry.CollectCredentialCandidates([]uint{groupID}, nil, time.Time{}) recorder := httptest.NewRecorder() request := httptest.NewRequest( @@ -1233,17 +1262,17 @@ func TestUpdateGroupSettingsEndpointRejectsTopLevelNullWithoutMutation(t *testin if !reflect.DeepEqual(afterGroup, beforeGroup) { t.Fatalf("persisted group changed: got=%#v want=%#v", afterGroup, beforeGroup) } - var afterKeys []models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&afterKeys).Error; err != nil { + var afterCredentials []models.Credential + if err := fixture.db.Where("group_id = ?", groupID).Order("id ASC").Find(&afterCredentials).Error; err != nil { t.Fatal(err) } - if !reflect.DeepEqual(afterKeys, beforeKeys) { - t.Fatalf("persisted keys changed: got=%#v want=%#v", afterKeys, beforeKeys) + if !reflect.DeepEqual(afterCredentials, beforeCredentials) { + t.Fatalf("persisted credentials changed: got=%#v want=%#v", afterCredentials, beforeCredentials) } if got := fixture.manager.Current().Revision; got != beforeRevision { t.Fatalf("snapshot revision = %d, want %d", got, beforeRevision) } - if got := fixture.registry.CollectCandidates([]uint{groupID}, nil, time.Time{}); !reflect.DeepEqual(got, beforeRegistry) { + if got := fixture.registry.CollectCredentialCandidates([]uint{groupID}, nil, time.Time{}); !reflect.DeepEqual(got, beforeRegistry) { t.Fatalf("Registry candidates changed: got=%#v want=%#v", got, beforeRegistry) } } @@ -1251,13 +1280,19 @@ func TestUpdateGroupSettingsEndpointRejectsTopLevelNullWithoutMutation(t *testin func TestUpdateGroupSettingsEndpointAllowsURLReuseAndPreservesI18nAndAuth(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - firstID := createGroupForKeyImport(t, fixture, "sk-update-first") + first, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: stringPointer("first-group"), ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://first.example.com/v1"}`), + Models: optionalGroupModels{Set: true}, Credentials: "sk-update-first", + }) + if err != nil { + t.Fatal(err) + } + firstID := first.GroupID second, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - Name: stringPointer("other-group"), - UpstreamURL: "https://conflict.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-update-second", + Name: stringPointer("other-group"), ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://conflict.example.com/v1"}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Credentials: "sk-update-second", }) if err != nil { t.Fatal(err) @@ -1267,7 +1302,7 @@ func TestUpdateGroupSettingsEndpointAllowsURLReuseAndPreservesI18nAndAuth(t *tes path := "/api/groups/" + strconv.FormatUint(uint64(firstID), 10) + "/settings" request := httptest.NewRequest(http.MethodPut, path, strings.NewReader( - `{"upstream_url":"https://unique.example.com/v1"}`, + `{"params":{"base_url":"https://unique.example.com/v1"}}`, )) request.Header.Set("Authorization", "Bearer test-auth-key") request.Header.Set("Content-Type", "application/json") @@ -1279,7 +1314,7 @@ func TestUpdateGroupSettingsEndpointAllowsURLReuseAndPreservesI18nAndAuth(t *tes } request = httptest.NewRequest(http.MethodPut, path, strings.NewReader( - `{"upstream_url":"https://conflict.example.com/v1/"}`, + `{"params":{"base_url":"https://conflict.example.com/v1/"}}`, )) request.Header.Set("Authorization", "Bearer test-auth-key") request.Header.Set("Content-Type", "application/json") @@ -1295,7 +1330,8 @@ func TestUpdateGroupSettingsEndpointAllowsURLReuseAndPreservesI18nAndAuth(t *tes if err := json.Unmarshal(recorder.Body.Bytes(), &duplicateSuccess); err != nil { t.Fatal(err) } - if duplicateSuccess.Code != 0 || duplicateSuccess.Data.UpstreamURL != "https://conflict.example.com/v1" { + if duplicateSuccess.Code != 0 || duplicateSuccess.Data.ChannelID != channel.OpenAICompatible || + string(duplicateSuccess.Data.Params) != `{"base_url":"https://conflict.example.com/v1"}` { t.Fatalf("duplicate URL success envelope = %#v", duplicateSuccess) } @@ -1309,7 +1345,7 @@ func TestUpdateGroupSettingsEndpointAllowsURLReuseAndPreservesI18nAndAuth(t *tes assertUpdateGroupErrorResponse(t, recorder, http.StatusConflict, app_errors.ErrDuplicateResource.Code, "分组名称已存在") request = httptest.NewRequest(http.MethodPut, path, strings.NewReader( - `{"upstream_url":" HTTPS://UNIQUE.example.com/v1/ ","enabled":false}`, + `{"params":{"base_url":" HTTPS://UNIQUE.example.com/v1/ "},"enabled":false}`, )) request.Header.Set("Authorization", "Bearer test-auth-key") request.Header.Set("Content-Type", "application/json") @@ -1328,7 +1364,7 @@ func TestUpdateGroupSettingsEndpointAllowsURLReuseAndPreservesI18nAndAuth(t *tes t.Fatal(err) } if success.Code != 0 || success.Message != "成功" || success.Data.Enabled || - success.Data.UpstreamURL != "https://unique.example.com/v1" { + string(success.Data.Params) != `{"base_url":"https://unique.example.com/v1"}` { t.Fatalf("success envelope = %#v", success) } @@ -1345,7 +1381,7 @@ func TestUpdateGroupSettingsEndpointAllowsURLReuseAndPreservesI18nAndAuth(t *tes func TestUpdateGroupSettingsEndpointRejectsOversizedJSON(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-update-limit") + groupID := createGroupForCredentialImport(t, fixture, "sk-update-limit") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) before := fixture.manager.Current().Revision @@ -1370,13 +1406,13 @@ func TestUpdateGroupModelsEndpointRejectsStrictInvalidBodiesWithoutMutation(t *t initControlI18n(t) fixture := newServiceFixture(t) created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://model-save-http-invalid.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://model-save-http-invalid.example.com/v1"}`), Models: optionalGroupModels{ Set: true, Values: []GroupModel{{ID: "provider-old", Alias: "old-public", AliasEnabled: true}}, }, - Keys: "sk-model-save-http-invalid", + Credentials: "sk-model-save-http-invalid", }) if err != nil { t.Fatal(err) @@ -1431,7 +1467,7 @@ func TestUpdateGroupModelsEndpointRejectsStrictInvalidBodiesWithoutMutation(t *t func TestGroupModelsHTTPReturnsStructuredConflictWithoutMutation(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-model-conflict-http") + groupID := createGroupForCredentialImport(t, fixture, "sk-model-conflict-http") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) beforeRevision := fixture.manager.Current().Revision @@ -1469,7 +1505,7 @@ func TestGroupModelsHTTPRejectsMissingAliasEnabledWithoutMutation(t *testing.T) initControlI18n(t) fixture := newServiceFixture(t) mustEnsureInitialPrices(t, fixture) - groupID := createGroupForKeyImport(t, fixture, "sk-model-alias-default") + groupID := createGroupForCredentialImport(t, fixture, "sk-model-alias-default") engine := gin.New() NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) beforeRevision := fixture.manager.Current().Revision @@ -1500,13 +1536,13 @@ func TestUpdateGroupModelsEndpointIDsAuthNotFoundAndSuccessDTO(t *testing.T) { fixture := newServiceFixture(t) mustEnsureInitialPrices(t, fixture) created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://model-save-http.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://model-save-http.example.com/v1"}`), Models: optionalGroupModels{ Set: true, Values: []GroupModel{{ID: "provider-old", Alias: "old-public", AliasEnabled: true}}, }, - Keys: "sk-model-save-http", + Credentials: "sk-model-save-http", }) if err != nil { t.Fatal(err) @@ -1712,13 +1748,13 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { initControlI18n(t) const authKey = "test-auth-key" - newServer := func(value *recordingDiscoveryDialect) (*Service, *gin.Engine) { + newServer := func(value *recordingDiscoveryExecutorTarget) (*Service, *gin.Engine) { t.Helper() fixture := newServiceFixture(t) if value == nil { - fixture.service.dialects = dialect.NewSet() + fixture.service.executor = nil } else { - fixture.service.dialects = dialect.NewSet(value) + fixture.service.executor = newRecordingDiscoveryExecutor(value) } engine := gin.New() NewServer(&config.Config{AuthKey: authKey}, fixture.service).RegisterRoutes(engine) @@ -1726,7 +1762,7 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { } t.Run("success preserves order", func(t *testing.T) { - _, engine := newServer(&recordingDiscoveryDialect{ + _, engine := newServer(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func( context.Context, @@ -1738,8 +1774,8 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }, }) recorder := serveDiscoveryRequest(t, engine, authKey, - `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],`+ - `"keys":"sk-upstream"}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},`+ + `"credentials":"sk-upstream"}`, ) if recorder.Code != http.StatusOK { t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) @@ -1765,7 +1801,7 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }) t.Run("empty list remains an array", func(t *testing.T) { - _, engine := newServer(&recordingDiscoveryDialect{ + _, engine := newServer(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func( context.Context, @@ -1777,8 +1813,8 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }, }) recorder := serveDiscoveryRequest(t, engine, authKey, - `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],`+ - `"keys":"sk-upstream"}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},`+ + `"credentials":"sk-upstream"}`, ) if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), `"models":[]`) { t.Fatalf("response = %d %s", recorder.Code, recorder.Body.String()) @@ -1786,7 +1822,7 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }) t.Run("authentication is inherited", func(t *testing.T) { - _, engine := newServer(&recordingDiscoveryDialect{ + _, engine := newServer(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { t.Fatal("ListModels called without valid management authentication") @@ -1795,8 +1831,8 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }) for _, token := range []string{"", "wrong-key"} { recorder := serveDiscoveryRequest(t, engine, token, - `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],`+ - `"keys":"sk-upstream"}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},`+ + `"credentials":"sk-upstream"}`, ) if recorder.Code != http.StatusUnauthorized || !strings.Contains(recorder.Body.String(), `"code":"UNAUTHORIZED"`) { t.Fatalf("auth %q response = %d %s", token, recorder.Code, recorder.Body.String()) @@ -1805,7 +1841,7 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }) t.Run("strict JSON", func(t *testing.T) { - _, engine := newServer(&recordingDiscoveryDialect{ + _, engine := newServer(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { t.Fatal("ListModels called for invalid JSON") @@ -1813,11 +1849,11 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }, }) for _, payload := range []string{ - `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],"keys":"sk-upstream","config":{}}`, - `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],"keys":"sk-upstream","unknown":true}`, - `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],"keys":"sk-upstream"}{}`, - `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],"keys":"sk-upstream"`, - `{"upstream_url":"https://api.example.com","protocol":"openai-completions","key":"sk-upstream"}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},"credentials":"sk-upstream","config":{}}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},"credentials":"sk-upstream","unknown":true}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},"credentials":"sk-upstream"}{}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},"credentials":"sk-upstream"`, + `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],"keys":"sk-upstream"}`, } { recorder := serveDiscoveryRequest(t, engine, authKey, payload) if recorder.Code != http.StatusBadRequest || !strings.Contains(recorder.Body.String(), `"code":"INVALID_JSON"`) { @@ -1833,8 +1869,8 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { status int code string }{ - {payload: `{"upstream_url":"/relative","protocols":["openai-completions"],"keys":"secret-key"}`, status: http.StatusBadRequest, code: "VALIDATION_FAILED"}, - {payload: `{"upstream_url":"https://api.example.com?token=query-secret","protocols":["openai-completions"],"keys":"secret-key"}`, status: http.StatusInternalServerError, code: "INTERNAL_SERVER_ERROR"}, + {payload: `{"channel_id":"openai_compatible","params":{"base_url":"/relative"},"credentials":"secret-key"}`, status: http.StatusBadRequest, code: "VALIDATION_FAILED"}, + {payload: `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com?token=query-secret"},"credentials":"secret-key"}`, status: http.StatusInternalServerError, code: "INTERNAL_SERVER_ERROR"}, } { recorder := serveDiscoveryRequest(t, engine, authKey, test.payload) if recorder.Code != test.status || !strings.Contains(recorder.Body.String(), `"code":"`+test.code+`"`) { @@ -1849,7 +1885,7 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }) t.Run("upstream failures map to localized bad gateway", func(t *testing.T) { - service, engine := newServer(&recordingDiscoveryDialect{ + service, engine := newServer(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { return nil, fmt.Errorf("raw upstream failure with secret-body") @@ -1864,8 +1900,8 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { {language: "zh-CN", message: "上游服务错误"}, } { recorder := serveDiscoveryRequestWithLanguage(t, engine, authKey, - `{"upstream_url":"https://api.example.com?token=query-secret",`+ - `"protocols":["openai-completions"],"keys":"secret-key"}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com?token=query-secret"},`+ + `"credentials":"secret-key"}`, test.language, ) if recorder.Code != http.StatusBadGateway || @@ -1882,7 +1918,7 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }) t.Run("timeout maps to bad gateway", func(t *testing.T) { - service, engine := newServer(&recordingDiscoveryDialect{ + service, engine := newServer(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(ctx context.Context, _ string, _ string, _ state.HeaderRules) ([]string, error) { <-ctx.Done() @@ -1891,8 +1927,8 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }) service.modelDiscoveryTimeout = 20 * time.Millisecond recorder := serveDiscoveryRequest(t, engine, authKey, - `{"upstream_url":"https://api.example.com","protocols":["openai-completions"],`+ - `"keys":"sk-upstream"}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},`+ + `"credentials":"sk-upstream"}`, ) if recorder.Code != http.StatusBadGateway || !strings.Contains(recorder.Body.String(), `"code":"BAD_GATEWAY"`) { t.Fatalf("response = %d %s", recorder.Code, recorder.Body.String()) @@ -1900,19 +1936,21 @@ func TestModelDiscoveryHTTPContract(t *testing.T) { }) t.Run("broken upstream JSON maps to bad gateway", func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - _, _ = writer.Write([]byte(`{"data":[`)) - })) - defer upstream.Close() - fixture := newServiceFixture(t) - fixture.service.dialects = dialect.NewSet(dialect.NewOpenAI(upstream.Client())) + fixture.service.executor = scriptedDiscoveryExecutor{execute: func( + context.Context, + execution.AttemptSpec, + ) execution.AttemptResult { + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusOK, Header: http.Header{}, Body: []byte(`{"data":[`), + } + }} engine := gin.New() NewServer(&config.Config{AuthKey: authKey}, fixture.service).RegisterRoutes(engine) recorder := serveDiscoveryRequest(t, engine, authKey, - `{"upstream_url":"`+upstream.URL+`","protocols":["openai-completions"],`+ - `"keys":"sk-upstream"}`, + `{"channel_id":"openai_compatible","params":{"base_url":"https://api.example.com"},`+ + `"credentials":"sk-upstream"}`, ) if recorder.Code != http.StatusBadGateway || !strings.Contains(recorder.Body.String(), `"code":"BAD_GATEWAY"`) { t.Fatalf("response = %d %s", recorder.Code, recorder.Body.String()) @@ -1932,7 +1970,7 @@ func TestServerDraftModelDiscoveryLogsOnlyMetadata(t *testing.T) { bodySecret = "distinctive-upstream-body" ) fixture := newServiceFixture(t) - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.Anthropic, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { return nil, fmt.Errorf("provider error: %s", bodySecret) @@ -1947,8 +1985,8 @@ func TestServerDraftModelDiscoveryLogsOnlyMetadata(t *testing.T) { t.Cleanup(func() { logrus.SetOutput(previousOutput) }) recorder := serveDiscoveryRequest(t, engine, authSecret, - `{"upstream_url":"https://api.example.com?token=`+querySecret+`",`+ - `"protocols":["anthropic"],"keys":"`+keySecret+`"}`, + `{"channel_id":"anthropic_compatible","params":{"base_url":"https://api.example.com?token=`+querySecret+`"},`+ + `"credentials":"`+keySecret+`"}`, ) if recorder.Code != http.StatusBadGateway { t.Fatalf("response = %d %s", recorder.Code, recorder.Body.String()) @@ -1977,22 +2015,22 @@ func TestServerGroupModelDiscoveryBodyContract(t *testing.T) { t.Helper() fixture := newServiceFixture(t) created, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - UpstreamURL: "https://persisted-server.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://persisted-server.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "persisted-server-key", + Credentials: "persisted-server-key", }) if err != nil { t.Fatalf("seed CreateGroup() error = %v", err) } if !activeKey { - if err := fixture.db.Model(&models.UpstreamKey{}). + if err := fixture.db.Model(&models.Credential{}). Where("group_id = ?", created.GroupID). - Update("status", models.UpstreamKeyStatusDisabled).Error; err != nil { + Update("status", models.CredentialStatusDisabled).Error; err != nil { t.Fatalf("disable persisted discovery key: %v", err) } } - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ + fixture.service.executor = newRecordingDiscoveryExecutor(&recordingDiscoveryExecutorTarget{ value: protocol.OpenAICompletions, listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { return []string{"z-model", "a-model"}, nil @@ -2079,15 +2117,15 @@ func TestServerGroupModelDiscoveryBodyContract(t *testing.T) { language string message string }{ - {language: "en-US", message: "No active upstream key is available for this group"}, - {language: "zh-CN", message: "该分组没有可用的活跃上游密钥"}, - {language: "ja-JP", message: "このグループには利用可能な有効なアップストリームキーがありません"}, + {language: "en-US", message: "No active credential is available for this group"}, + {language: "zh-CN", message: "该分组没有可用凭据"}, + {language: "ja-JP", message: "このグループには利用可能な認証情報がありません"}, } { recorder := serveGroupDiscoveryRequest( t, engine, authKey, test.language, groupID, nil, ) if recorder.Code != http.StatusConflict || - !strings.Contains(recorder.Body.String(), `"code":"NO_ACTIVE_UPSTREAM_KEY"`) || + !strings.Contains(recorder.Body.String(), `"code":"NO_ACTIVE_CREDENTIAL"`) || !strings.Contains(recorder.Body.String(), test.message) { t.Fatalf("%s response = %d %s", test.language, recorder.Code, recorder.Body.String()) } @@ -2362,7 +2400,7 @@ func TestSettingsHTTPFiltersPrivateRowsAndDoesNotLogValues(t *testing.T) { fixture.catalogRuntime, nil, fixture.encryption, - fixture.service.dialects, + fixture.service.executor, fixture.service.requestLogs, fixture.service.usageStats, fixture.service.homeStatistics, diff --git a/internal/control/service.go b/internal/control/service.go index 0104d0e0..0f8800b1 100644 --- a/internal/control/service.go +++ b/internal/control/service.go @@ -11,7 +11,8 @@ import ( "gorm.io/gorm" "gpt-load/internal/catalog" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/health" "gpt-load/internal/platform/config" "gpt-load/internal/platform/dbtx" @@ -31,28 +32,29 @@ const ( type Service struct { db *gorm.DB manager *state.Manager - registry *state.KeyRegistry - registrySnapshot func() []state.KeyRuntimeView + registry *state.CredentialRegistry + channelRegistry *channel.Registry + registrySnapshot func() []state.CredentialRuntimeView priceRuntime *PriceRuntime catalogRuntime *catalog.Runtime catalogSync *CatalogSyncCoordinator modelsDevAutoSyncOverride *bool encryption encryption.Service - dialects dialect.Set + executor execution.Executor requestLogs RequestLogReader usageStats UsageStatReader homeStatistics HomeStatisticsReader stats *health.StatsStore - mutations keyMutationCoordinator + mutations credentialMutationCoordinator requestLogStats RequestLogStatsReader modelDiscoveryTimeout time.Duration random io.Reader operationRandom io.Reader now func() time.Time publishSnapshot func(state.CompileInput) (*state.ConfigSnapshot, error) - reconcileRegistryGroup func(uint, []state.KeyEntry) (bool, error) - applyBatchRegistryMutation func(uint, []uint, GroupKeyBatchAction) error - restoreBatchRegistryEntries func(uint, []state.KeyEntry) error + reconcileRegistryGroup func(uint, []state.CredentialEntry) (bool, error) + applyBatchRegistryMutation func(uint, []uint, CredentialBatchAction) error + restoreBatchRegistryEntries func(uint, []state.CredentialEntry) error beforeAdvanceOperationStage func( context.Context, *models.ControlOperation, @@ -65,24 +67,33 @@ type Service struct { func NewService( db *gorm.DB, manager *state.Manager, - registry *state.KeyRegistry, + registry *state.CredentialRegistry, priceRuntime *PriceRuntime, catalogRuntime *catalog.Runtime, cfg *config.Config, encryptionService encryption.Service, - dialects dialect.Set, + executor execution.Executor, requestLogs RequestLogReader, usageStats UsageStatReader, homeStatistics HomeStatisticsReader, stats *health.StatsStore, mutations *health.MutationCoordinator, requestLogStats RequestLogStatsReader, + channelRegistries ...*channel.Registry, ) *Service { + channelRegistry := channel.NewRegistry() + for _, candidate := range channelRegistries { + if candidate != nil { + channelRegistry = candidate + break + } + } service := &Service{ db: db, manager: manager, registry: registry, - priceRuntime: priceRuntime, - catalogRuntime: catalogRuntime, - encryption: encryptionService, dialects: dialects, requestLogs: requestLogs, + channelRegistry: channelRegistry, + priceRuntime: priceRuntime, + catalogRuntime: catalogRuntime, + encryption: encryptionService, executor: executor, requestLogs: requestLogs, usageStats: usageStats, homeStatistics: homeStatistics, stats: stats, mutations: mutations, requestLogStats: requestLogStats, modelDiscoveryTimeout: defaultModelDiscoveryTimeout, @@ -97,8 +108,8 @@ func NewService( } service.publishSnapshot = manager.Publish service.reconcileRegistryGroup = registry.ReconcileGroup - service.applyBatchRegistryMutation = service.applyGroupKeyBatchRegistryMutation - service.restoreBatchRegistryEntries = registry.RestoreGroupKeyEntriesExact + service.applyBatchRegistryMutation = service.applyCredentialBatchRegistryMutation + service.restoreBatchRegistryEntries = registry.RestoreGroupCredentialEntriesExact service.registrySnapshot = registry.Snapshot return service } @@ -134,7 +145,7 @@ func (s *Service) writeGroupConfig( if err := cleanupUnreferencedAutomaticPrices(tx); err != nil { return err } - input, err := stateloader.BuildCompileInput(ctx, tx) + input, err := stateloader.BuildCompileInput(ctx, tx, s.channelRegistry) if err != nil { return err } @@ -182,7 +193,7 @@ func (s *Service) writeConfig( return err } var err error - input, err = stateloader.BuildCompileInput(ctx, tx) + input, err = stateloader.BuildCompileInput(ctx, tx, s.channelRegistry) if err != nil { return err } @@ -207,10 +218,10 @@ func (s *Service) writeConfig( return snapshot, nil } -func (s *Service) writeKeyConfig( +func (s *Service) writeCredentialConfig( ctx context.Context, groupID uint, - keyID uint, + credentialID uint, mutate func(*gorm.DB) error, afterCommit func() error, ) error { @@ -228,7 +239,7 @@ func (s *Service) writeKeyConfig( return withControlOperationContext( newControlOperationError(stageApplyCommittedRegistryMutation), groupID, - keyID, + credentialID, ) } } diff --git a/internal/control/service_integration_test.go b/internal/control/service_integration_test.go index 1792d46e..36aa3341 100644 --- a/internal/control/service_integration_test.go +++ b/internal/control/service_integration_test.go @@ -2,6 +2,7 @@ package control import ( "bytes" + "encoding/json" "io" "net/http" "net/http/httptest" @@ -11,12 +12,10 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/dialect" "gpt-load/internal/gateway" "gpt-load/internal/health" - platformhttp "gpt-load/internal/platform/httpclient" - "gpt-load/internal/platform/redact" - "gpt-load/internal/protocol" "gpt-load/internal/state" ) @@ -45,10 +44,11 @@ func TestControlWriteLockDoesNotBlockDataPlane(t *testing.T) { upstreamKey = "sk-data-plane-upstream" ) if _, err := fixture.manager.Publish(state.CompileInput{ + ChannelRegistry: fixture.channelRegistry, Groups: []state.GroupConfig{{ - ID: groupID, Name: "data-plane", UpstreamURL: upstream.URL, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "gpt-4o"}}, Enabled: true, + ID: groupID, Name: "data-plane", ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"` + upstream.URL + `/v1"}`), + Models: []state.ModelConfig{{ID: "gpt-4o"}}, Enabled: true, }}, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: fixture.encryption.Hash(accessKey), @@ -57,22 +57,25 @@ func TestControlWriteLockDoesNotBlockDataPlane(t *testing.T) { }); err != nil { t.Fatalf("Publish(data plane) error = %v", err) } - ciphertext, err := fixture.encryption.Encrypt(upstreamKey) + credentialData := `{"api_key":"` + upstreamKey + `"}` + ciphertext, err := fixture.encryption.Encrypt(credentialData) if err != nil { t.Fatalf("Encrypt(upstream key) error = %v", err) } - if err := fixture.registry.ApplyImport(groupID, []state.KeyEntry{{ - ID: upstreamKeyID, GroupID: groupID, Status: state.KeyStatusActive, EncryptedValue: ciphertext, + if err := fixture.registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: upstreamKeyID, GroupID: groupID, Version: 1, IdentityGeneration: 1, + Fingerprint: fixture.encryption.Hash(credentialData), Status: state.CredentialStatusActive, + EncryptedValue: ciphertext, }}); err != nil { - t.Fatalf("ApplyImport() error = %v", err) + t.Fatalf("ReplaceCredentials() error = %v", err) } - openAI := dialect.NewOpenAI(http.DefaultClient) + openAI := dialect.NewOpenAI() handler := gateway.NewHandler( fixture.manager, fixture.registry, fixture.encryption, - gateway.NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()), + gateway.NewExecutionForwarder(controlHTTPExecutor{}), dialect.NewSet(openAI), health.NewStatsStore(), health.NewMutationCoordinator(), @@ -132,7 +135,10 @@ func TestControlWriteLockDoesNotBlockDataPlane(t *testing.T) { t.Fatal("data-plane request did not complete while control writeMu was held") } } - if got, ok := fixture.registry.ActiveEncryptedValue(upstreamKeyID, groupID); !ok || got != ciphertext { + if got, ok := fixture.registry.ActiveEncryptedCredentialDataIfMatch(state.CredentialRef{ + ID: upstreamKeyID, GroupID: groupID, Version: 1, IdentityGeneration: 1, + Fingerprint: fixture.encryption.Hash(credentialData), EncryptedValue: ciphertext, + }); !ok || got != ciphertext { t.Fatalf("Registry value = %q, %t, want active ciphertext", got, ok) } diff --git a/internal/control/service_test.go b/internal/control/service_test.go index e2475887..7777c640 100644 --- a/internal/control/service_test.go +++ b/internal/control/service_test.go @@ -3,6 +3,7 @@ package control import ( "bytes" "context" + "encoding/json" "errors" "fmt" "io" @@ -18,13 +19,11 @@ import ( "github.com/sirupsen/logrus" "gorm.io/gorm" + "gpt-load/internal/channel" "gpt-load/internal/dialect" "gpt-load/internal/gateway" "gpt-load/internal/health" app_errors "gpt-load/internal/platform/errors" - platformhttp "gpt-load/internal/platform/httpclient" - "gpt-load/internal/platform/redact" - "gpt-load/internal/protocol" "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" "gpt-load/internal/storage" @@ -132,237 +131,6 @@ func TestReadSnapshotCancellationTakesPrecedenceAndReleasesConnection(t *testing } } -func TestWriteKeyConfigCommitsDatabaseThenAppliesRegistryWithoutPublishing(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-write-key-config") - var row models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Take(&row).Error; err != nil { - t.Fatal(err) - } - beforeSnapshot := fixture.manager.Current() - weight := 75 - callbackRan := false - - err := fixture.service.writeKeyConfig( - t.Context(), - groupID, - row.ID, - func(tx *gorm.DB) error { - return tx.Model(&models.UpstreamKey{}). - Where("id = ? AND group_id = ?", row.ID, groupID). - Updates(map[string]any{ - "status": models.UpstreamKeyStatusDisabled, - "weight_manual": weight, - }).Error - }, - func() error { - callbackRan = true - if fixture.service.writeMu.TryLock() { - fixture.service.writeMu.Unlock() - return errors.New("writeMu was not held through Registry convergence") - } - var committed models.UpstreamKey - if err := fixture.db.First(&committed, row.ID).Error; err != nil { - return err - } - if committed.Status != models.UpstreamKeyStatusDisabled || - committed.WeightManual == nil || *committed.WeightManual != weight { - return fmt.Errorf("callback observed uncommitted row: %#v", committed) - } - return fixture.registry.UpdateKeyConfig( - row.ID, - state.KeyStatusDisabled, - &weight, - ) - }, - ) - if err != nil { - t.Fatalf("writeKeyConfig() error = %v", err) - } - if !callbackRan { - t.Fatal("afterCommit callback did not run") - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("writeKeyConfig published Snapshot") - } - view := fixture.registry.Snapshot()[0] - if view.Status != state.KeyStatusDisabled || - view.WeightManual == nil || *view.WeightManual != weight { - t.Fatalf("Registry view = %#v", view) - } -} - -func TestWriteKeyConfigCommitFailureDoesNotApplyRegistry(t *testing.T) { - fixture, dsn := newFileServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-key-commit-failure") - var row models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Take(&row).Error; err != nil { - t.Fatal(err) - } - beforeSnapshot := fixture.manager.Current() - beforeRegistry := fixture.registry.Snapshot() - release := holdRollbackJournalReadLock(t, fixture.db, dsn) - - callbackRan := false - err := fixture.service.writeKeyConfig( - t.Context(), - groupID, - row.ID, - func(tx *gorm.DB) error { - return tx.Model(&models.UpstreamKey{}). - Where("id = ?", row.ID). - Update("status", models.UpstreamKeyStatusDisabled).Error - }, - func() error { - callbackRan = true - return nil - }, - ) - if err == nil { - t.Fatal("writeKeyConfig() error = nil") - } - if callbackRan { - t.Fatal("afterCommit callback ran after failed COMMIT") - } - if fixture.manager.Current() != beforeSnapshot || - !reflect.DeepEqual(fixture.registry.Snapshot(), beforeRegistry) { - t.Fatal("failed COMMIT changed runtime") - } - release() -} - -func TestWriteKeyConfigAfterCommitFailureIsTypedAndKeepsSnapshot(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-key-post-commit") - var row models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Take(&row).Error; err != nil { - t.Fatal(err) - } - beforeSnapshot := fixture.manager.Current() - const secretCause = "known-secret-must-not-be-logged" - err := fixture.service.writeKeyConfig( - t.Context(), - groupID, - row.ID, - func(tx *gorm.DB) error { - return tx.Model(&models.UpstreamKey{}). - Where("id = ?", row.ID). - Update("status", models.UpstreamKeyStatusDisabled).Error - }, - func() error { return errors.New(secretCause) }, - ) - if !errors.Is(err, app_errors.ErrInternalServer) { - t.Fatalf("error = %v", err) - } - var operationErr *controlOperationError - if !errors.As(err, &operationErr) || - operationErr.stage != stageApplyCommittedRegistryMutation || - operationErr.groupID != groupID || operationErr.keyID != row.ID || - operationErr.mismatchKind != "" { - t.Fatalf("operation error = %#v", operationErr) - } - if strings.Contains(err.Error(), secretCause) { - t.Fatalf("operation error leaked callback cause: %v", err) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("post-commit failure published Snapshot") - } - var committed models.UpstreamKey - if err := fixture.db.First(&committed, row.ID).Error; err != nil { - t.Fatal(err) - } - if committed.Status != models.UpstreamKeyStatusDisabled { - t.Fatalf("committed status = %q, want disabled", committed.Status) - } -} - -func TestWriteKeyConfigPostCommitCancellationDoesNotSkipRegistryCallback(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-key-cancel-after-commit") - var row models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Take(&row).Error; err != nil { - t.Fatal(err) - } - beforeSnapshot := fixture.manager.Current() - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - callbackRan := false - - err := fixture.service.writeKeyConfig( - ctx, - groupID, - row.ID, - func(tx *gorm.DB) error { - return tx.Model(&models.UpstreamKey{}). - Where("id = ?", row.ID). - Update("status", models.UpstreamKeyStatusDisabled).Error - }, - func() error { - callbackRan = true - cancel() - if !errors.Is(ctx.Err(), context.Canceled) { - return fmt.Errorf("context error = %v, want canceled", ctx.Err()) - } - return fixture.registry.UpdateKeyConfig(row.ID, state.KeyStatusDisabled, nil) - }, - ) - if err != nil { - t.Fatalf("writeKeyConfig() error = %v", err) - } - if !callbackRan { - t.Fatal("afterCommit callback did not run") - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("post-commit cancellation published Snapshot") - } - view := fixture.registry.Snapshot()[0] - if view.Status != state.KeyStatusDisabled { - t.Fatalf("Registry status = %q, want disabled", view.Status) - } - var committed models.UpstreamKey - if err := fixture.db.First(&committed, row.ID).Error; err != nil { - t.Fatal(err) - } - if committed.Status != models.UpstreamKeyStatusDisabled { - t.Fatalf("committed status = %q, want disabled", committed.Status) - } -} - -func TestWriteKeyConfigDoesNotCompileOrPublishUnrelatedConfiguration(t *testing.T) { - fixture := newServiceFixture(t) - groupID := createGroupForKeyImport(t, fixture, "sk-key-no-compile") - var row models.UpstreamKey - if err := fixture.db.Where("group_id = ?", groupID).Take(&row).Error; err != nil { - t.Fatal(err) - } - if err := fixture.db.Model(&models.Group{}). - Where("id = ?", groupID). - Update("protocols", models.JSON(`[]`)).Error; err != nil { - t.Fatal(err) - } - beforeSnapshot := fixture.manager.Current() - - err := fixture.service.writeKeyConfig( - t.Context(), - groupID, - row.ID, - func(tx *gorm.DB) error { - return tx.Model(&models.UpstreamKey{}). - Where("id = ?", row.ID). - Update("status", models.UpstreamKeyStatusDisabled).Error - }, - func() error { - return fixture.registry.UpdateKeyConfig(row.ID, state.KeyStatusDisabled, nil) - }, - ) - if err != nil { - t.Fatalf("writeKeyConfig() error = %v", err) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("writeKeyConfig published Snapshot") - } -} - func TestWriteConfigDiscardsConnectionAfterCommitBusy(t *testing.T) { fixture, dsn := newFileServiceFixture(t) beforeRevision := fixture.manager.Current().Revision @@ -416,9 +184,8 @@ func TestWriteConfigRollsBackWhenCompileRejectsCandidate(t *testing.T) { _, err := fixture.service.writeConfig(context.Background(), func(tx *gorm.DB) error { return tx.Create(&models.Group{ - Name: "invalid", UpstreamURL: "https://invalid.example", - Protocols: models.JSON(`[]`), - Models: models.JSON(`[]`), Config: models.JSON(`{}`), Enabled: true, + Name: "invalid", ChannelID: "unknown", Params: models.JSON(`{}`), + Models: models.JSON(`[]`), Overrides: models.JSON(`{}`), Enabled: true, }).Error }, nil) if err == nil { @@ -434,17 +201,17 @@ func TestWriteConfigAppliesRuntimeBeforePublishingSnapshot(t *testing.T) { fixture := newServiceFixture(t) beforeSnapshot := fixture.manager.Current() group := validControlGroup("registry-before-snapshot") - var key models.UpstreamKey + var credential models.Credential snapshot, err := fixture.service.writeConfig(t.Context(), func(tx *gorm.DB) error { if err := tx.Create(group).Error; err != nil { return err } - key = models.UpstreamKey{ - GroupID: group.ID, KeyValue: "ciphertext-runtime-order", - KeyHash: "hash-runtime-order", Status: models.UpstreamKeyStatusActive, + credential = models.Credential{ + GroupID: group.ID, Data: "ciphertext-runtime-order", + Fingerprint: "hash-runtime-order", Status: models.CredentialStatusActive, } - return tx.Create(&key).Error + return tx.Create(&credential).Error }, func() error { if fixture.service.writeMu.TryLock() { fixture.service.writeMu.Unlock() @@ -453,9 +220,12 @@ func TestWriteConfigAppliesRuntimeBeforePublishingSnapshot(t *testing.T) { if fixture.manager.Current() != beforeSnapshot { return fmt.Errorf("Snapshot published before Registry update") } - return fixture.registry.ApplyImport(group.ID, []state.KeyEntry{{ - ID: key.ID, GroupID: group.ID, Status: state.KeyStatusActive, - EncryptedValue: key.KeyValue, + return fixture.registry.ApplyCredentialImport(group.ID, []state.CredentialEntry{{ + ID: credential.ID, GroupID: group.ID, + Version: groupCollectionCredentialVersion(credential.UpdatedAtMS), + IdentityGeneration: groupCollectionCredentialIdentity(credential.Fingerprint), + Fingerprint: credential.Fingerprint, Status: state.CredentialStatusActive, + EncryptedValue: credential.Data, }}) }) if err != nil { @@ -467,7 +237,7 @@ func TestWriteConfigAppliesRuntimeBeforePublishingSnapshot(t *testing.T) { if _, ok := snapshot.Groups[group.ID]; !ok { t.Fatalf("Snapshot missing Group %d", group.ID) } - if got, ok := fixture.registry.EncryptedValue(key.ID); !ok || got != key.KeyValue { + if got, ok := fixture.registry.EncryptedCredentialData(credential.ID); !ok || got != credential.Data { t.Fatalf("Registry key = %q, %t", got, ok) } } @@ -495,12 +265,12 @@ func TestWriteConfigMakesCreatedGroupAndFirstKeyAtomicallyVisibleToDataPlane(t * })) defer upstream.Close() - dialects := dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)) + dialects := dialect.NewSet(dialect.NewOpenAI()) handler := gateway.NewHandler( fixture.manager, fixture.registry, fixture.encryption, - gateway.NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()), + gateway.NewExecutionForwarder(controlHTTPExecutor{}), dialects, health.NewStatsStore(), health.NewMutationCoordinator(), @@ -523,13 +293,15 @@ func TestWriteConfigMakesCreatedGroupAndFirstKeyAtomicallyVisibleToDataPlane(t * } group := validControlGroup("atomic-runtime-publication") - group.UpstreamURL = upstream.URL + "/v1" + group.ChannelID = string(channel.OpenAICompatible) + group.Params = models.JSON(`{"base_url":"` + upstream.URL + `/v1"}`) const providerKey = "sk-atomic-runtime-publication" - ciphertext, err := fixture.encryption.Encrypt(providerKey) + credentialData := `{"api_key":"` + providerKey + `"}` + ciphertext, err := fixture.encryption.Encrypt(credentialData) if err != nil { t.Fatalf("Encrypt(provider key) error = %v", err) } - var key models.UpstreamKey + var credential models.Credential runtimeApplied := make(chan struct{}) allowPublish := make(chan struct{}) var releaseOnce sync.Once @@ -546,16 +318,17 @@ func TestWriteConfigMakesCreatedGroupAndFirstKeyAtomicallyVisibleToDataPlane(t * if createErr := tx.Create(group).Error; createErr != nil { return createErr } - key = models.UpstreamKey{ - GroupID: group.ID, KeyValue: ciphertext, - KeyHash: fixture.encryption.Hash(providerKey), Status: models.UpstreamKeyStatusActive, + credential = models.Credential{ + GroupID: group.ID, Data: ciphertext, + Fingerprint: fixture.encryption.Hash(credentialData), Status: models.CredentialStatusActive, } - return tx.Create(&key).Error + return tx.Create(&credential).Error }, func() error { - if applyErr := fixture.registry.ApplyImport(group.ID, []state.KeyEntry{{ - ID: key.ID, GroupID: group.ID, Status: state.KeyStatusActive, - EncryptedValue: key.KeyValue, - }}); applyErr != nil { + entries, buildErr := stateloader.BuildGroupCredentialEntries(t.Context(), fixture.db, group.ID) + if buildErr != nil { + return buildErr + } + if applyErr := fixture.registry.ApplyCredentialImport(group.ID, entries); applyErr != nil { return applyErr } close(runtimeApplied) @@ -647,7 +420,7 @@ func TestWriteConfigRuntimeFailureKeepsOldSnapshot(t *testing.T) { for _, required := range []string{ `"stage":"apply_committed_registry_mutation"`, `"group_id":91`, - `"key_id":7`, + `"credential_id":7`, } { if !strings.Contains(logText, required) { t.Fatalf("log output missing %q: %s", required, logText) @@ -703,12 +476,12 @@ func TestConcurrentCreateGroupsPublishDatabaseTruth(t *testing.T) { fixture := newServiceFixture(t) before := fixture.manager.Current().Revision requests := []GroupCreateRequest{ - {UpstreamURL: "https://shared.example.com/v1", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Keys: "sk-shared-a", ConfirmSameUpstreamURL: true}, - {UpstreamURL: "https://shared.example.com/v1", Protocols: []protocol.Protocol{protocol.Anthropic}, Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Keys: "sk-shared-b", ConfirmSameUpstreamURL: true}, - {UpstreamURL: "https://one.example.com/v1", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Keys: "sk-one"}, - {UpstreamURL: "https://two.example.com/v1", Protocols: []protocol.Protocol{protocol.Gemini}, Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Keys: "sk-two"}, - {UpstreamURL: "https://three.example.com/v1", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Keys: "sk-three"}, - {UpstreamURL: "https://four.example.com/v1", Protocols: []protocol.Protocol{protocol.Anthropic}, Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Keys: "sk-four"}, + {ChannelID: channel.OpenAICompatible, Params: json.RawMessage(`{"base_url":"https://shared.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Credentials: "sk-shared-a", ConfirmSameTarget: true}, + {ChannelID: channel.OpenAICompatible, Params: json.RawMessage(`{"base_url":"https://shared.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Credentials: "sk-shared-b", ConfirmSameTarget: true}, + {ChannelID: channel.OpenAICompatible, Params: json.RawMessage(`{"base_url":"https://one.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Credentials: "sk-one"}, + {ChannelID: channel.OpenAICompatible, Params: json.RawMessage(`{"base_url":"https://two.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Credentials: "sk-two"}, + {ChannelID: channel.OpenAICompatible, Params: json.RawMessage(`{"base_url":"https://three.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Credentials: "sk-three"}, + {ChannelID: channel.OpenAICompatible, Params: json.RawMessage(`{"base_url":"https://four.example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, Credentials: "sk-four"}, } start := make(chan struct{}) @@ -751,9 +524,9 @@ func TestConcurrentCreateGroupsPublishDatabaseTruth(t *testing.T) { } func validControlGroup(name string) *models.Group { + baseURL := "https://" + name + ".example/v1" return &models.Group{ - Name: name, UpstreamURL: "https://" + name + ".example/v1", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[{"id":"gpt-4o"}]`), Config: models.JSON(`{}`), Enabled: true, + Name: name, ChannelID: "openai_compatible", Params: models.JSON(`{"base_url":"` + baseURL + `"}`), + Models: models.JSON(`[{"id":"gpt-4o"}]`), Overrides: models.JSON(`{}`), Enabled: true, } } diff --git a/internal/control/settings.go b/internal/control/settings.go index 258fab58..743da0a8 100644 --- a/internal/control/settings.go +++ b/internal/control/settings.go @@ -171,7 +171,7 @@ func (s *Service) UpdateSettingsIfMatch( updatedAutoSyncEnabled bool ) err = s.withControlTransaction(ctx, func(tx *gorm.DB) error { - currentInput, err := stateloader.BuildCompileInput(ctx, tx) + currentInput, err := stateloader.BuildCompileInput(ctx, tx, s.channelRegistry) if err != nil { return err } @@ -204,7 +204,7 @@ func (s *Service) UpdateSettingsIfMatch( if err := s.applySettingUpdates(tx, updates); err != nil { return err } - input, err = stateloader.BuildCompileInput(ctx, tx) + input, err = stateloader.BuildCompileInput(ctx, tx, s.channelRegistry) if err != nil { return err } diff --git a/internal/control/settings_occ_test.go b/internal/control/settings_occ_test.go index ac7194fb..58a36a46 100644 --- a/internal/control/settings_occ_test.go +++ b/internal/control/settings_occ_test.go @@ -10,8 +10,8 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" - "gpt-load/internal/protocol" "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" ) @@ -187,11 +187,10 @@ func TestSettingsETagIgnoresOtherResourcesAndSurvivesRuntimeReload(t *testing.T) groupName := "ETag unrelated group" if _, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ - Name: &groupName, - UpstreamURL: "https://example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + Name: &groupName, ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://example.com/v1"}`), Models: optionalGroupModels{Set: true, Values: []GroupModel{}}, - Keys: "sk-unrelated-etag-key", + Credentials: "sk-unrelated-etag-key", }); err != nil { t.Fatal(err) } @@ -204,7 +203,7 @@ func TestSettingsETagIgnoresOtherResourcesAndSurvivesRuntimeReload(t *testing.T) } reloadedManager := state.NewManager() - reloadedRegistry := state.NewKeyRegistry() + reloadedRegistry := state.NewCredentialRegistry() if err := stateloader.New(fixture.db, reloadedManager, reloadedRegistry).Load(t.Context()); err != nil { t.Fatal(err) } @@ -216,7 +215,7 @@ func TestSettingsETagIgnoresOtherResourcesAndSurvivesRuntimeReload(t *testing.T) fixture.catalogRuntime, nil, fixture.encryption, - fixture.service.dialects, + fixture.service.executor, fixture.service.requestLogs, fixture.service.usageStats, fixture.service.homeStatistics, diff --git a/internal/control/test_helpers_test.go b/internal/control/test_helpers_test.go index f07bc094..3d2f07cd 100644 --- a/internal/control/test_helpers_test.go +++ b/internal/control/test_helpers_test.go @@ -1,6 +1,7 @@ package control import ( + "encoding/json" "fmt" "net/http" "path/filepath" @@ -15,7 +16,7 @@ import ( "gorm.io/gorm/logger" "gpt-load/internal/catalog" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" "gpt-load/internal/gateway" "gpt-load/internal/health" "gpt-load/internal/platform/encryption" @@ -25,6 +26,7 @@ import ( "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" "gpt-load/internal/storage" + "gpt-load/internal/storage/models" ) var ( @@ -33,6 +35,48 @@ var ( testIdempotencySequence atomic.Uint64 ) +type blockingDecryptService struct { + encryption.Service + started chan<- struct{} + release <-chan struct{} +} + +func (service blockingDecryptService) Decrypt(ciphertext string) (string, error) { + close(service.started) + <-service.release + return service.Service.Decrypt(ciphertext) +} + +func loadCreatedGroupModels(t *testing.T, fixture serviceFixture, groupID uint) []GroupModel { + t.Helper() + var group models.Group + if err := fixture.db.First(&group, groupID).Error; err != nil { + t.Fatalf("query group %d: %v", groupID, err) + } + var result []GroupModel + if err := json.Unmarshal(group.Models, &result); err != nil { + t.Fatalf("decode group %d models: %v", groupID, err) + } + if result == nil { + result = make([]GroupModel, 0) + } + return result +} + +func createGroupWithCredentials(t *testing.T, fixture serviceFixture, credentials string) uint { + t.Helper() + name := fmt.Sprintf("credential-group-%d", testIdempotencySequence.Add(1)) + result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{ + Name: &name, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: optionalGroupModels{Set: true, Values: []GroupModel{{ID: "gpt-4o"}}}, + Credentials: credentials, + }) + if err != nil { + t.Fatalf("CreateGroup(%q) error = %v", name, err) + } + return result.GroupID +} + // RegisterRoutes keeps package tests concise while production registration // remains exclusively owned by the shared HTTP route registry. func (s *Server) RegisterRoutes(engine *gin.Engine) { @@ -67,7 +111,8 @@ func setRequiredTestIdempotencyHeader(request *http.Request) { type serviceFixture struct { db *gorm.DB manager *state.Manager - registry *state.KeyRegistry + registry *state.CredentialRegistry + channelRegistry *channel.Registry priceRuntime *PriceRuntime catalogRuntime *catalog.Runtime encryption encryption.Service @@ -134,7 +179,8 @@ func newServiceFixtureWithDSN(t *testing.T, dsn string) serviceFixture { t.Helper() db := openControlTestDBWithDSN(t, dsn) manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() + channelRegistry := channel.NewRegistry() keyService, err := encryption.NewService("control-test-master-key-material-2026") if err != nil { t.Fatalf("encryption.NewService() error = %v", err) @@ -148,7 +194,7 @@ func newServiceFixtureWithDSN(t *testing.T, dsn string) serviceFixture { priceRuntime := NewPriceRuntime() catalogRuntime := &catalog.Runtime{} return serviceFixture{ - db: db, manager: manager, registry: registry, encryption: keyService, + db: db, manager: manager, registry: registry, channelRegistry: channelRegistry, encryption: keyService, priceRuntime: priceRuntime, catalogRuntime: catalogRuntime, stats: stats, mutations: mutations, requestLogStats: requestLogStats, service: NewService( @@ -159,13 +205,14 @@ func newServiceFixtureWithDSN(t *testing.T, dsn string) serviceFixture { catalogRuntime, nil, keyService, - dialect.NewSet(), + controlHTTPExecutor{}, nil, nil, nil, stats, mutations, requestLogStats, + channelRegistry, ), } } diff --git a/internal/control/testdata/idempotency/group-key-import-k-twice.json b/internal/control/testdata/idempotency/group-key-import-k-twice.json deleted file mode 100644 index 39474f65..00000000 --- a/internal/control/testdata/idempotency/group-key-import-k-twice.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "name": "group-key-import-k-twice", - "operation_kind": "group_key_import", - "path_template": "/api/groups/:group_id/keys/import", - "resource_locator": "group:7", - "canonical_body": "{\"keys\":[\"K\",\"K\"]}", - "domain_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e6379", - "separator_hex": "00", - "framed_input_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e637900000000013100000004504f53540000001067726f75705f6b65795f696d706f7274000000212f6170692f67726f7570732f3a67726f75705f69642f6b6579732f696d706f72740000000767726f75703a3700000010636f6e74726f6c2d61646d696e2d7631000000127b226b657973223a5b224b222c224b225d7d", - "expected_sha256": "8e98a779323ab4b2652c7814126307c14aa106e435b5abdf5a3f98c8be1f264b" -} diff --git a/internal/control/testdata/idempotency/group-key-import-k.json b/internal/control/testdata/idempotency/group-key-import-k.json deleted file mode 100644 index fcfaf836..00000000 --- a/internal/control/testdata/idempotency/group-key-import-k.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "name": "group-key-import-k", - "operation_kind": "group_key_import", - "path_template": "/api/groups/:group_id/keys/import", - "resource_locator": "group:7", - "canonical_body": "{\"keys\":[\"K\"]}", - "domain_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e6379", - "separator_hex": "00", - "framed_input_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e637900000000013100000004504f53540000001067726f75705f6b65795f696d706f7274000000212f6170692f67726f7570732f3a67726f75705f69642f6b6579732f696d706f72740000000767726f75703a3700000010636f6e74726f6c2d61646d696e2d76310000000e7b226b657973223a5b224b225d7d", - "expected_sha256": "33ed46d786c43e583a1aab96104d834e26f8d46d402cf8fcacb83ec7e87373d4" -} diff --git a/internal/control/upstream_keys.go b/internal/control/upstream_keys.go deleted file mode 100644 index bf0272c4..00000000 --- a/internal/control/upstream_keys.go +++ /dev/null @@ -1,404 +0,0 @@ -package control - -import ( - "context" - "errors" - "fmt" - "time" - - "gorm.io/gorm" - - "gpt-load/internal/health" - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/state" - "gpt-load/internal/storage/models" -) - -type UpstreamKeyUpdateRequest struct { - Status optionalField[state.KeyStatus] `json:"status"` - WeightManual optionalField[int] `json:"weight_manual"` -} - -type GroupKeyRevealResult struct { - ID uint `json:"id"` - Key string `json:"key"` - RevealedAtMS int64 `json:"revealed_at_ms"` -} - -type groupKeysCapture struct { - group models.Group - rows []models.UpstreamKey - views []state.KeyRuntimeView - observedAt time.Time -} - -type groupKeysObservation struct { - group models.Group - rows []models.UpstreamKey - runtime map[uint]state.KeyRuntimeView - observedAt time.Time -} - -func equalOptionalWeight(left, right *int) bool { - if left == nil || right == nil { - return left == nil && right == nil - } - return *left == *right -} - -func (s *Service) captureGroupKeys( - ctx context.Context, - groupID uint, -) (groupKeysCapture, error) { - s.writeMu.RLock() - - var group models.Group - groupErr := s.db.WithContext(ctx). - Where("id = ?", groupID).Take(&group).Error - rows := make([]models.UpstreamKey, 0) - var rowsErr error - if groupErr == nil { - rowsErr = s.db.WithContext(ctx). - Where("group_id = ?", groupID). - Order("id ASC"). - Find(&rows).Error - } - var views []state.KeyRuntimeView - var observedAt time.Time - if groupErr == nil && rowsErr == nil { - views = s.registry.Snapshot() - observedAt = s.now().UTC() - } - s.writeMu.RUnlock() - - if groupErr != nil { - if errors.Is(groupErr, gorm.ErrRecordNotFound) { - return groupKeysCapture{}, groupNotFoundError() - } - return groupKeysCapture{}, app_errors.ParseDBError(groupErr) - } - if rowsErr != nil { - return groupKeysCapture{}, app_errors.ParseDBError(rowsErr) - } - return groupKeysCapture{ - group: group, rows: rows, views: views, observedAt: observedAt, - }, nil -} - -func validateGroupKeysCapture( - capture groupKeysCapture, -) (groupKeysObservation, error) { - groupID := capture.group.ID - byID := make(map[uint]state.KeyRuntimeView, len(capture.views)) - for _, view := range capture.views { - byID[view.ID] = view - } - for _, row := range capture.rows { - view, exists := byID[row.ID] - if !exists { - return groupKeysObservation{}, dbRegistryMismatch( - mismatchMissingRegistry, groupID, row.ID, - ) - } - if view.GroupID != groupID { - return groupKeysObservation{}, dbRegistryMismatch( - mismatchGroupID, groupID, row.ID, - ) - } - if view.Status != state.KeyStatus(row.Status) { - return groupKeysObservation{}, dbRegistryMismatch( - mismatchStatus, groupID, row.ID, - ) - } - if !equalOptionalWeight(view.WeightManual, row.WeightManual) { - return groupKeysObservation{}, dbRegistryMismatch( - mismatchWeightManual, groupID, row.ID, - ) - } - } - persisted := make(map[uint]struct{}, len(capture.rows)) - for _, row := range capture.rows { - persisted[row.ID] = struct{}{} - } - for _, view := range capture.views { - if view.GroupID != groupID { - continue - } - if _, exists := persisted[view.ID]; !exists { - return groupKeysObservation{}, dbRegistryMismatch( - mismatchExtraRegistry, groupID, view.ID, - ) - } - } - return groupKeysObservation{ - group: capture.group, rows: capture.rows, - runtime: byID, observedAt: capture.observedAt, - }, nil -} - -func (s *Service) mapGroupKeyItem( - row models.UpstreamKey, - view state.KeyRuntimeView, - group models.Group, - stats health.KeyStats, - observedAt time.Time, -) (GroupKeyItemResponse, error) { - plaintext, err := s.encryption.Decrypt(row.KeyValue) - if err != nil { - return GroupKeyItemResponse{}, fmt.Errorf( - "decrypt group key %d: %w", - row.ID, - app_errors.ErrInternalServer, - ) - } - mask, err := maskGroupKeyCollection(plaintext) - if err != nil { - return GroupKeyItemResponse{}, err - } - bucket := classifyHealthKey(state.GroupCatalogView{ - ID: group.ID, Name: group.Name, Enabled: group.Enabled, - WeightManual: cloneInt(group.WeightManual), - }, view, observedAt) - return mapGroupKeyCollectionItem(mask, row.ID, view, bucket, stats, observedAt) -} - -func normalizeUpstreamKeyUpdate( - request UpstreamKeyUpdateRequest, -) (status *state.KeyStatus, weight *int, weightSet bool, err error) { - if !request.Status.Set && !request.WeightManual.Set { - return nil, nil, false, app_errors.ErrBadRequest - } - if request.Status.Set { - if request.Status.Null || - (request.Status.Value != state.KeyStatusActive && - request.Status.Value != state.KeyStatusDisabled) { - return nil, nil, false, app_errors.ErrValidation - } - value := request.Status.Value - status = &value - } - if request.WeightManual.Set { - weightSet = true - if !request.WeightManual.Null { - if request.WeightManual.Value < 1 || - request.WeightManual.Value > state.MaxWeight { - return nil, nil, false, app_errors.ErrValidation - } - value := request.WeightManual.Value - weight = &value - } - } - return status, weight, weightSet, nil -} - -func findRuntimeKey( - views []state.KeyRuntimeView, - keyID uint, -) (state.KeyRuntimeView, bool) { - for _, view := range views { - if view.ID == keyID { - return view, true - } - } - return state.KeyRuntimeView{}, false -} - -func (s *Service) RevealGroupKey( - ctx context.Context, - groupID uint, - keyID uint, -) (GroupKeyRevealResult, error) { - if groupID == 0 || keyID == 0 { - return GroupKeyRevealResult{}, app_errors.ErrBadRequest - } - var group models.Group - if err := s.db.WithContext(ctx). - Select("id"). - Where("id = ?", groupID). - Take(&group).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return GroupKeyRevealResult{}, groupNotFoundError() - } - return GroupKeyRevealResult{}, app_errors.ParseDBError(err) - } - var row struct { - ID uint - KeyValue string - } - if err := s.db.WithContext(ctx). - Model(&models.UpstreamKey{}). - Select("id", "key_value"). - Where("id = ? AND group_id = ?", keyID, groupID). - Take(&row).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return GroupKeyRevealResult{}, keyNotFoundError() - } - return GroupKeyRevealResult{}, app_errors.ParseDBError(err) - } - plaintext, err := s.encryption.Decrypt(row.KeyValue) - if err != nil { - return GroupKeyRevealResult{}, fmt.Errorf( - "reveal group key %d: %w", - row.ID, - app_errors.ErrInternalServer, - ) - } - if _, err := maskGroupKeyCollection(plaintext); err != nil { - return GroupKeyRevealResult{}, fmt.Errorf( - "reveal group key %d: %w", - row.ID, - app_errors.ErrInternalServer, - ) - } - revealedAtMS, err := safeEpochMilliseconds(s.now()) - if err != nil { - return GroupKeyRevealResult{}, app_errors.ErrInternalServer - } - return GroupKeyRevealResult{ - ID: row.ID, - Key: plaintext, - RevealedAtMS: revealedAtMS, - }, nil -} - -func (s *Service) UpdateGroupKey( - ctx context.Context, - groupID uint, - keyID uint, - request UpstreamKeyUpdateRequest, -) (GroupKeyItemResponse, error) { - if groupID == 0 || keyID == 0 { - return GroupKeyItemResponse{}, app_errors.ErrBadRequest - } - status, weight, weightSet, err := normalizeUpstreamKeyUpdate(request) - if err != nil { - return GroupKeyItemResponse{}, err - } - - var committed models.UpstreamKey - err = s.writeKeyConfig( - ctx, - groupID, - keyID, - func(tx *gorm.DB) error { - var group models.Group - if err := tx.Select("id"). - Where("id = ?", groupID). - Take(&group).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return groupNotFoundError() - } - return app_errors.ParseDBError(err) - } - if err := tx.Where("id = ? AND group_id = ?", keyID, groupID). - Take(&committed).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return keyNotFoundError() - } - return app_errors.ParseDBError(err) - } - updates := make(map[string]any, 2) - if status != nil { - committed.Status = models.UpstreamKeyStatus(*status) - updates["status"] = committed.Status - } - if weightSet { - committed.WeightManual = cloneInt(weight) - updates["weight_manual"] = committed.WeightManual - } - if err := tx.Model(&models.UpstreamKey{}). - Where("id = ? AND group_id = ?", keyID, groupID). - Updates(updates).Error; err != nil { - return app_errors.ParseDBError(err) - } - return nil - }, - func() error { - view, exists := findRuntimeKey(s.registry.Snapshot(), keyID) - if exists && view.GroupID != groupID { - return fmt.Errorf("Registry key belongs to another Group") - } - if !exists { - return s.registry.ApplyImport(groupID, []state.KeyEntry{{ - ID: keyID, GroupID: groupID, - Status: state.KeyStatus(committed.Status), - WeightManual: cloneInt(committed.WeightManual), - WeightAuto: state.DefaultWeight, - EncryptedValue: committed.KeyValue, - }}) - } - return s.registry.UpdateKeyConfig( - keyID, - state.KeyStatus(committed.Status), - committed.WeightManual, - ) - }, - ) - if err != nil { - return GroupKeyItemResponse{}, err - } - - capture, err := s.captureGroupKeys(ctx, groupID) - if err != nil { - return GroupKeyItemResponse{}, err - } - observation, err := validateGroupKeysCapture(capture) - if err != nil { - return GroupKeyItemResponse{}, err - } - for _, row := range observation.rows { - if row.ID == keyID { - return s.mapGroupKeyItem( - row, - observation.runtime[keyID], - observation.group, - s.stats.Snapshot(keyID, observation.observedAt), - observation.observedAt, - ) - } - } - return GroupKeyItemResponse{}, keyNotFoundError() -} - -func (s *Service) DeleteGroupKey( - ctx context.Context, - groupID uint, - keyID uint, -) error { - if groupID == 0 || keyID == 0 { - return app_errors.ErrBadRequest - } - return s.writeKeyConfig( - ctx, - groupID, - keyID, - func(tx *gorm.DB) error { - var group models.Group - if err := tx.Select("id"). - Where("id = ?", groupID). - Take(&group).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return groupNotFoundError() - } - return app_errors.ParseDBError(err) - } - var row models.UpstreamKey - if err := tx.Select("id"). - Where("id = ? AND group_id = ?", keyID, groupID). - Take(&row).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return keyNotFoundError() - } - return app_errors.ParseDBError(err) - } - if err := tx.Delete(&row).Error; err != nil { - return app_errors.ParseDBError(err) - } - return nil - }, - func() error { - s.registry.RemoveKey(keyID) - return nil - }, - ) -} diff --git a/internal/control/upstream_keys_test.go b/internal/control/upstream_keys_test.go deleted file mode 100644 index eab37e5b..00000000 --- a/internal/control/upstream_keys_test.go +++ /dev/null @@ -1,1565 +0,0 @@ -package control - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "math/rand" - "net/http" - "net/http/httptest" - "reflect" - "strconv" - "strings" - "testing" - "time" - - "github.com/gin-gonic/gin" - "github.com/sirupsen/logrus" - - "gpt-load/internal/dialect" - "gpt-load/internal/platform/config" - "gpt-load/internal/platform/encryption" - app_errors "gpt-load/internal/platform/errors" - "gpt-load/internal/protocol" - "gpt-load/internal/scheduler" - "gpt-load/internal/state" - "gpt-load/internal/storage/models" -) - -func seedManagedUpstreamKey( - t *testing.T, - fixture serviceFixture, - groupID uint, - plaintext string, - status models.UpstreamKeyStatus, - weight *int, -) models.UpstreamKey { - t.Helper() - ciphertext, err := fixture.encryption.Encrypt(plaintext) - if err != nil { - t.Fatal(err) - } - row := models.UpstreamKey{ - GroupID: groupID, KeyValue: ciphertext, - KeyHash: fixture.encryption.Hash(plaintext), - Status: status, WeightManual: weight, - } - if err := fixture.db.Create(&row).Error; err != nil { - t.Fatal(err) - } - if err := fixture.registry.ApplyImport(groupID, []state.KeyEntry{{ - ID: row.ID, GroupID: groupID, EncryptedValue: row.KeyValue, - Status: state.KeyStatus(row.Status), WeightManual: row.WeightManual, - }}); err != nil { - t.Fatal(err) - } - return row -} - -func TestRevealGroupKeyReturnsPlaintextOnlyForOwningGroup(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("reveal-group-key") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - other := validControlGroup("reveal-group-key-other") - if err := fixture.db.Create(other).Error; err != nil { - t.Fatal(err) - } - const plaintext = "sk-reveal-group-key-secret" - row := seedManagedUpstreamKey( - t, - fixture, - group.ID, - plaintext, - models.UpstreamKeyStatusActive, - nil, - ) - now := time.Date(2026, time.August, 1, 9, 30, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - - revealed, err := fixture.service.RevealGroupKey(t.Context(), group.ID, row.ID) - if err != nil { - t.Fatalf("RevealGroupKey() error = %v", err) - } - if revealed.ID != row.ID || revealed.Key != plaintext || - revealed.RevealedAtMS != now.UnixMilli() { - t.Fatalf("RevealGroupKey() = %#v", revealed) - } - - if _, err := fixture.service.RevealGroupKey(t.Context(), other.ID, row.ID); !errors.Is(err, app_errors.ErrResourceNotFound) { - t.Fatalf("cross-Group RevealGroupKey() error = %v, want resource not found", err) - } -} - -func TestGroupKeyRevealHTTPUsesSecretResponseContract(t *testing.T) { - initControlI18n(t) - fixture := newServiceFixture(t) - group := validControlGroup("reveal-group-key-http") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - other := validControlGroup("reveal-group-key-http-other") - if err := fixture.db.Create(other).Error; err != nil { - t.Fatal(err) - } - const plaintext = "sk-reveal-group-key-http-secret" - row := seedManagedUpstreamKey( - t, - fixture, - group.ID, - plaintext, - models.UpstreamKeyStatusActive, - nil, - ) - engine := gin.New() - NewServer( - &config.Config{AuthKey: groupKeyHTTPAuth}, - fixture.service, - ).RegisterRoutes(engine) - - list := serveGroupKeyHTTPRequest( - t, - engine, - http.MethodGet, - fmt.Sprintf("/api/groups/%d/keys", group.ID), - "", - groupKeyHTTPAuth, - "en-US", - ) - if list.Code != http.StatusOK { - t.Fatalf("GET group keys = %d %s", list.Code, list.Body.String()) - } - assertGroupKeyHTTPDoesNotExpose( - t, - list.Body.String(), - plaintext, - row.KeyValue, - row.KeyHash, - ) - - reveal := serveGroupKeyHTTPRequest( - t, - engine, - http.MethodPost, - fmt.Sprintf("/api/groups/%d/keys/%d/reveal", group.ID, row.ID), - "", - groupKeyHTTPAuth, - "en-US", - ) - if reveal.Code != http.StatusOK || - reveal.Header().Get("Cache-Control") != "no-store" || - reveal.Header().Get("Pragma") != "no-cache" || - !strings.Contains(reveal.Body.String(), plaintext) { - t.Fatalf( - "POST group key reveal = %d headers=%v body=%s", - reveal.Code, - reveal.Header(), - reveal.Body.String(), - ) - } - - crossGroup := serveGroupKeyHTTPRequest( - t, - engine, - http.MethodPost, - fmt.Sprintf("/api/groups/%d/keys/%d/reveal", other.ID, row.ID), - "", - groupKeyHTTPAuth, - "en-US", - ) - if crossGroup.Code != http.StatusNotFound { - t.Fatalf("cross-Group reveal = %d %s", crossGroup.Code, crossGroup.Body.String()) - } - assertGroupKeyHTTPDoesNotExpose( - t, - crossGroup.Body.String(), - plaintext, - row.KeyValue, - row.KeyHash, - ) -} - -func TestUpdateGroupKeyCommittedRegistryFailureLogStage(t *testing.T) { - initControlI18n(t) - const ( - authKey = "known-update-auth-secret" - plaintext = "known-update-upstream-secret" - querySecret = "known-update-query-secret" - headerSecret = "known-update-header-secret" - ) - fixture := newServiceFixture(t) - group := validControlGroup("update-commit-log") - group.UpstreamURL += "?token=" + querySecret - group.Config = models.JSON( - `{"header_rules":{"set":{"X-Secret":"` + headerSecret + `"}}}`, - ) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - other := validControlGroup("update-commit-log-other") - if err := fixture.db.Create(other).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, plaintext, - models.UpstreamKeyStatusActive, nil, - ) - accessKey, err := fixture.service.CreateAccessKey( - t.Context(), - AccessKeyCreateRequest{Name: "update-commit-access-key"}, - ) - if err != nil { - t.Fatal(err) - } - accessKeyRow := loadAccessKeyRow(t, fixture.db, accessKey.ID) - if err := fixture.registry.Replace([]state.KeyEntry{{ - ID: row.ID, GroupID: other.ID, - Status: state.KeyStatusActive, WeightAuto: state.DefaultWeight, - EncryptedValue: row.KeyValue, - }}); err != nil { - t.Fatal(err) - } - beforeSnapshot := fixture.manager.Current() - - engine := gin.New() - NewServer(&config.Config{AuthKey: authKey}, fixture.service).RegisterRoutes(engine) - var logs bytes.Buffer - logger := logrus.StandardLogger() - previousOutput, previousFormatter := logger.Out, logger.Formatter - logrus.SetOutput(&logs) - logrus.SetFormatter(&logrus.JSONFormatter{DisableTimestamp: true}) - t.Cleanup(func() { - logrus.SetOutput(previousOutput) - logrus.SetFormatter(previousFormatter) - }) - - recorder := serveGroupKeyHTTPRequest( - t, - engine, - http.MethodPut, - fmt.Sprintf("/api/groups/%d/keys/%d", group.ID, row.ID), - `{"status":"disabled"}`, - authKey, - "en-US", - ) - if recorder.Code != http.StatusInternalServerError { - t.Fatalf("PUT response = %d %s, want 500", recorder.Code, recorder.Body.String()) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("post-commit failure published Snapshot") - } - var committed models.UpstreamKey - if err := fixture.db.First(&committed, row.ID).Error; err != nil { - t.Fatal(err) - } - if committed.Status != models.UpstreamKeyStatusDisabled { - t.Fatalf("committed status = %q, want disabled", committed.Status) - } - view, exists := findRuntimeKey(fixture.registry.Snapshot(), row.ID) - if !exists || view.GroupID != other.ID || view.Status != state.KeyStatusActive { - t.Fatalf("Registry view = %#v, exists=%t", view, exists) - } - - logText := logs.String() - for _, required := range []string{ - `"operation":"update_group_key"`, - `"stage":"apply_committed_registry_mutation"`, - `"group_id":` + strconv.FormatUint(uint64(group.ID), 10), - `"key_id":` + strconv.FormatUint(uint64(row.ID), 10), - } { - if !strings.Contains(logText, required) { - t.Fatalf("logs missing %q: %s", required, logText) - } - } - if strings.Contains(logText, `"mismatch_kind"`) { - t.Fatalf("post-commit log contains mismatch kind: %s", logText) - } - for _, forbidden := range []string{ - authKey, - accessKey.Key, - accessKeyRow.KeyValue, - accessKeyRow.KeyHash, - plaintext, - row.KeyValue, - row.KeyHash, - group.UpstreamURL, - querySecret, - headerSecret, - string(group.Config), - other.UpstreamURL, - } { - if strings.Contains(logText, forbidden) || - strings.Contains(recorder.Body.String(), forbidden) { - t.Fatalf( - "response/log exposes %q: response=%s logs=%s", - forbidden, - recorder.Body.String(), - logText, - ) - } - } -} - -func TestGroupKeyRoutesAreRegistered(t *testing.T) { - initControlI18n(t) - const authKey = "group-key-route-auth" - for _, test := range []struct { - name string - method string - suffix string - body string - }{ - {name: "list", method: http.MethodGet, suffix: "/keys"}, - { - name: "update", method: http.MethodPut, - suffix: "/keys/{key_id}", body: `{"status":"disabled"}`, - }, - {name: "reveal", method: http.MethodPost, suffix: "/keys/{key_id}/reveal"}, - {name: "delete", method: http.MethodDelete, suffix: "/keys/{key_id}"}, - } { - t.Run(test.name, func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-route-" + test.name) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-route-"+test.name, - models.UpstreamKeyStatusActive, nil, - ) - engine := gin.New() - NewServer( - &config.Config{AuthKey: authKey}, - fixture.service, - ).RegisterRoutes(engine) - path := "/api/groups/" + - strconv.FormatUint(uint64(group.ID), 10) + - strings.ReplaceAll( - test.suffix, - "{key_id}", - strconv.FormatUint(uint64(row.ID), 10), - ) - recorder := httptest.NewRecorder() - request := httptest.NewRequest( - test.method, - path, - strings.NewReader(test.body), - ) - request.Header.Set("Authorization", "Bearer "+authKey) - engine.ServeHTTP(recorder, request) - if recorder.Code != http.StatusOK { - t.Fatalf( - "%s %s status = %d, body=%s", - test.method, - path, - recorder.Code, - recorder.Body.String(), - ) - } - }) - } -} - -const groupKeyHTTPAuth = "group-key-http-auth" - -func serveGroupKeyHTTPRequest( - t *testing.T, - engine *gin.Engine, - method string, - path string, - body string, - auth string, - language string, -) *httptest.ResponseRecorder { - t.Helper() - recorder := httptest.NewRecorder() - request := httptest.NewRequest(method, path, strings.NewReader(body)) - if auth != "" { - request.Header.Set("Authorization", "Bearer "+auth) - } - if language != "" { - request.Header.Set("Accept-Language", language) - } - if body != "" { - request.Header.Set("Content-Type", "application/json") - } - engine.ServeHTTP(recorder, request) - return recorder -} - -func assertGroupKeyHTTPDoesNotExpose( - t *testing.T, - response string, - forbidden ...string, -) { - t.Helper() - for _, value := range forbidden { - if value != "" && strings.Contains(response, value) { - t.Fatalf("HTTP response exposes %q: %s", value, response) - } - } -} - -type groupKeyMutationState struct { - row models.UpstreamKey - runtime []state.KeyRuntimeView - snapshot *state.ConfigSnapshot -} - -func captureGroupKeyMutationState( - t *testing.T, - fixture serviceFixture, - keyID uint, -) groupKeyMutationState { - t.Helper() - var row models.UpstreamKey - if err := fixture.db.First(&row, keyID).Error; err != nil { - t.Fatalf("load UpstreamKey %d: %v", keyID, err) - } - return groupKeyMutationState{ - row: row, runtime: fixture.registry.Snapshot(), - snapshot: fixture.manager.Current(), - } -} - -func assertGroupKeyMutationStateUnchanged( - t *testing.T, - fixture serviceFixture, - keyID uint, - want groupKeyMutationState, -) { - t.Helper() - got := captureGroupKeyMutationState(t, fixture, keyID) - if !reflect.DeepEqual(got.row, want.row) { - t.Fatalf("persisted UpstreamKey changed:\ngot=%#v\nwant=%#v", got.row, want.row) - } - if !reflect.DeepEqual(got.runtime, want.runtime) { - t.Fatalf("Registry runtime changed:\ngot=%#v\nwant=%#v", got.runtime, want.runtime) - } - if got.snapshot != want.snapshot { - t.Fatalf( - "Snapshot changed: got=%p/%d want=%p/%d", - got.snapshot, - got.snapshot.Revision, - want.snapshot, - want.snapshot.Revision, - ) - } -} - -func TestGroupKeyHTTPRoutesRequireAuthentication(t *testing.T) { - initControlI18n(t) - for _, route := range []struct { - name string - method string - suffix string - body string - }{ - {name: "list", method: http.MethodGet, suffix: "/keys"}, - { - name: "update", method: http.MethodPut, - suffix: "/keys/{key_id}", body: `{"status":"disabled"}`, - }, - {name: "reveal", method: http.MethodPost, suffix: "/keys/{key_id}/reveal"}, - {name: "delete", method: http.MethodDelete, suffix: "/keys/{key_id}"}, - } { - for _, auth := range []struct { - name string - value string - }{ - {name: "missing"}, - {name: "wrong", value: groupKeyHTTPAuth + "-wrong"}, - } { - t.Run(route.name+"/"+auth.name, func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-auth-" + route.name + "-" + auth.name) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-auth-"+route.name+"-"+auth.name, - models.UpstreamKeyStatusActive, nil, - ) - before := captureGroupKeyMutationState(t, fixture, row.ID) - engine := gin.New() - NewServer( - &config.Config{AuthKey: groupKeyHTTPAuth}, - fixture.service, - ).RegisterRoutes(engine) - path := "/api/groups/" + - strconv.FormatUint(uint64(group.ID), 10) + - strings.ReplaceAll( - route.suffix, - "{key_id}", - strconv.FormatUint(uint64(row.ID), 10), - ) - recorder := serveGroupKeyHTTPRequest( - t, - engine, - route.method, - path, - route.body, - auth.value, - "en-US", - ) - if recorder.Code != http.StatusUnauthorized || - !strings.Contains(recorder.Body.String(), `"code":"UNAUTHORIZED"`) { - t.Fatalf( - "%s %s response = %d %s, want 401", - route.method, - path, - recorder.Code, - recorder.Body.String(), - ) - } - assertGroupKeyMutationStateUnchanged(t, fixture, row.ID, before) - }) - } - } -} - -func TestGroupKeyHTTPRejectsInvalidPathIDsWithoutMutation(t *testing.T) { - initControlI18n(t) - fixture := newServiceFixture(t) - group := validControlGroup("key-invalid-path") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-invalid-path", - models.UpstreamKeyStatusActive, nil, - ) - engine := gin.New() - NewServer( - &config.Config{AuthKey: groupKeyHTTPAuth}, - fixture.service, - ).RegisterRoutes(engine) - groupText := strconv.FormatUint(uint64(group.ID), 10) - keyText := strconv.FormatUint(uint64(row.ID), 10) - const overflow = "18446744073709551616" - - for _, test := range []struct { - name string - method string - path string - body string - }{ - {name: "list zero Group", method: http.MethodGet, path: "/api/groups/0/keys"}, - {name: "list invalid Group", method: http.MethodGet, path: "/api/groups/not-a-number/keys"}, - {name: "list overflow Group", method: http.MethodGet, path: "/api/groups/" + overflow + "/keys"}, - { - name: "update zero Group", method: http.MethodPut, - path: "/api/groups/0/keys/" + keyText, body: `{"status":"disabled"}`, - }, - { - name: "update invalid Key", method: http.MethodPut, - path: "/api/groups/" + groupText + "/keys/not-a-number", - body: `{"status":"disabled"}`, - }, - { - name: "update overflow Key", method: http.MethodPut, - path: "/api/groups/" + groupText + "/keys/" + overflow, - body: `{"status":"disabled"}`, - }, - { - name: "delete zero Key", method: http.MethodDelete, - path: "/api/groups/" + groupText + "/keys/0", - }, - { - name: "delete overflow Group", method: http.MethodDelete, - path: "/api/groups/" + overflow + "/keys/" + keyText, - }, - } { - t.Run(test.name, func(t *testing.T) { - before := captureGroupKeyMutationState(t, fixture, row.ID) - recorder := serveGroupKeyHTTPRequest( - t, - engine, - test.method, - test.path, - test.body, - groupKeyHTTPAuth, - "en-US", - ) - if recorder.Code != http.StatusBadRequest || - !strings.Contains(recorder.Body.String(), `"code":"BAD_REQUEST"`) { - t.Fatalf( - "%s %s response = %d %s, want 400", - test.method, - test.path, - recorder.Code, - recorder.Body.String(), - ) - } - assertGroupKeyMutationStateUnchanged(t, fixture, row.ID, before) - }) - } -} - -func TestUpdateGroupKeyHTTPStrictJSONValidation(t *testing.T) { - initControlI18n(t) - fixture := newServiceFixture(t) - group := validControlGroup("key-update-http-validation") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - weight := 33 - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-update-http-validation", - models.UpstreamKeyStatusActive, &weight, - ) - engine := gin.New() - NewServer( - &config.Config{AuthKey: groupKeyHTTPAuth}, - fixture.service, - ).RegisterRoutes(engine) - path := fmt.Sprintf("/api/groups/%d/keys/%d", group.ID, row.ID) - - for _, test := range []struct { - name string - body string - code string - }{ - {name: "empty body", code: app_errors.ErrInvalidJSON.Code}, - {name: "top-level null", body: `null`, code: app_errors.ErrInvalidJSON.Code}, - {name: "top-level array", body: `[]`, code: app_errors.ErrInvalidJSON.Code}, - {name: "unknown field", body: `{"unknown":true}`, code: app_errors.ErrInvalidJSON.Code}, - { - name: "multiple JSON values", - body: `{"status":"disabled"} {"status":"active"}`, - code: app_errors.ErrInvalidJSON.Code, - }, - {name: "empty object", body: `{}`, code: app_errors.ErrBadRequest.Code}, - {name: "null status", body: `{"status":null}`, code: app_errors.ErrValidation.Code}, - { - name: "invalid status", body: `{"status":"cooldown"}`, - code: app_errors.ErrValidation.Code, - }, - { - name: "negative weight", body: `{"weight_manual":-1}`, - code: app_errors.ErrValidation.Code, - }, - { - name: "large weight", body: `{"weight_manual":101}`, - code: app_errors.ErrValidation.Code, - }, - { - name: "fraction weight", body: `{"weight_manual":1.5}`, - code: app_errors.ErrInvalidJSON.Code, - }, - { - name: "string weight", body: `{"weight_manual":"1"}`, - code: app_errors.ErrInvalidJSON.Code, - }, - } { - t.Run(test.name, func(t *testing.T) { - before := captureGroupKeyMutationState(t, fixture, row.ID) - recorder := serveGroupKeyHTTPRequest( - t, - engine, - http.MethodPut, - path, - test.body, - groupKeyHTTPAuth, - "en-US", - ) - if recorder.Code != http.StatusBadRequest || - !strings.Contains( - recorder.Body.String(), - `"code":"`+test.code+`"`, - ) { - t.Fatalf( - "PUT validation response = %d %s, want %s", - recorder.Code, - recorder.Body.String(), - test.code, - ) - } - assertGroupKeyMutationStateUnchanged(t, fixture, row.ID, before) - assertGroupKeyHTTPDoesNotExpose( - t, - recorder.Body.String(), - "sk-update-http-validation", - row.KeyValue, - row.KeyHash, - group.UpstreamURL, - string(group.Config), - ) - }) - } -} - -func TestGroupKeyHTTPReturnsExactResourceNotFoundMessages(t *testing.T) { - initControlI18n(t) - for _, operation := range []struct { - name string - method string - body string - }{ - {name: "update", method: http.MethodPut, body: `{"status":"disabled"}`}, - {name: "delete", method: http.MethodDelete}, - } { - for _, resource := range []struct { - name string - groupMode string - keyMode string - message string - }{ - {name: "missing Group", groupMode: "missing", message: "Group not found"}, - {name: "missing Key", keyMode: "missing", message: "Key not found"}, - {name: "wrong owner", keyMode: "wrong-owner", message: "Key not found"}, - } { - t.Run(operation.name+"/"+resource.name, func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup( - "key-http-not-found-" + operation.name + "-" + - strings.ReplaceAll(resource.name, " ", "-"), - ) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - const plaintext = "sk-key-http-not-found" - row := seedManagedUpstreamKey( - t, fixture, group.ID, plaintext, - models.UpstreamKeyStatusActive, nil, - ) - groupID, keyID := group.ID, row.ID - var wrongOwner models.UpstreamKey - if resource.groupMode == "missing" { - groupID = 999999 - } - if resource.keyMode == "missing" { - keyID = 999999 - } - if resource.keyMode == "wrong-owner" { - other := validControlGroup("key-http-not-found-other") - if err := fixture.db.Create(other).Error; err != nil { - t.Fatal(err) - } - wrongOwner = seedManagedUpstreamKey( - t, fixture, other.ID, "sk-wrong-owner-http", - models.UpstreamKeyStatusActive, nil, - ) - keyID = wrongOwner.ID - } - before := captureGroupKeyMutationState(t, fixture, row.ID) - engine := gin.New() - NewServer( - &config.Config{AuthKey: groupKeyHTTPAuth}, - fixture.service, - ).RegisterRoutes(engine) - recorder := serveGroupKeyHTTPRequest( - t, - engine, - operation.method, - fmt.Sprintf("/api/groups/%d/keys/%d", groupID, keyID), - operation.body, - groupKeyHTTPAuth, - "en-US", - ) - if recorder.Code != http.StatusNotFound { - t.Fatalf( - "not-found response = %d %s", - recorder.Code, - recorder.Body.String(), - ) - } - var envelope struct { - Code string `json:"code"` - Message string `json:"message"` - Data json.RawMessage `json:"data"` - } - if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil { - t.Fatal(err) - } - if envelope.Code != app_errors.ErrResourceNotFound.Code || - envelope.Message != resource.message || - len(envelope.Data) != 0 { - t.Fatalf("not-found envelope = %#v", envelope) - } - assertGroupKeyMutationStateUnchanged(t, fixture, row.ID, before) - assertGroupKeyHTTPDoesNotExpose( - t, - recorder.Body.String(), - plaintext, - row.KeyValue, - row.KeyHash, - group.UpstreamURL, - string(group.Config), - wrongOwner.KeyValue, - wrongOwner.KeyHash, - ) - }) - } - } -} - -func TestDeleteGroupKeyHTTPSuccessAllowsLastKey(t *testing.T) { - initControlI18n(t) - fixture := newServiceFixture(t) - group := validControlGroup("key-delete-http") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - const plaintext = "sk-delete-http-last" - row := seedManagedUpstreamKey( - t, fixture, group.ID, plaintext, - models.UpstreamKeyStatusActive, nil, - ) - beforeSnapshot := fixture.manager.Current() - engine := gin.New() - NewServer( - &config.Config{AuthKey: groupKeyHTTPAuth}, - fixture.service, - ).RegisterRoutes(engine) - recorder := serveGroupKeyHTTPRequest( - t, - engine, - http.MethodDelete, - fmt.Sprintf("/api/groups/%d/keys/%d", group.ID, row.ID), - "", - groupKeyHTTPAuth, - "en-US", - ) - if recorder.Code != http.StatusOK { - t.Fatalf("DELETE response = %d %s", recorder.Code, recorder.Body.String()) - } - var envelope struct { - Code int `json:"code"` - Data json.RawMessage `json:"data"` - } - if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil { - t.Fatal(err) - } - if envelope.Code != 0 || len(envelope.Data) != 0 { - t.Fatalf("DELETE envelope = %#v, want code 0 and omitted data", envelope) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("HTTP delete published Snapshot") - } - if _, exists := fixture.registry.EncryptedValue(row.ID); exists { - t.Fatal("deleted Key remains in Registry") - } - var groupCount, keyCount int64 - if err := fixture.db.Model(&models.Group{}). - Where("id = ?", group.ID).Count(&groupCount).Error; err != nil { - t.Fatal(err) - } - if err := fixture.db.Model(&models.UpstreamKey{}). - Where("id = ?", row.ID).Count(&keyCount).Error; err != nil { - t.Fatal(err) - } - if groupCount != 1 || keyCount != 0 { - t.Fatalf("row counts = Group:%d Key:%d", groupCount, keyCount) - } - assertGroupKeyHTTPDoesNotExpose( - t, - recorder.Body.String(), - plaintext, - row.KeyValue, - row.KeyHash, - group.UpstreamURL, - string(group.Config), - ) -} - -func TestCaptureGroupKeysSeparatesLockedCaptureFromValidation(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-short-capture") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-short-capture", - models.UpstreamKeyStatusActive, nil, - ) - - capture, err := fixture.service.captureGroupKeys(t.Context(), group.ID) - if err != nil { - t.Fatal(err) - } - if len(capture.rows) != 1 || capture.rows[0].ID != row.ID || - len(capture.views) != 1 || capture.views[0].ID != row.ID { - t.Fatalf("raw capture = %#v", capture) - } - - fixture.service.writeMu.Lock() - validationDone := make(chan error, 1) - go func() { - _, err := validateGroupKeysCapture(capture) - validationDone <- err - }() - select { - case err := <-validationDone: - fixture.service.writeMu.Unlock() - if err != nil { - t.Fatalf("validateGroupKeysCapture() error = %v", err) - } - case <-time.After(2 * time.Second): - fixture.service.writeMu.Unlock() - t.Fatal("pair validation waited for writeMu") - } -} - -type blockingDecryptService struct { - encryption.Service - started chan<- struct{} - release <-chan struct{} -} - -func (service blockingDecryptService) Decrypt(ciphertext string) (string, error) { - close(service.started) - <-service.release - return service.Service.Decrypt(ciphertext) -} - -func TestUpdateGroupKeyChangesConfigAndPreservesAutomaticRuntime(t *testing.T) { - fixture := newServiceFixture(t) - now := time.Date(2026, time.July, 24, 12, 0, 0, 0, time.UTC) - fixture.service.now = func() time.Time { return now } - group := validControlGroup("key-update") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-update-secret", - models.UpstreamKeyStatusActive, nil, - ) - if !fixture.registry.SetAutoWeight(row.ID, 41) || - !fixture.registry.SetCooldown(row.ID, now.Add(time.Minute)) || - !fixture.registry.SetBlacklisted(row.ID) { - t.Fatal("seed runtime") - } - for range 3 { - if _, ok := fixture.registry.IncrFailure(row.ID); !ok { - t.Fatal("seed failure") - } - } - beforeSnapshot := fixture.manager.Current() - weight := 25 - got, err := fixture.service.UpdateGroupKey( - t.Context(), - group.ID, - row.ID, - UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{ - Set: true, Value: state.KeyStatusDisabled, - }, - WeightManual: optionalField[int]{Set: true, Value: weight}, - }, - ) - if err != nil { - t.Fatal(err) - } - if got.ConfiguredStatus != string(state.KeyStatusDisabled) || - got.EffectiveStatus != string(healthBucketDisabled) || - got.Weight != nil || got.Recovery.Mode != "manual" { - t.Fatalf("updated response = %#v", got) - } - view, exists := findRuntimeKey(fixture.registry.Snapshot(), row.ID) - if !exists || view.WeightManual == nil || *view.WeightManual != weight || - view.WeightAuto != 41 || !view.Blacklisted || view.FailureCount != 3 { - t.Fatalf("preserved runtime = %#v, exists=%t", view, exists) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("UpdateGroupKey published Snapshot") - } - - got, err = fixture.service.UpdateGroupKey( - t.Context(), - group.ID, - row.ID, - UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{ - Set: true, Value: state.KeyStatusActive, - }, - WeightManual: optionalField[int]{Set: true, Null: true}, - }, - ) - if err != nil { - t.Fatal(err) - } - if got.ConfiguredStatus != string(state.KeyStatusActive) || - got.EffectiveStatus != string(healthBucketBlacklisted) || - got.Weight != nil || got.Recovery.Mode != "probe" { - t.Fatalf("re-enabled response = %#v", got) - } - view, exists = findRuntimeKey(fixture.registry.Snapshot(), row.ID) - if !exists || view.WeightManual != nil || view.WeightAuto != 41 || - !view.Blacklisted || view.FailureCount != 3 { - t.Fatalf("re-enabled runtime = %#v, exists=%t", view, exists) - } - var persisted models.UpstreamKey - if err := fixture.db.First(&persisted, row.ID).Error; err != nil { - t.Fatal(err) - } - if persisted.Status != models.UpstreamKeyStatusActive || persisted.WeightManual != nil { - t.Fatalf("persisted key = %#v", persisted) - } -} - -func TestUpdateGroupKeyRepairsMissingRegistryEntryFromCommittedRow(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-repair") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - weight := 25 - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-repair", - models.UpstreamKeyStatusActive, &weight, - ) - fixture.registry.RemoveKey(row.ID) - beforeSnapshot := fixture.manager.Current() - - got, err := fixture.service.UpdateGroupKey( - t.Context(), - group.ID, - row.ID, - UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{ - Set: true, Value: state.KeyStatusDisabled, - }, - }, - ) - if err != nil { - t.Fatalf("UpdateGroupKey() error = %v", err) - } - if got.ConfiguredStatus != string(state.KeyStatusDisabled) || - got.EffectiveStatus != string(healthBucketDisabled) || - got.Weight != nil || got.Recovery.Mode != "manual" { - t.Fatalf("repaired response = %#v", got) - } - view, exists := findRuntimeKey(fixture.registry.Snapshot(), row.ID) - if !exists || view.WeightManual == nil || *view.WeightManual != 25 || - view.WeightAuto != state.DefaultWeight { - t.Fatalf("repaired runtime = %#v, exists=%t", view, exists) - } - value, exists := fixture.registry.EncryptedValue(row.ID) - if !exists || value != row.KeyValue { - t.Fatalf("repaired credential = %q, %t", value, exists) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("missing Registry repair published Snapshot") - } -} - -func TestUpdateGroupKeyConvergesStatusAndManualWeightMismatches(t *testing.T) { - for _, test := range []struct { - name string - mutate func(t *testing.T, fixture serviceFixture, row models.UpstreamKey) - }{ - { - name: "status", - mutate: func( - t *testing.T, - fixture serviceFixture, - row models.UpstreamKey, - ) { - if err := fixture.registry.UpdateKeyConfig( - row.ID, - state.KeyStatusDisabled, - row.WeightManual, - ); err != nil { - t.Fatal(err) - } - }, - }, - { - name: "manual weight", - mutate: func( - t *testing.T, - fixture serviceFixture, - row models.UpstreamKey, - ) { - wrongWeight := 77 - if err := fixture.registry.UpdateKeyConfig( - row.ID, - state.KeyStatusActive, - &wrongWeight, - ); err != nil { - t.Fatal(err) - } - }, - }, - } { - t.Run(test.name, func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-converge-" + test.name) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - persistedWeight := 25 - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-converge-"+test.name, - models.UpstreamKeyStatusActive, &persistedWeight, - ) - if !fixture.registry.SetAutoWeight(row.ID, 41) || - !fixture.registry.SetBlacklisted(row.ID) { - t.Fatal("seed automatic runtime") - } - if _, ok := fixture.registry.IncrFailure(row.ID); !ok { - t.Fatal("seed failure count") - } - test.mutate(t, fixture, row) - beforeSnapshot := fixture.manager.Current() - - got, err := fixture.service.UpdateGroupKey( - t.Context(), - group.ID, - row.ID, - UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{ - Set: true, Value: state.KeyStatusActive, - }, - }, - ) - if err != nil { - t.Fatalf("UpdateGroupKey() error = %v", err) - } - if got.ConfiguredStatus != string(state.KeyStatusActive) || - got.EffectiveStatus != string(healthBucketBlacklisted) || - got.Weight != nil || got.Recovery.Mode != "probe" { - t.Fatalf("converged response = %#v", got) - } - view, exists := findRuntimeKey( - fixture.registry.Snapshot(), - row.ID, - ) - if !exists || - view.Status != state.KeyStatusActive || - view.WeightManual == nil || - *view.WeightManual != persistedWeight || - view.WeightAuto != 41 || - !view.Blacklisted || - view.FailureCount != 1 { - t.Fatalf("converged Registry view = %#v, exists=%t", view, exists) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("mismatch convergence published Snapshot") - } - }) - } -} - -func TestUpdateGroupKeyLeavesExtraRegistryEntryFailLoud(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-update-extra") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-update-extra", - models.UpstreamKeyStatusActive, nil, - ) - extraID := row.ID + 1000 - if err := fixture.registry.ApplyImport(group.ID, []state.KeyEntry{{ - ID: extraID, GroupID: group.ID, - Status: state.KeyStatusActive, EncryptedValue: "cipher-extra", - }}); err != nil { - t.Fatal(err) - } - beforeSnapshot := fixture.manager.Current() - - got, err := fixture.service.UpdateGroupKey( - t.Context(), - group.ID, - row.ID, - UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{ - Set: true, Value: state.KeyStatusDisabled, - }, - }, - ) - if got != (GroupKeyItemResponse{}) { - t.Fatalf("UpdateGroupKey() response = %#v, want zero", got) - } - var operationErr *controlOperationError - if !errors.As(err, &operationErr) || - operationErr.stage != stageValidateDBRegistryPair || - operationErr.mismatchKind != mismatchExtraRegistry || - operationErr.groupID != group.ID || - operationErr.keyID != extraID { - t.Fatalf("UpdateGroupKey() error = %#v", err) - } - var committed models.UpstreamKey - if err := fixture.db.First(&committed, row.ID).Error; err != nil { - t.Fatal(err) - } - if committed.Status != models.UpstreamKeyStatusDisabled { - t.Fatalf("committed status = %q, want disabled", committed.Status) - } - views := fixture.registry.Snapshot() - updated, updatedExists := findRuntimeKey(views, row.ID) - extra, extraExists := findRuntimeKey(views, extraID) - if !updatedExists || updated.Status != state.KeyStatusDisabled || - !extraExists || extra.GroupID != group.ID { - t.Fatalf( - "Registry views = %#v, updated=%#v/%t extra=%#v/%t", - views, - updated, - updatedExists, - extra, - extraExists, - ) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("extra Registry mismatch published Snapshot") - } -} - -func TestDeleteGroupKeyAllowsLastKeyAndDoesNotPublish(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-delete") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-last", - models.UpstreamKeyStatusActive, nil, - ) - beforeSnapshot := fixture.manager.Current() - if err := fixture.service.DeleteGroupKey(t.Context(), group.ID, row.ID); err != nil { - t.Fatal(err) - } - if fixture.manager.Current() != beforeSnapshot { - t.Fatal("DeleteGroupKey published Snapshot") - } - if _, exists := fixture.registry.EncryptedValue(row.ID); exists { - t.Fatal("deleted key remains in Registry") - } - var groupCount, keyCount int64 - if err := fixture.db.Model(&models.Group{}). - Where("id = ?", group.ID).Count(&groupCount).Error; err != nil { - t.Fatal(err) - } - if err := fixture.db.Model(&models.UpstreamKey{}). - Where("id = ?", row.ID).Count(&keyCount).Error; err != nil { - t.Fatal(err) - } - if groupCount != 1 || keyCount != 0 { - t.Fatalf("counts = group:%d key:%d", groupCount, keyCount) - } -} - -func TestUpdateGroupKeyValidatesThreeStateFieldsWithoutMutation(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-update-validation") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - weight := 33 - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-update-validation", - models.UpstreamKeyStatusActive, &weight, - ) - for _, test := range []struct { - name string - request UpstreamKeyUpdateRequest - want *app_errors.APIError - }{ - {name: "empty", request: UpstreamKeyUpdateRequest{}, want: app_errors.ErrBadRequest}, - { - name: "null status", - request: UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{Set: true, Null: true}, - }, - want: app_errors.ErrValidation, - }, - { - name: "invalid status", - request: UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{ - Set: true, Value: state.KeyStatus("cooldown"), - }, - }, - want: app_errors.ErrValidation, - }, - { - name: "zero weight", - request: UpstreamKeyUpdateRequest{ - WeightManual: optionalField[int]{Set: true, Value: 0}, - }, - want: app_errors.ErrValidation, - }, - { - name: "negative weight", - request: UpstreamKeyUpdateRequest{ - WeightManual: optionalField[int]{Set: true, Value: -1}, - }, - want: app_errors.ErrValidation, - }, - { - name: "large weight", - request: UpstreamKeyUpdateRequest{ - WeightManual: optionalField[int]{Set: true, Value: 101}, - }, - want: app_errors.ErrValidation, - }, - } { - t.Run(test.name, func(t *testing.T) { - var beforeDB models.UpstreamKey - if err := fixture.db.First(&beforeDB, row.ID).Error; err != nil { - t.Fatal(err) - } - beforeRegistry := fixture.registry.Snapshot() - beforeSnapshot := fixture.manager.Current() - _, err := fixture.service.UpdateGroupKey( - t.Context(), group.ID, row.ID, test.request, - ) - if !errors.Is(err, test.want) { - t.Fatalf("UpdateGroupKey() error = %v, want %v", err, test.want) - } - var afterDB models.UpstreamKey - if err := fixture.db.First(&afterDB, row.ID).Error; err != nil { - t.Fatal(err) - } - if !reflect.DeepEqual(afterDB, beforeDB) || - !reflect.DeepEqual(fixture.registry.Snapshot(), beforeRegistry) || - fixture.manager.Current() != beforeSnapshot { - t.Fatal("invalid update mutated DB/Registry/Snapshot") - } - }) - } -} - -func TestUpdateGroupKeyAcceptsManualWeightBoundariesAndNull(t *testing.T) { - for _, test := range []struct { - name string - field optionalField[int] - want *int - initial *int - }{ - {name: "one", field: optionalField[int]{Set: true, Value: 1}, want: keyTestIntPointer(1)}, - {name: "max", field: optionalField[int]{Set: true, Value: 100}, want: keyTestIntPointer(100)}, - { - name: "null", field: optionalField[int]{Set: true, Null: true}, - initial: keyTestIntPointer(55), - }, - } { - t.Run(test.name, func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-weight-" + test.name) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-weight-"+test.name, - models.UpstreamKeyStatusActive, test.initial, - ) - got, err := fixture.service.UpdateGroupKey( - t.Context(), - group.ID, - row.ID, - UpstreamKeyUpdateRequest{WeightManual: test.field}, - ) - if err != nil { - t.Fatal(err) - } - if test.want == nil { - if got.WeightMode != "auto" || got.Weight == nil || *got.Weight != state.DefaultWeight { - t.Fatalf("automatic weight response = %#v", got) - } - } else if got.WeightMode != "manual" || got.Weight == nil || *got.Weight != *test.want { - t.Fatalf("manual weight response = %#v, want %d", got, *test.want) - } - }) - } -} - -func keyTestIntPointer(value int) *int { - return &value -} - -func TestUpdateAndDeleteGroupKeyReturnTypedResourceNotFound(t *testing.T) { - for _, operation := range []string{"update", "delete"} { - for _, test := range []struct { - name string - groupMissing bool - wrongOwner bool - wantResource string - }{ - {name: "missing Group", groupMissing: true, wantResource: "group"}, - {name: "missing Key", wantResource: "key"}, - {name: "wrong owner", wrongOwner: true, wantResource: "key"}, - } { - t.Run(operation+"/"+test.name, func(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-resource-" + operation + "-" + - strings.ReplaceAll(test.name, " ", "-")) - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - groupID := group.ID - keyID := uint(999) - if test.groupMissing { - groupID += 1000 - } - if test.wrongOwner { - other := validControlGroup("key-resource-other-" + operation) - if err := fixture.db.Create(other).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, other.ID, "sk-wrong-owner-"+operation, - models.UpstreamKeyStatusActive, nil, - ) - keyID = row.ID - } - var err error - if operation == "update" { - _, err = fixture.service.UpdateGroupKey( - t.Context(), - groupID, - keyID, - UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{ - Set: true, Value: state.KeyStatusDisabled, - }, - }, - ) - } else { - err = fixture.service.DeleteGroupKey(t.Context(), groupID, keyID) - } - var resourceErr *controlResourceNotFoundError - if !errors.As(err, &resourceErr) || - resourceErr.resource != test.wantResource || - !errors.Is(err, app_errors.ErrResourceNotFound) { - t.Fatalf("error = %#v, want %s not found", err, test.wantResource) - } - }) - } - } -} - -func TestDeleteGroupKeySucceedsWhenRegistryEntryIsMissing(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-delete-missing-runtime") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-delete-missing-runtime", - models.UpstreamKeyStatusActive, nil, - ) - fixture.registry.RemoveKey(row.ID) - beforeSnapshot := fixture.manager.Current() - - if err := fixture.service.DeleteGroupKey(t.Context(), group.ID, row.ID); err != nil { - t.Fatalf("DeleteGroupKey() error = %v", err) - } - var count int64 - if err := fixture.db.Model(&models.UpstreamKey{}). - Where("id = ?", row.ID).Count(&count).Error; err != nil { - t.Fatal(err) - } - if count != 0 || fixture.manager.Current() != beforeSnapshot { - t.Fatalf("delete missing Registry state = count:%d snapshotChanged:%t", - count, fixture.manager.Current() != beforeSnapshot) - } -} - -func TestUpdateAndDeleteGroupKeyNeverCallDiscovery(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-no-discovery") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - row := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-no-discovery", - models.UpstreamKeyStatusActive, nil, - ) - fixture.service.dialects = dialect.NewSet(&recordingDiscoveryDialect{ - value: protocol.OpenAICompletions, - listFn: func(context.Context, string, string, state.HeaderRules) ([]string, error) { - t.Fatal("Key mutation must not call model discovery or upstream") - return nil, nil - }, - }) - if _, err := fixture.service.UpdateGroupKey( - t.Context(), - group.ID, - row.ID, - UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{ - Set: true, Value: state.KeyStatusDisabled, - }, - }, - ); err != nil { - t.Fatal(err) - } - if err := fixture.service.DeleteGroupKey(t.Context(), group.ID, row.ID); err != nil { - t.Fatal(err) - } -} - -func TestUpdateAndDeleteGroupKeyChangeSchedulerEligibility(t *testing.T) { - fixture := newServiceFixture(t) - group := validControlGroup("key-candidates") - if err := fixture.db.Create(group).Error; err != nil { - t.Fatal(err) - } - disabled := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-candidate-disabled", - models.UpstreamKeyStatusActive, nil, - ) - remaining := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-candidate-remaining", - models.UpstreamKeyStatusActive, nil, - ) - deleted := seedManagedUpstreamKey( - t, fixture, group.ID, "sk-candidate-deleted", - models.UpstreamKeyStatusActive, nil, - ) - if got := fixture.registry.CollectCandidates( - []uint{group.ID}, nil, time.Time{}, - ); len(got) != 3 { - t.Fatalf("initial candidates = %#v", got) - } - if _, err := fixture.service.UpdateGroupKey( - t.Context(), - group.ID, - disabled.ID, - UpstreamKeyUpdateRequest{ - Status: optionalField[state.KeyStatus]{ - Set: true, Value: state.KeyStatusDisabled, - }, - }, - ); err != nil { - t.Fatal(err) - } - if err := fixture.service.DeleteGroupKey(t.Context(), group.ID, deleted.ID); err != nil { - t.Fatal(err) - } - rawPool := fixture.registry.CollectCandidates( - []uint{group.ID}, nil, time.Time{}, - ) - if len(rawPool) != 1 || rawPool[0].ID != remaining.ID { - t.Fatalf("raw candidates = %#v, want only remaining key %d", rawPool, remaining.ID) - } - snapshot, err := state.Compile(state.CompileInput{ - Groups: []state.GroupConfig{{ - ID: group.ID, Name: group.Name, UpstreamURL: group.UpstreamURL, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "gpt-4o"}}, - Enabled: true, - }}, - }) - if err != nil { - t.Fatal(err) - } - iterator := scheduler.New( - snapshot, - fixture.registry, - scheduler.Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: stringPointer("gpt-4o"), - }, - rand.New(rand.NewSource(1)), - ) - if selection, err := iterator.Next(); err != nil || selection.KeyID != remaining.ID { - t.Fatalf("Next() = (%#v, %v), want remaining key %d", selection, err, remaining.ID) - } -} diff --git a/internal/control/usage.go b/internal/control/usage.go index 5a80c709..27a23f8d 100644 --- a/internal/control/usage.go +++ b/internal/control/usage.go @@ -12,6 +12,7 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/platform/epochms" app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/platform/response" @@ -59,8 +60,10 @@ type usageSeriesResponse struct { } type usageBreakdownResponse struct { - GroupID uint `json:"group_id"` - Model string `json:"model"` + GroupID uint `json:"group_id"` + ChannelID *channel.ID `json:"channel_id"` + CredentialID *uint `json:"credential_id"` + Model string `json:"model"` usageAggregateResponse } @@ -113,7 +116,7 @@ func (server *Server) handleUsage(c *gin.Context) { return } if accessKeyID, scoped := currentAccessKeyID(c); scoped { - if query.GroupID != nil { + if query.GroupID != nil || query.ChannelID != "" || query.CredentialID != nil { writeServiceError(c, "usage", app_errors.ErrBadRequest) return } @@ -142,6 +145,8 @@ func parseUsageQuery(rawQuery string, observedAtMS int64) (requestlog.UsageQuery "from_ms": {}, "to_ms": {}, "group_id": {}, + "channel_id": {}, + "credential_id": {}, "upstream_model": {}, "breakdown_order": {}, } @@ -226,6 +231,20 @@ func parseUsageQuery(rawQuery string, observedAtMS int64) (requestlog.UsageQuery } query.GroupID = &groupID } + if value, ok := singleQueryValue(values, "channel_id"); ok { + channelID, apiErr := parseRequestLogChannelID(value) + if apiErr != nil { + return requestlog.UsageQuery{}, apiErr + } + query.ChannelID = channelID + } + if value, ok := singleQueryValue(values, "credential_id"); ok { + credentialID, apiErr := parseUsageGroupID(value) + if apiErr != nil { + return requestlog.UsageQuery{}, apiErr + } + query.CredentialID = &credentialID + } if value, ok := singleQueryValue(values, "upstream_model"); ok { if !validUsageModel(value) { return requestlog.UsageQuery{}, app_errors.ErrValidation @@ -394,15 +413,27 @@ func (service *Service) mapUsageResponse( if uint64(row.GroupID) > uint64(maxSafeInteger) { return usageResponse{}, fmt.Errorf("map usage breakdown: unsafe group") } + channelID, err := nullableRequestLogChannelID(row.ChannelID) + if err != nil { + return usageResponse{}, fmt.Errorf("map usage breakdown: %w", err) + } + credentialID, err := nullableRequestLogID(row.CredentialID, "credential") + if err != nil { + return usageResponse{}, fmt.Errorf("map usage breakdown: %w", err) + } aggregate, err := mapUsageAggregate(row.UsageAggregate) if err != nil { return usageResponse{}, err } result.Breakdown = append(result.Breakdown, usageBreakdownResponse{ - GroupID: row.GroupID, Model: row.Model, usageAggregateResponse: aggregate, + GroupID: row.GroupID, ChannelID: channelID, CredentialID: credentialID, + Model: row.Model, usageAggregateResponse: aggregate, }) if accessKeyScoped { - result.Breakdown[len(result.Breakdown)-1].GroupID = 0 + redacted := &result.Breakdown[len(result.Breakdown)-1] + redacted.GroupID = 0 + redacted.ChannelID = nil + redacted.CredentialID = nil } } return result, nil diff --git a/internal/control/usage_test.go b/internal/control/usage_test.go index 5b76154d..e17fd88c 100644 --- a/internal/control/usage_test.go +++ b/internal/control/usage_test.go @@ -14,6 +14,7 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" "gpt-load/internal/requestlog" "gpt-load/internal/storage/models" @@ -283,7 +284,8 @@ func TestUsageAPISelectsThirtyUTCDaysAndAppliesFilters(t *testing.T) { }, }}, Breakdown: []requestlog.UsageBreakdown{{ - GroupID: 9, Model: "upstream-model", + GroupID: 9, ChannelID: channel.OpenAI, CredentialID: 13, + Model: "upstream-model", UsageAggregate: requestlog.UsageAggregate{ RequestCount: 1, UncachedInputTokens: 2, CacheWriteUnknownTokens: 4, OutputTokens: 3, PricingPartialCount: 1, @@ -297,7 +299,7 @@ func TestUsageAPISelectsThirtyUTCDaysAndAppliesFilters(t *testing.T) { recorder := performUsageRequest( engine, "test-auth-key", - "range=30d&group_id=9&upstream_model=upstream-model&breakdown_order=cost", + "range=30d&group_id=9&channel_id=openai&credential_id=13&upstream_model=upstream-model&breakdown_order=cost", ) if recorder.Code != http.StatusOK { t.Fatalf("response = %d %s", recorder.Code, recorder.Body.String()) @@ -310,6 +312,8 @@ func TestUsageAPISelectsThirtyUTCDaysAndAppliesFilters(t *testing.T) { query.FromMS != time.Date(2026, time.June, 28, 0, 0, 0, 0, time.UTC).UnixMilli() || query.ToMS != time.Date(2026, time.July, 28, 0, 0, 0, 0, time.UTC).UnixMilli() || query.GroupID == nil || *query.GroupID != 9 || + query.ChannelID != channel.OpenAI || + query.CredentialID == nil || *query.CredentialID != 13 || query.UpstreamModel != "upstream-model" || query.Limit != 100 || query.BreakdownOrder != requestlog.UsageBreakdownOrderCost { t.Fatalf("30 day UsageQuery = %#v", query) @@ -337,6 +341,8 @@ func TestUsageAPISelectsThirtyUTCDaysAndAppliesFilters(t *testing.T) { } `json:"series"` Breakdown []struct { GroupID uint `json:"group_id"` + ChannelID string `json:"channel_id"` + CredentialID uint `json:"credential_id"` Model string `json:"model"` CacheWriteUnknownTokens int64 `json:"cache_write_unknown_tokens"` PricingPartialCount int64 `json:"pricing_partial_count"` @@ -364,12 +370,19 @@ func TestUsageAPISelectsThirtyUTCDaysAndAppliesFilters(t *testing.T) { envelope.Data.Series[0].PricingPartialCount != 1 || envelope.Data.Series[0].EstimatedCostNanoUSD != "1123456789012" || len(envelope.Data.Breakdown) != 1 || envelope.Data.Breakdown[0].GroupID != 9 || + envelope.Data.Breakdown[0].ChannelID != string(channel.OpenAI) || + envelope.Data.Breakdown[0].CredentialID != 13 || envelope.Data.Breakdown[0].Model != "upstream-model" || envelope.Data.Breakdown[0].CacheWriteUnknownTokens != 4 || envelope.Data.Breakdown[0].PricingPartialCount != 1 || envelope.Data.Breakdown[0].EstimatedCostNanoUSD != "250000000" { t.Fatalf("usage response = %#v", envelope.Data) } + for _, legacyField := range []string{`"key_id"`, `"upstream_key_id"`} { + if strings.Contains(recorder.Body.String(), legacyField) { + t.Fatalf("usage response exposes legacy field %s: %s", legacyField, recorder.Body.String()) + } + } } func TestUsageAPIValidatesModelAsUTF8BytesWithoutBoundaryWhitespaceOrControls(t *testing.T) { @@ -442,6 +455,15 @@ func TestUsageAPIRejectsStrictInvalidQueriesWithoutCallingReader(t *testing.T) { {query: "group_id=1&group_id=2", code: "BAD_REQUEST"}, {query: "group_id=9007199254740992", code: "VALIDATION_FAILED"}, {query: "group_id=-1", code: "BAD_REQUEST"}, + {query: "channel_id=", code: "VALIDATION_FAILED"}, + {query: "channel_id=unknown", code: "VALIDATION_FAILED"}, + {query: "channel_id=openai&channel_id=anthropic", code: "BAD_REQUEST"}, + {query: "credential_id=0", code: "VALIDATION_FAILED"}, + {query: "credential_id=01", code: "BAD_REQUEST"}, + {query: "credential_id=9007199254740992", code: "VALIDATION_FAILED"}, + {query: "credential_id=1&credential_id=2", code: "BAD_REQUEST"}, + {query: "key_id=1", code: "BAD_REQUEST"}, + {query: "upstream_key_id=1", code: "BAD_REQUEST"}, {query: "model=legacy", code: "BAD_REQUEST"}, {query: "upstream_model=", code: "VALIDATION_FAILED"}, {query: "breakdown_order=", code: "VALIDATION_FAILED"}, @@ -641,8 +663,8 @@ func TestUsageAPIBindsAccessKeyScopeAndRedactsProcessHealth(t *testing.T) { now := time.Date(2026, time.August, 8, 17, 30, 0, 0, time.UTC) reader := &recordingUsageStatReader{report: requestlog.UsageReport{ Breakdown: []requestlog.UsageBreakdown{{ - GroupID: 99, - Model: "allowed-model", + GroupID: 99, ChannelID: channel.OpenAI, CredentialID: 101, + Model: "allowed-model", }}, BreakdownCount: 1, }} @@ -664,7 +686,9 @@ func TestUsageAPIBindsAccessKeyScopeAndRedactsProcessHealth(t *testing.T) { var envelope struct { Data struct { Breakdown []struct { - GroupID uint `json:"group_id"` + GroupID uint `json:"group_id"` + ChannelID *string `json:"channel_id"` + CredentialID *uint `json:"credential_id"` } `json:"breakdown"` CollectionHealth usageCollectionHealthResponse `json:"collection_health"` } `json:"data"` @@ -673,6 +697,8 @@ func TestUsageAPIBindsAccessKeyScopeAndRedactsProcessHealth(t *testing.T) { t.Fatalf("decode AccessKey usage response: %v", err) } if len(envelope.Data.Breakdown) != 1 || envelope.Data.Breakdown[0].GroupID != 0 || + envelope.Data.Breakdown[0].ChannelID != nil || + envelope.Data.Breakdown[0].CredentialID != nil || envelope.Data.CollectionHealth.Scope != "access_key" || envelope.Data.CollectionHealth.DroppedTotal != 0 || envelope.Data.CollectionHealth.WriteFailureTotal != 0 || @@ -680,14 +706,17 @@ func TestUsageAPIBindsAccessKeyScopeAndRedactsProcessHealth(t *testing.T) { t.Fatalf("AccessKey usage redaction = %#v", envelope.Data) } - forbiddenFilter := performUsageRequest(engine, created.Key, "group_id=1") - if forbiddenFilter.Code != http.StatusBadRequest || len(reader.queries) != 1 { - t.Fatalf( - "AccessKey group filter = %d %s, calls=%d, want 400/no query", - forbiddenFilter.Code, - forbiddenFilter.Body.String(), - len(reader.queries), - ) + for _, filter := range []string{"group_id=1", "channel_id=openai", "credential_id=101"} { + forbidden := performUsageRequest(engine, created.Key, filter) + if forbidden.Code != http.StatusBadRequest || len(reader.queries) != 1 { + t.Fatalf( + "AccessKey internal filter %q = %d %s, calls=%d, want 400/no query", + filter, + forbidden.Code, + forbidden.Body.String(), + len(reader.queries), + ) + } } } diff --git a/internal/control/validation.go b/internal/control/validation.go index bb64ba81..69fbbe79 100644 --- a/internal/control/validation.go +++ b/internal/control/validation.go @@ -2,9 +2,11 @@ package control import ( "context" + cryptorand "crypto/rand" "crypto/sha256" "encoding/binary" "hash" + "net/http" "net/textproto" "sort" "strings" @@ -12,7 +14,8 @@ import ( "github.com/sirupsen/logrus" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/health" "gpt-load/internal/platform/encryption" "gpt-load/internal/platform/utils" @@ -27,8 +30,8 @@ type validationSweep interface { } type validationRegistry interface { - BlacklistedKeys() []state.KeyRef - RecoverIfMatch(ref state.KeyRef, weight int) bool + BlacklistedCredentials() []state.CredentialRef + RecoverIfMatch(ref state.CredentialRef, weight int) bool } type statsResetter interface { @@ -48,9 +51,10 @@ type validationWorker struct { snapshots snapshotSource registry validationRegistry stats statsResetter - mutations keyMutationCoordinator + mutations credentialMutationCoordinator decryptor credentialDecryptor - dialects dialect.Set + channels *channel.Registry + executor execution.Executor } type groupValidationSignature [sha256.Size]byte @@ -65,11 +69,12 @@ var _ validationSweep = (*validationWorker)(nil) func newValidationWorker( manager *state.Manager, - registry *state.KeyRegistry, + registry *state.CredentialRegistry, stats *health.StatsStore, mutations *health.MutationCoordinator, decryptor encryption.Service, - dialects dialect.Set, + channels *channel.Registry, + executor execution.Executor, ) *validationWorker { return &validationWorker{ snapshots: manager, @@ -77,7 +82,8 @@ func newValidationWorker( stats: stats, mutations: mutations, decryptor: decryptor, - dialects: dialects, + channels: channels, + executor: executor, } } @@ -89,14 +95,14 @@ func (worker *validationWorker) Validate(ctx context.Context) { if snapshot == nil { return } - refs := worker.registry.BlacklistedKeys() + refs := worker.registry.BlacklistedCredentials() if len(refs) == 0 { return } concurrency := min(validationConcurrency, len(refs)) - jobs := make(chan state.KeyRef) + jobs := make(chan state.CredentialRef) var workers sync.WaitGroup workers.Add(concurrency) for range concurrency { @@ -124,7 +130,7 @@ dispatch: func (worker *validationWorker) consumeValidationJobs( ctx context.Context, snapshot *state.ConfigSnapshot, - jobs <-chan state.KeyRef, + jobs <-chan state.CredentialRef, ) { for { select { @@ -139,7 +145,7 @@ func (worker *validationWorker) consumeValidationJobs( } } -func (worker *validationWorker) validateRef(ctx context.Context, snapshot *state.ConfigSnapshot, ref state.KeyRef) { +func (worker *validationWorker) validateRef(ctx context.Context, snapshot *state.ConfigSnapshot, ref state.CredentialRef) { if ctx.Err() != nil { return } @@ -149,18 +155,12 @@ func (worker *validationWorker) validateRef(ctx context.Context, snapshot *state return } target, ok := buildGroupValidationTarget(group) - if !ok && len(group.Protocols) == 0 { - logValidationFailure(ref, "", "missing_protocol") - return - } if !ok { - logValidationFailure(ref, string(group.Protocols[0]), "missing_model") + logValidationFailure(ref, "", "missing_target_or_model") return } - - selectedDialect, ok := worker.dialects[target.protocol] - if !ok || selectedDialect == nil { - logValidationFailure(ref, string(target.protocol), "missing_dialect") + if worker.executor == nil || worker.channels == nil { + logValidationFailure(ref, string(target.protocol), "missing_executor") return } if worker.decryptor == nil { @@ -171,20 +171,65 @@ func (worker *validationWorker) validateRef(ctx context.Context, snapshot *state logValidationFailure(ref, string(target.protocol), "conditional_recover") return } - apiKey, err := worker.decryptor.Decrypt(ref.EncryptedValue) + plaintext, err := worker.decryptor.Decrypt(ref.EncryptedValue) if err != nil { if ctx.Err() == nil { logValidationFailure(ref, string(target.protocol), "decrypt") } return } - if err := selectedDialect.Probe( - ctx, - group.UpstreamURL, - apiKey, - group.HeaderRules, - target.model, - ); err != nil { + credential, err := normalizeStoredCredential(worker.channels, group.ChannelID, plaintext) + if err != nil { + if ctx.Err() == nil { + logValidationFailure(ref, string(target.protocol), "credential") + } + return + } + apiKey, _ := credential.Value("api_key") + method, path, body, err := validationRequestShape(group, target) + if err != nil { + logValidationFailure(ref, string(target.protocol), "request") + return + } + requestID, err := newOperationID(cryptorand.Reader) + if err != nil { + logValidationFailure(ref, string(target.protocol), "request_identity") + return + } + attemptID, err := newOperationID(cryptorand.Reader) + if err != nil { + logValidationFailure(ref, string(target.protocol), "attempt_identity") + return + } + version := ref.Version + if version == 0 { + version = 1 + } + generation := ref.IdentityGeneration + if generation == 0 { + generation = 1 + } + spec := execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: requestID, AttemptID: attemptID, Sequence: 1, + ChannelID: string(group.ChannelID), ClientProtocol: target.protocol, + Operation: execution.OperationProbe, ClientModel: target.model, UpstreamModel: target.model, + Method: method, Path: path, Header: applyControlHeaderRules(group.HeaderRules, apiKey), Body: body, + TargetConfig: group.ResolvedTarget.TargetConfig, + Timeouts: executionTimeouts(group.Timeouts), + Credential: execution.NewCredentialSnapshot( + ref.ID, version, generation, credential.CanonicalJSON(), + ), + }) + if err := spec.Validate(); err != nil { + logValidationFailure(ref, string(target.protocol), "request") + return + } + result := worker.executor.Execute(ctx, spec) + if ctx.Err() != nil { + return + } + if result.Validate() != nil || result.Error != nil || + result.StatusCode < http.StatusOK || result.StatusCode >= http.StatusMultipleChoices { if ctx.Err() == nil { logValidationFailure(ref, string(target.protocol), "probe") } @@ -230,8 +275,8 @@ func (worker *validationWorker) validateRef(ctx context.Context, snapshot *state } func buildGroupValidationTarget(group state.GroupView) (groupValidationTarget, bool) { - selectedProtocol, ok := representativeProtocol(group.Protocols) - if !ok { + if group.ChannelID == "" || group.ResolvedTarget.ChannelID != group.ChannelID || + !group.ResolvedTarget.ProviderKind.Valid() { return groupValidationTarget{}, false } probeModel := strings.TrimSpace(group.ValidationModel) @@ -241,6 +286,14 @@ func buildGroupValidationTarget(group state.GroupView) (groupValidationTarget, b if probeModel == "" { return groupValidationTarget{}, false } + selectedProtocol, _, _, _, err := utilityRequestShape( + group.ResolvedTarget.ProviderKind, + execution.OperationProbe, + probeModel, + ) + if err != nil { + return groupValidationTarget{}, false + } return groupValidationTarget{ protocol: selectedProtocol, model: probeModel, @@ -248,19 +301,16 @@ func buildGroupValidationTarget(group state.GroupView) (groupValidationTarget, b }, true } -func representativeProtocol( - values []protocol.Protocol, -) (protocol.Protocol, bool) { - present := make(map[protocol.Protocol]struct{}, len(values)) - for _, value := range values { - present[value] = struct{}{} - } - for _, value := range protocol.DataPlaneProtocols() { - if _, exists := present[value]; exists { - return value, true - } - } - return "", false +func validationRequestShape( + group state.GroupView, + target groupValidationTarget, +) (string, string, []byte, error) { + _, method, path, body, err := utilityRequestShape( + group.ResolvedTarget.ProviderKind, + execution.OperationProbe, + target.model, + ) + return method, path, body, err } func computeGroupValidationSignature( @@ -270,7 +320,9 @@ func computeGroupValidationSignature( ) groupValidationSignature { hasher := sha256.New() writeValidationSignatureUint64(hasher, uint64(group.ID)) - writeValidationSignaturePart(hasher, []byte(group.UpstreamURL)) + writeValidationSignaturePart(hasher, []byte(group.ChannelID)) + writeValidationSignaturePart(hasher, []byte(group.ResolvedTarget.ProviderKind)) + writeValidationSignaturePart(hasher, group.ResolvedTarget.TargetConfig) writeValidationSignaturePart(hasher, []byte(selectedProtocol)) writeValidationSignaturePart(hasher, []byte(probeModel)) @@ -329,16 +381,16 @@ func normalizeValidationHeaderName(name string) string { return strings.ToLower(textproto.CanonicalMIMEHeaderKey(name)) } -func logValidationFailure(ref state.KeyRef, protocol, stage string) { +func logValidationFailure(ref state.CredentialRef, protocol, stage string) { utils.LogPlaneBestEffort( logrus.StandardLogger(), logrus.WarnLevel, utils.LogPlaneControl, logrus.Fields{ - "key_id": ref.ID, - "group_id": ref.GroupID, - "protocol": protocol, - "stage": stage, + "credential_id": ref.ID, + "group_id": ref.GroupID, + "protocol": protocol, + "stage": stage, }, "Validation failed", ) diff --git a/internal/control/validation_test.go b/internal/control/validation_test.go index b87c36a4..e5c11433 100644 --- a/internal/control/validation_test.go +++ b/internal/control/validation_test.go @@ -3,6 +3,7 @@ package control import ( "bytes" "context" + "encoding/json" "errors" "fmt" "net/http" @@ -15,7 +16,8 @@ import ( "github.com/sirupsen/logrus" - "gpt-load/internal/dialect" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/health" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -34,7 +36,7 @@ func TestValidationWorkerUsesExplicitModelAndCanonicalRepresentativeProtocol(t * nil, ), }), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, probes, ) @@ -55,7 +57,7 @@ func TestValidationWorkerFallsBackToFirstRealModelID(t *testing.T) { validationSnapshot(map[uint]state.GroupView{ 1: validationGroup([]protocol.Protocol{protocol.OpenAICompletions}, " \t", []state.ModelConfig{{ID: " real-model ", Alias: "external-model"}}), }), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, probes, ) @@ -78,14 +80,14 @@ func TestValidationWorkerUsesResponsesProbeForResponsesOnlyGroup(t *testing.T) { nil, ), }), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, probes, ) worker.Validate(context.Background()) want := []validationProbeCall{{ - protocol: protocol.OpenAIResponses, + protocol: protocol.OpenAICompletions, model: "gpt-5", apiKey: "plain-key-7", }} @@ -94,19 +96,64 @@ func TestValidationWorkerUsesResponsesProbeForResponsesOnlyGroup(t *testing.T) { } } +func TestValidationWorkerProbesStructuredCloudCredential(t *testing.T) { + group := validationGroup(nil, "anthropic.claude-test", nil) + setValidationChannel(&group, channel.AWSBedrock, json.RawMessage(`{"region":"us-east-1"}`)) + worker := newValidationWorkerForTest( + validationSnapshot(map[uint]state.GroupView{1: group}), + []state.CredentialRef{{ + ID: 7, GroupID: 1, Version: 1, IdentityGeneration: 7, + Fingerprint: "bedrock-fingerprint", EncryptedValue: "bedrock-cipher", + }}, + &validationProbeRecorder{}, + ) + worker.decryptor = fixedValidationDecryptor{ + plaintext: `{"access_key":"AKIA_TEST","secret_key":"bedrock-secret"}`, + } + var observed execution.AttemptSpec + worker.executor = scriptedDiscoveryExecutor{execute: func( + _ context.Context, + spec execution.AttemptSpec, + ) execution.AttemptResult { + observed = spec.Clone() + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: http.Header{}, + } + }} + + worker.Validate(context.Background()) + + if got, want := string(observed.Credential.Data()), + `{"access_key":"AKIA_TEST","secret_key":"bedrock-secret"}`; got != want { + t.Fatalf("credential = %s, want %s", got, want) + } + if observed.ChannelID != string(channel.AWSBedrock) || + observed.Operation != execution.OperationProbe || + observed.ClientProtocol != protocol.OpenAICompletions || + observed.UpstreamModel != "anthropic.claude-test" { + t.Fatalf("attempt = %#v", observed) + } + if got, want := worker.recorder.events(), []string{ + fmt.Sprintf("registry.weight:7:%d", state.DefaultWeight), "registry.recover:7", "stats.reset:7", + }; !reflect.DeepEqual(got, want) { + t.Fatalf("recovery events = %#v, want %#v", got, want) + } +} + func TestValidationSignatureUsesCanonicalLengthPrefixedEncoding(t *testing.T) { - group := state.GroupView{ - ID: 42, - UpstreamURL: "https://upstream.example.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - ValidationModel: " model-a ", - HeaderRules: state.HeaderRules{ - Set: map[string]string{ - "X-Zeta": "last", - "x-alpha": "first", - }, - Remove: []string{"X-Old", "x-beta"}, + group := validationGroup( + []protocol.Protocol{protocol.OpenAICompletions}, " model-a ", nil, + ) + group.ID = 42 + group.HeaderRules = state.HeaderRules{ + Set: map[string]string{ + "X-Zeta": "last", + "x-alpha": "first", }, + Remove: []string{"X-Old", "x-beta"}, } target, ok := buildGroupValidationTarget(group) @@ -116,7 +163,7 @@ func TestValidationSignatureUsesCanonicalLengthPrefixedEncoding(t *testing.T) { if target.protocol != protocol.OpenAICompletions || target.model != "model-a" { t.Fatalf("target protocol/model = %q/%q, want openai/model-a", target.protocol, target.model) } - const want = "7ece2253053165e5b6ee50d11ab75e86feffca6162a52f08361d6ee2010e0760" + const want = "96758d0d5d971626ff21a2e501edd86c6799c2aa17d7924cc40469d2e3a698b9" if got := fmt.Sprintf("%x", target.signature); got != want { t.Fatalf("validation signature = %s, want canonical digest %s", got, want) } @@ -132,11 +179,11 @@ func TestValidationSignatureChangesForEveryCoveredInput(t *testing.T) { {name: "group id", base: base, mutate: func(group *state.GroupView) { group.ID++ }}, - {name: "upstream URL", base: base, mutate: func(group *state.GroupView) { - group.UpstreamURL += "/changed" + {name: "channel target", base: base, mutate: func(group *state.GroupView) { + setValidationChannel(group, channel.OpenAICompatible, json.RawMessage(`{"base_url":"https://changed.example/v1"}`)) }}, - {name: "selected protocol", base: base, mutate: func(group *state.GroupView) { - group.Protocols[0] = protocol.Anthropic + {name: "provider family", base: base, mutate: func(group *state.GroupView) { + setValidationChannel(group, channel.Anthropic, json.RawMessage(`{}`)) }}, {name: "explicit validation model", base: base, mutate: func(group *state.GroupView) { group.ValidationModel = "model-b" @@ -215,12 +262,13 @@ func TestValidationWorkerSkipsMissingGroupProtocolModelAndDialect(t *testing.T) 2: validationGroup([]protocol.Protocol{protocol.OpenAICompletions}, " \t", nil), 3: validationGroup([]protocol.Protocol{protocol.Gemini}, "model", nil), }) - worker := newValidationWorkerForTest(snapshot, []state.KeyRef{ + worker := newValidationWorkerForTest(snapshot, []state.CredentialRef{ {ID: 1, GroupID: 9, EncryptedValue: "key-1"}, {ID: 2, GroupID: 1, EncryptedValue: "key-2"}, {ID: 3, GroupID: 2, EncryptedValue: "key-3"}, {ID: 4, GroupID: 3, EncryptedValue: "key-4"}, }, probes) + worker.executor = nil worker.Validate(context.Background()) @@ -236,7 +284,7 @@ func TestValidationWorkerKeepsKeyBlacklistedOnDecryptOrProbeFailure(t *testing.T probes := &validationProbeRecorder{errByKey: map[string]error{"plain-key-2": errors.New("probe failed")}} worker := newValidationWorkerForTest( validationSnapshot(map[uint]state.GroupView{1: validationGroup([]protocol.Protocol{protocol.OpenAICompletions}, "model", nil)}), - []state.KeyRef{ + []state.CredentialRef{ {ID: 1, GroupID: 1, EncryptedValue: "decrypt-fails"}, {ID: 2, GroupID: 1, EncryptedValue: "key-2"}, }, @@ -258,7 +306,7 @@ func TestValidationWorkerCoordinatesConditionalRecoveryAndStatsReset(t *testing. probes := &validationProbeRecorder{} worker := newValidationWorkerForTest( validationSnapshot(map[uint]state.GroupView{1: validationGroup([]protocol.Protocol{protocol.OpenAICompletions}, "model", nil)}), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, probes, ) coordinator := &barrierValidationMutationCoordinator{ @@ -293,7 +341,7 @@ func TestValidationWorkerCoordinatesConditionalRecoveryAndStatsReset(t *testing. t.Fatalf("snapshot reads = %d, want 1", got) } if got := worker.registry.(*validationRegistryRecorder).blacklistedCalls(); got != 1 { - t.Fatalf("BlacklistedKeys calls = %d, want 1", got) + t.Fatalf("BlacklistedCredentials calls = %d, want 1", got) } } @@ -301,7 +349,7 @@ func TestValidationWorkerRecoverIfMatchFailsDoesNotResetStats(t *testing.T) { probes := &validationProbeRecorder{} worker := newValidationWorkerForTest( validationSnapshot(map[uint]state.GroupView{1: validationGroup([]protocol.Protocol{protocol.OpenAICompletions}, "model", nil)}), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, probes, ) worker.registry.(*validationRegistryRecorder).recoveryOK = false @@ -315,9 +363,9 @@ func TestValidationWorkerRecoverIfMatchFailsDoesNotResetStats(t *testing.T) { func TestValidationWorkerFailureGenerationChangesDuringProbeRejectsRecovery(t *testing.T) { now := time.Date(2026, time.July, 27, 12, 0, 0, 0, time.UTC) - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 7, GroupID: 1, Status: state.KeyStatusActive, + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 7, GroupID: 1, Version: 1, IdentityGeneration: 7, Fingerprint: "test-7", Status: state.CredentialStatusActive, Blacklisted: true, FailureCount: 3, EncryptedValue: "key-7", }}); err != nil { t.Fatalf("Replace() error = %v", err) @@ -334,14 +382,14 @@ func TestValidationWorkerFailureGenerationChangesDuringProbeRejectsRecovery(t *t stats: stats, mutations: health.NewMutationCoordinator(), decryptor: validationDecryptor{}, - dialects: dialect.Set{protocol.OpenAICompletions: &validationTestDialect{ - protocol: protocol.OpenAICompletions, + channels: channel.NewRegistry(), + executor: &validationTestExecutor{ probes: &validationProbeRecorder{probe: func(context.Context, protocol.Protocol, string, string) error { close(probeStarted) <-releaseProbe return nil }}, - }}, + }, } done := make(chan struct{}) @@ -356,12 +404,13 @@ func TestValidationWorkerFailureGenerationChangesDuringProbeRejectsRecovery(t *t close(releaseProbe) awaitValidationDone(t, done) - if got, want := registry.BlacklistedKeys(), []state.KeyRef{{ - ID: 7, GroupID: 1, EncryptedValue: "key-7", FailureGeneration: 1, + if got, want := registry.BlacklistedCredentials(), []state.CredentialRef{{ + ID: 7, GroupID: 1, Version: 1, IdentityGeneration: 7, + Fingerprint: "test-7", EncryptedValue: "key-7", FailureGeneration: 1, }}; !reflect.DeepEqual(got, want) { t.Fatalf("blacklisted keys = %#v, want stale recovery rejected as %#v", got, want) } - if got, want := stats.Snapshot(7, now), (health.KeyStats{ + if got, want := stats.Snapshot(7, now), (health.CredentialStats{ Failure: 1, Problem: 1, ConsecutiveFailure: 1, ConsecutiveProblem: 1, }); got != want { t.Fatalf("stats after stale recovery = %#v, want %#v", got, want) @@ -375,11 +424,11 @@ func TestValidationSignatureChangesDuringProbeRejectRecovery(t *testing.T) { prepare func(*state.GroupView) mutate func(*state.GroupView) }{ - {name: "upstream URL", mutate: func(group *state.GroupView) { - group.UpstreamURL += "/changed" + {name: "channel target", mutate: func(group *state.GroupView) { + setValidationChannel(group, channel.OpenAICompatible, json.RawMessage(`{"base_url":"https://changed.example/v1"}`)) }}, - {name: "protocol", mutate: func(group *state.GroupView) { - group.Protocols[0] = protocol.Anthropic + {name: "provider family", mutate: func(group *state.GroupView) { + setValidationChannel(group, channel.Anthropic, json.RawMessage(`{}`)) }}, {name: "explicit model", mutate: func(group *state.GroupView) { group.ValidationModel = "model-b" @@ -405,13 +454,12 @@ func TestValidationSignatureChangesDuringProbeRejectRecovery(t *testing.T) { } worker := newValidationWorkerForTest( validationSnapshot(map[uint]state.GroupView{1: captured}), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, &validationProbeRecorder{}, ) current := cloneValidationGroup(captured) test.mutate(¤t) - worker.validationWorker.dialects[protocol.OpenAICompletions] = &validationTestDialect{ - protocol: protocol.OpenAICompletions, + worker.validationWorker.executor = &validationTestExecutor{ probes: &validationProbeRecorder{probe: func(context.Context, protocol.Protocol, string, string) error { worker.snapshots.(*validationSnapshotRecorder).setSnapshot( validationSnapshot(map[uint]state.GroupView{1: current}), @@ -436,21 +484,16 @@ func TestValidationWorkerUnrelatedSnapshotRevisionAllowsRecovery(t *testing.T) { Revision: 1, Groups: map[uint]state.GroupView{1: cloneValidationGroup(oldGroup)}, }, - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, &validationProbeRecorder{}, ) - worker.validationWorker.dialects[protocol.OpenAICompletions] = &validationTestDialect{ - protocol: protocol.OpenAICompletions, + worker.validationWorker.executor = &validationTestExecutor{ probes: &validationProbeRecorder{probe: func(context.Context, protocol.Protocol, string, string) error { worker.snapshots.(*validationSnapshotRecorder).setSnapshot(&state.ConfigSnapshot{ Revision: 2, Groups: map[uint]state.GroupView{ 1: cloneValidationGroup(oldGroup), - 2: { - ID: 2, UpstreamURL: "https://unrelated.example.com", - Protocols: []protocol.Protocol{protocol.Anthropic}, - Models: []state.ModelConfig{{ID: "unrelated"}}, - }, + 2: validationGroup([]protocol.Protocol{protocol.Anthropic}, "unrelated", []state.ModelConfig{{ID: "unrelated"}}), }, }) return nil @@ -473,11 +516,10 @@ func TestValidationSignatureRemovedOrDisabledGroupRejectsRecovery(t *testing.T) t.Run(name, func(t *testing.T) { worker := newValidationWorkerForTest( validationSnapshot(map[uint]state.GroupView{1: validationSignatureGroup()}), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, &validationProbeRecorder{}, ) - worker.validationWorker.dialects[protocol.OpenAICompletions] = &validationTestDialect{ - protocol: protocol.OpenAICompletions, + worker.validationWorker.executor = &validationTestExecutor{ probes: &validationProbeRecorder{probe: func(context.Context, protocol.Protocol, string, string) error { // Disabled Groups are omitted from ConfigSnapshot.Groups, as are removed Groups. worker.snapshots.(*validationSnapshotRecorder).setSnapshot(validationSnapshot(nil)) @@ -520,18 +562,18 @@ func (source *observableValidationSnapshotSource) active() bool { } type publicationValidationRegistry struct { - delegate *state.KeyRegistry + delegate *state.CredentialRegistry callbackActive func() bool recoverCallbackActive chan bool recoverEntered chan struct{} releaseRecover chan struct{} } -func (registry *publicationValidationRegistry) BlacklistedKeys() []state.KeyRef { - return registry.delegate.BlacklistedKeys() +func (registry *publicationValidationRegistry) BlacklistedCredentials() []state.CredentialRef { + return registry.delegate.BlacklistedCredentials() } -func (registry *publicationValidationRegistry) RecoverIfMatch(ref state.KeyRef, weight int) bool { +func (registry *publicationValidationRegistry) RecoverIfMatch(ref state.CredentialRef, weight int) bool { registry.recoverCallbackActive <- registry.callbackActive() close(registry.recoverEntered) <-registry.releaseRecover @@ -558,9 +600,9 @@ func TestValidationWorkerPublicationBoundaryBlocksPublishThroughRecoverAndReset( if _, err := manager.Publish(validationManagerCompileInput("https://upstream.example.com")); err != nil { t.Fatalf("initial Publish() error = %v", err) } - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 7, GroupID: 1, Status: state.KeyStatusActive, + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 7, GroupID: 1, Version: 1, IdentityGeneration: 7, Fingerprint: "test-7", Status: state.CredentialStatusActive, Blacklisted: true, FailureCount: 3, EncryptedValue: "key-7", }}); err != nil { t.Fatalf("Replace() error = %v", err) @@ -590,10 +632,8 @@ func TestValidationWorkerPublicationBoundaryBlocksPublishThroughRecoverAndReset( stats: blockingStats, mutations: health.NewMutationCoordinator(), decryptor: validationDecryptor{}, - dialects: dialect.Set{protocol.OpenAICompletions: &validationTestDialect{ - protocol: protocol.OpenAICompletions, - probes: probes, - }}, + channels: channel.NewRegistry(), + executor: &validationTestExecutor{probes: probes}, } validationDone := make(chan struct{}) @@ -634,10 +674,10 @@ func TestValidationWorkerPublicationBoundaryBlocksPublishThroughRecoverAndReset( t.Fatal("timed out waiting for Publish after validation callback returned") } - if got := registry.BlacklistedKeys(); len(got) != 0 { + if got := registry.BlacklistedCredentials(); len(got) != 0 { t.Fatalf("blacklisted keys after recovery = %#v, want none", got) } - if got := stats.Snapshot(7, now); got != (health.KeyStats{}) { + if got := stats.Snapshot(7, now); got != (health.CredentialStats{}) { t.Fatalf("stats after recovery = %#v, want reset", got) } if snapshots.active() { @@ -659,7 +699,7 @@ func TestValidationSignatureMismatchLogDoesNotLeakSensitiveInputs(t *testing.T) }) oldGroup := validationSignatureGroup() - oldGroup.UpstreamURL = "https://sensitive.example.com/path" + setValidationChannel(&oldGroup, channel.OpenAICompatible, json.RawMessage(`{"base_url":"https://sensitive.example.com/path"}`)) oldGroup.HeaderRules.Set["X-Secret"] = "sensitive-header-value" oldTarget, ok := buildGroupValidationTarget(oldGroup) if !ok { @@ -673,11 +713,10 @@ func TestValidationSignatureMismatchLogDoesNotLeakSensitiveInputs(t *testing.T) } worker := newValidationWorkerForTest( validationSnapshot(map[uint]state.GroupView{1: oldGroup}), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "cipher-secret"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "cipher-secret"}}, &validationProbeRecorder{}, ) - worker.validationWorker.dialects[protocol.OpenAICompletions] = &validationTestDialect{ - protocol: protocol.OpenAICompletions, + worker.validationWorker.executor = &validationTestExecutor{ probes: &validationProbeRecorder{probe: func(context.Context, protocol.Protocol, string, string) error { worker.snapshots.(*validationSnapshotRecorder).setSnapshot( validationSnapshot(map[uint]state.GroupView{1: currentGroup}), @@ -735,7 +774,7 @@ func (coordinator *barrierValidationMutationCoordinator) Do(_ uint, fn func()) { func TestValidationWorkerConditionalRecoveryFailureCompletesCoordinatorInterval(t *testing.T) { worker := newValidationWorkerForTest( validationSnapshot(map[uint]state.GroupView{1: validationGroup([]protocol.Protocol{protocol.OpenAICompletions}, "model", nil)}), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "key-7"}}, &validationProbeRecorder{}, ) worker.registry.(*validationRegistryRecorder).recoveryOK = false @@ -766,23 +805,23 @@ func TestValidationWorkerConditionalRecoveryFailureCompletesCoordinatorInterval( func TestValidationWorkerDoesNotRecoverDisabledOrReplacedKeyRef(t *testing.T) { tests := []struct { name string - mutate func(t *testing.T, registry *state.KeyRegistry) + mutate func(t *testing.T, registry *state.CredentialRegistry) expectedCipher string expectedEnabled bool }{ { - name: "disabled after sweep", expectedCipher: "cipher-original", mutate: func(t *testing.T, registry *state.KeyRegistry) { + name: "disabled after sweep", expectedCipher: "cipher-original", mutate: func(t *testing.T, registry *state.CredentialRegistry) { t.Helper() - if err := registry.SetKeyStatus(7, state.KeyStatusDisabled); err != nil { - t.Fatalf("SetKeyStatus() error = %v", err) + if err := registry.SetCredentialStatus(7, state.CredentialStatusDisabled); err != nil { + t.Fatalf("SetCredentialStatus() error = %v", err) } }, }, { - name: "replaced after sweep", expectedCipher: "cipher-replaced", expectedEnabled: true, mutate: func(t *testing.T, registry *state.KeyRegistry) { + name: "replaced after sweep", expectedCipher: "cipher-replaced", expectedEnabled: true, mutate: func(t *testing.T, registry *state.CredentialRegistry) { t.Helper() - if err := registry.Replace([]state.KeyEntry{{ - ID: 7, GroupID: 1, Status: state.KeyStatusActive, Blacklisted: true, + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 7, GroupID: 1, Version: 1, IdentityGeneration: 7, Fingerprint: "test-7", Status: state.CredentialStatusActive, Blacklisted: true, FailureCount: 5, WeightAuto: 17, EncryptedValue: "cipher-replaced", }}); err != nil { t.Fatalf("Replace() error = %v", err) @@ -793,9 +832,9 @@ func TestValidationWorkerDoesNotRecoverDisabledOrReplacedKeyRef(t *testing.T) { for _, test := range tests { t.Run(test.name, func(t *testing.T) { - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 7, GroupID: 1, Status: state.KeyStatusActive, Blacklisted: true, + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 7, GroupID: 1, Version: 1, IdentityGeneration: 7, Fingerprint: "test-7", Status: state.CredentialStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-original", }}); err != nil { t.Fatalf("Replace() error = %v", err) @@ -821,22 +860,23 @@ func TestValidationWorkerDoesNotRecoverDisabledOrReplacedKeyRef(t *testing.T) { close(releaseProbe) awaitValidationDone(t, done) - if got := len(registry.ActiveKeyIDs()); test.expectedEnabled && got != 1 { + if got := len(registry.ActiveCredentialIDs()); test.expectedEnabled && got != 1 { t.Fatalf("active key count = %d, want 1", got) } else if !test.expectedEnabled && got != 0 { t.Fatalf("active key count = %d, want 0", got) } if !test.expectedEnabled { - if err := registry.SetKeyStatus(7, state.KeyStatusActive); err != nil { - t.Fatalf("SetKeyStatus() error = %v", err) + if err := registry.SetCredentialStatus(7, state.CredentialStatusActive); err != nil { + t.Fatalf("SetCredentialStatus() error = %v", err) } } - if got, want := registry.BlacklistedKeys(), []state.KeyRef{{ - ID: 7, GroupID: 1, EncryptedValue: test.expectedCipher, FailureGeneration: 0, + if got, want := registry.BlacklistedCredentials(), []state.CredentialRef{{ + ID: 7, GroupID: 1, Version: 1, IdentityGeneration: 7, + Fingerprint: "test-7", EncryptedValue: test.expectedCipher, FailureGeneration: 0, }}; !reflect.DeepEqual(got, want) { t.Fatalf("blacklisted keys after stale recovery = %#v, want %#v", got, want) } - if got := registry.CollectCandidates([]uint{1}, nil, time.Time{}); len(got) != 0 { + if got := registry.CollectCredentialCandidates([]uint{1}, nil, time.Time{}); len(got) != 0 { t.Fatalf("candidates after stale recovery = %#v, want none", got) } }) @@ -857,12 +897,10 @@ func TestValidationWorkerFailureLogUsesSafeStructuredFields(t *testing.T) { }) worker := newValidationWorkerForTest( - validationSnapshot(map[uint]state.GroupView{1: { - UpstreamURL: "https://sensitive.example.com/path", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - ValidationModel: "model", - }}), - []state.KeyRef{{ID: 7, GroupID: 1, EncryptedValue: "cipher-secret"}}, + validationSnapshot(map[uint]state.GroupView{1: validationGroup( + []protocol.Protocol{protocol.OpenAICompletions}, "model", nil, + )}), + []state.CredentialRef{{ID: 7, GroupID: 1, EncryptedValue: "cipher-secret"}}, &validationProbeRecorder{}, ) worker.decryptor = validationDecryptor{errors: map[string]error{"cipher-secret": errors.New("plain-secret underlying failure")}} @@ -968,8 +1006,8 @@ func TestValidationWorkerDoesNotProbeQueuedJobAfterCancellation(t *testing.T) { nil, probes, ) - jobs := make(chan state.KeyRef, 1) - jobs <- state.KeyRef{ID: 7, GroupID: 1, EncryptedValue: "key-7"} + jobs := make(chan state.CredentialRef, 1) + jobs <- state.CredentialRef{ID: 7, GroupID: 1, EncryptedValue: "key-7"} close(jobs) ctx := &validationCancelAfterJobContext{ done: make(chan struct{}), @@ -1073,20 +1111,20 @@ func (recorder *validationEventRecorder) events() []string { type validationRegistryRecorder struct { mu sync.Mutex - refs []state.KeyRef + refs []state.CredentialRef blacklistedRead int recoveryOK bool recorder *validationEventRecorder } -func (registry *validationRegistryRecorder) BlacklistedKeys() []state.KeyRef { +func (registry *validationRegistryRecorder) BlacklistedCredentials() []state.CredentialRef { registry.mu.Lock() defer registry.mu.Unlock() registry.blacklistedRead++ - return append([]state.KeyRef(nil), registry.refs...) + return append([]state.CredentialRef(nil), registry.refs...) } -func (registry *validationRegistryRecorder) RecoverIfMatch(ref state.KeyRef, weight int) bool { +func (registry *validationRegistryRecorder) RecoverIfMatch(ref state.CredentialRef, weight int) bool { registry.mu.Lock() recoveryOK := registry.recoveryOK registry.mu.Unlock() @@ -1124,7 +1162,16 @@ func (decryptor validationDecryptor) Decrypt(value string) (string, error) { if err := decryptor.errors[value]; err != nil { return "", err } - return "plain-" + value, nil + encoded, err := json.Marshal(map[string]string{"api_key": "plain-" + value}) + return string(encoded), err +} + +type fixedValidationDecryptor struct { + plaintext string +} + +func (decryptor fixedValidationDecryptor) Decrypt(string) (string, error) { + return decryptor.plaintext, nil } type validationProbeCall struct { @@ -1173,37 +1220,50 @@ func (recorder *validationProbeRecorder) invoke(ctx context.Context, p protocol. return recorder.errByKey[apiKey] } -type validationTestDialect struct { - protocol protocol.Protocol - probes *validationProbeRecorder +type validationTestExecutor struct { + probes *validationProbeRecorder } -func (dialect *validationTestDialect) Protocol() protocol.Protocol { - return dialect.protocol +func (*validationTestExecutor) Capabilities() execution.CapabilitySet { + return execution.CapabilitySet{} } -func (*validationTestDialect) InspectRequest( - *dialect.ParsedRequest, -) (dialect.RequestMetadata, error) { - return dialect.RequestMetadata{}, nil +func (executor *validationTestExecutor) Execute( + ctx context.Context, + spec execution.AttemptSpec, +) execution.AttemptResult { + var credential struct { + APIKey string `json:"api_key"` + } + if err := json.Unmarshal(spec.Credential.Data(), &credential); err != nil { + return validationExecutionFailure(execution.ErrorKindInvalidRequest) + } + if err := executor.probes.invoke(ctx, spec.ClientProtocol, spec.UpstreamModel, credential.APIKey); err != nil { + kind := execution.ErrorKindProvider + if errors.Is(err, context.Canceled) { + kind = execution.ErrorKindCanceled + } + return validationExecutionFailure(kind) + } + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusOK, Header: http.Header{}, + } } -func (*validationTestDialect) BuildUpstreamURL(string, *dialect.ParsedRequest) (string, error) { - return "", nil +func (*validationTestExecutor) ExecuteStream( + context.Context, + execution.AttemptSpec, + execution.StreamSink, +) execution.StreamResult { + panic("unexpected stream execution") } -func (*validationTestDialect) InjectCredential(http.Header, string) {} - -func (*validationTestDialect) ListModels(context.Context, string, string, state.HeaderRules) ([]string, error) { - return nil, nil -} - -func (dialect *validationTestDialect) Probe(ctx context.Context, _ string, apiKey string, _ state.HeaderRules, model string) error { - return dialect.probes.invoke(ctx, dialect.protocol, model, apiKey) -} - -func (*validationTestDialect) ClassifyStatus(int, []byte) health.FailureCategory { - return health.FailureCategoryAmbiguous +func validationExecutionFailure(kind execution.ErrorKind) execution.AttemptResult { + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + Error: &execution.ErrorEvidence{Kind: kind, Summary: "test probe failed"}, + } } type validationTestWorker struct { @@ -1211,17 +1271,9 @@ type validationTestWorker struct { recorder *validationEventRecorder } -func newValidationWorkerForTest(snapshot *state.ConfigSnapshot, refs []state.KeyRef, probes *validationProbeRecorder) *validationTestWorker { +func newValidationWorkerForTest(snapshot *state.ConfigSnapshot, refs []state.CredentialRef, probes *validationProbeRecorder) *validationTestWorker { recorder := &validationEventRecorder{} registry := &validationRegistryRecorder{refs: refs, recoveryOK: true, recorder: recorder} - dialects := dialect.Set{} - for _, p := range []protocol.Protocol{ - protocol.OpenAICompletions, - protocol.OpenAIResponses, - protocol.Anthropic, - } { - dialects[p] = &validationTestDialect{protocol: p, probes: probes} - } return &validationTestWorker{ validationWorker: &validationWorker{ snapshots: &validationSnapshotRecorder{snapshot: snapshot}, @@ -1229,13 +1281,14 @@ func newValidationWorkerForTest(snapshot *state.ConfigSnapshot, refs []state.Key stats: &validationStatsRecorder{recorder: recorder}, mutations: health.NewMutationCoordinator(), decryptor: validationDecryptor{}, - dialects: dialects, + channels: channel.NewRegistry(), + executor: &validationTestExecutor{probes: probes}, }, recorder: recorder, } } -func newRealRegistryValidationWorker(registry *state.KeyRegistry, probes *validationProbeRecorder) *validationWorker { +func newRealRegistryValidationWorker(registry *state.CredentialRegistry, probes *validationProbeRecorder) *validationWorker { return &validationWorker{ snapshots: &validationSnapshotRecorder{snapshot: validationSnapshot(map[uint]state.GroupView{ 1: validationGroup([]protocol.Protocol{protocol.OpenAICompletions}, "model", nil), @@ -1244,7 +1297,8 @@ func newRealRegistryValidationWorker(registry *state.KeyRegistry, probes *valida stats: health.NewStatsStore(), mutations: health.NewMutationCoordinator(), decryptor: validationDecryptor{}, - dialects: dialect.Set{protocol.OpenAICompletions: &validationTestDialect{protocol: protocol.OpenAICompletions, probes: probes}}, + channels: channel.NewRegistry(), + executor: &validationTestExecutor{probes: probes}, } } @@ -1253,34 +1307,55 @@ func validationSnapshot(groups map[uint]state.GroupView) *state.ConfigSnapshot { } func validationGroup(protocols []protocol.Protocol, validationModel string, models []state.ModelConfig) state.GroupView { - return state.GroupView{ - UpstreamURL: "https://upstream.example.com", - Protocols: protocols, + group := state.GroupView{ ValidationModel: validationModel, Models: models, } + if len(protocols) == 0 { + return group + } + channelID := channel.OpenAI + switch protocols[0] { + case protocol.Anthropic: + channelID = channel.Anthropic + case protocol.Gemini: + channelID = channel.Gemini + } + setValidationChannel(&group, channelID, json.RawMessage(`{}`)) + return group +} + +func setValidationChannel(group *state.GroupView, channelID channel.ID, params json.RawMessage) { + resolved, err := channel.NewRegistry().Resolve(channelID, params) + if err != nil { + panic(err) + } + group.ChannelID = channelID + group.Params = append(json.RawMessage(nil), params...) + group.ResolvedTarget = resolved } func validationSignatureGroup() state.GroupView { - return state.GroupView{ - ID: 1, - UpstreamURL: "https://upstream.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - ValidationModel: "model-a", - Models: []state.ModelConfig{{ID: "model-a"}}, - HeaderRules: state.HeaderRules{ - Set: map[string]string{ - "X-Alpha": "first", - "X-Zeta": "last", - }, - Remove: []string{"X-Old"}, + group := validationGroup( + []protocol.Protocol{protocol.OpenAICompletions}, + "model-a", + []state.ModelConfig{{ID: "model-a"}}, + ) + group.ID = 1 + group.HeaderRules = state.HeaderRules{ + Set: map[string]string{ + "X-Alpha": "first", + "X-Zeta": "last", }, + Remove: []string{"X-Old"}, } + return group } func cloneValidationGroup(group state.GroupView) state.GroupView { cloned := group - cloned.Protocols = append([]protocol.Protocol(nil), group.Protocols...) + cloned.Params = append(json.RawMessage(nil), group.Params...) + cloned.ResolvedTarget.TargetConfig = append(json.RawMessage(nil), group.ResolvedTarget.TargetConfig...) cloned.Models = append([]state.ModelConfig(nil), group.Models...) cloned.HeaderRules.Set = make(map[string]string, len(group.HeaderRules.Set)) for name, value := range group.HeaderRules.Set { @@ -1291,22 +1366,22 @@ func cloneValidationGroup(group state.GroupView) state.GroupView { } func validationManagerCompileInput(upstreamURL string) state.CompileInput { - return state.CompileInput{Groups: []state.GroupConfig{{ + return state.CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ ID: 1, Name: "group", - UpstreamURL: upstreamURL, + ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"` + upstreamURL + `"}`), ValidationModel: "model-a", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: []state.ModelConfig{{ID: "model-a"}}, Enabled: true, }}} } -func validationRefs(count int) []state.KeyRef { - refs := make([]state.KeyRef, count) +func validationRefs(count int) []state.CredentialRef { + refs := make([]state.CredentialRef, count) for index := range refs { keyID := uint(index + 1) - refs[index] = state.KeyRef{ID: keyID, GroupID: 1, EncryptedValue: fmt.Sprintf("key-%d", keyID)} + refs[index] = state.CredentialRef{ID: keyID, GroupID: 1, EncryptedValue: fmt.Sprintf("key-%d", keyID)} } return refs } diff --git a/internal/control/wire_v3_contract_test.go b/internal/control/wire_v3_contract_test.go index 0fec592b..8abf2907 100644 --- a/internal/control/wire_v3_contract_test.go +++ b/internal/control/wire_v3_contract_test.go @@ -4,6 +4,8 @@ import ( "encoding/json" "strings" "testing" + + "gpt-load/internal/scheduler" ) func TestAccessKeyMillisWireOmitsLegacyTimestampKeys(t *testing.T) { @@ -38,7 +40,7 @@ func TestHealthMillisWireUsesNullableEpochMilliseconds(t *testing.T) { t, runtimeHealthResponse{ ObservedAtMS: 1_784_894_400_000, - CooldownKeys: []healthProblemKeyResponse{{ + CooldownCredentials: []healthProblemCredentialResponse{{ CooldownUntilMS: &cooldownUntilMS, Recovery: healthRecoveryResponse{ Automatic: true, @@ -52,7 +54,7 @@ func TestHealthMillisWireUsesNullableEpochMilliseconds(t *testing.T) { }, []string{ "observed_at_ms", - "cooldown_keys", + "cooldown_credentials", "request_log", }, ) @@ -65,7 +67,7 @@ func TestInspectMillisWireUsesEpochMilliseconds(t *testing.T) { routeInspectResponse{ ObservedAtMS: 1_784_894_400_000, Groups: []routeInspectGroupResponse{{ - Keys: []routeInspectKeyResponse{{ + Credentials: []routeInspectCredentialResponse{{ CooldownUntilMS: &cooldownUntilMS, }}, }}, @@ -74,6 +76,58 @@ func TestInspectMillisWireUsesEpochMilliseconds(t *testing.T) { ) } +func TestHealthAndRouteInspectionUseCredentialWireNames(t *testing.T) { + healthBody, err := json.Marshal(runtimeHealthResponse{ + Counts: healthCountsResponse{Credentials: 1, Available: 1}, + CooldownCredentials: []healthProblemCredentialResponse{{CredentialID: 7}}, + BlacklistedCredentials: []healthProblemCredentialResponse{}, + }) + if err != nil { + t.Fatalf("json.Marshal(health) error = %v", err) + } + for _, token := range []string{ + `"counts":{"credentials":1,"available":1,"cooldown":0,"blacklisted":0}`, + `"cooldown_credentials":[{"credential_id":7`, + `"blacklisted_credentials":[]`, + } { + if !strings.Contains(string(healthBody), token) { + t.Fatalf("health wire missing %s: %s", token, healthBody) + } + } + for _, legacy := range []string{ + `"total"`, `"disabled"`, `"cooldown_keys"`, `"blacklisted_keys"`, `"key_id"`, + } { + if strings.Contains(string(healthBody), legacy) { + t.Fatalf("health wire exposes legacy field %s: %s", legacy, healthBody) + } + } + + inspectBody, err := json.Marshal(routeInspectResponse{ + Groups: []routeInspectGroupResponse{{ + Credentials: []routeInspectCredentialResponse{{ + CredentialID: 9, + ReasonCode: optionalReason(scheduler.ReasonCredentialCooldown), + }}, + }}, + }) + if err != nil { + t.Fatalf("json.Marshal(route inspection) error = %v", err) + } + for _, token := range []string{ + `"credentials":[{"credential_id":9`, + `"reason_code":"credential_cooldown"`, + } { + if !strings.Contains(string(inspectBody), token) { + t.Fatalf("route inspection wire missing %s: %s", token, inspectBody) + } + } + for _, legacy := range []string{`"keys"`, `"key_id"`, `"key_cooldown"`} { + if strings.Contains(string(inspectBody), legacy) { + t.Fatalf("route inspection wire exposes legacy field %s: %s", legacy, inspectBody) + } + } +} + func TestRequestLogMillisCostWireUsesCursorV2Fields(t *testing.T) { assertManagementWireObject( t, diff --git a/internal/dialect/anthropic.go b/internal/dialect/anthropic.go index bee36c9d..e891ef89 100644 --- a/internal/dialect/anthropic.go +++ b/internal/dialect/anthropic.go @@ -1,192 +1,23 @@ package dialect import ( - "context" "fmt" - "net/http" - "net/url" - "strings" - "gpt-load/internal/health" - platformheader "gpt-load/internal/platform/httpheader" "gpt-load/internal/protocol" - "gpt-load/internal/state" ) -const ( - anthropicMessagesPath = "/v1/messages" - anthropicMessagesResourcePath = "/messages" - anthropicModelsResourcePath = "/models" - anthropicDefaultVersion = "2023-06-01" -) - -var anthropicFailureMarkers = failureMarkers{ - rateLimited: []string{"rate_limit_error", "rate limit"}, - modelUnavailable: []string{"model_not_found", "model not found", "model_not_supported", "model not supported", "no access to model"}, - invalidKey: []string{"authentication_error", "permission_error", "invalid x-api-key", "api key disabled", "api key banned"}, - upstreamHost: []string{"overloaded_error"}, -} - -type Anthropic struct { - client *http.Client -} - -type anthropicModelPage struct { - Data []struct { - ID string `json:"id"` - } `json:"data"` - HasMore bool `json:"has_more"` - LastID string `json:"last_id"` -} +type Anthropic struct{} var _ Dialect = (*Anthropic)(nil) -func NewAnthropic(client *http.Client) *Anthropic { - return &Anthropic{client: client} +func NewAnthropic() *Anthropic { + return &Anthropic{} } -func (d *Anthropic) Protocol() protocol.Protocol { +func (*Anthropic) Protocol() protocol.Protocol { return protocol.Anthropic } -func (d *Anthropic) InjectCredential(headers http.Header, apiKey string) { - if headers == nil { - return - } - headers.Set("X-Api-Key", apiKey) - if strings.TrimSpace(headers.Get("Anthropic-Version")) == "" { - headers.Set("Anthropic-Version", anthropicDefaultVersion) - } -} - -func (d *Anthropic) CredentialHeaderNames() []string { - return []string{"X-Api-Key"} -} - -func (d *Anthropic) BuildUpstreamURL(base string, req *ParsedRequest) (string, error) { - if req == nil { - return "", fmt.Errorf("parsed request is required") - } - if req.Path != anthropicMessagesPath { - return "", fmt.Errorf("invalid Anthropic Messages request path") - } - return resolveUpstreamURL(base, anthropicMessagesResourcePath, req.RawQuery) -} - -func (d *Anthropic) ListModels( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, -) ([]string, error) { - requestURL, err := resolveUpstreamURL(baseURL, anthropicModelsResourcePath, "") - if err != nil { - return nil, fmt.Errorf("build Anthropic model-list URL: %w", err) - } - parsed, err := url.Parse(requestURL) - if err != nil { - return nil, fmt.Errorf("parse Anthropic model-list URL: %w", err) - } - - collector := newModelListCollector() - seenCursors := make(map[string]struct{}) - afterID := "" - for pageNumber := 1; pageNumber <= maxModelListPages; pageNumber++ { - page, err := d.listModelsPage(ctx, parsed, apiKey, rules, afterID) - if err != nil { - return nil, err - } - pageModels := make([]string, 0, len(page.Data)) - for _, item := range page.Data { - pageModels = append(pageModels, item.ID) - } - if err := collector.Add(pageModels); err != nil { - return nil, err - } - if !page.HasMore { - return collector.Result(), nil - } - if strings.TrimSpace(page.LastID) == "" { - return nil, fmt.Errorf("Anthropic model-list cursor is empty") - } - if _, repeated := seenCursors[page.LastID]; repeated { - return nil, fmt.Errorf("Anthropic model-list cursor repeated") - } - if pageNumber == maxModelListPages || collector.Full() { - return nil, fmt.Errorf("Anthropic model-list pagination limit exceeded") - } - seenCursors[page.LastID] = struct{}{} - afterID = page.LastID - } - return nil, fmt.Errorf("Anthropic model-list pagination limit exceeded") -} - -func (d *Anthropic) listModelsPage( - ctx context.Context, - endpoint *url.URL, - apiKey string, - rules state.HeaderRules, - afterID string, -) (anthropicModelPage, error) { - pageEndpoint := *endpoint - query := pageEndpoint.Query() - query.Del("limit") - query.Set("limit", "1000") - query.Del("after_id") - if afterID != "" { - query.Set("after_id", afterID) - } - pageEndpoint.RawQuery = query.Encode() - - request, err := http.NewRequestWithContext(ctx, http.MethodGet, pageEndpoint.String(), nil) - if err != nil { - return anthropicModelPage{}, fmt.Errorf("create Anthropic model-list request: %w", err) - } - ApplyCredential(d, request.Header, apiKey, rules) - platformheader.NormalizeUpstreamRequestRepresentation(request, 0) - - response, err := d.client.Do(request) - if err != nil { - if contextErr := ctx.Err(); contextErr != nil { - return anthropicModelPage{}, fmt.Errorf("request Anthropic model list: %w", contextErr) - } - return anthropicModelPage{}, fmt.Errorf("request Anthropic model list failed") - } - defer response.Body.Close() - if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { - return anthropicModelPage{}, fmt.Errorf("request Anthropic model list: upstream status %d", response.StatusCode) - } - - var page anthropicModelPage - if err := decodeModelListPage(response, &page); err != nil { - return anthropicModelPage{}, fmt.Errorf("decode Anthropic model list: %w", err) - } - return page, nil -} - -func (d *Anthropic) Probe( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, - validationModel string, -) error { - if err := validateProbeModel(validationModel); err != nil { - return err - } - requestURL, err := d.BuildUpstreamURL(baseURL, &ParsedRequest{Method: http.MethodPost, Path: anthropicMessagesPath}) - if err != nil { - return fmt.Errorf("build %s probe URL failed", d.Protocol()) - } - return executeProbe(ctx, d.client, d, requestURL, apiKey, rules, struct { - Model string `json:"model"` - Messages []probeMessage `json:"messages"` - MaxTokens int `json:"max_tokens"` - }{ - Model: validationModel, - Messages: []probeMessage{{Role: "user", Content: "ping"}}, - MaxTokens: 1, - }) -} - func (d *Anthropic) InspectRequest(req *ParsedRequest) (RequestMetadata, error) { if req == nil { return RequestMetadata{}, fmt.Errorf("parsed request is required") @@ -198,13 +29,10 @@ func (d *Anthropic) InspectRequest(req *ParsedRequest) (RequestMetadata, error) } metadata.ObserveUsage = true metadata.Reasoning = inspectAnthropicReasoning(req.Body) + metadata.Operation, metadata.RequiredFeatures = chatExecutionMetadata( + req.Body, + metadata.Stream, + metadata.Reasoning, + ) return metadata, nil } - -func (d *Anthropic) ClassifyStatus(status int, body []byte) health.FailureCategory { - return classifyStatusWithMarkers(status, body, anthropicFailureMarkers) -} - -func (d *Anthropic) ClassifyProviderError(body []byte) health.FailureCategory { - return classifyProviderErrorWithMarkers(body, anthropicFailureMarkers) -} diff --git a/internal/dialect/anthropic_test.go b/internal/dialect/anthropic_test.go index 38155bf4..fd067840 100644 --- a/internal/dialect/anthropic_test.go +++ b/internal/dialect/anthropic_test.go @@ -1,25 +1,19 @@ package dialect import ( - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "net/http/httptest" - "strings" - "sync/atomic" "testing" - "time" - "gpt-load/internal/health" "gpt-load/internal/protocol" - "gpt-load/internal/state" ) +func TestAnthropicProtocol(t *testing.T) { + if got := NewAnthropic().Protocol(); got != protocol.Anthropic { + t.Fatalf("Protocol() = %q, want %q", got, protocol.Anthropic) + } +} + func TestAnthropicInspectRequest(t *testing.T) { - value := NewAnthropic(http.DefaultClient) + selected := NewAnthropic() tests := []struct { name string request *ParsedRequest @@ -37,536 +31,17 @@ func TestAnthropicInspectRequest(t *testing.T) { } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - metadata, err := value.InspectRequest(test.request) + metadata, err := selected.InspectRequest(test.request) if test.wantErr { if err == nil { - t.Fatalf("InspectRequest() = (%#v, nil), want error", metadata) + t.Fatalf("InspectRequest() = %#v, nil", metadata) } return } - if err != nil { - t.Fatalf("InspectRequest() error = %v", err) - } - if metadata.Model == nil || *metadata.Model != test.wantModel || + if err != nil || metadata.Model == nil || *metadata.Model != test.wantModel || metadata.Stream != test.wantStream || !metadata.ObserveUsage { - t.Fatalf("InspectRequest() = %#v, want model=%q stream=%t observe=true", metadata, test.wantModel, test.wantStream) + t.Fatalf("InspectRequest() = %#v, %v", metadata, err) } }) } } - -func TestAnthropicCredentialAndProtocol(t *testing.T) { - value := NewAnthropic(http.DefaultClient) - if got := value.Protocol(); got != protocol.Anthropic { - t.Fatalf("Protocol() = %q, want %q", got, protocol.Anthropic) - } - - for _, test := range []struct { - name string - headers http.Header - version string - }{ - {name: "missing version", headers: make(http.Header), version: anthropicDefaultVersion}, - {name: "blank version", headers: http.Header{"Anthropic-Version": {" "}}, version: anthropicDefaultVersion}, - {name: "existing version", headers: http.Header{"Anthropic-Version": {" 2024-01-01 "}}, version: " 2024-01-01 "}, - } { - t.Run(test.name, func(t *testing.T) { - value.InjectCredential(test.headers, "upstream-secret") - if got := test.headers.Get("X-Api-Key"); got != "upstream-secret" { - t.Fatalf("X-Api-Key = %q", got) - } - if got := test.headers.Get("Anthropic-Version"); got != test.version { - t.Fatalf("Anthropic-Version = %q, want %q", got, test.version) - } - }) - } - - headers := make(http.Header) - ApplyCredential(value, headers, "upstream-secret", state.HeaderRules{Set: map[string]string{ - "X-Api-Key": "override-${API_KEY}", - "Anthropic-Version": "2025-01-01", - }}) - if got := headers.Get("X-Api-Key"); got != "override-upstream-secret" { - t.Fatalf("overridden X-Api-Key = %q", got) - } - if got := headers.Get("Anthropic-Version"); got != "2025-01-01" { - t.Fatalf("overridden Anthropic-Version = %q", got) - } - - ApplyCredential(value, headers, "upstream-secret", state.HeaderRules{ - Set: map[string]string{ - "X-Api-Key": "must-be-removed", - "Anthropic-Version": "must-be-removed", - }, - Remove: []string{"x-api-key", "anthropic-version"}, - }) - if headers.Get("X-Api-Key") != "" || headers.Get("Anthropic-Version") != "" { - t.Fatalf("remove did not win: %#v", headers) - } - - value.InjectCredential(nil, "upstream-secret") - ApplyCredential(value, nil, "upstream-secret", state.HeaderRules{}) -} - -func TestAnthropicBuildUpstreamURL(t *testing.T) { - value := NewAnthropic(http.DefaultClient) - tests := []struct { - name string - base string - request *ParsedRequest - want string - wantErr bool - }{ - {name: "complete API prefix", base: "https://api.example.com/v1", request: &ParsedRequest{Path: "/v1/messages", RawQuery: "trace=true"}, want: "https://api.example.com/v1/messages?trace=true"}, - {name: "base path and queries", base: "https://api.example.com/compatible/v1/?tenant=one", request: &ParsedRequest{Path: "/v1/messages", RawQuery: "trace=true"}, want: "https://api.example.com/compatible/v1/messages?tenant=one&trace=true"}, - {name: "nil request", base: "https://api.example.com", wantErr: true}, - {name: "relative request path", base: "https://api.example.com", request: &ParsedRequest{Path: "v1/messages"}, wantErr: true}, - {name: "relative base", base: "api.example.com/v1", request: &ParsedRequest{Path: "/v1/messages"}, wantErr: true}, - {name: "unsupported base", base: "ftp://api.example.com/v1", request: &ParsedRequest{Path: "/v1/messages"}, wantErr: true}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - got, err := value.BuildUpstreamURL(test.base, test.request) - if test.wantErr { - if err == nil { - t.Fatalf("BuildUpstreamURL() = %q, want error", got) - } - return - } - if err != nil || got != test.want { - t.Fatalf("BuildUpstreamURL() = %q, %v, want %q", got, err, test.want) - } - }) - } -} - -func TestAnthropicClassifyStatus(t *testing.T) { - value := NewAnthropic(http.DefaultClient) - tests := []struct { - name string - status int - body string - want health.FailureCategory - }{ - {name: "success", status: http.StatusOK, body: `{"type":"authentication_error"}`, want: health.FailureCategoryOK}, - {name: "unauthorized", status: http.StatusUnauthorized, want: health.FailureCategoryInvalidKey}, - {name: "forbidden", status: http.StatusForbidden, want: health.FailureCategoryInvalidKey}, - {name: "not found", status: http.StatusNotFound, want: health.FailureCategoryModelUnavailable}, - {name: "rate limited", status: http.StatusTooManyRequests, want: health.FailureCategoryRateLimited}, - {name: "server error", status: http.StatusInternalServerError, want: health.FailureCategoryUpstreamHostError}, - {name: "provider 529 is host error", status: 529, want: health.FailureCategoryUpstreamHostError}, - {name: "authentication marker", status: http.StatusBadRequest, body: `{"type":"authentication_error"}`, want: health.FailureCategoryInvalidKey}, - {name: "permission marker", status: http.StatusBadRequest, body: `{"type":"permission_error"}`, want: health.FailureCategoryInvalidKey}, - {name: "rate marker", status: http.StatusBadRequest, body: `{"type":"rate_limit_error"}`, want: health.FailureCategoryRateLimited}, - {name: "overloaded marker", status: http.StatusBadRequest, body: `{"type":"overloaded_error"}`, want: health.FailureCategoryUpstreamHostError}, - {name: "model not found code", status: http.StatusBadRequest, body: `{"type":"model_not_found"}`, want: health.FailureCategoryModelUnavailable}, - {name: "model marker precedes key status", status: http.StatusForbidden, body: `{"type":"model_not_found"}`, want: health.FailureCategoryModelUnavailable}, - {name: "model not found message", status: http.StatusBadRequest, body: `{"message":"model not found"}`, want: health.FailureCategoryModelUnavailable}, - {name: "model not supported code", status: http.StatusBadRequest, body: `{"type":"model_not_supported"}`, want: health.FailureCategoryModelUnavailable}, - {name: "model not supported message", status: http.StatusBadRequest, body: `{"message":"model not supported"}`, want: health.FailureCategoryModelUnavailable}, - {name: "no model access", status: http.StatusBadRequest, body: `{"message":"no access to model"}`, want: health.FailureCategoryModelUnavailable}, - {name: "invalid request", status: http.StatusBadRequest, body: `{"type":"invalid_request_error"}`, want: health.FailureCategoryClientError}, - {name: "unsupported feature", status: http.StatusBadRequest, body: `{"message":"function calling is not supported with this model"}`, want: health.FailureCategoryClientError}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - if got := value.ClassifyStatus(test.status, []byte(test.body)); got != test.want { - t.Fatalf("ClassifyStatus() = %d, want %d", got, test.want) - } - }) - } -} - -func TestAnthropicListModelsPaginatesAndPreservesInputs(t *testing.T) { - var requestCount atomic.Int64 - const cursor = " cursor/+= ? " - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - pageNumber := requestCount.Add(1) - if request.Method != http.MethodGet || request.URL.Path != "/prefix/v1/models" { - t.Errorf("request = %s %s", request.Method, request.URL.Path) - } - if got := request.URL.Query()["limit"]; len(got) != 1 || got[0] != "1000" { - t.Errorf("limit = %#v, want one 1000", got) - } - if got := request.URL.Query().Get("tenant"); got != "one" { - t.Errorf("tenant = %q, want one", got) - } - if got := request.URL.Query().Get("fixed"); got != "preserved" { - t.Errorf("fixed = %q, want preserved", got) - } - if got := request.Header.Get("X-Api-Key"); got != "upstream-secret" { - t.Errorf("X-Api-Key = %q", got) - } - if got := request.Header.Get("Anthropic-Version"); got != anthropicDefaultVersion { - t.Errorf("Anthropic-Version = %q", got) - } - if got := request.Header.Get("X-Discovery-Rule"); got != "applied-rule" { - t.Errorf("X-Discovery-Rule = %q", got) - } - switch pageNumber { - case 1: - if _, exists := request.URL.Query()["after_id"]; exists { - t.Errorf("page 1 retained stale after_id: %q", request.URL.Query().Get("after_id")) - } - _, _ = writer.Write([]byte(`{"data":[{"id":"claude-a"},{"id":"shared"},{"id":"claude-a"}],"has_more":true,"last_id":"` + cursor + `"}`)) - case 2: - if got := request.URL.Query().Get("after_id"); got != cursor { - t.Errorf("page 2 after_id = %q, want byte-for-byte %q", got, cursor) - } - _, _ = writer.Write([]byte(`{"data":[{"id":"shared"},{"id":"claude-b"}],"has_more":false}`)) - default: - t.Errorf("unexpected page %d", pageNumber) - writer.WriteHeader(http.StatusInternalServerError) - } - })) - defer server.Close() - - models, err := NewAnthropic(server.Client()).ListModels( - context.Background(), - server.URL+"/prefix/v1?tenant=one&fixed=preserved&limit=1&limit=2&after_id=stale", - "upstream-secret", - state.HeaderRules{Set: map[string]string{"X-Discovery-Rule": "applied-rule"}}, - ) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - if strings.Join(models, ",") != "claude-a,shared,claude-b" { - t.Fatalf("models = %#v", models) - } - if got := requestCount.Load(); got != 2 { - t.Fatalf("request count = %d, want 2", got) - } -} - -func TestAnthropicListModelsIdentityRepresentation(t *testing.T) { - client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { - assertOutboundRequestRepresentation(t, request, 0, map[string]string{ - "X-Api-Key": "override-secret", - "Anthropic-Version": "2025-01-01", - "X-Business": "preserved", - }) - return anthropicTestResponse(http.StatusOK, `{"data":[]}`), nil - })} - models, err := NewAnthropic(client).ListModels( - context.Background(), - "https://api.example.com", - "secret", - legacyRepresentationHeaderRules(map[string]string{ - "X-Api-Key": "override-${API_KEY}", - "Anthropic-Version": "2025-01-01", - }), - ) - if err != nil || models == nil { - t.Fatalf("ListModels() = %#v, %v", models, err) - } -} - -func TestAnthropicListModelsRejectsBlankOrRepeatedCursor(t *testing.T) { - t.Run("blank cursor", func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - requestCount.Add(1) - return anthropicTestResponse(http.StatusOK, `{"data":[{"id":"claude-a"}],"has_more":true,"last_id":" \t "}`), nil - })} - models, err := NewAnthropic(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and error", models, err) - } - if got := requestCount.Load(); got != 1 { - t.Fatalf("request count = %d, want 1", got) - } - }) - - t.Run("repeated cursor", func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - requestCount.Add(1) - return anthropicTestResponse(http.StatusOK, `{"data":[{"id":"claude-a"}],"has_more":true,"last_id":"same-cursor"}`), nil - })} - models, err := NewAnthropic(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and error", models, err) - } - if got := requestCount.Load(); got != 2 { - t.Fatalf("request count = %d, want 2", got) - } - }) -} - -func TestAnthropicListModelsRejectsPageAndUniqueModelLimits(t *testing.T) { - t.Run("page 100 may end", func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - pageNumber := requestCount.Add(1) - return anthropicTestResponse(http.StatusOK, fmt.Sprintf( - `{"data":[{"id":"model-%d"}],"has_more":%t,"last_id":"cursor-%d"}`, - pageNumber, pageNumber < maxModelListPages, pageNumber, - )), nil - })} - models, err := NewAnthropic(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err != nil || len(models) != maxModelListPages { - t.Fatalf("ListModels() = %d models, %v, want %d models", len(models), err, maxModelListPages) - } - if got := requestCount.Load(); got != maxModelListPages { - t.Fatalf("request count = %d, want %d", got, maxModelListPages) - } - }) - - t.Run("page 100 may not continue", func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - pageNumber := requestCount.Add(1) - return anthropicTestResponse(http.StatusOK, fmt.Sprintf( - `{"data":[{"id":"model-%d"}],"has_more":true,"last_id":"cursor-%d"}`, - pageNumber, pageNumber, - )), nil - })} - models, err := NewAnthropic(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and error", models, err) - } - if got := requestCount.Load(); got != maxModelListPages { - t.Fatalf("request count = %d, want %d", got, maxModelListPages) - } - }) - - for _, test := range []struct { - name string - modelCount int - hasMore bool - wantErr bool - }{ - {name: "exact unique maximum may end", modelCount: maxUniqueModelListEntries}, - {name: "exact unique maximum may not continue", modelCount: maxUniqueModelListEntries, hasMore: true, wantErr: true}, - {name: "more than unique maximum is rejected", modelCount: maxUniqueModelListEntries + 1, wantErr: true}, - } { - t.Run(test.name, func(t *testing.T) { - body := mustAnthropicTestPayload(t, test.modelCount, test.hasMore, "next-cursor") - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - requestCount.Add(1) - return anthropicTestResponse(http.StatusOK, body), nil - })} - models, err := NewAnthropic(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if test.wantErr { - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and error", models, err) - } - } else if err != nil || len(models) != test.modelCount { - t.Fatalf("ListModels() = %d models, %v, want %d models", len(models), err, test.modelCount) - } - if got := requestCount.Load(); got != 1 { - t.Fatalf("request count = %d, want 1", got) - } - }) - } -} - -func TestAnthropicListModelsReturnsNoPartialResultAfterLaterPageFailure(t *testing.T) { - for _, test := range []struct { - name string - secondPage func() (*http.Response, error) - }{ - {name: "transport error", secondPage: func() (*http.Response, error) { - return nil, errors.New("later transport failure") - }}, - {name: "non-2xx", secondPage: func() (*http.Response, error) { - return anthropicTestResponse(http.StatusBadGateway, `{"error":"later failure"}`), nil - }}, - {name: "malformed JSON", secondPage: func() (*http.Response, error) { - return anthropicTestResponse(http.StatusOK, `{"data":[`), nil - }}, - } { - t.Run(test.name, func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - if requestCount.Add(1) == 1 { - return anthropicTestResponse(http.StatusOK, `{"data":[{"id":"first-page-model"}],"has_more":true,"last_id":"next"}`), nil - } - return test.secondPage() - })} - models, err := NewAnthropic(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and later-page error", models, err) - } - if got := requestCount.Load(); got != 2 { - t.Fatalf("request count = %d, want 2", got) - } - }) - } -} - -func TestAnthropicListModelsClosesBodyBeforeRequestingNextPage(t *testing.T) { - firstBody := &trackedReadCloser{Reader: strings.NewReader(`{"data":[{"id":"first"}],"has_more":true,"last_id":"next"}`)} - secondBody := &trackedReadCloser{Reader: strings.NewReader(`{"data":[{"id":"second"}],"has_more":false}`)} - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - switch requestCount.Add(1) { - case 1: - return &http.Response{StatusCode: http.StatusOK, Body: firstBody, Header: make(http.Header)}, nil - case 2: - if !firstBody.closed.Load() { - return nil, errors.New("first page body is still open") - } - return &http.Response{StatusCode: http.StatusOK, Body: secondBody, Header: make(http.Header)}, nil - default: - return nil, errors.New("unexpected extra request") - } - })} - models, err := NewAnthropic(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err != nil || strings.Join(models, ",") != "first,second" { - t.Fatalf("ListModels() = %#v, %v, want first,second", models, err) - } - if !secondBody.closed.Load() { - t.Fatal("last page body was not closed") - } -} - -func TestAnthropicListModelsRulesEmptyAndFailures(t *testing.T) { - t.Run("rules override defaults", func(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if got := request.Header.Get("X-Api-Key"); got != "override-upstream-secret" { - t.Errorf("X-Api-Key = %q", got) - } - if got := request.Header.Get("Anthropic-Version"); got != "2025-01-01" { - t.Errorf("Anthropic-Version = %q", got) - } - _, _ = writer.Write([]byte(`{"data":[]}`)) - })) - defer server.Close() - models, err := NewAnthropic(server.Client()).ListModels( - context.Background(), server.URL, "upstream-secret", - state.HeaderRules{Set: map[string]string{ - "X-Api-Key": "override-${API_KEY}", - "Anthropic-Version": "2025-01-01", - }}, - ) - if err != nil || models == nil || len(models) != 0 { - t.Fatalf("ListModels() = %#v, %v, want non-nil empty", models, err) - } - }) - - t.Run("rules remove defaults", func(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if request.Header.Get("X-Api-Key") != "" || request.Header.Get("Anthropic-Version") != "" { - t.Errorf("removed headers remain: %#v", request.Header) - } - _, _ = writer.Write([]byte(`{"data":[]}`)) - })) - defer server.Close() - _, err := NewAnthropic(server.Client()).ListModels( - context.Background(), server.URL, "upstream-secret", - state.HeaderRules{Remove: []string{"x-api-key", "anthropic-version"}}, - ) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - }) - - t.Run("non-success does not expose body or key", func(t *testing.T) { - const secret = "distinctive-upstream-secret" - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.WriteHeader(http.StatusUnauthorized) - _, _ = writer.Write([]byte(`{"error":"` + secret + ` in body"}`)) - })) - defer server.Close() - _, err := NewAnthropic(server.Client()).ListModels(context.Background(), server.URL, secret, state.HeaderRules{}) - if err == nil || strings.Contains(err.Error(), secret) || strings.Contains(err.Error(), server.URL) { - t.Fatalf("ListModels() error = %v", err) - } - }) - - t.Run("malformed JSON is contextual without URL or key", func(t *testing.T) { - const secret = "distinctive-upstream-secret" - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - _, _ = writer.Write([]byte(`{"data":[`)) - })) - defer server.Close() - _, err := NewAnthropic(server.Client()).ListModels( - context.Background(), server.URL+"?token=query-secret", secret, state.HeaderRules{}, - ) - if err == nil || !strings.Contains(err.Error(), "decode Anthropic model list") { - t.Fatalf("ListModels() error = %v", err) - } - for _, forbidden := range []string{secret, "query-secret", server.URL} { - if strings.Contains(err.Error(), forbidden) { - t.Fatalf("error exposes %q: %v", forbidden, err) - } - } - }) - - t.Run("caller context stops request", func(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { - <-request.Context().Done() - })) - defer server.Close() - ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) - defer cancel() - _, err := NewAnthropic(server.Client()).ListModels(ctx, server.URL, "upstream-secret", state.HeaderRules{}) - if !errors.Is(err, context.DeadlineExceeded) { - t.Fatalf("ListModels() error = %v, want deadline exceeded", err) - } - }) - - t.Run("response body closes", func(t *testing.T) { - body := &trackedReadCloser{Reader: strings.NewReader(`{"data":[]}`)} - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - return &http.Response{StatusCode: http.StatusOK, Body: body, Header: make(http.Header)}, nil - })} - _, err := NewAnthropic(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - if !body.closed.Load() { - t.Fatal("response body was not closed") - } - }) -} - -type roundTripFunc func(*http.Request) (*http.Response, error) - -func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { - return function(request) -} - -type trackedReadCloser struct { - io.Reader - closed atomic.Bool -} - -func (body *trackedReadCloser) Close() error { - body.closed.Store(true) - return nil -} - -func anthropicTestResponse(statusCode int, body string) *http.Response { - return &http.Response{ - StatusCode: statusCode, - Body: io.NopCloser(strings.NewReader(body)), - Header: make(http.Header), - } -} - -func mustAnthropicTestPayload(t *testing.T, modelCount int, hasMore bool, lastID string) string { - t.Helper() - type item struct { - ID string `json:"id"` - } - payload := struct { - Data []item `json:"data"` - HasMore bool `json:"has_more"` - LastID string `json:"last_id"` - }{ - Data: make([]item, modelCount), - HasMore: hasMore, - LastID: lastID, - } - for index := range payload.Data { - payload.Data[index].ID = fmt.Sprintf("model-%06d", index) - } - encoded, err := json.Marshal(payload) - if err != nil { - t.Fatalf("marshal Anthropic test payload: %v", err) - } - return string(encoded) -} diff --git a/internal/dialect/anthropic_usage_test.go b/internal/dialect/anthropic_usage_test.go index 5252adc2..505be662 100644 --- a/internal/dialect/anthropic_usage_test.go +++ b/internal/dialect/anthropic_usage_test.go @@ -2,7 +2,6 @@ package dialect import ( "fmt" - "net/http" "strings" "testing" @@ -11,7 +10,7 @@ import ( ) func TestUsageAnthropicCanonicalFixtures(t *testing.T) { - extractor, ok := any(NewAnthropic(http.DefaultClient)).(UsageExtractor) + extractor, ok := any(NewAnthropic()).(UsageExtractor) if !ok { t.Fatal("Anthropic does not expose UsageExtractor capability") } @@ -43,7 +42,7 @@ func TestUsageAnthropicCanonicalFixtures(t *testing.T) { } func TestUsageAnthropicCacheCreationMapping(t *testing.T) { - extractor := NewAnthropic(http.DefaultClient) + extractor := NewAnthropic() tests := []struct { name string body string @@ -112,7 +111,7 @@ func TestUsageAnthropicCacheCreationMapping(t *testing.T) { } func TestUsageAnthropicInvalidCacheCreationDoesNotFallback(t *testing.T) { - extractor := NewAnthropic(http.DefaultClient) + extractor := NewAnthropic() for _, cacheCreation := range []string{`[]`, `"unsupported"`} { result, err := extractor.ExtractUsage([]byte(`{"usage":{"input_tokens":80,"output_tokens":30,"cache_creation":` + cacheCreation + `,"cache_creation_input_tokens":12}}`)) if err != nil { @@ -129,7 +128,7 @@ func TestUsageAnthropicInvalidCacheCreationDoesNotFallback(t *testing.T) { } func TestUsageAnthropicDetailReconciliationRequiresUsableComponents(t *testing.T) { - extractor := NewAnthropic(http.DefaultClient) + extractor := NewAnthropic() tests := []struct { name string value string @@ -175,7 +174,7 @@ func TestUsageAnthropicDetailReconciliationRequiresUsableComponents(t *testing.T } func TestUsageAnthropicStrictNumbersAndRequiredFields(t *testing.T) { - extractor := NewAnthropic(http.DefaultClient) + extractor := NewAnthropic() tests := []struct { name string body string @@ -262,7 +261,7 @@ func TestUsageAnthropicStrictNumbersAndRequiredFields(t *testing.T) { } func TestUsageAnthropicRequiredPresenceControlsState(t *testing.T) { - extractor := NewAnthropic(http.DefaultClient) + extractor := NewAnthropic() tests := []struct { name string body string @@ -313,7 +312,7 @@ func TestUsageAnthropicStreamCompletesOnlyAtMessageStop(t *testing.T) { start := `{"type":"message_start","message":{"usage":{"input_tokens":80,"cache_read_input_tokens":20,"cache_creation":{"ephemeral_5m_input_tokens":5,"ephemeral_1h_input_tokens":7},"output_tokens":0}}}` delta := `{"type":"message_delta","usage":{"output_tokens":30}}` - partial := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + partial := NewAnthropic().NewUsageStreamExtractor() for _, event := range []string{start, delta} { if err := partial.Observe([]byte(event)); err != nil { t.Fatal(err) @@ -326,7 +325,7 @@ func TestUsageAnthropicStreamCompletesOnlyAtMessageStop(t *testing.T) { } requireUsageDiagnostics(t, partialResult.Diagnostics) - complete := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + complete := NewAnthropic().NewUsageStreamExtractor() for _, event := range []string{start, delta, `{"type":"message_stop"}`} { if err := complete.Observe([]byte(event)); err != nil { t.Fatal(err) @@ -340,7 +339,7 @@ func TestUsageAnthropicStreamCompletesOnlyAtMessageStop(t *testing.T) { } func TestUsageAnthropicStreamMergesCumulativeDeltaFields(t *testing.T) { - stream := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + stream := NewAnthropic().NewUsageStreamExtractor() for _, event := range []string{ `{"type":"message_start","message":{"usage":{"input_tokens":80,"cache_read_input_tokens":20,"cache_creation":{"ephemeral_5m_input_tokens":5,"ephemeral_1h_input_tokens":7},"output_tokens":0}}}`, `{"type":"message_delta","usage":{"output_tokens":10}}`, @@ -360,7 +359,7 @@ func TestUsageAnthropicStreamMergesCumulativeDeltaFields(t *testing.T) { } requireUsageDiagnostics(t, result.Diagnostics) - fallback := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + fallback := NewAnthropic().NewUsageStreamExtractor() for _, event := range []string{ `{"type":"message_start","message":{"usage":{"input_tokens":80,"cache_creation_input_tokens":12}}}`, `{"type":"message_stop"}`, @@ -425,7 +424,7 @@ func TestUsageAnthropicStreamReconcilesAggregateAfterCumulativeMerge(t *testing. for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - stream := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + stream := NewAnthropic().NewUsageStreamExtractor() for _, event := range tt.events { if err := stream.Observe([]byte(event)); err != nil { t.Fatal(err) @@ -442,7 +441,7 @@ func TestUsageAnthropicStreamReconcilesAggregateAfterCumulativeMerge(t *testing. func TestUsageAnthropicStreamRejectsOutOfOrderAndRegressingCumulatives(t *testing.T) { t.Run("delta and stop without valid start", func(t *testing.T) { - stream := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + stream := NewAnthropic().NewUsageStreamExtractor() for _, event := range []string{ `{"type":"message_delta","usage":{"output_tokens":30}}`, `{"type":"message_stop"}`, @@ -459,7 +458,7 @@ func TestUsageAnthropicStreamRejectsOutOfOrderAndRegressingCumulatives(t *testin }) t.Run("last good cumulatives survive invalid events", func(t *testing.T) { - stream := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + stream := NewAnthropic().NewUsageStreamExtractor() for _, event := range []string{ `{"type":"message_start","message":{"usage":{"input_tokens":80,"output_tokens":0}}}`, `{"type":"message_delta","usage":{"output_tokens":10,"cache_read_input_tokens":4}}`, @@ -490,7 +489,7 @@ func TestUsageAnthropicStreamRejectsOutOfOrderAndRegressingCumulatives(t *testin } func TestUsageAnthropicStreamStartStopWithoutDelta(t *testing.T) { - stream := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + stream := NewAnthropic().NewUsageStreamExtractor() for _, event := range []string{ `{"type":"message_start","message":{"usage":{"input_tokens":80,"cache_read_input_tokens":20,"output_tokens":0}}}`, `{"type":"message_stop"}`, @@ -508,7 +507,7 @@ func TestUsageAnthropicStreamStartStopWithoutDelta(t *testing.T) { } func TestUsageAnthropicStreamEOFRemainsPartial(t *testing.T) { - stream := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + stream := NewAnthropic().NewUsageStreamExtractor() for _, event := range []string{ `{"type":"message_start","message":{"usage":{"input_tokens":80}}}`, `{"type":"message_delta","usage":{"output_tokens":30}}`, @@ -526,7 +525,7 @@ func TestUsageAnthropicStreamEOFRemainsPartial(t *testing.T) { } func TestUsageAnthropicServerToolDetailKeepsKnownTokenPriceComplete(t *testing.T) { - stream := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + stream := NewAnthropic().NewUsageStreamExtractor() observeUsageJSONL(t, stream, readUsageFixture(t, "anthropic", "server-tool.jsonl")) result, finalized := stream.Finalize() if !finalized || result.State != usage.StateComplete || @@ -535,7 +534,7 @@ func TestUsageAnthropicServerToolDetailKeepsKnownTokenPriceComplete(t *testing.T } requireUsageDiagnostics(t, result.Diagnostics, usage.DiagnosticUnsupportedBillableDetail) - identity := pricing.Identity{ModelID: "claude-test"} + identity := pricing.Identity{ChannelID: "anthropic", ModelID: "claude-test"} table, err := pricing.NewTable([]pricing.Rule{{ Identity: identity, Prices: pricing.Prices{ @@ -554,7 +553,7 @@ func TestUsageAnthropicServerToolDetailKeepsKnownTokenPriceComplete(t *testing.T } func TestUsageAnthropicNonStreamServerToolDetailPricing(t *testing.T) { - identity := pricing.Identity{ModelID: "claude-test"} + identity := pricing.Identity{ChannelID: "anthropic", ModelID: "claude-test"} table, err := pricing.NewTable([]pricing.Rule{{ Identity: identity, Prices: pricing.Prices{ @@ -589,7 +588,7 @@ func TestUsageAnthropicNonStreamServerToolDetailPricing(t *testing.T) { `{"usage":{"input_tokens":80,"output_tokens":30,"server_tool_use":{"web_search_requests":%d}}}`, tt.count, ) - result, extractErr := NewAnthropic(http.DefaultClient).ExtractUsage([]byte(body)) + result, extractErr := NewAnthropic().ExtractUsage([]byte(body)) if extractErr != nil { t.Fatal(extractErr) } @@ -607,7 +606,7 @@ func TestUsageAnthropicNonStreamServerToolDetailPricing(t *testing.T) { } func TestUsageAnthropicStreamPatchState(t *testing.T) { - extractor := NewAnthropic(http.DefaultClient) + extractor := NewAnthropic() tests := []struct { name string steps []string @@ -687,7 +686,7 @@ func TestUsageAnthropicStreamPatchState(t *testing.T) { } func TestUsageAnthropicStreamPresenceAndInvalidStartHandling(t *testing.T) { - extractor := NewAnthropic(http.DefaultClient) + extractor := NewAnthropic() tests := []struct { name string steps []string @@ -767,7 +766,7 @@ func TestUsageAnthropicStreamPresenceAndInvalidStartHandling(t *testing.T) { } func TestUsageAnthropicMalformedPayloadPreservesStateAndDoesNotLeak(t *testing.T) { - stream := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor() + stream := NewAnthropic().NewUsageStreamExtractor() if err := stream.Observe([]byte(`{"type":"message_start","message":{"usage":{"input_tokens":80}}}`)); err != nil { t.Fatal(err) } @@ -791,7 +790,7 @@ func TestUsageAnthropicMalformedPayloadPreservesStateAndDoesNotLeak(t *testing.T } func TestUsageAnthropicInvalidUsageObjectIsDiagnosed(t *testing.T) { - extractor := NewAnthropic(http.DefaultClient) + extractor := NewAnthropic() result, err := extractor.ExtractUsage([]byte(`{"usage":[]}`)) if err != nil || result.State != usage.StateMissing || result.Tokens != (usage.Tokens{}) { t.Fatalf("invalid non-stream usage = %#v, %v", result, err) @@ -810,7 +809,7 @@ func TestUsageAnthropicInvalidUsageObjectIsDiagnosed(t *testing.T) { } func TestUsageAnthropicMalformedBodiesReturnSanitizedErrors(t *testing.T) { - extractor := NewAnthropic(http.DefaultClient) + extractor := NewAnthropic() for _, body := range [][]byte{ []byte(`[]`), []byte(`{"usage":{"input_tokens":80,"output_tokens":30}} {}`), diff --git a/internal/dialect/classify.go b/internal/dialect/classify.go deleted file mode 100644 index 146cbe95..00000000 --- a/internal/dialect/classify.go +++ /dev/null @@ -1,72 +0,0 @@ -package dialect - -import ( - "net/http" - "strings" - - "gpt-load/internal/health" -) - -type failureMarkers struct { - rateLimited []string - modelUnavailable []string - invalidKey []string - upstreamHost []string -} - -func classifyStatusWithMarkers( - status int, - body []byte, - markers failureMarkers, -) health.FailureCategory { - if status >= http.StatusOK && status < http.StatusMultipleChoices { - return health.FailureCategoryOK - } - if status == 0 { - return health.FailureCategoryAmbiguous - } - lowered := strings.ToLower(string(body)) - switch { - case status == http.StatusTooManyRequests || containsFailureMarker(lowered, markers.rateLimited): - return health.FailureCategoryRateLimited - case status == http.StatusNotFound || containsFailureMarker(lowered, markers.modelUnavailable): - return health.FailureCategoryModelUnavailable - case status == http.StatusUnauthorized || status == http.StatusForbidden || - containsFailureMarker(lowered, markers.invalidKey): - return health.FailureCategoryInvalidKey - case status >= http.StatusInternalServerError || containsFailureMarker(lowered, markers.upstreamHost): - return health.FailureCategoryUpstreamHostError - case status >= http.StatusBadRequest && status < http.StatusInternalServerError: - return health.FailureCategoryClientError - default: - return health.FailureCategoryAmbiguous - } -} - -func classifyProviderErrorWithMarkers( - body []byte, - markers failureMarkers, -) health.FailureCategory { - lowered := strings.ToLower(string(body)) - switch { - case containsFailureMarker(lowered, markers.rateLimited): - return health.FailureCategoryRateLimited - case containsFailureMarker(lowered, markers.modelUnavailable): - return health.FailureCategoryModelUnavailable - case containsFailureMarker(lowered, markers.invalidKey): - return health.FailureCategoryInvalidKey - case containsFailureMarker(lowered, markers.upstreamHost): - return health.FailureCategoryUpstreamHostError - default: - return health.FailureCategoryClientError - } -} - -func containsFailureMarker(body string, markers []string) bool { - for _, marker := range markers { - if strings.Contains(body, marker) { - return true - } - } - return false -} diff --git a/internal/dialect/classify_test.go b/internal/dialect/classify_test.go deleted file mode 100644 index 9e97d6f6..00000000 --- a/internal/dialect/classify_test.go +++ /dev/null @@ -1,90 +0,0 @@ -package dialect - -import ( - "net/http" - "testing" - - "gpt-load/internal/health" -) - -func TestClassifyProviderErrorIgnoresHTTP2xxShortcut(t *testing.T) { - type providerErrorClassifier interface { - ClassifyProviderError([]byte) health.FailureCategory - } - - dialects := []struct { - name string - value Dialect - cases []struct { - name string - body string - want health.FailureCategory - } - }{ - { - name: "openai", - value: NewOpenAI(http.DefaultClient), - cases: []struct { - name string - body string - want health.FailureCategory - }{ - {name: "rate limited", body: `{"error":{"type":"rate_limit_error"}}`, want: health.FailureCategoryRateLimited}, - {name: "model unavailable", body: `{"error":{"code":"model_not_found"}}`, want: health.FailureCategoryModelUnavailable}, - {name: "invalid key", body: `{"error":{"code":"invalid_api_key"}}`, want: health.FailureCategoryInvalidKey}, - {name: "host error", body: `{"error":{"type":"server_overloaded"}}`, want: health.FailureCategoryUpstreamHostError}, - {name: "unknown provider error", body: `{"error":{"type":"unexpected_provider_error"}}`, want: health.FailureCategoryClientError}, - }, - }, - { - name: "anthropic", - value: NewAnthropic(http.DefaultClient), - cases: []struct { - name string - body string - want health.FailureCategory - }{ - {name: "rate limited", body: `{"type":"error","error":{"type":"rate_limit_error"}}`, want: health.FailureCategoryRateLimited}, - {name: "model unavailable", body: `{"type":"error","error":{"type":"model_not_found"}}`, want: health.FailureCategoryModelUnavailable}, - {name: "invalid key", body: `{"type":"error","error":{"type":"authentication_error"}}`, want: health.FailureCategoryInvalidKey}, - {name: "host error", body: `{"type":"error","error":{"type":"overloaded_error"}}`, want: health.FailureCategoryUpstreamHostError}, - {name: "unknown provider error", body: `{"type":"error","error":{"type":"unexpected_provider_error"}}`, want: health.FailureCategoryClientError}, - }, - }, - { - name: "gemini", - value: NewGemini(http.DefaultClient), - cases: []struct { - name string - body string - want health.FailureCategory - }{ - {name: "rate limited", body: `{"error":{"status":"RESOURCE_EXHAUSTED"}}`, want: health.FailureCategoryRateLimited}, - {name: "model unavailable", body: `{"error":{"status":"MODEL_NOT_FOUND"}}`, want: health.FailureCategoryModelUnavailable}, - {name: "invalid key", body: `{"error":{"status":"UNAUTHENTICATED"}}`, want: health.FailureCategoryInvalidKey}, - {name: "host error", body: `{"error":{"status":"SERVICE_UNAVAILABLE"}}`, want: health.FailureCategoryUpstreamHostError}, - {name: "unknown provider error", body: `{"error":{"status":"UNEXPECTED_PROVIDER_ERROR"}}`, want: health.FailureCategoryClientError}, - }, - }, - } - - for _, dialectCase := range dialects { - t.Run(dialectCase.name, func(t *testing.T) { - classifier, ok := dialectCase.value.(providerErrorClassifier) - if !ok { - t.Fatalf("%T does not implement the optional provider-error classifier", dialectCase.value) - } - for _, test := range dialectCase.cases { - t.Run(test.name, func(t *testing.T) { - body := []byte(test.body) - if got := dialectCase.value.ClassifyStatus(http.StatusOK, body); got != health.FailureCategoryOK { - t.Fatalf("ClassifyStatus(200) = %v, want OK", got) - } - if got := classifier.ClassifyProviderError(body); got != test.want { - t.Fatalf("ClassifyProviderError() = %v, want %v", got, test.want) - } - }) - } - }) - } -} diff --git a/internal/dialect/dialect.go b/internal/dialect/dialect.go index 509eb51b..d61e277c 100644 --- a/internal/dialect/dialect.go +++ b/internal/dialect/dialect.go @@ -1,13 +1,11 @@ package dialect import ( - "context" "net/http" - "gpt-load/internal/health" + "gpt-load/internal/execution" "gpt-load/internal/protocol" "gpt-load/internal/reasoning" - "gpt-load/internal/state" "gpt-load/internal/usage" ) @@ -22,6 +20,8 @@ type ParsedRequest struct { type RequestMetadata struct { Model *string Stream bool + Operation execution.Operation + RequiredFeatures execution.FeatureSet ObserveUsage bool UsageDiagnostics usage.Diagnostics Reasoning reasoning.Config @@ -30,20 +30,6 @@ type RequestMetadata struct { type Dialect interface { Protocol() protocol.Protocol InspectRequest(req *ParsedRequest) (RequestMetadata, error) - BuildUpstreamURL(base string, req *ParsedRequest) (string, error) - InjectCredential(headers http.Header, apiKey string) - ListModels( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, - ) ([]string, error) - Probe( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, - validationModel string, - ) error - ClassifyStatus(status int, body []byte) health.FailureCategory } type ModelRewriter interface { @@ -58,7 +44,3 @@ type ModelRewriter interface { type ResponseModelInspector interface { InspectResponseModels(payload []byte) []string } - -type CredentialHeaderNamer interface { - CredentialHeaderNames() []string -} diff --git a/internal/dialect/gemini.go b/internal/dialect/gemini.go index a7c17311..20116f79 100644 --- a/internal/dialect/gemini.go +++ b/internal/dialect/gemini.go @@ -1,64 +1,30 @@ package dialect import ( - "context" "fmt" - "net/http" - "net/url" "strings" - "gpt-load/internal/health" - platformheader "gpt-load/internal/platform/httpheader" "gpt-load/internal/protocol" - "gpt-load/internal/state" ) const ( - geminiGenerationPrefix = "/v1beta/models/" - geminiGenerateSuffix = ":generateContent" - geminiStreamSuffix = ":streamGenerateContent" - geminiModelsResourcePath = "/models" + geminiGenerationPrefix = "/v1beta/models/" + geminiGenerateSuffix = ":generateContent" + geminiStreamSuffix = ":streamGenerateContent" ) -var geminiFailureMarkers = failureMarkers{ - rateLimited: []string{"resource_exhausted", "rate_limit", "rate limit", "quota_exceeded", "quota exceeded"}, - modelUnavailable: []string{"model_not_found", "model not found", "model_not_supported", "model not supported", "no access to model"}, - invalidKey: []string{"api_key_invalid", "unauthenticated", "permission_denied", "api key not valid", "invalid api key", "api key disabled", "api key banned"}, - upstreamHost: []string{"service_unavailable", "backend_error", "internal_error"}, -} - -type Gemini struct { - client *http.Client -} - -type geminiModelPage struct { - Models []struct { - Name string `json:"name"` - } `json:"models"` - NextPageToken string `json:"nextPageToken"` -} +type Gemini struct{} var _ Dialect = (*Gemini)(nil) -func NewGemini(client *http.Client) *Gemini { - return &Gemini{client: client} +func NewGemini() *Gemini { + return &Gemini{} } -func (d *Gemini) Protocol() protocol.Protocol { +func (*Gemini) Protocol() protocol.Protocol { return protocol.Gemini } -func (d *Gemini) InjectCredential(headers http.Header, apiKey string) { - if headers == nil { - return - } - headers.Set("X-Goog-Api-Key", apiKey) -} - -func (d *Gemini) CredentialHeaderNames() []string { - return []string{"X-Goog-Api-Key"} -} - func (d *Gemini) InspectRequest(req *ParsedRequest) (RequestMetadata, error) { if req == nil { return RequestMetadata{}, fmt.Errorf("parsed request is required") @@ -67,185 +33,18 @@ func (d *Gemini) InspectRequest(req *ParsedRequest) (RequestMetadata, error) { if err != nil { return RequestMetadata{}, err } - return RequestMetadata{ + metadata := RequestMetadata{ Model: &model, Stream: stream, ObserveUsage: true, Reasoning: inspectGeminiReasoning(req.Body), - }, nil -} - -func (d *Gemini) BuildUpstreamURL(base string, req *ParsedRequest) (string, error) { - if req == nil { - return "", fmt.Errorf("parsed request is required") } - _, stream, err := parseGeminiGenerationPath(req.Path) - if err != nil { - return "", err - } - resourcePath := strings.TrimPrefix(req.Path, "/v1beta") - upstream, err := resolveUpstreamURL(base, resourcePath, req.RawQuery) - if err != nil { - return "", err - } - parsed, err := url.Parse(upstream) - if err != nil { - return "", fmt.Errorf("parse Gemini upstream URL: %w", err) - } - query := parsed.Query() - query.Del("alt") - if stream { - query.Set("alt", "sse") - } - parsed.RawQuery = query.Encode() - return parsed.String(), nil -} - -func (d *Gemini) ListModels( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, -) ([]string, error) { - requestURL, err := resolveUpstreamURL(baseURL, geminiModelsResourcePath, "") - if err != nil { - return nil, fmt.Errorf("build Gemini model-list URL: %w", err) - } - parsed, err := url.Parse(requestURL) - if err != nil { - return nil, fmt.Errorf("parse Gemini model-list URL: %w", err) - } - - collector := newModelListCollector() - seenPageTokens := make(map[string]struct{}) - pageToken := "" - for pageNumber := 1; pageNumber <= maxModelListPages; pageNumber++ { - page, err := d.listModelsPage(ctx, parsed, apiKey, rules, pageToken) - if err != nil { - return nil, err - } - pageModels := make([]string, 0, len(page.Models)) - for _, item := range page.Models { - name := strings.TrimPrefix(item.Name, "models/") - if name != "" { - pageModels = append(pageModels, name) - } - } - if err := collector.Add(pageModels); err != nil { - return nil, err - } - if page.NextPageToken == "" { - return collector.Result(), nil - } - if strings.TrimSpace(page.NextPageToken) == "" { - return nil, fmt.Errorf("Gemini model-list page token is empty") - } - if _, repeated := seenPageTokens[page.NextPageToken]; repeated { - return nil, fmt.Errorf("Gemini model-list page token repeated") - } - if pageNumber == maxModelListPages || collector.Full() { - return nil, fmt.Errorf("Gemini model-list pagination limit exceeded") - } - seenPageTokens[page.NextPageToken] = struct{}{} - pageToken = page.NextPageToken - } - return nil, fmt.Errorf("Gemini model-list pagination limit exceeded") -} - -func (d *Gemini) listModelsPage( - ctx context.Context, - endpoint *url.URL, - apiKey string, - rules state.HeaderRules, - pageToken string, -) (geminiModelPage, error) { - pageEndpoint := *endpoint - query := pageEndpoint.Query() - query.Del("pageSize") - query.Set("pageSize", "1000") - query.Del("pageToken") - if pageToken != "" { - query.Set("pageToken", pageToken) - } - pageEndpoint.RawQuery = query.Encode() - - request, err := http.NewRequestWithContext(ctx, http.MethodGet, pageEndpoint.String(), nil) - if err != nil { - return geminiModelPage{}, fmt.Errorf("create Gemini model-list request: %w", err) - } - ApplyCredential(d, request.Header, apiKey, rules) - platformheader.NormalizeUpstreamRequestRepresentation(request, 0) - - response, err := d.client.Do(request) - if err != nil { - if contextErr := ctx.Err(); contextErr != nil { - return geminiModelPage{}, fmt.Errorf("request Gemini model list: %w", contextErr) - } - return geminiModelPage{}, fmt.Errorf("request Gemini model list failed") - } - defer response.Body.Close() - if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { - return geminiModelPage{}, fmt.Errorf("request Gemini model list: upstream status %d", response.StatusCode) - } - - var page geminiModelPage - if err := decodeModelListPage(response, &page); err != nil { - return geminiModelPage{}, fmt.Errorf("decode Gemini model list: %w", err) - } - return page, nil -} - -func (d *Gemini) Probe( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, - validationModel string, -) error { - if err := validateProbeModel(validationModel); err != nil { - return err - } - requestURL, err := d.BuildUpstreamURL(baseURL, &ParsedRequest{ - Method: http.MethodPost, - Path: geminiGenerationPrefix + validationModel + geminiGenerateSuffix, - }) - if err != nil { - return fmt.Errorf("build %s probe URL failed", d.Protocol()) - } - return executeProbe(ctx, d.client, d, requestURL, apiKey, rules, struct { - Contents []struct { - Role string `json:"role"` - Parts []struct { - Text string `json:"text"` - } `json:"parts"` - } `json:"contents"` - GenerationConfig struct { - MaxOutputTokens int `json:"maxOutputTokens"` - } `json:"generationConfig"` - }{ - Contents: []struct { - Role string `json:"role"` - Parts []struct { - Text string `json:"text"` - } `json:"parts"` - }{ - { - Role: "user", - Parts: []struct { - Text string `json:"text"` - }{{Text: "ping"}}, - }, - }, - GenerationConfig: struct { - MaxOutputTokens int `json:"maxOutputTokens"` - }{MaxOutputTokens: 1}, - }) -} - -func (d *Gemini) ClassifyStatus(status int, body []byte) health.FailureCategory { - return classifyStatusWithMarkers(status, body, geminiFailureMarkers) -} - -func (d *Gemini) ClassifyProviderError(body []byte) health.FailureCategory { - return classifyProviderErrorWithMarkers(body, geminiFailureMarkers) + metadata.Operation, metadata.RequiredFeatures = chatExecutionMetadata( + req.Body, + metadata.Stream, + metadata.Reasoning, + ) + return metadata, nil } func parseGeminiGenerationPath(path string) (model string, stream bool, err error) { diff --git a/internal/dialect/gemini_test.go b/internal/dialect/gemini_test.go index ae9781f6..e9f3804f 100644 --- a/internal/dialect/gemini_test.go +++ b/internal/dialect/gemini_test.go @@ -1,26 +1,19 @@ package dialect import ( - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "net/http/httptest" - "net/url" - "strings" - "sync/atomic" "testing" - "time" - "gpt-load/internal/health" "gpt-load/internal/protocol" - "gpt-load/internal/state" ) +func TestGeminiProtocol(t *testing.T) { + if got := NewGemini().Protocol(); got != protocol.Gemini { + t.Fatalf("Protocol() = %q, want %q", got, protocol.Gemini) + } +} + func TestGeminiInspectRequest(t *testing.T) { - value := NewGemini(http.DefaultClient) + selected := NewGemini() tests := []struct { name string request *ParsedRequest @@ -35,514 +28,21 @@ func TestGeminiInspectRequest(t *testing.T) { {name: "wrong suffix", request: &ParsedRequest{Path: "/v1beta/models/gemini:embedContent"}, wantErr: true}, {name: "empty model", request: &ParsedRequest{Path: "/v1beta/models/:generateContent"}, wantErr: true}, {name: "nested slash", request: &ParsedRequest{Path: "/v1beta/models/vendor/gemini:generateContent"}, wantErr: true}, - {name: "configured models prefix is not corrected", request: &ParsedRequest{Path: "/v1beta/models/models/custom:generateContent"}, wantErr: true}, {name: "boundary whitespace", request: &ParsedRequest{Path: "/v1beta/models/ gemini :generateContent"}, wantErr: true}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - metadata, err := value.InspectRequest(test.request) + metadata, err := selected.InspectRequest(test.request) if test.wantErr { if err == nil { - t.Fatalf("InspectRequest() = (%#v, nil), want error", metadata) + t.Fatalf("InspectRequest() = %#v, nil", metadata) } return } - if err != nil || metadata.Model == nil || - *metadata.Model != test.wantModel || - metadata.Stream != test.wantStream || - !metadata.ObserveUsage { - t.Fatalf("InspectRequest() = (%#v, %v), want model=%q stream=%t observe=true", metadata, err, test.wantModel, test.wantStream) + if err != nil || metadata.Model == nil || *metadata.Model != test.wantModel || + metadata.Stream != test.wantStream || !metadata.ObserveUsage { + t.Fatalf("InspectRequest() = %#v, %v", metadata, err) } }) } } - -func TestGeminiBuildUpstreamURL(t *testing.T) { - value := NewGemini(http.DefaultClient) - - streamURL, err := value.BuildUpstreamURL( - "https://api.example.com/proxy/v1beta/?alt=json&alt=proto&tenant=base", - &ParsedRequest{ - Path: "/v1beta/models/gemini-2.5-pro:streamGenerateContent", - RawQuery: "alt=xml&trace=true", - }, - ) - if err != nil { - t.Fatalf("stream BuildUpstreamURL() error = %v", err) - } - assertGeminiURL(t, streamURL, "/proxy/v1beta/models/gemini-2.5-pro:streamGenerateContent", url.Values{ - "alt": {"sse"}, - "tenant": {"base"}, - "trace": {"true"}, - }) - - nonStreamURL, err := value.BuildUpstreamURL( - "https://api.example.com/v1beta?alt=json&tenant=base", - &ParsedRequest{Path: "/v1beta/models/gemini-2.5-flash:generateContent", RawQuery: "alt=xml&trace=true"}, - ) - if err != nil { - t.Fatalf("non-stream BuildUpstreamURL() error = %v", err) - } - assertGeminiURL(t, nonStreamURL, "/v1beta/models/gemini-2.5-flash:generateContent", url.Values{ - "tenant": {"base"}, - "trace": {"true"}, - }) - - for _, request := range []*ParsedRequest{ - nil, - {Path: "/v1beta/models/gemini:embedContent"}, - {Path: "/v1beta/models/:generateContent"}, - } { - if got, err := value.BuildUpstreamURL("https://api.example.com", request); err == nil { - t.Fatalf("BuildUpstreamURL(%#v) = %q, want error", request, got) - } - } - if got, err := value.BuildUpstreamURL("ftp://api.example.com", &ParsedRequest{Path: "/v1beta/models/gemini:generateContent"}); err == nil { - t.Fatalf("BuildUpstreamURL() = %q, want invalid base error", got) - } -} - -func assertGeminiURL(t *testing.T, rawURL, wantPath string, wantQuery url.Values) { - t.Helper() - parsed, err := url.Parse(rawURL) - if err != nil { - t.Fatalf("parse URL %q: %v", rawURL, err) - } - if parsed.Path != wantPath { - t.Fatalf("path = %q, want %q", parsed.Path, wantPath) - } - if parsed.Query().Encode() != wantQuery.Encode() { - t.Fatalf("query = %q, want %q", parsed.Query().Encode(), wantQuery.Encode()) - } - if parsed.Query().Get("key") != "" { - t.Fatalf("unexpected key query: %q", parsed.RawQuery) - } -} - -func TestGeminiCredentialAndProtocol(t *testing.T) { - value := NewGemini(http.DefaultClient) - if got := value.Protocol(); got != protocol.Gemini { - t.Fatalf("Protocol() = %q, want %q", got, protocol.Gemini) - } - headers := make(http.Header) - value.InjectCredential(headers, "upstream-secret") - if got := headers.Get("X-Goog-Api-Key"); got != "upstream-secret" || len(headers) != 1 { - t.Fatalf("headers = %#v", headers) - } - - ApplyCredential(value, headers, "upstream-secret", state.HeaderRules{ - Set: map[string]string{"X-Goog-Api-Key": "override-${API_KEY}"}, - }) - if got := headers.Get("X-Goog-Api-Key"); got != "override-upstream-secret" { - t.Fatalf("overridden X-Goog-Api-Key = %q", got) - } - ApplyCredential(value, headers, "upstream-secret", state.HeaderRules{ - Set: map[string]string{"X-Goog-Api-Key": "must-be-removed"}, - Remove: []string{"x-goog-api-key"}, - }) - if got := headers.Get("X-Goog-Api-Key"); got != "" { - t.Fatalf("removed X-Goog-Api-Key = %q", got) - } - value.InjectCredential(nil, "upstream-secret") -} - -func TestGeminiClassifyStatus(t *testing.T) { - value := NewGemini(http.DefaultClient) - tests := []struct { - name string - status int - body string - want health.FailureCategory - }{ - {name: "success", status: http.StatusOK, body: `{"status":"api_key_invalid"}`, want: health.FailureCategoryOK}, - {name: "unauthorized", status: http.StatusUnauthorized, want: health.FailureCategoryInvalidKey}, - {name: "forbidden", status: http.StatusForbidden, want: health.FailureCategoryInvalidKey}, - {name: "not found", status: http.StatusNotFound, want: health.FailureCategoryModelUnavailable}, - {name: "rate limited", status: http.StatusTooManyRequests, want: health.FailureCategoryRateLimited}, - {name: "server error", status: http.StatusServiceUnavailable, want: health.FailureCategoryUpstreamHostError}, - {name: "redirect is ambiguous", status: http.StatusTemporaryRedirect, want: health.FailureCategoryAmbiguous}, - {name: "zero status with marker is ambiguous", status: 0, body: `{"status":"api_key_invalid"}`, want: health.FailureCategoryAmbiguous}, - {name: "invalid key marker", status: http.StatusBadRequest, body: `{"status":"api_key_invalid"}`, want: health.FailureCategoryInvalidKey}, - {name: "permission marker", status: http.StatusBadRequest, body: `{"status":"permission_denied"}`, want: health.FailureCategoryInvalidKey}, - {name: "quota marker", status: http.StatusBadRequest, body: `{"status":"resource_exhausted"}`, want: health.FailureCategoryRateLimited}, - {name: "model marker", status: http.StatusBadRequest, body: `{"message":"model not supported"}`, want: health.FailureCategoryModelUnavailable}, - {name: "invalid argument", status: http.StatusBadRequest, body: `{"status":"invalid_argument"}`, want: health.FailureCategoryClientError}, - {name: "ordinary unavailable text", status: http.StatusBadRequest, body: `{"message":"tool is unavailable for this request"}`, want: health.FailureCategoryClientError}, - {name: "ordinary client error", status: http.StatusBadRequest, body: `{"error":"invalid input"}`, want: health.FailureCategoryClientError}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - if got := value.ClassifyStatus(test.status, []byte(test.body)); got != test.want { - t.Fatalf("ClassifyStatus() = %d, want %d", got, test.want) - } - }) - } -} - -func TestGeminiListModelsPaginatesPreservesInputsAndParsesNames(t *testing.T) { - var requestCount atomic.Int64 - const pageToken = " token/+= ? " - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - pageNumber := requestCount.Add(1) - if request.Method != http.MethodGet || request.URL.Path != "/prefix/v1beta/models" { - t.Errorf("request = %s %s", request.Method, request.URL.Path) - } - if got := request.URL.Query()["pageSize"]; len(got) != 1 || got[0] != "1000" { - t.Errorf("pageSize = %#v, want one 1000", got) - } - if got := request.URL.Query().Get("tenant"); got != "one" { - t.Errorf("tenant = %q, want one", got) - } - if got := request.URL.Query().Get("fixed"); got != "preserved" { - t.Errorf("fixed = %q, want preserved", got) - } - if got := request.Header.Get("X-Goog-Api-Key"); got != "upstream-secret" { - t.Errorf("X-Goog-Api-Key = %q", got) - } - if got := request.Header.Get("X-Discovery-Rule"); got != "applied-rule" { - t.Errorf("X-Discovery-Rule = %q, want applied-rule", got) - } - switch pageNumber { - case 1: - if _, exists := request.URL.Query()["pageToken"]; exists { - t.Errorf("page 1 retained stale pageToken: %q", request.URL.Query().Get("pageToken")) - } - _, _ = writer.Write([]byte(`{ - "models":[ - {"name":"models/gemini-versioned","baseModelId":"gemini-base"}, - {"name":"models/models/vendor-model"}, - {"name":"shared"}, - {"baseModelId":"base-only"} - ], - "nextPageToken":"` + pageToken + `" - }`)) - case 2: - if got := request.URL.Query().Get("pageToken"); got != pageToken { - t.Errorf("page 2 pageToken = %q, want byte-for-byte %q", got, pageToken) - } - _, _ = writer.Write([]byte(`{ - "models":[ - {"name":"models/shared"}, - {"name":"models/gemini-next"}, - {"name":"models/"} - ] - }`)) - default: - t.Errorf("unexpected page %d", pageNumber) - writer.WriteHeader(http.StatusInternalServerError) - } - })) - defer server.Close() - - models, err := NewGemini(server.Client()).ListModels( - context.Background(), - server.URL+"/prefix/v1beta?tenant=one&fixed=preserved&pageSize=1&pageSize=2&pageToken=stale", - "upstream-secret", - state.HeaderRules{Set: map[string]string{"X-Discovery-Rule": "applied-rule"}}, - ) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - if strings.Join(models, ",") != "gemini-versioned,models/vendor-model,shared,gemini-next" { - t.Fatalf("models = %#v", models) - } - if got := requestCount.Load(); got != 2 { - t.Fatalf("request count = %d, want 2", got) - } -} - -func TestGeminiListModelsIdentityRepresentation(t *testing.T) { - client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { - assertOutboundRequestRepresentation(t, request, 0, map[string]string{ - "X-Goog-Api-Key": "override-secret", - "X-Business": "preserved", - }) - return geminiTestResponse(http.StatusOK, `{"models":[]}`), nil - })} - models, err := NewGemini(client).ListModels( - context.Background(), - "https://api.example.com", - "secret", - legacyRepresentationHeaderRules(map[string]string{"X-Goog-Api-Key": "override-${API_KEY}"}), - ) - if err != nil || models == nil { - t.Fatalf("ListModels() = %#v, %v", models, err) - } -} - -func TestGeminiListModelsRejectsBlankOrRepeatedPageToken(t *testing.T) { - t.Run("blank page token", func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - requestCount.Add(1) - return geminiTestResponse(http.StatusOK, `{"models":[{"name":"models/gemini-a"}],"nextPageToken":" \t "}`), nil - })} - models, err := NewGemini(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and error", models, err) - } - if got := requestCount.Load(); got != 1 { - t.Fatalf("request count = %d, want 1", got) - } - }) - - t.Run("repeated page token", func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - requestCount.Add(1) - return geminiTestResponse(http.StatusOK, `{"models":[{"name":"models/gemini-a"}],"nextPageToken":"same-token"}`), nil - })} - models, err := NewGemini(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and error", models, err) - } - if got := requestCount.Load(); got != 2 { - t.Fatalf("request count = %d, want 2", got) - } - }) -} - -func TestGeminiListModelsRejectsPageAndUniqueModelLimits(t *testing.T) { - t.Run("page 100 may end", func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - pageNumber := requestCount.Add(1) - nextPageToken := "" - if pageNumber < maxModelListPages { - nextPageToken = fmt.Sprintf("token-%d", pageNumber) - } - return geminiTestResponse(http.StatusOK, fmt.Sprintf( - `{"models":[{"name":"models/model-%d"}],"nextPageToken":%q}`, - pageNumber, nextPageToken, - )), nil - })} - models, err := NewGemini(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err != nil || len(models) != maxModelListPages { - t.Fatalf("ListModels() = %d models, %v, want %d models", len(models), err, maxModelListPages) - } - if got := requestCount.Load(); got != maxModelListPages { - t.Fatalf("request count = %d, want %d", got, maxModelListPages) - } - }) - - t.Run("page 100 may not continue", func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - pageNumber := requestCount.Add(1) - return geminiTestResponse(http.StatusOK, fmt.Sprintf( - `{"models":[{"name":"models/model-%d"}],"nextPageToken":"token-%d"}`, - pageNumber, pageNumber, - )), nil - })} - models, err := NewGemini(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and error", models, err) - } - if got := requestCount.Load(); got != maxModelListPages { - t.Fatalf("request count = %d, want %d", got, maxModelListPages) - } - }) - - for _, test := range []struct { - name string - modelCount int - continuation bool - wantErr bool - }{ - {name: "exact unique maximum may end", modelCount: maxUniqueModelListEntries}, - {name: "exact unique maximum may not continue", modelCount: maxUniqueModelListEntries, continuation: true, wantErr: true}, - {name: "more than unique maximum is rejected", modelCount: maxUniqueModelListEntries + 1, wantErr: true}, - } { - t.Run(test.name, func(t *testing.T) { - body := mustGeminiTestPayload(t, test.modelCount, test.continuation) - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - requestCount.Add(1) - return geminiTestResponse(http.StatusOK, body), nil - })} - models, err := NewGemini(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if test.wantErr { - if err == nil || models != nil { - t.Fatalf("ListModels() = %d models, %v, want nil models and error", len(models), err) - } - } else if err != nil || len(models) != test.modelCount { - t.Fatalf("ListModels() = %d models, %v, want %d models", len(models), err, test.modelCount) - } - if got := requestCount.Load(); got != 1 { - t.Fatalf("request count = %d, want 1", got) - } - }) - } -} - -func TestGeminiListModelsReturnsNoPartialResultAfterLaterPageFailure(t *testing.T) { - for _, test := range []struct { - name string - secondPage func() (*http.Response, error) - }{ - {name: "transport error", secondPage: func() (*http.Response, error) { - return nil, errors.New("later transport failure") - }}, - {name: "non-2xx", secondPage: func() (*http.Response, error) { - return geminiTestResponse(http.StatusBadGateway, `{"error":"later failure"}`), nil - }}, - {name: "malformed JSON", secondPage: func() (*http.Response, error) { - return geminiTestResponse(http.StatusOK, `{"models":[`), nil - }}, - } { - t.Run(test.name, func(t *testing.T) { - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - if requestCount.Add(1) == 1 { - return geminiTestResponse(http.StatusOK, `{"models":[{"name":"models/first-page-model"}],"nextPageToken":"next"}`), nil - } - return test.secondPage() - })} - models, err := NewGemini(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and later-page error", models, err) - } - if got := requestCount.Load(); got != 2 { - t.Fatalf("request count = %d, want 2", got) - } - }) - } -} - -func TestGeminiListModelsClosesBodyBeforeRequestingNextPage(t *testing.T) { - firstBody := &trackedReadCloser{Reader: strings.NewReader(`{"models":[{"name":"models/first"}],"nextPageToken":"next"}`)} - secondBody := &trackedReadCloser{Reader: strings.NewReader(`{"models":[{"name":"models/second"}]}`)} - var requestCount atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - switch requestCount.Add(1) { - case 1: - return &http.Response{StatusCode: http.StatusOK, Body: firstBody, Header: make(http.Header)}, nil - case 2: - if !firstBody.closed.Load() { - return nil, errors.New("first page body is still open") - } - return &http.Response{StatusCode: http.StatusOK, Body: secondBody, Header: make(http.Header)}, nil - default: - return nil, errors.New("unexpected extra request") - } - })} - models, err := NewGemini(client).ListModels(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}) - if err != nil || strings.Join(models, ",") != "first,second" { - t.Fatalf("ListModels() = %#v, %v, want first,second", models, err) - } - if !secondBody.closed.Load() { - t.Fatal("last page body was not closed") - } -} - -func TestGeminiListModelsRulesAndFailures(t *testing.T) { - t.Run("rules override and remove", func(t *testing.T) { - requests := atomic.Int64{} - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - call := requests.Add(1) - switch call { - case 1: - if got := request.Header.Get("X-Goog-Api-Key"); got != "override-upstream-secret" { - t.Errorf("overridden X-Goog-Api-Key = %q", got) - } - case 2: - if got := request.Header.Get("X-Goog-Api-Key"); got != "" { - t.Errorf("removed X-Goog-Api-Key = %q", got) - } - } - _, _ = writer.Write([]byte(`{"models":[]}`)) - })) - defer server.Close() - models, err := NewGemini(server.Client()).ListModels( - context.Background(), server.URL, "upstream-secret", - state.HeaderRules{Set: map[string]string{"X-Goog-Api-Key": "override-${API_KEY}"}}, - ) - if err != nil || models == nil || len(models) != 0 { - t.Fatalf("override ListModels() = %#v, %v", models, err) - } - _, err = NewGemini(server.Client()).ListModels( - context.Background(), server.URL, "upstream-secret", - state.HeaderRules{Remove: []string{"x-goog-api-key"}}, - ) - if err != nil { - t.Fatalf("remove ListModels() error = %v", err) - } - }) - - t.Run("non-success does not expose body or key", func(t *testing.T) { - const secret = "distinctive-gemini-secret" - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.WriteHeader(http.StatusForbidden) - _, _ = writer.Write([]byte(`{"error":"` + secret + ` in body"}`)) - })) - defer server.Close() - _, err := NewGemini(server.Client()).ListModels(context.Background(), server.URL, secret, state.HeaderRules{}) - if err == nil || strings.Contains(err.Error(), secret) || strings.Contains(err.Error(), server.URL) { - t.Fatalf("ListModels() error = %v", err) - } - }) - - t.Run("malformed JSON is contextual and key free", func(t *testing.T) { - const secret = "distinctive-gemini-secret" - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - _, _ = writer.Write([]byte(`{"models":[`)) - })) - defer server.Close() - _, err := NewGemini(server.Client()).ListModels( - context.Background(), server.URL+"?token=query-secret", secret, state.HeaderRules{}, - ) - if err == nil || !strings.Contains(err.Error(), "decode Gemini model list") { - t.Fatalf("ListModels() error = %v", err) - } - for _, forbidden := range []string{secret, "query-secret", server.URL} { - if strings.Contains(err.Error(), forbidden) { - t.Fatalf("error exposes %q: %v", forbidden, err) - } - } - }) - - t.Run("caller context stops request", func(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { - <-request.Context().Done() - })) - defer server.Close() - ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) - defer cancel() - _, err := NewGemini(server.Client()).ListModels(ctx, server.URL, "secret", state.HeaderRules{}) - if !errors.Is(err, context.DeadlineExceeded) { - t.Fatalf("ListModels() error = %v, want deadline exceeded", err) - } - }) -} - -func geminiTestResponse(statusCode int, body string) *http.Response { - return &http.Response{ - StatusCode: statusCode, - Body: io.NopCloser(strings.NewReader(body)), - Header: make(http.Header), - } -} - -func mustGeminiTestPayload(t *testing.T, modelCount int, continuation bool) string { - t.Helper() - type item struct { - Name string `json:"name"` - } - payload := struct { - Models []item `json:"models"` - NextPageToken string `json:"nextPageToken,omitempty"` - }{ - Models: make([]item, modelCount), - } - if continuation { - payload.NextPageToken = "next-token" - } - for index := range payload.Models { - payload.Models[index].Name = fmt.Sprintf("models/model-%06d", index) - } - encoded, err := json.Marshal(payload) - if err != nil { - t.Fatalf("marshal Gemini test payload: %v", err) - } - return string(encoded) -} diff --git a/internal/dialect/gemini_usage_test.go b/internal/dialect/gemini_usage_test.go index 540e2450..d02f95b6 100644 --- a/internal/dialect/gemini_usage_test.go +++ b/internal/dialect/gemini_usage_test.go @@ -1,7 +1,6 @@ package dialect import ( - "net/http" "strings" "testing" @@ -9,7 +8,7 @@ import ( ) func TestUsageGeminiCanonicalFixtures(t *testing.T) { - provider := Dialect(NewGemini(http.DefaultClient)) + provider := Dialect(NewGemini()) extractor, ok := provider.(UsageExtractor) if !ok { t.Fatal("Gemini does not implement UsageExtractor") @@ -36,7 +35,7 @@ func TestUsageGeminiCanonicalFixtures(t *testing.T) { } func TestUsageGeminiNonStreamOptionalFieldsAndArithmetic(t *testing.T) { - extractor := NewGemini(http.DefaultClient) + extractor := NewGemini() tests := []struct { name string body string @@ -109,7 +108,7 @@ func TestUsageGeminiNonStreamOptionalFieldsAndArithmetic(t *testing.T) { } func TestUsageGeminiRequiredFieldsAndStrictNumbers(t *testing.T) { - extractor := NewGemini(http.DefaultClient) + extractor := NewGemini() tests := []struct { name string body string @@ -179,7 +178,7 @@ func TestUsageGeminiRequiredFieldsAndStrictNumbers(t *testing.T) { } func TestUsageGeminiOptionalNullAndZeroPresence(t *testing.T) { - extractor := NewGemini(http.DefaultClient) + extractor := NewGemini() tests := []struct { name string body string @@ -240,7 +239,7 @@ func TestUsageGeminiOptionalNullAndZeroPresence(t *testing.T) { } func TestUsageGeminiStreamSnapshotsFinalityAndMalformedPayload(t *testing.T) { - extractor := NewGemini(http.DefaultClient) + extractor := NewGemini() tests := []struct { name string steps []string @@ -368,7 +367,7 @@ func TestUsageGeminiStreamSnapshotsFinalityAndMalformedPayload(t *testing.T) { } func TestUsageGeminiMalformedBodiesAreSanitized(t *testing.T) { - extractor := NewGemini(http.DefaultClient) + extractor := NewGemini() for _, body := range [][]byte{ []byte(`[]`), []byte(`{"usageMetadata":{"promptTokenCount":100}} {}`), diff --git a/internal/dialect/model_list.go b/internal/dialect/model_list.go deleted file mode 100644 index fa590301..00000000 --- a/internal/dialect/model_list.go +++ /dev/null @@ -1,117 +0,0 @@ -package dialect - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "net/http" - "reflect" - "strings" -) - -const ( - maxModelListPages = 100 - maxUniqueModelListEntries = 100_000 - maxModelListPageBytes = int64(8 << 20) - maxUniqueModelListJSONBytes = int64(16 << 20) -) - -type modelListCollector struct { - values []string - seen map[string]struct{} - jsonBytes int64 - maxJSONBytes int64 -} - -func newModelListCollector() *modelListCollector { - return &modelListCollector{ - values: make([]string, 0), - seen: make(map[string]struct{}), - jsonBytes: 2, - maxJSONBytes: maxUniqueModelListJSONBytes, - } -} - -func (collector *modelListCollector) Add(values []string) error { - for _, value := range values { - if _, duplicate := collector.seen[value]; duplicate { - continue - } - if len(collector.values) == maxUniqueModelListEntries { - return fmt.Errorf("model list unique-result limit exceeded") - } - encoded, err := json.Marshal(value) - if err != nil { - return fmt.Errorf("encode model ID: %w", err) - } - additional := int64(len(encoded)) - if len(collector.values) > 0 { - additional++ - } - if additional > collector.maxJSONBytes-collector.jsonBytes { - return fmt.Errorf("model list encoded-size limit exceeded") - } - collector.seen[value] = struct{}{} - collector.values = append(collector.values, value) - collector.jsonBytes += additional - } - return nil -} - -func (collector *modelListCollector) Full() bool { - return len(collector.values) == maxUniqueModelListEntries || - collector.jsonBytes == collector.maxJSONBytes -} - -func (collector *modelListCollector) Result() []string { - result := make([]string, len(collector.values)) - copy(result, collector.values) - return result -} - -func decodeModelListPage(response *http.Response, target any) error { - if response == nil || response.Body == nil || target == nil { - return fmt.Errorf("model-list response is incomplete") - } - targetValue := reflect.ValueOf(target) - if targetValue.Kind() != reflect.Ptr || targetValue.IsNil() { - return fmt.Errorf("model-list target must be a non-nil pointer") - } - - encodings := response.Header.Values("Content-Encoding") - if len(encodings) > 1 { - return fmt.Errorf("model-list response uses unsupported Content-Encoding") - } - if len(encodings) == 1 { - encoding := strings.TrimSpace(encodings[0]) - if encoding != "" && !strings.EqualFold(encoding, "identity") { - return fmt.Errorf("model-list response uses unsupported Content-Encoding") - } - } - if response.ContentLength > maxModelListPageBytes { - return fmt.Errorf("model-list response exceeds page limit") - } - body, err := io.ReadAll(io.LimitReader(response.Body, maxModelListPageBytes+1)) - if err != nil { - return fmt.Errorf("read model-list response: %w", err) - } - if int64(len(body)) > maxModelListPageBytes { - return fmt.Errorf("model-list response exceeds page limit") - } - - staging := reflect.New(targetValue.Elem().Type()) - decoder := json.NewDecoder(bytes.NewReader(body)) - if err := decoder.Decode(staging.Interface()); err != nil { - return err - } - var extra json.RawMessage - if err := decoder.Decode(&extra); err != io.EOF { - if err == nil { - return fmt.Errorf("model-list response contains multiple JSON values") - } - return fmt.Errorf("decode model-list response tail: %w", err) - } - targetValue.Elem().Set(staging.Elem()) - return nil -} diff --git a/internal/dialect/model_list_test.go b/internal/dialect/model_list_test.go deleted file mode 100644 index 80dbdcf5..00000000 --- a/internal/dialect/model_list_test.go +++ /dev/null @@ -1,179 +0,0 @@ -package dialect - -import ( - "fmt" - "io" - "net/http" - "strings" - "testing" -) - -type zeroReader struct{} - -func (zeroReader) Read(buffer []byte) (int, error) { - for index := range buffer { - buffer[index] = 0 - } - return len(buffer), nil -} - -func TestModelListCollectorPreservesFirstOccurrenceOrder(t *testing.T) { - collector := newModelListCollector() - if err := collector.Add([]string{"claude-a", "Shared", "claude-a"}); err != nil { - t.Fatalf("first Add() error = %v", err) - } - if err := collector.Add([]string{"shared", "Shared", "claude-b"}); err != nil { - t.Fatalf("second Add() error = %v", err) - } - - result := collector.Result() - if got, want := fmt.Sprint(result), "[claude-a Shared shared claude-b]"; got != want { - t.Fatalf("Result() = %s, want %s", got, want) - } - result[0] = "mutated" - if got := collector.Result()[0]; got != "claude-a" { - t.Fatalf("Result() exposed collector storage: first value = %q", got) - } -} - -func TestModelListCollectorRejectsMoreThanMaximumUniqueModels(t *testing.T) { - collector := newModelListCollector() - values := make([]string, maxUniqueModelListEntries) - for index := range values { - values[index] = fmt.Sprintf("model-%06d", index) - } - if err := collector.Add(values); err != nil { - t.Fatalf("Add(exact maximum) error = %v", err) - } - if !collector.Full() { - t.Fatal("Full() = false at exact maximum") - } - if err := collector.Add([]string{values[0]}); err != nil { - t.Fatalf("Add(duplicate at maximum) error = %v", err) - } - if err := collector.Add([]string{"one-too-many"}); err == nil { - t.Fatal("Add(100001st unique model) error = nil") - } - if got := len(collector.Result()); got != maxUniqueModelListEntries { - t.Fatalf("Result() length = %d, want %d", got, maxUniqueModelListEntries) - } -} - -func TestModelListCollectorBoundsJSONEncodedBytes(t *testing.T) { - if maxUniqueModelListJSONBytes != 16<<20 { - t.Fatalf("maxUniqueModelListJSONBytes = %d", maxUniqueModelListJSONBytes) - } - collector := newModelListCollector() - collector.maxJSONBytes = 12 // [] + "1234" + , + "x" = 12 - if err := collector.Add([]string{"1234", "x"}); err != nil { - t.Fatalf("Add(exact encoded limit) error = %v", err) - } - if err := collector.Add([]string{"x"}); err != nil { - t.Fatalf("duplicate changed encoded budget: %v", err) - } - if err := collector.Add([]string{"y"}); err == nil { - t.Fatal("Add(limit+1 model) error = nil") - } - - escaped := newModelListCollector() - escaped.maxJSONBytes = 8 // ["\n"] is exactly 6; another item exceeds 8. - if err := escaped.Add([]string{"\n"}); err != nil { - t.Fatal(err) - } - if err := escaped.Add([]string{"z"}); err == nil { - t.Fatal("escaped JSON bytes were counted as raw UTF-8") - } -} - -func TestDecodeModelListPageRequiresBoundedSingleIdentityJSON(t *testing.T) { - if maxModelListPageBytes != 8<<20 { - t.Fatalf("maxModelListPageBytes = %d", maxModelListPageBytes) - } - valid := `{"data":[{"id":"a"}]}` - tests := []struct { - name string - body string - encoding []string - wantError bool - }{ - {name: "missing encoding", body: valid}, - {name: "identity", body: valid + " ", encoding: []string{"identity"}}, - {name: "second JSON", body: valid + `{}`, wantError: true}, - {name: "gzip", body: valid, encoding: []string{"gzip"}, wantError: true}, - {name: "multiple encodings", body: valid, encoding: []string{"identity", "identity"}, wantError: true}, - {name: "blank encoding", body: valid, encoding: []string{" "}}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - response := &http.Response{Header: make(http.Header), Body: io.NopCloser(strings.NewReader(test.body)), ContentLength: int64(len(test.body))} - for _, value := range test.encoding { - response.Header.Add("Content-Encoding", value) - } - target := struct { - Data []struct { - ID string `json:"id"` - } `json:"data"` - Marker string `json:"marker"` - }{Marker: "unchanged"} - err := decodeModelListPage(response, &target) - if test.wantError == (err == nil) { - t.Fatalf("decodeModelListPage() error = %v", err) - } - if test.wantError && target.Marker != "unchanged" { - t.Fatalf("error returned partial target = %#v", target) - } - }) - } -} - -func TestDecodeModelListPageAcceptsExactKnownAndStreamingLimit(t *testing.T) { - body := `{}` + strings.Repeat(" ", int(maxModelListPageBytes)-2) - for _, contentLength := range []int64{maxModelListPageBytes, -1} { - response := &http.Response{ - Header: make(http.Header), - Body: io.NopCloser(strings.NewReader(body)), - ContentLength: contentLength, - } - var target struct{} - if err := decodeModelListPage(response, &target); err != nil { - t.Fatalf("decodeModelListPage(ContentLength=%d) error = %v", contentLength, err) - } - } -} - -func TestDecodeModelListPageRejectsKnownAndStreamingOverflow(t *testing.T) { - tests := []struct { - name string - contentLength int64 - body io.Reader - }{ - { - name: "known length", - contentLength: maxModelListPageBytes + 1, - body: strings.NewReader("{}"), - }, - { - name: "streaming length", - contentLength: -1, - body: io.LimitReader(zeroReader{}, maxModelListPageBytes+1), - }, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - target := struct { - Marker string `json:"marker"` - }{Marker: "unchanged"} - response := &http.Response{ - Header: make(http.Header), - Body: io.NopCloser(test.body), - ContentLength: test.contentLength, - } - if err := decodeModelListPage(response, &target); err == nil { - t.Fatal("decodeModelListPage() error = nil") - } - if target.Marker != "unchanged" { - t.Fatalf("partial target = %#v", target) - } - }) - } -} diff --git a/internal/dialect/model_rewrite_test.go b/internal/dialect/model_rewrite_test.go index 4d9c9b63..4cab1db5 100644 --- a/internal/dialect/model_rewrite_test.go +++ b/internal/dialect/model_rewrite_test.go @@ -18,13 +18,13 @@ func TestModelRewritersDeriveRequestsWithoutMutatingOriginal(t *testing.T) { wantModel string }{ { - name: "openai", rewriter: NewOpenAI(http.DefaultClient), upstream: "provider-openai", + name: "openai", rewriter: NewOpenAI(), upstream: "provider-openai", request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/chat/completions", Header: http.Header{"X-Test": {"one"}}, Body: []byte(`{"model":"public","metadata":{"n":9007199254740993}}`)}, wantPath: "/v1/chat/completions", wantModel: "provider-openai", }, { name: "openai responses", - rewriter: NewOpenAIResponses(http.DefaultClient), + rewriter: NewOpenAIResponses(), upstream: "provider-responses", request: &ParsedRequest{ Method: http.MethodPost, @@ -36,12 +36,12 @@ func TestModelRewritersDeriveRequestsWithoutMutatingOriginal(t *testing.T) { wantModel: "provider-responses", }, { - name: "anthropic", rewriter: NewAnthropic(http.DefaultClient), upstream: "provider-claude", + name: "anthropic", rewriter: NewAnthropic(), upstream: "provider-claude", request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/messages", Body: []byte(`{"model":"public","messages":[]}`)}, wantPath: "/v1/messages", wantModel: "provider-claude", }, { - name: "gemini", rewriter: NewGemini(http.DefaultClient), upstream: "provider-gemini", + name: "gemini", rewriter: NewGemini(), upstream: "provider-gemini", request: &ParsedRequest{Method: http.MethodPost, Path: "/v1beta/models/public:generateContent", Body: []byte(`{}`)}, wantPath: "/v1beta/models/provider-gemini:generateContent", }, @@ -94,12 +94,12 @@ func TestModelRewritersRewriteResponseFields(t *testing.T) { body string path []string }{ - {name: "openai", rewriter: NewOpenAI(http.DefaultClient), body: `{"model":"real","n":9007199254740993}`, path: []string{"model"}}, - {name: "openai responses object", rewriter: NewOpenAIResponses(http.DefaultClient), body: `{"type":"response","model":"real","n":9007199254740993}`, path: []string{"model"}}, - {name: "openai responses event", rewriter: NewOpenAIResponses(http.DefaultClient), body: `{"type":"response.completed","response":{"model":"real"},"n":9007199254740993}`, path: []string{"response", "model"}}, - {name: "anthropic response", rewriter: NewAnthropic(http.DefaultClient), body: `{"type":"message","model":"real"}`, path: []string{"model"}}, - {name: "anthropic message_start", rewriter: NewAnthropic(http.DefaultClient), body: `{"type":"message_start","message":{"model":"real"}}`, path: []string{"message", "model"}}, - {name: "gemini", rewriter: NewGemini(http.DefaultClient), body: `{"modelVersion":"real"}`, path: []string{"modelVersion"}}, + {name: "openai", rewriter: NewOpenAI(), body: `{"model":"real","n":9007199254740993}`, path: []string{"model"}}, + {name: "openai responses object", rewriter: NewOpenAIResponses(), body: `{"type":"response","model":"real","n":9007199254740993}`, path: []string{"model"}}, + {name: "openai responses event", rewriter: NewOpenAIResponses(), body: `{"type":"response.completed","response":{"model":"real"},"n":9007199254740993}`, path: []string{"response", "model"}}, + {name: "anthropic response", rewriter: NewAnthropic(), body: `{"type":"message","model":"real"}`, path: []string{"model"}}, + {name: "anthropic message_start", rewriter: NewAnthropic(), body: `{"type":"message_start","message":{"model":"real"}}`, path: []string{"message", "model"}}, + {name: "gemini", rewriter: NewGemini(), body: `{"modelVersion":"real"}`, path: []string{"modelVersion"}}, } for _, test := range tests { @@ -133,7 +133,7 @@ func TestModelRewritersRewriteResponseFields(t *testing.T) { func TestOpenAIResponsesResponseModelRewriteTreatsNullResponseAsAbsent(t *testing.T) { body := []byte(`{"type":"response.output_text.delta","response":null,"delta":"hello"}`) - got, err := NewOpenAIResponses(http.DefaultClient).RewriteResponseModel(body, "public") + got, err := NewOpenAIResponses().RewriteResponseModel(body, "public") if err != nil { t.Fatalf("RewriteResponseModel() error = %v", err) } @@ -148,10 +148,10 @@ func TestModelRewritersRejectInvalidTargetModels(t *testing.T) { value ModelRewriter req *ParsedRequest }{ - {name: "openai", value: NewOpenAI(http.DefaultClient), req: &ParsedRequest{Body: []byte(`{"model":"public"}`)}}, - {name: "openai responses", value: NewOpenAIResponses(http.DefaultClient), req: &ParsedRequest{Body: []byte(`{"model":"public"}`)}}, - {name: "anthropic", value: NewAnthropic(http.DefaultClient), req: &ParsedRequest{Body: []byte(`{"model":"public"}`)}}, - {name: "gemini", value: NewGemini(http.DefaultClient), req: &ParsedRequest{Path: "/v1beta/models/public:generateContent"}}, + {name: "openai", value: NewOpenAI(), req: &ParsedRequest{Body: []byte(`{"model":"public"}`)}}, + {name: "openai responses", value: NewOpenAIResponses(), req: &ParsedRequest{Body: []byte(`{"model":"public"}`)}}, + {name: "anthropic", value: NewAnthropic(), req: &ParsedRequest{Body: []byte(`{"model":"public"}`)}}, + {name: "gemini", value: NewGemini(), req: &ParsedRequest{Path: "/v1beta/models/public:generateContent"}}, } for _, test := range rewriters { t.Run(test.name, func(t *testing.T) { @@ -165,7 +165,7 @@ func TestModelRewritersRejectInvalidTargetModels(t *testing.T) { } }) } - gemini := NewGemini(http.DefaultClient) + gemini := NewGemini() if _, err := gemini.RewriteRequestModel(&ParsedRequest{Path: "/v1beta/models/public:generateContent"}, "models/provider"); err == nil { t.Fatal("Gemini RewriteRequestModel() accepted slash in model") } @@ -181,10 +181,10 @@ func TestModelRewritersRequireRequestModelAndCloneAbsentResponse(t *testing.T) { request *ParsedRequest response []byte }{ - {name: "openai", rewriter: NewOpenAI(http.DefaultClient), request: &ParsedRequest{Body: []byte(`{}`)}, response: []byte(`{"id":"one"}`)}, - {name: "openai responses", rewriter: NewOpenAIResponses(http.DefaultClient), request: &ParsedRequest{Body: []byte(`{}`)}, response: []byte(`{"id":"resp_one"}`)}, - {name: "anthropic", rewriter: NewAnthropic(http.DefaultClient), request: &ParsedRequest{Body: []byte(`{}`)}, response: []byte(`{"type":"content_block_delta"}`)}, - {name: "gemini", rewriter: NewGemini(http.DefaultClient), request: &ParsedRequest{Path: "/v1beta/models/:generateContent"}, response: []byte(`{"candidates":[]}`)}, + {name: "openai", rewriter: NewOpenAI(), request: &ParsedRequest{Body: []byte(`{}`)}, response: []byte(`{"id":"one"}`)}, + {name: "openai responses", rewriter: NewOpenAIResponses(), request: &ParsedRequest{Body: []byte(`{}`)}, response: []byte(`{"id":"resp_one"}`)}, + {name: "anthropic", rewriter: NewAnthropic(), request: &ParsedRequest{Body: []byte(`{}`)}, response: []byte(`{"type":"content_block_delta"}`)}, + {name: "gemini", rewriter: NewGemini(), request: &ParsedRequest{Path: "/v1beta/models/:generateContent"}, response: []byte(`{"candidates":[]}`)}, } { t.Run(test.name, func(t *testing.T) { if _, err := test.rewriter.RewriteRequestModel(test.request, "provider"); err == nil { diff --git a/internal/dialect/openai.go b/internal/dialect/openai.go index eaa8a5df..4b35fe94 100644 --- a/internal/dialect/openai.go +++ b/internal/dialect/openai.go @@ -1,146 +1,23 @@ package dialect import ( - "context" "fmt" - "net/http" - "gpt-load/internal/health" - platformheader "gpt-load/internal/platform/httpheader" "gpt-load/internal/protocol" - "gpt-load/internal/state" ) -const ( - openAICompletionsPath = "/v1/chat/completions" - openAICompletionsResourcePath = "/chat/completions" - openAIModelsResourcePath = "/models" -) - -var openAIFailureMarkers = failureMarkers{ - rateLimited: []string{"rate_limit", "rate limit", "insufficient_quota", "quota_exceeded", "quota exceeded"}, - modelUnavailable: []string{"model_not_found", "model not found", "model_not_supported", "model not supported", "no access to model"}, - invalidKey: []string{"invalid_api_key", "invalid api key", "incorrect api key", "api_key_disabled", "api key disabled", "account_deactivated", "api_key_banned", "api key banned"}, - upstreamHost: []string{"overloaded_error", "server_overloaded"}, -} - -type OpenAI struct { - client *http.Client -} +type OpenAI struct{} var _ Dialect = (*OpenAI)(nil) -func NewOpenAI(client *http.Client) *OpenAI { - return &OpenAI{client: client} +func NewOpenAI() *OpenAI { + return &OpenAI{} } func (d *OpenAI) Protocol() protocol.Protocol { return protocol.OpenAICompletions } -func (d *OpenAI) InjectCredential(headers http.Header, apiKey string) { - if headers == nil { - return - } - headers.Set("Authorization", "Bearer "+apiKey) -} - -func (d *OpenAI) CredentialHeaderNames() []string { - return []string{"Authorization"} -} - -func (d *OpenAI) BuildUpstreamURL(base string, req *ParsedRequest) (string, error) { - if req == nil { - return "", fmt.Errorf("parsed request is required") - } - if req.Path != openAICompletionsPath { - return "", fmt.Errorf("invalid OpenAI Completions request path") - } - return resolveUpstreamURL(base, openAICompletionsResourcePath, req.RawQuery) -} - -func (d *OpenAI) ListModels( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, -) ([]string, error) { - requestURL, err := resolveUpstreamURL(baseURL, openAIModelsResourcePath, "") - if err != nil { - return nil, fmt.Errorf("build OpenAI model-list URL: %w", err) - } - - req, err := http.NewRequestWithContext( - ctx, - http.MethodGet, - requestURL, - nil, - ) - if err != nil { - return nil, fmt.Errorf("create OpenAI model-list request: %w", err) - } - ApplyCredential(d, req.Header, apiKey, rules) - platformheader.NormalizeUpstreamRequestRepresentation(req, 0) - - resp, err := d.client.Do(req) - if err != nil { - if contextErr := ctx.Err(); contextErr != nil { - return nil, fmt.Errorf("request OpenAI model list: %w", contextErr) - } - return nil, fmt.Errorf("request OpenAI model list failed") - } - defer resp.Body.Close() - if resp.StatusCode < http.StatusOK || - resp.StatusCode >= http.StatusMultipleChoices { - return nil, fmt.Errorf( - "request OpenAI model list: upstream status %d", - resp.StatusCode, - ) - } - - var payload struct { - Data []struct { - ID string `json:"id"` - } `json:"data"` - } - if err := decodeModelListPage(resp, &payload); err != nil { - return nil, fmt.Errorf("decode OpenAI model list: %w", err) - } - - pageModels := make([]string, 0, len(payload.Data)) - for _, model := range payload.Data { - pageModels = append(pageModels, model.ID) - } - collector := newModelListCollector() - if err := collector.Add(pageModels); err != nil { - return nil, err - } - return collector.Result(), nil -} - -func (d *OpenAI) Probe( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, - validationModel string, -) error { - if err := validateProbeModel(validationModel); err != nil { - return err - } - requestURL, err := d.BuildUpstreamURL(baseURL, &ParsedRequest{Method: http.MethodPost, Path: openAICompletionsPath}) - if err != nil { - return fmt.Errorf("build %s probe URL failed", d.Protocol()) - } - return executeProbe(ctx, d.client, d, requestURL, apiKey, rules, struct { - Model string `json:"model"` - Messages []probeMessage `json:"messages"` - MaxTokens int `json:"max_tokens"` - }{ - Model: validationModel, - Messages: []probeMessage{{Role: "user", Content: "ping"}}, - MaxTokens: 1, - }) -} - func (d *OpenAI) InspectRequest(req *ParsedRequest) (RequestMetadata, error) { if req == nil { return RequestMetadata{}, fmt.Errorf("parsed request is required") @@ -153,13 +30,10 @@ func (d *OpenAI) InspectRequest(req *ParsedRequest) (RequestMetadata, error) { metadata.ObserveUsage = true metadata.UsageDiagnostics = openAIRequestPricingDiagnostics(req.Body) metadata.Reasoning = inspectOpenAICompletionsReasoning(req.Body) + metadata.Operation, metadata.RequiredFeatures = chatExecutionMetadata( + req.Body, + metadata.Stream, + metadata.Reasoning, + ) return metadata, nil } - -func (d *OpenAI) ClassifyStatus(status int, body []byte) health.FailureCategory { - return classifyStatusWithMarkers(status, body, openAIFailureMarkers) -} - -func (d *OpenAI) ClassifyProviderError(body []byte) health.FailureCategory { - return classifyProviderErrorWithMarkers(body, openAIFailureMarkers) -} diff --git a/internal/dialect/openai_responses.go b/internal/dialect/openai_responses.go index f5a92671..efd88732 100644 --- a/internal/dialect/openai_responses.go +++ b/internal/dialect/openai_responses.go @@ -2,27 +2,20 @@ package dialect import ( "bytes" - "context" "encoding/json" "fmt" "net/http" "net/url" - "strings" - "gpt-load/internal/health" "gpt-load/internal/protocol" - "gpt-load/internal/state" ) const ( - openAIResponsesPath = "/v1/responses" - openAIResponsesCompactPath = "/v1/responses/compact" - openAIResponsesResourceRoot = "/responses" + openAIResponsesPath = "/v1/responses" + openAIResponsesCompactPath = "/v1/responses/compact" ) -type OpenAIResponses struct { - openAI *OpenAI -} +type OpenAIResponses struct{} var ( _ Dialect = (*OpenAIResponses)(nil) @@ -30,8 +23,8 @@ var ( _ StreamEventClassifier = (*OpenAIResponses)(nil) ) -func NewOpenAIResponses(client *http.Client) *OpenAIResponses { - return &OpenAIResponses{openAI: NewOpenAI(client)} +func NewOpenAIResponses() *OpenAIResponses { + return &OpenAIResponses{} } func (d *OpenAIResponses) Protocol() protocol.Protocol { @@ -66,6 +59,11 @@ func (d *OpenAIResponses) InspectRequest(req *ParsedRequest) (RequestMetadata, e metadata.UsageDiagnostics = openAIRequestPricingDiagnostics(req.Body) metadata.Reasoning = inspectOpenAIResponsesReasoning(req.Body) } + metadata.Operation, metadata.RequiredFeatures = responsesExecutionMetadata( + req, + metadata.Stream, + metadata.Reasoning, + ) return metadata, nil } @@ -91,104 +89,6 @@ func inspectResponsesStreamQuery(rawQuery string) (bool, bool, error) { } } -func (d *OpenAIResponses) BuildUpstreamURL(base string, req *ParsedRequest) (string, error) { - if req == nil { - return "", fmt.Errorf("parsed request is required") - } - resourcePath, err := openAIResponsesResourcePath(req.Path) - if err != nil { - return "", err - } - return resolveUpstreamURL(base, resourcePath, req.RawQuery) -} - -func openAIResponsesResourcePath(path string) (string, error) { - switch { - case path == openAIResponsesPath: - return openAIResponsesResourceRoot, nil - case strings.HasPrefix(path, openAIResponsesPath+"/"): - return strings.TrimPrefix(path, "/v1"), nil - default: - return "", fmt.Errorf("invalid OpenAI Responses request path") - } -} - -func (d *OpenAIResponses) InjectCredential(headers http.Header, apiKey string) { - d.openAI.InjectCredential(headers, apiKey) -} - -func (d *OpenAIResponses) CredentialHeaderNames() []string { - return d.openAI.CredentialHeaderNames() -} - -func (d *OpenAIResponses) ListModels( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, -) ([]string, error) { - return d.openAI.ListModels(ctx, baseURL, apiKey, rules) -} - -func (d *OpenAIResponses) Probe( - ctx context.Context, - baseURL, apiKey string, - rules state.HeaderRules, - validationModel string, -) error { - if err := validateProbeModel(validationModel); err != nil { - return err - } - requestURL, err := d.BuildUpstreamURL(baseURL, &ParsedRequest{Method: http.MethodPost, Path: openAIResponsesPath}) - if err != nil { - return fmt.Errorf("build %s probe URL failed", d.Protocol()) - } - return executeProbe( - ctx, - d.openAI.client, - d, - requestURL, - apiKey, - rules, - struct { - Model string `json:"model"` - Input string `json:"input"` - MaxOutputTokens int `json:"max_output_tokens"` - Store bool `json:"store"` - }{ - Model: validationModel, - Input: "ping", - MaxOutputTokens: 16, - Store: false, - }, - ) -} - -func (d *OpenAIResponses) ClassifyStatus( - status int, - body []byte, -) health.FailureCategory { - if status == http.StatusNotFound { - lowered := strings.ToLower(string(body)) - switch { - case containsFailureMarker(lowered, openAIFailureMarkers.rateLimited): - return health.FailureCategoryRateLimited - case containsFailureMarker(lowered, openAIFailureMarkers.modelUnavailable): - return health.FailureCategoryModelUnavailable - case containsFailureMarker(lowered, openAIFailureMarkers.invalidKey): - return health.FailureCategoryInvalidKey - case containsFailureMarker(lowered, openAIFailureMarkers.upstreamHost): - return health.FailureCategoryUpstreamHostError - default: - return health.FailureCategoryClientError - } - } - return d.openAI.ClassifyStatus(status, body) -} - -func (d *OpenAIResponses) ClassifyProviderError(body []byte) health.FailureCategory { - return d.openAI.ClassifyProviderError(body) -} - func (*OpenAIResponses) RequiresTerminalEvent() bool { return true } diff --git a/internal/dialect/openai_responses_stream_test.go b/internal/dialect/openai_responses_stream_test.go index 202ecb77..36f35179 100644 --- a/internal/dialect/openai_responses_stream_test.go +++ b/internal/dialect/openai_responses_stream_test.go @@ -1,14 +1,13 @@ package dialect import ( - "net/http" "testing" ) func TestOpenAIResponsesClassifiesStreamEvents(t *testing.T) { t.Parallel() - classifier, ok := any(NewOpenAIResponses(http.DefaultClient)).(StreamEventClassifier) + classifier, ok := any(NewOpenAIResponses()).(StreamEventClassifier) if !ok { t.Fatal("OpenAI Responses does not expose StreamEventClassifier") } @@ -87,7 +86,7 @@ func TestOpenAIResponsesClassifiesStreamEvents(t *testing.T) { func TestOpenAIResponsesRejectsConflictingStreamEventNames(t *testing.T) { t.Parallel() - classifier := any(NewOpenAIResponses(http.DefaultClient)).(StreamEventClassifier) + classifier := any(NewOpenAIResponses()).(StreamEventClassifier) _, err := classifier.ClassifyStreamEvent(StreamEvent{ Name: "response.completed", Payload: []byte(`{"type":"response.failed","response":{"error":{"code":"server_error"}}}`), diff --git a/internal/dialect/openai_responses_test.go b/internal/dialect/openai_responses_test.go index b6a4f96f..df2db16f 100644 --- a/internal/dialect/openai_responses_test.go +++ b/internal/dialect/openai_responses_test.go @@ -1,24 +1,15 @@ package dialect import ( - "context" - "encoding/json" - "io" "net/http" - "net/http/httptest" "testing" - "time" - "gpt-load/internal/health" "gpt-load/internal/protocol" - "gpt-load/internal/state" "gpt-load/internal/usage" ) func TestOpenAIResponsesProtocolAndRequestMetadata(t *testing.T) { - t.Parallel() - - selected := NewOpenAIResponses(http.DefaultClient) + selected := NewOpenAIResponses() if got := selected.Protocol(); got != protocol.OpenAIResponses { t.Fatalf("Protocol() = %q, want %q", got, protocol.OpenAIResponses) } @@ -26,351 +17,56 @@ func TestOpenAIResponsesProtocolAndRequestMetadata(t *testing.T) { tests := []struct { name string request *ParsedRequest - wantModel *string + wantModel string wantStream bool wantObserve bool - wantError bool + wantErr bool }{ - { - name: "empty create body", - request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses"}, - wantObserve: true, - }, - { - name: "create with model and stream", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte(`{"model":"gpt-5","stream":true,"input":"ping"}`), - }, - wantModel: stringPointerForTest("gpt-5"), - wantStream: true, - wantObserve: true, - }, - { - name: "compact observes usage", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses/compact", - Body: []byte(`{"model":"gpt-5","input":"ping"}`), - }, - wantModel: stringPointerForTest("gpt-5"), - wantObserve: true, - }, - { - name: "input tokens does not observe usage", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses/input_tokens", - Body: []byte(`{"model":"gpt-5","input":"ping"}`), - }, - wantModel: stringPointerForTest("gpt-5"), - }, - { - name: "retrieve stream query true", - request: &ParsedRequest{ - Method: http.MethodGet, - Path: "/v1/responses/resp_123", - RawQuery: "include=output&stream=true", - }, - wantStream: true, - }, - { - name: "retrieve stream query false", - request: &ParsedRequest{ - Method: http.MethodGet, - Path: "/v1/responses/resp_123", - RawQuery: "stream=false", - }, - }, - { - name: "retrieve percent encoded stream query", - request: &ParsedRequest{ - Method: http.MethodGet, - Path: "/v1/responses/resp_123", - RawQuery: "%73tream=%74rue", - }, - wantStream: true, - }, - { - name: "nil request", - request: nil, - wantError: true, - }, - { - name: "non object", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte(`[]`), - }, - wantError: true, - }, - { - name: "multiple values", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte(`{} {}`), - }, - wantError: true, - }, - { - name: "invalid utf8", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte{0xff}, - }, - wantError: true, - }, - { - name: "blank model", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte(`{"model":""}`), - }, - wantError: true, - }, - { - name: "model boundary whitespace", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte(`{"model":" gpt-5"}`), - }, - wantError: true, - }, - { - name: "model case alias", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte(`{"Model":"gpt-5"}`), - }, - wantError: true, - }, - { - name: "duplicate model", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte(`{"model":"gpt-5","model":"gpt-5-mini"}`), - }, - wantError: true, - }, - { - name: "stream must be bool", - request: &ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte(`{"stream":"true"}`), - }, - wantError: true, - }, - { - name: "duplicate stream query", - request: &ParsedRequest{ - Method: http.MethodGet, - Path: "/v1/responses/resp_123", - RawQuery: "stream=true&stream=true", - }, - wantError: true, - }, - { - name: "mixed encoded duplicate stream query", - request: &ParsedRequest{ - Method: http.MethodGet, - Path: "/v1/responses/resp_123", - RawQuery: "stream=true&%73tream=false", - }, - wantError: true, - }, - { - name: "invalid stream query value", - request: &ParsedRequest{ - Method: http.MethodGet, - Path: "/v1/responses/resp_123", - RawQuery: "stream=1", - }, - wantError: true, - }, - { - name: "malformed stream query encoding", - request: &ParsedRequest{ - Method: http.MethodGet, - Path: "/v1/responses/resp_123", - RawQuery: "stream=%ZZ", - }, - wantError: true, - }, + {name: "empty create", request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses"}, wantObserve: true}, + {name: "create stream", request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses", Body: []byte(`{"model":"gpt-5","stream":true}`)}, wantModel: "gpt-5", wantStream: true, wantObserve: true}, + {name: "compact", request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses/compact", Body: []byte(`{"model":"gpt-5"}`)}, wantModel: "gpt-5", wantObserve: true}, + {name: "input tokens", request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses/input_tokens", Body: []byte(`{"model":"gpt-5"}`)}, wantModel: "gpt-5"}, + {name: "retrieve stream", request: &ParsedRequest{Method: http.MethodGet, Path: "/v1/responses/resp_123", RawQuery: "stream=true"}, wantStream: true}, + {name: "encoded stream", request: &ParsedRequest{Method: http.MethodGet, Path: "/v1/responses/resp_123", RawQuery: "%73tream=%74rue"}, wantStream: true}, + {name: "nil", wantErr: true}, + {name: "non object", request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses", Body: []byte(`[]`)}, wantErr: true}, + {name: "blank model", request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses", Body: []byte(`{"model":""}`)}, wantErr: true}, + {name: "duplicate stream query", request: &ParsedRequest{Method: http.MethodGet, Path: "/v1/responses/resp_123", RawQuery: "stream=true&stream=false"}, wantErr: true}, + {name: "invalid stream query", request: &ParsedRequest{Method: http.MethodGet, Path: "/v1/responses/resp_123", RawQuery: "stream=1"}, wantErr: true}, } - for _, test := range tests { t.Run(test.name, func(t *testing.T) { - t.Parallel() - got, err := selected.InspectRequest(test.request) - if test.wantError { + metadata, err := selected.InspectRequest(test.request) + if test.wantErr { if err == nil { - t.Fatalf("InspectRequest() = %#v, nil; want error", got) + t.Fatalf("InspectRequest() = %#v, nil", metadata) } return } - if err != nil { - t.Fatalf("InspectRequest() error = %v", err) + if err != nil || metadata.Stream != test.wantStream || metadata.ObserveUsage != test.wantObserve { + t.Fatalf("InspectRequest() = %#v, %v", metadata, err) } - if !sameOptionalString(got.Model, test.wantModel) || - got.Stream != test.wantStream || - got.ObserveUsage != test.wantObserve { - t.Fatalf( - "InspectRequest() = %#v, want model=%v stream=%t observe=%t", - got, - test.wantModel, - test.wantStream, - test.wantObserve, - ) + if test.wantModel == "" { + if metadata.Model != nil { + t.Fatalf("model = %v, want nil", metadata.Model) + } + } else if metadata.Model == nil || *metadata.Model != test.wantModel { + t.Fatalf("model = %v, want %q", metadata.Model, test.wantModel) } }) } } func TestOpenAIResponsesRequestMarksUnsupportedPricingModes(t *testing.T) { - t.Parallel() - selected := NewOpenAIResponses(http.DefaultClient) + selected := NewOpenAIResponses() for _, body := range []string{ `{"model":"gpt-5","service_tier":"priority"}`, `{"model":"gpt-5","speed":"fast"}`, `{"model":"gpt-5","reasoning":{"mode":"pro"}}`, } { - metadata, err := selected.InspectRequest(&ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/responses", - Body: []byte(body), - }) - if err != nil { - t.Fatalf("InspectRequest(%s) error = %v", body, err) - } - if !metadata.UsageDiagnostics.Has(usage.DiagnosticUnsupportedBillableDetail) { - t.Fatalf("InspectRequest(%s) diagnostics = %#v, want unsupported billable detail", body, metadata.UsageDiagnostics) + metadata, err := selected.InspectRequest(&ParsedRequest{Method: http.MethodPost, Path: "/v1/responses", Body: []byte(body)}) + if err != nil || !metadata.UsageDiagnostics.Has(usage.DiagnosticUnsupportedBillableDetail) { + t.Fatalf("InspectRequest(%s) = %#v, %v", body, metadata, err) } } } - -func TestOpenAIResponsesResource404DoesNotBlameKey(t *testing.T) { - t.Parallel() - - selected := NewOpenAIResponses(http.DefaultClient) - resourceBody := []byte( - `{"error":{"type":"invalid_request_error","message":"No response found with id resp_123"}}`, - ) - if got := selected.ClassifyStatus(http.StatusNotFound, resourceBody); got != health.FailureCategoryClientError { - t.Fatalf("resource 404 category = %v, want client error", got) - } - decision := health.Judge(selected, health.Attempt{ - StatusCode: http.StatusNotFound, - Body: resourceBody, - Now: time.Unix(1, 0), - }) - if decision.Category != health.FailureCategoryClientError || - decision.Action != health.ActionTerminate { - t.Fatalf("resource 404 decision = %#v", decision) - } - - modelBody := []byte(`{"error":{"code":"model_not_found"}}`) - if got := selected.ClassifyStatus(http.StatusNotFound, modelBody); got != health.FailureCategoryModelUnavailable { - t.Fatalf("model 404 category = %v, want model unavailable", got) - } -} - -func TestOpenAIResponsesProbeUsesMinimumValidRequest(t *testing.T) { - t.Parallel() - - var received map[string]any - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - request *http.Request, - ) { - if request.Method != http.MethodPost || - request.URL.Path != "/v1/responses" || - request.Header.Get("Authorization") != "Bearer sk-responses" { - t.Fatalf( - "request method/path/auth = %s/%s/%q", - request.Method, - request.URL.Path, - request.Header.Get("Authorization"), - ) - } - body, err := io.ReadAll(request.Body) - if err != nil { - t.Fatalf("read request body: %v", err) - } - if err := json.Unmarshal(body, &received); err != nil { - t.Fatalf("decode request body: %v", err) - } - writer.WriteHeader(http.StatusNoContent) - })) - defer upstream.Close() - - err := NewOpenAIResponses(upstream.Client()).Probe( - context.Background(), - upstream.URL+"/v1", - "sk-responses", - state.HeaderRules{}, - "gpt-5", - ) - if err != nil { - t.Fatalf("Probe() error = %v", err) - } - want := map[string]any{ - "model": "gpt-5", - "input": "ping", - "max_output_tokens": float64(16), - "store": false, - } - if len(received) != len(want) { - t.Fatalf("probe body = %#v, want %#v", received, want) - } - for name, value := range want { - if received[name] != value { - t.Fatalf("probe body[%q] = %#v, want %#v", name, received[name], value) - } - } -} - -func TestOpenAIResponsesUsesOpenAIURLAndCredentialRules(t *testing.T) { - t.Parallel() - - selected := NewOpenAIResponses(http.DefaultClient) - headers := make(http.Header) - selected.InjectCredential(headers, "sk-provider") - if got := headers.Get("Authorization"); got != "Bearer sk-provider" { - t.Fatalf("Authorization = %q, want Bearer provider key", got) - } - - got, err := selected.BuildUpstreamURL( - "https://api.example.com/compatible/v1?api-version=2026-01-01", - &ParsedRequest{ - Path: "/v1/responses/resp_123/input_items", - RawQuery: "limit=20&after=item_1", - }, - ) - if err != nil { - t.Fatalf("BuildUpstreamURL() error = %v", err) - } - const want = "https://api.example.com/compatible/v1/responses/resp_123/input_items?api-version=2026-01-01&limit=20&after=item_1" - if got != want { - t.Fatalf("BuildUpstreamURL() = %q, want %q", got, want) - } -} - -func stringPointerForTest(value string) *string { - return &value -} - -func sameOptionalString(left, right *string) bool { - if left == nil || right == nil { - return left == nil && right == nil - } - return *left == *right -} diff --git a/internal/dialect/openai_responses_usage_test.go b/internal/dialect/openai_responses_usage_test.go index 9744658f..0501c1e8 100644 --- a/internal/dialect/openai_responses_usage_test.go +++ b/internal/dialect/openai_responses_usage_test.go @@ -1,7 +1,6 @@ package dialect import ( - "net/http" "testing" "gpt-load/internal/usage" @@ -10,7 +9,7 @@ import ( func TestOpenAIResponsesUsageNonStreaming(t *testing.T) { t.Parallel() - extractor, ok := any(NewOpenAIResponses(http.DefaultClient)).(UsageExtractor) + extractor, ok := any(NewOpenAIResponses()).(UsageExtractor) if !ok { t.Fatal("OpenAI Responses does not expose UsageExtractor capability") } @@ -129,7 +128,7 @@ func TestOpenAIResponsesUsageNonStreaming(t *testing.T) { func TestOpenAIResponsesUsageStreamingFinality(t *testing.T) { t.Parallel() - extractor := NewOpenAIResponses(http.DefaultClient) + extractor := NewOpenAIResponses() tests := []struct { name string steps []string @@ -213,7 +212,7 @@ func TestOpenAIResponsesUsageStreamingFinality(t *testing.T) { func TestOpenAIResponsesUsageUsesExplicitSSEEventForFinality(t *testing.T) { t.Parallel() - stream := NewOpenAIResponses(http.DefaultClient).NewUsageStreamExtractor() + stream := NewOpenAIResponses().NewUsageStreamExtractor() observer, ok := stream.(UsageStreamEventObserver) if !ok { t.Fatal("OpenAI Responses usage stream does not observe SSE event names") diff --git a/internal/dialect/openai_test.go b/internal/dialect/openai_test.go index f45f2145..f8b7d571 100644 --- a/internal/dialect/openai_test.go +++ b/internal/dialect/openai_test.go @@ -1,157 +1,28 @@ package dialect import ( - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "reflect" - "strings" "testing" - "time" - "gpt-load/internal/health" "gpt-load/internal/protocol" - "gpt-load/internal/state" - "gpt-load/internal/testutil/fakeupstream" "gpt-load/internal/usage" ) -func TestParsedRequestCarriesForwardingInputs(t *testing.T) { - request := ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/chat/completions", - RawQuery: "trace=true", - Header: http.Header{"X-Test": {"value"}}, - Body: []byte(`{"model":"gpt-4o"}`), - } - - if request.Method != http.MethodPost || - request.Path != "/v1/chat/completions" || - request.RawQuery != "trace=true" || - request.Header.Get("X-Test") != "value" || - string(request.Body) != `{"model":"gpt-4o"}` { - t.Fatalf("ParsedRequest lost forwarding input: %#v", request) +func TestParsedRequestCarriesIngressInputs(t *testing.T) { + request := ParsedRequest{Method: "POST", Path: "/v1/chat/completions", RawQuery: "trace=true", Body: []byte(`{"model":"gpt-4o"}`)} + if request.Method != "POST" || request.Path != "/v1/chat/completions" || + request.RawQuery != "trace=true" || len(request.Body) == 0 { + t.Fatalf("ParsedRequest = %#v", request) } } func TestOpenAIProtocol(t *testing.T) { - dialect := NewOpenAI(http.DefaultClient) - - if got := dialect.Protocol(); got != protocol.OpenAICompletions { - t.Fatalf("OpenAI.Protocol() = %q, want %q", got, protocol.OpenAICompletions) - } -} - -func TestOpenAIInjectCredential(t *testing.T) { - dialect := NewOpenAI(http.DefaultClient) - headers := make(http.Header) - - dialect.InjectCredential(headers, "sk-upstream") - - if got := headers.Get("Authorization"); got != "Bearer sk-upstream" { - t.Fatalf("Authorization = %q, want Bearer credential", got) - } -} - -func TestOpenAIBuildUpstreamURL(t *testing.T) { - dialect := NewOpenAI(http.DefaultClient) - tests := []struct { - name string - base string - request *ParsedRequest - want string - wantErr bool - }{ - { - name: "complete API prefix", - base: "https://api.example.com/v1", - request: &ParsedRequest{ - Path: "/v1/chat/completions", - RawQuery: "trace=true", - }, - want: "https://api.example.com/v1/chat/completions?trace=true", - }, - { - name: "base path prefix", - base: "https://api.example.com/compatible/v1/", - request: &ParsedRequest{ - Path: "/v1/chat/completions", - }, - want: "https://api.example.com/compatible/v1/chat/completions", - }, - { - name: "base query is preserved", - base: "https://api.example.com/v1?api-version=2024-10-01", - request: &ParsedRequest{ - Path: "/v1/chat/completions", - }, - want: "https://api.example.com/v1/chat/completions?api-version=2024-10-01", - }, - { - name: "base and request queries are combined", - base: "https://api.example.com/v1?api-version=2024-10-01", - request: &ParsedRequest{ - Path: "/v1/chat/completions", - RawQuery: "trace=true", - }, - want: "https://api.example.com/v1/chat/completions?api-version=2024-10-01&trace=true", - }, - { - name: "nil request", - base: "https://api.example.com", - request: nil, - wantErr: true, - }, - { - name: "relative request path", - base: "https://api.example.com", - request: &ParsedRequest{ - Path: "v1/chat/completions", - }, - wantErr: true, - }, - { - name: "relative base", - base: "api.example.com/v1", - request: &ParsedRequest{ - Path: "/v1/chat/completions", - }, - wantErr: true, - }, - { - name: "unsupported base scheme", - base: "ftp://api.example.com/v1", - request: &ParsedRequest{ - Path: "/v1/chat/completions", - }, - wantErr: true, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got, err := dialect.BuildUpstreamURL(tt.base, tt.request) - if tt.wantErr { - if err == nil { - t.Fatalf("BuildUpstreamURL() = %q, want error", got) - } - return - } - if err != nil { - t.Fatalf("BuildUpstreamURL() error = %v", err) - } - if got != tt.want { - t.Fatalf("BuildUpstreamURL() = %q, want %q", got, tt.want) - } - }) + if got := NewOpenAI().Protocol(); got != protocol.OpenAICompletions { + t.Fatalf("Protocol() = %q, want %q", got, protocol.OpenAICompletions) } } func TestOpenAIInspectRequest(t *testing.T) { - dialect := NewOpenAI(http.DefaultClient) + selected := NewOpenAI() tests := []struct { name string request *ParsedRequest @@ -159,85 +30,42 @@ func TestOpenAIInspectRequest(t *testing.T) { wantStream bool wantErr bool }{ - { - name: "non-stream", - request: &ParsedRequest{Body: []byte(`{"model":"gpt-4o","messages":[]}`)}, - wantModel: "gpt-4o", - wantStream: false, - }, - { - name: "stream true", - request: &ParsedRequest{Body: []byte(`{"model":"gpt-4o-mini","stream":true}`)}, - wantModel: "gpt-4o-mini", - wantStream: true, - }, - { - name: "nil request", - request: nil, - wantErr: true, - }, - { - name: "invalid JSON", - request: &ParsedRequest{Body: []byte(`{"model":`)}, - wantErr: true, - }, - { - name: "missing model", - request: &ParsedRequest{Body: []byte(`{"stream":true}`)}, - wantErr: true, - }, - { - name: "blank model", - request: &ParsedRequest{Body: []byte(`{"model":" "}`)}, - wantErr: true, - }, - { - name: "model with boundary whitespace", - request: &ParsedRequest{Body: []byte(`{"model":" gpt-4o "}`)}, - wantErr: true, - }, + {name: "non-stream", request: &ParsedRequest{Body: []byte(`{"model":"gpt-4o","messages":[]}`)}, wantModel: "gpt-4o"}, + {name: "stream", request: &ParsedRequest{Body: []byte(`{"model":"gpt-4o-mini","stream":true}`)}, wantModel: "gpt-4o-mini", wantStream: true}, + {name: "nil", wantErr: true}, + {name: "invalid JSON", request: &ParsedRequest{Body: []byte(`{"model":`)}, wantErr: true}, + {name: "missing model", request: &ParsedRequest{Body: []byte(`{"stream":true}`)}, wantErr: true}, + {name: "blank model", request: &ParsedRequest{Body: []byte(`{"model":" "}`)}, wantErr: true}, + {name: "model boundary whitespace", request: &ParsedRequest{Body: []byte(`{"model":" gpt-4o "}`)}, wantErr: true}, } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - metadata, err := dialect.InspectRequest(tt.request) - if tt.wantErr { + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + metadata, err := selected.InspectRequest(test.request) + if test.wantErr { if err == nil { - t.Fatalf("InspectRequest() = (%#v, nil), want error", metadata) + t.Fatalf("InspectRequest() = %#v, nil", metadata) } return } - if err != nil { - t.Fatalf("InspectRequest() error = %v", err) - } - if metadata.Model == nil || *metadata.Model != tt.wantModel || - metadata.Stream != tt.wantStream || !metadata.ObserveUsage { - t.Fatalf( - "InspectRequest() = %#v, want model=%q stream=%t observe=true", - metadata, - tt.wantModel, - tt.wantStream, - ) + if err != nil || metadata.Model == nil || *metadata.Model != test.wantModel || + metadata.Stream != test.wantStream || !metadata.ObserveUsage { + t.Fatalf("InspectRequest() = %#v, %v", metadata, err) } }) } } func TestOpenAIInspectRequestMarksUnsupportedPricingModes(t *testing.T) { - t.Parallel() - dialect := NewOpenAI(http.DefaultClient) + selected := NewOpenAI() for _, body := range []string{ `{"model":"gpt-5","service_tier":"priority"}`, `{"model":"gpt-5","service_tier":"flex"}`, `{"model":"gpt-5","speed":"fast"}`, `{"model":"gpt-5","reasoning":{"mode":"pro"}}`, } { - metadata, err := dialect.InspectRequest(&ParsedRequest{Body: []byte(body)}) - if err != nil { - t.Fatalf("InspectRequest(%s) error = %v", body, err) - } - if !metadata.UsageDiagnostics.Has(usage.DiagnosticUnsupportedBillableDetail) { - t.Fatalf("InspectRequest(%s) diagnostics = %#v, want unsupported billable detail", body, metadata.UsageDiagnostics) + metadata, err := selected.InspectRequest(&ParsedRequest{Body: []byte(body)}) + if err != nil || !metadata.UsageDiagnostics.Has(usage.DiagnosticUnsupportedBillableDetail) { + t.Fatalf("InspectRequest(%s) = %#v, %v", body, metadata, err) } } for _, body := range []string{ @@ -245,270 +73,9 @@ func TestOpenAIInspectRequestMarksUnsupportedPricingModes(t *testing.T) { `{"model":"gpt-5","service_tier":"auto"}`, `{"model":"gpt-5","service_tier":"default"}`, } { - metadata, err := dialect.InspectRequest(&ParsedRequest{Body: []byte(body)}) - if err != nil { - t.Fatalf("InspectRequest(%s) error = %v", body, err) - } - if metadata.UsageDiagnostics.Has(usage.DiagnosticUnsupportedBillableDetail) { - t.Fatalf("InspectRequest(%s) unexpectedly marked unsupported", body) - } - } -} - -func TestOpenAIClassifyStatus(t *testing.T) { - dialect := NewOpenAI(http.DefaultClient) - tests := []struct { - name string - status int - body string - want health.FailureCategory - }{ - {name: "success is not retryable", status: http.StatusOK, want: health.FailureCategoryOK}, - { - name: "success body marker is not retryable", - status: http.StatusOK, - body: `{"choices":[{"message":{"content":"rate limit"}}]}`, - want: health.FailureCategoryOK, - }, - {name: "success marker ignored", status: http.StatusOK, body: `{"error":{"code":"rate_limit"}}`, want: health.FailureCategoryOK}, - {name: "generic bad request", status: http.StatusBadRequest, body: `{"error":"invalid input"}`, want: health.FailureCategoryClientError}, - {name: "payload too large", status: http.StatusRequestEntityTooLarge, want: health.FailureCategoryClientError}, - {name: "unprocessable entity", status: http.StatusUnprocessableEntity, want: health.FailureCategoryClientError}, - {name: "unauthorized", status: http.StatusUnauthorized, want: health.FailureCategoryInvalidKey}, - {name: "forbidden", status: http.StatusForbidden, want: health.FailureCategoryInvalidKey}, - {name: "model not found", status: http.StatusNotFound, want: health.FailureCategoryModelUnavailable}, - {name: "rate limited", status: http.StatusTooManyRequests, want: health.FailureCategoryRateLimited}, - {name: "upstream server error", status: http.StatusInternalServerError, want: health.FailureCategoryUpstreamHostError}, - { - name: "rate keyword overrides 400", - status: http.StatusBadRequest, - body: `{"error":{"code":"rate_limit_exceeded"}}`, - want: health.FailureCategoryRateLimited, - }, - {name: "rate marker precedes key status", status: http.StatusUnauthorized, body: `{"error":{"code":"rate_limit_exceeded"}}`, want: health.FailureCategoryRateLimited}, - { - name: "invalid key keyword overrides 400", - status: http.StatusBadRequest, - body: `{"error":{"code":"invalid_api_key"}}`, - want: health.FailureCategoryInvalidKey, - }, - { - name: "model keyword overrides 400", - status: http.StatusBadRequest, - body: `{"error":{"message":"model not supported"}}`, - want: health.FailureCategoryModelUnavailable, - }, - { - name: "unsupported parameter remains non retryable", - status: http.StatusBadRequest, - body: `{"error":{"message":"parameter response_format is not supported"}}`, - want: health.FailureCategoryClientError, - }, - {name: "quota field validation is client error", status: http.StatusBadRequest, body: `{"error":"quota field is invalid"}`, want: health.FailureCategoryClientError}, - {name: "disabled feature is client error", status: http.StatusBadRequest, body: `{"error":"feature is disabled"}`, want: health.FailureCategoryClientError}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := dialect.ClassifyStatus(tt.status, []byte(tt.body)); got != tt.want { - t.Fatalf("ClassifyStatus(%d) = %v, want %v", tt.status, got, tt.want) - } - }) - } -} - -func TestOpenAIListModelsUsesDefaultCredentialWithEmptyRules(t *testing.T) { - server := fakeupstream.New(fakeupstream.Step{ - Status: http.StatusOK, - Fixture: "openai/models.json", - }) - defer server.Close() - - dialect := NewOpenAI(server.Client()) - models, err := dialect.ListModels( - context.Background(), - server.URL+"/v1", - "sk-default-models", - state.HeaderRules{}, - ) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - want := []string{"gpt-4o", "gpt-4o-mini"} - if !reflect.DeepEqual(models, want) { - t.Fatalf("ListModels() = %#v, want %#v", models, want) - } - - requests := server.Requests() - if len(requests) != 1 { - t.Fatalf("upstream requests = %d, want 1", len(requests)) - } - if requests[0].Method != http.MethodGet || - requests[0].Path != "/v1/models" { - t.Fatalf("upstream request = %#v, want GET /v1/models", requests[0]) - } - if got := requests[0].Headers.Get("Authorization"); got != "Bearer sk-default-models" { - t.Fatalf("Authorization = %q, want default Bearer credential", got) - } -} - -func TestOpenAIListModelsAppliesHeaderRuleOverrides(t *testing.T) { - server := fakeupstream.New(fakeupstream.Step{ - Status: http.StatusOK, - Fixture: "openai/models.json", - }) - defer server.Close() - - dialect := NewOpenAI(server.Client()) - _, err := dialect.ListModels( - context.Background(), - server.URL+"/v1", - "sk-custom-models", - state.HeaderRules{ - Set: map[string]string{ - "Authorization": "Token ${API_KEY}", - "X-Custom-Key": "${API_KEY}", - }, - }, - ) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - - request := server.Requests()[0] - if got := request.Headers.Get("Authorization"); got != "Token sk-custom-models" { - t.Fatalf("Authorization = %q, want custom override", got) - } - if got := request.Headers.Get("X-Custom-Key"); got != "sk-custom-models" { - t.Fatalf("X-Custom-Key = %q, want expanded key", got) - } -} - -func TestOpenAIListModelsIdentityRepresentationAndCollectsUniqueModels(t *testing.T) { - client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { - assertOutboundRequestRepresentation(t, request, 0, map[string]string{ - "Authorization": "Token secret", - "X-Business": "preserved", - }) - return &http.Response{ - StatusCode: http.StatusOK, - Header: make(http.Header), - Body: io.NopCloser(strings.NewReader(`{"data":[{"id":"a"},{"id":"a"},{"id":"b"}]}`)), - }, nil - })} - - models, err := NewOpenAI(client).ListModels( - context.Background(), - "https://api.example.com", - "secret", - legacyRepresentationHeaderRules(map[string]string{"Authorization": "Token ${API_KEY}"}), - ) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - if !reflect.DeepEqual(models, []string{"a", "b"}) { - t.Fatalf("ListModels() = %#v, want [a b]", models) - } -} - -func TestOpenAIListModelsRejectsUniqueModelOverflow(t *testing.T) { - type item struct { - ID string `json:"id"` - } - payload := struct { - Data []item `json:"data"` - }{Data: make([]item, maxUniqueModelListEntries+1)} - for index := range payload.Data { - payload.Data[index].ID = fmt.Sprintf("model-%06d", index) - } - body, err := json.Marshal(payload) - if err != nil { - t.Fatalf("marshal OpenAI model payload: %v", err) - } - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Header: make(http.Header), - Body: io.NopCloser(strings.NewReader(string(body))), - ContentLength: int64(len(body)), - }, nil - })} - - models, err := NewOpenAI(client).ListModels( - context.Background(), "https://api.example.com", "secret", state.HeaderRules{}, - ) - if err == nil || models != nil { - t.Fatalf("ListModels() = %#v, %v, want nil models and error", models, err) - } -} - -func TestOpenAIListModelsHonorsContextTimeout(t *testing.T) { - server := fakeupstream.New(fakeupstream.Step{ - Status: http.StatusOK, - Fixture: "openai/models.json", - Delay: 200 * time.Millisecond, - }) - defer server.Close() - - ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) - defer cancel() - - dialect := NewOpenAI(server.Client()) - _, err := dialect.ListModels( - ctx, - server.URL+"/v1", - "sk-timeout", - state.HeaderRules{}, - ) - if !errors.Is(err, context.DeadlineExceeded) { - t.Fatalf("ListModels() error = %v, want context deadline exceeded", err) - } -} - -func TestOpenAIListModelsRejectsNonSuccessWithoutLeakingKey(t *testing.T) { - server := fakeupstream.New(fakeupstream.Step{ - Status: http.StatusUnauthorized, - Fixture: "openai/401.json", - }) - defer server.Close() - - const apiKey = "sk-must-not-leak" - dialect := NewOpenAI(server.Client()) - _, err := dialect.ListModels( - context.Background(), - server.URL+"/v1", - apiKey, - state.HeaderRules{}, - ) - if err == nil { - t.Fatal("ListModels() error = nil, want upstream status error") - } - if !strings.Contains(err.Error(), "status 401") { - t.Fatalf("ListModels() error = %v, want status 401", err) - } - if strings.Contains(err.Error(), apiKey) { - t.Fatalf("ListModels() error leaked API key: %v", err) - } -} - -func TestOpenAIListModelsTransportErrorDoesNotExposeURLOrKey(t *testing.T) { - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - return nil, errors.New("forced transport failure") - })} - - const apiKey = "sk-transport-canary" - _, err := NewOpenAI(client).ListModels( - context.Background(), - "https://api.example.com?token=query-canary", - apiKey, - state.HeaderRules{}, - ) - if err == nil { - t.Fatal("ListModels() error = nil, want transport error") - } - for _, forbidden := range []string{apiKey, "query-canary", "api.example.com"} { - if strings.Contains(err.Error(), forbidden) { - t.Fatalf("ListModels() error exposes %q: %v", forbidden, err) + metadata, err := selected.InspectRequest(&ParsedRequest{Body: []byte(body)}) + if err != nil || metadata.UsageDiagnostics.Has(usage.DiagnosticUnsupportedBillableDetail) { + t.Fatalf("InspectRequest(%s) = %#v, %v", body, metadata, err) } } } diff --git a/internal/dialect/openai_usage_options_test.go b/internal/dialect/openai_usage_options_test.go index a33e0f3f..e4b72712 100644 --- a/internal/dialect/openai_usage_options_test.go +++ b/internal/dialect/openai_usage_options_test.go @@ -41,7 +41,7 @@ func TestOpenAIInjectStreamUsage(t *testing.T) { } original := cloneParsedRequestForTest(request) - derived, err := NewOpenAI(http.DefaultClient).InjectStreamUsage(request) + derived, err := NewOpenAI().InjectStreamUsage(request) if err != nil { t.Fatalf("InjectStreamUsage() error = %v", err) } @@ -72,7 +72,7 @@ func TestOpenAIInjectStreamUsageRejectsInvalidStreamOptions(t *testing.T) { `{"stream_options":[]}`, } { request := &ParsedRequest{Body: []byte(body)} - if _, err := NewOpenAI(http.DefaultClient).InjectStreamUsage(request); err == nil { + if _, err := NewOpenAI().InjectStreamUsage(request); err == nil { t.Fatalf("InjectStreamUsage(%s) error = nil", body) } } @@ -80,11 +80,11 @@ func TestOpenAIInjectStreamUsageRejectsInvalidStreamOptions(t *testing.T) { func TestOpenAIInjectStreamUsageIsSemanticallyIdempotent(t *testing.T) { request := &ParsedRequest{Body: []byte(`{"model":"gpt-4o","stream_options":{"future":"keep"}}`)} - first, err := NewOpenAI(http.DefaultClient).InjectStreamUsage(request) + first, err := NewOpenAI().InjectStreamUsage(request) if err != nil { t.Fatalf("first InjectStreamUsage() error = %v", err) } - second, err := NewOpenAI(http.DefaultClient).InjectStreamUsage(first) + second, err := NewOpenAI().InjectStreamUsage(first) if err != nil { t.Fatalf("second InjectStreamUsage() error = %v", err) } @@ -99,10 +99,10 @@ func TestOnlyOpenAIImplementsStreamUsageInjector(t *testing.T) { value any want bool }{ - {name: "OpenAI", value: NewOpenAI(http.DefaultClient), want: true}, - {name: "OpenAI Responses", value: NewOpenAIResponses(http.DefaultClient), want: false}, - {name: "Anthropic", value: NewAnthropic(http.DefaultClient), want: false}, - {name: "Gemini", value: NewGemini(http.DefaultClient), want: false}, + {name: "OpenAI", value: NewOpenAI(), want: true}, + {name: "OpenAI Responses", value: NewOpenAIResponses(), want: false}, + {name: "Anthropic", value: NewAnthropic(), want: false}, + {name: "Gemini", value: NewGemini(), want: false}, {name: "dialect only", value: &usageDialectOnly{}, want: false}, } { t.Run(test.name, func(t *testing.T) { diff --git a/internal/dialect/openai_usage_test.go b/internal/dialect/openai_usage_test.go index 8dd29a80..05f8b3c0 100644 --- a/internal/dialect/openai_usage_test.go +++ b/internal/dialect/openai_usage_test.go @@ -1,7 +1,6 @@ package dialect import ( - "net/http" "strings" "testing" @@ -9,7 +8,7 @@ import ( ) func TestUsageOpenAICanonicalFixtures(t *testing.T) { - extractor := NewOpenAI(http.DefaultClient) + extractor := NewOpenAI() result, err := extractor.ExtractUsage(readUsageFixture(t, "openai", "nonstream.json")) if err != nil { t.Fatal(err) @@ -33,7 +32,7 @@ func TestUsageOpenAICanonicalFixtures(t *testing.T) { } func TestUsageOpenAINonStreamOptionalFields(t *testing.T) { - extractor := NewOpenAI(http.DefaultClient) + extractor := NewOpenAI() negativeTen := int64(-10) tests := []struct { name string @@ -117,7 +116,7 @@ func TestUsageOpenAINonStreamOptionalFields(t *testing.T) { } func TestUsageOpenAINonStreamRequiredAndStrictNumbers(t *testing.T) { - extractor := NewOpenAI(http.DefaultClient) + extractor := NewOpenAI() tests := []struct { name string body string @@ -225,7 +224,7 @@ func TestUsageOpenAINonStreamRequiredAndStrictNumbers(t *testing.T) { } func TestUsageOpenAIClampsCachedAndChecksTotal(t *testing.T) { - extractor := NewOpenAI(http.DefaultClient) + extractor := NewOpenAI() tests := []struct { name string body string @@ -293,7 +292,7 @@ func TestUsageOpenAIClampsCachedAndChecksTotal(t *testing.T) { } func TestUsageOpenAIInvalidUsageObjectIsDiagnosedAndDoesNotDiscardStreamState(t *testing.T) { - extractor := NewOpenAI(http.DefaultClient) + extractor := NewOpenAI() result, err := extractor.ExtractUsage([]byte(`{"usage":[]}`)) if err != nil || result.State != usage.StateMissing || result.Tokens != (usage.Tokens{}) { t.Fatalf("invalid non-stream usage = %#v, %v", result, err) @@ -333,7 +332,7 @@ func TestUsageOpenAIInvalidUsageObjectIsDiagnosedAndDoesNotDiscardStreamState(t } func TestUsageOpenAIStreamSnapshotsAndFinality(t *testing.T) { - extractor := NewOpenAI(http.DefaultClient) + extractor := NewOpenAI() tests := []struct { name string steps []string @@ -389,7 +388,7 @@ func TestUsageOpenAIStreamSnapshotsAndFinality(t *testing.T) { } func TestUsageOpenAIStreamReplaceSnapshotClearsAbsentTokensAndDiagnostics(t *testing.T) { - stream := NewOpenAI(http.DefaultClient).NewUsageStreamExtractor() + stream := NewOpenAI().NewUsageStreamExtractor() for _, payload := range []string{ `{"choices":[{"delta":{}}],"usage":{"prompt_tokens":100,"completion_tokens":10,"prompt_tokens_details":{"cached_tokens":20,"cache_write_tokens":1.5}}}`, `{"choices":[],"usage":{"prompt_tokens":100}}`, @@ -407,7 +406,7 @@ func TestUsageOpenAIStreamReplaceSnapshotClearsAbsentTokensAndDiagnostics(t *tes } func TestUsageOpenAIStreamMalformedPayloadKeepsDiagnosticAndState(t *testing.T) { - stream := NewOpenAI(http.DefaultClient).NewUsageStreamExtractor() + stream := NewOpenAI().NewUsageStreamExtractor() if err := stream.Observe([]byte(`{"choices":[{"delta":{}}],"usage":{"prompt_tokens":100,"completion_tokens":10}}`)); err != nil { t.Fatal(err) } @@ -428,7 +427,7 @@ func TestUsageOpenAIStreamMalformedPayloadKeepsDiagnosticAndState(t *testing.T) } func TestUsageOpenAIMalformedBodiesAndObserveOwnership(t *testing.T) { - extractor := NewOpenAI(http.DefaultClient) + extractor := NewOpenAI() for _, body := range [][]byte{[]byte(`{}`), []byte(`{"usage":null}`)} { result, err := extractor.ExtractUsage(body) if err != nil || result.State != usage.StateMissing || result.Tokens != (usage.Tokens{}) { diff --git a/internal/dialect/pipeline.go b/internal/dialect/pipeline.go deleted file mode 100644 index 85ff9a09..00000000 --- a/internal/dialect/pipeline.go +++ /dev/null @@ -1,31 +0,0 @@ -package dialect - -import ( - "net/http" - "strings" - - "gpt-load/internal/state" -) - -type credentialInjector interface { - InjectCredential(headers http.Header, apiKey string) -} - -func ApplyCredential( - injector credentialInjector, - headers http.Header, - apiKey string, - rules state.HeaderRules, -) { - if injector == nil || headers == nil { - return - } - - injector.InjectCredential(headers, apiKey) - for name, value := range rules.Set { - headers.Set(name, strings.ReplaceAll(value, "${API_KEY}", apiKey)) - } - for _, name := range rules.Remove { - headers.Del(name) - } -} diff --git a/internal/dialect/pipeline_test.go b/internal/dialect/pipeline_test.go deleted file mode 100644 index 0560d474..00000000 --- a/internal/dialect/pipeline_test.go +++ /dev/null @@ -1,70 +0,0 @@ -package dialect - -import ( - "net/http" - "testing" - - "gpt-load/internal/state" -) - -func TestApplyCredentialUsesDialectDefaultWithEmptyRules(t *testing.T) { - headers := make(http.Header) - - ApplyCredential( - NewOpenAI(http.DefaultClient), - headers, - "sk-default", - state.HeaderRules{}, - ) - - if got := headers.Get("Authorization"); got != "Bearer sk-default" { - t.Fatalf("Authorization = %q, want default Bearer credential", got) - } -} - -func TestApplyCredentialExpandsSetRulesAfterDefault(t *testing.T) { - headers := http.Header{"X-Remove-Me": {"old"}} - rules := state.HeaderRules{ - Set: map[string]string{ - "Authorization": "Token ${API_KEY}", - "X-Custom-Key": "prefix-${API_KEY}-${API_KEY}", - }, - Remove: []string{"X-Remove-Me"}, - } - - ApplyCredential( - NewOpenAI(http.DefaultClient), - headers, - "sk-custom", - rules, - ) - - if got := headers.Get("Authorization"); got != "Token sk-custom" { - t.Fatalf("Authorization = %q, want custom override", got) - } - if got := headers.Get("X-Custom-Key"); got != "prefix-sk-custom-sk-custom" { - t.Fatalf("X-Custom-Key = %q, want all templates expanded", got) - } - if got := headers.Get("X-Remove-Me"); got != "" { - t.Fatalf("X-Remove-Me = %q, want removed", got) - } -} - -func TestApplyCredentialRemoveWinsOverSet(t *testing.T) { - headers := make(http.Header) - rules := state.HeaderRules{ - Set: map[string]string{"Authorization": "Token ${API_KEY}"}, - Remove: []string{"Authorization"}, - } - - ApplyCredential( - NewOpenAI(http.DefaultClient), - headers, - "sk-removed", - rules, - ) - - if got := headers.Get("Authorization"); got != "" { - t.Fatalf("Authorization = %q, want final remove to win", got) - } -} diff --git a/internal/dialect/probe.go b/internal/dialect/probe.go deleted file mode 100644 index 87038428..00000000 --- a/internal/dialect/probe.go +++ /dev/null @@ -1,69 +0,0 @@ -package dialect - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "net/http" - "strings" - - platformheader "gpt-load/internal/platform/httpheader" - "gpt-load/internal/state" -) - -type probeMessage struct { - Role string `json:"role"` - Content string `json:"content"` -} - -func validateProbeModel(validationModel string) error { - if strings.TrimSpace(validationModel) == "" { - return fmt.Errorf("probe validation model is required") - } - return nil -} - -func executeProbe( - ctx context.Context, - client *http.Client, - dialect Dialect, - requestURL, apiKey string, - rules state.HeaderRules, - payload any, -) error { - if ctx == nil { - return fmt.Errorf("probe context is required") - } - if client == nil { - return fmt.Errorf("request %s probe failed", dialect.Protocol()) - } - - body, err := json.Marshal(payload) - if err != nil { - return fmt.Errorf("encode %s probe request failed", dialect.Protocol()) - } - request, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL, bytes.NewReader(body)) - if err != nil { - return fmt.Errorf("create %s probe request failed", dialect.Protocol()) - } - request.Header.Set("Content-Type", "application/json") - ApplyCredential(dialect, request.Header, apiKey, rules) - platformheader.NormalizeUpstreamRequestRepresentation(request, int64(len(body))) - - response, err := client.Do(request) - if err != nil { - if contextErr := ctx.Err(); contextErr != nil { - return contextErr - } - return fmt.Errorf("request %s probe failed", dialect.Protocol()) - } - if response == nil || response.Body == nil { - return fmt.Errorf("request %s probe failed", dialect.Protocol()) - } - defer response.Body.Close() - if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { - return fmt.Errorf("request %s probe: upstream status %d", dialect.Protocol(), response.StatusCode) - } - return nil -} diff --git a/internal/dialect/probe_test.go b/internal/dialect/probe_test.go deleted file mode 100644 index f5520f08..00000000 --- a/internal/dialect/probe_test.go +++ /dev/null @@ -1,255 +0,0 @@ -package dialect - -import ( - "context" - "encoding/json" - "errors" - "io" - "net/http" - "net/http/httptest" - "strings" - "sync/atomic" - "testing" - - "gpt-load/internal/state" -) - -func TestOpenAIProbeSendsMinimalAuthenticatedRequest(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if request.Method != http.MethodPost || request.URL.Path != openAICompletionsPath { - t.Errorf("request = %s %s, want POST %s", request.Method, request.URL.Path, openAICompletionsPath) - } - if got := request.Header.Get("Authorization"); got != "Bearer secret" { - t.Errorf("Authorization = %q, want bearer credential", got) - } - assertOpenAIProbePayload(t, request.Body, "gpt-test") - writer.WriteHeader(http.StatusNoContent) - })) - defer server.Close() - - if err := NewOpenAI(server.Client()).Probe( - context.Background(), server.URL+"/v1", "secret", state.HeaderRules{}, "gpt-test", - ); err != nil { - t.Fatalf("Probe() error = %v", err) - } -} - -func TestAnthropicProbeSendsMinimalAuthenticatedRequest(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if request.Method != http.MethodPost || request.URL.Path != anthropicMessagesPath { - t.Errorf("request = %s %s, want POST %s", request.Method, request.URL.Path, anthropicMessagesPath) - } - if got := request.Header.Get("X-Api-Key"); got != "secret" { - t.Errorf("X-Api-Key = %q, want credential", got) - } - if got := request.Header.Get("Anthropic-Version"); got != anthropicDefaultVersion { - t.Errorf("Anthropic-Version = %q, want %q", got, anthropicDefaultVersion) - } - assertOpenAIProbePayload(t, request.Body, "claude-test") - writer.WriteHeader(http.StatusAccepted) - })) - defer server.Close() - - if err := NewAnthropic(server.Client()).Probe( - context.Background(), server.URL+"/v1", "secret", state.HeaderRules{}, "claude-test", - ); err != nil { - t.Fatalf("Probe() error = %v", err) - } -} - -func TestGeminiProbeSendsMinimalAuthenticatedRequest(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if request.Method != http.MethodPost || request.URL.Path != "/v1beta/models/gemini-test:generateContent" { - t.Errorf("request = %s %s, want POST Gemini generation path", request.Method, request.URL.Path) - } - if got := request.Header.Get("X-Goog-Api-Key"); got != "secret" { - t.Errorf("X-Goog-Api-Key = %q, want credential", got) - } - if got := request.URL.Query().Get("key"); got != "" { - t.Errorf("key query = %q, want empty", got) - } - assertGeminiProbePayload(t, request.Body, "gemini-test") - writer.WriteHeader(http.StatusOK) - })) - defer server.Close() - - if err := NewGemini(server.Client()).Probe( - context.Background(), server.URL+"/v1beta?tenant=one", "secret", state.HeaderRules{}, "gemini-test", - ); err != nil { - t.Fatalf("Probe() error = %v", err) - } -} - -func TestProbeIdentityRepresentationPreservesRulesAndCredential(t *testing.T) { - client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { - assertProbeRequestRepresentation(t, request, map[string]string{ - "Authorization": "Token secret", - "Content-Type": "application/json", - "X-Business": "preserved", - "X-Probe": "prefix-secret", - }) - return &http.Response{ - StatusCode: http.StatusNoContent, - Body: io.NopCloser(strings.NewReader("ignored")), - Header: make(http.Header), - }, nil - })} - - rules := legacyRepresentationHeaderRules(map[string]string{ - "Authorization": "Token ${API_KEY}", - "X-Probe": "prefix-${API_KEY}", - }) - err := NewOpenAI(client).Probe(context.Background(), "https://api.example.com", "secret", rules, "gpt-test") - if err != nil { - t.Fatalf("Probe() error = %v", err) - } -} - -func TestProbeAcceptsOnly2xxAndRedactsTransportDetails(t *testing.T) { - t.Run("accepts every 2xx and closes response", func(t *testing.T) { - body := &trackedReadCloser{Reader: strings.NewReader("ignored")} - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - return &http.Response{StatusCode: 299, Body: body, Header: make(http.Header)}, nil - })} - if err := NewOpenAI(client).Probe(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}, "gpt-test"); err != nil { - t.Fatalf("Probe() error = %v", err) - } - if !body.closed.Load() { - t.Fatal("response body was not closed") - } - }) - - t.Run("non 2xx contains only protocol and status", func(t *testing.T) { - const secret = "distinctive-probe-secret" - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.WriteHeader(http.StatusForbidden) - _, _ = writer.Write([]byte(`{"error":"` + secret + `"}`)) - })) - defer server.Close() - - err := NewOpenAI(server.Client()).Probe(context.Background(), server.URL+"?token=query-secret", secret, state.HeaderRules{}, "gpt-test") - if err == nil || !strings.Contains(err.Error(), "openai") || !strings.Contains(err.Error(), "403") { - t.Fatalf("Probe() error = %v, want protocol and status", err) - } - for _, forbidden := range []string{secret, "query-secret", server.URL} { - if strings.Contains(err.Error(), forbidden) { - t.Fatalf("Probe() error exposes %q: %v", forbidden, err) - } - } - }) - - t.Run("transport error is redacted", func(t *testing.T) { - const secret = "distinctive-probe-secret" - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - return nil, errors.New("transport failed for https://bad.example.com?token=" + secret) - })} - err := NewOpenAI(client).Probe(context.Background(), "https://api.example.com", secret, state.HeaderRules{}, "gpt-test") - if err == nil || !strings.Contains( - err.Error(), - "request openai-completions probe failed", - ) { - t.Fatalf("Probe() error = %v, want redacted transport failure", err) - } - for _, forbidden := range []string{secret, "bad.example.com", "api.example.com"} { - if strings.Contains(err.Error(), forbidden) { - t.Fatalf("Probe() error exposes %q: %v", forbidden, err) - } - } - }) - - t.Run("nil client returns a safe error", func(t *testing.T) { - err := NewOpenAI(nil).Probe(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}, "gpt-test") - if err == nil || !strings.Contains( - err.Error(), - "request openai-completions probe failed", - ) || strings.Contains(err.Error(), "secret") { - t.Fatalf("Probe() error = %v, want redacted nil-client failure", err) - } - }) -} - -func TestProbeHonorsContextCancellation(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - err := NewOpenAI(http.DefaultClient).Probe(ctx, "https://api.example.com", "secret", state.HeaderRules{}, "gpt-test") - if !errors.Is(err, context.Canceled) { - t.Fatalf("Probe() error = %v, want context cancellation", err) - } -} - -func TestProbeRejectsBlankValidationModelWithoutRequest(t *testing.T) { - var requests atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - requests.Add(1) - return nil, errors.New("unexpected request") - })} - - for _, value := range []Dialect{NewOpenAI(client), NewAnthropic(client), NewGemini(client)} { - if err := value.Probe(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}, " \t "); err == nil { - t.Fatalf("%s Probe() error = nil for blank validation model", value.Protocol()) - } - } - if got := requests.Load(); got != 0 { - t.Fatalf("requests = %d, want none", got) - } -} - -func TestGeminiProbeDoesNotNormalizeConfiguredModel(t *testing.T) { - var requests atomic.Int64 - client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { - requests.Add(1) - return nil, errors.New("unexpected request") - })} - - err := NewGemini(client).Probe(context.Background(), "https://api.example.com", "secret", state.HeaderRules{}, "models/gemini-test") - if err == nil { - t.Fatal("Probe() error = nil for an unnormalized models/ prefix") - } - if got := requests.Load(); got != 0 { - t.Fatalf("requests = %d, want none", got) - } -} - -func assertOpenAIProbePayload(t *testing.T, body io.Reader, wantModel string) { - t.Helper() - var payload struct { - Model string `json:"model"` - Messages []struct { - Role string `json:"role"` - Content string `json:"content"` - } `json:"messages"` - MaxTokens int `json:"max_tokens"` - } - if err := json.NewDecoder(body).Decode(&payload); err != nil { - t.Fatalf("decode Probe payload: %v", err) - } - if payload.Model != wantModel || payload.MaxTokens != 1 || len(payload.Messages) != 1 || - payload.Messages[0].Role != "user" || payload.Messages[0].Content != "ping" { - t.Fatalf("Probe payload = %#v, want model %q and one user ping with max_tokens 1", payload, wantModel) - } -} - -func assertGeminiProbePayload(t *testing.T, body io.Reader, wantModel string) { - t.Helper() - var payload struct { - Contents []struct { - Role string `json:"role"` - Parts []struct { - Text string `json:"text"` - } `json:"parts"` - } `json:"contents"` - GenerationConfig struct { - MaxOutputTokens int `json:"maxOutputTokens"` - } `json:"generationConfig"` - } - if err := json.NewDecoder(body).Decode(&payload); err != nil { - t.Fatalf("decode Probe payload: %v", err) - } - if len(payload.Contents) != 1 || payload.Contents[0].Role != "user" || - len(payload.Contents[0].Parts) != 1 || payload.Contents[0].Parts[0].Text != "ping" || - payload.GenerationConfig.MaxOutputTokens != 1 { - t.Fatalf("Probe payload = %#v, want one user ping with maxOutputTokens 1", payload) - } - _ = wantModel -} diff --git a/internal/dialect/request_execution.go b/internal/dialect/request_execution.go new file mode 100644 index 00000000..0652e2a8 --- /dev/null +++ b/internal/dialect/request_execution.go @@ -0,0 +1,243 @@ +package dialect + +import ( + "bytes" + "encoding/json" + "net/http" + "strings" + + "gpt-load/internal/execution" + "gpt-load/internal/reasoning" +) + +func chatExecutionMetadata( + body []byte, + stream bool, + reasoningConfig reasoning.Config, +) (execution.Operation, execution.FeatureSet) { + return execution.OperationChatCompletion, + inspectRequiredFeatures(body, stream, reasoningConfig.RequiresCapability(), false) +} + +func responsesExecutionMetadata( + request *ParsedRequest, + stream bool, + reasoningConfig reasoning.Config, +) (execution.Operation, execution.FeatureSet) { + operation := responsesOperation(request) + nativeResource := operation == execution.OperationResponsesRetrieve || + operation == execution.OperationResponsesDelete || + operation == execution.OperationResponsesCancel || + operation == execution.OperationResponsesInputItems || + operation == execution.OperationResponsesPassthrough + if operation == execution.OperationResponsesCreate { + nativeResource = responsesCreateStoresResource(request.Body) + } + return operation, inspectRequiredFeatures( + request.Body, + stream, + reasoningConfig.RequiresCapability(), + nativeResource, + ) +} + +func responsesOperation(request *ParsedRequest) execution.Operation { + if request == nil { + return execution.OperationResponsesPassthrough + } + switch { + case request.Method == http.MethodPost && request.Path == openAIResponsesPath: + return execution.OperationResponsesCreate + case request.Method == http.MethodPost && request.Path == openAIResponsesCompactPath: + return execution.OperationResponsesCompact + case request.Method == http.MethodPost && request.Path == openAIResponsesPath+"/input_tokens": + return execution.OperationResponsesInputTokens + } + + resourcePath := strings.TrimPrefix(request.Path, openAIResponsesPath+"/") + segments := strings.Split(resourcePath, "/") + if resourcePath == request.Path || len(segments) == 0 || segments[0] == "" { + return execution.OperationResponsesPassthrough + } + switch { + case len(segments) == 1 && request.Method == http.MethodGet: + return execution.OperationResponsesRetrieve + case len(segments) == 1 && request.Method == http.MethodDelete: + return execution.OperationResponsesDelete + case len(segments) == 2 && segments[1] == "cancel" && request.Method == http.MethodPost: + return execution.OperationResponsesCancel + case len(segments) == 2 && segments[1] == "input_items" && request.Method == http.MethodGet: + return execution.OperationResponsesInputItems + default: + return execution.OperationResponsesPassthrough + } +} + +func inspectRequiredFeatures( + body []byte, + stream bool, + reasoningPresent bool, + nativeResource bool, +) execution.FeatureSet { + features := make([]execution.Feature, 0, 6) + if stream { + features = append(features, execution.FeatureStreaming) + } + if reasoningPresent { + features = append(features, execution.FeatureReasoning) + } + + root, ok := decodeExecutionFeatureObject(body) + if ok { + if hasMeaningfulField(root, "tools") || + hasMeaningfulField(root, "tool_choice") || + hasMeaningfulField(root, "toolChoice") { + features = append(features, execution.FeatureTools) + } + if containsMultimodalValue(root) { + features = append(features, execution.FeatureMultimodal) + } + if containsStructuredOutput(root) { + features = append(features, execution.FeatureStructuredOutput) + } + } + if nativeResource { + features = append(features, execution.FeatureNativeResourceSemantics) + } + result, err := execution.NewFeatureSet(features...) + if err != nil { + return execution.FeatureSet{} + } + return result +} + +func decodeExecutionFeatureObject(body []byte) (map[string]any, bool) { + if len(bytes.TrimSpace(body)) == 0 { + return nil, false + } + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + var object map[string]any + if err := decoder.Decode(&object); err != nil || object == nil { + return nil, false + } + return object, true +} + +func hasMeaningfulField(object map[string]any, field string) bool { + value, ok := object[field] + if !ok || value == nil { + return false + } + switch typed := value.(type) { + case []any: + return len(typed) > 0 + case map[string]any: + return len(typed) > 0 + case string: + return typed != "" + default: + return true + } +} + +func containsMultimodalValue(value any) bool { + switch typed := value.(type) { + case map[string]any: + for key, child := range typed { + normalizedKey := normalizeFeatureName(key) + if normalizedKey == "inlinedata" || + normalizedKey == "filedata" || + normalizedKey == "imageurl" { + return true + } + if normalizedKey == "type" { + if text, ok := child.(string); ok && multimodalType(text) { + return true + } + } + if normalizedKey == "modalities" && containsNonTextModality(child) { + return true + } + if containsMultimodalValue(child) { + return true + } + } + case []any: + for _, child := range typed { + if containsMultimodalValue(child) { + return true + } + } + } + return false +} + +func multimodalType(value string) bool { + switch normalizeFeatureName(value) { + case "image", "imageurl", "inputimage", "inputaudio", "file", "inputfile", "document": + return true + default: + return false + } +} + +func containsNonTextModality(value any) bool { + items, ok := value.([]any) + if !ok { + return false + } + for _, item := range items { + text, ok := item.(string) + if ok && !strings.EqualFold(strings.TrimSpace(text), "text") { + return true + } + } + return false +} + +func containsStructuredOutput(root map[string]any) bool { + if hasMeaningfulField(root, "response_format") { + return true + } + for _, parentName := range []string{ + "text", + "output_config", + "outputConfig", + "generation_config", + "generationConfig", + } { + parent, ok := root[parentName].(map[string]any) + if !ok { + continue + } + for key, value := range parent { + switch normalizeFeatureName(key) { + case "format", "responseschema", "responsejsonschema", "responsemimetype": + if value != nil { + return true + } + } + } + } + return false +} + +func responsesCreateStoresResource(body []byte) bool { + root, ok := decodeExecutionFeatureObject(body) + if !ok { + return true + } + value, exists := root["store"] + if !exists || value == nil { + return true + } + store, ok := value.(bool) + return !ok || store +} + +func normalizeFeatureName(value string) string { + value = strings.ToLower(value) + value = strings.ReplaceAll(value, "_", "") + return strings.ReplaceAll(value, "-", "") +} diff --git a/internal/dialect/request_execution_test.go b/internal/dialect/request_execution_test.go new file mode 100644 index 00000000..15d72e81 --- /dev/null +++ b/internal/dialect/request_execution_test.go @@ -0,0 +1,188 @@ +package dialect + +import ( + "net/http" + "testing" + + "gpt-load/internal/execution" +) + +func TestInspectRequestIdentifiesExecutionOperationAndFeatures(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + dialect Dialect + request *ParsedRequest + operation execution.Operation + features []execution.Feature + noFeatures []execution.Feature + }{ + { + name: "OpenAI Chat", + dialect: NewOpenAI(), + request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/chat/completions", Body: []byte(`{ + "model":"gpt-4o","stream":true,"reasoning_effort":"high", + "tools":[{"type":"function","function":{"name":"lookup"}}], + "messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":"https://example.test/image.png"}}]}], + "response_format":{"type":"json_schema","json_schema":{"name":"answer","schema":{"type":"object"}}} + }`)}, + operation: execution.OperationChatCompletion, + features: []execution.Feature{ + execution.FeatureStreaming, + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + }, + }, + { + name: "Anthropic Messages", + dialect: NewAnthropic(), + request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/messages", Body: []byte(`{ + "model":"claude-sonnet-4","stream":true, + "thinking":{"type":"enabled","budget_tokens":1024}, + "tools":[{"name":"lookup","input_schema":{"type":"object"}}], + "messages":[{"role":"user","content":[{"type":"document","source":{"type":"base64","media_type":"application/pdf","data":"AA=="}}]}], + "output_config":{"format":{"type":"json_schema","schema":{"type":"object"}}} + }`)}, + operation: execution.OperationChatCompletion, + features: []execution.Feature{ + execution.FeatureStreaming, + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + }, + }, + { + name: "Gemini generateContent", + dialect: NewGemini(), + request: &ParsedRequest{Method: http.MethodPost, Path: "/v1beta/models/gemini-2.5-pro:streamGenerateContent", Body: []byte(`{ + "contents":[{"role":"user","parts":[{"inlineData":{"mimeType":"image/png","data":"AA=="}}]}], + "tools":[{"functionDeclarations":[{"name":"lookup"}]}], + "generationConfig":{"thinkingConfig":{"thinkingLevel":"high"},"responseSchema":{"type":"OBJECT"}} + }`)}, + operation: execution.OperationChatCompletion, + features: []execution.Feature{ + execution.FeatureStreaming, + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + }, + }, + { + name: "Responses create stored", + dialect: NewOpenAIResponses(), + request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses", Body: []byte(`{ + "model":"gpt-5","stream":true,"reasoning":{"effort":"high"}, + "tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}], + "input":[{"role":"user","content":[{"type":"input_file","file_id":"file_123"}]}], + "text":{"format":{"type":"json_schema","name":"answer","schema":{"type":"object"}}} + }`)}, + operation: execution.OperationResponsesCreate, + features: []execution.Feature{ + execution.FeatureStreaming, + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + execution.FeatureNativeResourceSemantics, + }, + }, + { + name: "Responses create transient", + dialect: NewOpenAIResponses(), + request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses", Body: []byte(`{"model":"gpt-5","store":false}`)}, + operation: execution.OperationResponsesCreate, + noFeatures: []execution.Feature{execution.FeatureNativeResourceSemantics}, + }, + { + name: "Responses reasoning none is not a capability requirement", + dialect: NewOpenAIResponses(), + request: &ParsedRequest{Method: http.MethodPost, Path: "/v1/responses", Body: []byte(`{"model":"gpt-5","reasoning":{"effort":"none"},"store":false}`)}, + operation: execution.OperationResponsesCreate, + noFeatures: []execution.Feature{execution.FeatureReasoning, execution.FeatureNativeResourceSemantics}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + metadata, err := test.dialect.InspectRequest(test.request) + if err != nil { + t.Fatalf("InspectRequest() error = %v", err) + } + if metadata.Operation != test.operation { + t.Fatalf("Operation = %q, want %q", metadata.Operation, test.operation) + } + for _, feature := range test.features { + if !metadata.RequiredFeatures.Has(feature) { + t.Errorf("RequiredFeatures missing %q", feature) + } + } + for _, feature := range test.noFeatures { + if metadata.RequiredFeatures.Has(feature) { + t.Errorf("RequiredFeatures unexpectedly contains %q", feature) + } + } + }) + } +} + +func TestOpenAIResponsesInspectRequestIdentifiesLifecycleOperations(t *testing.T) { + t.Parallel() + + selected := NewOpenAIResponses() + tests := []struct { + name string + method string + path string + body string + operation execution.Operation + model string + }{ + {name: "create", method: http.MethodPost, path: "/v1/responses", body: `{"model":"gpt-5"}`, operation: execution.OperationResponsesCreate, model: "gpt-5"}, + {name: "compact", method: http.MethodPost, path: "/v1/responses/compact", body: `{"model":"gpt-5","input":[]}`, operation: execution.OperationResponsesCompact, model: "gpt-5"}, + {name: "input tokens", method: http.MethodPost, path: "/v1/responses/input_tokens", body: `{"model":"gpt-5","input":[]}`, operation: execution.OperationResponsesInputTokens, model: "gpt-5"}, + {name: "retrieve", method: http.MethodGet, path: "/v1/responses/resp_123", operation: execution.OperationResponsesRetrieve}, + {name: "delete", method: http.MethodDelete, path: "/v1/responses/resp_123", operation: execution.OperationResponsesDelete}, + {name: "cancel", method: http.MethodPost, path: "/v1/responses/resp_123/cancel", operation: execution.OperationResponsesCancel}, + {name: "input items", method: http.MethodGet, path: "/v1/responses/resp_123/input_items", operation: execution.OperationResponsesInputItems}, + {name: "vendor extension", method: http.MethodPatch, path: "/v1/responses/vendor-extension/nested", body: `{"vendor":true}`, operation: execution.OperationResponsesPassthrough}, + {name: "resource head", method: http.MethodHead, path: "/v1/responses/resp_123", operation: execution.OperationResponsesPassthrough}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + metadata, err := selected.InspectRequest(&ParsedRequest{ + Method: test.method, + Path: test.path, + Body: []byte(test.body), + }) + if err != nil { + t.Fatalf("InspectRequest() error = %v", err) + } + if metadata.Operation != test.operation { + t.Fatalf("Operation = %q, want %q", metadata.Operation, test.operation) + } + gotModel := "" + if metadata.Model != nil { + gotModel = *metadata.Model + } + if got := gotModel; got != test.model { + t.Fatalf("Model = %q, want %q", got, test.model) + } + if !metadata.RequiredFeatures.Has(execution.FeatureNativeResourceSemantics) && + (test.operation == execution.OperationResponsesRetrieve || + test.operation == execution.OperationResponsesDelete || + test.operation == execution.OperationResponsesCancel || + test.operation == execution.OperationResponsesInputItems || + test.operation == execution.OperationResponsesPassthrough) { + t.Fatal("resource operation must require native resource semantics") + } + }) + } +} diff --git a/internal/dialect/request_reasoning_test.go b/internal/dialect/request_reasoning_test.go index 1d67d740..5243473e 100644 --- a/internal/dialect/request_reasoning_test.go +++ b/internal/dialect/request_reasoning_test.go @@ -21,7 +21,7 @@ func TestInspectRequestCapturesExplicitReasoningConfiguration(t *testing.T) { }{ { name: "OpenAI completions effort", - value: NewOpenAI(http.DefaultClient), + value: NewOpenAI(), request: &ParsedRequest{ Body: []byte(`{"model":"gpt-5","reasoning_effort":"HIGH"}`), }, @@ -29,7 +29,7 @@ func TestInspectRequestCapturesExplicitReasoningConfiguration(t *testing.T) { }, { name: "OpenAI Responses mode and effort", - value: NewOpenAIResponses(http.DefaultClient), + value: NewOpenAIResponses(), request: &ParsedRequest{ Method: http.MethodPost, Path: "/v1/responses", @@ -39,7 +39,7 @@ func TestInspectRequestCapturesExplicitReasoningConfiguration(t *testing.T) { }, { name: "Anthropic adaptive effort", - value: NewAnthropic(http.DefaultClient), + value: NewAnthropic(), request: &ParsedRequest{ Body: []byte(`{"model":"claude-sonnet-4-6","thinking":{"type":"adaptive"},"output_config":{"effort":"HIGH"}}`), }, @@ -47,7 +47,7 @@ func TestInspectRequestCapturesExplicitReasoningConfiguration(t *testing.T) { }, { name: "Anthropic manual budget", - value: NewAnthropic(http.DefaultClient), + value: NewAnthropic(), request: &ParsedRequest{ Body: []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"enabled","budget_tokens":4096}}`), }, @@ -55,7 +55,7 @@ func TestInspectRequestCapturesExplicitReasoningConfiguration(t *testing.T) { }, { name: "Gemini level and dynamic budget", - value: NewGemini(http.DefaultClient), + value: NewGemini(), request: &ParsedRequest{ Path: "/v1beta/models/gemini-3-flash:generateContent", Body: []byte(`{"generationConfig":{"thinkingConfig":{"thinkingLevel":"HIGH","thinkingBudget":-1}}}`), @@ -88,12 +88,12 @@ func TestInspectRequestIgnoresInvalidReasoningMetadata(t *testing.T) { }{ { name: "OpenAI completions wrong effort type", - value: NewOpenAI(http.DefaultClient), + value: NewOpenAI(), request: &ParsedRequest{Body: []byte(`{"model":"gpt-5","reasoning_effort":3}`)}, }, { name: "OpenAI Responses unknown reasoning object", - value: NewOpenAIResponses(http.DefaultClient), + value: NewOpenAIResponses(), request: &ParsedRequest{ Method: http.MethodPost, Path: "/v1/responses", @@ -102,13 +102,13 @@ func TestInspectRequestIgnoresInvalidReasoningMetadata(t *testing.T) { }, { name: "Anthropic wrong budget type", - value: NewAnthropic(http.DefaultClient), + value: NewAnthropic(), request: &ParsedRequest{Body: []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"enabled","budget_tokens":"4096"}}`)}, want: reasoning.Config{Mode: "enabled"}, }, { name: "Gemini invalid body remains pass through metadata", - value: NewGemini(http.DefaultClient), + value: NewGemini(), request: &ParsedRequest{ Path: "/v1beta/models/gemini-3-flash:generateContent", Body: []byte("{"), diff --git a/internal/dialect/response_model_test.go b/internal/dialect/response_model_test.go index 271b4386..be86ef45 100644 --- a/internal/dialect/response_model_test.go +++ b/internal/dialect/response_model_test.go @@ -1,7 +1,6 @@ package dialect import ( - "net/http" "reflect" "testing" ) @@ -15,43 +14,43 @@ func TestResponseModelInspectorsObserveProtocolFields(t *testing.T) { }{ { name: "openai chat completion", - dialect: NewOpenAI(http.DefaultClient), + dialect: NewOpenAI(), payload: `{"id":"chatcmpl-1","model":"gpt-4.1"}`, want: []string{"gpt-4.1"}, }, { name: "openai responses object", - dialect: NewOpenAIResponses(http.DefaultClient), + dialect: NewOpenAIResponses(), payload: `{"id":"resp-1","model":"gpt-5"}`, want: []string{"gpt-5"}, }, { name: "openai responses stream event", - dialect: NewOpenAIResponses(http.DefaultClient), + dialect: NewOpenAIResponses(), payload: `{"type":"response.completed","response":{"id":"resp-1","model":"gpt-5.1"}}`, want: []string{"gpt-5.1"}, }, { name: "openai responses preserves conflicting declarations", - dialect: NewOpenAIResponses(http.DefaultClient), + dialect: NewOpenAIResponses(), payload: `{"model":"gpt-5","response":{"model":"gpt-5.1"}}`, want: []string{"gpt-5", "gpt-5.1"}, }, { name: "anthropic message", - dialect: NewAnthropic(http.DefaultClient), + dialect: NewAnthropic(), payload: `{"type":"message","model":"claude-sonnet-4"}`, want: []string{"claude-sonnet-4"}, }, { name: "anthropic message start", - dialect: NewAnthropic(http.DefaultClient), + dialect: NewAnthropic(), payload: `{"type":"message_start","message":{"model":"claude-sonnet-4-20250514"}}`, want: []string{"claude-sonnet-4-20250514"}, }, { name: "gemini generate content", - dialect: NewGemini(http.DefaultClient), + dialect: NewGemini(), payload: `{"modelVersion":"gemini-2.5-pro"}`, want: []string{"gemini-2.5-pro"}, }, @@ -78,13 +77,13 @@ func TestResponseModelInspectorsIgnoreMissingInvalidAndEmptyModels(t *testing.T) dialect Dialect payload string }{ - {name: "openai missing", dialect: NewOpenAI(http.DefaultClient), payload: `{"id":"chatcmpl-1"}`}, - {name: "openai whitespace", dialect: NewOpenAI(http.DefaultClient), payload: `{"model":" \t\n"}`}, - {name: "responses null", dialect: NewOpenAIResponses(http.DefaultClient), payload: `{"model":null}`}, - {name: "anthropic empty", dialect: NewAnthropic(http.DefaultClient), payload: `{"model":""}`}, - {name: "gemini wrong type", dialect: NewGemini(http.DefaultClient), payload: `{"modelVersion":42}`}, - {name: "gemini whitespace", dialect: NewGemini(http.DefaultClient), payload: `{"modelVersion":"　"}`}, - {name: "malformed", dialect: NewOpenAI(http.DefaultClient), payload: `{"model":`}, + {name: "openai missing", dialect: NewOpenAI(), payload: `{"id":"chatcmpl-1"}`}, + {name: "openai whitespace", dialect: NewOpenAI(), payload: `{"model":" \t\n"}`}, + {name: "responses null", dialect: NewOpenAIResponses(), payload: `{"model":null}`}, + {name: "anthropic empty", dialect: NewAnthropic(), payload: `{"model":""}`}, + {name: "gemini wrong type", dialect: NewGemini(), payload: `{"modelVersion":42}`}, + {name: "gemini whitespace", dialect: NewGemini(), payload: `{"modelVersion":"　"}`}, + {name: "malformed", dialect: NewOpenAI(), payload: `{"model":`}, } for _, test := range tests { diff --git a/internal/dialect/set_test.go b/internal/dialect/set_test.go index 1cc1b6c5..843e497e 100644 --- a/internal/dialect/set_test.go +++ b/internal/dialect/set_test.go @@ -1,14 +1,13 @@ package dialect import ( - "net/http" "testing" "gpt-load/internal/protocol" ) func TestNewSetIndexesDialectsByProtocol(t *testing.T) { - openAI := NewOpenAI(http.DefaultClient) + openAI := NewOpenAI() got := NewSet(openAI) if len(got) != 1 || got[protocol.OpenAICompletions] != openAI { t.Fatalf("NewSet(OpenAI) = %#v, want OpenAI indexed", got) diff --git a/internal/dialect/upstream_url.go b/internal/dialect/upstream_url.go deleted file mode 100644 index 4c48ee4b..00000000 --- a/internal/dialect/upstream_url.go +++ /dev/null @@ -1,31 +0,0 @@ -package dialect - -import ( - "fmt" - "net/url" - "strings" -) - -func resolveUpstreamURL(baseURL, resourcePath, rawQuery string) (string, error) { - if resourcePath == "" || !strings.HasPrefix(resourcePath, "/") || strings.HasPrefix(resourcePath, "//") { - return "", fmt.Errorf("upstream resource path must start with one slash") - } - - parsed, err := url.Parse(baseURL) - if err != nil || parsed.Opaque != "" || parsed.Host == "" || - (!strings.EqualFold(parsed.Scheme, "http") && - !strings.EqualFold(parsed.Scheme, "https")) || - parsed.User != nil || parsed.Fragment != "" { - return "", fmt.Errorf("invalid upstream base URL") - } - - parsed.Path = strings.TrimRight(parsed.Path, "/") + resourcePath - parsed.RawPath = "" - if parsed.RawQuery == "" { - parsed.RawQuery = rawQuery - } else if rawQuery != "" { - parsed.RawQuery += "&" + rawQuery - } - parsed.Fragment = "" - return parsed.String(), nil -} diff --git a/internal/dialect/upstream_url_test.go b/internal/dialect/upstream_url_test.go deleted file mode 100644 index c669acb0..00000000 --- a/internal/dialect/upstream_url_test.go +++ /dev/null @@ -1,91 +0,0 @@ -package dialect - -import ( - "strings" - "testing" -) - -func TestResolveUpstreamURL(t *testing.T) { - tests := []struct { - name string - baseURL string - resourcePath string - rawQuery string - want string - wantErr bool - }{ - { - name: "complete API prefix", - baseURL: "https://api.openai.com/v1", - resourcePath: "/chat/completions", - rawQuery: "trace=true", - want: "https://api.openai.com/v1/chat/completions?trace=true", - }, - { - name: "custom path prefix and trailing slash", - baseURL: "https://proxy.example.com/tenant/openai/v1/", - resourcePath: "/responses/resp_123", - want: "https://proxy.example.com/tenant/openai/v1/responses/resp_123", - }, - { - name: "root base URL", - baseURL: "https://api.example.com", - resourcePath: "/chat/completions", - want: "https://api.example.com/chat/completions", - }, - { - name: "base and client queries keep current order", - baseURL: "https://proxy.example.com/v1?api-version=2026-01-01", - resourcePath: "/chat/completions", - rawQuery: "trace=true&filter=a%2Fb", - want: "https://proxy.example.com/v1/chat/completions?api-version=2026-01-01&trace=true&filter=a%2Fb", - }, - { - name: "empty request query preserves base query", - baseURL: "https://proxy.example.com/v1?api-version=2026-01-01", - resourcePath: "/models", - want: "https://proxy.example.com/v1/models?api-version=2026-01-01", - }, - {name: "empty resource", baseURL: "https://api.example.com/v1", wantErr: true}, - {name: "relative resource", baseURL: "https://api.example.com/v1", resourcePath: "models", wantErr: true}, - {name: "double slash resource", baseURL: "https://api.example.com/v1", resourcePath: "//models", wantErr: true}, - {name: "relative base", baseURL: "api.example.com/v1", resourcePath: "/models", wantErr: true}, - {name: "unsupported scheme", baseURL: "ftp://api.example.com/v1", resourcePath: "/models", wantErr: true}, - {name: "userinfo", baseURL: "https://user:secret@api.example.com/v1", resourcePath: "/models", wantErr: true}, - {name: "fragment", baseURL: "https://api.example.com/v1#secret", resourcePath: "/models", wantErr: true}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got, err := resolveUpstreamURL(tt.baseURL, tt.resourcePath, tt.rawQuery) - if tt.wantErr { - if err == nil { - t.Fatalf("resolveUpstreamURL() = %q, want error", got) - } - return - } - if err != nil { - t.Fatalf("resolveUpstreamURL() error = %v", err) - } - if got != tt.want { - t.Fatalf("resolveUpstreamURL() = %q, want %q", got, tt.want) - } - }) - } -} - -func TestResolveUpstreamURLDoesNotExposeMalformedBaseURL(t *testing.T) { - const secret = "query-secret" - - _, err := resolveUpstreamURL( - "https://api.example.com/v1/%zz?token="+secret, - "/models", - "", - ) - if err == nil { - t.Fatal("resolveUpstreamURL() error = nil, want malformed base URL error") - } - if strings.Contains(err.Error(), secret) { - t.Fatalf("resolveUpstreamURL() error exposes query secret: %v", err) - } -} diff --git a/internal/dialect/usage_test.go b/internal/dialect/usage_test.go index 5ee5d78a..87ad5b0b 100644 --- a/internal/dialect/usage_test.go +++ b/internal/dialect/usage_test.go @@ -3,15 +3,11 @@ package dialect import ( "bufio" "bytes" - "context" - "net/http" "os" "path/filepath" "testing" - "gpt-load/internal/health" "gpt-load/internal/protocol" - "gpt-load/internal/state" "gpt-load/internal/usage" ) @@ -40,17 +36,17 @@ func TestUsageProviderNonStreamConformance(t *testing.T) { }{ { name: "OpenAI", - extractor: NewOpenAI(http.DefaultClient), + extractor: NewOpenAI(), body: `{"usage":{"prompt_tokens":100,"completion_tokens":30,"prompt_tokens_details":{"cached_tokens":20}}}`, }, { name: "Anthropic", - extractor: NewAnthropic(http.DefaultClient), + extractor: NewAnthropic(), body: `{"usage":{"input_tokens":80,"cache_read_input_tokens":20,"output_tokens":30}}`, }, { name: "Gemini", - extractor: NewGemini(http.DefaultClient), + extractor: NewGemini(), body: `{"usageMetadata":{"promptTokenCount":100,"cachedContentTokenCount":20,"candidatesTokenCount":30}}`, }, } @@ -131,19 +127,3 @@ func (d *usageDialectOnly) Protocol() protocol.Protocol { return protocol.OpenAI func (d *usageDialectOnly) InspectRequest(*ParsedRequest) (RequestMetadata, error) { return RequestMetadata{}, nil } - -func (d *usageDialectOnly) BuildUpstreamURL(string, *ParsedRequest) (string, error) { return "", nil } - -func (d *usageDialectOnly) InjectCredential(http.Header, string) {} - -func (d *usageDialectOnly) ListModels(context.Context, string, string, state.HeaderRules) ([]string, error) { - return nil, nil -} - -func (d *usageDialectOnly) Probe(context.Context, string, string, state.HeaderRules, string) error { - return nil -} - -func (d *usageDialectOnly) ClassifyStatus(int, []byte) health.FailureCategory { - return health.FailureCategoryOK -} diff --git a/internal/execution/bifrost/cloud_providers_test.go b/internal/execution/bifrost/cloud_providers_test.go new file mode 100644 index 00000000..b9b181bd --- /dev/null +++ b/internal/execution/bifrost/cloud_providers_test.go @@ -0,0 +1,86 @@ +package bifrost + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "gpt-load/internal/execution" +) + +func TestAzureAPIKeyUsesSelectedDirectCredentialAndNormalizesResult(t *testing.T) { + t.Parallel() + + const azureKey = "azure-secret-do-not-leak" + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + if request.URL.Path != "/openai/v1/chat/completions" { + t.Errorf("path = %q", request.URL.Path) + } + if request.Header.Get("Api-Key") != azureKey || request.Header.Get("Authorization") != "" { + t.Errorf("auth headers = %#v", request.Header) + } + body, err := io.ReadAll(request.Body) + if err != nil { + t.Errorf("read body: %v", err) + return + } + var payload map[string]any + if json.Unmarshal(body, &payload) != nil || payload["model"] != "upstream-model" { + t.Errorf("payload = %s", body) + } + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Ms-Request-Id", "azure-request") + _, _ = io.WriteString(writer, `{"id":"chatcmpl-azure","object":"chat.completion","created":1,"model":"served-model","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}}`) + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec("https://unused.example") + spec.ChannelID = "azure_openai" + spec.TargetConfig = json.RawMessage(`{"endpoint":"` + server.URL + `"}`) + spec.Credential = execution.NewCredentialSnapshot(31, 1, 1, []byte(`{"api_key":"`+azureKey+`"}`)) + result := runtime.Execute(context.Background(), spec) + + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if calls.Load() != 1 || result.StatusCode != http.StatusOK || result.UpstreamRequestID != "azure-request" || result.Model != "served-model" { + t.Fatalf("calls/result = %d/%+v", calls.Load(), result) + } + assertNoPrivateLeak(t, result, azureKey) +} + +func TestAzureErrorDoesNotExposeStructuredCredentialSecrets(t *testing.T) { + t.Parallel() + + const azureKey = "azure-error-secret" + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + calls.Add(1) + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusUnauthorized) + _, _ = io.WriteString(writer, `{"error":{"type":"invalid_api_key","code":"bad_key","message":"rejected `+azureKey+`"}}`) + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec("https://unused.example") + spec.ChannelID = "azure_openai" + spec.TargetConfig = json.RawMessage(`{"endpoint":"` + server.URL + `"}`) + spec.Credential = execution.NewCredentialSnapshot(32, 1, 1, []byte(`{"api_key":"`+azureKey+`"}`)) + result := runtime.Execute(context.Background(), spec) + + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if calls.Load() != 1 || result.Error == nil || result.Error.Kind != execution.ErrorKindHTTP || result.Error.StatusCode != http.StatusUnauthorized { + t.Fatalf("calls/result = %d/%+v", calls.Load(), result) + } + assertNoPrivateLeak(t, result, azureKey) +} diff --git a/internal/execution/bifrost/compatible_matrix_test.go b/internal/execution/bifrost/compatible_matrix_test.go new file mode 100644 index 00000000..92f23639 --- /dev/null +++ b/internal/execution/bifrost/compatible_matrix_test.go @@ -0,0 +1,330 @@ +package bifrost + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestCompatibleAnthropicAndGeminiNativeUnaryStream(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID channel.ID + prefix string + unaryPath string + streamPath string + unaryBody string + streamBody string + unarySpec func() execution.AttemptSpec + streamSpec func() execution.AttemptSpec + credential string + streamQuery string + }{ + { + name: "Anthropic", channelID: channel.AnthropicCompatible, prefix: "/tenant/v1", + unaryPath: "/tenant/v1/messages", streamPath: "/tenant/v1/messages", credential: "X-Api-Key", + unaryBody: anthropicResponsesConvertedFixture, streamBody: anthropicResponsesStreamFixture, + unarySpec: func() execution.AttemptSpec { return anthropicSpec(false) }, + streamSpec: func() execution.AttemptSpec { return anthropicSpec(true) }, + }, + { + name: "Gemini", channelID: channel.GeminiCompatible, prefix: "/tenant/v1beta", + unaryPath: "/tenant/v1beta/models/upstream-gemini:generateContent", + streamPath: "/tenant/v1beta/models/upstream-gemini:streamGenerateContent", credential: "X-Goog-Api-Key", + unaryBody: geminiResponsesConvertedFixture, streamBody: geminiResponsesStreamFixture, streamQuery: "sse", + unarySpec: func() execution.AttemptSpec { return geminiSpec(false) }, + streamSpec: func() execution.AttemptSpec { return geminiSpec(true) }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Header.Get(test.credential) != testAPIKey || request.Header.Get("Authorization") != "" { + t.Errorf("credential headers = %#v", request.Header) + } + body, _ := io.ReadAll(request.Body) + var payload map[string]any + _ = json.Unmarshal(body, &payload) + isStream := request.URL.Path == test.streamPath && (test.streamPath != test.unaryPath || payload["stream"] == true) + if isStream { + if test.streamQuery != "" && request.URL.Query().Get("alt") != test.streamQuery { + t.Errorf("stream query = %q", request.URL.RawQuery) + } + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(writer, test.streamBody) + return + } + if request.URL.Path != test.unaryPath { + t.Errorf("unary path = %q", request.URL.Path) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, test.unaryBody) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true}) + target, _ := json.Marshal(map[string]string{"base_url": server.URL + test.prefix}) + unarySpec := test.unarySpec() + unarySpec.ChannelID = string(test.channelID) + unarySpec.TargetConfig = target + unarySpec.ClientModel = unarySpec.UpstreamModel + unary := runtime.Execute(context.Background(), unarySpec) + if err := unary.Validate(); err != nil || unary.Error != nil || string(unary.Body) != test.unaryBody || unary.Usage == nil { + t.Fatalf("unary = %+v err=%v body=%s", unary, err, unary.Body) + } + + streamSpec := test.streamSpec() + streamSpec.ChannelID = string(test.channelID) + streamSpec.TargetConfig = target + streamSpec.ClientModel = streamSpec.UpstreamModel + if test.streamQuery != "" { + streamSpec.Query.Set("alt", test.streamQuery) + } + var data bytes.Buffer + stream := runtime.ExecuteStream(context.Background(), streamSpec, func(event execution.StreamEvent) error { + if event.Kind == execution.StreamEventData { + data.Write(event.Data) + } + return nil + }) + if err := stream.Validate(); err != nil || stream.Error != nil || data.String() != test.streamBody || stream.Usage == nil { + t.Fatalf("stream = %+v err=%v data=%s", stream, err, data.String()) + } + }) + } +} + +func TestCompatibleOpenAIResponsesCreateStreamAndLifecycle(t *testing.T) { + t.Parallel() + + const unaryResponse = `{"id":"resp_1","object":"response","status":"completed","model":"served","output":[],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7}}` + const streamResponse = "event: response.completed\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"model\":\"served\",\"usage\":{\"input_tokens\":5,\"output_tokens\":2,\"total_tokens\":7}}}\n\n" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if !strings.HasPrefix(request.URL.Path, "/tenant/v1/responses") || request.Header.Get("Authorization") != "Bearer "+testAPIKey { + t.Errorf("request = %s %#v", request.URL.Path, request.Header) + } + if request.Method == http.MethodGet { + writer.WriteHeader(http.StatusPartialContent) + _, _ = io.WriteString(writer, `{"id":"resp_1","status":"completed"}`) + return + } + body, _ := io.ReadAll(request.Body) + var payload map[string]any + _ = json.Unmarshal(body, &payload) + if payload["stream"] == true { + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(writer, streamResponse) + return + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, unaryResponse) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true}) + target, _ := json.Marshal(map[string]string{"base_url": server.URL + "/tenant/v1"}) + create := openAIResponsesSpec(execution.OperationResponsesCreate, http.MethodPost, "/v1/responses") + create.ChannelID = string(channel.OpenAICompatible) + create.TargetConfig = target + create.ClientModel = create.UpstreamModel + unary := runtime.Execute(context.Background(), create) + if err := unary.Validate(); err != nil || unary.Error != nil || string(unary.Body) != unaryResponse { + t.Fatalf("unary = %+v err=%v body=%s", unary, err, unary.Body) + } + var data bytes.Buffer + stream := runtime.ExecuteStream(context.Background(), create, func(event execution.StreamEvent) error { + if event.Kind == execution.StreamEventData { + data.Write(event.Data) + } + return nil + }) + if err := stream.Validate(); err != nil || stream.Error != nil || data.String() != streamResponse { + t.Fatalf("stream = %+v err=%v data=%s", stream, err, data.String()) + } + retrieve := openAIResponsesSpec(execution.OperationResponsesRetrieve, http.MethodGet, "/v1/responses/resp_1") + retrieve.ChannelID = string(channel.OpenAICompatible) + retrieve.TargetConfig = target + retrieve.ClientModel, retrieve.UpstreamModel, retrieve.Body = "", "", nil + lifecycle := runtime.Execute(context.Background(), retrieve) + if err := lifecycle.Validate(); err != nil || lifecycle.StatusCode != http.StatusPartialContent || lifecycle.Error != nil { + t.Fatalf("lifecycle = %+v err=%v", lifecycle, err) + } +} + +func TestCompatibleListModelsAndProbeUseConfiguredPrefix(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID channel.ID + protocol protocol.Protocol + baseSuffix string + modelsPath string + probePath string + probeBody []byte + }{ + {name: "OpenAI", channelID: channel.OpenAICompatible, protocol: protocol.OpenAICompletions, baseSuffix: "/custom/v1", modelsPath: "/custom/v1/models", probePath: "/custom/v1/chat/completions", probeBody: []byte(`{"model":"client","messages":[{"role":"user","content":"ping"}]}`)}, + {name: "Anthropic", channelID: channel.AnthropicCompatible, protocol: protocol.Anthropic, baseSuffix: "/custom/v1", modelsPath: "/custom/v1/models", probePath: "/custom/v1/messages", probeBody: []byte(`{"model":"client","max_tokens":1,"messages":[{"role":"user","content":"ping"}]}`)}, + {name: "Gemini", channelID: channel.GeminiCompatible, protocol: protocol.Gemini, baseSuffix: "/custom/v1beta", modelsPath: "/custom/v1beta/models", probePath: "/custom/v1beta/models/probe-upstream:generateContent", probeBody: []byte(`{"contents":[{"parts":[{"text":"ping"}]}]}`)}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case test.modelsPath: + if request.Method != http.MethodGet || request.URL.Query().Get("page") != "next" { + t.Errorf("models target = %s %s?%s", request.Method, request.URL.Path, request.URL.RawQuery) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"data":[]}`) + case test.probePath: + if request.Method != http.MethodPost { + t.Errorf("probe method = %s", request.Method) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"ok":true}`) + default: + t.Errorf("unexpected path = %s", request.URL.Path) + writer.WriteHeader(http.StatusNotFound) + } + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true}) + target, _ := json.Marshal(map[string]string{"base_url": server.URL + test.baseSuffix}) + modelsPath := "/v1/models" + if test.protocol == protocol.Gemini { + modelsPath = "/v1beta/models" + } + models := utilitySpec(test.channelID, test.protocol, execution.OperationListModels, http.MethodGet, modelsPath, nil) + models.TargetConfig = target + models.Query.Set("page", "next") + modelsResult := runtime.Execute(context.Background(), models) + if err := modelsResult.Validate(); err != nil || modelsResult.Error != nil { + t.Fatalf("models = %+v err=%v", modelsResult, err) + } + + probeClientPath := "/v1/chat/completions" + if test.protocol == protocol.Anthropic { + probeClientPath = "/v1/messages" + } + if test.protocol == protocol.Gemini { + probeClientPath = "/v1beta/models/client:generateContent" + } + probe := utilitySpec(test.channelID, test.protocol, execution.OperationProbe, http.MethodPost, probeClientPath, test.probeBody) + probe.TargetConfig = target + probe.ClientModel, probe.UpstreamModel = "probe-client", "probe-upstream" + probeResult := runtime.Execute(context.Background(), probe) + if err := probeResult.Validate(); err != nil || probeResult.Error != nil { + t.Fatalf("probe = %+v err=%v", probeResult, err) + } + }) + } +} + +func TestOpenAICompatibleNonV1PrefixKeepsListModelsAndProbeFunctional(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "application/json") + switch request.URL.Path { + case "/vendor/api/v4/models": + _, _ = io.WriteString(writer, `{"object":"list","data":[{"id":"model-1","object":"model","created":1,"owned_by":"vendor"}]}`) + case "/vendor/api/v4/chat/completions": + _, _ = io.WriteString(writer, `{"id":"chat_1","object":"chat.completion","created":1,"model":"served","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`) + default: + t.Errorf("path = %s", request.URL.Path) + writer.WriteHeader(http.StatusNotFound) + } + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true}) + target, _ := json.Marshal(map[string]string{"base_url": server.URL + "/vendor/api/v4"}) + models := utilitySpec(channel.OpenAICompatible, protocol.OpenAICompletions, execution.OperationListModels, http.MethodGet, "/v1/models", nil) + models.TargetConfig = target + modelsResult := runtime.Execute(context.Background(), models) + if err := modelsResult.Validate(); err != nil || modelsResult.Error != nil || !bytes.Contains(modelsResult.Body, []byte(`"model-1"`)) { + t.Fatalf("models = %+v err=%v body=%s", modelsResult, err, modelsResult.Body) + } + + probe := utilitySpec(channel.OpenAICompatible, protocol.OpenAICompletions, execution.OperationProbe, http.MethodPost, "/v1/chat/completions", []byte(`{"model":"client","messages":[{"role":"user","content":"ping"}]}`)) + probe.TargetConfig = target + probe.ClientModel, probe.UpstreamModel = "probe-client", "probe-upstream" + probeResult := runtime.Execute(context.Background(), probe) + if err := probeResult.Validate(); err != nil || probeResult.Error != nil || !bytes.Contains(probeResult.Body, []byte(`"chat_1"`)) { + t.Fatalf("probe = %+v err=%v body=%s", probeResult, err, probeResult.Body) + } +} + +func TestOpenAICompatibleConvertedResponsesUsesChatCompletions(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Method != http.MethodPost || request.URL.Path != "/vendor/api/v4/chat/completions" { + t.Errorf("request = %s %s", request.Method, request.URL.Path) + writer.WriteHeader(http.StatusNotFound) + return + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"id":"chat_1","object":"chat.completion","created":1,"model":"served","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true}) + target, _ := json.Marshal(map[string]string{"base_url": server.URL + "/vendor/api/v4"}) + spec := openAIResponsesSpec(execution.OperationResponsesCreate, http.MethodPost, "/v1/responses") + spec.ChannelID = string(channel.OpenAICompatible) + spec.TargetConfig = target + spec.ClientModel = "client-model" + spec.UpstreamModel = "upstream-model" + spec.Body = []byte(`{"model":"client-model","input":"hello","store":false}`) + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil || result.Error != nil { + t.Fatalf("result = %+v err=%v body=%s", result, err, result.Body) + } + if !bytes.Contains(result.Body, []byte(`"object":"response"`)) { + t.Fatalf("converted Responses body = %s", result.Body) + } +} + +func TestListModelsConvertsProviderResultToClientProtocol(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/v1/models" { + t.Errorf("path = %s", request.URL.Path) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"object":"list","data":[{"id":"model-1","object":"model","created":1,"owned_by":"vendor"}]}`) + })) + defer server.Close() + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: server.URL}) + spec := utilitySpec(channel.OpenAI, protocol.Anthropic, execution.OperationListModels, http.MethodGet, "/v1/models", nil) + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil || result.Error != nil { + t.Fatalf("result = %+v err=%v", result, err) + } + var body map[string]any + if err := json.Unmarshal(result.Body, &body); err != nil { + t.Fatal(err) + } + if _, exists := body["has_more"]; !exists { + t.Fatalf("Anthropic model list shape missing: %s", result.Body) + } + if strings.Contains(string(result.Body), "gptload-custom-") || strings.Contains(string(result.Body), "openai/model-1") { + t.Fatalf("internal provider leaked: %s", result.Body) + } +} diff --git a/internal/execution/bifrost/converted.go b/internal/execution/bifrost/converted.go new file mode 100644 index 00000000..15b62c71 --- /dev/null +++ b/internal/execution/bifrost/converted.go @@ -0,0 +1,523 @@ +package bifrost + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + + "github.com/maximhq/bifrost/core/providers/anthropic" + "github.com/maximhq/bifrost/core/providers/gemini" + "github.com/maximhq/bifrost/core/providers/openai" + "github.com/maximhq/bifrost/core/schemas" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +type responsesUnarySDKResult struct { + response *schemas.BifrostResponsesResponse + err *schemas.BifrostError +} + +type responsesStreamSDKResult struct { + stream chan *schemas.BifrostStreamChunk + err *schemas.BifrostError +} + +func buildConvertedResponsesRequest(spec execution.AttemptSpec, provider schemas.ModelProvider) (*schemas.BifrostResponsesRequest, error) { + conversionContext := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + defer conversionContext.Cancel() + + var request *schemas.BifrostResponsesRequest + switch spec.ClientProtocol { + case protocol.OpenAIResponses: + if spec.Operation != execution.OperationResponsesCreate { + return nil, fmt.Errorf("converted Responses operation is not supported") + } + var wire openai.OpenAIResponsesRequest + if err := json.Unmarshal(spec.Body, &wire); err != nil { + return nil, fmt.Errorf("invalid OpenAI Responses request body") + } + request = wire.ToBifrostResponsesRequest(conversionContext) + case protocol.Anthropic: + if spec.Operation != execution.OperationChatCompletion && spec.Operation != execution.OperationProbe { + return nil, fmt.Errorf("converted Anthropic operation is not supported") + } + var wire anthropic.AnthropicMessageRequest + if err := json.Unmarshal(spec.Body, &wire); err != nil { + return nil, fmt.Errorf("invalid Anthropic request body") + } + request = wire.ToBifrostResponsesRequest(conversionContext) + case protocol.Gemini: + if spec.Operation != execution.OperationChatCompletion && spec.Operation != execution.OperationProbe { + return nil, fmt.Errorf("converted Gemini operation is not supported") + } + var wire gemini.GeminiGenerationRequest + if err := json.Unmarshal(spec.Body, &wire); err != nil { + return nil, fmt.Errorf("invalid Gemini request body") + } + wire.Model = spec.ClientModel + request = wire.ToBifrostResponsesRequest(conversionContext) + default: + return nil, fmt.Errorf("converted protocol is not supported") + } + if request == nil || len(request.Input) == 0 { + return nil, fmt.Errorf("converted request input is required") + } + request.Provider = provider + request.Model = spec.UpstreamModel + request.Fallbacks = nil + request.RawRequestBody = nil + stripResponsesControlParams(request.Params) + return request, nil +} + +func stripResponsesControlParams(params *schemas.ResponsesParameters) { + if params == nil || len(params.ExtraParams) == 0 { + return + } + for key := range params.ExtraParams { + normalized := strings.ReplaceAll(strings.ToLower(key), "-", "_") + switch normalized { + case "provider", "fallback", "fallbacks", "authorization", "proxy_authorization", + "api_key", "apikey", "x_api_key", "x_goog_api_key": + delete(params.ExtraParams, key) + } + } +} + +func (r *Runtime) executeConvertedResponses( + parent context.Context, + spec execution.AttemptSpec, + prepared preparedAttempt, +) execution.AttemptResult { + requestContext, requestCancel := boundedRequestContext(parent, spec.Timeouts.Request) + defer requestCancel() + callContext, callCancel := context.WithCancel(requestContext) + defer callCancel() + preResponse := startPreResponseGate(callCancel, spec.Timeouts) + defer preResponse.stop() + + bifrostContext := newSDKContext(callContext, spec, prepared.directKey) + setTypedRequestURL(bifrostContext, prepared.typedURL) + outcomeChannel := make(chan responsesUnarySDKResult, 1) + go func() { + response, bifrostError := r.core.ResponsesRequest(bifrostContext, prepared.responsesRequest) + outcomeChannel <- responsesUnarySDKResult{response: response, err: bifrostError} + }() + + var outcome responsesUnarySDKResult + select { + case outcome = <-outcomeChannel: + if preResponse.expired() || requestContext.Err() != nil { + return unaryContextFailure(requestContext, preResponse.expired()) + } + case <-callContext.Done(): + return unaryContextFailure(requestContext, preResponse.expired()) + } + preResponse.stop() + + if outcome.err != nil { + return unaryErrorResult(outcome.err, bifrostContext, prepared.secrets) + } + if outcome.response == nil { + return attemptedUnaryFailure(execution.ErrorKindInternal, "execution runtime returned no response") + } + headers := responseHeaders(outcome.response.ExtraFields.ProviderResponseHeaders, bifrostContext, false) + body, usageEvidence, err := encodeConvertedResponsesResponse(prepared.clientProtocol, bifrostContext, outcome.response) + if err != nil { + return startedUnaryFailure(http.StatusOK, headers, execution.ErrorKindInternal, "encode converted response") + } + if needsClientModelAlias(spec) { + body, err = rewriteClientResponseModel(spec.ClientProtocol, body, spec.ClientModel) + if err != nil { + return startedUnaryFailure(http.StatusOK, headers, execution.ErrorKindInternal, "rewrite converted response model") + } + } + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: headers, + Body: body, + Model: outcome.response.Model, + UpstreamRequestID: upstreamRequestID(headers), + Usage: usageEvidence, + } +} + +func encodeConvertedResponsesResponse( + clientProtocol protocol.Protocol, + ctx *schemas.BifrostContext, + response *schemas.BifrostResponsesResponse, +) ([]byte, *execution.UsageEvidence, error) { + if response == nil { + return nil, nil, fmt.Errorf("response is nil") + } + usageEvidence, err := usageEvidenceFromResponses(response.Usage) + if err != nil { + return nil, nil, err + } + var wire any + switch clientProtocol { + case protocol.OpenAIResponses: + wire = response.WithDefaults() + case protocol.Anthropic: + wire = anthropic.ToAnthropicResponsesResponse(ctx, response) + case protocol.Gemini: + wire = gemini.ToGeminiResponsesResponse(response) + default: + return nil, nil, fmt.Errorf("unsupported converted response protocol") + } + if wire == nil { + return nil, nil, fmt.Errorf("response conversion returned nil") + } + body, err := marshalClientWire(clientProtocol, wire) + if err != nil { + return nil, nil, err + } + return body, usageEvidence, nil +} + +func marshalClientWire(clientProtocol protocol.Protocol, value any) ([]byte, error) { + body, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("marshal converted response") + } + if clientProtocol != protocol.OpenAIResponses { + return body, nil + } + var object map[string]json.RawMessage + if err := json.Unmarshal(body, &object); err != nil { + return nil, fmt.Errorf("sanitize converted OpenAI response") + } + delete(object, "extra_fields") + delete(object, "provider_extra_fields") + if rawUsage := object["usage"]; len(rawUsage) > 0 { + var usageObject map[string]json.RawMessage + if json.Unmarshal(rawUsage, &usageObject) == nil { + delete(usageObject, "cost") + object["usage"], _ = json.Marshal(usageObject) + } + } + body, err = json.Marshal(object) + if err != nil { + return nil, fmt.Errorf("marshal sanitized OpenAI response") + } + return body, nil +} + +func convertedTypedTarget( + providerKind channel.ProviderKind, + baseURL string, + model string, + responses bool, + stream bool, + rawQuery string, +) (string, error) { + var resourcePath string + switch providerKind { + case channel.ProviderOpenAI, channel.ProviderOpenAICompatible: + resourcePath = "/chat/completions" + // Compatible targets are configured with Chat Completions only. Bifrost + // translates a converted Responses request back to that operation; the + // Responses path is reserved for native OpenAI targets. + if responses && providerKind == channel.ProviderOpenAI { + resourcePath = "/responses" + } + if providerKind == channel.ProviderOpenAI { + resourcePath = "/v1" + resourcePath + } + case channel.ProviderAnthropic, channel.ProviderAnthropicCompatible: + resourcePath = "/messages" + if providerKind == channel.ProviderAnthropic { + resourcePath = "/v1" + resourcePath + } + case channel.ProviderGemini, channel.ProviderGeminiCompatible: + action := "generateContent" + if stream { + action = "streamGenerateContent" + } + resourcePath = "/models/" + url.PathEscape(model) + ":" + action + if stream { + rawQuery = setRawQueryValue(rawQuery, "alt", "sse") + } + case channel.ProviderAzureOpenAI, channel.ProviderAWSBedrock, channel.ProviderGoogleVertex: + if rawQuery != "" { + return "", fmt.Errorf("provider-specific query is not supported") + } + return "", nil + default: + return "", fmt.Errorf("unsupported provider kind") + } + if baseURL != "" && providerKind != channel.ProviderGeminiCompatible { + return resolveTypedTargetURL(baseURL, resourcePath, rawQuery) + } + return appendTypedQuery(resourcePath, rawQuery), nil +} + +func convertedListModelsTarget(providerKind channel.ProviderKind, baseURL string, rawQuery string) (string, error) { + resourcePath := "/models" + switch providerKind { + case channel.ProviderOpenAI, channel.ProviderAnthropic: + resourcePath = "/v1/models" + case channel.ProviderGemini: + resourcePath = "/models" + case channel.ProviderOpenAICompatible, channel.ProviderAnthropicCompatible: + if baseURL == "" { + return "", fmt.Errorf("compatible target is required") + } + return resolveTypedTargetURL(baseURL, resourcePath, rawQuery) + case channel.ProviderGeminiCompatible: + resourcePath = "/models" + case channel.ProviderAzureOpenAI, channel.ProviderAWSBedrock, channel.ProviderGoogleVertex: + if rawQuery != "" { + return "", fmt.Errorf("provider-specific query is not supported") + } + return "", nil + default: + return "", fmt.Errorf("unsupported provider kind") + } + return appendTypedQuery(resourcePath, rawQuery), nil +} + +func appendTypedQuery(path, rawQuery string) string { + if rawQuery == "" { + return path + } + return path + "?" + rawQuery +} + +func setRawQueryValue(source, key, value string) string { + kept := removeRawQueryValue(source, key) + if kept == "" { + return url.QueryEscape(key) + "=" + url.QueryEscape(value) + } + return kept + "&" + url.QueryEscape(key) + "=" + url.QueryEscape(value) +} + +func removeRawQueryValue(source, key string) string { + segments := strings.Split(source, "&") + kept := make([]string, 0, len(segments)) + for _, segment := range segments { + rawKey := segment + if separator := strings.IndexByte(rawKey, '='); separator >= 0 { + rawKey = rawKey[:separator] + } + decodedKey, err := url.QueryUnescape(rawKey) + if err == nil && strings.EqualFold(decodedKey, key) { + continue + } + if source != "" || segment != "" { + kept = append(kept, segment) + } + } + return strings.Join(kept, "&") +} + +func (r *Runtime) executeConvertedResponsesStream( + parent context.Context, + spec execution.AttemptSpec, + prepared preparedAttempt, + sink execution.StreamSink, +) execution.StreamResult { + requestContext, requestCancel := boundedRequestContext(parent, spec.Timeouts.Request) + defer requestCancel() + callContext, callCancel := context.WithCancel(requestContext) + defer callCancel() + preResponse := startPreResponseGate(callCancel, spec.Timeouts) + defer preResponse.stop() + + bifrostContext := newSDKContext(callContext, spec, prepared.directKey) + setTypedRequestURL(bifrostContext, prepared.typedURL) + outcomeChannel := make(chan responsesStreamSDKResult, 1) + go func() { + stream, bifrostError := r.core.ResponsesStreamRequest(bifrostContext, prepared.responsesRequest) + outcomeChannel <- responsesStreamSDKResult{stream: stream, err: bifrostError} + }() + + var outcome responsesStreamSDKResult + select { + case outcome = <-outcomeChannel: + if preResponse.expired() || requestContext.Err() != nil { + return streamContextFailure(requestContext, preResponse.expired(), false, nil, "", "", nil) + } + case <-callContext.Done(): + return streamContextFailure(requestContext, preResponse.expired(), false, nil, "", "", nil) + } + preResponse.stop() + if outcome.err != nil { + return streamErrorResult(outcome.err, bifrostContext, prepared.secrets, false, 0, nil, "", nil) + } + if outcome.stream == nil { + return attemptedStreamFailure(execution.ErrorKindInternal, "execution runtime returned no stream") + } + + headers := responseHeaders(nil, bifrostContext, true) + requestID := upstreamRequestID(headers) + if err := sink(execution.StreamEvent{ + Sequence: 1, + Kind: execution.StreamEventReady, + StatusCode: http.StatusOK, + Header: headers.Clone(), + UpstreamRequestID: requestID, + }); err != nil { + callCancel() + return streamSinkFailure(headers, requestID, "", nil) + } + + encoder := newConvertedResponsesStreamEncoder(prepared.clientProtocol) + sequence := uint64(1) + model := spec.UpstreamModel + var usageEvidence *execution.UsageEvidence + idleTimer := newIdleTimer(spec.Timeouts.StreamIdle) + defer idleTimer.stop() + + for { + select { + case <-requestContext.Done(): + callCancel() + return streamContextFailure(requestContext, false, true, headers, requestID, model, usageEvidence) + case <-idleTimer.channel(): + callCancel() + return streamContextFailure(requestContext, true, true, headers, requestID, model, usageEvidence) + case chunk, open := <-outcome.stream: + idleTimer.pause() + if !open { + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: headers, + Model: model, + UpstreamRequestID: requestID, + Usage: usageEvidence, + } + } + if chunk == nil || chunk.BifrostResponsesStreamResponse == nil { + callCancel() + if chunk != nil && chunk.BifrostError != nil { + return streamErrorResult(chunk.BifrostError, bifrostContext, prepared.secrets, true, http.StatusOK, headers, model, usageEvidence) + } + return streamErrorResult(nil, bifrostContext, prepared.secrets, true, http.StatusOK, headers, model, usageEvidence) + } + + response := chunk.BifrostResponsesStreamResponse + var chunkUsage *execution.UsageEvidence + if response.Response != nil { + if response.Response.Model != "" { + model = response.Response.Model + } + var err error + chunkUsage, err = usageEvidenceFromResponses(response.Response.Usage) + if err != nil { + callCancel() + return streamErrorResult(nil, bifrostContext, prepared.secrets, true, http.StatusOK, headers, model, usageEvidence) + } + } + frames, err := encoder.encode(bifrostContext, response) + if err != nil { + callCancel() + return streamErrorResult(nil, bifrostContext, prepared.secrets, true, http.StatusOK, headers, model, usageEvidence) + } + for _, frame := range frames { + if needsClientModelAlias(spec) { + frame, err = rewriteClientSSEEvent(frame, spec.ClientProtocol, spec.ClientModel) + if err != nil { + callCancel() + return streamAliasFailure(http.StatusOK, headers, requestID, model, usageEvidence) + } + } + sequence++ + if err := sink(execution.StreamEvent{Sequence: sequence, Kind: execution.StreamEventData, Data: frame}); err != nil { + callCancel() + return streamSinkFailure(headers, requestID, model, usageEvidence) + } + } + if chunkUsage != nil { + usageEvidence = cloneUsage(chunkUsage) + sequence++ + if err := sink(execution.StreamEvent{Sequence: sequence, Kind: execution.StreamEventUsage, Usage: cloneUsage(chunkUsage)}); err != nil { + callCancel() + return streamSinkFailure(headers, requestID, model, usageEvidence) + } + } + idleTimer.resume() + } + } +} + +type convertedResponsesStreamEncoder struct { + clientProtocol protocol.Protocol + geminiState *gemini.BifrostToGeminiStreamState +} + +func newConvertedResponsesStreamEncoder(clientProtocol protocol.Protocol) *convertedResponsesStreamEncoder { + encoder := &convertedResponsesStreamEncoder{clientProtocol: clientProtocol} + if clientProtocol == protocol.Gemini { + encoder.geminiState = gemini.NewBifrostToGeminiStreamState() + } + return encoder +} + +func (e *convertedResponsesStreamEncoder) encode( + ctx *schemas.BifrostContext, + response *schemas.BifrostResponsesStreamResponse, +) ([][]byte, error) { + switch e.clientProtocol { + case protocol.OpenAIResponses: + wire := response.WithDefaults() + if wire == nil { + return nil, nil + } + body, err := marshalClientWire(protocol.OpenAIResponses, wire) + if err != nil { + return nil, err + } + return [][]byte{frameNamedSSE(string(wire.Type), body)}, nil + case protocol.Anthropic: + events := anthropic.ToAnthropicResponsesStreamResponse(ctx, response) + frames := make([][]byte, 0, len(events)) + for _, event := range events { + if event == nil { + continue + } + body, err := json.Marshal(event) + if err != nil { + return nil, fmt.Errorf("marshal Anthropic stream event") + } + frames = append(frames, frameNamedSSE(string(event.Type), body)) + } + return frames, nil + case protocol.Gemini: + wire := gemini.ToGeminiResponsesStreamResponse(response, e.geminiState) + if wire == nil { + return nil, nil + } + body, err := json.Marshal(wire) + if err != nil { + return nil, fmt.Errorf("marshal Gemini stream response") + } + return [][]byte{frameSSE(body)}, nil + default: + return nil, fmt.Errorf("unsupported converted stream protocol") + } +} + +func frameNamedSSE(eventName string, payload []byte) []byte { + if eventName == "" { + return frameSSE(payload) + } + framed := make([]byte, 0, len(eventName)+len(payload)+16) + framed = append(framed, "event: "...) + framed = append(framed, eventName...) + framed = append(framed, '\n') + framed = append(framed, "data: "...) + framed = append(framed, payload...) + framed = append(framed, '\n', '\n') + return framed +} diff --git a/internal/execution/bifrost/converted_protocols_test.go b/internal/execution/bifrost/converted_protocols_test.go new file mode 100644 index 00000000..ac2b9e17 --- /dev/null +++ b/internal/execution/bifrost/converted_protocols_test.go @@ -0,0 +1,601 @@ +package bifrost + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestConvertedResponsesUnaryUsesCanonicalSDKConverters(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID channel.ID + clientProtocol protocol.Protocol + operation execution.Operation + clientPath string + clientBody []byte + upstreamPath string + credentialName string + upstreamBody string + modelInPath bool + assertClient func(*testing.T, map[string]any) + runtime func(*testing.T, string) *Runtime + }{ + { + name: "Anthropic client to OpenAI", + channelID: channel.OpenAI, + clientProtocol: protocol.Anthropic, + operation: execution.OperationChatCompletion, + clientPath: "/v1/messages", + clientBody: []byte(`{"model":"client-model","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`), + upstreamPath: "/v1/responses", + credentialName: "Authorization", + upstreamBody: openAIResponsesConvertedFixture, + assertClient: func(t *testing.T, body map[string]any) { + if body["type"] != "message" || body["role"] != "assistant" || body["model"] != "client-model" { + t.Errorf("Anthropic response = %#v", body) + } + }, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: base}) + }, + }, + { + name: "Anthropic client to Gemini", + channelID: channel.Gemini, + clientProtocol: protocol.Anthropic, + operation: execution.OperationChatCompletion, + clientPath: "/v1/messages", + clientBody: []byte(`{"model":"client-model","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`), + upstreamPath: "/v1beta/models/upstream-model:generateContent", + credentialName: "X-Goog-Api-Key", + upstreamBody: geminiResponsesConvertedFixture, + modelInPath: true, + assertClient: func(t *testing.T, body map[string]any) { + if body["type"] != "message" || body["role"] != "assistant" || body["model"] != "client-model" { + t.Errorf("Anthropic response = %#v", body) + } + }, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, geminiBaseURL: base + "/v1beta"}) + }, + }, + { + name: "Gemini client to OpenAI", + channelID: channel.OpenAI, + clientProtocol: protocol.Gemini, + operation: execution.OperationChatCompletion, + clientPath: "/v1beta/models/client-model:generateContent", + clientBody: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + upstreamPath: "/v1/responses", + credentialName: "Authorization", + upstreamBody: openAIResponsesConvertedFixture, + assertClient: func(t *testing.T, body map[string]any) { + candidates, _ := body["candidates"].([]any) + if len(candidates) != 1 || body["modelVersion"] != "client-model" { + t.Errorf("Gemini response = %#v", body) + } + }, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: base}) + }, + }, + { + name: "Gemini client to Anthropic", + channelID: channel.Anthropic, + clientProtocol: protocol.Gemini, + operation: execution.OperationChatCompletion, + clientPath: "/v1beta/models/client-model:generateContent", + clientBody: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + upstreamPath: "/v1/messages", + credentialName: "X-Api-Key", + upstreamBody: anthropicResponsesConvertedFixture, + assertClient: func(t *testing.T, body map[string]any) { + candidates, _ := body["candidates"].([]any) + if len(candidates) != 1 || body["modelVersion"] != "client-model" { + t.Errorf("Gemini response = %#v", body) + } + }, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, anthropicBaseURL: base}) + }, + }, + { + name: "OpenAI Responses client to Anthropic", + channelID: channel.Anthropic, + clientProtocol: protocol.OpenAIResponses, + operation: execution.OperationResponsesCreate, + clientPath: "/v1/responses", + clientBody: []byte(`{"model":"client-model","input":"hello","max_output_tokens":16}`), + upstreamPath: "/v1/messages", + credentialName: "X-Api-Key", + upstreamBody: anthropicResponsesConvertedFixture, + assertClient: func(t *testing.T, body map[string]any) { + if body["object"] != "response" || body["model"] != "client-model" { + t.Errorf("OpenAI Responses response = %#v", body) + } + }, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, anthropicBaseURL: base}) + }, + }, + { + name: "OpenAI Responses client to Gemini", + channelID: channel.Gemini, + clientProtocol: protocol.OpenAIResponses, + operation: execution.OperationResponsesCreate, + clientPath: "/v1/responses", + clientBody: []byte(`{"model":"client-model","input":"hello","max_output_tokens":16}`), + upstreamPath: "/v1beta/models/upstream-model:generateContent", + credentialName: "X-Goog-Api-Key", + upstreamBody: geminiResponsesConvertedFixture, + modelInPath: true, + assertClient: func(t *testing.T, body map[string]any) { + if body["object"] != "response" || body["model"] != "client-model" { + t.Errorf("OpenAI Responses response = %#v", body) + } + }, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, geminiBaseURL: base + "/v1beta"}) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Method != http.MethodPost || request.URL.Path != test.upstreamPath || request.URL.Query().Get("trace") != "kept" { + t.Errorf("upstream target = %s %s?%s", request.Method, request.URL.Path, request.URL.RawQuery) + } + wantCredential := testAPIKey + if test.credentialName == "Authorization" { + wantCredential = "Bearer " + testAPIKey + } + if request.Header.Get(test.credentialName) != wantCredential || request.Header.Get("X-Api-Key") == "client-secret" { + t.Errorf("credential headers = %#v", request.Header) + } + requestBody, _ := io.ReadAll(request.Body) + var upstreamRequest map[string]any + if err := json.Unmarshal(requestBody, &upstreamRequest); err != nil { + t.Fatalf("decode upstream request: %v", err) + } + if (!test.modelInPath && upstreamRequest["model"] != "upstream-model") || upstreamRequest["provider"] != nil || upstreamRequest["fallbacks"] != nil { + t.Errorf("upstream request = %#v", upstreamRequest) + } + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Request-Id", "converted-request") + _, _ = io.WriteString(writer, test.upstreamBody) + })) + defer server.Close() + + runtime := test.runtime(t, server.URL) + spec := convertedSpec(test.channelID, test.clientProtocol, test.operation, test.clientPath, test.clientBody) + spec.Query.Set("trace", "kept") + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.Error != nil || result.StatusCode != http.StatusOK || result.UpstreamRequestID != "converted-request" { + t.Fatalf("result = %+v, body=%s", result, result.Body) + } + var clientBody map[string]any + if err := json.Unmarshal(result.Body, &clientBody); err != nil { + t.Fatalf("decode client response: %v; body=%s", err, result.Body) + } + test.assertClient(t, clientBody) + if result.Usage == nil { + t.Fatal("normalized usage is missing") + } + }) + } +} + +func TestConvertedOpenAIChatUnaryUsesSelectedProvider(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID channel.ID + upstreamPath string + credentialName string + upstreamBody string + modelInPath bool + runtime func(*testing.T, string) *Runtime + }{ + { + name: "Anthropic", channelID: channel.Anthropic, upstreamPath: "/v1/messages", + credentialName: "X-Api-Key", upstreamBody: anthropicResponsesConvertedFixture, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, anthropicBaseURL: base}) + }, + }, + { + name: "Gemini", channelID: channel.Gemini, upstreamPath: "/v1beta/models/upstream-model:generateContent", + credentialName: "X-Goog-Api-Key", upstreamBody: geminiResponsesConvertedFixture, modelInPath: true, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, geminiBaseURL: base + "/v1beta"}) + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != test.upstreamPath || request.URL.Query().Get("trace") != "kept" { + t.Errorf("target = %s?%s", request.URL.Path, request.URL.RawQuery) + } + if request.Header.Get(test.credentialName) != testAPIKey { + t.Errorf("credential %s = %q", test.credentialName, request.Header.Get(test.credentialName)) + } + body, _ := io.ReadAll(request.Body) + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("decode request: %v", err) + } + if !test.modelInPath && payload["model"] != "upstream-model" { + t.Errorf("request model = %#v", payload["model"]) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, test.upstreamBody) + })) + defer server.Close() + + runtime := test.runtime(t, server.URL) + spec := convertedSpec( + test.channelID, + protocol.OpenAICompletions, + execution.OperationChatCompletion, + "/v1/chat/completions", + []byte(`{"model":"client-model","messages":[{"role":"user","content":"hello"}]}`), + ) + spec.Query.Set("trace", "kept") + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + var body map[string]any + if err := json.Unmarshal(result.Body, &body); err != nil { + t.Fatalf("decode response: %v", err) + } + choices, _ := body["choices"].([]any) + if result.Error != nil || len(choices) != 1 || result.Usage == nil || body["model"] != "client-model" { + t.Fatalf("result/body = %+v/%#v", result, body) + } + }) + } +} + +func TestConvertedResponsesStreamUsesClientProtocolFraming(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID channel.ID + clientProtocol protocol.Protocol + operation execution.Operation + clientPath string + clientBody []byte + upstreamPath string + credentialName string + upstreamStream string + upstreamModel string + streamInBody bool + wantFragments []string + runtime func(*testing.T, string) *Runtime + }{ + { + name: "Anthropic client from OpenAI stream", channelID: channel.OpenAI, + clientProtocol: protocol.Anthropic, operation: execution.OperationChatCompletion, + clientPath: "/v1/messages", clientBody: []byte(`{"model":"client-model","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`), + upstreamPath: "/v1/responses", credentialName: "Authorization", upstreamStream: openAIResponsesStreamFixture, + upstreamModel: "gpt-upstream", + streamInBody: true, + wantFragments: []string{"event: message_start", `"type":"message_start"`, "hello", "event: message_stop"}, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: base}) + }, + }, + { + name: "Gemini client from OpenAI stream", channelID: channel.OpenAI, + clientProtocol: protocol.Gemini, operation: execution.OperationChatCompletion, + clientPath: "/v1beta/models/client-model:streamGenerateContent", clientBody: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + upstreamPath: "/v1/responses", credentialName: "Authorization", upstreamStream: openAIResponsesStreamFixture, + upstreamModel: "gpt-upstream", + streamInBody: true, + wantFragments: []string{"data: ", "hello", "usageMetadata"}, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: base}) + }, + }, + { + name: "OpenAI Responses client from Anthropic stream", channelID: channel.Anthropic, + clientProtocol: protocol.OpenAIResponses, operation: execution.OperationResponsesCreate, + clientPath: "/v1/responses", clientBody: []byte(`{"model":"client-model","input":"hello","max_output_tokens":16}`), + upstreamPath: "/v1/messages", credentialName: "X-Api-Key", upstreamStream: anthropicResponsesStreamFixture, + upstreamModel: "claude-upstream", + streamInBody: true, + wantFragments: []string{"event: response.created", "response.output_text.delta", "hello", "event: response.completed"}, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, anthropicBaseURL: base}) + }, + }, + { + name: "OpenAI Responses client from Gemini stream", channelID: channel.Gemini, + clientProtocol: protocol.OpenAIResponses, operation: execution.OperationResponsesCreate, + clientPath: "/v1/responses", clientBody: []byte(`{"model":"client-model","input":"hello","max_output_tokens":16}`), + upstreamPath: "/v1beta/models/upstream-model:streamGenerateContent", credentialName: "X-Goog-Api-Key", upstreamStream: geminiResponsesStreamFixture, + upstreamModel: "gemini-upstream", + wantFragments: []string{"event: response.created", "response.output_text.delta", "hello", "event: response.completed"}, + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, geminiBaseURL: base + "/v1beta"}) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != test.upstreamPath || request.URL.Query().Get("trace") != "kept" { + t.Errorf("target = %s?%s", request.URL.Path, request.URL.RawQuery) + } + wantCredential := testAPIKey + if test.credentialName == "Authorization" { + wantCredential = "Bearer " + testAPIKey + } + if request.Header.Get(test.credentialName) != wantCredential { + t.Errorf("credential = %q", request.Header.Get(test.credentialName)) + } + body, _ := io.ReadAll(request.Body) + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil || (test.streamInBody && payload["stream"] != true) { + t.Errorf("stream request = %s err=%v", body, err) + } + writer.Header().Set("Content-Type", "text/event-stream") + writer.Header().Set("X-Request-Id", "converted-stream") + _, _ = io.WriteString(writer, test.upstreamStream) + })) + defer server.Close() + + runtime := test.runtime(t, server.URL) + spec := convertedSpec(test.channelID, test.clientProtocol, test.operation, test.clientPath, test.clientBody) + spec.Query.Set("trace", "kept") + var data bytes.Buffer + ready := 0 + usageEvents := 0 + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + if err := event.Validate(); err != nil { + t.Fatalf("event validation: %v; event=%+v", err, event) + } + switch event.Kind { + case execution.StreamEventReady: + ready++ + if data.Len() != 0 || event.UpstreamRequestID != "converted-stream" { + t.Errorf("ready ordering/id = %d/%q", data.Len(), event.UpstreamRequestID) + } + case execution.StreamEventData: + data.Write(event.Data) + case execution.StreamEventUsage: + usageEvents++ + } + return nil + }) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v data=%s", err, result, data.String()) + } + if result.Error != nil || ready != 1 || usageEvents == 0 || result.Usage == nil { + t.Fatalf("result/events = %+v ready=%d usage=%d data=%s", result, ready, usageEvents, data.String()) + } + for _, fragment := range test.wantFragments { + if !strings.Contains(data.String(), fragment) { + t.Errorf("stream missing %q: %s", fragment, data.String()) + } + } + if !strings.Contains(data.String(), `"client-model"`) || strings.Contains(data.String(), `"`+test.upstreamModel+`"`) { + t.Errorf("client model alias was not applied: %s", data.String()) + } + }) + } +} + +func TestConvertedOpenAIChatStreamUsesSelectedProvider(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID channel.ID + upstreamPath string + credentialName string + upstreamStream string + upstreamModel string + runtime func(*testing.T, string) *Runtime + }{ + { + name: "Anthropic", channelID: channel.Anthropic, upstreamPath: "/v1/messages", + credentialName: "X-Api-Key", upstreamStream: anthropicResponsesStreamFixture, upstreamModel: "claude-upstream", + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, anthropicBaseURL: base}) + }, + }, + { + name: "Gemini", channelID: channel.Gemini, upstreamPath: "/v1beta/models/upstream-model:streamGenerateContent", + credentialName: "X-Goog-Api-Key", upstreamStream: geminiResponsesStreamFixture, upstreamModel: "gemini-upstream", + runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, geminiBaseURL: base + "/v1beta"}) + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != test.upstreamPath || request.URL.Query().Get("trace") != "kept" { + t.Errorf("target = %s?%s", request.URL.Path, request.URL.RawQuery) + } + if request.Header.Get(test.credentialName) != testAPIKey { + t.Errorf("credential = %q", request.Header.Get(test.credentialName)) + } + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(writer, test.upstreamStream) + })) + defer server.Close() + + runtime := test.runtime(t, server.URL) + spec := convertedSpec( + test.channelID, + protocol.OpenAICompletions, + execution.OperationChatCompletion, + "/v1/chat/completions", + []byte(`{"model":"client-model","stream":true,"messages":[{"role":"user","content":"hello"}]}`), + ) + spec.Query.Set("trace", "kept") + var data bytes.Buffer + usageEvents := 0 + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + if event.Kind == execution.StreamEventData { + data.Write(event.Data) + } + if event.Kind == execution.StreamEventUsage { + usageEvents++ + } + return nil + }) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v data=%s", err, result, data.String()) + } + if result.Error != nil || !strings.Contains(data.String(), "hello") || !strings.Contains(data.String(), "data: [DONE]\n\n") || + !strings.Contains(data.String(), `"client-model"`) || strings.Contains(data.String(), `"`+test.upstreamModel+`"`) || usageEvents == 0 { + t.Fatalf("result/data/usage = %+v/%s/%d", result, data.String(), usageEvents) + } + }) + } +} + +func TestConvertedResponsesStreamHonorsCancellationAndUpstreamError(t *testing.T) { + t.Parallel() + + t.Run("context cancellation", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(writer, "event: response.created\n"+ + "data: {\"type\":\"response.created\",\"sequence_number\":0,\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":123,\"status\":\"in_progress\",\"model\":\"gpt-upstream\",\"output\":[]}}\n\n") + if flusher, ok := writer.(http.Flusher); ok { + flusher.Flush() + } + <-request.Context().Done() + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: server.URL}) + spec := convertedSpec( + channel.OpenAI, + protocol.Anthropic, + execution.OperationChatCompletion, + "/v1/messages", + []byte(`{"model":"client-model","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`), + ) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var data bytes.Buffer + result := runtime.ExecuteStream(ctx, spec, func(event execution.StreamEvent) error { + if event.Kind == execution.StreamEventData { + data.Write(event.Data) + cancel() + } + return nil + }) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.Error == nil || result.Error.Kind != execution.ErrorKindCanceled || !result.ResponseStarted { + t.Fatalf("canceled result = %+v", result) + } + if strings.Contains(data.String(), "message_stop") { + t.Fatalf("terminal event emitted after cancellation: %s", data.String()) + } + }) + + t.Run("upstream HTTP error", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Request-Id", "converted-429") + writer.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(writer, `{"error":{"type":"rate_limit_error","code":"rate_limited","message":"rejected `+testAPIKey+`"}}`) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: server.URL}) + spec := convertedSpec( + channel.OpenAI, + protocol.Anthropic, + execution.OperationChatCompletion, + "/v1/messages", + []byte(`{"model":"client-model","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`), + ) + var events []execution.StreamEvent + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + return nil + }) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.Error == nil || result.Error.Kind != execution.ErrorKindHTTP || result.StatusCode != http.StatusTooManyRequests { + t.Fatalf("HTTP error result = %+v", result) + } + if len(events) != 0 { + t.Fatalf("converted error emitted client data: %+v", events) + } + assertNoPrivateLeak(t, result, testAPIKey, "gptload-custom-") + }) +} + +func convertedSpec(channelID channel.ID, clientProtocol protocol.Protocol, operation execution.Operation, path string, body []byte) execution.AttemptSpec { + return execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: "converted-request", AttemptID: "converted-attempt", Sequence: 1, + ChannelID: string(channelID), ClientProtocol: clientProtocol, Operation: operation, + ClientModel: "client-model", UpstreamModel: "upstream-model", + Method: http.MethodPost, Path: path, Query: make(map[string][]string), + Header: http.Header{"Authorization": {"Bearer client-secret"}, "X-Api-Key": {"client-secret"}, "X-Test": {"kept"}}, + Body: body, TargetConfig: json.RawMessage(`{}`), + Credential: execution.NewCredentialSnapshot(12, 1, 1, []byte(`{"api_key":"`+testAPIKey+`"}`)), + }) +} + +const openAIResponsesConvertedFixture = `{"id":"resp_1","object":"response","created_at":123,"status":"completed","model":"gpt-upstream","output":[{"id":"msg_1","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"hello","annotations":[]}]}],"usage":{"input_tokens":4,"input_tokens_details":{"cached_tokens":1},"output_tokens":2,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":6}}` + +const anthropicResponsesConvertedFixture = `{"id":"msg_1","type":"message","role":"assistant","model":"claude-upstream","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":4,"cache_read_input_tokens":1,"output_tokens":2}}` + +const geminiResponsesConvertedFixture = `{"candidates":[{"content":{"parts":[{"text":"hello"}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":4,"cachedContentTokenCount":1,"candidatesTokenCount":2,"totalTokenCount":6},"modelVersion":"gemini-upstream","responseId":"resp_1"}` + +const openAIResponsesStreamFixture = "event: response.created\n" + + "data: {\"type\":\"response.created\",\"sequence_number\":0,\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":123,\"status\":\"in_progress\",\"model\":\"gpt-upstream\",\"output\":[]}}\n\n" + + "event: response.output_text.delta\n" + + "data: {\"type\":\"response.output_text.delta\",\"sequence_number\":1,\"output_index\":0,\"content_index\":0,\"item_id\":\"msg_1\",\"delta\":\"hello\"}\n\n" + + "event: response.completed\n" + + "data: {\"type\":\"response.completed\",\"sequence_number\":2,\"response\":" + openAIResponsesConvertedFixture + "}\n\n" + + "data: [DONE]\n\n" + +const anthropicResponsesStreamFixture = "event: message_start\n" + + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-upstream\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":4,\"cache_read_input_tokens\":1,\"output_tokens\":0}}}\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\":\"hello\"}}\n\n" + + "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\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":2}}\n\n" + + "event: message_stop\n" + + "data: {\"type\":\"message_stop\"}\n\n" + +const geminiResponsesStreamFixture = "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hello\"}],\"role\":\"model\"},\"index\":0}],\"modelVersion\":\"gemini-upstream\",\"responseId\":\"resp_1\"}\n\n" + + "data: {\"candidates\":[{\"content\":{\"parts\":[],\"role\":\"model\"},\"finishReason\":\"STOP\",\"index\":0}],\"usageMetadata\":{\"promptTokenCount\":4,\"cachedContentTokenCount\":1,\"candidatesTokenCount\":2,\"totalTokenCount\":6},\"modelVersion\":\"gemini-upstream\",\"responseId\":\"resp_1\"}\n\n" diff --git a/internal/execution/bifrost/executor.go b/internal/execution/bifrost/executor.go new file mode 100644 index 00000000..e571ad7f --- /dev/null +++ b/internal/execution/bifrost/executor.go @@ -0,0 +1,1135 @@ +package bifrost + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/url" + "strings" + "time" + + "github.com/maximhq/bifrost/core/providers/openai" + "github.com/maximhq/bifrost/core/schemas" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +const ( + openAIChatPath = "/v1/chat/completions" + defaultRequestBudget = 5 * time.Minute +) + +type preparedAttempt struct { + provider schemas.ModelProvider + mode channel.RouteMode + request *schemas.BifrostChatRequest + responsesRequest *schemas.BifrostResponsesRequest + listModelsRequest *schemas.BifrostListModelsRequest + passthrough *schemas.BifrostPassthroughRequest + typedURL string + clientProtocol protocol.Protocol + directKey schemas.Key + secrets []string +} + +type unarySDKResult struct { + response *schemas.BifrostChatResponse + err *schemas.BifrostError +} + +type streamSDKResult struct { + stream chan *schemas.BifrostStreamChunk + err *schemas.BifrostError +} + +// Execute executes one non-streaming attempt. +func (r *Runtime) Execute(parent context.Context, spec execution.AttemptSpec) execution.AttemptResult { + prepared, preflightError := r.prepare(spec, false) + if preflightError != nil { + return *preflightError + } + if prepared.mode == channel.RouteNative { + return r.executeNative(parent, spec, prepared) + } + if prepared.responsesRequest != nil { + return r.executeConvertedResponses(parent, spec, prepared) + } + if prepared.listModelsRequest != nil { + return r.executeListModels(parent, spec, prepared) + } + + requestContext, requestCancel := boundedRequestContext(parent, spec.Timeouts.Request) + defer requestCancel() + callContext, callCancel := context.WithCancel(requestContext) + defer callCancel() + preResponse := startPreResponseGate(callCancel, spec.Timeouts) + defer preResponse.stop() + + bifrostContext := newSDKContext(callContext, spec, prepared.directKey) + setTypedRequestURL(bifrostContext, prepared.typedURL) + outcomeChannel := make(chan unarySDKResult, 1) + go func() { + response, bifrostError := r.core.ChatCompletionRequest(bifrostContext, prepared.request) + outcomeChannel <- unarySDKResult{response: response, err: bifrostError} + }() + + var outcome unarySDKResult + select { + case outcome = <-outcomeChannel: + if preResponse.expired() || requestContext.Err() != nil { + return unaryContextFailure(requestContext, preResponse.expired()) + } + case <-callContext.Done(): + return unaryContextFailure(requestContext, preResponse.expired()) + } + preResponse.stop() + + if outcome.err != nil { + return unaryErrorResult(outcome.err, bifrostContext, prepared.secrets) + } + if outcome.response == nil { + return attemptedUnaryFailure(execution.ErrorKindInternal, "execution runtime returned no response") + } + body, usageEvidence, err := encodeChatResponse(outcome.response) + if err != nil { + return startedUnaryFailure( + http.StatusOK, + responseHeaders(outcome.response.ExtraFields.ProviderResponseHeaders, bifrostContext, false), + execution.ErrorKindInternal, + "encode upstream response", + ) + } + headers := responseHeaders(outcome.response.ExtraFields.ProviderResponseHeaders, bifrostContext, false) + if needsClientModelAlias(spec) { + body, err = rewriteClientResponseModel(spec.ClientProtocol, body, spec.ClientModel) + if err != nil { + return startedUnaryFailure(http.StatusOK, headers, execution.ErrorKindInternal, "rewrite converted response model") + } + } + requestID := upstreamRequestID(headers) + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: headers, + Body: body, + Model: outcome.response.Model, + UpstreamRequestID: requestID, + Usage: usageEvidence, + } +} + +// ExecuteStream executes one streaming attempt and emits ordered events synchronously. +func (r *Runtime) ExecuteStream( + parent context.Context, + spec execution.AttemptSpec, + sink execution.StreamSink, +) execution.StreamResult { + if sink == nil { + return notSentStreamFailure(execution.ErrorKindInvalidRequest, "stream sink is required") + } + prepared, preflightError := r.prepare(spec, true) + if preflightError != nil { + return streamFromAttemptFailure(*preflightError) + } + if prepared.mode == channel.RouteNative { + return r.executeNativeStream(parent, spec, prepared, sink) + } + if prepared.responsesRequest != nil { + return r.executeConvertedResponsesStream(parent, spec, prepared, sink) + } + + requestContext, requestCancel := boundedRequestContext(parent, spec.Timeouts.Request) + defer requestCancel() + callContext, callCancel := context.WithCancel(requestContext) + defer callCancel() + preResponse := startPreResponseGate(callCancel, spec.Timeouts) + defer preResponse.stop() + + bifrostContext := newSDKContext(callContext, spec, prepared.directKey) + setTypedRequestURL(bifrostContext, prepared.typedURL) + outcomeChannel := make(chan streamSDKResult, 1) + go func() { + stream, bifrostError := r.core.ChatCompletionStreamRequest(bifrostContext, prepared.request) + outcomeChannel <- streamSDKResult{stream: stream, err: bifrostError} + }() + + var outcome streamSDKResult + select { + case outcome = <-outcomeChannel: + if preResponse.expired() || requestContext.Err() != nil { + return streamContextFailure(requestContext, preResponse.expired(), false, nil, "", "", nil) + } + case <-callContext.Done(): + return streamContextFailure(requestContext, preResponse.expired(), false, nil, "", "", nil) + } + preResponse.stop() + if outcome.err != nil { + return streamErrorResult(outcome.err, bifrostContext, prepared.secrets, false, 0, nil, "", nil) + } + if outcome.stream == nil { + return attemptedStreamFailure(execution.ErrorKindInternal, "execution runtime returned no stream") + } + + headers := responseHeaders(nil, bifrostContext, true) + requestID := upstreamRequestID(headers) + ready := execution.StreamEvent{ + Sequence: 1, + Kind: execution.StreamEventReady, + StatusCode: http.StatusOK, + Header: headers.Clone(), + UpstreamRequestID: requestID, + } + if err := sink(ready); err != nil { + callCancel() + return streamSinkFailure(headers, requestID, "", nil) + } + + sequence := uint64(1) + model := "" + var usageEvidence *execution.UsageEvidence + idleTimer := newIdleTimer(spec.Timeouts.StreamIdle) + defer idleTimer.stop() + + for { + select { + case <-requestContext.Done(): + callCancel() + return streamContextFailure(requestContext, false, true, headers, requestID, model, usageEvidence) + case <-idleTimer.channel(): + callCancel() + return streamContextFailure(requestContext, true, true, headers, requestID, model, usageEvidence) + case chunk, open := <-outcome.stream: + idleTimer.pause() + if !open { + sequence++ + if err := sink(execution.StreamEvent{ + Sequence: sequence, + Kind: execution.StreamEventData, + Data: []byte("data: [DONE]\n\n"), + }); err != nil { + callCancel() + return streamSinkFailure(headers, requestID, model, usageEvidence) + } + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: headers, + Model: model, + UpstreamRequestID: requestID, + Usage: usageEvidence, + } + } + if chunk == nil { + callCancel() + return streamErrorResult(nil, bifrostContext, prepared.secrets, true, http.StatusOK, headers, model, usageEvidence) + } + if chunk.BifrostError != nil { + callCancel() + return streamErrorResult(chunk.BifrostError, bifrostContext, prepared.secrets, true, http.StatusOK, headers, model, usageEvidence) + } + if chunk.BifrostChatResponse == nil { + callCancel() + return streamErrorResult(nil, bifrostContext, prepared.secrets, true, http.StatusOK, headers, model, usageEvidence) + } + + response := chunk.BifrostChatResponse + if response.Model != "" { + model = response.Model + } + payload, chunkUsage, err := encodeChatResponse(response) + if err != nil { + callCancel() + return streamErrorResult(nil, bifrostContext, prepared.secrets, true, http.StatusOK, headers, model, usageEvidence) + } + if needsClientModelAlias(spec) { + payload, err = rewriteClientResponseModel(spec.ClientProtocol, payload, spec.ClientModel) + if err != nil { + callCancel() + return streamAliasFailure(http.StatusOK, headers, requestID, model, usageEvidence) + } + } + sequence++ + if err := sink(execution.StreamEvent{ + Sequence: sequence, + Kind: execution.StreamEventData, + Data: frameSSE(payload), + }); err != nil { + callCancel() + return streamSinkFailure(headers, requestID, model, usageEvidence) + } + if chunkUsage != nil { + usageEvidence = cloneUsage(chunkUsage) + sequence++ + if err := sink(execution.StreamEvent{ + Sequence: sequence, + Kind: execution.StreamEventUsage, + Usage: cloneUsage(chunkUsage), + }); err != nil { + callCancel() + return streamSinkFailure(headers, requestID, model, usageEvidence) + } + } + idleTimer.resume() + } + } +} + +func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAttempt, *execution.AttemptResult) { + if r == nil || r.core == nil || r.registry == nil { + failure := notSentUnaryFailure(execution.ErrorKindInternal, "execution runtime is unavailable") + return preparedAttempt{}, &failure + } + if r.closed.Load() { + failure := notSentUnaryFailure(execution.ErrorKindInternal, "execution runtime is shut down") + return preparedAttempt{}, &failure + } + if err := spec.Validate(); err != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid execution attempt: "+safeValidationReason(err)) + return preparedAttempt{}, &failure + } + if !supportedRequestShape(spec, stream) { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "unsupported protocol or operation") + return preparedAttempt{}, &failure + } + + channelID := channel.ID(spec.ChannelID) + if len(bytes.TrimSpace(spec.TargetConfig)) == 0 { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "channel target is required") + return preparedAttempt{}, &failure + } + resolved, err := r.registry.ResolveExecutionTarget(channelID, spec.TargetConfig) + if err != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid channel target: "+safeValidationReason(err)) + return preparedAttempt{}, &failure + } + mode, ok := resolved.Mode(spec.ClientProtocol, spec.Operation) + if !ok { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "channel does not support the requested route") + return preparedAttempt{}, &failure + } + if spec.Operation == execution.OperationResponsesPassthrough && + resolved.ProviderKind != channel.ProviderOpenAI && + resolved.ProviderKind != channel.ProviderOpenAICompatible { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "Responses passthrough requires an OpenAI provider") + return preparedAttempt{}, &failure + } + var provider schemas.ModelProvider + targetBaseURL := "" + targetBaseQuery := "" + nativePassthroughEligible := true + switch resolved.ProviderKind { + case channel.ProviderOpenAI: + provider = schemas.OpenAI + case channel.ProviderAnthropic: + provider = schemas.Anthropic + case channel.ProviderGemini: + provider = schemas.Gemini + case channel.ProviderAzureOpenAI: + provider = schemas.Azure + case channel.ProviderAWSBedrock: + provider = schemas.Bedrock + case channel.ProviderGoogleVertex: + provider = schemas.Vertex + case channel.ProviderOpenAICompatible, channel.ProviderAnthropicCompatible, channel.ProviderGeminiCompatible: + var target struct { + BaseURL string `json:"base_url"` + } + if err := json.Unmarshal(resolved.TargetConfig, &target); err != nil || target.BaseURL == "" { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid compatible channel target") + return preparedAttempt{}, &failure + } + baseProvider := schemas.OpenAI + switch resolved.ProviderKind { + case channel.ProviderAnthropicCompatible: + baseProvider = schemas.Anthropic + case channel.ProviderGeminiCompatible: + baseProvider = schemas.Gemini + } + var targetErr error + targetBaseURL, targetBaseQuery, targetErr = splitCompatibleBaseURL(target.BaseURL) + if targetErr != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid compatible channel target") + return preparedAttempt{}, &failure + } + provider = customProviderKey(baseProvider, target.BaseURL) + r.account.setConfig(provider, providerConfig(targetBaseURL, true, baseProvider, r.allowPrivate)) + // Core v1.7.7 hardcodes /v1 for OpenAI passthrough. A complete + // compatible prefix with any other suffix must keep the typed path so + // RequestPathOverrides can address the configured endpoint exactly. + if baseProvider == schemas.OpenAI { + nativePassthroughEligible = strings.HasSuffix(targetBaseURL, "/v1") + } + default: + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "unsupported channel") + return preparedAttempt{}, &failure + } + if mode == channel.RouteNative && + (spec.Operation == execution.OperationListModels || spec.Operation == execution.OperationProbe) && + !providerKindNativeForClient(resolved.ProviderKind, spec.ClientProtocol) { + mode = channel.RouteConverted + } + if mode == channel.RouteNative && !nativePassthroughEligible { + mode = channel.RouteConverted + } + safeQuery := mergeRawQueries(targetBaseQuery, safeAttemptQuery(spec)) + if resolved.ProviderKind == channel.ProviderGemini || + resolved.ProviderKind == channel.ProviderGeminiCompatible { + safeQuery = removeRawQueryValue(safeQuery, "alt") + if stream { + safeQuery = setRawQueryValue(safeQuery, "alt", "sse") + } + } + + credential, err := r.registry.ValidateCredential(channelID, spec.Credential.Data()) + if err != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid credential: "+safeValidationReason(err)) + return preparedAttempt{}, &failure + } + directKey, secrets, err := directKeyForAttempt(spec, resolved, credential) + if err != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid credential") + return preparedAttempt{}, &failure + } + if mode == channel.RouteNative { + body, err := sanitizeNativeRequestBody(spec, stream) + if err != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid native request body") + return preparedAttempt{}, &failure + } + passthroughPath, err := nativePassthroughPath(spec, resolved.ProviderKind) + if err != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid native request path") + return preparedAttempt{}, &failure + } + return preparedAttempt{ + provider: provider, + mode: mode, + passthrough: &schemas.BifrostPassthroughRequest{ + Provider: provider, + Model: spec.UpstreamModel, + Method: spec.Method, + Path: passthroughPath, + RawQuery: safeQuery, + Body: body, + SafeHeaders: safePassthroughHeaders(spec.Header), + }, + directKey: directKey, + secrets: secrets, + }, nil + } + if spec.Operation == execution.OperationListModels { + typedURL, targetErr := convertedListModelsTarget(resolved.ProviderKind, targetBaseURL, safeQuery) + if targetErr != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid compatible channel target") + return preparedAttempt{}, &failure + } + return preparedAttempt{ + provider: provider, + mode: mode, + listModelsRequest: &schemas.BifrostListModelsRequest{Provider: provider}, + typedURL: typedURL, + clientProtocol: spec.ClientProtocol, + directKey: directKey, + secrets: secrets, + }, nil + } + if spec.ClientProtocol != protocol.OpenAICompletions || + (spec.Operation != execution.OperationChatCompletion && spec.Operation != execution.OperationProbe) { + request, conversionErr := buildConvertedResponsesRequest(spec, provider) + if conversionErr != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, conversionErr.Error()) + return preparedAttempt{}, &failure + } + typedURL, targetErr := convertedTypedTarget(resolved.ProviderKind, targetBaseURL, spec.UpstreamModel, true, stream, safeQuery) + if targetErr != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid compatible channel target") + return preparedAttempt{}, &failure + } + return preparedAttempt{ + provider: provider, + mode: mode, + responsesRequest: request, + typedURL: typedURL, + clientProtocol: spec.ClientProtocol, + directKey: directKey, + secrets: secrets, + }, nil + } + + var openAIRequest openai.OpenAIChatRequest + if err := json.Unmarshal(spec.Body, &openAIRequest); err != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid OpenAI chat request body") + return preparedAttempt{}, &failure + } + conversionContext := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + request := openAIRequest.ToBifrostChatRequest(conversionContext) + conversionContext.Cancel() + if request == nil || request.Input == nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "OpenAI chat messages are required") + return preparedAttempt{}, &failure + } + request.Provider = provider + request.Model = spec.UpstreamModel + request.Fallbacks = nil + request.RawRequestBody = nil + if request.Params != nil && request.Params.ExtraParams != nil { + delete(request.Params.ExtraParams, "provider") + delete(request.Params.ExtraParams, "fallback") + delete(request.Params.ExtraParams, "fallbacks") + } + typedURL, err := convertedTypedTarget(resolved.ProviderKind, targetBaseURL, spec.UpstreamModel, false, stream, safeQuery) + if err != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid compatible channel target") + return preparedAttempt{}, &failure + } + return preparedAttempt{ + provider: provider, mode: mode, request: request, typedURL: typedURL, + clientProtocol: spec.ClientProtocol, directKey: directKey, secrets: secrets, + }, nil +} + +func providerKindNativeForClient(providerKind channel.ProviderKind, clientProtocol protocol.Protocol) bool { + switch providerKind { + case channel.ProviderOpenAI, channel.ProviderOpenAICompatible: + return clientProtocol == protocol.OpenAICompletions || clientProtocol == protocol.OpenAIResponses + case channel.ProviderAnthropic, channel.ProviderAnthropicCompatible: + return clientProtocol == protocol.Anthropic + case channel.ProviderGemini, channel.ProviderGeminiCompatible: + return clientProtocol == protocol.Gemini + default: + return false + } +} + +func supportedRequestShape(spec execution.AttemptSpec, stream bool) bool { + if spec.Operation == execution.OperationListModels { + if stream || spec.Method != http.MethodGet { + return false + } + switch spec.ClientProtocol { + case protocol.OpenAICompletions, protocol.OpenAIResponses, protocol.Anthropic: + return spec.Path == "/v1/models" + case protocol.Gemini: + return spec.Path == "/v1beta/models" + default: + return false + } + } + if spec.Operation == execution.OperationProbe { + if stream || spec.Method != http.MethodPost || strings.TrimSpace(spec.UpstreamModel) == "" { + return false + } + switch spec.ClientProtocol { + case protocol.OpenAICompletions: + return spec.Path == openAIChatPath + case protocol.OpenAIResponses: + return spec.Path == "/v1/responses" + case protocol.Anthropic: + return spec.Path == "/v1/messages" + case protocol.Gemini: + return validGeminiGeneratePath(spec.Path, "generateContent") + default: + return false + } + } + switch spec.ClientProtocol { + case protocol.OpenAICompletions: + return spec.Operation == execution.OperationChatCompletion && + spec.Method == http.MethodPost && spec.Path == openAIChatPath + case protocol.OpenAIResponses: + switch spec.Operation { + case execution.OperationResponsesCreate: + return spec.Method == http.MethodPost && spec.Path == "/v1/responses" + case execution.OperationResponsesRetrieve: + return spec.Method == http.MethodGet && validResponsesResourcePath(spec.Path, "") + case execution.OperationResponsesDelete: + return spec.Method == http.MethodDelete && validResponsesResourcePath(spec.Path, "") + case execution.OperationResponsesCancel: + return spec.Method == http.MethodPost && validResponsesResourcePath(spec.Path, "cancel") + case execution.OperationResponsesInputItems: + return spec.Method == http.MethodGet && validResponsesResourcePath(spec.Path, "input_items") + case execution.OperationResponsesCompact: + return !stream && spec.Method == http.MethodPost && spec.Path == "/v1/responses/compact" + case execution.OperationResponsesInputTokens: + return !stream && spec.Method == http.MethodPost && spec.Path == "/v1/responses/input_tokens" + case execution.OperationResponsesPassthrough: + return validResponsesPassthroughShape(spec, stream) + } + case protocol.Anthropic: + return spec.Operation == execution.OperationChatCompletion && spec.Method == http.MethodPost && spec.Path == "/v1/messages" + case protocol.Gemini: + if spec.Operation != execution.OperationChatCompletion || spec.Method != http.MethodPost { + return false + } + action := "generateContent" + if stream { + action = "streamGenerateContent" + } + return validGeminiGeneratePath(spec.Path, action) + } + return false +} + +func validResponsesPassthroughShape(spec execution.AttemptSpec, stream bool) bool { + remainder, ok := strings.CutPrefix(spec.Path, "/v1/responses/") + if !ok || remainder == "" { + return false + } + switch spec.Method { + case http.MethodOptions, http.MethodConnect, http.MethodTrace: + return false + case http.MethodHead: + return !stream + default: + return true + } +} + +func validGeminiGeneratePath(path, action string) bool { + prefix := "/v1beta/models/" + remainder, ok := strings.CutPrefix(path, prefix) + if !ok { + return false + } + model, suffix, ok := strings.Cut(remainder, ":") + return ok && suffix == action && validResourceID(model) && !strings.Contains(model, "/") +} + +func nativePassthroughPath(spec execution.AttemptSpec, providerKind channel.ProviderKind) (string, error) { + switch providerKind { + case channel.ProviderOpenAI, channel.ProviderOpenAICompatible: + return spec.Path, nil + case channel.ProviderAnthropic: + return spec.Path, nil + case channel.ProviderAnthropicCompatible: + path, ok := strings.CutPrefix(spec.Path, "/v1") + if !ok || path == "" { + return "", fmt.Errorf("invalid Anthropic path") + } + return path, nil + case channel.ProviderGemini, channel.ProviderGeminiCompatible: + path, ok := strings.CutPrefix(spec.Path, "/v1beta") + if !ok || path == "" { + return "", fmt.Errorf("invalid Gemini path") + } + if spec.Operation == execution.OperationListModels { + if path != "/models" { + return "", fmt.Errorf("invalid Gemini models path") + } + return path, nil + } + marker := "/models/" + modelStart := strings.Index(path, marker) + colon := strings.LastIndex(path, ":") + if modelStart < 0 || colon <= modelStart+len(marker) { + return "", fmt.Errorf("invalid Gemini model path") + } + return path[:modelStart+len(marker)] + url.PathEscape(spec.UpstreamModel) + path[colon:], nil + default: + return "", fmt.Errorf("unsupported native channel") + } +} + +func validResponsesResourcePath(path, action string) bool { + remainder, ok := strings.CutPrefix(path, "/v1/responses/") + if !ok || remainder == "" { + return false + } + parts := strings.Split(remainder, "/") + if action == "" { + return len(parts) == 1 && validResourceID(parts[0]) + } + return len(parts) == 2 && validResourceID(parts[0]) && parts[1] == action +} + +func validResourceID(value string) bool { + return value != "" && value != "." && value != ".." && !strings.ContainsAny(value, "\\?#") +} + +func customProviderKey(baseProvider schemas.ModelProvider, baseURL string) schemas.ModelProvider { + digest := sha256.Sum256([]byte(string(baseProvider) + "\x00" + baseURL)) + return schemas.ModelProvider("gptload-custom-" + hex.EncodeToString(digest[:])) +} + +func directKeyForAttempt( + spec execution.AttemptSpec, + resolved channel.ResolvedTarget, + credential channel.Credential, +) (schemas.Key, []string, error) { + key := schemas.Key{ + ID: fmt.Sprintf("%d:%d:%d", spec.Credential.ID, spec.Credential.IdentityGeneration, spec.Credential.Version), + Name: "selected credential", + Models: schemas.WhiteList{"*"}, + Weight: 1, + } + secrets := credentialSecrets(credential) + apiKey, _ := credential.Value("api_key") + switch resolved.ProviderKind { + case channel.ProviderOpenAI, channel.ProviderAnthropic, channel.ProviderGemini, + channel.ProviderOpenAICompatible, channel.ProviderAnthropicCompatible, channel.ProviderGeminiCompatible: + if apiKey == "" { + return schemas.Key{}, nil, fmt.Errorf("api_key is required") + } + key.Value = plainSecret(apiKey) + case channel.ProviderAzureOpenAI: + var target struct { + Endpoint string `json:"endpoint"` + } + if json.Unmarshal(resolved.TargetConfig, &target) != nil || target.Endpoint == "" { + return schemas.Key{}, nil, fmt.Errorf("azure endpoint is required") + } + key.AzureKeyConfig = &schemas.AzureKeyConfig{Endpoint: plainSecret(target.Endpoint)} + if apiKey != "" { + key.Value = plainSecret(apiKey) + } else { + clientID, clientOK := credential.Value("client_id") + clientSecret, secretOK := credential.Value("client_secret") + tenantID, tenantOK := credential.Value("tenant_id") + if !clientOK || !secretOK || !tenantOK { + return schemas.Key{}, nil, fmt.Errorf("complete Entra credentials are required") + } + key.AzureKeyConfig.ClientID = plainSecretPtr(clientID) + key.AzureKeyConfig.ClientSecret = plainSecretPtr(clientSecret) + key.AzureKeyConfig.TenantID = plainSecretPtr(tenantID) + } + case channel.ProviderAWSBedrock: + var target struct { + Region string `json:"region"` + } + if json.Unmarshal(resolved.TargetConfig, &target) != nil || target.Region == "" { + return schemas.Key{}, nil, fmt.Errorf("Bedrock region is required") + } + config := &schemas.BedrockKeyConfig{Region: plainSecretPtr(target.Region)} + if apiKey != "" { + key.Value = plainSecret(apiKey) + } else { + if value, ok := credential.Value("access_key"); ok { + config.AccessKey = plainSecret(value) + } + if value, ok := credential.Value("secret_key"); ok { + config.SecretKey = plainSecret(value) + } + config.SessionToken = optionalPlainSecret(credential, "session_token") + config.RoleARN = optionalPlainSecret(credential, "role_arn") + config.ExternalID = optionalPlainSecret(credential, "external_id") + config.RoleSessionName = optionalPlainSecret(credential, "session_name") + } + key.BedrockKeyConfig = config + case channel.ProviderGoogleVertex: + var target struct { + ProjectID string `json:"project_id"` + ProjectNumber string `json:"project_number"` + Location string `json:"location"` + } + if json.Unmarshal(resolved.TargetConfig, &target) != nil || target.ProjectID == "" || target.Location == "" { + return schemas.Key{}, nil, fmt.Errorf("Vertex project and location are required") + } + serviceAccount, ok := credential.Value("service_account_json") + if !ok || serviceAccount == "" { + return schemas.Key{}, nil, fmt.Errorf("Vertex service account is required") + } + key.VertexKeyConfig = &schemas.VertexKeyConfig{ + ProjectID: plainSecret(target.ProjectID), + ProjectNumber: plainSecret(target.ProjectNumber), + Region: plainSecret(target.Location), + AuthCredentials: plainSecret(serviceAccount), + } + default: + return schemas.Key{}, nil, fmt.Errorf("unsupported provider kind") + } + return key, secrets, nil +} + +func plainSecret(value string) schemas.SecretVar { + return schemas.SecretVar{Val: value, SecretType: schemas.SecretTypePlainText} +} + +func plainSecretPtr(value string) *schemas.SecretVar { + secret := plainSecret(value) + return &secret +} + +func optionalPlainSecret(credential channel.Credential, field string) *schemas.SecretVar { + value, ok := credential.Value(field) + if !ok || value == "" { + return nil + } + return plainSecretPtr(value) +} + +func credentialSecrets(credential channel.Credential) []string { + fields := []string{ + "api_key", "client_id", "client_secret", "tenant_id", "access_key", "secret_key", + "session_token", "role_arn", "external_id", "session_name", "service_account_json", + } + seen := make(map[string]struct{}, len(fields)) + result := make([]string, 0, len(fields)) + appendSecret := func(value string) { + if value == "" { + return + } + if _, exists := seen[value]; exists { + return + } + seen[value] = struct{}{} + result = append(result, value) + } + for _, field := range fields { + value, ok := credential.Value(field) + if !ok { + continue + } + appendSecret(value) + if field == "service_account_json" { + var object map[string]json.RawMessage + if json.Unmarshal([]byte(value), &object) == nil { + for _, nestedField := range []string{"private_key", "private_key_id", "client_email"} { + var nested string + if json.Unmarshal(object[nestedField], &nested) == nil { + appendSecret(nested) + } + } + } + } + } + return result +} + +func newSDKContext(parent context.Context, spec execution.AttemptSpec, directKey schemas.Key) *schemas.BifrostContext { + bifrostContext := schemas.NewBifrostContext(parent, schemas.NoDeadline) + bifrostContext.SetValue(schemas.BifrostContextKeyRequestID, spec.RequestID) + bifrostContext.SetValue(schemas.BifrostContextKeyDirectKey, directKey) + if spec.Timeouts.StreamIdle > 0 { + bifrostContext.SetValue(schemas.BifrostContextKeyStreamIdleTimeout, spec.Timeouts.StreamIdle) + } + if headers := safeRequestHeaders(spec.Header); len(headers) > 0 { + bifrostContext.SetValue(schemas.BifrostContextKeyExtraHeaders, headers) + } + return bifrostContext +} + +func safeRequestHeaders(source http.Header) map[string][]string { + if len(source) == 0 { + return nil + } + result := make(map[string][]string) + for name, values := range source { + canonical := http.CanonicalHeaderKey(name) + if !safeUpstreamRequestHeader(canonical) { + continue + } + result[canonical] = append([]string(nil), values...) + } + return result +} + +func safePassthroughHeaders(source http.Header) map[string]string { + if len(source) == 0 { + return nil + } + result := make(map[string]string) + for name, values := range source { + canonical := http.CanonicalHeaderKey(name) + if !safeUpstreamRequestHeader(canonical) || len(values) == 0 { + continue + } + result[canonical] = strings.Join(values, ", ") + } + return result +} + +func safeUpstreamQuery(source url.Values) url.Values { + if len(source) == 0 { + return nil + } + result := make(url.Values) + for key, values := range source { + if unsafeUpstreamQueryKey(key) { + continue + } + result[key] = append([]string(nil), values...) + } + return result +} + +func safeRawUpstreamQuery(source string) string { + if source == "" { + return "" + } + segments := strings.Split(source, "&") + kept := make([]string, 0, len(segments)) + for _, segment := range segments { + key := segment + if separator := strings.IndexByte(key, '='); separator >= 0 { + key = key[:separator] + } + if decoded, err := url.QueryUnescape(key); err == nil { + key = decoded + } + if unsafeUpstreamQueryKey(key) { + continue + } + kept = append(kept, segment) + } + return strings.Join(kept, "&") +} + +func safeAttemptQuery(spec execution.AttemptSpec) string { + if spec.RawQuery != "" { + return safeRawUpstreamQuery(spec.RawQuery) + } + return safeUpstreamQuery(spec.Query).Encode() +} + +func unsafeUpstreamQueryKey(key string) bool { + normalized := strings.ReplaceAll(strings.ToLower(key), "-", "_") + switch normalized { + case "provider", "fallback", "fallbacks", "authorization", "proxy_authorization", + "api_key", "apikey", "x_api_key", "x_goog_api_key", "key", "access_token": + return true + default: + return false + } +} + +func resolveTypedTargetURL(baseURL, resourcePath, rawQuery string) (string, error) { + parsed, err := url.Parse(baseURL) + if err != nil || parsed == nil || !parsed.IsAbs() || parsed.Host == "" || parsed.User != nil || parsed.Fragment != "" { + return "", fmt.Errorf("invalid target URL") + } + parsed.Path = strings.TrimRight(parsed.Path, "/") + resourcePath + parsed.RawPath = "" + if rawQuery != "" { + parsed.RawQuery = rawQuery + } + return parsed.String(), nil +} + +func splitCompatibleBaseURL(value string) (string, string, error) { + parsed, err := url.Parse(value) + if err != nil || parsed == nil || !parsed.IsAbs() || parsed.Host == "" || parsed.User != nil || parsed.Fragment != "" { + return "", "", fmt.Errorf("invalid target URL") + } + query := safeRawUpstreamQuery(parsed.RawQuery) + if query != parsed.RawQuery { + return "", "", fmt.Errorf("unsafe target query") + } + parsed.RawQuery = "" + parsed.ForceQuery = false + return parsed.String(), query, nil +} + +func mergeRawQueries(baseQuery, requestQuery string) string { + if baseQuery == "" { + return requestQuery + } + if requestQuery == "" { + return baseQuery + } + return baseQuery + "&" + requestQuery +} + +func setTypedRequestURL(ctx *schemas.BifrostContext, requestURL string) { + if ctx != nil && requestURL != "" { + ctx.SetValue(schemas.BifrostContextKeyURLPath, requestURL) + } +} + +func safeUpstreamRequestHeader(canonical string) bool { + switch canonical { + case "Authorization", "Proxy-Authorization", "Api-Key", "X-Api-Key", "X-Goog-Api-Key", + "Host", "Content-Length", "Connection", "Proxy-Connection", "Keep-Alive", "Te", "Trailer", + "Transfer-Encoding", "Upgrade", "Cookie", "Set-Cookie": + return false + } + return !strings.HasPrefix(strings.ToLower(canonical), "x-bf-") +} + +func boundedRequestContext(parent context.Context, budget time.Duration) (context.Context, context.CancelFunc) { + if parent == nil { + parent = context.Background() + } + if budget > 0 { + return context.WithTimeout(parent, budget) + } + if _, hasDeadline := parent.Deadline(); hasDeadline { + return context.WithCancel(parent) + } + // Bifrost's fasthttp operation may outlive a canceled caller while it releases + // pooled objects. A finite default keeps that SDK cleanup goroutine bounded. + return context.WithTimeout(parent, defaultRequestBudget) +} + +type responseGate struct { + timer *time.Timer + fired chan struct{} + stopOne chan struct{} +} + +func startPreResponseGate(cancel context.CancelFunc, timeouts execution.AttemptTimeouts) *responseGate { + budget := timeouts.FirstByte + gate := &responseGate{fired: make(chan struct{}), stopOne: make(chan struct{})} + if budget <= 0 { + return gate + } + // Core exposes neither an independent dial hook nor transport configuration + // per attempt. Applying Connect here would incorrectly turn it into a whole + // response budget. Request remains the total budget; FirstByte is the only + // safe pre-response approximation (and covers the buffered unary call). + gate.timer = time.AfterFunc(budget, func() { + select { + case <-gate.stopOne: + return + default: + } + close(gate.fired) + cancel() + }) + return gate +} + +func (g *responseGate) stop() { + if g == nil { + return + } + select { + case <-g.stopOne: + return + default: + close(g.stopOne) + } + if g.timer != nil { + g.timer.Stop() + } +} + +func (g *responseGate) expired() bool { + if g == nil { + return false + } + select { + case <-g.fired: + return true + default: + return false + } +} + +func minimumPositive(values ...time.Duration) time.Duration { + var result time.Duration + for _, value := range values { + if value > 0 && (result == 0 || value < result) { + result = value + } + } + return result +} + +type streamIdleTimer struct { + duration time.Duration + timer *time.Timer + ch <-chan time.Time +} + +func newIdleTimer(duration time.Duration) *streamIdleTimer { + timer := &streamIdleTimer{duration: duration} + if duration > 0 { + timer.timer = time.NewTimer(duration) + timer.ch = timer.timer.C + } + return timer +} + +func (t *streamIdleTimer) channel() <-chan time.Time { + if t == nil { + return nil + } + return t.ch +} + +func (t *streamIdleTimer) pause() { + if t == nil || t.timer == nil { + return + } + if !t.timer.Stop() { + select { + case <-t.timer.C: + default: + } + } + t.ch = nil +} + +func (t *streamIdleTimer) resume() { + if t == nil || t.timer == nil { + return + } + t.timer.Reset(t.duration) + t.ch = t.timer.C +} + +func (t *streamIdleTimer) stop() { + if t != nil { + t.pause() + } +} + +func safeValidationReason(err error) string { + if err == nil { + return "validation failed" + } + message := strings.TrimSpace(err.Error()) + if message == "" || strings.ContainsAny(message, "\r\n\x00") { + return "validation failed" + } + if len(message) > 192 { + return message[:192] + } + return message +} + +func unaryContextFailure(ctx context.Context, preResponseExpired bool) execution.AttemptResult { + kind, summary := contextFailure(ctx, preResponseExpired) + return attemptedUnaryFailure(kind, summary) +} + +func streamContextFailure( + ctx context.Context, + preResponseExpired bool, + started bool, + headers http.Header, + requestID string, + model string, + usageEvidence *execution.UsageEvidence, +) execution.StreamResult { + kind, summary := contextFailure(ctx, preResponseExpired) + if !started { + return attemptedStreamFailure(kind, summary) + } + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: headers.Clone(), + Model: model, + UpstreamRequestID: requestID, + Usage: cloneUsage(usageEvidence), + Error: &execution.ErrorEvidence{Kind: kind, Summary: summary, RequestID: requestID}, + } +} + +func contextFailure(ctx context.Context, preResponseExpired bool) (execution.ErrorKind, string) { + if preResponseExpired || errors.Is(ctx.Err(), context.DeadlineExceeded) { + return execution.ErrorKindTimeout, "upstream request timed out" + } + return execution.ErrorKindCanceled, "upstream request canceled" +} diff --git a/internal/execution/bifrost/executor_test.go b/internal/execution/bifrost/executor_test.go new file mode 100644 index 00000000..750de0ed --- /dev/null +++ b/internal/execution/bifrost/executor_test.go @@ -0,0 +1,1065 @@ +package bifrost + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" + "gpt-load/internal/usage" +) + +const testAPIKey = "sk-test-do-not-leak" + +func TestRuntimeUnaryUsesSelectedCredentialModelAndEndpoint(t *testing.T) { + t.Parallel() + + const responseBody = `{"id": "chatcmpl-1", "object":"chat.completion", "created":123, "model":"served-model", "choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}], "usage":{"prompt_tokens":10,"completion_tokens":3,"total_tokens":13,"prompt_tokens_details":{"cached_tokens":4}}, "vendor_response":{"precise":1.2300}}` + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + if request.URL.Path != "/v1/chat/completions" { + t.Errorf("path = %q", request.URL.Path) + } + if got := request.Header.Get("Authorization"); got != "Bearer "+testAPIKey { + t.Errorf("authorization = %q", got) + } + for _, name := range []string{"Proxy-Authorization", "Api-Key", "X-Api-Key", "X-Goog-Api-Key", "Connection", "Proxy-Connection"} { + if value := request.Header.Get(name); value != "" { + t.Errorf("sensitive/transport header %s reached upstream: %q", name, value) + } + } + if got := request.Header.Get("X-Test-Header"); got != "forward-me" { + t.Errorf("safe business header = %q", got) + } + body, err := io.ReadAll(request.Body) + if err != nil { + t.Errorf("read request: %v", err) + return + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Errorf("decode request: %v", err) + return + } + if payload["model"] != "upstream-model" { + t.Errorf("model = %#v", payload["model"]) + } + if _, exists := payload["stream"]; exists { + t.Errorf("stream = %#v, want omitted", payload["stream"]) + } + for _, field := range []string{"fallbacks", "provider", "authorization", "api_key", "x-api-key"} { + if _, exists := payload[field]; exists { + t.Errorf("client control field %q reached upstream", field) + } + } + vendor, ok := payload["vendor_extension"].(map[string]any) + if !ok || vendor["precise"] != float64(1.23) { + t.Errorf("vendor extension was not preserved: %#v", payload["vendor_extension"]) + } + if !bytes.Contains(body, []byte(`"precise":1.2300`)) { + t.Errorf("vendor extension raw number was rewritten: %s", body) + } + + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Request-Id", "upstream-request-1") + writer.Header().Set("X-Upstream-Secret", "must-not-forward") + _, _ = io.WriteString(writer, responseBody) + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + // Equal aliases must retain the native byte-for-byte fast path. + spec.ClientModel = spec.UpstreamModel + result := runtime.Execute(context.Background(), spec) + + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.StatusCode != http.StatusOK || result.DispatchState != execution.DispatchMaybeSent || !result.ResponseStarted { + t.Fatalf("unexpected result metadata: %+v", result) + } + if result.UpstreamRequestID != "upstream-request-1" { + t.Fatalf("upstream request ID = %q", result.UpstreamRequestID) + } + if result.Header.Get("X-Upstream-Secret") != "" { + t.Fatal("unsafe response header was forwarded") + } + if result.Model != "served-model" { + t.Fatalf("result model = %q", result.Model) + } + if string(result.Body) != responseBody { + t.Fatalf("native response bytes changed:\n got: %s\nwant: %s", result.Body, responseBody) + } + assertUsage(t, result.Usage, usage.Tokens{UncachedInput: 6, CacheRead: 4, Output: 3}) + if calls.Load() != 1 { + t.Fatalf("calls = %d, want 1", calls.Load()) + } + if runtime.account.keyPoolCalls.Load() != 0 { + t.Fatalf("key-pool calls = %d, want 0", runtime.account.keyPoolCalls.Load()) + } + assertNoPrivateLeak(t, result, testAPIKey, "gptload-custom-") +} + +func TestNativePassthroughRewritesModelAliasInUnaryAndStreamResponses(t *testing.T) { + t.Parallel() + + t.Run("unary JSON", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"id":"chat_1","object":"chat.completion","created":1,"model":"provider-model","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`) + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.ClientModel = "public-model" + spec.UpstreamModel = "provider-model" + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.Model != "provider-model" || !bytes.Contains(result.Body, []byte(`"model":"public-model"`)) || bytes.Contains(result.Body, []byte(`"model":"provider-model"`)) { + t.Fatalf("model/body = %q/%s", result.Model, result.Body) + } + }) + + t.Run("SSE split across upstream writes", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "text/event-stream") + flusher := writer.(http.Flusher) + _, _ = io.WriteString(writer, `data: {"id":"chat_1","object":"chat.completion.chunk","created":1,"model":"provider-`) + flusher.Flush() + time.Sleep(5 * time.Millisecond) + _, _ = io.WriteString(writer, "model\",\"choices\":[]}\n\ndata: [DONE]\n\n") + flusher.Flush() + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.ClientModel = "public-model" + spec.UpstreamModel = "provider-model" + var data bytes.Buffer + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + if event.Kind == execution.StreamEventData { + data.Write(event.Data) + } + return nil + }) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.Model != "provider-model" || !bytes.Contains(data.Bytes(), []byte(`"model":"public-model"`)) || bytes.Contains(data.Bytes(), []byte(`"model":"provider-model"`)) { + t.Fatalf("model/data = %q/%s", result.Model, data.String()) + } + }) +} + +func TestProductionRuntimeAllowsConfiguredPrivateCompatibleEndpoint(t *testing.T) { + t.Parallel() + + server := responseServer(t, "private-endpoint") + defer server.Close() + runtime, err := NewRuntime(context.Background()) + if err != nil { + t.Fatalf("new production runtime: %v", err) + } + defer runtime.Shutdown() + + result := runtime.Execute(context.Background(), compatibleSpec(server.URL)) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.UpstreamRequestID != "private-endpoint" { + t.Fatalf("request ID = %q", result.UpstreamRequestID) + } +} + +func TestRuntimeUnaryHTTPErrorIsSingleLogicalCallAndSanitized(t *testing.T) { + t.Parallel() + + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + writer.Header().Set("X-Request-Id", "request-401") + writer.Header().Set("Retry-After", "2") + writer.WriteHeader(http.StatusUnauthorized) + _, _ = io.WriteString(writer, `{"error":{"type":"invalid_api_key","code":"bad_key","message":"rejected `+testAPIKey+`"}}`) + })) + defer server.Close() + + runtime := newTestRuntime(t) + result := runtime.Execute(context.Background(), compatibleSpec(server.URL)) + + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if calls.Load() != 1 { + t.Fatalf("calls = %d, want MaxRetries=0 to produce one call", calls.Load()) + } + if result.Error == nil || result.Error.Kind != execution.ErrorKindHTTP || result.Error.StatusCode != http.StatusUnauthorized { + t.Fatalf("unexpected error evidence: %+v", result.Error) + } + if result.Error.Type != "invalid_api_key" || result.Error.Code != "bad_key" { + t.Fatalf("type/code = %q/%q", result.Error.Type, result.Error.Code) + } + if result.Error.RetryAfter != 2*time.Second { + t.Fatalf("retry after = %s", result.Error.RetryAfter) + } + if result.UpstreamRequestID != "request-401" { + t.Fatalf("request ID = %q", result.UpstreamRequestID) + } + assertNoPrivateLeak(t, result, testAPIKey, "gptload-custom-") +} + +func TestRuntimeRejectsInvalidAttemptBeforeDispatch(t *testing.T) { + t.Parallel() + + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + calls.Add(1) + })) + defer server.Close() + + tests := []struct { + name string + mutate func(*execution.AttemptSpec) + }{ + {name: "unsupported channel", mutate: func(spec *execution.AttemptSpec) { spec.ChannelID = "other" }}, + {name: "official config has field", mutate: func(spec *execution.AttemptSpec) { + spec.ChannelID = string(channel.OpenAI) + spec.TargetConfig = json.RawMessage(`{"base_url":"` + server.URL + `"}`) + }}, + {name: "official target missing", mutate: func(spec *execution.AttemptSpec) { + spec.ChannelID = string(channel.OpenAI) + spec.TargetConfig = nil + }}, + {name: "compatible missing base URL", mutate: func(spec *execution.AttemptSpec) { spec.TargetConfig = json.RawMessage(`{}`) }}, + {name: "compatible unknown target field", mutate: func(spec *execution.AttemptSpec) { + spec.TargetConfig = json.RawMessage(`{"base_url":"` + server.URL + `","token":"secret"}`) + }}, + {name: "credential unknown field", mutate: func(spec *execution.AttemptSpec) { + spec.Credential = execution.NewCredentialSnapshot(1, 1, 1, []byte(`{"api_key":"`+testAPIKey+`","other":"x"}`)) + }}, + {name: "credential empty", mutate: func(spec *execution.AttemptSpec) { + spec.Credential = execution.NewCredentialSnapshot(1, 1, 1, []byte(`{"api_key":""}`)) + }}, + {name: "wrong protocol", mutate: func(spec *execution.AttemptSpec) { spec.ClientProtocol = protocol.Anthropic }}, + {name: "wrong operation", mutate: func(spec *execution.AttemptSpec) { spec.Operation = execution.OperationListModels }}, + {name: "wrong method", mutate: func(spec *execution.AttemptSpec) { spec.Method = http.MethodGet }}, + {name: "wrong path", mutate: func(spec *execution.AttemptSpec) { spec.Path = "/v1/responses" }}, + } + + runtime := newTestRuntime(t) + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + spec := compatibleSpec(server.URL) + test.mutate(&spec) + result := runtime.Execute(context.Background(), spec) + if result.DispatchState != execution.DispatchNotSent || result.Error == nil || result.Error.Kind != execution.ErrorKindInvalidRequest { + t.Fatalf("unexpected result: %+v", result) + } + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v", err) + } + assertNoPrivateLeak(t, result, testAPIKey) + }) + } + if calls.Load() != 0 { + t.Fatalf("upstream calls = %d, want 0", calls.Load()) + } +} + +func TestRuntimeAcceptsCanonicalOfficialTargetWithoutChangingCredential(t *testing.T) { + t.Parallel() + + runtime := newTestRuntime(t) + spec := compatibleSpec("https://unused.example") + spec.ChannelID = string(channel.OpenAI) + spec.TargetConfig = json.RawMessage(`{}`) + prepared, failure := runtime.prepare(spec, false) + if failure != nil { + t.Fatalf("official preflight failed: %+v", failure) + } + if prepared.provider != "openai" || prepared.directKey.Value.GetValue() != testAPIKey { + t.Fatalf("prepared provider/key mismatch: provider=%q key_matches=%t", prepared.provider, prepared.directKey.Value.GetValue() == testAPIKey) + } +} + +func TestRuntimePreparesStructuredCloudCredentialsForSelectedProvider(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID string + targetConfig string + credential string + wantProvider string + }{ + { + name: "Azure API key", channelID: "azure_openai", + targetConfig: `{"endpoint":"https://resource.openai.azure.com"}`, + credential: `{"api_key":"azure-secret"}`, wantProvider: "azure", + }, + { + name: "Azure Entra", channelID: "azure_openai", + targetConfig: `{"endpoint":"https://resource.services.ai.azure.com"}`, + credential: `{"client_id":"client","client_secret":"secret","tenant_id":"tenant"}`, wantProvider: "azure", + }, + { + name: "Bedrock API key", channelID: "aws_bedrock", + targetConfig: `{"region":"us-east-1"}`, + credential: `{"api_key":"bedrock-secret"}`, wantProvider: "bedrock", + }, + { + name: "Bedrock SigV4 role", channelID: "aws_bedrock", + targetConfig: `{"region":"eu-west-1"}`, + credential: `{"access_key":"access","secret_key":"secret","session_token":"token","role_arn":"arn:aws:iam::123456789012:role/test"}`, wantProvider: "bedrock", + }, + { + name: "Vertex service account", channelID: "google_vertex", + targetConfig: `{"location":"us-central1","project_id":"project-one"}`, + credential: `{"service_account_json":"{\"type\":\"service_account\",\"project_id\":\"project-one\",\"client_email\":\"svc@example.iam.gserviceaccount.com\",\"private_key\":\"secret\"}"}`, wantProvider: "vertex", + }, + } + runtime := newTestRuntime(t) + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + spec := compatibleSpec("https://unused.example") + spec.ChannelID = test.channelID + spec.TargetConfig = json.RawMessage(test.targetConfig) + spec.Credential = execution.NewCredentialSnapshot(23, 5, 7, []byte(test.credential)) + prepared, failure := runtime.prepare(spec, false) + if failure != nil { + t.Fatalf("prepare() failure = %+v", failure) + } + if string(prepared.provider) != test.wantProvider { + t.Fatalf("provider = %q, want %q", prepared.provider, test.wantProvider) + } + if prepared.directKey.ID != "23:7:5" || prepared.directKey.Weight != 1 || len(prepared.directKey.Models) != 1 || prepared.directKey.Models[0] != "*" { + t.Fatalf("direct key identity = %#v", prepared.directKey) + } + switch test.name { + case "Azure API key": + if prepared.directKey.Value.GetValue() != "azure-secret" || prepared.directKey.AzureKeyConfig == nil || prepared.directKey.AzureKeyConfig.Endpoint.GetValue() != "https://resource.openai.azure.com" { + t.Fatalf("Azure key = %#v", prepared.directKey) + } + case "Azure Entra": + config := prepared.directKey.AzureKeyConfig + if config == nil || prepared.directKey.Value.GetValue() != "" || config.ClientID == nil || config.ClientID.GetValue() != "client" || config.ClientSecret == nil || config.ClientSecret.GetValue() != "secret" || config.TenantID == nil || config.TenantID.GetValue() != "tenant" { + t.Fatalf("Azure Entra key = %#v", prepared.directKey) + } + case "Bedrock API key": + if prepared.directKey.Value.GetValue() != "bedrock-secret" || prepared.directKey.BedrockKeyConfig == nil || prepared.directKey.BedrockKeyConfig.Region == nil || prepared.directKey.BedrockKeyConfig.Region.GetValue() != "us-east-1" { + t.Fatalf("Bedrock API key = %#v", prepared.directKey) + } + case "Bedrock SigV4 role": + config := prepared.directKey.BedrockKeyConfig + if config == nil || config.AccessKey.GetValue() != "access" || config.SecretKey.GetValue() != "secret" || config.SessionToken == nil || config.SessionToken.GetValue() != "token" || config.RoleARN == nil || config.RoleARN.GetValue() == "" || config.Region == nil || config.Region.GetValue() != "eu-west-1" { + t.Fatalf("Bedrock SigV4 key = %#v", prepared.directKey) + } + case "Vertex service account": + config := prepared.directKey.VertexKeyConfig + if config == nil || config.ProjectID.GetValue() != "project-one" || config.Region.GetValue() != "us-central1" || !strings.Contains(config.AuthCredentials.GetValue(), `"private_key":"secret"`) { + t.Fatalf("Vertex key = %#v", prepared.directKey) + } + } + }) + } +} + +func TestRuntimeIsolatesCompatibleEndpointsWithinOneCore(t *testing.T) { + t.Parallel() + + serverA := responseServer(t, "a") + defer serverA.Close() + serverB := responseServer(t, "b") + defer serverB.Close() + + runtime := newTestRuntime(t) + resultA := runtime.Execute(context.Background(), compatibleSpec(serverA.URL)) + resultB := runtime.Execute(context.Background(), compatibleSpec(serverB.URL)) + if resultA.UpstreamRequestID != "a" || resultB.UpstreamRequestID != "b" { + t.Fatalf("endpoint results = %q/%q", resultA.UpstreamRequestID, resultB.UpstreamRequestID) + } + assertNoPrivateLeak(t, resultA, "gptload-custom-") + assertNoPrivateLeak(t, resultB, "gptload-custom-") +} + +func TestRuntimeKeepsExactNonV1CompatiblePrefixViaTypedFallback(t *testing.T) { + t.Parallel() + + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + if request.URL.Path != "/tenant/openai/chat/completions" { + t.Errorf("path = %q", request.URL.Path) + } + writeSuccess(writer, "compatible-prefix") + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.TargetConfig = json.RawMessage(`{"base_url":"` + server.URL + `/tenant/openai"}`) + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if calls.Load() != 1 || result.Error != nil || result.StatusCode != http.StatusOK { + t.Fatalf("calls/result = %d/%+v", calls.Load(), result) + } +} + +func TestRuntimeStreamFramesReadyDataUsageAndDone(t *testing.T) { + t.Parallel() + + const rawStream = "data: {\"id\": \"chatcmpl-s\", \"object\":\"chat.completion.chunk\",\"created\":123,\"model\":\"served-model\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hi\"},\"finish_reason\":null}],\"vendor_chunk\":{\"precise\":1.2300}}\n\n" + + "data: {\"id\":\"chatcmpl-s\",\"object\":\"chat.completion.chunk\",\"created\":123,\"model\":\"served-model\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n" + + "data: {\"id\":\"chatcmpl-s\",\"object\":\"chat.completion.chunk\",\"created\":123,\"model\":\"served-model\",\"choices\":[],\"usage\":{\"prompt_tokens\":8,\"completion_tokens\":2,\"total_tokens\":10,\"prompt_tokens_details\":{\"cached_tokens\":3}}}\n\n" + + "data: [DONE]\n\n" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if got := request.Header.Get("Authorization"); got != "Bearer "+testAPIKey { + t.Errorf("authorization = %q", got) + } + body, err := io.ReadAll(request.Body) + if err != nil { + t.Errorf("read request: %v", err) + return + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Errorf("decode request: %v", err) + return + } + if payload["model"] != "upstream-model" || payload["stream"] != true { + t.Errorf("stream request controls = %#v", payload) + } + if _, exists := payload["stream_options"]; exists { + t.Errorf("adapter injected stream_options without an explicit request option: %#v", payload["stream_options"]) + } + if _, exists := payload["provider"]; exists { + t.Error("provider control reached streaming upstream") + } + if _, exists := payload["fallbacks"]; exists { + t.Error("fallback control reached streaming upstream") + } + if !bytes.Contains(body, []byte(`"precise":1.2300`)) { + t.Errorf("stream vendor extension was rewritten: %s", body) + } + writer.Header().Set("Content-Type", "text/event-stream") + writer.Header().Set("X-Request-Id", "stream-request") + writer.WriteHeader(http.StatusOK) + flusher := writer.(http.Flusher) + _, _ = io.WriteString(writer, rawStream) + flusher.Flush() + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + // Equal aliases must retain the native byte-for-byte fast path. + spec.ClientModel = spec.UpstreamModel + var events []execution.StreamEvent + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + if err := event.Validate(); err != nil { + t.Fatalf("event validation: %v; event=%+v", err, event) + } + events = append(events, event.Clone()) + return nil + }) + + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if len(events) < 3 || events[0].Kind != execution.StreamEventReady { + t.Fatalf("events = %+v", events) + } + if events[0].UpstreamRequestID != "stream-request" { + t.Fatalf("ready request ID = %q", events[0].UpstreamRequestID) + } + var sawUsage, sawDone bool + var rawData bytes.Buffer + for _, event := range events[1:] { + if event.Kind == execution.StreamEventUsage { + sawUsage = true + assertUsage(t, event.Usage, usage.Tokens{UncachedInput: 5, CacheRead: 3, Output: 2}) + } + if event.Kind == execution.StreamEventData && bytes.Contains(event.Data, []byte("data: [DONE]\n\n")) { + sawDone = true + } + if event.Kind == execution.StreamEventData { + rawData.Write(event.Data) + } + assertNoPrivateLeak(t, event, "gptload-openai-", testAPIKey) + } + if !sawUsage || !sawDone { + t.Fatalf("usage/done = %t/%t; events=%+v", sawUsage, sawDone, events) + } + if rawData.String() != rawStream { + t.Fatalf("native SSE bytes changed:\n got: %q\nwant: %q", rawData.String(), rawStream) + } + assertUsage(t, result.Usage, usage.Tokens{UncachedInput: 5, CacheRead: 3, Output: 2}) +} + +func TestSanitizeNativeChatBodyPreservesExplicitStreamOptions(t *testing.T) { + t.Parallel() + + body, err := sanitizeNativeChatBody( + []byte(`{"model":"client","messages":[],"stream_options":{"include_usage":false,"vendor_option":"keep"}}`), + "upstream", + true, + ) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatal(err) + } + options, ok := payload["stream_options"].(map[string]any) + if !ok || options["include_usage"] != false || options["vendor_option"] != "keep" { + t.Fatalf("stream_options changed: %#v", payload["stream_options"]) + } +} + +func TestRuntimePreservesSafeQueryForNativeAndTypedCompatibleTargets(t *testing.T) { + t.Parallel() + + for _, test := range []struct { + name string + basePrefix string + }{ + {name: "native passthrough", basePrefix: "/v1"}, + {name: "typed fallback", basePrefix: "/tenant/openai"}, + } { + t.Run(test.name, func(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + wantPath := test.basePrefix + "/chat/completions" + if request.URL.Path != wantPath { + t.Errorf("path = %q, want %q", request.URL.Path, wantPath) + } + if request.URL.Query().Get("trace") != "keep" { + t.Errorf("safe query = %q", request.URL.RawQuery) + } + for _, key := range []string{"provider", "fallbacks", "api_key"} { + if request.URL.Query().Has(key) { + t.Errorf("control query %q reached upstream: %q", key, request.URL.RawQuery) + } + } + writeSuccess(writer, "query") + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.TargetConfig = json.RawMessage(`{"base_url":"` + server.URL + test.basePrefix + `"}`) + spec.Query.Set("trace", "keep") + spec.Query.Set("provider", "injected") + spec.Query.Set("fallbacks", "injected") + spec.Query.Set("api_key", "injected") + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.Error != nil || calls.Load() != 1 { + t.Fatalf("calls/result = %d/%+v", calls.Load(), result) + } + }) + } +} + +func TestSafeRawUpstreamQueryPreservesOpaqueAllowedSegments(t *testing.T) { + t.Parallel() + + const raw = "first=%2F&provider=injected&broken=%ZZ&X-Api-Key=injected&&first=%2f&empty=&%61pi_key=injected&fallbacks=x" + const want = "first=%2F&broken=%ZZ&&first=%2f&empty=" + if got := safeRawUpstreamQuery(raw); got != want { + t.Fatalf("safeRawUpstreamQuery() = %q, want %q", got, want) + } + + runtime := newTestRuntime(t) + spec := compatibleSpec("https://upstream.example") + spec.Query = nil + spec.RawQuery = raw + prepared, failure := runtime.prepare(spec, false) + if failure != nil || prepared.passthrough == nil { + t.Fatalf("prepare() failure = %+v; prepared=%+v", failure, prepared) + } + if prepared.passthrough.RawQuery != want { + t.Fatalf("prepared RawQuery = %q, want %q", prepared.passthrough.RawQuery, want) + } +} + +func TestRuntimePreservesRawQueryBytesForNativeAndTypedCompatibleTargets(t *testing.T) { + t.Parallel() + + for _, test := range []struct { + name string + basePrefix string + }{ + {name: "native passthrough", basePrefix: "/v1"}, + {name: "typed fallback", basePrefix: "/tenant/openai"}, + } { + t.Run(test.name, func(t *testing.T) { + const raw = "first=%2F&provider=injected&second=%2f&&first=again&api%5Fkey=injected" + const want = "first=%2F&second=%2f&&first=again" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.RawQuery != want { + t.Errorf("RawQuery = %q, want %q", request.URL.RawQuery, want) + } + writeSuccess(writer, "raw-query") + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.TargetConfig = json.RawMessage(`{"base_url":"` + server.URL + test.basePrefix + `"}`) + spec.Query = nil + spec.RawQuery = raw + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.Error != nil { + t.Fatalf("result = %+v", result) + } + }) + } +} + +func TestRuntimeMergesCompatibleBaseAndRequestQueries(t *testing.T) { + t.Parallel() + + for _, test := range []struct { + name string + basePrefix string + }{ + {name: "native passthrough", basePrefix: "/v1"}, + {name: "typed fallback", basePrefix: "/tenant/openai"}, + } { + t.Run(test.name, func(t *testing.T) { + const want = "api-version=2025-01-01&tenant=base&trace=%2F&trace=%2f" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.RawQuery != want { + t.Errorf("RawQuery = %q, want %q", request.URL.RawQuery, want) + } + writeSuccess(writer, "merged-query") + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.TargetConfig = json.RawMessage(`{"base_url":"` + server.URL + test.basePrefix + `?api-version=2025-01-01&tenant=base"}`) + spec.Query = nil + spec.RawQuery = "trace=%2F&api_key=injected&trace=%2f" + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.Error != nil { + t.Fatalf("result = %+v", result) + } + }) + } +} + +func TestRuntimeStreamStopsWithoutDoneOnSinkFailure(t *testing.T) { + t.Parallel() + + server := slowStreamServer(t, 200*time.Millisecond) + defer server.Close() + runtime := newTestRuntime(t) + sinkErr := errors.New("sink closed") + var events []execution.StreamEvent + result := runtime.ExecuteStream(context.Background(), compatibleSpec(server.URL), func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + if event.Kind == execution.StreamEventData { + return sinkErr + } + return nil + }) + if result.Error == nil || result.Error.Kind != execution.ErrorKindCanceled { + t.Fatalf("unexpected result: %+v", result) + } + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v", err) + } + for _, event := range events { + if string(event.Data) == "data: [DONE]\n\n" { + t.Fatal("DONE emitted after sink failure") + } + } +} + +func TestRuntimeStreamCancellationAndHTTPErrorDoNotEmitDone(t *testing.T) { + t.Parallel() + + t.Run("context cancellation", func(t *testing.T) { + server := slowStreamServer(t, 150*time.Millisecond) + defer server.Close() + runtime := newTestRuntime(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var events []execution.StreamEvent + result := runtime.ExecuteStream(ctx, compatibleSpec(server.URL), func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + if event.Kind == execution.StreamEventData { + cancel() + } + return nil + }) + if result.Error == nil || result.Error.Kind != execution.ErrorKindCanceled { + t.Fatalf("unexpected result: %+v", result) + } + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v", err) + } + assertNoDone(t, events) + }) + + t.Run("upstream HTTP error", func(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + calls.Add(1) + writer.Header().Set("X-Request-Id", "stream-401") + writer.WriteHeader(http.StatusUnauthorized) + _, _ = io.WriteString(writer, `{"error":{"type":"invalid_api_key","message":"`+testAPIKey+`"}}`) + })) + defer server.Close() + runtime := newTestRuntime(t) + var events []execution.StreamEvent + result := runtime.ExecuteStream(context.Background(), compatibleSpec(server.URL), func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + return nil + }) + if calls.Load() != 1 || result.Error == nil || result.Error.Kind != execution.ErrorKindHTTP || result.StatusCode != http.StatusUnauthorized { + t.Fatalf("calls/result = %d/%+v", calls.Load(), result) + } + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v", err) + } + if len(events) != 2 || events[0].Kind != execution.StreamEventReady || events[0].StatusCode != http.StatusUnauthorized || events[1].Kind != execution.StreamEventData { + t.Fatalf("rejected stream events = %+v", events) + } + if bytes.Contains(events[1].Data, []byte(testAPIKey)) || !bytes.Contains(events[1].Data, []byte("[REDACTED]")) { + t.Fatalf("rejected stream body was not safely redacted: %s", events[1].Data) + } + assertNoPrivateLeak(t, result, testAPIKey, "gptload-custom-") + }) +} + +func TestRuntimeFirstByteAndStreamIdleTimeouts(t *testing.T) { + t.Parallel() + + t.Run("first byte", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + time.Sleep(150 * time.Millisecond) + writeSuccess(writer, "late") + })) + defer server.Close() + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.Timeouts.FirstByte = 20 * time.Millisecond + spec.Timeouts.Request = time.Second + started := time.Now() + result := runtime.Execute(context.Background(), spec) + if elapsed := time.Since(started); elapsed > 120*time.Millisecond { + t.Fatalf("first-byte gate returned after %s", elapsed) + } + if result.Error == nil || result.Error.Kind != execution.ErrorKindTimeout || result.ResponseStarted { + t.Fatalf("unexpected first-byte result: %+v", result) + } + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v", err) + } + }) + + t.Run("unary connect budget does not replace first-byte budget", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + time.Sleep(150 * time.Millisecond) + writeSuccess(writer, "late") + })) + defer server.Close() + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.Timeouts.Connect = 20 * time.Millisecond + spec.Timeouts.FirstByte = 500 * time.Millisecond + spec.Timeouts.Request = time.Second + result := runtime.Execute(context.Background(), spec) + if result.Error != nil || !result.ResponseStarted || result.StatusCode != http.StatusOK { + t.Fatalf("unexpected connect-gate result: %+v", result) + } + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v", err) + } + }) + + t.Run("stream connect budget does not replace first-byte budget", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + time.Sleep(80 * time.Millisecond) + writer.Header().Set("Content-Type", "text/event-stream") + writer.WriteHeader(http.StatusOK) + _, _ = io.WriteString(writer, "data: [DONE]\n\n") + writer.(http.Flusher).Flush() + })) + defer server.Close() + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.Timeouts.Connect = 20 * time.Millisecond + spec.Timeouts.FirstByte = 500 * time.Millisecond + spec.Timeouts.Request = time.Second + var events []execution.StreamEvent + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + return nil + }) + if result.Error != nil || !result.ResponseStarted || result.StatusCode != http.StatusOK { + t.Fatalf("unexpected connect-gate result: %+v", result) + } + if len(events) < 2 || events[0].Kind != execution.StreamEventReady { + t.Fatalf("events = %+v", events) + } + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v", err) + } + }) + + t.Run("stream first byte requires a complete data event", func(t *testing.T) { + finished := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(writer, "data: partial\n") + writer.(http.Flusher).Flush() + // Core v1.7.7 returns cancellation to its caller but does not + // guarantee that fasthttp closes an already-open body stream. Keep + // the fake finite while the assertion below verifies our logical + // first-event timeout boundary. + select { + case <-request.Context().Done(): + case <-time.After(200 * time.Millisecond): + } + close(finished) + })) + defer server.Close() + + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.Timeouts.FirstByte = 20 * time.Millisecond + spec.Timeouts.Request = time.Second + started := time.Now() + result := runtime.ExecuteStream(context.Background(), spec, func(execution.StreamEvent) error { return nil }) + if elapsed := time.Since(started); elapsed > 120*time.Millisecond { + t.Fatalf("first-event gate returned after %s", elapsed) + } + if result.Error == nil || result.Error.Kind != execution.ErrorKindTimeout || result.ResponseStarted { + t.Fatalf("unexpected first-event result: %+v", result) + } + select { + case <-finished: + case <-time.After(time.Second): + t.Fatal("timed out waiting for finite upstream fixture") + } + }) + + t.Run("stream idle", func(t *testing.T) { + server := slowStreamServer(t, 150*time.Millisecond) + defer server.Close() + runtime := newTestRuntime(t) + spec := compatibleSpec(server.URL) + spec.Timeouts.StreamIdle = 20 * time.Millisecond + spec.Timeouts.Request = time.Second + var events []execution.StreamEvent + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + return nil + }) + if result.Error == nil || result.Error.Kind != execution.ErrorKindTimeout || !result.ResponseStarted { + t.Fatalf("unexpected idle result: %+v", result) + } + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v", err) + } + if len(events) < 2 || events[0].Kind != execution.StreamEventReady { + t.Fatalf("events = %+v", events) + } + for _, event := range events { + if string(event.Data) == "data: [DONE]\n\n" { + t.Fatal("DONE emitted after idle timeout") + } + } + }) +} + +func TestRuntimeShutdownIsIdempotentAndClosesExecution(t *testing.T) { + t.Parallel() + + runtime := newTestRuntime(t) + runtime.Shutdown() + runtime.Shutdown() + result := runtime.Execute(context.Background(), compatibleSpec("https://example.com")) + if result.DispatchState != execution.DispatchNotSent || result.Error == nil || result.Error.Kind != execution.ErrorKindInternal { + t.Fatalf("unexpected result after shutdown: %+v", result) + } +} + +func TestRuntimeCapabilities(t *testing.T) { + t.Parallel() + + runtime := newTestRuntime(t) + required, err := execution.NewFeatureSet(execution.FeatureStreaming, execution.FeatureTools) + if err != nil { + t.Fatal(err) + } + if !runtime.Capabilities().Supports(execution.OperationChatCompletion, required) { + t.Fatal("chat streaming/tools capability missing") + } + if !runtime.Capabilities().Supports(execution.OperationResponsesCreate, required) { + t.Fatal("Responses streaming/tools capability missing") + } + resourceFeatures, err := execution.NewFeatureSet(execution.FeatureNativeResourceSemantics) + if err != nil { + t.Fatal(err) + } + if !runtime.Capabilities().Supports(execution.OperationResponsesRetrieve, resourceFeatures) { + t.Fatal("Responses resource capability missing") + } + if !runtime.Capabilities().Supports(execution.OperationResponsesCompact, mustTestFeatures(t, execution.FeatureTools)) { + t.Fatal("Responses compact tools capability missing") + } + if runtime.Capabilities().Supports(execution.OperationResponsesCompact, mustTestFeatures(t, execution.FeatureStreaming)) || + runtime.Capabilities().Supports(execution.OperationResponsesInputTokens, resourceFeatures) { + t.Fatal("non-streaming Responses utility capability was overstated") + } +} + +func mustTestFeatures(t *testing.T, features ...execution.Feature) execution.FeatureSet { + t.Helper() + set, err := execution.NewFeatureSet(features...) + if err != nil { + t.Fatal(err) + } + return set +} + +func newTestRuntime(t *testing.T) *Runtime { + t.Helper() + runtime, err := newRuntime(context.Background(), runtimeOptions{allowPrivateNetwork: true}) + if err != nil { + t.Fatalf("new runtime: %v", err) + } + t.Cleanup(runtime.Shutdown) + return runtime +} + +func compatibleSpec(baseURL string) execution.AttemptSpec { + baseURL = strings.TrimRight(baseURL, "/") + "/v1" + return execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: "request-1", + AttemptID: "attempt-1", + Sequence: 1, + ChannelID: string(channel.OpenAICompatible), + ClientProtocol: protocol.OpenAICompletions, + Operation: execution.OperationChatCompletion, + ClientModel: "client-model", + UpstreamModel: "upstream-model", + Method: http.MethodPost, + Path: "/v1/chat/completions", + Query: make(map[string][]string), + Header: http.Header{ + "Authorization": []string{"Bearer client-injected"}, + "Proxy-Authorization": []string{"Basic client-injected"}, + "Api-Key": []string{"client-injected"}, + "X-Api-Key": []string{"client-injected"}, + "X-Goog-Api-Key": []string{"client-injected"}, + "Connection": []string{"keep-alive"}, + "Proxy-Connection": []string{"keep-alive"}, + "X-Test-Header": []string{"forward-me"}, + }, + Body: []byte(`{"model":"injected-model","provider":"injected-provider","fallbacks":["other/model"],"authorization":"body-injected","api_key":"body-injected","x-api-key":"body-injected","messages":[{"role":"user","content":"hello"}],"vendor_extension":{"precise":1.2300,"nested":{"keep":"yes"}}}`), + TargetConfig: json.RawMessage(`{"base_url":"` + baseURL + `"}`), + Timeouts: execution.AttemptTimeouts{ + Connect: time.Second, + FirstByte: time.Second, + Request: 2 * time.Second, + StreamIdle: time.Second, + }, + Credential: execution.NewCredentialSnapshot(7, 3, 2, []byte(`{"api_key":"`+testAPIKey+`"}`)), + }) +} + +func responseServer(t *testing.T, requestID string) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("X-Request-Id", requestID) + writeSuccess(writer, requestID) + })) +} + +func writeSuccess(writer http.ResponseWriter, content string) { + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"id":"chatcmpl","object":"chat.completion","created":1,"model":"served","choices":[{"index":0,"message":{"role":"assistant","content":"`+content+`"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`) +} + +func slowStreamServer(t *testing.T, delay time.Duration) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "text/event-stream") + writer.WriteHeader(http.StatusOK) + flusher := writer.(http.Flusher) + _, _ = io.WriteString(writer, "data: {\"id\":\"chatcmpl-slow\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"served\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"a\"},\"finish_reason\":null}]}\n\n") + flusher.Flush() + time.Sleep(delay) + _, _ = io.WriteString(writer, "data: [DONE]\n\n") + flusher.Flush() + })) +} + +func assertUsage(t *testing.T, evidence *execution.UsageEvidence, want usage.Tokens) { + t.Helper() + if evidence == nil { + t.Fatal("usage evidence is nil") + } + if evidence.Normalized.Tokens != want || evidence.Normalized.State != usage.StateComplete { + t.Fatalf("usage = %+v, want tokens=%+v complete", evidence.Normalized, want) + } + if len(evidence.Raw) == 0 || !json.Valid(evidence.Raw) { + t.Fatalf("raw usage is not JSON: %q", evidence.Raw) + } +} + +func assertNoPrivateLeak(t *testing.T, value any, forbidden ...string) { + t.Helper() + encoded, err := json.Marshal(value) + if err != nil { + t.Fatalf("marshal value: %v", err) + } + for _, secret := range forbidden { + if secret != "" && strings.Contains(string(encoded), secret) { + t.Fatalf("serialized value leaked %q: %s", secret, encoded) + } + } +} + +func assertNoDone(t *testing.T, events []execution.StreamEvent) { + t.Helper() + for _, event := range events { + if string(event.Data) == "data: [DONE]\n\n" { + t.Fatal("DONE emitted after stream failure") + } + } +} diff --git a/internal/execution/bifrost/list_models.go b/internal/execution/bifrost/list_models.go new file mode 100644 index 00000000..0ce7eb49 --- /dev/null +++ b/internal/execution/bifrost/list_models.go @@ -0,0 +1,112 @@ +package bifrost + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + + "github.com/maximhq/bifrost/core/providers/anthropic" + "github.com/maximhq/bifrost/core/providers/gemini" + "github.com/maximhq/bifrost/core/providers/openai" + "github.com/maximhq/bifrost/core/schemas" + + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +type listModelsSDKResult struct { + response *schemas.BifrostListModelsResponse + err *schemas.BifrostError +} + +func (r *Runtime) executeListModels( + parent context.Context, + spec execution.AttemptSpec, + prepared preparedAttempt, +) execution.AttemptResult { + requestContext, requestCancel := boundedRequestContext(parent, spec.Timeouts.Request) + defer requestCancel() + callContext, callCancel := context.WithCancel(requestContext) + defer callCancel() + preResponse := startPreResponseGate(callCancel, spec.Timeouts) + defer preResponse.stop() + + bifrostContext := newSDKContext(callContext, spec, prepared.directKey) + setTypedRequestURL(bifrostContext, prepared.typedURL) + outcomeChannel := make(chan listModelsSDKResult, 1) + go func() { + response, bifrostError := r.core.ListModelsRequest(bifrostContext, prepared.listModelsRequest) + outcomeChannel <- listModelsSDKResult{response: response, err: bifrostError} + }() + + var outcome listModelsSDKResult + select { + case outcome = <-outcomeChannel: + if preResponse.expired() || requestContext.Err() != nil { + return unaryContextFailure(requestContext, preResponse.expired()) + } + case <-callContext.Done(): + return unaryContextFailure(requestContext, preResponse.expired()) + } + preResponse.stop() + if outcome.err != nil { + return unaryErrorResult(outcome.err, bifrostContext, prepared.secrets) + } + if outcome.response == nil { + return attemptedUnaryFailure(execution.ErrorKindInternal, "execution runtime returned no model list") + } + body, err := encodeListModelsResponse(prepared.clientProtocol, prepared.provider, outcome.response) + if err != nil { + return startedUnaryFailure(http.StatusOK, responseHeaders(outcome.response.ExtraFields.ProviderResponseHeaders, bifrostContext, false), execution.ErrorKindInternal, "encode model list") + } + headers := responseHeaders(outcome.response.ExtraFields.ProviderResponseHeaders, bifrostContext, false) + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: headers, + Body: body, + UpstreamRequestID: upstreamRequestID(headers), + } +} + +func encodeListModelsResponse(clientProtocol protocol.Protocol, provider schemas.ModelProvider, response *schemas.BifrostListModelsResponse) ([]byte, error) { + response = sanitizeListModelsResponse(provider, response) + var wire any + switch clientProtocol { + case protocol.OpenAICompletions, protocol.OpenAIResponses: + converted := openai.ToOpenAIListModelsResponse(response) + if converted != nil { + converted.Object = "list" + } + wire = converted + case protocol.Anthropic: + wire = anthropic.ToAnthropicListModelsResponse(response) + case protocol.Gemini: + wire = gemini.ToGeminiListModelsResponse(response) + default: + return nil, fmt.Errorf("unsupported model list protocol") + } + body, err := json.Marshal(wire) + if err != nil { + return nil, fmt.Errorf("marshal model list") + } + return body, nil +} + +func sanitizeListModelsResponse(provider schemas.ModelProvider, response *schemas.BifrostListModelsResponse) *schemas.BifrostListModelsResponse { + if response == nil { + return nil + } + clone := *response + clone.Data = append([]schemas.Model(nil), response.Data...) + prefix := string(provider) + "/" + for index := range clone.Data { + clone.Data[index].ID = strings.TrimPrefix(clone.Data[index].ID, prefix) + } + clone.ExtraFields = schemas.BifrostResponseExtraFields{} + clone.KeyStatuses = nil + return &clone +} diff --git a/internal/execution/bifrost/mapping.go b/internal/execution/bifrost/mapping.go new file mode 100644 index 00000000..2b088b6e --- /dev/null +++ b/internal/execution/bifrost/mapping.go @@ -0,0 +1,694 @@ +package bifrost + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "strconv" + "strings" + "time" + "unicode/utf8" + + "github.com/maximhq/bifrost/core/schemas" + + "gpt-load/internal/dialect" + "gpt-load/internal/execution" + platformheader "gpt-load/internal/platform/httpheader" +) + +type openAIChatResponse struct { + ID string `json:"id"` + Choices []schemas.BifrostResponseChoice `json:"choices"` + Created int `json:"created"` + Model string `json:"model"` + Object string `json:"object"` + ServiceTier *schemas.BifrostServiceTier `json:"service_tier,omitempty"` + SystemFingerprint string `json:"system_fingerprint,omitempty"` + Usage *openAIUsage `json:"usage,omitempty"` +} + +type openAIUsage struct { + PromptTokens int `json:"prompt_tokens"` + PromptTokensDetails *openAIPromptTokensDetails `json:"prompt_tokens_details,omitempty"` + CompletionTokens int `json:"completion_tokens"` + CompletionTokensDetails *schemas.ChatCompletionTokensDetails `json:"completion_tokens_details,omitempty"` + TotalTokens int `json:"total_tokens"` +} + +type openAIPromptTokensDetails struct { + AudioTokens int `json:"audio_tokens,omitempty"` + CachedTokens int `json:"cached_tokens,omitempty"` + CacheWriteTokens int `json:"cache_write_tokens,omitempty"` +} + +func encodeChatResponse(response *schemas.BifrostChatResponse) ([]byte, *execution.UsageEvidence, error) { + if response == nil { + return nil, nil, fmt.Errorf("response is nil") + } + wireUsage := toOpenAIUsage(response.Usage) + wire := openAIChatResponse{ + ID: response.ID, + Choices: response.Choices, + Created: response.Created, + Model: response.Model, + Object: response.Object, + ServiceTier: response.ServiceTier, + SystemFingerprint: response.SystemFingerprint, + Usage: wireUsage, + } + body, err := json.Marshal(wire) + if err != nil { + return nil, nil, fmt.Errorf("marshal OpenAI response") + } + if wireUsage == nil { + return body, nil, nil + } + rawUsage, err := json.Marshal(wireUsage) + if err != nil { + return nil, nil, fmt.Errorf("marshal OpenAI usage") + } + normalized, err := dialect.NewOpenAI().ExtractUsage(body) + if err != nil { + return nil, nil, fmt.Errorf("normalize OpenAI usage") + } + return body, &execution.UsageEvidence{Normalized: normalized, Raw: rawUsage}, nil +} + +func toOpenAIUsage(source *schemas.BifrostLLMUsage) *openAIUsage { + if source == nil { + return nil + } + result := &openAIUsage{ + PromptTokens: source.PromptTokens, + CompletionTokens: source.CompletionTokens, + CompletionTokensDetails: source.CompletionTokensDetails, + TotalTokens: source.TotalTokens, + } + if source.PromptTokensDetails != nil { + result.PromptTokensDetails = &openAIPromptTokensDetails{ + AudioTokens: source.PromptTokensDetails.AudioTokens, + CachedTokens: source.PromptTokensDetails.CachedReadTokens, + CacheWriteTokens: source.PromptTokensDetails.CachedWriteTokens, + } + } + return result +} + +func usageEvidenceFromPassthrough(source *schemas.BifrostPassthroughUsage) (*execution.UsageEvidence, error) { + if source == nil || source.LLMUsage == nil { + return nil, nil + } + wireUsage := toOpenAIUsage(source.LLMUsage) + rawUsage, err := json.Marshal(wireUsage) + if err != nil { + return nil, fmt.Errorf("marshal OpenAI usage") + } + syntheticBody, err := json.Marshal(struct { + Usage *openAIUsage `json:"usage"` + }{Usage: wireUsage}) + if err != nil { + return nil, fmt.Errorf("marshal OpenAI usage response") + } + normalized, err := dialect.NewOpenAI().ExtractUsage(syntheticBody) + if err != nil { + return nil, fmt.Errorf("normalize OpenAI usage") + } + return &execution.UsageEvidence{Normalized: normalized, Raw: rawUsage}, nil +} + +func usageEvidenceFromResponses(source *schemas.ResponsesResponseUsage) (*execution.UsageEvidence, error) { + if source == nil { + return nil, nil + } + rawUsage, err := json.Marshal(source) + if err != nil { + return nil, fmt.Errorf("marshal Responses usage") + } + syntheticBody, err := json.Marshal(struct { + Usage *schemas.ResponsesResponseUsage `json:"usage"` + }{Usage: source}) + if err != nil { + return nil, fmt.Errorf("marshal Responses usage response") + } + normalized, err := dialect.NewOpenAIResponses().ExtractUsage(syntheticBody) + if err != nil { + return nil, fmt.Errorf("normalize Responses usage") + } + return &execution.UsageEvidence{Normalized: normalized, Raw: rawUsage}, nil +} + +func openAIResponseModel(body []byte, fallback string) string { + var response struct { + Model string `json:"model"` + } + if err := json.Unmarshal(body, &response); err == nil && strings.TrimSpace(response.Model) != "" { + return response.Model + } + return fallback +} + +func frameSSE(payload []byte) []byte { + framed := make([]byte, 0, len(payload)+8) + framed = append(framed, "data: "...) + framed = append(framed, payload...) + framed = append(framed, '\n', '\n') + return framed +} + +func responseHeaders(primary map[string]string, bifrostContext *schemas.BifrostContext, stream bool) http.Header { + combined := make(map[string]string, len(primary)+4) + for key, value := range primary { + combined[key] = value + } + if bifrostContext != nil { + if contextHeaders, ok := bifrostContext.Value(schemas.BifrostContextKeyProviderResponseHeaders).(map[string]string); ok { + for key, value := range contextHeaders { + combined[key] = value + } + } + } + connectionTokens := make(map[string]struct{}) + for key, value := range combined { + if !strings.EqualFold(key, "Connection") { + continue + } + for _, token := range strings.Split(value, ",") { + if normalized := strings.ToLower(strings.TrimSpace(token)); normalized != "" { + connectionTokens[normalized] = struct{}{} + } + } + } + secrets := directKeySecrets(bifrostContext) + headers := make(http.Header) + for key, value := range combined { + if strings.ContainsAny(value, "\r\n\x00") { + continue + } + canonical := http.CanonicalHeaderKey(key) + if unsafeProviderResponseHeader(canonical, value, secrets, connectionTokens) { + continue + } + headers.Set(canonical, value) + } + if stream { + headers.Set("Content-Type", "text/event-stream") + } else if headers.Get("Content-Type") == "" { + headers.Set("Content-Type", "application/json") + } + for _, canonical := range []string{ + "Openai-Request-Id", "X-Ms-Request-Id", "X-Amzn-Requestid", "X-Goog-Request-Id", + } { + if headers.Get("X-Request-Id") == "" && headers.Get(canonical) != "" { + headers.Set("X-Request-Id", headers.Get(canonical)) + } + } + return headers +} + +func directKeySecrets(bifrostContext *schemas.BifrostContext) []string { + if bifrostContext == nil { + return nil + } + switch key := bifrostContext.Value(schemas.BifrostContextKeyDirectKey).(type) { + case schemas.Key: + keyCopy := key + return keySecrets(&keyCopy) + case *schemas.Key: + if key != nil { + return keySecrets(key) + } + } + return nil +} + +func unsafeProviderResponseHeader( + canonical string, + value string, + secrets []string, + connectionTokens map[string]struct{}, +) bool { + lowered := strings.ToLower(canonical) + if _, found := connectionTokens[strings.ToLower(canonical)]; found { + return true + } + if platformheader.IsCredentialName(canonical) || + strings.EqualFold(canonical, "Set-Cookie") || + strings.EqualFold(canonical, "Set-Cookie2") || + strings.Contains(lowered, "secret") || + strings.HasPrefix(strings.ToLower(canonical), "x-bf-") || + strings.HasPrefix(strings.ToLower(canonical), "x-gptload-") || + containsSecret(value, secrets) { + return true + } + switch canonical { + case "Connection", "Proxy-Connection", "Keep-Alive", "Proxy-Authenticate", + "Proxy-Authorization", "Te", "Trailer", "Transfer-Encoding", "Upgrade": + return true + default: + return false + } +} + +func passthroughHTTPError(status int, headers http.Header, body []byte, secrets []string) *execution.ErrorEvidence { + typeValue, codeValue := openAIErrorTypeCode(body) + typeValue = sanitizeEvidenceValue(typeValue, secrets) + codeValue = sanitizeEvidenceValue(codeValue, secrets) + requestID := upstreamRequestID(headers) + return &execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, + StatusCode: status, + Type: typeValue, + Code: codeValue, + Summary: errorSummary(execution.ErrorKindHTTP, status), + RequestID: requestID, + RetryAfter: retryAfter(headers), + Header: evidenceHeaders(headers), + } +} + +func openAIErrorTypeCode(body []byte) (string, string) { + var payload struct { + Error struct { + Type json.RawMessage `json:"type"` + Code json.RawMessage `json:"code"` + } `json:"error"` + } + if err := json.Unmarshal(body, &payload); err != nil { + return "", "" + } + return evidenceScalar(payload.Error.Type), evidenceScalar(payload.Error.Code) +} + +func evidenceScalar(raw json.RawMessage) string { + if len(raw) == 0 || bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + return "" + } + var text string + if err := json.Unmarshal(raw, &text); err == nil { + return text + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var number json.Number + if err := decoder.Decode(&number); err == nil { + return number.String() + } + return "" +} + +func redactSecrets(body []byte, secrets []string) []byte { + clone := append([]byte(nil), body...) + for _, secret := range secrets { + if secret != "" { + clone = bytes.ReplaceAll(clone, []byte(secret), []byte("[REDACTED]")) + } + } + return clone +} + +func evidenceHeaders(headers http.Header) http.Header { + if len(headers) == 0 { + return nil + } + result := make(http.Header) + for name, values := range headers { + canonical := http.CanonicalHeaderKey(name) + if canonical == "Request-Id" || canonical == "X-Request-Id" || canonical == "Retry-After" || + canonical == "Traceparent" || strings.HasPrefix(canonical, "Ratelimit-") || strings.HasPrefix(canonical, "X-Ratelimit-") { + result[canonical] = append([]string(nil), values...) + } + } + if len(result) == 0 { + return nil + } + return result +} + +func upstreamRequestID(headers http.Header) string { + if headers == nil { + return "" + } + if value := headers.Get("X-Request-Id"); value != "" { + return value + } + return headers.Get("Request-Id") +} + +func retryAfter(headers http.Header) time.Duration { + value := strings.TrimSpace(headers.Get("Retry-After")) + if value == "" { + return 0 + } + if seconds, err := strconv.ParseInt(value, 10, 64); err == nil && seconds >= 0 { + return time.Duration(seconds) * time.Second + } + if deadline, err := http.ParseTime(value); err == nil { + duration := time.Until(deadline) + if duration > 0 { + return duration + } + } + return 0 +} + +func unaryErrorResult(bifrostError *schemas.BifrostError, bifrostContext *schemas.BifrostContext, secrets []string) execution.AttemptResult { + headers := responseHeaders(nil, bifrostContext, false) + evidence, started, status := errorEvidence(bifrostError, bifrostContext, secrets, false, 0, headers) + result := execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: started, + StatusCode: status, + Header: nil, + UpstreamRequestID: "", + Error: evidence, + } + if started { + result.Header = headers + result.UpstreamRequestID = upstreamRequestID(headers) + result.Body = encodeSafeErrorBody(evidence) + } + return result +} + +func streamErrorResult( + bifrostError *schemas.BifrostError, + bifrostContext *schemas.BifrostContext, + secrets []string, + alreadyStarted bool, + initialStatus int, + headers http.Header, + model string, + usageEvidence *execution.UsageEvidence, +) execution.StreamResult { + if headers == nil { + headers = responseHeaders(nil, bifrostContext, true) + } + evidence, started, status := errorEvidence(bifrostError, bifrostContext, secrets, alreadyStarted, initialStatus, headers) + result := execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + Model: model, + Usage: cloneUsage(usageEvidence), + Error: evidence, + } + if started { + result.ResponseStarted = true + result.StatusCode = status + result.Header = headers.Clone() + result.UpstreamRequestID = upstreamRequestID(headers) + } + return result +} + +func errorEvidence( + bifrostError *schemas.BifrostError, + bifrostContext *schemas.BifrostContext, + secrets []string, + alreadyStarted bool, + initialStatus int, + headers http.Header, +) (*execution.ErrorEvidence, bool, int) { + if bifrostError == nil { + evidence := &execution.ErrorEvidence{Kind: execution.ErrorKindInternal, Summary: "execution runtime returned an invalid response"} + if alreadyStarted { + evidence.RequestID = upstreamRequestID(headers) + return evidence, true, initialStatus + } + return evidence, false, 0 + } + typeValue := errorType(bifrostError) + codeValue := "" + if bifrostError.Error != nil && bifrostError.Error.Code != nil { + codeValue = *bifrostError.Error.Code + } + typeValue = sanitizeEvidenceValue(typeValue, secrets) + codeValue = sanitizeEvidenceValue(codeValue, secrets) + synthetic := typeValue == schemas.RequestTimedOut || typeValue == schemas.RequestCancelled || typeValue == schemas.ProviderConnectionFailed + status := 0 + if bifrostError.StatusCode != nil && *bifrostError.StatusCode >= 100 && *bifrostError.StatusCode <= 599 { + status = *bifrostError.StatusCode + } + started := alreadyStarted || (status > 0 && !synthetic && !bifrostError.IsBifrostError) + resultStatus := status + if alreadyStarted { + resultStatus = initialStatus + } else if !started { + resultStatus = 0 + } + kind := classifyError(bifrostError, typeValue, started) + requestID := "" + if started { + requestID = upstreamRequestID(headers) + } + evidenceStatus := resultStatus + if alreadyStarted && status != initialStatus { + evidenceStatus = 0 + } + evidence := &execution.ErrorEvidence{ + Kind: kind, + StatusCode: evidenceStatus, + Type: typeValue, + Code: codeValue, + Summary: errorSummary(kind, errorSummaryStatus(status, resultStatus)), + RequestID: requestID, + RetryAfter: retryAfter(headers), + Header: evidenceHeaders(headers), + } + if !started { + evidence.StatusCode = 0 + evidence.RequestID = "" + evidence.RetryAfter = 0 + evidence.Header = nil + } + return evidence, started, resultStatus +} + +func errorSummaryStatus(upstreamStatus, resultStatus int) int { + if upstreamStatus > 0 { + return upstreamStatus + } + return resultStatus +} + +func errorType(bifrostError *schemas.BifrostError) string { + if bifrostError == nil { + return "" + } + if bifrostError.Error != nil && bifrostError.Error.Type != nil { + return *bifrostError.Error.Type + } + if bifrostError.Type != nil { + return *bifrostError.Type + } + return "" +} + +func classifyError(bifrostError *schemas.BifrostError, typeValue string, started bool) execution.ErrorKind { + switch typeValue { + case schemas.RequestTimedOut: + return execution.ErrorKindTimeout + case schemas.RequestCancelled: + return execution.ErrorKindCanceled + case schemas.ProviderConnectionFailed: + return execution.ErrorKindTransport + case "invalid_request_error": + if !started && bifrostError != nil && bifrostError.IsBifrostError { + return execution.ErrorKindInvalidRequest + } + } + if started && bifrostError != nil && bifrostError.StatusCode != nil && *bifrostError.StatusCode >= 400 { + return execution.ErrorKindHTTP + } + if bifrostError != nil && bifrostError.IsBifrostError { + return execution.ErrorKindInternal + } + return execution.ErrorKindProvider +} + +func errorSummary(kind execution.ErrorKind, status int) string { + switch kind { + case execution.ErrorKindTransport: + return "upstream transport failed" + case execution.ErrorKindTimeout: + return "upstream request timed out" + case execution.ErrorKindCanceled: + return "upstream request canceled" + case execution.ErrorKindHTTP: + return fmt.Sprintf("upstream returned HTTP %d", status) + case execution.ErrorKindProvider: + return "upstream provider failed" + case execution.ErrorKindInvalidRequest: + return "upstream rejected the request" + default: + return "execution adapter failed" + } +} + +func sanitizeEvidenceValue(value string, secrets []string) string { + value = strings.TrimSpace(value) + if value == "" || containsSecret(value, secrets) || strings.ContainsAny(value, "\r\n\x00") { + return "" + } + const limit = 128 + if utf8.RuneCountInString(value) <= limit { + return value + } + runes := []rune(value) + return string(runes[:limit]) +} + +func containsSecret(value string, secrets []string) bool { + for _, secret := range secrets { + if secret != "" && strings.Contains(value, secret) { + return true + } + } + return false +} + +func keySecrets(key *schemas.Key) []string { + if key == nil { + return nil + } + seen := make(map[string]struct{}) + result := make([]string, 0, 12) + appendSecret := func(secret *schemas.SecretVar) { + if secret == nil { + return + } + value := secret.GetValue() + if value == "" { + return + } + if _, exists := seen[value]; exists { + return + } + seen[value] = struct{}{} + result = append(result, value) + } + appendSecret(&key.Value) + if config := key.AzureKeyConfig; config != nil { + appendSecret(config.ClientID) + appendSecret(config.ClientSecret) + appendSecret(config.TenantID) + } + if config := key.BedrockKeyConfig; config != nil { + appendSecret(&config.AccessKey) + appendSecret(&config.SecretKey) + appendSecret(config.SessionToken) + appendSecret(config.RoleARN) + appendSecret(config.ExternalID) + appendSecret(config.RoleSessionName) + } + if config := key.VertexKeyConfig; config != nil { + appendSecret(&config.AuthCredentials) + } + return result +} + +func encodeSafeErrorBody(evidence *execution.ErrorEvidence) []byte { + if evidence == nil { + return nil + } + body, err := json.Marshal(struct { + Error struct { + Type string `json:"type,omitempty"` + Code string `json:"code,omitempty"` + Message string `json:"message"` + } `json:"error"` + }{ + Error: struct { + Type string `json:"type,omitempty"` + Code string `json:"code,omitempty"` + Message string `json:"message"` + }{Type: evidence.Type, Code: evidence.Code, Message: evidence.Summary}, + }) + if err != nil { + return nil + } + return body +} + +func notSentUnaryFailure(kind execution.ErrorKind, summary string) execution.AttemptResult { + return execution.AttemptResult{ + DispatchState: execution.DispatchNotSent, + Error: &execution.ErrorEvidence{Kind: kind, Summary: summary}, + } +} + +func attemptedUnaryFailure(kind execution.ErrorKind, summary string) execution.AttemptResult { + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + Error: &execution.ErrorEvidence{Kind: kind, Summary: summary}, + } +} + +func startedUnaryFailure(status int, headers http.Header, kind execution.ErrorKind, summary string) execution.AttemptResult { + requestID := upstreamRequestID(headers) + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: status, + Header: headers.Clone(), + UpstreamRequestID: requestID, + Error: &execution.ErrorEvidence{ + Kind: kind, + Summary: summary, + RequestID: requestID, + }, + } +} + +func notSentStreamFailure(kind execution.ErrorKind, summary string) execution.StreamResult { + return execution.StreamResult{ + DispatchState: execution.DispatchNotSent, + Error: &execution.ErrorEvidence{Kind: kind, Summary: summary}, + } +} + +func attemptedStreamFailure(kind execution.ErrorKind, summary string) execution.StreamResult { + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + Error: &execution.ErrorEvidence{Kind: kind, Summary: summary}, + } +} + +func streamFromAttemptFailure(failure execution.AttemptResult) execution.StreamResult { + return execution.StreamResult{ + DispatchState: failure.DispatchState, + Error: cloneError(failure.Error), + } +} + +func streamSinkFailure(headers http.Header, requestID, model string, usageEvidence *execution.UsageEvidence) execution.StreamResult { + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: headers.Clone(), + Model: model, + UpstreamRequestID: requestID, + Usage: cloneUsage(usageEvidence), + Error: &execution.ErrorEvidence{ + Kind: execution.ErrorKindCanceled, + Summary: "stream consumer stopped", + RequestID: requestID, + }, + } +} + +func cloneUsage(source *execution.UsageEvidence) *execution.UsageEvidence { + if source == nil { + return nil + } + clone := source.Clone() + return &clone +} + +func cloneError(source *execution.ErrorEvidence) *execution.ErrorEvidence { + if source == nil { + return nil + } + clone := source.Clone() + return &clone +} diff --git a/internal/execution/bifrost/model_alias.go b/internal/execution/bifrost/model_alias.go new file mode 100644 index 00000000..e8a9d6e8 --- /dev/null +++ b/internal/execution/bifrost/model_alias.go @@ -0,0 +1,279 @@ +package bifrost + +import ( + "bytes" + "fmt" + + "gpt-load/internal/dialect" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +const maxNativeAliasSSEEventBytes = 10 << 20 + +const maxNativeFirstSSEEventBytes = 1 << 20 + +type nativeFirstSSEEventGate struct { + pending []byte + scanStart int + ready bool +} + +func (g *nativeFirstSSEEventGate) push(chunk []byte) ([]byte, error) { + if g == nil || g.ready { + return append([]byte(nil), chunk...), nil + } + g.pending = append(g.pending, chunk...) + if len(g.pending) > maxNativeFirstSSEEventBytes { + return nil, fmt.Errorf("first native SSE event exceeds limit") + } + for { + eventEnd, found := firstCompleteNativeSSEEvent(g.pending, g.scanStart) + if !found { + return nil, nil + } + event := g.pending[g.scanStart:eventEnd] + g.scanStart = eventEnd + if !nativeSSEEventHasData(event) { + continue + } + g.ready = true + output := append([]byte(nil), g.pending...) + g.pending = nil + g.scanStart = 0 + return output, nil + } +} + +func (g *nativeFirstSSEEventGate) finish() error { + if g == nil || g.ready { + return nil + } + return fmt.Errorf("native SSE stream ended before the first data event") +} + +func firstCompleteNativeSSEEvent(data []byte, start int) (int, bool) { + if start < 0 || start > len(data) { + return 0, false + } + lineStart := start + for index := start; index < len(data); { + if data[index] != '\r' && data[index] != '\n' { + index++ + continue + } + terminatorStart := index + index++ + if data[terminatorStart] == '\r' && index < len(data) && data[index] == '\n' { + index++ + } + if terminatorStart == lineStart { + return index, true + } + lineStart = index + } + return 0, false +} + +func nativeSSEEventHasData(event []byte) bool { + for _, line := range splitNativeSSELines(event) { + isData, value := parseNativeSSEDataLine(line.content) + if isData && len(value) > 0 { + return true + } + } + return false +} + +func needsClientModelAlias(spec execution.AttemptSpec) bool { + return spec.ClientModel != "" && spec.UpstreamModel != "" && spec.ClientModel != spec.UpstreamModel +} + +func rewriteClientResponseModel(clientProtocol protocol.Protocol, body []byte, clientModel string) ([]byte, error) { + var rewriter dialect.ModelRewriter + switch clientProtocol { + case protocol.OpenAICompletions: + rewriter = dialect.NewOpenAI() + case protocol.OpenAIResponses: + rewriter = dialect.NewOpenAIResponses() + case protocol.Anthropic: + rewriter = dialect.NewAnthropic() + case protocol.Gemini: + rewriter = dialect.NewGemini() + default: + return nil, fmt.Errorf("unsupported client protocol") + } + rewritten, err := rewriter.RewriteResponseModel(body, clientModel) + if err != nil { + return nil, fmt.Errorf("rewrite native response model") + } + return rewritten, nil +} + +type nativeAliasSSERewriter struct { + clientProtocol protocol.Protocol + clientModel string + pending []byte +} + +func newNativeAliasSSERewriter(spec execution.AttemptSpec) *nativeAliasSSERewriter { + if !needsClientModelAlias(spec) { + return nil + } + return &nativeAliasSSERewriter{ + clientProtocol: spec.ClientProtocol, + clientModel: spec.ClientModel, + } +} + +func (r *nativeAliasSSERewriter) push(chunk []byte) ([]byte, error) { + if r == nil || len(chunk) == 0 { + return append([]byte(nil), chunk...), nil + } + r.pending = append(r.pending, chunk...) + var output bytes.Buffer + for { + index, delimiterLength := firstNativeSSEDelimiter(r.pending) + if index < 0 { + if len(r.pending) > maxNativeAliasSSEEventBytes { + return nil, fmt.Errorf("native SSE event exceeds limit") + } + return output.Bytes(), nil + } + eventEnd := index + delimiterLength + if eventEnd > maxNativeAliasSSEEventBytes { + return nil, fmt.Errorf("native SSE event exceeds limit") + } + event := append([]byte(nil), r.pending[:eventEnd]...) + r.pending = r.pending[eventEnd:] + rewritten, err := rewriteClientSSEEvent(event, r.clientProtocol, r.clientModel) + if err != nil { + return nil, err + } + _, _ = output.Write(rewritten) + } +} + +func (r *nativeAliasSSERewriter) finish() ([]byte, error) { + if r == nil || len(r.pending) == 0 { + return nil, nil + } + if len(r.pending) > maxNativeAliasSSEEventBytes { + return nil, fmt.Errorf("native SSE event exceeds limit") + } + event := append([]byte(nil), r.pending...) + r.pending = nil + return rewriteClientSSEEvent(event, r.clientProtocol, r.clientModel) +} + +func firstNativeSSEDelimiter(data []byte) (int, int) { + indexLF := bytes.Index(data, []byte("\n\n")) + indexCRLF := bytes.Index(data, []byte("\r\n\r\n")) + switch { + case indexLF < 0 && indexCRLF < 0: + return -1, 0 + case indexLF < 0: + return indexCRLF, 4 + case indexCRLF < 0: + return indexLF, 2 + case indexCRLF < indexLF: + return indexCRLF, 4 + default: + return indexLF, 2 + } +} + +type nativeSSELine struct { + content []byte + terminator []byte + isData bool + data []byte +} + +func rewriteClientSSEEvent(event []byte, clientProtocol protocol.Protocol, clientModel string) ([]byte, error) { + lines := splitNativeSSELines(event) + dataValues := make([][]byte, 0, 1) + firstDataLine := -1 + for index := range lines { + if !lines[index].isData { + continue + } + if firstDataLine < 0 { + firstDataLine = index + } + dataValues = append(dataValues, lines[index].data) + } + if firstDataLine < 0 { + return append([]byte(nil), event...), nil + } + payload := bytes.Join(dataValues, []byte{'\n'}) + if len(payload) == 0 || bytes.Equal(bytes.TrimSpace(payload), []byte("[DONE]")) { + return append([]byte(nil), event...), nil + } + rewritten, err := rewriteClientResponseModel(clientProtocol, payload, clientModel) + if err != nil { + return nil, err + } + if bytes.Equal(rewritten, payload) { + return append([]byte(nil), event...), nil + } + + var output bytes.Buffer + output.Grow(len(event) - len(payload) + len(rewritten)) + for index, line := range lines { + switch { + case index == firstDataLine: + _, _ = output.WriteString("data: ") + _, _ = output.Write(rewritten) + _, _ = output.Write(line.terminator) + case line.isData: + continue + default: + _, _ = output.Write(line.content) + _, _ = output.Write(line.terminator) + } + } + return output.Bytes(), nil +} + +func splitNativeSSELines(event []byte) []nativeSSELine { + lines := make([]nativeSSELine, 0, 4) + for start := 0; start < len(event); { + end := start + for end < len(event) && event[end] != '\n' && event[end] != '\r' { + end++ + } + terminatorEnd := end + if terminatorEnd < len(event) { + terminatorEnd++ + if event[end] == '\r' && terminatorEnd < len(event) && event[terminatorEnd] == '\n' { + terminatorEnd++ + } + } + line := nativeSSELine{content: event[start:end], terminator: event[end:terminatorEnd]} + line.isData, line.data = parseNativeSSEDataLine(line.content) + lines = append(lines, line) + start = terminatorEnd + } + return lines +} + +func parseNativeSSEDataLine(line []byte) (bool, []byte) { + if len(line) == 0 || line[0] == ':' { + return false, nil + } + separator := bytes.IndexByte(line, ':') + field := line + var value []byte + if separator >= 0 { + field = line[:separator] + value = line[separator+1:] + if len(value) > 0 && value[0] == ' ' { + value = value[1:] + } + } + if !bytes.Equal(field, []byte("data")) { + return false, nil + } + return true, value +} diff --git a/internal/execution/bifrost/model_alias_test.go b/internal/execution/bifrost/model_alias_test.go new file mode 100644 index 00000000..09bb56ac --- /dev/null +++ b/internal/execution/bifrost/model_alias_test.go @@ -0,0 +1,92 @@ +package bifrost + +import ( + "bytes" + "testing" + + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestRewriteNativeResponseModelCoversNativeProtocolShapes(t *testing.T) { + t.Parallel() + + for _, test := range []struct { + name string + clientProtocol protocol.Protocol + body string + }{ + {name: "OpenAI chat", clientProtocol: protocol.OpenAICompletions, body: `{"model":"provider-model","choices":[]}`}, + {name: "OpenAI Responses object", clientProtocol: protocol.OpenAIResponses, body: `{"type":"response","model":"provider-model","output":[]}`}, + {name: "OpenAI Responses event", clientProtocol: protocol.OpenAIResponses, body: `{"type":"response.completed","response":{"model":"provider-model"}}`}, + {name: "Anthropic message", clientProtocol: protocol.Anthropic, body: `{"type":"message","model":"provider-model","content":[]}`}, + {name: "Anthropic message_start", clientProtocol: protocol.Anthropic, body: `{"type":"message_start","message":{"model":"provider-model"}}`}, + {name: "Gemini response", clientProtocol: protocol.Gemini, body: `{"modelVersion":"provider-model","candidates":[]}`}, + } { + t.Run(test.name, func(t *testing.T) { + rewritten, err := rewriteClientResponseModel(test.clientProtocol, []byte(test.body), "public-model") + if err != nil { + t.Fatalf("rewriteClientResponseModel() error = %v", err) + } + if !bytes.Contains(rewritten, []byte(`"public-model"`)) || bytes.Contains(rewritten, []byte(`"provider-model"`)) { + t.Fatalf("rewritten response = %s", rewritten) + } + }) + } +} + +func TestNativeAliasSSERewriterCoversNativeProtocolShapes(t *testing.T) { + t.Parallel() + + for _, test := range []struct { + name string + clientProtocol protocol.Protocol + payload string + }{ + {name: "OpenAI chat", clientProtocol: protocol.OpenAICompletions, payload: `{"model":"provider-model","choices":[]}`}, + {name: "OpenAI Responses", clientProtocol: protocol.OpenAIResponses, payload: `{"type":"response.completed","response":{"model":"provider-model"}}`}, + {name: "Anthropic", clientProtocol: protocol.Anthropic, payload: `{"type":"message_start","message":{"model":"provider-model"}}`}, + {name: "Gemini", clientProtocol: protocol.Gemini, payload: `{"modelVersion":"provider-model","candidates":[]}`}, + } { + t.Run(test.name, func(t *testing.T) { + rewriter := newNativeAliasSSERewriter(execution.AttemptSpec{ + ClientProtocol: test.clientProtocol, + ClientModel: "public-model", + UpstreamModel: "provider-model", + }) + raw := []byte("event: update\r\ndata: " + test.payload + "\r\n\r\ndata: [DONE]\r\n\r\n") + split := len(raw) / 2 + first, err := rewriter.push(raw[:split]) + if err != nil { + t.Fatalf("first push error = %v", err) + } + second, err := rewriter.push(raw[split:]) + if err != nil { + t.Fatalf("second push error = %v", err) + } + tail, err := rewriter.finish() + if err != nil { + t.Fatalf("finish error = %v", err) + } + output := bytes.Join([][]byte{first, second, tail}, nil) + if !bytes.Contains(output, []byte(`"public-model"`)) || bytes.Contains(output, []byte(`"provider-model"`)) { + t.Fatalf("rewritten SSE = %s", output) + } + if !bytes.Contains(output, []byte("data: [DONE]\r\n\r\n")) { + t.Fatalf("terminal SSE marker changed: %q", output) + } + }) + } +} + +func TestNativeAliasSSERewriterIsDisabledForEqualAliases(t *testing.T) { + t.Parallel() + + if rewriter := newNativeAliasSSERewriter(execution.AttemptSpec{ + ClientProtocol: protocol.OpenAICompletions, + ClientModel: "same-model", + UpstreamModel: "same-model", + }); rewriter != nil { + t.Fatal("equal aliases enabled native SSE rewriting") + } +} diff --git a/internal/execution/bifrost/native_protocols_test.go b/internal/execution/bifrost/native_protocols_test.go new file mode 100644 index 00000000..1ce2aebb --- /dev/null +++ b/internal/execution/bifrost/native_protocols_test.go @@ -0,0 +1,245 @@ +package bifrost + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" + "gpt-load/internal/usage" +) + +func TestOfficialAnthropicMessagesNativeUnary(t *testing.T) { + t.Parallel() + + const responseBody = `{"id": "msg_1", "type":"message", "role":"assistant", "model":"served-claude", "content":[{"type":"text","text":"ok"}], "stop_reason":"end_turn", "usage":{"input_tokens":6,"cache_read_input_tokens":4,"output_tokens":3}, "vendor_response":{"precise":1.2300}}` + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Method != http.MethodPost || request.URL.Path != "/v1/messages" { + t.Errorf("request target = %s %s", request.Method, request.URL.Path) + } + if request.Header.Get("X-Api-Key") != testAPIKey || request.Header.Get("Authorization") != "" { + t.Errorf("credential headers = x-api-key:%q authorization:%q", request.Header.Get("X-Api-Key"), request.Header.Get("Authorization")) + } + if request.Header.Get("Anthropic-Version") == "" { + t.Error("anthropic-version was not injected") + } + body, err := io.ReadAll(request.Body) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatal(err) + } + if payload["model"] != "upstream-claude" || payload["stream"] != false { + t.Errorf("request controls = %#v", payload) + } + for _, field := range []string{"provider", "fallbacks", "authorization", "api_key"} { + if _, exists := payload[field]; exists { + t.Errorf("control field %q reached upstream", field) + } + } + if !bytes.Contains(body, []byte(`"precise":1.2300`)) { + t.Errorf("unknown field changed: %s", body) + } + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("Request-Id", "anthropic-unary") + _, _ = io.WriteString(writer, responseBody) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, anthropicBaseURL: server.URL}) + spec := anthropicSpec(false) + spec.ClientModel = spec.UpstreamModel + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if string(result.Body) != responseBody || result.UpstreamRequestID != "anthropic-unary" { + t.Fatalf("result/body = %+v/%s", result, result.Body) + } + assertUsage(t, result.Usage, usage.Tokens{UncachedInput: 6, CacheRead: 4, Output: 3}) +} + +func TestOfficialAnthropicMessagesNativeStream(t *testing.T) { + t.Parallel() + + const rawStream = "event: message_start\n" + + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_stream\",\"model\":\"served-claude\",\"usage\":{\"input_tokens\":6,\"cache_read_input_tokens\":4,\"output_tokens\":0}}}\n\n" + + "event: content_block_delta\n" + + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"},\"vendor\":{\"precise\":1.2300}}\n\n" + + "event: message_delta\n" + + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":3}}\n\n" + + "event: message_stop\n" + + "data: {\"type\":\"message_stop\"}\n\n" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + body, _ := io.ReadAll(request.Body) + var payload map[string]any + _ = json.Unmarshal(body, &payload) + if payload["model"] != "upstream-claude" || payload["stream"] != true { + t.Errorf("request controls = %#v", payload) + } + writer.Header().Set("Content-Type", "text/event-stream") + writer.Header().Set("Request-Id", "anthropic-stream") + writer.WriteHeader(http.StatusOK) + _, _ = io.WriteString(writer, rawStream) + writer.(http.Flusher).Flush() + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, anthropicBaseURL: server.URL}) + spec := anthropicSpec(true) + spec.ClientModel = spec.UpstreamModel + var events []execution.StreamEvent + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + return nil + }) + assertRawStream(t, result, events, rawStream, usage.Tokens{UncachedInput: 6, CacheRead: 4, Output: 3}) +} + +func TestOfficialGeminiGenerateNativeUnary(t *testing.T) { + t.Parallel() + + const responseBody = `{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}], "modelVersion":"served-gemini", "usageMetadata":{"promptTokenCount":10,"cachedContentTokenCount":4,"candidatesTokenCount":3,"totalTokenCount":13}, "vendorResponse":{"precise":1.2300}}` + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Method != http.MethodPost || request.URL.Path != "/v1beta/models/upstream-gemini:generateContent" { + t.Errorf("request target = %s %s", request.Method, request.URL.Path) + } + if request.Header.Get("X-Goog-Api-Key") != testAPIKey || request.Header.Get("Authorization") != "" { + t.Errorf("credential headers = goog:%q auth:%q", request.Header.Get("X-Goog-Api-Key"), request.Header.Get("Authorization")) + } + body, err := io.ReadAll(request.Body) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatal(err) + } + for _, field := range []string{"provider", "fallbacks", "api_key"} { + if _, exists := payload[field]; exists { + t.Errorf("body control field %q reached upstream", field) + } + } + if payload["model"] != "client-gemini" || payload["stream"] != true { + t.Errorf("native Gemini body fields changed: %#v", payload) + } + if !bytes.Contains(body, []byte(`"precise":1.2300`)) { + t.Errorf("unknown field changed: %s", body) + } + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Goog-Request-Id", "gemini-unary") + _, _ = io.WriteString(writer, responseBody) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, geminiBaseURL: server.URL + "/v1beta"}) + spec := geminiSpec(false) + spec.ClientModel = spec.UpstreamModel + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if string(result.Body) != responseBody || result.UpstreamRequestID != "gemini-unary" { + t.Fatalf("result/body = %+v/%s", result, result.Body) + } + assertUsage(t, result.Usage, usage.Tokens{UncachedInput: 6, CacheRead: 4, Output: 3}) +} + +func TestOfficialGeminiGenerateNativeStream(t *testing.T) { + t.Parallel() + + const rawStream = "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"hi\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":10,\"cachedContentTokenCount\":4,\"candidatesTokenCount\":3,\"totalTokenCount\":13},\"vendor\":{\"precise\":1.2300}}\n\n" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/v1beta/models/upstream-gemini:streamGenerateContent" || request.URL.Query().Get("alt") != "sse" { + t.Errorf("request target = %s?%s", request.URL.Path, request.URL.RawQuery) + } + writer.Header().Set("Content-Type", "text/event-stream") + writer.Header().Set("X-Goog-Request-Id", "gemini-stream") + writer.WriteHeader(http.StatusOK) + _, _ = io.WriteString(writer, rawStream) + writer.(http.Flusher).Flush() + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, geminiBaseURL: server.URL + "/v1beta"}) + spec := geminiSpec(true) + spec.ClientModel = spec.UpstreamModel + spec.Query.Set("alt", "sse") + var events []execution.StreamEvent + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + return nil + }) + assertRawStream(t, result, events, rawStream, usage.Tokens{UncachedInput: 6, CacheRead: 4, Output: 3}) +} + +func anthropicSpec(stream bool) execution.AttemptSpec { + return execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: "anthropic-request", AttemptID: "anthropic-attempt", Sequence: 1, + ChannelID: string(channel.Anthropic), ClientProtocol: protocol.Anthropic, + Operation: execution.OperationChatCompletion, ClientModel: "client-claude", UpstreamModel: "upstream-claude", + Method: http.MethodPost, Path: "/v1/messages", Query: make(map[string][]string), + Header: http.Header{"Authorization": {"Bearer client"}, "X-Api-Key": {"client"}, "X-Test-Header": {"forward-me"}}, + Body: []byte(`{"model":"client-claude","stream":true,"provider":"injected","fallbacks":["other"],"authorization":"body","api_key":"body","max_tokens":64,"messages":[{"role":"user","content":"hello"}],"vendor_extension":{"precise":1.2300}}`), + TargetConfig: json.RawMessage(`{}`), Credential: execution.NewCredentialSnapshot(8, 1, 1, []byte(`{"api_key":"`+testAPIKey+`"}`)), + }) +} + +func geminiSpec(stream bool) execution.AttemptSpec { + action := "generateContent" + if stream { + action = "streamGenerateContent" + } + return execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: "gemini-request", AttemptID: "gemini-attempt", Sequence: 1, + ChannelID: string(channel.Gemini), ClientProtocol: protocol.Gemini, + Operation: execution.OperationChatCompletion, ClientModel: "client-gemini", UpstreamModel: "upstream-gemini", + Method: http.MethodPost, Path: "/v1beta/models/client-gemini:" + action, Query: make(map[string][]string), + Header: http.Header{"Authorization": {"Bearer client"}, "X-Goog-Api-Key": {"client"}, "X-Test-Header": {"forward-me"}}, + Body: []byte(`{"model":"client-gemini","stream":true,"provider":"injected","fallbacks":["other"],"api_key":"body","contents":[{"role":"user","parts":[{"text":"hello"}]}],"vendor_extension":{"precise":1.2300}}`), + TargetConfig: json.RawMessage(`{}`), Credential: execution.NewCredentialSnapshot(9, 1, 1, []byte(`{"api_key":"`+testAPIKey+`"}`)), + }) +} + +func newProtocolTestRuntime(t *testing.T, options runtimeOptions) *Runtime { + t.Helper() + runtime, err := newRuntime(context.Background(), options) + if err != nil { + t.Fatalf("new runtime: %v", err) + } + t.Cleanup(runtime.Shutdown) + return runtime +} + +func assertRawStream(t *testing.T, result execution.StreamResult, events []execution.StreamEvent, raw string, want usage.Tokens) { + t.Helper() + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if len(events) < 3 || events[0].Kind != execution.StreamEventReady { + t.Fatalf("events = %+v", events) + } + var data bytes.Buffer + var sawUsage bool + for _, event := range events[1:] { + if event.Kind == execution.StreamEventData { + data.Write(event.Data) + } + if event.Kind == execution.StreamEventUsage { + sawUsage = true + assertUsage(t, event.Usage, want) + } + } + if data.String() != raw || !sawUsage { + t.Fatalf("stream data/usage = %q/%t", data.String(), sawUsage) + } + assertUsage(t, result.Usage, want) +} diff --git a/internal/execution/bifrost/passthrough.go b/internal/execution/bifrost/passthrough.go new file mode 100644 index 00000000..50097ddd --- /dev/null +++ b/internal/execution/bifrost/passthrough.go @@ -0,0 +1,676 @@ +package bifrost + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strconv" + "strings" + + "github.com/maximhq/bifrost/core/schemas" + + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +const maxStreamErrorEvidenceBytes = 64 << 10 + +type passthroughUnarySDKResult struct { + response *schemas.BifrostPassthroughResponse + err *schemas.BifrostError +} + +type passthroughStreamSDKResult struct { + stream chan *schemas.BifrostStreamChunk + err *schemas.BifrostError +} + +func sanitizeNativeChatBody(body []byte, upstreamModel string, stream ...bool) ([]byte, error) { + object, err := decodeNativeJSONObject(body) + if err != nil { + return nil, err + } + changed := stripNativeControlFields(object) + if !jsonStringEquals(object["model"], upstreamModel) { + if err := forceJSONModel(object, upstreamModel); err != nil { + return nil, err + } + changed = true + } + if len(stream) > 0 { + _, exists := object["stream"] + if (exists && !jsonBoolEquals(object["stream"], stream[0])) || (!exists && stream[0]) { + object["stream"] = json.RawMessage(strconv.FormatBool(stream[0])) + changed = true + } + } + if !changed { + return append([]byte(nil), body...), nil + } + return encodeNativeJSONObject(object) +} + +func sanitizeNativeRequestBody(spec execution.AttemptSpec, stream bool) ([]byte, error) { + switch { + case spec.Operation == execution.OperationListModels: + return nil, nil + case (spec.ClientProtocol == protocol.OpenAICompletions || spec.ClientProtocol == protocol.Anthropic) && + (spec.Operation == execution.OperationChatCompletion || spec.Operation == execution.OperationProbe): + body, err := sanitizeNativeChatBody(spec.Body, spec.UpstreamModel, stream) + if err != nil { + return nil, err + } + if spec.ClientProtocol == protocol.OpenAICompletions && stream && spec.IncludeUsage { + body, err = includeOpenAIStreamUsage(body) + if err != nil { + return nil, err + } + } + if spec.Operation != execution.OperationProbe { + return body, nil + } + object, err := decodeNativeJSONObject(body) + if err != nil { + return nil, err + } + if _, exists := object["stream"]; !exists { + object["stream"] = json.RawMessage("false") + return encodeNativeJSONObject(object) + } + return body, nil + case spec.ClientProtocol == protocol.Gemini && + (spec.Operation == execution.OperationChatCompletion || spec.Operation == execution.OperationProbe): + object, err := decodeNativeJSONObject(spec.Body) + if err != nil { + // Preserve malformed/opaque provider payloads for the upstream to + // validate, while valid JSON is scrubbed below. + return append([]byte(nil), spec.Body...), nil + } + changed := stripNativeControlFields(object) + if !changed { + return append([]byte(nil), spec.Body...), nil + } + return encodeNativeJSONObject(object) + case spec.ClientProtocol == protocol.OpenAIResponses && + (spec.Operation == execution.OperationResponsesCreate || spec.Operation == execution.OperationProbe): + object, err := decodeNativeJSONObject(spec.Body) + if err != nil { + return nil, err + } + changed := stripNativeControlFields(object) + if !jsonStringEquals(object["model"], spec.UpstreamModel) { + if err := forceJSONModel(object, spec.UpstreamModel); err != nil { + return nil, err + } + changed = true + } + _, exists := object["stream"] + if (exists && !jsonBoolEquals(object["stream"], stream)) || (!exists && stream) { + object["stream"] = json.RawMessage(strconv.FormatBool(stream)) + changed = true + } + if !changed { + return append([]byte(nil), spec.Body...), nil + } + return encodeNativeJSONObject(object) + case spec.ClientProtocol == protocol.OpenAIResponses && + (spec.Operation == execution.OperationResponsesCompact || spec.Operation == execution.OperationResponsesInputTokens): + object, err := decodeNativeJSONObject(spec.Body) + if err != nil { + return nil, err + } + changed := stripNativeControlFields(object) + if !jsonStringEquals(object["model"], spec.UpstreamModel) { + if err := forceJSONModel(object, spec.UpstreamModel); err != nil { + return nil, err + } + changed = true + } + if _, exists := object["stream"]; exists { + delete(object, "stream") + changed = true + } + if !changed { + return append([]byte(nil), spec.Body...), nil + } + return encodeNativeJSONObject(object) + case spec.ClientProtocol == protocol.OpenAIResponses && spec.Operation == execution.OperationResponsesPassthrough: + if len(bytes.TrimSpace(spec.Body)) == 0 { + return nil, nil + } + object, err := decodeNativeJSONObject(spec.Body) + if err != nil { + return nil, err + } + changed := stripNativeControlFields(object) + if _, hasModel := object["model"]; hasModel && strings.TrimSpace(spec.UpstreamModel) != "" { + if !jsonStringEquals(object["model"], spec.UpstreamModel) { + if err := forceJSONModel(object, spec.UpstreamModel); err != nil { + return nil, err + } + changed = true + } + } + if !changed { + return append([]byte(nil), spec.Body...), nil + } + return encodeNativeJSONObject(object) + default: + if len(bytes.TrimSpace(spec.Body)) == 0 { + return nil, nil + } + object, err := decodeNativeJSONObject(spec.Body) + if err != nil { + return nil, err + } + stripNativeControlFields(object) + return encodeNativeJSONObject(object) + } +} + +func decodeNativeJSONObject(body []byte) (map[string]json.RawMessage, error) { + decoder := json.NewDecoder(bytes.NewReader(body)) + var object map[string]json.RawMessage + if err := decoder.Decode(&object); err != nil || object == nil { + return nil, fmt.Errorf("decode request object") + } + var trailing json.RawMessage + if err := decoder.Decode(&trailing); err != io.EOF { + return nil, fmt.Errorf("request body has trailing data") + } + return object, nil +} + +func stripNativeControlFields(object map[string]json.RawMessage) bool { + changed := false + for key := range object { + normalized := strings.ReplaceAll(strings.ToLower(key), "-", "_") + switch normalized { + case "provider", "fallback", "fallbacks", "authorization", "proxy_authorization", + "api_key", "apikey", "x_api_key", "x_goog_api_key": + delete(object, key) + changed = true + } + } + return changed +} + +func jsonStringEquals(raw json.RawMessage, expected string) bool { + var value string + return json.Unmarshal(raw, &value) == nil && value == expected +} + +func jsonBoolEquals(raw json.RawMessage, expected bool) bool { + var value bool + return json.Unmarshal(raw, &value) == nil && value == expected +} + +func includeOpenAIStreamUsage(body []byte) ([]byte, error) { + object, err := decodeNativeJSONObject(body) + if err != nil { + return nil, err + } + options := make(map[string]json.RawMessage) + if raw, exists := object["stream_options"]; exists && + !bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + if err := json.Unmarshal(raw, &options); err != nil || options == nil { + return nil, fmt.Errorf("stream_options must be an object or null") + } + } + options["include_usage"] = json.RawMessage("true") + encoded, err := json.Marshal(options) + if err != nil { + return nil, fmt.Errorf("encode stream_options") + } + object["stream_options"] = encoded + return encodeNativeJSONObject(object) +} + +func forceJSONModel(object map[string]json.RawMessage, upstreamModel string) error { + model, err := json.Marshal(upstreamModel) + if err != nil { + return fmt.Errorf("encode selected model") + } + object["model"] = model + return nil +} + +func encodeNativeJSONObject(object map[string]json.RawMessage) ([]byte, error) { + encoded, err := json.Marshal(object) + if err != nil { + return nil, fmt.Errorf("encode request object") + } + return encoded, nil +} + +func (r *Runtime) executeNative( + parent context.Context, + spec execution.AttemptSpec, + prepared preparedAttempt, +) execution.AttemptResult { + requestContext, requestCancel := boundedRequestContext(parent, spec.Timeouts.Request) + defer requestCancel() + callContext, callCancel := context.WithCancel(requestContext) + defer callCancel() + preResponse := startPreResponseGate(callCancel, spec.Timeouts) + defer preResponse.stop() + + bifrostContext := newSDKContext(callContext, spec, prepared.directKey) + outcomeChannel := make(chan passthroughUnarySDKResult, 1) + go func() { + response, bifrostError := r.core.Passthrough(bifrostContext, prepared.provider, prepared.passthrough) + outcomeChannel <- passthroughUnarySDKResult{response: response, err: bifrostError} + }() + + var outcome passthroughUnarySDKResult + select { + case outcome = <-outcomeChannel: + if preResponse.expired() || requestContext.Err() != nil { + return unaryContextFailure(requestContext, preResponse.expired()) + } + case <-callContext.Done(): + return unaryContextFailure(requestContext, preResponse.expired()) + } + preResponse.stop() + + if outcome.err != nil { + return unaryErrorResult(outcome.err, bifrostContext, prepared.secrets) + } + if outcome.response == nil { + return attemptedUnaryFailure(execution.ErrorKindInternal, "execution runtime returned no response") + } + status := outcome.response.StatusCode + if !validUpstreamStatus(status) { + return attemptedUnaryFailure(execution.ErrorKindInternal, "execution runtime returned an invalid status") + } + headers := responseHeaders(outcome.response.Headers, bifrostContext, false) + requestID := upstreamRequestID(headers) + body := append([]byte(nil), outcome.response.Body...) + model := openAIResponseModel(body, spec.UpstreamModel) + usageEvidence, err := usageEvidenceFromPassthrough(outcome.response.PassthroughUsage) + if err != nil { + return startedUnaryFailure(status, headers, execution.ErrorKindInternal, "normalize upstream usage") + } + if status >= http.StatusOK && status < http.StatusMultipleChoices { + if needsClientModelAlias(spec) { + body, err = rewriteClientResponseModel(spec.ClientProtocol, body, spec.ClientModel) + if err != nil { + return startedUnaryFailure(status, headers, execution.ErrorKindInternal, "rewrite native response model") + } + } + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: status, + Header: headers, + Body: body, + Model: model, + UpstreamRequestID: requestID, + Usage: usageEvidence, + } + } + evidence := passthroughHTTPError(status, headers, body, prepared.secrets) + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: status, + Header: headers, + Body: redactSecrets(body, prepared.secrets), + Model: model, + UpstreamRequestID: requestID, + Error: evidence, + } +} + +func (r *Runtime) executeNativeStream( + parent context.Context, + spec execution.AttemptSpec, + prepared preparedAttempt, + sink execution.StreamSink, +) execution.StreamResult { + requestContext, requestCancel := boundedRequestContext(parent, spec.Timeouts.Request) + defer requestCancel() + callContext, callCancel := context.WithCancel(requestContext) + bifrostContext := newSDKContext(callContext, spec, prepared.directKey) + cancelCall := func() { + callCancel() + bifrostContext.Cancel() + } + defer cancelCall() + preResponse := startPreResponseGate(cancelCall, spec.Timeouts) + defer preResponse.stop() + + outcomeChannel := make(chan passthroughStreamSDKResult, 1) + go func() { + stream, bifrostError := r.core.PassthroughStream(bifrostContext, prepared.provider, prepared.passthrough) + outcomeChannel <- passthroughStreamSDKResult{stream: stream, err: bifrostError} + }() + + var outcome passthroughStreamSDKResult + select { + case outcome = <-outcomeChannel: + if preResponse.expired() || requestContext.Err() != nil { + return streamContextFailure(requestContext, preResponse.expired(), false, nil, "", "", nil) + } + case <-callContext.Done(): + return streamContextFailure(requestContext, preResponse.expired(), false, nil, "", "", nil) + } + if outcome.err != nil { + return streamErrorResult(outcome.err, bifrostContext, prepared.secrets, false, 0, nil, "", nil) + } + if outcome.stream == nil { + return attemptedStreamFailure(execution.ErrorKindInternal, "execution runtime returned no stream") + } + + sequence := uint64(0) + responseObserved := false + started := false + status := 0 + var headers http.Header + requestID := "" + model := spec.UpstreamModel + var usageEvidence *execution.UsageEvidence + var errorBody bytes.Buffer + firstEventGate := &nativeFirstSSEEventGate{} + aliasRewriter := newNativeAliasSSERewriter(spec) + idleTimer := newIdleTimer(spec.Timeouts.StreamIdle) + defer idleTimer.stop() + emitReady := func() error { + if started { + return nil + } + preResponse.stop() + started = true + sequence = 1 + return sink(execution.StreamEvent{ + Sequence: sequence, + Kind: execution.StreamEventReady, + StatusCode: status, + Header: headers.Clone(), + UpstreamRequestID: requestID, + }) + } + + for { + select { + case <-callContext.Done(): + return nativeStreamContextFailure(requestContext, preResponse.expired(), started, status, headers, requestID, model, usageEvidence) + case <-requestContext.Done(): + cancelCall() + return nativeStreamContextFailure(requestContext, false, started, status, headers, requestID, model, usageEvidence) + case <-idleTimer.channel(): + cancelCall() + return nativeStreamContextFailure(requestContext, true, started, status, headers, requestID, model, usageEvidence) + case chunk, open := <-outcome.stream: + idleTimer.pause() + if callContext.Err() != nil || requestContext.Err() != nil || preResponse.expired() { + return nativeStreamContextFailure(requestContext, preResponse.expired(), started, status, headers, requestID, model, usageEvidence) + } + if !open { + if !responseObserved { + return attemptedStreamFailure(execution.ErrorKindInternal, "execution runtime returned an empty stream") + } + if status >= http.StatusOK && status < http.StatusMultipleChoices { + if err := firstEventGate.finish(); err != nil { + return attemptedStreamFailure(execution.ErrorKindInternal, "invalid upstream SSE stream") + } + } + if aliasRewriter != nil && status >= http.StatusOK && status < http.StatusMultipleChoices { + data, err := aliasRewriter.finish() + if err != nil { + return streamAliasFailure(status, headers, requestID, model, usageEvidence) + } + if len(data) > 0 { + sequence++ + if err := sink(execution.StreamEvent{Sequence: sequence, Kind: execution.StreamEventData, Data: data}); err != nil { + return nativeStreamSinkFailure(status, headers, requestID, model, usageEvidence) + } + } + } + return finishNativeStream(status, headers, requestID, model, usageEvidence, errorBody.Bytes(), prepared.secrets) + } + if chunk == nil { + cancelCall() + return nativeStreamRuntimeFailure(bifrostContext, prepared.secrets, started, status, headers, model, usageEvidence) + } + if chunk.BifrostError != nil { + cancelCall() + return streamErrorResult(chunk.BifrostError, bifrostContext, prepared.secrets, started, status, headers, model, usageEvidence) + } + response := chunk.BifrostPassthroughResponse + if response == nil || !validUpstreamStatus(response.StatusCode) { + cancelCall() + return nativeStreamRuntimeFailure(bifrostContext, prepared.secrets, started, status, headers, model, usageEvidence) + } + if !responseObserved { + responseObserved = true + status = response.StatusCode + headers = responseHeaders(response.Headers, bifrostContext, true) + requestID = upstreamRequestID(headers) + if status < http.StatusOK || status >= http.StatusMultipleChoices { + if err := emitReady(); err != nil { + callCancel() + return nativeStreamSinkFailure(status, headers, requestID, model, usageEvidence) + } + } + } else if response.StatusCode != status { + cancelCall() + return nativeStreamRuntimeFailure(bifrostContext, prepared.secrets, started, status, headers, model, usageEvidence) + } + + if len(response.Body) > 0 { + data := append([]byte(nil), response.Body...) + if status < http.StatusOK || status >= http.StatusMultipleChoices { + appendBounded(&errorBody, data, maxStreamErrorEvidenceBytes) + data = redactSecrets(data, prepared.secrets) + } else { + var err error + data, err = firstEventGate.push(data) + if err != nil { + cancelCall() + if !started { + return attemptedStreamFailure(execution.ErrorKindInternal, "invalid upstream SSE stream") + } + return nativeStreamProtocolFailure(status, headers, requestID, model, usageEvidence) + } + if aliasRewriter != nil && len(data) > 0 { + data, err = aliasRewriter.push(data) + if err != nil { + cancelCall() + if !started { + return attemptedStreamFailure(execution.ErrorKindInternal, "rewrite client response model") + } + return streamAliasFailure(status, headers, requestID, model, usageEvidence) + } + } + } + if len(data) > 0 { + if err := emitReady(); err != nil { + cancelCall() + return nativeStreamSinkFailure(status, headers, requestID, model, usageEvidence) + } + sequence++ + if err := sink(execution.StreamEvent{Sequence: sequence, Kind: execution.StreamEventData, Data: data}); err != nil { + cancelCall() + return nativeStreamSinkFailure(status, headers, requestID, model, usageEvidence) + } + } + } + if response.PassthroughUsage != nil { + chunkUsage, err := usageEvidenceFromPassthrough(response.PassthroughUsage) + if err != nil { + cancelCall() + return nativeStreamRuntimeFailure(bifrostContext, prepared.secrets, started, status, headers, model, usageEvidence) + } + if chunkUsage != nil { + usageEvidence = cloneUsage(chunkUsage) + if !started { + idleTimer.resume() + continue + } + sequence++ + if err := sink(execution.StreamEvent{ + Sequence: sequence, + Kind: execution.StreamEventUsage, + Usage: cloneUsage(chunkUsage), + }); err != nil { + cancelCall() + return nativeStreamSinkFailure(status, headers, requestID, model, usageEvidence) + } + } + } + idleTimer.resume() + } + } +} + +func nativeStreamRuntimeFailure( + bifrostContext *schemas.BifrostContext, + secrets []string, + started bool, + status int, + headers http.Header, + model string, + usageEvidence *execution.UsageEvidence, +) execution.StreamResult { + return streamErrorResult(nil, bifrostContext, secrets, started, status, headers, model, usageEvidence) +} + +func nativeStreamContextFailure( + ctx context.Context, + preResponseExpired bool, + started bool, + status int, + headers http.Header, + requestID string, + model string, + usageEvidence *execution.UsageEvidence, +) execution.StreamResult { + if !started { + return streamContextFailure(ctx, preResponseExpired, false, nil, "", model, usageEvidence) + } + kind, summary := contextFailure(ctx, preResponseExpired) + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: status, + Header: headers.Clone(), + Model: model, + UpstreamRequestID: requestID, + Usage: cloneUsage(usageEvidence), + Error: &execution.ErrorEvidence{ + Kind: kind, + Summary: summary, + RequestID: requestID, + }, + } +} + +func nativeStreamSinkFailure( + status int, + headers http.Header, + requestID string, + model string, + usageEvidence *execution.UsageEvidence, +) execution.StreamResult { + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: status, + Header: headers.Clone(), + Model: model, + UpstreamRequestID: requestID, + Usage: cloneUsage(usageEvidence), + Error: &execution.ErrorEvidence{ + Kind: execution.ErrorKindCanceled, + Summary: "stream consumer stopped", + RequestID: requestID, + }, + } +} + +func streamAliasFailure( + status int, + headers http.Header, + requestID string, + model string, + usageEvidence *execution.UsageEvidence, +) execution.StreamResult { + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: status, + Header: headers.Clone(), + Model: model, + UpstreamRequestID: requestID, + Usage: cloneUsage(usageEvidence), + Error: &execution.ErrorEvidence{ + Kind: execution.ErrorKindInternal, + Summary: "rewrite client response model", + RequestID: requestID, + }, + } +} + +func nativeStreamProtocolFailure( + status int, + headers http.Header, + requestID string, + model string, + usageEvidence *execution.UsageEvidence, +) execution.StreamResult { + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: status, + Header: headers.Clone(), + Model: model, + UpstreamRequestID: requestID, + Usage: cloneUsage(usageEvidence), + Error: &execution.ErrorEvidence{ + Kind: execution.ErrorKindInternal, + Summary: "invalid upstream SSE stream", + RequestID: requestID, + }, + } +} + +func finishNativeStream( + status int, + headers http.Header, + requestID string, + model string, + usageEvidence *execution.UsageEvidence, + errorBody []byte, + secrets []string, +) execution.StreamResult { + result := execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: status, + Header: headers.Clone(), + Model: model, + UpstreamRequestID: requestID, + Usage: cloneUsage(usageEvidence), + } + if status < http.StatusOK || status >= http.StatusMultipleChoices { + result.Error = passthroughHTTPError(status, headers, errorBody, secrets) + } + return result +} + +func appendBounded(destination *bytes.Buffer, source []byte, limit int) { + if destination == nil || limit <= 0 || destination.Len() >= limit { + return + } + remaining := limit - destination.Len() + if len(source) > remaining { + source = source[:remaining] + } + _, _ = destination.Write(source) +} + +func validUpstreamStatus(status int) bool { return status >= 100 && status <= 599 } diff --git a/internal/execution/bifrost/responses_passthrough_test.go b/internal/execution/bifrost/responses_passthrough_test.go new file mode 100644 index 00000000..83d59423 --- /dev/null +++ b/internal/execution/bifrost/responses_passthrough_test.go @@ -0,0 +1,129 @@ +package bifrost + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestOpenAIResponsesUnknownNamespaceUsesExplicitNativePassthrough(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID channel.ID + method string + path string + body []byte + status int + response string + target func(string) json.RawMessage + newRuntime func(*testing.T, string) *Runtime + }{ + { + name: "official vendor extension", channelID: channel.OpenAI, + method: http.MethodPatch, path: "/v1/responses/vendor_extension", + body: []byte(`{"model":"client-model","provider":"injected","api_key":"injected","vendor":{"precise":1.2300}}`), + status: http.StatusMultiStatus, + response: `{"object":"vendor.result","vendor":{"precise":1.2300}}`, + target: func(string) json.RawMessage { return json.RawMessage(`{}`) }, + newRuntime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: base}) + }, + }, + { + name: "compatible complete v1 prefix", channelID: channel.OpenAICompatible, + method: http.MethodPatch, path: "/v1/responses/vendor_extension", + body: []byte(`{"model":"client-model","vendor":{"precise":1.2300}}`), + status: http.StatusOK, + response: `{"object":"vendor.result"}`, + target: func(base string) json.RawMessage { + encoded, _ := json.Marshal(map[string]string{"base_url": base + "/v1"}) + return encoded + }, + newRuntime: func(t *testing.T, _ string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true}) + }, + }, + { + name: "model optional head resource", channelID: channel.OpenAI, + method: http.MethodHead, path: "/v1/responses/resp_123/vendor_state", + status: http.StatusNoContent, + target: func(string) json.RawMessage { return json.RawMessage(`{}`) }, + newRuntime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: base}) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + if request.Method != test.method || request.URL.Path != test.path || request.URL.Query().Get("vendor") != "kept" { + t.Errorf("target = %s %s?%s", request.Method, request.URL.Path, request.URL.RawQuery) + } + if request.Header.Get("Authorization") != "Bearer "+testAPIKey || request.Header.Get("Api-Key") != "" { + t.Errorf("credential headers = %#v", request.Header) + } + body, _ := io.ReadAll(request.Body) + if len(test.body) == 0 { + if len(body) != 0 { + t.Errorf("unexpected body = %s", body) + } + } else { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("decode body: %v", err) + } + if payload["model"] != "upstream-model" || payload["provider"] != nil || payload["api_key"] != nil || payload["vendor"] == nil { + t.Errorf("sanitized body = %#v", payload) + } + } + writer.Header().Set("X-Request-Id", "vendor-request") + writer.WriteHeader(test.status) + _, _ = io.WriteString(writer, test.response) + })) + defer server.Close() + + runtime := test.newRuntime(t, server.URL) + spec := convertedSpec(test.channelID, protocol.OpenAIResponses, execution.OperationResponsesPassthrough, test.path, test.body) + spec.Method = test.method + spec.TargetConfig = test.target(server.URL) + spec.Query.Set("vendor", "kept") + if len(test.body) == 0 { + spec.ClientModel = "" + spec.UpstreamModel = "" + } + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if calls.Load() != 1 || result.Error != nil || result.StatusCode != test.status || string(result.Body) != test.response { + t.Fatalf("calls/result = %d/%+v body=%s", calls.Load(), result, result.Body) + } + }) + } +} + +func TestResponsesPassthroughFailsClosedForNonOpenAIProvider(t *testing.T) { + t.Parallel() + + runtime := newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true}) + spec := convertedSpec(channel.Anthropic, protocol.OpenAIResponses, execution.OperationResponsesPassthrough, "/v1/responses/vendor_extension", []byte(`{"vendor":true}`)) + spec.ClientModel = "" + spec.UpstreamModel = "" + result := runtime.Execute(context.Background(), spec) + if result.DispatchState != execution.DispatchNotSent || result.ResponseStarted || result.Error == nil || result.Error.Kind != execution.ErrorKindInvalidRequest { + t.Fatalf("result = %+v", result) + } +} diff --git a/internal/execution/bifrost/responses_test.go b/internal/execution/bifrost/responses_test.go new file mode 100644 index 00000000..c593ee96 --- /dev/null +++ b/internal/execution/bifrost/responses_test.go @@ -0,0 +1,298 @@ +package bifrost + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" + "gpt-load/internal/usage" +) + +func TestOfficialOpenAIResponsesCreateNativePassthrough(t *testing.T) { + t.Parallel() + + const responseBody = `{"id": "resp_1", "object":"response", "status":"completed", "model":"served-model", "output":[], "usage":{"input_tokens":10,"input_tokens_details":{"cached_tokens":4},"output_tokens":3,"total_tokens":13}, "vendor_response":{"precise":1.2300}}` + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Method != http.MethodPost || request.URL.Path != "/v1/responses" { + t.Errorf("request target = %s %s", request.Method, request.URL.Path) + } + if request.URL.Query().Get("include") != "output_text" { + t.Errorf("query = %q", request.URL.RawQuery) + } + if request.Header.Get("Authorization") != "Bearer "+testAPIKey { + t.Errorf("authorization = %q", request.Header.Get("Authorization")) + } + if request.Header.Get("X-Test-Header") != "forward-me" { + t.Errorf("business header = %q", request.Header.Get("X-Test-Header")) + } + body, err := io.ReadAll(request.Body) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatal(err) + } + if payload["model"] != "upstream-model" || payload["stream"] != false { + t.Errorf("request controls = %#v", payload) + } + for _, field := range []string{"provider", "fallbacks", "authorization", "api_key"} { + if _, exists := payload[field]; exists { + t.Errorf("control field %q reached upstream", field) + } + } + if !bytes.Contains(body, []byte(`"precise":1.2300`)) { + t.Errorf("unknown request field changed: %s", body) + } + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Request-Id", "responses-create") + _, _ = io.WriteString(writer, responseBody) + })) + defer server.Close() + + runtime := newOfficialOpenAITestRuntime(t, server.URL) + spec := openAIResponsesSpec(execution.OperationResponsesCreate, http.MethodPost, "/v1/responses") + spec.ClientModel = spec.UpstreamModel + spec.Query.Set("include", "output_text") + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.StatusCode != http.StatusOK || result.UpstreamRequestID != "responses-create" || result.Model != "served-model" { + t.Fatalf("result metadata = %+v", result) + } + if string(result.Body) != responseBody { + t.Fatalf("response bytes changed:\n got: %s\nwant: %s", result.Body, responseBody) + } + assertUsage(t, result.Usage, usage.Tokens{UncachedInput: 6, CacheRead: 4, Output: 3}) +} + +func TestOfficialOpenAIResponsesCreateNativeStream(t *testing.T) { + t.Parallel() + + const rawStream = "event: response.output_text.delta\n" + + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\",\"vendor\":{\"precise\":1.2300}}\n\n" + + "event: response.completed\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_stream\",\"model\":\"served-model\",\"usage\":{\"input_tokens\":8,\"input_tokens_details\":{\"cached_tokens\":3},\"output_tokens\":2,\"total_tokens\":10}}}\n\n" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + body, err := io.ReadAll(request.Body) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatal(err) + } + if payload["model"] != "upstream-model" || payload["stream"] != true { + t.Errorf("request controls = %#v", payload) + } + if _, exists := payload["stream_options"]; exists { + t.Errorf("unexpected stream_options = %#v", payload["stream_options"]) + } + writer.Header().Set("Content-Type", "text/event-stream") + writer.Header().Set("X-Request-Id", "responses-stream") + writer.WriteHeader(http.StatusOK) + _, _ = io.WriteString(writer, rawStream) + writer.(http.Flusher).Flush() + })) + defer server.Close() + + runtime := newOfficialOpenAITestRuntime(t, server.URL) + spec := openAIResponsesSpec(execution.OperationResponsesCreate, http.MethodPost, "/v1/responses") + spec.ClientModel = spec.UpstreamModel + var events []execution.StreamEvent + result := runtime.ExecuteStream( + context.Background(), + spec, + func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + return nil + }, + ) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if len(events) < 3 || events[0].Kind != execution.StreamEventReady || events[0].UpstreamRequestID != "responses-stream" { + t.Fatalf("events = %+v", events) + } + var data bytes.Buffer + var sawUsage bool + for _, event := range events[1:] { + if event.Kind == execution.StreamEventData { + data.Write(event.Data) + } + if event.Kind == execution.StreamEventUsage { + sawUsage = true + assertUsage(t, event.Usage, usage.Tokens{UncachedInput: 5, CacheRead: 3, Output: 2}) + } + } + if data.String() != rawStream || !sawUsage { + t.Fatalf("stream data/usage = %q/%t", data.String(), sawUsage) + } + assertUsage(t, result.Usage, usage.Tokens{UncachedInput: 5, CacheRead: 3, Output: 2}) +} + +func TestOfficialOpenAIResponsesLifecyclePreservesTargetAndStatus(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + operation execution.Operation + method string + path string + status int + body string + }{ + {name: "retrieve", operation: execution.OperationResponsesRetrieve, method: http.MethodGet, path: "/v1/responses/resp_123", status: http.StatusOK, body: `{"id":"resp_123","object":"response","status":"completed"}`}, + {name: "delete", operation: execution.OperationResponsesDelete, method: http.MethodDelete, path: "/v1/responses/resp_123", status: http.StatusOK, body: `{"id":"resp_123","object":"response.deleted","deleted":true}`}, + {name: "cancel", operation: execution.OperationResponsesCancel, method: http.MethodPost, path: "/v1/responses/resp_123/cancel", status: http.StatusAccepted, body: `{"id":"resp_123","object":"response","status":"cancelling"}`}, + {name: "input items", operation: execution.OperationResponsesInputItems, method: http.MethodGet, path: "/v1/responses/resp_123/input_items", status: http.StatusPartialContent, body: `{"object":"list","data":[]}`}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Method != test.method || request.URL.Path != test.path { + t.Errorf("request target = %s %s", request.Method, request.URL.Path) + } + if request.URL.Query().Get("after") != "item_1" { + t.Errorf("query = %q", request.URL.RawQuery) + } + if request.Header.Get("Authorization") != "Bearer "+testAPIKey { + t.Errorf("authorization = %q", request.Header.Get("Authorization")) + } + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Request-Id", "lifecycle-"+test.name) + writer.WriteHeader(test.status) + _, _ = io.WriteString(writer, test.body) + })) + defer server.Close() + + runtime := newOfficialOpenAITestRuntime(t, server.URL) + spec := openAIResponsesSpec(test.operation, test.method, test.path) + spec.ClientModel = "" + spec.UpstreamModel = "" + spec.Body = nil + spec.Query.Set("after", "item_1") + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.StatusCode != test.status || string(result.Body) != test.body || result.Error != nil { + t.Fatalf("result = %+v, body=%s", result, result.Body) + } + }) + } +} + +func TestOfficialOpenAIResponsesUtilityOperationsPreserveNativeSemantics(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + operation execution.Operation + path string + body []byte + response string + }{ + { + name: "compact", + operation: execution.OperationResponsesCompact, + path: "/v1/responses/compact", + body: []byte(`{"model":"client-model","input":[{"role":"user","content":"hello"}],"provider":"injected","fallbacks":["other/model"],"vendor_extension":{"precise":1.2300}}`), + response: `{"id":"cmp_123","object":"response.compaction","output":[],"vendor":{"precise":1.2300}}`, + }, + { + name: "input tokens", + operation: execution.OperationResponsesInputTokens, + path: "/v1/responses/input_tokens", + body: []byte(`{"model":"client-model","input":"hello","provider":"injected","api_key":"injected","vendor_extension":{"precise":1.2300}}`), + response: `{"object":"response.input_tokens","input_tokens":7,"vendor":{"precise":1.2300}}`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + if request.Method != http.MethodPost || request.URL.Path != test.path || request.URL.Query().Get("trace") != "kept" { + t.Errorf("target = %s %s?%s", request.Method, request.URL.Path, request.URL.RawQuery) + } + if request.Header.Get("Authorization") != "Bearer "+testAPIKey || request.Header.Get("Api-Key") != "" { + t.Errorf("credential headers = authorization:%q api-key:%q", request.Header.Get("Authorization"), request.Header.Get("Api-Key")) + } + body, _ := io.ReadAll(request.Body) + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("decode request body: %v", err) + } + if payload["model"] != "upstream-model" || payload["provider"] != nil || payload["fallbacks"] != nil || payload["api_key"] != nil { + t.Errorf("sanitized request = %#v", payload) + } + if _, ok := payload["vendor_extension"]; !ok { + t.Errorf("vendor extension was dropped: %#v", payload) + } + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Request-Id", "utility-"+test.name) + _, _ = io.WriteString(writer, test.response) + })) + defer server.Close() + + runtime := newOfficialOpenAITestRuntime(t, server.URL) + spec := openAIResponsesSpec(test.operation, http.MethodPost, test.path) + spec.Body = test.body + spec.Query.Set("trace", "kept") + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if calls.Load() != 1 || result.StatusCode != http.StatusOK || result.Error != nil || string(result.Body) != test.response { + t.Fatalf("calls/result = %d/%+v body=%s", calls.Load(), result, result.Body) + } + }) + } +} + +func newOfficialOpenAITestRuntime(t *testing.T, baseURL string) *Runtime { + t.Helper() + runtime, err := newRuntime(context.Background(), runtimeOptions{ + allowPrivateNetwork: true, + openAIBaseURL: baseURL, + }) + if err != nil { + t.Fatalf("new runtime: %v", err) + } + t.Cleanup(runtime.Shutdown) + return runtime +} + +func openAIResponsesSpec(operation execution.Operation, method, path string) execution.AttemptSpec { + return execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: "responses-request", + AttemptID: "responses-attempt", + Sequence: 1, + ChannelID: string(channel.OpenAI), + ClientProtocol: protocol.OpenAIResponses, + Operation: operation, + ClientModel: "client-model", + UpstreamModel: "upstream-model", + Method: method, + Path: path, + Query: make(map[string][]string), + Header: http.Header{ + "Authorization": []string{"Bearer client-injected"}, + "Api-Key": []string{"client-injected"}, + "X-Test-Header": []string{"forward-me"}, + }, + Body: []byte(`{"model":"client-model","stream":true,"provider":"injected","fallbacks":["other/model"],"authorization":"body-injected","api_key":"body-injected","input":"hello","vendor_extension":{"precise":1.2300}}`), + TargetConfig: json.RawMessage(`{}`), + Credential: execution.NewCredentialSnapshot(7, 3, 2, []byte(`{"api_key":"`+testAPIKey+`"}`)), + }) +} diff --git a/internal/execution/bifrost/runtime.go b/internal/execution/bifrost/runtime.go new file mode 100644 index 00000000..464dda88 --- /dev/null +++ b/internal/execution/bifrost/runtime.go @@ -0,0 +1,302 @@ +// Package bifrost adapts the official Bifrost Core SDK to the neutral execution contract. +package bifrost + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "sync/atomic" + + core "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" +) + +const ( + providerConcurrency = 32 + providerBufferSize = 256 + providerTimeoutSecs = 300 +) + +var errKeyPoolDisabled = errors.New("Bifrost key pool is disabled; DirectKey is required") + +// Runtime owns one long-lived Bifrost Core and executes attempts with explicit credentials. +type Runtime struct { + core *core.Bifrost + account *directAccount + registry *channel.Registry + capabilities execution.CapabilitySet + allowPrivate bool + shutdownOnce sync.Once + closed atomic.Bool +} + +type runtimeOptions struct { + allowPrivateNetwork bool + openAIBaseURL string + anthropicBaseURL string + geminiBaseURL string +} + +// NewRuntime initializes one production runtime. Administrators may configure private-network +// compatible endpoints; Bifrost still rejects link-local and otherwise unsafe destinations. +func NewRuntime(ctx context.Context) (*Runtime, error) { + return newRuntime(ctx, runtimeOptions{allowPrivateNetwork: true}) +} + +func newRuntime(ctx context.Context, options runtimeOptions) (*Runtime, error) { + account := newDirectAccount() + account.setConfig(schemas.OpenAI, providerConfig(options.openAIBaseURL, false, schemas.OpenAI, options.allowPrivateNetwork)) + account.setConfig(schemas.Anthropic, providerConfig(options.anthropicBaseURL, false, schemas.Anthropic, options.allowPrivateNetwork)) + account.setConfig(schemas.Gemini, providerConfig(options.geminiBaseURL, false, schemas.Gemini, options.allowPrivateNetwork)) + account.setConfig(schemas.Azure, providerConfig("", false, schemas.Azure, options.allowPrivateNetwork)) + account.setConfig(schemas.Bedrock, providerConfig("", false, schemas.Bedrock, options.allowPrivateNetwork)) + account.setConfig(schemas.Vertex, providerConfig("", false, schemas.Vertex, options.allowPrivateNetwork)) + + bifrostCore, err := core.Init(ctx, schemas.BifrostConfig{ + Account: account, + LLMPlugins: []schemas.LLMPlugin{}, + MCPPlugins: []schemas.MCPPlugin{}, + Logger: core.NewNoOpLogger(), + InitialPoolSize: 64, + }) + if err != nil { + return nil, fmt.Errorf("initialize execution runtime: %w", err) + } + features, err := execution.NewFeatureSet( + execution.FeatureStreaming, + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + ) + if err != nil { + bifrostCore.Shutdown() + return nil, fmt.Errorf("initialize execution capabilities: %w", err) + } + resourceFeatures, err := execution.NewFeatureSet(execution.FeatureNativeResourceSemantics) + if err != nil { + bifrostCore.Shutdown() + return nil, fmt.Errorf("initialize resource capabilities: %w", err) + } + responsesFeatures, err := execution.NewFeatureSet( + execution.FeatureStreaming, + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + execution.FeatureNativeResourceSemantics, + ) + if err != nil { + bifrostCore.Shutdown() + return nil, fmt.Errorf("initialize responses capabilities: %w", err) + } + passthroughFeatures, err := execution.NewFeatureSet( + execution.FeatureStreaming, + execution.FeatureNativeResourceSemantics, + ) + if err != nil { + bifrostCore.Shutdown() + return nil, fmt.Errorf("initialize passthrough capabilities: %w", err) + } + responsesUtilityFeatures, err := execution.NewFeatureSet( + execution.FeatureTools, + execution.FeatureReasoning, + execution.FeatureMultimodal, + execution.FeatureStructuredOutput, + ) + if err != nil { + bifrostCore.Shutdown() + return nil, fmt.Errorf("initialize Responses utility capabilities: %w", err) + } + capabilities, err := execution.NewCapabilitySet( + execution.Capability{Operation: execution.OperationChatCompletion, Features: features}, + execution.Capability{Operation: execution.OperationResponsesCreate, Features: responsesFeatures}, + execution.Capability{Operation: execution.OperationResponsesRetrieve, Features: resourceFeatures}, + execution.Capability{Operation: execution.OperationResponsesDelete, Features: resourceFeatures}, + execution.Capability{Operation: execution.OperationResponsesCancel, Features: resourceFeatures}, + execution.Capability{Operation: execution.OperationResponsesInputItems, Features: resourceFeatures}, + execution.Capability{Operation: execution.OperationResponsesCompact, Features: responsesUtilityFeatures}, + execution.Capability{Operation: execution.OperationResponsesInputTokens, Features: responsesUtilityFeatures}, + execution.Capability{Operation: execution.OperationResponsesPassthrough, Features: passthroughFeatures}, + execution.Capability{Operation: execution.OperationListModels}, + execution.Capability{Operation: execution.OperationProbe}, + ) + if err != nil { + bifrostCore.Shutdown() + return nil, fmt.Errorf("initialize execution capabilities: %w", err) + } + return &Runtime{ + core: bifrostCore, + account: account, + registry: channel.NewRegistry(), + capabilities: capabilities, + allowPrivate: options.allowPrivateNetwork, + }, nil +} + +// Capabilities returns an independent capability snapshot. +func (r *Runtime) Capabilities() execution.CapabilitySet { + if r == nil { + return execution.CapabilitySet{} + } + return r.capabilities.Clone() +} + +// Shutdown releases the shared SDK runtime. It is safe to call more than once. +func (r *Runtime) Shutdown() { + if r == nil { + return + } + r.shutdownOnce.Do(func() { + r.closed.Store(true) + if r.core != nil { + r.core.Shutdown() + } + }) +} + +func providerConfig(baseURL string, custom bool, baseProvider schemas.ModelProvider, allowPrivateNetwork bool) *schemas.ProviderConfig { + networkBaseURL := baseURL + if custom && baseProvider == schemas.OpenAI { + // OpenAI passthrough always appends /v1 before the request path. Channel + // base_url is a complete API prefix, so remove only that exact suffix. + networkBaseURL = strings.TrimSuffix(baseURL, "/v1") + } + config := &schemas.ProviderConfig{ + NetworkConfig: schemas.NetworkConfig{ + BaseURL: networkBaseURL, + DefaultRequestTimeoutInSeconds: providerTimeoutSecs, + MaxRetries: 0, + StreamIdleTimeoutInSeconds: 120, + AllowPrivateNetwork: allowPrivateNetwork, + }, + ConcurrencyAndBufferSize: schemas.ConcurrencyAndBufferSize{ + Concurrency: providerConcurrency, + BufferSize: providerBufferSize, + }, + SendBackRawRequest: false, + SendBackRawResponse: false, + StoreRawRequestResponse: false, + } + if custom { + customConfig := &schemas.CustomProviderConfig{ + BaseProviderType: baseProvider, + AllowedRequests: &schemas.AllowedRequests{ + ListModels: true, + ChatCompletion: true, + ChatCompletionStream: true, + Passthrough: true, + PassthroughStream: true, + }, + } + if baseProvider == schemas.OpenAI { + // Channel base_url is the complete API prefix (commonly ending in /v1), + // while Bifrost's OpenAI default assumes a provider origin and appends /v1. + customConfig.RequestPathOverrides = map[schemas.RequestType]string{ + schemas.ChatCompletionRequest: baseURL + "/chat/completions", + schemas.ChatCompletionStreamRequest: baseURL + "/chat/completions", + } + } + config.CustomProviderConfig = customConfig + } + return config +} + +type directAccount struct { + mu sync.RWMutex + configs map[schemas.ModelProvider]*schemas.ProviderConfig + keyPoolCalls atomic.Uint64 +} + +func newDirectAccount() *directAccount { + return &directAccount{configs: make(map[schemas.ModelProvider]*schemas.ProviderConfig)} +} + +func (a *directAccount) GetConfiguredProviders() ([]schemas.ModelProvider, error) { + a.mu.RLock() + defer a.mu.RUnlock() + providers := make([]schemas.ModelProvider, 0, len(a.configs)) + for provider := range a.configs { + providers = append(providers, provider) + } + return providers, nil +} + +func (a *directAccount) GetKeysForProvider(context.Context, schemas.ModelProvider) ([]schemas.Key, error) { + a.keyPoolCalls.Add(1) + return nil, errKeyPoolDisabled +} + +func (a *directAccount) GetConfigForProvider(provider schemas.ModelProvider) (*schemas.ProviderConfig, error) { + a.mu.RLock() + defer a.mu.RUnlock() + config, ok := a.configs[provider] + if !ok { + return nil, fmt.Errorf("provider config is not registered") + } + return cloneProviderConfig(config), nil +} + +func (a *directAccount) setConfig(provider schemas.ModelProvider, config *schemas.ProviderConfig) { + a.mu.Lock() + defer a.mu.Unlock() + if _, exists := a.configs[provider]; !exists { + a.configs[provider] = cloneProviderConfig(config) + } +} + +func cloneProviderConfig(source *schemas.ProviderConfig) *schemas.ProviderConfig { + if source == nil { + return nil + } + clone := *source + clone.NetworkConfig.ExtraHeaders = cloneStringMap(source.NetworkConfig.ExtraHeaders) + clone.NetworkConfig.BetaHeaderOverrides = cloneBoolMap(source.NetworkConfig.BetaHeaderOverrides) + if source.CustomProviderConfig != nil { + custom := *source.CustomProviderConfig + if source.CustomProviderConfig.AllowedRequests != nil { + allowed := *source.CustomProviderConfig.AllowedRequests + custom.AllowedRequests = &allowed + } + custom.RequestPathOverrides = make(map[schemas.RequestType]string, len(source.CustomProviderConfig.RequestPathOverrides)) + for key, value := range source.CustomProviderConfig.RequestPathOverrides { + custom.RequestPathOverrides[key] = value + } + clone.CustomProviderConfig = &custom + } + if source.OpenAIConfig != nil { + openAIConfig := *source.OpenAIConfig + clone.OpenAIConfig = &openAIConfig + } + return &clone +} + +func cloneStringMap(source map[string]string) map[string]string { + if source == nil { + return nil + } + clone := make(map[string]string, len(source)) + for key, value := range source { + clone[key] = value + } + return clone +} + +func cloneBoolMap(source map[string]bool) map[string]bool { + if source == nil { + return nil + } + clone := make(map[string]bool, len(source)) + for key, value := range source { + clone[key] = value + } + return clone +} + +var _ execution.Executor = (*Runtime)(nil) +var _ schemas.Account = (*directAccount)(nil) diff --git a/internal/execution/bifrost/utility_operations_test.go b/internal/execution/bifrost/utility_operations_test.go new file mode 100644 index 00000000..05f705e3 --- /dev/null +++ b/internal/execution/bifrost/utility_operations_test.go @@ -0,0 +1,150 @@ +package bifrost + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestOfficialNativeListModelsPreservesWireResponse(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID channel.ID + protocol protocol.Protocol + path string + body string + credential string + runtime func(*testing.T, string) *Runtime + wantUpstream string + }{ + {name: "openai", channelID: channel.OpenAI, protocol: protocol.OpenAICompletions, path: "/v1/models", body: `{"object":"list","data":[{"id":"gpt-test","object":"model"}],"vendor":{"precise":1.2300}}`, credential: "Authorization", wantUpstream: "/v1/models", runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: base}) + }}, + {name: "anthropic", channelID: channel.Anthropic, protocol: protocol.Anthropic, path: "/v1/models", body: `{"data":[{"id":"claude-test","type":"model"}],"has_more":false,"vendor":{"precise":1.2300}}`, credential: "X-Api-Key", wantUpstream: "/v1/models", runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, anthropicBaseURL: base}) + }}, + {name: "gemini", channelID: channel.Gemini, protocol: protocol.Gemini, path: "/v1beta/models", body: `{"models":[{"name":"models/gemini-test"}],"vendor":{"precise":1.2300}}`, credential: "X-Goog-Api-Key", wantUpstream: "/v1beta/models", runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, geminiBaseURL: base + "/v1beta"}) + }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + if request.Method != http.MethodGet || request.URL.Path != test.wantUpstream || request.URL.Query().Get("page") != "next" { + t.Errorf("target = %s %s?%s", request.Method, request.URL.Path, request.URL.RawQuery) + } + wantCredential := testAPIKey + if test.credential == "Authorization" { + wantCredential = "Bearer " + testAPIKey + } + if request.Header.Get(test.credential) != wantCredential { + t.Errorf("credential header %s = %q", test.credential, request.Header.Get(test.credential)) + } + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Request-Id", "models-"+test.name) + _, _ = io.WriteString(writer, test.body) + })) + defer server.Close() + + runtime := test.runtime(t, server.URL) + spec := utilitySpec(test.channelID, test.protocol, execution.OperationListModels, http.MethodGet, test.path, nil) + spec.Query.Set("page", "next") + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if string(result.Body) != test.body || result.UpstreamRequestID != "models-"+test.name { + t.Fatalf("result/body = %+v/%s", result, result.Body) + } + if calls.Load() != 1 { + t.Fatalf("calls = %d", calls.Load()) + } + }) + } +} + +func TestOfficialNativeProbeUsesSelectedModelAndProtocolTarget(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + channelID channel.ID + protocol protocol.Protocol + path string + body []byte + runtime func(*testing.T, string) *Runtime + assert func(*testing.T, *http.Request, map[string]any) + }{ + {name: "openai", channelID: channel.OpenAI, protocol: protocol.OpenAICompletions, path: "/v1/chat/completions", body: []byte(`{"model":"client","messages":[{"role":"user","content":"ping"}]}`), runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, openAIBaseURL: base}) + }, assert: func(t *testing.T, request *http.Request, payload map[string]any) { + if request.URL.Path != "/v1/chat/completions" || payload["model"] != "probe-upstream" || payload["stream"] != false { + t.Errorf("OpenAI probe = %s %#v", request.URL.Path, payload) + } + }}, + {name: "anthropic", channelID: channel.Anthropic, protocol: protocol.Anthropic, path: "/v1/messages", body: []byte(`{"model":"client","max_tokens":1,"messages":[{"role":"user","content":"ping"}]}`), runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, anthropicBaseURL: base}) + }, assert: func(t *testing.T, request *http.Request, payload map[string]any) { + if request.URL.Path != "/v1/messages" || payload["model"] != "probe-upstream" || payload["stream"] != false { + t.Errorf("Anthropic probe = %s %#v", request.URL.Path, payload) + } + }}, + {name: "gemini", channelID: channel.Gemini, protocol: protocol.Gemini, path: "/v1beta/models/client:generateContent", body: []byte(`{"contents":[{"parts":[{"text":"ping"}]}]}`), runtime: func(t *testing.T, base string) *Runtime { + return newProtocolTestRuntime(t, runtimeOptions{allowPrivateNetwork: true, geminiBaseURL: base + "/v1beta"}) + }, assert: func(t *testing.T, request *http.Request, payload map[string]any) { + if request.URL.Path != "/v1beta/models/probe-upstream:generateContent" { + t.Errorf("Gemini probe path = %s", request.URL.Path) + } + }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + body, _ := io.ReadAll(request.Body) + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Errorf("decode probe: %v", err) + } + test.assert(t, request, payload) + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"ok":true}`) + })) + defer server.Close() + + runtime := test.runtime(t, server.URL) + spec := utilitySpec(test.channelID, test.protocol, execution.OperationProbe, http.MethodPost, test.path, test.body) + spec.ClientModel = "probe-client" + spec.UpstreamModel = "probe-upstream" + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if calls.Load() != 1 || result.Error != nil || result.StatusCode != http.StatusOK { + t.Fatalf("calls/result = %d/%+v", calls.Load(), result) + } + }) + } +} + +func utilitySpec(channelID channel.ID, clientProtocol protocol.Protocol, operation execution.Operation, method, path string, body []byte) execution.AttemptSpec { + return execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: "utility-request", AttemptID: "utility-attempt", Sequence: 1, + ChannelID: string(channelID), ClientProtocol: clientProtocol, Operation: operation, + Method: method, Path: path, Query: make(map[string][]string), Header: http.Header{"Authorization": {"Bearer client"}}, Body: body, + TargetConfig: json.RawMessage(`{}`), Credential: execution.NewCredentialSnapshot(10, 1, 1, []byte(`{"api_key":"`+testAPIKey+`"}`)), + }) +} diff --git a/internal/execution/contracts.go b/internal/execution/contracts.go new file mode 100644 index 00000000..d321da14 --- /dev/null +++ b/internal/execution/contracts.go @@ -0,0 +1,562 @@ +// Package execution defines provider-neutral contracts for one selected upstream attempt. +package execution + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "sort" + "time" + + "gpt-load/internal/protocol" + "gpt-load/internal/usage" +) + +// ValidationError identifies one invalid contract field without embedding secret values. +type ValidationError struct { + Field string + Reason string +} + +// Error implements error. +func (e *ValidationError) Error() string { + return fmt.Sprintf("%s: %s", e.Field, e.Reason) +} + +// Operation identifies one logical upstream operation. +type Operation string + +const ( + OperationChatCompletion Operation = "chat_completion" + OperationResponsesCreate Operation = "responses_create" + OperationResponsesRetrieve Operation = "responses_retrieve" + OperationResponsesDelete Operation = "responses_delete" + OperationResponsesCancel Operation = "responses_cancel" + OperationResponsesInputItems Operation = "responses_input_items" + OperationResponsesCompact Operation = "responses_compact" + OperationResponsesInputTokens Operation = "responses_input_tokens" + OperationResponsesPassthrough Operation = "responses_passthrough" + OperationListModels Operation = "list_models" + OperationProbe Operation = "probe" +) + +// Valid reports whether the operation is supported by the execution contract. +func (o Operation) Valid() bool { + switch o { + case OperationChatCompletion, + OperationResponsesCreate, + OperationResponsesRetrieve, + OperationResponsesDelete, + OperationResponsesCancel, + OperationResponsesInputItems, + OperationResponsesCompact, + OperationResponsesInputTokens, + OperationResponsesPassthrough, + OperationListModels, + OperationProbe: + return true + default: + return false + } +} + +// Feature identifies an optional execution capability. +type Feature string + +const ( + FeatureStreaming Feature = "streaming" + FeatureTools Feature = "tools" + FeatureReasoning Feature = "reasoning" + FeatureMultimodal Feature = "multimodal" + FeatureStructuredOutput Feature = "structured_output" + FeatureNativeResourceSemantics Feature = "native_resource_semantics" +) + +// Valid reports whether the feature is recognized. +func (f Feature) Valid() bool { + switch f { + case FeatureStreaming, + FeatureTools, + FeatureReasoning, + FeatureMultimodal, + FeatureStructuredOutput, + FeatureNativeResourceSemantics: + return true + default: + return false + } +} + +// FeatureSet is an immutable set of optional execution capabilities. +type FeatureSet struct { + features map[Feature]struct{} +} + +// NewFeatureSet constructs a feature set. +func NewFeatureSet(features ...Feature) (FeatureSet, error) { + set := FeatureSet{features: make(map[Feature]struct{}, len(features))} + for _, feature := range features { + if !feature.Valid() { + return FeatureSet{}, &ValidationError{Field: "feature", Reason: "unsupported value"} + } + set.features[feature] = struct{}{} + } + return set, nil +} + +// Has reports whether feature is present. +func (s FeatureSet) Has(feature Feature) bool { + _, ok := s.features[feature] + return ok +} + +// Features returns a stable snapshot of the set. +func (s FeatureSet) Features() []Feature { + features := make([]Feature, 0, len(s.features)) + for feature := range s.features { + features = append(features, feature) + } + sort.Slice(features, func(i, j int) bool { + return features[i] < features[j] + }) + return features +} + +// Clone returns an independent feature set. +func (s FeatureSet) Clone() FeatureSet { + clone := FeatureSet{features: make(map[Feature]struct{}, len(s.features))} + for feature := range s.features { + clone.features[feature] = struct{}{} + } + return clone +} + +// Validate validates all features in the set. +func (s FeatureSet) Validate() error { + for feature := range s.features { + if !feature.Valid() { + return &ValidationError{Field: "feature", Reason: "unsupported value"} + } + } + return nil +} + +// Capability binds an operation to the optional features supported for it. +type Capability struct { + Operation Operation + Features FeatureSet +} + +// CapabilitySet is an immutable set of supported operations and features. +type CapabilitySet struct { + capabilities map[Operation]FeatureSet +} + +// NewCapabilitySet constructs a capability set. +func NewCapabilitySet(capabilities ...Capability) (CapabilitySet, error) { + set := CapabilitySet{capabilities: make(map[Operation]FeatureSet, len(capabilities))} + for _, capability := range capabilities { + if !capability.Operation.Valid() { + return CapabilitySet{}, &ValidationError{Field: "operation", Reason: "unsupported value"} + } + if err := capability.Features.Validate(); err != nil { + return CapabilitySet{}, err + } + features, ok := set.capabilities[capability.Operation] + if !ok { + features = FeatureSet{features: make(map[Feature]struct{})} + } + for feature := range capability.Features.features { + features.features[feature] = struct{}{} + } + set.capabilities[capability.Operation] = features + } + return set, nil +} + +// Has reports whether operation is supported. +func (s CapabilitySet) Has(operation Operation) bool { + _, ok := s.capabilities[operation] + return ok +} + +// Supports reports whether operation and all required features are supported. +func (s CapabilitySet) Supports(operation Operation, required FeatureSet) bool { + if !operation.Valid() || required.Validate() != nil { + return false + } + available, ok := s.capabilities[operation] + if !ok { + return false + } + for feature := range required.features { + if !available.Has(feature) { + return false + } + } + return true +} + +// Features returns an independent set of features supported for operation. +func (s CapabilitySet) Features(operation Operation) FeatureSet { + features, ok := s.capabilities[operation] + if !ok { + return FeatureSet{features: make(map[Feature]struct{})} + } + return features.Clone() +} + +// Operations returns a stable snapshot of supported operations. +func (s CapabilitySet) Operations() []Operation { + operations := make([]Operation, 0, len(s.capabilities)) + for operation := range s.capabilities { + operations = append(operations, operation) + } + sort.Slice(operations, func(i, j int) bool { + return operations[i] < operations[j] + }) + return operations +} + +// Clone returns an independent capability set. +func (s CapabilitySet) Clone() CapabilitySet { + clone := CapabilitySet{capabilities: make(map[Operation]FeatureSet, len(s.capabilities))} + for operation, features := range s.capabilities { + clone.capabilities[operation] = features.Clone() + } + return clone +} + +// Validate validates all operations and features in the set. +func (s CapabilitySet) Validate() error { + for operation, features := range s.capabilities { + if !operation.Valid() { + return &ValidationError{Field: "operation", Reason: "unsupported value"} + } + if err := features.Validate(); err != nil { + return err + } + } + return nil +} + +// CredentialSnapshot is the exact logical credential selected for an attempt. +// Secret data is private and is never included in JSON. +type CredentialSnapshot struct { + ID uint `json:"id"` + Version uint64 `json:"version"` + IdentityGeneration uint64 `json:"identity_generation"` + + data []byte +} + +// NewCredentialSnapshot copies the credential data into a snapshot. +func NewCredentialSnapshot(id uint, version, identityGeneration uint64, data []byte) CredentialSnapshot { + return CredentialSnapshot{ + ID: id, + Version: version, + IdentityGeneration: identityGeneration, + data: cloneBytes(data), + } +} + +// Data returns an independent copy of the secret credential data. +func (c CredentialSnapshot) Data() []byte { + return cloneBytes(c.data) +} + +// Clone returns an independent credential snapshot. +func (c CredentialSnapshot) Clone() CredentialSnapshot { + clone := c + clone.data = cloneBytes(c.data) + return clone +} + +// MarshalJSON exposes only non-secret credential identity metadata. +func (c CredentialSnapshot) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + ID uint `json:"id"` + Version uint64 `json:"version"` + IdentityGeneration uint64 `json:"identity_generation"` + }{ + ID: c.ID, + Version: c.Version, + IdentityGeneration: c.IdentityGeneration, + }) +} + +// AttemptTimeouts contains per-attempt timeout overrides. Zero inherits the caller default. +type AttemptTimeouts struct { + Connect time.Duration `json:"connect,omitempty"` + FirstByte time.Duration `json:"first_byte,omitempty"` + Request time.Duration `json:"request,omitempty"` + StreamIdle time.Duration `json:"stream_idle,omitempty"` +} + +// AttemptSpec is a fully selected, provider-neutral upstream attempt. +// NewAttemptSpec or Clone must be used at ownership boundaries because Query, +// Header, Body, TargetConfig, and Credential contain reference-backed values. +type AttemptSpec struct { + RequestID string `json:"request_id"` + AttemptID string `json:"attempt_id"` + Sequence uint32 `json:"sequence"` + ChannelID string `json:"channel_id"` + ClientProtocol protocol.Protocol `json:"client_protocol"` + Operation Operation `json:"operation"` + ClientModel string `json:"client_model,omitempty"` + UpstreamModel string `json:"upstream_model,omitempty"` + Method string `json:"method"` + Path string `json:"path"` + Query url.Values `json:"query,omitempty"` + // RawQuery preserves the original query bytes when exact forwarding matters. + // It is mutually exclusive with Query and intentionally is not URL-decoded. + RawQuery string `json:"raw_query,omitempty"` + Header http.Header `json:"header,omitempty"` + Body []byte `json:"body,omitempty"` + // IncludeUsage asks the executor to request provider usage details when the + // selected operation supports an explicit wire option. + IncludeUsage bool `json:"include_usage,omitempty"` + // TargetConfig is non-secret configuration resolved by the channel registry. + TargetConfig json.RawMessage `json:"target_config,omitempty"` + Timeouts AttemptTimeouts `json:"timeouts"` + Credential CredentialSnapshot `json:"credential"` +} + +// NewAttemptSpec takes ownership of an independent clone of spec. +func NewAttemptSpec(spec AttemptSpec) AttemptSpec { + return spec.Clone() +} + +// Clone returns an independent attempt specification. +func (s AttemptSpec) Clone() AttemptSpec { + clone := s + clone.Query = cloneValues(s.Query) + clone.Header = cloneHeader(s.Header) + clone.Body = cloneBytes(s.Body) + clone.TargetConfig = cloneRawMessage(s.TargetConfig) + clone.Credential = s.Credential.Clone() + return clone +} + +// DispatchState describes whether an attempt may have reached the upstream. +type DispatchState string + +const ( + DispatchNotSent DispatchState = "not_sent" + DispatchMaybeSent DispatchState = "maybe_sent" +) + +// Valid reports whether the dispatch state is recognized. +func (s DispatchState) Valid() bool { + return s == DispatchNotSent || s == DispatchMaybeSent +} + +// MaxErrorSummaryLength bounds sanitized error summaries retained by the contract. +const MaxErrorSummaryLength = 512 + +// ErrorKind is the stable category used for retry and health decisions. +type ErrorKind string + +const ( + ErrorKindTransport ErrorKind = "transport" + ErrorKindTimeout ErrorKind = "timeout" + ErrorKindCanceled ErrorKind = "canceled" + ErrorKindHTTP ErrorKind = "http" + ErrorKindProvider ErrorKind = "provider" + ErrorKindInvalidRequest ErrorKind = "invalid_request" + ErrorKindInternal ErrorKind = "internal" +) + +// Valid reports whether the error kind is recognized. +func (k ErrorKind) Valid() bool { + switch k { + case ErrorKindTransport, + ErrorKindTimeout, + ErrorKindCanceled, + ErrorKindHTTP, + ErrorKindProvider, + ErrorKindInvalidRequest, + ErrorKindInternal: + return true + default: + return false + } +} + +// ErrorEvidence contains provider-neutral evidence for an unsuccessful attempt. +// Summary must be bounded and sanitized; raw upstream error messages are not retained. +type ErrorEvidence struct { + Kind ErrorKind `json:"kind"` + StatusCode int `json:"status_code,omitempty"` + Type string `json:"type,omitempty"` + Code string `json:"code,omitempty"` + Summary string `json:"summary,omitempty"` + RequestID string `json:"request_id,omitempty"` + RetryAfter time.Duration `json:"retry_after,omitempty"` + Header http.Header `json:"-"` +} + +// Clone returns an independent error evidence value. +func (e ErrorEvidence) Clone() ErrorEvidence { + clone := e + clone.Header = cloneHeader(e.Header) + return clone +} + +// UsageEvidence contains normalized usage and the raw provider evidence used to derive it. +type UsageEvidence struct { + Normalized usage.Result `json:"normalized"` + Raw []byte `json:"raw,omitempty"` +} + +// Clone returns an independent usage evidence value. +func (e UsageEvidence) Clone() UsageEvidence { + clone := e + clone.Raw = cloneBytes(e.Raw) + return clone +} + +// AttemptResult is the terminal result of a non-streaming attempt. +// The result owns Header, Body, Usage, and Error after return from Executor. +type AttemptResult struct { + DispatchState DispatchState `json:"dispatch_state"` + ResponseStarted bool `json:"response_started"` + StatusCode int `json:"status_code,omitempty"` + Header http.Header `json:"header,omitempty"` + Body []byte `json:"body,omitempty"` + Model string `json:"model,omitempty"` + UpstreamRequestID string `json:"upstream_request_id,omitempty"` + Usage *UsageEvidence `json:"usage,omitempty"` + Error *ErrorEvidence `json:"error,omitempty"` +} + +// Clone returns an independent attempt result. +func (r AttemptResult) Clone() AttemptResult { + clone := r + clone.Header = cloneHeader(r.Header) + clone.Body = cloneBytes(r.Body) + if r.Usage != nil { + usage := r.Usage.Clone() + clone.Usage = &usage + } + if r.Error != nil { + evidence := r.Error.Clone() + clone.Error = &evidence + } + return clone +} + +// StreamEventKind identifies one streaming event category. +type StreamEventKind string + +const ( + StreamEventReady StreamEventKind = "ready" + StreamEventData StreamEventKind = "data" + StreamEventUsage StreamEventKind = "usage" +) + +// Valid reports whether the streaming event kind is recognized. +func (k StreamEventKind) Valid() bool { + return k == StreamEventReady || k == StreamEventData || k == StreamEventUsage +} + +// StreamEvent is one ordered item emitted by a streaming attempt. +// Ownership of Header, Data, and Usage transfers to the stream sink. Executors +// must not mutate or reuse reference-backed event values after the sink call. +type StreamEvent struct { + Sequence uint64 `json:"sequence"` + Kind StreamEventKind `json:"kind"` + StatusCode int `json:"status_code,omitempty"` + Header http.Header `json:"header,omitempty"` + UpstreamRequestID string `json:"upstream_request_id,omitempty"` + Data []byte `json:"data,omitempty"` + Usage *UsageEvidence `json:"usage,omitempty"` +} + +// Clone returns an independent stream event. +func (e StreamEvent) Clone() StreamEvent { + clone := e + clone.Header = cloneHeader(e.Header) + clone.Data = cloneBytes(e.Data) + if e.Usage != nil { + usage := e.Usage.Clone() + clone.Usage = &usage + } + return clone +} + +// StreamResult is the terminal metadata returned after streaming ends. +// The result owns Header, Usage, and Error after return from Executor. +type StreamResult struct { + DispatchState DispatchState `json:"dispatch_state"` + ResponseStarted bool `json:"response_started"` + StatusCode int `json:"status_code,omitempty"` + Header http.Header `json:"header,omitempty"` + Model string `json:"model,omitempty"` + UpstreamRequestID string `json:"upstream_request_id,omitempty"` + Usage *UsageEvidence `json:"usage,omitempty"` + Error *ErrorEvidence `json:"error,omitempty"` +} + +// Clone returns an independent streaming result. +func (r StreamResult) Clone() StreamResult { + clone := r + clone.Header = cloneHeader(r.Header) + if r.Usage != nil { + usage := r.Usage.Clone() + clone.Usage = &usage + } + if r.Error != nil { + evidence := r.Error.Clone() + clone.Error = &evidence + } + return clone +} + +// StreamSink consumes events synchronously. Returning an error stops the stream. +type StreamSink func(StreamEvent) error + +// Executor executes exactly the selected attempt described by AttemptSpec. +// Implementations must return an independent capability snapshot and must not +// mutate spec or retain reference-backed values from it. +type Executor interface { + Capabilities() CapabilitySet + Execute(context.Context, AttemptSpec) AttemptResult + ExecuteStream(context.Context, AttemptSpec, StreamSink) StreamResult +} + +var _ json.Marshaler = CredentialSnapshot{} + +func cloneBytes(value []byte) []byte { + if value == nil { + return nil + } + return append([]byte(nil), value...) +} + +func cloneRawMessage(value json.RawMessage) json.RawMessage { + if value == nil { + return nil + } + return json.RawMessage(cloneBytes(value)) +} + +func cloneHeader(value http.Header) http.Header { + if value == nil { + return nil + } + return value.Clone() +} + +func cloneValues(value url.Values) url.Values { + if value == nil { + return nil + } + clone := make(url.Values, len(value)) + for key, values := range value { + clone[key] = append([]string(nil), values...) + } + return clone +} diff --git a/internal/execution/contracts_test.go b/internal/execution/contracts_test.go new file mode 100644 index 00000000..7b159088 --- /dev/null +++ b/internal/execution/contracts_test.go @@ -0,0 +1,548 @@ +package execution + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/url" + "strings" + "testing" + "time" + + "gpt-load/internal/protocol" + "gpt-load/internal/usage" +) + +func TestOperationAndDispatchEnums(t *testing.T) { + operations := []Operation{ + OperationChatCompletion, + OperationResponsesCreate, + OperationResponsesRetrieve, + OperationResponsesDelete, + OperationResponsesCancel, + OperationResponsesInputItems, + OperationResponsesCompact, + OperationResponsesInputTokens, + OperationResponsesPassthrough, + OperationListModels, + OperationProbe, + } + for _, operation := range operations { + if !operation.Valid() { + t.Fatalf("expected operation %q to be valid", operation) + } + } + if Operation("unknown").Valid() { + t.Fatal("expected unknown operation to be invalid") + } + for _, feature := range []Feature{ + FeatureStreaming, + FeatureTools, + FeatureReasoning, + FeatureMultimodal, + FeatureStructuredOutput, + FeatureNativeResourceSemantics, + } { + if !feature.Valid() { + t.Fatalf("expected feature %q to be valid", feature) + } + } + if Feature("unknown").Valid() { + t.Fatal("expected unknown feature to be invalid") + } + + for _, state := range []DispatchState{DispatchNotSent, DispatchMaybeSent} { + if !state.Valid() { + t.Fatalf("expected dispatch state %q to be valid", state) + } + } + if DispatchState("sent").Valid() { + t.Fatal("expected unknown dispatch state to be invalid") + } + for _, kind := range []ErrorKind{ + ErrorKindTransport, + ErrorKindTimeout, + ErrorKindCanceled, + ErrorKindHTTP, + ErrorKindProvider, + ErrorKindInvalidRequest, + ErrorKindInternal, + } { + if !kind.Valid() { + t.Fatalf("expected error kind %q to be valid", kind) + } + } + if ErrorKind("unknown").Valid() { + t.Fatal("expected unknown error kind to be invalid") + } +} + +func TestCapabilitySetIsValidatedAndDefensivelyCopied(t *testing.T) { + chatFeatures, err := NewFeatureSet(FeatureStreaming, FeatureTools, FeatureStreaming) + if err != nil { + t.Fatalf("NewFeatureSet() error = %v", err) + } + set, err := NewCapabilitySet( + Capability{Operation: OperationResponsesCreate}, + Capability{Operation: OperationChatCompletion, Features: chatFeatures}, + Capability{Operation: OperationResponsesCreate}, + ) + if err != nil { + t.Fatalf("NewCapabilitySet() error = %v", err) + } + if err := set.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + if !set.Has(OperationChatCompletion) || !set.Has(OperationResponsesCreate) { + t.Fatal("expected configured operations to be present") + } + if set.Has(OperationProbe) { + t.Fatal("expected an unconfigured operation to be absent") + } + requiredStreaming, err := NewFeatureSet(FeatureStreaming) + if err != nil { + t.Fatalf("NewFeatureSet() error = %v", err) + } + requiredReasoning, err := NewFeatureSet(FeatureReasoning) + if err != nil { + t.Fatalf("NewFeatureSet() error = %v", err) + } + if !set.Supports(OperationChatCompletion, requiredStreaming) { + t.Fatal("expected chat completion streaming to be supported") + } + if set.Supports(OperationChatCompletion, requiredReasoning) { + t.Fatal("expected unsupported reasoning feature to be rejected") + } + if !set.Features(OperationChatCompletion).Has(FeatureTools) { + t.Fatal("expected operation features to be returned") + } + + want := []Operation{OperationChatCompletion, OperationResponsesCreate} + got := set.Operations() + if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] { + t.Fatalf("Operations() = %v, want %v", got, want) + } + got[0] = OperationProbe + if !set.Has(OperationChatCompletion) || set.Has(OperationProbe) { + t.Fatal("mutating Operations() result changed the set") + } + + clone := set.Clone() + clone.capabilities[OperationProbe] = FeatureSet{} + clone.capabilities[OperationChatCompletion].features[FeatureReasoning] = struct{}{} + if set.Has(OperationProbe) { + t.Fatal("mutating clone changed the original set") + } + if set.Features(OperationChatCompletion).Has(FeatureReasoning) { + t.Fatal("mutating cloned features changed the original set") + } + + if _, err := NewCapabilitySet(Capability{Operation: Operation("unknown")}); err == nil { + t.Fatal("expected an invalid operation to be rejected") + } + if _, err := NewFeatureSet(Feature("unknown")); err == nil { + t.Fatal("expected an invalid feature to be rejected") + } +} + +func TestCredentialSnapshotProtectsSecretDataAndJSON(t *testing.T) { + secret := []byte(`{"api_key":"sk-secret-value"}`) + credential := NewCredentialSnapshot(41, 7, 3, secret) + secret[0] = 'X' + if got := string(credential.Data()); got != `{"api_key":"sk-secret-value"}` { + t.Fatalf("Data() = %q after input mutation", got) + } + + copyOfData := credential.Data() + copyOfData[0] = 'Y' + if got := string(credential.Data()); got != `{"api_key":"sk-secret-value"}` { + t.Fatalf("Data() = %q after returned copy mutation", got) + } + + clone := credential.Clone() + clone.data[0] = 'Z' + if got := string(credential.Data()); got != `{"api_key":"sk-secret-value"}` { + t.Fatalf("Data() = %q after clone mutation", got) + } + + encoded, err := json.Marshal(credential) + if err != nil { + t.Fatalf("json.Marshal() error = %v", err) + } + if strings.Contains(string(encoded), "sk-secret-value") || strings.Contains(string(encoded), "api_key") { + t.Fatalf("credential JSON leaked secret data: %s", encoded) + } + if !strings.Contains(string(encoded), `"id":41`) || !strings.Contains(string(encoded), `"version":7`) || !strings.Contains(string(encoded), `"identity_generation":3`) { + t.Fatalf("credential JSON omitted identity metadata: %s", encoded) + } +} + +func TestAttemptSpecOwnsReferenceBackedValues(t *testing.T) { + credentialData := []byte(`{"api_key":"sk-secret-value"}`) + raw := validAttemptSpec(credentialData) + owned := NewAttemptSpec(raw) + + raw.Header.Set("X-Test", "mutated") + raw.Query.Set("api-version", "mutated") + raw.Body[0] = 'X' + raw.TargetConfig[0] = 'Y' + credentialData[0] = 'Y' + if got := owned.Header.Get("X-Test"); got != "original" { + t.Fatalf("owned header = %q, want original", got) + } + if got := owned.Query.Get("api-version"); got != "2026-01-01" { + t.Fatalf("owned query = %q, want original", got) + } + if got := string(owned.Body); got != `{"model":"client-model"}` { + t.Fatalf("owned body = %q", got) + } + if got := string(owned.TargetConfig); got != `{"base_url":"https://upstream.example"}` { + t.Fatalf("owned target config = %q", got) + } + if got := string(owned.Credential.Data()); got != `{"api_key":"sk-secret-value"}` { + t.Fatalf("owned credential = %q", got) + } + + clone := owned.Clone() + clone.Header.Set("X-Test", "clone") + clone.Query.Set("api-version", "clone") + clone.Body[0] = 'Z' + clone.TargetConfig[0] = 'Q' + clone.Credential.data[0] = 'Q' + if owned.Header.Get("X-Test") != "original" || owned.Query.Get("api-version") != "2026-01-01" { + t.Fatal("mutating clone changed original maps") + } + if string(owned.Body) != `{"model":"client-model"}` || string(owned.Credential.Data()) != `{"api_key":"sk-secret-value"}` { + t.Fatal("mutating clone changed original byte slices") + } + if string(owned.TargetConfig) != `{"base_url":"https://upstream.example"}` { + t.Fatal("mutating clone changed original raw messages") + } + + encoded, err := json.Marshal(owned) + if err != nil { + t.Fatalf("json.Marshal() error = %v", err) + } + if strings.Contains(string(encoded), "sk-secret-value") || strings.Contains(string(encoded), "api_key") { + t.Fatalf("attempt JSON leaked credential data: %s", encoded) + } +} + +func TestAttemptSpecRawQueryPreservesOpaqueBytesAndRejectsUnsafeShapes(t *testing.T) { + t.Parallel() + + spec := validAttemptSpec([]byte(`{"api_key":"sk-secret-value"}`)) + spec.Query = nil + spec.RawQuery = "first=%2F&broken=%ZZ&first=%2f&&empty=" + owned := NewAttemptSpec(spec) + if owned.RawQuery != spec.RawQuery || owned.Clone().RawQuery != spec.RawQuery { + t.Fatalf("RawQuery was changed: owned=%q clone=%q", owned.RawQuery, owned.Clone().RawQuery) + } + if err := owned.Validate(); err != nil { + t.Fatalf("opaque RawQuery Validate() error = %v", err) + } + + for _, test := range []struct { + name string + rawQuery string + }{ + {name: "fragment", rawQuery: "trace=kept#fragment"}, + {name: "newline", rawQuery: "trace=kept\nInjected=true"}, + {name: "nul", rawQuery: "trace=kept\x00"}, + } { + t.Run(test.name, func(t *testing.T) { + invalid := spec.Clone() + invalid.RawQuery = test.rawQuery + if err := invalid.Validate(); err == nil { + t.Fatalf("RawQuery %q was accepted", test.rawQuery) + } else if validationErr, ok := err.(*ValidationError); !ok || validationErr.Field != "raw_query" { + t.Fatalf("Validate() error = %v, want raw_query validation error", err) + } + }) + } + + conflict := spec.Clone() + conflict.Query = url.Values{"trace": {"kept"}} + if err := conflict.Validate(); err == nil { + t.Fatal("RawQuery and Query were accepted together") + } else if validationErr, ok := err.(*ValidationError); !ok || validationErr.Field != "raw_query" { + t.Fatalf("Validate() error = %v, want raw_query validation error", err) + } +} + +func TestAttemptAndStreamResultsAreDefensivelyCopied(t *testing.T) { + result := AttemptResult{ + DispatchState: DispatchMaybeSent, + StatusCode: http.StatusOK, + Header: http.Header{"X-Request-Id": {"request-1"}}, + Body: []byte("response"), + Usage: &UsageEvidence{ + Normalized: usage.Result{State: usage.StateComplete}, + Raw: []byte(`{"input_tokens":1}`), + }, + Error: &ErrorEvidence{ + Kind: ErrorKindProvider, + StatusCode: http.StatusTooManyRequests, + Type: "rate_limit", + Summary: "request rate limited", + Header: http.Header{"Retry-After": {"1"}}, + }, + } + clone := result.Clone() + clone.Header.Set("X-Request-ID", "changed") + clone.Body[0] = 'X' + clone.Usage.Raw[0] = 'Y' + clone.Error.Header.Set("Retry-After", "2") + if result.Header.Get("X-Request-ID") != "request-1" || string(result.Body) != "response" { + t.Fatal("mutating attempt result clone changed original response data") + } + if string(result.Usage.Raw) != `{"input_tokens":1}` || result.Error.Header.Get("Retry-After") != "1" { + t.Fatal("mutating attempt result clone changed original evidence") + } + + event := StreamEvent{ + Sequence: 1, + Kind: StreamEventUsage, + Data: []byte("event"), + Usage: &UsageEvidence{Raw: []byte("usage")}, + } + eventClone := event.Clone() + eventClone.Data[0] = 'X' + eventClone.Usage.Raw[0] = 'Y' + if string(event.Data) != "event" || string(event.Usage.Raw) != "usage" { + t.Fatal("mutating stream event clone changed original") + } + + streamResult := StreamResult{ + DispatchState: DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: http.Header{"X-Request-Id": {"request-1"}}, + UpstreamRequestID: "upstream-1", + Usage: &UsageEvidence{Raw: []byte("usage")}, + Error: &ErrorEvidence{Kind: ErrorKindProvider, Type: "stream", Summary: "stream ended", Header: http.Header{"Retry-After": {"1"}}}, + } + streamClone := streamResult.Clone() + streamClone.Header.Set("X-Request-ID", "changed") + streamClone.Usage.Raw[0] = 'X' + streamClone.Error.Header.Set("Retry-After", "2") + if streamResult.Header.Get("X-Request-ID") != "request-1" || string(streamResult.Usage.Raw) != "usage" || streamResult.Error.Header.Get("Retry-After") != "1" { + t.Fatal("mutating stream result clone changed original") + } +} + +func TestValidationAcceptsValidContractsAndRejectsInvalidFields(t *testing.T) { + spec := validAttemptSpec([]byte(`{"api_key":"sk-secret-value"}`)) + if err := spec.Validate(); err != nil { + t.Fatalf("valid AttemptSpec.Validate() error = %v", err) + } + + tests := []struct { + name string + mutate func(*AttemptSpec) + field string + }{ + {name: "request id", mutate: func(s *AttemptSpec) { s.RequestID = "" }, field: "request_id"}, + {name: "attempt id", mutate: func(s *AttemptSpec) { s.AttemptID = "" }, field: "attempt_id"}, + {name: "sequence", mutate: func(s *AttemptSpec) { s.Sequence = 0 }, field: "sequence"}, + {name: "channel", mutate: func(s *AttemptSpec) { s.ChannelID = "" }, field: "channel_id"}, + {name: "protocol", mutate: func(s *AttemptSpec) { s.ClientProtocol = protocol.Protocol("unknown") }, field: "client_protocol"}, + {name: "operation", mutate: func(s *AttemptSpec) { s.Operation = Operation("unknown") }, field: "operation"}, + {name: "model", mutate: func(s *AttemptSpec) { s.UpstreamModel = "" }, field: "upstream_model"}, + {name: "method", mutate: func(s *AttemptSpec) { s.Method = "" }, field: "method"}, + {name: "path", mutate: func(s *AttemptSpec) { s.Path = "https://upstream.example/v1/chat/completions" }, field: "path"}, + {name: "header", mutate: func(s *AttemptSpec) { s.Header.Set("X-Test", "bad\nvalue") }, field: "header"}, + {name: "connect timeout", mutate: func(s *AttemptSpec) { s.Timeouts.Connect = -time.Second }, field: "timeouts.connect"}, + {name: "first byte timeout", mutate: func(s *AttemptSpec) { s.Timeouts.FirstByte = -time.Second }, field: "timeouts.first_byte"}, + {name: "timeout", mutate: func(s *AttemptSpec) { s.Timeouts.Request = -time.Second }, field: "timeouts.request"}, + {name: "stream idle timeout", mutate: func(s *AttemptSpec) { s.Timeouts.StreamIdle = -time.Second }, field: "timeouts.stream_idle"}, + {name: "target config", mutate: func(s *AttemptSpec) { s.TargetConfig = json.RawMessage("not-json") }, field: "target_config"}, + {name: "credential", mutate: func(s *AttemptSpec) { s.Credential = NewCredentialSnapshot(0, 7, 3, []byte("secret")) }, field: "credential.id"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + invalid := spec.Clone() + tt.mutate(&invalid) + err := invalid.Validate() + if err == nil { + t.Fatalf("Validate() error = nil, want field %q", tt.field) + } + var validationErr *ValidationError + if !errors.As(err, &validationErr) { + t.Fatalf("Validate() error type = %T, want *ValidationError", err) + } + if validationErr.Field != tt.field { + t.Fatalf("ValidationError.Field = %q, want %q", validationErr.Field, tt.field) + } + if strings.Contains(err.Error(), "sk-secret-value") { + t.Fatalf("validation error leaked credential data: %v", err) + } + }) + } + + for _, operation := range []Operation{ + OperationResponsesRetrieve, + OperationResponsesDelete, + OperationResponsesCancel, + OperationResponsesInputItems, + OperationResponsesPassthrough, + OperationListModels, + OperationProbe, + } { + modelOptional := spec.Clone() + modelOptional.Operation = operation + modelOptional.ClientModel = "" + modelOptional.UpstreamModel = "" + modelOptional.Method = http.MethodGet + modelOptional.Path = "/v1/resources" + if err := modelOptional.Validate(); err != nil { + t.Fatalf("model-optional operation %q Validate() error = %v", operation, err) + } + } + responsesCreate := spec.Clone() + responsesCreate.Operation = OperationResponsesCreate + responsesCreate.UpstreamModel = "" + if err := responsesCreate.Validate(); err == nil { + t.Fatal("expected responses create without a model to be rejected") + } + for _, operation := range []Operation{OperationResponsesCompact, OperationResponsesInputTokens} { + modelRequired := spec.Clone() + modelRequired.Operation = operation + modelRequired.UpstreamModel = "" + if err := modelRequired.Validate(); err == nil { + t.Fatalf("expected %q without a model to be rejected", operation) + } + } + + if err := (AttemptResult{DispatchState: DispatchMaybeSent, ResponseStarted: true, StatusCode: http.StatusOK}).Validate(); err != nil { + t.Fatalf("valid AttemptResult.Validate() error = %v", err) + } + if err := (StreamResult{DispatchState: DispatchNotSent, Error: &ErrorEvidence{Kind: ErrorKindTransport, Type: "transport", Summary: "dial failed"}}).Validate(); err != nil { + t.Fatalf("valid StreamResult.Validate() error = %v", err) + } + if err := (AttemptResult{DispatchState: DispatchState("unknown")}).Validate(); err == nil { + t.Fatal("expected invalid result dispatch state to be rejected") + } + invalidResults := []AttemptResult{ + {DispatchState: DispatchMaybeSent, StatusCode: http.StatusOK}, + {DispatchState: DispatchMaybeSent, ResponseStarted: true}, + {DispatchState: DispatchNotSent, ResponseStarted: true, StatusCode: http.StatusOK}, + {DispatchState: DispatchNotSent, Body: []byte("body"), Error: &ErrorEvidence{Kind: ErrorKindTransport, Summary: "not sent"}}, + {DispatchState: DispatchMaybeSent, ResponseStarted: true, StatusCode: http.StatusInternalServerError}, + } + for i, result := range invalidResults { + if err := result.Validate(); err == nil { + t.Fatalf("invalid result %d was accepted: %+v", i, result) + } + } + failedResult := AttemptResult{ + DispatchState: DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusTooManyRequests, + Error: &ErrorEvidence{Kind: ErrorKindHTTP, StatusCode: http.StatusTooManyRequests, Summary: "rate limited"}, + } + if err := failedResult.Validate(); err != nil { + t.Fatalf("failed result with evidence Validate() error = %v", err) + } + if err := (ErrorEvidence{Kind: ErrorKindHTTP, RetryAfter: -time.Second, Summary: "retry"}).Validate(); err == nil { + t.Fatal("expected negative retry-after to be rejected") + } + if err := (ErrorEvidence{Kind: ErrorKindProvider, Summary: strings.Repeat("x", MaxErrorSummaryLength+1)}).Validate(); err == nil { + t.Fatal("expected overlong error summary to be rejected") + } + if err := (ErrorEvidence{Kind: ErrorKindHTTP, Summary: "safe", Header: http.Header{"Authorization": {"Bearer secret"}}}).Validate(); err == nil { + t.Fatal("expected unsafe error evidence header to be rejected") + } + safeEvidence := ErrorEvidence{ + Kind: ErrorKindHTTP, + Summary: "safe summary", + Header: http.Header{"Retry-After": {"1"}}, + } + encodedEvidence, err := json.Marshal(safeEvidence) + if err != nil { + t.Fatalf("json.Marshal(error evidence) error = %v", err) + } + if strings.Contains(string(encodedEvidence), "Retry-After") { + t.Fatalf("error evidence JSON exposed headers: %s", encodedEvidence) + } +} + +func TestStreamEventValidation(t *testing.T) { + ready := StreamEvent{ + Sequence: 1, + Kind: StreamEventReady, + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + UpstreamRequestID: "upstream-1", + } + if err := ready.Validate(); err != nil { + t.Fatalf("valid ready event Validate() error = %v", err) + } + readyClone := ready.Clone() + readyClone.Header.Set("Content-Type", "changed") + if ready.Header.Get("Content-Type") != "text/event-stream" { + t.Fatal("mutating ready event clone changed original header") + } + if err := (StreamEvent{Sequence: 2, Kind: StreamEventData, Data: []byte("data")}).Validate(); err != nil { + t.Fatalf("valid data event Validate() error = %v", err) + } + if err := (StreamEvent{Sequence: 3, Kind: StreamEventUsage, Usage: &UsageEvidence{Normalized: usage.Result{State: usage.StateComplete}}}).Validate(); err != nil { + t.Fatalf("valid usage event Validate() error = %v", err) + } + if err := (StreamEvent{Sequence: 0, Kind: StreamEventData, Data: []byte("data")}).Validate(); err == nil { + t.Fatal("expected zero sequence to be rejected") + } + if err := (StreamEvent{Sequence: 1, Kind: StreamEventKind("unknown")}).Validate(); err == nil { + t.Fatal("expected unknown event kind to be rejected") + } + if err := (StreamEvent{Sequence: 2, Kind: StreamEventData, StatusCode: http.StatusOK, Data: []byte("data")}).Validate(); err == nil { + t.Fatal("expected data event with ready metadata to be rejected") + } + if err := (StreamEvent{Sequence: 1, Kind: StreamEventReady, StatusCode: http.StatusOK, Header: http.Header{}, Data: []byte("data")}).Validate(); err == nil { + t.Fatal("expected ready event with data to be rejected") + } +} + +func validAttemptSpec(credentialData []byte) AttemptSpec { + return AttemptSpec{ + RequestID: "request-1", + AttemptID: "attempt-1", + Sequence: 1, + ChannelID: "openai", + ClientProtocol: protocol.OpenAICompletions, + Operation: OperationChatCompletion, + ClientModel: "client-model", + UpstreamModel: "upstream-model", + Method: http.MethodPost, + Path: "/v1/chat/completions", + Query: url.Values{"api-version": {"2026-01-01"}}, + Header: http.Header{"X-Test": {"original"}}, + Body: []byte(`{"model":"client-model"}`), + TargetConfig: json.RawMessage(`{"base_url":"https://upstream.example"}`), + Timeouts: AttemptTimeouts{ + Connect: 5 * time.Second, + FirstByte: 10 * time.Second, + Request: 30 * time.Second, + StreamIdle: 15 * time.Second, + }, + Credential: NewCredentialSnapshot(41, 7, 3, credentialData), + } +} + +type contractExecutor struct{} + +func (contractExecutor) Capabilities() CapabilitySet { + capabilities, _ := NewCapabilitySet(Capability{Operation: OperationChatCompletion}) + return capabilities +} + +func (contractExecutor) Execute(context.Context, AttemptSpec) AttemptResult { + return AttemptResult{} +} + +func (contractExecutor) ExecuteStream(context.Context, AttemptSpec, StreamSink) StreamResult { + return StreamResult{} +} + +var _ Executor = contractExecutor{} diff --git a/internal/execution/validation.go b/internal/execution/validation.go new file mode 100644 index 00000000..53f1b18e --- /dev/null +++ b/internal/execution/validation.go @@ -0,0 +1,372 @@ +package execution + +import ( + "encoding/json" + "net/http" + "net/url" + "strings" + "unicode" + "unicode/utf8" +) + +// Validate validates the credential identity and secret data. +func (c CredentialSnapshot) Validate() error { + if c.ID == 0 { + return validationError("credential.id", "must be greater than zero") + } + if c.Version == 0 { + return validationError("credential.version", "must be greater than zero") + } + if c.IdentityGeneration == 0 { + return validationError("credential.identity_generation", "must be greater than zero") + } + if len(c.data) == 0 { + return validationError("credential.data", "must not be empty") + } + return nil +} + +// Validate validates the timeout values. +func (t AttemptTimeouts) Validate() error { + if t.Connect < 0 { + return validationError("timeouts.connect", "must not be negative") + } + if t.FirstByte < 0 { + return validationError("timeouts.first_byte", "must not be negative") + } + if t.Request < 0 { + return validationError("timeouts.request", "must not be negative") + } + if t.StreamIdle < 0 { + return validationError("timeouts.stream_idle", "must not be negative") + } + return nil +} + +// Validate validates the complete attempt specification. +func (s AttemptSpec) Validate() error { + if strings.TrimSpace(s.RequestID) == "" { + return validationError("request_id", "must not be empty") + } + if strings.TrimSpace(s.AttemptID) == "" { + return validationError("attempt_id", "must not be empty") + } + if s.Sequence == 0 { + return validationError("sequence", "must be greater than zero") + } + if strings.TrimSpace(s.ChannelID) == "" { + return validationError("channel_id", "must not be empty") + } + if !s.ClientProtocol.Valid() { + return validationError("client_protocol", "unsupported value") + } + if !s.Operation.Valid() { + return validationError("operation", "unsupported value") + } + if operationRequiresModel(s.Operation) { + if strings.TrimSpace(s.ClientModel) == "" { + return validationError("client_model", "must not be empty for this operation") + } + if strings.TrimSpace(s.UpstreamModel) == "" { + return validationError("upstream_model", "must not be empty for this operation") + } + } + if !validHTTPToken(s.Method) { + return validationError("method", "must be a valid HTTP method") + } + if err := validateRequestPath(s.Path); err != nil { + return validationError("path", "must be an absolute request path without query or fragment") + } + if s.RawQuery != "" { + if len(s.Query) != 0 { + return validationError("raw_query", "must not be combined with query") + } + if strings.Contains(s.RawQuery, "#") || containsControl(s.RawQuery) { + return validationError("raw_query", "contains a fragment or control characters") + } + } + if err := validateValues(s.Query); err != nil { + return validationError("query", "contains invalid control characters") + } + if err := validateHeader(s.Header); err != nil { + return validationError("header", "contains an invalid name or value") + } + if len(s.TargetConfig) > 0 && !json.Valid(s.TargetConfig) { + return validationError("target_config", "must contain valid JSON") + } + if err := s.Timeouts.Validate(); err != nil { + return err + } + return s.Credential.Validate() +} + +// Validate validates error evidence. +func (e ErrorEvidence) Validate() error { + if !e.Kind.Valid() { + return validationError("error.kind", "unsupported value") + } + if !validOptionalHTTPStatus(e.StatusCode) { + return validationError("error.status_code", "must be zero or a valid HTTP status") + } + if strings.TrimSpace(e.Summary) == "" { + return validationError("error.summary", "must not be empty") + } + if utf8.RuneCountInString(e.Summary) > MaxErrorSummaryLength { + return validationError("error.summary", "exceeds maximum length") + } + if containsControl(e.Summary) { + return validationError("error.summary", "must be a single safe line") + } + if containsControl(e.Type) { + return validationError("error.type", "contains control characters") + } + if containsControl(e.Code) { + return validationError("error.code", "contains control characters") + } + if containsControl(e.RequestID) { + return validationError("error.request_id", "contains control characters") + } + if e.RetryAfter < 0 { + return validationError("error.retry_after", "must not be negative") + } + if err := validateSafeEvidenceHeader(e.Header); err != nil { + return validationError("error.header", "contains a non-safe header") + } + return nil +} + +// Validate validates a terminal attempt result. +func (r AttemptResult) Validate() error { + if err := validateResultMetadata( + r.DispatchState, + r.ResponseStarted, + r.StatusCode, + r.Header, + r.UpstreamRequestID, + r.Error, + ); err != nil { + return err + } + if r.DispatchState == DispatchNotSent && len(r.Body) > 0 { + return validationError("body", "must be empty when dispatch_state is not_sent") + } + if !r.ResponseStarted && len(r.Body) > 0 { + return validationError("body", "requires response_started") + } + return nil +} + +// Validate validates a stream event. +func (e StreamEvent) Validate() error { + if !e.Kind.Valid() { + return validationError("stream_event.kind", "unsupported value") + } + if e.Sequence == 0 { + return validationError("stream_event.sequence", "must be greater than zero") + } + switch e.Kind { + case StreamEventReady: + if e.Sequence != 1 { + return validationError("stream_event.sequence", "ready event must be first") + } + if !validHTTPStatus(e.StatusCode) { + return validationError("stream_event.status_code", "ready event requires a valid HTTP status") + } + if e.Header == nil { + return validationError("stream_event.header", "ready event requires initial headers") + } + if err := validateHeader(e.Header); err != nil { + return validationError("stream_event.header", "contains an invalid name or value") + } + if len(e.Data) > 0 || e.Usage != nil { + return validationError("stream_event.ready", "must not contain data or usage") + } + case StreamEventData: + if e.Sequence == 1 { + return validationError("stream_event.sequence", "data must follow a ready event") + } + if len(e.Data) == 0 { + return validationError("stream_event.data", "must not be empty") + } + if e.Usage != nil || hasReadyMetadata(e) { + return validationError("stream_event.data", "must not contain ready metadata or usage") + } + case StreamEventUsage: + if e.Sequence == 1 { + return validationError("stream_event.sequence", "usage must follow a ready event") + } + if e.Usage == nil { + return validationError("stream_event.usage", "must not be empty") + } + if len(e.Data) > 0 || hasReadyMetadata(e) { + return validationError("stream_event.usage", "must not contain ready metadata or data") + } + } + return nil +} + +// Validate validates a terminal streaming result. +func (r StreamResult) Validate() error { + return validateResultMetadata( + r.DispatchState, + r.ResponseStarted, + r.StatusCode, + r.Header, + r.UpstreamRequestID, + r.Error, + ) +} + +func validateResultMetadata( + dispatchState DispatchState, + responseStarted bool, + statusCode int, + header http.Header, + upstreamRequestID string, + errorEvidence *ErrorEvidence, +) error { + if !dispatchState.Valid() { + return validationError("dispatch_state", "unsupported value") + } + if statusCode > 0 && !responseStarted { + return validationError("response_started", "must be true when status_code is present") + } + if !responseStarted && (len(header) > 0 || upstreamRequestID != "") { + return validationError("response_started", "must be true when upstream response metadata is present") + } + if responseStarted && !validHTTPStatus(statusCode) { + return validationError("status_code", "response_started requires a valid HTTP status") + } + if dispatchState == DispatchNotSent { + if responseStarted || statusCode != 0 || len(header) > 0 || upstreamRequestID != "" { + return validationError("dispatch_state", "not_sent cannot contain upstream response metadata") + } + } + if responseStarted && dispatchState != DispatchMaybeSent { + return validationError("dispatch_state", "response_started requires maybe_sent") + } + if err := validateHeader(header); err != nil { + return validationError("header", "contains an invalid name or value") + } + if containsControl(upstreamRequestID) { + return validationError("upstream_request_id", "contains control characters") + } + if errorEvidence == nil { + if !responseStarted || statusCode < http.StatusOK || statusCode >= http.StatusMultipleChoices { + return validationError("error", "is required unless a started response has a 2xx status") + } + return nil + } + if err := errorEvidence.Validate(); err != nil { + return err + } + if errorEvidence.StatusCode > 0 && !responseStarted { + return validationError("error.status_code", "requires response_started") + } + if responseStarted && errorEvidence.StatusCode != 0 && errorEvidence.StatusCode != statusCode { + return validationError("error.status_code", "must match result status_code") + } + return nil +} + +func operationRequiresModel(operation Operation) bool { + return operation == OperationChatCompletion || + operation == OperationResponsesCreate || + operation == OperationResponsesCompact || + operation == OperationResponsesInputTokens +} + +func validateRequestPath(path string) error { + if !strings.HasPrefix(path, "/") || strings.HasPrefix(path, "//") { + return validationError("path", "invalid prefix") + } + parsed, err := url.ParseRequestURI(path) + if err != nil || parsed.IsAbs() || parsed.Host != "" || parsed.RawQuery != "" || parsed.ForceQuery || parsed.Fragment != "" { + return validationError("path", "invalid request path") + } + return nil +} + +func validateValues(values url.Values) error { + for key, items := range values { + if key == "" || containsControl(key) { + return validationError("query", "invalid key") + } + for _, item := range items { + if containsControl(item) { + return validationError("query", "invalid value") + } + } + } + return nil +} + +func validateHeader(header http.Header) error { + for name, values := range header { + if !validHTTPToken(name) { + return validationError("header", "invalid name") + } + for _, value := range values { + if strings.ContainsAny(value, "\r\n\x00") { + return validationError("header", "invalid value") + } + } + } + return nil +} + +func validateSafeEvidenceHeader(header http.Header) error { + if err := validateHeader(header); err != nil { + return err + } + for name := range header { + canonical := http.CanonicalHeaderKey(name) + if canonical == "Retry-After" || + canonical == "Request-Id" || + canonical == "X-Request-Id" || + canonical == "Traceparent" || + strings.HasPrefix(canonical, "Ratelimit-") || + strings.HasPrefix(canonical, "X-Ratelimit-") { + continue + } + return validationError("header", "not allowlisted") + } + return nil +} + +func validOptionalHTTPStatus(statusCode int) bool { + return statusCode == 0 || validHTTPStatus(statusCode) +} + +func validHTTPStatus(statusCode int) bool { + return statusCode >= 100 && statusCode <= 599 +} + +func validHTTPToken(value string) bool { + if value == "" { + return false + } + for _, r := range value { + if r <= 31 || r >= 127 || strings.ContainsRune("()<>@,;:\\\"/[]?={} \t", r) { + return false + } + } + return true +} + +func containsControl(value string) bool { + for _, r := range value { + if unicode.IsControl(r) { + return true + } + } + return false +} + +func hasReadyMetadata(event StreamEvent) bool { + return event.StatusCode != 0 || event.Header != nil || event.UpstreamRequestID != "" +} + +func validationError(field, reason string) *ValidationError { + return &ValidationError{Field: field, Reason: reason} +} diff --git a/internal/gateway/channel_test_helpers_test.go b/internal/gateway/channel_test_helpers_test.go new file mode 100644 index 00000000..4c0f317e --- /dev/null +++ b/internal/gateway/channel_test_helpers_test.go @@ -0,0 +1,79 @@ +package gateway + +import ( + "context" + "encoding/json" + "strconv" + "testing" + + "gpt-load/internal/channel" + bifrostexecutor "gpt-load/internal/execution/bifrost" + "gpt-load/internal/platform/encryption" + "gpt-load/internal/protocol" + "gpt-load/internal/state" +) + +func testChannelConfig( + t *testing.T, + selectedProtocol protocol.Protocol, + baseURL string, +) (channel.ID, json.RawMessage) { + t.Helper() + var channelID channel.ID + switch selectedProtocol { + case protocol.OpenAICompletions, protocol.OpenAIResponses: + channelID = channel.OpenAICompatible + case protocol.Anthropic: + channelID = channel.AnthropicCompatible + case protocol.Gemini: + channelID = channel.GeminiCompatible + default: + t.Fatalf("unsupported test protocol %q", selectedProtocol) + } + params, err := json.Marshal(map[string]string{"base_url": baseURL}) + if err != nil { + t.Fatalf("encode channel params: %v", err) + } + return channelID, params +} + +func testCredentialConfig(id, groupID uint) state.CredentialConfig { + return state.CredentialConfig{ + ID: id, GroupID: groupID, Status: state.CredentialStatusActive, + Version: 1, IdentityGeneration: uint64(id), + Fingerprint: "test-credential-" + strconv.FormatUint(uint64(id), 10), + } +} + +func testCredentialEntry( + t *testing.T, + service encryption.Service, + id uint, + groupID uint, + apiKey string, +) state.CredentialEntry { + t.Helper() + payload, err := json.Marshal(map[string]string{"api_key": apiKey}) + if err != nil { + t.Fatalf("encode credential %d: %v", id, err) + } + encrypted, err := service.Encrypt(string(payload)) + if err != nil { + t.Fatalf("encrypt credential %d: %v", id, err) + } + return state.CredentialEntry{ + ID: id, GroupID: groupID, Status: state.CredentialStatusActive, + Version: 1, IdentityGeneration: uint64(id), + Fingerprint: "test-credential-" + strconv.FormatUint(uint64(id), 10), EncryptedValue: encrypted, + } +} + +func newTestExecutionForwarder(t *testing.T) AttemptForwarder { + t.Helper() + runtime, err := bifrostexecutor.NewRuntime(context.Background()) + if err != nil { + t.Fatalf("initialize execution runtime: %v", err) + } + t.Cleanup(runtime.Shutdown) + return NewExecutionForwarder(runtime) +} diff --git a/internal/gateway/dialects_integration_test.go b/internal/gateway/dialects_integration_test.go index d8de6804..42c9c0bd 100644 --- a/internal/gateway/dialects_integration_test.go +++ b/internal/gateway/dialects_integration_test.go @@ -18,12 +18,11 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/dialect" "gpt-load/internal/health" "gpt-load/internal/platform/contentcoding" "gpt-load/internal/platform/encryption" - platformhttp "gpt-load/internal/platform/httpclient" - "gpt-load/internal/platform/redact" "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/testutil/fakeupstream" @@ -62,7 +61,7 @@ func newDialectGatewayEngine( model string, dialects dialect.Set, groups ...dialectGatewayGroup, -) (*gin.Engine, *state.KeyRegistry) { +) (*gin.Engine, *state.CredentialRegistry) { t.Helper() gin.SetMode(gin.TestMode) keyService, err := encryption.NewService("dialect-gateway-test-master-key") @@ -71,34 +70,31 @@ func newDialectGatewayEngine( } configs := make([]state.GroupConfig, 0, len(groups)) - entries := make([]state.KeyEntry, 0) - keyID := uint(1) + entries := make([]state.CredentialEntry, 0) + credentialConfigs := make([]state.CredentialConfig, 0) + credentialID := uint(1) for _, group := range groups { models := group.models if len(models) == 0 { models = []state.ModelConfig{{ID: model}} } + baseURL := testUpstreamBaseURL(group.upstreamURL, selectedProtocol) + channelID, params := testChannelConfig(t, selectedProtocol, baseURL) configs = append(configs, state.GroupConfig{ - ID: group.id, Name: group.name, UpstreamURL: testUpstreamBaseURL(group.upstreamURL, selectedProtocol), - Protocols: []protocol.Protocol{selectedProtocol}, - Models: models, Enabled: true, + ID: group.id, Name: group.name, ChannelID: channelID, Params: params, + Models: models, Enabled: true, }) for _, apiKey := range group.apiKeys { - ciphertext, encryptErr := keyService.Encrypt(apiKey) - if encryptErr != nil { - t.Fatalf("Encrypt(group %d key) error = %v", group.id, encryptErr) - } - entries = append(entries, state.KeyEntry{ - ID: keyID, GroupID: group.id, Status: state.KeyStatusActive, - EncryptedValue: ciphertext, - }) - keyID++ + entries = append(entries, testCredentialEntry(t, keyService, credentialID, group.id, apiKey)) + credentialConfigs = append(credentialConfigs, testCredentialConfig(credentialID, group.id)) + credentialID++ } } manager := state.NewManager() snapshot, err := manager.Publish(state.CompileInput{ - Groups: configs, + ChannelRegistry: channel.NewRegistry(), Groups: configs, + Credentials: credentialConfigs, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: keyService.Hash("gl-client"), Status: state.AccessKeyStatusActive, @@ -116,15 +112,15 @@ func newDialectGatewayEngine( snapshot.Groups[group.id] = view } - registry := state.NewKeyRegistry() - if err := registry.Replace(entries); err != nil { - t.Fatalf("Replace() error = %v", err) + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials(entries); err != nil { + t.Fatalf("ReplaceCredentials() error = %v", err) } handler := NewHandler( manager, registry, keyService, - NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()), + newTestExecutionForwarder(t), dialects, health.NewStatsStore(), health.NewMutationCoordinator(), @@ -148,22 +144,22 @@ func TestGatewayDecodesSupportedClientContentCodings(t *testing.T) { }{ { name: "OpenAI completions", protocol: protocol.OpenAICompletions, - dialects: dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewOpenAI()), path: "/v1/chat/completions", plaintext: []byte(`{"model":"public-model","stream":false}`), }, { name: "Anthropic messages", protocol: protocol.Anthropic, - dialects: dialect.NewSet(dialect.NewAnthropic(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewAnthropic()), path: "/v1/messages", plaintext: []byte(`{"model":"public-model","max_tokens":1,"messages":[]}`), }, { name: "Gemini generation", protocol: protocol.Gemini, - dialects: dialect.NewSet(dialect.NewGemini(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewGemini()), path: "/v1beta/models/public-model:generateContent", plaintext: []byte(`{"contents":[{"role":"user","parts":[{"text":"ping"}]}]}`), }, { name: "Responses create", protocol: protocol.OpenAIResponses, - dialects: dialect.NewSet(dialect.NewOpenAIResponses(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewOpenAIResponses()), path: "/v1/responses", plaintext: []byte(`{"model":"public-model","input":"ping"}`), }, } @@ -256,7 +252,7 @@ func TestResponsesNamespaceRoutesOrdinaryMethodsAndRejectsDangerousMethodsLocall t, protocol.OpenAIResponses, "public-model", - dialect.NewSet(dialect.NewOpenAIResponses(http.DefaultClient)), + dialect.NewSet(dialect.NewOpenAIResponses()), dialectGatewayGroup{ id: 1, name: "responses", upstreamURL: upstream.URL, apiKeys: []string{"sk-responses"}, @@ -469,7 +465,7 @@ func TestResponsesNamespaceRejectsNormalizationEscapesBeforeForward(t *testing.T t, protocol.OpenAIResponses, "public-model", - dialect.NewSet(dialect.NewOpenAIResponses(http.DefaultClient)), + dialect.NewSet(dialect.NewOpenAIResponses()), dialectGatewayGroup{ id: 1, name: "responses", upstreamURL: upstream.URL, apiKeys: []string{"sk-responses"}, @@ -520,7 +516,7 @@ func TestResponsesNamespaceRejectsDecodedDuplicateStreamQueryBeforeForward(t *te t, protocol.OpenAIResponses, "public-model", - dialect.NewSet(dialect.NewOpenAIResponses(http.DefaultClient)), + dialect.NewSet(dialect.NewOpenAIResponses()), dialectGatewayGroup{ id: 1, name: "responses", upstreamURL: upstream.URL, apiKeys: []string{"sk-responses"}, @@ -557,7 +553,7 @@ func TestGatewayRewritesEachAttemptFromOriginal(t *testing.T) { defer second.Close() engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, "public", - dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), + dialect.NewSet(dialect.NewOpenAI()), dialectGatewayGroup{ id: 1, name: "first", upstreamURL: first.URL, apiKeys: []string{"sk-first"}, models: []state.ModelConfig{{ID: "provider-one", Alias: "public"}}, @@ -608,7 +604,7 @@ func TestHandlerHostFailureSkipsGroupForCurrentRequestOnly(t *testing.T) { defer backup.Close() engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, "gpt-4o", - dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), + dialect.NewSet(dialect.NewOpenAI()), dialectGatewayGroup{id: 1, name: "primary", upstreamURL: primary.URL, apiKeys: []string{"sk-primary-one", "sk-primary-two"}}, dialectGatewayGroup{id: 2, name: "backup", upstreamURL: backup.URL, @@ -641,7 +637,7 @@ func TestHandlerReturnsLastHostErrorWhenSkippedGroupHasNoBackup(t *testing.T) { defer upstream.Close() engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, "gpt-4o", - dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), + dialect.NewSet(dialect.NewOpenAI()), dialectGatewayGroup{id: 1, name: "only", upstreamURL: upstream.URL, apiKeys: []string{"sk-one", "sk-two"}}, ) @@ -670,28 +666,28 @@ func TestForwarderRewritesAliasedNonStreamingResponses(t *testing.T) { }{ { name: "OpenAI", value: protocol.OpenAICompletions, - dialects: dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewOpenAI()), path: "/v1/chat/completions", requestBody: `{"model":"public-model"}`, upstreamResponse: `{"id":"chatcmpl-1","model":"provider-response"}`, responseField: "model", }, { name: "OpenAI Responses", value: protocol.OpenAIResponses, - dialects: dialect.NewSet(dialect.NewOpenAIResponses(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewOpenAIResponses()), path: "/v1/responses", requestBody: `{"model":"public-model","input":"ping"}`, upstreamResponse: `{"id":"resp-1","object":"response","model":"provider-response"}`, responseField: "model", }, { name: "Anthropic", value: protocol.Anthropic, - dialects: dialect.NewSet(dialect.NewAnthropic(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewAnthropic()), path: "/v1/messages", requestBody: `{"model":"public-model"}`, upstreamResponse: `{"type":"message","model":"provider-response"}`, responseField: "model", }, { name: "Gemini", value: protocol.Gemini, - dialects: dialect.NewSet(dialect.NewGemini(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewGemini()), path: "/v1beta/models/public-model:generateContent", requestBody: `{}`, upstreamResponse: `{"modelVersion":"provider-response","candidates":[]}`, responseField: "modelVersion", @@ -772,7 +768,7 @@ func TestTransparentModelRoutePreservesWire(t *testing.T) { defer upstream.Close() engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, "same-model", - dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), + dialect.NewSet(dialect.NewOpenAI()), dialectGatewayGroup{ id: 1, name: "transparent", upstreamURL: upstream.URL, apiKeys: []string{"provider-key"}, models: []state.ModelConfig{{ID: "same-model"}}, @@ -822,14 +818,14 @@ func TestGatewayRewritesAliasedStreams(t *testing.T) { }{ { name: "OpenAI", value: protocol.OpenAICompletions, - dialects: dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewOpenAI()), path: "/v1/chat/completions", requestBody: `{"model":"public-model","stream":true}`, streamBody: "data: {\"id\":\"1\",\"model\":\"provider-model\",\"choices\":[]}\n\ndata: [DONE]\n\n", want: `"model":"public-model"`, unchanged: "data: [DONE]\n\n", }, { name: "Anthropic", value: protocol.Anthropic, - dialects: dialect.NewSet(dialect.NewAnthropic(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewAnthropic()), path: "/v1/messages", requestBody: `{"model":"public-model","stream":true}`, streamBody: "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"model\":\"provider-model\"}}\n\n" + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"unchanged\"}}\n\n", @@ -837,7 +833,7 @@ func TestGatewayRewritesAliasedStreams(t *testing.T) { }, { name: "Gemini", value: protocol.Gemini, - dialects: dialect.NewSet(dialect.NewGemini(http.DefaultClient)), + dialects: dialect.NewSet(dialect.NewGemini()), path: "/v1beta/models/public-model:streamGenerateContent", requestBody: `{}`, streamBody: "data: {\"modelVersion\":\"provider-model\",\"candidates\":[]}\n\n", want: `"modelVersion":"public-model"`, @@ -910,13 +906,13 @@ func TestAnthropicGatewayNonStreamAuthAndForwarding(t *testing.T) { wantVersion string }{ {name: "Bearer remains Anthropic", authHeader: "Bearer gl-client", wantVersion: anthropicDefaultVersionForTest}, - {name: "x-api-key carrier", apiKey: "gl-client", version: "2024-01-01", wantVersion: "2024-01-01"}, + {name: "x-api-key carrier", apiKey: "gl-client", version: "2024-01-01", wantVersion: anthropicDefaultVersionForTest}, } { t.Run(test.name, func(t *testing.T) { upstream := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "success.json"}) defer upstream.Close() engine, _ := newDialectGatewayEngine(t, protocol.Anthropic, "claude-3-5-sonnet", - dialect.NewSet(dialect.NewAnthropic(http.DefaultClient)), + dialect.NewSet(dialect.NewAnthropic()), dialectGatewayGroup{id: 1, name: "anthropic", upstreamURL: upstream.URL, apiKeys: []string{"sk-anthropic-upstream"}}, ) body := `{"model":"claude-3-5-sonnet","messages":[{"role":"user","content":"ping"}]}` @@ -960,7 +956,7 @@ func TestAnthropicGatewayFailover(t *testing.T) { ) defer upstream.Close() engine, _ := newDialectGatewayEngine(t, protocol.Anthropic, "claude-3-5-sonnet", - dialect.NewSet(dialect.NewAnthropic(http.DefaultClient)), + dialect.NewSet(dialect.NewAnthropic()), dialectGatewayGroup{id: 1, name: "anthropic", upstreamURL: upstream.URL, apiKeys: []string{"sk-anthropic-one", "sk-anthropic-two"}}, ) request := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-3-5-sonnet"}`)) @@ -990,7 +986,7 @@ func TestAnthropicGatewayFailover(t *testing.T) { })) defer clientError.Close() engine, _ = newDialectGatewayEngine(t, protocol.Anthropic, "claude-3-5-sonnet", - dialect.NewSet(dialect.NewAnthropic(http.DefaultClient)), + dialect.NewSet(dialect.NewAnthropic()), dialectGatewayGroup{id: 1, name: "anthropic", upstreamURL: clientError.URL, apiKeys: []string{"sk-one", "sk-two"}}, ) request = httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-3-5-sonnet"}`)) @@ -1024,7 +1020,7 @@ func TestAnthropicGatewayStream(t *testing.T) { defer backup.Close() engine, _ := newDialectGatewayEngine(t, protocol.Anthropic, "claude-3-5-sonnet", - dialect.NewSet(dialect.NewAnthropic(http.DefaultClient)), + dialect.NewSet(dialect.NewAnthropic()), dialectGatewayGroup{ id: 1, name: "primary", upstreamURL: primary.URL, apiKeys: []string{"sk-primary"}, headerRules: state.HeaderRules{Set: map[string]string{"Accept-Encoding": "gzip"}}, @@ -1101,7 +1097,7 @@ func TestGeminiGatewayNonStreamAuthAndQuery(t *testing.T) { upstream := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "success.json"}) defer upstream.Close() engine, _ := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro", - dialect.NewSet(dialect.NewGemini(http.DefaultClient)), + dialect.NewSet(dialect.NewGemini()), dialectGatewayGroup{ id: 1, name: "gemini", upstreamURL: upstream.URL + "?tenant=base&alt=proto", apiKeys: []string{"gemini-upstream-key"}, @@ -1144,7 +1140,7 @@ func TestGeminiGatewayFailover(t *testing.T) { ) defer upstream.Close() engine, _ := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro", - dialect.NewSet(dialect.NewGemini(http.DefaultClient)), + dialect.NewSet(dialect.NewGemini()), dialectGatewayGroup{id: 1, name: "gemini", upstreamURL: upstream.URL, apiKeys: []string{"gemini-key-one", "gemini-key-two"}}, ) request := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:generateContent?trace=true", strings.NewReader(`{}`)) @@ -1166,7 +1162,7 @@ func TestGeminiGatewayStream(t *testing.T) { upstream := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "stream.sse", Stream: true}) defer upstream.Close() engine, _ := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro", - dialect.NewSet(dialect.NewGemini(http.DefaultClient)), + dialect.NewSet(dialect.NewGemini()), dialectGatewayGroup{ id: 1, name: "gemini-stream", upstreamURL: upstream.URL + "?tenant=base&alt=json", apiKeys: []string{"gemini-stream-key"}, @@ -1200,7 +1196,7 @@ func TestGeminiGatewayStream(t *testing.T) { } func TestGeminiGatewayCompressed(t *testing.T) { - t.Run("request-written compressed response terminates before retry", func(t *testing.T) { + t.Run("SDK-normalized response remains one logical attempt", func(t *testing.T) { compressed := fakeupstream.New(fakeupstream.Step{ Status: http.StatusOK, Fixture: "stream.sse", Stream: true, Headers: http.Header{"Content-Encoding": {"gzip"}}, @@ -1209,7 +1205,7 @@ func TestGeminiGatewayCompressed(t *testing.T) { backup := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "stream.sse", Stream: true}) defer backup.Close() engine, _ := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro", - dialect.NewSet(dialect.NewGemini(http.DefaultClient)), + dialect.NewSet(dialect.NewGemini()), dialectGatewayGroup{id: 1, name: "compressed", upstreamURL: compressed.URL, apiKeys: []string{"bad-key"}}, dialectGatewayGroup{id: 2, name: "backup", upstreamURL: backup.URL, apiKeys: []string{"good-key"}}, ) @@ -1217,14 +1213,14 @@ func TestGeminiGatewayCompressed(t *testing.T) { request.Header.Set("Authorization", "Bearer gl-client") recorder := httptest.NewRecorder() engine.ServeHTTP(recorder, request) - if recorder.Code != http.StatusBadGateway || recorder.Header().Get(debugHeaderAttempts) != "1" || - !strings.Contains(recorder.Body.String(), reasonUpstreamProtocol.Code) || + if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "1" || + !strings.Contains(recorder.Body.String(), "data:") || len(compressed.Requests()) != 1 || len(backup.Requests()) != 0 { t.Fatalf("response = %d headers=%v compressed=%d backup=%d body=%s", recorder.Code, recorder.Header(), len(compressed.Requests()), len(backup.Requests()), recorder.Body.String()) } }) - t.Run("all bad candidates return fixed protocol error", func(t *testing.T) { + t.Run("transport encoding metadata does not trigger candidate retry", func(t *testing.T) { first := fakeupstream.New(fakeupstream.Step{ Status: http.StatusOK, Fixture: "stream.sse", Stream: true, Headers: http.Header{"Content-Encoding": {"gzip"}}, @@ -1236,7 +1232,7 @@ func TestGeminiGatewayCompressed(t *testing.T) { }) defer second.Close() engine, registry := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro", - dialect.NewSet(dialect.NewGemini(http.DefaultClient)), + dialect.NewSet(dialect.NewGemini()), dialectGatewayGroup{id: 1, name: "first", upstreamURL: first.URL, apiKeys: []string{"secret-one"}}, dialectGatewayGroup{id: 2, name: "second", upstreamURL: second.URL, apiKeys: []string{"secret-two"}}, ) @@ -1244,12 +1240,13 @@ func TestGeminiGatewayCompressed(t *testing.T) { request.Header.Set("Authorization", "Bearer gl-client") recorder := httptest.NewRecorder() engine.ServeHTTP(recorder, request) - if recorder.Code != http.StatusBadGateway || recorder.Header().Get(debugHeaderAttempts) != "1" || - !strings.Contains(recorder.Body.String(), reasonUpstreamProtocol.Code) || - len(registry.CollectCandidates([]uint{1, 2}, nil, time.Time{})) != 2 { - t.Fatalf("response = %d headers=%v candidates=%d body=%s", recorder.Code, recorder.Header(), len(registry.CollectCandidates([]uint{1, 2}, nil, time.Time{})), recorder.Body.String()) + if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "1" || + !strings.Contains(recorder.Body.String(), "data:") || + len(first.Requests()) != 1 || len(second.Requests()) != 0 || + len(registry.CollectCredentialCandidates([]uint{1, 2}, nil, time.Time{})) != 2 { + t.Fatalf("response = %d headers=%v candidates=%d body=%s", recorder.Code, recorder.Header(), len(registry.CollectCredentialCandidates([]uint{1, 2}, nil, time.Time{})), recorder.Body.String()) } - for _, forbidden := range []string{"secret-one", "secret-two", "data:"} { + for _, forbidden := range []string{"secret-one", "secret-two"} { if strings.Contains(recorder.Body.String(), forbidden) { t.Fatalf("response exposes %q: %s", forbidden, recorder.Body.String()) } diff --git a/internal/gateway/execution_boundary.go b/internal/gateway/execution_boundary.go new file mode 100644 index 00000000..494ee8d3 --- /dev/null +++ b/internal/gateway/execution_boundary.go @@ -0,0 +1,154 @@ +package gateway + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/health" +) + +type normalizedCredential struct { + apiKey string + payload []byte + secrets []string +} + +func normalizeChannelCredential( + registry *channel.Registry, + channelID channel.ID, + decrypted string, +) (normalizedCredential, error) { + if registry == nil || channelID == "" { + return normalizedCredential{}, fmt.Errorf("credential channel is unavailable") + } + trimmed := strings.TrimSpace(decrypted) + if trimmed == "" { + return normalizedCredential{}, fmt.Errorf("credential is empty") + } + raw := []byte(trimmed) + if !json.Valid(raw) { + encoded, err := json.Marshal(map[string]string{"api_key": trimmed}) + if err != nil { + return normalizedCredential{}, fmt.Errorf("encode legacy credential") + } + raw = encoded + } + validated, err := registry.ValidateCredential(channelID, raw) + if err != nil { + return normalizedCredential{}, fmt.Errorf("validate credential: %w", err) + } + payload := bytes.Clone(validated.CanonicalJSON()) + apiKey, _ := validated.Value("api_key") + return normalizedCredential{ + apiKey: apiKey, payload: payload, secrets: credentialSecretValues(payload), + }, nil +} + +func credentialSecretValues(payload []byte) []string { + var values map[string]string + if json.Unmarshal(payload, &values) != nil { + return nil + } + fields := []string{ + "api_key", "client_id", "client_secret", "tenant_id", "access_key", "secret_key", + "session_token", "role_arn", "external_id", "session_name", "service_account_json", + } + seen := make(map[string]struct{}, len(fields)) + result := make([]string, 0, len(fields)) + appendValue := func(value string) { + if value == "" { + return + } + if _, exists := seen[value]; exists { + return + } + seen[value] = struct{}{} + result = append(result, value) + } + for _, field := range fields { + value := values[field] + appendValue(value) + if field != "service_account_json" || value == "" { + continue + } + var nested map[string]json.RawMessage + if json.Unmarshal([]byte(value), &nested) != nil { + continue + } + for _, nestedField := range []string{"private_key", "private_key_id", "client_email"} { + var nestedValue string + if json.Unmarshal(nested[nestedField], &nestedValue) == nil { + appendValue(nestedValue) + } + } + } + return result +} + +func judgeUpstreamResult( + result UpstreamResult, + now time.Time, +) health.Result { + result = normalizeUpstreamResultEvidence(result) + downstreamErr := result.Err + if result.ExecutionError != nil && !result.Committed { + downstreamErr = nil + } + return health.JudgeExecution(health.ExecutionAttempt{ + DispatchState: result.DispatchState, + StatusCode: result.StatusCode, + Header: result.Header, + Evidence: result.ExecutionError, + DownstreamCommitted: result.Committed, + DownstreamErr: downstreamErr, + Now: now, + }) +} + +func normalizeUpstreamResultEvidence(result UpstreamResult) UpstreamResult { + if result.DispatchState.Valid() { + return result + } + result.DispatchState, result.ExecutionError = inferAttemptEvidence(result) + result.ResponseStarted = result.ResponseStarted || result.StatusCode != 0 + return result +} + +// inferAttemptEvidence keeps the AttemptForwarder boundary fail-safe when a +// custom implementation supplies HTTP or transport evidence without the +// normalized execution fields used by the built-in adapter. +func inferAttemptEvidence(result UpstreamResult) (execution.DispatchState, *execution.ErrorEvidence) { + dispatchState := execution.DispatchNotSent + if result.RequestWritten || result.StatusCode != 0 || result.ResponseStarted { + dispatchState = execution.DispatchMaybeSent + } + if result.Err != nil { + kind := execution.ErrorKindTransport + switch { + case errors.Is(result.Err, context.Canceled): + kind = execution.ErrorKindCanceled + case errors.Is(result.Err, context.DeadlineExceeded): + kind = execution.ErrorKindTimeout + } + return dispatchState, &execution.ErrorEvidence{Kind: kind, Summary: result.Err.Error()} + } + if result.StatusCode >= 200 && result.StatusCode < 300 && !result.ProviderErrorBeforeCommit { + return dispatchState, nil + } + summary := strings.TrimSpace(string(result.ClassificationBody)) + if summary == "" { + summary = strings.TrimSpace(string(result.Body)) + } + return dispatchState, &execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, + StatusCode: result.StatusCode, + Summary: summary, + } +} diff --git a/internal/gateway/execution_boundary_test.go b/internal/gateway/execution_boundary_test.go new file mode 100644 index 00000000..6c1f96d2 --- /dev/null +++ b/internal/gateway/execution_boundary_test.go @@ -0,0 +1,99 @@ +package gateway + +import ( + "net/http" + "reflect" + "testing" + "time" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/health" +) + +func TestNormalizeChannelCredentialAcceptsLegacyAndTypedStorage(t *testing.T) { + t.Parallel() + + for _, test := range []struct { + name string + decrypted string + wantAPIKey string + wantPayload string + }{ + { + name: "legacy raw key", decrypted: " sk-legacy ", + wantAPIKey: "sk-legacy", wantPayload: `{"api_key":"sk-legacy"}`, + }, + { + name: "typed credential", decrypted: ` {"api_key":"sk-typed"} `, + wantAPIKey: "sk-typed", wantPayload: `{"api_key":"sk-typed"}`, + }, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + credential, err := normalizeChannelCredential(channel.NewRegistry(), channel.OpenAI, test.decrypted) + if err != nil { + t.Fatalf("normalizeChannelCredential() error = %v", err) + } + if credential.apiKey != test.wantAPIKey || string(credential.payload) != test.wantPayload { + t.Fatalf("normalizeChannelCredential() = %q %s", credential.apiKey, credential.payload) + } + }) + } + + for _, invalid := range []string{"", " ", `[]`, `{}`, `{"api_key":""}`} { + if credential, err := normalizeChannelCredential(channel.NewRegistry(), channel.OpenAI, invalid); err == nil { + t.Fatalf("normalizeChannelCredential(%q) = %#v, nil", invalid, credential) + } + } +} + +func TestNormalizeChannelCredentialPreservesStructuredCloudSecrets(t *testing.T) { + t.Parallel() + + got, err := normalizeChannelCredential( + channel.NewRegistry(), + channel.AWSBedrock, + ` {"access_key":"AKIA_TEST","secret_key":"bedrock-secret","session_token":"bedrock-session"} `, + ) + if err != nil { + t.Fatalf("normalizeChannelCredential() error = %v", err) + } + if got.apiKey != "" || string(got.payload) != + `{"access_key":"AKIA_TEST","secret_key":"bedrock-secret","session_token":"bedrock-session"}` { + t.Fatalf("normalized credential = %#v / %s", got, got.payload) + } + if want := []string{"AKIA_TEST", "bedrock-secret", "bedrock-session"}; !reflect.DeepEqual(got.secrets, want) { + t.Fatalf("secrets = %#v, want %#v", got.secrets, want) + } +} + +func TestJudgeUpstreamResultUsesNeutralExecutionEvidence(t *testing.T) { + t.Parallel() + + now := time.Unix(1_800_000_000, 0) + result := judgeUpstreamResult(UpstreamResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusTooManyRequests, + ExecutionError: &execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, StatusCode: http.StatusTooManyRequests, + Summary: "upstream rejected request", RetryAfter: 7 * time.Second, + }, + }, now) + if result.Category != health.FailureCategoryRateLimited || + result.Action != health.ActionCooldownCredential || + !result.CooldownUntil.Equal(now.Add(7*time.Second)) { + t.Fatalf("JudgeExecution() = %#v", result) + } + + result = judgeUpstreamResult(UpstreamResult{ + DispatchState: execution.DispatchNotSent, + ExecutionError: &execution.ErrorEvidence{ + Kind: execution.ErrorKindTransport, Summary: "connection failed", + }, + }, now) + if result.Category != health.FailureCategoryUpstreamHostError || result.Action != health.ActionSkipGroup { + t.Fatalf("JudgeExecution(not sent) = %#v", result) + } +} diff --git a/internal/gateway/execution_forward.go b/internal/gateway/execution_forward.go new file mode 100644 index 00000000..c5c22cef --- /dev/null +++ b/internal/gateway/execution_forward.go @@ -0,0 +1,413 @@ +package gateway + +import ( + "context" + "fmt" + "io" + "net/http" + "strings" + "time" + + "gpt-load/internal/execution" + "gpt-load/internal/platform/redact" + "gpt-load/internal/usage" +) + +// ExecutionForwarder adapts the provider-neutral executor to the gateway's +// downstream commit boundary while the handler orchestration remains intact. +type ExecutionForwarder struct { + executor execution.Executor + representation *responseProcessor + writeTimeout time.Duration +} + +func NewExecutionForwarder(executor execution.Executor) *ExecutionForwarder { + return &ExecutionForwarder{ + executor: executor, representation: &responseProcessor{redactor: redact.New()}, + writeTimeout: downstreamWriteTimeout, + } +} + +func (forwarder *ExecutionForwarder) Forward( + ctx context.Context, + input ForwardInput, +) UpstreamResult { + spec, err := newExecutionAttemptSpec(input) + if err != nil || forwarder == nil || forwarder.executor == nil { + return executionInputFailure(err) + } + result := upstreamFromExecutionResult(ctx, input, forwarder.executor.Execute(ctx, spec)) + return forwarder.prepareBufferedResult(input, result) +} + +func (forwarder *ExecutionForwarder) ForwardStream( + ctx context.Context, + input ForwardInput, + downstream http.ResponseWriter, +) UpstreamResult { + spec, err := newExecutionAttemptSpec(input) + if err != nil || forwarder == nil || forwarder.executor == nil || downstream == nil { + return executionInputFailure(err) + } + + writeTimeout := forwarder.writeTimeout + if writeTimeout <= 0 { + writeTimeout = downstreamWriteTimeout + } + controller := newStreamWriteController(downstream, writeTimeout) + defer func() { _ = controller.clear() }() + + var ( + ready *execution.StreamEvent + committed bool + downstreamErr error + errorBody []byte + streamUsage *execution.UsageEvidence + ) + sink := func(event execution.StreamEvent) error { + if err := event.Validate(); err != nil { + downstreamErr = fmt.Errorf("%w: invalid execution stream event", ErrUpstreamProtocol) + return downstreamErr + } + switch event.Kind { + case execution.StreamEventReady: + copy := event.Clone() + if copy.StatusCode >= http.StatusOK && copy.StatusCode < http.StatusMultipleChoices { + if err := validateStreamContentEncoding(copy.Header); err != nil { + downstreamErr = err + return err + } + copy.Header = normalizeStreamResponseHeaders( + sanitizeForwardResponseHeaders(copy.Header, input), + ) + } else { + copy.Header = sanitizeForwardResponseHeaders(copy.Header, input) + } + ready = © + return nil + case execution.StreamEventUsage: + if event.Usage != nil { + copy := event.Usage.Clone() + streamUsage = © + } + return nil + case execution.StreamEventData: + if ready == nil { + downstreamErr = fmt.Errorf("%w: execution data arrived before response metadata", ErrUpstreamProtocol) + return downstreamErr + } + if ready.StatusCode < http.StatusOK || ready.StatusCode >= http.StatusMultipleChoices { + errorBody = appendExecutionErrorBody(errorBody, event.Data) + return nil + } + if !committed { + committed = true + if err := commitStream(controller, ready.StatusCode, ready.Header, event.Data); err != nil { + downstreamErr = err + return err + } + if input.OnStreamReady != nil { + input.OnStreamReady() + } + if input.OnFirstResponse != nil { + input.OnFirstResponse() + } + return nil + } + written, err := controller.write(event.Data) + if err != nil { + downstreamErr = &streamFailure{ + kind: streamFailureDownstreamWrite, + err: fmt.Errorf("write execution stream: %w", err), + } + return downstreamErr + } + if written != len(event.Data) { + downstreamErr = &streamFailure{ + kind: streamFailureDownstreamWrite, + err: fmt.Errorf("write execution stream: %w", io.ErrShortWrite), + } + return downstreamErr + } + if err := controller.flush(); err != nil { + downstreamErr = &streamFailure{ + kind: streamFailureDownstreamWrite, + err: fmt.Errorf("flush execution stream: %w", err), + } + return downstreamErr + } + return nil + default: + downstreamErr = fmt.Errorf("%w: unsupported execution stream event", ErrUpstreamProtocol) + return downstreamErr + } + } + + terminal := forwarder.executor.ExecuteStream(ctx, spec, sink) + result := upstreamFromExecutionStreamResult(ctx, input, terminal, streamUsage) + result.Committed = committed + if len(errorBody) > 0 { + result.Body = append([]byte(nil), errorBody...) + result.ClassificationBody = append([]byte(nil), errorBody...) + } + if downstreamErr != nil { + result.Err = downstreamErr + } + if committed { + result.Stream = executionStreamObservation(ctx, terminal, downstreamErr) + } else if ready != nil { + result.StatusCode = ready.StatusCode + result.Header = ready.Header.Clone() + result.ResponseStarted = true + result.UpstreamRequestID = ready.UpstreamRequestID + } + if !committed && result.HasResponse() { + result = forwarder.prepareBufferedResult(input, result) + } + return result +} + +func (forwarder *ExecutionForwarder) prepareBufferedResult( + input ForwardInput, + result UpstreamResult, +) UpstreamResult { + if forwarder == nil || !result.HasResponse() { + return result + } + representation := forwarder.representation + if representation == nil { + representation = &responseProcessor{redactor: redact.New()} + } + secrets := append([]string(nil), input.CredentialSecrets...) + if input.APIKey != "" { + secrets = append(secrets, input.APIKey) + } + if result.StatusCode >= http.StatusOK && result.StatusCode < http.StatusMultipleChoices { + policy := classifyResponseBody(input.Request.Method, result.StatusCode) + if !policy.readBody { + if err := validateStreamContentEncoding(result.Header); err != nil || + (result.StatusCode == http.StatusPartialContent && !hasSafeHeadContentRange(result.Header)) { + return executionRepresentationFailure(result, fmt.Errorf("%w: invalid bodyless response representation", ErrUpstreamProtocol)) + } + headers := sanitizeForwardResponseHeaders(result.Header, input, secrets...) + headers, _ = normalizeBufferedResponse(input.Request.Method, result.StatusCode, headers, nil) + result.Header = headers + result.Body = nil + return result + } + prepared, err := representation.prepareSuccessRepresentation( + input, result.StatusCode, result.Header, result.Body, secrets, + ) + if err != nil { + return executionRepresentationFailure(result, err) + } + result.Header = prepared.headers + result.Body = prepared.downstream + result.ClassificationBody = prepared.inspectable + return result + } + prepared := representation.prepareErrorRepresentation(input, result.Header, result.Body, secrets) + result.Header = prepared.headers + result.Body = prepared.downstream + result.ClassificationBody = prepared.classification + return result +} + +func executionRepresentationFailure(result UpstreamResult, err error) UpstreamResult { + return UpstreamResult{ + Err: err, RequestWritten: result.RequestWritten, DispatchState: result.DispatchState, + } +} + +func newExecutionAttemptSpec(input ForwardInput) (execution.AttemptSpec, error) { + if input.Request == nil { + return execution.AttemptSpec{}, fmt.Errorf("request is required") + } + headers := cloneEndToEndHeaders(input.Request.Header) + removeDownstreamCredentials(headers) + for name, value := range input.Group.HeaderRules.Set { + headers.Set(name, strings.ReplaceAll(value, "${API_KEY}", input.APIKey)) + } + for _, name := range input.Group.HeaderRules.Remove { + headers.Del(name) + } + sanitizeUpstreamRequestHeaders(headers) + headers.Set("Accept-Encoding", "identity") + spec := execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: input.RequestID, + AttemptID: input.AttemptID, + Sequence: input.AttemptSequence, + ChannelID: input.ChannelID, + ClientProtocol: input.ClientProtocol, + Operation: input.Operation, + ClientModel: input.ExternalModel, + UpstreamModel: input.UpstreamModelID, + Method: input.Request.Method, + Path: input.Request.Path, + RawQuery: input.Request.RawQuery, + Header: headers, + Body: input.Request.Body, + IncludeUsage: input.ObserveUsage && input.Group.InjectUsageOptions, + TargetConfig: input.TargetConfig, + Timeouts: execution.AttemptTimeouts{ + Connect: input.Group.Timeouts.Connect, + FirstByte: input.Group.Timeouts.FirstByte, + Request: input.Group.Timeouts.Request, + StreamIdle: input.Group.Timeouts.StreamIdle, + }, + Credential: input.Credential, + }) + if err := spec.Validate(); err != nil { + return execution.AttemptSpec{}, err + } + return spec, nil +} + +func upstreamFromExecutionResult( + ctx context.Context, + input ForwardInput, + result execution.AttemptResult, +) UpstreamResult { + upstream := baseExecutionResult(input, result.DispatchState, result.ResponseStarted, + result.StatusCode, result.Header, result.Model, result.UpstreamRequestID, + result.Usage, result.Error) + upstream.Body = append([]byte(nil), result.Body...) + upstream.ClassificationBody = append([]byte(nil), result.Body...) + if !result.ResponseStarted && result.Error != nil { + upstream.Err = executionFailureError(ctx, result.Error) + } + return upstream +} + +func upstreamFromExecutionStreamResult( + ctx context.Context, + input ForwardInput, + result execution.StreamResult, + streamUsage *execution.UsageEvidence, +) UpstreamResult { + usageEvidence := result.Usage + if usageEvidence == nil { + usageEvidence = streamUsage + } + upstream := baseExecutionResult(input, result.DispatchState, result.ResponseStarted, + result.StatusCode, result.Header, result.Model, result.UpstreamRequestID, + usageEvidence, result.Error) + if !result.ResponseStarted && result.Error != nil { + upstream.Err = executionFailureError(ctx, result.Error) + } + return upstream +} + +func baseExecutionResult( + input ForwardInput, + dispatchState execution.DispatchState, + responseStarted bool, + statusCode int, + header http.Header, + model string, + upstreamRequestID string, + usageEvidence *execution.UsageEvidence, + errorEvidence *execution.ErrorEvidence, +) UpstreamResult { + result := UpstreamResult{ + StatusCode: statusCode, + Header: cloneEndToEndHeaders(header), + UpstreamReportedModel: model, + ResponseModelObserved: model != "", + RequestWritten: dispatchState == execution.DispatchMaybeSent, + DispatchState: dispatchState, + ResponseStarted: responseStarted, + UpstreamRequestID: upstreamRequestID, + } + if usageEvidence != nil { + result.Usage = usageEvidence.Normalized + } else if input.ObserveUsage { + result.Usage = usage.Result{State: usage.StateMissing} + } else { + result.Usage = usage.Result{State: usage.StateNotApplicable} + } + if errorEvidence != nil { + copy := errorEvidence.Clone() + result.ExecutionError = © + result.ErrorSummary = copy.Summary + result.ProviderErrorBeforeCommit = responseStarted && + statusCode >= http.StatusOK && statusCode < http.StatusMultipleChoices + } + return result +} + +func executionFailureError(ctx context.Context, evidence *execution.ErrorEvidence) error { + if ctx != nil && ctx.Err() != nil { + return ctx.Err() + } + if evidence == nil { + return fmt.Errorf("upstream execution failed") + } + switch evidence.Kind { + case execution.ErrorKindCanceled: + return context.Canceled + case execution.ErrorKindTimeout: + return context.DeadlineExceeded + case execution.ErrorKindInvalidRequest, execution.ErrorKindProvider, execution.ErrorKindInternal: + return fmt.Errorf("%w: upstream execution failed", ErrUpstreamProtocol) + default: + return fmt.Errorf("upstream execution failed") + } +} + +func executionInputFailure(cause error) UpstreamResult { + if cause == nil { + cause = fmt.Errorf("execution runtime is unavailable") + } + evidence := execution.ErrorEvidence{ + Kind: execution.ErrorKindInvalidRequest, + Summary: "invalid execution attempt", + } + return UpstreamResult{ + Err: fmt.Errorf("%w: %v", ErrUpstreamProtocol, cause), + DispatchState: execution.DispatchNotSent, + ExecutionError: &evidence, + ErrorSummary: evidence.Summary, + } +} + +func appendExecutionErrorBody(current, chunk []byte) []byte { + remaining := maxStreamingErrorBodyBytes - len(current) + if remaining <= 0 { + return current + } + if len(chunk) > remaining { + chunk = chunk[:remaining] + } + return append(current, chunk...) +} + +func executionStreamObservation( + ctx context.Context, + terminal execution.StreamResult, + downstreamErr error, +) StreamObservation { + if downstreamErr != nil { + return prioritizeStreamObservation(ctx, downstreamErr, StreamObservation{}) + } + if terminal.Error == nil { + return StreamObservation{EndReason: StreamEndCleanEOF} + } + switch terminal.Error.Kind { + case execution.ErrorKindCanceled: + return streamTerminalObservation(StreamEndClientCanceled) + case execution.ErrorKindTimeout: + return streamTerminalObservation(StreamEndIdleTimeout) + case execution.ErrorKindProvider, execution.ErrorKindHTTP: + return StreamObservation{ + EndReason: StreamEndSSEError, + ErrorSummary: terminal.Error.Summary, + } + case execution.ErrorKindInvalidRequest, execution.ErrorKindInternal: + return streamTerminalObservation(StreamEndUpstreamProtocolError) + default: + return streamTerminalObservation(StreamEndUpstreamTerminated) + } +} + +var _ AttemptForwarder = (*ExecutionForwarder)(nil) diff --git a/internal/gateway/execution_forward_test.go b/internal/gateway/execution_forward_test.go new file mode 100644 index 00000000..27a21c26 --- /dev/null +++ b/internal/gateway/execution_forward_test.go @@ -0,0 +1,278 @@ +package gateway + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "reflect" + "testing" + "time" + + "gpt-load/internal/dialect" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" + "gpt-load/internal/state" + "gpt-load/internal/usage" +) + +type fakeExecutionExecutor struct { + unary func(context.Context, execution.AttemptSpec) execution.AttemptResult + stream func(context.Context, execution.AttemptSpec, execution.StreamSink) execution.StreamResult +} + +func (fakeExecutionExecutor) Capabilities() execution.CapabilitySet { + return execution.CapabilitySet{} +} + +func (value fakeExecutionExecutor) Execute( + ctx context.Context, + spec execution.AttemptSpec, +) execution.AttemptResult { + return value.unary(ctx, spec) +} + +func (value fakeExecutionExecutor) ExecuteStream( + ctx context.Context, + spec execution.AttemptSpec, + sink execution.StreamSink, +) execution.StreamResult { + return value.stream(ctx, spec, sink) +} + +func TestExecutionForwarderBuildsFrozenAttemptAndMapsUnaryResult(t *testing.T) { + t.Parallel() + + wantUsage := usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{UncachedInput: 7, Output: 3}} + executor := fakeExecutionExecutor{unary: func( + _ context.Context, + spec execution.AttemptSpec, + ) execution.AttemptResult { + if spec.RequestID != "request-1" || spec.AttemptID != "attempt-1" || spec.Sequence != 2 || + spec.ChannelID != "openai" || spec.ClientProtocol != protocol.OpenAICompletions || + spec.Operation != execution.OperationChatCompletion || spec.ClientModel != "public" || + spec.UpstreamModel != "upstream" || spec.Method != http.MethodPost || + spec.Path != "/v1/chat/completions" || spec.RawQuery != "trace=kept" || + string(spec.Body) != `{"model":"public"}` || spec.Credential.ID != 7 { + t.Fatalf("attempt spec = %#v", spec) + } + if spec.Header.Get("Authorization") != "" || spec.Header.Get("X-Remove") != "" || + spec.Header.Get("X-Template") != "Bearer secret" || spec.Header.Get("X-Test") != "kept" { + t.Fatalf("attempt headers = %#v", spec.Header) + } + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, + ResponseStarted: true, + StatusCode: http.StatusOK, + Header: http.Header{"X-Request-Id": {"upstream-1"}}, + Body: []byte(`{"ok":true}`), + Model: "upstream", + UpstreamRequestID: "upstream-1", + Usage: &execution.UsageEvidence{Normalized: wantUsage}, + } + }} + + result := NewExecutionForwarder(executor).Forward(context.Background(), executionForwardInput()) + if result.Err != nil || result.StatusCode != http.StatusOK || + string(result.Body) != `{"ok":true}` || !reflect.DeepEqual(result.Usage, wantUsage) || + result.DispatchState != execution.DispatchMaybeSent || !result.ResponseStarted || + result.UpstreamRequestID != "upstream-1" || result.UpstreamReportedModel != "upstream" { + t.Fatalf("Forward() = %#v", result) + } +} + +func TestExecutionForwarderKeepsHTTPFailureAsUncommittedResponse(t *testing.T) { + t.Parallel() + + evidence := execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, StatusCode: http.StatusTooManyRequests, + Summary: "upstream rejected request", RetryAfter: 3 * time.Second, + } + executor := fakeExecutionExecutor{unary: func(context.Context, execution.AttemptSpec) execution.AttemptResult { + return execution.AttemptResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": {"3"}}, + Body: []byte(`{"error":{"type":"rate_limit"}}`), + Error: &evidence, + } + }} + + result := NewExecutionForwarder(executor).Forward(context.Background(), executionForwardInput()) + if result.Err != nil || !result.HasResponse() || result.ExecutionError == nil || + result.ExecutionError.Kind != execution.ErrorKindHTTP || result.ErrorSummary != evidence.Summary { + t.Fatalf("Forward() = %#v", result) + } + result.ExecutionError.Summary = "changed" + if evidence.Summary != "upstream rejected request" { + t.Fatal("Forward() retained executor-owned error evidence") + } +} + +func TestExecutionForwarderCommitsSuccessfulStreamOnlyOnFirstData(t *testing.T) { + t.Parallel() + + readyObserved := false + firstResponseObserved := false + executor := fakeExecutionExecutor{stream: func( + _ context.Context, + _ execution.AttemptSpec, + sink execution.StreamSink, + ) execution.StreamResult { + if err := sink(execution.StreamEvent{ + Sequence: 1, Kind: execution.StreamEventReady, StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}, "X-Request-Id": {"stream-1"}}, + UpstreamRequestID: "stream-1", + }); err != nil { + t.Fatalf("ready sink: %v", err) + } + if readyObserved || firstResponseObserved { + t.Fatal("ready metadata committed downstream before first data") + } + if err := sink(execution.StreamEvent{Sequence: 2, Kind: execution.StreamEventData, Data: []byte("data: one\n\n")}); err != nil { + t.Fatalf("data sink: %v", err) + } + if !readyObserved || !firstResponseObserved { + t.Fatal("first data did not cross commit boundary") + } + if err := sink(execution.StreamEvent{Sequence: 3, Kind: execution.StreamEventData, Data: []byte("data: [DONE]\n\n")}); err != nil { + t.Fatalf("done sink: %v", err) + } + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, + UpstreamRequestID: "stream-1", + } + }} + input := executionForwardInput() + input.OnStreamReady = func() { readyObserved = true } + input.OnFirstResponse = func() { firstResponseObserved = true } + recorder := httptest.NewRecorder() + result := NewExecutionForwarder(executor).ForwardStream(context.Background(), input, recorder) + if result.Err != nil || !result.Committed || result.Stream.EndReason != StreamEndCleanEOF || + recorder.Code != http.StatusOK || recorder.Body.String() != "data: one\n\ndata: [DONE]\n\n" { + t.Fatalf("ForwardStream() result=%#v response=%d %q", result, recorder.Code, recorder.Body.String()) + } +} + +func TestExecutionForwarderBuffersRejectedStreamWithoutCommitting(t *testing.T) { + t.Parallel() + + evidence := execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, StatusCode: http.StatusUnauthorized, + Summary: "upstream rejected request", + } + executor := fakeExecutionExecutor{stream: func( + _ context.Context, + _ execution.AttemptSpec, + sink execution.StreamSink, + ) execution.StreamResult { + if err := sink(execution.StreamEvent{ + Sequence: 1, Kind: execution.StreamEventReady, + StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": {"application/json"}}, + }); err != nil { + t.Fatalf("ready sink: %v", err) + } + if err := sink(execution.StreamEvent{Sequence: 2, Kind: execution.StreamEventData, Data: []byte(`{"error":"invalid"}`)}); err != nil { + t.Fatalf("data sink: %v", err) + } + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": {"application/json"}}, + Error: &evidence, + } + }} + + recorder := httptest.NewRecorder() + result := NewExecutionForwarder(executor).ForwardStream( + context.Background(), executionForwardInput(), recorder, + ) + if result.Committed || recorder.Code != http.StatusOK || recorder.Body.Len() != 0 || + result.StatusCode != http.StatusUnauthorized || string(result.Body) != `{"error":"invalid"}` || + result.ExecutionError == nil { + t.Fatalf("ForwardStream() result=%#v response=%d %q", result, recorder.Code, recorder.Body.String()) + } +} + +func TestExecutionForwarderDoesNotRetryOrCommitAfterDownstreamFailure(t *testing.T) { + t.Parallel() + + downstreamErr := errors.New("downstream closed") + executor := fakeExecutionExecutor{stream: func( + _ context.Context, + _ execution.AttemptSpec, + sink execution.StreamSink, + ) execution.StreamResult { + _ = sink(execution.StreamEvent{ + Sequence: 1, Kind: execution.StreamEventReady, StatusCode: http.StatusOK, + Header: make(http.Header), + }) + if err := sink(execution.StreamEvent{Sequence: 2, Kind: execution.StreamEventData, Data: []byte("data: one\n\n")}); err == nil { + t.Fatal("data sink unexpectedly succeeded") + } + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusOK, Header: make(http.Header), + Error: &execution.ErrorEvidence{Kind: execution.ErrorKindCanceled, Summary: "stream sink stopped"}, + } + }} + writer := &failingExecutionResponseWriter{header: make(http.Header), err: downstreamErr} + result := NewExecutionForwarder(executor).ForwardStream( + context.Background(), executionForwardInput(), writer, + ) + if !result.Committed || !errors.Is(result.Err, downstreamErr) || result.Stream.EndReason != StreamEndDownstreamWriteFailure { + t.Fatalf("ForwardStream() = %#v", result) + } +} + +type failingExecutionResponseWriter struct { + header http.Header + err error +} + +func (writer *failingExecutionResponseWriter) Header() http.Header { return writer.header } +func (*failingExecutionResponseWriter) WriteHeader(int) {} +func (writer *failingExecutionResponseWriter) Write([]byte) (int, error) { + return 0, writer.err +} +func (*failingExecutionResponseWriter) FlushError() error { return nil } + +func executionForwardInput() ForwardInput { + return ForwardInput{ + Dialect: dialect.NewOpenAI(), + RequestID: "request-1", + AttemptID: "attempt-1", + AttemptSequence: 2, + ClientProtocol: protocol.OpenAICompletions, + Operation: execution.OperationChatCompletion, + ChannelID: "openai", + TargetConfig: []byte(`{}`), + APIKey: "secret", + Credential: execution.NewCredentialSnapshot( + 7, 2, 3, []byte(`{"api_key":"secret"}`), + ), + Request: &dialect.ParsedRequest{ + Method: http.MethodPost, + Path: "/v1/chat/completions", + RawQuery: "trace=kept", + Header: http.Header{ + "Authorization": {"Bearer downstream"}, + "X-Remove": {"remove"}, + "X-Test": {"kept"}, + }, + Body: []byte(`{"model":"public"}`), + }, + ExternalModel: "public", + UpstreamModelID: "upstream", + Group: state.GroupView{ + Timeouts: state.TimeoutConfig{ + Connect: time.Second, FirstByte: time.Second, + Request: time.Second, StreamIdle: time.Second, + }, + HeaderRules: state.HeaderRules{ + Set: map[string]string{"X-Template": "Bearer ${API_KEY}"}, + Remove: []string{"X-Remove"}, + }, + }, + } +} diff --git a/internal/gateway/forward.go b/internal/gateway/forward.go index 2c986c59..df27d6b0 100644 --- a/internal/gateway/forward.go +++ b/internal/gateway/forward.go @@ -6,40 +6,47 @@ import ( "encoding/json" "errors" "fmt" - "io" - "math" "mime" "net" "net/http" - "net/http/httptrace" "net/textproto" - "sort" "strings" - "sync" - "sync/atomic" - "time" "gpt-load/internal/dialect" + "gpt-load/internal/execution" "gpt-load/internal/platform/contentcoding" - platformhttp "gpt-load/internal/platform/httpclient" platformheader "gpt-load/internal/platform/httpheader" - "gpt-load/internal/platform/redact" + "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/usage" ) +// ForwardInput is the frozen logical-attempt input shared by the gateway +// orchestrator and the provider-neutral execution adapter. type ForwardInput struct { - Dialect dialect.Dialect - ObserveUsage bool - Group state.GroupView - APIKey string - Request *dialect.ParsedRequest - ExternalModel string - UpstreamModelID string - OnStreamReady func() - OnFirstResponse func() + Dialect dialect.Dialect + ObserveUsage bool + Group state.GroupView + APIKey string + CredentialSecrets []string + Request *dialect.ParsedRequest + ExternalModel string + UpstreamModelID string + OnStreamReady func() + OnFirstResponse func() + + RequestID string + AttemptID string + AttemptSequence uint32 + ClientProtocol protocol.Protocol + Operation execution.Operation + ChannelID string + TargetConfig json.RawMessage + Credential execution.CredentialSnapshot } +// UpstreamResult is the gateway's stable view of one logical execution +// attempt. SDK-internal transport recovery remains inside this result. type UpstreamResult struct { StatusCode int Header http.Header @@ -52,23 +59,19 @@ type UpstreamResult struct { Err error RequestWritten bool Committed bool - RetryableBeforeCommit bool ProviderErrorBeforeCommit bool Stream StreamObservation Usage usage.Result + DispatchState execution.DispatchState + ResponseStarted bool + UpstreamRequestID string + ExecutionError *execution.ErrorEvidence } func (result UpstreamResult) HasResponse() bool { return result.StatusCode != 0 && result.Err == nil } -type Forwarder struct { - clients *platformhttp.HTTPClientManager - redactor *redact.Redactor - streamWriteTimeout time.Duration - usageCapture *usageCaptureBoundary -} - const ( maxNonStreamingResponseBodyBytes = int64(32 << 20) maxErrorResponseBodyBytes = int64(64 << 10) @@ -78,569 +81,14 @@ const ( var ErrUpstreamProtocol = errors.New("upstream protocol error") -func NewForwarder(clients *platformhttp.HTTPClientManager, redactor *redact.Redactor) *Forwarder { - return &Forwarder{ - clients: clients, redactor: redactor, - streamWriteTimeout: downstreamWriteTimeout, - usageCapture: newUsageCaptureBoundary(), - } -} - -func (forwarder *Forwarder) Forward(ctx context.Context, input ForwardInput) UpstreamResult { - if forwarder == nil || forwarder.clients == nil || forwarder.redactor == nil || input.Dialect == nil || input.Request == nil { - return UpstreamResult{Err: fmt.Errorf("forward input is incomplete")} - } - request, wroteRequest, _, err := forwarder.newUpstreamRequest(ctx, input, false) - if err != nil { - return UpstreamResult{Err: err} - } - knownSecrets := resolvedCredentialSecretValues(input, request.Header) - summarySecrets := resolvedErrorSummarySecretValues( - input.APIKey, - input.Group.HeaderRules, - knownSecrets..., - ) - response, err := forwarder.clients.GetClient(nonStreamingClientConfig(input.Group.Timeouts)).Do(request) - if err != nil { - return UpstreamResult{Err: fmt.Errorf("perform upstream request: %w", err), RequestWritten: wroteRequest.Load()} - } - defer response.Body.Close() - - representationHeaders := response.Header.Clone() - headers := cloneEndToEndHeaders(response.Header) - policy := classifyResponseBody(input.Request.Method, response.StatusCode) - if !policy.readBody { - encoding, parseErr := contentcoding.ParseContentEncoding( - headerFieldValues(representationHeaders, "Content-Encoding"), - ) - if parseErr != nil || encoding != contentcoding.Identity { - return UpstreamResult{ - Err: fmt.Errorf("%w: bodyless response has non-identity Content-Encoding", ErrUpstreamProtocol), - RequestWritten: true, - } - } - if response.StatusCode == http.StatusPartialContent && - !hasSafeHeadContentRange(representationHeaders) { - return UpstreamResult{ - Err: fmt.Errorf("%w: invalid HEAD partial Content-Range", ErrUpstreamProtocol), - RequestWritten: true, - } - } - safeHeaders := sanitizeForwardResponseHeaders(representationHeaders, input, knownSecrets...) - safeHeaders, _ = normalizeBufferedResponse( - input.Request.Method, - response.StatusCode, - safeHeaders, - nil, - ) - return UpstreamResult{ - StatusCode: response.StatusCode, - Header: safeHeaders, - RequestWritten: true, - Usage: usage.Result{State: usage.StateNotApplicable}, - } - } - - success := response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices - limit := maxErrorResponseBodyBytes - if success { - limit = maxNonStreamingResponseBodyBytes - } - overflow := response.ContentLength > limit - var body []byte - if !overflow { - body, overflow, err = readBodyAtMost(response.Body, limit) - } - if err != nil { - return UpstreamResult{Err: fmt.Errorf("read upstream response: %w", err), RequestWritten: true} - } - if overflow && success { - return UpstreamResult{ - Err: fmt.Errorf("%w: non-streaming response body exceeds limit", ErrUpstreamProtocol), - RequestWritten: true, - } - } - - result := UpstreamResult{ - StatusCode: response.StatusCode, - Body: body, - RequestWritten: true, - Usage: usage.Result{State: usage.StateNotApplicable}, - } - if success { - prepared, prepareErr := forwarder.prepareSuccessRepresentation( - input, - response.StatusCode, - representationHeaders, - body, - knownSecrets, - ) - if prepareErr != nil { - return UpstreamResult{Err: prepareErr, RequestWritten: true} - } - result.Body = prepared.downstream - result.Header = prepared.headers - applyResponseModelObservation(&result, prepared.modelObservation) - if input.ObserveUsage && forwarder.usageCapture != nil { - result.Usage = forwarder.usageCapture.extractNonStreamingPlain( - input.Dialect, - prepared.inspectable, - ) - } - return result - } - if !success { - prepared := forwarder.prepareErrorRepresentation( - input, - representationHeaders, - body, - knownSecrets, - ) - if overflow { - prepared = forwarder.failClosedErrorRepresentation(input, headers, knownSecrets) - } - result.Body = prepared.downstream - result.ClassificationBody = prepared.classification - result.ErrorSummary = summarizeErrorBody( - forwarder.redactor, prepared.classification, "", summarySecrets..., - ) - result.Header = prepared.headers - return result - } - return result -} - -func (forwarder *Forwarder) ForwardStream( - ctx context.Context, - input ForwardInput, - downstream http.ResponseWriter, -) (result UpstreamResult) { - if forwarder == nil || forwarder.clients == nil || forwarder.redactor == nil || - input.Dialect == nil || input.Request == nil || downstream == nil { - return UpstreamResult{Err: fmt.Errorf("stream forward input is incomplete")} - } - - deadline := newFirstEventDeadline(ctx, input.Group.Timeouts.FirstByte) - defer deadline.stop() - request, wroteRequest, replay, err := forwarder.newUpstreamRequest(deadline.ctx, input, true) - if err != nil { - return UpstreamResult{Err: err} - } - knownSecrets := resolvedCredentialSecretValues(input, request.Header) - summarySecrets := resolvedErrorSummarySecretValues( - input.APIKey, - input.Group.HeaderRules, - knownSecrets..., - ) - response, err := forwarder.clients.GetClient(streamingClientConfig(input.Group.Timeouts)).Do(request) - if err != nil { - requestWritten := wroteRequest.Load() - return UpstreamResult{ - Err: streamAttemptError(ctx, deadline.ctx, fmt.Errorf("perform upstream stream request: %w", err)), - RequestWritten: requestWritten, RetryableBeforeCommit: retryableBeforeCommit(ctx, requestWritten), - } - } - streamBody := response.Body - defer func() { _ = streamBody.Close() }() - - representationHeaders := response.Header.Clone() - headers := cloneEndToEndHeaders(response.Header) - result = UpstreamResult{ - StatusCode: response.StatusCode, - RequestWritten: true, - } - if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { - body, overflow, readErr := readStreamingErrorBody(streamBody) - if readErr != nil { - result.Err = streamAttemptError(ctx, deadline.ctx, fmt.Errorf("read upstream stream error response: %w", readErr)) - result.RetryableBeforeCommit = retryableBeforeCommit(ctx, result.RequestWritten) - return result - } - prepared := forwarder.prepareErrorRepresentation(input, representationHeaders, body, knownSecrets) - if overflow { - prepared = forwarder.failClosedErrorRepresentation(input, headers, knownSecrets) - } - result.Body = prepared.downstream - result.ClassificationBody = prepared.classification - result.ErrorSummary = summarizeErrorBody( - forwarder.redactor, - result.ClassificationBody, - "", - summarySecrets..., - ) - result.Header = prepared.headers - return result - } - streamEvents := newStreamEventObserver( - input.Dialect, - forwarder.usageCapture.newStreamForRequest( - input.Dialect, - input.ObserveUsage, - ), - ) - modelTracker := newResponseModelTracker(input.Dialect, input.UpstreamModelID) - defer func() { - result.Usage = streamEvents.finalizeUsage() - if result.Stream.EndReason == StreamEndCleanEOF { - applyResponseModelObservation(&result, modelTracker.observation()) - } - }() - - if err := validateStreamContentEncoding(representationHeaders); err != nil { - result.Err = err - result.RetryableBeforeCommit = retryableBeforeCommit(ctx, result.RequestWritten) - return result - } - - rewriteModel := needsModelRewrite(input) - var rewriter dialect.ModelRewriter - if rewriteModel { - var ok bool - rewriter, ok = input.Dialect.(dialect.ModelRewriter) - if !ok { - result.Err = fmt.Errorf("%w: dialect does not support model rewrite", ErrUpstreamProtocol) - result.RetryableBeforeCommit = retryableBeforeCommit(ctx, result.RequestWritten) - return result - } - } - firstPayloadPending := true - rewrittenStream := newSSEEventRewriteStream(streamBody, func( - event dialect.StreamEvent, - errorEvent bool, - ) (sseEventRewriteResult, error) { - firstPayload := firstPayloadPending - firstPayloadPending = false - modelTracker.observe(event.Payload) - safePayload := event.Payload - for _, secret := range knownSecrets { - var ok bool - safePayload, ok = rewriteBoundedLiteral( - safePayload, - secret, - redact.Placeholder, - int64(maxSSEEventBytes), - ) - if !ok { - return sseEventRewriteResult{}, fmt.Errorf("%w: redact upstream SSE credential", ErrUpstreamProtocol) - } - } - safeEvent := dialect.StreamEvent{ - Name: event.Name, - Payload: safePayload, - } - wasTerminal := streamEvents.sawTerminal - providerError, err := streamEvents.classify( - safeEvent, - errorEvent, - ) - if err != nil { - return sseEventRewriteResult{}, err - } - terminal := !wasTerminal && streamEvents.sawTerminal - if !firstPayload { - streamEvents.observeUsageEvent(safeEvent) - } - if providerError { - observationPayload := forwarder.redactor.Bytes(safePayload) - streamEvents.observeError( - summarizeErrorBody( - forwarder.redactor, - observationPayload, - fixedErrorSummary("upstream_sse_error"), - summarySecrets..., - ), - ) - } - if !rewriteModel { - return sseEventRewriteResult{body: safePayload, terminal: terminal}, nil - } - if providerError { - var ok bool - safePayload, ok = rewriteBoundedLiteral( - safePayload, - input.UpstreamModelID, - input.ExternalModel, - int64(maxSSEEventBytes), - ) - if !ok { - return sseEventRewriteResult{}, fmt.Errorf("%w: rewrite upstream SSE model literal", ErrUpstreamProtocol) - } - return sseEventRewriteResult{body: safePayload, terminal: terminal}, nil - } - rewritten, err := rewriter.RewriteResponseModel(safePayload, input.ExternalModel) - if err != nil { - return sseEventRewriteResult{}, fmt.Errorf("%w: rewrite upstream response model: %v", ErrUpstreamProtocol, err) - } - return sseEventRewriteResult{body: rewritten, terminal: terminal}, nil - }) - streamBody = rewrittenStream - invalidateRewrittenBodyHeaders(headers) - - firstEvent, err := bufferFirstSSEEvent(streamBody) - if err != nil { - if !rewriteModel && errors.Is(err, errSSEEventTooLarge) { - err = errFirstSSEEventTooLarge - } - if errors.Is(err, errFirstSSEEventTooLarge) || errors.Is(err, errSSEEventTooLarge) { - err = fmt.Errorf("%w: %w", ErrUpstreamProtocol, err) - } - result.Err = streamAttemptError(ctx, deadline.ctx, err) - result.RetryableBeforeCommit = retryableBeforeCommit(ctx, result.RequestWritten) - return result - } - if ctxErr := ctx.Err(); ctxErr != nil { - result.Err = ctxErr - return result - } - - if firstEvent.IsProviderError || streamEvents.firstProviderError { - streamEvents.observeUsageEvent(dialect.StreamEvent{ - Name: firstEvent.Name, - Payload: firstEvent.Payload, - }) - result.ClassificationBody = forwarder.safeProviderErrorPayload( - firstEvent.Payload, - knownSecrets, - ) - result.ErrorSummary = fixedErrorSummary("upstream_sse_error") - result.Header = providerErrorRateLimitHeaders( - sanitizeForwardResponseHeaders(headers, input, knownSecrets...), - ) - result.ProviderErrorBeforeCommit = true - result.Usage = usage.Result{State: usage.StateMissing} - return result - } - if !deadline.disarm() { - result.Err = streamAttemptError(ctx, deadline.ctx, context.DeadlineExceeded) - result.RetryableBeforeCommit = retryableBeforeCommit(ctx, result.RequestWritten) - return result - } - streamEvents.observeUsageEvent(dialect.StreamEvent{ - Name: firstEvent.Name, - Payload: firstEvent.Payload, - }) - - if input.OnStreamReady != nil { - input.OnStreamReady() - } - - result.Header = normalizeStreamResponseHeaders( - sanitizeForwardResponseHeaders(headers, input, knownSecrets...), - ) - streamWriter := newStreamWriteController(downstream, forwarder.streamWriteTimeout) - defer func() { _ = streamWriter.clear() }() - - result.Committed = true - releaseCommittedRequestReplay(input.Request, replay) - if err := commitStream(streamWriter, response.StatusCode, result.Header, firstEvent.Prefix); err != nil { - result.Err = err - result.Stream = observeStreamTermination(ctx, err, streamEvents) - return result - } - if input.OnFirstResponse != nil { - input.OnFirstResponse() - } - if rewrittenStream.consumeReadTerminalBoundary() { - streamEvents.markTerminalForwarded() - } - if err := pumpStreamWithFlushObserver( - deadline.ctx, - streamBody, - streamWriter, - input.Group.Timeouts.StreamIdle, - func() { - if rewrittenStream.consumeReadTerminalBoundary() { - streamEvents.markTerminalForwarded() - } - }, - ); err != nil { - result.Err = err - } - if result.Err == nil { - result.Err = streamEvents.validateEOF() - if result.Err == nil && rewrittenStream.consumeReadTerminalBoundary() { - streamEvents.markTerminalForwarded() - } - } - result.Stream = observeStreamTermination(ctx, result.Err, streamEvents) - return result -} - -func retryableBeforeCommit(parent context.Context, requestWritten bool) bool { - return !requestWritten && parent != nil && parent.Err() == nil -} - -func providerErrorRateLimitHeaders(source http.Header) http.Header { - filtered := make(http.Header) - for name, values := range source { - lowered := strings.ToLower(name) - // Keep this allowlist synchronized with health.ParseRateLimitReset. - if lowered != "retry-after" && - !(strings.HasPrefix(lowered, "anthropic-ratelimit-") && - strings.HasSuffix(lowered, "-reset")) && - lowered != "x-ratelimit-reset" && - !strings.HasPrefix(lowered, "x-ratelimit-reset-") { - continue - } - filtered[name] = append([]string(nil), values...) - } - return filtered -} - -func (forwarder *Forwarder) safeProviderErrorPayload( - payload []byte, - knownSecrets []string, -) []byte { - safe := bytes.Clone(payload) - for _, secret := range knownSecrets { - var ok bool - safe, ok = rewriteBoundedLiteral( - safe, - secret, - redact.Placeholder, - int64(maxFirstSSEEventBytes), - ) - if !ok { - return []byte(redact.Placeholder) - } - } - safe = forwarder.redactor.Bytes(safe) - if len(safe) > maxFirstSSEEventBytes { - return []byte(redact.Placeholder) - } - return safe -} - -func releaseCommittedRequestReplay(parsed *dialect.ParsedRequest, replay *requestReplay) { - if parsed != nil { - parsed.Body = nil - } - if replay != nil { - replay.release() - } -} - -func (forwarder *Forwarder) newUpstreamRequest( - ctx context.Context, - input ForwardInput, - stream bool, -) (*http.Request, *atomic.Bool, *requestReplay, error) { - parsed := input.Request - rewrite := needsModelRewrite(input) - if rewrite { - rewriter, ok := input.Dialect.(dialect.ModelRewriter) - if !ok { - return nil, nil, nil, fmt.Errorf("dialect does not support model rewriting") - } - derived, err := rewriter.RewriteRequestModel(input.Request, input.UpstreamModelID) - if err != nil { - return nil, nil, nil, fmt.Errorf("rewrite upstream request model: %w", err) - } - if derived == nil { - return nil, nil, nil, fmt.Errorf("rewrite upstream request model returned nil request") - } - if int64(len(derived.Body)) > maxRequestBodyBytes { - return nil, nil, nil, fmt.Errorf("%w: rewritten request body exceeds limit", errRequestTooLarge) - } - parsed = derived - } - if stream && input.Group.InjectUsageOptions && forwarder != nil && forwarder.usageCapture != nil { - parsed = forwarder.usageCapture.injectStreamUsage(input.Dialect, parsed) - } - upstreamURL, err := input.Dialect.BuildUpstreamURL(input.Group.UpstreamURL, parsed) - if err != nil { - return nil, nil, nil, fmt.Errorf("build upstream URL: %w", err) - } - replay := newRequestReplay(parsed.Body) - request, err := http.NewRequestWithContext(ctx, parsed.Method, upstreamURL, nil) - if err != nil { - return nil, nil, nil, fmt.Errorf("create upstream request: %w", err) - } - request.Body = replay.open() - request.ContentLength = int64(len(parsed.Body)) - request.GetBody = nil - request.Header = cloneEndToEndHeaders(parsed.Header) - removeDownstreamCredentials(request.Header) - dialect.ApplyCredential(input.Dialect, request.Header, input.APIKey, input.Group.HeaderRules) - sanitizeUpstreamRequestHeaders(request.Header) - platformheader.NormalizeUpstreamRequestRepresentation(request, int64(len(parsed.Body))) - if _, exists := request.Header["User-Agent"]; !exists { - request.Header["User-Agent"] = nil - } - - wroteRequest := &atomic.Bool{} - trace := &httptrace.ClientTrace{WroteRequest: func(httptrace.WroteRequestInfo) { wroteRequest.Store(true) }} - request = request.WithContext(httptrace.WithClientTrace(request.Context(), trace)) - return request, wroteRequest, replay, nil -} - func needsModelRewrite(input ForwardInput) bool { return input.ExternalModel != "" && input.UpstreamModelID != "" && input.ExternalModel != input.UpstreamModelID } -func isResolvedCredentialHeaderName(selected dialect.Dialect, name string) bool { - if platformheader.IsCredentialName(name) { - return true - } - namer, ok := selected.(dialect.CredentialHeaderNamer) - if !ok { - return false - } - identity := canonicalHeaderIdentity(name) - if identity == "" { - return false - } - for _, dialectName := range namer.CredentialHeaderNames() { - if identity == canonicalHeaderIdentity(dialectName) { - return true - } - } - return false -} - -func resolvedCredentialSecretValues(input ForwardInput, finalHeaders http.Header) []string { - if finalHeaders == nil { - finalHeaders = make(http.Header) - if input.Request != nil { - finalHeaders = cloneEndToEndHeaders(input.Request.Header) - } - removeDownstreamCredentials(finalHeaders) - dialect.ApplyCredential( - input.Dialect, - finalHeaders, - input.APIKey, - input.Group.HeaderRules, - ) - sanitizeUpstreamRequestHeaders(finalHeaders) - } - - secrets := make([]string, 0, 8) - seen := make(map[string]struct{}) - appendSecret := func(value string) { - if value == "" { - return - } - if _, exists := seen[value]; exists { - return - } - seen[value] = struct{}{} - secrets = append(secrets, value) - } - appendSecret(input.APIKey) - for name, values := range finalHeaders { - if !isResolvedCredentialHeaderName(input.Dialect, name) { - continue - } - for _, value := range values { - appendSecret(value) - } - } - sort.SliceStable(secrets, func(left, right int) bool { - return len(secrets[left]) > len(secrets[right]) - }) - return secrets +func isResolvedCredentialHeaderName(name string) bool { + return platformheader.IsCredentialName(name) } func rewriteBoundedLiteral(body []byte, literal, replacement string, limit int64) ([]byte, bool) { @@ -648,9 +96,7 @@ func rewriteBoundedLiteral(body []byte, literal, replacement string, limit int64 return body, int64(len(body)) <= limit } if !json.Valid(body) { - return replaceAllBounded( - body, []byte(literal), []byte(replacement), limit, - ) + return replaceAllBounded(body, []byte(literal), []byte(replacement), limit) } decoder := json.NewDecoder(bytes.NewReader(body)) @@ -683,10 +129,7 @@ func rewriteJSONLiteralStrings( ) (any, bool, bool) { switch typed := value.(type) { case string: - rewritten, changed, ok := replaceStringBounded( - typed, literal, replacement, limit, budget, - ) - return rewritten, changed, ok + return replaceStringBounded(typed, literal, replacement, limit, budget) case []any: changed := false for index, item := range typed { @@ -809,42 +252,6 @@ func invalidateRewrittenBodyHeaders(headers http.Header) { } } -func nonStreamingClientConfig(timeouts state.TimeoutConfig) *platformhttp.Config { - return &platformhttp.Config{ - ConnectTimeout: timeouts.Connect, - RequestTimeout: timeouts.Request, - IdleConnTimeout: 90 * time.Second, - MaxIdleConns: 100, - MaxIdleConnsPerHost: 20, - ResponseHeaderTimeout: timeouts.FirstByte, - DisableCompression: true, - WriteBufferSize: 32 * 1024, - ReadBufferSize: 32 * 1024, - ForceAttemptHTTP2: true, - TLSHandshakeTimeout: timeouts.Connect, - ExpectContinueTimeout: time.Second, - DisableRedirects: true, - } -} - -func streamingClientConfig(timeouts state.TimeoutConfig) *platformhttp.Config { - return &platformhttp.Config{ - ConnectTimeout: timeouts.Connect, - RequestTimeout: 0, - IdleConnTimeout: 90 * time.Second, - MaxIdleConns: 100, - MaxIdleConnsPerHost: 20, - ResponseHeaderTimeout: timeouts.FirstByte, - DisableCompression: true, - WriteBufferSize: 32 * 1024, - ReadBufferSize: 32 * 1024, - ForceAttemptHTTP2: true, - TLSHandshakeTimeout: timeouts.Connect, - ExpectContinueTimeout: time.Second, - DisableRedirects: true, - } -} - func validateStreamContentEncoding(headers http.Header) error { encoding, err := contentcoding.ParseContentEncoding( headerFieldValues(headers, "Content-Encoding"), @@ -855,87 +262,6 @@ func validateStreamContentEncoding(headers http.Header) error { return nil } -func readStreamingErrorBody(body io.Reader) ([]byte, bool, error) { - return readBodyAtMost(body, maxErrorResponseBodyBytes) -} - -func readBodyAtMost(reader io.Reader, limit int64) ([]byte, bool, error) { - if reader == nil || limit < 0 { - return nil, false, fmt.Errorf("response reader/limit is invalid") - } - var limited io.Reader = reader - if limit < math.MaxInt64 { - limited = io.LimitReader(reader, limit+1) - } - body, err := io.ReadAll(limited) - if err != nil { - return nil, false, err - } - if int64(len(body)) > limit { - return nil, true, nil - } - return body, false, nil -} - -func streamAttemptError(parent, attempt context.Context, fallback error) error { - if err := parent.Err(); err != nil { - return err - } - if cause := context.Cause(attempt); cause != nil { - return cause - } - return fallback -} - -type firstEventDeadline struct { - ctx context.Context - cancel context.CancelCauseFunc - mu sync.Mutex - timer *time.Timer - disarmed bool - expired bool -} - -func newFirstEventDeadline(parent context.Context, timeout time.Duration) *firstEventDeadline { - ctx, cancel := context.WithCancelCause(parent) - deadline := &firstEventDeadline{ctx: ctx, cancel: cancel} - deadline.timer = time.AfterFunc(timeout, deadline.expire) - return deadline -} - -func (deadline *firstEventDeadline) expire() { - deadline.mu.Lock() - defer deadline.mu.Unlock() - if deadline.disarmed { - return - } - deadline.expired = true - deadline.cancel(context.DeadlineExceeded) -} - -func (deadline *firstEventDeadline) disarm() bool { - deadline.mu.Lock() - defer deadline.mu.Unlock() - if deadline.expired { - return false - } - deadline.disarmed = true - if deadline.timer != nil { - deadline.timer.Stop() - } - return true -} - -func (deadline *firstEventDeadline) stop() { - deadline.mu.Lock() - defer deadline.mu.Unlock() - deadline.disarmed = true - if deadline.timer != nil { - deadline.timer.Stop() - } - deadline.cancel(context.Canceled) -} - var hopByHopHeaders = map[string]struct{}{ "Connection": {}, "Proxy-Connection": {}, "Keep-Alive": {}, "Proxy-Authenticate": {}, "Proxy-Authorization": {}, "Te": {}, @@ -1013,23 +339,6 @@ func headerValuesContainLiteral(values []string, literal string) bool { return false } -func sanitizeUpstreamResponseHeaders(source http.Header, apiKey string) http.Header { - headers := cloneEndToEndHeaders(source) - namesToDelete := make([]string, 0) - for actualName, values := range headers { - if platformheader.IsCredentialName(actualName) || - strings.EqualFold(actualName, "Set-Cookie") || - strings.EqualFold(actualName, "Set-Cookie2") || - headerValuesContainLiteral(values, apiKey) { - namesToDelete = append(namesToDelete, actualName) - } - } - for _, name := range namesToDelete { - deleteHeaderField(headers, name) - } - return headers -} - func sanitizeForwardResponseHeaders( source http.Header, input ForwardInput, @@ -1040,11 +349,12 @@ func sanitizeForwardResponseHeaders( deleted := endToEndHeaderCleanupDeletedField(source, headers) namesToDelete := make([]string, 0) for actualName, values := range headers { - deleteField := isResolvedCredentialHeaderName(input.Dialect, actualName) || + deleteField := isResolvedCredentialHeaderName(actualName) || + strings.HasPrefix(strings.ToLower(actualName), "x-upstream-") || strings.EqualFold(actualName, "Set-Cookie") || strings.EqualFold(actualName, "Set-Cookie2") || headerValuesContainLiteral(values, input.APIKey) - for _, secret := range additionalSecrets { + for _, secret := range append(append([]string(nil), input.CredentialSecrets...), additionalSecrets...) { if deleteField || secret == "" || secret == input.APIKey { continue } @@ -1063,8 +373,7 @@ func sanitizeForwardResponseHeaders( } if !needsModelRewrite(input) { if deleted { - deleteHeaderField(headers, "Signature") - deleteHeaderField(headers, "Signature-Input") + stripDownstreamResponseSignatures(headers) } return headers } @@ -1084,17 +393,11 @@ func sanitizeForwardResponseHeaders( } continue } - if nameContainsModel { - deleteHeaderField(headers, name) - deleted = true - continue - } deleteHeaderField(headers, name) deleted = true } if deleted { - deleteHeaderField(headers, "Signature") - deleteHeaderField(headers, "Signature-Input") + stripDownstreamResponseSignatures(headers) } return headers } diff --git a/internal/gateway/forward_test.go b/internal/gateway/forward_test.go deleted file mode 100644 index ed5de098..00000000 --- a/internal/gateway/forward_test.go +++ /dev/null @@ -1,4115 +0,0 @@ -package gateway - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "math" - mathrand "math/rand" - "net/http" - "net/http/httptest" - "reflect" - "slices" - "strconv" - "strings" - "sync" - "sync/atomic" - "testing" - "time" - - "github.com/klauspost/compress/zstd" - - "gpt-load/internal/dialect" - "gpt-load/internal/health" - "gpt-load/internal/platform/contentcoding" - platformhttp "gpt-load/internal/platform/httpclient" - platformheader "gpt-load/internal/platform/httpheader" - "gpt-load/internal/platform/redact" - "gpt-load/internal/protocol" - "gpt-load/internal/state" - "gpt-load/internal/usage" -) - -type dialectSpecificCredentialHeaders struct { - dialect.Dialect -} - -func (dialectSpecificCredentialHeaders) CredentialHeaderNames() []string { - return []string{"X-Dialect-Credential"} -} - -func TestReadBodyAtMostDoesNotReturnPartialBody(t *testing.T) { - if body, overflow, err := readBodyAtMost(strings.NewReader("1234"), 4); err != nil || overflow || string(body) != "1234" { - t.Fatalf("exact = %q, %t, %v", body, overflow, err) - } - if body, overflow, err := readBodyAtMost(strings.NewReader("12345"), 4); err != nil || !overflow || body != nil { - t.Fatalf("overflow = %q, %t, %v", body, overflow, err) - } - - wantErr := errors.New("read failed") - body, overflow, err := readBodyAtMost(io.MultiReader(strings.NewReader("12"), responseReadError{err: wantErr}), 4) - if !errors.Is(err, wantErr) || overflow || body != nil { - t.Fatalf("read error = %q, %t, %v", body, overflow, err) - } - - body, overflow, err = readBodyAtMost(strings.NewReader("x"), math.MaxInt64) - if err != nil || overflow || string(body) != "x" { - t.Fatalf("max int limit = %q, %t, %v", body, overflow, err) - } - - body, overflow, err = readBodyAtMost(strings.NewReader("x"), -1) - if err == nil || overflow || body != nil { - t.Fatalf("negative limit = %q, %t, %v", body, overflow, err) - } -} - -func TestForwardCredentialSecretsExcludeOrdinaryHeaderRuleValues(t *testing.T) { - const apiKey = "fake-provider-key-for-redaction" - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Group: state.GroupView{ - HeaderRules: state.HeaderRules{Set: map[string]string{ - "Authorization": "Bearer ${API_KEY}", - "Accept-Encoding": "gzip", - "X-Custom": "ordinary", - }}, - }, - APIKey: apiKey, - Request: &dialect.ParsedRequest{ - Header: make(http.Header), - }, - } - - got := resolvedCredentialSecretValues(input, nil) - want := []string{"Bearer " + apiKey, apiKey} - if !reflect.DeepEqual(got, want) { - t.Fatalf("resolvedCredentialSecretValues() = %#v, want %#v", got, want) - } - got[0] = "mutated" - if next := resolvedCredentialSecretValues(input, nil); !reflect.DeepEqual(next, want) { - t.Fatalf("resolvedCredentialSecretValues() returned aliased values: %#v", next) - } - - for _, selected := range []dialect.Dialect{ - dialect.NewOpenAI(http.DefaultClient), - dialect.NewAnthropic(http.DefaultClient), - dialect.NewGemini(http.DefaultClient), - } { - namer, ok := selected.(dialect.CredentialHeaderNamer) - if !ok { - t.Fatalf("%T does not implement CredentialHeaderNamer", selected) - } - names := namer.CredentialHeaderNames() - if len(names) == 0 { - t.Fatalf("%T returned no credential Header names", selected) - } - original := names[0] - names[0] = "mutated" - if next := namer.CredentialHeaderNames(); len(next) == 0 || next[0] != original { - t.Fatalf("%T returned an aliased credential Header slice: %#v", selected, next) - } - } -} - -func TestGatewayCredentialPolicyRecognizesSharedBaseAndDialectSpecificNames(t *testing.T) { - baseNames := []string{ - "Authorization", - "Proxy-Authorization", - "Api-Key", - "X-Api-Key", - "X-Goog-Api-Key", - } - selected := dialectSpecificCredentialHeaders{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - } - source := http.Header{"X-Safe": {"kept"}} - for index, name := range baseNames { - if !platformheader.IsCredentialName(name) { - t.Fatalf("shared policy does not recognize base credential Header %q", name) - } - source.Set(name, fmt.Sprintf("base-secret-%d", index)) - } - source.Set("X-Dialect-Credential", "dialect-secret") - input := ForwardInput{ - Dialect: selected, - APIKey: "provider-api-key", - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - - sanitized := sanitizeForwardResponseHeaders(source, input) - for _, name := range append(baseNames, "X-Dialect-Credential") { - if values := sanitized.Values(name); values != nil { - t.Errorf("credential response Header %q survived: %#v", name, values) - } - } - if got := sanitized.Get("X-Safe"); got != "kept" { - t.Fatalf("safe response Header = %q, want kept", got) - } - - finalHeaders := http.Header{ - "Authorization": {"base-final-secret"}, - "X-Dialect-Credential": {"dialect-final-secret"}, - } - secrets := resolvedCredentialSecretValues(input, finalHeaders) - for _, want := range []string{ - "provider-api-key", - "base-final-secret", - "dialect-final-secret", - } { - if !slices.Contains(secrets, want) { - t.Errorf("resolved credential secrets = %#v, want %q", secrets, want) - } - } -} - -func TestForwardRequestRulesCannotRestoreReservedHeaders(t *testing.T) { - const apiKey = "fake-request-sanitizer-provider-key" - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Group: state.GroupView{ - UpstreamURL: "https://api.example.test", - HeaderRules: state.HeaderRules{Set: map[string]string{ - "Authorization": "Token ${API_KEY}", - "Connection": "X-Rule-Hop", - "X-Rule-Hop": "restored", - "Cookie": "session=fake", - "Cookie2": "legacy=fake", - "Proxy-Authorization": "Bearer proxy-fake", - "Proxy-Custom": "proxy-fake", - requestIDHeader: "restored-request-id", - debugHeaderGroup: "restored-group", - debugHeaderKey: "restored-key", - debugHeaderAttempts: "restored-attempts", - "Accept-Encoding": "gzip", - "Content-Encoding": "br", - }, Remove: []string{"Accept-Encoding", "Content-Encoding"}}, - }, - APIKey: apiKey, - Request: &dialect.ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/chat/completions", - Header: http.Header{ - "Connection": {"X-Client-Hop"}, - "X-Client-Hop": {"client-hop"}, - "Cookie": {"client=fake"}, - "Proxy-Client-Hop": {"proxy-client"}, - "accept-encoding": {"zstd"}, - "content-encoding": {"gzip"}, - }, - Body: []byte(`{"model":"gpt-test"}`), - }, - } - - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - request, _, replay, err := forwarder.newUpstreamRequest(t.Context(), input, false) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { - _ = request.Body.Close() - replay.release() - }) - - if got := request.Header.Get("Authorization"); got != "Token "+apiKey { - t.Fatalf("Authorization = %q, want configured credential", got) - } - if values := headerFieldValues(request.Header, "Accept-Encoding"); !reflect.DeepEqual(values, []string{"identity"}) { - t.Fatalf("Accept-Encoding values = %#v, want [identity]", values) - } - for _, name := range []string{ - "Connection", - "X-Rule-Hop", - "X-Client-Hop", - "Cookie", - "Cookie2", - "Proxy-Authorization", - "Proxy-Custom", - "Proxy-Client-Hop", - requestIDHeader, - debugHeaderGroup, - debugHeaderKey, - debugHeaderAttempts, - "Content-Encoding", - } { - if values := headerFieldValues(request.Header, name); values != nil { - t.Errorf("upstream request Header %s survived: %#v", name, values) - } - } -} - -func TestForwardAllRequestsUseIdentityRepresentation(t *testing.T) { - tests := []struct { - name string - stream bool - configure func(*ForwardInput) - assert func(*testing.T, []byte) - }{ - { - name: "normal request", - assert: func(t *testing.T, body []byte) { - t.Helper() - if string(body) != `{"model":"public-model","stream":true}` { - t.Fatalf("upstream body = %s, want unchanged plaintext", body) - } - }, - }, - { - name: "alias rewrite", - configure: func(input *ForwardInput) { - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - }, - assert: func(t *testing.T, body []byte) { - t.Helper() - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil || payload["model"] != "provider-model" { - t.Fatalf("upstream alias body = %s, error = %v", body, err) - } - }, - }, - { - name: "stream usage injection", - stream: true, - configure: func(input *ForwardInput) { - input.Group.InjectUsageOptions = true - }, - assert: func(t *testing.T, body []byte) { - t.Helper() - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - t.Fatalf("decode upstream stream body: %v", err) - } - options, ok := payload["stream_options"].(map[string]any) - if !ok || options["include_usage"] != true { - t.Fatalf("upstream stream body = %s, want include_usage", body) - } - }, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - input := streamForwardInput("https://api.example.test") - input.Request.Header = http.Header{ - "Accept-Encoding": {"gzip"}, - "accept-encoding": {"br"}, - "Content-Encoding": {"gzip"}, - "content-encoding": {"zstd"}, - "Content-Type": {"application/json"}, - "Accept": {"application/json"}, - "Idempotency-Key": {"request-1"}, - } - input.Request.Body = []byte(`{"model":"public-model","stream":true}`) - if test.configure != nil { - test.configure(&input) - } - - request, _, replay, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).newUpstreamRequest(t.Context(), input, test.stream) - if err != nil { - t.Fatalf("newUpstreamRequest() error = %v", err) - } - t.Cleanup(func() { - _ = request.Body.Close() - replay.release() - }) - body, err := io.ReadAll(request.Body) - if err != nil { - t.Fatalf("read upstream request body: %v", err) - } - test.assert(t, body) - if values := headerFieldValues(request.Header, "Accept-Encoding"); !reflect.DeepEqual(values, []string{"identity"}) { - t.Fatalf("Accept-Encoding values = %#v, want [identity]", values) - } - if values := headerFieldValues(request.Header, "Content-Encoding"); values != nil { - t.Fatalf("Content-Encoding values = %#v, want absent", values) - } - if request.ContentLength != int64(len(body)) { - t.Fatalf("ContentLength = %d, want %d", request.ContentLength, len(body)) - } - for name, want := range map[string]string{ - "Content-Type": "application/json", - "Accept": "application/json", - "Idempotency-Key": "request-1", - "Authorization": "Bearer sk-upstream-secret", - } { - if got := request.Header.Get(name); got != want { - t.Errorf("%s = %q, want %q", name, got, want) - } - } - }) - } -} - -func TestForwardNormalizesFinalRequestMetadata(t *testing.T) { - input := streamForwardInput("https://api.example.test") - input.Request.Header = http.Header{ - "Content-Type": {"application/json"}, - "Accept": {"application/json"}, - "Idempotency-Key": {"request-1"}, - "Content-Length": {"999"}, - "content-length": {"998"}, - "Content-Encoding": {"gzip"}, - "content-encoding": {"br"}, - "ETag": {`"client"`}, - "Digest": {"sha-256=client"}, - "Content-MD5": {"client-md5"}, - "Content-Range": {"bytes 0-1/2"}, - "Content-Digest": {"sha-256=:Y2xpZW50:"}, - "Repr-Digest": {"sha-256=:cmVwcg==:"}, - "Signature": {"client-signature"}, - "Signature-Input": {"client-signature-input"}, - "accept-encoding": {"gzip"}, - "Accept-Encoding": {"br"}, - } - input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ - "Authorization": "Token ${API_KEY}", - "Digest": "sha-256=rule", - "Signature": "rule-signature", - "Content-Encoding": "zstd", - "Accept-Encoding": "deflate", - }} - - request, _, replay, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).newUpstreamRequest(t.Context(), input, false) - if err != nil { - t.Fatalf("newUpstreamRequest() error = %v", err) - } - t.Cleanup(func() { - _ = request.Body.Close() - replay.release() - }) - body, err := io.ReadAll(request.Body) - if err != nil { - t.Fatalf("read upstream body: %v", err) - } - if request.ContentLength != int64(len(body)) { - t.Fatalf("ContentLength = %d, want %d", request.ContentLength, len(body)) - } - for _, name := range []string{ - "Content-Length", - "Content-Encoding", - "ETag", - "Digest", - "Content-MD5", - "Content-Range", - "Content-Digest", - "Repr-Digest", - "Signature", - "Signature-Input", - } { - if values := headerFieldValues(request.Header, name); values != nil { - t.Errorf("%s values = %#v, want removed", name, values) - } - } - if values := headerFieldValues(request.Header, "Accept-Encoding"); !reflect.DeepEqual(values, []string{"identity"}) { - t.Fatalf("Accept-Encoding values = %#v, want [identity]", values) - } - for name, want := range map[string]string{ - "Content-Type": "application/json", - "Accept": "application/json", - "Idempotency-Key": "request-1", - "Authorization": "Token sk-upstream-secret", - } { - if got := request.Header.Get(name); got != want { - t.Errorf("%s = %q, want %q", name, got, want) - } - } -} - -func TestNewUpstreamRequestDoesNotExposeGetBody(t *testing.T) { - const payload = `{"model":"gpt-test","messages":[]}` - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Group: state.GroupView{ - UpstreamURL: "https://example.com", - }, - APIKey: "fake-provider-key", - Request: &dialect.ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/chat/completions", - Header: http.Header{ - "Idempotency-Key": {"client-idempotency"}, - "X-Idempotency-Key": {"client-x-idempotency"}, - }, - Body: []byte(payload), - }, - } - - request, _, replay, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).newUpstreamRequest(t.Context(), input, false) - if err != nil { - t.Fatalf("newUpstreamRequest() error = %v", err) - } - t.Cleanup(func() { - _ = request.Body.Close() - replay.release() - }) - - body, err := io.ReadAll(request.Body) - if err != nil { - t.Fatalf("read upstream Body: %v", err) - } - if string(body) != payload || request.ContentLength != int64(len(payload)) { - t.Fatalf( - "upstream Body/ContentLength = %q/%d, want %q/%d", - body, - request.ContentLength, - payload, - len(payload), - ) - } - if request.GetBody != nil { - t.Fatal("newUpstreamRequest() exposed GetBody to net/http.Transport") - } - if request.Header.Get("Idempotency-Key") != "client-idempotency" || - request.Header.Get("X-Idempotency-Key") != "client-x-idempotency" { - t.Fatalf("idempotency headers = %#v", request.Header) - } -} - -func TestForwarderBoundsNonStreamingBodies(t *testing.T) { - if maxNonStreamingResponseBodyBytes != 32<<20 || maxErrorResponseBodyBytes != 64<<10 { - t.Fatalf("response limits = %d/%d", maxNonStreamingResponseBodyBytes, maxErrorResponseBodyBytes) - } - tests := []struct { - name, key string - status int - size int64 - wantProtocolErr bool - wantPlaceholder bool - }{ - {name: "success exact", key: "key", status: http.StatusOK, size: maxNonStreamingResponseBodyBytes}, - {name: "success plus one", key: "key", status: http.StatusOK, size: maxNonStreamingResponseBodyBytes + 1, wantProtocolErr: true}, - {name: "error exact", key: "key", status: http.StatusUnauthorized, size: maxErrorResponseBodyBytes}, - {name: "error plus one", key: "key", status: http.StatusUnauthorized, size: maxErrorResponseBodyBytes + 1, wantPlaceholder: true}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.WriteHeader(test.status) - writer.(http.Flusher).Flush() - _, _ = io.CopyN(writer, repeatingByteReader('x'), test.size) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, test.key, 10*time.Second) - if test.wantProtocolErr { - if !errors.Is(result.Err, ErrUpstreamProtocol) || !result.RequestWritten || result.StatusCode != 0 || - result.Body != nil || result.ClassificationBody != nil { - t.Fatalf("result = %#v", result) - } - return - } - if result.Err != nil || result.StatusCode != test.status { - t.Fatalf("result = %#v", result) - } - if test.wantPlaceholder { - if string(result.Body) != redact.Placeholder || string(result.ClassificationBody) != redact.Placeholder { - t.Fatalf("safe bodies = %q/%q", result.Body, result.ClassificationBody) - } - return - } - if int64(len(result.Body)) != test.size { - t.Fatalf("body length = %d, want %d", len(result.Body), test.size) - } - }) - } -} - -func TestForwarderBoundsDecompressedErrorBodies(t *testing.T) { - for _, encoding := range []string{"gzip", "br", "deflate", "zstd"} { - t.Run(encoding, func(t *testing.T) { - for _, size := range []int{1 << 20, 1<<20 + 1} { - plain := bytes.Repeat([]byte("x"), size) - wire := compressResponseWithBoundedZstdWindow(t, encoding, plain) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", encoding) - writer.WriteHeader(http.StatusUnauthorized) - _, _ = writer.Write(wire) - })) - - result := testForward(t, upstream.URL, "key", 10*time.Second) - upstream.Close() - if result.Err != nil || result.StatusCode != http.StatusUnauthorized { - t.Fatalf("size %d result = %#v", size, result) - } - if size == 1<<20 { - if !bytes.Equal(result.Body, plain) || !bytes.Equal(result.ClassificationBody, plain) || - len(headerFieldValues(result.Header, "Content-Encoding")) != 0 || - result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { - t.Fatalf("exact limit result body lengths = %d/%d, headers=%#v", len(result.Body), len(result.ClassificationBody), result.Header) - } - } else if string(result.Body) != redact.Placeholder || - string(result.ClassificationBody) != redact.Placeholder || result.Header.Get("Content-Encoding") != "" { - t.Fatalf("overflow result = %#v", result) - } - } - }) - } -} - -func TestForwarderFailsClosedWhenRedactionExpandsErrorBeyondBounds(t *testing.T) { - tests := []struct { - name string - encoding string - plain []byte - wantPlaceholder bool - }{ - { - name: "identity redaction remains within decoded limit", - plain: bytes.Repeat([]byte("a"), int(maxErrorResponseBodyBytes)), - }, - { - name: "gzip classification body exceeds limit after redaction", - encoding: "gzip", - plain: bytes.Repeat([]byte("a"), int(maxDecompressedErrorBodyBytes)), - wantPlaceholder: true, - }, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - wire := test.plain - if test.encoding != "" { - wire = encodeContentCodingForGatewayTest(t, contentcoding.Encoding(test.encoding), test.plain) - } - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - if test.encoding != "" { - writer.Header().Set("Content-Encoding", test.encoding) - } - writer.WriteHeader(http.StatusUnauthorized) - _, _ = writer.Write(wire) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "a", time.Second) - if result.Err != nil || result.StatusCode != http.StatusUnauthorized || - result.Header.Get("Content-Encoding") != "" { - t.Fatalf( - "result status=%d body=%d classification=%d encoding=%q err=%v", - result.StatusCode, len(result.Body), len(result.ClassificationBody), - result.Header.Get("Content-Encoding"), result.Err, - ) - } - if test.wantPlaceholder { - if string(result.Body) != redact.Placeholder || string(result.ClassificationBody) != redact.Placeholder { - t.Fatalf("overflow result = %#v", result) - } - return - } - if len(result.Body) != len(result.ClassificationBody) || int64(len(result.Body)) > maxDecompressedErrorBodyBytes || bytes.Contains(result.Body, []byte("a")) { - t.Fatalf("safe identity result = %#v", result) - } - }) - } -} - -func TestPrepareErrorRepresentationFailsClosedForOversizedWire(t *testing.T) { - plain := make([]byte, 100<<10) - const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_." - random := mathrand.New(mathrand.NewSource(1)) - for index := range plain { - plain[index] = alphabet[random.Intn(len(alphabet))] - } - gzipWire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) - if len(gzipWire) <= int(maxErrorResponseBodyBytes) { - t.Fatalf("gzip fixture length = %d, want oversized wire", len(gzipWire)) - } - - tests := []struct { - name string - headers http.Header - wire []byte - }{ - { - name: "identity", - headers: make(http.Header), - wire: bytes.Repeat([]byte("x"), int(maxErrorResponseBodyBytes)+1), - }, - { - name: "gzip", - headers: http.Header{"Content-Encoding": {"gzip"}}, - wire: gzipWire, - }, - } - forwarder := &Forwarder{redactor: redact.New()} - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - headers := test.headers.Clone() - prepared := forwarder.prepareErrorRepresentation(ForwardInput{}, headers, test.wire, nil) - if string(prepared.downstream) != redact.Placeholder || - string(prepared.classification) != redact.Placeholder || - len(headerFieldValues(prepared.headers, "Content-Encoding")) != 0 || - !reflect.DeepEqual(headers, test.headers) { - t.Fatalf( - "safe representation = %#v, source headers = %#v", - prepared, headers, - ) - } - }) - } -} - -func compressResponseWithBoundedZstdWindow(t *testing.T, encoding string, plain []byte) []byte { - t.Helper() - if encoding != "zstd" { - return encodeContentCodingForGatewayTest(t, contentcoding.Encoding(encoding), plain) - } - encoder, err := zstd.NewWriter( - nil, - zstd.WithWindowSize(int(maxDecompressedErrorBodyBytes)), - zstd.WithEncoderConcurrency(1), - ) - if err != nil { - t.Fatal(err) - } - defer encoder.Close() - return encoder.EncodeAll(plain, nil) -} - -func TestForwarderDoesNotReturnPartialBodyOnReadError(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - connection, buffer, err := writer.(http.Hijacker).Hijack() - if err != nil { - t.Errorf("hijack response: %v", err) - return - } - defer connection.Close() - _, _ = fmt.Fprint(buffer, "HTTP/1.1 200 OK\r\nContent-Length: 10\r\n\r\n1234") - _ = buffer.Flush() - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "key", time.Second) - if result.Err == nil || !result.RequestWritten || result.StatusCode != 0 || result.Body != nil || result.ClassificationBody != nil { - t.Fatalf("result = %#v", result) - } -} - -type repeatingByteReader byte - -func (reader repeatingByteReader) Read(destination []byte) (int, error) { - for index := range destination { - destination[index] = byte(reader) - } - return len(destination), nil -} - -type responseReadError struct { - err error -} - -func (reader responseReadError) Read([]byte) (int, error) { - return 0, reader.err -} - -func TestSanitizeUpstreamResponseHeadersRemovesCurrentCredential(t *testing.T) { - const secret = "provider-secret" - source := http.Header{ - "Authorization": {"Bearer unrelated"}, - "Proxy-Authorization": {"unrelated"}, - "Api-Key": {"unrelated"}, - "X-Api-Key": {"unrelated"}, - "X-Goog-Api-Key": {"unrelated"}, - "X-Echo": {"prefix-" + secret + "-suffix"}, - "X-Multi": {"safe", secret}, - "X-Safe": {"kept"}, - } - before := source.Clone() - got := sanitizeUpstreamResponseHeaders(source, secret) - for _, name := range []string{ - "Authorization", "Proxy-Authorization", "Api-Key", - "X-Api-Key", "X-Goog-Api-Key", "X-Echo", "X-Multi", - } { - if got.Values(name) != nil { - t.Fatalf("Header %s survived: %#v", name, got.Values(name)) - } - } - if got.Get("X-Safe") != "kept" || !reflect.DeepEqual(source, before) { - t.Fatalf("safe/source headers = %#v / %#v", got, source) - } -} - -func TestSanitizeUpstreamResponseHeadersHandlesNonCanonicalFieldNames(t *testing.T) { - const secret = "provider-secret-noncanonical" - source := http.Header{ - "authorization": {"Bearer unrelated"}, - "pRoXy-AuThOrIzAtIoN": {"unrelated"}, - "aPi-kEy": {"unrelated"}, - "x-aPi-kEy": {"unrelated"}, - "x-gOoG-aPi-kEy": {"unrelated"}, - "X-Echo": {"prefix-" + secret}, - "x-echo": {"suffix-" + secret}, - "X-Safe": {"kept"}, - } - before := source.Clone() - - got := sanitizeUpstreamResponseHeaders(source, secret) - - credentialNames := []string{ - "Authorization", "Proxy-Authorization", "Api-Key", - "X-Api-Key", "X-Goog-Api-Key", - } - for actualName, values := range got { - for _, credentialName := range credentialNames { - if strings.EqualFold(actualName, credentialName) { - t.Fatalf("credential Header %q survived: %#v", actualName, values) - } - } - if headerValuesContainLiteral(values, secret) { - t.Fatalf("Header %q retained current key: %#v", actualName, values) - } - } - if safe := got["X-Safe"]; len(safe) != 1 || safe[0] != "kept" || !reflect.DeepEqual(source, before) { - t.Fatalf("safe/source headers = %#v / %#v", got, source) - } -} - -func TestSanitizeUpstreamResponseHeadersRemovesAllCasingsOfMatchedField(t *testing.T) { - const secret = "provider-secret-duplicate-casing" - source := http.Header{ - "X-Echo": {"safe"}, - "x-echo": {"prefix-" + secret}, - "X-Safe": {"kept"}, - } - before := source.Clone() - - got := sanitizeUpstreamResponseHeaders(source, secret) - - for actualName := range got { - if strings.EqualFold(actualName, "X-Echo") { - t.Fatalf("logical Header X-Echo survived as %q: %#v", actualName, got[actualName]) - } - } - if safe := got["X-Safe"]; len(safe) != 1 || safe[0] != "kept" || !reflect.DeepEqual(source, before) { - t.Fatalf("safe/source headers = %#v / %#v", got, source) - } -} - -func TestSanitizeForwardResponseHeadersPreservesRepresentationHeaderLiteralCollisions(t *testing.T) { - for _, test := range []struct { - name, upstreamModel, headerName, headerValue string - }{ - {name: "content encoding", upstreamModel: "gzip", headerName: "Content-Encoding", headerValue: "gzip"}, - {name: "content type", upstreamModel: "json", headerName: "Content-Type", headerValue: "application/json"}, - {name: "content type preserves case", upstreamModel: "JSON", headerName: "Content-Type", headerValue: "application/JSON"}, - {name: "problem json content type", upstreamModel: "problem", headerName: "Content-Type", headerValue: "application/problem+json"}, - {name: "ndjson content type", upstreamModel: "ndjson", headerName: "Content-Type", headerValue: "application/x-ndjson"}, - {name: "event stream content type", upstreamModel: "event", headerName: "Content-Type", headerValue: "text/event-stream"}, - {name: "plain text content type", upstreamModel: "plain", headerName: "Content-Type", headerValue: "text/plain"}, - {name: "content length", upstreamModel: "42", headerName: "Content-Length", headerValue: "42"}, - } { - t.Run(test.name, func(t *testing.T) { - source := http.Header{ - test.headerName: {test.headerValue}, - "X-Upstream-Model": {test.upstreamModel}, - } - got := sanitizeForwardResponseHeaders(source, ForwardInput{ - ExternalModel: "public-model", UpstreamModelID: test.upstreamModel, - }) - - if got.Get(test.headerName) != test.headerValue { - t.Fatalf("%s = %q, want %q", test.headerName, got.Get(test.headerName), test.headerValue) - } - if got.Get("X-Upstream-Model") != "" { - t.Fatalf("X-Upstream-Model survived: %#v", got) - } - }) - } -} - -func TestSanitizeForwardResponseHeadersRemovesAliasedModelFromFieldNames(t *testing.T) { - const upstreamModel = "provider-model" - source := http.Header{ - "X-Provider-Model-Quota": {"safe"}, - "x-provider-model-quota": {"also-safe"}, - "X-Safe": {"kept"}, - } - got := sanitizeForwardResponseHeaders(source, ForwardInput{ - ExternalModel: "public-model", UpstreamModelID: upstreamModel, - }) - - for name := range got { - if strings.Contains(strings.ToLower(name), upstreamModel) { - t.Fatalf("Header field name leaked aliased model as %q: %#v", name, got[name]) - } - } - if got.Get("X-Safe") != "kept" { - t.Fatalf("safe header changed: %#v", got) - } -} - -func TestSanitizeForwardResponseHeadersPreservesSafeContentTypeFieldNameCollisions(t *testing.T) { - for _, test := range []struct { - name string - upstreamModel string - contentType string - }{ - {name: "content in field name", upstreamModel: "content", contentType: "application/json"}, - {name: "type in field name", upstreamModel: "type", contentType: "text/event-stream"}, - } { - t.Run(test.name, func(t *testing.T) { - customHeader := "X-" + test.upstreamModel + "-Quota" - source := http.Header{ - "Content-Type": {test.contentType}, - customHeader: {"safe"}, - } - got := sanitizeForwardResponseHeaders(source, ForwardInput{ - ExternalModel: "public-model", UpstreamModelID: test.upstreamModel, - }) - - if got.Get("Content-Type") != test.contentType { - t.Fatalf("Content-Type = %q, want %q", got.Get("Content-Type"), test.contentType) - } - if got.Get(customHeader) != "" { - t.Fatalf("custom Header field-name collision survived: %#v", got) - } - }) - } -} - -func TestSanitizeForwardResponseHeadersAlwaysDropsSignatures(t *testing.T) { - const upstreamModel = "provider-model" - for _, test := range []struct { - name string - source http.Header - }{ - { - name: "custom field name", - source: http.Header{ - "X-Provider-Model-Quota": {"safe"}, - }, - }, - { - name: "custom field value", - source: http.Header{ - "X-Upstream": {"selected=" + upstreamModel}, - }, - }, - { - name: "content type", - source: http.Header{ - "Content-Type": {"application/vnd.provider-model+json"}, - }, - }, - } { - t.Run(test.name, func(t *testing.T) { - source := test.source.Clone() - source["Signature"] = []string{"sig-canonical"} - source["sIgNaTuRe"] = []string{"sig-noncanonical"} - source["Signature-Input"] = []string{"input-canonical"} - source["sIgNaTuRe-InPuT"] = []string{"input-noncanonical"} - - got := sanitizeForwardResponseHeaders(source, ForwardInput{ - ExternalModel: "public-model", UpstreamModelID: upstreamModel, - }) - - for name := range got { - if strings.EqualFold(name, "Signature") || strings.EqualFold(name, "Signature-Input") { - t.Fatalf("invalidated signature Header survived as %q: %#v", name, got[name]) - } - } - }) - } - - t.Run("alias without deletion still removes signatures", func(t *testing.T) { - source := http.Header{ - "X-Safe": {"kept"}, - "Signature": {"sig"}, - "Signature-Input": {"input"}, - } - got := sanitizeForwardResponseHeaders(source, ForwardInput{ - ExternalModel: "public-model", UpstreamModelID: upstreamModel, - }) - if got.Get("Signature") != "" || got.Get("Signature-Input") != "" { - t.Fatalf("unchanged Header set retained signatures: %#v", got) - } - }) - - t.Run("no alias still removes signatures", func(t *testing.T) { - source := http.Header{ - "X-Upstream": {"selected=" + upstreamModel}, - "Signature": {"sig"}, - "Signature-Input": {"input"}, - } - got := sanitizeForwardResponseHeaders(source, ForwardInput{ - ExternalModel: upstreamModel, UpstreamModelID: upstreamModel, - }) - if got.Get("Signature") != "" || got.Get("Signature-Input") != "" { - t.Fatalf("non-alias Header set retained signatures: %#v", got) - } - }) -} - -func TestSanitizeForwardResponseHeadersRemovesAliasedModelFromContentTypeParameters(t *testing.T) { - const upstreamModel = "provider-model" - for _, test := range []struct { - name, headerName, headerValue, upstreamModel string - }{ - { - name: "parameter value", headerName: "Content-Type", - headerValue: `application/json; model=provider-model`, upstreamModel: upstreamModel, - }, - { - name: "parameter name", headerName: "Content-Type", - headerValue: `application/json; provider-model=safe`, upstreamModel: upstreamModel, - }, - { - name: "malformed non-canonical field", headerName: "cOnTeNt-TyPe", - headerValue: `application/json; model=provider-model; broken`, upstreamModel: upstreamModel, - }, - { - name: "malformed without parameter delimiter", headerName: "Content-Type", - headerValue: `application/json provider-model`, upstreamModel: upstreamModel, - }, - { - name: "vendor media type", headerName: "Content-Type", - headerValue: `application/vnd.provider-model+json`, upstreamModel: upstreamModel, - }, - { - name: "model crosses parameter delimiter", headerName: "Content-Type", - headerValue: `application/json; model=safe`, upstreamModel: "json; model", - }, - } { - t.Run(test.name, func(t *testing.T) { - source := http.Header{ - test.headerName: {test.headerValue}, - "X-Safe": {"kept"}, - } - got := sanitizeForwardResponseHeaders(source, ForwardInput{ - ExternalModel: "public-model", UpstreamModelID: test.upstreamModel, - }) - - for name := range got { - if strings.EqualFold(name, "Content-Type") { - t.Fatalf("logical Content-Type survived as %q: %#v", name, got[name]) - } - } - if got.Get("X-Safe") != "kept" { - t.Fatalf("safe header changed: %#v", got) - } - }) - } -} - -func TestSanitizeForwardResponseHeadersRemovesCaseSensitiveModelFromRawContentTypeParameters(t *testing.T) { - source := http.Header{ - "Content-Type": {`application/JSON; JSON=safe`}, - "X-Safe": {"kept"}, - } - got := sanitizeForwardResponseHeaders(source, ForwardInput{ - ExternalModel: "public-model", UpstreamModelID: "JSON", - }) - - if got.Get("Content-Type") != "" { - t.Fatalf("Content-Type survived: %#v", got) - } - if got.Get("X-Safe") != "kept" { - t.Fatalf("safe header changed: %#v", got) - } -} - -func TestForwarderDoesNotTreatShortKeyAsContentCodingFailure(t *testing.T) { - tests := []struct { - name string - status int - wantBody string - }{ - {name: "success coding header is terminated", status: http.StatusOK, wantBody: `{"error":"safe"}`}, - {name: "error coding header is terminated", status: http.StatusUnauthorized, wantBody: `{"error":"safe"}`}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - wire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, []byte(`{"error":"safe"}`)) - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Encoding", "gzip") - w.WriteHeader(test.status) - _, _ = w.Write(wire) - })) - defer upstream.Close() - result := testForward(t, upstream.URL, "gzip", time.Second) - if result.Err != nil || result.StatusCode != test.status || - string(result.Body) != test.wantBody || - len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { - t.Fatalf("result = %#v", result) - } - }) - } -} - -func TestErrorRepresentationRejectsAllContentEncodingCasings(t *testing.T) { - headers := http.Header{ - "Content-Encoding": {"identity"}, - "cOnTeNt-EnCoDiNg": {"gzip"}, - } - prepared := (&Forwarder{redactor: redact.New()}).prepareErrorRepresentation( - ForwardInput{}, headers, []byte("opaque"), nil, - ) - if string(prepared.downstream) != redact.Placeholder || - string(prepared.classification) != redact.Placeholder { - t.Fatalf("prepared = %#v", prepared) - } - for actualName := range prepared.headers { - if strings.EqualFold(actualName, "Content-Encoding") { - t.Fatalf("Content-Encoding survived as %q: %#v", actualName, prepared.headers[actualName]) - } - } -} - -func TestForwarderSanitizesResponseHeadersOnAllPaths(t *testing.T) { - for _, test := range []struct { - name string - stream bool - status int - }{ - {name: "nonstream success", status: http.StatusOK}, - {name: "nonstream error", status: http.StatusUnauthorized}, - {name: "stream success", stream: true, status: http.StatusOK}, - {name: "stream error", stream: true, status: http.StatusUnauthorized}, - } { - t.Run(test.name, func(t *testing.T) { - const ( - secret = "provider-secret-all-paths" - ordinaryRuleValue = "independent-header-rule-value" - ) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - writer.Header().Set("X-Echo", "prefix-"+secret) - writer.Header().Set("X-Rule-Echo", "prefix-"+request.Header.Get("X-Custom-Credential")) - writer.Header().Set("X-Safe", "kept") - writer.WriteHeader(test.status) - if test.stream && test.status == http.StatusOK { - _, _ = writer.Write([]byte("data: ok\n\n")) - return - } - _, _ = writer.Write([]byte(`{"error":"safe"}`)) - })) - defer upstream.Close() - - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - input := streamForwardInput(upstream.URL) - input.APIKey = secret - input.Group.HeaderRules = state.HeaderRules{ - Set: map[string]string{"X-Custom-Credential": ordinaryRuleValue}, - } - downstream := newRecordingResponseWriter() - var result UpstreamResult - if test.stream { - result = forwarder.ForwardStream(t.Context(), input, downstream) - } else { - result = forwarder.Forward(t.Context(), input) - } - if result.Err != nil || - result.Header.Get("X-Echo") != "" || - result.Header.Get("X-Rule-Echo") != "prefix-"+ordinaryRuleValue || - result.Header.Get("X-Safe") != "kept" { - t.Fatalf("result = %#v", result) - } - if test.stream && test.status == http.StatusOK && - (downstream.header.Get("X-Echo") != "" || - downstream.header.Get("X-Rule-Echo") != "prefix-"+ordinaryRuleValue || - downstream.header.Get("X-Safe") != "kept") { - t.Fatalf("downstream headers = %#v", downstream.header) - } - }) - } -} - -func TestForwarderSanitizesResponseHeadersWithCurrentAttemptKey(t *testing.T) { - for _, secret := range []string{"secret-one", "secret-two"} { - t.Run(secret, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("X-Echo", "prefix-"+secret) - writer.Header().Set("X-Safe", "kept") - _, _ = writer.Write([]byte(`{"ok":true}`)) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, secret, time.Second) - if result.Err != nil || result.Header.Get("X-Echo") != "" || - result.Header.Get("X-Safe") != "kept" { - t.Fatalf("result = %#v", result) - } - }) - } -} - -func TestForwardStreamForcesIdentityAfterHeaderRules(t *testing.T) { - tests := []struct { - name string - rules state.HeaderRules - }{ - { - name: "set cannot override identity", - rules: state.HeaderRules{Set: map[string]string{ - "Accept-Encoding": "gzip", - "X-Custom": "prefix-${API_KEY}", - }}, - }, - { - name: "remove cannot delete identity", - rules: state.HeaderRules{Remove: []string{"Accept-Encoding"}}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - received := make(chan http.Header, 1) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - received <- request.Header.Clone() - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: ok\n\n")) - })) - defer upstream.Close() - - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - downstream := newRecordingResponseWriter() - input := streamForwardInput(upstream.URL) - input.Group.HeaderRules = tt.rules - result := forwarder.ForwardStream(context.Background(), input, downstream) - - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - headers := <-received - if got := headers.Get("Accept-Encoding"); got != "identity" { - t.Fatalf("Accept-Encoding = %q, want identity", got) - } - if got := headers.Get("Authorization"); got != "Bearer sk-upstream-secret" { - t.Fatalf("Authorization = %q", got) - } - if tt.rules.Set != nil && headers.Get("X-Custom") != "prefix-sk-upstream-secret" { - t.Fatalf("X-Custom = %q", headers.Get("X-Custom")) - } - }) - } -} - -func TestForwardAliasForcesIdentityAfterHeaderRules(t *testing.T) { - var receivedHeader http.Header - var receivedBody []byte - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - receivedHeader = request.Header.Clone() - receivedBody, _ = io.ReadAll(request.Body) - writer.Header().Set("Content-Type", "application/json") - _, _ = writer.Write([]byte(`{"ok":true}`)) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public"}`) - input.ExternalModel = "public" - input.UpstreamModelID = "provider" - input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{"Accept-Encoding": "gzip"}} - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - result := forwarder.Forward(context.Background(), input) - - if result.Err != nil || result.StatusCode != http.StatusOK { - t.Fatalf("Forward() result = %#v", result) - } - if receivedHeader.Get("Accept-Encoding") != "identity" { - t.Fatalf("Accept-Encoding = %q, want identity", receivedHeader.Get("Accept-Encoding")) - } - if string(receivedBody) != `{"model":"provider"}` { - t.Fatalf("upstream body = %s, want provider model", receivedBody) - } -} - -func TestForwarderRewritesAliasedNonStreamingResponsesAtBounds(t *testing.T) { - t.Run("unsupported response encoding fails closed", func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - writer.Header().Set("Content-Encoding", "gzip") - _, _ = writer.Write([]byte(`{"model":"provider"}`)) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public"}`) - input.ExternalModel = "public" - input.UpstreamModelID = "provider" - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - - if !errors.Is(result.Err, ErrUpstreamProtocol) || !result.RequestWritten || result.StatusCode != 0 { - t.Fatalf("Forward() result = %#v, want uninspectable alias response protocol error", result) - } - }) - - t.Run("rewrite expansion cannot exceed response bound", func(t *testing.T) { - const prefix = `{"model":"p","padding":"` - const suffix = `"}` - padding := strings.Repeat("x", int(maxNonStreamingResponseBodyBytes)-len(prefix)-len(suffix)) - responseBody := prefix + padding + suffix - if int64(len(responseBody)) != maxNonStreamingResponseBodyBytes { - t.Fatalf("test response size = %d", len(responseBody)) - } - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - writer.Header().Set("Content-Length", strconv.Itoa(len(responseBody))) - _, _ = io.WriteString(writer, responseBody) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public"}`) - input.ExternalModel = "public" - input.UpstreamModelID = "p" - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - - if !errors.Is(result.Err, ErrUpstreamProtocol) || !result.RequestWritten || result.StatusCode != 0 { - t.Fatalf("Forward() result = %#v, want rewritten response overflow", result) - } - if strings.Contains(result.Err.Error(), "%!w()") { - t.Fatalf("overflow error formats nil rewrite error: %v", result.Err) - } - }) -} - -type readyGuardWriter struct { - *recordingResponseWriter - ready *atomic.Bool - committedTooSoon atomic.Bool -} - -func (writer *readyGuardWriter) WriteHeader(status int) { - if !writer.ready.Load() { - writer.committedTooSoon.Store(true) - } - writer.recordingResponseWriter.WriteHeader(status) -} - -func TestForwardStreamCallsReadyBeforeCommit(t *testing.T) { - releaseUpstream := make(chan struct{}) - var releaseOnce sync.Once - release := func() { releaseOnce.Do(func() { close(releaseUpstream) }) } - - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, "data: ok\n\n") - writer.(http.Flusher).Flush() - <-releaseUpstream - })) - defer func() { - release() - upstream.Close() - }() - - var ready atomic.Bool - var calls atomic.Int32 - readyCalled := make(chan struct{}) - var readyOnce sync.Once - input := streamForwardInput(upstream.URL) - input.OnStreamReady = func() { - calls.Add(1) - ready.Store(true) - readyOnce.Do(func() { close(readyCalled) }) - } - downstream := &readyGuardWriter{ - recordingResponseWriter: newRecordingResponseWriter(), - ready: &ready, - } - done := make(chan UpstreamResult, 1) - go func() { - done <- NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - }() - - waitForSignal(t, readyCalled, "stream-ready callback") - if calls.Load() != 1 { - t.Fatalf("stream-ready calls while upstream is active = %d, want 1", calls.Load()) - } - select { - case result := <-done: - t.Fatalf("ForwardStream() returned before upstream release: %#v", result) - default: - } - release() - - select { - case result := <-done: - if result.Err != nil || !result.Committed || calls.Load() != 1 || - downstream.committedTooSoon.Load() { - t.Fatalf("result=%#v calls=%d committedTooSoon=%t", - result, calls.Load(), downstream.committedTooSoon.Load()) - } - case <-time.After(time.Second): - t.Fatal("ForwardStream() did not finish after upstream release") - } -} - -func providerErrorBeforeCommit(result UpstreamResult) bool { - return result.ProviderErrorBeforeCommit -} - -func TestForwardStreamReturnsFirstProviderErrorBeforeCommit(t *testing.T) { - const secret = "sk-obviously-fake-provider-error" - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - writer.Header().Set("Retry-After", "9") - _, _ = io.WriteString( - writer, - "event: error\ndata: {\"error\":{\"type\":\"rate_limit_error\",\"message\":\""+secret+"\"}}\n\n", - ) - })) - defer upstream.Close() - - var ready atomic.Int32 - input := streamForwardInput(upstream.URL) - input.APIKey = secret - input.OnStreamReady = func() { ready.Add(1) } - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream(context.Background(), input, downstream) - - if result.Err != nil || result.StatusCode != http.StatusOK || result.Committed || - !providerErrorBeforeCommit(result) || !result.RequestWritten { - t.Fatalf("ForwardStream() result = %#v", result) - } - if result.Body != nil || bytes.Contains(result.ClassificationBody, []byte(secret)) || - !bytes.Contains(result.ClassificationBody, []byte("rate_limit_error")) { - t.Fatalf("body/classification = %q/%q", result.Body, result.ClassificationBody) - } - if result.ErrorSummary != fixedErrorSummary("upstream_sse_error") { - t.Fatalf("ErrorSummary = %q", result.ErrorSummary) - } - if result.Header.Get("Retry-After") != "9" || ready.Load() != 0 || - downstream.status != 0 || downstream.body.Len() != 0 { - t.Fatalf("header/ready/downstream = %#v/%d/%d/%q", - result.Header, ready.Load(), downstream.status, downstream.body.String()) - } - if result.Usage.State != usage.StateMissing { - t.Fatalf("Usage = %#v, want missing", result.Usage) - } -} - -func TestForwardStreamBoundsAndRedactsFirstProviderError(t *testing.T) { - const ( - secret = "sk-obviously-fake-bounded-error" - headerSecret = "obviously-fake-header-rule-value" - ) - padding := strings.Repeat("x", maxFirstSSEEventBytes-256) - payload := `{"error":{"type":"server_overloaded","message":"` + - secret + ` ` + headerSecret + ` ` + padding + `"}}` - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, "data: "+payload+"\n\n") - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.APIKey = secret - input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ - "X-Provider-Trace": headerSecret, - }} - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream(context.Background(), input, newRecordingResponseWriter()) - - if result.Err != nil || !providerErrorBeforeCommit(result) { - t.Fatalf("ForwardStream() result = %#v", result) - } - if len(result.ClassificationBody) > maxFirstSSEEventBytes || - bytes.Contains(result.ClassificationBody, []byte(secret)) || - !bytes.Contains(result.ClassificationBody, []byte(headerSecret)) { - t.Fatalf("unsafe ClassificationBody length/content = %d/%q", - len(result.ClassificationBody), result.ClassificationBody) - } - if strings.Contains(result.ErrorSummary, secret) || - strings.Contains(result.ErrorSummary, headerSecret) || - len(result.ErrorSummary) > maxRequestLogSummaryBytes { - t.Fatalf("unsafe ErrorSummary = %q", result.ErrorSummary) - } -} - -func TestForwardStreamClassifiesAliasedNonObjectProviderErrorBeforeRewrite(t *testing.T) { - for _, test := range []struct { - name string - payload string - }{ - {name: "plain text", payload: "rate_limit_error provider-model"}, - {name: "JSON array", payload: `["rate_limit_error","provider-model"]`}, - } { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, "event: error\ndata: "+test.payload+"\n\n") - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - input.Request.Body = []byte(`{"model":"public-model","stream":true}`) - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream(context.Background(), input, downstream) - - if result.Err != nil || !result.ProviderErrorBeforeCommit || - result.Committed || result.StatusCode != http.StatusOK { - t.Fatalf("ForwardStream() result = %#v", result) - } - if bytes.Contains(result.ClassificationBody, []byte("provider-model")) || - !bytes.Contains(result.ClassificationBody, []byte("public-model")) || - !bytes.Contains(result.ClassificationBody, []byte("rate_limit_error")) { - t.Fatalf("ClassificationBody = %q", result.ClassificationBody) - } - decision := health.Judge(input.Dialect, health.Attempt{ - StatusCode: result.StatusCode, - Body: result.ClassificationBody, - Header: result.Header, - Now: time.Unix(1, 0), - ProviderErrorBeforeCommit: true, - }) - if decision.Category != health.FailureCategoryRateLimited || - decision.Action != health.ActionCooldownKey { - t.Fatalf("health decision = %#v", decision) - } - if downstream.status != 0 || downstream.body.Len() != 0 { - t.Fatalf("downstream status/body = %d/%q", downstream.status, downstream.body.String()) - } - }) - } -} - -func TestForwardStreamProviderErrorReturnsOnlyRateLimitHeaders(t *testing.T) { - const ( - apiKey = "sk-obviously-fake-provider-header" - ruleEcho = "ordinary-header-rule-literal" - retryAfter = "9" - ) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - writer.Header().Set("Retry-After", retryAfter) - writer.Header().Set("Anthropic-Ratelimit-Tokens-Reset", "30s") - writer.Header().Set("X-Ratelimit-Reset-Requests", "45s") - writer.Header().Set("X-Rule-Echo", ruleEcho) - writer.Header().Set("Authorization", "Bearer "+apiKey) - writer.Header().Set("X-Api-Key", apiKey) - _, _ = io.WriteString(writer, "event: error\ndata: rate_limit_error\n\n") - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.APIKey = apiKey - input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ - "X-Rule-Echo": ruleEcho, - }} - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream(context.Background(), input, newRecordingResponseWriter()) - - if result.Err != nil || !result.ProviderErrorBeforeCommit { - t.Fatalf("ForwardStream() result = %#v", result) - } - want := http.Header{ - "Retry-After": {retryAfter}, - "Anthropic-Ratelimit-Tokens-Reset": {"30s"}, - "X-Ratelimit-Reset-Requests": {"45s"}, - } - if !reflect.DeepEqual(result.Header, want) { - t.Fatalf("provider error Header = %#v, want %#v", result.Header, want) - } - now := time.Date(2026, time.July, 27, 12, 0, 0, 0, time.UTC) - decision := health.Judge(input.Dialect, health.Attempt{ - StatusCode: result.StatusCode, - Body: result.ClassificationBody, - Header: result.Header, - Now: now, - ProviderErrorBeforeCommit: true, - }) - if decision.Category != health.FailureCategoryRateLimited || - decision.Action != health.ActionCooldownKey || - !decision.CooldownUntil.Equal(now.Add(9*time.Second)) { - t.Fatalf("health decision = %#v", decision) - } -} - -func TestForwardStreamObservesCleanEOFAndFirstSSEError(t *testing.T) { - t.Run("clean EOF", func(t *testing.T) { - const wire = "data: {\"model\":\"gpt-4o\"}\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, wire) - })) - defer upstream.Close() - - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - context.Background(), - streamForwardInput(upstream.URL), - downstream, - ) - - if result.Err != nil || !result.Committed || - result.Stream.EndReason != StreamEndCleanEOF || - result.Stream.ErrorSummary != "" { - t.Fatalf("ForwardStream() clean result = %#v", result) - } - if downstream.status != http.StatusOK || downstream.body.String() != wire { - t.Fatalf( - "clean downstream status/body = %d/%q, want %d/%q", - downstream.status, - downstream.body.String(), - http.StatusOK, - wire, - ) - } - }) - - t.Run("first safe SSE error summary", func(t *testing.T) { - const ( - apiKey = "opaque-provider-secret" - resolvedAuth = "Token " + apiKey - resolvedRule = "opaque literal\tvalue" - ordinaryEncoding = "gzip" - globalSecret = "sk-global-secret-123456789" - upstreamModel = "upstream-private-model" - externalModel = "public-model" - disallowedSummary = "must-not-persist" - ) - firstMessage := strings.Join([]string{ - apiKey, - resolvedAuth, - resolvedRule, - ordinaryEncoding, - globalSecret, - upstreamModel, - }, " ") - firstPayload := `{"error":{"message":` + strconv.Quote(firstMessage) + - `},"debug":"` + disallowedSummary + `"}` - secondPayload := `{"error":{"message":"second error must not replace first"}}` - wire := "data: {\"model\":\"" + upstreamModel + "\",\"ready\":true}\r\n\r\n" + - "event: error\r\ndata: " + firstPayload + "\r\n\r\n" + - "event: error\ndata: " + secondPayload + "\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, wire) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.APIKey = apiKey - input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ - "Authorization": "Token ${API_KEY}", - "X-Custom": resolvedRule, - "Accept-Encoding": ordinaryEncoding, - }} - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - input.Request.Body = []byte(`{"model":"` + externalModel + `","stream":true}`) - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream(context.Background(), input, downstream) - - wantSummary := strings.Repeat(redact.Placeholder+" ", 5) + upstreamModel - if result.Err != nil || !result.Committed || - result.Stream.EndReason != StreamEndSSEError || - result.Stream.ErrorSummary != wantSummary { - t.Fatalf( - "ForwardStream() SSE observation = %#v, want summary %q", - result, - wantSummary, - ) - } - output := downstream.body.String() - for _, secret := range []string{apiKey, resolvedAuth, upstreamModel} { - if strings.Contains(output, secret) { - t.Fatalf("downstream SSE leaked %q: %q", secret, output) - } - } - firstBoundary := strings.Index(output, "\r\n\r\n") - if firstBoundary < 0 { - t.Fatalf("downstream SSE lacks first CRLF event boundary: %q", output) - } - secondBoundaryOffset := strings.Index(output[firstBoundary+len("\r\n\r\n"):], "\r\n\r\n") - if secondBoundaryOffset < 0 { - t.Fatalf("downstream SSE lacks committed error event boundary: %q", output) - } - secondStart := firstBoundary + len("\r\n\r\n") - secondEvent := output[secondStart : secondStart+secondBoundaryOffset] - dataIndex := strings.Index(secondEvent, "data: ") - if dataIndex < 0 { - t.Fatalf("first downstream SSE error event lacks data: %q", secondEvent) - } - var decodedFirstPayload map[string]any - if err := json.Unmarshal( - []byte(secondEvent[dataIndex+len("data: "):]), - &decodedFirstPayload, - ); err != nil { - t.Fatalf("decode first downstream SSE payload: %v", err) - } - firstError, ok := decodedFirstPayload["error"].(map[string]any) - if !ok { - t.Fatalf("first downstream SSE error = %#v", decodedFirstPayload["error"]) - } - downstreamMessage, ok := firstError["message"].(string) - if !ok { - t.Fatalf("first downstream SSE error message = %#v", firstError["message"]) - } - for _, ordinary := range []string{resolvedRule, ordinaryEncoding} { - if !strings.Contains(downstreamMessage, ordinary) { - t.Fatalf( - "ordinary HeaderRule value %q was changed in downstream SSE message: %q", - ordinary, - downstreamMessage, - ) - } - if strings.Contains(result.Stream.ErrorSummary, ordinary) { - t.Fatalf("stream ErrorSummary leaked HeaderRules value %q: %q", ordinary, result.Stream.ErrorSummary) - } - } - if strings.Contains(result.Stream.ErrorSummary, "opaque literal value") { - t.Fatalf("stream ErrorSummary leaked normalized HeaderRules value: %q", result.Stream.ErrorSummary) - } - if !strings.Contains(output, globalSecret) { - t.Fatalf("observation redactor changed downstream SSE: %q", output) - } - if strings.Contains(result.Stream.ErrorSummary, disallowedSummary) || - strings.Contains(result.Stream.ErrorSummary, "second error") || - strings.Contains(result.Stream.ErrorSummary, globalSecret) { - t.Fatalf("unsafe/non-first SSE summary = %q", result.Stream.ErrorSummary) - } - if !strings.Contains(output, externalModel) || - !strings.Contains(output, "second error must not replace first") || - !strings.Contains(output, "\r\n\r\n") || - strings.Count(output, "event: error") != 2 { - t.Fatalf("rewritten SSE framing/body = %q", output) - } - }) -} - -type cancelingStreamResponseWriter struct { - *recordingResponseWriter - cancel context.CancelFunc - err error -} - -func (writer *cancelingStreamResponseWriter) Write([]byte) (int, error) { - writer.cancel() - return 0, writer.err -} - -func TestForwardStreamPrioritizesCancellationAndDownstreamFailure(t *testing.T) { - const wire = "data: first\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, wire) - })) - defer upstream.Close() - - t.Run("client cancellation wins over write failure", func(t *testing.T) { - wantErr := errors.New("downstream write failed after cancellation") - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - downstream := &cancelingStreamResponseWriter{ - recordingResponseWriter: newRecordingResponseWriter(), - cancel: cancel, - err: wantErr, - } - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream(ctx, streamForwardInput(upstream.URL), downstream) - - if !result.Committed || - result.Stream.EndReason != StreamEndClientCanceled || - !errors.Is(result.Err, wantErr) { - t.Fatalf("ForwardStream() cancellation-priority result = %#v", result) - } - if downstream.status != http.StatusOK || downstream.body.Len() != 0 { - t.Fatalf( - "downstream status/body = %d/%q, want committed 200/empty", - downstream.status, - downstream.body.String(), - ) - } - }) - - t.Run("ordinary write failure", func(t *testing.T) { - wantErr := errors.New("downstream write failed") - downstream := newRecordingResponseWriter() - downstream.writeErr = wantErr - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - context.Background(), - streamForwardInput(upstream.URL), - downstream, - ) - - if !result.Committed || - result.Stream.EndReason != StreamEndDownstreamWriteFailure || - !errors.Is(result.Err, wantErr) { - t.Fatalf("ForwardStream() write-failure result = %#v", result) - } - if downstream.status != http.StatusOK || downstream.body.Len() != 0 { - t.Fatalf( - "downstream status/body = %d/%q, want committed 200/empty", - downstream.status, - downstream.body.String(), - ) - } - }) -} - -func TestForwardStreamRejectsUnsupportedSuccessEncodingBeforeCommit(t *testing.T) { - tests := []struct { - name string - encodings []string - wantCommit bool - }{ - {name: "missing encoding", wantCommit: true}, - {name: "empty encoding", encodings: []string{""}, wantCommit: true}, - {name: "identity", encodings: []string{" identity "}, wantCommit: true}, - {name: "gzip", encodings: []string{"gzip"}}, - {name: "encoding list", encodings: []string{"identity, gzip"}}, - {name: "multiple values", encodings: []string{"identity", "gzip"}}, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - for _, encoding := range test.encodings { - writer.Header().Add("Content-Encoding", encoding) - } - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: ok\n\n")) - })) - defer upstream.Close() - - var calls atomic.Int32 - input := streamForwardInput(upstream.URL) - input.OnStreamReady = func() { calls.Add(1) } - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - downstream := newRecordingResponseWriter() - result := forwarder.ForwardStream(context.Background(), input, downstream) - - if test.wantCommit { - if result.Err != nil || !result.Committed || - downstream.body.String() != "data: ok\n\n" || calls.Load() != 1 { - t.Fatalf("ForwardStream() valid result = %#v, body=%q calls=%d", - result, downstream.body.String(), calls.Load()) - } - return - } - if !errors.Is(result.Err, ErrUpstreamProtocol) || result.Committed || - result.RetryableBeforeCommit || calls.Load() != 0 { - t.Fatalf("ForwardStream() protocol result = %#v, calls=%d", result, calls.Load()) - } - if downstream.status != 0 || downstream.body.Len() != 0 || downstream.flushes != 0 { - t.Fatalf("downstream was touched before protocol rejection: %#v", downstream) - } - }) - } -} - -func TestForwardStreamRejectsOversizedFirstEventAsProtocolError(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte(":" + strings.Repeat("x", maxFirstSSEEventBytes))) - })) - defer upstream.Close() - - var calls atomic.Int32 - input := streamForwardInput(upstream.URL) - input.OnStreamReady = func() { calls.Add(1) } - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - downstream := newRecordingResponseWriter() - result := forwarder.ForwardStream(context.Background(), input, downstream) - - if !errors.Is(result.Err, ErrUpstreamProtocol) || - !errors.Is(result.Err, errFirstSSEEventTooLarge) || - result.Committed || result.RetryableBeforeCommit || calls.Load() != 0 { - t.Fatalf("ForwardStream() result = %#v calls=%d, want terminal pre-commit protocol error", - result, calls.Load()) - } - if downstream.status != 0 || downstream.body.Len() != 0 || downstream.flushes != 0 { - t.Fatalf("downstream was touched before oversized event rejection: %#v", downstream) - } -} - -func TestForwardStreamRejectsOversizedAliasedEventAsProtocolError(t *testing.T) { - event := "data: " + strings.Repeat("x", maxSSEEventBytes-len("data: \n\n")+1) + "\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, event) - writer.(http.Flusher).Flush() - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public","stream":true}`) - input.ExternalModel = "public" - input.UpstreamModelID = "provider" - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if !errors.Is(result.Err, ErrUpstreamProtocol) || !errors.Is(result.Err, errSSEEventTooLarge) || - result.Committed || result.RetryableBeforeCommit { - t.Fatalf("ForwardStream() result = %#v", result) - } -} - -func TestForwardStreamRejectsMalformedAliasedEventAsProtocolError(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, "data: not-json\n\n") - writer.(http.Flusher).Flush() - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public","stream":true}`) - input.ExternalModel = "public" - input.UpstreamModelID = "provider" - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if !errors.Is(result.Err, ErrUpstreamProtocol) || result.Committed || result.RetryableBeforeCommit { - t.Fatalf("ForwardStream() result = %#v, want terminal pre-commit protocol error", result) - } - if downstream.status != 0 || downstream.body.Len() != 0 || downstream.flushes != 0 { - t.Fatalf("downstream was touched before malformed event rejection: %#v", downstream) - } -} - -func TestForwardStreamSanitizesAliasedErrorEventPayloads(t *testing.T) { - const ( - upstreamModel = "org/model" - externalModel = "public-model" - secret = "stream/secret" - ) - stream := `data: {"ready":true}` + "\n\n" + - "event: error\n" + - `data: {"model":"org/model","org\/model":"org\u002fmodel failed","credential":"stream\u002fsecret"}` + "\n\n" + - "event: error\n" + - `data: {"message":"later org\/model stream\u002fsecret"}` + "\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, stream) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - input.APIKey = secret - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - payloads := decodeSSEJSONPayloads(t, downstream.body.Bytes()) - if len(payloads) != 3 { - t.Fatalf("decoded payload count = %d, wire=%q", len(payloads), downstream.body.String()) - } - if payloads[1]["model"] != externalModel || - payloads[1][externalModel] != externalModel+" failed" || - payloads[1]["credential"] != redact.Placeholder { - t.Fatalf("first error payload = %#v", payloads[1]) - } - if payloads[2]["message"] != "later "+externalModel+" "+redact.Placeholder { - t.Fatalf("later error payload = %#v", payloads[2]) - } -} - -func TestForwardStreamSanitizesDataOnlyAliasedErrorPayloads(t *testing.T) { - const ( - upstreamModel = "provider-model" - externalModel = "public-model" - ) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, `data: {"ready":true} - -data: {"error":{"message":"provider-model failed"}} - -`) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - payloads := decodeSSEJSONPayloads(t, downstream.body.Bytes()) - if len(payloads) != 2 { - t.Fatalf("decoded payloads = %#v, wire=%q", payloads, downstream.body.String()) - } - errorObject := payloads[1]["error"].(map[string]any) - if errorObject["message"] != externalModel+" failed" { - t.Fatalf("error message = %#v, want %q", errorObject["message"], externalModel+" failed") - } -} - -func TestForwardStreamSanitizesTypeErrorAliasedPayloads(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, `data: {"ready":true} - -data: {"type":"error","message":"provider-model failed"} - -`) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - payloads := decodeSSEJSONPayloads(t, downstream.body.Bytes()) - if len(payloads) != 2 || payloads[1]["message"] != "public-model failed" { - t.Fatalf("sanitized payloads = %#v, wire=%q", payloads, downstream.body.String()) - } -} - -func TestForwardStreamAliasOnlyRewritesModelFieldsInSuccessfulEvents(t *testing.T) { - const ( - upstreamModel = "provider-model" - externalModel = "public-model" - secret = "provider-secret" - content = "provider-model " + redact.Placeholder - ) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, `data: {"model":"provider-model","choices":[{"delta":{"content":"provider-model provider-secret","tool_calls":[{"function":{"arguments":"{\"model\":\"provider-model\",\"key\":\"provider-secret\"}"}}]}}]} - -`) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - input.APIKey = secret - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - payloads := decodeSSEJSONPayloads(t, downstream.body.Bytes()) - if len(payloads) != 1 || payloads[0]["model"] != externalModel { - t.Fatalf("rewritten payloads = %#v, wire=%q", payloads, downstream.body.String()) - } - choices := payloads[0]["choices"].([]any) - delta := choices[0].(map[string]any)["delta"].(map[string]any) - if delta["content"] != content { - t.Fatalf("generated content = %#v, want %q", delta["content"], content) - } - toolCalls := delta["tool_calls"].([]any) - arguments := toolCalls[0].(map[string]any)["function"].(map[string]any)["arguments"] - if arguments != `{"model":"provider-model","key":"[REDACTED]"}` { - t.Fatalf("tool arguments = %#v, want provider model preserved and credential redacted", arguments) - } -} - -func TestForwardStreamUsesFinalSSEEventTypeForAliasSanitization(t *testing.T) { - for _, test := range []struct { - name string - eventLines string - }{ - {name: "later message event", eventLines: "event: error\nevent: message\n"}, - {name: "later empty event", eventLines: "event: error\nevent\n"}, - } { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, test.eventLines+`data: {"model":"provider-model","content":"provider-model is generated content"}`+"\n\n") - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - payloads := decodeSSEJSONPayloads(t, downstream.body.Bytes()) - if len(payloads) != 1 || payloads[0]["model"] != "public-model" || - payloads[0]["content"] != "provider-model is generated content" { - t.Fatalf("rewritten payloads = %#v, wire=%q", payloads, downstream.body.String()) - } - }) - } -} - -func TestForwardStreamSanitizesUnaliasedErrorEventPayloads(t *testing.T) { - const secret = "stream/secret" - stream := `data: {"ready":true}` + "\n\n" + - "event: error\n" + - `data: {"credential":"stream\u002fsecret","raw":"stream/secret"}` + "\n\n" + - "event: error\n" + - `data: {"message":"later stream\/secret"}` + "\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - setRepresentationMetadata(writer.Header()) - _, _ = io.WriteString(writer, stream) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.APIKey = secret - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - payloads := decodeSSEJSONPayloads(t, downstream.body.Bytes()) - if len(payloads) != 3 || payloads[1]["credential"] != redact.Placeholder || - payloads[1]["raw"] != redact.Placeholder || - payloads[2]["message"] != "later "+redact.Placeholder { - t.Fatalf("sanitized unaliased payloads = %#v, wire=%q", payloads, downstream.body.String()) - } - assertRepresentationMetadata(t, result.Header, false) -} - -func TestForwardStreamPreservesHeuristicCredentialLikeContent(t *testing.T) { - for _, test := range []struct { - name string - externalModel string - upstreamModel string - wantModel string - }{ - {name: "no alias"}, - { - name: "alias", externalModel: "public-model", - upstreamModel: "provider-model", wantModel: "public-model", - }, - } { - t.Run(test.name, func(t *testing.T) { - payload := `{"token":"demo-value","api_key":"example"}` - if test.upstreamModel != "" { - payload = `{"model":"provider-model","token":"demo-value","api_key":"example"}` - } - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, "data: "+payload+"\n\n") - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.APIKey = "actual/secret" - input.ExternalModel = test.externalModel - input.UpstreamModelID = test.upstreamModel - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - payloads := decodeSSEJSONPayloads(t, downstream.body.Bytes()) - if len(payloads) != 1 || payloads[0]["token"] != "demo-value" || payloads[0]["api_key"] != "example" { - t.Fatalf("preserved payloads = %#v, wire=%q", payloads, downstream.body.String()) - } - if test.wantModel != "" && payloads[0]["model"] != test.wantModel { - t.Fatalf("rewritten model = %#v, want %q", payloads[0]["model"], test.wantModel) - } - }) - } -} - -func TestForwardStreamUnaliasedCredentialRewriteFailureRespectsCommitBoundary(t *testing.T) { - const ( - secret = "secret" - collision = `{"secret":"first","[REDACTED]":"second"}` - ) - for _, test := range []struct { - name string - stream string - wantCommit bool - wantBody string - }{ - { - name: "first error event fails before commit", - stream: "event: error\ndata: " + collision + "\n\n", - }, - { - name: "later error event terminates committed stream", - stream: `data: {"ok":true}` + "\n\n" + - "event: error\ndata: " + collision + "\n\n", - wantCommit: true, - wantBody: `data: {"ok":true}` + "\n\n", - }, - } { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, test.stream) - })) - defer upstream.Close() - - var calls atomic.Int32 - input := streamForwardInput(upstream.URL) - input.APIKey = secret - input.OnStreamReady = func() { calls.Add(1) } - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - wantCalls := int32(0) - if test.wantCommit { - wantCalls = 1 - } - if calls.Load() != wantCalls { - t.Fatalf("stream-ready calls = %d, want %d", calls.Load(), wantCalls) - } - if !errors.Is(result.Err, ErrUpstreamProtocol) || result.Committed != test.wantCommit { - t.Fatalf("ForwardStream() result = %#v, want protocol error committed=%t", result, test.wantCommit) - } - if !test.wantCommit && result.RetryableBeforeCommit { - t.Fatalf("request-written pre-commit failure is retryable: %#v", result) - } - if downstream.body.String() != test.wantBody { - t.Fatalf("downstream body = %q, want %q", downstream.body.String(), test.wantBody) - } - }) - } -} - -func TestForwardStreamRedactsEscapedAPIKeyBeforeOverlappingModelLiteral(t *testing.T) { - const ( - upstreamModel = "provider" - externalModel = "public" - secret = "provider/secret" - ) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, `data: {"ready":true} - -event: error -data: {"message":"provider\u002fsecret"} - -`) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - input.APIKey = secret - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - payloads := decodeSSEJSONPayloads(t, downstream.body.Bytes()) - if len(payloads) != 2 || payloads[1]["message"] != redact.Placeholder { - t.Fatalf("sanitized overlapping payloads = %#v, wire=%q", payloads, downstream.body.String()) - } -} - -func TestForwardStreamAliasPayloadRewriteFailureRespectsCommitBoundary(t *testing.T) { - const ( - upstreamModel = "provider-model" - externalModel = "public-model" - collision = `{"provider-model":"first","public-model":"second"}` - ) - for _, test := range []struct { - name string - stream string - wantCommit bool - wantBody string - }{ - { - name: "first error event fails before commit", - stream: "event: error\ndata: " + collision + "\n\n", - }, - { - name: "later error event terminates committed stream", - stream: `data: {"model":"provider-model","delta":"ok"}` + "\n\n" + - "event: error\ndata: " + collision + "\n\n", - wantCommit: true, - wantBody: `data: {"delta":"ok","model":"public-model"}` + "\n\n", - }, - } { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, test.stream) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if !errors.Is(result.Err, ErrUpstreamProtocol) || result.Committed != test.wantCommit { - t.Fatalf("ForwardStream() result = %#v, want protocol error committed=%t", result, test.wantCommit) - } - if !test.wantCommit && result.RetryableBeforeCommit { - t.Fatalf("request-written pre-commit failure is retryable: %#v", result) - } - if downstream.body.String() != test.wantBody { - t.Fatalf("downstream body = %q, want %q", downstream.body.String(), test.wantBody) - } - }) - } -} - -func TestForwardStreamReturnsSafeBoundedNonSuccessResponse(t *testing.T) { - const secret = "custom-upstream-secret" - tests := []struct { - name string - body string - wantBody string - }{ - {name: "inspectable", body: `{"error":{"api_key":"` + secret + `"}}`, wantBody: `{"error":{"api_key":"[REDACTED]"}}`}, - {name: "over limit", body: strings.Repeat("x", maxStreamingErrorBodyBytes) + secret, wantBody: redact.Placeholder}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - setRepresentationMetadata(writer.Header()) - writer.WriteHeader(http.StatusUnauthorized) - _, _ = writer.Write([]byte(tt.body)) - })) - defer upstream.Close() - - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - downstream := newRecordingResponseWriter() - input := streamForwardInput(upstream.URL) - input.APIKey = secret - result := forwarder.ForwardStream(context.Background(), input, downstream) - - if result.Err != nil || result.Committed || result.StatusCode != http.StatusUnauthorized { - t.Fatalf("ForwardStream() result = %#v", result) - } - if string(result.Body) != tt.wantBody || string(result.ClassificationBody) != tt.wantBody { - t.Fatalf("safe bodies = %q / %q, want %q", result.Body, result.ClassificationBody, tt.wantBody) - } - if bytes.Contains(result.Body, []byte(secret)) || bytes.Contains(result.ClassificationBody, []byte(secret)) { - t.Fatal("streaming error result leaked plaintext key") - } - if downstream.status != 0 || downstream.body.Len() != 0 { - t.Fatal("ForwardStream() wrote non-success response before Handler verdict") - } - if tt.name == "over limit" { - if result.Header.Get("Content-Length") != strconv.Itoa(len(redact.Placeholder)) { - t.Fatalf("Content-Length = %q", result.Header.Get("Content-Length")) - } - assertRepresentationMetadata(t, result.Header, false) - } - }) - } -} - -func TestForwardStreamBuffersCodedNonSuccessRepresentation(t *testing.T) { - plain := []byte(`{"error":{"code":"rate_limited"}}`) - for _, encoding := range []contentcoding.Encoding{ - contentcoding.Identity, - contentcoding.Gzip, - contentcoding.Brotli, - contentcoding.Deflate, - contentcoding.Zstd, - } { - t.Run(string(encoding), func(t *testing.T) { - wire := encodeContentCodingForGatewayTest(t, encoding, plain) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", string(encoding)) - writer.WriteHeader(http.StatusUnauthorized) - _, _ = writer.Write(wire) - })) - defer upstream.Close() - - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), streamForwardInput(upstream.URL), downstream, - ) - if result.Err != nil || result.Committed || result.StatusCode != http.StatusUnauthorized || - !bytes.Equal(result.Body, plain) || !bytes.Equal(result.ClassificationBody, plain) || - len(headerFieldValues(result.Header, "Content-Encoding")) != 0 || - result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { - t.Fatalf("ForwardStream() result = %#v, want buffered plaintext error", result) - } - if downstream.status != 0 || downstream.body.Len() != 0 { - t.Fatalf("ForwardStream() wrote non-success response before Handler verdict: %d/%q", downstream.status, downstream.body.String()) - } - }) - } - - t.Run("multiple content-encoding values fail closed", func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Add("Content-Encoding", "identity") - writer.Header().Add("Content-Encoding", "gzip") - writer.WriteHeader(http.StatusBadGateway) - _, _ = writer.Write([]byte("opaque")) - })) - defer upstream.Close() - - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), streamForwardInput(upstream.URL), newRecordingResponseWriter(), - ) - if result.Err != nil || result.Committed || result.StatusCode != http.StatusBadGateway || - string(result.Body) != redact.Placeholder || string(result.ClassificationBody) != redact.Placeholder || - len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { - t.Fatalf("ForwardStream() result = %#v, want fail-closed buffered error", result) - } - }) -} - -func TestForwardStreamContentCodingBoundary(t *testing.T) { - for _, test := range []struct { - name string - values []string - wantError bool - }{ - {name: "absent"}, - {name: "empty", values: []string{""}}, - {name: "identity", values: []string{"IDENTITY"}}, - {name: "gzip", values: []string{"gzip"}, wantError: true}, - {name: "brotli", values: []string{"br"}, wantError: true}, - {name: "deflate", values: []string{"deflate"}, wantError: true}, - {name: "zstd", values: []string{"zstd"}, wantError: true}, - {name: "stacked", values: []string{"gzip, br"}, wantError: true}, - {name: "unknown", values: []string{"compress"}, wantError: true}, - {name: "multiple", values: []string{"identity", "gzip"}, wantError: true}, - } { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - for _, value := range test.values { - writer.Header().Add("Content-Encoding", value) - } - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: first\n\n")) - })) - defer upstream.Close() - - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), streamForwardInput(upstream.URL), downstream, - ) - if test.wantError { - if !errors.Is(result.Err, ErrUpstreamProtocol) || result.Committed || - downstream.status != 0 || downstream.body.Len() != 0 { - t.Fatalf("ForwardStream() result/downstream = %#v/%d/%q", result, downstream.status, downstream.body.String()) - } - return - } - if result.Err != nil || !result.Committed || downstream.status != http.StatusOK || - downstream.body.String() != "data: first\n\n" { - t.Fatalf("ForwardStream() result/downstream = %#v/%d/%q", result, downstream.status, downstream.body.String()) - } - }) - } - - t.Run("case-colliding header map is rejected", func(t *testing.T) { - err := validateStreamContentEncoding(http.Header{ - "Content-Encoding": {"identity"}, - "content-encoding": {"identity"}, - }) - if !errors.Is(err, ErrUpstreamProtocol) { - t.Fatalf("validateStreamContentEncoding() error = %v, want upstream protocol error", err) - } - }) -} - -func TestForwardStreamTimesOutBeforeCompleteFirstEvent(t *testing.T) { - upstreamCanceled := make(chan struct{}) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - writer.WriteHeader(http.StatusOK) - _, _ = writer.Write([]byte("data: partial\n")) - writer.(http.Flusher).Flush() - <-request.Context().Done() - close(upstreamCanceled) - })) - defer upstream.Close() - - var calls atomic.Int32 - input := streamForwardInput(upstream.URL) - input.Group.Timeouts.FirstByte = 25 * time.Millisecond - input.OnStreamReady = func() { calls.Add(1) } - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - downstream := newRecordingResponseWriter() - result := forwarder.ForwardStream(context.Background(), input, downstream) - - if !errors.Is(result.Err, context.DeadlineExceeded) || result.Committed || - result.RetryableBeforeCommit || calls.Load() != 0 { - t.Fatalf("ForwardStream() timeout result = %#v calls=%d", result, calls.Load()) - } - if downstream.status != 0 || downstream.body.Len() != 0 { - t.Fatalf("partial event reached downstream: status/body=%d/%q", downstream.status, downstream.body.String()) - } - select { - case <-upstreamCanceled: - case <-time.After(time.Second): - t.Fatal("first-event timeout did not cancel upstream request") - } -} - -func TestForwardStreamDoesNotRetryParentDeadline(t *testing.T) { - ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second)) - defer cancel() - - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - result := forwarder.ForwardStream(ctx, streamForwardInput("http://127.0.0.1:1"), newRecordingResponseWriter()) - - if !errors.Is(result.Err, context.DeadlineExceeded) { - t.Fatalf("ForwardStream() error = %v, want parent deadline", result.Err) - } - if result.RetryableBeforeCommit { - t.Fatalf("ForwardStream() marked parent deadline retryable: %#v", result) - } -} - -func TestReleaseCommittedRequestReplayReleasesParsedBodyWithoutMutatingHTTPRequest(t *testing.T) { - input := streamForwardInput("https://example.com") - wantBody := bytes.Clone(input.Request.Body) - request, _, replay, err := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).newUpstreamRequest(context.Background(), input, true) - if err != nil { - t.Fatalf("newUpstreamRequest() error = %v", err) - } - if request.GetBody != nil { - t.Fatal("newUpstreamRequest() exposed GetBody before commit") - } - originalBody := request.Body - - releaseCommittedRequestReplay(input.Request, replay) - - if input.Request.Body != nil { - t.Fatal("ParsedRequest.Body still retains the replay buffer") - } - if request.Body != originalBody || request.GetBody != nil { - t.Fatal("committed release mutated fields owned by the HTTP transport") - } - activeBody, err := io.ReadAll(request.Body) - if err != nil { - t.Fatalf("read active request body: %v", err) - } - if !bytes.Equal(activeBody, wantBody) { - t.Fatalf("active request body = %q, want %q", activeBody, wantBody) - } -} - -func TestForwardStreamPreservesParsedRequestBeforeCommit(t *testing.T) { - requestStarted := make(chan struct{}) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - _, _ = io.Copy(io.Discard, request.Body) - close(requestStarted) - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: partial\n")) - writer.(http.Flusher).Flush() - <-request.Context().Done() - })) - defer upstream.Close() - - var calls atomic.Int32 - input := streamForwardInput(upstream.URL) - input.OnStreamReady = func() { calls.Add(1) } - input.Request.RawQuery = "trace=true" - input.Request.Header.Set("X-Test", "one") - want := cloneParsedRequestForGatewayTest(input.Request) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - done := make(chan UpstreamResult, 1) - go func() { - done <- NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - ctx, input, newRecordingResponseWriter(), - ) - }() - - waitForSignal(t, requestStarted, "upstream request") - if !reflect.DeepEqual(input.Request, want) { - t.Fatalf("ForwardStream() mutated ParsedRequest before commit:\n got %#v\nwant %#v", input.Request, want) - } - cancel() - select { - case result := <-done: - if result.Committed { - t.Fatalf("ForwardStream() committed partial event: %#v", result) - } - if calls.Load() != 0 { - t.Fatalf("stream-ready calls after pre-commit cancellation = %d, want 0", calls.Load()) - } - if !reflect.DeepEqual(input.Request, want) { - t.Fatalf("ForwardStream() mutated ParsedRequest after pre-commit return:\n got %#v\nwant %#v", input.Request, want) - } - case <-time.After(time.Second): - t.Fatal("ForwardStream() did not finish after pre-commit cancellation") - } -} - -func TestForwardStreamReleasesParsedBodyAfterCommit(t *testing.T) { - requestBodyRead := make(chan struct{}) - firstSSEEvent := make(chan struct{}) - downstream := newBarrierStreamResponseWriter() - - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - _, _ = io.Copy(io.Discard, request.Body) - close(requestBodyRead) - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: ready\n\n")) - writer.(http.Flusher).Flush() - close(firstSSEEvent) - <-downstream.release - })) - defer func() { - downstream.unblock() - upstream.Close() - }() - - input := streamForwardInput(upstream.URL) - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - done := make(chan UpstreamResult, 1) - go func() { - done <- forwarder.ForwardStream(context.Background(), input, downstream) - }() - - waitForSignal(t, requestBodyRead, "upstream request body read") - waitForSignal(t, firstSSEEvent, "first upstream SSE event") - waitForSignal(t, downstream.firstFlush, "first downstream SSE event flush") - if input.Request.Body != nil { - t.Fatalf("committed stream still retains %d request body bytes", len(input.Request.Body)) - } - downstream.unblock() - - select { - case result := <-done: - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() result = %#v", result) - } - case <-time.After(time.Second): - t.Fatal("ForwardStream() did not finish after upstream release") - } -} - -func cloneParsedRequestForGatewayTest(request *dialect.ParsedRequest) *dialect.ParsedRequest { - clone := *request - clone.Header = request.Header.Clone() - clone.Body = bytes.Clone(request.Body) - return &clone -} - -func TestStreamingClientConfigHasNoTotalTimeout(t *testing.T) { - timeouts := state.TimeoutConfig{ - Connect: 2 * time.Second, FirstByte: 3 * time.Second, - Request: 4 * time.Second, StreamIdle: 5 * time.Second, - } - config := streamingClientConfig(timeouts) - - if config.ConnectTimeout != timeouts.Connect || config.ResponseHeaderTimeout != timeouts.FirstByte { - t.Fatalf("stream connect/header timeouts = %s/%s", config.ConnectTimeout, config.ResponseHeaderTimeout) - } - if config.RequestTimeout != 0 { - t.Fatalf("stream RequestTimeout = %s, want 0", config.RequestTimeout) - } - if !config.DisableCompression { - t.Fatal("stream DisableCompression = false; transport must not implicitly decode while request normalization explicitly negotiates identity") - } -} - -func TestNonStreamingClientConfigDisablesImplicitCompression(t *testing.T) { - config := nonStreamingClientConfig(state.TimeoutConfig{}) - if !config.DisableCompression { - t.Fatal("non-stream DisableCompression = false; transport must not implicitly decode while request normalization explicitly negotiates identity") - } -} - -func TestGatewayClientConfigsDisableRedirects(t *testing.T) { - timeouts := state.TimeoutConfig{ - Connect: 2 * time.Second, FirstByte: 3 * time.Second, - Request: 4 * time.Second, StreamIdle: 5 * time.Second, - } - for name, config := range map[string]*platformhttp.Config{ - "non-streaming": nonStreamingClientConfig(timeouts), - "streaming": streamingClientConfig(timeouts), - } { - if !config.DisableRedirects { - t.Errorf("%s client DisableRedirects = false, want true", name) - } - } -} - -func TestForwarderPreservesEndToEndRequestAndSuccessfulResponse(t *testing.T) { - var received *http.Request - var receivedBody []byte - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - received = request.Clone(request.Context()) - received.Header = request.Header.Clone() - receivedBody, _ = io.ReadAll(request.Body) - writer.Header().Set("Content-Type", "application/json") - writer.Header().Set("X-Upstream", "kept") - writer.Header().Set("Connection", "X-Upstream-Hop") - writer.Header().Add("Connection", "X-Upstream-Hop-Second") - writer.Header().Set("X-Upstream-Hop", "drop") - writer.Header().Set("X-Upstream-Hop-Second", "drop") - writer.WriteHeader(http.StatusOK) - _, _ = writer.Write([]byte(`{"id":"chatcmpl-1"}`)) - })) - defer upstream.Close() - - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - result := forwarder.Forward(context.Background(), ForwardInput{ - Dialect: dialect.NewOpenAI(upstream.Client()), - Group: state.GroupView{ - ID: 1, Name: "openai", UpstreamURL: upstream.URL, - Timeouts: state.TimeoutConfig{Connect: time.Second, FirstByte: time.Second, Request: time.Second}, - HeaderRules: state.HeaderRules{Set: map[string]string{"X-Custom": "prefix-${API_KEY}"}}, - }, - APIKey: "sk-upstream-secret", - Request: &dialect.ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/chat/completions", - RawQuery: "trace=true", - Header: http.Header{ - "Content-Type": {"application/json"}, - "X-Passthrough": {"kept"}, - "Authorization": {"Bearer downstream-key"}, - "Connection": {"X-Drop", "X-Drop-Second"}, - "X-Drop": {"drop"}, - "X-Drop-Second": {"drop"}, - }, - Body: []byte(`{"model":"gpt-4o"}`), - }, - }) - - if result.Err != nil || result.StatusCode != http.StatusOK || !result.RequestWritten { - t.Fatalf("Forward() result = %#v", result) - } - if string(result.Body) != `{"id":"chatcmpl-1"}` || len(result.ClassificationBody) != 0 { - t.Fatalf("response bodies = wire %q classify %q", result.Body, result.ClassificationBody) - } - if result.Header.Get("X-Upstream") != "kept" || - result.Header.Get("X-Upstream-Hop") != "" || - result.Header.Get("X-Upstream-Hop-Second") != "" { - t.Fatalf("response headers = %#v", result.Header) - } - if received.URL.RawQuery != "trace=true" || string(receivedBody) != `{"model":"gpt-4o"}` { - t.Fatalf("upstream request URL/body = %s?%s %q", received.URL.Path, received.URL.RawQuery, receivedBody) - } - if received.Header.Get("Authorization") != "Bearer sk-upstream-secret" || - received.Header.Get("X-Custom") != "prefix-sk-upstream-secret" || - received.Header.Get("X-Passthrough") != "kept" { - t.Fatalf("upstream headers = %#v", received.Header) - } - if received.Header.Get("X-Drop") != "" || - received.Header.Get("X-Drop-Second") != "" || - strings.Contains(received.Header.Get("Authorization"), "downstream-key") { - t.Fatalf("upstream retained forbidden header: %#v", received.Header) - } - if got := received.Header.Get("User-Agent"); got != "" { - t.Fatalf("upstream User-Agent = %q, want downstream absence preserved", got) - } -} - -func TestForwarderIsolatesAliasedModelFromNonStreamingErrors(t *testing.T) { - const ( - externalModel = "public-model" - upstreamModel = "provider-model" - secret = "custom-upstream-secret" - ) - plain := []byte(`{"error":{"message":"model provider-model rejected custom-upstream-secret"}}`) - - for _, encoding := range []string{"", "gzip"} { - name := "identity" - if encoding != "" { - name = encoding - } - t.Run(name, func(t *testing.T) { - wire := plain - if encoding != "" { - wire = encodeContentCodingForGatewayTest(t, contentcoding.Encoding(encoding), plain) - } - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - writer.Header().Set("X-Upstream-Model", upstreamModel) - if encoding != "" { - writer.Header().Set("Content-Encoding", encoding) - } - setRepresentationMetadata(writer.Header()) - writer.WriteHeader(http.StatusBadRequest) - _, _ = writer.Write(wire) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public-model"}`) - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - input.APIKey = secret - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - - if result.Err != nil || result.StatusCode != http.StatusBadRequest { - t.Fatalf("Forward() result = %#v", result) - } - if !bytes.Contains(result.ClassificationBody, []byte(upstreamModel)) || - bytes.Contains(result.ClassificationBody, []byte(secret)) || - !bytes.Contains(result.ClassificationBody, []byte(redact.Placeholder)) { - t.Fatalf("ClassificationBody = %q", result.ClassificationBody) - } - downstreamBody := result.Body - if bytes.Contains(downstreamBody, []byte(upstreamModel)) || - !bytes.Contains(downstreamBody, []byte(externalModel)) || - bytes.Contains(downstreamBody, []byte(secret)) || - !bytes.Contains(downstreamBody, []byte(redact.Placeholder)) { - t.Fatalf("downstream body = %q", downstreamBody) - } - if result.Header.Get("Content-Type") != "application/json" || - result.Header.Get("Content-Encoding") != "" || - result.Header.Get("Content-Length") != strconv.Itoa(len(result.Body)) { - t.Fatalf("representation headers = %#v", result.Header) - } - assertHeadersDoNotContain(t, result.Header, upstreamModel) - assertRepresentationMetadata(t, result.Header, false) - }) - } -} - -func TestForwarderRewritesEscapedAliasedModelInJSONErrors(t *testing.T) { - const upstreamModel = "org/model" - externalModel := "public\"\\model" - plain := []byte(`{"org\/model":"org\u002fmodel unavailable","nested":["org/model"],"number":9007199254740993}`) - - for _, encoding := range []string{"", "gzip"} { - name := "identity" - wire := plain - if encoding != "" { - name = encoding - wire = encodeContentCodingForGatewayTest(t, contentcoding.Encoding(encoding), plain) - } - t.Run(name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - if encoding != "" { - writer.Header().Set("Content-Encoding", encoding) - } - writer.WriteHeader(http.StatusBadRequest) - _, _ = writer.Write(wire) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - if result.Err != nil || result.StatusCode != http.StatusBadRequest { - t.Fatalf("Forward() result = %#v", result) - } - if !bytes.Equal(result.ClassificationBody, plain) { - t.Fatalf("ClassificationBody = %q, want original safe JSON %q", result.ClassificationBody, plain) - } - downstreamBody := result.Body - if !json.Valid(downstreamBody) { - t.Fatalf("downstream body is invalid JSON: %q", downstreamBody) - } - if !bytes.Contains(downstreamBody, []byte("9007199254740993")) { - t.Fatalf("downstream body lost JSON number precision: %q", downstreamBody) - } - var decoded map[string]any - if err := json.Unmarshal(downstreamBody, &decoded); err != nil { - t.Fatalf("decode downstream JSON: %v", err) - } - if decoded[externalModel] != externalModel+" unavailable" { - t.Fatalf("rewritten object entry = %#v", decoded) - } - nested, ok := decoded["nested"].([]any) - if !ok || len(nested) != 1 || nested[0] != externalModel { - t.Fatalf("rewritten nested value = %#v", decoded["nested"]) - } - }) - } -} - -func TestForwarderRedactsEscapedAPIKeyInJSONErrors(t *testing.T) { - const secret = "secret/key" - plain := []byte(`{"secret\/key":"secret\u002fkey rejected"}`) - - for _, encoding := range []string{"", "gzip"} { - name := "identity" - wire := plain - if encoding != "" { - name = encoding - wire = encodeContentCodingForGatewayTest(t, contentcoding.Encoding(encoding), plain) - } - t.Run(name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - if encoding != "" { - writer.Header().Set("Content-Encoding", encoding) - } - writer.WriteHeader(http.StatusUnauthorized) - _, _ = writer.Write(wire) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.APIKey = secret - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - if result.Err != nil || result.StatusCode != http.StatusUnauthorized { - t.Fatalf("Forward() result = %#v", result) - } - downstreamBody := result.Body - for _, body := range [][]byte{downstreamBody, result.ClassificationBody} { - var decoded map[string]any - if err := json.Unmarshal(body, &decoded); err != nil { - t.Fatalf("decode safe body %q: %v", body, err) - } - if decoded[redact.Placeholder] != redact.Placeholder+" rejected" { - t.Fatalf("safe decoded body = %#v", decoded) - } - } - }) - } -} - -func TestPrepareErrorRepresentationFailsClosedWhenEscapedAPIKeyRewriteCollides(t *testing.T) { - const secret = "secret/key" - plain := []byte(`{"secret\/key":"leak","[REDACTED]":"safe"}`) - forwarder := &Forwarder{redactor: redact.New()} - headers := http.Header{"Content-Type": {"application/json"}} - prepared := forwarder.prepareErrorRepresentation(ForwardInput{APIKey: secret}, headers, plain, []string{secret}) - - if string(prepared.downstream) != redact.Placeholder || string(prepared.classification) != redact.Placeholder { - t.Fatalf("collision result = %#v", prepared) - } -} - -func TestPrepareErrorRepresentationFailsClosedBeforeRawAPIKeyCollisionRedaction(t *testing.T) { - const secret = "secret" - plain := []byte(`{"secret":"leak","[REDACTED]":"safe"}`) - forwarder := &Forwarder{redactor: redact.New()} - headers := http.Header{"Content-Type": {"application/json"}} - prepared := forwarder.prepareErrorRepresentation(ForwardInput{APIKey: secret}, headers, plain, []string{secret}) - - if string(prepared.downstream) != redact.Placeholder || string(prepared.classification) != redact.Placeholder { - t.Fatalf("collision result = %#v", prepared) - } -} - -func TestForwarderFailsClosedWhenJSONKeyRewriteCollides(t *testing.T) { - plain := []byte(`{"provider-model":"first","public-model":"second"}`) - forwarder := &Forwarder{redactor: redact.New()} - headers := http.Header{"Content-Type": {"application/json"}} - prepared := forwarder.prepareErrorRepresentation(ForwardInput{ - ExternalModel: "public-model", UpstreamModelID: "provider-model", - }, headers, plain, nil) - - if string(prepared.downstream) != redact.Placeholder || string(prepared.classification) != redact.Placeholder { - t.Fatalf("collision result = %#v", prepared) - } -} - -func TestPrepareErrorRepresentationFailsClosedWhenModelExpansionExceedsLimit(t *testing.T) { - plain := bytes.Repeat([]byte("x"), 64<<10) - forwarder := &Forwarder{redactor: redact.New()} - headers := make(http.Header) - prepared := forwarder.prepareErrorRepresentation(ForwardInput{ - ExternalModel: strings.Repeat("a", 32), UpstreamModelID: "x", - }, headers, plain, nil) - - if string(prepared.downstream) != redact.Placeholder || string(prepared.classification) != redact.Placeholder { - t.Fatalf("expansion result = %#v", prepared) - } -} - -func TestForwardStreamIsolatesAliasedModelFromNonSuccessResponse(t *testing.T) { - const ( - externalModel = "public-model" - upstreamModel = "provider-model" - secret = "custom-upstream-secret" - ) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - writer.Header().Set("X-Upstream-Model", upstreamModel) - setRepresentationMetadata(writer.Header()) - writer.WriteHeader(http.StatusTooManyRequests) - _, _ = writer.Write([]byte(`{"error":"provider-model rate limited custom-upstream-secret"}`)) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public-model","stream":true}`) - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - input.APIKey = secret - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || result.Committed || result.StatusCode != http.StatusTooManyRequests { - t.Fatalf("ForwardStream() result = %#v", result) - } - if bytes.Contains(result.Body, []byte(upstreamModel)) || - !bytes.Contains(result.Body, []byte(externalModel)) || - bytes.Contains(result.Body, []byte(secret)) { - t.Fatalf("downstream body = %q", result.Body) - } - if !bytes.Contains(result.ClassificationBody, []byte(upstreamModel)) || - bytes.Contains(result.ClassificationBody, []byte(secret)) { - t.Fatalf("ClassificationBody = %q", result.ClassificationBody) - } - assertHeadersDoNotContain(t, result.Header, upstreamModel) - assertRepresentationMetadata(t, result.Header, false) - if downstream.status != 0 || downstream.body.Len() != 0 { - t.Fatalf("ForwardStream() wrote error before Handler verdict: %#v", downstream) - } -} - -func TestForwarderPreservesErrorModelWithoutAlias(t *testing.T) { - const upstreamModel = "provider-model" - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("X-Upstream-Model", upstreamModel) - writer.WriteHeader(http.StatusBadRequest) - _, _ = writer.Write([]byte(upstreamModel)) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"provider-model"}`) - input.ExternalModel = upstreamModel - input.UpstreamModelID = upstreamModel - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - - if string(result.Body) != upstreamModel || string(result.ClassificationBody) != upstreamModel || - result.Header.Get("X-Upstream-Model") != upstreamModel { - t.Fatalf("non-alias response changed: %#v", result) - } -} - -func TestForwarderRedactsCompressedErrorAsPlaintext(t *testing.T) { - const secret = "custom-upstream-secret" - plain := []byte(`{"error":{"api_key":"` + secret + `","code":"invalid_api_key"}}`) - encoded := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", "gzip") - setRepresentationMetadata(writer.Header()) - writer.WriteHeader(http.StatusUnauthorized) - _, _ = writer.Write(encoded) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, secret, time.Second) - if result.Err != nil || result.StatusCode != http.StatusUnauthorized { - t.Fatalf("Forward() result = %#v", result) - } - if result.Header.Get("Content-Encoding") != "" { - t.Fatalf("Content-Encoding = %q, want absent", result.Header.Get("Content-Encoding")) - } - for _, body := range [][]byte{result.Body, result.ClassificationBody} { - if bytes.Contains(body, []byte(secret)) || !bytes.Contains(body, []byte(redact.Placeholder)) { - t.Fatalf("safe body = %q, want placeholder and no secret", body) - } - } - if result.Header.Get("Content-Length") != strconv.Itoa(len(result.Body)) { - t.Fatalf("Content-Length = %q, body length = %d", result.Header.Get("Content-Length"), len(result.Body)) - } - assertRepresentationMetadata(t, result.Header, false) -} - -func TestForwarderErrorRepresentationsArePlain(t *testing.T) { - plain := []byte(`{"error":{"code":"rate_limited"}}`) - for _, encoding := range []contentcoding.Encoding{ - contentcoding.Identity, - contentcoding.Gzip, - contentcoding.Brotli, - contentcoding.Deflate, - contentcoding.Zstd, - } { - t.Run(string(encoding), func(t *testing.T) { - wire := encodeContentCodingForGatewayTest(t, encoding, plain) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - if encoding != contentcoding.Identity { - writer.Header().Set("Content-Encoding", string(encoding)) - } - writer.WriteHeader(http.StatusTooManyRequests) - _, _ = writer.Write(wire) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "safe-upstream-key", time.Second) - if result.Err != nil || result.StatusCode != http.StatusTooManyRequests || - !bytes.Equal(result.Body, plain) || !bytes.Equal(result.ClassificationBody, plain) || - len(headerFieldValues(result.Header, "Content-Encoding")) != 0 || - result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { - t.Fatalf("Forward() result = %#v, want plaintext error representation", result) - } - }) - } -} - -func TestForwarderErrorRepresentationMetadataPolicy(t *testing.T) { - plain := []byte(`{"error":"safe"}`) - for _, test := range []struct { - name string - encoding contentcoding.Encoding - wantMetadataKept bool - }{ - {name: "unchanged identity", encoding: contentcoding.Identity, wantMetadataKept: true}, - {name: "decoded gzip", encoding: contentcoding.Gzip, wantMetadataKept: false}, - } { - t.Run(test.name, func(t *testing.T) { - wire := encodeContentCodingForGatewayTest(t, test.encoding, plain) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", string(test.encoding)) - setRepresentationMetadata(writer.Header()) - writer.WriteHeader(http.StatusBadGateway) - _, _ = writer.Write(wire) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "safe-upstream-key", time.Second) - if result.Err != nil || !bytes.Equal(result.Body, plain) { - t.Fatalf("Forward() result = %#v", result) - } - if test.wantMetadataKept { - for name, want := range map[string]string{ - "ETag": `"wire-v1"`, - "Digest": "sha-256=wire-digest", - "Content-MD5": "d2lyZQ==", - "Content-Range": "bytes 0-9/10", - "Content-Digest": "sha-256=:d2lyZQ==:", - "Repr-Digest": "sha-256=:cmVwcg==:", - } { - if got := result.Header.Get(name); got != want { - t.Fatalf("%s = %q, want %q", name, got, want) - } - } - if result.Header.Get("Signature") != "" || result.Header.Get("Signature-Input") != "" { - t.Fatalf("normalized identity representation retained stale signatures: %#v", result.Header) - } - } else { - assertRepresentationMetadata(t, result.Header, false) - } - }) - } -} - -func TestForwarderErrorContentCodingFailsClosed(t *testing.T) { - for _, test := range []struct { - name string - coding string - body []byte - apiKey string - }{ - {name: "unknown", coding: "unknown", body: []byte("opaque")}, - {name: "stacked", coding: "gzip, br", body: []byte("opaque")}, - {name: "corrupt gzip", coding: "gzip", body: []byte("not-gzip")}, - {name: "short key no special bypass", coding: "gzip", body: encodeContentCodingForGatewayTest(t, contentcoding.Gzip, []byte(`{"error":"safe"}`)), apiKey: "gzip"}, - } { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", test.coding) - writer.WriteHeader(http.StatusUnauthorized) - _, _ = writer.Write(test.body) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, test.apiKey, time.Second) - want := []byte(redact.Placeholder) - if test.name == "short key no special bypass" { - want = []byte(`{"error":"safe"}`) - } - if result.Err != nil || result.StatusCode != http.StatusUnauthorized || - !bytes.Equal(result.Body, want) || !bytes.Equal(result.ClassificationBody, want) || - len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { - t.Fatalf("Forward() result = %#v, want safe error representation", result) - } - }) - } -} - -func TestForwarderReturnsUnchangedCompressedErrorAsPlaintext(t *testing.T) { - plain := []byte(`{"error":{"code":"rate_limited"}}`) - encoded := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", "gzip") - setRepresentationMetadata(writer.Header()) - writer.WriteHeader(http.StatusTooManyRequests) - _, _ = writer.Write(encoded) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "custom-upstream-secret", time.Second) - if result.Err != nil || result.StatusCode != http.StatusTooManyRequests { - t.Fatalf("Forward() result = %#v", result) - } - if !bytes.Equal(result.Body, plain) { - t.Fatalf("downstream body = %q, want %q", result.Body, plain) - } - if !bytes.Equal(result.ClassificationBody, plain) { - t.Fatalf("ClassificationBody = %q, want %q", result.ClassificationBody, plain) - } - if len(headerFieldValues(result.Header, "Content-Encoding")) != 0 || - result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { - t.Fatalf("plaintext response headers = %#v", result.Header) - } - assertRepresentationMetadata(t, result.Header, false) -} - -func TestForwarderFailsClosedForUndecodableError(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", "unsupported") - setRepresentationMetadata(writer.Header()) - writer.WriteHeader(http.StatusBadGateway) - _, _ = writer.Write([]byte("opaque-secret-body")) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "opaque-secret-body", time.Second) - if result.Err != nil || result.StatusCode != http.StatusBadGateway { - t.Fatalf("Forward() result = %#v", result) - } - if result.Header.Get("Content-Encoding") != "" || - result.Header.Get("Content-Type") != "text/plain; charset=utf-8" || - result.Header.Get("Content-Length") != strconv.Itoa(len(redact.Placeholder)) || - string(result.Body) != redact.Placeholder || - string(result.ClassificationBody) != redact.Placeholder { - t.Fatalf("fail-closed result headers/body = %#v %q", result.Header, result.Body) - } - assertRepresentationMetadata(t, result.Header, false) -} - -func TestForwarderFailsClosedForMalformedEncoding(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", "gzip, br") - writer.WriteHeader(http.StatusBadGateway) - _, _ = writer.Write([]byte("opaque-body")) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "sk-secret", time.Second) - if result.Header.Get("Content-Encoding") != "" || - string(result.Body) != redact.Placeholder || - string(result.ClassificationBody) != redact.Placeholder { - t.Fatalf("fail-closed result headers/body = %#v %q", result.Header, result.Body) - } -} - -func TestForwarderFailsClosedForMultipleEncodingFieldValues(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Add("Content-Encoding", "identity") - writer.Header().Add("Content-Encoding", "gzip") - writer.WriteHeader(http.StatusBadGateway) - _, _ = writer.Write([]byte("opaque-multi-value-body")) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "sk-secret", time.Second) - if result.Err != nil || result.StatusCode != http.StatusBadGateway { - t.Fatalf("Forward() result = %#v", result) - } - if len(result.Header.Values("Content-Encoding")) != 0 || - string(result.Body) != redact.Placeholder || - string(result.ClassificationBody) != redact.Placeholder { - t.Fatalf("fail-closed result headers/body = %#v %q", result.Header, result.Body) - } -} - -func TestForwarderFailsClosedForUnsupportedEmptyBody(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", "unsupported") - writer.WriteHeader(http.StatusBadGateway) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "sk-secret", time.Second) - if result.Err != nil || result.StatusCode != http.StatusBadGateway { - t.Fatalf("Forward() result = %#v", result) - } - if result.Header.Get("Content-Encoding") != "" || - string(result.Body) != redact.Placeholder || - string(result.ClassificationBody) != redact.Placeholder { - t.Fatalf("fail-closed result headers/body = %#v %q", result.Header, result.Body) - } -} - -func TestForwarderFailsClosedForGzipEmptyBody(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", "gzip") - writer.WriteHeader(http.StatusBadGateway) - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "sk-secret", time.Second) - if result.Err != nil || result.StatusCode != http.StatusBadGateway { - t.Fatalf("Forward() result = %#v", result) - } - if result.Header.Get("Content-Encoding") != "" || - string(result.Body) != redact.Placeholder || - string(result.ClassificationBody) != redact.Placeholder { - t.Fatalf("fail-closed result headers/body = %#v %q", result.Header, result.Body) - } -} - -func TestForwarderBodylessResponseSemantics(t *testing.T) { - tests := []struct { - name string - method string - status int - contentEncoding string - contentLength string - wantErr bool - wantContentLength string - }{ - { - name: "head preserves oversized representation length without reading it", - method: http.MethodHead, - status: http.StatusOK, - contentEncoding: "identity", - contentLength: strconv.FormatInt(maxNonStreamingResponseBodyBytes+1, 10), - wantContentLength: strconv.FormatInt(maxNonStreamingResponseBodyBytes+1, 10), - }, - { - name: "not modified is bodyless", - method: http.MethodGet, - status: http.StatusNotModified, - contentEncoding: "identity", - contentLength: "321", - }, - { - name: "no content removes length", - method: http.MethodGet, - status: http.StatusNoContent, - contentEncoding: "identity", - contentLength: "321", - }, - { - name: "reset content keeps zero length only", - method: http.MethodGet, - status: http.StatusResetContent, - contentEncoding: "identity", - contentLength: "0", - wantContentLength: "0", - }, - { - name: "bodyless response rejects compressed encoding", - method: http.MethodHead, - status: http.StatusOK, - contentEncoding: "gzip", - wantErr: true, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if request.Method != test.method { - t.Errorf("request method = %s, want %s", request.Method, test.method) - } - if test.contentEncoding != "" { - writer.Header().Set("Content-Encoding", test.contentEncoding) - } - if test.contentLength != "" { - writer.Header().Set("Content-Length", test.contentLength) - } - writer.Header().Set("ETag", `"safe"`) - writer.WriteHeader(test.status) - _, _ = writer.Write([]byte("must not be read")) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ObserveUsage = false - input.Request.Method = test.method - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - - if test.wantErr { - if !errors.Is(result.Err, ErrUpstreamProtocol) || !result.RequestWritten { - t.Fatalf("Forward() result = %#v, want written upstream protocol error", result) - } - return - } - if result.Err != nil || result.StatusCode != test.status || !result.RequestWritten || len(result.Body) != 0 || len(result.ClassificationBody) != 0 { - t.Fatalf("Forward() result = %#v, want bodyless response", result) - } - if values := headerFieldValues(result.Header, "Content-Encoding"); len(values) != 0 { - t.Fatalf("Content-Encoding values = %#v, want absent", values) - } - if got := result.Header.Get("Content-Length"); got != test.wantContentLength { - t.Fatalf("Content-Length = %q, want %q", got, test.wantContentLength) - } - if result.Header.Get("ETag") != `"safe"` { - t.Fatalf("ETag = %q, want safe validator preserved", result.Header.Get("ETag")) - } - }) - } -} - -func TestForwarderHeadPartialContentRejectsInvalidContentRange(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Length", "3") - writer.Header().Set("Content-Range", "bytes 10-5/100") - writer.WriteHeader(http.StatusPartialContent) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ObserveUsage = false - input.Request.Method = http.MethodHead - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - if !errors.Is(result.Err, ErrUpstreamProtocol) || !result.RequestWritten { - t.Fatalf("Forward() result = %#v, want written upstream protocol error", result) - } -} - -func TestForwarderBodylessResponseInvalidatesSignaturesAfterHeaderNormalization(t *testing.T) { - tests := []struct { - name string - method string - status int - contentEncoding string - contentLength string - credentialHeader bool - }{ - { - name: "head removes identity coding", - method: http.MethodHead, - status: http.StatusOK, - contentEncoding: "identity", - contentLength: "12", - }, - { - name: "not modified removes identity coding", - method: http.MethodGet, - status: http.StatusNotModified, - contentEncoding: "identity", - contentLength: "12", - }, - { - name: "reset content removes nonzero length", - method: http.MethodGet, - status: http.StatusResetContent, - contentLength: "12", - }, - { - name: "head sanitizes credential header", - method: http.MethodHead, - status: http.StatusOK, - contentLength: "12", - credentialHeader: true, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - if test.contentEncoding != "" { - writer.Header().Set("Content-Encoding", test.contentEncoding) - } - if test.contentLength != "" { - writer.Header().Set("Content-Length", test.contentLength) - } - writer.Header().Set("Signature", "sig") - writer.Header().Set("Signature-Input", `sig1=("content-encoding" "content-length")`) - if test.credentialHeader { - writer.Header().Set("Authorization", "Bearer upstream-secret") - } - writer.WriteHeader(test.status) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.ObserveUsage = false - input.Request.Method = test.method - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - if result.Err != nil || result.StatusCode != test.status || !result.RequestWritten || - len(headerFieldValues(result.Header, "Signature")) != 0 || - len(headerFieldValues(result.Header, "Signature-Input")) != 0 { - t.Fatalf("Forward() result = %#v, want bodyless result without stale signatures", result) - } - }) - } -} - -func TestForwarderMarksConnectionFailureAsNotWritten(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) - url := upstream.URL - upstream.Close() - - result := testForward(t, url, "sk-secret", 200*time.Millisecond) - if result.Err == nil || result.RequestWritten { - t.Fatalf("connection failure result = %#v, want error before write", result) - } -} - -func TestForwarderMarksTimeoutAfterRequestWrite(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - _, _ = io.Copy(io.Discard, request.Body) - <-request.Context().Done() - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, "sk-secret", 25*time.Millisecond) - if result.Err == nil || !result.RequestWritten || !isTimeoutError(result.Err) { - t.Fatalf("post-write timeout result = %#v", result) - } -} - -func TestForwarderDecodesContentEncodingListedAsConnectionToken(t *testing.T) { - plain := []byte(`{"id":"safe"}`) - clients := platformhttp.NewHTTPClientManager() - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Group: state.GroupView{ - ID: 1, UpstreamURL: "https://api.example.test", - Timeouts: state.TimeoutConfig{Connect: time.Second, FirstByte: time.Second, Request: time.Second}, - }, - APIKey: "sk-upstream-secret", - Request: &dialect.ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/chat/completions", - Header: make(http.Header), - Body: []byte(`{"model":"gpt-4o"}`), - }, - } - clients.GetClient(nonStreamingClientConfig(input.Group.Timeouts)).Transport = forwarderRoundTripper(func(request *http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{ - "Connection": {"Content-Encoding"}, - "Content-Encoding": {"gzip"}, - }, - Body: io.NopCloser(bytes.NewReader(encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain))), - Request: request, - }, nil - }) - - result := NewForwarder(clients, redact.New()).Forward(context.Background(), input) - if result.Err != nil || !bytes.Equal(result.Body, plain) || len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { - t.Fatalf("Forward() result = %#v, want plaintext response without Content-Encoding", result) - } -} - -func TestForwarderBodylessResponseInvalidatesSignaturesAfterConnectionTokenCleanup(t *testing.T) { - clients := platformhttp.NewHTTPClientManager() - input := streamForwardInput("https://api.example.test") - input.ObserveUsage = false - input.Request.Method = http.MethodHead - clients.GetClient(nonStreamingClientConfig(input.Group.Timeouts)).Transport = forwarderRoundTripper(func(request *http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{ - "Connection": {"X-Signed"}, - "X-Signed": {"upstream-only"}, - "Content-Length": {"12"}, - "Signature": {"sig"}, - "Signature-Input": {`sig1=("x-signed" "content-length")`}, - }, - Body: io.NopCloser(strings.NewReader("")), - Request: request, - }, nil - }) - - result := NewForwarder(clients, redact.New()).Forward(context.Background(), input) - if result.Err != nil || result.StatusCode != http.StatusOK || !result.RequestWritten { - t.Fatalf("Forward() result = %#v, want bodyless forwarded response", result) - } - if len(headerFieldValues(result.Header, "X-Signed")) != 0 || - len(headerFieldValues(result.Header, "Signature")) != 0 || - len(headerFieldValues(result.Header, "Signature-Input")) != 0 { - t.Fatalf("bodyless response retained hop-by-hop cleanup or stale signatures: %#v", result.Header) - } - if result.Header.Get("Content-Length") != "12" { - t.Fatalf("Content-Length = %q, want preserved bodyless representation length", result.Header.Get("Content-Length")) - } -} - -func TestForwardStreamRejectsContentEncodingListedAsConnectionToken(t *testing.T) { - plain := []byte("data: first\n\n") - wire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) - clients := platformhttp.NewHTTPClientManager() - input := streamForwardInput("https://api.example.test") - clients.GetClient(streamingClientConfig(input.Group.Timeouts)).Transport = forwarderRoundTripper(func(request *http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{ - "Connection": {"Content-Encoding"}, - "Content-Encoding": {"gzip"}, - }, - Body: io.NopCloser(bytes.NewReader(wire)), - ContentLength: int64(len(wire)), - Request: request, - }, nil - }) - - downstream := newRecordingResponseWriter() - result := NewForwarder(clients, redact.New()).ForwardStream(context.Background(), input, downstream) - if !errors.Is(result.Err, ErrUpstreamProtocol) || result.Committed || result.RetryableBeforeCommit { - t.Fatalf("ForwardStream() result = %#v, want terminal pre-commit protocol error", result) - } - if downstream.status != 0 || downstream.body.Len() != 0 || downstream.flushes != 0 { - t.Fatalf("downstream was touched before protocol rejection: %#v", downstream) - } -} - -func TestForwarderDropsSignaturesFromUnchangedIdentityResponse(t *testing.T) { - plain := []byte(`{"id":"safe"}`) - clients := platformhttp.NewHTTPClientManager() - input := streamForwardInput("https://api.example.test") - input.ObserveUsage = false - input.Request.Body = []byte(`{"model":"gpt-4o"}`) - clients.GetClient(nonStreamingClientConfig(input.Group.Timeouts)).Transport = forwarderRoundTripper(func(request *http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{ - "Content-Length": {strconv.Itoa(len(plain))}, - "ETag": {`"wire-v1"`}, - "Signature": {"sig"}, - "Signature-Input": {`sig1=("etag")`}, - }, - Body: io.NopCloser(bytes.NewReader(plain)), - ContentLength: int64(len(plain)), - Request: request, - }, nil - }) - - result := NewForwarder(clients, redact.New()).Forward(context.Background(), input) - if result.Err != nil || !bytes.Equal(result.Body, plain) || !reflect.DeepEqual(headerFieldValues(result.Header, "ETag"), []string{`"wire-v1"`}) { - t.Fatalf("Forward() result = %#v, bodyEqual=%t etag=%#v want unchanged identity representation with safe validator", result, bytes.Equal(result.Body, plain), headerFieldValues(result.Header, "ETag")) - } - if len(headerFieldValues(result.Header, "Signature")) != 0 || - len(headerFieldValues(result.Header, "Signature-Input")) != 0 { - t.Fatalf("unchanged identity response retained signatures: %#v", result.Header) - } -} - -func TestForwarderBodylessResponseDropsUnchangedIdentitySignatures(t *testing.T) { - clients := platformhttp.NewHTTPClientManager() - input := streamForwardInput("https://api.example.test") - input.ObserveUsage = false - input.Request.Method = http.MethodHead - clients.GetClient(nonStreamingClientConfig(input.Group.Timeouts)).Transport = forwarderRoundTripper(func(request *http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{ - "Content-Length": {"12"}, - "ETag": {`"wire-v1"`}, - "Signature": {"sig"}, - "Signature-Input": {`sig1=("etag")`}, - }, - Body: io.NopCloser(strings.NewReader("")), - Request: request, - }, nil - }) - - result := NewForwarder(clients, redact.New()).Forward(context.Background(), input) - if result.Err != nil || result.StatusCode != http.StatusOK || !reflect.DeepEqual(headerFieldValues(result.Header, "ETag"), []string{`"wire-v1"`}) { - t.Fatalf("Forward() result = %#v, etag=%#v want unchanged bodyless representation with safe validator", result, headerFieldValues(result.Header, "ETag")) - } - if len(headerFieldValues(result.Header, "Signature")) != 0 || - len(headerFieldValues(result.Header, "Signature-Input")) != 0 { - t.Fatalf("unchanged bodyless response retained signatures: %#v", result.Header) - } -} - -type forwarderRoundTripper func(*http.Request) (*http.Response, error) - -func (roundTrip forwarderRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { - return roundTrip(request) -} - -func testForward(t *testing.T, upstreamURL, apiKey string, timeout time.Duration) UpstreamResult { - t.Helper() - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - return forwarder.Forward(context.Background(), ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Group: state.GroupView{ - ID: 1, UpstreamURL: testUpstreamBaseURL(upstreamURL, protocol.OpenAICompletions), - Timeouts: state.TimeoutConfig{Connect: timeout, FirstByte: timeout, Request: timeout}, - }, - APIKey: apiKey, - Request: &dialect.ParsedRequest{ - Method: http.MethodPost, Path: "/v1/chat/completions", - Header: make(http.Header), Body: []byte(`{"model":"gpt-4o"}`), - }, - }) -} - -func streamForwardInput(upstreamURL string) ForwardInput { - return ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - ObserveUsage: true, - Group: state.GroupView{ - ID: 1, Name: "openai", UpstreamURL: testUpstreamBaseURL(upstreamURL, protocol.OpenAICompletions), - Timeouts: state.TimeoutConfig{ - Connect: time.Second, FirstByte: time.Second, - Request: time.Second, StreamIdle: time.Second, - }, - }, - APIKey: "sk-upstream-secret", - Request: &dialect.ParsedRequest{ - Method: http.MethodPost, Path: "/v1/chat/completions", - Header: make(http.Header), Body: []byte(`{"model":"gpt-4o","stream":true}`), - }, - } -} - -type streamUsageInjectorDialect struct { - *dialect.OpenAI - inject func(*dialect.ParsedRequest) (*dialect.ParsedRequest, error) -} - -func (d streamUsageInjectorDialect) InjectStreamUsage(request *dialect.ParsedRequest) (*dialect.ParsedRequest, error) { - return d.inject(request) -} - -type dialectWithoutStreamUsageInjector struct{ dialect.Dialect } - -func TestForwardNeverInjectsNonStreamingRequest(t *testing.T) { - var received map[string]any - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if err := json.NewDecoder(request.Body).Decode(&received); err != nil { - t.Error(err) - } - writer.Header().Set("Content-Type", "application/json") - _, _ = writer.Write([]byte(`{"id":"response"}`)) - })) - defer upstream.Close() - - injectCalls := 0 - input := streamForwardInput(upstream.URL) - input.Group.InjectUsageOptions = true - input.Request.Body = []byte(`{"model":"gpt-4o"}`) - input.Dialect = streamUsageInjectorDialect{ - OpenAI: dialect.NewOpenAI(http.DefaultClient), - inject: func(request *dialect.ParsedRequest) (*dialect.ParsedRequest, error) { - injectCalls++ - return dialect.NewOpenAI(http.DefaultClient).InjectStreamUsage(request) - }, - } - - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward(context.Background(), input) - if result.Err != nil || injectCalls != 0 { - t.Fatalf("Forward() result/inject calls = %#v/%d", result, injectCalls) - } - if _, exists := received["stream_options"]; exists { - t.Fatalf("non-stream upstream request was injected: %#v", received) - } -} - -func TestForwardStreamSkipsDialectWithoutStreamUsageInjector(t *testing.T) { - var received map[string]any - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if err := json.NewDecoder(request.Body).Decode(&received); err != nil { - t.Error(err) - } - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: {\"choices\":[]}\n\n")) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Group.InjectUsageOptions = true - input.Dialect = dialectWithoutStreamUsageInjector{Dialect: dialect.NewOpenAI(http.DefaultClient)} - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - result := forwarder.ForwardStream(context.Background(), input, newRecordingResponseWriter()) - if result.Err != nil || !result.Committed || forwarder.usageCapture.failureTotal.Load() != 0 { - t.Fatalf("ForwardStream() result/failures = %#v/%d", result, forwarder.usageCapture.failureTotal.Load()) - } - if _, exists := received["stream_options"]; exists { - t.Fatalf("capability-less dialect injected usage: %#v", received) - } -} - -func TestForwardStreamInjectsOpenAIUsageWhenEffective(t *testing.T) { - var received map[string]any - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if err := json.NewDecoder(request.Body).Decode(&received); err != nil { - t.Error(err) - } - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: {\"choices\":[]}\n\n")) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Group.InjectUsageOptions = true - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, newRecordingResponseWriter(), - ) - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() = %#v", result) - } - options, ok := received["stream_options"].(map[string]any) - if !ok || options["include_usage"] != true { - t.Fatalf("upstream request = %#v, want stream_options.include_usage=true", received) - } -} - -func TestForwardStreamInjectsUsageAfterModelAliasRewrite(t *testing.T) { - var received map[string]any - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if err := json.NewDecoder(request.Body).Decode(&received); err != nil { - t.Error(err) - } - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: {\"choices\":[]}\n\n")) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Group.InjectUsageOptions = true - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - input.Request.Body = []byte(`{"model":"public-model","stream":true}`) - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, newRecordingResponseWriter(), - ) - if result.Err != nil || !result.Committed || received["model"] != "provider-model" { - t.Fatalf("ForwardStream() result/request = %#v/%#v", result, received) - } - options, ok := received["stream_options"].(map[string]any) - if !ok || options["include_usage"] != true { - t.Fatalf("rewritten request did not retain injected usage: %#v", received) - } -} - -func TestForwardStreamNormalizesFinalRequestMetadata(t *testing.T) { - type capturedRequest struct { - body []byte - headers http.Header - contentLength int64 - } - tests := []struct { - name string - configure func(*ForwardInput) - wantBodyChanged bool - }{ - { - name: "usage injection invalidates stale metadata", - configure: func(input *ForwardInput) { - input.Group.InjectUsageOptions = true - }, - wantBodyChanged: true, - }, - { - name: "disabled injection still removes stale metadata", - }, - { - name: "failed injection still removes stale metadata when body is unchanged", - configure: func(input *ForwardInput) { - input.Group.InjectUsageOptions = true - input.Dialect = streamUsageInjectorDialect{ - OpenAI: dialect.NewOpenAI(http.DefaultClient), - inject: func(*dialect.ParsedRequest) (*dialect.ParsedRequest, error) { - return nil, errors.New("inject failed") - }, - } - }, - }, - { - name: "model alias rewrite invalidates stale metadata", - configure: func(input *ForwardInput) { - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - input.Request.Body = []byte(`{"model":"public-model","stream":true}`) - }, - wantBodyChanged: true, - }, - { - name: "group header rules cannot restore stale metadata", - configure: func(input *ForwardInput) { - input.Group.InjectUsageOptions = true - input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ - "Digest": "sha-256=group-digest", - "Signature": "group-signature", - }} - }, - wantBodyChanged: true, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - received := make(chan capturedRequest, 1) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - body, err := io.ReadAll(request.Body) - if err != nil { - t.Error(err) - } - received <- capturedRequest{ - body: body, - headers: request.Header.Clone(), - contentLength: request.ContentLength, - } - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: {\"choices\":[]}\n\n")) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - setRepresentationMetadata(input.Request.Header) - if test.configure != nil { - test.configure(&input) - } - originalBody := bytes.Clone(input.Request.Body) - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, newRecordingResponseWriter(), - ) - if result.Err != nil || !result.Committed { - t.Fatalf("ForwardStream() = %#v", result) - } - got := <-received - if changed := !bytes.Equal(got.body, originalBody); changed != test.wantBodyChanged { - t.Fatalf("body changed = %t, want %t: original=%s upstream=%s", changed, test.wantBodyChanged, originalBody, got.body) - } - if got.contentLength != int64(len(got.body)) { - t.Fatalf("ContentLength = %d, want %d", got.contentLength, len(got.body)) - } - - assertRepresentationMetadata(t, got.headers, false) - }) - } -} - -func TestForwardStreamDoesNotInjectWhenEffectiveFalse(t *testing.T) { - for _, body := range []string{ - `{"model":"gpt-4o","stream":true}`, - `{"model":"gpt-4o","stream":true,"stream_options":{"include_usage":false}}`, - `{"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true}}`, - } { - t.Run(body, func(t *testing.T) { - var received []byte - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - received, _ = io.ReadAll(request.Body) - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: {\"choices\":[]}\n\n")) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Group.InjectUsageOptions = false - input.Request.Body = []byte(body) - want := bytes.Clone(input.Request.Body) - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, newRecordingResponseWriter(), - ) - if result.Err != nil || !bytes.Equal(received, want) { - t.Fatalf("result/body = %#v/%s, want original %s", result, received, want) - } - }) - } -} - -func TestForwardStreamUsageInjectionFailureKeepsRewrittenRequest(t *testing.T) { - tests := []struct { - name string - body string - inject func(*dialect.ParsedRequest) (*dialect.ParsedRequest, error) - }{ - {name: "error", inject: func(*dialect.ParsedRequest) (*dialect.ParsedRequest, error) { return nil, errors.New("inject failed") }}, - {name: "panic", inject: func(*dialect.ParsedRequest) (*dialect.ParsedRequest, error) { panic("inject panic") }}, - {name: "nil", inject: func(*dialect.ParsedRequest) (*dialect.ParsedRequest, error) { return nil, nil }}, - {name: "oversize", inject: func(request *dialect.ParsedRequest) (*dialect.ParsedRequest, error) { - return &dialect.ParsedRequest{Method: request.Method, Path: request.Path, RawQuery: request.RawQuery, Header: request.Header.Clone(), Body: bytes.Repeat([]byte("x"), int(maxRequestBodyBytes)+1)}, nil - }}, - {name: "invalid options", body: `{"model":"public-model","stream":true,"stream_options":false}`, inject: dialect.NewOpenAI(http.DefaultClient).InjectStreamUsage}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - var received map[string]any - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - if err := json.NewDecoder(request.Body).Decode(&received); err != nil { - t.Error(err) - } - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = writer.Write([]byte("data: {\"choices\":[]}\n\n")) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Group.InjectUsageOptions = true - input.ExternalModel = "public-model" - input.UpstreamModelID = "upstream-model" - if test.body != "" { - input.Request.Body = []byte(test.body) - } else { - input.Request.Body = []byte(`{"model":"public-model","stream":true}`) - } - input.Dialect = streamUsageInjectorDialect{OpenAI: dialect.NewOpenAI(http.DefaultClient), inject: test.inject} - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - result := forwarder.ForwardStream(context.Background(), input, newRecordingResponseWriter()) - if result.Err != nil || !result.Committed || received["model"] != "upstream-model" { - t.Fatalf("result/request = %#v/%#v", result, received) - } - if options, exists := received["stream_options"].(map[string]any); exists && options["include_usage"] == true { - t.Fatalf("fail-open request unexpectedly injected usage: %#v", received) - } - if forwarder.usageCapture.failureTotal.Load() != 1 { - t.Fatalf("failure total = %d, want 1", forwarder.usageCapture.failureTotal.Load()) - } - }) - } -} - -func setRepresentationMetadata(headers http.Header) { - headers.Set("ETag", `"wire-v1"`) - headers.Set("Digest", "sha-256=wire-digest") - headers.Set("Content-MD5", "d2lyZQ==") - headers.Set("Content-Range", "bytes 0-9/10") - headers.Set("Content-Digest", "sha-256=:d2lyZQ==:") - headers.Set("Repr-Digest", "sha-256=:cmVwcg==:") - headers.Set("Signature", "sig1=:c2lnbmF0dXJl:") - headers.Set("Signature-Input", `sig1=("content-digest");created=1`) -} - -func assertRepresentationMetadata(t *testing.T, headers http.Header, wantPreserved bool) { - t.Helper() - want := map[string]string{ - "ETag": `"wire-v1"`, - "Digest": "sha-256=wire-digest", - "Content-MD5": "d2lyZQ==", - "Content-Range": "bytes 0-9/10", - "Content-Digest": "sha-256=:d2lyZQ==:", - "Repr-Digest": "sha-256=:cmVwcg==:", - } - for name, value := range want { - got := headers.Get(name) - if wantPreserved && got != value { - t.Errorf("%s = %q, want preserved value %q", name, got, value) - } - if !wantPreserved && got != "" { - t.Errorf("%s = %q, want removed after body rewrite", name, got) - } - } - for _, name := range []string{"Signature", "Signature-Input"} { - if got := headers.Get(name); got != "" { - t.Errorf("%s = %q, want removed from downstream response", name, got) - } - } -} - -func TestInvalidateRewrittenBodyHeadersRemovesRepresentationMetadata(t *testing.T) { - headers := make(http.Header) - headers.Set("Content-Length", "123") - setRepresentationMetadata(headers) - for _, name := range []string{ - "content-length", "eTAG", "dIGEST", "content-md5", "content-range", - "content-digest", "repr-digest", "signature", "signature-input", - } { - headers[name] = []string{"case-colliding-stale-value"} - } - - invalidateRewrittenBodyHeaders(headers) - - for actualName := range headers { - if isRepresentationMetadataHeader(actualName) && - !strings.EqualFold(actualName, "Content-Encoding") { - t.Fatalf("stale metadata survived as %q: %#v", actualName, headers[actualName]) - } - } - assertRepresentationMetadata(t, headers, false) -} - -func assertHeadersDoNotContain(t *testing.T, headers http.Header, literal string) { - t.Helper() - for name, values := range headers { - for _, value := range values { - if strings.Contains(value, literal) { - t.Fatalf("header %q leaked %q in value %q", name, literal, value) - } - } - } -} - -func decodeSSEJSONPayloads(t *testing.T, stream []byte) []map[string]any { - t.Helper() - remaining := bytes.Clone(stream) - payloads := make([]map[string]any, 0) - for len(remaining) > 0 { - boundary := bytes.Index(remaining, []byte("\n\n")) - if boundary < 0 { - t.Fatalf("incomplete SSE output: %q", remaining) - } - event := remaining[:boundary+2] - remaining = remaining[boundary+2:] - var values [][]byte - for _, line := range splitSSEEventLines(event) { - if line.isData { - values = append(values, line.data) - } - } - if len(values) == 0 { - continue - } - var decoded map[string]any - if err := json.Unmarshal(bytes.Join(values, []byte{'\n'}), &decoded); err != nil { - t.Fatalf("decode SSE JSON payload: %v", err) - } - payloads = append(payloads, decoded) - } - return payloads -} diff --git a/internal/gateway/handler.go b/internal/gateway/handler.go index 62e1ab5f..c054efc2 100644 --- a/internal/gateway/handler.go +++ b/internal/gateway/handler.go @@ -15,7 +15,9 @@ import ( "github.com/gin-gonic/gin" "github.com/sirupsen/logrus" + "gpt-load/internal/channel" "gpt-load/internal/dialect" + "gpt-load/internal/execution" "gpt-load/internal/health" "gpt-load/internal/httplifecycle" "gpt-load/internal/platform/contentcoding" @@ -64,28 +66,29 @@ type PriceTableProvider interface { Load() *pricing.Table } -type keyMutationCoordinator interface { +type credentialMutationCoordinator interface { Do(uint, func()) } -type runtimeKeyRegistry interface { - scheduler.KeySource - CaptureActiveKeyRefs(groupIDs []uint) []state.KeyRef - ActiveEncryptedValueIfMatch(ref state.KeyRef) (string, bool) - SetCooldown(keyID uint, until time.Time) bool - IncrFailure(keyID uint) (int, bool) - SetBlacklisted(keyID uint) bool - ClearFailure(keyID uint) bool +type runtimeCredentialRegistry interface { + scheduler.CredentialSource + CaptureActiveCredentialRefs(groupIDs []uint) []state.CredentialRef + ActiveEncryptedCredentialDataIfMatch(ref state.CredentialRef) (string, bool) + SetCooldown(credentialID uint, until time.Time) bool + IncrFailure(credentialID uint) (int, bool) + SetBlacklisted(credentialID uint) bool + ClearFailure(credentialID uint) bool } type Handler struct { manager *state.Manager - registry runtimeKeyRegistry + channels *channel.Registry + registry runtimeCredentialRegistry encryption encryption.Service forwarder AttemptForwarder dialects dialect.Set stats *health.StatsStore - mutations keyMutationCoordinator + mutations credentialMutationCoordinator limiter AccessKeyRPMLimiter requestLogSink telemetry.RequestLogSink priceTables PriceTableProvider @@ -106,6 +109,7 @@ func (handler *Handler) freezeAttemptPricing( applicable bool, ) frozenAttemptPricing { frozen := frozenAttemptPricing{ + channelID: string(selection.ChannelID), groupID: selection.GroupID, upstreamModel: optionalModelValue(selection.UpstreamModelID), applicable: applicable, @@ -118,7 +122,7 @@ func (handler *Handler) freezeAttemptPricing( func NewHandler( manager *state.Manager, - registry *state.KeyRegistry, + registry *state.CredentialRegistry, encryptionService encryption.Service, forwarder AttemptForwarder, dialects dialect.Set, @@ -135,7 +139,7 @@ func NewHandler( requestLogSink = telemetry.NoopRequestLogSink{} } return &Handler{ - manager: manager, registry: registry, encryption: encryptionService, + manager: manager, channels: channel.NewRegistry(), registry: registry, encryption: encryptionService, forwarder: forwarder, dialects: dialects, stats: stats, mutations: mutations, limiter: limiter, requestLogSink: requestLogSink, priceTables: priceTables, newRandom: func() *rand.Rand { return rand.New(rand.NewSource(rand.Int63())) }, @@ -161,7 +165,8 @@ func NewHandler( // focused gateway tests. func NewHandlerWithLifecycle( manager *state.Manager, - registry *state.KeyRegistry, + registry *state.CredentialRegistry, + channelRegistry *channel.Registry, encryptionService encryption.Service, forwarder AttemptForwarder, dialects dialect.Set, @@ -184,6 +189,9 @@ func NewHandlerWithLifecycle( requestLogSink, priceTables, ) + if channelRegistry != nil { + handler.channels = channelRegistry + } handler.lifecycle = lifecycle return handler } @@ -194,47 +202,47 @@ func (unlimitedAccessKeyRPMLimiter) Allow(uint, int64) ratelimit.LimitDecision { return ratelimit.LimitDecision{Allowed: true} } -func (handler *Handler) applyKeyAction( - keyID uint, +func (handler *Handler) applyCredentialAction( + credentialID uint, decision health.Result, statusCode int, attemptNow time.Time, ) { switch decision.Action { - case health.ActionCooldownKey: + case health.ActionCooldownCredential: mutate := func() { until := decision.CooldownUntil if decision.UseFixed { until = attemptNow.Add(fixedCooldown) } - if handler.registry.SetCooldown(keyID, until) { - handler.stats.RecordProblem(keyID, decision.Category, statusCode, attemptNow) + if handler.registry.SetCooldown(credentialID, until) { + handler.stats.RecordProblem(credentialID, decision.Category, statusCode, attemptNow) } } if handler.mutations == nil { mutate() } else { - handler.mutations.Do(keyID, mutate) + handler.mutations.Do(credentialID, mutate) } - case health.ActionFailKey: - handler.mutations.Do(keyID, func() { - count, ok := handler.registry.IncrFailure(keyID) + case health.ActionFailCredential: + handler.mutations.Do(credentialID, func() { + count, ok := handler.registry.IncrFailure(credentialID) if !ok { return } if count >= blacklistFailureThreshold && - !handler.registry.SetBlacklisted(keyID) { + !handler.registry.SetBlacklisted(credentialID) { return } - handler.stats.RecordFailure(keyID, decision.Category, statusCode, attemptNow) + handler.stats.RecordFailure(credentialID, decision.Category, statusCode, attemptNow) }) } } -func (handler *Handler) recordSuccess(keyID uint, at time.Time) { - handler.mutations.Do(keyID, func() { - if handler.registry.ClearFailure(keyID) { - handler.stats.RecordSuccess(keyID, at) +func (handler *Handler) recordCredentialSuccess(credentialID uint, at time.Time) { + handler.mutations.Do(credentialID, func() { + if handler.registry.ClearFailure(credentialID) { + handler.stats.RecordSuccess(credentialID, at) } }) } @@ -392,39 +400,42 @@ func (handler *Handler) Handle(ginContext *gin.Context) { recorder.completeCanceled(0, -1) return } - candidateGroupIDs := scheduler.CandidateGroupIDs( - snapshot, - selectedRoute.Protocol, - accessKey, - ) - capturedRefs := handler.registry.CaptureActiveKeyRefs(candidateGroupIDs) - allowedKeyRefs := make(map[uint]state.KeyRef, len(capturedRefs)) + query := scheduler.Query{ + ClientProtocol: selectedRoute.Protocol, + Operation: metadata.Operation, + RequiredFeatures: metadata.RequiredFeatures, + ExternalModel: metadata.Model, + AccessKey: accessKey, + } + candidateGroupIDs := scheduler.CandidateGroupIDsForQuery(snapshot, query) + capturedRefs := handler.registry.CaptureActiveCredentialRefs(candidateGroupIDs) + allowedCredentialRefs := make(map[uint]state.CredentialRef, len(capturedRefs)) for _, ref := range capturedRefs { - allowedKeyRefs[ref.ID] = ref + allowedCredentialRefs[ref.ID] = ref } recorder.setClientModel(model) + recorder.setOperation(metadata.Operation) recorder.setStream(metadata.Stream) recorder.setReasoning(metadata.Reasoning) recorder.setUsageApplicable(metadata.ObserveUsage) recorder.setUsageDiagnostics(metadata.UsageDiagnostics) - allowedKeyIDs := make(map[uint]struct{}, len(allowedKeyRefs)) - for keyID := range allowedKeyRefs { - allowedKeyIDs[keyID] = struct{}{} + allowedCredentialIDs := make(map[uint]struct{}, len(allowedCredentialRefs)) + for credentialID := range allowedCredentialRefs { + allowedCredentialIDs[credentialID] = struct{}{} } - iterator := scheduler.New(snapshot, handler.registry, scheduler.Query{ - Protocol: selectedRoute.Protocol, ExternalModel: metadata.Model, AccessKey: accessKey, - AllowedKeyIDs: allowedKeyIDs, - }, handler.newRandom()) + query.AllowedCredentialIDs = allowedCredentialIDs + iterator := scheduler.New(snapshot, handler.registry, query, handler.newRandom()) handler.executeAttempts( ginContext, iterator, - allowedKeyRefs, + allowedCredentialRefs, selectedDialect, parsed, model, metadata.Stream, metadata.ObserveUsage, + metadata.Operation, recorder, ) } @@ -524,12 +535,13 @@ func headerFieldValues(headers http.Header, name string) []string { func (handler *Handler) executeAttempts( ginContext *gin.Context, iterator *scheduler.Iterator, - allowedKeyRefs map[uint]state.KeyRef, + allowedCredentialRefs map[uint]state.CredentialRef, selectedDialect dialect.Dialect, parsed *dialect.ParsedRequest, externalModel string, stream bool, observeUsage bool, + operation execution.Operation, recorder *requestRecorder, ) { type deferredAttempt struct { @@ -555,29 +567,55 @@ func (handler *Handler) executeAttempts( if err != nil { break } - ref, allowed := allowedKeyRefs[selection.KeyID] + ref, allowed := allowedCredentialRefs[selection.CredentialID] if !allowed || ref.GroupID != selection.GroupID { continue } - encrypted, active := handler.registry.ActiveEncryptedValueIfMatch(ref) + encrypted, active := handler.registry.ActiveEncryptedCredentialDataIfMatch(ref) if !active { continue } - apiKey, err := handler.encryption.Decrypt(encrypted) + decryptedCredential, err := handler.encryption.Decrypt(encrypted) + if err != nil { + continue + } + normalizedCredential, err := normalizeChannelCredential( + handler.channels, + selection.ChannelID, + decryptedCredential, + ) if err != nil { continue } attempts++ updateDebugHeaders(ginContext.Writer.Header(), selection.Group.Name, attempts) - selectedKeyID := selection.KeyID + selectedCredentialID := selection.CredentialID + executionRequestID := "untracked" + if recorder != nil { + executionRequestID = recorder.requestID + } input := ForwardInput{ Dialect: selectedDialect, ObserveUsage: observeUsage, - Group: selection.Group, APIKey: apiKey, Request: parsed, + Group: selection.Group, APIKey: normalizedCredential.apiKey, + CredentialSecrets: normalizedCredential.secrets, Request: parsed, ExternalModel: externalModel, UpstreamModelID: optionalModelValue(selection.UpstreamModelID), + RequestID: executionRequestID, + AttemptID: executionRequestID + ":" + strconv.Itoa(attempts), + AttemptSequence: uint32(attempts), + ClientProtocol: selectedDialect.Protocol(), + Operation: operation, + ChannelID: string(selection.ChannelID), + TargetConfig: selection.ResolvedTarget.TargetConfig, + Credential: execution.NewCredentialSnapshot( + selection.CredentialID, + ref.Version, + ref.IdentityGeneration, + normalizedCredential.payload, + ), OnStreamReady: func() { - handler.recordSuccess(selectedKeyID, handler.now()) + handler.recordCredentialSuccess(selectedCredentialID, handler.now()) }, OnFirstResponse: func() { recorder.recordFirstResponse() @@ -595,6 +633,7 @@ func (handler *Handler) executeAttempts( } else { result = handler.forwarder.Forward(ginContext.Request.Context(), input) } + result = normalizeUpstreamResultEvidence(result) attemptCompleted := time.Time{} if recorder != nil { attemptCompleted = recorder.now() @@ -610,7 +649,7 @@ func (handler *Handler) executeAttempts( ) } recordedAttempt := recorder.recordStreamAttempt( - selection, apiKey, result, attemptStarted, attemptCompleted, + selection, normalizedCredential.secrets, result, attemptStarted, attemptCompleted, ) recorder.completeStream(result, optionalModelValue(selection.UpstreamModelID), recordedAttempt) } @@ -620,7 +659,7 @@ func (handler *Handler) executeAttempts( if recorder != nil { recordedAttempt := recorder.recordAttempt( selection, - apiKey, + normalizedCredential.secrets, result, health.Result{ Category: health.FailureCategoryDownstreamCancel, @@ -637,20 +676,14 @@ func (handler *Handler) executeAttempts( if !stream && !result.ProviderErrorBeforeCommit && result.HasResponse() && result.StatusCode >= http.StatusOK && result.StatusCode < http.StatusMultipleChoices { - handler.recordSuccess(selection.KeyID, attemptNow) + handler.recordCredentialSuccess(selection.CredentialID, attemptNow) } - decision := health.Judge(selectedDialect, health.Attempt{ - StatusCode: result.StatusCode, Body: result.ClassificationBody, - Header: result.Header, Now: attemptNow, - Err: result.Err, RequestWritten: result.RequestWritten, - Committed: result.Committed, RetryableBeforeCommit: result.RetryableBeforeCommit, - ProviderErrorBeforeCommit: result.ProviderErrorBeforeCommit, - }) + decision := judgeUpstreamResult(result, attemptNow) recordedAttempt := recorder.recordAttempt( - selection, apiKey, result, decision, attemptStarted, attemptCompleted, + selection, normalizedCredential.secrets, result, decision, attemptStarted, attemptCompleted, ) lastAttemptIndex = recordedAttempt - handler.applyKeyAction(selection.KeyID, decision, result.StatusCode, attemptNow) + handler.applyCredentialAction(selection.CredentialID, decision, result.StatusCode, attemptNow) if decision.Action == health.ActionSkipGroup { iterator.SkipGroup(selection.GroupID) } diff --git a/internal/gateway/handler_test.go b/internal/gateway/handler_test.go index d00d63bb..f01ee0ca 100644 --- a/internal/gateway/handler_test.go +++ b/internal/gateway/handler_test.go @@ -20,13 +20,13 @@ import ( "github.com/gin-gonic/gin" + "gpt-load/internal/channel" "gpt-load/internal/dialect" "gpt-load/internal/health" "gpt-load/internal/platform/config" "gpt-load/internal/platform/contentcoding" "gpt-load/internal/platform/encryption" platformhttp "gpt-load/internal/platform/httpclient" - "gpt-load/internal/platform/redact" "gpt-load/internal/platform/utils" "gpt-load/internal/protocol" "gpt-load/internal/ratelimit" @@ -40,6 +40,68 @@ func withProviderErrorBeforeCommit(result UpstreamResult) UpstreamResult { return result } +func TestHandlerForwardsStructuredCloudCredentialWithoutAPIKeyAssumption(t *testing.T) { + forwarder := &scriptedForwarder{results: []UpstreamResult{{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: []byte(`{"id":"ok","model":"anthropic.claude-test"}`), + }}} + handler, manager, registry := newHandlerForTest(t, forwarder) + if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []state.GroupConfig{{ + ID: 1, Name: "bedrock", ChannelID: channel.AWSBedrock, + Params: json.RawMessage(`{"region":"us-east-1"}`), + Models: []state.ModelConfig{{ID: "anthropic.claude-test"}}, Enabled: true, + }}, + Credentials: []state.CredentialConfig{{ + ID: 1, GroupID: 1, Status: state.CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "bedrock-credential", + }}, + AccessKeys: []state.AccessKeyConfig{{ + ID: 1, Name: "client", KeyHash: handler.encryption.Hash("gl-client"), + Status: state.AccessKeyStatusActive, + }}, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + const canonical = `{"access_key":"AKIA_TEST","secret_key":"bedrock-secret","session_token":"bedrock-session"}` + encrypted, err := handler.encryption.Encrypt(canonical) + if err != nil { + t.Fatalf("Encrypt() error = %v", err) + } + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Status: state.CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "bedrock-credential", + EncryptedValue: encrypted, + }}); err != nil { + t.Fatalf("ReplaceCredentials() error = %v", err) + } + engine := gin.New() + bindGatewayRoutesForTest(t, engine, handler) + request := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + bytes.NewBufferString(`{"model":"anthropic.claude-test","messages":[{"role":"user","content":"ping"}]}`), + ) + request.Header.Set("Authorization", "Bearer gl-client") + response := httptest.NewRecorder() + + engine.ServeHTTP(response, request) + + if response.Code != http.StatusOK || len(forwarder.inputs) != 1 { + t.Fatalf("status/inputs = %d/%d; body=%s", response.Code, len(forwarder.inputs), response.Body.String()) + } + input := forwarder.inputs[0] + if input.ChannelID != string(channel.AWSBedrock) || input.APIKey != "" || + string(input.Credential.Data()) != canonical { + t.Fatalf("forward input = %#v; credential=%s", input, input.Credential.Data()) + } + if want := []string{"AKIA_TEST", "bedrock-secret", "bedrock-session"}; !reflect.DeepEqual(input.CredentialSecrets, want) { + t.Fatalf("credential secrets = %#v, want %#v", input.CredentialSecrets, want) + } +} + type scriptedForwarder struct { results []UpstreamResult streamResults []UpstreamResult @@ -110,24 +172,24 @@ func (*blockingRequestBody) Close() error { } type mutatingRuntimeRegistry struct { - *state.KeyRegistry + *state.CredentialRegistry mutate func() mutated bool } type recordingRuntimeRegistry struct { - *state.KeyRegistry - cooldownKeyID uint - cooldownUntil time.Time - cooldownCalls int - incrFailureCalls int - lastFailureCount int - blacklistCalls int - clearCalls int + *state.CredentialRegistry + cooldownCredentialID uint + cooldownUntil time.Time + cooldownCalls int + incrFailureCalls int + lastFailureCount int + blacklistCalls int + clearCalls int } type captureCountingRuntimeRegistry struct { - runtimeKeyRegistry + runtimeCredentialRegistry captureCalls int } @@ -146,11 +208,11 @@ func (value *cancelingSuccessfulInspectDialect) InspectRequest( return metadata, err } -func (registry *captureCountingRuntimeRegistry) CaptureActiveKeyRefs( +func (registry *captureCountingRuntimeRegistry) CaptureActiveCredentialRefs( groupIDs []uint, -) []state.KeyRef { +) []state.CredentialRef { registry.captureCalls++ - return registry.runtimeKeyRegistry.CaptureActiveKeyRefs(groupIDs) + return registry.runtimeCredentialRegistry.CaptureActiveCredentialRefs(groupIDs) } type gatewayMutationObservation struct { @@ -159,14 +221,14 @@ type gatewayMutationObservation struct { incrFailureCalls int lastFailureCount int blacklistCalls int - stats health.KeyStats + stats health.CredentialStats } func TestHandlerCoordinatesCooldownMutation(t *testing.T) { now := time.Date(2026, time.August, 1, 12, 0, 0, 0, time.UTC) - registry := &recordingRuntimeRegistry{KeyRegistry: state.NewKeyRegistry()} - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "cipher", + registry := &recordingRuntimeRegistry{CredentialRegistry: state.NewCredentialRegistry()} + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, EncryptedValue: "cipher", }}); err != nil { t.Fatal(err) } @@ -180,20 +242,20 @@ func TestHandlerCoordinatesCooldownMutation(t *testing.T) { handler := &Handler{registry: registry, stats: stats, mutations: coordinator} done := make(chan struct{}) go func() { - handler.applyKeyAction(1, health.Result{ + handler.applyCredentialAction(1, health.Result{ Category: health.FailureCategoryRateLimited, - Action: health.ActionCooldownKey, + Action: health.ActionCooldownCredential, }, http.StatusTooManyRequests, now) close(done) }() receiveTestSignal(t, coordinator.entered, "cooldown coordinator entry") - if registry.cooldownCalls != 0 || stats.Snapshot(1, now) != (health.KeyStats{}) { + if registry.cooldownCalls != 0 || stats.Snapshot(1, now) != (health.CredentialStats{}) { t.Fatal("cooldown bundle changed state before coordinator callback") } close(coordinator.releaseEntry) observed := receiveTestSignal(t, coordinator.observed, "cooldown mutation observation") - if observed.cooldownCalls != 1 || observed.stats != (health.KeyStats{ + if observed.cooldownCalls != 1 || observed.stats != (health.CredentialStats{ Problem: 1, ConsecutiveProblem: 1, LastFailureCategory: health.FailureCategoryRateLimited, @@ -214,18 +276,18 @@ type barrierGatewayMutationCoordinator struct { } type gatewayFailureOrderRegistry struct { - *state.KeyRegistry + *state.CredentialRegistry failureEntered chan struct{} releaseFailure chan struct{} enterOnce sync.Once } -func (registry *gatewayFailureOrderRegistry) IncrFailure(keyID uint) (int, bool) { +func (registry *gatewayFailureOrderRegistry) IncrFailure(credentialID uint) (int, bool) { registry.enterOnce.Do(func() { close(registry.failureEntered) <-registry.releaseFailure }) - return registry.KeyRegistry.IncrFailure(keyID) + return registry.CredentialRegistry.IncrFailure(credentialID) } func newBarrierGatewayMutationCoordinator( @@ -258,35 +320,35 @@ func receiveTestSignal[T any](t *testing.T, signal <-chan T, name string) T { } } -func (registry *recordingRuntimeRegistry) SetCooldown(keyID uint, until time.Time) bool { - registry.cooldownKeyID = keyID +func (registry *recordingRuntimeRegistry) SetCooldown(credentialID uint, until time.Time) bool { + registry.cooldownCredentialID = credentialID registry.cooldownUntil = until registry.cooldownCalls++ - return registry.KeyRegistry.SetCooldown(keyID, until) + return registry.CredentialRegistry.SetCooldown(credentialID, until) } -func (registry *recordingRuntimeRegistry) IncrFailure(keyID uint) (int, bool) { +func (registry *recordingRuntimeRegistry) IncrFailure(credentialID uint) (int, bool) { registry.incrFailureCalls++ - count, ok := registry.KeyRegistry.IncrFailure(keyID) + count, ok := registry.CredentialRegistry.IncrFailure(credentialID) registry.lastFailureCount = count return count, ok } -func (registry *recordingRuntimeRegistry) SetBlacklisted(keyID uint) bool { +func (registry *recordingRuntimeRegistry) SetBlacklisted(credentialID uint) bool { registry.blacklistCalls++ - return registry.KeyRegistry.SetBlacklisted(keyID) + return registry.CredentialRegistry.SetBlacklisted(credentialID) } -func (registry *recordingRuntimeRegistry) ClearFailure(keyID uint) bool { +func (registry *recordingRuntimeRegistry) ClearFailure(credentialID uint) bool { registry.clearCalls++ - return registry.KeyRegistry.ClearFailure(keyID) + return registry.CredentialRegistry.ClearFailure(credentialID) } func TestHandlerCoordinatesSuccessMutation(t *testing.T) { now := time.Date(2026, time.July, 27, 12, 0, 0, 0, time.UTC) - registry := &recordingRuntimeRegistry{KeyRegistry: state.NewKeyRegistry()} - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, FailureCount: 2, EncryptedValue: "cipher", + registry := &recordingRuntimeRegistry{CredentialRegistry: state.NewCredentialRegistry()} + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, FailureCount: 2, EncryptedValue: "cipher", }}); err != nil { t.Fatalf("Replace() error = %v", err) } @@ -301,16 +363,16 @@ func TestHandlerCoordinatesSuccessMutation(t *testing.T) { done := make(chan struct{}) go func() { - handler.recordSuccess(1, now) + handler.recordCredentialSuccess(1, now) close(done) }() receiveTestSignal(t, coordinator.entered, "success coordinator entry") - if registry.clearCalls != 0 || stats.Snapshot(1, now) != (health.KeyStats{}) { + if registry.clearCalls != 0 || stats.Snapshot(1, now) != (health.CredentialStats{}) { t.Fatalf("success bundle changed state before coordinator callback") } close(coordinator.releaseEntry) observed := receiveTestSignal(t, coordinator.observed, "success mutation observation") - if observed.clearCalls != 1 || observed.stats != (health.KeyStats{Success: 1}) { + if observed.clearCalls != 1 || observed.stats != (health.CredentialStats{Success: 1}) { t.Fatalf("coordinator callback observation = %#v", observed) } select { @@ -324,9 +386,9 @@ func TestHandlerCoordinatesSuccessMutation(t *testing.T) { func TestHandlerRecordsCooldownFailureContext(t *testing.T) { now := time.Date(2026, time.July, 27, 12, 0, 0, 0, time.UTC) - registry := &recordingRuntimeRegistry{KeyRegistry: state.NewKeyRegistry()} - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "cipher", + registry := &recordingRuntimeRegistry{CredentialRegistry: state.NewCredentialRegistry()} + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, EncryptedValue: "cipher", }}); err != nil { t.Fatalf("Replace() error = %v", err) } @@ -334,22 +396,22 @@ func TestHandlerRecordsCooldownFailureContext(t *testing.T) { handler := &Handler{registry: registry, stats: stats} until := now.Add(30 * time.Second) - handler.applyKeyAction(1, health.Result{ + handler.applyCredentialAction(1, health.Result{ Category: health.FailureCategoryRateLimited, - Action: health.ActionCooldownKey, + Action: health.ActionCooldownCredential, CooldownUntil: until, }, http.StatusTooManyRequests, now) - if registry.cooldownCalls != 1 || registry.cooldownKeyID != 1 || + if registry.cooldownCalls != 1 || registry.cooldownCredentialID != 1 || !registry.cooldownUntil.Equal(until) { t.Fatalf( "cooldown mutation = calls:%d key:%d until:%s", registry.cooldownCalls, - registry.cooldownKeyID, + registry.cooldownCredentialID, registry.cooldownUntil, ) } - if got, want := stats.Snapshot(1, now), (health.KeyStats{ + if got, want := stats.Snapshot(1, now), (health.CredentialStats{ Problem: 1, ConsecutiveProblem: 1, LastFailureCategory: health.FailureCategoryRateLimited, @@ -368,37 +430,37 @@ func TestHandlerSkipsStatsWhenRegistryKeyWasDeletedBeforeCompletion(t *testing.T { name: "cooldown", mutate: func(handler *Handler) { - handler.applyKeyAction(1, health.Result{ + handler.applyCredentialAction(1, health.Result{ Category: health.FailureCategoryRateLimited, - Action: health.ActionCooldownKey, + Action: health.ActionCooldownCredential, }, http.StatusTooManyRequests, now) }, }, { name: "attributable failure", mutate: func(handler *Handler) { - handler.applyKeyAction(1, health.Result{ + handler.applyCredentialAction(1, health.Result{ Category: health.FailureCategoryInvalidKey, - Action: health.ActionFailKey, + Action: health.ActionFailCredential, }, http.StatusUnauthorized, now) }, }, { name: "success", mutate: func(handler *Handler) { - handler.recordSuccess(1, now) + handler.recordCredentialSuccess(1, now) }, }, } { t.Run(test.name, func(t *testing.T) { stats := health.NewStatsStore() handler := &Handler{ - registry: state.NewKeyRegistry(), + registry: state.NewCredentialRegistry(), stats: stats, mutations: health.NewMutationCoordinator(), } test.mutate(handler) - if got := stats.Snapshot(1, now); got != (health.KeyStats{}) { + if got := stats.Snapshot(1, now); got != (health.CredentialStats{}) { t.Fatalf("stats after deleted-key completion = %#v, want zero", got) } }) @@ -407,9 +469,9 @@ func TestHandlerSkipsStatsWhenRegistryKeyWasDeletedBeforeCompletion(t *testing.T func TestHandlerCoordinatesAttributableFailureMutation(t *testing.T) { now := time.Date(2026, time.July, 27, 12, 0, 0, 0, time.UTC) - registry := &recordingRuntimeRegistry{KeyRegistry: state.NewKeyRegistry()} - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, FailureCount: 2, EncryptedValue: "cipher", + registry := &recordingRuntimeRegistry{CredentialRegistry: state.NewCredentialRegistry()} + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, FailureCount: 2, EncryptedValue: "cipher", }}); err != nil { t.Fatalf("Replace() error = %v", err) } @@ -426,22 +488,22 @@ func TestHandlerCoordinatesAttributableFailureMutation(t *testing.T) { done := make(chan struct{}) go func() { - handler.applyKeyAction(1, health.Result{ + handler.applyCredentialAction(1, health.Result{ Category: health.FailureCategoryInvalidKey, - Action: health.ActionFailKey, + Action: health.ActionFailCredential, }, http.StatusUnauthorized, now) close(done) }() receiveTestSignal(t, coordinator.entered, "failure coordinator entry") if registry.incrFailureCalls != 0 || registry.blacklistCalls != 0 || - stats.Snapshot(1, now) != (health.KeyStats{}) { + stats.Snapshot(1, now) != (health.CredentialStats{}) { t.Fatalf("failure bundle changed state before coordinator callback") } close(coordinator.releaseEntry) observed := receiveTestSignal(t, coordinator.observed, "failure mutation observation") if observed.incrFailureCalls != 1 || observed.lastFailureCount != 3 || observed.blacklistCalls != 1 || - observed.stats != (health.KeyStats{ + observed.stats != (health.CredentialStats{ Failure: 1, Problem: 1, ConsecutiveFailure: 1, @@ -462,18 +524,18 @@ func TestHandlerCoordinatesAttributableFailureMutation(t *testing.T) { func TestGatewayFailureAndValidationRecoveryFailureFirstKeepsRegistryAndStatsFailed(t *testing.T) { now := time.Date(2026, time.July, 27, 12, 0, 0, 0, time.UTC) - baseRegistry := state.NewKeyRegistry() - if err := baseRegistry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + baseRegistry := state.NewCredentialRegistry() + if err := baseRegistry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, Blacklisted: true, FailureCount: 3, EncryptedValue: "cipher", }}); err != nil { t.Fatalf("Replace() error = %v", err) } - ref := baseRegistry.BlacklistedKeys()[0] + ref := baseRegistry.BlacklistedCredentials()[0] registry := &gatewayFailureOrderRegistry{ - KeyRegistry: baseRegistry, - failureEntered: make(chan struct{}), - releaseFailure: make(chan struct{}), + CredentialRegistry: baseRegistry, + failureEntered: make(chan struct{}), + releaseFailure: make(chan struct{}), } stats := health.NewStatsStore() stats.RecordFailure(1, health.FailureCategoryAmbiguous, 0, now) @@ -482,7 +544,7 @@ func TestGatewayFailureAndValidationRecoveryFailureFirstKeepsRegistryAndStatsFai failureDone := make(chan struct{}) go func() { - handler.applyKeyAction(1, health.Result{Action: health.ActionFailKey}, 0, now) + handler.applyCredentialAction(1, health.Result{Action: health.ActionFailCredential}, 0, now) close(failureDone) }() receiveTestSignal(t, registry.failureEntered, "gateway failure mutation") @@ -507,12 +569,13 @@ func TestGatewayFailureAndValidationRecoveryFailureFirstKeepsRegistryAndStatsFai t.Fatal("stale validation recovery = true after gateway failure, want false") } - if got, want := baseRegistry.BlacklistedKeys(), []state.KeyRef{{ - ID: 1, GroupID: 1, EncryptedValue: "cipher", FailureGeneration: 1, + if got, want := baseRegistry.BlacklistedCredentials(), []state.CredentialRef{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, + Fingerprint: "test-1", EncryptedValue: "cipher", FailureGeneration: 1, }}; !reflect.DeepEqual(got, want) { t.Fatalf("blacklisted keys = %#v, want %#v", got, want) } - if got, want := stats.Snapshot(1, now), (health.KeyStats{ + if got, want := stats.Snapshot(1, now), (health.CredentialStats{ Failure: 2, Problem: 2, ConsecutiveFailure: 2, ConsecutiveProblem: 2, }); got != want { t.Fatalf("stats = %#v, want %#v", got, want) @@ -521,14 +584,14 @@ func TestGatewayFailureAndValidationRecoveryFailureFirstKeepsRegistryAndStatsFai func TestGatewayFailureAndValidationRecoveryRecoveryFirstLeavesNewFailure(t *testing.T) { now := time.Date(2026, time.July, 27, 12, 0, 0, 0, time.UTC) - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, Blacklisted: true, FailureCount: 3, EncryptedValue: "cipher", }}); err != nil { t.Fatalf("Replace() error = %v", err) } - ref := registry.BlacklistedKeys()[0] + ref := registry.BlacklistedCredentials()[0] stats := health.NewStatsStore() stats.RecordFailure(1, health.FailureCategoryAmbiguous, 0, now) mutations := health.NewMutationCoordinator() @@ -555,7 +618,7 @@ func TestGatewayFailureAndValidationRecoveryRecoveryFirstLeavesNewFailure(t *tes failureDone := make(chan struct{}) go func() { close(failureAttempted) - handler.applyKeyAction(1, health.Result{Action: health.ActionFailKey}, 0, now) + handler.applyCredentialAction(1, health.Result{Action: health.ActionFailCredential}, 0, now) close(failureDone) }() receiveTestSignal(t, failureAttempted, "gateway failure attempt") @@ -565,25 +628,25 @@ func TestGatewayFailureAndValidationRecoveryRecoveryFirstLeavesNewFailure(t *tes } receiveTestSignal(t, failureDone, "gateway failure completion") - if got := registry.BlacklistedKeys(); len(got) != 0 { + if got := registry.BlacklistedCredentials(); len(got) != 0 { t.Fatalf("blacklisted keys = %#v, want recovered key with new sub-threshold failure", got) } - if refs := registry.CaptureActiveKeyRefs([]uint{1}); len(refs) != 1 || refs[0].FailureGeneration != 2 { + if refs := registry.CaptureActiveCredentialRefs([]uint{1}); len(refs) != 1 || refs[0].FailureGeneration != 2 { t.Fatalf("active refs = %#v, want generation 2 after recovery then failure", refs) } - if got, want := stats.Snapshot(1, now), (health.KeyStats{ + if got, want := stats.Snapshot(1, now), (health.CredentialStats{ Failure: 1, Problem: 1, ConsecutiveFailure: 1, ConsecutiveProblem: 1, }); got != want { t.Fatalf("stats = %#v, want %#v", got, want) } } -func (registry *mutatingRuntimeRegistry) CollectCandidates( +func (registry *mutatingRuntimeRegistry) CollectCredentialCandidates( groupIDs []uint, excluded func(uint) bool, now time.Time, -) []state.KeyMeta { - candidates := registry.KeyRegistry.CollectCandidates(groupIDs, excluded, now) +) []state.CredentialMeta { + candidates := registry.CredentialRegistry.CollectCredentialCandidates(groupIDs, excluded, now) if !registry.mutated && len(candidates) > 0 { registry.mutated = true registry.mutate() @@ -650,14 +713,14 @@ func TestHandlerRecordsNonStreamingResultStatsByAction(t *testing.T) { tests := []struct { name string result UpstreamResult - want health.KeyStats + want health.CredentialStats }{ { name: "2xx success", result: UpstreamResult{ StatusCode: http.StatusOK, Header: make(http.Header), Body: []byte(`{"ok":true}`), RequestWritten: true, }, - want: health.KeyStats{Success: 1}, + want: health.CredentialStats{Success: 1}, }, { name: "invalid key", @@ -665,7 +728,7 @@ func TestHandlerRecordsNonStreamingResultStatsByAction(t *testing.T) { StatusCode: http.StatusUnauthorized, Header: make(http.Header), Body: []byte(`{"error":"invalid key"}`), ClassificationBody: []byte(`{"error":"invalid key"}`), RequestWritten: true, }, - want: health.KeyStats{ + want: health.CredentialStats{ Failure: 1, Problem: 1, ConsecutiveFailure: 1, @@ -680,7 +743,7 @@ func TestHandlerRecordsNonStreamingResultStatsByAction(t *testing.T) { StatusCode: http.StatusBadRequest, Header: make(http.Header), Body: []byte(`{"error":"invalid input"}`), ClassificationBody: []byte(`{"error":"invalid input"}`), RequestWritten: true, }, - want: health.KeyStats{}, + want: health.CredentialStats{}, }, { name: "rate limited", @@ -688,7 +751,7 @@ func TestHandlerRecordsNonStreamingResultStatsByAction(t *testing.T) { StatusCode: http.StatusTooManyRequests, Header: make(http.Header), Body: []byte(`{"error":"rate limit"}`), ClassificationBody: []byte(`{"error":"rate limit"}`), RequestWritten: true, }, - want: health.KeyStats{ + want: health.CredentialStats{ Problem: 1, ConsecutiveProblem: 1, LastFailureCategory: health.FailureCategoryRateLimited, @@ -701,7 +764,7 @@ func TestHandlerRecordsNonStreamingResultStatsByAction(t *testing.T) { StatusCode: http.StatusNotFound, Header: make(http.Header), Body: []byte(`{"error":"model not found"}`), ClassificationBody: []byte(`{"error":"model not found"}`), RequestWritten: true, }, - want: health.KeyStats{ + want: health.CredentialStats{ Problem: 1, ConsecutiveProblem: 1, LastFailureCategory: health.FailureCategoryModelUnavailable, @@ -714,21 +777,21 @@ func TestHandlerRecordsNonStreamingResultStatsByAction(t *testing.T) { StatusCode: http.StatusInternalServerError, Header: make(http.Header), Body: []byte(`{"error":"overloaded"}`), ClassificationBody: []byte(`{"error":"overloaded"}`), RequestWritten: true, }, - want: health.KeyStats{}, + want: health.CredentialStats{}, }, { name: "pre-write transport", result: UpstreamResult{ Err: errors.New("dial upstream"), }, - want: health.KeyStats{}, + want: health.CredentialStats{}, }, { name: "context canceled", result: UpstreamResult{ Err: context.Canceled, }, - want: health.KeyStats{}, + want: health.CredentialStats{}, }, } @@ -776,7 +839,7 @@ func TestHandlerRecordsInvalidKeyPerAttempt(t *testing.T) { request.Header.Set("Authorization", "Bearer gl-client") engine.ServeHTTP(httptest.NewRecorder(), request) - if got := stats.Snapshot(1, now); got != (health.KeyStats{ + if got := stats.Snapshot(1, now); got != (health.CredentialStats{ Failure: 1, Problem: 1, ConsecutiveFailure: 1, @@ -786,7 +849,7 @@ func TestHandlerRecordsInvalidKeyPerAttempt(t *testing.T) { }) { t.Fatalf("first key stats = %#v, want one failure", got) } - if got := stats.Snapshot(2, now); got != (health.KeyStats{Success: 1}) { + if got := stats.Snapshot(2, now); got != (health.CredentialStats{Success: 1}) { t.Fatalf("second key stats = %#v, want one success", got) } } @@ -829,19 +892,19 @@ func TestHandlerRecordsStreamSuccessAtReadyTime(t *testing.T) { }() receiveTestSignal(t, forwarder.ready, "stream-ready callback") - if got := stats.Snapshot(1, now); got != (health.KeyStats{Success: 1}) { + if got := stats.Snapshot(1, now); got != (health.CredentialStats{Success: 1}) { t.Fatalf("stats before forward returns = %#v, want one success", got) } forwarder.Release() receiveTestSignal(t, done, "stream request completion") - if got := stats.Snapshot(1, now); got != (health.KeyStats{Success: 1}) { + if got := stats.Snapshot(1, now); got != (health.CredentialStats{Success: 1}) { t.Fatalf("stats after forward returns = %#v, want one success", got) } }) } preCommitForwarder := &scriptedForwarder{streamResults: []UpstreamResult{{ - Err: errors.New("first stream event failed"), RequestWritten: true, RetryableBeforeCommit: true, + Err: errors.New("first stream event failed"), RequestWritten: true, }}} engine, handler, _, stats := newStatsHandlerTestRuntime(t, preCommitForwarder, "sk-one") handler.now = func() time.Time { return now } @@ -852,7 +915,7 @@ func TestHandlerRecordsStreamSuccessAtReadyTime(t *testing.T) { ) request.Header.Set("Authorization", "Bearer gl-client") engine.ServeHTTP(httptest.NewRecorder(), request) - if got := stats.Snapshot(1, now); got != (health.KeyStats{}) { + if got := stats.Snapshot(1, now); got != (health.CredentialStats{}) { t.Fatalf("pre-commit stream stats = %#v, want zero", got) } } @@ -881,7 +944,7 @@ func TestHandlerDoesNotRecordCanceledAttempt(t *testing.T) { request.Header.Set("Authorization", "Bearer gl-client") engine.ServeHTTP(httptest.NewRecorder(), request) - if got := stats.Snapshot(1, now); got != (health.KeyStats{}) { + if got := stats.Snapshot(1, now); got != (health.CredentialStats{}) { t.Fatalf("canceled attempt stats = %#v, want zero", got) } } @@ -906,7 +969,7 @@ func TestHandlerDoesNotRecordCommittedNonStreamingAttempt(t *testing.T) { request.Header.Set("Authorization", "Bearer gl-client") engine.ServeHTTP(httptest.NewRecorder(), request) - if got := stats.Snapshot(1, now); got != (health.KeyStats{}) { + if got := stats.Snapshot(1, now); got != (health.CredentialStats{}) { t.Fatalf("committed attempt stats = %#v, want zero", got) } } @@ -1054,15 +1117,16 @@ func TestHandlerEnforcesModelUTF8ByteLimitBeforeAttempt(t *testing.T) { "sk-one", ) if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ - ID: 1, Name: "openai", UpstreamURL: "http://upstream.invalid", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 1, Name: "openai", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ ID: "gpt-4o", Alias: test.model, }}, Enabled: true, }}, + Credentials: []state.CredentialConfig{testCredentialConfig(1, 1)}, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", @@ -1166,7 +1230,7 @@ func TestHandlerServesLocalModelEndpoints(t *testing.T) { }, { name: "Gemini with query key", method: http.MethodGet, target: "/v1beta/models?key=gl-client", - expected: `{"models":[{"name":"models/alpha"},{"name":"models/zeta"}]}`, + expected: `{"models":[{"name":"models/alpha"},{"name":"models/beta"},{"name":"models/zeta"}]}`, }, } for _, test := range tests { @@ -1236,10 +1300,10 @@ func TestHandlerModelEndpointHasNoDataPlaneSideEffects(t *testing.T) { } manager := state.NewManager() if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ - ID: 1, Name: "openai", UpstreamURL: "https://unused.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "alpha"}}, Enabled: true, + ID: 1, Name: "openai", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "alpha"}}, Enabled: true, }}, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: keyService.Hash("gl-client"), Status: state.AccessKeyStatusActive, @@ -1249,7 +1313,7 @@ func TestHandlerModelEndpointHasNoDataPlaneSideEffects(t *testing.T) { } spyEncryption := &decryptPanicEncryption{Service: keyService} handler := NewHandler( - manager, state.NewKeyRegistry(), spyEncryption, panicForwarder{}, dialect.NewSet(), health.NewStatsStore(), + manager, state.NewCredentialRegistry(), spyEncryption, panicForwarder{}, dialect.NewSet(), health.NewStatsStore(), health.NewMutationCoordinator(), nil, nil, nil, ) @@ -1303,16 +1367,15 @@ func newModelListHandlerEngineWithLimit( } manager := state.NewManager() if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{ { - ID: 1, Name: "multi", UpstreamURL: "https://multi.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic, protocol.Gemini}, - Models: []state.ModelConfig{{ID: "zeta"}, {ID: "alpha"}}, Enabled: true, + ID: 1, Name: "multi", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "zeta"}, {ID: "alpha"}}, Enabled: true, }, { - ID: 2, Name: "openai", UpstreamURL: "https://openai.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - Models: []state.ModelConfig{{ID: "beta"}}, Enabled: true, + ID: 2, Name: "openai", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "beta"}}, Enabled: true, }, }, AccessKeys: []state.AccessKeyConfig{{ @@ -1323,7 +1386,7 @@ func newModelListHandlerEngineWithLimit( t.Fatalf("Publish() error = %v", err) } handler := NewHandler( - manager, state.NewKeyRegistry(), keyService, &scriptedForwarder{}, dialect.NewSet(), health.NewStatsStore(), + manager, state.NewCredentialRegistry(), keyService, &scriptedForwarder{}, dialect.NewSet(), health.NewStatsStore(), health.NewMutationCoordinator(), nil, nil, nil, ) @@ -1335,19 +1398,19 @@ func newModelListHandlerEngineWithLimit( type panicRuntimeRegistry struct{} -func (panicRuntimeRegistry) CollectCandidates([]uint, func(uint) bool, time.Time) []state.KeyMeta { +func (panicRuntimeRegistry) CollectCredentialCandidates([]uint, func(uint) bool, time.Time) []state.CredentialMeta { panic("model endpoint collected upstream candidates") } -func (panicRuntimeRegistry) ActiveEncryptedValue(uint, uint) (string, bool) { +func (panicRuntimeRegistry) ActiveEncryptedCredentialData(uint, uint) (string, bool) { panic("model endpoint read an upstream key") } -func (panicRuntimeRegistry) CaptureActiveKeyRefs([]uint) []state.KeyRef { +func (panicRuntimeRegistry) CaptureActiveCredentialRefs([]uint) []state.CredentialRef { panic("model endpoint captured upstream keys") } -func (panicRuntimeRegistry) ActiveEncryptedValueIfMatch(state.KeyRef) (string, bool) { +func (panicRuntimeRegistry) ActiveEncryptedCredentialDataIfMatch(state.CredentialRef) (string, bool) { panic("model endpoint matched an upstream key") } @@ -1719,7 +1782,7 @@ func TestHandlerCapturesKeyIdentityOnlyAfterDecodedInspection(t *testing.T) { Body: []byte(`{"ok":true}`), RequestWritten: true, }}} handler, _, registry := newHandlerForTest(t, forwarder, "sk-upstream") - counting := &captureCountingRuntimeRegistry{runtimeKeyRegistry: registry} + counting := &captureCountingRuntimeRegistry{runtimeCredentialRegistry: registry} handler.registry = counting engine := gin.New() bindGatewayRoutesForTest(t, engine, handler) @@ -1754,7 +1817,7 @@ func TestHandlerCapturesKeyIdentityOnlyAfterDecodedInspection(t *testing.T) { t.Run("canceled body read", func(t *testing.T) { forwarder := &scriptedForwarder{} handler, _, registry := newHandlerForTest(t, forwarder, "sk-upstream") - counting := &captureCountingRuntimeRegistry{runtimeKeyRegistry: registry} + counting := &captureCountingRuntimeRegistry{runtimeCredentialRegistry: registry} handler.registry = counting engine := gin.New() bindGatewayRoutesForTest(t, engine, handler) @@ -1790,7 +1853,7 @@ func TestHandlerCapturesKeyIdentityOnlyAfterDecodedInspection(t *testing.T) { t.Run("canceled after successful inspection", func(t *testing.T) { forwarder := &scriptedForwarder{} handler, _, registry := newHandlerForTest(t, forwarder, "sk-upstream") - counting := &captureCountingRuntimeRegistry{runtimeKeyRegistry: registry} + counting := &captureCountingRuntimeRegistry{runtimeCredentialRegistry: registry} handler.registry = counting requestContext, cancel := context.WithCancel(context.Background()) baseDialect := handler.dialects[protocol.OpenAICompletions] @@ -2295,7 +2358,7 @@ func TestHandlerTerminatesRequestWrittenStreamFailuresBeforeCommit(t *testing.T) { name: "protocol error does not reach committed success", results: []UpstreamResult{ - {Err: protocolFailure, RequestWritten: true, RetryableBeforeCommit: true}, + {Err: protocolFailure, RequestWritten: true}, {StatusCode: http.StatusOK, RequestWritten: true, Committed: true}, }, wantStatus: http.StatusBadGateway, wantCode: reasonUpstreamProtocol.Code, wantAttempts: 1, @@ -2303,7 +2366,7 @@ func TestHandlerTerminatesRequestWrittenStreamFailuresBeforeCommit(t *testing.T) { name: "first-event timeout does not reach committed success", results: []UpstreamResult{ - {Err: context.DeadlineExceeded, RequestWritten: true, RetryableBeforeCommit: true}, + {Err: context.DeadlineExceeded, RequestWritten: true}, {StatusCode: http.StatusOK, RequestWritten: true, Committed: true}, }, wantStatus: http.StatusGatewayTimeout, wantCode: reasonUpstreamTimeout.Code, wantAttempts: 1, @@ -2311,24 +2374,24 @@ func TestHandlerTerminatesRequestWrittenStreamFailuresBeforeCommit(t *testing.T) { name: "protocol errors exhausted", results: []UpstreamResult{ - {Err: protocolFailure, RequestWritten: true, RetryableBeforeCommit: true}, - {Err: protocolFailure, RequestWritten: true, RetryableBeforeCommit: true}, + {Err: protocolFailure, RequestWritten: true}, + {Err: protocolFailure, RequestWritten: true}, }, wantStatus: http.StatusBadGateway, wantCode: reasonUpstreamProtocol.Code, wantAttempts: 1, }, { name: "first-event timeouts exhausted", results: []UpstreamResult{ - {Err: context.DeadlineExceeded, RequestWritten: true, RetryableBeforeCommit: true}, - {Err: context.DeadlineExceeded, RequestWritten: true, RetryableBeforeCommit: true}, + {Err: context.DeadlineExceeded, RequestWritten: true}, + {Err: context.DeadlineExceeded, RequestWritten: true}, }, wantStatus: http.StatusGatewayTimeout, wantCode: reasonUpstreamTimeout.Code, wantAttempts: 1, }, { name: "transport failures exhausted", results: []UpstreamResult{ - {Err: errors.New("stream disconnected"), RequestWritten: true, RetryableBeforeCommit: true}, - {Err: errors.New("stream disconnected"), RequestWritten: true, RetryableBeforeCommit: true}, + {Err: errors.New("stream disconnected"), RequestWritten: true}, + {Err: errors.New("stream disconnected"), RequestWritten: true}, }, wantStatus: http.StatusBadGateway, wantCode: reasonUpstreamConnect.Code, wantAttempts: 1, }, @@ -2360,31 +2423,12 @@ func TestHandlerTerminatesRequestWrittenStreamFailuresBeforeCommit(t *testing.T) func TestHandlerRetriesClassifiedFirstProviderErrorAndReturns502OnExhaustion(t *testing.T) { const marker = "rate_limit_error" - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - writer.Header().Set("Retry-After", "1") - _, _ = io.WriteString( - writer, - `event: error`+"\n"+`data: {"error":{"type":"`+marker+`"}}`+"\n\n", - ) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ - "X-Ordinary-Marker": marker, - }} - providerError := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream(t.Context(), input, newRecordingResponseWriter()) - if providerError.Err != nil || - !providerError.ProviderErrorBeforeCommit || - !bytes.Contains(providerError.ClassificationBody, []byte(marker)) { - t.Fatalf( - "ordinary HeaderRule changed Provider classification evidence: %#v", - providerError, - ) + providerError := UpstreamResult{ + StatusCode: http.StatusOK, + Header: http.Header{"Retry-After": []string{"1"}}, + ClassificationBody: []byte(`{"error":{"type":"` + marker + `"}}`), + ProviderErrorBeforeCommit: true, + RequestWritten: true, } t.Run("retryable category changes key", func(t *testing.T) { @@ -2465,7 +2509,7 @@ func TestHandlerFirstProviderErrorDoesNotCommitOrRecordSuccess(t *testing.T) { if recorder.Code != http.StatusBadGateway || len(forwarder.streamInputs) != 1 { t.Fatalf("status/attempts = %d/%d", recorder.Code, len(forwarder.streamInputs)) } - if got := stats.Snapshot(1, now); got != (health.KeyStats{}) { + if got := stats.Snapshot(1, now); got != (health.CredentialStats{}) { t.Fatalf("stats = %#v, want no success/failure record", got) } count, ok := registry.IncrFailure(1) @@ -2479,11 +2523,11 @@ func TestHandlerDoesNotRetryRequestWrittenAmbiguousPreCommitFailures(t *testing. name string result UpstreamResult }{ - {name: "timeout", result: UpstreamResult{Err: context.DeadlineExceeded, RequestWritten: true, RetryableBeforeCommit: true}}, - {name: "read error", result: UpstreamResult{Err: errors.New("read failed"), RequestWritten: true, RetryableBeforeCommit: true}}, - {name: "clean EOF", result: UpstreamResult{Err: errIncompleteSSEEvent, RequestWritten: true, RetryableBeforeCommit: true}}, - {name: "framing", result: UpstreamResult{Err: fmt.Errorf("%w: framing", ErrUpstreamProtocol), RequestWritten: true, RetryableBeforeCommit: true}}, - {name: "encoding", result: UpstreamResult{Err: fmt.Errorf("%w: encoding", ErrUpstreamProtocol), RequestWritten: true, RetryableBeforeCommit: true}}, + {name: "timeout", result: UpstreamResult{Err: context.DeadlineExceeded, RequestWritten: true}}, + {name: "read error", result: UpstreamResult{Err: errors.New("read failed"), RequestWritten: true}}, + {name: "clean EOF", result: UpstreamResult{Err: errIncompleteSSEEvent, RequestWritten: true}}, + {name: "framing", result: UpstreamResult{Err: fmt.Errorf("%w: framing", ErrUpstreamProtocol), RequestWritten: true}}, + {name: "encoding", result: UpstreamResult{Err: fmt.Errorf("%w: encoding", ErrUpstreamProtocol), RequestWritten: true}}, } for _, test := range tests { @@ -2508,9 +2552,9 @@ func TestHandlerDoesNotRetryRequestWrittenAmbiguousPreCommitFailures(t *testing. } } -func TestHandlerRetriesRequestNotWrittenTransportFailure(t *testing.T) { +func TestHandlerSkipsGroupAfterRequestNotWrittenTransportFailure(t *testing.T) { forwarder := &scriptedForwarder{streamResults: []UpstreamResult{ - {Err: errors.New("dial failed"), RequestWritten: false, RetryableBeforeCommit: true}, + {Err: errors.New("dial failed"), RequestWritten: false}, {StatusCode: http.StatusOK, RequestWritten: true, Committed: true}, }} engine, _, _ := newHandlerTestRuntime(t, forwarder, "sk-first", "sk-second") @@ -2523,7 +2567,7 @@ func TestHandlerRetriesRequestNotWrittenTransportFailure(t *testing.T) { recorder := httptest.NewRecorder() engine.ServeHTTP(recorder, request) - if recorder.Code != http.StatusOK || len(forwarder.streamInputs) != 2 { + if recorder.Code != http.StatusBadGateway || len(forwarder.streamInputs) != 1 { t.Fatalf("status/attempts = %d/%d, body=%s", recorder.Code, len(forwarder.streamInputs), recorder.Body.String()) } @@ -2545,7 +2589,7 @@ func TestHandlerStopsAtStreamingTerminalBoundaries(t *testing.T) { { name: "downstream cancellation", result: UpstreamResult{ - Err: context.Canceled, RequestWritten: true, RetryableBeforeCommit: true, + Err: context.Canceled, RequestWritten: true, }, }, } @@ -2588,7 +2632,7 @@ func TestHandlerDoesNotRetryDownstreamWriteDeadline(t *testing.T) { forwarder := &scriptedForwarder{streamResults: []UpstreamResult{ { Err: deadlineErr, RequestWritten: true, - Committed: true, RetryableBeforeCommit: true, + Committed: true, }, {StatusCode: http.StatusOK, RequestWritten: true, Committed: true}, }} @@ -2619,7 +2663,7 @@ func TestHandlerDoesNotRetryDownstreamWriteDeadline(t *testing.T) { func TestHandlerDoesNotAdvanceCandidatesAfterRequestDeadline(t *testing.T) { forwarder := &scriptedForwarder{streamResults: []UpstreamResult{ - {Err: context.DeadlineExceeded, RequestWritten: true, RetryableBeforeCommit: true}, + {Err: context.DeadlineExceeded, RequestWritten: true}, {StatusCode: http.StatusOK, RequestWritten: true, Committed: true}, }} engine, _, _ := newHandlerTestRuntime(t, forwarder, "sk-one", "sk-two") @@ -2767,7 +2811,7 @@ func TestHandlerAppliesExactCooldownDeadline(t *testing.T) { t.Run(test.name, func(t *testing.T) { forwarder := &scriptedForwarder{results: []UpstreamResult{test.result}} handler, _, registry := newHandlerForTest(t, forwarder, "sk-one") - recording := &recordingRuntimeRegistry{KeyRegistry: registry} + recording := &recordingRuntimeRegistry{CredentialRegistry: registry} handler.registry = recording handler.now = func() time.Time { return attemptNow } engine := gin.New() @@ -2779,10 +2823,10 @@ func TestHandlerAppliesExactCooldownDeadline(t *testing.T) { recorder := httptest.NewRecorder() engine.ServeHTTP(recorder, request) - if recording.cooldownCalls != 1 || recording.cooldownKeyID != 1 || + if recording.cooldownCalls != 1 || recording.cooldownCredentialID != 1 || !recording.cooldownUntil.Equal(test.want) { t.Fatalf("cooldown = calls:%d key:%d until:%v, want 1/1/%v", - recording.cooldownCalls, recording.cooldownKeyID, + recording.cooldownCalls, recording.cooldownCredentialID, recording.cooldownUntil, test.want) } }) @@ -2838,7 +2882,7 @@ func TestHandlerBlacklistsKeyOnThirdInvalidFailure(t *testing.T) { if recorder.Code != http.StatusUnauthorized { t.Fatalf("attempt %d response = %d %s", attempt, recorder.Code, recorder.Body.String()) } - blacklisted := registry.BlacklistedKeys() + blacklisted := registry.BlacklistedCredentials() if (attempt < 3 && len(blacklisted) != 0) || (attempt == 3 && (len(blacklisted) != 1 || blacklisted[0].ID != 1)) { t.Fatalf("attempt %d blacklisted = %#v", attempt, blacklisted) @@ -2934,9 +2978,9 @@ func TestHandlerLeavesKeyRegistryUnchangedForNonKeyActions(t *testing.T) { health.Action(255), } { t.Run(fmt.Sprintf("action_%d", action), func(t *testing.T) { - recording := &recordingRuntimeRegistry{KeyRegistry: state.NewKeyRegistry()} + recording := &recordingRuntimeRegistry{CredentialRegistry: state.NewCredentialRegistry()} handler := &Handler{registry: recording} - handler.applyKeyAction(1, health.Result{Action: action}, 0, time.Time{}) + handler.applyCredentialAction(1, health.Result{Action: action}, 0, time.Time{}) if recording.cooldownCalls != 0 || recording.incrFailureCalls != 0 || recording.blacklistCalls != 0 || recording.clearCalls != 0 { t.Fatalf("mutation calls = cooldown:%d failure:%d blacklist:%d clear:%d", @@ -3015,10 +3059,12 @@ func TestHandlerReturnsModelRequiredByFilterForProtocolOnlyRequest(t *testing.T) t.Fatalf("NewService() error = %v", err) } if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ ID: 1, Name: "responses", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, + ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Enabled: true, }}, AccessKeys: []state.AccessKeyConfig{{ @@ -3034,7 +3080,7 @@ func TestHandlerReturnsModelRequiredByFilterForProtocolOnlyRequest(t *testing.T) t.Fatalf("Publish() error = %v", err) } handler.dialects = dialect.NewSet( - dialect.NewOpenAIResponses(http.DefaultClient), + dialect.NewOpenAIResponses(), ) engine := gin.New() bindGatewayRoutesForTest(t, engine, handler) @@ -3169,7 +3215,7 @@ func TestHandlerDoesNotExposeAliasedUpstreamModelWhenRetryBudgetIsExhausted(t *t defer upstream.Close() engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, externalModel, - dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), + dialect.NewSet(dialect.NewOpenAI()), dialectGatewayGroup{id: 1, name: "openai-1", upstreamURL: upstream.URL, apiKeys: []string{"sk-one"}, models: []state.ModelConfig{{ID: upstreamModel, Alias: externalModel}}}, @@ -3216,11 +3262,16 @@ func TestHandlerKeepsFrozenSnapshotAndInjectUsageSettingAcrossRetry(t *testing.T return } if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ - ID: 1, Name: "openai", UpstreamURL: "http://upstream.invalid", Enabled: true, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: []state.ModelConfig{{ID: "gpt-4o"}}, + ID: 1, Name: "openai", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Enabled: true, + Models: []state.ModelConfig{{ID: "gpt-4o"}}, Settings: config.Settings{state.SettingInjectUsageOptions: false}, }}, + Credentials: []state.CredentialConfig{ + testCredentialConfig(1, 1), + testCredentialConfig(2, 1), + }, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: keyService.Hash("gl-client"), Status: state.AccessKeyStatusActive, }}, @@ -3254,18 +3305,18 @@ func TestHandlerKeepsFrozenSnapshotAndInjectUsageSettingAcrossRetry(t *testing.T func TestHandlerSkipsCandidateChangedAfterCollection(t *testing.T) { tests := []struct { name string - mutate func(*testing.T, *state.KeyRegistry, encryption.Service) + mutate func(*testing.T, *state.CredentialRegistry, encryption.Service) }{ { name: "key moved to another group", - mutate: func(t *testing.T, registry *state.KeyRegistry, keyService encryption.Service) { + mutate: func(t *testing.T, registry *state.CredentialRegistry, keyService encryption.Service) { t.Helper() encrypted, err := keyService.Encrypt("sk-group-two") if err != nil { t.Fatalf("Encrypt(group two key) error = %v", err) } - if err := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 2, Status: state.KeyStatusActive, EncryptedValue: encrypted, + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 2, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, EncryptedValue: encrypted, }}); err != nil { t.Fatalf("Replace(moved key) error = %v", err) } @@ -3273,10 +3324,10 @@ func TestHandlerSkipsCandidateChangedAfterCollection(t *testing.T) { }, { name: "key disabled", - mutate: func(t *testing.T, registry *state.KeyRegistry, _ encryption.Service) { + mutate: func(t *testing.T, registry *state.CredentialRegistry, _ encryption.Service) { t.Helper() - if err := registry.SetKeyStatus(1, state.KeyStatusDisabled); err != nil { - t.Fatalf("SetKeyStatus(disabled) error = %v", err) + if err := registry.SetCredentialStatus(1, state.CredentialStatusDisabled); err != nil { + t.Fatalf("SetCredentialStatus(disabled) error = %v", err) } }, }, @@ -3293,10 +3344,10 @@ func TestHandlerSkipsCandidateChangedAfterCollection(t *testing.T) { t.Fatalf("NewService() error = %v", err) } runtimeRegistry := &mutatingRuntimeRegistry{ - KeyRegistry: registry, - mutate: func() { tt.mutate(t, registry, keyService) }, + CredentialRegistry: registry, + mutate: func() { tt.mutate(t, registry, keyService) }, } - openAI := dialect.NewOpenAI(http.DefaultClient) + openAI := dialect.NewOpenAI() handler := NewHandler( manager, registry, keyService, forwarder, dialect.NewSet(openAI), health.NewStatsStore(), health.NewMutationCoordinator(), @@ -3332,32 +3383,32 @@ func TestHandlerFreezesKeyIdentityAfterInspectingBody(t *testing.T) { tests := []struct { name string seedPlain string - seedStatus state.KeyStatus + seedStatus state.CredentialStatus currentPlain string - mutate func(*state.KeyRegistry, string) error + mutate func(*state.CredentialRegistry, string) error }{ { name: "new import", currentPlain: "sk-imported", - mutate: func(registry *state.KeyRegistry, encrypted string) error { - return registry.ApplyImport(1, []state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + mutate: func(registry *state.CredentialRegistry, encrypted string) error { + return registry.ApplyCredentialImport(1, []state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, EncryptedValue: encrypted, }}) }, }, { name: "disabled becomes active", seedPlain: "sk-enabled", - seedStatus: state.KeyStatusDisabled, currentPlain: "sk-enabled", - mutate: func(registry *state.KeyRegistry, _ string) error { - return registry.SetKeyStatus(1, state.KeyStatusActive) + seedStatus: state.CredentialStatusDisabled, currentPlain: "sk-enabled", + mutate: func(registry *state.CredentialRegistry, _ string) error { + return registry.SetCredentialStatus(1, state.CredentialStatusActive) }, }, { name: "same ID gets new ciphertext", seedPlain: "sk-old", - seedStatus: state.KeyStatusActive, currentPlain: "sk-replaced", - mutate: func(registry *state.KeyRegistry, encrypted string) error { - return registry.ApplyImport(1, []state.KeyEntry{{ - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + seedStatus: state.CredentialStatusActive, currentPlain: "sk-replaced", + mutate: func(registry *state.CredentialRegistry, encrypted string) error { + return registry.ApplyCredentialImport(1, []state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, EncryptedValue: encrypted, }}) }, @@ -3397,8 +3448,8 @@ func TestHandlerFreezesKeyIdentityAfterInspectingBody(t *testing.T) { t.Fatalf("Encrypt(seed key) error = %v", err) } } - if replaceErr := registry.Replace([]state.KeyEntry{{ - ID: 1, GroupID: 1, Status: test.seedStatus, + if replaceErr := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: test.seedStatus, EncryptedValue: seedCiphertext, }}); replaceErr != nil { t.Fatalf("Replace(seed key) error = %v", replaceErr) @@ -3473,18 +3524,18 @@ func TestHandlerFreezesKeyIdentityAfterInspectingBody(t *testing.T) { func TestHandlerAllowsCapturedUnavailableIdentityAfterRecovery(t *testing.T) { tests := []struct { name string - makeUnavailable func(*testing.T, *state.KeyRegistry) - recover func(*testing.T, *state.KeyRegistry, string) + makeUnavailable func(*testing.T, *state.CredentialRegistry) + recover func(*testing.T, *state.CredentialRegistry, string) }{ { name: "blacklisted", - makeUnavailable: func(t *testing.T, registry *state.KeyRegistry) { + makeUnavailable: func(t *testing.T, registry *state.CredentialRegistry) { t.Helper() if ok := registry.SetBlacklisted(3); !ok { t.Fatal("SetBlacklisted(3) = false") } }, - recover: func(t *testing.T, registry *state.KeyRegistry, _ string) { + recover: func(t *testing.T, registry *state.CredentialRegistry, _ string) { t.Helper() if ok := registry.Recover(3); !ok { t.Fatal("Recover(3) = false") @@ -3493,16 +3544,16 @@ func TestHandlerAllowsCapturedUnavailableIdentityAfterRecovery(t *testing.T) { }, { name: "cooldown", - makeUnavailable: func(t *testing.T, registry *state.KeyRegistry) { + makeUnavailable: func(t *testing.T, registry *state.CredentialRegistry) { t.Helper() if ok := registry.SetCooldown(3, time.Now().Add(time.Hour)); !ok { t.Fatal("SetCooldown(3) = false") } }, - recover: func(t *testing.T, registry *state.KeyRegistry, encrypted string) { + recover: func(t *testing.T, registry *state.CredentialRegistry, encrypted string) { t.Helper() - if err := registry.ApplyImport(1, []state.KeyEntry{{ - ID: 3, GroupID: 1, Status: state.KeyStatusActive, + if err := registry.ApplyCredentialImport(1, []state.CredentialEntry{{ + ID: 3, GroupID: 1, Version: 1, IdentityGeneration: 3, Fingerprint: "test-3", Status: state.CredentialStatusActive, EncryptedValue: encrypted, }}); err != nil { t.Fatalf("ApplyImport(recovered cooldown key) error = %v", err) @@ -3553,17 +3604,17 @@ func TestHandlerAllowsCapturedUnavailableIdentityAfterRecovery(t *testing.T) { return encrypted } recoverableCiphertext := encrypt("sk-recoverable") - if err := registry.Replace([]state.KeyEntry{ + if err := registry.ReplaceCredentials([]state.CredentialEntry{ { - ID: 1, GroupID: 1, Status: state.KeyStatusActive, + ID: 1, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-1", Status: state.CredentialStatusActive, EncryptedValue: encrypt("sk-first"), }, { - ID: 2, GroupID: 1, Status: state.KeyStatusDisabled, + ID: 2, GroupID: 1, Version: 1, IdentityGeneration: 2, Fingerprint: "test-2", Status: state.CredentialStatusDisabled, EncryptedValue: encrypt("sk-newly-enabled"), }, { - ID: 3, GroupID: 1, Status: state.KeyStatusActive, + ID: 3, GroupID: 1, Version: 1, IdentityGeneration: 3, Fingerprint: "test-3", Status: state.CredentialStatusActive, EncryptedValue: recoverableCiphertext, }, }); err != nil { @@ -3585,8 +3636,8 @@ func TestHandlerAllowsCapturedUnavailableIdentityAfterRecovery(t *testing.T) { }() receiveTestSignal(t, firstForward, "first forward") - if err := registry.SetKeyStatus(2, state.KeyStatusActive); err != nil { - t.Fatalf("SetKeyStatus(2, active) error = %v", err) + if err := registry.SetCredentialStatus(2, state.CredentialStatusActive); err != nil { + t.Fatalf("SetCredentialStatus(2, active) error = %v", err) } test.recover(t, registry, recoverableCiphertext) close(releaseForward) @@ -3618,12 +3669,19 @@ func newRealGatewayEngine(t *testing.T, upstreamURL string, upstreamKeys ...stri t.Fatalf("NewService() error = %v", err) } manager := state.NewManager() + baseURL := testUpstreamBaseURL(upstreamURL, protocol.OpenAICompletions) + channelID, params := testChannelConfig(t, protocol.OpenAICompletions, baseURL) + credentialConfigs := make([]state.CredentialConfig, 0, len(upstreamKeys)) + for index := range upstreamKeys { + credentialConfigs = append(credentialConfigs, testCredentialConfig(uint(index+1), 1)) + } if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ - ID: 1, Name: "openai", UpstreamURL: testUpstreamBaseURL(upstreamURL, protocol.OpenAICompletions), - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "gpt-4o"}}, Enabled: true, + ID: 1, Name: "openai", ChannelID: channelID, Params: params, + Models: []state.ModelConfig{{ID: "gpt-4o"}}, Enabled: true, }}, + Credentials: credentialConfigs, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: keyService.Hash("gl-client"), Status: state.AccessKeyStatusActive, @@ -3631,28 +3689,21 @@ func newRealGatewayEngine(t *testing.T, upstreamURL string, upstreamKeys ...stri }); err != nil { t.Fatalf("Publish() error = %v", err) } - registry := state.NewKeyRegistry() - entries := make([]state.KeyEntry, 0, len(upstreamKeys)) + registry := state.NewCredentialRegistry() + entries := make([]state.CredentialEntry, 0, len(upstreamKeys)) for index, plaintext := range upstreamKeys { - encrypted, err := keyService.Encrypt(plaintext) - if err != nil { - t.Fatalf("Encrypt() error = %v", err) - } - entries = append(entries, state.KeyEntry{ - ID: uint(index + 1), GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: encrypted, - }) + entries = append(entries, testCredentialEntry(t, keyService, uint(index+1), 1, plaintext)) } - if err := registry.Replace(entries); err != nil { - t.Fatalf("Replace() error = %v", err) + if err := registry.ReplaceCredentials(entries); err != nil { + t.Fatalf("ReplaceCredentials() error = %v", err) } - clients := platformhttp.NewHTTPClientManager() - openAI := dialect.NewOpenAI(clients.GetClient(testDialectClientConfig())) + openAI := dialect.NewOpenAI() handler := NewHandler( manager, registry, keyService, - NewForwarder(clients, redact.New()), + newTestExecutionForwarder(t), dialect.NewSet(openAI), health.NewStatsStore(), health.NewMutationCoordinator(), @@ -3679,7 +3730,7 @@ func newHandlerTestRuntime( t *testing.T, forwarder AttemptForwarder, upstreamKeys ...string, -) (*gin.Engine, *state.Manager, *state.KeyRegistry) { +) (*gin.Engine, *state.Manager, *state.CredentialRegistry) { t.Helper() handler, manager, registry := newHandlerForTest(t, forwarder, upstreamKeys...) engine := gin.New() @@ -3691,7 +3742,7 @@ func newStatsHandlerTestRuntime( t *testing.T, forwarder AttemptForwarder, upstreamKeys ...string, -) (*gin.Engine, *Handler, *state.KeyRegistry, *health.StatsStore) { +) (*gin.Engine, *Handler, *state.CredentialRegistry, *health.StatsStore) { t.Helper() stats := health.NewStatsStore() handler, _, registry := newHandlerForTestWithStats(t, forwarder, stats, upstreamKeys...) @@ -3704,7 +3755,7 @@ func newHandlerForTest( t *testing.T, forwarder AttemptForwarder, upstreamKeys ...string, -) (*Handler, *state.Manager, *state.KeyRegistry) { +) (*Handler, *state.Manager, *state.CredentialRegistry) { return newHandlerForTestWithStats(t, forwarder, health.NewStatsStore(), upstreamKeys...) } @@ -3713,7 +3764,7 @@ func newHandlerForTestWithStats( forwarder AttemptForwarder, stats *health.StatsStore, upstreamKeys ...string, -) (*Handler, *state.Manager, *state.KeyRegistry) { +) (*Handler, *state.Manager, *state.CredentialRegistry) { t.Helper() gin.SetMode(gin.TestMode) keyService, err := encryption.NewService("handler-test-master-key") @@ -3721,12 +3772,25 @@ func newHandlerForTestWithStats( t.Fatalf("NewService() error = %v", err) } manager := state.NewManager() + credentialConfigs := make([]state.CredentialConfig, 0, len(upstreamKeys)) + for index := range upstreamKeys { + credentialConfigs = append(credentialConfigs, state.CredentialConfig{ + ID: uint(index + 1), + GroupID: 1, + Status: state.CredentialStatusActive, + Version: 1, + IdentityGeneration: uint64(index + 1), + Fingerprint: fmt.Sprintf("credential-%d", index+1), + }) + } if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ - ID: 1, Name: "openai", UpstreamURL: "http://upstream.invalid", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "gpt-4o"}}, Enabled: true, + ID: 1, Name: "openai", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "gpt-4o"}}, Enabled: true, }}, + Credentials: credentialConfigs, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: keyService.Hash("gl-client"), Status: state.AccessKeyStatusActive, @@ -3734,22 +3798,29 @@ func newHandlerForTestWithStats( }); err != nil { t.Fatalf("Publish() error = %v", err) } - registry := state.NewKeyRegistry() - entries := make([]state.KeyEntry, 0, len(upstreamKeys)) + registry := state.NewCredentialRegistry() + entries := make([]state.CredentialEntry, 0, len(upstreamKeys)) for index, plaintext := range upstreamKeys { - encrypted, err := keyService.Encrypt(plaintext) + credential, err := json.Marshal(map[string]string{"api_key": plaintext}) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + encrypted, err := keyService.Encrypt(string(credential)) if err != nil { t.Fatalf("Encrypt() error = %v", err) } - entries = append(entries, state.KeyEntry{ - ID: uint(index + 1), GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: encrypted, + entries = append(entries, state.CredentialEntry{ + ID: uint(index + 1), GroupID: 1, + Version: 1, IdentityGeneration: uint64(index + 1), + Fingerprint: fmt.Sprintf("credential-%d", index+1), + Status: state.CredentialStatusActive, EncryptedValue: encrypted, }) } - if err := registry.Replace(entries); err != nil { - t.Fatalf("Replace() error = %v", err) + if err := registry.ReplaceCredentials(entries); err != nil { + t.Fatalf("ReplaceCredentials() error = %v", err) } - openAI := dialect.NewOpenAI(http.DefaultClient) + openAI := dialect.NewOpenAI() handler := NewHandler( manager, registry, keyService, forwarder, dialect.NewSet(openAI), stats, health.NewMutationCoordinator(), diff --git a/internal/gateway/models.go b/internal/gateway/models.go index 98dec42c..6ed096dc 100644 --- a/internal/gateway/models.go +++ b/internal/gateway/models.go @@ -91,16 +91,21 @@ func collectVisibleModelIDs( continue } } - for modelID, targets := range snapshot.Candidates[selectedProtocol] { - if len(accessKey.Filters.Models) > 0 { - if _, ok := accessKey.Filters.Models[modelID]; !ok { + for _, byModel := range snapshot.ExecutionCandidates[selectedProtocol] { + for modelID, targets := range byModel { + if modelID == state.NoModelRouteKey { continue } + if len(accessKey.Filters.Models) > 0 { + if _, ok := accessKey.Filters.Models[modelID]; !ok { + continue + } + } + if !anyVisibleTarget(targets, accessKey.Filters.Groups) { + continue + } + visible[modelID] = struct{}{} } - if !anyVisibleTarget(targets, accessKey.Filters.Groups) { - continue - } - visible[modelID] = struct{}{} } } result := make([]string, 0, len(visible)) diff --git a/internal/gateway/models_test.go b/internal/gateway/models_test.go index c596b8c8..93736e1f 100644 --- a/internal/gateway/models_test.go +++ b/internal/gateway/models_test.go @@ -7,33 +7,32 @@ import ( "strings" "testing" + "gpt-load/internal/channel" "gpt-load/internal/protocol" "gpt-load/internal/state" ) func TestVisibleModelIDs(t *testing.T) { snapshot, err := state.Compile(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{ { - ID: 1, Name: "first", UpstreamURL: "https://first.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic, protocol.Gemini}, + ID: 1, Name: "first", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{ {ID: "zeta"}, {ID: "shared", Alias: "first-alias"}, {ID: "alpha"}, }, Enabled: true, }, { - ID: 2, Name: "second", UpstreamURL: "https://second.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, + ID: 2, Name: "second", ChannelID: channel.Anthropic, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{ {ID: "shared", Alias: "second-alias"}, {ID: "beta"}, }, Enabled: true, }, { - ID: 3, Name: "disabled", UpstreamURL: "https://disabled.example.com", - Protocols: []protocol.Protocol{protocol.Gemini}, - Models: []state.ModelConfig{{ID: "hidden"}}, Enabled: false, + ID: 3, Name: "disabled", ChannelID: channel.Gemini, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "hidden"}}, Enabled: false, }, }, }) @@ -55,8 +54,8 @@ func TestVisibleModelIDs(t *testing.T) { {name: "group filter keeps any matching target", snapshot: snapshot, accessKey: state.AccessKeyView{Filters: state.FilterSet{Groups: map[uint]struct{}{2: {}}}}, value: protocol.OpenAICompletions, want: []string{"beta", "second-alias"}}, {name: "joint filters", snapshot: snapshot, accessKey: state.AccessKeyView{Filters: state.FilterSet{Protocols: map[protocol.Protocol]struct{}{protocol.Anthropic: {}}, Models: map[string]struct{}{"beta": {}, "second-alias": {}, "shared": {}}, Groups: map[uint]struct{}{2: {}}}}, value: protocol.Anthropic, want: []string{"beta", "second-alias"}}, {name: "dangling group filter", snapshot: snapshot, accessKey: state.AccessKeyView{Filters: state.FilterSet{Groups: map[uint]struct{}{99: {}}}}, value: protocol.OpenAICompletions, want: []string{}}, - {name: "disabled group model absent", snapshot: snapshot, value: protocol.Gemini, want: []string{"alpha", "first-alias", "zeta"}}, - {name: "missing protocol", snapshot: snapshot, value: protocol.OpenAIResponses, want: []string{}}, + {name: "disabled group model absent", snapshot: snapshot, value: protocol.Gemini, want: []string{"alpha", "beta", "first-alias", "second-alias", "zeta"}}, + {name: "responses uses channel capabilities", snapshot: snapshot, value: protocol.OpenAIResponses, want: []string{"alpha", "beta", "first-alias", "second-alias", "zeta"}}, {name: "nil snapshot", snapshot: nil, value: protocol.OpenAICompletions, want: []string{}}, } for _, test := range tests { @@ -72,10 +71,9 @@ func TestVisibleModelIDs(t *testing.T) { func TestVisibleOpenAIModelIDsUnionsChatAndResponses(t *testing.T) { t.Parallel() - snapshot, err := state.Compile(state.CompileInput{Groups: []state.GroupConfig{ + snapshot, err := state.Compile(state.CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{ { - ID: 1, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{ {ID: "chat-only"}, {ID: "shared"}, @@ -83,8 +81,7 @@ func TestVisibleOpenAIModelIDsUnionsChatAndResponses(t *testing.T) { Enabled: true, }, { - ID: 2, - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, + ID: 2, ChannelID: channel.Anthropic, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{ {ID: "responses-only"}, {ID: "shared"}, @@ -92,11 +89,7 @@ func TestVisibleOpenAIModelIDsUnionsChatAndResponses(t *testing.T) { Enabled: true, }, { - ID: 3, - Protocols: []protocol.Protocol{ - protocol.OpenAICompletions, - protocol.OpenAIResponses, - }, + ID: 3, ChannelID: channel.Gemini, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "both"}}, Enabled: true, }, @@ -121,7 +114,7 @@ func TestVisibleOpenAIModelIDsUnionsChatAndResponses(t *testing.T) { protocol.OpenAIResponses: {}, }, }}, - want: []string{"both", "responses-only", "shared"}, + want: []string{"both", "chat-only", "responses-only", "shared"}, }, { name: "chat protocol filter", @@ -130,7 +123,7 @@ func TestVisibleOpenAIModelIDsUnionsChatAndResponses(t *testing.T) { protocol.OpenAICompletions: {}, }, }}, - want: []string{"both", "chat-only", "shared"}, + want: []string{"both", "chat-only", "responses-only", "shared"}, }, } @@ -202,16 +195,15 @@ func TestMarshalModelList(t *testing.T) { func TestBuildVisibleModelListBoundsFinalAnthropicJSONAcrossGroups(t *testing.T) { snapshot, err := state.Compile(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{ { - ID: 1, Name: "first", UpstreamURL: "https://first.example.com", - Protocols: []protocol.Protocol{protocol.Anthropic}, - Models: []state.ModelConfig{{ID: "zeta"}}, Enabled: true, + ID: 1, Name: "first", ChannelID: channel.Anthropic, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "zeta"}}, Enabled: true, }, { - ID: 2, Name: "second", UpstreamURL: "https://second.example.com", - Protocols: []protocol.Protocol{protocol.Anthropic}, - Models: []state.ModelConfig{{ID: "alpha"}}, Enabled: true, + ID: 2, Name: "second", ChannelID: channel.Anthropic, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "alpha"}}, Enabled: true, }, }, }) diff --git a/internal/gateway/openai_responses_stream_test.go b/internal/gateway/openai_responses_stream_test.go deleted file mode 100644 index 80839d0d..00000000 --- a/internal/gateway/openai_responses_stream_test.go +++ /dev/null @@ -1,545 +0,0 @@ -package gateway - -import ( - "context" - "errors" - "io" - "net/http" - "net/http/httptest" - "strings" - "sync" - "testing" - "time" - - "gpt-load/internal/dialect" - "gpt-load/internal/health" - platformhttp "gpt-load/internal/platform/httpclient" - "gpt-load/internal/platform/redact" - "gpt-load/internal/usage" -) - -type cancelAfterFlushResponseWriter struct { - *recordingResponseWriter - cancel context.CancelFunc - cancelAt int -} - -func (writer *cancelAfterFlushResponseWriter) FlushError() error { - writer.flushes++ - if writer.flushes == writer.cancelAt { - writer.cancel() - } - return nil -} - -type cancelAfterFlushWaitForCloseResponseWriter struct { - *recordingResponseWriter - cancel context.CancelFunc - closed <-chan struct{} - cancelAt int -} - -func (writer *cancelAfterFlushWaitForCloseResponseWriter) FlushError() error { - writer.flushes++ - if writer.flushes == writer.cancelAt { - writer.cancel() - select { - case <-writer.closed: - case <-time.After(time.Second): - return errors.New("timed out waiting for stream close") - } - } - return nil -} - -type closeNotifyingReadCloser struct { - reader io.Reader - closed chan struct{} - closeOnce sync.Once -} - -func (body *closeNotifyingReadCloser) Read(target []byte) (int, error) { - return body.reader.Read(target) -} - -func (body *closeNotifyingReadCloser) Close() error { - body.closeOnce.Do(func() { close(body.closed) }) - return nil -} - -type releaseAfterFlushResponseWriter struct { - *recordingResponseWriter - release func() - releaseOnce sync.Once -} - -func (writer *releaseAfterFlushResponseWriter) FlushError() error { - writer.flushes++ - writer.releaseOnce.Do(writer.release) - return nil -} - -func TestResponsesSSETerminalLifecycleAndUsage(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - wire string - wantReason StreamEndReason - wantUsage usage.Result - wantErr bool - }{ - { - name: "completed explicit event", - wire: "event: response.completed\n" + - "data: {\"response\":{\"usage\":{\"input_tokens\":100,\"input_tokens_details\":{\"cached_tokens\":20},\"output_tokens\":30,\"total_tokens\":130}}}\n\n", - wantReason: StreamEndCleanEOF, - wantUsage: usage.Result{ - State: usage.StateComplete, - Tokens: usage.Tokens{ - UncachedInput: 80, - CacheRead: 20, - Output: 30, - }, - }, - }, - { - name: "incomplete terminal", - wire: "data: {\"type\":\"response.incomplete\",\"response\":{\"usage\":{\"input_tokens\":50,\"output_tokens\":5,\"total_tokens\":55}}}\n\n", - wantReason: StreamEndProviderIncomplete, - wantUsage: usage.Result{ - State: usage.StateComplete, - Tokens: usage.Tokens{UncachedInput: 50, Output: 5}, - }, - }, - { - name: "missing terminal after partial usage", - wire: "data: {\"type\":\"response.output_item.added\",\"response\":{\"usage\":{\"input_tokens\":10,\"output_tokens\":1}}}\n\n", - wantReason: StreamEndUpstreamProtocolError, - wantUsage: usage.Result{ - State: usage.StatePartial, - Tokens: usage.Tokens{UncachedInput: 10, Output: 1}, - }, - wantErr: true, - }, - } - - for _, test := range tests { - test := test - t.Run(test.name, func(t *testing.T) { - t.Parallel() - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - _ *http.Request, - ) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, test.wire) - })) - defer upstream.Close() - - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - context.Background(), - responsesStreamForwardInput(upstream.URL), - downstream, - ) - - if result.Committed != true || - result.Stream.EndReason != test.wantReason || - result.Usage != test.wantUsage { - t.Fatalf("ForwardStream() result = %#v", result) - } - if test.wantErr != errors.Is(result.Err, ErrUpstreamProtocol) { - t.Fatalf("ForwardStream() error = %v, want protocol error %t", result.Err, test.wantErr) - } - if downstream.status != http.StatusOK || - downstream.body.String() != test.wire { - t.Fatalf( - "downstream status/body = %d/%q, want %d/%q", - downstream.status, - downstream.body.String(), - http.StatusOK, - test.wire, - ) - } - }) - } -} - -func TestResponsesSSECompletedFlushWinsOverLaterClientCancellation(t *testing.T) { - t.Parallel() - - first := "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" - terminal := "event: response.completed\n" + - "data: {\"response\":{\"usage\":{\"input_tokens\":100,\"output_tokens\":30,\"total_tokens\":130}}}\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - request *http.Request, - ) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, first+terminal) - writer.(http.Flusher).Flush() - <-request.Context().Done() - })) - defer upstream.Close() - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - downstream := &cancelAfterFlushResponseWriter{ - recordingResponseWriter: newRecordingResponseWriter(), - cancel: cancel, - cancelAt: 2, - } - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - ctx, - responsesStreamForwardInput(upstream.URL), - downstream, - ) - - wantUsage := usage.Result{ - State: usage.StateComplete, - Tokens: usage.Tokens{UncachedInput: 100, Output: 30}, - } - if !errors.Is(result.Err, context.Canceled) || !result.Committed || - result.Stream.EndReason != StreamEndCleanEOF || - result.Usage != wantUsage || - downstream.status != http.StatusOK || - downstream.body.String() != first+terminal { - t.Fatalf("ForwardStream() result/downstream = %#v / %#v", result, downstream) - } -} - -func TestResponsesSSECompletedFlushKeepsTerminalBoundaryWhenCloseRacesAfterFlush(t *testing.T) { - first := "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" - terminal := "event: response.completed\n" + - "data: {\"response\":{\"usage\":{\"input_tokens\":100,\"output_tokens\":30,\"total_tokens\":130}}}\n\n" - closed := make(chan struct{}) - body := &closeNotifyingReadCloser{ - reader: strings.NewReader(first + terminal), - closed: closed, - } - clients := platformhttp.NewHTTPClientManager() - input := responsesStreamForwardInput("https://api.example.test") - clients.GetClient(streamingClientConfig(input.Group.Timeouts)).Transport = forwarderRoundTripper(func(request *http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{ - "Content-Type": {"text/event-stream"}, - }, - Body: body, - Request: request, - }, nil - }) - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - downstream := &cancelAfterFlushWaitForCloseResponseWriter{ - recordingResponseWriter: newRecordingResponseWriter(), - cancel: cancel, - closed: closed, - cancelAt: 2, - } - result := NewForwarder(clients, redact.New()).ForwardStream(ctx, input, downstream) - - if !errors.Is(result.Err, context.Canceled) || !result.Committed || - result.Stream.EndReason != StreamEndCleanEOF || - downstream.body.String() != first+terminal { - t.Fatalf("ForwardStream() result/downstream = %#v / %#v", result, downstream) - } -} - -func TestResponsesSSELargeTerminalAfterPureCREventDoesNotMarkPartialTerminalForwarded(t *testing.T) { - first := "event: response.output_text.delta\r" + - "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\r\r" - terminal := "event: response.completed\n" + - "data: {\"response\":{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"" + - strings.Repeat("x", streamReadBufferSize*2) + - "\"}]}],\"usage\":{\"input_tokens\":100,\"output_tokens\":30,\"total_tokens\":130}}}\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - _ *http.Request, - ) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, first+terminal) - })) - defer upstream.Close() - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - downstream := &cancelAfterFlushResponseWriter{ - recordingResponseWriter: newRecordingResponseWriter(), - cancel: cancel, - cancelAt: 2, - } - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - ctx, - responsesStreamForwardInput(upstream.URL), - downstream, - ) - - if !errors.Is(result.Err, context.Canceled) || - result.Stream.EndReason != StreamEndClientCanceled || - downstream.body.Len() <= len(first) || - downstream.body.Len() >= len(first+terminal) { - t.Fatalf("ForwardStream() result = %#v, want client cancellation before terminal forwarding", result) - } -} - -func TestResponsesSSEForwardsCompleteTerminalEventBeforeContinuing(t *testing.T) { - t.Parallel() - - first := "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" - terminal := "event: response.completed\n" + - "data: {\"response\":{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"" + - strings.Repeat("x", streamReadBufferSize*2) + - "\"}]}],\"usage\":{\"input_tokens\":100,\"output_tokens\":30,\"total_tokens\":130}}}\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - _ *http.Request, - ) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, first+terminal) - })) - defer upstream.Close() - - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - context.Background(), - responsesStreamForwardInput(upstream.URL), - downstream, - ) - - wantUsage := usage.Result{ - State: usage.StateComplete, - Tokens: usage.Tokens{UncachedInput: 100, Output: 30}, - } - if result.Err != nil || !result.Committed || - result.Stream.EndReason != StreamEndCleanEOF || - result.Usage != wantUsage || - downstream.status != http.StatusOK || - downstream.body.String() != first+terminal { - t.Fatalf( - "ForwardStream() result/body length = %#v / %d, want complete terminal length %d", - result, - downstream.body.Len(), - len(first+terminal), - ) - } -} - -func TestResponsesSSEPreservesSplitCRLFAfterTerminal(t *testing.T) { - t.Parallel() - - first := "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" - terminalPrefix := "event: response.completed\r\n" + - "data: {\"response\":{\"usage\":{\"input_tokens\":100,\"output_tokens\":30,\"total_tokens\":130}}}\r\n" + - "\r" - release := make(chan struct{}) - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - _ *http.Request, - ) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, first+terminalPrefix) - writer.(http.Flusher).Flush() - <-release - _, _ = io.WriteString(writer, "\n") - writer.(http.Flusher).Flush() - })) - defer upstream.Close() - - downstream := &releaseAfterFlushResponseWriter{ - recordingResponseWriter: newRecordingResponseWriter(), - release: func() { close(release) }, - } - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - context.Background(), - responsesStreamForwardInput(upstream.URL), - downstream, - ) - - wantUsage := usage.Result{ - State: usage.StateComplete, - Tokens: usage.Tokens{UncachedInput: 100, Output: 30}, - } - if result.Err != nil || !result.Committed || - result.Stream.EndReason != StreamEndCleanEOF || - result.Usage != wantUsage || - downstream.status != http.StatusOK || - downstream.body.String() != first+terminalPrefix+"\n" { - t.Fatalf("ForwardStream() result/body = %#v / %q", result, downstream.body.String()) - } -} - -func TestResponsesSSEFailedBeforeAndAfterCommit(t *testing.T) { - t.Parallel() - - t.Run("first failed event remains pre-commit provider error", func(t *testing.T) { - t.Parallel() - wire := "event: response.failed\n" + - "data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"code\":\"rate_limit_error\"},\"usage\":{\"input_tokens\":25,\"output_tokens\":2,\"total_tokens\":27}}}\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - _ *http.Request, - ) { - writer.Header().Set("Content-Type", "text/event-stream") - writer.Header().Set("Retry-After", "1") - _, _ = io.WriteString(writer, wire) - })) - defer upstream.Close() - - input := responsesStreamForwardInput(upstream.URL) - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream(context.Background(), input, downstream) - - if result.Err != nil || result.Committed || - !result.ProviderErrorBeforeCommit || - downstream.status != 0 || downstream.body.Len() != 0 { - t.Fatalf("ForwardStream() result/downstream = %#v / %#v", result, downstream) - } - if want := (usage.Result{ - State: usage.StateComplete, - Tokens: usage.Tokens{UncachedInput: 25, Output: 2}, - }); result.Usage != want { - t.Fatalf("ForwardStream() usage = %#v, want %#v", result.Usage, want) - } - decision := health.Judge(input.Dialect, health.Attempt{ - StatusCode: result.StatusCode, - Body: result.ClassificationBody, - Header: result.Header, - Now: time.Unix(1, 0), - ProviderErrorBeforeCommit: true, - }) - if decision.Category != health.FailureCategoryRateLimited || - decision.Action != health.ActionCooldownKey { - t.Fatalf("health decision = %#v", decision) - } - }) - - t.Run("failed after first event cannot cross commit boundary", func(t *testing.T) { - t.Parallel() - wire := "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" + - "event: response.failed\n" + - "data: {\"response\":{\"error\":{\"code\":\"server_error\"}}}\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - _ *http.Request, - ) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, wire) - })) - defer upstream.Close() - - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - context.Background(), - responsesStreamForwardInput(upstream.URL), - downstream, - ) - - if result.Err != nil || !result.Committed || - result.ProviderErrorBeforeCommit || - result.Stream.EndReason != StreamEndSSEError || - downstream.status != http.StatusOK || - downstream.body.String() != wire { - t.Fatalf("ForwardStream() result/downstream = %#v / %#v", result, downstream) - } - }) -} - -func TestResponsesSSEEventNameConflictIsProtocolError(t *testing.T) { - t.Parallel() - - wire := "event: response.completed\n" + - "data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"code\":\"server_error\"}}}\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - _ *http.Request, - ) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, wire) - })) - defer upstream.Close() - - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - context.Background(), - responsesStreamForwardInput(upstream.URL), - downstream, - ) - if !errors.Is(result.Err, ErrUpstreamProtocol) || - result.Committed || result.ProviderErrorBeforeCommit || - downstream.status != 0 || downstream.body.Len() != 0 { - t.Fatalf("ForwardStream() result/downstream = %#v / %#v", result, downstream) - } -} - -func TestResponsesSSERejectsEventsAfterTerminalWithoutReplacingUsage(t *testing.T) { - t.Parallel() - - terminal := "event: response.completed\n" + - "data: {\"response\":{\"usage\":{\"input_tokens\":100,\"output_tokens\":30,\"total_tokens\":130}}}\n\n" - lateEvent := "event: response.output_item.added\n" + - "data: {\"response\":{\"usage\":{\"input_tokens\":900,\"output_tokens\":90,\"total_tokens\":990}}}\n\n" - upstream := httptest.NewServer(http.HandlerFunc(func( - writer http.ResponseWriter, - _ *http.Request, - ) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, terminal+lateEvent) - })) - defer upstream.Close() - - downstream := newRecordingResponseWriter() - result := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).ForwardStream( - context.Background(), - responsesStreamForwardInput(upstream.URL), - downstream, - ) - wantUsage := usage.Result{ - State: usage.StateComplete, - Tokens: usage.Tokens{UncachedInput: 100, Output: 30}, - } - if !errors.Is(result.Err, ErrUpstreamProtocol) || !result.Committed || - result.Stream.EndReason != StreamEndUpstreamProtocolError || - result.Usage != wantUsage || - downstream.body.String() != terminal { - t.Fatalf("ForwardStream() result/downstream = %#v / %#v", result, downstream) - } -} - -func responsesStreamForwardInput(upstreamURL string) ForwardInput { - input := streamForwardInput(upstreamURL) - input.Dialect = dialect.NewOpenAIResponses(http.DefaultClient) - input.Request.Path = "/v1/responses" - input.Request.Body = []byte(`{"model":"gpt-5","stream":true}`) - input.Group.InjectUsageOptions = true - return input -} diff --git a/internal/gateway/representation_test_helpers_test.go b/internal/gateway/representation_test_helpers_test.go new file mode 100644 index 00000000..d4cbe61c --- /dev/null +++ b/internal/gateway/representation_test_helpers_test.go @@ -0,0 +1,55 @@ +package gateway + +import ( + "net/http" + "strings" + "testing" +) + +func setRepresentationMetadata(headers http.Header) { + headers.Set("ETag", `"wire-v1"`) + headers.Set("Digest", "sha-256=wire-digest") + headers.Set("Content-MD5", "d2lyZQ==") + headers.Set("Content-Range", "bytes 0-9/10") + headers.Set("Content-Digest", "sha-256=:d2lyZQ==:") + headers.Set("Repr-Digest", "sha-256=:cmVwcg==:") + headers.Set("Signature", "sig1=:c2lnbmF0dXJl:") + headers.Set("Signature-Input", `sig1=("content-digest");created=1`) +} + +func assertHeadersDoNotContain(t *testing.T, headers http.Header, literal string) { + t.Helper() + for name, values := range headers { + for _, value := range values { + if strings.Contains(value, literal) { + t.Fatalf("header %q leaked %q in value %q", name, literal, value) + } + } + } +} + +func assertRepresentationMetadata(t *testing.T, headers http.Header, wantPreserved bool) { + t.Helper() + want := map[string]string{ + "ETag": `"wire-v1"`, + "Digest": "sha-256=wire-digest", + "Content-MD5": "d2lyZQ==", + "Content-Range": "bytes 0-9/10", + "Content-Digest": "sha-256=:d2lyZQ==:", + "Repr-Digest": "sha-256=:cmVwcg==:", + } + for name, value := range want { + got := headers.Get(name) + if wantPreserved && got != value { + t.Errorf("%s = %q, want preserved value %q", name, got, value) + } + if !wantPreserved && got != "" { + t.Errorf("%s = %q, want removed after body rewrite", name, got) + } + } + for _, name := range []string{"Signature", "Signature-Input"} { + if got := headers.Get(name); got != "" { + t.Errorf("%s = %q, want removed from downstream response", name, got) + } + } +} diff --git a/internal/gateway/request_id_test.go b/internal/gateway/request_id_test.go index 2c6719e5..7ed9007d 100644 --- a/internal/gateway/request_id_test.go +++ b/internal/gateway/request_id_test.go @@ -2,7 +2,6 @@ package gateway import ( "bytes" - "context" "encoding/json" "errors" "net/http" @@ -12,11 +11,7 @@ import ( "github.com/sirupsen/logrus" - "gpt-load/internal/dialect" - platformhttp "gpt-load/internal/platform/httpclient" - "gpt-load/internal/platform/redact" "gpt-load/internal/ratelimit" - "gpt-load/internal/state" ) func TestHandlerGeneratesCanonicalRequestIDAndRejectsSpoofing(t *testing.T) { @@ -124,32 +119,3 @@ func TestHandlerRequestIDGenerationFailurePreservesDataPlaneAndSkipsEmit(t *test t.Fatalf("request ID warning contains redundant plane field: %#v", entry) } } - -func TestNewUpstreamRequestRemovesReservedRequestIDAfterHeaderRules(t *testing.T) { - openAI := dialect.NewOpenAI(http.DefaultClient) - input := ForwardInput{ - Dialect: openAI, - Group: state.GroupView{ - UpstreamURL: "https://upstream.example.com", - HeaderRules: state.HeaderRules{ - Set: map[string]string{requestIDHeader: "group-spoof-${API_KEY}"}, - }, - }, - APIKey: "opaque-provider-secret", - Request: &dialect.ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/chat/completions", - Header: http.Header{requestIDHeader: {"client-spoof"}}, - Body: []byte(`{"model":"gpt-4o"}`), - }, - } - - request, _, replay, err := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).newUpstreamRequest(context.Background(), input, false) - if err != nil { - t.Fatalf("newUpstreamRequest() error = %v", err) - } - t.Cleanup(replay.release) - if got := request.Header.Values(requestIDHeader); len(got) != 0 { - t.Fatalf("upstream %s = %#v, want reserved header removed", requestIDHeader, got) - } -} diff --git a/internal/gateway/request_log.go b/internal/gateway/request_log.go index aa14a71c..94cbb865 100644 --- a/internal/gateway/request_log.go +++ b/internal/gateway/request_log.go @@ -12,6 +12,7 @@ import ( "time" "unicode/utf8" + "gpt-load/internal/execution" "gpt-load/internal/health" "gpt-load/internal/platform/redact" "gpt-load/internal/pricing" @@ -40,6 +41,7 @@ type requestOutcome struct { } type frozenAttemptPricing struct { + channelID string groupID uint upstreamModel string table *pricing.Table @@ -52,6 +54,7 @@ type requestRecorder struct { startedAt time.Time accessKeyID uint protocol protocol.Protocol + operation execution.Operation clientModel string stream bool firstResponseMs *int64 @@ -159,6 +162,12 @@ func (recorder *requestRecorder) setClientModel(model string) { } } +func (recorder *requestRecorder) setOperation(operation execution.Operation) { + if recorder != nil { + recorder.operation = operation + } +} + func (recorder *requestRecorder) setStream(stream bool) { if recorder != nil { recorder.stream = stream @@ -213,7 +222,7 @@ func (recorder *requestRecorder) beforeForward() time.Time { func (recorder *requestRecorder) recordAttempt( selection scheduler.Selection, - apiKey string, + credentialSecrets []string, result UpstreamResult, decision health.Result, startedAt time.Time, @@ -227,8 +236,9 @@ func (recorder *requestRecorder) recordAttempt( rules = selection.Group.HeaderRules } summarySecrets := resolvedErrorSummarySecretValues( - apiKey, + "", rules, + credentialSecrets..., ) return recorder.appendAttempt( selection, @@ -244,7 +254,7 @@ func (recorder *requestRecorder) recordAttempt( func (recorder *requestRecorder) recordStreamAttempt( selection scheduler.Selection, - apiKey string, + credentialSecrets []string, result UpstreamResult, startedAt time.Time, completedAt time.Time, @@ -258,8 +268,9 @@ func (recorder *requestRecorder) recordStreamAttempt( rules = selection.Group.HeaderRules } summarySecrets := resolvedErrorSummarySecretValues( - apiKey, + "", rules, + credentialSecrets..., ) return recorder.appendAttempt( selection, @@ -295,18 +306,24 @@ func (recorder *requestRecorder) appendAttempt( duration = 0 } attempt := telemetry.Attempt{ - Sequence: len(recorder.attempts) + 1, - GroupID: selection.GroupID, - GroupName: selection.Group.Name, - KeyID: selection.KeyID, - UpstreamModel: optionalModelValue(selection.UpstreamModelID), - StatusCode: result.StatusCode, - DurationMs: duration.Milliseconds(), - FailureCategory: category, - Action: action, - ErrorCode: errorCode, - ErrorSummary: errorSummary, - Committed: result.Committed, + Sequence: len(recorder.attempts) + 1, + GroupID: selection.GroupID, + GroupName: selection.Group.Name, + ChannelID: selection.ChannelID, + CredentialID: selection.CredentialID, + Operation: recorder.operation, + RouteMode: selection.RouteMode, + UpstreamModel: optionalModelValue(selection.UpstreamModelID), + UpstreamRequestID: result.UpstreamRequestID, + DispatchState: result.DispatchState, + ResponseStarted: result.ResponseStarted, + StatusCode: result.StatusCode, + DurationMs: duration.Milliseconds(), + FailureCategory: category, + Action: action, + ErrorCode: errorCode, + ErrorSummary: errorSummary, + Committed: result.Committed, } if attempt.ErrorCode != "" && attempt.ErrorSummary == "" { attempt.ErrorSummary = fixedErrorSummary(attempt.ErrorCode) @@ -444,7 +461,7 @@ func (recorder *requestRecorder) bindUsage( return } attempt := recorder.attempts[attemptIndex] - if attempt.GroupID == 0 || attempt.KeyID == 0 || attempt.Sequence < 1 { + if attempt.GroupID == 0 || attempt.ChannelID == "" || attempt.CredentialID == 0 || attempt.Sequence < 1 { return } if !applicable || !recorder.usageApplicable { @@ -462,7 +479,8 @@ func (recorder *requestRecorder) bindUsage( recorder.usage = telemetry.UsageObservation{ Result: result, GroupID: attempt.GroupID, - KeyID: attempt.KeyID, + ChannelID: attempt.ChannelID, + CredentialID: attempt.CredentialID, AttemptSequence: attempt.Sequence, Pricing: pricingObservation, } @@ -485,7 +503,10 @@ func quoteFrozenAttempt( if frozen.table == nil || frozen.upstreamModel == "" { return observation } - quote, receipt := frozen.table.QuoteWithReceipt(pricing.Identity{ModelID: frozen.upstreamModel}, result) + quote, receipt := frozen.table.QuoteWithReceipt(pricing.Identity{ + ChannelID: frozen.channelID, + ModelID: frozen.upstreamModel, + }, result) observation.CostState = string(quote.State) observation.PricingCompleteness = string(quote.Completeness) observation.EstimatedCostNanoUSD = int64(quote.EstimatedCostNanoUSD) @@ -579,10 +600,10 @@ func telemetryAction(value health.Action) telemetry.Action { switch value { case health.ActionRetry: return telemetry.ActionRetry - case health.ActionCooldownKey: - return telemetry.ActionCooldownKey - case health.ActionFailKey: - return telemetry.ActionFailKey + case health.ActionCooldownCredential: + return telemetry.ActionCooldownCredential + case health.ActionFailCredential: + return telemetry.ActionFailCredential case health.ActionSkipGroup: return telemetry.ActionSkipGroup default: diff --git a/internal/gateway/request_log_test.go b/internal/gateway/request_log_test.go index 10d6c9c2..38b830aa 100644 --- a/internal/gateway/request_log_test.go +++ b/internal/gateway/request_log_test.go @@ -3,6 +3,7 @@ package gateway import ( "bytes" "context" + "encoding/json" "errors" "fmt" "io" @@ -10,7 +11,6 @@ import ( "net/http" "net/http/httptest" "reflect" - "strconv" "strings" "sync" "testing" @@ -20,10 +20,10 @@ import ( "github.com/gin-gonic/gin" "github.com/sirupsen/logrus" + "gpt-load/internal/channel" "gpt-load/internal/dialect" "gpt-load/internal/health" "gpt-load/internal/platform/encryption" - platformhttp "gpt-load/internal/platform/httpclient" "gpt-load/internal/platform/redact" "gpt-load/internal/pricing" "gpt-load/internal/protocol" @@ -35,6 +35,14 @@ import ( "gpt-load/internal/usage" ) +func requestLogSelection(groupID, credentialID uint, name string) scheduler.Selection { + return scheduler.Selection{ + GroupID: groupID, ChannelID: channel.OpenAI, CredentialID: credentialID, + RouteMode: channel.RouteNative, + Group: state.GroupView{ID: groupID, Name: name, ChannelID: channel.OpenAI}, + } +} + func TestRequestRecorderCapturesFirstResponseAfterDownstreamCommit(t *testing.T) { startedAt := time.Date(2026, time.August, 5, 8, 0, 0, 0, time.UTC) sink := &recordingRequestLogSink{} @@ -203,7 +211,7 @@ func TestNewRequestRecorderInitializesUsageNotApplicable(t *testing.T) { if event.Usage.Result.State != usage.StateNotApplicable || event.Usage.Result.Tokens != (usage.Tokens{}) || event.Usage.GroupID != 0 || - event.Usage.KeyID != 0 || + event.Usage.CredentialID != 0 || event.Usage.AttemptSequence != 0 { t.Fatalf("initial Usage = %#v", event.Usage) } @@ -225,13 +233,12 @@ func TestHandlerRecordsProtocolOnlyResponsesResourceWithoutFabricatingModels(t * "sk-first", ) if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ - ID: 1, - Name: "responses", - UpstreamURL: "http://upstream.invalid", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Enabled: true, + ID: 1, Name: "responses", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Enabled: true, }}, + Credentials: []state.CredentialConfig{testCredentialConfig(1, 1)}, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", @@ -241,7 +248,7 @@ func TestHandlerRecordsProtocolOnlyResponsesResourceWithoutFabricatingModels(t * }); err != nil { t.Fatalf("Publish() error = %v", err) } - handler.dialects = dialect.NewSet(dialect.NewOpenAIResponses(http.DefaultClient)) + handler.dialects = dialect.NewSet(dialect.NewOpenAIResponses()) request := httptest.NewRequest(http.MethodGet, "/v1/responses/resp_123", nil) request.Header.Set("Authorization", "Bearer gl-client") @@ -351,7 +358,7 @@ func TestHandlerRecordsModelConsistencyOnlyForSuccessfulModeledRequests(t *testi } } -func TestHandlerBillableResponsesUsageWithoutModelIsUnpriced(t *testing.T) { +func TestHandlerBillableResponsesUsageWithoutPriceIsUnpriced(t *testing.T) { forwarder := &scriptedForwarder{results: []UpstreamResult{{ StatusCode: http.StatusOK, Header: make(http.Header), @@ -367,10 +374,13 @@ func TestHandlerBillableResponsesUsageWithoutModelIsUnpriced(t *testing.T) { t, forwarder, &recordingAccessKeyRPMLimiter{}, sink, "sk-first", ) if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ - ID: 1, Name: "responses", UpstreamURL: "http://upstream.invalid", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, Enabled: true, + ID: 1, Name: "responses", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "unpriced-model"}}, Enabled: true, }}, + Credentials: []state.CredentialConfig{testCredentialConfig(1, 1)}, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: handler.encryption.Hash("gl-client"), Status: state.AccessKeyStatusActive, @@ -378,16 +388,16 @@ func TestHandlerBillableResponsesUsageWithoutModelIsUnpriced(t *testing.T) { }); err != nil { t.Fatal(err) } - handler.dialects = dialect.NewSet(dialect.NewOpenAIResponses(http.DefaultClient)) + handler.dialects = dialect.NewSet(dialect.NewOpenAIResponses()) handler.priceTables = &mutableGatewayPriceTableProvider{table: mustGatewayPriceTable(t, 1, true)} - request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"input":"ping"}`)) + request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"unpriced-model","input":"ping"}`)) request.Header.Set("Authorization", "Bearer gl-client") engine.ServeHTTP(httptest.NewRecorder(), request) events := sink.snapshot() if len(events) != 1 || withoutPricingReceipt(events[0].Usage.Pricing) != (telemetry.PricingObservation{ - CostState: "unpriced", PricingCompleteness: "unavailable", + UpstreamModel: "unpriced-model", CostState: "unpriced", PricingCompleteness: "unavailable", }) { t.Fatalf("events = %#v", events) } @@ -406,17 +416,17 @@ func TestRequestRecorderBindsUsageToRecordedAttempt(t *testing.T) { t.Run(test.name, func(t *testing.T) { sink := &recordingRequestLogSink{} recorder := newRequestRecorder(sink, "req-usage-bind", time.Unix(100, 0), 9, protocol.OpenAICompletions, func() time.Time { return time.Unix(101, 0) }) - recorder.appendAttempt(scheduler.Selection{GroupID: 11, Group: state.GroupView{Name: "first"}, KeyID: 21}, UpstreamResult{}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(100, 0)) - second := recorder.appendAttempt(scheduler.Selection{GroupID: 12, Group: state.GroupView{Name: "second"}, KeyID: 22}, UpstreamResult{StatusCode: http.StatusOK}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(100, 0)) + recorder.appendAttempt(requestLogSelection(11, 21, "first"), UpstreamResult{}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(100, 0)) + second := recorder.appendAttempt(requestLogSelection(12, 22, "second"), UpstreamResult{StatusCode: http.StatusOK}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(100, 0)) recorder.completeResponse(UpstreamResult{StatusCode: http.StatusOK, Usage: test.result}, health.Result{}, "provider", second) recorder.emit() event := sink.events[0] - if event.Usage.Result != test.result || event.Usage.GroupID != 12 || event.Usage.KeyID != 22 || event.Usage.AttemptSequence != 2 { + if event.Usage.Result != test.result || event.Usage.GroupID != 12 || event.Usage.ChannelID != channel.OpenAI || event.Usage.CredentialID != 22 || event.Usage.AttemptSequence != 2 { t.Fatalf("Usage = %#v", event.Usage) } - if event.Attempts[1].GroupID != event.Usage.GroupID || event.Attempts[1].KeyID != event.Usage.KeyID || event.Attempts[1].Sequence != event.Usage.AttemptSequence { + if event.Attempts[1].GroupID != event.Usage.GroupID || event.Attempts[1].ChannelID != event.Usage.ChannelID || event.Attempts[1].CredentialID != event.Usage.CredentialID || event.Attempts[1].Sequence != event.Usage.AttemptSequence { t.Fatalf("attempt/Usage identity = %#v/%#v", event.Attempts[1], event.Usage) } }) @@ -426,12 +436,12 @@ func TestRequestRecorderBindsUsageToRecordedAttempt(t *testing.T) { func TestRequestRecorderNon2xxKeepsAttributionButNotApplicable(t *testing.T) { sink := &recordingRequestLogSink{} recorder := newRequestRecorder(sink, "req-usage-429", time.Unix(100, 0), 9, protocol.OpenAICompletions, func() time.Time { return time.Unix(101, 0) }) - index := recorder.appendAttempt(scheduler.Selection{GroupID: 12, Group: state.GroupView{Name: "second"}, KeyID: 22}, UpstreamResult{StatusCode: http.StatusTooManyRequests}, telemetry.FailureCategoryRateLimited, telemetry.ActionRetry, "", "", time.Unix(100, 0), time.Unix(100, 0)) + index := recorder.appendAttempt(requestLogSelection(12, 22, "second"), UpstreamResult{StatusCode: http.StatusTooManyRequests}, telemetry.FailureCategoryRateLimited, telemetry.ActionRetry, "", "", time.Unix(100, 0), time.Unix(100, 0)) recorder.completeResponse(UpstreamResult{StatusCode: http.StatusTooManyRequests, Usage: usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{Output: 30}}}, health.Result{Category: health.FailureCategoryRateLimited}, "provider", index) recorder.emit() event := sink.events[0] - if event.Usage.Result != (usage.Result{State: usage.StateNotApplicable}) || event.Usage.GroupID != 12 || event.Usage.KeyID != 22 || event.Usage.AttemptSequence != 1 { + if event.Usage.Result != (usage.Result{State: usage.StateNotApplicable}) || event.Usage.GroupID != 12 || event.Usage.ChannelID != channel.OpenAI || event.Usage.CredentialID != 22 || event.Usage.AttemptSequence != 1 { t.Fatalf("Usage = %#v", event.Usage) } } @@ -441,7 +451,7 @@ func TestRequestRecorderInvalidAttemptIndexDoesNotForgeUsageAttribution(t *testi t.Run(fmt.Sprintf("index_%d", index), func(t *testing.T) { sink := &recordingRequestLogSink{} recorder := newRequestRecorder(sink, "req-usage-invalid", time.Unix(100, 0), 9, protocol.OpenAICompletions, func() time.Time { return time.Unix(101, 0) }) - recorder.appendAttempt(scheduler.Selection{GroupID: 12, Group: state.GroupView{Name: "second"}, KeyID: 22}, UpstreamResult{}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(100, 0)) + recorder.appendAttempt(requestLogSelection(12, 22, "second"), UpstreamResult{}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(100, 0)) recorder.bindUsage(index, usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{Output: 30}}, true) recorder.emit() @@ -461,14 +471,14 @@ func TestRequestRecorderInvalidAttemptIndexDoesNotForgeUsageAttribution(t *testi func TestRequestRecorderDownstreamFailureKeepsBoundUsage(t *testing.T) { sink := &recordingRequestLogSink{} recorder := newRequestRecorder(sink, "req-usage-write", time.Unix(100, 0), 9, protocol.OpenAICompletions, func() time.Time { return time.Unix(101, 0) }) - index := recorder.appendAttempt(scheduler.Selection{GroupID: 12, Group: state.GroupView{Name: "second"}, KeyID: 22}, UpstreamResult{StatusCode: http.StatusOK}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(100, 0)) + index := recorder.appendAttempt(requestLogSelection(12, 22, "second"), UpstreamResult{StatusCode: http.StatusOK}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(100, 0)) result := usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}} recorder.completeResponse(UpstreamResult{StatusCode: http.StatusOK, Usage: result}, health.Result{}, "provider", index) recorder.completeDownstreamWrite(http.StatusOK) recorder.emit() event := sink.events[0] - if event.Status != telemetry.RequestStatusIncomplete || event.Usage.Result != result || event.Usage.GroupID != 12 || event.Usage.KeyID != 22 || event.Usage.AttemptSequence != 1 { + if event.Status != telemetry.RequestStatusIncomplete || event.Usage.Result != result || event.Usage.GroupID != 12 || event.Usage.ChannelID != channel.OpenAI || event.Usage.CredentialID != 22 || event.Usage.AttemptSequence != 1 { t.Fatalf("event = %#v", event) } } @@ -486,7 +496,7 @@ func TestHandlerPublishesUsageForActualNonStreamingResponseAttempt(t *testing.T) engine.ServeHTTP(httptest.NewRecorder(), request) events := sink.snapshot() - if len(events) != 1 || events[0].Usage.Result != (usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}}) || events[0].Usage.GroupID != 1 || events[0].Usage.KeyID != 2 || events[0].Usage.AttemptSequence != 2 { + if len(events) != 1 || events[0].Usage.Result != (usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}}) || events[0].Usage.GroupID != 1 || events[0].Usage.CredentialID != 2 || events[0].Usage.AttemptSequence != 2 { t.Fatalf("events = %#v", events) } } @@ -546,13 +556,13 @@ func TestRequestRecorderPreservesKnownCostWhenUsedOutputPriceIsUnavailable(t *te ) model := "gpt-4o" recorder.freezeNextAttemptPricing(frozenAttemptPricing{ - groupID: 1, upstreamModel: model, table: table, applicable: true, + channelID: string(channel.OpenAI), groupID: 1, + upstreamModel: model, table: table, applicable: true, }) + selection := requestLogSelection(1, 2, "group") + selection.UpstreamModelID = &model index := recorder.appendAttempt( - scheduler.Selection{ - GroupID: 1, Group: state.GroupView{Name: "group"}, KeyID: 2, - UpstreamModelID: &model, - }, + selection, UpstreamResult{StatusCode: http.StatusOK}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(101, 0), @@ -585,13 +595,13 @@ func TestRequestRecorderMergesRequestPricingDiagnosticsIntoKnownCost(t *testing. recorder.setUsageDiagnostics(diagnostics) model := "gpt-4o" recorder.freezeNextAttemptPricing(frozenAttemptPricing{ - groupID: 1, upstreamModel: model, table: table, applicable: true, + channelID: string(channel.OpenAI), groupID: 1, + upstreamModel: model, table: table, applicable: true, }) + selection := requestLogSelection(1, 2, "group") + selection.UpstreamModelID = &model index := recorder.appendAttempt( - scheduler.Selection{ - GroupID: 1, Group: state.GroupView{Name: "group"}, KeyID: 2, - UpstreamModelID: &model, - }, + selection, UpstreamResult{StatusCode: http.StatusOK}, telemetry.FailureCategoryOK, telemetry.ActionTerminate, "", "", time.Unix(100, 0), time.Unix(101, 0), @@ -617,10 +627,9 @@ func TestRequestRecorderMergesRequestPricingDiagnosticsIntoKnownCost(t *testing. } func TestHandlerUsesSelectedProviderModelForPricingInsteadOfAliasOrBodyModel(t *testing.T) { - providerID := "openai" model := "provider-model" table, err := pricing.NewTable([]pricing.Rule{{ - Identity: pricing.Identity{ModelID: model}, + Identity: pricing.Identity{ChannelID: string(channel.OpenAI), ModelID: model}, Prices: pricing.Prices{Input: pricing.Price{ NanoUSDPerMillion: 3_000_000_000, Set: true, }}, @@ -638,12 +647,13 @@ func TestHandlerUsesSelectedProviderModelForPricingInsteadOfAliasOrBodyModel(t * t, forwarder, &recordingAccessKeyRPMLimiter{}, sink, "sk-first", ) if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ - ID: 1, Name: "provider", ProviderID: &providerID, - UpstreamURL: "http://upstream.invalid", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: model, Alias: "client-alias"}}, Enabled: true, + ID: 1, Name: "provider", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: model, Alias: "client-alias"}}, Enabled: true, }}, + Credentials: []state.CredentialConfig{testCredentialConfig(1, 1)}, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: handler.encryption.Hash("gl-client"), Status: state.AccessKeyStatusActive, @@ -718,7 +728,7 @@ func mustGatewayPriceTable(t *testing.T, input int64, outputSet bool) *pricing.T prices.Output = pricing.Price{NanoUSDPerMillion: pricing.NanoUSD(input), Set: true} } table, err := pricing.NewTable([]pricing.Rule{{ - Identity: pricing.Identity{ModelID: "gpt-4o"}, + Identity: pricing.Identity{ChannelID: string(channel.OpenAI), ModelID: "gpt-4o"}, Prices: prices, }}) if err != nil { @@ -753,11 +763,11 @@ func TestHandlerDiscardsPreCommitStreamUsageOnRetry(t *testing.T) { } event := events[0] if event.Usage.Result != second || event.Usage.Result == first || - event.Usage.GroupID != 1 || event.Usage.KeyID != 2 || event.Usage.AttemptSequence != 2 { + event.Usage.GroupID != 1 || event.Usage.CredentialID != 2 || event.Usage.AttemptSequence != 2 { t.Fatalf("event Usage = %#v, want second attempt Usage", event.Usage) } if len(event.Attempts) != 2 || event.Attempts[1].GroupID != event.Usage.GroupID || - event.Attempts[1].KeyID != event.Usage.KeyID || event.Attempts[1].Sequence != event.Usage.AttemptSequence { + event.Attempts[1].CredentialID != event.Usage.CredentialID || event.Attempts[1].Sequence != event.Usage.AttemptSequence { t.Fatalf("attempts/Usage = %#v/%#v", event.Attempts, event.Usage) } } @@ -821,7 +831,7 @@ func TestHandlerFirstProviderErrorRequestLogContract(t *testing.T) { }, }) || event.Usage.GroupID != 1 || - event.Usage.KeyID != 1 || + event.Usage.CredentialID != 1 || event.Usage.AttemptSequence != 1 || len(event.Attempts) != 1 { t.Fatalf("response/event = %d/%#v", response.Code, event) @@ -829,7 +839,7 @@ func TestHandlerFirstProviderErrorRequestLogContract(t *testing.T) { attempt := event.Attempts[0] if attempt.StatusCode != http.StatusOK || attempt.FailureCategory != telemetry.FailureCategoryRateLimited || - attempt.Action != telemetry.ActionCooldownKey || + attempt.Action != telemetry.ActionCooldownCredential || attempt.Committed || attempt.ErrorSummary != fixedErrorSummary("upstream_sse_error") { t.Fatalf("attempt = %#v", attempt) @@ -866,7 +876,7 @@ func TestHandlerRetryExhaustionUsesProviderErrorAttemptAndItsFrozenPrice(t *test ClassificationBody: []byte(`{"error":"invalid key"}`), RequestWritten: true, }, - {Err: errors.New("dial failed"), RetryableBeforeCommit: true}, + {Err: errors.New("dial failed")}, }} forwarder.onStreamCall = func(index int, _ http.ResponseWriter) { switch index { @@ -904,7 +914,7 @@ func TestHandlerRetryExhaustionUsesProviderErrorAttemptAndItsFrozenPrice(t *test t.Fatalf("response/events = %d/%#v", response.Code, events) } event := events[0] - if event.Usage.GroupID != 1 || event.Usage.KeyID != 1 || + if event.Usage.GroupID != 1 || event.Usage.CredentialID != 1 || event.Usage.AttemptSequence != 1 || withoutPricingReceipt(event.Usage.Pricing) != (telemetry.PricingObservation{ UpstreamModel: "gpt-4o", CostState: string(pricing.CostStatePriced), @@ -923,8 +933,8 @@ func TestHandlerRememberedLastResponseKeepsOriginalUsageAttempt(t *testing.T) { sink := &recordingRequestLogSink{} forwarder := &scriptedForwarder{results: []UpstreamResult{ {StatusCode: http.StatusTooManyRequests, Header: make(http.Header), ClassificationBody: []byte(`{"error":{"type":"rate_limit_error"}}`), RequestWritten: true}, - {Err: errors.New("transport two"), RetryableBeforeCommit: true}, - {Err: errors.New("transport three"), RetryableBeforeCommit: true}, + {Err: errors.New("transport two")}, + {Err: errors.New("transport three")}, }} engine, handler, _, _ := newRequestLogHandlerTestRuntime(t, forwarder, &recordingAccessKeyRPMLimiter{}, sink, "sk-first", "sk-second", "sk-third") handler.newRandom = func() *rand.Rand { return rand.New(zeroSource{}) } @@ -933,7 +943,7 @@ func TestHandlerRememberedLastResponseKeepsOriginalUsageAttempt(t *testing.T) { engine.ServeHTTP(httptest.NewRecorder(), request) events := sink.snapshot() - if len(events) != 1 || events[0].Usage.Result != (usage.Result{State: usage.StateNotApplicable}) || events[0].Usage.GroupID != 1 || events[0].Usage.KeyID != 1 || events[0].Usage.AttemptSequence != 1 { + if len(events) != 1 || events[0].Usage.Result != (usage.Result{State: usage.StateNotApplicable}) || events[0].Usage.GroupID != 1 || events[0].Usage.CredentialID != 1 || events[0].Usage.AttemptSequence != 1 { t.Fatalf("events = %#v", events) } } @@ -949,7 +959,7 @@ func TestHandlerFinalNon2xxPublishesNotApplicableWithResponseAttribution(t *test engine.ServeHTTP(httptest.NewRecorder(), request) events := sink.snapshot() - if len(events) != 1 || events[0].Usage.Result != (usage.Result{State: usage.StateNotApplicable}) || events[0].Usage.GroupID != 1 || events[0].Usage.KeyID != 1 || events[0].Usage.AttemptSequence != 1 { + if len(events) != 1 || events[0].Usage.Result != (usage.Result{State: usage.StateNotApplicable}) || events[0].Usage.GroupID != 1 || events[0].Usage.CredentialID != 1 || events[0].Usage.AttemptSequence != 1 { t.Fatalf("events = %#v", events) } } @@ -961,7 +971,7 @@ func TestHandlerTerminalAttemptUsageKeepsRouteAttribution(t *testing.T) { upstreamKeys []string wantStatus telemetry.RequestStatus wantGroupID uint - wantKeyID uint + wantCredentialID uint wantAttemptSequence int wantUpstreamModel string }{ @@ -971,7 +981,7 @@ func TestHandlerTerminalAttemptUsageKeepsRouteAttribution(t *testing.T) { upstreamKeys: []string{"sk-first"}, wantStatus: telemetry.RequestStatusError, wantGroupID: 1, - wantKeyID: 1, + wantCredentialID: 1, wantAttemptSequence: 1, wantUpstreamModel: "gpt-4o", }, @@ -981,22 +991,22 @@ func TestHandlerTerminalAttemptUsageKeepsRouteAttribution(t *testing.T) { upstreamKeys: []string{"sk-first"}, wantStatus: telemetry.RequestStatusCanceled, wantGroupID: 1, - wantKeyID: 1, + wantCredentialID: 1, wantAttemptSequence: 1, wantUpstreamModel: "gpt-4o", }, { - name: "retried transports use the final attempt", + name: "transport failure keeps the skipped group attempt", forwarder: &scriptedForwarder{results: []UpstreamResult{ - {Err: errors.New("transport one"), RetryableBeforeCommit: true}, - {Err: errors.New("transport two"), RetryableBeforeCommit: true}, + {Err: errors.New("transport one")}, + {Err: errors.New("transport two")}, {Err: errors.New("transport three"), RequestWritten: true}, }}, upstreamKeys: []string{"sk-first", "sk-second", "sk-third"}, wantStatus: telemetry.RequestStatusError, wantGroupID: 1, - wantKeyID: 3, - wantAttemptSequence: 3, + wantCredentialID: 1, + wantAttemptSequence: 1, wantUpstreamModel: "gpt-4o", }, { @@ -1013,6 +1023,10 @@ func TestHandlerTerminalAttemptUsageKeepsRouteAttribution(t *testing.T) { request.Header.Set("Authorization", "Bearer gl-client") engine.ServeHTTP(httptest.NewRecorder(), request) events := sink.snapshot() + var wantChannelID channel.ID + if test.wantGroupID != 0 { + wantChannelID = channel.OpenAI + } wantPricing := telemetry.PricingObservation{ CostState: "not_applicable", PricingCompleteness: "not_applicable", @@ -1025,7 +1039,8 @@ func TestHandlerTerminalAttemptUsageKeepsRouteAttribution(t *testing.T) { events[0].Usage != (telemetry.UsageObservation{ Result: usage.Result{State: usage.StateNotApplicable}, GroupID: test.wantGroupID, - KeyID: test.wantKeyID, + ChannelID: wantChannelID, + CredentialID: test.wantCredentialID, AttemptSequence: test.wantAttemptSequence, Pricing: wantPricing, }) { @@ -1253,19 +1268,6 @@ func (reader cancelingErrorReadCloser) Read([]byte) (int, error) { func (cancelingErrorReadCloser) Close() error { return nil } -type countingHealthDialect struct { - dialect.Dialect - classifyCalls int -} - -func (value *countingHealthDialect) ClassifyStatus( - status int, - body []byte, -) health.FailureCategory { - value.classifyCalls++ - return value.Dialect.ClassifyStatus(status, body) -} - type cancelingExtractDialect struct { dialect.Dialect cancel context.CancelFunc @@ -1316,7 +1318,7 @@ func newRequestLogHandlerTestRuntime( limiter AccessKeyRPMLimiter, sink telemetry.RequestLogSink, upstreamKeys ...string, -) (*gin.Engine, *Handler, *state.Manager, *state.KeyRegistry) { +) (*gin.Engine, *Handler, *state.Manager, *state.CredentialRegistry) { t.Helper() handler, manager, registry := newHandlerForTest(t, forwarder, upstreamKeys...) handler.limiter = limiter @@ -1390,10 +1392,10 @@ func TestHandlerUsesFrozenRPMLimitAcrossSnapshotPublish(t *testing.T) { publish := func(limit int64) { t.Helper() if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{{ - ID: 1, Name: "openai", UpstreamURL: "https://unused.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "gpt-4o"}}, Enabled: true, + ID: 1, Name: "openai", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "gpt-4o"}}, Enabled: true, }}, AccessKeys: []state.AccessKeyConfig{{ ID: 1, Name: "client", KeyHash: keyService.Hash("gl-client"), @@ -1412,7 +1414,7 @@ func TestHandlerUsesFrozenRPMLimitAcrossSnapshotPublish(t *testing.T) { } handler := NewHandler( manager, - state.NewKeyRegistry(), + state.NewCredentialRegistry(), keyService, &scriptedForwarder{}, dialect.NewSet(), @@ -1688,13 +1690,10 @@ func TestHandlerPrioritizesClientCancellationOverLocalInferenceErrors(t *testing ) stats := health.NewStatsStore() handler.stats = stats - runtimeRegistry := &recordingRuntimeRegistry{KeyRegistry: registry} + runtimeRegistry := &recordingRuntimeRegistry{CredentialRegistry: registry} handler.registry = runtimeRegistry - healthDialect := &countingHealthDialect{ - Dialect: handler.dialects[protocol.OpenAICompletions], - } extractDialect := &cancelingExtractDialect{ - Dialect: healthDialect, + Dialect: handler.dialects[protocol.OpenAICompletions], err: sentinel, } if test.cancelOnExtract { @@ -1748,19 +1747,13 @@ func TestHandlerPrioritizesClientCancellationOverLocalInferenceErrors(t *testing if len(forwarder.inputs)+len(forwarder.streamInputs) != 0 { t.Fatal("canceled local failure reached upstream") } - if healthDialect.classifyCalls != 0 { - t.Fatalf( - "health classifier calls = %d, want zero", - healthDialect.classifyCalls, - ) - } if runtimeRegistry.cooldownCalls != 0 || runtimeRegistry.incrFailureCalls != 0 || runtimeRegistry.blacklistCalls != 0 || runtimeRegistry.clearCalls != 0 { t.Fatalf("Registry side effects = %#v", runtimeRegistry) } - if got := stats.Snapshot(1, time.Now()); got != (health.KeyStats{}) { + if got := stats.Snapshot(1, time.Now()); got != (health.CredentialStats{}) { t.Fatalf("StatsStore side effects = %#v, want zero", got) } }) @@ -1813,7 +1806,7 @@ func TestHandlerRecordsNonStreamingRetryChain(t *testing.T) { first, second := event.Attempts[0], event.Attempts[1] if first.Sequence != 1 || !first.WillRetry || first.FailureCategory != telemetry.FailureCategoryInvalidKey || - first.Action != telemetry.ActionFailKey || + first.Action != telemetry.ActionFailCredential || first.ErrorCode != "upstream_invalid_key" || first.ErrorSummary != "first invalid key" || first.DurationMs <= 0 { @@ -2026,15 +2019,10 @@ func TestHandlerRecordsCommittedStreamTerminalMatrix(t *testing.T) { sink, "sk-one", ) - runtimeRegistry := &recordingRuntimeRegistry{KeyRegistry: registry} + runtimeRegistry := &recordingRuntimeRegistry{CredentialRegistry: registry} handler.registry = runtimeRegistry now := time.Date(2026, time.July, 24, 13, 0, 0, 0, time.UTC) handler.now = func() time.Time { return now } - selectedDialect := &countingHealthDialect{ - Dialect: handler.dialects[protocol.OpenAICompletions], - } - handler.dialects[protocol.OpenAICompletions] = selectedDialect - request := httptest.NewRequest( http.MethodPost, "/v1/chat/completions", @@ -2079,16 +2067,13 @@ func TestHandlerRecordsCommittedStreamTerminalMatrix(t *testing.T) { downstreamBody, ) } - if selectedDialect.classifyCalls != 0 { - t.Fatalf("health classifier calls = %d, want 0", selectedDialect.classifyCalls) - } if runtimeRegistry.cooldownCalls != 0 || runtimeRegistry.incrFailureCalls != 0 || runtimeRegistry.blacklistCalls != 0 || runtimeRegistry.clearCalls != 1 { t.Fatalf("Registry side effects = %#v", runtimeRegistry) } - if got := handler.stats.Snapshot(1, now); got != (health.KeyStats{Success: 1}) { + if got := handler.stats.Snapshot(1, now); got != (health.CredentialStats{Success: 1}) { t.Fatalf("StatsStore side effects = %#v, want existing ready success only", got) } }) @@ -2158,12 +2143,8 @@ func TestHandlerStreamTelemetryDoesNotChangeHealthOrRetrySideEffects(t *testing. "sk-one", "sk-two", ) - runtimeRegistry := &recordingRuntimeRegistry{KeyRegistry: registry} + runtimeRegistry := &recordingRuntimeRegistry{CredentialRegistry: registry} handler.registry = runtimeRegistry - selectedDialect := &countingHealthDialect{ - Dialect: handler.dialects[protocol.OpenAICompletions], - } - handler.dialects[protocol.OpenAICompletions] = selectedDialect now := time.Date(2026, time.July, 24, 14, 0, 0, 0, time.UTC) handler.now = func() time.Time { return now } @@ -2178,16 +2159,13 @@ func TestHandlerStreamTelemetryDoesNotChangeHealthOrRetrySideEffects(t *testing. if len(forwarder.streamInputs) != 1 { t.Fatalf("stream attempts = %d, want observation-only terminal return", len(forwarder.streamInputs)) } - if selectedDialect.classifyCalls != 0 { - t.Fatalf("health classifier calls = %d, want 0", selectedDialect.classifyCalls) - } if runtimeRegistry.cooldownCalls != 0 || runtimeRegistry.incrFailureCalls != 0 || runtimeRegistry.blacklistCalls != 0 || runtimeRegistry.clearCalls != 0 { t.Fatalf("Registry side effects = %#v", runtimeRegistry) } - if got := handler.stats.Snapshot(1, now); got != (health.KeyStats{}) { + if got := handler.stats.Snapshot(1, now); got != (health.CredentialStats{}) { t.Fatalf("StatsStore side effects = %#v, want zero", got) } events := sink.snapshot() @@ -2266,17 +2244,13 @@ func TestHandlerRecordsCanceledSuccessfulResponseWithoutHealthSideEffects(t *tes handler, _, registry := newHandlerForTestWithStats( t, forwarder, stats, "sk-one", ) - runtimeRegistry := &recordingRuntimeRegistry{KeyRegistry: registry} + runtimeRegistry := &recordingRuntimeRegistry{CredentialRegistry: registry} handler.registry = runtimeRegistry handler.limiter = &recordingAccessKeyRPMLimiter{} sink := &recordingRequestLogSink{} handler.requestLogSink = sink handler.newRequestID = func() (string, error) { return fixedRequestID, nil } handler.requestNow = newSteppingRequestClock() - selectedDialect := &countingHealthDialect{ - Dialect: handler.dialects[protocol.OpenAICompletions], - } - handler.dialects[protocol.OpenAICompletions] = selectedDialect engine := gin.New() bindGatewayRoutesForTest(t, engine, handler) @@ -2300,7 +2274,7 @@ func TestHandlerRecordsCanceledSuccessfulResponseWithoutHealthSideEffects(t *tes t.Fatalf("event = %#v, want uncommitted client cancellation", event) } if event.Usage.Result != (usage.Result{State: usage.StateNotApplicable}) || - event.Usage.GroupID != 1 || event.Usage.KeyID != 1 || + event.Usage.GroupID != 1 || event.Usage.CredentialID != 1 || event.Usage.AttemptSequence != 1 { t.Fatalf("event Usage = %#v, want canceled attempt route attribution", event.Usage) } @@ -2310,19 +2284,13 @@ func TestHandlerRecordsCanceledSuccessfulResponseWithoutHealthSideEffects(t *tes attempt.WillRetry { t.Fatalf("attempt = %#v, want read-only downstream_cancel/terminate", attempt) } - if selectedDialect.classifyCalls != 0 { - t.Fatalf( - "health classifier calls = %d, canceled early branch must not judge again", - selectedDialect.classifyCalls, - ) - } if runtimeRegistry.cooldownCalls != 0 || runtimeRegistry.incrFailureCalls != 0 || runtimeRegistry.blacklistCalls != 0 || runtimeRegistry.clearCalls != 0 { t.Fatalf("Registry side effects = %#v", runtimeRegistry) } - if got := stats.Snapshot(1, time.Now()); got != (health.KeyStats{}) { + if got := stats.Snapshot(1, time.Now()); got != (health.CredentialStats{}) { t.Fatalf("StatsStore side effects = %#v", got) } } @@ -2438,7 +2406,7 @@ func TestHandlerPrioritizesClientCancellationOverDownstreamWriteFailure(t *testi event.StatusCode != test.wantHTTP || event.UpstreamModel != "gpt-4o" || event.Usage.GroupID != 1 || - event.Usage.KeyID != 1 || + event.Usage.CredentialID != 1 || event.Usage.AttemptSequence != 1 { t.Fatalf( "event status/code/http = %q/%q/%d, want %q/%q/%d with route attribution: %#v", @@ -2455,82 +2423,6 @@ func TestHandlerPrioritizesClientCancellationOverDownstreamWriteFailure(t *testi } } -func TestForwarderRedactsResolvedCredentialSecretsBeforeClassification(t *testing.T) { - const ( - apiKey = "opaque-provider-secret" - authorization = "Token " + apiKey - apiKeyHeader = "Secondary " + apiKey - customRule = "opaque literal\tvalue" - ordinaryEncoding = "gzip" - disallowedBody = "resolved-opaque-disallowed" - ) - server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - writer.Header().Set("Content-Type", "application/json") - writer.WriteHeader(http.StatusBadRequest) - message := strings.Join( - []string{apiKey, authorization, apiKeyHeader, customRule, ordinaryEncoding}, - " ", - ) - _, _ = writer.Write([]byte( - `{"error":{"message":` + strconv.Quote(message) + - `},"debug":"` + disallowedBody + `"}`, - )) - })) - defer server.Close() - - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Group: state.GroupView{ - UpstreamURL: server.URL, - HeaderRules: state.HeaderRules{Set: map[string]string{ - "Authorization": "Token ${API_KEY}", - "Api-Key": "Secondary ${API_KEY}", - "X-Custom-Rule": customRule, - "Accept-Encoding": ordinaryEncoding, - }}, - }, - APIKey: apiKey, - Request: &dialect.ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/chat/completions", - Header: make(http.Header), - Body: []byte(`{"model":"gpt-4o"}`), - }, - } - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - result := forwarder.Forward(context.Background(), input) - - for _, secret := range []string{apiKey, authorization, apiKeyHeader} { - if strings.Contains(string(result.ClassificationBody), secret) || - strings.Contains(result.ErrorSummary, secret) { - t.Fatalf("resolved credential %q leaked: %#v", secret, result) - } - } - classificationSummary := allowedErrorSummary(result.ClassificationBody) - for _, ordinary := range []string{customRule, ordinaryEncoding} { - if !strings.Contains(classificationSummary, ordinary) { - t.Fatalf( - "ordinary HeaderRule value %q was changed in ClassificationBody summary %q: %#v", - ordinary, - classificationSummary, - result, - ) - } - if strings.Contains(result.ErrorSummary, ordinary) { - t.Fatalf("ErrorSummary leaked HeaderRules value %q: %#v", ordinary, result) - } - } - if strings.Contains(result.ErrorSummary, "opaque literal value") { - t.Fatalf("ErrorSummary leaked normalized HeaderRules value: %#v", result) - } - if result.ErrorSummary == "" || !strings.Contains(result.ErrorSummary, redact.Placeholder) { - t.Fatalf("ErrorSummary = %q, want non-empty redacted allowed-path summary", result.ErrorSummary) - } - if strings.Contains(result.ErrorSummary, disallowedBody) { - t.Fatalf("ErrorSummary used disallowed JSON path: %q", result.ErrorSummary) - } -} - func TestRequestRecorderRedactsHeaderRuleLiteralsFromNonStreamingAndSSESummaries(t *testing.T) { const ( apiKey = "fake-recorder-provider-key" @@ -2538,15 +2430,15 @@ func TestRequestRecorderRedactsHeaderRuleLiteralsFromNonStreamingAndSSESummaries ordinaryEncoding = "gzip" ) selection := scheduler.Selection{ - GroupID: 11, + GroupID: 11, ChannelID: channel.OpenAI, CredentialID: 21, + RouteMode: channel.RouteNative, Group: state.GroupView{ - ID: 11, + ID: 11, ChannelID: channel.OpenAI, HeaderRules: state.HeaderRules{Set: map[string]string{ "X-Custom": customRule, "Accept-Encoding": ordinaryEncoding, }}, }, - KeyID: 21, } for _, test := range []struct { name string @@ -2563,7 +2455,7 @@ func TestRequestRecorderRedactsHeaderRuleLiteralsFromNonStreamingAndSSESummaries record: func(recorder *requestRecorder, result UpstreamResult) int { return recorder.recordAttempt( selection, - apiKey, + []string{apiKey}, result, health.Result{ Category: health.FailureCategoryClientError, @@ -2595,7 +2487,7 @@ func TestRequestRecorderRedactsHeaderRuleLiteralsFromNonStreamingAndSSESummaries record: func(recorder *requestRecorder, result UpstreamResult) int { return recorder.recordStreamAttempt( selection, - apiKey, + []string{apiKey}, result, time.Unix(100, 0), time.Unix(101, 0), @@ -2649,14 +2541,14 @@ func TestRequestRecorderRedactsHeaderRuleLiteralsFromNonStreamingAndSSESummaries func TestRequestRecorderPreservesFixedTerminalSummariesWithShortHeaderRuleLiteral(t *testing.T) { const shortLiteral = "a" selection := scheduler.Selection{ - GroupID: 11, + GroupID: 11, ChannelID: channel.OpenAI, CredentialID: 21, + RouteMode: channel.RouteNative, Group: state.GroupView{ - ID: 11, + ID: 11, ChannelID: channel.OpenAI, HeaderRules: state.HeaderRules{Set: map[string]string{ "X-Short": shortLiteral, }}, }, - KeyID: 21, } sink := &recordingRequestLogSink{} recorder := newRequestRecorder( @@ -2676,7 +2568,7 @@ func TestRequestRecorderPreservesFixedTerminalSummariesWithShortHeaderRuleLitera } attempt := recorder.recordStreamAttempt( selection, - "fake-fixed-summary-provider-key", + []string{"fake-fixed-summary-provider-key"}, result, time.Unix(100, 0), time.Unix(101, 0), @@ -2702,15 +2594,15 @@ func TestRequestRecorderPreservesFixedTerminalSummariesWithShortHeaderRuleLitera func TestRequestRecorderPreservesProviderErrorFixedSummaryWithOrdinaryHeaderRules(t *testing.T) { const marker = "rate_limit_error" selection := scheduler.Selection{ - GroupID: 11, + GroupID: 11, ChannelID: channel.OpenAI, CredentialID: 21, + RouteMode: channel.RouteNative, Group: state.GroupView{ - ID: 11, + ID: 11, ChannelID: channel.OpenAI, HeaderRules: state.HeaderRules{Set: map[string]string{ "X-Ordinary-Marker": marker, "X-Short-Literal": "a", }}, }, - KeyID: 21, } sink := &recordingRequestLogSink{} recorder := newRequestRecorder( @@ -2731,12 +2623,12 @@ func TestRequestRecorderPreservesProviderErrorFixedSummaryWithOrdinaryHeaderRule } decision := health.Result{ Category: health.FailureCategoryRateLimited, - Action: health.ActionCooldownKey, + Action: health.ActionCooldownCredential, } recorder.recordAttempt( selection, - "fake-provider-fixed-summary-key", + []string{"fake-provider-fixed-summary-key"}, result, decision, time.Unix(100, 0), diff --git a/internal/gateway/response_model_observation_test.go b/internal/gateway/response_model_observation_test.go deleted file mode 100644 index e52ff412..00000000 --- a/internal/gateway/response_model_observation_test.go +++ /dev/null @@ -1,178 +0,0 @@ -package gateway - -import ( - "bytes" - "context" - "errors" - "io" - "net/http" - "net/http/httptest" - "strings" - "testing" - "unicode/utf8" - - platformhttp "gpt-load/internal/platform/httpclient" - "gpt-load/internal/platform/redact" -) - -func TestForwardObservesUpstreamModelBeforeAliasRewrite(t *testing.T) { - tests := []struct { - name string - responseModel string - wantMismatch bool - }{ - {name: "match", responseModel: "provider-model"}, - {name: "mismatch", responseModel: "different-provider-model", wantMismatch: true}, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - _, _ = io.WriteString(writer, `{"id":"chatcmpl-1","model":"`+test.responseModel+`"}`) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public-model"}`) - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - - if result.Err != nil || result.StatusCode != http.StatusOK { - t.Fatalf("Forward() result = %#v", result) - } - if !result.ResponseModelObserved || result.UpstreamReportedModel != test.responseModel || - result.ResponseModelMismatch != test.wantMismatch { - t.Fatalf("model observation = %#v", result) - } - if bytes.Contains(result.Body, []byte(test.responseModel)) || - !bytes.Contains(result.Body, []byte("public-model")) { - t.Fatalf("downstream alias rewrite changed = %q", result.Body) - } - }) - } -} - -func TestForwardMarksSuccessfulResponseWithoutModelAsUnobserved(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - _, _ = io.WriteString(writer, `{"id":"chatcmpl-1"}`) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"provider-model"}`) - input.ExternalModel = "provider-model" - input.UpstreamModelID = "provider-model" - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - - if result.Err != nil || result.StatusCode != http.StatusOK || - result.ResponseModelObserved || result.UpstreamReportedModel != "" || - result.ResponseModelMismatch { - t.Fatalf("Forward() result = %#v", result) - } -} - -func TestForwardBoundsObservedModelAfterRawComparison(t *testing.T) { - responseModel := strings.Repeat("模型", 128) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "application/json") - _, _ = io.WriteString(writer, `{"id":"chatcmpl-1","model":"`+responseModel+`"}`) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public-model"}`) - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), input, - ) - - if result.Err != nil || result.StatusCode != http.StatusOK || - !result.ResponseModelObserved || !result.ResponseModelMismatch { - t.Fatalf("Forward() result = %#v", result) - } - if len(result.UpstreamReportedModel) > 255 || !utf8.ValidString(result.UpstreamReportedModel) { - t.Fatalf("bounded response model length/validity = %d/%t", len(result.UpstreamReportedModel), utf8.ValidString(result.UpstreamReportedModel)) - } - if result.UpstreamReportedModel == responseModel { - t.Fatal("Forward() retained the unbounded response model") - } -} - -func TestResponseModelTrackerPreservesMismatchWhenProjectionCollidesWithExpectedModel(t *testing.T) { - const truncatedMarker = "...[truncated]" - rawModel := strings.Repeat("x", 512) - expected := strings.Repeat("x", 255-len(truncatedMarker)) + truncatedMarker - tracker := newResponseModelTracker(streamForwardInput("http://example.com").Dialect, expected) - - tracker.observe([]byte(`{"model":"` + rawModel + `"}`)) - observation := tracker.observation() - - if !observation.observed || !observation.mismatch { - t.Fatalf("model observation = %#v", observation) - } - if len(observation.reported) > 255 || observation.reported == expected { - t.Fatalf("bounded collision model = %q (%d bytes), expected a distinct bounded value", observation.reported, len(observation.reported)) - } -} - -func TestForwardStreamKeepsMismatchingDeclarationObservedBeforeLaterMatches(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString( - writer, - "data: {\"id\":\"chunk-1\",\"model\":\"different-provider-model\"}\n\n"+ - "data: {\"id\":\"chunk-1\",\"model\":\"provider-model\"}\n\n", - ) - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"public-model","stream":true}`) - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - downstream := newRecordingResponseWriter() - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Err != nil || !result.Committed || !result.ResponseModelObserved || - !result.ResponseModelMismatch || result.UpstreamReportedModel != "different-provider-model" { - t.Fatalf("ForwardStream() result = %#v", result) - } - if bytes.Contains(downstream.body.Bytes(), []byte("provider-model")) || - !bytes.Contains(downstream.body.Bytes(), []byte("public-model")) { - t.Fatalf("downstream stream = %q", downstream.body.String()) - } -} - -func TestForwardStreamDiscardsModelObservationWhenRequestDoesNotCompleteNormally(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Type", "text/event-stream") - _, _ = io.WriteString(writer, "data: {\"model\":\"different-provider-model\"}\n\n") - })) - defer upstream.Close() - - input := streamForwardInput(upstream.URL) - input.Request.Body = []byte(`{"model":"provider-model","stream":true}`) - input.ExternalModel = "provider-model" - input.UpstreamModelID = "provider-model" - downstream := newRecordingResponseWriter() - downstream.writeErr = errors.New("downstream write failed") - result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( - context.Background(), input, downstream, - ) - - if result.Stream.EndReason != StreamEndDownstreamWriteFailure || - result.ResponseModelObserved || result.ResponseModelMismatch || - result.UpstreamReportedModel != "" { - t.Fatalf("ForwardStream() result = %#v", result) - } -} diff --git a/internal/gateway/response_representation.go b/internal/gateway/response_representation.go index 2b0272c3..f5a3ec40 100644 --- a/internal/gateway/response_representation.go +++ b/internal/gateway/response_representation.go @@ -31,7 +31,11 @@ type preparedErrorRepresentation struct { changed bool } -func (forwarder *Forwarder) prepareSuccessRepresentation( +type responseProcessor struct { + redactor *redact.Redactor +} + +func (forwarder *responseProcessor) prepareSuccessRepresentation( input ForwardInput, status int, headers http.Header, @@ -94,7 +98,11 @@ func (forwarder *Forwarder) prepareSuccessRepresentation( } } - changed := encoding != contentcoding.Identity || + // The execution adapter may already have projected the selected upstream + // model back to the client alias. Conservatively treat every aliased route + // as a representation rewrite so upstream validators/digests are never + // forwarded for bytes whose identity changed before this boundary. + changed := encoding != contentcoding.Identity || needsModelRewrite(input) || !bytes.Equal(originalPlain, downstreamPlain) if status == http.StatusPartialContent && changed { return preparedSuccessRepresentation{}, successRepresentationProtocolError("changed partial response") @@ -141,7 +149,7 @@ func (forwarder *Forwarder) prepareSuccessRepresentation( }, nil } -func (forwarder *Forwarder) prepareErrorRepresentation( +func (forwarder *responseProcessor) prepareErrorRepresentation( input ForwardInput, headers http.Header, wire []byte, @@ -232,7 +240,7 @@ func (forwarder *Forwarder) prepareErrorRepresentation( } } -func (forwarder *Forwarder) failClosedErrorRepresentation( +func (forwarder *responseProcessor) failClosedErrorRepresentation( input ForwardInput, headers http.Header, secrets []string, diff --git a/internal/gateway/response_representation_test.go b/internal/gateway/response_representation_test.go deleted file mode 100644 index dfd7d569..00000000 --- a/internal/gateway/response_representation_test.go +++ /dev/null @@ -1,1163 +0,0 @@ -package gateway - -import ( - "bytes" - "compress/gzip" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "net/http/httptest" - "strconv" - "testing" - "time" - - "gpt-load/internal/dialect" - "gpt-load/internal/platform/contentcoding" - platformhttp "gpt-load/internal/platform/httpclient" - "gpt-load/internal/platform/redact" - "gpt-load/internal/protocol" - "gpt-load/internal/state" - "gpt-load/internal/usage" -) - -type successRepresentationEncoding struct { - name string - encoding string - wire []byte -} - -func assertRepresentationValidatorsPreservedWithoutSignatures(t *testing.T, headers http.Header) { - t.Helper() - for name, want := range map[string]string{ - "ETag": `"wire-v1"`, - "Digest": "sha-256=wire-digest", - "Content-MD5": "d2lyZQ==", - "Content-Range": "bytes 0-9/10", - "Content-Digest": "sha-256=:d2lyZQ==:", - "Repr-Digest": "sha-256=:cmVwcg==:", - } { - if got := headers.Get(name); got != want { - t.Errorf("%s = %q, want preserved value %q", name, got, want) - } - } - if headers.Get("Signature") != "" || headers.Get("Signature-Input") != "" { - t.Errorf("stale signatures survived representation normalization: %#v", headers) - } -} - -func TestPrepareSuccessRepresentationReturnsPlainIdentity(t *testing.T) { - plain := []byte(`{"id":"safe-response","usage":{"prompt_tokens":100,"completion_tokens":30}}`) - encodings := []contentcoding.Encoding{ - contentcoding.Identity, - contentcoding.Gzip, - contentcoding.Brotli, - contentcoding.Deflate, - contentcoding.Zstd, - } - for _, encoding := range encodings { - t.Run(string(encoding), func(t *testing.T) { - result := forwardSuccessRepresentation( - t, - successRepresentationEncoding{ - name: string(encoding), - encoding: string(encoding), - wire: encodeContentCodingForGatewayTest(t, encoding, plain), - }, - "fake-provider-secret-plaintext-boundary", - nil, - ) - - if result.Err != nil || result.StatusCode != http.StatusOK { - t.Fatalf("Forward() result = %#v, want successful response", result) - } - if !bytes.Equal(result.Body, plain) { - t.Fatalf("downstream body = %x, want plaintext %x", result.Body, plain) - } - if values := headerFieldValues(result.Header, "Content-Encoding"); len(values) != 0 { - t.Fatalf("Content-Encoding values = %#v, want absent", values) - } - if got := result.Header.Get("Content-Length"); got != strconv.Itoa(len(plain)) { - t.Fatalf("Content-Length = %q, want %d", got, len(plain)) - } - }) - } -} - -func TestPrepareSuccessRepresentationMetadataPolicy(t *testing.T) { - plain := []byte(`{"id":"safe-response"}`) - - identity := forwardSuccessRepresentation( - t, - successRepresentationEncoding{ - name: "identity", - encoding: "identity", - wire: bytes.Clone(plain), - }, - "fake-provider-secret-metadata-identity", - nil, - ) - if identity.Err != nil || !bytes.Equal(identity.Body, plain) { - t.Fatalf("identity Forward() result = %#v", identity) - } - assertRepresentationValidatorsPreservedWithoutSignatures(t, identity.Header) - - gzipResult := forwardSuccessRepresentation( - t, - successRepresentationEncoding{ - name: "gzip", - encoding: "gzip", - wire: encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain), - }, - "fake-provider-secret-metadata-gzip", - nil, - ) - if gzipResult.Err != nil || !bytes.Equal(gzipResult.Body, plain) { - t.Fatalf("gzip Forward() result = %#v", gzipResult) - } - assertRepresentationMetadata(t, gzipResult.Header, false) - if values := headerFieldValues(gzipResult.Header, "Content-Encoding"); len(values) != 0 { - t.Fatalf("Content-Encoding values = %#v, want absent", values) - } -} - -func TestPrepareSuccessRepresentationRebuildsContentLengthAcrossHeaderCasings(t *testing.T) { - plain := []byte(`{"id":"safe-response"}`) - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - headers := http.Header{ - "Content-Encoding": {"identity"}, - "Content-Length": {"999"}, - "content-length": {"998"}, - } - - result, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).prepareSuccessRepresentation(input, http.StatusOK, headers, plain, nil) - if err != nil { - t.Fatalf("prepareSuccessRepresentation() error = %v", err) - } - if values := headerFieldValues(result.headers, "Content-Length"); len(values) != 1 || values[0] != strconv.Itoa(len(plain)) { - t.Fatalf("Content-Length values = %#v, want exactly [%d]", values, len(plain)) - } -} - -func TestPrepareSuccessRepresentationDropsSignaturesAfterRepresentationHeaderNormalization(t *testing.T) { - plain := []byte(`{"id":"safe-response"}`) - headers := http.Header{ - "Content-Encoding": {"identity"}, - "Content-Length": {strconv.Itoa(len(plain))}, - "Signature": {"sig1=:c2lnOg==:"}, - "Signature-Input": {`sig1=("content-encoding" "content-length");created=1`}, - } - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - - result, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).prepareSuccessRepresentation(input, http.StatusOK, headers, plain, nil) - if err != nil || result.changed || result.headers.Get("Signature") != "" || result.headers.Get("Signature-Input") != "" { - t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want normalized representation without signatures", result, err) - } -} - -func TestPrepareSuccessRepresentationRejectsPartialContentWithoutSafeContentRange(t *testing.T) { - plain := []byte(`{"id":"partial"}`) - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - ExternalModel: "public-model", - UpstreamModelID: "provider-model", - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - headers := http.Header{ - "Content-Encoding": {"identity"}, - "Content-Length": {strconv.Itoa(len(plain))}, - "Content-Range": {"bytes provider-model/100"}, - } - - result, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).prepareSuccessRepresentation(input, http.StatusPartialContent, headers, plain, nil) - if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { - t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want unsafe partial response rejection", result, err) - } -} - -func TestPrepareSuccessRepresentationValidatesPartialContentRangeAgainstPlaintext(t *testing.T) { - plain := []byte("abc") - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - tests := []struct { - name string - headers http.Header - wantErr bool - }{ - { - name: "matching decimal range", - headers: http.Header{ - "Content-Length": {"3"}, - "Content-Range": {"bytes 10-12/100"}, - }, - }, - { - name: "range length differs from plaintext", - headers: http.Header{ - "Content-Length": {"3"}, - "Content-Range": {"bytes 10-13/100"}, - }, - wantErr: true, - }, - { - name: "plus-prefixed range bound", - headers: http.Header{ - "Content-Length": {"3"}, - "Content-Range": {"bytes +10-12/100"}, - }, - wantErr: true, - }, - { - name: "plus-prefixed complete length", - headers: http.Header{ - "Content-Length": {"3"}, - "Content-Range": {"bytes 10-12/+100"}, - }, - wantErr: true, - }, - { - name: "case-colliding duplicate range", - headers: http.Header{ - "Content-Length": {"3"}, - "Content-Range": {"bytes 10-12/100"}, - "content-range": {"bytes 10-12/100"}, - }, - wantErr: true, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - result, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).prepareSuccessRepresentation(input, http.StatusPartialContent, test.headers, plain, nil) - if test.wantErr { - if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { - t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want protocol error", result, err) - } - return - } - if err != nil || !bytes.Equal(result.downstream, plain) || - result.headers.Get("Content-Range") != "bytes 10-12/100" { - t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want valid partial response", result, err) - } - }) - } -} - -func TestHasSafeHeadContentRange(t *testing.T) { - tests := []struct { - name string - headers http.Header - want bool - }{ - { - name: "matching content length", - headers: http.Header{ - "Content-Length": {"3"}, - "Content-Range": {"bytes 10-12/100"}, - }, - want: true, - }, - { - name: "wildcard total", - headers: http.Header{ - "Content-Length": {"3"}, - "Content-Range": {"bytes 10-12/*"}, - }, - want: true, - }, - { - name: "content length absent", - headers: http.Header{ - "Content-Range": {"bytes 10-12/100"}, - }, - want: true, - }, - { - name: "content length mismatch", - headers: http.Header{ - "Content-Length": {"4"}, - "Content-Range": {"bytes 10-12/100"}, - }, - }, - { - name: "invalid bounds", - headers: http.Header{ - "Content-Length": {"3"}, - "Content-Range": {"bytes 12-10/100"}, - }, - }, - { - name: "case colliding content range", - headers: http.Header{ - "Content-Length": {"3"}, - "Content-Range": {"bytes 10-12/100"}, - "content-range": {"bytes 10-12/100"}, - }, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - if got := hasSafeHeadContentRange(test.headers); got != test.want { - t.Fatalf("hasSafeHeadContentRange() = %t, want %t", got, test.want) - } - }) - } -} - -func TestPrepareSuccessRepresentationRejectsChangedPartialContent(t *testing.T) { - const ( - upstreamModel = "provider-model" - externalModel = "public-model" - ) - plain := []byte(`{"model":"provider-model","id":"partial"}`) - forwardPartial := func( - t *testing.T, - encoding contentcoding.Encoding, - configure func(*ForwardInput), - ) UpstreamResult { - t.Helper() - wire := encodeContentCodingForGatewayTest(t, encoding, plain) - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - writer.Header().Set("Content-Encoding", string(encoding)) - writer.Header().Set("Content-Range", fmt.Sprintf("bytes 0-%d/100", len(plain)-1)) - writer.WriteHeader(http.StatusPartialContent) - _, _ = writer.Write(wire) - })) - defer upstream.Close() - - input := usageForwardInput(upstream.URL, dialect.NewOpenAI(http.DefaultClient)) - input.Request.Body = []byte(`{"model":"public-model"}`) - if configure != nil { - configure(&input) - } - return NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), - input, - ) - } - - t.Run("identity alias rewrite", func(t *testing.T) { - result := forwardPartial(t, contentcoding.Identity, func(input *ForwardInput) { - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - }) - if !errors.Is(result.Err, ErrUpstreamProtocol) || result.StatusCode != 0 || len(result.Body) != 0 { - t.Fatalf("Forward() result = %#v, want changed 206 protocol error", result) - } - }) - - t.Run("compressed representation", func(t *testing.T) { - result := forwardPartial(t, contentcoding.Gzip, nil) - if !errors.Is(result.Err, ErrUpstreamProtocol) || result.StatusCode != 0 || len(result.Body) != 0 { - t.Fatalf("Forward() result = %#v, want decoded 206 protocol error", result) - } - }) - - t.Run("identity unchanged", func(t *testing.T) { - result := forwardPartial(t, contentcoding.Identity, nil) - if result.Err != nil || result.StatusCode != http.StatusPartialContent || - !bytes.Equal(result.Body, plain) || result.Header.Get("Content-Range") != fmt.Sprintf("bytes 0-%d/100", len(plain)-1) || - len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { - t.Fatalf("Forward() result = %#v, want unchanged legal 206", result) - } - }) -} - -func TestPrepareSuccessRepresentationBoundsWireAndDecodedBodies(t *testing.T) { - forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - - caseCollidingHeaders := http.Header{ - "Content-Encoding": {"identity"}, - "content-encoding": {"gzip"}, - } - result, err := forwarder.prepareSuccessRepresentation( - input, - http.StatusOK, - caseCollidingHeaders, - []byte(`{"id":"safe-response"}`), - nil, - ) - if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { - t.Fatalf("case-colliding Content-Encoding result/error = %#v/%v, want protocol error", result, err) - } - - overflow := gzipRepeatedByte(t, 'x', maxNonStreamingResponseBodyBytes+1) - result, err = forwarder.prepareSuccessRepresentation( - input, - http.StatusOK, - http.Header{"Content-Encoding": {"gzip"}}, - overflow, - nil, - ) - if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { - t.Fatalf("decoded overflow result/error = %#v/%v, want protocol error", result, err) - } - - result, err = forwarder.prepareSuccessRepresentation( - input, - http.StatusOK, - http.Header{"Content-Encoding": {"identity"}}, - bytes.Repeat([]byte("w"), int(maxNonStreamingResponseBodyBytes)+1), - nil, - ) - if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { - t.Fatalf("wire overflow result/error = %#v/%v, want protocol error", result, err) - } -} - -func gzipRepeatedByte(t *testing.T, value byte, count int64) []byte { - t.Helper() - var wire bytes.Buffer - compressor := gzip.NewWriter(&wire) - if _, err := io.CopyN(compressor, repeatingByteReader(value), count); err != nil { - t.Fatalf("stream gzip fixture: %v", err) - } - if err := compressor.Close(); err != nil { - t.Fatalf("close gzip fixture: %v", err) - } - return wire.Bytes() -} - -func successRepresentationEncodings(t *testing.T, plain []byte) []successRepresentationEncoding { - t.Helper() - encodings := []string{"", "identity", "gzip", "br", "deflate", "zstd"} - result := make([]successRepresentationEncoding, 0, len(encodings)) - for _, encoding := range encodings { - fixtureEncoding := contentcoding.Identity - if encoding != "" { - fixtureEncoding = contentcoding.Encoding(encoding) - } - wire := encodeContentCodingForGatewayTest(t, fixtureEncoding, plain) - name := encoding - if name == "" { - name = "empty" - } - result = append(result, successRepresentationEncoding{ - name: name, - encoding: encoding, - wire: wire, - }) - } - return result -} - -func forwardSuccessRepresentation( - t *testing.T, - representation successRepresentationEncoding, - apiKey string, - configureInput func(*ForwardInput), -) UpstreamResult { - t.Helper() - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - if representation.encoding != "" { - writer.Header().Set("Content-Encoding", representation.encoding) - } - writer.Header().Set("Content-Length", strconv.Itoa(len(representation.wire))) - setRepresentationMetadata(writer.Header()) - _, _ = writer.Write(representation.wire) - })) - defer upstream.Close() - - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - Group: state.GroupView{ - ID: 1, - UpstreamURL: testUpstreamBaseURL(upstream.URL, protocol.OpenAICompletions), - Timeouts: state.TimeoutConfig{ - Connect: time.Second, - FirstByte: time.Second, - Request: 5 * time.Second, - }, - }, - APIKey: apiKey, - Request: &dialect.ParsedRequest{ - Method: http.MethodPost, - Path: "/v1/chat/completions", - Header: make(http.Header), - Body: []byte(`{"model":"public-model"}`), - }, - } - if configureInput != nil { - configureInput(&input) - } - return NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( - context.Background(), - input, - ) -} - -func TestForwardSuccessRepresentationReturnsPlaintextAndConditionalMetadata(t *testing.T) { - plain := []byte(`{"id":"safe-response","usage":{"prompt_tokens":100,"completion_tokens":30}}`) - for _, representation := range successRepresentationEncodings(t, plain) { - t.Run(representation.name, func(t *testing.T) { - result := forwardSuccessRepresentation( - t, - representation, - "fake-provider-secret-unchanged", - nil, - ) - - if result.Err != nil || result.StatusCode != http.StatusOK || !result.RequestWritten { - t.Fatalf("Forward() result = %#v, want successful written response", result) - } - if !bytes.Equal(result.Body, plain) { - t.Fatalf("downstream body = %x, want plaintext %x", result.Body, plain) - } - if values := headerFieldValues(result.Header, "Content-Encoding"); len(values) != 0 { - t.Fatalf("Content-Encoding values = %#v, want absent", values) - } - if result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { - t.Fatalf("Content-Length = %q, want %d", result.Header.Get("Content-Length"), len(plain)) - } - switch representation.encoding { - case "": - assertRepresentationMetadata(t, result.Header, true) - case "identity": - assertRepresentationValidatorsPreservedWithoutSignatures(t, result.Header) - default: - assertRepresentationMetadata(t, result.Header, false) - } - }) - } -} - -func TestPrepareSuccessRepresentationSanitizesUnchangedCredentialMetadataCollision(t *testing.T) { - const credential = "fake-metadata-credential" - wire := []byte(`{"id":"safe-response"}`) - contentLength := strconv.Itoa(len(wire)) - tests := []struct { - name string - value string - secret string - }{ - {name: "Content-Encoding", value: "identity", secret: "identity"}, - {name: "Content-Length", value: contentLength, secret: contentLength}, - {name: "ETag", value: `"fake-metadata-credential"`, secret: credential}, - {name: "Digest", value: "sha-256=fake-metadata-credential", secret: credential}, - {name: "Content-MD5", value: "fake-metadata-credential", secret: credential}, - {name: "Content-Range", value: "bytes fake-metadata-credential/100", secret: credential}, - {name: "Content-Digest", value: "sha-256=:fake-metadata-credential:", secret: credential}, - {name: "Repr-Digest", value: "sha-256=:fake-metadata-credential:", secret: credential}, - {name: "Signature", value: "sig1=:fake-metadata-credential:", secret: credential}, - {name: "Signature-Input", value: `sig1=("fake-metadata-credential")`, secret: credential}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - headers := http.Header{ - "Content-Encoding": {"identity"}, - "Content-Length": {contentLength}, - } - headers.Set(test.name, test.value) - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - APIKey: test.secret, - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - - result, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).prepareSuccessRepresentation(input, http.StatusOK, headers, wire, []string{test.secret}) - - if test.name == "Content-Length" { - if !errors.Is(err, ErrUpstreamProtocol) || - result.downstream != nil || - result.headers != nil { - t.Fatalf( - "prepareSuccessRepresentation() = %#v, %v; want fail-closed regenerated length collision", - result, - err, - ) - } - return - } - if err != nil || !bytes.Equal(result.downstream, wire) || - len(headerFieldValues(result.headers, "Content-Encoding")) != 0 || - result.headers.Get(test.name) != "" { - t.Fatalf( - "prepareSuccessRepresentation() = %#v, %v; want sanitized plaintext response", - result, - err, - ) - } - }) - } - - headers := http.Header{ - "Content-Encoding": {"identity"}, - "Content-Length": {contentLength}, - } - setRepresentationMetadata(headers) - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - APIKey: credential, - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - result, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).prepareSuccessRepresentation(input, http.StatusOK, headers, wire, []string{credential}) - if err != nil || - result.changed || - !bytes.Equal(result.downstream, wire) || - len(headerFieldValues(result.headers, "Content-Encoding")) != 0 || - result.headers.Get("Content-Length") != contentLength { - t.Fatalf( - "non-collision representation = %#v, %v; want unchanged wire and framing", - result, - err, - ) - } - assertRepresentationValidatorsPreservedWithoutSignatures(t, result.headers) -} - -func TestPrepareSuccessRepresentationRebuildsChangedBodyDespiteStaleCredentialMetadata(t *testing.T) { - const credential = "fake-stale-metadata-credential" - plain := []byte(`{"echo":"fake-stale-metadata-credential"}`) - wire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) - headers := http.Header{ - "Content-Encoding": {"gzip"}, - "Content-Length": {strconv.Itoa(len(wire))}, - } - setRepresentationMetadata(headers) - headers.Set("ETag", `"fake-stale-metadata-credential"`) - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - APIKey: credential, - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - - result, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).prepareSuccessRepresentation(input, http.StatusOK, headers, wire, []string{credential}) - - if err != nil || !result.changed { - t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want rebuilt response", result, err) - } - gotPlain := result.downstream - if string(gotPlain) != `{"echo":"[REDACTED]"}` { - t.Fatalf("rebuilt plain = %q", gotPlain) - } - if len(headerFieldValues(result.headers, "Content-Encoding")) != 0 || - result.headers.Get("Content-Length") != strconv.Itoa(len(result.downstream)) { - t.Fatalf("rebuilt framing headers = %#v", result.headers) - } - assertRepresentationMetadata(t, result.headers, false) -} - -func TestPrepareSuccessRepresentationRemovesAliasedModelFromUnchangedMetadata(t *testing.T) { - const ( - upstreamModel = "provider-private-model" - externalModel = "public-model" - ) - wire := []byte(`{"id":"safe-response"}`) - headers := http.Header{ - "Content-Encoding": {"identity"}, - "Content-Length": {strconv.Itoa(len(wire))}, - "X-Safe": {"kept"}, - } - for _, name := range []string{ - "ETag", "Digest", "Content-MD5", "Content-Range", "Content-Digest", - "Repr-Digest", "Signature", "Signature-Input", - } { - headers.Set(name, "contains="+upstreamModel) - } - input := ForwardInput{ - Dialect: dialect.NewOpenAI(http.DefaultClient), - ExternalModel: externalModel, - UpstreamModelID: upstreamModel, - Request: &dialect.ParsedRequest{Header: make(http.Header)}, - } - - result, err := NewForwarder( - platformhttp.NewHTTPClientManager(), - redact.New(), - ).prepareSuccessRepresentation(input, http.StatusOK, headers, wire, nil) - - if err != nil || result.changed || !bytes.Equal(result.downstream, wire) { - t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want unchanged wire", result, err) - } - if len(headerFieldValues(result.headers, "Content-Encoding")) != 0 || - result.headers.Get("Content-Length") != strconv.Itoa(len(wire)) || - result.headers.Get("X-Safe") != "kept" { - t.Fatalf("unchanged safe headers = %#v", result.headers) - } - for _, name := range []string{ - "ETag", "Digest", "Content-MD5", "Content-Range", "Content-Digest", - "Repr-Digest", "Signature", "Signature-Input", - } { - if value := result.headers.Get(name); value != "" { - t.Fatalf("%s leaked aliased upstream model: %#v", name, result.headers) - } - } -} - -func TestForwardSuccessRepresentationRedactsAndRebuildsMetadata(t *testing.T) { - const ( - apiKey = "fake-provider-secret-redaction" - upstreamModel = "provider-model" - externalModel = "public-model" - ) - tests := []struct { - name string - plain []byte - configureInput func(*ForwardInput) - wantPlain []byte - }{ - { - name: "credential redaction", - plain: []byte(`{"echo":"Bearer fake-provider-secret-redaction fake-provider-secret-redaction"}`), - wantPlain: []byte(`{"echo":"[REDACTED] [REDACTED]"}`), - }, - { - name: "model alias rewrite", - plain: []byte(`{"model":"provider-model","id":"safe-response"}`), - configureInput: func(input *ForwardInput) { - input.ExternalModel = externalModel - input.UpstreamModelID = upstreamModel - }, - wantPlain: []byte(`{"id":"safe-response","model":"public-model"}`), - }, - } - - for _, test := range tests { - for _, representation := range successRepresentationEncodings(t, test.plain) { - t.Run(test.name+"/"+representation.name, func(t *testing.T) { - result := forwardSuccessRepresentation( - t, - representation, - apiKey, - test.configureInput, - ) - - if result.Err != nil || result.StatusCode != http.StatusOK || !result.RequestWritten { - t.Fatalf("Forward() result = %#v, want successful rewritten response", result) - } - plain := result.Body - if !bytes.Equal(plain, test.wantPlain) { - t.Fatalf("downstream plain = %q, want %q", plain, test.wantPlain) - } - if bytes.Contains(plain, []byte(apiKey)) { - t.Fatalf("downstream plain leaked credential: %q", plain) - } - if values := headerFieldValues(result.Header, "Content-Encoding"); len(values) != 0 { - t.Fatalf("Content-Encoding values = %#v, want absent", values) - } - if result.Header.Get("Content-Length") != strconv.Itoa(len(result.Body)) { - t.Fatalf("Content-Length = %q, want %d", result.Header.Get("Content-Length"), len(result.Body)) - } - assertRepresentationMetadata(t, result.Header, false) - }) - } - } -} - -func TestForwardSuccessRepresentationRejectsUnknownStackedMalformedAndPlainOverflow(t *testing.T) { - compressedOverflow := gzipRepeatedByte(t, 'a', maxNonStreamingResponseBodyBytes+1) - tests := []struct { - name string - apiKey string - encodings []string - wire []byte - contentLen string - writeBody func(io.Writer) - }{ - {name: "unknown", apiKey: "fake-provider-secret-unknown", encodings: []string{"compress"}, wire: []byte("opaque-wire")}, - {name: "stacked", apiKey: "fake-provider-secret-stacked", encodings: []string{"gzip, br"}, wire: []byte("opaque-wire")}, - {name: "multiple fields", apiKey: "fake-provider-secret-multiple", encodings: []string{"identity", "gzip"}, wire: []byte("opaque-wire")}, - {name: "malformed gzip", apiKey: "fake-provider-secret-gzip", encodings: []string{"gzip"}, wire: []byte("not-a-gzip-stream")}, - {name: "malformed br", apiKey: "fake-provider-secret-br", encodings: []string{"br"}, wire: []byte{0xff, 0xff, 0xff}}, - {name: "malformed deflate", apiKey: "fake-provider-secret-deflate", encodings: []string{"deflate"}, wire: []byte("not-a-deflate-stream")}, - {name: "malformed zstd", apiKey: "fake-provider-secret-zstd", encodings: []string{"zstd"}, wire: []byte("not-a-zstd-stream")}, - {name: "empty gzip", apiKey: "fake-provider-secret-empty-gzip", encodings: []string{"gzip"}}, - {name: "empty br", apiKey: "fake-provider-secret-empty-br", encodings: []string{"br"}}, - {name: "empty deflate", apiKey: "fake-provider-secret-empty-deflate", encodings: []string{"deflate"}}, - {name: "empty zstd", apiKey: "fake-provider-secret-empty-zstd", encodings: []string{"zstd"}}, - { - name: "wire overflow", - apiKey: "fake-provider-secret-wire-overflow", - encodings: nil, - writeBody: func(writer io.Writer) { - _, _ = io.CopyN(writer, repeatingByteReader('w'), maxNonStreamingResponseBodyBytes+1) - }, - }, - { - name: "plain overflow", - apiKey: "fake-provider-secret-plain-overflow", - encodings: []string{"gzip"}, - wire: compressedOverflow, - }, - { - name: "Content-Length credential collision", - apiKey: "42", - wire: bytes.Repeat([]byte("x"), 42), - contentLen: "42", - }, - { - name: "rebuilt Content-Length credential collision", - apiKey: "21", - wire: []byte(`{"echo":"21"}`), - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { - for _, encoding := range test.encodings { - writer.Header().Add("Content-Encoding", encoding) - } - if test.contentLen != "" { - writer.Header().Set("Content-Length", test.contentLen) - } - setRepresentationMetadata(writer.Header()) - if test.writeBody != nil { - test.writeBody(writer) - } else { - _, _ = writer.Write(test.wire) - } - })) - defer upstream.Close() - - result := testForward(t, upstream.URL, test.apiKey, 10*time.Second) - if !errors.Is(result.Err, ErrUpstreamProtocol) || - !result.RequestWritten || - result.StatusCode != 0 || - len(result.Body) != 0 || - len(result.ClassificationBody) != 0 { - t.Fatalf("Forward() result = %#v, want fail-closed post-write protocol error", result) - } - }) - } -} - -func TestForwardSuccessRepresentationDropsContentEncodingCredentialCollision(t *testing.T) { - plain := []byte(`{"id":"safe-response"}`) - result := forwardSuccessRepresentation( - t, - successRepresentationEncoding{ - name: "gzip", - encoding: "gzip", - wire: encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain), - }, - "gzip", - nil, - ) - if result.Err != nil || !bytes.Equal(result.Body, plain) || - len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { - t.Fatalf("Forward() result = %#v, want plaintext with coding header removed", result) - } -} - -func TestForwardSuccessRepresentationFailsClosedOnResidualCredentialTokens(t *testing.T) { - tests := []struct { - name string - apiKey string - body string - }{ - {name: "numeric", apiKey: "42", body: `{"echo":42}`}, - {name: "boolean", apiKey: "true", body: `{"echo":true}`}, - {name: "null", apiKey: "null", body: `{"echo":null}`}, - {name: "placeholder collision", apiKey: redact.Placeholder, body: `{"echo":"[REDACTED]"}`}, - {name: "rewritten key collision", apiKey: "fake-key", body: `{"fake-key":1,"[REDACTED]":2}`}, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - result := forwardSuccessRepresentation( - t, - successRepresentationEncoding{name: "identity", encoding: "identity", wire: []byte(test.body)}, - test.apiKey, - nil, - ) - - if !errors.Is(result.Err, ErrUpstreamProtocol) || - !result.RequestWritten || - result.StatusCode != 0 || - len(result.Body) != 0 || - len(result.ClassificationBody) != 0 { - t.Fatalf("Forward() result = %#v, want fail-closed residual credential protocol error", result) - } - }) - } -} - -func TestForwardSuccessRepresentationRedactsEscapedAndFinalCredentialValues(t *testing.T) { - tests := []struct { - name string - apiKey string - body string - configureInput func(*ForwardInput) - }{ - { - name: "escaped JSON credential", - apiKey: "fake-escaped-credential", - body: `{"echo":"fake-\u0065scaped-credential"}`, - }, - { - name: "final custom credential Header value", - apiKey: "fake-provider-key-for-custom-credential", - body: `{"echo":"final-custom-fake-credential"}`, - configureInput: func(input *ForwardInput) { - input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ - "Authorization": "final-custom-fake-credential", - }} - }, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - result := forwardSuccessRepresentation( - t, - successRepresentationEncoding{name: "identity", encoding: "identity", wire: []byte(test.body)}, - test.apiKey, - test.configureInput, - ) - if result.Err != nil || result.StatusCode != http.StatusOK || - string(result.Body) != `{"echo":"[REDACTED]"}` { - t.Fatalf("Forward() result = %#v, want redacted JSON credential string", result) - } - }) - } -} - -func TestForwardSuccessRepresentationRedactsMultipleCredentialsInSingleJSONPass(t *testing.T) { - const ( - apiKey = "fake-overlap-credential" - finalCredential = "prefix-fake-overlap-credential" - ) - plain := []byte( - `{"long":"prefix-fake-overlap-credential","short":"fake-overlap-credential",` + - `"escaped":"fake-\u006fverlap-credential"}`, - ) - result := forwardSuccessRepresentation( - t, - successRepresentationEncoding{name: "identity", encoding: "identity", wire: plain}, - apiKey, - func(input *ForwardInput) { - input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ - "Authorization": finalCredential, - }} - }, - ) - - var got map[string]string - if err := json.Unmarshal(result.Body, &got); result.Err != nil || - result.StatusCode != http.StatusOK || - err != nil || - got["long"] != redact.Placeholder || - got["short"] != redact.Placeholder || - got["escaped"] != redact.Placeholder { - t.Fatalf("Forward() error/status/body/decoded = %v/%d/%q/%#v, decode error=%v", result.Err, result.StatusCode, result.Body, got, err) - } -} - -type overflowingResponseDialect struct { - *dialect.OpenAI -} - -func (dialect overflowingResponseDialect) RewriteResponseModel([]byte, string) ([]byte, error) { - return bytes.Repeat([]byte("x"), int(maxNonStreamingResponseBodyBytes)+1), nil -} - -type nonJSONAliasResponseDialect struct { - *dialect.OpenAI -} - -func (dialect nonJSONAliasResponseDialect) RewriteResponseModel([]byte, string) ([]byte, error) { - return []byte("alias=fake-non-json-alias-key"), nil -} - -func TestForwardSuccessRepresentationRejectsChangedOutputWireOverflow(t *testing.T) { - result := forwardSuccessRepresentation( - t, - successRepresentationEncoding{ - name: "identity", - encoding: "identity", - wire: []byte(`{"model":"provider-model"}`), - }, - "fake-provider-secret-output-overflow", - func(input *ForwardInput) { - input.Dialect = overflowingResponseDialect{OpenAI: dialect.NewOpenAI(http.DefaultClient)} - input.ExternalModel = "public-model" - input.UpstreamModelID = "provider-model" - }, - ) - - if !errors.Is(result.Err, ErrUpstreamProtocol) || - !result.RequestWritten || - result.StatusCode != 0 || - len(result.Body) != 0 || - len(result.ClassificationBody) != 0 { - t.Fatalf("Forward() result = %#v, want fail-closed changed output overflow", result) - } -} - -func TestForwardSuccessRepresentationRejectsCredentialIntroducedByAliasRewrite(t *testing.T) { - const providerBody = `{"model":"provider-model","id":"safe-response"}` - tests := []struct { - name string - encoding string - apiKey string - finalCredential string - externalModel string - nonJSON bool - }{ - { - name: "identity escaped API key", - encoding: "identity", - apiKey: "fake= http.StatusInternalServerError && statusCode <= 599: + return FailureCategoryUpstreamHostError + case statusCode >= http.StatusBadRequest && statusCode <= 499: + return FailureCategoryClientError + case attempt.Evidence != nil && attempt.Evidence.Kind == execution.ErrorKindInvalidRequest: + return FailureCategoryClientError + default: + return FailureCategoryAmbiguous + } +} + +func resultForExecutionCategory(category FailureCategory, attempt ExecutionAttempt) Result { + switch category { + case FailureCategoryRateLimited: + if attempt.Evidence != nil && attempt.Evidence.RetryAfter > 0 { + return Result{ + Category: category, + Action: ActionCooldownCredential, + CooldownUntil: attempt.Now.Add(attempt.Evidence.RetryAfter), + } + } + header := attempt.Header + if len(header) == 0 && attempt.Evidence != nil { + header = attempt.Evidence.Header + } + if until, ok := ParseRateLimitReset(header, attempt.Now); ok { + return Result{Category: category, Action: ActionCooldownCredential, CooldownUntil: until} + } + return Result{Category: category, Action: ActionCooldownCredential, UseFixed: true} + case FailureCategoryModelUnavailable: + return Result{ + Category: category, + Action: ActionCooldownCredential, + CooldownUntil: attempt.Now.Add(time.Hour), + } + case FailureCategoryInvalidKey: + return Result{Category: category, Action: ActionFailCredential} + case FailureCategoryUpstreamHostError: + return Result{Category: category, Action: ActionSkipGroup} + default: + return Result{Category: category, Action: ActionTerminate} + } +} + +func isSuccessStatus(statusCode int) bool { + return statusCode >= http.StatusOK && statusCode < http.StatusMultipleChoices +} + +func containsAny(value string, markers ...string) bool { + for _, marker := range markers { + if strings.Contains(value, marker) { + return true + } + } + return false +} diff --git a/internal/health/execution_judge_test.go b/internal/health/execution_judge_test.go new file mode 100644 index 00000000..8a26fdfd --- /dev/null +++ b/internal/health/execution_judge_test.go @@ -0,0 +1,208 @@ +package health + +import ( + "context" + "net/http" + "testing" + "time" + + "gpt-load/internal/execution" +) + +func TestJudgeExecutionUsesNeutralEvidenceAndReplayBoundary(t *testing.T) { + now := time.Date(2026, time.August, 9, 12, 0, 0, 0, time.UTC) + tests := []struct { + name string + attempt ExecutionAttempt + want Result + }{ + { + name: "downstream cancellation always terminates", + attempt: ExecutionAttempt{ + DownstreamErr: context.Canceled, + Evidence: evidence(execution.ErrorKindTransport, 0, "dial failed"), + }, + want: Result{Category: FailureCategoryDownstreamCancel, Action: ActionTerminate}, + }, + { + name: "committed response never retries", + attempt: ExecutionAttempt{ + DownstreamCommitted: true, + Evidence: evidence(execution.ErrorKindHTTP, http.StatusTooManyRequests, "rate limited"), + }, + want: Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, + }, + { + name: "clean committed response is successful", + attempt: ExecutionAttempt{ + DownstreamCommitted: true, + StatusCode: http.StatusOK, + }, + want: Result{Category: FailureCategoryOK, Action: ActionTerminate}, + }, + { + name: "not sent transport error skips group", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchNotSent, + Evidence: evidence(execution.ErrorKindTransport, 0, "dial failed"), + }, + want: Result{Category: FailureCategoryUpstreamHostError, Action: ActionSkipGroup}, + }, + { + name: "not sent timeout skips group", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchNotSent, + Evidence: evidence(execution.ErrorKindTimeout, 0, "connect timeout"), + }, + want: Result{Category: FailureCategoryUpstreamHostError, Action: ActionSkipGroup}, + }, + { + name: "maybe sent transport error is ambiguous", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + Evidence: evidence(execution.ErrorKindTransport, 0, "connection reset"), + }, + want: Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, + }, + { + name: "maybe sent timeout is ambiguous", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + Evidence: evidence(execution.ErrorKindTimeout, 0, "response timeout"), + }, + want: Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, + }, + { + name: "rate limit cools credential using evidence duration", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusTooManyRequests, + Now: now, + Evidence: &execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, + StatusCode: http.StatusTooManyRequests, + Summary: "rate limited", + RetryAfter: 12 * time.Second, + }, + }, + want: Result{ + Category: FailureCategoryRateLimited, + Action: ActionCooldownCredential, + CooldownUntil: now.Add(12 * time.Second), + }, + }, + { + name: "rate limit falls back to safe response header", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": {"30"}}, + Now: now, + Evidence: evidence(execution.ErrorKindHTTP, http.StatusTooManyRequests, "rate limited"), + }, + want: Result{ + Category: FailureCategoryRateLimited, + Action: ActionCooldownCredential, + CooldownUntil: now.Add(30 * time.Second), + }, + }, + { + name: "unauthorized fails credential", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusUnauthorized, + Evidence: evidence(execution.ErrorKindHTTP, http.StatusUnauthorized, "invalid credential"), + }, + want: Result{Category: FailureCategoryInvalidKey, Action: ActionFailCredential}, + }, + { + name: "payment required fails credential", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusPaymentRequired, + Evidence: evidence(execution.ErrorKindHTTP, http.StatusPaymentRequired, "billing disabled"), + }, + want: Result{Category: FailureCategoryInvalidKey, Action: ActionFailCredential}, + }, + { + name: "model not found cools credential", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusNotFound, + Now: now, + Evidence: &execution.ErrorEvidence{ + Kind: execution.ErrorKindProvider, + StatusCode: http.StatusNotFound, + Code: "model_not_found", + Summary: "model unavailable", + }, + }, + want: Result{ + Category: FailureCategoryModelUnavailable, + Action: ActionCooldownCredential, + CooldownUntil: now.Add(time.Hour), + }, + }, + { + name: "generic not found terminates", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusNotFound, + Evidence: evidence(execution.ErrorKindHTTP, http.StatusNotFound, "endpoint not found"), + }, + want: Result{Category: FailureCategoryClientError, Action: ActionTerminate}, + }, + { + name: "provider rate limit under success status still cools credential", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusOK, + Now: now, + Evidence: &execution.ErrorEvidence{ + Kind: execution.ErrorKindProvider, + StatusCode: http.StatusOK, + Type: "rate_limit_error", + Summary: "rate limited", + }, + }, + want: Result{Category: FailureCategoryRateLimited, Action: ActionCooldownCredential, UseFixed: true}, + }, + { + name: "server error skips group", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusServiceUnavailable, + Evidence: evidence(execution.ErrorKindHTTP, http.StatusServiceUnavailable, "overloaded"), + }, + want: Result{Category: FailureCategoryUpstreamHostError, Action: ActionSkipGroup}, + }, + { + name: "invalid request terminates before dispatch", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchNotSent, + Evidence: evidence(execution.ErrorKindInvalidRequest, 0, "unsupported request"), + }, + want: Result{Category: FailureCategoryClientError, Action: ActionTerminate}, + }, + { + name: "clean success terminates", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusOK, + }, + want: Result{Category: FailureCategoryOK, Action: ActionTerminate}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := JudgeExecution(test.attempt); got != test.want { + t.Fatalf("JudgeExecution() = %#v, want %#v", got, test.want) + } + }) + } +} + +func evidence(kind execution.ErrorKind, statusCode int, summary string) *execution.ErrorEvidence { + return &execution.ErrorEvidence{Kind: kind, StatusCode: statusCode, Summary: summary} +} diff --git a/internal/health/judge.go b/internal/health/judge.go deleted file mode 100644 index 03ef4dc7..00000000 --- a/internal/health/judge.go +++ /dev/null @@ -1,145 +0,0 @@ -package health - -import ( - "context" - "errors" - "net/http" - "time" -) - -type StatusClassifier interface { - ClassifyStatus(status int, body []byte) FailureCategory -} - -type ProviderErrorClassifier interface { - ClassifyProviderError(body []byte) FailureCategory -} - -type Attempt struct { - StatusCode int - Body []byte - Header http.Header - Now time.Time - Err error - RequestWritten bool - Committed bool - RetryableBeforeCommit bool - ProviderErrorBeforeCommit bool -} - -type Action uint8 - -const ( - ActionTerminate Action = iota - ActionRetry - ActionCooldownKey - ActionFailKey - ActionSkipGroup -) - -type Result struct { - Category FailureCategory - Action Action - CooldownUntil time.Time - UseFixed bool -} - -func (result Result) ShouldRetry() bool { - switch result.Action { - case ActionRetry, ActionCooldownKey, ActionFailKey, ActionSkipGroup: - return true - default: - return false - } -} - -type Rule func(StatusClassifier, Attempt) (Result, bool) - -var judgeRules = []Rule{ - terminalBoundaryRule, - explicitPreCommitRetryRule, - transportRule, - providerErrorRule, - statusRule, -} - -func Judge(classifier StatusClassifier, attempt Attempt) Result { - for _, rule := range judgeRules { - if result, matched := rule(classifier, attempt); matched { - return result - } - } - return Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate} -} - -func terminalBoundaryRule(_ StatusClassifier, attempt Attempt) (Result, bool) { - switch { - case errors.Is(attempt.Err, context.Canceled): - return Result{Category: FailureCategoryDownstreamCancel, Action: ActionTerminate}, true - case attempt.Committed && attempt.Err == nil: - return Result{Category: FailureCategoryOK, Action: ActionTerminate}, true - case attempt.Committed: - return Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, true - case attempt.Err != nil && attempt.RequestWritten: - return Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, true - default: - return Result{}, false - } -} - -func explicitPreCommitRetryRule(_ StatusClassifier, attempt Attempt) (Result, bool) { - if attempt.Err == nil || attempt.RequestWritten || !attempt.RetryableBeforeCommit { - return Result{}, false - } - return Result{Category: FailureCategoryAmbiguous, Action: ActionRetry}, true -} - -func transportRule(_ StatusClassifier, attempt Attempt) (Result, bool) { - if attempt.Err == nil { - return Result{}, false - } - if !attempt.RequestWritten { - return Result{Category: FailureCategoryUpstreamHostError, Action: ActionSkipGroup}, true - } - return Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, true -} - -func providerErrorRule(classifier StatusClassifier, attempt Attempt) (Result, bool) { - if !attempt.ProviderErrorBeforeCommit { - return Result{}, false - } - providerClassifier, ok := classifier.(ProviderErrorClassifier) - if !ok { - return Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, true - } - return resultForCategory( - providerClassifier.ClassifyProviderError(attempt.Body), - attempt, - ), true -} - -func statusRule(classifier StatusClassifier, attempt Attempt) (Result, bool) { - if classifier == nil || attempt.StatusCode == 0 { - return Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, true - } - category := classifier.ClassifyStatus(attempt.StatusCode, attempt.Body) - return resultForCategory(category, attempt), true -} - -func resultForCategory(category FailureCategory, attempt Attempt) Result { - switch category { - case FailureCategoryRateLimited: - if until, ok := ParseRateLimitReset(attempt.Header, attempt.Now); ok { - return Result{Category: category, Action: ActionCooldownKey, CooldownUntil: until} - } - return Result{Category: category, Action: ActionCooldownKey, UseFixed: true} - case FailureCategoryModelUnavailable: - return Result{Category: category, Action: ActionCooldownKey, CooldownUntil: attempt.Now.Add(time.Hour)} - case FailureCategoryInvalidKey: - return Result{Category: category, Action: ActionFailKey} - case FailureCategoryUpstreamHostError: - return Result{Category: category, Action: ActionSkipGroup} - default: - return Result{Category: category, Action: ActionTerminate} - } -} diff --git a/internal/health/judge_test.go b/internal/health/judge_test.go deleted file mode 100644 index ec2e4ebf..00000000 --- a/internal/health/judge_test.go +++ /dev/null @@ -1,252 +0,0 @@ -package health - -import ( - "context" - "errors" - "net/http" - "testing" - "time" -) - -type classifierFunc func(int, []byte) FailureCategory - -func (function classifierFunc) ClassifyStatus(status int, body []byte) FailureCategory { - return function(status, body) -} - -func fixedCategory(category FailureCategory) StatusClassifier { - return classifierFunc(func(int, []byte) FailureCategory { return category }) -} - -type providerClassifier struct { - status FailureCategory - provider FailureCategory -} - -func (classifier providerClassifier) ClassifyStatus(int, []byte) FailureCategory { - return classifier.status -} - -func (classifier providerClassifier) ClassifyProviderError([]byte) FailureCategory { - return classifier.provider -} - -func TestClassifyProviderErrorIgnoresHTTP2xxShortcut(t *testing.T) { - now := time.Date(2026, time.July, 27, 12, 0, 0, 0, time.UTC) - attempt := Attempt{ - StatusCode: http.StatusOK, - Body: []byte(`{"error":{"type":"rate_limit_error"}}`), - Header: http.Header{"Retry-After": {"15"}}, - Now: now, - ProviderErrorBeforeCommit: true, - } - got := Judge(providerClassifier{ - status: FailureCategoryOK, - provider: FailureCategoryRateLimited, - }, attempt) - want := Result{ - Category: FailureCategoryRateLimited, - Action: ActionCooldownKey, - CooldownUntil: now.Add(15 * time.Second), - } - if got != want { - t.Fatalf("Judge() = %#v, want %#v", got, want) - } - - withoutCapability := Attempt{ - StatusCode: http.StatusOK, - Body: attempt.Body, - Now: now, - ProviderErrorBeforeCommit: true, - } - if got := Judge(fixedCategory(FailureCategoryOK), withoutCapability); got != (Result{ - Category: FailureCategoryAmbiguous, - Action: ActionTerminate, - }) { - t.Fatalf("Judge() without optional classifier = %#v, want fail-closed ambiguous", got) - } -} - -func TestJudgeAppliesRulesInSafetyOrder(t *testing.T) { - now := time.Date(2026, time.July, 20, 12, 0, 0, 0, time.UTC) - - tests := []struct { - name string - classifier StatusClassifier - attempt Attempt - want Result - }{ - { - name: "committed overrides rate limit", - classifier: classifierFunc(func(int, []byte) FailureCategory { - return FailureCategoryRateLimited - }), - attempt: Attempt{Committed: true, StatusCode: http.StatusTooManyRequests, Now: now}, - want: Result{Category: FailureCategoryOK, Action: ActionTerminate}, - }, - { - name: "downstream cancellation terminates", - attempt: Attempt{Err: context.Canceled, RetryableBeforeCommit: true}, - want: Result{Category: FailureCategoryDownstreamCancel, Action: ActionTerminate}, - }, - { - name: "request-written ambiguous error terminates", - attempt: Attempt{Err: errors.New("connection reset"), RequestWritten: true}, - want: Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, - }, - { - name: "request-written explicit pre-commit signal terminates", - attempt: Attempt{Err: errors.New("first event timeout"), RequestWritten: true, RetryableBeforeCommit: true}, - want: Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, - }, - { - name: "pre-write transport skips group", - attempt: Attempt{Err: errors.New("dial tcp failed")}, - want: Result{Category: FailureCategoryUpstreamHostError, Action: ActionSkipGroup}, - }, - { - name: "rate limit uses parsed reset", - classifier: fixedCategory(FailureCategoryRateLimited), - attempt: Attempt{StatusCode: http.StatusTooManyRequests, - Header: http.Header{"Retry-After": {"30"}}, Now: now}, - want: Result{Category: FailureCategoryRateLimited, Action: ActionCooldownKey, CooldownUntil: now.Add(30 * time.Second)}, - }, - { - name: "rate limit without valid reset requests fixed fallback", - classifier: fixedCategory(FailureCategoryRateLimited), - attempt: Attempt{StatusCode: http.StatusTooManyRequests, Now: now}, - want: Result{Category: FailureCategoryRateLimited, Action: ActionCooldownKey, UseFixed: true}, - }, - { - name: "model unavailable cools for one hour", - classifier: fixedCategory(FailureCategoryModelUnavailable), - attempt: Attempt{StatusCode: http.StatusNotFound, Now: now}, - want: Result{Category: FailureCategoryModelUnavailable, Action: ActionCooldownKey, CooldownUntil: now.Add(time.Hour)}, - }, - { - name: "invalid key fails key", - classifier: fixedCategory(FailureCategoryInvalidKey), - attempt: Attempt{StatusCode: http.StatusUnauthorized, Now: now}, - want: Result{Category: FailureCategoryInvalidKey, Action: ActionFailKey}, - }, - { - name: "host error skips group", - classifier: fixedCategory(FailureCategoryUpstreamHostError), - attempt: Attempt{StatusCode: http.StatusServiceUnavailable, Now: now}, - want: Result{Category: FailureCategoryUpstreamHostError, Action: ActionSkipGroup}, - }, - { - name: "client error terminates", - classifier: fixedCategory(FailureCategoryClientError), - attempt: Attempt{StatusCode: http.StatusBadRequest, Now: now}, - want: Result{Category: FailureCategoryClientError, Action: ActionTerminate}, - }, - { - name: "nil classifier terminates", - attempt: Attempt{StatusCode: http.StatusInternalServerError}, - want: Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, - }, - { - name: "empty attempt terminates", - classifier: fixedCategory(FailureCategoryRateLimited), - want: Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, - }, - { - name: "OK terminates", - classifier: fixedCategory(FailureCategoryOK), - attempt: Attempt{StatusCode: http.StatusOK}, - want: Result{Category: FailureCategoryOK, Action: ActionTerminate}, - }, - { - name: "ambiguous terminates", - classifier: fixedCategory(FailureCategoryAmbiguous), - attempt: Attempt{StatusCode: http.StatusTemporaryRedirect}, - want: Result{Category: FailureCategoryAmbiguous, Action: ActionTerminate}, - }, - { - name: "downstream category terminates", - classifier: fixedCategory(FailureCategoryDownstreamCancel), - attempt: Attempt{StatusCode: http.StatusBadRequest}, - want: Result{Category: FailureCategoryDownstreamCancel, Action: ActionTerminate}, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - if got := Judge(test.classifier, test.attempt); got != test.want { - t.Fatalf("Judge() = %#v, want %#v", got, test.want) - } - }) - } -} - -func TestResultShouldRetryIsDerivedFromAction(t *testing.T) { - var zero Result - if zero.Action != ActionTerminate { - t.Fatalf("zero Result Action = %d, want ActionTerminate", zero.Action) - } - if zero.ShouldRetry() { - t.Fatal("zero Result ShouldRetry() = true, want false") - } - - tests := []struct { - action Action - want bool - }{ - {action: ActionTerminate, want: false}, - {action: ActionRetry, want: true}, - {action: ActionCooldownKey, want: true}, - {action: ActionFailKey, want: true}, - {action: ActionSkipGroup, want: true}, - {action: Action(255), want: false}, - } - for _, test := range tests { - if got := (Result{Action: test.action}).ShouldRetry(); got != test.want { - t.Fatalf("Action %d ShouldRetry() = %t, want %t", test.action, got, test.want) - } - } -} - -func TestJudgeReportsCategoryWithoutChangingAction(t *testing.T) { - cases := []struct { - name string - attempt Attempt - class FailureCategory - want Result - }{ - { - name: "client canceled", - attempt: Attempt{Err: context.Canceled, Committed: true}, - want: Result{Category: FailureCategoryDownstreamCancel, Action: ActionTerminate}, - }, - { - name: "committed clean", - attempt: Attempt{StatusCode: http.StatusOK, Committed: true}, - want: Result{Category: FailureCategoryOK, Action: ActionTerminate}, - }, - { - name: "retryable pre-commit", - attempt: Attempt{Err: errors.New("bootstrap"), RetryableBeforeCommit: true}, - want: Result{Category: FailureCategoryAmbiguous, Action: ActionRetry}, - }, - { - name: "transport before write", - attempt: Attempt{Err: errors.New("dial"), RequestWritten: false}, - want: Result{Category: FailureCategoryUpstreamHostError, Action: ActionSkipGroup}, - }, - { - name: "rate limited", - attempt: Attempt{StatusCode: http.StatusTooManyRequests, Now: time.Unix(1, 0)}, - class: FailureCategoryRateLimited, - want: Result{Category: FailureCategoryRateLimited, Action: ActionCooldownKey, UseFixed: true}, - }, - } - - for _, test := range cases { - t.Run(test.name, func(t *testing.T) { - if got := Judge(fixedCategory(test.class), test.attempt); got != test.want { - t.Fatalf("Judge() = %#v, want %#v", got, test.want) - } - }) - } -} diff --git a/internal/health/mutation.go b/internal/health/mutation.go index abdd5a6e..380c1813 100644 --- a/internal/health/mutation.go +++ b/internal/health/mutation.go @@ -15,24 +15,24 @@ func NewMutationCoordinator() *MutationCoordinator { return &MutationCoordinator{} } -func (coordinator *MutationCoordinator) Do(keyID uint, fn func()) { +func (coordinator *MutationCoordinator) Do(credentialID uint, fn func()) { if fn == nil { return } - index := int(keyID % uint(len(coordinator.stripes))) + index := int(credentialID % uint(len(coordinator.stripes))) coordinator.stripes[index].Lock() defer coordinator.stripes[index].Unlock() fn() } -func (coordinator *MutationCoordinator) DoMany(keyIDs []uint, fn func()) { +func (coordinator *MutationCoordinator) DoMany(credentialIDs []uint, fn func()) { if fn == nil { return } - selected := make(map[int]struct{}, len(keyIDs)) - indexes := make([]int, 0, len(keyIDs)) - for _, keyID := range keyIDs { - index := int(keyID % uint(len(coordinator.stripes))) + selected := make(map[int]struct{}, len(credentialIDs)) + indexes := make([]int, 0, len(credentialIDs)) + for _, credentialID := range credentialIDs { + index := int(credentialID % uint(len(coordinator.stripes))) if _, exists := selected[index]; exists { continue } diff --git a/internal/health/runtime_checkpoint.go b/internal/health/runtime_checkpoint.go index 385ebc27..8c4a07e3 100644 --- a/internal/health/runtime_checkpoint.go +++ b/internal/health/runtime_checkpoint.go @@ -14,9 +14,9 @@ type StatsBucketCheckpoint struct { Problem uint64 `json:"problem"` } -// StatsRuntimeCheckpoint contains the raw rolling-window state for one key. +// StatsRuntimeCheckpoint contains the raw rolling-window state for one credential. type StatsRuntimeCheckpoint struct { - KeyID uint `json:"key_id"` + CredentialID uint `json:"credential_id"` Buckets []StatsBucketCheckpoint `json:"buckets"` ConsecutiveFailure uint64 `json:"consecutive_failure"` ConsecutiveProblem uint64 `json:"consecutive_problem"` @@ -27,11 +27,11 @@ type StatsRuntimeCheckpoint struct { } // CaptureRuntimeCheckpoint returns detached rolling-window state in stable -// key order. It does not include request queues or access-key rate limits. +// credential order. It does not include request queues or access-key rate limits. func (store *StatsStore) CaptureRuntimeCheckpoint() []StatsRuntimeCheckpoint { store.mu.Lock() checkpoints := make([]StatsRuntimeCheckpoint, 0, len(store.windows)) - for keyID, window := range store.windows { + for credentialID, window := range store.windows { buckets := make([]StatsBucketCheckpoint, 0, len(window.buckets)) for _, bucket := range window.buckets { buckets = append(buckets, StatsBucketCheckpoint{ @@ -43,7 +43,7 @@ func (store *StatsStore) CaptureRuntimeCheckpoint() []StatsRuntimeCheckpoint { }) } checkpoints = append(checkpoints, StatsRuntimeCheckpoint{ - KeyID: keyID, + CredentialID: credentialID, Buckets: buckets, ConsecutiveFailure: window.consecutiveFailure, ConsecutiveProblem: window.consecutiveProblem, @@ -54,12 +54,14 @@ func (store *StatsStore) CaptureRuntimeCheckpoint() []StatsRuntimeCheckpoint { }) } store.mu.Unlock() - sort.Slice(checkpoints, func(i, j int) bool { return checkpoints[i].KeyID < checkpoints[j].KeyID }) + sort.Slice(checkpoints, func(i, j int) bool { + return checkpoints[i].CredentialID < checkpoints[j].CredentialID + }) return checkpoints } -// RestoreRuntimeCheckpoint replaces the supplied keys' rolling-window state. -// Invalid key IDs and malformed bucket lists are ignored because checkpoint +// RestoreRuntimeCheckpoint replaces the supplied credentials' rolling-window state. +// Invalid credential IDs and malformed bucket lists are ignored because checkpoint // data is explicitly best-effort. func (store *StatsStore) RestoreRuntimeCheckpoint(checkpoints []StatsRuntimeCheckpoint) int { store.mu.Lock() @@ -67,10 +69,10 @@ func (store *StatsStore) RestoreRuntimeCheckpoint(checkpoints []StatsRuntimeChec restored := 0 for _, checkpoint := range checkpoints { - if checkpoint.KeyID == 0 || len(checkpoint.Buckets) > statsBucketCount { + if checkpoint.CredentialID == 0 || len(checkpoint.Buckets) > statsBucketCount { continue } - window := &keyStatsWindow{ + window := &credentialStatsWindow{ consecutiveFailure: checkpoint.ConsecutiveFailure, consecutiveProblem: checkpoint.ConsecutiveProblem, lastEventAt: checkpoint.LastEventAt, @@ -91,7 +93,7 @@ func (store *StatsStore) RestoreRuntimeCheckpoint(checkpoints []StatsRuntimeChec problem: bucket.Problem, } } - store.windows[checkpoint.KeyID] = window + store.windows[checkpoint.CredentialID] = window restored++ } return restored diff --git a/internal/health/stats.go b/internal/health/stats.go index 488a54e2..11676000 100644 --- a/internal/health/stats.go +++ b/internal/health/stats.go @@ -11,7 +11,7 @@ const ( StatsWindow = time.Duration(statsBucketCount) * statsBucketWidth ) -type KeyStats struct { +type CredentialStats struct { Success uint64 Failure uint64 Problem uint64 @@ -23,7 +23,7 @@ type KeyStats struct { type StatsStore struct { mu sync.Mutex - windows map[uint]*keyStatsWindow + windows map[uint]*credentialStatsWindow } type statsBucket struct { @@ -34,7 +34,7 @@ type statsBucket struct { problem uint64 } -type keyStatsWindow struct { +type credentialStatsWindow struct { buckets [statsBucketCount]statsBucket consecutiveFailure uint64 consecutiveProblem uint64 @@ -53,39 +53,39 @@ const ( ) func NewStatsStore() *StatsStore { - return &StatsStore{windows: make(map[uint]*keyStatsWindow)} + return &StatsStore{windows: make(map[uint]*credentialStatsWindow)} } -func (store *StatsStore) RecordSuccess(keyID uint, at time.Time) { - store.record(keyID, statsEventSuccess, FailureCategoryOK, 0, at) +func (store *StatsStore) RecordSuccess(credentialID uint, at time.Time) { + store.record(credentialID, statsEventSuccess, FailureCategoryOK, 0, at) } func (store *StatsStore) RecordFailure( - keyID uint, + credentialID uint, category FailureCategory, statusCode int, at time.Time, ) { - store.record(keyID, statsEventFailure, category, statusCode, at) + store.record(credentialID, statsEventFailure, category, statusCode, at) } func (store *StatsStore) RecordProblem( - keyID uint, + credentialID uint, category FailureCategory, statusCode int, at time.Time, ) { - store.record(keyID, statsEventProblem, category, statusCode, at) + store.record(credentialID, statsEventProblem, category, statusCode, at) } func (store *StatsStore) record( - keyID uint, + credentialID uint, event statsEvent, category FailureCategory, statusCode int, at time.Time, ) { - if keyID == 0 { + if credentialID == 0 { return } @@ -95,10 +95,10 @@ func (store *StatsStore) record( store.mu.Lock() defer store.mu.Unlock() - window := store.windows[keyID] + window := store.windows[credentialID] if window == nil { - window = &keyStatsWindow{} - store.windows[keyID] = window + window = &credentialStatsWindow{} + store.windows[credentialID] = window } bucket := &window.buckets[slot] @@ -137,23 +137,23 @@ func (store *StatsStore) record( window.lastStatusCode = statusCode } -func (store *StatsStore) Reset(keyID uint) { - if keyID == 0 { +func (store *StatsStore) Reset(credentialID uint) { + if credentialID == 0 { return } store.mu.Lock() - delete(store.windows, keyID) + delete(store.windows, credentialID) store.mu.Unlock() } -func (store *StatsStore) ClearProblemState(keyID uint) { - if keyID == 0 { +func (store *StatsStore) ClearProblemState(credentialID uint) { + if credentialID == 0 { return } store.mu.Lock() defer store.mu.Unlock() - window := store.windows[keyID] + window := store.windows[credentialID] if window == nil { return } @@ -163,21 +163,21 @@ func (store *StatsStore) ClearProblemState(keyID uint) { window.lastStatusCode = 0 } -func (store *StatsStore) Snapshot(keyID uint, now time.Time) KeyStats { - if keyID == 0 { - return KeyStats{} +func (store *StatsStore) Snapshot(credentialID uint, now time.Time) CredentialStats { + if credentialID == 0 { + return CredentialStats{} } currentMinute := now.UnixNano() / int64(statsBucketWidth) store.mu.Lock() defer store.mu.Unlock() - window := store.windows[keyID] + window := store.windows[credentialID] if window == nil { - return KeyStats{} + return CredentialStats{} } - stats := KeyStats{ + stats := CredentialStats{ ConsecutiveFailure: window.consecutiveFailure, ConsecutiveProblem: window.consecutiveProblem, LastFailureCategory: window.lastFailureCategory, diff --git a/internal/health/stats_test.go b/internal/health/stats_test.go index f07b61ca..ddd8ba54 100644 --- a/internal/health/stats_test.go +++ b/internal/health/stats_test.go @@ -28,13 +28,13 @@ func TestStatsStoreSnapshotUnknownKeyReturnsZero(t *testing.T) { now := statsBase() for _, keyID := range []uint{0, 1} { - if got := store.Snapshot(keyID, now); got != (KeyStats{}) { + if got := store.Snapshot(keyID, now); got != (CredentialStats{}) { t.Fatalf("Snapshot(%d) = %#v, want zero value", keyID, got) } } store.RecordFailure(0, FailureCategoryAmbiguous, 0, now) - if got := store.Snapshot(0, now); got != (KeyStats{}) { + if got := store.Snapshot(0, now); got != (CredentialStats{}) { t.Fatalf("Snapshot(0) after Record = %#v, want zero value", got) } } @@ -49,7 +49,7 @@ func TestStatsStoreRecordAggregatesRollingWindow(t *testing.T) { store.RecordFailure(1, FailureCategoryAmbiguous, 0, base.Add(-5*time.Minute)) got := store.Snapshot(1, base) - want := KeyStats{Success: 2, Failure: 1, Problem: 1} + want := CredentialStats{Success: 2, Failure: 1, Problem: 1} if got != want { t.Fatalf("Snapshot() = %#v, want %#v", got, want) } @@ -64,7 +64,7 @@ func TestStatsStoreRecordDiscardsOlderSlotCollision(t *testing.T) { store.RecordSuccess(1, base.Add(5*time.Minute)) got := store.Snapshot(1, base.Add(5*time.Minute)) - want := KeyStats{Success: 1} + want := CredentialStats{Success: 1} if got != want { t.Fatalf("Snapshot() = %#v, want %#v", got, want) } @@ -76,7 +76,7 @@ func TestStatsStoreRecordDiscardsOlderSlotCollision(t *testing.T) { store.RecordSuccess(1, base) got := store.Snapshot(1, base.Add(5*time.Minute)) - want := KeyStats{Failure: 1, Problem: 1, ConsecutiveFailure: 1, ConsecutiveProblem: 1} + want := CredentialStats{Failure: 1, Problem: 1, ConsecutiveFailure: 1, ConsecutiveProblem: 1} if got != want { t.Fatalf("Snapshot() = %#v, want %#v", got, want) } @@ -91,7 +91,7 @@ func TestStatsStoreSnapshotExcludesExpiredAndFutureBuckets(t *testing.T) { store.RecordSuccess(1, base.Add(time.Minute)) got := store.Snapshot(1, base) - want := KeyStats{ConsecutiveFailure: 0} + want := CredentialStats{ConsecutiveFailure: 0} if got != want { t.Fatalf("Snapshot() = %#v, want %#v", got, want) } @@ -103,14 +103,14 @@ func TestStatsStoreConsecutiveFailureLifecycle(t *testing.T) { store.RecordFailure(1, FailureCategoryAmbiguous, 0, base.Add(-5*time.Minute)) store.RecordFailure(1, FailureCategoryAmbiguous, 0, base.Add(-4*time.Minute)) - if got, want := store.Snapshot(1, base), (KeyStats{ + if got, want := store.Snapshot(1, base), (CredentialStats{ Failure: 1, Problem: 1, ConsecutiveFailure: 2, ConsecutiveProblem: 2, }); got != want { t.Fatalf("after failures Snapshot() = %#v, want %#v", got, want) } store.RecordSuccess(1, base) - if got, want := store.Snapshot(1, base), (KeyStats{Success: 1, Failure: 1, Problem: 1}); got != want { + if got, want := store.Snapshot(1, base), (CredentialStats{Success: 1, Failure: 1, Problem: 1}); got != want { t.Fatalf("after success Snapshot() = %#v, want %#v", got, want) } } @@ -120,7 +120,7 @@ func TestStatsStoreProblemContextLifecycle(t *testing.T) { base := statsBase() store.RecordProblem(7, FailureCategoryRateLimited, 429, base) - if got, want := store.Snapshot(7, base), (KeyStats{ + if got, want := store.Snapshot(7, base), (CredentialStats{ Problem: 1, ConsecutiveProblem: 1, LastFailureCategory: FailureCategoryRateLimited, @@ -130,7 +130,7 @@ func TestStatsStoreProblemContextLifecycle(t *testing.T) { } store.RecordFailure(7, FailureCategoryInvalidKey, 401, base.Add(time.Second)) - if got, want := store.Snapshot(7, base), (KeyStats{ + if got, want := store.Snapshot(7, base), (CredentialStats{ Failure: 1, Problem: 2, ConsecutiveFailure: 1, @@ -142,7 +142,7 @@ func TestStatsStoreProblemContextLifecycle(t *testing.T) { } store.RecordSuccess(7, base.Add(2*time.Second)) - if got, want := store.Snapshot(7, base), (KeyStats{ + if got, want := store.Snapshot(7, base), (CredentialStats{ Success: 1, Failure: 1, Problem: 2, @@ -159,7 +159,7 @@ func TestStatsStoreOlderEventDoesNotReplaceLatestFailureContext(t *testing.T) { store.RecordProblem(7, FailureCategoryRateLimited, 429, base.Add(-time.Minute)) got := store.Snapshot(7, base) - want := KeyStats{ + want := CredentialStats{ Failure: 1, Problem: 2, ConsecutiveFailure: 1, @@ -179,7 +179,7 @@ func TestStatsStoreSnapshotReturnsValueCopy(t *testing.T) { snapshot := store.Snapshot(1, base) snapshot.Success = 99 - if got, want := store.Snapshot(1, base), (KeyStats{Success: 1}); got != want { + if got, want := store.Snapshot(1, base), (CredentialStats{Success: 1}); got != want { t.Fatalf("Snapshot() after mutating prior result = %#v, want %#v", got, want) } } @@ -192,7 +192,7 @@ func TestStatsStoreResetClearsWindowAndStreak(t *testing.T) { store.RecordFailure(7, FailureCategoryAmbiguous, 0, base) store.Reset(7) - if got := store.Snapshot(7, base); got != (KeyStats{}) { + if got := store.Snapshot(7, base); got != (CredentialStats{}) { t.Fatalf("Snapshot() after Reset = %#v, want zero", got) } } @@ -207,7 +207,7 @@ func TestStatsStoreClearProblemStatePreservesRollingBuckets(t *testing.T) { store.ClearProblemState(7) got := store.Snapshot(7, base) - want := KeyStats{Success: 1, Failure: 2, Problem: 2} + want := CredentialStats{Success: 1, Failure: 2, Problem: 2} if got != want { t.Fatalf("Snapshot() after ClearProblemState = %#v, want %#v", got, want) } @@ -224,13 +224,13 @@ func TestStatsStoreResetUnknownAndZeroKeyIsNoop(t *testing.T) { store.Reset(0) store.Reset(7) - if got := store.Snapshot(7, base); got != (KeyStats{}) { + if got := store.Snapshot(7, base); got != (CredentialStats{}) { t.Fatalf("Snapshot() after Reset unknown key = %#v, want zero", got) } store.RecordFailure(7, FailureCategoryAmbiguous, 0, base) store.Reset(0) - if got, want := store.Snapshot(7, base), (KeyStats{ + if got, want := store.Snapshot(7, base), (CredentialStats{ Failure: 1, Problem: 1, ConsecutiveFailure: 1, ConsecutiveProblem: 1, }); got != want { t.Fatalf("Snapshot() after Reset(0) = %#v, want %#v", got, want) diff --git a/internal/platform/errors/errors.go b/internal/platform/errors/errors.go index 4af6cc6d..3f9d0dee 100644 --- a/internal/platform/errors/errors.go +++ b/internal/platform/errors/errors.go @@ -30,16 +30,15 @@ var ( ErrDuplicateResource = &APIError{HTTPStatus: http.StatusConflict, Code: "DUPLICATE_RESOURCE", Message: "Resource already exists"} ErrResourceNotFound = &APIError{HTTPStatus: http.StatusNotFound, Code: "NOT_FOUND", Message: "Resource not found"} ErrGroupInUse = &APIError{HTTPStatus: http.StatusConflict, Code: "GROUP_IN_USE", Message: "Group is referenced by access keys"} - ErrInvalidKeyState = &APIError{HTTPStatus: http.StatusConflict, Code: "INVALID_KEY_STATE", Message: "Key cannot be restored from its current state"} + ErrInvalidCredentialState = &APIError{HTTPStatus: http.StatusConflict, Code: "INVALID_CREDENTIAL_STATE", Message: "Credential cannot be restored from its current state"} ErrInternalServer = &APIError{HTTPStatus: http.StatusInternalServerError, Code: "INTERNAL_SERVER_ERROR", Message: "An unexpected error occurred"} ErrDatabase = &APIError{HTTPStatus: http.StatusInternalServerError, Code: "DATABASE_ERROR", Message: "Database operation failed"} ErrUnauthorized = &APIError{HTTPStatus: http.StatusUnauthorized, Code: "UNAUTHORIZED", Message: "Authentication failed"} ErrForbidden = &APIError{HTTPStatus: http.StatusForbidden, Code: "FORBIDDEN", Message: "Access is forbidden"} ErrAuthLocked = &APIError{HTTPStatus: http.StatusTooManyRequests, Code: "AUTH_LOCKED", Message: "Authentication is temporarily locked"} - ErrUpstreamURLConflict = &APIError{HTTPStatus: http.StatusConflict, Code: "UPSTREAM_URL_CONFLICT", Message: "Upstream URL conflicts with an existing group"} + ErrChannelTargetConflict = &APIError{HTTPStatus: http.StatusConflict, Code: "CHANNEL_TARGET_CONFLICT", Message: "Channel target conflicts with an existing group"} ErrModelNameConflict = &APIError{HTTPStatus: http.StatusConflict, Code: "MODEL_NAME_CONFLICT", Message: "Client model names conflict within the group"} - ErrUpstreamURLChangeConfirmationRequired = &APIError{HTTPStatus: http.StatusConflict, Code: "UPSTREAM_URL_CHANGE_CONFIRMATION_REQUIRED", Message: "Upstream URL change requires explicit confirmation"} - ErrNoActiveUpstreamKey = &APIError{HTTPStatus: http.StatusConflict, Code: "NO_ACTIVE_UPSTREAM_KEY", Message: "No active upstream key available for this group"} + ErrNoActiveCredential = &APIError{HTTPStatus: http.StatusConflict, Code: "NO_ACTIVE_CREDENTIAL", Message: "No active credential is available for this group"} ErrBadGateway = &APIError{HTTPStatus: http.StatusBadGateway, Code: "BAD_GATEWAY", Message: "Upstream service error"} ErrIdempotencyKeyRequired = &APIError{HTTPStatus: http.StatusPreconditionRequired, Code: "IDEMPOTENCY_KEY_REQUIRED", Message: "Idempotency-Key is required"} ErrInvalidIdempotencyKey = &APIError{HTTPStatus: http.StatusBadRequest, Code: "INVALID_IDEMPOTENCY_KEY", Message: "Idempotency-Key must be a canonical UUID v4"} diff --git a/internal/platform/errors/errors_test.go b/internal/platform/errors/errors_test.go index 9f1a0a1b..c8c84023 100644 --- a/internal/platform/errors/errors_test.go +++ b/internal/platform/errors/errors_test.go @@ -29,10 +29,10 @@ func TestParseDBErrorUsesDatabaseIndependentGORMErrors(t *testing.T) { } } -func TestUpstreamURLChangeConfirmationRequiredContract(t *testing.T) { - err := ErrUpstreamURLChangeConfirmationRequired - if err.HTTPStatus != http.StatusConflict || err.Code != "UPSTREAM_URL_CHANGE_CONFIRMATION_REQUIRED" { - t.Fatalf("ErrUpstreamURLChangeConfirmationRequired = %#v", err) +func TestChannelTargetConflictContract(t *testing.T) { + err := ErrChannelTargetConflict + if err.HTTPStatus != http.StatusConflict || err.Code != "CHANNEL_TARGET_CONFLICT" { + t.Fatalf("ErrChannelTargetConflict = %#v", err) } } @@ -64,11 +64,11 @@ func TestParseDBErrorRecognizesNativeMySQLAndPostgreSQLUniqueConstraints(t *test func TestNewAPIErrorWithDataDoesNotMutateBase(t *testing.T) { data := map[string]any{"id": 12} - got := NewAPIErrorWithData(ErrUpstreamURLConflict, data) - if got == ErrUpstreamURLConflict || got.Data == nil { + got := NewAPIErrorWithData(ErrChannelTargetConflict, data) + if got == ErrChannelTargetConflict || got.Data == nil { t.Fatalf("NewAPIErrorWithData() = %#v", got) } - if ErrUpstreamURLConflict.Data != nil { - t.Fatalf("base Data = %#v, want nil", ErrUpstreamURLConflict.Data) + if ErrChannelTargetConflict.Data != nil { + t.Fatalf("base Data = %#v, want nil", ErrChannelTargetConflict.Data) } } diff --git a/internal/platform/i18n/i18n_test.go b/internal/platform/i18n/i18n_test.go index 2fec58b0..91618608 100644 --- a/internal/platform/i18n/i18n_test.go +++ b/internal/platform/i18n/i18n_test.go @@ -51,7 +51,7 @@ func TestRequestTooLargeTranslations(t *testing.T) { } } -func TestUpstreamURLChangeConfirmationTranslations(t *testing.T) { +func TestChannelTargetConflictTranslations(t *testing.T) { if err := Init(); err != nil { t.Fatalf("Init() error = %v", err) } @@ -60,9 +60,9 @@ func TestUpstreamURLChangeConfirmationTranslations(t *testing.T) { language string message string }{ - {language: "zh-CN", message: "修改上游地址需要明确确认"}, - {language: "en-US", message: "Changing the upstream URL requires explicit confirmation"}, - {language: "ja-JP", message: "アップストリームURLの変更には明示的な確認が必要です"}, + {language: "zh-CN", message: "已有分组使用相同渠道目标"}, + {language: "en-US", message: "An existing group already uses this channel target"}, + {language: "ja-JP", message: "同じチャネル接続先を使用するグループが既に存在します"}, } { t.Run(test.language, func(t *testing.T) { gin.SetMode(gin.TestMode) @@ -71,7 +71,7 @@ func TestUpstreamURLChangeConfirmationTranslations(t *testing.T) { context.Request.Header.Set("Accept-Language", test.language) Middleware()(context) - if got := Message(context, "group.upstream_url_change_confirmation_required"); got != test.message { + if got := Message(context, "group.channel_target_conflict"); got != test.message { t.Fatalf("Message() = %q, want %q", got, test.message) } }) diff --git a/internal/platform/i18n/locales/en-US.go b/internal/platform/i18n/locales/en-US.go index bab6bc63..0b860390 100644 --- a/internal/platform/i18n/locales/en-US.go +++ b/internal/platform/i18n/locales/en-US.go @@ -2,34 +2,33 @@ package locales // MessagesEnUS contains English (US) control-plane translations. var MessagesEnUS = map[string]string{ - "common.success": "Success", - "route.not_found": "Route not found", - "route.method_not_allowed": "Method not allowed", - "bad_request": "Bad request", - "request_too_large": "Request body is too large", - "bad_gateway": "Upstream service error", - "internal_error": "Internal error", - "idempotency.required": "Idempotency-Key is required", - "idempotency.invalid": "Idempotency-Key must be a canonical UUID v4", - "idempotency.reused": "Idempotency-Key was already used for another request", - "idempotency.expired": "The idempotent result retention period expired", - "control.operation_incomplete": "The resource was committed but runtime recovery is incomplete", - "control.recovery_pending": "An earlier committed operation is still recovering", - "settings.precondition_required": "If-Match is required", - "settings.version_conflict": "Settings changed since they were loaded", - "model_price.not_found": "Model price not found", - "model_price.unpriced_confirmation_required": "Marking a model price as unpriced requires explicit confirmation", - "model_price.referenced": "The model price is still referenced by Groups", - "model_price.automatic_delete_forbidden": "Automatic model prices cannot be deleted manually", - "auth.invalid_key": "Invalid authorization key", - "auth.forbidden": "The current key cannot access this feature", - "auth.locked": "Too many authentication attempts; try again later", - "group.not_found": "Group not found", - "group.name_exists": "Group name already exists", - "group.in_use": "The group is still referenced by access keys", - "group.upstream_url_conflict": "An existing group already uses this upstream URL", - "group.model_name_conflict": "Client model names must be unique within the group", - "group.upstream_url_change_confirmation_required": "Changing the upstream URL requires explicit confirmation", - "group.no_active_upstream_key": "No active upstream key is available for this group", - "key.not_found": "Key not found", + "common.success": "Success", + "route.not_found": "Route not found", + "route.method_not_allowed": "Method not allowed", + "bad_request": "Bad request", + "request_too_large": "Request body is too large", + "bad_gateway": "Upstream service error", + "internal_error": "Internal error", + "idempotency.required": "Idempotency-Key is required", + "idempotency.invalid": "Idempotency-Key must be a canonical UUID v4", + "idempotency.reused": "Idempotency-Key was already used for another request", + "idempotency.expired": "The idempotent result retention period expired", + "control.operation_incomplete": "The resource was committed but runtime recovery is incomplete", + "control.recovery_pending": "An earlier committed operation is still recovering", + "settings.precondition_required": "If-Match is required", + "settings.version_conflict": "Settings changed since they were loaded", + "model_price.not_found": "Model price not found", + "model_price.unpriced_confirmation_required": "Marking a model price as unpriced requires explicit confirmation", + "model_price.referenced": "The model price is still referenced by Groups", + "model_price.automatic_delete_forbidden": "Automatic model prices cannot be deleted manually", + "auth.invalid_key": "Invalid authorization key", + "auth.forbidden": "The current key cannot access this feature", + "auth.locked": "Too many authentication attempts; try again later", + "group.not_found": "Group not found", + "group.name_exists": "Group name already exists", + "group.in_use": "The group is still referenced by access keys", + "group.channel_target_conflict": "An existing group already uses this channel target", + "group.model_name_conflict": "Client model names must be unique within the group", + "group.no_active_credential": "No active credential is available for this group", + "credential.not_found": "Credential not found", } diff --git a/internal/platform/i18n/locales/ja-JP.go b/internal/platform/i18n/locales/ja-JP.go index 52fb6eb9..cf1b83e3 100644 --- a/internal/platform/i18n/locales/ja-JP.go +++ b/internal/platform/i18n/locales/ja-JP.go @@ -2,34 +2,33 @@ package locales // MessagesJaJP contains Japanese control-plane translations. var MessagesJaJP = map[string]string{ - "common.success": "成功", - "route.not_found": "ルートが見つかりません", - "route.method_not_allowed": "許可されていないHTTPメソッドです", - "bad_request": "不正なリクエスト", - "request_too_large": "リクエストボディが大きすぎます", - "bad_gateway": "上流サービスエラー", - "internal_error": "内部エラー", - "idempotency.required": "Idempotency-Key が必要です", - "idempotency.invalid": "Idempotency-Key は正規 UUID v4 である必要があります", - "idempotency.reused": "Idempotency-Key は別のリクエストで使用済みです", - "idempotency.expired": "冪等結果の保持期間が終了しました", - "control.operation_incomplete": "リソースは確定しましたが、実行時の復旧が未完了です", - "control.recovery_pending": "以前に確定した操作を復旧中です", - "settings.precondition_required": "If-Match が必要です", - "settings.version_conflict": "読み込み後に設定が変更されました", - "model_price.not_found": "モデル価格が存在しません", - "model_price.unpriced_confirmation_required": "モデル価格を未設定としてマークするには明示的な確認が必要です", - "model_price.referenced": "モデル価格はグループから参照されています", - "model_price.automatic_delete_forbidden": "自動モデル価格は手動で削除できません", - "auth.invalid_key": "無効な認証キー", - "auth.forbidden": "現在のキーにはこの機能へのアクセス権がありません", - "auth.locked": "認証試行回数が多すぎます。しばらくしてから再試行してください", - "group.not_found": "グループが存在しません", - "group.name_exists": "グループ名が既に存在します", - "group.in_use": "グループはアクセスキーから参照されています", - "group.upstream_url_conflict": "このアップストリームURLは既存のグループで使用されています", - "group.model_name_conflict": "グループ内のクライアントモデル名は一意である必要があります", - "group.upstream_url_change_confirmation_required": "アップストリームURLの変更には明示的な確認が必要です", - "group.no_active_upstream_key": "このグループには利用可能な有効なアップストリームキーがありません", - "key.not_found": "キーが存在しません", + "common.success": "成功", + "route.not_found": "ルートが見つかりません", + "route.method_not_allowed": "許可されていないHTTPメソッドです", + "bad_request": "不正なリクエスト", + "request_too_large": "リクエストボディが大きすぎます", + "bad_gateway": "上流サービスエラー", + "internal_error": "内部エラー", + "idempotency.required": "Idempotency-Key が必要です", + "idempotency.invalid": "Idempotency-Key は正規 UUID v4 である必要があります", + "idempotency.reused": "Idempotency-Key は別のリクエストで使用済みです", + "idempotency.expired": "冪等結果の保持期間が終了しました", + "control.operation_incomplete": "リソースは確定しましたが、実行時の復旧が未完了です", + "control.recovery_pending": "以前に確定した操作を復旧中です", + "settings.precondition_required": "If-Match が必要です", + "settings.version_conflict": "読み込み後に設定が変更されました", + "model_price.not_found": "モデル価格が存在しません", + "model_price.unpriced_confirmation_required": "モデル価格を未設定としてマークするには明示的な確認が必要です", + "model_price.referenced": "モデル価格はグループから参照されています", + "model_price.automatic_delete_forbidden": "自動モデル価格は手動で削除できません", + "auth.invalid_key": "無効な認証キー", + "auth.forbidden": "現在のキーにはこの機能へのアクセス権がありません", + "auth.locked": "認証試行回数が多すぎます。しばらくしてから再試行してください", + "group.not_found": "グループが存在しません", + "group.name_exists": "グループ名が既に存在します", + "group.in_use": "グループはアクセスキーから参照されています", + "group.channel_target_conflict": "同じチャネル接続先を使用するグループが既に存在します", + "group.model_name_conflict": "グループ内のクライアントモデル名は一意である必要があります", + "group.no_active_credential": "このグループには利用可能な認証情報がありません", + "credential.not_found": "認証情報が存在しません", } diff --git a/internal/platform/i18n/locales/zh-CN.go b/internal/platform/i18n/locales/zh-CN.go index c3e75d7a..63d29875 100644 --- a/internal/platform/i18n/locales/zh-CN.go +++ b/internal/platform/i18n/locales/zh-CN.go @@ -2,34 +2,33 @@ package locales // MessagesZhCN contains Simplified Chinese control-plane translations. var MessagesZhCN = map[string]string{ - "common.success": "操作成功", - "route.not_found": "路由不存在", - "route.method_not_allowed": "请求方法不被允许", - "bad_request": "请求错误", - "request_too_large": "请求体过大", - "bad_gateway": "上游服务错误", - "internal_error": "内部错误", - "idempotency.required": "必须提供 Idempotency-Key", - "idempotency.invalid": "Idempotency-Key 必须是规范 UUID v4", - "idempotency.reused": "Idempotency-Key 已用于其它请求", - "idempotency.expired": "幂等结果保留期已过", - "control.operation_incomplete": "资源已提交，但运行时恢复尚未完成", - "control.recovery_pending": "较早提交的操作仍在恢复", - "settings.precondition_required": "必须提供 If-Match", - "settings.version_conflict": "设置已在加载后发生变化", - "model_price.not_found": "模型价格不存在", - "model_price.unpriced_confirmation_required": "将模型价格标记为未定价需要明确确认", - "model_price.referenced": "模型价格仍被分组引用", - "model_price.automatic_delete_forbidden": "自动模型价格不能手动删除", - "auth.invalid_key": "无效的授权密钥", - "auth.forbidden": "当前密钥无权访问此功能", - "auth.locked": "认证尝试过多，请稍后重试", - "group.not_found": "分组不存在", - "group.name_exists": "分组名称已存在", - "group.in_use": "分组仍被访问密钥引用", - "group.upstream_url_conflict": "已有分组使用该上游地址", - "group.model_name_conflict": "分组内客户端模型名称必须唯一", - "group.upstream_url_change_confirmation_required": "修改上游地址需要明确确认", - "group.no_active_upstream_key": "该分组没有可用的活跃上游密钥", - "key.not_found": "密钥不存在", + "common.success": "操作成功", + "route.not_found": "路由不存在", + "route.method_not_allowed": "请求方法不被允许", + "bad_request": "请求错误", + "request_too_large": "请求体过大", + "bad_gateway": "上游服务错误", + "internal_error": "内部错误", + "idempotency.required": "必须提供 Idempotency-Key", + "idempotency.invalid": "Idempotency-Key 必须是规范 UUID v4", + "idempotency.reused": "Idempotency-Key 已用于其它请求", + "idempotency.expired": "幂等结果保留期已过", + "control.operation_incomplete": "资源已提交，但运行时恢复尚未完成", + "control.recovery_pending": "较早提交的操作仍在恢复", + "settings.precondition_required": "必须提供 If-Match", + "settings.version_conflict": "设置已在加载后发生变化", + "model_price.not_found": "模型价格不存在", + "model_price.unpriced_confirmation_required": "将模型价格标记为未定价需要明确确认", + "model_price.referenced": "模型价格仍被分组引用", + "model_price.automatic_delete_forbidden": "自动模型价格不能手动删除", + "auth.invalid_key": "无效的授权密钥", + "auth.forbidden": "当前密钥无权访问此功能", + "auth.locked": "认证尝试过多，请稍后重试", + "group.not_found": "分组不存在", + "group.name_exists": "分组名称已存在", + "group.in_use": "分组仍被访问密钥引用", + "group.channel_target_conflict": "已有分组使用相同渠道目标", + "group.model_name_conflict": "分组内客户端模型名称必须唯一", + "group.no_active_credential": "该分组没有可用凭据", + "credential.not_found": "凭据不存在", } diff --git a/internal/platform/response/response_test.go b/internal/platform/response/response_test.go index ae5d1d44..1fd8490a 100644 --- a/internal/platform/response/response_test.go +++ b/internal/platform/response/response_test.go @@ -49,10 +49,10 @@ func TestErrorI18nFromAPIErrorIncludesOptionalData(t *testing.T) { context.Request.Header.Set("Accept-Language", "en-US") i18n.Middleware()(context) - apiErr := app_errors.NewAPIErrorWithData(app_errors.ErrUpstreamURLConflict, map[string]any{ + apiErr := app_errors.NewAPIErrorWithData(app_errors.ErrChannelTargetConflict, map[string]any{ "groups": []string{"group-a", "group-b"}, }) - ErrorI18nFromAPIError(context, apiErr, "group.upstream_url_conflict") + ErrorI18nFromAPIError(context, apiErr, "group.channel_target_conflict") var body struct { Code string `json:"code"` @@ -63,7 +63,7 @@ func TestErrorI18nFromAPIErrorIncludesOptionalData(t *testing.T) { if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil { t.Fatalf("decode response: %v", err) } - if body.Code != app_errors.ErrUpstreamURLConflict.Code || len(body.Data.Groups) != 2 { + if body.Code != app_errors.ErrChannelTargetConflict.Code || len(body.Data.Groups) != 2 { t.Fatalf("response = %#v", body) } } diff --git a/internal/pricing/compile.go b/internal/pricing/compile.go index 8cab5a7f..8098e811 100644 --- a/internal/pricing/compile.go +++ b/internal/pricing/compile.go @@ -9,6 +9,7 @@ import ( const ( maxProviderIDBytes = 246 + maxChannelIDBytes = 64 maxModelIDBytes = 255 ) @@ -53,7 +54,11 @@ func NewTable(rules []Rule) (*Table, error) { return nil, err } if _, exists := table.rules[rule.Identity]; exists { - return nil, fmt.Errorf("duplicate pricing model ID %q", rule.Identity.ModelID) + return nil, fmt.Errorf( + "duplicate pricing identity channel_id=%q model_id=%q", + rule.Identity.ChannelID, + rule.Identity.ModelID, + ) } table.rules[rule.Identity] = cloneRule(rule) } @@ -99,13 +104,35 @@ func validateRule(rule Rule) error { } func validateIdentity(identity Identity) error { - if len(identity.ModelID) == 0 || len(identity.ModelID) > maxModelIDBytes { + if err := validateChannelID(identity.ChannelID); err != nil { + return err + } + return validateModelID(identity.ModelID) +} + +func validateChannelID(channelID string) error { + if len(channelID) == 0 || len(channelID) > maxChannelIDBytes { + return fmt.Errorf("channel ID must be 1 through %d bytes", maxChannelIDBytes) + } + if strings.TrimSpace(channelID) != channelID { + return fmt.Errorf("channel ID must not have surrounding whitespace") + } + for _, character := range channelID { + if unicode.IsControl(character) { + return fmt.Errorf("channel ID must not contain control characters") + } + } + return nil +} + +func validateModelID(modelID string) error { + if len(modelID) == 0 || len(modelID) > maxModelIDBytes { return fmt.Errorf("model ID must be 1 through %d bytes", maxModelIDBytes) } - if strings.TrimSpace(identity.ModelID) != identity.ModelID { + if strings.TrimSpace(modelID) != modelID { return fmt.Errorf("model ID must not have surrounding whitespace") } - for _, character := range identity.ModelID { + for _, character := range modelID { if unicode.IsControl(character) { return fmt.Errorf("model ID must not contain control characters") } @@ -113,20 +140,38 @@ func validateIdentity(identity Identity) error { return nil } -func validateReceiptRule(rule ReceiptRule, requireScope bool) error { - if err := validateIdentity(Identity{ModelID: rule.ModelID}); err != nil { - return err - } - if rule.ScopeKey == "" { - if requireScope { +func validateReceiptRule(rule ReceiptRule, schemaVersion int) error { + switch schemaVersion { + case 1: + if err := validateModelID(rule.ModelID); err != nil { + return err + } + if rule.ChannelID != "" { + return fmt.Errorf("legacy v1 receipt must not contain a channel ID") + } + if rule.ScopeKey == "" { return fmt.Errorf("legacy receipt scope key is required") } + return validateScopeKey(rule.ScopeKey) + case 2: + if err := validateModelID(rule.ModelID); err != nil { + return err + } + if rule.ScopeKey != "" { + return fmt.Errorf("global receipt must not contain a scope key") + } + if rule.ChannelID != "" { + return fmt.Errorf("global receipt must not contain a channel ID") + } return nil + case 3: + if rule.ScopeKey != "" { + return fmt.Errorf("channel receipt must not contain a scope key") + } + return validateIdentity(Identity{ChannelID: rule.ChannelID, ModelID: rule.ModelID}) + default: + return fmt.Errorf("unsupported receipt rule schema version") } - if !requireScope { - return fmt.Errorf("global receipt must not contain a scope key") - } - return validateScopeKey(rule.ScopeKey) } func validateScopeKey(scopeKey string) error { diff --git a/internal/pricing/compile_test.go b/internal/pricing/compile_test.go index 23cb1c37..c7139ab0 100644 --- a/internal/pricing/compile_test.go +++ b/internal/pricing/compile_test.go @@ -43,9 +43,9 @@ func TestProviderScopeKeyRejectsNonCanonicalProviderIDs(t *testing.T) { } } -func TestNewTableRequiresExactGlobalModelIdentity(t *testing.T) { - firstRule := Rule{Identity: Identity{ModelID: "openai/gpt-4.1"}} - secondRule := Rule{Identity: Identity{ModelID: "gpt-4.1"}} +func TestNewTableRequiresExactChannelModelIdentity(t *testing.T) { + firstRule := Rule{Identity: Identity{ChannelID: "openai", ModelID: "shared-model"}} + secondRule := Rule{Identity: Identity{ChannelID: "azure-openai", ModelID: "shared-model"}} table := mustTable(t, firstRule, secondRule) for _, identity := range []Identity{firstRule.Identity, secondRule.Identity} { @@ -54,8 +54,8 @@ func TestNewTableRequiresExactGlobalModelIdentity(t *testing.T) { } } for _, identity := range []Identity{ - {ModelID: "openai/gpt-4"}, - {ModelID: "gpt-4.1-preview"}, + {ChannelID: "anthropic", ModelID: "shared-model"}, + {ChannelID: "openai", ModelID: "shared-model-preview"}, } { if got, ok := table.Lookup(identity); ok { t.Fatalf("Lookup(%#v) = %#v, true, want exact miss", identity, got) @@ -64,14 +64,19 @@ func TestNewTableRequiresExactGlobalModelIdentity(t *testing.T) { } func TestNewTableRejectsInvalidIdentityAndDuplicates(t *testing.T) { - valid := Identity{ModelID: "gpt-4.1"} + valid := Identity{ChannelID: "openai", ModelID: "gpt-4.1"} invalid := []Identity{ {}, - {ModelID: ""}, - {ModelID: " gpt-4.1"}, - {ModelID: "gpt-4.1 "}, - {ModelID: "gpt-4.1\n"}, - {ModelID: strings.Repeat("m", 256)}, + {ChannelID: "", ModelID: "gpt-4.1"}, + {ChannelID: " openai", ModelID: "gpt-4.1"}, + {ChannelID: "openai ", ModelID: "gpt-4.1"}, + {ChannelID: "openai\n", ModelID: "gpt-4.1"}, + {ChannelID: strings.Repeat("c", 65), ModelID: "gpt-4.1"}, + {ChannelID: "openai", ModelID: ""}, + {ChannelID: "openai", ModelID: " gpt-4.1"}, + {ChannelID: "openai", ModelID: "gpt-4.1 "}, + {ChannelID: "openai", ModelID: "gpt-4.1\n"}, + {ChannelID: "openai", ModelID: strings.Repeat("m", 256)}, } for _, identity := range invalid { t.Run(identity.ModelID, func(t *testing.T) { @@ -83,21 +88,43 @@ func TestNewTableRejectsInvalidIdentityAndDuplicates(t *testing.T) { if _, err := NewTable([]Rule{{Identity: valid}, {Identity: valid, IsManual: true}}); err == nil { t.Fatal("NewTable() accepted duplicate exact identity") + } else if !strings.Contains(err.Error(), valid.ChannelID) || !strings.Contains(err.Error(), valid.ModelID) { + t.Fatalf("duplicate error %q does not identify channel and model", err) } } -func TestNewTableRejectsRepeatedModelID(t *testing.T) { +func TestNewTableAllowsRepeatedModelIDAcrossChannels(t *testing.T) { + rules := []Rule{ + {Identity: Identity{ChannelID: "openai", ModelID: "shared-model"}}, + {Identity: Identity{ChannelID: "azure-openai", ModelID: "shared-model"}}, + } + table, err := NewTable(rules) + if err != nil { + t.Fatalf("NewTable() error = %v", err) + } + for _, rule := range rules { + if got, ok := table.Lookup(rule.Identity); !ok || got.Identity != rule.Identity { + t.Fatalf("Lookup(%#v) = %#v, %t", rule.Identity, got, ok) + } + } + if _, ok := table.Lookup(Identity{ChannelID: "anthropic", ModelID: "shared-model"}); ok { + t.Fatal("Lookup() matched a model from another channel") + } +} + +func TestNewTableRejectsDuplicateChannelModelIdentity(t *testing.T) { + identity := Identity{ChannelID: "openai", ModelID: "shared-model"} _, err := NewTable([]Rule{ - {Identity: Identity{ModelID: "shared-model"}}, - {Identity: Identity{ModelID: "shared-model"}}, + {Identity: identity}, + {Identity: identity}, }) if err == nil { - t.Fatal("NewTable() accepted two rules for the same upstream model") + t.Fatal("NewTable() accepted duplicate channel/model identity") } } func TestNewTableValidatesPricesAndContextTiers(t *testing.T) { - identity := Identity{ModelID: "gpt-4.1"} + identity := Identity{ChannelID: "openai", ModelID: "gpt-4.1"} tests := []struct { name string rule Rule @@ -131,7 +158,7 @@ func TestNewTableValidatesPricesAndContextTiers(t *testing.T) { func TestTableDeepClonesConstructionAndLookup(t *testing.T) { rules := []Rule{{ - Identity: Identity{ModelID: "claude:3_5/model"}, + Identity: Identity{ChannelID: "anthropic", ModelID: "claude:3_5/model"}, Prices: Prices{Input: fixedPrice(1)}, ContextTiers: []ContextTier{{ InputThresholdTokens: 100, diff --git a/internal/pricing/quote.go b/internal/pricing/quote.go index 95dc32fb..d1366541 100644 --- a/internal/pricing/quote.go +++ b/internal/pricing/quote.go @@ -67,11 +67,14 @@ func (table *Table) QuoteWithReceipt( {code: "output", tokens: result.Tokens.Output, price: prices.Output, multiplier: directPriceMultiplier}, } receipt := &Receipt{ - SchemaVersion: 2, - Method: ReceiptMethodUnitRateSum, - MethodVersion: 1, - Currency: "USD", - Rule: ReceiptRule{ModelID: identity.ModelID}, + SchemaVersion: 3, + Method: ReceiptMethodUnitRateSum, + MethodVersion: 1, + Currency: "USD", + Rule: ReceiptRule{ + ChannelID: identity.ChannelID, + ModelID: identity.ModelID, + }, ContextThresholdTokens: selectedThreshold, LineItems: make([]ReceiptLine, 0, len(components)+1), } diff --git a/internal/pricing/quote_test.go b/internal/pricing/quote_test.go index aa486cb2..69eaa8ad 100644 --- a/internal/pricing/quote_test.go +++ b/internal/pricing/quote_test.go @@ -28,7 +28,7 @@ func TestQuoteMapsUsageStatesBeforeLookup(t *testing.T) { } func TestQuoteUsesHighestEligibleTierAtInclusiveBoundary(t *testing.T) { - identity := Identity{ModelID: "gpt-4.1"} + identity := Identity{ChannelID: "openai", ModelID: "gpt-4.1"} table := mustTable(t, Rule{ Identity: identity, Prices: Prices{Input: fixedPrice(1)}, @@ -49,7 +49,7 @@ func TestQuoteUsesHighestEligibleTierAtInclusiveBoundary(t *testing.T) { } func TestQuoteTierReplacesBaseSlotsWithoutFallback(t *testing.T) { - identity := Identity{ModelID: "claude-sonnet"} + identity := Identity{ChannelID: "anthropic", ModelID: "claude-sonnet"} table := mustTable(t, Rule{ Identity: identity, Prices: Prices{Input: fixedPrice(10), Output: fixedPrice(20)}, @@ -70,7 +70,7 @@ func TestQuoteTierReplacesBaseSlotsWithoutFallback(t *testing.T) { } func TestQuotePricesCacheWriteFiveMinutesDirectlyAndOneHourAtExactEightFifths(t *testing.T) { - identity := Identity{ModelID: "claude-sonnet"} + identity := Identity{ChannelID: "anthropic", ModelID: "claude-sonnet"} table := mustTable(t, Rule{Identity: identity, Prices: Prices{CacheWrite: fixedPrice(5)}}) quote := table.Quote(identity, usage.Result{ Tokens: usage.Tokens{CacheWrite5M: 1_000_000, CacheWrite1H: 1_000_000}, @@ -83,7 +83,7 @@ func TestQuotePricesCacheWriteFiveMinutesDirectlyAndOneHourAtExactEightFifths(t } func TestQuoteIncludesUnknownCacheWriteInTierButNeverChargesIt(t *testing.T) { - identity := Identity{ModelID: "claude-sonnet"} + identity := Identity{ChannelID: "anthropic", ModelID: "claude-sonnet"} table := mustTable(t, Rule{ Identity: identity, Prices: Prices{Input: fixedPrice(1), CacheWrite: fixedPrice(math.MaxInt64)}, @@ -103,7 +103,7 @@ func TestQuoteIncludesUnknownCacheWriteInTierButNeverChargesIt(t *testing.T) { } func TestQuotePreservesKnownCostWhenAnotherPositiveComponentIsUnavailable(t *testing.T) { - identity := Identity{ModelID: "gpt-4.1"} + identity := Identity{ChannelID: "openai", ModelID: "gpt-4.1"} table := mustTable(t, Rule{Identity: identity, Prices: Prices{Input: fixedPrice(7)}}) quote := table.Quote(identity, usage.Result{ Tokens: usage.Tokens{UncachedInput: 1_000_000, Output: 1}, @@ -116,7 +116,7 @@ func TestQuotePreservesKnownCostWhenAnotherPositiveComponentIsUnavailable(t *tes } func TestQuoteCompleteUsageDiagnosticsDoNotMakeKnownCostPartial(t *testing.T) { - identity := Identity{ModelID: "gpt-4.1"} + identity := Identity{ChannelID: "openai", ModelID: "gpt-4.1"} table := mustTable(t, Rule{Identity: identity, Prices: Prices{Input: fixedPrice(9)}}) tests := []struct { name string @@ -151,7 +151,7 @@ func TestQuoteCompleteUsageDiagnosticsDoNotMakeKnownCostPartial(t *testing.T) { } func TestQuoteSeparatesPricedStateFromCompleteness(t *testing.T) { - identity := Identity{ModelID: "gpt-4.1"} + identity := Identity{ChannelID: "openai", ModelID: "gpt-4.1"} tests := []struct { name string result usage.Result @@ -171,9 +171,12 @@ func TestQuoteSeparatesPricedStateFromCompleteness(t *testing.T) { }) } - if got := allNull.Quote(Identity{ModelID: "missing"}, usage.Result{State: usage.StateComplete}); got != (Quote{State: CostStateUnpriced, Completeness: CompletenessUnavailable}) { + if got := allNull.Quote(Identity{ChannelID: identity.ChannelID, ModelID: "missing"}, usage.Result{State: usage.StateComplete}); got != (Quote{State: CostStateUnpriced, Completeness: CompletenessUnavailable}) { t.Fatalf("missing exact rule Quote() = %#v, want unpriced unavailable", got) } + if got := allNull.Quote(Identity{ChannelID: "other-channel", ModelID: identity.ModelID}, usage.Result{State: usage.StateComplete}); got != (Quote{State: CostStateUnpriced, Completeness: CompletenessUnavailable}) { + t.Fatalf("wrong-channel rule Quote() = %#v, want unpriced unavailable", got) + } } explicitZero := mustTable(t, Rule{Identity: identity, Prices: Prices{Output: fixedPrice(0)}}) @@ -184,7 +187,7 @@ func TestQuoteSeparatesPricedStateFromCompleteness(t *testing.T) { } func TestQuoteRejectsNegativeTokensAndOverflowFailClosed(t *testing.T) { - identity := Identity{ModelID: "gpt-4.1"} + identity := Identity{ChannelID: "openai", ModelID: "gpt-4.1"} table := mustTable(t, Rule{Identity: identity, Prices: Prices{ Input: fixedPrice(math.MaxInt64), Output: fixedPrice(math.MaxInt64), CacheRead: fixedPrice(0), CacheWrite: fixedPrice(0), @@ -214,7 +217,7 @@ func TestQuoteRejectsNegativeTokensAndOverflowFailClosed(t *testing.T) { } func TestQuotePartialUsageKeepsKnownCostAndPartialCompleteness(t *testing.T) { - identity := Identity{ModelID: "gemini-2.5-pro"} + identity := Identity{ChannelID: "gemini", ModelID: "gemini-2.5-pro"} table := mustTable(t, Rule{Identity: identity, Prices: Prices{Output: fixedPrice(4)}}) var diagnostics usage.Diagnostics diagnostics.Add(usage.DiagnosticInvalidPayload) @@ -230,7 +233,7 @@ func TestQuotePartialUsageKeepsKnownCostAndPartialCompleteness(t *testing.T) { } func TestQuoteWithReceiptFreezesExactTieredCalculation(t *testing.T) { - identity := Identity{ModelID: "claude-sonnet"} + identity := Identity{ChannelID: "anthropic", ModelID: "claude-sonnet"} threshold := int64(1_000) table := mustTable(t, Rule{ Identity: identity, @@ -263,12 +266,15 @@ func TestQuoteWithReceiptFreezesExactTieredCalculation(t *testing.T) { }) { t.Fatalf("QuoteWithReceipt() quote = %#v", quote) } - if receipt == nil || receipt.SchemaVersion != 2 || receipt.Method != "unit_rate_sum" || + if receipt == nil || receipt.SchemaVersion != 3 || receipt.Method != "unit_rate_sum" || receipt.MethodVersion != 1 || receipt.Currency != "USD" || - receipt.Rule != (ReceiptRule{ModelID: identity.ModelID}) || receipt.ContextThresholdTokens == nil || + receipt.Rule != (ReceiptRule{ChannelID: identity.ChannelID, ModelID: identity.ModelID}) || receipt.ContextThresholdTokens == nil || *receipt.ContextThresholdTokens != threshold || receipt.TotalNanoUSD != 27 { t.Fatalf("QuoteWithReceipt() receipt = %#v", receipt) } + if err := ValidateReceipt(*receipt); err != nil { + t.Fatalf("generated v3 receipt is invalid: %v", err) + } if got, want := len(receipt.LineItems), 5; got != want { t.Fatalf("receipt line items = %d, want %d: %#v", got, want, receipt.LineItems) } @@ -280,7 +286,7 @@ func TestQuoteWithReceiptFreezesExactTieredCalculation(t *testing.T) { } func TestQuoteWithReceiptPreservesUnpricedPositiveComponents(t *testing.T) { - identity := Identity{ModelID: "gpt-4.1"} + identity := Identity{ChannelID: "openai", ModelID: "gpt-4.1"} table := mustTable(t, Rule{Identity: identity, Prices: Prices{Input: fixedPrice(7)}}) quote, receipt := table.QuoteWithReceipt(identity, usage.Result{ diff --git a/internal/pricing/receipt.go b/internal/pricing/receipt.go index 9060145d..476c070d 100644 --- a/internal/pricing/receipt.go +++ b/internal/pricing/receipt.go @@ -5,11 +5,11 @@ import "fmt" // ValidateReceipt verifies a persisted request-time receipt without consulting // the mutable current pricing table. func ValidateReceipt(receipt Receipt) error { - if (receipt.SchemaVersion != 1 && receipt.SchemaVersion != 2) || + if (receipt.SchemaVersion != 1 && receipt.SchemaVersion != 2 && receipt.SchemaVersion != 3) || receipt.Method != ReceiptMethodUnitRateSum || receipt.MethodVersion != 1 || receipt.Currency != "USD" { return fmt.Errorf("unsupported pricing receipt contract") } - if err := validateReceiptRule(receipt.Rule, receipt.SchemaVersion == 1); err != nil { + if err := validateReceiptRule(receipt.Rule, receipt.SchemaVersion); err != nil { return fmt.Errorf("invalid pricing receipt rule: %w", err) } if receipt.ContextThresholdTokens != nil && *receipt.ContextThresholdTokens < 0 { diff --git a/internal/pricing/receipt_test.go b/internal/pricing/receipt_test.go new file mode 100644 index 00000000..85e2114d --- /dev/null +++ b/internal/pricing/receipt_test.go @@ -0,0 +1,88 @@ +package pricing + +import ( + "strings" + "testing" +) + +func TestValidateReceiptPreservesVersionedRuleIdentitySemantics(t *testing.T) { + tests := []struct { + name string + receipt Receipt + }{ + { + name: "v1 requires and accepts legacy scope", + receipt: validReceipt(1, ReceiptRule{ + ScopeKey: "provider:openai", + ModelID: "gpt-4.1", + }), + }, + { + name: "v2 remains global", + receipt: validReceipt(2, ReceiptRule{ModelID: "gpt-4.1"}), + }, + { + name: "v3 freezes channel and model", + receipt: validReceipt(3, ReceiptRule{ + ChannelID: "openai", + ModelID: "gpt-4.1", + }), + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if err := ValidateReceipt(test.receipt); err != nil { + t.Fatalf("ValidateReceipt() error = %v", err) + } + }) + } +} + +func TestValidateReceiptRejectsCrossVersionIdentityReinterpretation(t *testing.T) { + tests := []struct { + name string + receipt Receipt + }{ + {name: "v1 missing scope", receipt: validReceipt(1, ReceiptRule{ModelID: "gpt-4.1"})}, + {name: "v1 with channel", receipt: validReceipt(1, ReceiptRule{ScopeKey: "provider:openai", ChannelID: "openai", ModelID: "gpt-4.1"})}, + {name: "v2 with legacy scope", receipt: validReceipt(2, ReceiptRule{ScopeKey: "provider:openai", ModelID: "gpt-4.1"})}, + {name: "v2 with channel", receipt: validReceipt(2, ReceiptRule{ChannelID: "openai", ModelID: "gpt-4.1"})}, + {name: "v3 missing channel", receipt: validReceipt(3, ReceiptRule{ModelID: "gpt-4.1"})}, + {name: "v3 with legacy scope", receipt: validReceipt(3, ReceiptRule{ScopeKey: "provider:openai", ChannelID: "openai", ModelID: "gpt-4.1"})}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if err := ValidateReceipt(test.receipt); err == nil { + t.Fatalf("ValidateReceipt() accepted incompatible rule %#v", test.receipt.Rule) + } + }) + } +} + +func TestValidateReceiptV3UsesStrictChannelIDValidation(t *testing.T) { + for _, channelID := range []string{ + "", + " openai", + "openai ", + "openai\n", + strings.Repeat("c", 65), + } { + t.Run(channelID, func(t *testing.T) { + receipt := validReceipt(3, ReceiptRule{ChannelID: channelID, ModelID: "gpt-4.1"}) + if err := ValidateReceipt(receipt); err == nil { + t.Fatalf("ValidateReceipt() accepted channel ID %q", channelID) + } + }) + } +} + +func validReceipt(schemaVersion int, rule ReceiptRule) Receipt { + return Receipt{ + SchemaVersion: schemaVersion, + Method: ReceiptMethodUnitRateSum, + MethodVersion: 1, + Currency: "USD", + Rule: rule, + LineItems: []ReceiptLine{}, + } +} diff --git a/internal/pricing/types.go b/internal/pricing/types.go index 0fc480fe..de98fc98 100644 --- a/internal/pricing/types.go +++ b/internal/pricing/types.go @@ -15,18 +15,19 @@ type Prices struct { CacheWrite Price } -// Identity is the exact global pricing identity of a real upstream model. -// Routing groups and providers deliberately do not participate in matching. +// Identity is the exact channel and upstream-model pricing identity. type Identity struct { - ModelID string `json:"model_id"` + ChannelID string `json:"channel_id"` + ModelID string `json:"model_id"` } // ReceiptRule is the frozen model identity written into a request-time cost -// receipt. ScopeKey is retained only to read historical v1 receipts. New v2 -// receipts leave it empty because prices are global. +// receipt. ScopeKey is retained only to read historical v1 receipts; ChannelID +// is populated only by v3 and later receipts. type ReceiptRule struct { - ScopeKey string `json:"scope_key,omitempty"` - ModelID string `json:"model_id"` + ScopeKey string `json:"scope_key,omitempty"` + ChannelID string `json:"channel_id,omitempty"` + ModelID string `json:"model_id"` } // ContextTier replaces all base prices once its inclusive threshold is met. @@ -35,7 +36,7 @@ type ContextTier struct { Prices Prices } -// Rule is one exact scope and upstream-model price definition. +// Rule is one exact channel and upstream-model price definition. type Rule struct { Identity Identity Prices Prices diff --git a/internal/reasoning/config.go b/internal/reasoning/config.go index 0e6ac445..21fad2c0 100644 --- a/internal/reasoning/config.go +++ b/internal/reasoning/config.go @@ -11,3 +11,16 @@ type Config struct { func (config Config) Present() bool { return config.Mode != "" || config.Effort != "" || config.BudgetTokens != nil } + +// RequiresCapability reports whether the request asks the upstream to perform +// reasoning. An explicit "none" setting is a valid opt-out and must not make +// routing require a reasoning-capable target. +func (config Config) RequiresCapability() bool { + if config.BudgetTokens != nil { + return true + } + if config.Mode != "" && config.Mode != "none" { + return true + } + return config.Effort != "" && config.Effort != "none" +} diff --git a/internal/requestlog/channel_identity_test.go b/internal/requestlog/channel_identity_test.go new file mode 100644 index 00000000..0b67b686 --- /dev/null +++ b/internal/requestlog/channel_identity_test.go @@ -0,0 +1,328 @@ +package requestlog + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/platform/redact" + "gpt-load/internal/pricing" + "gpt-load/internal/storage/models" + "gpt-load/internal/telemetry" + "gpt-load/internal/usage" +) + +func TestMapEventFreezesChannelCredentialAndExecutionAttempt(t *testing.T) { + event := channelScopedEvent(t, "00000000-0000-4000-8000-000000007001") + + row, err := mapEvent(redact.New(), event) + if err != nil { + t.Fatalf("mapEvent() error = %v", err) + } + if row.ChannelID != string(channel.OpenAI) || row.CredentialID != 8 { + t.Fatalf("request attribution = %q/%d", row.ChannelID, row.CredentialID) + } + if len(row.AttemptRows) != 1 { + t.Fatalf("attempt rows = %#v", row.AttemptRows) + } + attempt := row.AttemptRows[0] + if attempt.ChannelID != string(channel.OpenAI) || attempt.CredentialID != 8 || + attempt.Operation != string(execution.OperationChatCompletion) || + attempt.RouteMode != string(channel.RouteNative) || + attempt.DispatchState != string(execution.DispatchMaybeSent) || + !attempt.ResponseStarted || !attempt.Committed || + attempt.UpstreamRequestID != "upstream-request-1" { + t.Fatalf("frozen attempt = %#v", attempt) + } +} + +func TestMapEventV3ReceiptMustMatchChannelAndModel(t *testing.T) { + event := channelScopedEvent(t, "00000000-0000-4000-8000-000000007002") + receipt := decodeReceiptJSON(t, event.Usage.Pricing.ReceiptJSON) + receipt.Rule.ChannelID = string(channel.Anthropic) + event.Usage.Pricing.ReceiptJSON = encodeReceiptJSON(t, receipt) + if _, err := mapEvent(redact.New(), event); err == nil || + !strings.Contains(err.Error(), "frozen pricing observation") { + t.Fatalf("mapEvent() error = %v, want v3 channel/model mismatch", err) + } + + event = channelScopedEvent(t, "00000000-0000-4000-8000-000000007003") + receipt = decodeReceiptJSON(t, event.Usage.Pricing.ReceiptJSON) + receipt.Rule.ModelID = "other-model" + event.Usage.Pricing.ReceiptJSON = encodeReceiptJSON(t, receipt) + if _, err := mapEvent(redact.New(), event); err == nil || + !strings.Contains(err.Error(), "frozen pricing observation") { + t.Fatalf("mapEvent() error = %v, want v3 channel/model mismatch", err) + } +} + +func TestMapEventDoesNotAcceptHistoricalReceiptSchemasForNewWrites(t *testing.T) { + event := channelScopedEvent(t, "00000000-0000-4000-8000-000000007004") + event.Usage.Pricing.ReceiptJSON = encodeReceiptJSON(t, emptyReceipt(2, pricing.ReceiptRule{ + ModelID: event.UpstreamModel, + })) + if _, err := mapEvent(redact.New(), event); err == nil { + t.Fatal("mapEvent() accepted a historical v2 receipt for a new write") + } +} + +func TestMapEventSanitizesUpstreamRequestID(t *testing.T) { + event := channelScopedEvent(t, "00000000-0000-4000-8000-000000007005") + event.Attempts[0].UpstreamRequestID = "sk-secret-upstream-request" + row, err := mapEvent(redact.New(), event) + if err != nil { + t.Fatalf("mapEvent() error = %v", err) + } + if len(row.AttemptRows) != 1 || strings.Contains(row.AttemptRows[0].UpstreamRequestID, "secret") || + row.AttemptRows[0].UpstreamRequestID != redact.Placeholder { + t.Fatalf("sanitized upstream request ID = %#v", row.AttemptRows) + } +} + +func TestDecodeAttemptRowsAcceptsHistoricalReceiptsButChecksV3Identity(t *testing.T) { + for _, schemaVersion := range []int{1, 2} { + rule := pricing.ReceiptRule{ModelID: "model-a"} + if schemaVersion == 1 { + rule.ScopeKey = "provider:openai" + } + rows := []models.RequestLogAttempt{{ + ChannelID: string(channel.OpenAI), + CredentialID: 9, + UpstreamModel: "model-a", + PricingReceipt: models.JSON(encodeReceiptJSON(t, emptyReceipt(schemaVersion, rule))), + }} + if _, err := decodeAttemptRows(rows); err != nil { + t.Fatalf("decodeAttemptRows(v%d) error = %v", schemaVersion, err) + } + } + + v3 := emptyReceipt(3, pricing.ReceiptRule{ + ChannelID: string(channel.Anthropic), + ModelID: "model-a", + }) + _, err := decodeAttemptRows([]models.RequestLogAttempt{{ + ChannelID: string(channel.OpenAI), + CredentialID: 9, + UpstreamModel: "model-a", + PricingReceipt: models.JSON(encodeReceiptJSON(t, v3)), + }}) + if err == nil || !strings.Contains(err.Error(), "identity") { + t.Fatalf("decodeAttemptRows(v3 mismatch) error = %v", err) + } +} + +func TestUsageAggregationKeepsChannelAndCredentialIdentitiesSeparate(t *testing.T) { + completedAt := time.Date(2026, time.August, 9, 1, 0, 0, 0, time.UTC) + first := aggregationRow("channel-a", completedAt, 7, "shared-model") + first.ChannelID = string(channel.OpenAI) + first.CredentialID = 11 + second := aggregationRow("channel-b", completedAt, 7, "shared-model") + second.ChannelID = string(channel.Anthropic) + second.CredentialID = 11 + third := aggregationRow("credential-b", completedAt, 7, "shared-model") + third.ChannelID = string(channel.OpenAI) + third.CredentialID = 12 + + deltas, err := buildUsageStatDeltas([]models.RequestLog{first, second, third}) + if err != nil { + t.Fatalf("buildUsageStatDeltas() error = %v", err) + } + if len(deltas) != 3 { + t.Fatalf("usage identities = %#v, want three channel/credential identities", deltas) + } + + journals, err := buildUsageAggregationJournals([]models.RequestLog{first}) + if err != nil { + t.Fatalf("buildUsageAggregationJournals() error = %v", err) + } + if len(journals) != 1 || journals[0].ChannelID != first.ChannelID || + journals[0].CredentialID != first.CredentialID { + t.Fatalf("journal attribution = %#v", journals) + } +} + +func TestUsageWriterPersistsSeparateChannelCredentialStats(t *testing.T) { + db := openRequestLogQueryDB(t) + completedAt := time.Date(2026, time.August, 9, 1, 30, 0, 0, time.UTC) + first := aggregationRow("channel-write-a", completedAt, 7, "shared-model") + first.ChannelID = string(channel.OpenAI) + first.CredentialID = 11 + second := aggregationRow("channel-write-b", completedAt, 7, "shared-model") + second.ChannelID = string(channel.Anthropic) + second.CredentialID = 11 + third := aggregationRow("credential-write-b", completedAt, 7, "shared-model") + third.ChannelID = string(channel.OpenAI) + third.CredentialID = 12 + + if err := (&gormBatchWriter{db: db}).WriteBatch(t.Context(), []models.RequestLog{ + first, + second, + third, + }); err != nil { + t.Fatalf("WriteBatch() error = %v", err) + } + var stats []models.UsageStat + if err := db.Order("channel_id ASC").Order("credential_id ASC").Find(&stats).Error; err != nil { + t.Fatalf("query usage stats: %v", err) + } + if len(stats) != 3 { + t.Fatalf("usage stats = %#v, want three exact identities", stats) + } +} + +func TestQueryUsageFiltersAndProjectsChannelCredential(t *testing.T) { + db := openRequestLogQueryDB(t) + start := time.Date(2026, time.August, 9, 2, 0, 0, 0, time.UTC) + first := usageStat(start, 7, "shared-model", 2) + first.ChannelID = string(channel.OpenAI) + first.CredentialID = 11 + second := usageStat(start, 7, "shared-model", 5) + second.ChannelID = string(channel.Anthropic) + second.CredentialID = 12 + createUsageStats(t, db, first, second) + + credentialID := uint(11) + report, err := newRequestLogTestService(db).QueryUsage(t.Context(), UsageQuery{ + FromMS: start.UnixMilli(), + ToMS: start.Add(time.Hour).UnixMilli(), + Granularity: UsageGranularityHour, + ChannelID: channel.OpenAI, + CredentialID: &credentialID, + Limit: 100, + BreakdownOrder: UsageBreakdownOrderRequests, + }) + if err != nil { + t.Fatalf("QueryUsage() error = %v", err) + } + if report.Summary.RequestCount != 2 || len(report.Breakdown) != 1 || + report.Breakdown[0].ChannelID != channel.OpenAI || + report.Breakdown[0].CredentialID != 11 { + t.Fatalf("scoped usage report = %#v", report) + } +} + +func TestAccessScopedUsageBreakdownNeverMergesChannelsOrCredentials(t *testing.T) { + db := openRequestLogQueryDB(t) + start := time.Date(2026, time.August, 9, 2, 30, 0, 0, time.UTC).Truncate(time.Hour) + rows := []models.UsageStat{ + usageStat(start, 7, "shared-model", 2), + usageStat(start, 8, "shared-model", 3), + usageStat(start, 9, "shared-model", 5), + } + for index := range rows { + rows[index].AccessKeyID = 41 + rows[index].ChannelID = string(channel.OpenAI) + rows[index].CredentialID = uint(11 + index) + } + rows[2].ChannelID = string(channel.Anthropic) + createUsageStats(t, db, rows...) + + accessKeyID := uint(41) + report, err := newRequestLogTestService(db).QueryUsage(t.Context(), UsageQuery{ + FromMS: start.UnixMilli(), + ToMS: start.Add(time.Hour).UnixMilli(), + Granularity: UsageGranularityHour, + AccessKeyID: &accessKeyID, + Limit: 100, + BreakdownOrder: UsageBreakdownOrderRequests, + }) + if err != nil { + t.Fatalf("QueryUsage() error = %v", err) + } + if report.BreakdownCount != 3 || len(report.Breakdown) != 3 { + t.Fatalf("access-scoped breakdown = %#v", report) + } + for _, breakdown := range report.Breakdown { + if breakdown.ChannelID == "" || breakdown.CredentialID == 0 { + t.Fatalf("breakdown lost exact identity: %#v", breakdown) + } + } +} + +func TestListFiltersAttemptsByChannelAndCredential(t *testing.T) { + db := openRequestLogQueryDB(t) + completedAt := time.Date(2026, time.August, 9, 3, 0, 0, 0, time.UTC) + first := requestLogQueryRow("channel-list-a", completedAt, 1, "model-a", []Attempt{{ + GroupID: 7, ChannelID: channel.OpenAI, CredentialID: 11, + }}) + first.GroupID = 7 + first.ChannelID = string(channel.OpenAI) + first.CredentialID = 11 + second := requestLogQueryRow("channel-list-b", completedAt.Add(-time.Second), 1, "model-a", []Attempt{{ + GroupID: 7, ChannelID: channel.Anthropic, CredentialID: 12, + }}) + createRequestLogQueryRow(t, db, first) + createRequestLogQueryRow(t, db, second) + + credentialID := uint(11) + page, err := newRequestLogTestService(db).List(context.Background(), ListQuery{ + ChannelID: channel.OpenAI, CredentialID: &credentialID, Limit: 50, + }) + if err != nil { + t.Fatalf("List() error = %v", err) + } + if len(page.Items) != 1 || page.Items[0].RequestID != first.ID || + page.Items[0].ChannelID != channel.OpenAI || page.Items[0].CredentialID != 11 { + t.Fatalf("filtered page = %#v", page) + } +} + +func channelScopedEvent(t *testing.T, requestID string) telemetry.RequestEvent { + t.Helper() + event := testEvent(requestID) + event.Attempts[0].ChannelID = channel.OpenAI + event.Attempts[0].CredentialID = 8 + event.Attempts[0].Operation = execution.OperationChatCompletion + event.Attempts[0].RouteMode = channel.RouteNative + event.Attempts[0].UpstreamRequestID = "upstream-request-1" + event.Attempts[0].DispatchState = execution.DispatchMaybeSent + event.Attempts[0].ResponseStarted = true + event.Attempts[0].Committed = true + event.Usage.ChannelID = channel.OpenAI + event.Usage.CredentialID = 8 + event.Usage.Result = usage.Result{State: usage.StateComplete} + event.Usage.Pricing = telemetry.PricingObservation{ + UpstreamModel: event.UpstreamModel, + CostState: string(pricing.CostStatePriced), + PricingCompleteness: string(pricing.CompletenessComplete), + EstimatedCostNanoUSD: 0, + ReceiptJSON: encodeReceiptJSON(t, emptyReceipt(3, pricing.ReceiptRule{ + ChannelID: string(channel.OpenAI), + ModelID: event.UpstreamModel, + })), + } + return event +} + +func emptyReceipt(schemaVersion int, rule pricing.ReceiptRule) pricing.Receipt { + return pricing.Receipt{ + SchemaVersion: schemaVersion, + Method: pricing.ReceiptMethodUnitRateSum, + MethodVersion: 1, + Currency: "USD", + Rule: rule, + LineItems: []pricing.ReceiptLine{}, + } +} + +func encodeReceiptJSON(t *testing.T, receipt pricing.Receipt) string { + t.Helper() + encoded, err := json.Marshal(receipt) + if err != nil { + t.Fatalf("json.Marshal(receipt) error = %v", err) + } + return string(encoded) +} + +func decodeReceiptJSON(t *testing.T, encoded string) pricing.Receipt { + t.Helper() + var receipt pricing.Receipt + if err := json.Unmarshal([]byte(encoded), &receipt); err != nil { + t.Fatalf("json.Unmarshal(receipt) error = %v", err) + } + return receipt +} diff --git a/internal/requestlog/external_database_integration_test.go b/internal/requestlog/external_database_integration_test.go index 6f17c311..1eefdfc3 100644 --- a/internal/requestlog/external_database_integration_test.go +++ b/internal/requestlog/external_database_integration_test.go @@ -43,8 +43,11 @@ func TestExternalDatabaseRequestLogLifecycle(t *testing.T) { unique := uint64(time.Now().UnixNano()) & 0xffffffffffff requestID := fmt.Sprintf("00000000-0000-4000-8000-%012x", unique) groupID := uint(unique % 1_000_000) + credentialID := groupID + 1 modelID := fmt.Sprintf("external-request-log-%012x", unique) row := aggregationRow(requestID, now.Add(-40*24*time.Hour), groupID, modelID) + row.ChannelID = "openai" + row.CredentialID = credentialID writer := &gormBatchWriter{db: db} if err := writer.WriteBatch(context.Background(), []models.RequestLog{row}); err != nil { t.Fatalf("first WriteBatch() error = %v", err) @@ -55,12 +58,32 @@ func TestExternalDatabaseRequestLogLifecycle(t *testing.T) { assertExternalRequestLogCount(t, db, "id = ?", []any{requestID}, 1) // aggregationRow is deliberately placed on an exact hour, so its completed // timestamp is the usage aggregate bucket as well. - assertExternalUsageStatCount(t, db, row.CompletedAtMS, row.AccessKeyID, groupID, modelID, 1) + assertExternalUsageStatCount( + t, + db, + row.CompletedAtMS, + row.AccessKeyID, + row.ChannelID, + groupID, + credentialID, + modelID, + 1, + ) service := NewService(db, redact.New(), staticRetentionPolicy{days: 1}) service.Sweep(context.Background(), now) assertExternalRequestLogCount(t, db, "id = ?", []any{requestID}, 0) - assertExternalUsageStatCount(t, db, row.CompletedAtMS, row.AccessKeyID, groupID, modelID, 0) + assertExternalUsageStatCount( + t, + db, + row.CompletedAtMS, + row.AccessKeyID, + row.ChannelID, + groupID, + credentialID, + modelID, + 0, + ) var journals int64 if err := db.Model(&models.UsageAggregationJournal{}). Where("request_id = ?", requestID). @@ -88,7 +111,9 @@ func assertExternalUsageStatCount( db *gorm.DB, bucketStartMS int64, accessKeyID uint, + channelID string, groupID uint, + credentialID uint, modelID string, want int64, ) { @@ -96,10 +121,12 @@ func assertExternalUsageStatCount( var count int64 if err := db.Model(&models.UsageStat{}). Where( - "bucket_start_ms = ? AND access_key_id = ? AND group_id = ? AND model = ?", + "bucket_start_ms = ? AND access_key_id = ? AND channel_id = ? AND group_id = ? AND credential_id = ? AND model = ?", bucketStartMS, accessKeyID, + channelID, groupID, + credentialID, modelID, ). Count(&count).Error; err != nil { diff --git a/internal/requestlog/home_statistics_test.go b/internal/requestlog/home_statistics_test.go index 0848d669..814b45ff 100644 --- a/internal/requestlog/home_statistics_test.go +++ b/internal/requestlog/home_statistics_test.go @@ -359,11 +359,11 @@ func homeStatisticsStat( func createHomeStatisticsGroup(t *testing.T, db *gorm.DB, name string) models.Group { t.Helper() row := models.Group{ - Name: name, - UpstreamURL: "https://example.invalid", - Protocols: models.JSON(`[]`), - Models: models.JSON(`[]`), - Enabled: true, + Name: name, + ChannelID: "openai", + Params: models.JSON(`{}`), + Models: models.JSON(`[]`), + Enabled: true, } if err := db.Create(&row).Error; err != nil { t.Fatalf("create Group %q: %v", name, err) diff --git a/internal/requestlog/mapper.go b/internal/requestlog/mapper.go index 3771c44b..6459e50d 100644 --- a/internal/requestlog/mapper.go +++ b/internal/requestlog/mapper.go @@ -37,7 +37,7 @@ func mapEvent( return models.RequestLog{}, fmt.Errorf("map request event usage/pricing: %w", err) } - receipt, err := canonicalPricingReceipt(event.Usage.Pricing) + receipt, err := canonicalPricingReceipt(event.Usage.Pricing, string(event.Usage.ChannelID)) if err != nil { return models.RequestLog{}, fmt.Errorf("map request event pricing receipt: %w", err) } @@ -48,21 +48,28 @@ func mapEvent( attemptReceipt = receipt } attempts = append(attempts, models.RequestLogAttempt{ - RequestID: event.RequestID, - Sequence: attempt.Sequence, - CompletedAtMS: completedAtMS, - GroupID: attempt.GroupID, - GroupName: redactIdentityValue(redactor, attempt.GroupName), - KeyID: attempt.KeyID, - UpstreamModel: redactIdentityValue(redactor, projectModel(attempt.UpstreamModel)), - StatusCode: attempt.StatusCode, - DurationMs: attempt.DurationMs, - FailureCategory: string(attempt.FailureCategory), - Action: string(attempt.Action), - WillRetry: attempt.WillRetry, - ErrorCode: attempt.ErrorCode, - ErrorSummary: sanitizeSummary(redactor, attempt.ErrorSummary), - PricingReceipt: attemptReceipt, + RequestID: event.RequestID, + Sequence: attempt.Sequence, + CompletedAtMS: completedAtMS, + GroupID: attempt.GroupID, + GroupName: redactIdentityValue(redactor, attempt.GroupName), + ChannelID: string(attempt.ChannelID), + CredentialID: attempt.CredentialID, + Operation: string(attempt.Operation), + RouteMode: string(attempt.RouteMode), + UpstreamModel: redactIdentityValue(redactor, projectModel(attempt.UpstreamModel)), + UpstreamRequestID: redactIdentityValue(redactor, projectModel(attempt.UpstreamRequestID)), + DispatchState: string(attempt.DispatchState), + ResponseStarted: attempt.ResponseStarted, + StatusCode: attempt.StatusCode, + DurationMs: attempt.DurationMs, + FailureCategory: string(attempt.FailureCategory), + Action: string(attempt.Action), + WillRetry: attempt.WillRetry, + ErrorCode: attempt.ErrorCode, + ErrorSummary: sanitizeSummary(redactor, attempt.ErrorSummary), + Committed: attempt.Committed, + PricingReceipt: attemptReceipt, }) } @@ -74,6 +81,8 @@ func mapEvent( CompletedAtMS: completedAtMS, AccessKeyID: event.AccessKeyID, GroupID: event.Usage.GroupID, + ChannelID: string(event.Usage.ChannelID), + CredentialID: event.Usage.CredentialID, Protocol: string(event.Protocol), ClientModel: redactIdentityValue(redactor, projectModel(event.ClientModel)), UpstreamModel: redactIdentityValue(redactor, projectModel(event.UpstreamModel)), @@ -146,7 +155,10 @@ func validateModelObservation(event telemetry.RequestEvent) error { return nil } -func canonicalPricingReceipt(observation telemetry.PricingObservation) (models.JSON, error) { +func canonicalPricingReceipt( + observation telemetry.PricingObservation, + channelID string, +) (models.JSON, error) { if observation.ReceiptJSON == "" { return nil, nil } @@ -157,7 +169,10 @@ func canonicalPricingReceipt(observation telemetry.PricingObservation) (models.J if err := pricing.ValidateReceipt(receipt); err != nil { return nil, err } - if receipt.SchemaVersion != 2 || receipt.Rule != (pricing.ReceiptRule{ModelID: observation.UpstreamModel}) || + if receipt.SchemaVersion != 3 || receipt.Rule != (pricing.ReceiptRule{ + ChannelID: channelID, + ModelID: observation.UpstreamModel, + }) || receipt.TotalNanoUSD != observation.EstimatedCostNanoUSD { return nil, fmt.Errorf("receipt does not match frozen pricing observation") } @@ -213,8 +228,15 @@ func validateFrozenObservation(event telemetry.RequestEvent) error { pricingObservation.UpstreamModel == "" { return fmt.Errorf("priced observation requires exact identity") } + for _, attempt := range event.Attempts { + if _, err := validatedAttemptObservation(attempt); err != nil { + return err + } + } - usageBound := event.Usage.GroupID != 0 || event.Usage.KeyID != 0 || + usageCredentialID := event.Usage.CredentialID + usageBound := event.Usage.GroupID != 0 || usageCredentialID != 0 || + event.Usage.ChannelID != "" || event.Usage.AttemptSequence != 0 if !usageBound { if result.State != usage.StateNotApplicable { @@ -225,16 +247,21 @@ func validateFrozenObservation(event telemetry.RequestEvent) error { } return nil } - if event.Usage.GroupID == 0 || event.Usage.KeyID == 0 || event.Usage.AttemptSequence < 1 { + if event.Usage.GroupID == 0 || usageCredentialID == 0 || event.Usage.AttemptSequence < 1 { return fmt.Errorf("partial usage attribution") } + if !validChannelID(string(event.Usage.ChannelID)) { + return fmt.Errorf("invalid usage channel ID") + } matchingAttempts := 0 boundModel := "" for _, attempt := range event.Attempts { + attemptCredentialID, _ := validatedAttemptObservation(attempt) if attempt.Sequence == event.Usage.AttemptSequence && attempt.GroupID == event.Usage.GroupID && - attempt.KeyID == event.Usage.KeyID { + attempt.ChannelID == event.Usage.ChannelID && + attemptCredentialID == usageCredentialID { matchingAttempts++ boundModel = attempt.UpstreamModel } @@ -252,6 +279,26 @@ func validateFrozenObservation(event telemetry.RequestEvent) error { return nil } +func validatedAttemptObservation(attempt telemetry.Attempt) (uint, error) { + credentialID := attempt.CredentialID + if (credentialID == 0) != (attempt.ChannelID == "") { + return 0, fmt.Errorf("invalid attempt %d identity", attempt.Sequence) + } + if !validRawModelOrEmpty(attempt.UpstreamRequestID) { + return 0, fmt.Errorf("invalid attempt %d upstream request ID", attempt.Sequence) + } + if (credentialID != 0 || attempt.ChannelID != "") && + (!validChannelID(string(attempt.ChannelID)) || !attempt.Operation.Valid() || + !attempt.RouteMode.Valid() || !attempt.DispatchState.Valid()) { + return 0, fmt.Errorf("invalid attempt %d execution observation", attempt.Sequence) + } + return credentialID, nil +} + +func validChannelID(channelID string) bool { + return channelID != "" && len(channelID) <= 64 && validRawModel(channelID) +} + func validateFrozenPricingState( usageState usage.State, costState pricing.CostState, diff --git a/internal/requestlog/mapper_test.go b/internal/requestlog/mapper_test.go index 25f2ef4b..8244de89 100644 --- a/internal/requestlog/mapper_test.go +++ b/internal/requestlog/mapper_test.go @@ -38,7 +38,7 @@ func TestMapEventPersistsFrozenUsagePricingAndAttribution(t *testing.T) { }, } event.Usage.AttemptSequence = 1 - event.Usage.KeyID = 8 + event.Usage.CredentialID = 8 event.Usage.Pricing = telemetry.PricingObservation{ UpstreamModel: "upstream-model", CostState: string(pricing.CostStatePriced), @@ -69,7 +69,7 @@ func TestMapEventPersistsFrozenUsagePricingAndAttribution(t *testing.T) { } if len(row.AttemptRows) != 1 || row.AttemptRows[0].GroupID != 7 || - row.AttemptRows[0].KeyID != 8 { + row.AttemptRows[0].CredentialID != 8 { t.Fatalf("attempts = %+v", row.AttemptRows) } } @@ -93,7 +93,7 @@ func TestMapEventPersistsEveryValidFrozenPricingState(t *testing.T) { event := testEvent("state-" + test.name) event.Usage.Result.State = test.usageState event.Usage.AttemptSequence = 1 - event.Usage.KeyID = 8 + event.Usage.CredentialID = 8 event.Usage.Pricing = telemetry.PricingObservation{ UpstreamModel: event.UpstreamModel, CostState: string(test.costState), @@ -114,7 +114,7 @@ func TestMapEventRejectsInvalidFrozenObservationAtomically(t *testing.T) { event := testEvent("invalid-frozen") event.Usage.Result = usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{Output: 1}} event.Usage.AttemptSequence = 1 - event.Usage.KeyID = 8 + event.Usage.CredentialID = 8 event.Usage.Pricing = telemetry.PricingObservation{ UpstreamModel: event.UpstreamModel, CostState: string(pricing.CostStatePriced), PricingCompleteness: string(pricing.CompletenessComplete), diff --git a/internal/requestlog/phase4_v3_test.go b/internal/requestlog/phase4_v3_test.go index 00f34abb..8434fd5f 100644 --- a/internal/requestlog/phase4_v3_test.go +++ b/internal/requestlog/phase4_v3_test.go @@ -4,6 +4,8 @@ import ( "testing" "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/redact" "gpt-load/internal/pricing" "gpt-load/internal/telemetry" @@ -22,10 +24,12 @@ func TestMapEventV3PersistsMillisecondLedgerAndQuotesUpstreamModel(t *testing.T) Status: telemetry.RequestStatusSuccess, StatusCode: 200, Attempts: []telemetry.Attempt{{ - Sequence: 1, GroupID: 73, KeyID: 9, UpstreamModel: "provider-model", + Sequence: 1, GroupID: 73, ChannelID: channel.OpenAI, CredentialID: 9, + Operation: execution.OperationChatCompletion, RouteMode: channel.RouteNative, + UpstreamModel: "provider-model", DispatchState: execution.DispatchMaybeSent, }}, Usage: telemetry.UsageObservation{ - GroupID: 73, KeyID: 9, AttemptSequence: 1, + GroupID: 73, ChannelID: channel.OpenAI, CredentialID: 9, AttemptSequence: 1, Result: usage.Result{ State: usage.StateComplete, Tokens: usage.Tokens{Output: 1_000_000}, diff --git a/internal/requestlog/process_log.go b/internal/requestlog/process_log.go index 1f927654..e050f474 100644 --- a/internal/requestlog/process_log.go +++ b/internal/requestlog/process_log.go @@ -27,7 +27,7 @@ func projectProcessLog( if err != nil { return logrus.InfoLevel, nil, false } - groupID, keyID := attributedAttempt(event) + groupID, channelID, credentialID := attributedAttempt(event) level := logrus.WarnLevel if event.Status == telemetry.RequestStatusSuccess || @@ -54,8 +54,11 @@ func projectProcessLog( if groupID > 0 { fields["group"] = groupID } - if keyID > 0 { - fields["kid"] = keyID + if channelID != "" { + fields["channel"] = channelID + } + if credentialID > 0 { + fields["credential"] = credentialID } if attemptCount := len(event.Attempts); attemptCount != 1 { fields["attempts"] = attemptCount @@ -113,15 +116,17 @@ func processInputTokens(row models.RequestLog) (int64, bool) { return total, true } -func attributedAttempt(event telemetry.RequestEvent) (uint, uint) { +func attributedAttempt(event telemetry.RequestEvent) (uint, string, uint) { + usageCredentialID := event.Usage.CredentialID for _, attempt := range event.Attempts { if attempt.Sequence == event.Usage.AttemptSequence && attempt.GroupID == event.Usage.GroupID && - attempt.KeyID == event.Usage.KeyID { - return attempt.GroupID, attempt.KeyID + attempt.ChannelID == event.Usage.ChannelID && + attempt.CredentialID == usageCredentialID { + return attempt.GroupID, string(attempt.ChannelID), usageCredentialID } } - return 0, 0 + return 0, "", 0 } func (service *Service) logCompletedRequest( diff --git a/internal/requestlog/process_log_test.go b/internal/requestlog/process_log_test.go index 16cc07b3..905af5a0 100644 --- a/internal/requestlog/process_log_test.go +++ b/internal/requestlog/process_log_test.go @@ -54,7 +54,7 @@ func TestProjectProcessLogUsesFrozenPricingAndInputOutputTotals(t *testing.T) { event.ErrorSummary = "provider failed" event.DurationMs = 58_773 event.Usage.AttemptSequence = 1 - event.Usage.KeyID = 8 + event.Usage.CredentialID = 8 event.Usage.Result = usage.Result{ State: usage.StatePartial, Tokens: usage.Tokens{ @@ -77,7 +77,7 @@ func TestProjectProcessLogUsesFrozenPricingAndInputOutputTotals(t *testing.T) { "status": "error", "http": http.StatusBadGateway, "proto": string(protocol.OpenAICompletions), "ak_id": uint(42), "model": "client-model", "up_model": "upstream-model", - "group": uint(7), "kid": uint(8), "duration": "58.8s", + "group": uint(7), "channel": "openai", "credential": uint(8), "duration": "58.8s", "in_tokens": int64(15), "out_tokens": int64(6), "usage": "partial", "cost_usd": "0.000055", "err": "upstream_error", "err_msg": "provider failed", @@ -127,7 +127,7 @@ func TestProjectProcessLogOmitsDefaultAndZeroValueFields(t *testing.T) { want := logrus.Fields{ "event": "request_completed", "status": "success", "proto": "anthropic", "ak_id": uint(42), - "model": "client-model", "group": uint(7), "kid": uint(8), + "model": "client-model", "group": uint(7), "channel": "openai", "credential": uint(8), "duration": "25ms", } assertProcessFields(t, fields, want) @@ -176,14 +176,14 @@ func TestProjectProcessLogOmitsMissingAttribution(t *testing.T) { if !ok { t.Fatal("projection skipped") } - for _, name := range []string{"group", "kid", "group_name", "up_model"} { + for _, name := range []string{"group", "credential", "group_name", "up_model"} { if value, exists := fields[name]; exists { t.Fatalf("%s = %#v, want omitted", name, value) } } } -func TestProjectProcessLogKeepsKeyIDVisibleAfterRedaction(t *testing.T) { +func TestProjectProcessLogKeepsCredentialIDVisibleAfterRedaction(t *testing.T) { event := testEvent("key-id-visible") _, fields, ok := projectProcessLog(redact.New(), event) if !ok { @@ -195,8 +195,8 @@ func TestProjectProcessLogKeepsKeyIDVisibleAfterRedaction(t *testing.T) { if err := redact.NewHook(redact.New()).Fire(entry); err != nil { t.Fatalf("redaction hook error = %v", err) } - if got := entry.Data["kid"]; got != uint(8) { - t.Fatalf("kid = %#v, want 8", got) + if got := entry.Data["credential"]; got != uint(8) { + t.Fatalf("credential = %#v, want 8", got) } } diff --git a/internal/requestlog/query.go b/internal/requestlog/query.go index ccf37af9..a719af4d 100644 --- a/internal/requestlog/query.go +++ b/internal/requestlog/query.go @@ -8,6 +8,8 @@ import ( "gorm.io/gorm" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/pricing" "gpt-load/internal/protocol" "gpt-load/internal/reasoning" @@ -148,15 +150,19 @@ func applyNullableRange[T int | int64]( } func applyAttemptFilters(query *gorm.DB, input ListQuery) *gorm.DB { - conditions := make([]string, 0, 6) - arguments := make([]any, 0, 6) + conditions := make([]string, 0, 8) + arguments := make([]any, 0, 8) if input.GroupID != nil { conditions = append(conditions, "attempt.group_id = ?") arguments = append(arguments, *input.GroupID) } - if input.UpstreamKeyID != nil { - conditions = append(conditions, "attempt.key_id = ?") - arguments = append(arguments, *input.UpstreamKeyID) + if input.ChannelID != "" { + conditions = append(conditions, "attempt.channel_id = ?") + arguments = append(arguments, input.ChannelID) + } + if input.CredentialID != nil { + conditions = append(conditions, "attempt.credential_id = ?") + arguments = append(arguments, *input.CredentialID) } if input.UpstreamModel != "" { conditions = append(conditions, "attempt.upstream_model = ?") @@ -233,22 +239,35 @@ func decodeAttemptRows(rows []models.RequestLogAttempt) ([]Attempt, error) { if err := pricing.ValidateReceipt(decoded); err != nil { return nil, fmt.Errorf("decode request log pricing receipt: %w", err) } + if decoded.SchemaVersion == 3 && decoded.Rule != (pricing.ReceiptRule{ + ChannelID: row.ChannelID, + ModelID: row.UpstreamModel, + }) { + return nil, fmt.Errorf("decode request log pricing receipt: v3 identity does not match attempt") + } receipt = &decoded } attempts = append(attempts, Attempt{ - Sequence: row.Sequence, - GroupID: row.GroupID, - GroupName: row.GroupName, - KeyID: row.KeyID, - UpstreamModel: row.UpstreamModel, - StatusCode: row.StatusCode, - DurationMs: row.DurationMs, - FailureCategory: telemetry.FailureCategory(row.FailureCategory), - Action: telemetry.Action(row.Action), - WillRetry: row.WillRetry, - ErrorCode: row.ErrorCode, - ErrorSummary: row.ErrorSummary, - PricingReceipt: receipt, + Sequence: row.Sequence, + GroupID: row.GroupID, + GroupName: row.GroupName, + ChannelID: channel.ID(row.ChannelID), + CredentialID: row.CredentialID, + Operation: execution.Operation(row.Operation), + RouteMode: channel.RouteMode(row.RouteMode), + UpstreamModel: row.UpstreamModel, + UpstreamRequestID: row.UpstreamRequestID, + DispatchState: execution.DispatchState(row.DispatchState), + ResponseStarted: row.ResponseStarted, + StatusCode: row.StatusCode, + DurationMs: row.DurationMs, + FailureCategory: telemetry.FailureCategory(row.FailureCategory), + Action: telemetry.Action(row.Action), + WillRetry: row.WillRetry, + ErrorCode: row.ErrorCode, + ErrorSummary: row.ErrorSummary, + Committed: row.Committed, + PricingReceipt: receipt, }) } return attempts, nil @@ -288,6 +307,8 @@ func decodeRequestLogRows(rows []models.RequestLog) ([]Record, error) { }, Attempts: nil, GroupID: row.GroupID, + ChannelID: channel.ID(row.ChannelID), + CredentialID: row.CredentialID, UsageState: usage.State(row.UsageState), CostState: pricing.CostState(row.CostState), PricingCompleteness: pricing.Completeness(row.PricingCompleteness), diff --git a/internal/requestlog/query_test.go b/internal/requestlog/query_test.go index 9d30f6a9..fa985065 100644 --- a/internal/requestlog/query_test.go +++ b/internal/requestlog/query_test.go @@ -10,6 +10,8 @@ import ( "gorm.io/gorm" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/redact" "gpt-load/internal/pricing" "gpt-load/internal/protocol" @@ -76,7 +78,7 @@ func TestDecodeAttemptRowsExposesOnlyNormalizedSafeFields(t *testing.T) { Sequence: 1, GroupID: 7, GroupName: "Primary", - KeyID: 11, + CredentialID: 11, UpstreamModel: "model", StatusCode: 200, DurationMs: 10, @@ -98,6 +100,10 @@ func TestDecodeAttemptRowsExposesOnlyNormalizedSafeFields(t *testing.T) { t.Fatalf("decoded attempt retains historical key material %q: %s", forbidden, encoded) } } + if bytes.Contains(encoded, []byte(`"key_id"`)) || + !bytes.Contains(encoded, []byte(`"credential_id":11`)) { + t.Fatalf("decoded attempt identity JSON = %s, want credential_id only", encoded) + } } func TestDecodeRequestLogRowsRejectsInvalidUsageCostValues(t *testing.T) { @@ -431,13 +437,13 @@ func TestServiceListGroupFilterUsesAnyAttemptWhileAttributionUsesFinalGroup(t *t db := openRequestLogQueryDB(t) event := testEvent("00000000-0000-4000-8000-000000000210") event.Attempts = []telemetry.Attempt{ - {Sequence: 1, GroupID: 12, GroupName: "retry", KeyID: 1, UpstreamModel: "retry-model", FailureCategory: telemetry.FailureCategoryRateLimited, Action: telemetry.ActionRetry, WillRetry: true}, - {Sequence: 2, GroupID: 13, GroupName: "final", KeyID: 2, UpstreamModel: "final-model", FailureCategory: telemetry.FailureCategoryOK, Action: telemetry.ActionTerminate}, + {Sequence: 1, GroupID: 12, GroupName: "retry", ChannelID: channel.OpenAI, CredentialID: 1, Operation: execution.OperationChatCompletion, RouteMode: channel.RouteNative, UpstreamModel: "retry-model", DispatchState: execution.DispatchMaybeSent, FailureCategory: telemetry.FailureCategoryRateLimited, Action: telemetry.ActionRetry, WillRetry: true}, + {Sequence: 2, GroupID: 13, GroupName: "final", ChannelID: channel.OpenAI, CredentialID: 2, Operation: execution.OperationChatCompletion, RouteMode: channel.RouteNative, UpstreamModel: "final-model", DispatchState: execution.DispatchMaybeSent, FailureCategory: telemetry.FailureCategoryOK, Action: telemetry.ActionTerminate}, } event.UpstreamModel = "final-model" event.UpstreamReportedModel = "final-model" event.Usage = telemetry.UsageObservation{ - GroupID: 13, KeyID: 2, AttemptSequence: 2, + GroupID: 13, ChannelID: channel.OpenAI, CredentialID: 2, AttemptSequence: 2, Result: usage.Result{ State: usage.StateNotApplicable, }, @@ -476,12 +482,12 @@ func TestServiceListAttemptFiltersMustMatchTheSameAttempt(t *testing.T) { "client-model", []Attempt{ { - Sequence: 1, GroupID: 12, GroupName: "retry", KeyID: 1, + Sequence: 1, GroupID: 12, GroupName: "retry", CredentialID: 1, StatusCode: 429, FailureCategory: telemetry.FailureCategoryRateLimited, Action: telemetry.ActionRetry, WillRetry: true, }, { - Sequence: 2, GroupID: 13, GroupName: "final", KeyID: 2, + Sequence: 2, GroupID: 13, GroupName: "final", CredentialID: 2, StatusCode: 200, FailureCategory: telemetry.FailureCategoryOK, Action: telemetry.ActionTerminate, }, @@ -523,13 +529,13 @@ func TestServiceListZeroRetryRangeIncludesRequestsWithoutUpstreamAttempts(t *tes ) single := requestLogQueryRow( "00000000-0000-4000-8000-000000000213", base.Add(time.Second), 71, "single", - []Attempt{{Sequence: 1, GroupID: 12, GroupName: "single", KeyID: 1}}, + []Attempt{{Sequence: 1, GroupID: 12, GroupName: "single", CredentialID: 1}}, ) retried := requestLogQueryRow( "00000000-0000-4000-8000-000000000214", base.Add(2*time.Second), 71, "retried", []Attempt{ - {Sequence: 1, GroupID: 12, GroupName: "retry", KeyID: 1, WillRetry: true}, - {Sequence: 2, GroupID: 12, GroupName: "final", KeyID: 2}, + {Sequence: 1, GroupID: 12, GroupName: "retry", CredentialID: 1, WillRetry: true}, + {Sequence: 2, GroupID: 12, GroupName: "final", CredentialID: 2}, }, ) for _, row := range []models.RequestLog{zero, single, retried} { @@ -636,7 +642,7 @@ func TestServiceListOmitsAttemptsAndDetailLoadsThem(t *testing.T) { CompletedAtMS: row.CompletedAtMS, GroupID: 7, GroupName: "Primary", - KeyID: 9, + CredentialID: 9, StatusCode: 200, DurationMs: 10, FailureCategory: string(telemetry.FailureCategoryOK), @@ -695,9 +701,12 @@ func requestLogQueryRow( if groupID == 0 { groupID = 1 } - keyID := attempt.KeyID - if keyID == 0 { - keyID = 1 + credentialID := attempt.CredentialID + if credentialID == 0 { + credentialID = attempt.CredentialID + } + if credentialID == 0 { + credentialID = 1 } sequence := attempt.Sequence if sequence == 0 { @@ -712,20 +721,27 @@ func requestLogQueryRow( action = telemetry.ActionTerminate } attemptRows = append(attemptRows, models.RequestLogAttempt{ - RequestID: id, - Sequence: sequence, - CompletedAtMS: completedAt.UTC().UnixMilli(), - GroupID: groupID, - GroupName: attempt.GroupName, - KeyID: keyID, - UpstreamModel: attempt.UpstreamModel, - StatusCode: attempt.StatusCode, - DurationMs: attempt.DurationMs, - FailureCategory: string(failureCategory), - Action: string(action), - WillRetry: attempt.WillRetry, - ErrorCode: attempt.ErrorCode, - ErrorSummary: attempt.ErrorSummary, + RequestID: id, + Sequence: sequence, + CompletedAtMS: completedAt.UTC().UnixMilli(), + GroupID: groupID, + GroupName: attempt.GroupName, + ChannelID: string(attempt.ChannelID), + CredentialID: credentialID, + Operation: string(attempt.Operation), + RouteMode: string(attempt.RouteMode), + UpstreamModel: attempt.UpstreamModel, + UpstreamRequestID: attempt.UpstreamRequestID, + DispatchState: string(attempt.DispatchState), + ResponseStarted: attempt.ResponseStarted, + StatusCode: attempt.StatusCode, + DurationMs: attempt.DurationMs, + FailureCategory: string(failureCategory), + Action: string(action), + WillRetry: attempt.WillRetry, + ErrorCode: attempt.ErrorCode, + ErrorSummary: attempt.ErrorSummary, + Committed: attempt.Committed, }) } return models.RequestLog{ diff --git a/internal/requestlog/service_test.go b/internal/requestlog/service_test.go index 2910ff62..7106dcb7 100644 --- a/internal/requestlog/service_test.go +++ b/internal/requestlog/service_test.go @@ -477,7 +477,7 @@ func TestEmitPersistsGatewayFrozenQuoteWithoutRepricing(t *testing.T) { Tokens: usage.Tokens{Output: 1_000_000}, } eventA.Usage.AttemptSequence = 1 - eventA.Usage.KeyID = 8 + eventA.Usage.CredentialID = 8 eventA.Usage.Pricing = telemetry.PricingObservation{ UpstreamModel: "snapshot-model", CostState: string(pricing.CostStatePriced), PricingCompleteness: string(pricing.CompletenessComplete), @@ -500,7 +500,7 @@ func TestEmitPersistsGatewayFrozenQuoteWithoutRepricing(t *testing.T) { Tokens: usage.Tokens{Output: 1_000_000}, } eventB.Usage.AttemptSequence = 1 - eventB.Usage.KeyID = 8 + eventB.Usage.CredentialID = 8 eventB.Usage.Pricing = telemetry.PricingObservation{ UpstreamModel: "snapshot-model", CostState: string(pricing.CostStatePriced), PricingCompleteness: string(pricing.CompletenessComplete), diff --git a/internal/requestlog/test_helpers_test.go b/internal/requestlog/test_helpers_test.go index 5a51ec54..dbebf77d 100644 --- a/internal/requestlog/test_helpers_test.go +++ b/internal/requestlog/test_helpers_test.go @@ -13,6 +13,8 @@ import ( "gorm.io/gorm" "gorm.io/gorm/logger" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/redact" "gpt-load/internal/pricing" "gpt-load/internal/protocol" @@ -226,8 +228,12 @@ func testEvent(id string) telemetry.RequestEvent { Sequence: 1, GroupID: 7, GroupName: "primary", - KeyID: 8, + ChannelID: channel.OpenAI, + CredentialID: 8, + Operation: execution.OperationChatCompletion, + RouteMode: channel.RouteNative, UpstreamModel: "upstream-model", + DispatchState: execution.DispatchMaybeSent, StatusCode: 200, DurationMs: 20, FailureCategory: telemetry.FailureCategoryOK, @@ -235,7 +241,8 @@ func testEvent(id string) telemetry.RequestEvent { }}, Usage: telemetry.UsageObservation{ GroupID: 7, - KeyID: 8, + ChannelID: channel.OpenAI, + CredentialID: 8, AttemptSequence: 1, Result: usage.Result{ State: usage.StateNotApplicable, diff --git a/internal/requestlog/types.go b/internal/requestlog/types.go index c60da79c..031168f3 100644 --- a/internal/requestlog/types.go +++ b/internal/requestlog/types.go @@ -4,6 +4,8 @@ import ( "errors" "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/pricing" "gpt-load/internal/protocol" "gpt-load/internal/reasoning" @@ -26,19 +28,26 @@ const ( ) type Attempt struct { - Sequence int `json:"sequence"` - GroupID uint `json:"group_id"` - GroupName string `json:"group_name"` - KeyID uint `json:"key_id"` - UpstreamModel string `json:"upstream_model"` - StatusCode int `json:"status_code"` - DurationMs int64 `json:"duration_ms"` - FailureCategory telemetry.FailureCategory `json:"failure_category"` - Action telemetry.Action `json:"action"` - WillRetry bool `json:"will_retry"` - ErrorCode string `json:"error_code"` - ErrorSummary string `json:"error_summary"` - PricingReceipt *pricing.Receipt `json:"pricing_receipt,omitempty"` + Sequence int `json:"sequence"` + GroupID uint `json:"group_id"` + GroupName string `json:"group_name"` + ChannelID channel.ID `json:"channel_id"` + CredentialID uint `json:"credential_id"` + Operation execution.Operation `json:"operation"` + RouteMode channel.RouteMode `json:"route_mode"` + UpstreamModel string `json:"upstream_model"` + UpstreamRequestID string `json:"upstream_request_id"` + DispatchState execution.DispatchState `json:"dispatch_state"` + ResponseStarted bool `json:"response_started"` + StatusCode int `json:"status_code"` + DurationMs int64 `json:"duration_ms"` + FailureCategory telemetry.FailureCategory `json:"failure_category"` + Action telemetry.Action `json:"action"` + WillRetry bool `json:"will_retry"` + ErrorCode string `json:"error_code"` + ErrorSummary string `json:"error_summary"` + Committed bool `json:"committed"` + PricingReceipt *pricing.Receipt `json:"pricing_receipt,omitempty"` } type RetryState string @@ -57,6 +66,7 @@ type ListQuery struct { FromMS *int64 ToMS *int64 GroupID *uint + ChannelID channel.ID ClientModel string UpstreamModel string AccessKeyID *uint @@ -69,7 +79,7 @@ type ListQuery struct { CostState pricing.CostState PricingCompleteness pricing.Completeness CachePresent *bool - UpstreamKeyID *uint + CredentialID *uint AttemptStatusCode *int FailureCategory telemetry.FailureCategory AttemptErrorCode string @@ -117,6 +127,8 @@ type Record struct { Reasoning reasoning.Config Attempts []Attempt GroupID uint + ChannelID channel.ID + CredentialID uint UsageState usage.State CostState pricing.CostState PricingCompleteness pricing.Completeness @@ -155,6 +167,8 @@ type UsageQuery struct { BucketWidthMS int64 AccessKeyID *uint GroupID *uint + ChannelID channel.ID + CredentialID *uint UpstreamModel string Limit int BreakdownOrder UsageBreakdownOrder @@ -184,8 +198,10 @@ type UsageSeriesPoint struct { } type UsageBreakdown struct { - GroupID uint - Model string + GroupID uint + ChannelID channel.ID + CredentialID uint + Model string UsageAggregate } diff --git a/internal/requestlog/usage_query.go b/internal/requestlog/usage_query.go index 1856907e..887c5b6e 100644 --- a/internal/requestlog/usage_query.go +++ b/internal/requestlog/usage_query.go @@ -7,6 +7,7 @@ import ( "gorm.io/gorm" + "gpt-load/internal/channel" "gpt-load/internal/platform/dbtx" "gpt-load/internal/platform/epochms" "gpt-load/internal/pricing" @@ -99,6 +100,12 @@ func validateUsageQuery(input UsageQuery) (int64, error) { if input.AccessKeyID != nil && *input.AccessKeyID == 0 { return 0, fmt.Errorf("query usage: invalid access key scope") } + if input.ChannelID != "" && !validChannelID(string(input.ChannelID)) { + return 0, fmt.Errorf("query usage: invalid channel scope") + } + if input.CredentialID != nil && *input.CredentialID == 0 { + return 0, fmt.Errorf("query usage: invalid credential scope") + } bucketWidthMS := input.BucketWidthMS switch input.Granularity { case UsageGranularityHour: @@ -150,6 +157,12 @@ func usageStatScope(db *gorm.DB, input UsageQuery) *gorm.DB { if input.GroupID != nil { scope = scope.Where("group_id = ?", *input.GroupID) } + if input.ChannelID != "" { + scope = scope.Where("channel_id = ?", input.ChannelID) + } + if input.CredentialID != nil { + scope = scope.Where("credential_id = ?", *input.CredentialID) + } if input.AccessKeyID != nil { scope = scope.Where("access_key_id = ?", *input.AccessKeyID) } @@ -241,9 +254,11 @@ func queryUsageSeries(scope *gorm.DB, bucketWidthMS int64) ([]UsageSeriesPoint, func queryUsageBreakdownCount(scope *gorm.DB, collapseByModel bool) (int64, error) { var grouped *gorm.DB if collapseByModel { - grouped = scope.Select("model").Group("model") + grouped = scope.Select("channel_id, credential_id, model"). + Group("channel_id, credential_id, model") } else { - grouped = scope.Select("group_id, model").Group("group_id, model") + grouped = scope.Select("channel_id, group_id, credential_id, model"). + Group("channel_id, group_id, credential_id, model") } var count int64 if err := scope.Session(&gorm.Session{NewDB: true}). @@ -266,11 +281,11 @@ func queryUsageBreakdown( var rows []usageBreakdownRow var query *gorm.DB if collapseByModel { - query = scope.Select("0 AS group_id, model, " + usageAggregateSelect). - Group("model") + query = scope.Select("0 AS group_id, channel_id, credential_id, model, " + usageAggregateSelect). + Group("channel_id, credential_id, model") } else { - query = scope.Select("group_id, model, " + usageAggregateSelect). - Group("group_id, model") + query = scope.Select("group_id, channel_id, credential_id, model, " + usageAggregateSelect). + Group("channel_id, group_id, credential_id, model") } switch order { case UsageBreakdownOrderRequests: @@ -288,7 +303,9 @@ func queryUsageBreakdown( ) } if err := query. + Order("channel_id ASC"). Order("group_id ASC"). + Order("credential_id ASC"). Order("model ASC"). Limit(limit + 1). Find(&rows).Error; err != nil { @@ -306,6 +323,8 @@ func queryUsageBreakdown( } breakdown = append(breakdown, UsageBreakdown{ GroupID: row.GroupID, + ChannelID: channel.ID(row.ChannelID), + CredentialID: row.CredentialID, Model: row.Model, UsageAggregate: row.UsageAggregate, }) @@ -463,7 +482,9 @@ type usageStatIntegrity struct { } type usageBreakdownRow struct { - GroupID uint - Model string + GroupID uint + ChannelID string + CredentialID uint + Model string UsageAggregate } diff --git a/internal/requestlog/worker.go b/internal/requestlog/worker.go index c73c8c29..008c4a30 100644 --- a/internal/requestlog/worker.go +++ b/internal/requestlog/worker.go @@ -121,7 +121,9 @@ func stageUsageAggregationJournals( type usageStatKey struct { BucketStartMS int64 AccessKeyID uint + ChannelID string GroupID uint + CredentialID uint Model string } @@ -191,7 +193,9 @@ func applyUsageJournalBatch( stat := existingStats[key] stat.BucketStartMS = key.BucketStartMS stat.AccessKeyID = key.AccessKeyID + stat.ChannelID = key.ChannelID stat.GroupID = key.GroupID + stat.CredentialID = key.CredentialID stat.Model = key.Model total, err := checkedUsageStatTotal(stat, deltas[key]) if err != nil { @@ -230,7 +234,9 @@ func usageStatUpsertClause() clause.OnConflict { Columns: []clause.Column{ {Name: "bucket_start_ms"}, {Name: "access_key_id"}, + {Name: "channel_id"}, {Name: "group_id"}, + {Name: "credential_id"}, {Name: "model"}, }, DoUpdates: clause.AssignmentColumns([]string{ @@ -274,7 +280,9 @@ func buildUsageAggregationJournals( RequestID: row.ID, BucketStartMS: key.BucketStartMS, AccessKeyID: key.AccessKeyID, + ChannelID: key.ChannelID, GroupID: key.GroupID, + CredentialID: key.CredentialID, Model: key.Model, RequestCount: delta.RequestCount, SuccessCount: delta.SuccessCount, @@ -304,7 +312,9 @@ func buildUsageJournalDeltas( key := usageStatKey{ BucketStartMS: journal.BucketStartMS, AccessKeyID: journal.AccessKeyID, + ChannelID: journal.ChannelID, GroupID: journal.GroupID, + CredentialID: journal.CredentialID, Model: journal.Model, } delta := deltas[key] @@ -359,7 +369,9 @@ func buildUsageStatDeltas(rows []models.RequestLog) (map[usageStatKey]usageStatD key := usageStatKey{ BucketStartMS: bucketStartMS, AccessKeyID: row.AccessKeyID, + ChannelID: row.ChannelID, GroupID: row.GroupID, + CredentialID: row.CredentialID, Model: row.UpstreamModel, } delta := deltas[key] @@ -506,9 +518,15 @@ func sortedUsageStatKeys(deltas map[usageStatKey]usageStatDelta) []usageStatKey if keys[left].AccessKeyID != keys[right].AccessKeyID { return keys[left].AccessKeyID < keys[right].AccessKeyID } + if keys[left].ChannelID != keys[right].ChannelID { + return keys[left].ChannelID < keys[right].ChannelID + } if keys[left].GroupID != keys[right].GroupID { return keys[left].GroupID < keys[right].GroupID } + if keys[left].CredentialID != keys[right].CredentialID { + return keys[left].CredentialID < keys[right].CredentialID + } return keys[left].Model < keys[right].Model }) return keys @@ -522,19 +540,23 @@ func queryExistingUsageStats( for index, key := range keys { if index == 0 { query = query.Where( - "bucket_start_ms = ? AND access_key_id = ? AND group_id = ? AND model = ?", + "bucket_start_ms = ? AND access_key_id = ? AND channel_id = ? AND group_id = ? AND credential_id = ? AND model = ?", key.BucketStartMS, key.AccessKeyID, + key.ChannelID, key.GroupID, + key.CredentialID, key.Model, ) continue } query = query.Or( - "bucket_start_ms = ? AND access_key_id = ? AND group_id = ? AND model = ?", + "bucket_start_ms = ? AND access_key_id = ? AND channel_id = ? AND group_id = ? AND credential_id = ? AND model = ?", key.BucketStartMS, key.AccessKeyID, + key.ChannelID, key.GroupID, + key.CredentialID, key.Model, ) } @@ -547,7 +569,9 @@ func queryExistingUsageStats( key := usageStatKey{ BucketStartMS: row.BucketStartMS, AccessKeyID: row.AccessKeyID, + ChannelID: row.ChannelID, GroupID: row.GroupID, + CredentialID: row.CredentialID, Model: row.Model, } existing[key] = row diff --git a/internal/requestlog/worker_test.go b/internal/requestlog/worker_test.go index e86ecdae..03da3ff3 100644 --- a/internal/requestlog/worker_test.go +++ b/internal/requestlog/worker_test.go @@ -18,6 +18,7 @@ import ( "gorm.io/gorm" "gorm.io/gorm/logger" + "gpt-load/internal/channel" "gpt-load/internal/platform/redact" "gpt-load/internal/pricing" "gpt-load/internal/storage/models" @@ -1095,7 +1096,7 @@ func TestWorkerCountsDuplicateReplayAsSuccessfulDeliveryWithoutReaggregation(t * event.UpstreamReportedModel = "gpt-4o" event.Attempts[0].UpstreamModel = "gpt-4o" event.Usage = telemetry.UsageObservation{ - GroupID: 14, KeyID: 8, AttemptSequence: 1, + GroupID: 14, ChannelID: channel.OpenAI, CredentialID: 8, AttemptSequence: 1, Result: usage.Result{ State: usage.StateComplete, Tokens: usage.Tokens{UncachedInput: 1_000_000, Output: 1_000_000}, diff --git a/internal/scheduler/channel_scheduler_test.go b/internal/scheduler/channel_scheduler_test.go new file mode 100644 index 00000000..f4ce8981 --- /dev/null +++ b/internal/scheduler/channel_scheduler_test.go @@ -0,0 +1,354 @@ +package scheduler + +import ( + "encoding/json" + "errors" + "math/rand" + "net/http" + "slices" + "strings" + "testing" + "time" + + "gpt-load/internal/channel" + "gpt-load/internal/dialect" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" + "gpt-load/internal/state" +) + +func TestIteratorExhaustsNativeTierBeforeConvertedTier(t *testing.T) { + t.Parallel() + + snapshot := channelSchedulerSnapshot(t) + convertedWeight := 100 + nativeWeight := 1 + converted := snapshot.Groups[1] + converted.WeightManual = &convertedWeight + snapshot.Groups[1] = converted + native := snapshot.Groups[2] + native.WeightManual = &nativeWeight + snapshot.Groups[2] = native + + iterator := New(snapshot, fakeCredentialSource{keys: []state.CredentialMeta{ + {ID: 11, GroupID: 1, WeightAuto: state.DefaultWeight}, + {ID: 21, GroupID: 2, WeightAuto: state.DefaultWeight}, + }}, Query{ + ClientProtocol: protocol.OpenAICompletions, + Operation: execution.OperationChatCompletion, + ExternalModel: modelPointer("public"), + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, + }, rand.New(zeroRandSource{})) + + first, err := iterator.Next() + if err != nil { + t.Fatalf("first Next() error = %v", err) + } + if first.CredentialID != 21 || first.GroupID != 2 || + first.ChannelID != channel.OpenAI || first.RouteMode != channel.RouteNative || + first.ResolvedTarget.ChannelID != channel.OpenAI || first.UpstreamModelID == nil || *first.UpstreamModelID != "native-model" { + t.Fatalf("first Selection = %#v, want native OpenAI credential", first) + } + + second, err := iterator.Next() + if err != nil { + t.Fatalf("second Next() error = %v", err) + } + if second.CredentialID != 11 || second.GroupID != 1 || second.ChannelID != channel.Anthropic || + second.RouteMode != channel.RouteConverted || second.UpstreamModelID == nil || *second.UpstreamModelID != "converted-model" { + t.Fatalf("second Selection = %#v, want converted Anthropic credential", second) + } + if _, err := iterator.Next(); !errors.Is(err, ErrExhausted) { + t.Fatalf("third Next() error = %v, want ErrExhausted", err) + } +} + +func TestIteratorSkipGroupAndAllowedCredentialIDsApplyAcrossRouteTiers(t *testing.T) { + t.Parallel() + + allowed := map[uint]struct{}{11: {}, 21: {}} + iterator := New(channelSchedulerSnapshot(t), fakeCredentialSource{keys: []state.CredentialMeta{ + {ID: 11, GroupID: 1}, {ID: 12, GroupID: 1}, {ID: 21, GroupID: 2}, + }}, Query{ + ClientProtocol: protocol.OpenAICompletions, + Operation: execution.OperationChatCompletion, + ExternalModel: modelPointer("public"), + AllowedCredentialIDs: allowed, + }, rand.New(zeroRandSource{})) + delete(allowed, 11) + allowed[12] = struct{}{} + iterator.SkipGroup(2) + + selection, err := iterator.Next() + if err != nil || selection.CredentialID != 11 || selection.GroupID != 1 || selection.RouteMode != channel.RouteConverted { + t.Fatalf("Next() = (%#v, %v), want frozen allowed converted credential 11", selection, err) + } + if _, err := iterator.Next(); !errors.Is(err, ErrExhausted) { + t.Fatalf("second Next() error = %v, want ErrExhausted", err) + } +} + +func TestIteratorRoutesModelLessResponsesResourceOperation(t *testing.T) { + t.Parallel() + + snapshot, err := state.Compile(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []state.GroupConfig{{ + ID: 7, Name: "responses", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Enabled: true, + }}, + }) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + selection, err := New(snapshot, fakeCredentialSource{keys: []state.CredentialMeta{{ID: 71, GroupID: 7}}}, Query{ + ClientProtocol: protocol.OpenAIResponses, + Operation: execution.OperationResponsesRetrieve, + ExternalModel: nil, + }, rand.New(zeroRandSource{})).Next() + if err != nil { + t.Fatalf("Next() error = %v", err) + } + if selection.CredentialID != 71 || selection.GroupID != 7 || selection.UpstreamModelID != nil || + selection.RouteMode != channel.RouteNative { + t.Fatalf("Selection = %#v", selection) + } +} + +func TestCandidateGroupIDsForQueryUsesExecutionRoutesAndAccessKeyFilters(t *testing.T) { + t.Parallel() + + snapshot := channelSchedulerSnapshot(t) + query := Query{ + ClientProtocol: protocol.OpenAICompletions, + Operation: execution.OperationChatCompletion, + ExternalModel: modelPointer("public"), + AccessKey: state.AccessKeyView{ + Status: state.AccessKeyStatusActive, + Filters: state.FilterSet{ + Groups: map[uint]struct{}{1: {}, 2: {}}, + Protocols: map[protocol.Protocol]struct{}{protocol.OpenAICompletions: {}}, + Models: map[string]struct{}{"public": {}}, + }, + }, + } + if got := CandidateGroupIDsForQuery(snapshot, query); !slices.Equal(got, []uint{2, 1}) { + t.Fatalf("CandidateGroupIDsForQuery() = %#v, want native then converted [2 1]", got) + } + + query.AccessKey.Filters.Groups = map[uint]struct{}{1: {}} + if got := CandidateGroupIDsForQuery(snapshot, query); !slices.Equal(got, []uint{1}) { + t.Fatalf("CandidateGroupIDsForQuery() with group filter = %#v, want [1]", got) + } + + query.AccessKey.Filters.Protocols = map[protocol.Protocol]struct{}{protocol.Anthropic: {}} + if got := CandidateGroupIDsForQuery(snapshot, query); len(got) != 0 { + t.Fatalf("CandidateGroupIDsForQuery() with protocol filter = %#v, want empty", got) + } +} + +func TestCandidateGroupIDsForQuerySupportsModelLessResourceOperation(t *testing.T) { + t.Parallel() + + snapshot, err := state.Compile(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []state.GroupConfig{ + { + ID: 7, Name: "official", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Enabled: true, + }, + { + ID: 8, Name: "compatible", ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://compatible.example/v1"}`), Enabled: true, + }, + }, + }) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + required := schedulerFeatureSet(t, execution.FeatureNativeResourceSemantics) + got := CandidateGroupIDsForQuery(snapshot, Query{ + ClientProtocol: protocol.OpenAIResponses, + Operation: execution.OperationResponsesRetrieve, + RequiredFeatures: required, + ExternalModel: nil, + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, + }) + if !slices.Equal(got, []uint{7, 8}) { + t.Fatalf("CandidateGroupIDsForQuery() = %#v, want model-less Responses groups [7 8]", got) + } +} + +func TestOpenRouterRoutesResponsesWithReasoningOptOut(t *testing.T) { + t.Parallel() + + snapshot, err := state.Compile(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []state.GroupConfig{{ + ID: 9, Name: "openrouter", ChannelID: channel.OpenRouter, + Params: json.RawMessage(`{}`), Enabled: true, + Models: []state.ModelConfig{{ID: "openai/gpt-5.6-luna", Alias: "gpt-5.6-luna"}}, + }}, + }) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + request := &dialect.ParsedRequest{ + Method: http.MethodPost, + Path: "/v1/responses", + Body: []byte(`{"model":"gpt-5.6-luna","reasoning":{"effort":"none"},"store":false}`), + } + metadata, err := dialect.NewOpenAIResponses().InspectRequest(request) + if err != nil { + t.Fatalf("InspectRequest() error = %v", err) + } + if metadata.RequiredFeatures.Has(execution.FeatureReasoning) { + t.Fatal("reasoning effort none must not require reasoning capability") + } + selection, err := New(snapshot, fakeCredentialSource{keys: []state.CredentialMeta{{ + ID: 91, GroupID: 9, + }}}, Query{ + ClientProtocol: protocol.OpenAIResponses, + Operation: metadata.Operation, + RequiredFeatures: metadata.RequiredFeatures, + ExternalModel: modelPointer("gpt-5.6-luna"), + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, + }, rand.New(zeroRandSource{})).Next() + if err != nil { + t.Fatalf("Next() error = %v", err) + } + if selection.GroupID != 9 || selection.ChannelID != channel.OpenRouter || + selection.RouteMode != channel.RouteConverted || selection.UpstreamModelID == nil || + *selection.UpstreamModelID != "openai/gpt-5.6-luna" { + t.Fatalf("Selection = %#v, want converted OpenRouter target", selection) + } +} + +func TestCandidateGroupIDsForQueryRejectsUnsupportedCapability(t *testing.T) { + t.Parallel() + + got := CandidateGroupIDsForQuery(channelSchedulerSnapshot(t), Query{ + ClientProtocol: protocol.OpenAICompletions, + Operation: execution.OperationChatCompletion, + RequiredFeatures: schedulerFeatureSet(t, execution.FeatureReasoning), + ExternalModel: modelPointer("converted-only"), + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, + }) + if len(got) != 0 { + t.Fatalf("CandidateGroupIDsForQuery() = %#v, want empty for unsupported capability", got) + } +} + +func TestCapabilityUnsupportedIsStableAndInspectionIsNeutral(t *testing.T) { + t.Parallel() + + snapshot := channelSchedulerSnapshot(t) + required := schedulerFeatureSet(t, execution.FeatureReasoning) + query := Query{ + ClientProtocol: protocol.OpenAICompletions, + Operation: execution.OperationChatCompletion, + RequiredFeatures: required, + ExternalModel: modelPointer("converted-only"), + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, + } + iterator := New(snapshot, fakeCredentialSource{keys: []state.CredentialMeta{{ID: 11, GroupID: 1}}}, query, rand.New(zeroRandSource{})) + if iterator.StaticReason() != ReasonCapabilityUnsupported { + t.Fatalf("StaticReason() = %q, want %q", iterator.StaticReason(), ReasonCapabilityUnsupported) + } + if _, err := iterator.Next(); !errors.Is(err, ErrExhausted) { + t.Fatalf("Next() error = %v, want ErrExhausted", err) + } + + inspection, err := Inspect(snapshot, []state.CredentialRuntimeView{{ + ID: 11, GroupID: 1, Status: state.CredentialStatusActive, WeightAuto: state.DefaultWeight, + }}, query, time.Unix(100, 0)) + if err != nil { + t.Fatalf("Inspect() error = %v", err) + } + if inspection.Routable || inspection.Reason != ReasonCapabilityUnsupported || + inspection.ClientProtocol != protocol.OpenAICompletions || inspection.Operation != execution.OperationChatCompletion || + len(inspection.RequiredFeatures) != 1 || inspection.RequiredFeatures[0] != execution.FeatureReasoning || len(inspection.Groups) != 1 { + t.Fatalf("Inspection = %#v", inspection) + } + group := inspection.Groups[0] + if group.ChannelID != channel.Anthropic || group.RouteMode != channel.RouteConverted || group.CapabilitySupported || + group.Reason != ReasonCapabilityUnsupported || len(group.Credentials) != 0 { + t.Fatalf("GroupInspection = %#v", group) + } + encoded, err := json.Marshal(inspection) + if err != nil { + t.Fatalf("json.Marshal(Inspection) error = %v", err) + } + implementationName := "bif" + "rost" + if strings.Contains(string(encoded), "target_config") || strings.Contains(string(encoded), "base_url") || + strings.Contains(strings.ToLower(string(encoded)), implementationName) { + t.Fatalf("Inspection leaked execution implementation details: %s", encoded) + } +} + +func TestInspectionExplainsAllowedCredentialScope(t *testing.T) { + t.Parallel() + + inspection, err := Inspect(channelSchedulerSnapshot(t), []CredentialRuntimeView{ + {ID: 11, GroupID: 1, Status: state.CredentialStatusActive}, + {ID: 12, GroupID: 1, Status: state.CredentialStatusActive}, + {ID: 21, GroupID: 2, Status: state.CredentialStatusActive}, + }, Query{ + ClientProtocol: protocol.OpenAICompletions, + Operation: execution.OperationChatCompletion, + ExternalModel: modelPointer("public"), + AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, + AllowedCredentialIDs: map[uint]struct{}{11: {}}, + }, time.Unix(100, 0)) + if err != nil { + t.Fatalf("Inspect() error = %v", err) + } + if !inspection.Routable || len(inspection.Groups) != 2 { + t.Fatalf("Inspection = %#v", inspection) + } + native, converted := inspection.Groups[0], inspection.Groups[1] + if native.RouteMode != channel.RouteNative || native.Routable || + native.Reason != ReasonNoAvailableCredential || len(native.Credentials) != 1 || + native.Credentials[0].Reason != ReasonCredentialNotAllowed { + t.Fatalf("native GroupInspection = %#v", native) + } + if converted.RouteMode != channel.RouteConverted || !converted.Routable || len(converted.Credentials) != 2 || + !converted.Credentials[0].Available || converted.Credentials[0].CredentialID != 11 || + converted.Credentials[1].Reason != ReasonCredentialNotAllowed { + t.Fatalf("converted GroupInspection = %#v", converted) + } +} + +func channelSchedulerSnapshot(t *testing.T) *state.ConfigSnapshot { + t.Helper() + snapshot, err := state.Compile(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []state.GroupConfig{ + { + ID: 1, Name: "converted", ChannelID: channel.Anthropic, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{ + {ID: "converted-model", Alias: "public"}, + {ID: "converted-only", Alias: "converted-only"}, + }, + Enabled: true, + }, + { + ID: 2, Name: "native", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "native-model", Alias: "public"}}, Enabled: true, + }, + }, + }) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + return snapshot +} + +func schedulerFeatureSet(t *testing.T, features ...execution.Feature) execution.FeatureSet { + t.Helper() + set, err := execution.NewFeatureSet(features...) + if err != nil { + t.Fatalf("NewFeatureSet() error = %v", err) + } + return set +} diff --git a/internal/scheduler/doc.go b/internal/scheduler/doc.go index 45a78fca..cc49bff7 100644 --- a/internal/scheduler/doc.go +++ b/internal/scheduler/doc.go @@ -1,2 +1,2 @@ -// Package scheduler contains the pure routing and upstream-key selection domain. +// Package scheduler contains the pure channel routing and credential selection domain. package scheduler diff --git a/internal/scheduler/inspect.go b/internal/scheduler/inspect.go index 68eb30bf..85552dd6 100644 --- a/internal/scheduler/inspect.go +++ b/internal/scheduler/inspect.go @@ -6,6 +6,8 @@ import ( "sort" "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/protocol" "gpt-load/internal/state" ) @@ -19,38 +21,47 @@ const ( ReasonProtocolFiltered ReasonCode = "protocol_filtered" ReasonModelFiltered ReasonCode = "model_filtered" ReasonModelRequiredByFilter ReasonCode = "model_required_by_filter" + ReasonCapabilityUnsupported ReasonCode = "capability_unsupported" ReasonNoRouteTarget ReasonCode = "no_route_target" ReasonGroupDisabled ReasonCode = "group_disabled" ReasonGroupFiltered ReasonCode = "group_filtered" ReasonNoAvailableGroup ReasonCode = "no_available_group" - ReasonNoKeys ReasonCode = "no_keys" + ReasonNoCredentials ReasonCode = "no_credentials" ReasonGroupWeightZero ReasonCode = "group_weight_zero" - ReasonKeyDisabled ReasonCode = "key_disabled" - ReasonKeyBlacklisted ReasonCode = "key_blacklisted" - ReasonKeyCooldown ReasonCode = "key_cooldown" - ReasonKeyWeightZero ReasonCode = "key_weight_zero" - ReasonNoAvailableKey ReasonCode = "no_available_key" + ReasonCredentialDisabled ReasonCode = "credential_disabled" + ReasonCredentialBlacklisted ReasonCode = "credential_blacklisted" + ReasonCredentialCooldown ReasonCode = "credential_cooldown" + ReasonCredentialWeightZero ReasonCode = "credential_weight_zero" + ReasonCredentialNotAllowed ReasonCode = "credential_not_allowed" + ReasonNoAvailableCredential ReasonCode = "no_available_credential" ) type Inspection struct { - Routable bool - Reason ReasonCode - Groups []GroupInspection + ClientProtocol protocol.Protocol + Operation execution.Operation + RequiredFeatures []execution.Feature + ExternalModel *string + Routable bool + Reason ReasonCode + Groups []GroupInspection } type GroupInspection struct { - GroupID uint - GroupName string - UpstreamModelID *string - WeightManual *int - Included bool - Routable bool - Reason ReasonCode - Keys []KeyInspection + GroupID uint + GroupName string + ChannelID channel.ID + RouteMode channel.RouteMode + CapabilitySupported bool + UpstreamModelID *string + WeightManual *int + Included bool + Routable bool + Reason ReasonCode + Credentials []CredentialInspection } -type KeyInspection struct { - KeyID uint +type CredentialInspection struct { + CredentialID uint Available bool Reason ReasonCode WeightManual *int @@ -59,16 +70,15 @@ type KeyInspection struct { CooldownUntil time.Time } -type targetDecision struct { - target evaluationTarget - group state.GroupCatalogView - included bool - reason ReasonCode -} +// CredentialRuntimeView is the scheduler's neutral view of runtime health. +type CredentialRuntimeView = state.CredentialRuntimeView -type evaluationTarget struct { - GroupID uint - UpstreamModelID *string +type targetDecision struct { + target state.RouteTarget + group state.GroupCatalogView + capabilityOK bool + included bool + reason ReasonCode } func cloneWeight(weight *int) *int { @@ -81,27 +91,44 @@ func cloneWeight(weight *int) *int { func evaluateTargets( snapshot *state.ConfigSnapshot, - modelIndex map[protocol.Protocol]map[string][]state.RouteTarget, - protocolIndex map[protocol.Protocol][]uint, - query Query, + index state.ExecutionCandidateIndex, + query normalizedQuery, ) ([]targetDecision, ReasonCode, error) { if snapshot == nil { return nil, "", fmt.Errorf("%w: nil ConfigSnapshot", ErrInconsistentSnapshot) } - if len(query.AccessKey.Filters.Protocols) > 0 { - if _, allowed := query.AccessKey.Filters.Protocols[query.Protocol]; !allowed { + if query.accessKey.Status == state.AccessKeyStatusDisabled { + return []targetDecision{}, ReasonAccessKeyDisabled, nil + } + if len(query.accessKey.Filters.Protocols) > 0 { + if _, allowed := query.accessKey.Filters.Protocols[query.clientProtocol]; !allowed { return []targetDecision{}, ReasonProtocolFiltered, nil } } - if len(query.AccessKey.Filters.Models) > 0 { - if query.ExternalModel == nil { + if len(query.accessKey.Filters.Models) > 0 { + if query.externalModel == nil { return []targetDecision{}, ReasonModelRequiredByFilter, nil } - if _, allowed := query.AccessKey.Filters.Models[*query.ExternalModel]; !allowed { + if _, allowed := query.accessKey.Filters.Models[*query.externalModel]; !allowed { return []targetDecision{}, ReasonModelFiltered, nil } } - routes := evaluationTargets(modelIndex, protocolIndex, query) + if !query.clientProtocol.Valid() || !query.operation.Valid() || query.requiredFeatures.Validate() != nil { + return []targetDecision{}, ReasonCapabilityUnsupported, nil + } + byOperation := index[query.clientProtocol] + if len(byOperation) == 0 { + return []targetDecision{}, ReasonNoRouteTarget, nil + } + byModel, operationSupported := byOperation[query.operation] + if !operationSupported { + return []targetDecision{}, ReasonCapabilityUnsupported, nil + } + modelKey := state.NoModelRouteKey + if query.externalModel != nil { + modelKey = *query.externalModel + } + routes := byModel[modelKey] if len(routes) == 0 { return []targetDecision{}, ReasonNoRouteTarget, nil } @@ -109,29 +136,43 @@ func evaluateTargets( decisions := make([]targetDecision, 0, len(routes)) seenGroups := make(map[uint]struct{}, len(routes)) included := 0 - for _, target := range routes { - if _, duplicate := seenGroups[target.GroupID]; duplicate { + for _, route := range routes { + if _, duplicate := seenGroups[route.GroupID]; duplicate { continue } - seenGroups[target.GroupID] = struct{}{} - group, exists := snapshot.GroupCatalog[target.GroupID] + seenGroups[route.GroupID] = struct{}{} + group, exists := snapshot.GroupCatalog[route.GroupID] if !exists { return nil, "", fmt.Errorf( "%w: route target group %d missing from catalog", ErrInconsistentSnapshot, - target.GroupID, + route.GroupID, ) } - decision := targetDecision{target: target, group: group, included: true} + capabilityOK := route.ResolvedTarget.Supports( + query.clientProtocol, + query.operation, + query.requiredFeatures, + ) + decision := targetDecision{ + target: cloneRouteTarget(route), group: group, + capabilityOK: capabilityOK, included: true, + } + groupFiltered := false + if len(query.accessKey.Filters.Groups) > 0 { + _, allowed := query.accessKey.Filters.Groups[route.GroupID] + groupFiltered = !allowed + } switch { case !group.Enabled: decision.included = false decision.reason = ReasonGroupDisabled - case len(query.AccessKey.Filters.Groups) > 0: - if _, allowed := query.AccessKey.Filters.Groups[target.GroupID]; !allowed { - decision.included = false - decision.reason = ReasonGroupFiltered - } + case groupFiltered: + decision.included = false + decision.reason = ReasonGroupFiltered + case !capabilityOK: + decision.included = false + decision.reason = ReasonCapabilityUnsupported } if decision.included { included++ @@ -150,31 +191,6 @@ func evaluateTargets( return decisions, reason, nil } -func evaluationTargets( - modelIndex map[protocol.Protocol]map[string][]state.RouteTarget, - protocolIndex map[protocol.Protocol][]uint, - query Query, -) []evaluationTarget { - if query.ExternalModel == nil { - groupIDs := protocolIndex[query.Protocol] - result := make([]evaluationTarget, 0, len(groupIDs)) - for _, groupID := range groupIDs { - result = append(result, evaluationTarget{GroupID: groupID}) - } - return result - } - routes := modelIndex[query.Protocol][*query.ExternalModel] - result := make([]evaluationTarget, 0, len(routes)) - for _, route := range routes { - upstreamModelID := route.UpstreamModelID - result = append(result, evaluationTarget{ - GroupID: route.GroupID, - UpstreamModelID: &upstreamModelID, - }) - } - return result -} - func accessKeyAllowsGroup(accessKey state.AccessKeyView, groupID uint) bool { if len(accessKey.Filters.Groups) == 0 { return true @@ -190,54 +206,62 @@ func normalizedAutoWeight(weight int) int { return weight } -func effectiveWeight(groupManual, keyManual *int, keyAuto int) int64 { +func effectiveWeight(groupManual, credentialManual *int, credentialAuto int) int64 { groupWeight := state.DefaultWeight if groupManual != nil { groupWeight = *groupManual } - keyWeight := normalizedAutoWeight(keyAuto) - if keyManual != nil { - keyWeight = *keyManual + credentialWeight := normalizedAutoWeight(credentialAuto) + if credentialManual != nil { + credentialWeight = *credentialManual } - if groupWeight <= 0 || keyWeight <= 0 { + if groupWeight <= 0 || credentialWeight <= 0 { return 0 } - return int64(groupWeight) * int64(keyWeight) + return int64(groupWeight) * int64(credentialWeight) } -func inspectKey( +func inspectCredential( group state.GroupCatalogView, - key state.KeyRuntimeView, + credential CredentialRuntimeView, + allowedCredentialIDs map[uint]struct{}, now time.Time, -) KeyInspection { - result := KeyInspection{ - KeyID: key.ID, WeightManual: cloneWeight(key.WeightManual), - WeightAuto: normalizedAutoWeight(key.WeightAuto), +) CredentialInspection { + result := CredentialInspection{ + CredentialID: credential.ID, + WeightManual: cloneWeight(credential.WeightManual), + WeightAuto: normalizedAutoWeight(credential.WeightAuto), } if group.WeightManual != nil && *group.WeightManual == 0 { result.Reason = ReasonGroupWeightZero return result } - if key.Status != state.KeyStatusActive { - result.Reason = ReasonKeyDisabled + if allowedCredentialIDs != nil { + if _, allowed := allowedCredentialIDs[credential.ID]; !allowed { + result.Reason = ReasonCredentialNotAllowed + return result + } + } + if credential.Status != state.CredentialStatusActive { + result.Reason = ReasonCredentialDisabled return result } - if key.WeightManual != nil && *key.WeightManual == 0 { - result.Reason = ReasonKeyWeightZero + if credential.WeightManual != nil && *credential.WeightManual == 0 { + result.Reason = ReasonCredentialWeightZero return result } - switch key.RuntimeState(now) { - case state.KeyRuntimeBlacklisted: - result.Reason = ReasonKeyBlacklisted - case state.KeyRuntimeCooldown: - result.Reason = ReasonKeyCooldown - result.CooldownUntil = key.CooldownUntil + switch credential.RuntimeState(now) { + case state.CredentialRuntimeBlacklisted: + result.Reason = ReasonCredentialBlacklisted + case state.CredentialRuntimeCooldown: + result.Reason = ReasonCredentialCooldown + result.CooldownUntil = credential.CooldownUntil default: result.Available = true result.EffectiveWeight = effectiveWeight( group.WeightManual, - key.WeightManual, - key.WeightAuto, + credential.WeightManual, + credential.WeightAuto, ) } return result @@ -245,44 +269,50 @@ func inspectKey( func Inspect( snapshot *state.ConfigSnapshot, - keys []state.KeyRuntimeView, + credentials []CredentialRuntimeView, query Query, now time.Time, ) (Inspection, error) { - result := Inspection{Groups: []GroupInspection{}} + normalized := normalizeQuery(query) + result := Inspection{ + ClientProtocol: normalized.clientProtocol, + Operation: normalized.operation, + RequiredFeatures: normalized.requiredFeatures.Features(), + ExternalModel: cloneString(normalized.externalModel), + Groups: []GroupInspection{}, + } if snapshot == nil { return Inspection{}, fmt.Errorf("%w: nil ConfigSnapshot", ErrInconsistentSnapshot) } - if query.AccessKey.Status == state.AccessKeyStatusDisabled { + if normalized.accessKey.Status == state.AccessKeyStatusDisabled { result.Reason = ReasonAccessKeyDisabled return result, nil } - keysByGroup := make(map[uint][]state.KeyRuntimeView) - for _, key := range keys { - if _, exists := snapshot.GroupCatalog[key.GroupID]; !exists { + credentialsByGroup := make(map[uint][]CredentialRuntimeView) + for _, credential := range credentials { + if _, exists := snapshot.GroupCatalog[credential.GroupID]; !exists { return Inspection{}, fmt.Errorf( - "%w: Registry key %d group %d missing from catalog", + "%w: registry credential %d group %d missing from catalog", ErrInconsistentSnapshot, - key.ID, - key.GroupID, + credential.ID, + credential.GroupID, ) } - cloned := key - cloned.WeightManual = cloneWeight(key.WeightManual) - keysByGroup[key.GroupID] = append(keysByGroup[key.GroupID], cloned) + cloned := credential + cloned.WeightManual = cloneWeight(credential.WeightManual) + credentialsByGroup[credential.GroupID] = append(credentialsByGroup[credential.GroupID], cloned) } - for groupID := range keysByGroup { - sort.Slice(keysByGroup[groupID], func(i, j int) bool { - return keysByGroup[groupID][i].ID < keysByGroup[groupID][j].ID + for groupID := range credentialsByGroup { + sort.Slice(credentialsByGroup[groupID], func(i, j int) bool { + return credentialsByGroup[groupID][i].ID < credentialsByGroup[groupID][j].ID }) } decisions, staticReason, err := evaluateTargets( snapshot, - snapshot.RouteCatalog, - snapshot.ProtocolRouteCatalog, - query, + snapshot.ExecutionRouteCatalog, + normalized, ) if err != nil { return Inspection{}, err @@ -290,32 +320,40 @@ func Inspect( for _, decision := range decisions { groupResult := GroupInspection{ GroupID: decision.group.ID, GroupName: decision.group.Name, - UpstreamModelID: cloneString(decision.target.UpstreamModelID), - WeightManual: cloneWeight(decision.group.WeightManual), - Included: decision.included, Reason: decision.reason, - Keys: []KeyInspection{}, + ChannelID: decision.target.ResolvedTarget.ChannelID, + RouteMode: decision.target.Mode, + CapabilitySupported: decision.capabilityOK, + UpstreamModelID: optionalModel(decision.target.UpstreamModelID), + WeightManual: cloneWeight(decision.group.WeightManual), + Included: decision.included, Reason: decision.reason, + Credentials: []CredentialInspection{}, } if !decision.included { result.Groups = append(result.Groups, groupResult) continue } - groupKeys := keysByGroup[decision.group.ID] + groupCredentials := credentialsByGroup[decision.group.ID] groupWeightZero := decision.group.WeightManual != nil && *decision.group.WeightManual == 0 - for _, key := range groupKeys { - keyResult := inspectKey(decision.group, key, now) - if keyResult.Available && keyResult.EffectiveWeight > 0 { + for _, credential := range groupCredentials { + credentialResult := inspectCredential( + decision.group, + credential, + normalized.allowedCredentialIDs, + now, + ) + if credentialResult.Available && credentialResult.EffectiveWeight > 0 { groupResult.Routable = true } - groupResult.Keys = append(groupResult.Keys, keyResult) + groupResult.Credentials = append(groupResult.Credentials, credentialResult) } switch { case groupWeightZero: groupResult.Reason = ReasonGroupWeightZero - case len(groupKeys) == 0: - groupResult.Reason = ReasonNoKeys + case len(groupCredentials) == 0: + groupResult.Reason = ReasonNoCredentials case !groupResult.Routable: - groupResult.Reason = ReasonNoAvailableKey + groupResult.Reason = ReasonNoAvailableCredential } if groupResult.Routable { result.Routable = true @@ -328,7 +366,7 @@ func Inspect( if staticReason != "" { result.Reason = staticReason } else { - result.Reason = ReasonNoAvailableKey + result.Reason = ReasonNoAvailableCredential } return result, nil } diff --git a/internal/scheduler/inspect_test.go b/internal/scheduler/inspect_test.go index 1b79ff91..968c2650 100644 --- a/internal/scheduler/inspect_test.go +++ b/internal/scheduler/inspect_test.go @@ -1,6 +1,7 @@ package scheduler import ( + "encoding/json" "errors" "fmt" "math/rand" @@ -8,6 +9,8 @@ import ( "testing" "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/protocol" "gpt-load/internal/state" ) @@ -20,19 +23,20 @@ func inspectSnapshot(t *testing.T) *state.ConfigSnapshot { t.Helper() zero := 0 snapshot, err := state.Compile(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []state.GroupConfig{ { - ID: 2, Name: "disabled", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 2, Name: "disabled", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "provider-disabled", Alias: "public"}}, Enabled: false, }, { - ID: 1, Name: "active", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 1, Name: "active", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "provider-active", Alias: "public"}}, Enabled: true, }, { - ID: 3, Name: "weight-zero", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + ID: 3, Name: "weight-zero", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []state.ModelConfig{{ID: "provider-zero", Alias: "public"}}, WeightManual: &zero, Enabled: true, }, @@ -54,7 +58,7 @@ func TestInspectAppliesTopLevelReasonPriority(t *testing.T) { { name: "access key disabled", query: Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("public"), + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("public"), AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusDisabled}, }, reason: ReasonAccessKeyDisabled, @@ -62,7 +66,7 @@ func TestInspectAppliesTopLevelReasonPriority(t *testing.T) { { name: "protocol filtered before model", query: Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("public"), + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("public"), AccessKey: state.AccessKeyView{ Status: state.AccessKeyStatusActive, Filters: state.FilterSet{ @@ -76,7 +80,7 @@ func TestInspectAppliesTopLevelReasonPriority(t *testing.T) { { name: "model filtered", query: Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("public"), + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("public"), AccessKey: state.AccessKeyView{ Status: state.AccessKeyStatusActive, Filters: state.FilterSet{Models: map[string]struct{}{"other": {}}}, @@ -87,7 +91,7 @@ func TestInspectAppliesTopLevelReasonPriority(t *testing.T) { { name: "no route target", query: Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("missing"), + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("missing"), AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, }, reason: ReasonNoRouteTarget, @@ -110,7 +114,7 @@ func TestInspectRequiresModelWhenAccessKeyHasModelFilter(t *testing.T) { t.Parallel() got, err := Inspect(inspectSnapshot(t), nil, Query{ - Protocol: protocol.OpenAICompletions, + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: nil, AccessKey: state.AccessKeyView{ Status: state.AccessKeyStatusActive, @@ -136,27 +140,27 @@ func TestInspectExplainsGroupsAndKeysInStableOrder(t *testing.T) { now := inspectNow() snapshot := inspectSnapshot(t) zero := 0 - keys := []state.KeyRuntimeView{ + keys := []state.CredentialRuntimeView{ { - ID: 32, GroupID: 3, Status: state.KeyStatusActive, + ID: 32, GroupID: 3, Status: state.CredentialStatusActive, WeightAuto: state.DefaultWeight, }, { - ID: 12, GroupID: 1, Status: state.KeyStatusActive, + ID: 12, GroupID: 1, Status: state.CredentialStatusActive, WeightAuto: state.DefaultWeight, CooldownUntil: now.Add(time.Minute), }, { - ID: 11, GroupID: 1, Status: state.KeyStatusActive, + ID: 11, GroupID: 1, Status: state.CredentialStatusActive, WeightManual: &zero, WeightAuto: state.DefaultWeight, }, { - ID: 13, GroupID: 1, Status: state.KeyStatusActive, + ID: 13, GroupID: 1, Status: state.CredentialStatusActive, WeightAuto: 40, }, } - before := append([]state.KeyRuntimeView(nil), keys...) + before := append([]state.CredentialRuntimeView(nil), keys...) got, err := Inspect(snapshot, keys, Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("public"), + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("public"), AccessKey: state.AccessKeyView{ Status: state.AccessKeyStatusActive, Filters: state.FilterSet{Groups: map[uint]struct{}{1: {}, 3: {}}}, @@ -176,22 +180,22 @@ func TestInspectExplainsGroupsAndKeysInStableOrder(t *testing.T) { t.Fatalf("group order = %#v", got.Groups) } if got.Groups[1].Included || got.Groups[1].Reason != ReasonGroupDisabled || - len(got.Groups[1].Keys) != 0 { + len(got.Groups[1].Credentials) != 0 { t.Fatalf("disabled group = %#v", got.Groups[1]) } active := got.Groups[0] - if !active.Included || !active.Routable || len(active.Keys) != 3 || - active.Keys[0].KeyID != 11 || active.Keys[1].KeyID != 12 || - active.Keys[2].KeyID != 13 { + if !active.Included || !active.Routable || len(active.Credentials) != 3 || + active.Credentials[0].CredentialID != 11 || active.Credentials[1].CredentialID != 12 || + active.Credentials[2].CredentialID != 13 { t.Fatalf("active group = %#v", active) } - if active.Keys[0].Reason != ReasonKeyWeightZero || - active.Keys[1].Reason != ReasonKeyCooldown || - !active.Keys[2].Available || active.Keys[2].EffectiveWeight != 50*40 { - t.Fatalf("key explanations = %#v", active.Keys) + if active.Credentials[0].Reason != ReasonCredentialWeightZero || + active.Credentials[1].Reason != ReasonCredentialCooldown || + !active.Credentials[2].Available || active.Credentials[2].EffectiveWeight != 50*40 { + t.Fatalf("key explanations = %#v", active.Credentials) } if got.Groups[2].Reason != ReasonGroupWeightZero || - got.Groups[2].Keys[0].Reason != ReasonGroupWeightZero { + got.Groups[2].Credentials[0].Reason != ReasonGroupWeightZero { t.Fatalf("weight-zero group = %#v", got.Groups[2]) } } @@ -201,36 +205,36 @@ func TestInspectEligiblePoolMatchesIteratorInitialWeightedPool(t *testing.T) { snapshot := inspectSnapshot(t) manual := 25 zero := 0 - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{ - {ID: 11, GroupID: 1, Status: state.KeyStatusActive, WeightAuto: 40, EncryptedValue: "one"}, + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{ + {ID: 11, GroupID: 1, Status: state.CredentialStatusActive, WeightAuto: 40, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "one"}, { - ID: 12, GroupID: 1, Status: state.KeyStatusActive, - WeightManual: &manual, WeightAuto: 80, EncryptedValue: "two", + ID: 12, GroupID: 1, Status: state.CredentialStatusActive, + WeightManual: &manual, WeightAuto: 80, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "two", }, { - ID: 13, GroupID: 1, Status: state.KeyStatusActive, - CooldownUntil: now.Add(time.Minute), EncryptedValue: "cooldown", + ID: 13, GroupID: 1, Status: state.CredentialStatusActive, + CooldownUntil: now.Add(time.Minute), Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cooldown", }, { - ID: 14, GroupID: 1, Status: state.KeyStatusActive, - Blacklisted: true, EncryptedValue: "blacklisted", + ID: 14, GroupID: 1, Status: state.CredentialStatusActive, + Blacklisted: true, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "blacklisted", }, { - ID: 15, GroupID: 1, Status: state.KeyStatusDisabled, - EncryptedValue: "disabled", + ID: 15, GroupID: 1, Status: state.CredentialStatusDisabled, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "disabled", }, { - ID: 16, GroupID: 1, Status: state.KeyStatusActive, - WeightManual: &zero, EncryptedValue: "key-zero", + ID: 16, GroupID: 1, Status: state.CredentialStatusActive, + WeightManual: &zero, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "key-zero", }, - {ID: 21, GroupID: 2, Status: state.KeyStatusActive, EncryptedValue: "disabled-group"}, - {ID: 31, GroupID: 3, Status: state.KeyStatusActive, EncryptedValue: "zero-group"}, + {ID: 21, GroupID: 2, Status: state.CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "disabled-group"}, + {ID: 31, GroupID: 3, Status: state.CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "zero-group"}, }); err != nil { t.Fatalf("Replace() error = %v", err) } query := Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("public"), + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("public"), AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, } inspection, err := Inspect(snapshot, registry.Snapshot(), query, now) @@ -239,9 +243,9 @@ func TestInspectEligiblePoolMatchesIteratorInitialWeightedPool(t *testing.T) { } inspectPool := make(map[uint]int64) for _, group := range inspection.Groups { - for _, key := range group.Keys { + for _, key := range group.Credentials { if key.Available && key.EffectiveWeight > 0 { - inspectPool[key.KeyID] = key.EffectiveWeight + inspectPool[key.CredentialID] = key.EffectiveWeight } } } @@ -257,7 +261,7 @@ func TestInspectEligiblePoolMatchesIteratorInitialWeightedPool(t *testing.T) { rand.New(rand.NewSource(1)), func() time.Time { return now }, ) - weighted, _ := iterator.weightedPool(now) + weighted, _ := iterator.weightedPoolForMode(channel.RouteNative, now) iteratorPool := make(map[uint]int64, len(weighted)) for _, key := range weighted { iteratorPool[key.meta.ID] = key.weight @@ -266,10 +270,10 @@ func TestInspectEligiblePoolMatchesIteratorInitialWeightedPool(t *testing.T) { t.Fatalf("Inspector pool = %#v, Iterator pool = %#v", inspectPool, iteratorPool) } for _, group := range inspection.Groups { - for _, key := range group.Keys { + for _, key := range group.Credentials { if key.Reason != "" { - if _, selected := iteratorPool[key.KeyID]; selected { - t.Fatalf("excluded key %d entered Iterator pool", key.KeyID) + if _, selected := iteratorPool[key.CredentialID]; selected { + t.Fatalf("excluded key %d entered Iterator pool", key.CredentialID) } } } @@ -277,10 +281,10 @@ func TestInspectEligiblePoolMatchesIteratorInitialWeightedPool(t *testing.T) { } func TestInspectRejectsCatalogRegistryMismatch(t *testing.T) { - _, err := Inspect(inspectSnapshot(t), []state.KeyRuntimeView{{ - ID: 99, GroupID: 999, Status: state.KeyStatusActive, + _, err := Inspect(inspectSnapshot(t), []state.CredentialRuntimeView{{ + ID: 99, GroupID: 999, Status: state.CredentialStatusActive, }}, Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("public"), + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("public"), AccessKey: state.AccessKeyView{Status: state.AccessKeyStatusActive}, }, inspectNow()) if !errors.Is(err, ErrInconsistentSnapshot) { @@ -290,7 +294,7 @@ func TestInspectRejectsCatalogRegistryMismatch(t *testing.T) { func TestInspectReportsNoKeysForIncludedGroup(t *testing.T) { got, err := Inspect(inspectSnapshot(t), nil, Query{ - Protocol: protocol.OpenAICompletions, + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("public"), AccessKey: state.AccessKeyView{ Status: state.AccessKeyStatusActive, @@ -300,12 +304,12 @@ func TestInspectReportsNoKeysForIncludedGroup(t *testing.T) { if err != nil { t.Fatalf("Inspect() error = %v", err) } - if got.Routable || got.Reason != ReasonNoAvailableKey || + if got.Routable || got.Reason != ReasonNoAvailableCredential || len(got.Groups) != 3 || !got.Groups[0].Included || - got.Groups[0].Reason != ReasonNoKeys || - got.Groups[0].Keys == nil || - len(got.Groups[0].Keys) != 0 { + got.Groups[0].Reason != ReasonNoCredentials || + got.Groups[0].Credentials == nil || + len(got.Groups[0].Credentials) != 0 { t.Fatalf("no-key Inspection = %#v", got) } } @@ -314,27 +318,27 @@ func TestInspectUsesExactKeyReasonPriority(t *testing.T) { now := inspectNow() zero := 0 snapshot := inspectSnapshot(t) - keys := []state.KeyRuntimeView{ + keys := []state.CredentialRuntimeView{ { - ID: 11, GroupID: 1, Status: state.KeyStatusDisabled, + ID: 11, GroupID: 1, Status: state.CredentialStatusDisabled, WeightManual: &zero, Blacklisted: true, CooldownUntil: now.Add(time.Minute), }, { - ID: 12, GroupID: 1, Status: state.KeyStatusActive, + ID: 12, GroupID: 1, Status: state.CredentialStatusActive, WeightManual: &zero, Blacklisted: true, }, { - ID: 13, GroupID: 1, Status: state.KeyStatusActive, + ID: 13, GroupID: 1, Status: state.CredentialStatusActive, Blacklisted: true, CooldownUntil: now.Add(time.Minute), }, { - ID: 14, GroupID: 1, Status: state.KeyStatusActive, + ID: 14, GroupID: 1, Status: state.CredentialStatusActive, CooldownUntil: now.Add(time.Minute), }, } got, err := Inspect(snapshot, keys, Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("public"), + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("public"), AccessKey: state.AccessKeyView{ Status: state.AccessKeyStatusActive, Filters: state.FilterSet{Groups: map[uint]struct{}{1: {}}}, @@ -344,18 +348,18 @@ func TestInspectUsesExactKeyReasonPriority(t *testing.T) { t.Fatalf("Inspect() error = %v", err) } want := []ReasonCode{ - ReasonKeyDisabled, - ReasonKeyWeightZero, - ReasonKeyBlacklisted, - ReasonKeyCooldown, + ReasonCredentialDisabled, + ReasonCredentialWeightZero, + ReasonCredentialBlacklisted, + ReasonCredentialCooldown, } for index, reason := range want { - if got.Groups[0].Keys[index].Reason != reason { - t.Fatalf("key %d reason = %q, want %q", index, got.Groups[0].Keys[index].Reason, reason) + if got.Groups[0].Credentials[index].Reason != reason { + t.Fatalf("key %d reason = %q, want %q", index, got.Groups[0].Credentials[index].Reason, reason) } } - if got.Routable || got.Reason != ReasonNoAvailableKey || - got.Groups[0].Reason != ReasonNoAvailableKey { + if got.Routable || got.Reason != ReasonNoAvailableCredential || + got.Groups[0].Reason != ReasonNoAvailableCredential { t.Fatalf("unavailable Inspection = %#v", got) } } @@ -364,9 +368,9 @@ func TestInspectSummarizesStaticGroupExclusions(t *testing.T) { group := func(id uint, enabled bool) state.GroupConfig { return state.GroupConfig{ ID: id, Name: fmt.Sprintf("group-%d", id), - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: fmt.Sprintf("provider-%d", id), Alias: "public"}}, - Enabled: enabled, + ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: fmt.Sprintf("provider-%d", id), Alias: "public"}}, + Enabled: enabled, } } tests := []struct { @@ -395,12 +399,15 @@ func TestInspectSummarizesStaticGroupExclusions(t *testing.T) { } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - snapshot, err := state.Compile(state.CompileInput{Groups: test.groups}) + snapshot, err := state.Compile(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: test.groups, + }) if err != nil { t.Fatalf("Compile() error = %v", err) } result, err := Inspect(snapshot, nil, Query{ - Protocol: protocol.OpenAICompletions, + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("public"), AccessKey: state.AccessKeyView{ Status: state.AccessKeyStatusActive, diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index 44388643..8b87424a 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -1,4 +1,4 @@ -// Package scheduler selects upstream keys without IO or persistence access. +// Package scheduler selects channel targets and credentials without IO or persistence access. package scheduler import ( @@ -6,76 +6,90 @@ import ( "math/rand" "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/protocol" "gpt-load/internal/state" ) var ErrExhausted = errors.New("scheduler exhausted") -type KeySource interface { - CollectCandidates(groupIDs []uint, excluded func(uint) bool, now time.Time) []state.KeyMeta +type CredentialSource interface { + CollectCredentialCandidates(groupIDs []uint, excluded func(uint) bool, now time.Time) []state.CredentialMeta } type Query struct { - Protocol protocol.Protocol - ExternalModel *string - AccessKey state.AccessKeyView - AllowedKeyIDs map[uint]struct{} + ClientProtocol protocol.Protocol + Operation execution.Operation + RequiredFeatures execution.FeatureSet + ExternalModel *string + AccessKey state.AccessKeyView + AllowedCredentialIDs map[uint]struct{} } type Selection struct { - KeyID uint + CredentialID uint GroupID uint + ChannelID channel.ID + ResolvedTarget channel.ResolvedTarget + RouteMode channel.RouteMode UpstreamModelID *string Group state.GroupView } type candidateTarget struct { - target evaluationTarget + target state.RouteTarget group state.GroupView } -type weightedKey struct { - meta state.KeyMeta +type weightedCredential struct { + meta state.CredentialMeta weight int64 } type Iterator struct { - keys KeySource - random *rand.Rand - targets map[uint]candidateTarget - groupIDs []uint - allowedKeyIDs map[uint]struct{} - tried map[uint]struct{} - skippedGroups map[uint]struct{} - staticReason ReasonCode - now func() time.Time + credentials CredentialSource + random *rand.Rand + targetsByMode map[channel.RouteMode]map[uint]candidateTarget + groupIDsByMode map[channel.RouteMode][]uint + allowedCredentialIDs map[uint]struct{} + tried map[uint]struct{} + skippedGroups map[uint]struct{} + staticReason ReasonCode + now func() time.Time } -func New(snapshot *state.ConfigSnapshot, keys KeySource, query Query, random *rand.Rand) *Iterator { - return newWithClock(snapshot, keys, query, random, time.Now) +type normalizedQuery struct { + clientProtocol protocol.Protocol + operation execution.Operation + requiredFeatures execution.FeatureSet + externalModel *string + accessKey state.AccessKeyView + allowedCredentialIDs map[uint]struct{} } -func CandidateGroupIDs( - snapshot *state.ConfigSnapshot, - selectedProtocol protocol.Protocol, - accessKey state.AccessKeyView, -) []uint { +func New(snapshot *state.ConfigSnapshot, credentials CredentialSource, query Query, random *rand.Rand) *Iterator { + return newWithClock(snapshot, credentials, query, random, time.Now) +} + +// CandidateGroupIDsForQuery returns the frozen credential-capture scope for a +// fully classified execution query. +func CandidateGroupIDsForQuery(snapshot *state.ConfigSnapshot, query Query) []uint { if snapshot == nil { return nil } - groupIDs := make([]uint, 0, len(snapshot.Groups)) - if accessKey.Status == state.AccessKeyStatusDisabled { - return groupIDs + decisions, _, err := evaluateTargets( + snapshot, + snapshot.ExecutionCandidates, + normalizeQuery(query), + ) + if err != nil { + return []uint{} } - if len(accessKey.Filters.Protocols) > 0 { - if _, allowed := accessKey.Filters.Protocols[selectedProtocol]; !allowed { - return groupIDs - } - } - for _, groupID := range snapshot.ProtocolCandidates[selectedProtocol] { - if accessKeyAllowsGroup(accessKey, groupID) { - groupIDs = append(groupIDs, groupID) + groupIDs := make([]uint, 0, len(decisions)) + for _, decision := range decisions { + if decision.included { + groupIDs = append(groupIDs, decision.target.GroupID) } } return groupIDs @@ -83,25 +97,30 @@ func CandidateGroupIDs( func newWithClock( snapshot *state.ConfigSnapshot, - keys KeySource, + credentials CredentialSource, query Query, random *rand.Rand, now func() time.Time, ) *Iterator { iterator := &Iterator{ - keys: keys, - random: random, - targets: make(map[uint]candidateTarget), - allowedKeyIDs: cloneAllowedKeyIDs(query.AllowedKeyIDs), - tried: make(map[uint]struct{}), - skippedGroups: make(map[uint]struct{}), - now: now, + credentials: credentials, + random: random, + targetsByMode: make(map[channel.RouteMode]map[uint]candidateTarget), + groupIDsByMode: make(map[channel.RouteMode][]uint), + allowedCredentialIDs: cloneAllowedCredentialIDs(query), + tried: make(map[uint]struct{}), + skippedGroups: make(map[uint]struct{}), + now: now, } targets, staticReason := filterTargetsWithReason(snapshot, query) iterator.staticReason = staticReason for _, target := range targets { - iterator.targets[target.target.GroupID] = target - iterator.groupIDs = append(iterator.groupIDs, target.target.GroupID) + mode := target.target.Mode + if iterator.targetsByMode[mode] == nil { + iterator.targetsByMode[mode] = make(map[uint]candidateTarget) + } + iterator.targetsByMode[mode][target.target.GroupID] = target + iterator.groupIDsByMode[mode] = append(iterator.groupIDsByMode[mode], target.target.GroupID) } return iterator } @@ -113,13 +132,14 @@ func (iterator *Iterator) StaticReason() ReasonCode { return iterator.staticReason } -func cloneAllowedKeyIDs(source map[uint]struct{}) map[uint]struct{} { +func cloneAllowedCredentialIDs(query Query) map[uint]struct{} { + source := query.AllowedCredentialIDs if source == nil { return nil } cloned := make(map[uint]struct{}, len(source)) - for keyID := range source { - cloned[keyID] = struct{}{} + for credentialID := range source { + cloned[credentialID] = struct{}{} } return cloned } @@ -134,38 +154,42 @@ func (iterator *Iterator) SkipGroup(groupID uint) { iterator.skippedGroups[groupID] = struct{}{} } -func (iterator *Iterator) weightedPool(now time.Time) ([]weightedKey, int64) { - if iterator == nil || iterator.keys == nil || len(iterator.groupIDs) == 0 { +func (iterator *Iterator) weightedPoolForMode(mode channel.RouteMode, now time.Time) ([]weightedCredential, int64) { + if iterator == nil || iterator.credentials == nil { return nil, 0 } - pool := iterator.keys.CollectCandidates(iterator.groupIDs, func(keyID uint) bool { - _, tried := iterator.tried[keyID] + groupIDs := iterator.groupIDsByMode[mode] + if len(groupIDs) == 0 { + return nil, 0 + } + pool := iterator.credentials.CollectCredentialCandidates(groupIDs, func(credentialID uint) bool { + _, tried := iterator.tried[credentialID] return tried }, now) - weighted := make([]weightedKey, 0, len(pool)) + weighted := make([]weightedCredential, 0, len(pool)) var total int64 - for _, key := range pool { - if iterator.allowedKeyIDs != nil { - if _, allowed := iterator.allowedKeyIDs[key.ID]; !allowed { + for _, credential := range pool { + if iterator.allowedCredentialIDs != nil { + if _, allowed := iterator.allowedCredentialIDs[credential.ID]; !allowed { continue } } - if _, skipped := iterator.skippedGroups[key.GroupID]; skipped { + if _, skipped := iterator.skippedGroups[credential.GroupID]; skipped { continue } - target, ok := iterator.targets[key.GroupID] + target, ok := iterator.targetsByMode[mode][credential.GroupID] if !ok { continue } weight := effectiveWeight( target.group.WeightManual, - key.WeightManual, - key.WeightAuto, + credential.WeightManual, + credential.WeightAuto, ) if weight <= 0 { continue } - weighted = append(weighted, weightedKey{meta: key, weight: weight}) + weighted = append(weighted, weightedCredential{meta: credential, weight: weight}) total += weight } return weighted, total @@ -175,28 +199,26 @@ func (iterator *Iterator) Next() (Selection, error) { if iterator == nil || iterator.random == nil || iterator.now == nil { return Selection{}, ErrExhausted } - weighted, total := iterator.weightedPool(iterator.now()) - if total <= 0 { - return Selection{}, ErrExhausted - } - - ticket := iterator.random.Int63n(total) - selected := weighted[len(weighted)-1].meta - for _, candidate := range weighted { - if ticket < candidate.weight { - selected = candidate.meta - break + for _, mode := range []channel.RouteMode{channel.RouteNative, channel.RouteConverted} { + weighted, total := iterator.weightedPoolForMode(mode, iterator.now()) + if total <= 0 { + continue } - ticket -= candidate.weight + + ticket := iterator.random.Int63n(total) + selected := weighted[len(weighted)-1].meta + for _, candidate := range weighted { + if ticket < candidate.weight { + selected = candidate.meta + break + } + ticket -= candidate.weight + } + iterator.tried[selected.ID] = struct{}{} + target := iterator.targetsByMode[mode][selected.GroupID] + return newSelection(selected, target), nil } - iterator.tried[selected.ID] = struct{}{} - target := iterator.targets[selected.GroupID] - return Selection{ - KeyID: selected.ID, - GroupID: selected.GroupID, - UpstreamModelID: cloneString(target.target.UpstreamModelID), - Group: target.group, - }, nil + return Selection{}, ErrExhausted } func filterTargetsWithReason( @@ -208,9 +230,8 @@ func filterTargetsWithReason( } decisions, staticReason, err := evaluateTargets( snapshot, - snapshot.Candidates, - snapshot.ProtocolCandidates, - query, + snapshot.ExecutionCandidates, + normalizeQuery(query), ) if err != nil { return nil, "" @@ -225,13 +246,86 @@ func filterTargetsWithReason( continue } targets = append(targets, candidateTarget{ - target: decision.target, + target: cloneRouteTarget(decision.target), group: group, }) } return targets, staticReason } +func normalizeQuery(query Query) normalizedQuery { + clientProtocol := query.ClientProtocol + operation := query.Operation + if operation == "" { + if clientProtocol == protocol.OpenAIResponses { + if query.ExternalModel == nil { + operation = execution.OperationResponsesRetrieve + } else { + operation = execution.OperationResponsesCreate + } + } else { + operation = execution.OperationChatCompletion + } + } + return normalizedQuery{ + clientProtocol: clientProtocol, + operation: operation, + requiredFeatures: query.RequiredFeatures.Clone(), + externalModel: cloneString(query.ExternalModel), + accessKey: query.AccessKey, + allowedCredentialIDs: cloneAllowedCredentialIDs(query), + } +} + +func newSelection(credential state.CredentialMeta, target candidateTarget) Selection { + upstreamModelID := optionalModel(target.target.UpstreamModelID) + resolvedTarget := target.target.ResolvedTarget + resolvedTarget.TargetConfig = append([]byte(nil), resolvedTarget.TargetConfig...) + return Selection{ + CredentialID: credential.ID, + GroupID: credential.GroupID, + ChannelID: resolvedTarget.ChannelID, + ResolvedTarget: resolvedTarget, + RouteMode: target.target.Mode, + UpstreamModelID: upstreamModelID, + Group: cloneGroupView(target.group), + } +} + +func optionalModel(value string) *string { + if value == "" { + return nil + } + return cloneString(&value) +} + +func cloneRouteTarget(target state.RouteTarget) state.RouteTarget { + target.ResolvedTarget.TargetConfig = append([]byte(nil), target.ResolvedTarget.TargetConfig...) + return target +} + +func cloneGroupView(group state.GroupView) state.GroupView { + group.Params = append([]byte(nil), group.Params...) + group.ClientProtocols = append([]protocol.Protocol(nil), group.ClientProtocols...) + group.Models = append([]state.ModelConfig(nil), group.Models...) + group.WeightManual = cloneWeight(group.WeightManual) + group.HeaderRules.Set = cloneStringMap(group.HeaderRules.Set) + group.HeaderRules.Remove = append([]string(nil), group.HeaderRules.Remove...) + group.ResolvedTarget.TargetConfig = append([]byte(nil), group.ResolvedTarget.TargetConfig...) + return group +} + +func cloneStringMap(source map[string]string) map[string]string { + if source == nil { + return nil + } + cloned := make(map[string]string, len(source)) + for key, value := range source { + cloned[key] = value + } + return cloned +} + func cloneString(value *string) *string { if value == nil { return nil diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go index 0a58a6dd..11d0d115 100644 --- a/internal/scheduler/scheduler_test.go +++ b/internal/scheduler/scheduler_test.go @@ -1,6 +1,7 @@ package scheduler import ( + "encoding/json" "errors" "math/rand" "reflect" @@ -8,6 +9,8 @@ import ( "testing" "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/protocol" "gpt-load/internal/state" ) @@ -67,7 +70,7 @@ func TestFilterTargetsAppliesAccessKeyDimensions(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { targets, _ := filterTargetsWithReason(snapshot, Query{ - Protocol: tt.protocol, + ClientProtocol: tt.protocol, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer(tt.model), AccessKey: state.AccessKeyView{ID: 10, Filters: tt.filters}, }) @@ -88,7 +91,7 @@ func TestFilterTargetsSkipsCandidateWithoutGroupView(t *testing.T) { got, _ := filterTargetsWithReason( snapshot, Query{ - Protocol: protocol.OpenAICompletions, + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o"), }, ) @@ -97,114 +100,20 @@ func TestFilterTargetsSkipsCandidateWithoutGroupView(t *testing.T) { } } -func TestCandidateGroupIDsUsesFrozenAccessKeyProtocolAndFilters(t *testing.T) { +func TestIteratorSelectsProtocolOnlyGroupWithoutModel(t *testing.T) { + t.Parallel() + snapshot, err := state.Compile(state.CompileInput{ - Groups: []state.GroupConfig{ - { - ID: 4, Name: "disabled-openai", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "disabled-model"}}, - Enabled: false, - }, - { - ID: 3, Name: "multi", - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.Anthropic}, - Models: []state.ModelConfig{{ID: "multi-model"}}, - Enabled: true, - }, - { - ID: 2, Name: "anthropic", - Protocols: []protocol.Protocol{protocol.Anthropic}, - Models: []state.ModelConfig{{ID: "anthropic-model"}}, - Enabled: true, - }, - { - ID: 1, Name: "openai", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []state.ModelConfig{{ID: "openai-model"}}, - Enabled: true, - }, - }, + ChannelRegistry: channel.NewRegistry(), + Groups: []state.GroupConfig{{ + ID: 7, Name: "responses", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Enabled: true, + }}, }) if err != nil { t.Fatalf("Compile() error = %v", err) } - - tests := []struct { - name string - selected protocol.Protocol - accessKey state.AccessKeyView - want []uint - }{ - { - name: "enabled protocol groups sorted without model", - selected: protocol.OpenAICompletions, - accessKey: state.AccessKeyView{ - Status: state.AccessKeyStatusActive, - Filters: state.FilterSet{ - Models: map[string]struct{}{"model-known-after-body": {}}, - }, - }, - want: []uint{1, 3}, - }, - { - name: "group filter", - selected: protocol.OpenAICompletions, - accessKey: state.AccessKeyView{ - Status: state.AccessKeyStatusActive, - Filters: state.FilterSet{ - Groups: map[uint]struct{}{2: {}, 3: {}}, - }, - }, - want: []uint{3}, - }, - { - name: "allowed protocol filter", - selected: protocol.Anthropic, - accessKey: state.AccessKeyView{ - Status: state.AccessKeyStatusActive, - Filters: state.FilterSet{ - Protocols: map[protocol.Protocol]struct{}{protocol.Anthropic: {}}, - }, - }, - want: []uint{2, 3}, - }, - { - name: "denied protocol filter", - selected: protocol.OpenAICompletions, - accessKey: state.AccessKeyView{ - Status: state.AccessKeyStatusActive, - Filters: state.FilterSet{ - Protocols: map[protocol.Protocol]struct{}{protocol.Anthropic: {}}, - }, - }, - want: []uint{}, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - got := CandidateGroupIDs(snapshot, test.selected, test.accessKey) - if !reflect.DeepEqual(got, test.want) { - t.Fatalf("CandidateGroupIDs() = %#v, want %#v", got, test.want) - } - }) - } -} - -func TestIteratorSelectsProtocolOnlyGroupWithoutModel(t *testing.T) { - t.Parallel() - - snapshot, err := state.Compile(state.CompileInput{Groups: []state.GroupConfig{{ - ID: 7, - Name: "responses", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Enabled: true, - }}}) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - source := fakeKeySource{keys: []state.KeyMeta{{ + source := fakeCredentialSource{keys: []state.CredentialMeta{{ ID: 71, GroupID: 7, WeightAuto: state.DefaultWeight, }}} @@ -212,8 +121,9 @@ func TestIteratorSelectsProtocolOnlyGroupWithoutModel(t *testing.T) { snapshot, source, Query{ - Protocol: protocol.OpenAIResponses, - ExternalModel: nil, + ClientProtocol: protocol.OpenAIResponses, + Operation: execution.OperationResponsesRetrieve, + ExternalModel: nil, AccessKey: state.AccessKeyView{ Status: state.AccessKeyStatusActive, }, @@ -223,40 +133,14 @@ func TestIteratorSelectsProtocolOnlyGroupWithoutModel(t *testing.T) { if err != nil { t.Fatalf("Next() error = %v", err) } - if selection.GroupID != 7 || selection.KeyID != 71 || + if selection.GroupID != 7 || selection.CredentialID != 71 || selection.UpstreamModelID != nil { t.Fatalf("Selection = %#v, want protocol-only group/key and nil model", selection) } } -func TestCandidateGroupIDsUsesProtocolProjectionWithoutModels(t *testing.T) { - t.Parallel() - - snapshot, err := state.Compile(state.CompileInput{Groups: []state.GroupConfig{{ - ID: 8, - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Enabled: true, - }}}) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - got := CandidateGroupIDs( - snapshot, - protocol.OpenAIResponses, - state.AccessKeyView{ - Status: state.AccessKeyStatusActive, - Filters: state.FilterSet{ - Models: map[string]struct{}{"must-not-apply-before-body": {}}, - }, - }, - ) - if !reflect.DeepEqual(got, []uint{8}) { - t.Fatalf("CandidateGroupIDs() = %#v, want [8]", got) - } -} - -type fakeKeySource struct { - keys []state.KeyMeta +type fakeCredentialSource struct { + keys []state.CredentialMeta } type zeroRandSource struct{} @@ -264,12 +148,12 @@ type zeroRandSource struct{} func (zeroRandSource) Int63() int64 { return 0 } func (zeroRandSource) Seed(int64) {} -func (source fakeKeySource) CollectCandidates(groupIDs []uint, excluded func(uint) bool, _ time.Time) []state.KeyMeta { +func (source fakeCredentialSource) CollectCredentialCandidates(groupIDs []uint, excluded func(uint) bool, _ time.Time) []state.CredentialMeta { allowed := make(map[uint]struct{}, len(groupIDs)) for _, groupID := range groupIDs { allowed[groupID] = struct{}{} } - result := make([]state.KeyMeta, 0, len(source.keys)) + result := make([]state.CredentialMeta, 0, len(source.keys)) for _, key := range source.keys { if _, ok := allowed[key.GroupID]; !ok || (excluded != nil && excluded(key.ID)) { continue @@ -282,15 +166,15 @@ func (source fakeKeySource) CollectCandidates(groupIDs []uint, excluded func(uin func TestIteratorUsesInjectedTimeForCandidateEligibility(t *testing.T) { now := time.Date(2026, time.July, 20, 12, 0, 0, 0, time.UTC) - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 11, GroupID: 1, Status: state.KeyStatusActive, - CooldownUntil: now.Add(time.Second), EncryptedValue: "cipher", + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 11, GroupID: 1, Status: state.CredentialStatusActive, + CooldownUntil: now.Add(time.Second), Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher", }}); err != nil { t.Fatalf("Replace() error = %v", err) } - query := Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")} + query := Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")} cooling := newWithClock(schedulerSnapshot(), registry, query, rand.New(rand.NewSource(1)), func() time.Time { return now }) @@ -302,26 +186,26 @@ func TestIteratorUsesInjectedTimeForCandidateEligibility(t *testing.T) { return now.Add(time.Second) }) selection, err := expired.Next() - if err != nil || selection.KeyID != 11 { + if err != nil || selection.CredentialID != 11 { t.Fatalf("Next() at cooldown boundary = (%#v, %v), want key 11", selection, err) } } func TestIteratorSkipGroupExcludesWholeGroup(t *testing.T) { - source := fakeKeySource{keys: []state.KeyMeta{ + source := fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 11, GroupID: 1}, {ID: 12, GroupID: 1}, {ID: 21, GroupID: 2}, }} iterator := New(schedulerSnapshot(), source, - Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, + Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, rand.New(zeroRandSource{})) first, err := iterator.Next() - if err != nil || first.KeyID != 11 { + if err != nil || first.CredentialID != 11 { t.Fatalf("first Next() = (%#v, %v), want key 11", first, err) } iterator.SkipGroup(1) iterator.SkipGroup(1) second, err := iterator.Next() - if err != nil || second.KeyID != 21 { + if err != nil || second.CredentialID != 21 { t.Fatalf("second Next() = (%#v, %v), want key 21", second, err) } if _, err := iterator.Next(); !errors.Is(err, ErrExhausted) { @@ -330,10 +214,10 @@ func TestIteratorSkipGroupExcludesWholeGroup(t *testing.T) { } func TestIteratorSkipGroupIsRequestLocal(t *testing.T) { - source := fakeKeySource{keys: []state.KeyMeta{ + source := fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 11, GroupID: 1}, {ID: 21, GroupID: 2}, }} - query := Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")} + query := Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")} first := New(schedulerSnapshot(), source, query, rand.New(zeroRandSource{})) first.SkipGroup(1) selection, err := first.Next() @@ -351,9 +235,9 @@ func TestIteratorSkipGroupIgnoresNilReceiverAndZeroID(t *testing.T) { var nilIterator *Iterator nilIterator.SkipGroup(1) - iterator := New(schedulerSnapshot(), fakeKeySource{keys: []state.KeyMeta{ + iterator := New(schedulerSnapshot(), fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 11, GroupID: 1}, {ID: 21, GroupID: 2}, - }}, Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, rand.New(zeroRandSource{})) + }}, Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, rand.New(zeroRandSource{})) iterator.SkipGroup(0) selection, err := iterator.Next() if err != nil || selection.GroupID != 1 { @@ -362,24 +246,24 @@ func TestIteratorSkipGroupIgnoresNilReceiverAndZeroID(t *testing.T) { } func TestIteratorSkipGroupExcludesKeysAddedAfterSkip(t *testing.T) { - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{ - {ID: 11, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "one"}, - {ID: 21, GroupID: 2, Status: state.KeyStatusActive, EncryptedValue: "two"}, + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{ + {ID: 11, GroupID: 1, Status: state.CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "one"}, + {ID: 21, GroupID: 2, Status: state.CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "two"}, }); err != nil { t.Fatalf("Replace() error = %v", err) } iterator := New(schedulerSnapshot(), registry, - Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, + Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, rand.New(zeroRandSource{})) iterator.SkipGroup(1) - if err := registry.ApplyImport(1, []state.KeyEntry{{ - ID: 12, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "new", + if err := registry.ApplyCredentialImport(1, []state.CredentialEntry{{ + ID: 12, GroupID: 1, Status: state.CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "new", }}); err != nil { t.Fatalf("ApplyImport() error = %v", err) } selection, err := iterator.Next() - if err != nil || selection.KeyID != 21 { + if err != nil || selection.CredentialID != 21 { t.Fatalf("Next() = (%#v, %v), want group 2 key 21", selection, err) } if _, err := iterator.Next(); !errors.Is(err, ErrExhausted) { @@ -388,7 +272,7 @@ func TestIteratorSkipGroupExcludesKeysAddedAfterSkip(t *testing.T) { } func TestIteratorNextNeverRepeatsAndExhausts(t *testing.T) { - source := fakeKeySource{keys: []state.KeyMeta{ + source := fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 11, GroupID: 1}, {ID: 12, GroupID: 1}, {ID: 21, GroupID: 2}, @@ -396,7 +280,7 @@ func TestIteratorNextNeverRepeatsAndExhausts(t *testing.T) { iterator := New( schedulerSnapshot(), source, - Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, + Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, rand.New(rand.NewSource(7)), ) @@ -406,10 +290,10 @@ func TestIteratorNextNeverRepeatsAndExhausts(t *testing.T) { if err != nil { t.Fatalf("Next() error = %v", err) } - if _, duplicate := seen[selection.KeyID]; duplicate { - t.Fatalf("key %d selected twice", selection.KeyID) + if _, duplicate := seen[selection.CredentialID]; duplicate { + t.Fatalf("key %d selected twice", selection.CredentialID) } - seen[selection.KeyID] = struct{}{} + seen[selection.CredentialID] = struct{}{} if selection.Group.ID != selection.GroupID || selection.UpstreamModelID == nil || *selection.UpstreamModelID == "" { @@ -431,19 +315,19 @@ func TestIteratorUsesEffectiveWeights(t *testing.T) { group.WeightManual = &lightGroup snapshot.Groups[2] = group - source := fakeKeySource{keys: []state.KeyMeta{ + source := fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 1, GroupID: 1, WeightAuto: 100}, {ID: 2, GroupID: 2, WeightAuto: 100}, }} counts := map[uint]int{} random := rand.New(rand.NewSource(99)) for range 12000 { - iterator := New(snapshot, source, Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, random) + iterator := New(snapshot, source, Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, random) selection, err := iterator.Next() if err != nil { t.Fatalf("Next() error = %v", err) } - counts[selection.KeyID]++ + counts[selection.CredentialID]++ } ratio := float64(counts[1]) / float64(counts[2]) if ratio < 1.85 || ratio > 2.15 { @@ -453,7 +337,7 @@ func TestIteratorUsesEffectiveWeights(t *testing.T) { func TestIteratorUsesKeyWeights(t *testing.T) { manualWeight := 100 - source := fakeKeySource{keys: []state.KeyMeta{ + source := fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 1, GroupID: 1, WeightManual: &manualWeight, WeightAuto: 1}, {ID: 2, GroupID: 1, WeightAuto: 50}, }} @@ -463,14 +347,14 @@ func TestIteratorUsesKeyWeights(t *testing.T) { iterator := New( schedulerSnapshot(), source, - Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, + Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, random, ) selection, err := iterator.Next() if err != nil { t.Fatalf("Next() error = %v", err) } - counts[selection.KeyID]++ + counts[selection.CredentialID]++ } ratio := float64(counts[1]) / float64(counts[2]) if ratio < 1.85 || ratio > 2.15 { @@ -485,14 +369,14 @@ func TestEffectiveWeightCombinesGroupAndKeyWeights(t *testing.T) { tests := []struct { name string group state.GroupView - key state.KeyMeta + key state.CredentialMeta want int64 }{ - {name: "defaults", group: state.GroupView{}, key: state.KeyMeta{}, want: 50 * 50}, - {name: "group and auto key", group: state.GroupView{WeightManual: &groupManual}, key: state.KeyMeta{WeightAuto: 30}, want: 20 * 30}, - {name: "manual key overrides auto", group: state.GroupView{WeightManual: &groupManual}, key: state.KeyMeta{WeightManual: &keyManual, WeightAuto: 90}, want: 20 * 40}, - {name: "zero group", group: state.GroupView{WeightManual: &zero}, key: state.KeyMeta{WeightAuto: 30}}, - {name: "zero key", group: state.GroupView{WeightManual: &groupManual}, key: state.KeyMeta{WeightManual: &zero, WeightAuto: 30}}, + {name: "defaults", group: state.GroupView{}, key: state.CredentialMeta{}, want: 50 * 50}, + {name: "group and auto key", group: state.GroupView{WeightManual: &groupManual}, key: state.CredentialMeta{WeightAuto: 30}, want: 20 * 30}, + {name: "manual key overrides auto", group: state.GroupView{WeightManual: &groupManual}, key: state.CredentialMeta{WeightManual: &keyManual, WeightAuto: 90}, want: 20 * 40}, + {name: "zero group", group: state.GroupView{WeightManual: &zero}, key: state.CredentialMeta{WeightAuto: 30}}, + {name: "zero key", group: state.GroupView{WeightManual: &groupManual}, key: state.CredentialMeta{WeightManual: &zero, WeightAuto: 30}}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { @@ -514,16 +398,16 @@ func TestIteratorExcludesZeroManualWeights(t *testing.T) { group := snapshot.Groups[1] group.WeightManual = &zero snapshot.Groups[1] = group - source := fakeKeySource{keys: []state.KeyMeta{ + source := fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 11, GroupID: 1, WeightAuto: state.DefaultWeight}, {ID: 21, GroupID: 2, WeightAuto: state.DefaultWeight}, }} random := rand.New(rand.NewSource(1)) for range 200 { - iterator := New(snapshot, source, Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, random) + iterator := New(snapshot, source, Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, random) selection, err := iterator.Next() - if err != nil || selection.KeyID != 21 { + if err != nil || selection.CredentialID != 21 { t.Fatalf("Next() = (%#v, %v), want enabled group key 21", selection, err) } } @@ -531,16 +415,16 @@ func TestIteratorExcludesZeroManualWeights(t *testing.T) { t.Run("key", func(t *testing.T) { zero := 0 - source := fakeKeySource{keys: []state.KeyMeta{ + source := fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 11, GroupID: 1, WeightManual: &zero, WeightAuto: state.DefaultWeight}, {ID: 12, GroupID: 1, WeightAuto: state.DefaultWeight}, }} random := rand.New(rand.NewSource(1)) for range 200 { - iterator := New(schedulerSnapshot(), source, Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, random) + iterator := New(schedulerSnapshot(), source, Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, random) selection, err := iterator.Next() - if err != nil || selection.KeyID != 12 { + if err != nil || selection.CredentialID != 12 { t.Fatalf("Next() = (%#v, %v), want enabled key 12", selection, err) } } @@ -558,8 +442,8 @@ func TestIteratorExhaustsWhenEffectiveWeightPoolIsEmpty(t *testing.T) { } iterator := New( snapshot, - fakeKeySource{keys: []state.KeyMeta{{ID: 11, GroupID: 1, WeightAuto: state.DefaultWeight}}}, - Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, + fakeCredentialSource{keys: []state.CredentialMeta{{ID: 11, GroupID: 1, WeightAuto: state.DefaultWeight}}}, + Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, rand.New(rand.NewSource(1)), ) if _, err := iterator.Next(); !errors.Is(err, ErrExhausted) { @@ -571,11 +455,11 @@ func TestIteratorExhaustsWhenEffectiveWeightPoolIsEmpty(t *testing.T) { zero := 0 iterator := New( schedulerSnapshot(), - fakeKeySource{keys: []state.KeyMeta{ + fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 11, GroupID: 1, WeightManual: &zero, WeightAuto: state.DefaultWeight}, {ID: 21, GroupID: 2, WeightManual: &zero, WeightAuto: state.DefaultWeight}, }}, - Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, + Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, rand.New(rand.NewSource(1)), ) if _, err := iterator.Next(); !errors.Is(err, ErrExhausted) { @@ -587,49 +471,49 @@ func TestIteratorExhaustsWhenEffectiveWeightPoolIsEmpty(t *testing.T) { func TestIteratorUsesDefaultWeights(t *testing.T) { iterator := New( schedulerSnapshot(), - fakeKeySource{keys: []state.KeyMeta{{ID: 11, GroupID: 1}}}, - Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, + fakeCredentialSource{keys: []state.CredentialMeta{{ID: 11, GroupID: 1}}}, + Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, rand.New(rand.NewSource(1)), ) selection, err := iterator.Next() - if err != nil || selection.KeyID != 11 { + if err != nil || selection.CredentialID != 11 { t.Fatalf("Next() with default weights = (%#v, %v), want key 11", selection, err) } } func TestIteratorReadsRegistryChangesBetweenNextCalls(t *testing.T) { - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 11, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "cipher-one", + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 11, GroupID: 1, Status: state.CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }}); err != nil { t.Fatalf("Replace() error = %v", err) } iterator := New( schedulerSnapshot(), registry, - Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, + Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, rand.New(rand.NewSource(1)), ) first, err := iterator.Next() - if err != nil || first.KeyID != 11 { + if err != nil || first.CredentialID != 11 { t.Fatalf("first Next() = (%#v, %v)", first, err) } - if err := registry.ApplyImport(1, []state.KeyEntry{{ - ID: 12, GroupID: 1, Status: state.KeyStatusActive, EncryptedValue: "cipher-two", + if err := registry.ApplyCredentialImport(1, []state.CredentialEntry{{ + ID: 12, GroupID: 1, Status: state.CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two", }}); err != nil { t.Fatalf("ApplyImport() error = %v", err) } second, err := iterator.Next() - if err != nil || second.KeyID != 12 { + if err != nil || second.CredentialID != 12 { t.Fatalf("second Next() = (%#v, %v), want newly added key 12", second, err) } } -func TestIteratorRestrictsLiveCandidatesToAllowedKeyIDs(t *testing.T) { - registry := state.NewKeyRegistry() - if err := registry.Replace([]state.KeyEntry{{ - ID: 11, GroupID: 1, Status: state.KeyStatusActive, - EncryptedValue: "cipher-captured", +func TestIteratorRestrictsLiveCandidatesToAllowedCredentialIDs(t *testing.T) { + registry := state.NewCredentialRegistry() + if err := registry.ReplaceCredentials([]state.CredentialEntry{{ + ID: 11, GroupID: 1, Status: state.CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-captured", }}); err != nil { t.Fatalf("Replace() error = %v", err) } @@ -638,23 +522,23 @@ func TestIteratorRestrictsLiveCandidatesToAllowedKeyIDs(t *testing.T) { schedulerSnapshot(), registry, Query{ - Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o"), - AllowedKeyIDs: allowed, + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o"), + AllowedCredentialIDs: allowed, }, rand.New(zeroRandSource{}), ) delete(allowed, 11) allowed[12] = struct{}{} - if err := registry.ApplyImport(1, []state.KeyEntry{{ - ID: 12, GroupID: 1, Status: state.KeyStatusActive, - EncryptedValue: "cipher-imported-after-iterator", + if err := registry.ApplyCredentialImport(1, []state.CredentialEntry{{ + ID: 12, GroupID: 1, Status: state.CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-imported-after-iterator", }}); err != nil { t.Fatalf("ApplyImport() error = %v", err) } first, err := iterator.Next() - if err != nil || first.KeyID != 11 { + if err != nil || first.CredentialID != 11 { t.Fatalf("first Next() = (%#v, %v), want captured key 11", first, err) } if _, err := iterator.Next(); !errors.Is(err, ErrExhausted) { @@ -664,7 +548,7 @@ func TestIteratorRestrictsLiveCandidatesToAllowedKeyIDs(t *testing.T) { func TestIteratorPropertyNeverEscapesAccessFilters(t *testing.T) { snapshot := schedulerSnapshot() - source := fakeKeySource{keys: []state.KeyMeta{ + source := fakeCredentialSource{keys: []state.CredentialMeta{ {ID: 11, GroupID: 1}, {ID: 12, GroupID: 1}, {ID: 21, GroupID: 2}, {ID: 22, GroupID: 2}, {ID: 31, GroupID: 3}, @@ -684,7 +568,7 @@ func TestIteratorPropertyNeverEscapesAccessFilters(t *testing.T) { filters.Models = map[string]struct{}{"gpt-4o": {}} } query := Query{ - Protocol: protocol.OpenAICompletions, + ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o"), AccessKey: state.AccessKeyView{ID: uint(caseIndex + 1), Filters: filters}, } @@ -733,9 +617,9 @@ func TestIteratorExhaustsForNilOrEmptyDependencies(t *testing.T) { name string iterator *Iterator }{ - {name: "nil snapshot", iterator: New(nil, fakeKeySource{}, Query{}, rand.New(rand.NewSource(1)))}, - {name: "nil key source", iterator: New(schedulerSnapshot(), nil, Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, rand.New(rand.NewSource(1)))}, - {name: "nil random", iterator: New(schedulerSnapshot(), fakeKeySource{}, Query{Protocol: protocol.OpenAICompletions, ExternalModel: modelPointer("gpt-4o")}, nil)}, + {name: "nil snapshot", iterator: New(nil, fakeCredentialSource{}, Query{}, rand.New(rand.NewSource(1)))}, + {name: "nil key source", iterator: New(schedulerSnapshot(), nil, Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, rand.New(rand.NewSource(1)))}, + {name: "nil random", iterator: New(schedulerSnapshot(), fakeCredentialSource{}, Query{ClientProtocol: protocol.OpenAICompletions, Operation: execution.OperationChatCompletion, ExternalModel: modelPointer("gpt-4o")}, nil)}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -747,22 +631,22 @@ func TestIteratorExhaustsForNilOrEmptyDependencies(t *testing.T) { } func schedulerSnapshot() *state.ConfigSnapshot { - return &state.ConfigSnapshot{ - Candidates: map[protocol.Protocol]map[string][]state.RouteTarget{ - protocol.OpenAICompletions: { - "gpt-4o": { - {GroupID: 1, UpstreamModelID: "gpt-4o"}, - {GroupID: 2, UpstreamModelID: "provider-gpt-4o"}, - }, + snapshot, err := state.Compile(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []state.GroupConfig{ + { + ID: 1, Name: "one", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "gpt-4o", Alias: "gpt-4o"}}, Enabled: true, + }, + { + ID: 2, Name: "two", ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://two.example/v1"}`), + Models: []state.ModelConfig{{ID: "provider-gpt-4o", Alias: "gpt-4o"}}, Enabled: true, }, }, - Groups: map[uint]state.GroupView{ - 1: {ID: 1, Name: "one", UpstreamURL: "https://one.example"}, - 2: {ID: 2, Name: "two", UpstreamURL: "https://two.example"}, - }, - GroupCatalog: map[uint]state.GroupCatalogView{ - 1: {ID: 1, Name: "one", Enabled: true}, - 2: {ID: 2, Name: "two", Enabled: true}, - }, + }) + if err != nil { + panic(err) } + return snapshot } diff --git a/internal/state/channel_snapshot_test.go b/internal/state/channel_snapshot_test.go new file mode 100644 index 00000000..52e30b67 --- /dev/null +++ b/internal/state/channel_snapshot_test.go @@ -0,0 +1,262 @@ +package state + +import ( + "encoding/json" + "strings" + "testing" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestCompileBuildsOperationAwareChannelCandidates(t *testing.T) { + t.Parallel() + + input := CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []GroupConfig{{ + ID: 7, Name: "anthropic", ChannelID: channel.Anthropic, + Params: json.RawMessage(`{}`), + Models: []ModelConfig{{ID: "claude-upstream", Alias: "claude-public"}}, + Enabled: true, + }}, + Credentials: []CredentialConfig{{ + ID: 31, GroupID: 7, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "fingerprint-31", + }}, + } + + snapshot, err := Compile(input) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + + tests := []struct { + clientProtocol protocol.Protocol + operation execution.Operation + wantMode channel.RouteMode + }{ + {protocol.Anthropic, execution.OperationChatCompletion, channel.RouteNative}, + {protocol.OpenAICompletions, execution.OperationChatCompletion, channel.RouteConverted}, + {protocol.Gemini, execution.OperationChatCompletion, channel.RouteConverted}, + {protocol.OpenAIResponses, execution.OperationResponsesCreate, channel.RouteConverted}, + } + for _, test := range tests { + got := snapshot.ExecutionCandidates[test.clientProtocol][test.operation]["claude-public"] + if len(got) != 1 { + t.Fatalf("ExecutionCandidates[%q][%q] = %#v, want one", test.clientProtocol, test.operation, got) + } + if got[0].GroupID != 7 || got[0].UpstreamModelID != "claude-upstream" || got[0].Mode != test.wantMode { + t.Errorf("candidate = %#v, want group 7, upstream model and mode %q", got[0], test.wantMode) + } + if got[0].ResolvedTarget.ChannelID != channel.Anthropic || got[0].ResolvedTarget.CatalogProviderID != "anthropic" { + t.Errorf("candidate target = %#v", got[0].ResolvedTarget) + } + } + + view := snapshot.Groups[7] + if view.ChannelID != channel.Anthropic || string(view.Params) != `{}` || view.ResolvedTarget.ChannelID != channel.Anthropic { + t.Fatalf("GroupView channel fields = %#v", view) + } + if len(view.ClientProtocols) != 4 { + t.Fatalf("GroupView.ClientProtocols = %#v, want registry-derived protocols", view.ClientProtocols) + } +} + +func TestCompileIndexesNativeResponsesResourceOperationsWithoutModels(t *testing.T) { + t.Parallel() + + snapshot, err := Compile(CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []GroupConfig{ + {ID: 1, Name: "openai", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Enabled: true}, + {ID: 2, Name: "compatible", ChannelID: channel.OpenAICompatible, Params: json.RawMessage(`{"base_url":"https://proxy.example/v1"}`), Enabled: true}, + }, + }) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + + for _, operation := range []execution.Operation{ + execution.OperationResponsesRetrieve, + execution.OperationResponsesDelete, + execution.OperationResponsesCancel, + execution.OperationResponsesInputItems, + } { + got := snapshot.ExecutionCandidates[protocol.OpenAIResponses][operation][NoModelRouteKey] + if len(got) != 2 || got[0].GroupID != 1 || got[1].GroupID != 2 || + got[0].Mode != channel.RouteNative || got[1].Mode != channel.RouteNative { + t.Errorf("resource candidates for %q = %#v, want native OpenAI and OpenAI-compatible", operation, got) + } + } + if got := snapshot.ExecutionCandidates[protocol.OpenAIResponses][execution.OperationResponsesCreate]; len(got) != 0 { + t.Fatalf("model-less Responses create candidates = %#v, want none", got) + } +} + +func TestCompileIndexesAllNativeResponsesExtensions(t *testing.T) { + t.Parallel() + + snapshot, err := Compile(CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []GroupConfig{ + {ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []ModelConfig{{ID: "official-upstream", Alias: "public"}}, Enabled: true}, + {ID: 2, ChannelID: channel.OpenAICompatible, Params: json.RawMessage(`{"base_url":"https://proxy.example/v1"}`), Models: []ModelConfig{{ID: "compatible-upstream", Alias: "public"}}, Enabled: true}, + }, + }) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + + for _, operation := range []execution.Operation{ + execution.OperationResponsesCompact, + execution.OperationResponsesInputTokens, + } { + got := snapshot.ExecutionCandidates[protocol.OpenAIResponses][operation]["public"] + if len(got) != 2 || got[0].UpstreamModelID != "official-upstream" || + got[1].UpstreamModelID != "compatible-upstream" { + t.Errorf("model candidates for %q = %#v", operation, got) + } + } + + noModel := snapshot.ExecutionCandidates[protocol.OpenAIResponses][execution.OperationResponsesPassthrough][NoModelRouteKey] + if len(noModel) != 2 || noModel[0].GroupID != 1 || noModel[1].GroupID != 2 { + t.Fatalf("model-less passthrough candidates = %#v", noModel) + } + withModel := snapshot.ExecutionCandidates[protocol.OpenAIResponses][execution.OperationResponsesPassthrough]["public"] + if len(withModel) != 2 || withModel[0].UpstreamModelID != "official-upstream" || + withModel[1].UpstreamModelID != "compatible-upstream" { + t.Fatalf("model passthrough candidates = %#v", withModel) + } +} + +func TestCompileOrdersNativeTargetsBeforeConvertedTargets(t *testing.T) { + t.Parallel() + + snapshot, err := Compile(CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []GroupConfig{ + {ID: 1, ChannelID: channel.Anthropic, Params: json.RawMessage(`{}`), Models: []ModelConfig{{ID: "converted", Alias: "public"}}, Enabled: true}, + {ID: 2, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []ModelConfig{{ID: "native", Alias: "public"}}, Enabled: true}, + }, + }) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + targets := snapshot.ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion]["public"] + if len(targets) != 2 || targets[0].GroupID != 2 || targets[0].Mode != channel.RouteNative || + targets[1].GroupID != 1 || targets[1].Mode != channel.RouteConverted { + t.Fatalf("targets = %#v, want native before converted", targets) + } +} + +func TestCompileChannelSnapshotOwnsInputData(t *testing.T) { + t.Parallel() + + weight := 12 + params := json.RawMessage(`{"base_url":"https://proxy.example/v1/"}`) + input := CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []GroupConfig{{ + ID: 1, ChannelID: channel.OpenAICompatible, Params: params, + Models: []ModelConfig{{ID: "upstream", Alias: "public"}}, + WeightManual: &weight, Enabled: true, + }}, + } + snapshot, err := Compile(input) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + + params[2] = 'X' + input.Groups[0].Models[0] = ModelConfig{ID: "changed", Alias: "changed"} + weight = 99 + + view := snapshot.Groups[1] + if string(view.Params) != `{"base_url":"https://proxy.example/v1"}` { + t.Fatalf("GroupView.Params = %s", view.Params) + } + if len(view.Models) != 1 || view.Models[0].ID != "upstream" || view.Models[0].Alias != "public" { + t.Fatalf("GroupView.Models = %#v", view.Models) + } + if view.WeightManual == nil || *view.WeightManual != 12 { + t.Fatalf("GroupView.WeightManual = %v", view.WeightManual) + } + target := snapshot.ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion]["public"][0] + if string(target.ResolvedTarget.TargetConfig) != `{"base_url":"https://proxy.example/v1"}` { + t.Fatalf("candidate TargetConfig = %s", target.ResolvedTarget.TargetConfig) + } + if target.UpstreamModelID != "upstream" { + t.Fatalf("candidate upstream model = %q", target.UpstreamModelID) + } +} + +func TestCompileRejectsInvalidChannelAndCredentialConfiguration(t *testing.T) { + t.Parallel() + + invalidWeight := -1 + tests := []struct { + name string + input CompileInput + wantErr string + }{ + { + name: "missing registry", + input: CompileInput{Groups: []GroupConfig{{ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`)}}}, + wantErr: "channel registry is required", + }, + { + name: "unknown channel", + input: CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{{ID: 1, ChannelID: channel.ID("unknown"), Params: json.RawMessage(`{}`)}}}, + wantErr: "unknown channel", + }, + { + name: "invalid params", + input: CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{{ID: 1, ChannelID: channel.OpenAICompatible, Params: json.RawMessage(`{}`)}}}, + wantErr: "params.base_url", + }, + { + name: "credential unknown group", + input: CompileInput{ChannelRegistry: channel.NewRegistry(), Credentials: []CredentialConfig{{ + ID: 1, GroupID: 99, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "fingerprint-1", + }}}, + wantErr: "unknown group", + }, + { + name: "duplicate credential", + input: CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{{ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`)}}, Credentials: []CredentialConfig{ + {ID: 1, GroupID: 1, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "fingerprint-1"}, + {ID: 1, GroupID: 1, Status: CredentialStatusDisabled, Version: 2, IdentityGeneration: 2, Fingerprint: "fingerprint-2"}, + }}, + wantErr: "duplicate credential id", + }, + { + name: "invalid credential status", + input: CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{{ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`)}}, Credentials: []CredentialConfig{{ + ID: 1, GroupID: 1, Status: CredentialStatus("revoked"), + Version: 1, IdentityGeneration: 1, Fingerprint: "fingerprint-1", + }}}, + wantErr: "invalid status", + }, + { + name: "invalid credential weight", + input: CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{{ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`)}}, Credentials: []CredentialConfig{{ + ID: 1, GroupID: 1, Status: CredentialStatusActive, WeightManual: &invalidWeight, + Version: 1, IdentityGeneration: 1, Fingerprint: "fingerprint-1", + }}}, + wantErr: "manual weight", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + _, err := Compile(test.input) + if err == nil || !strings.Contains(err.Error(), test.wantErr) { + t.Fatalf("Compile() error = %v, want substring %q", err, test.wantErr) + } + }) + } +} diff --git a/internal/state/credential_registry_test.go b/internal/state/credential_registry_test.go new file mode 100644 index 00000000..bab1b8f6 --- /dev/null +++ b/internal/state/credential_registry_test.go @@ -0,0 +1,79 @@ +package state + +import "testing" + +func TestCredentialRegistryKeepsIdentityAndHealthGenerationsIndependent(t *testing.T) { + t.Parallel() + + registry := NewCredentialRegistry() + entry := CredentialEntry{ + ID: 11, GroupID: 7, Version: 42, IdentityGeneration: 99, + Fingerprint: "fingerprint", Status: CredentialStatusActive, + WeightAuto: DefaultWeight, EncryptedValue: "encrypted-data", + } + if err := ValidateCredentialEntries([]CredentialEntry{entry}); err != nil { + t.Fatalf("ValidateCredentialEntries() error = %v", err) + } + if err := registry.ReplaceCredentials([]CredentialEntry{entry}); err != nil { + t.Fatalf("ReplaceCredentials() error = %v", err) + } + + refs := registry.CaptureActiveCredentialRefs([]uint{7}) + if len(refs) != 1 { + t.Fatalf("CaptureActiveCredentialRefs() = %#v", refs) + } + ref := refs[0] + if ref.Version != 42 || ref.IdentityGeneration != 99 || ref.Fingerprint != "fingerprint" { + t.Fatalf("credential ref = %#v", ref) + } + + if _, ok := registry.IncrFailure(11); !ok { + t.Fatal("IncrFailure() = false") + } + if got, ok := registry.ActiveEncryptedCredentialDataIfMatch(ref); !ok || got != "encrypted-data" { + t.Fatalf("health generation invalidated credential identity: %q, %t", got, ok) + } + + changed := ref + changed.Fingerprint = "other-fingerprint" + if _, ok := registry.ActiveEncryptedCredentialDataIfMatch(changed); ok { + t.Fatal("mismatched full fingerprint was accepted") + } + changed = ref + changed.IdentityGeneration++ + if _, ok := registry.ActiveEncryptedCredentialDataIfMatch(changed); ok { + t.Fatal("mismatched identity generation was accepted") + } + changed = ref + changed.Version++ + if _, ok := registry.ActiveEncryptedCredentialDataIfMatch(changed); ok { + t.Fatal("mismatched version was accepted") + } +} + +func TestValidateCredentialEntriesRequiresDurableIdentityEvidence(t *testing.T) { + t.Parallel() + + base := CredentialEntry{ + ID: 1, GroupID: 2, Version: 3, IdentityGeneration: 4, + Fingerprint: "fingerprint", Status: CredentialStatusActive, + WeightAuto: DefaultWeight, EncryptedValue: "encrypted", + } + tests := []struct { + name string + mutate func(*CredentialEntry) + }{ + {name: "version", mutate: func(entry *CredentialEntry) { entry.Version = 0 }}, + {name: "identity generation", mutate: func(entry *CredentialEntry) { entry.IdentityGeneration = 0 }}, + {name: "fingerprint", mutate: func(entry *CredentialEntry) { entry.Fingerprint = " " }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + entry := base + test.mutate(&entry) + if err := ValidateCredentialEntries([]CredentialEntry{entry}); err == nil { + t.Fatal("ValidateCredentialEntries() error = nil") + } + }) + } +} diff --git a/internal/state/loader/channel_loader_test.go b/internal/state/loader/channel_loader_test.go new file mode 100644 index 00000000..41181bd2 --- /dev/null +++ b/internal/state/loader/channel_loader_test.go @@ -0,0 +1,140 @@ +package loader_test + +import ( + "encoding/json" + "fmt" + "strings" + "testing" + "time" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" + "gpt-load/internal/state" + "gpt-load/internal/state/loader" + "gpt-load/internal/storage/models" +) + +func TestBuildCompileInputMapsChannelAndCredentialMetadata(t *testing.T) { + t.Parallel() + + db := openMigratedDatabase(t) + group := models.Group{ + Name: "compatible", ChannelID: string(channel.OpenAICompatible), + Params: models.JSON(`{"base_url":"https://proxy.example/v1/"}`), + Models: models.JSON(`[{"id":"upstream-model","alias":"public-model"}]`), + Overrides: models.JSON(`{}`), Enabled: true, + } + mustCreate(t, db, &group) + credential := models.Credential{ + GroupID: group.ID, Data: "encrypted-channel-secret", Fingerprint: "fingerprint-one", + Status: models.CredentialStatusActive, UpdatedAtMS: 42, + } + mustCreate(t, db, &credential) + + input, err := loader.BuildCompileInput(t.Context(), db, channel.NewRegistry()) + if err != nil { + t.Fatalf("BuildCompileInput() error = %v", err) + } + if len(input.Groups) != 1 || input.Groups[0].ChannelID != channel.OpenAICompatible || + string(input.Groups[0].Params) != `{"base_url":"https://proxy.example/v1/"}` { + t.Fatalf("CompileInput.Groups = %#v", input.Groups) + } + if len(input.Credentials) != 1 { + t.Fatalf("CompileInput.Credentials = %#v", input.Credentials) + } + metadata := input.Credentials[0] + if metadata.ID != credential.ID || metadata.GroupID != group.ID || metadata.Version != 42 || + metadata.IdentityGeneration == 0 || metadata.Fingerprint != credential.Fingerprint { + t.Fatalf("credential metadata = %#v", metadata) + } + if strings.Contains(fmt.Sprintf("%#v", input), credential.Data) { + t.Fatal("CompileInput contains encrypted credential data") + } + + snapshot, err := state.Compile(input) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + targets := snapshot.ExecutionCandidates[protocol.OpenAIResponses][execution.OperationResponsesCreate]["public-model"] + if len(targets) != 1 || targets[0].GroupID != group.ID || targets[0].Mode != channel.RouteNative || + targets[0].UpstreamModelID != "upstream-model" { + t.Fatalf("Responses candidates = %#v", targets) + } +} + +func TestBuildGroupCredentialEntriesUsesStableCredentialIdentity(t *testing.T) { + t.Parallel() + + db := openMigratedDatabase(t) + group := models.Group{ + Name: "credentials", ChannelID: string(channel.OpenAI), Params: models.JSON(`{}`), + Models: models.JSON(`[]`), Overrides: models.JSON(`{}`), Enabled: true, + } + mustCreate(t, db, &group) + weight := 9 + credentials := []models.Credential{ + {GroupID: group.ID, Data: "cipher-one", Fingerprint: "fingerprint-one", Status: models.CredentialStatusActive, UpdatedAtMS: 0}, + {GroupID: group.ID, Data: "cipher-two", Fingerprint: "fingerprint-two", Status: models.CredentialStatusDisabled, WeightManual: &weight, UpdatedAtMS: 99}, + } + for index := range credentials { + mustCreate(t, db, &credentials[index]) + } + + entries, err := loader.BuildGroupCredentialEntries(t.Context(), db, group.ID) + if err != nil { + t.Fatalf("BuildGroupCredentialEntries() error = %v", err) + } + if len(entries) != 2 || entries[0].ID != credentials[0].ID || entries[1].ID != credentials[1].ID { + t.Fatalf("entries = %#v", entries) + } + if entries[0].Version == 0 || entries[0].IdentityGeneration == 0 || + entries[0].Fingerprint != credentials[0].Fingerprint || entries[0].EncryptedValue != credentials[0].Data { + t.Fatalf("first entry identity = %#v", entries[0]) + } + if entries[1].Version != 99 || entries[1].IdentityGeneration == entries[0].IdentityGeneration || + entries[1].Status != state.CredentialStatusDisabled || entries[1].WeightManual == nil || *entries[1].WeightManual != weight { + t.Fatalf("second entry identity = %#v", entries[1]) + } +} + +func TestLoaderUsesCredentialsAndKeepsCiphertextOutOfSnapshot(t *testing.T) { + t.Parallel() + + db := openMigratedDatabase(t) + group := models.Group{ + Name: "runtime", ChannelID: string(channel.Anthropic), Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"claude-upstream","alias":"claude-public"}]`), + Overrides: models.JSON(`{}`), Enabled: true, + } + mustCreate(t, db, &group) + credential := models.Credential{ + GroupID: group.ID, Data: "new-credential-cipher", Fingerprint: "new-fingerprint", + Status: models.CredentialStatusActive, + } + mustCreate(t, db, &credential) + manager := state.NewManager() + registry := state.NewCredentialRegistry() + if err := loader.New(db, manager, registry, channel.NewRegistry()).Load(t.Context()); err != nil { + t.Fatalf("Load() error = %v", err) + } + + candidates := registry.CollectCredentialCandidates([]uint{group.ID}, nil, time.Time{}) + if len(candidates) != 1 || candidates[0].ID != credential.ID || candidates[0].Version == 0 || candidates[0].IdentityGeneration == 0 { + t.Fatalf("credential candidates = %#v", candidates) + } + if got, ok := registry.EncryptedCredentialData(credential.ID); !ok || got != credential.Data { + t.Fatalf("EncryptedValue(%d) = %q, %t", credential.ID, got, ok) + } + snapshot := manager.Current() + if snapshot == nil { + t.Fatal("Current() = nil") + } + encoded, err := json.Marshal(snapshot) + if err != nil { + t.Fatalf("json.Marshal(snapshot) error = %v", err) + } + if strings.Contains(string(encoded), credential.Data) || strings.Contains(fmt.Sprintf("%#v", snapshot), credential.Data) { + t.Fatal("snapshot exposes encrypted credential data") + } +} diff --git a/internal/state/loader/loader.go b/internal/state/loader/loader.go index 92c32577..3ff28ac4 100644 --- a/internal/state/loader/loader.go +++ b/internal/state/loader/loader.go @@ -7,12 +7,14 @@ import ( "encoding/json" "errors" "fmt" + "hash/fnv" "io" "strings" "gorm.io/gorm" "gorm.io/gorm/clause" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -22,15 +24,17 @@ import ( ) type Loader struct { - db *gorm.DB - manager *state.Manager - registry *state.KeyRegistry + db *gorm.DB + manager *state.Manager + registry *state.CredentialRegistry + channelRegistry *channel.Registry } type compileRows struct { - settings []models.SystemSetting - groups []models.Group - accessKeys []models.AccessKey + settings []models.SystemSetting + groups []models.Group + credentials []models.Credential + accessKeys []models.AccessKey } type modelDTO struct { @@ -67,8 +71,16 @@ func (f filterDTO) toState() state.FilterSet { return filters } -func New(db *gorm.DB, manager *state.Manager, registry *state.KeyRegistry) *Loader { - return &Loader{db: db, manager: manager, registry: registry} +func New( + db *gorm.DB, + manager *state.Manager, + registry *state.CredentialRegistry, + registries ...*channel.Registry, +) *Loader { + return &Loader{ + db: db, manager: manager, registry: registry, + channelRegistry: selectChannelRegistry(registries), + } } func (l *Loader) Load(ctx context.Context) error { @@ -76,8 +88,8 @@ func (l *Loader) Load(ctx context.Context) error { if err != nil { return fmt.Errorf("read runtime state: %w", err) } - if err := state.ValidateKeyEntries(entries); err != nil { - return fmt.Errorf("validate upstream keys: %w", err) + if err := state.ValidateCredentialEntries(entries); err != nil { + return fmt.Errorf("validate credentials: %w", err) } snapshot, err := l.manager.Publish(input) if err != nil { @@ -87,13 +99,13 @@ func (l *Loader) Load(ctx context.Context) error { "event": "startup.config_snapshot_publish", "revision": snapshot.Revision, }).Info("config snapshot published") - if err := l.registry.Replace(entries); err != nil { - return fmt.Errorf("replace key registry: %w", err) + if err := l.registry.ReplaceCredentials(entries); err != nil { + return fmt.Errorf("replace credential registry: %w", err) } logrus.WithFields(logrus.Fields{ - "event": "startup.key_registry_load", - "keys": len(entries), - }).Info("key registry loaded") + "event": "startup.credential_registry_load", + "credentials": len(entries), + }).Info("credential registry loaded") return nil } @@ -108,6 +120,12 @@ func queryCompileRows(ctx context.Context, db *gorm.DB) (compileRows, error) { if err := db.Order("id ASC").Find(&rows.groups).Error; err != nil { return compileRows{}, fmt.Errorf("query groups: %w", err) } + if err := db. + Select("id", "group_id", "fingerprint", "status", "weight_manual", "updated_at_ms"). + Order("id ASC"). + Find(&rows.credentials).Error; err != nil { + return compileRows{}, fmt.Errorf("query credential metadata: %w", err) + } if err := db. Select("id", "name", "key_hash", "status", "filters", "rpm_limit"). Order("id ASC"). @@ -117,16 +135,20 @@ func queryCompileRows(ctx context.Context, db *gorm.DB) (compileRows, error) { return rows, nil } -func queryUpstreamKeys(ctx context.Context, db *gorm.DB) ([]models.UpstreamKey, error) { - var rows []models.UpstreamKey +func queryCredentials(ctx context.Context, db *gorm.DB) ([]models.Credential, error) { + var rows []models.Credential if err := db.WithContext(ctx).Order("id ASC").Find(&rows).Error; err != nil { - return nil, fmt.Errorf("query upstream keys: %w", err) + return nil, fmt.Errorf("query credentials: %w", err) } return rows, nil } // BuildCompileInput maps persisted configuration rows into a runtime compiler input. -func BuildCompileInput(ctx context.Context, db *gorm.DB) (state.CompileInput, error) { +func BuildCompileInput( + ctx context.Context, + db *gorm.DB, + registries ...*channel.Registry, +) (state.CompileInput, error) { rows, err := queryCompileRows(ctx, db) if err != nil { return state.CompileInput{}, err @@ -135,6 +157,8 @@ func BuildCompileInput(ctx context.Context, db *gorm.DB) (state.CompileInput, er if err != nil { return state.CompileInput{}, err } + input.ChannelRegistry = selectChannelRegistry(registries) + input.Credentials = mapCredentialConfigs(rows.credentials) input.AccessKeys, err = mapAccessKeys(rows.accessKeys) if err != nil { return state.CompileInput{}, err @@ -142,13 +166,13 @@ func BuildCompileInput(ctx context.Context, db *gorm.DB) (state.CompileInput, er return input, nil } -// BuildGroupKeyEntries maps the persisted key configuration for one group. +// BuildGroupCredentialEntries maps persisted credentials for one group. // Runtime health fields are initialized to their durable baseline. -func BuildGroupKeyEntries( +func BuildGroupCredentialEntries( ctx context.Context, db *gorm.DB, groupID uint, -) ([]state.KeyEntry, error) { +) ([]state.CredentialEntry, error) { if groupID == 0 { return nil, fmt.Errorf("group id is required") } @@ -160,30 +184,39 @@ func BuildGroupKeyEntries( Take(&group).Error; err != nil { return nil, fmt.Errorf("query group %d: %w", groupID, err) } - var rows []models.UpstreamKey + var rows []models.Credential if err := db.WithContext(ctx). Where("group_id = ?", groupID). Order("id ASC"). Find(&rows).Error; err != nil { - return nil, fmt.Errorf("query group %d upstream keys: %w", groupID, err) + return nil, fmt.Errorf("query group %d credentials: %w", groupID, err) } - entries := mapUpstreamKeys(rows) - if err := state.ValidateKeyEntries(entries); err != nil { - return nil, fmt.Errorf("validate group %d upstream keys: %w", groupID, err) + entries := mapCredentials(rows) + if err := state.ValidateCredentialEntries(entries); err != nil { + return nil, fmt.Errorf("validate group %d credentials: %w", groupID, err) } return entries, nil } -func (l *Loader) read(ctx context.Context) (state.CompileInput, []state.KeyEntry, error) { - input, err := BuildCompileInput(ctx, l.db) +func (l *Loader) read(ctx context.Context) (state.CompileInput, []state.CredentialEntry, error) { + input, err := BuildCompileInput(ctx, l.db, l.channelRegistry) if err != nil { return state.CompileInput{}, nil, err } - upstreamKeys, err := queryUpstreamKeys(ctx, l.db) + credentials, err := queryCredentials(ctx, l.db) if err != nil { return state.CompileInput{}, nil, err } - return input, mapUpstreamKeys(upstreamKeys), nil + return input, mapCredentials(credentials), nil +} + +func selectChannelRegistry(registries []*channel.Registry) *channel.Registry { + for _, registry := range registries { + if registry != nil { + return registry + } + } + return channel.NewRegistry() } func decodeJSON(raw models.JSON, target any) error { @@ -276,17 +309,13 @@ func mapSystemAndGroups(rows compileRows) (state.CompileInput, error) { input.SystemSettings[row.Key] = value } for _, row := range rows.groups { - var protocols []protocol.Protocol - if err := decodeJSON(row.Protocols, &protocols); err != nil { - return state.CompileInput{}, fmt.Errorf("decode group %d protocols: %w", row.ID, err) - } var storedModels []modelDTO if err := decodeJSON(row.Models, &storedModels); err != nil { return state.CompileInput{}, fmt.Errorf("decode group %d models: %w", row.ID, err) } settings := make(config.Settings) - if err := decodeJSON(row.Config, &settings); err != nil { - return state.CompileInput{}, fmt.Errorf("decode group %d config: %w", row.ID, err) + if err := decodeJSON(row.Overrides, &settings); err != nil { + return state.CompileInput{}, fmt.Errorf("decode group %d overrides: %w", row.ID, err) } validationModel := "" if row.ValidationModel != nil { @@ -297,30 +326,22 @@ func mapSystemAndGroups(rows compileRows) (state.CompileInput, error) { for _, model := range storedModels { runtimeModels = append(runtimeModels, state.ModelConfig{ID: model.ID, Alias: model.Alias}) } - input.Groups = append(input.Groups, state.GroupConfig{ + group := state.GroupConfig{ ID: row.ID, Name: row.Name, - ProviderID: cloneStringPointer(row.ProviderID), - UpstreamURL: row.UpstreamURL, + ChannelID: channel.ID(row.ChannelID), + Params: append(json.RawMessage(nil), row.Params...), ValidationModel: validationModel, - Protocols: append([]protocol.Protocol(nil), protocols...), Models: runtimeModels, Settings: settings, - WeightManual: row.WeightManual, + WeightManual: cloneWeight(row.WeightManual), Enabled: row.Enabled, - }) + } + input.Groups = append(input.Groups, group) } return input, nil } -func cloneStringPointer(value *string) *string { - if value == nil { - return nil - } - cloned := *value - return &cloned -} - func mapAccessKeys(rows []models.AccessKey) ([]state.AccessKeyConfig, error) { result := make([]state.AccessKeyConfig, 0, len(rows)) for _, row := range rows { @@ -336,14 +357,56 @@ func mapAccessKeys(rows []models.AccessKey) ([]state.AccessKeyConfig, error) { return result, nil } -func mapUpstreamKeys(rows []models.UpstreamKey) []state.KeyEntry { - result := make([]state.KeyEntry, 0, len(rows)) +func mapCredentialConfigs(rows []models.Credential) []state.CredentialConfig { + result := make([]state.CredentialConfig, 0, len(rows)) for _, row := range rows { - result = append(result, state.KeyEntry{ - ID: row.ID, GroupID: row.GroupID, WeightManual: row.WeightManual, - WeightAuto: state.DefaultWeight, - Status: state.KeyStatus(row.Status), EncryptedValue: row.KeyValue, + result = append(result, state.CredentialConfig{ + ID: row.ID, GroupID: row.GroupID, WeightManual: cloneWeight(row.WeightManual), + Status: state.CredentialStatus(row.Status), + Version: credentialVersion(row.UpdatedAtMS), + IdentityGeneration: credentialIdentityGeneration(row.Fingerprint), + Fingerprint: row.Fingerprint, }) } return result } + +func mapCredentials(rows []models.Credential) []state.CredentialEntry { + result := make([]state.CredentialEntry, 0, len(rows)) + for _, row := range rows { + result = append(result, state.CredentialEntry{ + ID: row.ID, GroupID: row.GroupID, + Version: credentialVersion(row.UpdatedAtMS), + IdentityGeneration: credentialIdentityGeneration(row.Fingerprint), + Fingerprint: row.Fingerprint, WeightManual: cloneWeight(row.WeightManual), + WeightAuto: state.DefaultWeight, + Status: state.CredentialStatus(row.Status), EncryptedValue: row.Data, + }) + } + return result +} + +func credentialVersion(updatedAtMS int64) uint64 { + if updatedAtMS < 1 { + return 1 + } + return uint64(updatedAtMS) +} + +func credentialIdentityGeneration(fingerprint string) uint64 { + hasher := fnv.New64a() + _, _ = hasher.Write([]byte(fingerprint)) + generation := hasher.Sum64() + if generation == 0 { + return 1 + } + return generation +} + +func cloneWeight(value *int) *int { + if value == nil { + return nil + } + cloned := *value + return &cloned +} diff --git a/internal/state/loader/loader_test.go b/internal/state/loader/loader_test.go index 08cfe8c4..5b1fea15 100644 --- a/internal/state/loader/loader_test.go +++ b/internal/state/loader/loader_test.go @@ -12,6 +12,8 @@ import ( "gorm.io/gorm" "gorm.io/gorm/clause" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/config" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -31,9 +33,9 @@ func TestLoadSystemSettingsDecodesPersistedValues(t *testing.T) { mustCreate(t, db, &row) } mustCreate(t, db, &models.Group{ - Name: "unrelated", UpstreamURL: "https://unrelated.example.com", - Protocols: models.JSON(`["openai-completions"]`), Models: models.JSON(`[]`), - Config: models.JSON(`{}`), Enabled: true, + Name: "unrelated", ChannelID: string(channel.OpenAI), Params: models.JSON(`{}`), + Models: models.JSON(`[]`), + Overrides: models.JSON(`{}`), Enabled: true, }) var orderColumns []clause.OrderByColumn @@ -180,7 +182,7 @@ func TestLoadSystemSettingsExcludesInternalSystemSettings(t *testing.T) { func TestLoaderPublishesDefaultsFromEmptyDatabase(t *testing.T) { db := openMigratedDatabase(t) manager := state.NewManager() - if err := loader.New(db, manager, state.NewKeyRegistry()).Load(context.Background()); err != nil { + if err := loader.New(db, manager, state.NewCredentialRegistry()).Load(context.Background()); err != nil { t.Fatal(err) } got := manager.Current().Settings @@ -197,7 +199,7 @@ func TestLoaderRejectsUnknownPublicSystemSetting(t *testing.T) { db := openMigratedDatabase(t) mustCreate(t, db, &models.SystemSetting{Key: "unknown_public", Value: "true"}) manager := state.NewManager() - err := loader.New(db, manager, state.NewKeyRegistry()).Load(context.Background()) + err := loader.New(db, manager, state.NewCredentialRegistry()).Load(context.Background()) if err == nil || manager.Current() != nil { t.Fatalf("Load() error/current = %v/%#v", err, manager.Current()) } @@ -218,7 +220,7 @@ func TestLoaderRejectsInvalidKnownPublicSystemSettingsWithoutPublishing(t *testi mustCreate(t, db, &models.SystemSetting{Key: test.key, Value: test.value}) manager := state.NewManager() - err := loader.New(db, manager, state.NewKeyRegistry()).Load(context.Background()) + err := loader.New(db, manager, state.NewCredentialRegistry()).Load(context.Background()) if err == nil { t.Fatalf("Load() accepted %s=%s", test.key, test.value) } @@ -232,7 +234,7 @@ func TestLoaderRejectsInvalidKnownPublicSystemSettingsWithoutPublishing(t *testi func TestLoaderLoadsEmptyMigratedDatabase(t *testing.T) { db := openMigratedDatabase(t) manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) @@ -245,8 +247,8 @@ func TestLoaderLoadsEmptyMigratedDatabase(t *testing.T) { if snapshot.Revision != 1 { t.Errorf("snapshot revision = %d, want 1", snapshot.Revision) } - if snapshot.Candidates == nil || len(snapshot.Candidates) != 0 { - t.Errorf("snapshot candidates = %#v, want initialized empty map", snapshot.Candidates) + if snapshot.ExecutionCandidates == nil || len(snapshot.ExecutionCandidates) != 0 { + t.Errorf("snapshot candidates = %#v, want initialized empty map", snapshot.ExecutionCandidates) } if snapshot.Groups == nil || len(snapshot.Groups) != 0 { t.Errorf("snapshot groups = %#v, want initialized empty map", snapshot.Groups) @@ -254,7 +256,7 @@ func TestLoaderLoadsEmptyMigratedDatabase(t *testing.T) { if snapshot.AccessKeysByHash == nil || len(snapshot.AccessKeysByHash) != 0 { t.Errorf("snapshot access keys = %#v, want initialized empty map", snapshot.AccessKeysByHash) } - if got := registry.CollectCandidates([]uint{1}, nil, time.Time{}); len(got) != 0 { + if got := registry.CollectCredentialCandidates([]uint{1}, nil, time.Time{}); len(got) != 0 { t.Errorf("registry candidates = %#v, want empty", got) } } @@ -270,9 +272,8 @@ func TestBuildCompileInputReadsUncommittedTransactionState(t *testing.T) { }) group := models.Group{ - Name: "pending", UpstreamURL: "https://pending.example", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[{"id":"gpt-pending"}]`), Config: models.JSON(`{}`), Enabled: true, + Name: "pending", ChannelID: string(channel.OpenAI), Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"gpt-pending"}]`), Overrides: models.JSON(`{}`), Enabled: true, } mustCreate(t, tx, &group) mustCreate(t, tx, &models.AccessKey{ @@ -294,8 +295,8 @@ func TestBuildCompileInputReturnsIndependentData(t *testing.T) { db := openMigratedDatabase(t) mustCreate(t, db, &models.SystemSetting{Key: "connect_timeout", Value: "20"}) group := createRuntimeGroup(t, db, "owned", protocol.OpenAICompletions, "gpt-owned") - if err := db.Model(&group).Update("config", models.JSON(`{"request_timeout":30}`)).Error; err != nil { - t.Fatalf("update group config: %v", err) + if err := db.Model(&group).Update("overrides", models.JSON(`{"request_timeout":30}`)).Error; err != nil { + t.Fatalf("update group overrides: %v", err) } mustCreate(t, db, &models.AccessKey{ Name: "owned", KeyValue: "cipher", KeyHash: "owned-hash", Status: "active", @@ -308,7 +309,7 @@ func TestBuildCompileInputReturnsIndependentData(t *testing.T) { t.Fatalf("first BuildCompileInput() error = %v", err) } first.SystemSettings["connect_timeout"] = 99 - first.Groups[0].Protocols[0] = protocol.Anthropic + first.Groups[0].Params[0] = '[' first.Groups[0].Settings["request_timeout"] = 99 first.AccessKeys[0].Filters.Groups[999] = struct{}{} first.AccessKeys[0].Filters.Models["mutated"] = struct{}{} @@ -320,8 +321,8 @@ func TestBuildCompileInputReturnsIndependentData(t *testing.T) { if got := fmt.Sprint(second.SystemSettings["connect_timeout"]); got != "20" { t.Fatalf("connect_timeout = %q, want 20", got) } - if second.Groups[0].Protocols[0] != protocol.OpenAICompletions { - t.Fatalf("protocol = %q, want openai", second.Groups[0].Protocols[0]) + if second.Groups[0].ChannelID != channel.OpenAI || string(second.Groups[0].Params) != `{}` { + t.Fatalf("channel group = %#v, want OpenAI with canonical params", second.Groups[0]) } if got := fmt.Sprint(second.Groups[0].Settings["request_timeout"]); got != "30" { t.Fatalf("request_timeout = %q, want 30", got) @@ -334,39 +335,6 @@ func TestBuildCompileInputReturnsIndependentData(t *testing.T) { } } -func TestBuildCompileInputDoesNotQueryUpstreamKeys(t *testing.T) { - db := openMigratedDatabase(t) - group := createRuntimeGroup(t, db, "query-boundary", protocol.OpenAICompletions, "gpt-query") - mustCreate(t, db, &models.AccessKey{ - Name: "query-boundary", KeyValue: "cipher", KeyHash: "query-boundary-hash", - KeySuffix: "0002", Status: "active", Filters: models.JSON(`{}`), - }) - mustCreate(t, db, &models.UpstreamKey{ - GroupID: group.ID, KeyValue: "upstream-cipher", KeyHash: "upstream-query-hash", - Status: models.UpstreamKeyStatusActive, - }) - - const callbackName = "test:build_compile_input_tables" - seen := make(map[string]int) - if err := db.Callback().Query().After("gorm:query").Register(callbackName, func(tx *gorm.DB) { - seen[tx.Statement.Table]++ - }); err != nil { - t.Fatalf("register query callback: %v", err) - } - - if _, err := loader.BuildCompileInput(context.Background(), db); err != nil { - t.Fatalf("BuildCompileInput() error = %v", err) - } - if seen["upstream_keys"] != 0 { - t.Fatalf("upstream_keys query count = %d, want 0", seen["upstream_keys"]) - } - for _, table := range []string{"system_settings", "groups", "access_keys"} { - if seen[table] != 1 { - t.Errorf("%s query count = %d, want 1; all=%#v", table, seen[table], seen) - } - } -} - func TestLoaderMapsSystemAndGroupRows(t *testing.T) { db := openMigratedDatabase(t) mustCreate(t, db, &models.SystemSetting{Key: "connect_timeout", Value: "20"}) @@ -376,15 +344,15 @@ func TestLoaderMapsSystemAndGroupRows(t *testing.T) { }) enabled := models.Group{ - Name: "enabled", - UpstreamURL: "https://enabled.example.com/v1", - Protocols: models.JSON(`["openai-completions"]`), + Name: "enabled", + ChannelID: string(channel.OpenAI), + Params: models.JSON(`{}`), Models: models.JSON(`[ {"id":"gpt-4o","alias":"Primary"}, {"id":"gpt-4o","alias":"Secondary"}, {"id":"gpt-4.1","alias":"Other"} ]`), - Config: models.JSON(`{ + Overrides: models.JSON(`{ "request_timeout":30, "header_rules":{"set":{"X-Group":"group"},"remove":["X-Group-Remove"]} }`), @@ -392,12 +360,12 @@ func TestLoaderMapsSystemAndGroupRows(t *testing.T) { } mustCreate(t, db, &enabled) disabled := models.Group{ - Name: "disabled", - UpstreamURL: "https://disabled.example.com/v1", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[{"id":"hidden","alias":"Hidden"}]`), - Config: models.JSON(`{}`), - Enabled: true, + Name: "disabled", + ChannelID: string(channel.OpenAI), + Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"hidden","alias":"Hidden"}]`), + Overrides: models.JSON(`{}`), + Enabled: true, } mustCreate(t, db, &disabled) if err := db.Model(&disabled).Update("enabled", false).Error; err != nil { @@ -405,7 +373,7 @@ func TestLoaderMapsSystemAndGroupRows(t *testing.T) { } manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } @@ -457,7 +425,7 @@ func TestLoaderMapsSystemAndGroupRows(t *testing.T) { t.Errorf("group header remove rules = %#v, want group override", view.HeaderRules.Remove) } - openAICandidates := snapshot.Candidates[protocol.OpenAICompletions] + openAICandidates := snapshot.ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion] if len(openAICandidates) != 3 { t.Fatalf("OpenAI candidates = %#v, want three external model names", openAICandidates) } @@ -473,7 +441,7 @@ func TestLoaderMapsSystemAndGroupRows(t *testing.T) { if _, ok := openAICandidates["hidden"]; ok { t.Fatal("disabled group model hidden is present in candidates") } - if got := snapshot.RouteCatalog[protocol.OpenAICompletions]["Hidden"]; len(got) != 1 || + if got := snapshot.ExecutionRouteCatalog[protocol.OpenAICompletions][execution.OperationChatCompletion]["Hidden"]; len(got) != 1 || got[0].GroupID != disabled.ID || got[0].UpstreamModelID != "hidden" { t.Fatalf("disabled RouteCatalog entry = %#v", got) } @@ -494,17 +462,17 @@ func TestLoaderMapsValidationModelIntoRuntimeSnapshot(t *testing.T) { db := openMigratedDatabase(t) group := models.Group{ Name: "validation-" + test.name, - UpstreamURL: "https://validation.example.com/v1", - Protocols: models.JSON(`["openai-completions"]`), + ChannelID: string(channel.OpenAI), + Params: models.JSON(`{}`), Models: models.JSON(`[{"id":"real-model","alias":"public-model"}]`), ValidationModel: test.validationModel, - Config: models.JSON(`{}`), + Overrides: models.JSON(`{}`), Enabled: true, } mustCreate(t, db, &group) manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } @@ -513,7 +481,7 @@ func TestLoaderMapsValidationModelIntoRuntimeSnapshot(t *testing.T) { if got := snapshot.Groups[group.ID].ValidationModel; got != test.want { t.Fatalf("ValidationModel = %q, want %q", got, test.want) } - if got := snapshot.Candidates[protocol.OpenAICompletions]["public-model"][0].UpstreamModelID; got != "real-model" { + if got := snapshot.ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion]["public-model"][0].UpstreamModelID; got != "real-model" { t.Fatalf("candidate upstream model = %q, want real-model", got) } }) @@ -523,30 +491,30 @@ func TestLoaderMapsValidationModelIntoRuntimeSnapshot(t *testing.T) { func TestLoaderRejectsInvalidGroupRowsWithoutPublishing(t *testing.T) { tests := []struct { name string - protocols models.JSON + channelID string + params models.JSON models models.JSON - config models.JSON + overrides models.JSON wantError string }{ - {name: "protocols object", protocols: models.JSON(`{}`), models: models.JSON(`[]`), config: models.JSON(`{}`), wantError: "protocols"}, - {name: "models object", protocols: models.JSON(`["openai-completions"]`), models: models.JSON(`{}`), config: models.JSON(`{}`), wantError: "models"}, - {name: "config array", protocols: models.JSON(`["openai-completions"]`), models: models.JSON(`[]`), config: models.JSON(`[]`), wantError: "config"}, - {name: "unknown group setting", protocols: models.JSON(`["openai-completions"]`), models: models.JSON(`[{"id":"gpt-4o"}]`), config: models.JSON(`{"unknown":true}`), wantError: "unknown group setting"}, - {name: "duplicate protocol", protocols: models.JSON(`["openai-completions","openai-completions"]`), models: models.JSON(`[{"id":"gpt-4o"}]`), config: models.JSON(`{}`), wantError: "duplicate protocol"}, - {name: "blank model id", protocols: models.JSON(`["openai-completions"]`), models: models.JSON(`[{"id":" "}]`), config: models.JSON(`{}`), wantError: "model id is required"}, + {name: "unknown channel", channelID: "unknown", params: models.JSON(`{}`), models: models.JSON(`[]`), overrides: models.JSON(`{}`), wantError: "unknown channel"}, + {name: "invalid channel params", channelID: string(channel.OpenAICompatible), params: models.JSON(`{}`), models: models.JSON(`[]`), overrides: models.JSON(`{}`), wantError: "params.base_url"}, + {name: "models object", channelID: string(channel.OpenAI), params: models.JSON(`{}`), models: models.JSON(`{}`), overrides: models.JSON(`{}`), wantError: "models"}, + {name: "overrides array", channelID: string(channel.OpenAI), params: models.JSON(`{}`), models: models.JSON(`[]`), overrides: models.JSON(`[]`), wantError: "overrides"}, + {name: "unknown group setting", channelID: string(channel.OpenAI), params: models.JSON(`{}`), models: models.JSON(`[{"id":"gpt-4o"}]`), overrides: models.JSON(`{"unknown":true}`), wantError: "unknown group setting"}, + {name: "blank model id", channelID: string(channel.OpenAI), params: models.JSON(`{}`), models: models.JSON(`[{"id":" "}]`), overrides: models.JSON(`{}`), wantError: "model id is required"}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { db := openMigratedDatabase(t) group := models.Group{ - Name: "invalid", UpstreamURL: "https://invalid.example.com/v1", - Protocols: test.protocols, - Models: test.models, Config: test.config, Enabled: true, + Name: "invalid", ChannelID: test.channelID, Params: test.params, + Models: test.models, Overrides: test.overrides, Enabled: true, } mustCreate(t, db, &group) manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() err := loader.New(db, manager, registry).Load(context.Background()) if err == nil { @@ -558,14 +526,14 @@ func TestLoaderRejectsInvalidGroupRowsWithoutPublishing(t *testing.T) { if manager.Current() != nil { t.Fatalf("Current() = %#v after failed load, want nil", manager.Current()) } - if got := registry.CollectCandidates([]uint{group.ID}, nil, time.Time{}); len(got) != 0 { + if got := registry.CollectCredentialCandidates([]uint{group.ID}, nil, time.Time{}); len(got) != 0 { t.Fatalf("registry candidates after failed load = %#v, want empty", got) } }) } } -func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) { +func TestLoaderMapsAccessAndCredentials(t *testing.T) { db := openMigratedDatabase(t) firstGroup := createRuntimeGroup(t, db, "first", protocol.OpenAICompletions, "gpt-4o") if err := db.Model(&firstGroup).Update("models", models.JSON(`[{"id":"gpt-4o","alias":"Primary"}]`)).Error; err != nil { @@ -589,18 +557,18 @@ func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) { mustCreate(t, db, &disabledAccess) firstWeight := 7 - keys := []models.UpstreamKey{ - {GroupID: firstGroup.ID, KeyValue: "upstream-cipher-one", KeyHash: "upstream-hash-one", Status: models.UpstreamKeyStatusActive, WeightManual: &firstWeight}, - {GroupID: firstGroup.ID, KeyValue: "upstream-cipher-two", KeyHash: "upstream-hash-two", Status: models.UpstreamKeyStatusDisabled}, - {GroupID: secondGroup.ID, KeyValue: "upstream-cipher-three", KeyHash: "upstream-hash-three", Status: models.UpstreamKeyStatusActive}, - {GroupID: secondGroup.ID, KeyValue: "upstream-cipher-four", KeyHash: "upstream-hash-four", Status: models.UpstreamKeyStatusDisabled}, + credentials := []models.Credential{ + {GroupID: firstGroup.ID, Data: "credential-cipher-one", Fingerprint: "credential-fingerprint-one", Status: models.CredentialStatusActive, WeightManual: &firstWeight}, + {GroupID: firstGroup.ID, Data: "credential-cipher-two", Fingerprint: "credential-fingerprint-two", Status: models.CredentialStatusDisabled}, + {GroupID: secondGroup.ID, Data: "credential-cipher-three", Fingerprint: "credential-fingerprint-three", Status: models.CredentialStatusActive}, + {GroupID: secondGroup.ID, Data: "credential-cipher-four", Fingerprint: "credential-fingerprint-four", Status: models.CredentialStatusDisabled}, } - for index := range keys { - mustCreate(t, db, &keys[index]) + for index := range credentials { + mustCreate(t, db, &credentials[index]) } manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } @@ -627,7 +595,7 @@ func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) { got.Name != secondGroup.Name || !got.Enabled { t.Fatalf("second GroupCatalog entry = %#v", got) } - if len(snapshot.RouteCatalog) == 0 { + if len(snapshot.ExecutionRouteCatalog) == 0 { t.Fatal("RouteCatalog was not compiled from persisted groups") } if _, ok := access.Filters.Groups[firstGroup.ID]; !ok { @@ -646,27 +614,27 @@ func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) { t.Errorf("access filters models = %#v, must not expose hidden upstream id", access.Filters.Models) } - candidates := registry.CollectCandidates([]uint{firstGroup.ID, secondGroup.ID}, nil, time.Time{}) + candidates := registry.CollectCredentialCandidates([]uint{firstGroup.ID, secondGroup.ID}, nil, time.Time{}) if len(candidates) != 2 { - t.Fatalf("registry candidates = %#v, want two active keys", candidates) + t.Fatalf("registry candidates = %#v, want two active credentials", candidates) } - if candidates[0].ID != keys[0].ID || candidates[0].GroupID != firstGroup.ID || candidates[0].WeightManual == nil || *candidates[0].WeightManual != firstWeight { - t.Errorf("first candidate = %#v, want active weighted key %d", candidates[0], keys[0].ID) + if candidates[0].ID != credentials[0].ID || candidates[0].GroupID != firstGroup.ID || candidates[0].WeightManual == nil || *candidates[0].WeightManual != firstWeight { + t.Errorf("first candidate = %#v, want active weighted credential %d", candidates[0], credentials[0].ID) } - if candidates[1].ID != keys[2].ID || candidates[1].GroupID != secondGroup.ID { - t.Errorf("second candidate = %#v, want active key %d", candidates[1], keys[2].ID) + if candidates[1].ID != credentials[2].ID || candidates[1].GroupID != secondGroup.ID { + t.Errorf("second candidate = %#v, want active credential %d", candidates[1], credentials[2].ID) } - for _, key := range keys { - got, ok := registry.EncryptedValue(key.ID) - if !ok || got != key.KeyValue { - t.Errorf("EncryptedValue(%d) = %q, %t, want %q, true", key.ID, got, ok, key.KeyValue) + for _, credential := range credentials { + got, ok := registry.EncryptedCredentialData(credential.ID) + if !ok || got != credential.Data { + t.Errorf("EncryptedValue(%d) = %q, %t, want %q, true", credential.ID, got, ok, credential.Data) } } snapshotText := fmt.Sprintf("%#v", snapshot) for _, secret := range []string{ activeAccess.KeyValue, disabledAccess.KeyValue, - keys[0].KeyValue, keys[1].KeyValue, keys[2].KeyValue, keys[3].KeyValue, + credentials[0].Data, credentials[1].Data, credentials[2].Data, credentials[3].Data, } { if strings.Contains(snapshotText, secret) { t.Errorf("snapshot exposes credential material %q", secret) @@ -674,48 +642,48 @@ func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) { } } -func TestBuildGroupKeyEntriesReadsOnlyRequestedGroupInStableOrder(t *testing.T) { +func TestBuildGroupCredentialEntriesReadsOnlyRequestedGroupInStableOrder(t *testing.T) { db := openMigratedDatabase(t) firstGroup := createRuntimeGroup(t, db, "first-entries", protocol.OpenAICompletions, "gpt-4o") secondGroup := createRuntimeGroup(t, db, "second-entries", protocol.Anthropic, "claude") weight := 9 - keys := []models.UpstreamKey{ + credentials := []models.Credential{ { - GroupID: secondGroup.ID, KeyValue: "other-cipher", - KeyHash: "other-hash", Status: models.UpstreamKeyStatusActive, + GroupID: secondGroup.ID, Data: "other-cipher", + Fingerprint: "other-fingerprint", Status: models.CredentialStatusActive, }, { - GroupID: firstGroup.ID, KeyValue: "first-cipher", - KeyHash: "first-hash", Status: models.UpstreamKeyStatusActive, + GroupID: firstGroup.ID, Data: "first-cipher", + Fingerprint: "first-fingerprint", Status: models.CredentialStatusActive, }, { - GroupID: firstGroup.ID, KeyValue: "second-cipher", - KeyHash: "second-hash", Status: models.UpstreamKeyStatusDisabled, + GroupID: firstGroup.ID, Data: "second-cipher", + Fingerprint: "second-fingerprint", Status: models.CredentialStatusDisabled, WeightManual: &weight, }, } - for index := range keys { - mustCreate(t, db, &keys[index]) + for index := range credentials { + mustCreate(t, db, &credentials[index]) } - got, err := loader.BuildGroupKeyEntries(t.Context(), db, firstGroup.ID) + got, err := loader.BuildGroupCredentialEntries(t.Context(), db, firstGroup.ID) if err != nil { - t.Fatalf("BuildGroupKeyEntries() error = %v", err) + t.Fatalf("BuildGroupCredentialEntries() error = %v", err) } - if len(got) != 2 || got[0].ID != keys[1].ID || got[1].ID != keys[2].ID { - t.Fatalf("BuildGroupKeyEntries() = %#v", got) + if len(got) != 2 || got[0].ID != credentials[1].ID || got[1].ID != credentials[2].ID { + t.Fatalf("BuildGroupCredentialEntries() = %#v", got) } if got[1].WeightManual == nil || *got[1].WeightManual != weight || - got[1].Status != state.KeyStatusDisabled || + got[1].Status != state.CredentialStatusDisabled || got[1].EncryptedValue != "second-cipher" { t.Fatalf("second entry = %#v", got[1]) } } -func TestBuildGroupKeyEntriesRejectsMissingGroup(t *testing.T) { +func TestBuildGroupCredentialEntriesRejectsMissingGroup(t *testing.T) { db := openMigratedDatabase(t) - if _, err := loader.BuildGroupKeyEntries(t.Context(), db, 999); err == nil { - t.Fatal("BuildGroupKeyEntries(missing group) error = nil") + if _, err := loader.BuildGroupCredentialEntries(t.Context(), db, 999); err == nil { + t.Fatal("BuildGroupCredentialEntries(missing group) error = nil") } } @@ -728,7 +696,7 @@ func TestLoaderMapsAccessKeyRPMLimit(t *testing.T) { mustCreate(t, db, &accessKey) manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } @@ -795,11 +763,11 @@ func TestLoaderRejectsInvalidCredentialRowsWithoutPublishing(t *testing.T) { wantError: "unknown field", }, { - name: "empty upstream ciphertext", + name: "empty credential ciphertext", insert: func(t *testing.T, db *gorm.DB, group models.Group) { - mustCreate(t, db, &models.UpstreamKey{ - GroupID: group.ID, KeyValue: "", KeyHash: "empty-cipher-hash", - Status: models.UpstreamKeyStatusActive, + mustCreate(t, db, &models.Credential{ + GroupID: group.ID, Data: "", Fingerprint: "empty-cipher-fingerprint", + Status: models.CredentialStatusActive, }) }, wantError: "encrypted value is required", @@ -812,7 +780,7 @@ func TestLoaderRejectsInvalidCredentialRowsWithoutPublishing(t *testing.T) { group := createRuntimeGroup(t, db, "valid", protocol.OpenAICompletions, "gpt-4o") test.insert(t, db, group) manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() err := loader.New(db, manager, registry).Load(context.Background()) if err == nil { @@ -824,7 +792,7 @@ func TestLoaderRejectsInvalidCredentialRowsWithoutPublishing(t *testing.T) { if manager.Current() != nil { t.Fatalf("Current() = %#v after failed load, want nil", manager.Current()) } - if got := registry.CollectCandidates([]uint{group.ID}, nil, time.Time{}); len(got) != 0 { + if got := registry.CollectCredentialCandidates([]uint{group.ID}, nil, time.Time{}); len(got) != 0 { t.Fatalf("registry candidates after failed load = %#v, want empty", got) } }) @@ -862,10 +830,16 @@ func mustCreate(t *testing.T, db *gorm.DB, value any) { func createRuntimeGroup(t *testing.T, db *gorm.DB, name string, p protocol.Protocol, model string) models.Group { t.Helper() + channelID := channel.OpenAI + switch p { + case protocol.Anthropic: + channelID = channel.Anthropic + case protocol.Gemini: + channelID = channel.Gemini + } group := models.Group{ - Name: name, UpstreamURL: "https://" + name + ".example.com/v1", - Protocols: models.JSON(fmt.Sprintf(`[%q]`, p)), - Models: models.JSON(fmt.Sprintf(`[{"id":%q}]`, model)), Config: models.JSON(`{}`), Enabled: true, + Name: name, ChannelID: string(channelID), Params: models.JSON(`{}`), + Models: models.JSON(fmt.Sprintf(`[{"id":%q}]`, model)), Overrides: models.JSON(`{}`), Enabled: true, } mustCreate(t, db, &group) return group diff --git a/internal/state/loader/provider_id_test.go b/internal/state/loader/provider_id_test.go deleted file mode 100644 index ace1f9f3..00000000 --- a/internal/state/loader/provider_id_test.go +++ /dev/null @@ -1,32 +0,0 @@ -package loader - -import ( - "testing" - - "gpt-load/internal/storage/models" -) - -func TestMapSystemAndGroupsDeepClonesProviderID(t *testing.T) { - providerID := "openai" - rows := compileRows{groups: []models.Group{{ - ID: 1, - Name: "provider-group", - ProviderID: &providerID, - UpstreamURL: "https://api.openai.com/v1", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[{"id":"gpt-4o","alias":"public"}]`), - Config: models.JSON(`{}`), - Enabled: true, - }}} - - input, err := mapSystemAndGroups(rows) - if err != nil { - t.Fatalf("mapSystemAndGroups() error = %v", err) - } - providerID = "mutated-source" - *rows.groups[0].ProviderID = "mutated-row" - - if got := input.Groups[0].ProviderID; got == nil || *got != "openai" { - t.Fatalf("GroupConfig.ProviderID = %v, want independent openai", got) - } -} diff --git a/internal/state/loader/weights_test.go b/internal/state/loader/weights_test.go index 21360903..a7d3b479 100644 --- a/internal/state/loader/weights_test.go +++ b/internal/state/loader/weights_test.go @@ -6,6 +6,8 @@ import ( "testing" "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/state/loader" @@ -16,22 +18,22 @@ func TestLoaderMapsSchedulingWeights(t *testing.T) { db := openMigratedDatabase(t) groupWeight := 25 group := models.Group{ - Name: "weighted", UpstreamURL: "https://weighted.example.com", - Protocols: models.JSON(`["openai-completions"]`), Models: models.JSON(`[{"id":"gpt-weighted"}]`), - Config: models.JSON(`{}`), Enabled: true, WeightManual: &groupWeight, + Name: "weighted", ChannelID: string(channel.OpenAI), Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"gpt-weighted"}]`), + Overrides: models.JSON(`{}`), Enabled: true, WeightManual: &groupWeight, } mustCreate(t, db, &group) keyWeight := 30 - keys := []models.UpstreamKey{ - {GroupID: group.ID, KeyValue: "cipher-manual", KeyHash: "hash-manual", Status: models.UpstreamKeyStatusActive, WeightManual: &keyWeight}, - {GroupID: group.ID, KeyValue: "cipher-default", KeyHash: "hash-default", Status: models.UpstreamKeyStatusActive}, + credentials := []models.Credential{ + {GroupID: group.ID, Data: "cipher-manual", Fingerprint: "fingerprint-manual", Status: models.CredentialStatusActive, WeightManual: &keyWeight}, + {GroupID: group.ID, Data: "cipher-default", Fingerprint: "fingerprint-default", Status: models.CredentialStatusActive}, } - for index := range keys { - mustCreate(t, db, &keys[index]) + for index := range credentials { + mustCreate(t, db, &credentials[index]) } manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } @@ -40,19 +42,19 @@ func TestLoaderMapsSchedulingWeights(t *testing.T) { t.Fatalf("GroupView.WeightManual = %v, want %d", view.WeightManual, groupWeight) } - candidates := registry.CollectCandidates([]uint{group.ID}, nil, time.Time{}) + candidates := registry.CollectCredentialCandidates([]uint{group.ID}, nil, time.Time{}) if len(candidates) != 2 { - t.Fatalf("CollectCandidates() = %#v, want two keys", candidates) + t.Fatalf("CollectCandidates() = %#v, want two credentials", candidates) } if candidates[0].WeightManual == nil || *candidates[0].WeightManual != keyWeight { t.Fatalf("first WeightManual = %v, want %d", candidates[0].WeightManual, keyWeight) } for _, candidate := range candidates { if candidate.WeightAuto != state.DefaultWeight { - t.Errorf("key %d WeightAuto = %d, want %d", candidate.ID, candidate.WeightAuto, state.DefaultWeight) + t.Errorf("credential %d WeightAuto = %d, want %d", candidate.ID, candidate.WeightAuto, state.DefaultWeight) } } - if got := manager.Current().Candidates[protocol.OpenAICompletions]["gpt-weighted"]; len(got) != 1 { + if got := manager.Current().ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion]["gpt-weighted"]; len(got) != 1 { t.Fatalf("route candidates = %#v, want one group", got) } } @@ -62,19 +64,19 @@ func TestLoaderPreservesManualWeightBoundaries(t *testing.T) { t.Run(fmt.Sprintf("weight_%d", weight), func(t *testing.T) { db := openMigratedDatabase(t) group := models.Group{ - Name: "boundary", UpstreamURL: "https://boundary.example.com", - Protocols: models.JSON(`["openai-completions"]`), Models: models.JSON(`[{"id":"gpt-boundary"}]`), - Config: models.JSON(`{}`), Enabled: true, WeightManual: &weight, + Name: "boundary", ChannelID: string(channel.OpenAI), Params: models.JSON(`{}`), + Models: models.JSON(`[{"id":"gpt-boundary"}]`), + Overrides: models.JSON(`{}`), Enabled: true, WeightManual: &weight, } mustCreate(t, db, &group) - key := models.UpstreamKey{ - GroupID: group.ID, KeyValue: "cipher", KeyHash: "hash", - Status: models.UpstreamKeyStatusActive, WeightManual: &weight, + credential := models.Credential{ + GroupID: group.ID, Data: "cipher", Fingerprint: "fingerprint", + Status: models.CredentialStatusActive, WeightManual: &weight, } - mustCreate(t, db, &key) + mustCreate(t, db, &credential) manager := state.NewManager() - registry := state.NewKeyRegistry() + registry := state.NewCredentialRegistry() if err := loader.New(db, manager, registry).Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } @@ -83,7 +85,7 @@ func TestLoaderPreservesManualWeightBoundaries(t *testing.T) { t.Fatalf("GroupView.WeightManual = %v, want explicit %d", view.WeightManual, weight) } - candidates := registry.CollectCandidates([]uint{group.ID}, nil, time.Time{}) + candidates := registry.CollectCredentialCandidates([]uint{group.ID}, nil, time.Time{}) if len(candidates) != 1 || candidates[0].WeightManual == nil || *candidates[0].WeightManual != weight { t.Fatalf("CollectCandidates() = %#v, want explicit manual weight %d", candidates, weight) } diff --git a/internal/state/manager_test.go b/internal/state/manager_test.go index cc6966bf..b481b611 100644 --- a/internal/state/manager_test.go +++ b/internal/state/manager_test.go @@ -1,11 +1,12 @@ package state import ( + "encoding/json" "sync" "testing" "time" - "gpt-load/internal/protocol" + "gpt-load/internal/channel" ) func TestManagerPublishStartsAtRevisionOne(t *testing.T) { @@ -56,7 +57,7 @@ func TestManagerPublishFailureKeepsCurrentSnapshot(t *testing.T) { } invalid := managerCompileInput(2) - invalid.Groups[0].Protocols = []protocol.Protocol{protocol.Protocol("invalid")} + invalid.Groups[0].ChannelID = channel.ID("invalid") if _, err := manager.Publish(invalid); err == nil { t.Fatal("Publish() error = nil, want invalid protocol error") } @@ -94,7 +95,7 @@ func TestManagerMatchesCompiledInputWithoutChangingRevision(t *testing.T) { } invalid := managerCompileInput(1) - invalid.Groups[0].Protocols = []protocol.Protocol{"reserved"} + invalid.Groups[0].Params = json.RawMessage(`{"unexpected":true}`) if _, err := manager.Matches(invalid); err == nil { t.Fatal("Matches(invalid) error = nil, want compile rejection") } @@ -248,12 +249,11 @@ func awaitManagerSignal(t *testing.T, signal <-chan struct{}, name string) { } func managerCompileInput(groupID uint) CompileInput { - return CompileInput{Groups: []GroupConfig{{ - ID: groupID, - Name: "group", - UpstreamURL: "https://upstream.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "model"}}, - Enabled: true, - }}} + return CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []GroupConfig{{ + ID: groupID, Name: "group", ChannelID: channel.OpenAI, + Params: json.RawMessage(`{}`), Models: []ModelConfig{{ID: "model"}}, Enabled: true, + }}, + } } diff --git a/internal/state/registry.go b/internal/state/registry.go index b184dd62..9235e92d 100644 --- a/internal/state/registry.go +++ b/internal/state/registry.go @@ -3,198 +3,216 @@ package state import ( "fmt" "sort" + "strings" "sync" "time" ) -type KeyStatus string +type CredentialStatus string const ( - KeyStatusActive KeyStatus = "active" - KeyStatusDisabled KeyStatus = "disabled" + CredentialStatusActive CredentialStatus = "active" + CredentialStatusDisabled CredentialStatus = "disabled" ) -type KeyEntry struct { - ID uint - GroupID uint - WeightManual *int - WeightAuto int - Status KeyStatus - CooldownUntil time.Time - Blacklisted bool - FailureCount int - FailureGeneration uint64 - EncryptedValue string +type CredentialEntry struct { + ID uint + GroupID uint + Version uint64 + IdentityGeneration uint64 + Fingerprint string + WeightManual *int + WeightAuto int + Status CredentialStatus + CooldownUntil time.Time + Blacklisted bool + FailureCount int + FailureGeneration uint64 + EncryptedValue string } -type KeyMeta struct { - ID uint - GroupID uint - WeightManual *int - WeightAuto int +type CredentialMeta struct { + ID uint + GroupID uint + Version uint64 + IdentityGeneration uint64 + WeightManual *int + WeightAuto int } -type KeyRef struct { - ID uint - GroupID uint - EncryptedValue string - FailureGeneration uint64 +type CredentialRef struct { + ID uint + GroupID uint + Version uint64 + IdentityGeneration uint64 + Fingerprint string + EncryptedValue string + FailureGeneration uint64 } -type KeyRegistry struct { - mu sync.RWMutex - buckets map[uint]map[uint]*KeyEntry - keyGroups map[uint]uint +type CredentialRegistry struct { + mu sync.RWMutex + buckets map[uint]map[uint]*CredentialEntry + credentialGroups map[uint]uint } -func NewKeyRegistry() *KeyRegistry { - return &KeyRegistry{ - buckets: make(map[uint]map[uint]*KeyEntry), - keyGroups: make(map[uint]uint), +func NewCredentialRegistry() *CredentialRegistry { + return &CredentialRegistry{ + buckets: make(map[uint]map[uint]*CredentialEntry), + credentialGroups: make(map[uint]uint), } } -func ValidateKeyEntries(entries []KeyEntry) error { +func ValidateCredentialEntries(entries []CredentialEntry) error { seen := make(map[uint]struct{}, len(entries)) for _, entry := range entries { if entry.ID == 0 { - return fmt.Errorf("key id is required") + return fmt.Errorf("credential id is required") } if entry.GroupID == 0 { - return fmt.Errorf("key %d group id is required", entry.ID) + return fmt.Errorf("credential %d group id is required", entry.ID) } - if entry.Status != KeyStatusActive && entry.Status != KeyStatusDisabled { - return fmt.Errorf("key %d has invalid status %q", entry.ID, entry.Status) + if entry.Status != CredentialStatusActive && entry.Status != CredentialStatusDisabled { + return fmt.Errorf("credential %d has invalid status %q", entry.ID, entry.Status) } - if err := validateManualWeight(fmt.Sprintf("key %d", entry.ID), entry.WeightManual); err != nil { + if err := validateManualWeight(fmt.Sprintf("credential %d", entry.ID), entry.WeightManual); err != nil { return err } if entry.WeightAuto < 0 || entry.WeightAuto > MaxWeight { - return fmt.Errorf("key %d auto weight must be between 1 and %d", entry.ID, MaxWeight) + return fmt.Errorf("credential %d auto weight must be between 0 and %d", entry.ID, MaxWeight) } if entry.EncryptedValue == "" { - return fmt.Errorf("key %d encrypted value is required", entry.ID) + return fmt.Errorf("credential %d encrypted value is required", entry.ID) + } + if entry.Version == 0 { + return fmt.Errorf("credential %d version is required", entry.ID) + } + if entry.IdentityGeneration == 0 { + return fmt.Errorf("credential %d identity generation is required", entry.ID) + } + if strings.TrimSpace(entry.Fingerprint) == "" { + return fmt.Errorf("credential %d fingerprint is required", entry.ID) } if _, duplicate := seen[entry.ID]; duplicate { - return fmt.Errorf("duplicate key id %d", entry.ID) + return fmt.Errorf("duplicate credential id %d", entry.ID) } seen[entry.ID] = struct{}{} } return nil } -func (r *KeyRegistry) Replace(entries []KeyEntry) error { - if err := ValidateKeyEntries(entries); err != nil { +func (r *CredentialRegistry) ReplaceCredentials(entries []CredentialEntry) error { + if err := ValidateCredentialEntries(entries); err != nil { return err } - buckets := make(map[uint]map[uint]*KeyEntry) - keyGroups := make(map[uint]uint, len(entries)) + buckets := make(map[uint]map[uint]*CredentialEntry) + credentialGroups := make(map[uint]uint, len(entries)) for _, entry := range entries { if buckets[entry.GroupID] == nil { - buckets[entry.GroupID] = make(map[uint]*KeyEntry) + buckets[entry.GroupID] = make(map[uint]*CredentialEntry) } - cloned := cloneKeyEntry(entry) + cloned := cloneCredentialEntry(entry) buckets[entry.GroupID][entry.ID] = &cloned - keyGroups[entry.ID] = entry.GroupID + credentialGroups[entry.ID] = entry.GroupID } r.mu.Lock() r.buckets = buckets - r.keyGroups = keyGroups + r.credentialGroups = credentialGroups r.mu.Unlock() return nil } -func (r *KeyRegistry) ApplyImport(groupID uint, entries []KeyEntry) error { +func (r *CredentialRegistry) ApplyCredentialImport(groupID uint, entries []CredentialEntry) error { if groupID == 0 { return fmt.Errorf("group id is required") } - if err := ValidateKeyEntries(entries); err != nil { + if err := ValidateCredentialEntries(entries); err != nil { return err } for _, entry := range entries { if entry.GroupID != groupID { - return fmt.Errorf("key %d belongs to group %d, want %d", entry.ID, entry.GroupID, groupID) + return fmt.Errorf("credential %d belongs to group %d, want %d", entry.ID, entry.GroupID, groupID) } } r.mu.Lock() defer r.mu.Unlock() for _, entry := range entries { - if existingGroupID, exists := r.keyGroups[entry.ID]; exists && existingGroupID != groupID { - return fmt.Errorf("key %d already belongs to group %d", entry.ID, existingGroupID) + if existingGroupID, exists := r.credentialGroups[entry.ID]; exists && existingGroupID != groupID { + return fmt.Errorf("credential %d already belongs to group %d", entry.ID, existingGroupID) } } for _, entry := range entries { if r.buckets[groupID] == nil { - r.buckets[groupID] = make(map[uint]*KeyEntry) + r.buckets[groupID] = make(map[uint]*CredentialEntry) } - cloned := cloneKeyEntry(entry) + cloned := cloneCredentialEntry(entry) r.buckets[groupID][entry.ID] = &cloned - r.keyGroups[entry.ID] = groupID + r.credentialGroups[entry.ID] = groupID } return nil } -// SnapshotGroupKeyEntriesExact returns detached selected entries without -// applying the configuration replacement resets used by cloneKeyEntry. -func (r *KeyRegistry) SnapshotGroupKeyEntriesExact( +// SnapshotGroupCredentialEntriesExact returns detached selected entries without +// applying the configuration replacement resets used by cloneCredentialEntry. +func (r *CredentialRegistry) SnapshotGroupCredentialEntriesExact( groupID uint, - keyIDs []uint, -) ([]KeyEntry, error) { + credentialIDs []uint, +) ([]CredentialEntry, error) { r.mu.RLock() defer r.mu.RUnlock() - if err := r.validateGroupKeyIDsLocked(groupID, keyIDs); err != nil { + if err := r.validateGroupCredentialIDsLocked(groupID, credentialIDs); err != nil { return nil, err } - entries := make([]KeyEntry, 0, len(keyIDs)) - for _, keyID := range keyIDs { - entries = append(entries, detachKeyEntryExact(*r.buckets[groupID][keyID])) + entries := make([]CredentialEntry, 0, len(credentialIDs)) + for _, credentialID := range credentialIDs { + entries = append(entries, detachCredentialEntryExact(*r.buckets[groupID][credentialID])) } return entries, nil } -// RestoreGroupKeyEntriesExact atomically restores only selected entries, +// RestoreGroupCredentialEntriesExact atomically restores only selected entries, // including their runtime failure generation. Unselected entries are untouched. -func (r *KeyRegistry) RestoreGroupKeyEntriesExact(groupID uint, entries []KeyEntry) error { +func (r *CredentialRegistry) RestoreGroupCredentialEntriesExact(groupID uint, entries []CredentialEntry) error { if groupID == 0 { return fmt.Errorf("group id is required") } if len(entries) == 0 { - return fmt.Errorf("key entries are required") + return fmt.Errorf("credential entries are required") } - if err := ValidateKeyEntries(entries); err != nil { + if err := ValidateCredentialEntries(entries); err != nil { return err } for _, entry := range entries { if entry.GroupID != groupID { - return fmt.Errorf("key %d belongs to group %d, want %d", entry.ID, entry.GroupID, groupID) + return fmt.Errorf("credential %d belongs to group %d, want %d", entry.ID, entry.GroupID, groupID) } } r.mu.Lock() defer r.mu.Unlock() for _, entry := range entries { - if existingGroupID, exists := r.keyGroups[entry.ID]; exists && existingGroupID != groupID { - return fmt.Errorf("key %d already belongs to group %d", entry.ID, existingGroupID) + if existingGroupID, exists := r.credentialGroups[entry.ID]; exists && existingGroupID != groupID { + return fmt.Errorf("credential %d already belongs to group %d", entry.ID, existingGroupID) } } if r.buckets[groupID] == nil { - r.buckets[groupID] = make(map[uint]*KeyEntry) + r.buckets[groupID] = make(map[uint]*CredentialEntry) } for _, entry := range entries { - detached := detachKeyEntryExact(entry) + detached := detachCredentialEntryExact(entry) r.buckets[groupID][entry.ID] = &detached - r.keyGroups[entry.ID] = groupID + r.credentialGroups[entry.ID] = groupID } return nil } // MatchesGroup compares the persisted configuration-owned fields for one // group. Runtime health state is intentionally excluded. -func (r *KeyRegistry) MatchesGroup(groupID uint, entries []KeyEntry) bool { - if groupID == 0 || ValidateKeyEntries(entries) != nil { +func (r *CredentialRegistry) MatchesGroup(groupID uint, entries []CredentialEntry) bool { + if groupID == 0 || ValidateCredentialEntries(entries) != nil { return false } for _, entry := range entries { @@ -210,17 +228,17 @@ func (r *KeyRegistry) MatchesGroup(groupID uint, entries []KeyEntry) bool { // ReconcileGroup makes one group match DB truth while preserving runtime // health state for entries whose persisted configuration is already equal. -func (r *KeyRegistry) ReconcileGroup(groupID uint, entries []KeyEntry) (bool, error) { +func (r *CredentialRegistry) ReconcileGroup(groupID uint, entries []CredentialEntry) (bool, error) { if groupID == 0 { return false, fmt.Errorf("group id is required") } - if err := ValidateKeyEntries(entries); err != nil { + if err := ValidateCredentialEntries(entries); err != nil { return false, err } for _, entry := range entries { if entry.GroupID != groupID { return false, fmt.Errorf( - "key %d belongs to group %d, want %d", + "credential %d belongs to group %d, want %d", entry.ID, entry.GroupID, groupID, @@ -231,10 +249,10 @@ func (r *KeyRegistry) ReconcileGroup(groupID uint, entries []KeyEntry) (bool, er r.mu.Lock() defer r.mu.Unlock() for _, entry := range entries { - if existingGroupID, exists := r.keyGroups[entry.ID]; exists && + if existingGroupID, exists := r.credentialGroups[entry.ID]; exists && existingGroupID != groupID { return false, fmt.Errorf( - "key %d already belongs to group %d", + "credential %d already belongs to group %d", entry.ID, existingGroupID, ) @@ -245,47 +263,50 @@ func (r *KeyRegistry) ReconcileGroup(groupID uint, entries []KeyEntry) (bool, er } previous := r.buckets[groupID] - next := make(map[uint]*KeyEntry, len(entries)) + next := make(map[uint]*CredentialEntry, len(entries)) for _, desired := range entries { if existing := previous[desired.ID]; existing != nil && - samePersistedKeyConfig(*existing, desired) { + samePersistedCredentialConfig(*existing, desired) { next[desired.ID] = existing continue } - cloned := cloneKeyEntry(desired) + cloned := cloneCredentialEntry(desired) next[desired.ID] = &cloned } - for keyID := range previous { - delete(r.keyGroups, keyID) + for credentialID := range previous { + delete(r.credentialGroups, credentialID) } if len(next) == 0 { delete(r.buckets, groupID) } else { r.buckets[groupID] = next - for keyID := range next { - r.keyGroups[keyID] = groupID + for credentialID := range next { + r.credentialGroups[credentialID] = groupID } } return true, nil } -func (r *KeyRegistry) matchesGroupLocked(groupID uint, entries []KeyEntry) bool { +func (r *CredentialRegistry) matchesGroupLocked(groupID uint, entries []CredentialEntry) bool { current := r.buckets[groupID] if len(current) != len(entries) { return false } for _, desired := range entries { existing := current[desired.ID] - if existing == nil || !samePersistedKeyConfig(*existing, desired) { + if existing == nil || !samePersistedCredentialConfig(*existing, desired) { return false } } return true } -func samePersistedKeyConfig(left, right KeyEntry) bool { +func samePersistedCredentialConfig(left, right CredentialEntry) bool { if left.ID != right.ID || left.GroupID != right.GroupID || + left.Version != right.Version || + left.IdentityGeneration != right.IdentityGeneration || + left.Fingerprint != right.Fingerprint || left.Status != right.Status || left.EncryptedValue != right.EncryptedValue { return false @@ -296,49 +317,49 @@ func samePersistedKeyConfig(left, right KeyEntry) bool { return *left.WeightManual == *right.WeightManual } -func (r *KeyRegistry) RemoveKey(keyID uint) bool { +func (r *CredentialRegistry) RemoveCredential(credentialID uint) bool { r.mu.Lock() defer r.mu.Unlock() - groupID, ok := r.keyGroups[keyID] + groupID, ok := r.credentialGroups[credentialID] if !ok { return false } - delete(r.buckets[groupID], keyID) + delete(r.buckets[groupID], credentialID) if len(r.buckets[groupID]) == 0 { delete(r.buckets, groupID) } - delete(r.keyGroups, keyID) + delete(r.credentialGroups, credentialID) return true } -func (r *KeyRegistry) UpdateGroupKeyStatuses( +func (r *CredentialRegistry) UpdateGroupCredentialStatuses( groupID uint, - keyIDs []uint, - status KeyStatus, + credentialIDs []uint, + status CredentialStatus, ) error { - if status != KeyStatusActive && status != KeyStatusDisabled { - return fmt.Errorf("invalid key status %q", status) + if status != CredentialStatusActive && status != CredentialStatusDisabled { + return fmt.Errorf("invalid credential status %q", status) } r.mu.Lock() defer r.mu.Unlock() - if err := r.validateGroupKeyIDsLocked(groupID, keyIDs); err != nil { + if err := r.validateGroupCredentialIDsLocked(groupID, credentialIDs); err != nil { return err } - for _, keyID := range keyIDs { - r.buckets[groupID][keyID].Status = status + for _, credentialID := range credentialIDs { + r.buckets[groupID][credentialID].Status = status } return nil } -func (r *KeyRegistry) RemoveGroupKeys(groupID uint, keyIDs []uint) error { +func (r *CredentialRegistry) RemoveGroupCredentials(groupID uint, credentialIDs []uint) error { r.mu.Lock() defer r.mu.Unlock() - if err := r.validateGroupKeyIDsLocked(groupID, keyIDs); err != nil { + if err := r.validateGroupCredentialIDsLocked(groupID, credentialIDs); err != nil { return err } - for _, keyID := range keyIDs { - delete(r.buckets[groupID], keyID) - delete(r.keyGroups, keyID) + for _, credentialID := range credentialIDs { + delete(r.buckets[groupID], credentialID) + delete(r.credentialGroups, credentialID) } if len(r.buckets[groupID]) == 0 { delete(r.buckets, groupID) @@ -346,37 +367,37 @@ func (r *KeyRegistry) RemoveGroupKeys(groupID uint, keyIDs []uint) error { return nil } -func (r *KeyRegistry) validateGroupKeyIDsLocked(groupID uint, keyIDs []uint) error { +func (r *CredentialRegistry) validateGroupCredentialIDsLocked(groupID uint, credentialIDs []uint) error { if groupID == 0 { return fmt.Errorf("group id is required") } - if len(keyIDs) == 0 { - return fmt.Errorf("key ids are required") + if len(credentialIDs) == 0 { + return fmt.Errorf("credential ids are required") } - seen := make(map[uint]struct{}, len(keyIDs)) - for _, keyID := range keyIDs { - if keyID == 0 { - return fmt.Errorf("key id is required") + seen := make(map[uint]struct{}, len(credentialIDs)) + for _, credentialID := range credentialIDs { + if credentialID == 0 { + return fmt.Errorf("credential id is required") } - if _, duplicate := seen[keyID]; duplicate { - return fmt.Errorf("duplicate key id %d", keyID) + if _, duplicate := seen[credentialID]; duplicate { + return fmt.Errorf("duplicate credential id %d", credentialID) } - seen[keyID] = struct{}{} - actualGroupID, exists := r.keyGroups[keyID] + seen[credentialID] = struct{}{} + actualGroupID, exists := r.credentialGroups[credentialID] if !exists { - return fmt.Errorf("key %d not found", keyID) + return fmt.Errorf("credential %d not found", credentialID) } if actualGroupID != groupID { - return fmt.Errorf("key %d belongs to group %d, want %d", keyID, actualGroupID, groupID) + return fmt.Errorf("credential %d belongs to group %d, want %d", credentialID, actualGroupID, groupID) } - if r.buckets[groupID][keyID] == nil { - return fmt.Errorf("key %d not found in group %d", keyID, groupID) + if r.buckets[groupID][credentialID] == nil { + return fmt.Errorf("credential %d not found in group %d", credentialID, groupID) } } return nil } -func (r *KeyRegistry) RemoveGroup(groupID uint) bool { +func (r *CredentialRegistry) RemoveGroup(groupID uint) bool { if groupID == 0 { return false } @@ -386,75 +407,77 @@ func (r *KeyRegistry) RemoveGroup(groupID uint) bool { if !exists { return false } - for keyID := range bucket { - delete(r.keyGroups, keyID) + for credentialID := range bucket { + delete(r.credentialGroups, credentialID) } delete(r.buckets, groupID) return true } -func (r *KeyRegistry) SetKeyStatus(keyID uint, status KeyStatus) error { - if status != KeyStatusActive && status != KeyStatusDisabled { - return fmt.Errorf("invalid key status %q", status) +func (r *CredentialRegistry) SetCredentialStatus(credentialID uint, status CredentialStatus) error { + if status != CredentialStatusActive && status != CredentialStatusDisabled { + return fmt.Errorf("invalid credential status %q", status) } r.mu.Lock() defer r.mu.Unlock() - groupID, ok := r.keyGroups[keyID] + groupID, ok := r.credentialGroups[credentialID] if !ok { - return fmt.Errorf("key %d not found", keyID) + return fmt.Errorf("credential %d not found", credentialID) } - r.buckets[groupID][keyID].Status = status + r.buckets[groupID][credentialID].Status = status return nil } -func (r *KeyRegistry) UpdateKeyConfig( - keyID uint, - status KeyStatus, +func (r *CredentialRegistry) UpdateCredentialConfig( + credentialID uint, + status CredentialStatus, weightManual *int, ) error { - if status != KeyStatusActive && status != KeyStatusDisabled { - return fmt.Errorf("invalid key status %q", status) + if status != CredentialStatusActive && status != CredentialStatusDisabled { + return fmt.Errorf("invalid credential status %q", status) } - if err := validateManualWeight(fmt.Sprintf("key %d", keyID), weightManual); err != nil { + if err := validateManualWeight(fmt.Sprintf("credential %d", credentialID), weightManual); err != nil { return err } clonedWeight := cloneWeight(weightManual) r.mu.Lock() defer r.mu.Unlock() - entry, ok := r.entryLocked(keyID) + entry, ok := r.entryLocked(credentialID) if !ok { - return fmt.Errorf("key %d not found", keyID) + return fmt.Errorf("credential %d not found", credentialID) } entry.Status = status entry.WeightManual = clonedWeight return nil } -func (r *KeyRegistry) EncryptedValue(keyID uint) (string, bool) { +// EncryptedCredentialData returns encrypted credential data for the selected +// stable credential ID. Decryption belongs to the gateway execution boundary. +func (r *CredentialRegistry) EncryptedCredentialData(credentialID uint) (string, bool) { r.mu.RLock() defer r.mu.RUnlock() - groupID, ok := r.keyGroups[keyID] + groupID, ok := r.credentialGroups[credentialID] if !ok { return "", false } - return r.buckets[groupID][keyID].EncryptedValue, true + return r.buckets[groupID][credentialID].EncryptedValue, true } -func (r *KeyRegistry) ActiveEncryptedValue(keyID, expectedGroupID uint) (string, bool) { +func (r *CredentialRegistry) ActiveEncryptedCredentialData(credentialID, expectedGroupID uint) (string, bool) { r.mu.RLock() defer r.mu.RUnlock() - groupID, ok := r.keyGroups[keyID] + groupID, ok := r.credentialGroups[credentialID] if !ok || groupID != expectedGroupID { return "", false } - entry, ok := r.buckets[groupID][keyID] - if !ok || entry.Status != KeyStatusActive { + entry, ok := r.buckets[groupID][credentialID] + if !ok || entry.Status != CredentialStatusActive { return "", false } return entry.EncryptedValue, true } -func (r *KeyRegistry) CaptureActiveKeyRefs(groupIDs []uint) []KeyRef { +func (r *CredentialRegistry) CaptureActiveCredentialRefs(groupIDs []uint) []CredentialRef { selectedGroups := make(map[uint]struct{}, len(groupIDs)) for _, groupID := range groupIDs { if groupID != 0 { @@ -463,14 +486,16 @@ func (r *KeyRegistry) CaptureActiveKeyRefs(groupIDs []uint) []KeyRef { } r.mu.RLock() - refs := make([]KeyRef, 0) + refs := make([]CredentialRef, 0) for groupID := range selectedGroups { for _, entry := range r.buckets[groupID] { - if entry.Status != KeyStatusActive { + if entry.Status != CredentialStatusActive { continue } - refs = append(refs, KeyRef{ - ID: entry.ID, GroupID: entry.GroupID, EncryptedValue: entry.EncryptedValue, + refs = append(refs, CredentialRef{ + ID: entry.ID, GroupID: entry.GroupID, + Version: entry.Version, IdentityGeneration: entry.IdentityGeneration, + Fingerprint: entry.Fingerprint, EncryptedValue: entry.EncryptedValue, FailureGeneration: entry.FailureGeneration, }) } @@ -485,10 +510,10 @@ func (r *KeyRegistry) CaptureActiveKeyRefs(groupIDs []uint) []KeyRef { return refs } -func (r *KeyRegistry) ActiveEncryptedValueIfMatch(ref KeyRef) (string, bool) { +func (r *CredentialRegistry) ActiveEncryptedCredentialDataIfMatch(ref CredentialRef) (string, bool) { r.mu.RLock() defer r.mu.RUnlock() - groupID, ok := r.keyGroups[ref.ID] + groupID, ok := r.credentialGroups[ref.ID] if !ok || groupID != ref.GroupID { return "", false } @@ -496,8 +521,11 @@ func (r *KeyRegistry) ActiveEncryptedValueIfMatch(ref KeyRef) (string, bool) { if !ok || entry.ID != ref.ID || entry.GroupID != ref.GroupID || + entry.Version != ref.Version || + entry.IdentityGeneration != ref.IdentityGeneration || + entry.Fingerprint != ref.Fingerprint || entry.EncryptedValue != ref.EncryptedValue || - entry.Status != KeyStatusActive { + entry.Status != CredentialStatusActive { return "", false } // FailureGeneration is intentionally excluded: failure accounting must not @@ -505,12 +533,12 @@ func (r *KeyRegistry) ActiveEncryptedValueIfMatch(ref KeyRef) (string, bool) { return entry.EncryptedValue, true } -func (r *KeyRegistry) ActiveKeyIDs() []uint { +func (r *CredentialRegistry) ActiveCredentialIDs() []uint { r.mu.RLock() - ids := make([]uint, 0, len(r.keyGroups)) + ids := make([]uint, 0, len(r.credentialGroups)) for _, bucket := range r.buckets { for _, entry := range bucket { - if entry.Status == KeyStatusActive { + if entry.Status == CredentialStatusActive { ids = append(ids, entry.ID) } } @@ -520,17 +548,19 @@ func (r *KeyRegistry) ActiveKeyIDs() []uint { return ids } -func (r *KeyRegistry) CollectCandidates(groupIDs []uint, excluded func(uint) bool, now time.Time) []KeyMeta { +// CollectCredentialCandidates returns currently schedulable credentials. +func (r *CredentialRegistry) CollectCredentialCandidates(groupIDs []uint, excluded func(uint) bool, now time.Time) []CredentialMeta { r.mu.RLock() - metas := make([]KeyMeta, 0) + metas := make([]CredentialMeta, 0) for _, groupID := range groupIDs { for _, entry := range r.buckets[groupID] { view := runtimeView(entry) - if view.RuntimeState(now) != KeyRuntimeAvailable { + if view.RuntimeState(now) != CredentialRuntimeAvailable { continue } - meta := KeyMeta{ + meta := CredentialMeta{ ID: view.ID, GroupID: view.GroupID, + Version: view.Version, IdentityGeneration: view.IdentityGeneration, WeightManual: cloneWeight(view.WeightManual), WeightAuto: view.WeightAuto, } metas = append(metas, meta) @@ -554,10 +584,10 @@ func (r *KeyRegistry) CollectCandidates(groupIDs []uint, excluded func(uint) boo return filtered } -func (r *KeyRegistry) SetCooldown(keyID uint, until time.Time) bool { +func (r *CredentialRegistry) SetCooldown(credentialID uint, until time.Time) bool { r.mu.Lock() defer r.mu.Unlock() - entry, ok := r.entryLocked(keyID) + entry, ok := r.entryLocked(credentialID) if !ok { return false } @@ -567,13 +597,13 @@ func (r *KeyRegistry) SetCooldown(keyID uint, until time.Time) bool { return true } -func (r *KeyRegistry) SetAutoWeight(keyID uint, weight int) bool { +func (r *CredentialRegistry) SetAutoWeight(credentialID uint, weight int) bool { if weight < 1 || weight > MaxWeight { return false } r.mu.Lock() defer r.mu.Unlock() - entry, ok := r.entryLocked(keyID) + entry, ok := r.entryLocked(credentialID) if !ok { return false } @@ -581,13 +611,13 @@ func (r *KeyRegistry) SetAutoWeight(keyID uint, weight int) bool { return true } -func (r *KeyRegistry) RestoreRuntimeState(keyID uint, weight int) bool { +func (r *CredentialRegistry) RestoreRuntimeState(credentialID uint, weight int) bool { if weight < 1 || weight > MaxWeight { return false } r.mu.Lock() defer r.mu.Unlock() - entry, ok := r.entryLocked(keyID) + entry, ok := r.entryLocked(credentialID) if !ok { return false } @@ -599,10 +629,10 @@ func (r *KeyRegistry) RestoreRuntimeState(keyID uint, weight int) bool { return true } -func (r *KeyRegistry) SetBlacklisted(keyID uint) bool { +func (r *CredentialRegistry) SetBlacklisted(credentialID uint) bool { r.mu.Lock() defer r.mu.Unlock() - entry, ok := r.entryLocked(keyID) + entry, ok := r.entryLocked(credentialID) if !ok { return false } @@ -613,10 +643,10 @@ func (r *KeyRegistry) SetBlacklisted(keyID uint) bool { return true } -func (r *KeyRegistry) IncrFailure(keyID uint) (int, bool) { +func (r *CredentialRegistry) IncrFailure(credentialID uint) (int, bool) { r.mu.Lock() defer r.mu.Unlock() - entry, ok := r.entryLocked(keyID) + entry, ok := r.entryLocked(credentialID) if !ok { return 0, false } @@ -625,10 +655,10 @@ func (r *KeyRegistry) IncrFailure(keyID uint) (int, bool) { return entry.FailureCount, true } -func (r *KeyRegistry) ClearFailure(keyID uint) bool { +func (r *CredentialRegistry) ClearFailure(credentialID uint) bool { r.mu.Lock() defer r.mu.Unlock() - entry, ok := r.entryLocked(keyID) + entry, ok := r.entryLocked(credentialID) if !ok { return false } @@ -639,10 +669,10 @@ func (r *KeyRegistry) ClearFailure(keyID uint) bool { return true } -func (r *KeyRegistry) Recover(keyID uint) bool { +func (r *CredentialRegistry) Recover(credentialID uint) bool { r.mu.Lock() defer r.mu.Unlock() - entry, ok := r.entryLocked(keyID) + entry, ok := r.entryLocked(credentialID) if !ok { return false } @@ -654,18 +684,22 @@ func (r *KeyRegistry) Recover(keyID uint) bool { return true } -func (r *KeyRegistry) RecoverIfMatch(ref KeyRef, weight int) bool { +func (r *CredentialRegistry) RecoverIfMatch(ref CredentialRef, weight int) bool { r.mu.Lock() defer r.mu.Unlock() if weight < 1 || weight > MaxWeight { return false } - groupID, ok := r.keyGroups[ref.ID] + groupID, ok := r.credentialGroups[ref.ID] if !ok || groupID != ref.GroupID { return false } entry, ok := r.buckets[groupID][ref.ID] - if !ok || entry.Status != KeyStatusActive || !entry.Blacklisted || entry.GroupID != ref.GroupID || entry.EncryptedValue != ref.EncryptedValue || entry.FailureGeneration != ref.FailureGeneration { + if !ok || entry.Status != CredentialStatusActive || !entry.Blacklisted || + entry.GroupID != ref.GroupID || entry.Version != ref.Version || + entry.IdentityGeneration != ref.IdentityGeneration || + entry.Fingerprint != ref.Fingerprint || entry.EncryptedValue != ref.EncryptedValue || + entry.FailureGeneration != ref.FailureGeneration { return false } entry.WeightAuto = weight @@ -675,16 +709,18 @@ func (r *KeyRegistry) RecoverIfMatch(ref KeyRef, weight int) bool { return true } -func (r *KeyRegistry) BlacklistedKeys() []KeyRef { +func (r *CredentialRegistry) BlacklistedCredentials() []CredentialRef { r.mu.RLock() - refs := make([]KeyRef, 0) + refs := make([]CredentialRef, 0) for _, bucket := range r.buckets { for _, entry := range bucket { - if entry.Status != KeyStatusActive || !entry.Blacklisted { + if entry.Status != CredentialStatusActive || !entry.Blacklisted { continue } - refs = append(refs, KeyRef{ - ID: entry.ID, GroupID: entry.GroupID, EncryptedValue: entry.EncryptedValue, + refs = append(refs, CredentialRef{ + ID: entry.ID, GroupID: entry.GroupID, + Version: entry.Version, IdentityGeneration: entry.IdentityGeneration, + Fingerprint: entry.Fingerprint, EncryptedValue: entry.EncryptedValue, FailureGeneration: entry.FailureGeneration, }) } @@ -699,16 +735,16 @@ func (r *KeyRegistry) BlacklistedKeys() []KeyRef { return refs } -func (r *KeyRegistry) entryLocked(keyID uint) (*KeyEntry, bool) { - groupID, ok := r.keyGroups[keyID] +func (r *CredentialRegistry) entryLocked(credentialID uint) (*CredentialEntry, bool) { + groupID, ok := r.credentialGroups[credentialID] if !ok { return nil, false } - entry, ok := r.buckets[groupID][keyID] + entry, ok := r.buckets[groupID][credentialID] return entry, ok } -func cloneKeyEntry(entry KeyEntry) KeyEntry { +func cloneCredentialEntry(entry CredentialEntry) CredentialEntry { entry.WeightManual = cloneWeight(entry.WeightManual) entry.FailureGeneration = 0 if entry.WeightAuto == 0 { @@ -717,7 +753,7 @@ func cloneKeyEntry(entry KeyEntry) KeyEntry { return entry } -func detachKeyEntryExact(entry KeyEntry) KeyEntry { +func detachCredentialEntryExact(entry CredentialEntry) CredentialEntry { entry.WeightManual = cloneWeight(entry.WeightManual) return entry } diff --git a/internal/state/registry_test.go b/internal/state/registry_test.go index e4d558e4..e51b1ee7 100644 --- a/internal/state/registry_test.go +++ b/internal/state/registry_test.go @@ -10,32 +10,32 @@ import ( ) func TestKeyRegistryReplaceAndEncryptedValue(t *testing.T) { - registry := NewKeyRegistry() + registry := NewCredentialRegistry() weight := 7 - entries := []KeyEntry{ + entries := []CredentialEntry{ { ID: 1, GroupID: 10, WeightManual: &weight, - Status: KeyStatusActive, EncryptedValue: "cipher-one", + Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }, { ID: 2, GroupID: 20, - Status: KeyStatusDisabled, EncryptedValue: "cipher-two", + Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two", }, } - if err := registry.Replace(entries); err != nil { + if err := registry.ReplaceCredentials(entries); err != nil { t.Fatalf("Replace() error = %v", err) } for _, test := range []struct { - keyID uint - want string + credentialID uint + want string }{ - {keyID: 1, want: "cipher-one"}, - {keyID: 2, want: "cipher-two"}, + {credentialID: 1, want: "cipher-one"}, + {credentialID: 2, want: "cipher-two"}, } { - got, ok := registry.EncryptedValue(test.keyID) + got, ok := registry.EncryptedCredentialData(test.credentialID) if !ok || got != test.want { - t.Errorf("EncryptedValue(%d) = %q, %t, want %q, true", test.keyID, got, ok, test.want) + t.Errorf("EncryptedValue(%d) = %q, %t, want %q, true", test.credentialID, got, ok, test.want) } } @@ -47,29 +47,29 @@ func TestKeyRegistryReplaceAndEncryptedValue(t *testing.T) { t.Fatalf("stored WeightManual = %d after caller mutation, want 7", storedWeight) } - if err := registry.Replace([]KeyEntry{{ + if err := registry.ReplaceCredentials([]CredentialEntry{{ ID: 3, GroupID: 20, - Status: KeyStatusActive, EncryptedValue: "cipher-three", + Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three", }}); err != nil { t.Fatalf("second Replace() error = %v", err) } for _, removedID := range []uint{1, 2} { - if got, ok := registry.EncryptedValue(removedID); ok || got != "" { + if got, ok := registry.EncryptedCredentialData(removedID); ok || got != "" { t.Errorf("EncryptedValue(%d) after replacement = %q, %t, want empty, false", removedID, got, ok) } } - if got, ok := registry.EncryptedValue(3); !ok || got != "cipher-three" { + if got, ok := registry.EncryptedCredentialData(3); !ok || got != "cipher-three" { t.Fatalf("EncryptedValue(3) = %q, %t, want %q, true", got, ok, "cipher-three") } } func TestKeyRegistryRestoreRuntimeState(t *testing.T) { - registry := NewKeyRegistry() + registry := NewCredentialRegistry() now := time.Date(2026, time.August, 1, 10, 0, 0, 0, time.UTC) - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, CooldownUntil: now.Add(time.Hour), Blacklisted: true, - FailureCount: 3, EncryptedValue: "cipher-one", + FailureCount: 3, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }}) if ok := registry.RestoreRuntimeState(1, 37); !ok { @@ -85,16 +85,16 @@ func TestKeyRegistryRestoreRuntimeState(t *testing.T) { before := registryEntry(t, registry, 1) for _, test := range []struct { - keyID uint - weight int + credentialID uint + weight int }{ - {keyID: 0, weight: 37}, - {keyID: 99, weight: 37}, - {keyID: 1, weight: 0}, - {keyID: 1, weight: MaxWeight + 1}, + {credentialID: 0, weight: 37}, + {credentialID: 99, weight: 37}, + {credentialID: 1, weight: 0}, + {credentialID: 1, weight: MaxWeight + 1}, } { - if registry.RestoreRuntimeState(test.keyID, test.weight) { - t.Fatalf("RestoreRuntimeState(%d, %d) = true, want false", test.keyID, test.weight) + if registry.RestoreRuntimeState(test.credentialID, test.weight) { + t.Fatalf("RestoreRuntimeState(%d, %d) = true, want false", test.credentialID, test.weight) } } if after := registryEntry(t, registry, 1); !reflect.DeepEqual(after, before) { @@ -103,26 +103,26 @@ func TestKeyRegistryRestoreRuntimeState(t *testing.T) { } func TestKeyRegistryBatchMutationsAreAllOrNothing(t *testing.T) { - newRegistry := func(t *testing.T) *KeyRegistry { + newRegistry := func(t *testing.T) *CredentialRegistry { t.Helper() - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, - {ID: 2, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "cipher-two"}, - {ID: 3, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-three"}, + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + {ID: 2, GroupID: 10, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, + {ID: 3, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three"}, }) return registry } t.Run("update statuses", func(t *testing.T) { registry := newRegistry(t) - if err := registry.UpdateGroupKeyStatuses(10, []uint{1, 2}, KeyStatusDisabled); err != nil { - t.Fatalf("UpdateGroupKeyStatuses() error = %v", err) + if err := registry.UpdateGroupCredentialStatuses(10, []uint{1, 2}, CredentialStatusDisabled); err != nil { + t.Fatalf("UpdateGroupCredentialStatuses() error = %v", err) } - if got := registryEntry(t, registry, 1).Status; got != KeyStatusDisabled { + if got := registryEntry(t, registry, 1).Status; got != CredentialStatusDisabled { t.Fatalf("key 1 status = %q, want disabled", got) } - if got := registryEntry(t, registry, 2).Status; got != KeyStatusDisabled { + if got := registryEntry(t, registry, 2).Status; got != CredentialStatusDisabled { t.Fatalf("key 2 status = %q, want disabled", got) } }) @@ -139,8 +139,8 @@ func TestKeyRegistryBatchMutationsAreAllOrNothing(t *testing.T) { t.Run("update rejects "+test.name, func(t *testing.T) { registry := newRegistry(t) before := registry.Snapshot() - if err := registry.UpdateGroupKeyStatuses(10, test.ids, KeyStatusDisabled); err == nil { - t.Fatal("UpdateGroupKeyStatuses() error = nil, want error") + if err := registry.UpdateGroupCredentialStatuses(10, test.ids, CredentialStatusDisabled); err == nil { + t.Fatal("UpdateGroupCredentialStatuses() error = nil, want error") } if after := registry.Snapshot(); !reflect.DeepEqual(after, before) { t.Fatalf("registry after rejected update = %#v, want %#v", after, before) @@ -149,8 +149,8 @@ func TestKeyRegistryBatchMutationsAreAllOrNothing(t *testing.T) { t.Run("remove rejects "+test.name, func(t *testing.T) { registry := newRegistry(t) before := registry.Snapshot() - if err := registry.RemoveGroupKeys(10, test.ids); err == nil { - t.Fatal("RemoveGroupKeys() error = nil, want error") + if err := registry.RemoveGroupCredentials(10, test.ids); err == nil { + t.Fatal("RemoveGroupCredentials() error = nil, want error") } if after := registry.Snapshot(); !reflect.DeepEqual(after, before) { t.Fatalf("registry after rejected remove = %#v, want %#v", after, before) @@ -159,126 +159,126 @@ func TestKeyRegistryBatchMutationsAreAllOrNothing(t *testing.T) { } registry := newRegistry(t) - if err := registry.RemoveGroupKeys(10, []uint{2, 1}); err != nil { - t.Fatalf("RemoveGroupKeys() error = %v", err) + if err := registry.RemoveGroupCredentials(10, []uint{2, 1}); err != nil { + t.Fatalf("RemoveGroupCredentials() error = %v", err) } - if _, ok := registry.EncryptedValue(1); ok { + if _, ok := registry.EncryptedCredentialData(1); ok { t.Fatal("key 1 remains after batch removal") } - if _, ok := registry.EncryptedValue(2); ok { + if _, ok := registry.EncryptedCredentialData(2); ok { t.Fatal("key 2 remains after batch removal") } - if got, ok := registry.EncryptedValue(3); !ok || got != "cipher-three" { + if got, ok := registry.EncryptedCredentialData(3); !ok || got != "cipher-three" { t.Fatalf("other-group key = %q/%t", got, ok) } } func TestKeyRegistryActiveEncryptedValueRequiresExpectedGroupAndActiveStatus(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-active"}, - {ID: 2, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "cipher-disabled"}, + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-active"}, + {ID: 2, GroupID: 10, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-disabled"}, }) tests := []struct { - name string - keyID uint - groupID uint - want string - wantOK bool + name string + credentialID uint + groupID uint + want string + wantOK bool }{ - {name: "matching active key", keyID: 1, groupID: 10, want: "cipher-active", wantOK: true}, - {name: "group mismatch", keyID: 1, groupID: 20}, - {name: "disabled key", keyID: 2, groupID: 10}, - {name: "missing key", keyID: 99, groupID: 10}, + {name: "matching active key", credentialID: 1, groupID: 10, want: "cipher-active", wantOK: true}, + {name: "group mismatch", credentialID: 1, groupID: 20}, + {name: "disabled key", credentialID: 2, groupID: 10}, + {name: "missing key", credentialID: 99, groupID: 10}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got, ok := registry.ActiveEncryptedValue(tt.keyID, tt.groupID) + got, ok := registry.ActiveEncryptedCredentialData(tt.credentialID, tt.groupID) if got != tt.want || ok != tt.wantOK { - t.Fatalf("ActiveEncryptedValue(%d, %d) = %q, %t, want %q, %t", tt.keyID, tt.groupID, got, ok, tt.want, tt.wantOK) + t.Fatalf("ActiveEncryptedCredentialData(%d, %d) = %q, %t, want %q, %t", tt.credentialID, tt.groupID, got, ok, tt.want, tt.wantOK) } }) } - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-moved", + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-moved", }}) - if got, ok := registry.ActiveEncryptedValue(1, 10); got != "" || ok { - t.Fatalf("ActiveEncryptedValue(1, old group) = %q, %t, want empty, false", got, ok) + if got, ok := registry.ActiveEncryptedCredentialData(1, 10); got != "" || ok { + t.Fatalf("ActiveEncryptedCredentialData(1, old group) = %q, %t, want empty, false", got, ok) } - if got, ok := registry.ActiveEncryptedValue(1, 20); got != "cipher-moved" || !ok { - t.Fatalf("ActiveEncryptedValue(1, new group) = %q, %t, want cipher-moved, true", got, ok) + if got, ok := registry.ActiveEncryptedCredentialData(1, 20); got != "cipher-moved" || !ok { + t.Fatalf("ActiveEncryptedCredentialData(1, new group) = %q, %t, want cipher-moved, true", got, ok) } } func TestKeyRegistryCaptureActiveKeyRefsIncludesTemporarilyUnavailableKeys(t *testing.T) { now := time.Date(2026, time.July, 27, 12, 0, 0, 0, time.UTC) - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ { - ID: 21, GroupID: 2, Status: KeyStatusActive, - CooldownUntil: now.Add(time.Hour), EncryptedValue: "cipher-cooldown", + ID: 21, GroupID: 2, Status: CredentialStatusActive, + CooldownUntil: now.Add(time.Hour), Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-cooldown", }, { - ID: 12, GroupID: 1, Status: KeyStatusDisabled, - EncryptedValue: "cipher-disabled", + ID: 12, GroupID: 1, Status: CredentialStatusDisabled, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-disabled", }, { - ID: 11, GroupID: 1, Status: KeyStatusActive, - Blacklisted: true, EncryptedValue: "cipher-blacklisted", + ID: 11, GroupID: 1, Status: CredentialStatusActive, + Blacklisted: true, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-blacklisted", }, { - ID: 22, GroupID: 2, Status: KeyStatusActive, - EncryptedValue: "cipher-available", + ID: 22, GroupID: 2, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-available", }, { - ID: 31, GroupID: 3, Status: KeyStatusActive, - EncryptedValue: "cipher-other-group", + ID: 31, GroupID: 3, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-other-group", }, }) - want := []KeyRef{ - {ID: 11, GroupID: 1, EncryptedValue: "cipher-blacklisted"}, - {ID: 21, GroupID: 2, EncryptedValue: "cipher-cooldown"}, - {ID: 22, GroupID: 2, EncryptedValue: "cipher-available"}, + want := []CredentialRef{ + {ID: 11, GroupID: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-blacklisted"}, + {ID: 21, GroupID: 2, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-cooldown"}, + {ID: 22, GroupID: 2, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-available"}, } - got := registry.CaptureActiveKeyRefs([]uint{2, 1}) + got := registry.CaptureActiveCredentialRefs([]uint{2, 1}) if !reflect.DeepEqual(got, want) { - t.Fatalf("CaptureActiveKeyRefs() = %#v, want %#v", got, want) + t.Fatalf("CaptureActiveCredentialRefs() = %#v, want %#v", got, want) } if _, ok := registry.IncrFailure(11); !ok { t.Fatal("IncrFailure(11) = false, want true") } want[0].FailureGeneration = 1 - if got := registry.CaptureActiveKeyRefs([]uint{1, 2}); !reflect.DeepEqual(got, want) { - t.Fatalf("CaptureActiveKeyRefs() after failure = %#v, want %#v", got, want) + if got := registry.CaptureActiveCredentialRefs([]uint{1, 2}); !reflect.DeepEqual(got, want) { + t.Fatalf("CaptureActiveCredentialRefs() after failure = %#v, want %#v", got, want) } got[0].EncryptedValue = "caller-mutated" - if again := registry.CaptureActiveKeyRefs([]uint{1, 2}); !reflect.DeepEqual(again, want) { - t.Fatalf("CaptureActiveKeyRefs() aliases caller result: %#v", again) + if again := registry.CaptureActiveCredentialRefs([]uint{1, 2}); !reflect.DeepEqual(again, want) { + t.Fatalf("CaptureActiveCredentialRefs() aliases caller result: %#v", again) } } func TestKeyRegistryActiveEncryptedValueIfMatchRejectsIdentityChanges(t *testing.T) { - ref := KeyRef{ID: 1, GroupID: 10, EncryptedValue: "cipher-original"} - newRegistry := func(t *testing.T) *KeyRegistry { + ref := CredentialRef{ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-original"} + newRegistry := func(t *testing.T) *CredentialRegistry { t.Helper() - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, - EncryptedValue: "cipher-original", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-original", }}) return registry } t.Run("matching active identity", func(t *testing.T) { registry := newRegistry(t) - got, ok := registry.ActiveEncryptedValueIfMatch(ref) + got, ok := registry.ActiveEncryptedCredentialDataIfMatch(ref) if !ok || got != "cipher-original" { t.Fatalf( - "ActiveEncryptedValueIfMatch() = %q, %t, want cipher-original, true", + "ActiveEncryptedCredentialDataIfMatch() = %q, %t, want cipher-original, true", got, ok, ) @@ -290,9 +290,9 @@ func TestKeyRegistryActiveEncryptedValueIfMatchRejectsIdentityChanges(t *testing if _, ok := registry.IncrFailure(1); !ok { t.Fatal("IncrFailure(1) = false, want true") } - got, ok := registry.ActiveEncryptedValueIfMatch(ref) + got, ok := registry.ActiveEncryptedCredentialDataIfMatch(ref) if !ok || got != "cipher-original" { - t.Fatalf("ActiveEncryptedValueIfMatch() = %q, %t, want cipher-original, true", got, ok) + t.Fatalf("ActiveEncryptedCredentialDataIfMatch() = %q, %t, want cipher-original, true", got, ok) } }) @@ -300,53 +300,53 @@ func TestKeyRegistryActiveEncryptedValueIfMatchRejectsIdentityChanges(t *testing registry := newRegistry(t) mismatched := ref mismatched.ID = 2 - if got, ok := registry.ActiveEncryptedValueIfMatch(mismatched); ok || got != "" { - t.Fatalf("ActiveEncryptedValueIfMatch(wrong ID) = %q, %t, want empty, false", got, ok) + if got, ok := registry.ActiveEncryptedCredentialDataIfMatch(mismatched); ok || got != "" { + t.Fatalf("ActiveEncryptedCredentialDataIfMatch(wrong ID) = %q, %t, want empty, false", got, ok) } }) t.Run("replaced ciphertext", func(t *testing.T) { registry := newRegistry(t) - if err := registry.ApplyImport(10, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, - EncryptedValue: "cipher-replaced", + if err := registry.ApplyCredentialImport(10, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-replaced", }}); err != nil { t.Fatalf("ApplyImport(replacement) error = %v", err) } - if got, ok := registry.ActiveEncryptedValueIfMatch(ref); ok || got != "" { - t.Fatalf("ActiveEncryptedValueIfMatch(replaced) = %q, %t, want empty, false", got, ok) + if got, ok := registry.ActiveEncryptedCredentialDataIfMatch(ref); ok || got != "" { + t.Fatalf("ActiveEncryptedCredentialDataIfMatch(replaced) = %q, %t, want empty, false", got, ok) } }) t.Run("moved group", func(t *testing.T) { registry := newRegistry(t) - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 20, Status: KeyStatusActive, - EncryptedValue: "cipher-original", + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 20, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-original", }}) - if got, ok := registry.ActiveEncryptedValueIfMatch(ref); ok || got != "" { - t.Fatalf("ActiveEncryptedValueIfMatch(moved) = %q, %t, want empty, false", got, ok) + if got, ok := registry.ActiveEncryptedCredentialDataIfMatch(ref); ok || got != "" { + t.Fatalf("ActiveEncryptedCredentialDataIfMatch(moved) = %q, %t, want empty, false", got, ok) } }) t.Run("disabled", func(t *testing.T) { registry := newRegistry(t) - if err := registry.SetKeyStatus(1, KeyStatusDisabled); err != nil { - t.Fatalf("SetKeyStatus(disabled) error = %v", err) + if err := registry.SetCredentialStatus(1, CredentialStatusDisabled); err != nil { + t.Fatalf("SetCredentialStatus(disabled) error = %v", err) } - if got, ok := registry.ActiveEncryptedValueIfMatch(ref); ok || got != "" { - t.Fatalf("ActiveEncryptedValueIfMatch(disabled) = %q, %t, want empty, false", got, ok) + if got, ok := registry.ActiveEncryptedCredentialDataIfMatch(ref); ok || got != "" { + t.Fatalf("ActiveEncryptedCredentialDataIfMatch(disabled) = %q, %t, want empty, false", got, ok) } }) } func TestKeyRegistryActiveKeyIDs(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 9, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-nine"}, - {ID: 5, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "cipher-five"}, - {ID: 3, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-three"}, - {ID: 7, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-seven"}, + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 9, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-nine"}, + {ID: 5, GroupID: 10, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-five"}, + {ID: 3, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three"}, + {ID: 7, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-seven"}, }) if ok := registry.SetCooldown(3, time.Now().Add(time.Hour)); !ok { t.Fatal("SetCooldown(3) = false, want true") @@ -356,34 +356,34 @@ func TestKeyRegistryActiveKeyIDs(t *testing.T) { } want := []uint{3, 7, 9} - got := registry.ActiveKeyIDs() + got := registry.ActiveCredentialIDs() if !reflect.DeepEqual(got, want) { - t.Fatalf("ActiveKeyIDs() = %v, want %v", got, want) + t.Fatalf("ActiveCredentialIDs() = %v, want %v", got, want) } got[0] = 99 - if again := registry.ActiveKeyIDs(); !reflect.DeepEqual(again, want) { - t.Fatalf("ActiveKeyIDs() after caller mutation = %v, want %v", again, want) + if again := registry.ActiveCredentialIDs(); !reflect.DeepEqual(again, want) { + t.Fatalf("ActiveCredentialIDs() after caller mutation = %v, want %v", again, want) } } func TestKeyRegistrySnapshotIsSortedDetachedAndCredentialFree(t *testing.T) { weight := 17 now := time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC) - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ { - ID: 22, GroupID: 2, Status: KeyStatusDisabled, - EncryptedValue: "cipher-disabled", + ID: 22, GroupID: 2, Status: CredentialStatusDisabled, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-disabled", }, { ID: 11, GroupID: 1, WeightManual: &weight, WeightAuto: 23, - Status: KeyStatusActive, CooldownUntil: now.Add(time.Minute), - Blacklisted: true, FailureCount: 3, EncryptedValue: "cipher-secret", + Status: CredentialStatusActive, CooldownUntil: now.Add(time.Minute), + Blacklisted: true, FailureCount: 3, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-secret", }, { - ID: 12, GroupID: 1, Status: KeyStatusActive, - EncryptedValue: "cipher-default-weight", + ID: 12, GroupID: 1, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-default-weight", }, }) @@ -407,28 +407,28 @@ func TestKeyRegistrySnapshotIsSortedDetachedAndCredentialFree(t *testing.T) { t.Fatalf("Snapshot aliases Registry: %#v", again[0]) } - typ := reflect.TypeOf(KeyRuntimeView{}) + typ := reflect.TypeOf(CredentialRuntimeView{}) for _, forbidden := range []string{ "EncryptedValue", "KeyHash", "Hash", "Mask", "HeaderRules", } { if _, exists := typ.FieldByName(forbidden); exists { - t.Fatalf("KeyRuntimeView exposes forbidden field %s", forbidden) + t.Fatalf("CredentialRuntimeView exposes forbidden field %s", forbidden) } } } func TestKeyRuntimeViewDoesNotExposeFailureGeneration(t *testing.T) { - typ := reflect.TypeOf(KeyRuntimeView{}) + typ := reflect.TypeOf(CredentialRuntimeView{}) if _, exists := typ.FieldByName("FailureGeneration"); exists { - t.Fatal("KeyRuntimeView exposes forbidden field FailureGeneration") + t.Fatal("CredentialRuntimeView exposes forbidden field FailureGeneration") } } func TestKeyRegistryFailureGenerationTracksActualFailureStateChanges(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, - FailureGeneration: 99, EncryptedValue: "cipher-one", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, + FailureGeneration: 99, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }}) assertGeneration := func(want uint64) { @@ -476,50 +476,50 @@ func TestKeyRegistryFailureGenerationTracksActualFailureStateChanges(t *testing. } assertGeneration(4) - if err := registry.ApplyImport(10, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, - FailureGeneration: 73, EncryptedValue: "cipher-imported", + if err := registry.ApplyCredentialImport(10, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, + FailureGeneration: 73, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-imported", }}); err != nil { t.Fatalf("ApplyImport() error = %v", err) } assertGeneration(0) - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, - FailureGeneration: 42, EncryptedValue: "cipher-replaced", + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, + FailureGeneration: 42, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-replaced", }}) assertGeneration(0) } func TestKeyRegistryReplaceFailurePreservesRegistry(t *testing.T) { - invalidBatches := map[string][]KeyEntry{ + invalidBatches := map[string][]CredentialEntry{ "duplicate ids": { - {ID: 2, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-two"}, - {ID: 2, GroupID: 20, Status: KeyStatusDisabled, EncryptedValue: "duplicate"}, + {ID: 2, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, + {ID: 2, GroupID: 20, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "duplicate"}, }, "invalid status": { - {ID: 3, GroupID: 30, Status: KeyStatus("cooldown"), EncryptedValue: "cipher-three"}, + {ID: 3, GroupID: 30, Status: CredentialStatus("cooldown"), Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three"}, }, } for name, batch := range invalidBatches { t.Run(name, func(t *testing.T) { - registry := NewKeyRegistry() - if err := registry.Replace([]KeyEntry{{ + registry := NewCredentialRegistry() + if err := registry.ReplaceCredentials([]CredentialEntry{{ ID: 1, GroupID: 10, - Status: KeyStatusActive, EncryptedValue: "original-cipher", + Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "original-cipher", }}); err != nil { t.Fatalf("seed Replace() error = %v", err) } - if err := registry.Replace(batch); err == nil { + if err := registry.ReplaceCredentials(batch); err == nil { t.Fatal("invalid Replace() error = nil, want error") } - if got, ok := registry.EncryptedValue(1); !ok || got != "original-cipher" { + if got, ok := registry.EncryptedCredentialData(1); !ok || got != "original-cipher" { t.Errorf("original EncryptedValue(1) = %q, %t after failed replacement, want %q, true", got, ok, "original-cipher") } for _, entry := range batch { - if got, ok := registry.EncryptedValue(entry.ID); ok || got != "" { + if got, ok := registry.EncryptedCredentialData(entry.ID); ok || got != "" { t.Errorf("invalid EncryptedValue(%d) = %q, %t after failed replacement, want empty, false", entry.ID, got, ok) } } @@ -528,16 +528,16 @@ func TestKeyRegistryReplaceFailurePreservesRegistry(t *testing.T) { } func TestKeyRegistryApplyImportUpsertsOnlyRequestedGroup(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "old-one"}, - {ID: 2, GroupID: 20, Status: KeyStatusDisabled, EncryptedValue: "cipher-two"}, + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "old-one"}, + {ID: 2, GroupID: 20, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, }) weight := 12 - if err := registry.ApplyImport(10, []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "new-one"}, - {ID: 3, GroupID: 10, WeightManual: &weight, Status: KeyStatusActive, EncryptedValue: "cipher-three"}, + if err := registry.ApplyCredentialImport(10, []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "new-one"}, + {ID: 3, GroupID: 10, WeightManual: &weight, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three"}, }); err != nil { t.Fatalf("ApplyImport() error = %v", err) } @@ -556,23 +556,23 @@ func TestKeyRegistryApplyImportUpsertsOnlyRequestedGroup(t *testing.T) { } func TestKeyRegistryReconcileGroupUsesDBTruthAndPreservesSatisfiedRuntimeState(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ { - ID: 1, GroupID: 10, Status: KeyStatusActive, - EncryptedValue: "cipher-one", + ID: 1, GroupID: 10, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }, { - ID: 2, GroupID: 10, Status: KeyStatusActive, - EncryptedValue: "stale-cipher", + ID: 2, GroupID: 10, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "stale-cipher", }, { - ID: 3, GroupID: 10, Status: KeyStatusDisabled, - EncryptedValue: "removed-from-db", + ID: 3, GroupID: 10, Status: CredentialStatusDisabled, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "removed-from-db", }, { - ID: 4, GroupID: 20, Status: KeyStatusActive, - EncryptedValue: "other-group", + ID: 4, GroupID: 20, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "other-group", }, }) if !registry.SetCooldown(1, time.Now().Add(time.Hour)) { @@ -582,14 +582,14 @@ func TestKeyRegistryReconcileGroupUsesDBTruthAndPreservesSatisfiedRuntimeState(t t.Fatal("IncrFailure(1) = false") } - wantGroup := []KeyEntry{ + wantGroup := []CredentialEntry{ { - ID: 1, GroupID: 10, Status: KeyStatusActive, - EncryptedValue: "cipher-one", + ID: 1, GroupID: 10, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }, { - ID: 2, GroupID: 10, Status: KeyStatusDisabled, - EncryptedValue: "cipher-two", + ID: 2, GroupID: 10, Status: CredentialStatusDisabled, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two", }, } if registry.MatchesGroup(10, wantGroup) { @@ -609,14 +609,14 @@ func TestKeyRegistryReconcileGroupUsesDBTruthAndPreservesSatisfiedRuntimeState(t got.CooldownUntil.IsZero() { t.Fatalf("satisfied key runtime state was reset: %#v", got) } - if got := registryEntry(t, registry, 2); got.Status != KeyStatusDisabled || + if got := registryEntry(t, registry, 2); got.Status != CredentialStatusDisabled || got.EncryptedValue != "cipher-two" || got.FailureCount != 0 { t.Fatalf("stale key was not rebuilt from DB truth: %#v", got) } - if _, ok := registry.EncryptedValue(3); ok { + if _, ok := registry.EncryptedCredentialData(3); ok { t.Fatal("key absent from DB truth remains in reconciled group") } - if got, ok := registry.EncryptedValue(4); !ok || got != "other-group" { + if got, ok := registry.EncryptedCredentialData(4); !ok || got != "other-group" { t.Fatalf("other group changed: %q, %t", got, ok) } @@ -633,43 +633,43 @@ func TestKeyRegistryReconcileGroupUsesDBTruthAndPreservesSatisfiedRuntimeState(t } func TestKeyRegistryReconcileGroupRejectsCrossGroupIdentityWithoutMutation(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 20, Status: KeyStatusActive, - EncryptedValue: "other-group", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 20, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "other-group", }}) - _, err := registry.ReconcileGroup(10, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, - EncryptedValue: "attempted-move", + _, err := registry.ReconcileGroup(10, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "attempted-move", }}) if err == nil { t.Fatal("ReconcileGroup(cross-group key) error = nil") } - if got, ok := registry.EncryptedValue(1); !ok || got != "other-group" { + if got, ok := registry.EncryptedCredentialData(1); !ok || got != "other-group" { t.Fatalf("failed reconciliation mutated registry: %q, %t", got, ok) } } func TestKeyRegistryRemoveKey(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, - {ID: 2, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-two"}, - {ID: 3, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-three"}, + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + {ID: 2, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, + {ID: 3, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three"}, }) - if removed := registry.RemoveKey(1); !removed { - t.Fatal("RemoveKey(1) = false, want true") + if removed := registry.RemoveCredential(1); !removed { + t.Fatal("RemoveCredential(1) = false, want true") } assertEncryptedValue(t, registry, 1, "", false) assertEncryptedValue(t, registry, 2, "cipher-two", true) - if removed := registry.RemoveKey(1); removed { - t.Fatal("second RemoveKey(1) = true, want false") + if removed := registry.RemoveCredential(1); removed { + t.Fatal("second RemoveCredential(1) = true, want false") } - if removed := registry.RemoveKey(2); !removed { - t.Fatal("RemoveKey(2) = false, want true") + if removed := registry.RemoveCredential(2); !removed { + t.Fatal("RemoveCredential(2) = false, want true") } registry.mu.RLock() _, emptyBucketRetained := registry.buckets[10] @@ -681,22 +681,22 @@ func TestKeyRegistryRemoveKey(t *testing.T) { } func TestKeyRegistryRemoveGroupClearsBucketAndReverseIndexesAtomically(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, - {ID: 2, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "cipher-two"}, - {ID: 3, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-three"}, + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + {ID: 2, GroupID: 10, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, + {ID: 3, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three"}, }) if removed := registry.RemoveGroup(10); !removed { t.Fatal("RemoveGroup(10) = false, want true") } for _, id := range []uint{1, 2} { - if value, ok := registry.EncryptedValue(id); ok || value != "" { + if value, ok := registry.EncryptedCredentialData(id); ok || value != "" { t.Fatalf("removed key %d = %q, %t", id, value, ok) } } - if value, ok := registry.EncryptedValue(3); !ok || value != "cipher-three" { + if value, ok := registry.EncryptedCredentialData(3); !ok || value != "cipher-three" { t.Fatalf("other group key = %q, %t", value, ok) } if removed := registry.RemoveGroup(10); removed { @@ -708,12 +708,12 @@ func TestKeyRegistryRemoveGroupClearsBucketAndReverseIndexesAtomically(t *testin } func TestKeyRegistryRemoveGroupIsRaceSafeWithRuntimeReaders(t *testing.T) { - registry := NewKeyRegistry() - entries := make([]KeyEntry, 0, 100) + registry := NewCredentialRegistry() + entries := make([]CredentialEntry, 0, 100) for id := uint(1); id <= 100; id++ { - entries = append(entries, KeyEntry{ - ID: id, GroupID: 10, Status: KeyStatusActive, - EncryptedValue: fmt.Sprintf("cipher-%d", id), + entries = append(entries, CredentialEntry{ + ID: id, GroupID: 10, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: fmt.Sprintf("cipher-%d", id), }) } mustReplaceKeyEntries(t, registry, entries) @@ -727,8 +727,8 @@ func TestKeyRegistryRemoveGroupIsRaceSafeWithRuntimeReaders(t *testing.T) { <-start for range 100 { _ = registry.Snapshot() - _ = registry.CollectCandidates([]uint{10}, nil, time.Time{}) - _ = registry.ActiveKeyIDs() + _ = registry.CollectCredentialCandidates([]uint{10}, nil, time.Time{}) + _ = registry.ActiveCredentialIDs() } }() } @@ -741,47 +741,47 @@ func TestKeyRegistryRemoveGroupIsRaceSafeWithRuntimeReaders(t *testing.T) { } func TestKeyRegistrySetKeyStatus(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }}) - if err := registry.SetKeyStatus(1, KeyStatusDisabled); err != nil { - t.Fatalf("SetKeyStatus(disabled) error = %v", err) + if err := registry.SetCredentialStatus(1, CredentialStatusDisabled); err != nil { + t.Fatalf("SetCredentialStatus(disabled) error = %v", err) } - if got := keyStatus(t, registry, 1); got != KeyStatusDisabled { - t.Fatalf("key status = %q, want %q", got, KeyStatusDisabled) + if got := keyStatus(t, registry, 1); got != CredentialStatusDisabled { + t.Fatalf("key status = %q, want %q", got, CredentialStatusDisabled) } - if err := registry.SetKeyStatus(1, KeyStatus("cooldown")); err == nil { - t.Fatal("SetKeyStatus(invalid) error = nil, want error") + if err := registry.SetCredentialStatus(1, CredentialStatus("cooldown")); err == nil { + t.Fatal("SetCredentialStatus(invalid) error = nil, want error") } - if got := keyStatus(t, registry, 1); got != KeyStatusDisabled { - t.Fatalf("key status after invalid update = %q, want %q", got, KeyStatusDisabled) + if got := keyStatus(t, registry, 1); got != CredentialStatusDisabled { + t.Fatalf("key status after invalid update = %q, want %q", got, CredentialStatusDisabled) } - if err := registry.SetKeyStatus(99, KeyStatusActive); err == nil { - t.Fatal("SetKeyStatus(missing key) error = nil, want error") + if err := registry.SetCredentialStatus(99, CredentialStatusActive); err == nil { + t.Fatal("SetCredentialStatus(missing key) error = nil, want error") } } func TestKeyRegistryCollectCandidatesFiltersStatusAndExcluded(t *testing.T) { - registry := NewKeyRegistry() + registry := NewCredentialRegistry() weight := 7 - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 4, GroupID: 30, Status: KeyStatusActive, EncryptedValue: "unselected"}, + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 4, GroupID: 30, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "unselected"}, { - ID: 3, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-three", + ID: 3, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three", CooldownUntil: time.Date(2026, time.July, 16, 12, 0, 0, 0, time.UTC), Blacklisted: true, FailureCount: 9, }, - {ID: 2, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "disabled"}, - {ID: 5, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "excluded"}, - {ID: 1, GroupID: 10, WeightManual: &weight, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, + {ID: 2, GroupID: 10, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "disabled"}, + {ID: 5, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "excluded"}, + {ID: 1, GroupID: 10, WeightManual: &weight, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, }) excluded := map[uint]bool{5: true} - got := registry.CollectCandidates([]uint{20, 10}, func(keyID uint) bool { - return excluded[keyID] + got := registry.CollectCredentialCandidates([]uint{20, 10}, func(credentialID uint) bool { + return excluded[credentialID] }, time.Time{}) if len(got) != 1 { t.Fatalf("CollectCandidates() length = %d, want 1: %#v", len(got), got) @@ -794,23 +794,23 @@ func TestKeyRegistryCollectCandidatesFiltersStatusAndExcluded(t *testing.T) { } *got[0].WeightManual = 99 - again := registry.CollectCandidates([]uint{10}, nil, time.Time{}) + again := registry.CollectCredentialCandidates([]uint{10}, nil, time.Time{}) if len(again) != 1 || again[0].WeightManual == nil || *again[0].WeightManual != 7 { t.Fatalf("CollectCandidates() after caller mutation = %#v, want isolated weight 7", again) } - typ := reflect.TypeOf(KeyMeta{}) + typ := reflect.TypeOf(CredentialMeta{}) for _, field := range []string{"EncryptedValue", "CooldownUntil", "Blacklisted", "FailureCount"} { if _, ok := typ.FieldByName(field); ok { - t.Fatalf("KeyMeta exposes private field %s", field) + t.Fatalf("CredentialMeta exposes private field %s", field) } } } func TestKeyRegistryReservedRuntimeMutationsAreAtomic(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }}) until := time.Date(2026, time.July, 16, 12, 0, 0, 0, time.UTC) @@ -858,7 +858,7 @@ func TestKeyRegistryConcurrentMutationsAndCollection(t *testing.T) { operations = 48 ) - registry := NewKeyRegistry() + registry := NewCredentialRegistry() groupIDs := make([]uint, writerCount) for writer := 0; writer < writerCount; writer++ { groupIDs[writer] = uint(writer + 1) @@ -875,45 +875,45 @@ func TestKeyRegistryConcurrentMutationsAndCollection(t *testing.T) { defer wg.Done() <-start for operation := 0; operation < operations; operation++ { - keyID := keyBase + uint(operation) - entry := KeyEntry{ - ID: keyID, GroupID: groupID, Status: KeyStatusActive, - EncryptedValue: fmt.Sprintf("cipher-%d", keyID), + credentialID := keyBase + uint(operation) + entry := CredentialEntry{ + ID: credentialID, GroupID: groupID, Status: CredentialStatusActive, + Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: fmt.Sprintf("cipher-%d", credentialID), } - if err := registry.ApplyImport(groupID, []KeyEntry{entry}); err != nil { - errors <- fmt.Errorf("ApplyImport(group %d, key %d): %w", groupID, keyID, err) + if err := registry.ApplyCredentialImport(groupID, []CredentialEntry{entry}); err != nil { + errors <- fmt.Errorf("ApplyImport(group %d, key %d): %w", groupID, credentialID, err) continue } - if err := registry.SetKeyStatus(keyID, KeyStatusDisabled); err != nil { - errors <- fmt.Errorf("SetKeyStatus(%d, disabled): %w", keyID, err) + if err := registry.SetCredentialStatus(credentialID, CredentialStatusDisabled); err != nil { + errors <- fmt.Errorf("SetCredentialStatus(%d, disabled): %w", credentialID, err) } - if ok := registry.SetAutoWeight(keyID, operation%MaxWeight+1); !ok { - errors <- fmt.Errorf("SetAutoWeight(%d) = false", keyID) + if ok := registry.SetAutoWeight(credentialID, operation%MaxWeight+1); !ok { + errors <- fmt.Errorf("SetAutoWeight(%d) = false", credentialID) } - if ok := registry.SetCooldown(keyID, time.Unix(int64(operation+1), 0)); !ok { - errors <- fmt.Errorf("SetCooldown(%d) = false", keyID) + if ok := registry.SetCooldown(credentialID, time.Unix(int64(operation+1), 0)); !ok { + errors <- fmt.Errorf("SetCooldown(%d) = false", credentialID) } - if ok := registry.SetBlacklisted(keyID); !ok { - errors <- fmt.Errorf("SetBlacklisted(%d) = false", keyID) + if ok := registry.SetBlacklisted(credentialID); !ok { + errors <- fmt.Errorf("SetBlacklisted(%d) = false", credentialID) } - if _, ok := registry.IncrFailure(keyID); !ok { - errors <- fmt.Errorf("IncrFailure(%d) = false", keyID) + if _, ok := registry.IncrFailure(credentialID); !ok { + errors <- fmt.Errorf("IncrFailure(%d) = false", credentialID) } - if ok := registry.ClearFailure(keyID); !ok { - errors <- fmt.Errorf("ClearFailure(%d) = false", keyID) + if ok := registry.ClearFailure(credentialID); !ok { + errors <- fmt.Errorf("ClearFailure(%d) = false", credentialID) } - if ok := registry.Recover(keyID); !ok { - errors <- fmt.Errorf("Recover(%d) = false", keyID) + if ok := registry.Recover(credentialID); !ok { + errors <- fmt.Errorf("Recover(%d) = false", credentialID) } - if err := registry.SetKeyStatus(keyID, KeyStatusActive); err != nil { - errors <- fmt.Errorf("SetKeyStatus(%d, active): %w", keyID, err) + if err := registry.SetCredentialStatus(credentialID, CredentialStatusActive); err != nil { + errors <- fmt.Errorf("SetCredentialStatus(%d, active): %w", credentialID, err) } if operation%3 == 0 { - if removed := registry.RemoveKey(keyID); !removed { - errors <- fmt.Errorf("RemoveKey(%d) = false", keyID) + if removed := registry.RemoveCredential(credentialID); !removed { + errors <- fmt.Errorf("RemoveCredential(%d) = false", credentialID) } - } else if _, ok := registry.EncryptedValue(keyID); !ok { - errors <- fmt.Errorf("EncryptedValue(%d) missing after import", keyID) + } else if _, ok := registry.EncryptedCredentialData(credentialID); !ok { + errors <- fmt.Errorf("EncryptedValue(%d) missing after import", credentialID) } } }() @@ -926,10 +926,10 @@ func TestKeyRegistryConcurrentMutationsAndCollection(t *testing.T) { <-start for operation := 0; operation < writerCount*operations; operation++ { writer := operation % writerCount - keyID := uint((writer+1)*1000 + operation/writerCount%operations) - registry.EncryptedValue(keyID) - registry.BlacklistedKeys() - candidates := registry.CollectCandidates(groupIDs, func(candidateID uint) bool { + credentialID := uint((writer+1)*1000 + operation/writerCount%operations) + registry.EncryptedCredentialData(credentialID) + registry.BlacklistedCredentials() + candidates := registry.CollectCredentialCandidates(groupIDs, func(candidateID uint) bool { return candidateID%19 == uint(reader)%19 }, time.Now()) for index := 1; index < len(candidates); index++ { @@ -962,16 +962,16 @@ func TestKeyRegistryConcurrentMutationsAndCollection(t *testing.T) { } finalWeight := 13 - finalEntries := []KeyEntry{ - {ID: 42, GroupID: 4, Status: KeyStatusActive, EncryptedValue: "final-forty-two"}, - {ID: 43, GroupID: 4, Status: KeyStatusDisabled, EncryptedValue: "final-disabled"}, - {ID: 7, GroupID: 2, WeightManual: &finalWeight, Status: KeyStatusActive, EncryptedValue: "final-seven"}, + finalEntries := []CredentialEntry{ + {ID: 42, GroupID: 4, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "final-forty-two"}, + {ID: 43, GroupID: 4, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "final-disabled"}, + {ID: 7, GroupID: 2, WeightManual: &finalWeight, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "final-seven"}, } - if err := registry.Replace(finalEntries); err != nil { + if err := registry.ReplaceCredentials(finalEntries); err != nil { t.Fatalf("final Replace() error = %v", err) } - got := registry.CollectCandidates([]uint{4, 2}, nil, time.Time{}) + got := registry.CollectCredentialCandidates([]uint{4, 2}, nil, time.Time{}) if len(got) != 2 { t.Fatalf("final CollectCandidates() length = %d, want 2: %#v", len(got), got) } @@ -982,12 +982,12 @@ func TestKeyRegistryConcurrentMutationsAndCollection(t *testing.T) { t.Errorf("final CollectCandidates()[1] = %#v, want group 4 key 42 without manual weight", got[1]) } - for keyID, want := range map[uint]string{ + for credentialID, want := range map[uint]string{ 7: "final-seven", 42: "final-forty-two", 43: "final-disabled", } { - assertEncryptedValue(t, registry, keyID, want, true) + assertEncryptedValue(t, registry, credentialID, want, true) } for writer := 0; writer < writerCount; writer++ { assertEncryptedValue(t, registry, uint((writer+1)*1000), "", false) @@ -997,56 +997,56 @@ func TestKeyRegistryConcurrentMutationsAndCollection(t *testing.T) { func TestValidateKeyEntriesRejectsMalformedEntries(t *testing.T) { tests := []struct { name string - entries []KeyEntry + entries []CredentialEntry wantInError string }{ { name: "zero id", - entries: []KeyEntry{{GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher"}}, + entries: []CredentialEntry{{GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher"}}, wantInError: "id is required", }, { name: "zero group id", - entries: []KeyEntry{{ID: 1, Status: KeyStatusActive, EncryptedValue: "cipher"}}, + entries: []CredentialEntry{{ID: 1, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher"}}, wantInError: "group id is required", }, { name: "invalid status", - entries: []KeyEntry{{ID: 1, GroupID: 10, Status: KeyStatus("cooldown"), EncryptedValue: "cipher"}}, + entries: []CredentialEntry{{ID: 1, GroupID: 10, Status: CredentialStatus("cooldown"), Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher"}}, wantInError: "invalid status", }, { name: "empty ciphertext", - entries: []KeyEntry{{ID: 1, GroupID: 10, Status: KeyStatusActive}}, + entries: []CredentialEntry{{ID: 1, GroupID: 10, Status: CredentialStatusActive}}, wantInError: "encrypted value is required", }, { name: "duplicate id", - entries: []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, - {ID: 1, GroupID: 20, Status: KeyStatusDisabled, EncryptedValue: "cipher-two"}, + entries: []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + {ID: 1, GroupID: 20, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, }, - wantInError: "duplicate key id 1", + wantInError: "duplicate credential id 1", }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - err := ValidateKeyEntries(test.entries) + err := ValidateCredentialEntries(test.entries) if err == nil { - t.Fatal("ValidateKeyEntries() error = nil, want error") + t.Fatal("ValidateCredentialEntries() error = nil, want error") } if !strings.Contains(err.Error(), test.wantInError) { - t.Fatalf("ValidateKeyEntries() error = %q, want substring %q", err, test.wantInError) + t.Fatalf("ValidateCredentialEntries() error = %q, want substring %q", err, test.wantInError) } }) } - if err := ValidateKeyEntries([]KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one"}, - {ID: 2, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "cipher-two"}, + if err := ValidateCredentialEntries([]CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + {ID: 2, GroupID: 10, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, }); err != nil { - t.Fatalf("ValidateKeyEntries(valid) error = %v", err) + t.Fatalf("ValidateCredentialEntries(valid) error = %v", err) } } @@ -1054,40 +1054,40 @@ func TestKeyRegistryApplyImportFailureDoesNotMutateRegistry(t *testing.T) { tests := []struct { name string groupID uint - batch []KeyEntry + batch []CredentialEntry }{ { name: "invalid entry after valid entry", groupID: 10, - batch: []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "updated-one"}, - {ID: 2, GroupID: 10, Status: KeyStatusActive}, + batch: []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "updated-one"}, + {ID: 2, GroupID: 10, Status: CredentialStatusActive}, }, }, { name: "entry group differs from requested group", groupID: 10, - batch: []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "updated-one"}, - {ID: 3, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-three"}, + batch: []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "updated-one"}, + {ID: 3, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three"}, }, }, { name: "zero requested group", - batch: []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "updated-one"}, + batch: []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "updated-one"}, }, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "original-one", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "original-one", }}) - if err := registry.ApplyImport(test.groupID, test.batch); err == nil { + if err := registry.ApplyCredentialImport(test.groupID, test.batch); err == nil { t.Fatal("ApplyImport(invalid batch) error = nil, want error") } assertEncryptedValue(t, registry, 1, "original-one", true) @@ -1101,14 +1101,14 @@ func TestKeyRegistryApplyImportFailureDoesNotMutateRegistry(t *testing.T) { } func TestKeyRegistryApplyImportRejectsExistingIDFromAnotherGroup(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "original-one", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "original-one", }}) - err := registry.ApplyImport(20, []KeyEntry{ - {ID: 2, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "cipher-two"}, - {ID: 1, GroupID: 20, Status: KeyStatusActive, EncryptedValue: "moved-one"}, + err := registry.ApplyCredentialImport(20, []CredentialEntry{ + {ID: 2, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, + {ID: 1, GroupID: 20, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "moved-one"}, }) if err == nil { t.Fatal("ApplyImport(cross-group key id) error = nil, want error") @@ -1124,21 +1124,21 @@ func TestKeyRegistryUpdateKeyConfigAtomicallyPreservesRuntimeState(t *testing.T) oldWeight := 20 newWeight := 80 cooldown := time.Date(2026, time.July, 24, 12, 30, 0, 0, time.UTC) - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ ID: 11, GroupID: 7, WeightManual: &oldWeight, WeightAuto: 42, - Status: KeyStatusActive, CooldownUntil: cooldown, - Blacklisted: true, FailureCount: 3, EncryptedValue: "cipher-secret", + Status: CredentialStatusActive, CooldownUntil: cooldown, + Blacklisted: true, FailureCount: 3, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-secret", }}) - if err := registry.UpdateKeyConfig(11, KeyStatusDisabled, &newWeight); err != nil { - t.Fatalf("UpdateKeyConfig() error = %v", err) + if err := registry.UpdateCredentialConfig(11, CredentialStatusDisabled, &newWeight); err != nil { + t.Fatalf("UpdateCredentialConfig() error = %v", err) } newWeight = 99 registry.mu.RLock() - got := cloneKeyEntry(*registry.buckets[7][11]) + got := cloneCredentialEntry(*registry.buckets[7][11]) registry.mu.RUnlock() - if got.Status != KeyStatusDisabled || got.WeightManual == nil || *got.WeightManual != 80 { + if got.Status != CredentialStatusDisabled || got.WeightManual == nil || *got.WeightManual != 80 { t.Fatalf("config = %#v", got) } if got.WeightAuto != 42 || !got.CooldownUntil.Equal(cooldown) || @@ -1147,11 +1147,11 @@ func TestKeyRegistryUpdateKeyConfigAtomicallyPreservesRuntimeState(t *testing.T) t.Fatalf("runtime fields changed = %#v", got) } - if err := registry.UpdateKeyConfig(11, KeyStatusActive, nil); err != nil { + if err := registry.UpdateCredentialConfig(11, CredentialStatusActive, nil); err != nil { t.Fatal(err) } view := registry.Snapshot()[0] - if view.Status != KeyStatusActive || view.WeightManual != nil || + if view.Status != CredentialStatusActive || view.WeightManual != nil || view.WeightAuto != 42 || !view.Blacklisted || view.FailureCount != 3 { t.Fatalf("cleared manual state = %#v", view) } @@ -1159,27 +1159,27 @@ func TestKeyRegistryUpdateKeyConfigAtomicallyPreservesRuntimeState(t *testing.T) func TestKeyRegistryUpdateKeyConfigRejectsInvalidInputWithoutPartialMutation(t *testing.T) { weight := 30 - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ ID: 1, GroupID: 10, WeightManual: &weight, - Status: KeyStatusActive, EncryptedValue: "cipher-one", + Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }}) before := registry.Snapshot() for _, test := range []struct { - name string - keyID uint - status KeyStatus - weight *int + name string + credentialID uint + status CredentialStatus + weight *int }{ - {name: "missing", keyID: 99, status: KeyStatusDisabled}, - {name: "invalid status", keyID: 1, status: KeyStatus("cooldown")}, - {name: "negative weight", keyID: 1, status: KeyStatusDisabled, weight: intPointer(-1)}, - {name: "large weight", keyID: 1, status: KeyStatusDisabled, weight: intPointer(101)}, + {name: "missing", credentialID: 99, status: CredentialStatusDisabled}, + {name: "invalid status", credentialID: 1, status: CredentialStatus("cooldown")}, + {name: "negative weight", credentialID: 1, status: CredentialStatusDisabled, weight: intPointer(-1)}, + {name: "large weight", credentialID: 1, status: CredentialStatusDisabled, weight: intPointer(101)}, } { t.Run(test.name, func(t *testing.T) { - if err := registry.UpdateKeyConfig(test.keyID, test.status, test.weight); err == nil { - t.Fatal("UpdateKeyConfig() error = nil") + if err := registry.UpdateCredentialConfig(test.credentialID, test.status, test.weight); err == nil { + t.Fatal("UpdateCredentialConfig() error = nil") } if got := registry.Snapshot(); !reflect.DeepEqual(got, before) { t.Fatalf("Registry mutated:\ngot=%#v\nwant=%#v", got, before) @@ -1189,9 +1189,9 @@ func TestKeyRegistryUpdateKeyConfigRejectsInvalidInputWithoutPartialMutation(t * } func TestKeyRegistryUpdateKeyConfigIsRaceSafeWithRuntimeMutations(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-one", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }}) start := make(chan struct{}) var workers sync.WaitGroup @@ -1203,7 +1203,7 @@ func TestKeyRegistryUpdateKeyConfigIsRaceSafeWithRuntimeMutations(t *testing.T) <-start for index := 0; index < 100; index++ { weight := (worker + index) % (MaxWeight + 1) - _ = registry.UpdateKeyConfig(1, KeyStatusActive, &weight) + _ = registry.UpdateCredentialConfig(1, CredentialStatusActive, &weight) _ = registry.SetAutoWeight(1, (index%MaxWeight)+1) _ = registry.SetCooldown(1, time.Unix(int64(index), 0)) _, _ = registry.IncrFailure(1) @@ -1213,10 +1213,10 @@ func TestKeyRegistryUpdateKeyConfigIsRaceSafeWithRuntimeMutations(t *testing.T) close(start) workers.Wait() view := registry.Snapshot()[0] - if view.GroupID != 10 || view.Status != KeyStatusActive { + if view.GroupID != 10 || view.Status != CredentialStatusActive { t.Fatalf("final safe view = %#v", view) } - if value, ok := registry.EncryptedValue(1); !ok || value != "cipher-one" { + if value, ok := registry.EncryptedCredentialData(1); !ok || value != "cipher-one" { t.Fatalf("credential = %q, %t", value, ok) } } @@ -1225,28 +1225,28 @@ func intPointer(value int) *int { return &value } -func mustReplaceKeyEntries(t *testing.T, registry *KeyRegistry, entries []KeyEntry) { +func mustReplaceKeyEntries(t *testing.T, registry *CredentialRegistry, entries []CredentialEntry) { t.Helper() - if err := registry.Replace(entries); err != nil { + if err := registry.ReplaceCredentials(entries); err != nil { t.Fatalf("Replace() error = %v", err) } } -func assertEncryptedValue(t *testing.T, registry *KeyRegistry, keyID uint, want string, wantOK bool) { +func assertEncryptedValue(t *testing.T, registry *CredentialRegistry, credentialID uint, want string, wantOK bool) { t.Helper() - got, ok := registry.EncryptedValue(keyID) + got, ok := registry.EncryptedCredentialData(credentialID) if got != want || ok != wantOK { - t.Errorf("EncryptedValue(%d) = %q, %t, want %q, %t", keyID, got, ok, want, wantOK) + t.Errorf("EncryptedValue(%d) = %q, %t, want %q, %t", credentialID, got, ok, want, wantOK) } } -func keyStatus(t *testing.T, registry *KeyRegistry, keyID uint) KeyStatus { +func keyStatus(t *testing.T, registry *CredentialRegistry, credentialID uint) CredentialStatus { t.Helper() registry.mu.RLock() defer registry.mu.RUnlock() - groupID, ok := registry.keyGroups[keyID] + groupID, ok := registry.credentialGroups[credentialID] if !ok { - t.Fatalf("key %d does not exist", keyID) + t.Fatalf("key %d does not exist", credentialID) } - return registry.buckets[groupID][keyID].Status + return registry.buckets[groupID][credentialID].Status } diff --git a/internal/state/runtime_checkpoint.go b/internal/state/runtime_checkpoint.go index 7b887bb6..e4c3f1c3 100644 --- a/internal/state/runtime_checkpoint.go +++ b/internal/state/runtime_checkpoint.go @@ -5,10 +5,10 @@ import ( "time" ) -// KeyRuntimeCheckpoint contains only the mutable health state that is safe to +// CredentialRuntimeCheckpoint contains only the mutable health state that is safe to // carry across a process restart. Persisted key configuration remains owned by // SQLite and is matched by ID plus group ID during restore. -type KeyRuntimeCheckpoint struct { +type CredentialRuntimeCheckpoint struct { ID uint `json:"id"` GroupID uint `json:"group_id"` WeightAuto int `json:"weight_auto"` @@ -19,12 +19,12 @@ type KeyRuntimeCheckpoint struct { // CaptureRuntimeCheckpoint returns detached runtime health state in stable // order. It intentionally excludes credentials and failure generations. -func (r *KeyRegistry) CaptureRuntimeCheckpoint() []KeyRuntimeCheckpoint { +func (r *CredentialRegistry) CaptureRuntimeCheckpoint() []CredentialRuntimeCheckpoint { r.mu.RLock() - checkpoints := make([]KeyRuntimeCheckpoint, 0, len(r.keyGroups)) + checkpoints := make([]CredentialRuntimeCheckpoint, 0, len(r.credentialGroups)) for _, bucket := range r.buckets { for _, entry := range bucket { - checkpoints = append(checkpoints, KeyRuntimeCheckpoint{ + checkpoints = append(checkpoints, CredentialRuntimeCheckpoint{ ID: entry.ID, GroupID: entry.GroupID, WeightAuto: entry.WeightAuto, @@ -46,7 +46,7 @@ func (r *KeyRegistry) CaptureRuntimeCheckpoint() []KeyRuntimeCheckpoint { // RestoreRuntimeCheckpoint applies matching runtime health state and skips // keys that were removed or moved to another group while the process was down. -func (r *KeyRegistry) RestoreRuntimeCheckpoint(checkpoints []KeyRuntimeCheckpoint) int { +func (r *CredentialRegistry) RestoreRuntimeCheckpoint(checkpoints []CredentialRuntimeCheckpoint) int { r.mu.Lock() defer r.mu.Unlock() @@ -57,7 +57,7 @@ func (r *KeyRegistry) RestoreRuntimeCheckpoint(checkpoints []KeyRuntimeCheckpoin checkpoint.FailureCount < 0 { continue } - groupID, ok := r.keyGroups[checkpoint.ID] + groupID, ok := r.credentialGroups[checkpoint.ID] if !ok || groupID != checkpoint.GroupID { continue } diff --git a/internal/state/runtime_settings_test.go b/internal/state/runtime_settings_test.go index 33d037a6..2950e51b 100644 --- a/internal/state/runtime_settings_test.go +++ b/internal/state/runtime_settings_test.go @@ -7,8 +7,8 @@ import ( "testing" "time" + "gpt-load/internal/channel" "gpt-load/internal/platform/config" - "gpt-load/internal/protocol" ) func TestCompilePublishesDefaultRuntimeSettingsWithoutGroups(t *testing.T) { @@ -74,8 +74,8 @@ func TestCompileRejectsInvalidGlobalSettingsWithoutGroups(t *testing.T) { } func TestCompileValidatesDisabledGroupSettings(t *testing.T) { - _, err := Compile(CompileInput{Groups: []GroupConfig{{ - ID: 1, Protocols: []protocol.Protocol{protocol.OpenAICompletions}, + _, err := Compile(CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{{ + ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Settings: config.Settings{"request_log_retention_days": 30}, Enabled: false, }}}) diff --git a/internal/state/runtime_view.go b/internal/state/runtime_view.go index c22723d1..2726daea 100644 --- a/internal/state/runtime_view.go +++ b/internal/state/runtime_view.go @@ -5,53 +5,57 @@ import ( "time" ) -type KeyRuntimeView struct { - ID uint - GroupID uint - WeightManual *int - WeightAuto int - Status KeyStatus - CooldownUntil time.Time - Blacklisted bool - FailureCount int +type CredentialRuntimeView struct { + ID uint + GroupID uint + Version uint64 + IdentityGeneration uint64 + WeightManual *int + WeightAuto int + Status CredentialStatus + CooldownUntil time.Time + Blacklisted bool + FailureCount int } -type KeyRuntimeState string +type CredentialRuntimeState string const ( - KeyRuntimeAvailable KeyRuntimeState = "available" - KeyRuntimeDisabled KeyRuntimeState = "disabled" - KeyRuntimeBlacklisted KeyRuntimeState = "blacklisted" - KeyRuntimeCooldown KeyRuntimeState = "cooldown" + CredentialRuntimeAvailable CredentialRuntimeState = "available" + CredentialRuntimeDisabled CredentialRuntimeState = "disabled" + CredentialRuntimeBlacklisted CredentialRuntimeState = "blacklisted" + CredentialRuntimeCooldown CredentialRuntimeState = "cooldown" ) -func (view KeyRuntimeView) RuntimeState(now time.Time) KeyRuntimeState { - if view.Status != KeyStatusActive { - return KeyRuntimeDisabled +func (view CredentialRuntimeView) RuntimeState(now time.Time) CredentialRuntimeState { + if view.Status != CredentialStatusActive { + return CredentialRuntimeDisabled } if view.Blacklisted { - return KeyRuntimeBlacklisted + return CredentialRuntimeBlacklisted } if view.CooldownUntil.After(now) { - return KeyRuntimeCooldown + return CredentialRuntimeCooldown } - return KeyRuntimeAvailable + return CredentialRuntimeAvailable } -func runtimeView(entry *KeyEntry) KeyRuntimeView { - return KeyRuntimeView{ - ID: entry.ID, - GroupID: entry.GroupID, - WeightManual: cloneWeight(entry.WeightManual), - WeightAuto: entry.WeightAuto, - Status: entry.Status, - CooldownUntil: entry.CooldownUntil, - Blacklisted: entry.Blacklisted, - FailureCount: entry.FailureCount, +func runtimeView(entry *CredentialEntry) CredentialRuntimeView { + return CredentialRuntimeView{ + ID: entry.ID, + GroupID: entry.GroupID, + Version: entry.Version, + IdentityGeneration: entry.IdentityGeneration, + WeightManual: cloneWeight(entry.WeightManual), + WeightAuto: entry.WeightAuto, + Status: entry.Status, + CooldownUntil: entry.CooldownUntil, + Blacklisted: entry.Blacklisted, + FailureCount: entry.FailureCount, } } -func sortRuntimeViews(views []KeyRuntimeView) { +func sortRuntimeViews(views []CredentialRuntimeView) { sort.Slice(views, func(i, j int) bool { if views[i].GroupID != views[j].GroupID { return views[i].GroupID < views[j].GroupID @@ -60,9 +64,9 @@ func sortRuntimeViews(views []KeyRuntimeView) { }) } -func (r *KeyRegistry) Snapshot() []KeyRuntimeView { +func (r *CredentialRegistry) Snapshot() []CredentialRuntimeView { r.mu.RLock() - views := make([]KeyRuntimeView, 0, len(r.keyGroups)) + views := make([]CredentialRuntimeView, 0, len(r.credentialGroups)) for _, bucket := range r.buckets { for _, entry := range bucket { views = append(views, runtimeView(entry)) diff --git a/internal/state/snapshot.go b/internal/state/snapshot.go index 603f798d..8c7fea2c 100644 --- a/internal/state/snapshot.go +++ b/internal/state/snapshot.go @@ -2,35 +2,50 @@ package state import ( + "encoding/json" "fmt" "sort" "strings" "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/config" - "gpt-load/internal/pricing" "gpt-load/internal/protocol" ) type CompileInput struct { - SystemSettings config.Settings - Groups []GroupConfig - AccessKeys []AccessKeyConfig + SystemSettings config.Settings + ChannelRegistry *channel.Registry + Groups []GroupConfig + Credentials []CredentialConfig + AccessKeys []AccessKeyConfig } type GroupConfig struct { ID uint Name string - ProviderID *string - UpstreamURL string + ChannelID channel.ID + Params json.RawMessage ValidationModel string - Protocols []protocol.Protocol Models []ModelConfig Settings config.Settings WeightManual *int Enabled bool } +// CredentialConfig contains only non-secret credential metadata required to +// validate a runtime configuration publication. +type CredentialConfig struct { + ID uint + GroupID uint + Status CredentialStatus + WeightManual *int + Version uint64 + IdentityGeneration uint64 + Fingerprint string +} + type ModelConfig struct { ID string Alias string @@ -68,8 +83,18 @@ type FilterSet struct { type RouteTarget struct { GroupID uint UpstreamModelID string + Mode channel.RouteMode + ResolvedTarget channel.ResolvedTarget } +// NoModelRouteKey identifies operations whose upstream resource ID, rather +// than a model, determines the target after affinity resolution. +const NoModelRouteKey = "" + +// ExecutionCandidateIndex indexes targets by client protocol, logical +// operation, and external model. Resource operations use NoModelRouteKey. +type ExecutionCandidateIndex map[protocol.Protocol]map[execution.Operation]map[string][]RouteTarget + type TimeoutConfig struct { Connect time.Duration FirstByte time.Duration @@ -85,10 +110,11 @@ type HeaderRules struct { type GroupView struct { ID uint Name string - ProviderID *string - UpstreamURL string + ChannelID channel.ID + Params json.RawMessage + ResolvedTarget channel.ResolvedTarget ValidationModel string - Protocols []protocol.Protocol + ClientProtocols []protocol.Protocol Models []ModelConfig Timeouts TimeoutConfig HeaderRules HeaderRules @@ -112,16 +138,14 @@ type AccessKeyView struct { } type ConfigSnapshot struct { - Revision uint64 - Settings RuntimeSettings - Candidates map[protocol.Protocol]map[string][]RouteTarget - ProtocolCandidates map[protocol.Protocol][]uint - Groups map[uint]GroupView - AccessKeysByHash map[string]AccessKeyView - RouteCatalog map[protocol.Protocol]map[string][]RouteTarget - ProtocolRouteCatalog map[protocol.Protocol][]uint - GroupCatalog map[uint]GroupCatalogView - AccessKeysByID map[uint]AccessKeyView + Revision uint64 + Settings RuntimeSettings + ExecutionCandidates ExecutionCandidateIndex + ExecutionRouteCatalog ExecutionCandidateIndex + Groups map[uint]GroupView + AccessKeysByHash map[string]AccessKeyView + GroupCatalog map[uint]GroupCatalogView + AccessKeysByID map[uint]AccessKeyView } func Compile(input CompileInput) (*ConfigSnapshot, error) { @@ -134,15 +158,13 @@ func Compile(input CompileInput) (*ConfigSnapshot, error) { } snapshot := &ConfigSnapshot{ - Settings: runtimeSettings, - Candidates: make(map[protocol.Protocol]map[string][]RouteTarget), - ProtocolCandidates: make(map[protocol.Protocol][]uint), - Groups: make(map[uint]GroupView), - AccessKeysByHash: make(map[string]AccessKeyView), - RouteCatalog: make(map[protocol.Protocol]map[string][]RouteTarget), - ProtocolRouteCatalog: make(map[protocol.Protocol][]uint), - GroupCatalog: make(map[uint]GroupCatalogView), - AccessKeysByID: make(map[uint]AccessKeyView), + Settings: runtimeSettings, + ExecutionCandidates: make(ExecutionCandidateIndex), + ExecutionRouteCatalog: make(ExecutionCandidateIndex), + Groups: make(map[uint]GroupView), + AccessKeysByHash: make(map[string]AccessKeyView), + GroupCatalog: make(map[uint]GroupCatalogView), + AccessKeysByID: make(map[uint]AccessKeyView), } for _, group := range input.Groups { @@ -150,8 +172,9 @@ func Compile(input CompileInput) (*ConfigSnapshot, error) { ID: group.ID, Name: group.Name, Enabled: group.Enabled, WeightManual: cloneWeight(group.WeightManual), } - appendRouteTarget(snapshot.RouteCatalog, group) - appendProtocolGroup(snapshot.ProtocolRouteCatalog, group) + if err := appendExecutionTargets(snapshot.ExecutionRouteCatalog, input.ChannelRegistry, group); err != nil { + return nil, err + } resolved, err := ResolveGroupRuntimeSettings(runtimeSettings, group.Settings) if err != nil { return nil, fmt.Errorf("compile group %d settings: %w", group.ID, err) @@ -159,22 +182,34 @@ func Compile(input CompileInput) (*ConfigSnapshot, error) { if !group.Enabled { continue } + view := GroupView{ ID: group.ID, Name: group.Name, - ProviderID: cloneStringPointer(group.ProviderID), - UpstreamURL: group.UpstreamURL, ValidationModel: strings.TrimSpace(group.ValidationModel), - Protocols: append([]protocol.Protocol(nil), group.Protocols...), Models: append([]ModelConfig(nil), group.Models...), Timeouts: resolved.Timeouts, HeaderRules: resolved.HeaderRules, InjectUsageOptions: resolved.InjectUsageOptions, WeightManual: cloneWeight(group.WeightManual), } + params, err := input.ChannelRegistry.ValidateParams(group.ChannelID, group.Params) + if err != nil { + return nil, fmt.Errorf("compile group %d params: %w", group.ID, err) + } + target, err := input.ChannelRegistry.Resolve(group.ChannelID, group.Params) + if err != nil { + return nil, fmt.Errorf("compile group %d channel: %w", group.ID, err) + } + descriptor, _ := input.ChannelRegistry.Get(group.ChannelID) + view.ChannelID = group.ChannelID + view.Params = params.CanonicalJSON() + view.ResolvedTarget = cloneResolvedTarget(target) + view.ClientProtocols = append([]protocol.Protocol(nil), descriptor.ClientProtocols...) + if err := appendExecutionTargets(snapshot.ExecutionCandidates, input.ChannelRegistry, group); err != nil { + return nil, err + } snapshot.Groups[group.ID] = view - appendRouteTarget(snapshot.Candidates, group) - appendProtocolGroup(snapshot.ProtocolCandidates, group) } for _, accessKey := range input.AccessKeys { @@ -184,10 +219,8 @@ func Compile(input CompileInput) (*ConfigSnapshot, error) { } } - sortRouteIndex(snapshot.Candidates) - sortRouteIndex(snapshot.RouteCatalog) - sortProtocolGroupIndex(snapshot.ProtocolCandidates) - sortProtocolGroupIndex(snapshot.ProtocolRouteCatalog) + sortExecutionRouteIndex(snapshot.ExecutionCandidates) + sortExecutionRouteIndex(snapshot.ExecutionRouteCatalog) return snapshot, nil } @@ -198,54 +231,99 @@ func newAccessKeyView(input AccessKeyConfig) AccessKeyView { } } -func cloneStringPointer(value *string) *string { - if value == nil { - return nil +func appendExecutionTargets( + index ExecutionCandidateIndex, + registry *channel.Registry, + group GroupConfig, +) error { + target, err := registry.Resolve(group.ChannelID, group.Params) + if err != nil { + return fmt.Errorf("compile group %d channel: %w", group.ID, err) } - cloned := *value - return &cloned -} - -func appendRouteTarget(index map[protocol.Protocol]map[string][]RouteTarget, group GroupConfig) { - for _, model := range group.Models { - target := RouteTarget{ - GroupID: group.ID, - UpstreamModelID: strings.TrimSpace(model.ID), - } - external := externalModelName(model) - for _, value := range group.Protocols { - if index[value] == nil { - index[value] = make(map[string][]RouteTarget) + descriptor, ok := registry.Get(group.ChannelID) + if !ok { + return fmt.Errorf("compile group %d channel: unknown channel %q", group.ID, group.ChannelID) + } + for _, clientProtocol := range descriptor.ClientProtocols { + capabilities := target.Capabilities(clientProtocol) + for _, operation := range capabilities.Operations() { + if operation == execution.OperationListModels || operation == execution.OperationProbe { + continue + } + mode, ok := target.Mode(clientProtocol, operation) + if !ok { + return fmt.Errorf("compile group %d channel has no route mode for %q/%q", group.ID, clientProtocol, operation) + } + switch operation { + case execution.OperationResponsesRetrieve, + execution.OperationResponsesDelete, + execution.OperationResponsesCancel, + execution.OperationResponsesInputItems: + appendExecutionTarget(index, clientProtocol, operation, NoModelRouteKey, RouteTarget{ + GroupID: group.ID, Mode: mode, ResolvedTarget: cloneResolvedTarget(target), + }) + case execution.OperationResponsesPassthrough: + appendExecutionTarget(index, clientProtocol, operation, NoModelRouteKey, RouteTarget{ + GroupID: group.ID, Mode: mode, ResolvedTarget: cloneResolvedTarget(target), + }) + fallthrough + case execution.OperationChatCompletion, + execution.OperationResponsesCreate, + execution.OperationResponsesCompact, + execution.OperationResponsesInputTokens: + for _, model := range group.Models { + appendExecutionTarget(index, clientProtocol, operation, externalModelName(model), RouteTarget{ + GroupID: group.ID, UpstreamModelID: strings.TrimSpace(model.ID), + Mode: mode, ResolvedTarget: cloneResolvedTarget(target), + }) + } + default: + return fmt.Errorf("compile group %d channel has unsupported routable operation %q", group.ID, operation) } - index[value][external] = append(index[value][external], target) } } + return nil } -func appendProtocolGroup(index map[protocol.Protocol][]uint, group GroupConfig) { - for _, value := range group.Protocols { - index[value] = append(index[value], group.ID) +func appendExecutionTarget( + index ExecutionCandidateIndex, + clientProtocol protocol.Protocol, + operation execution.Operation, + externalModel string, + target RouteTarget, +) { + if index[clientProtocol] == nil { + index[clientProtocol] = make(map[execution.Operation]map[string][]RouteTarget) } -} - -func sortProtocolGroupIndex(index map[protocol.Protocol][]uint) { - for value := range index { - sort.Slice(index[value], func(i, j int) bool { - return index[value][i] < index[value][j] - }) + if index[clientProtocol][operation] == nil { + index[clientProtocol][operation] = make(map[string][]RouteTarget) } + index[clientProtocol][operation][externalModel] = append( + index[clientProtocol][operation][externalModel], + target, + ) } -func sortRouteIndex(index map[protocol.Protocol]map[string][]RouteTarget) { - for _, byModel := range index { - for model := range byModel { - sort.Slice(byModel[model], func(i, j int) bool { - left, right := byModel[model][i], byModel[model][j] - if left.GroupID != right.GroupID { - return left.GroupID < right.GroupID - } - return left.UpstreamModelID < right.UpstreamModelID - }) +func cloneResolvedTarget(target channel.ResolvedTarget) channel.ResolvedTarget { + target.TargetConfig = append(json.RawMessage(nil), target.TargetConfig...) + return target +} + +func sortExecutionRouteIndex(index ExecutionCandidateIndex) { + for _, byOperation := range index { + for _, byModel := range byOperation { + for model := range byModel { + sort.Slice(byModel[model], func(i, j int) bool { + left, right := byModel[model][i], byModel[model][j] + if left.Mode != right.Mode { + return left.Mode == channel.RouteNative + } + if left.GroupID != right.GroupID { + return left.GroupID < right.GroupID + } + return left.UpstreamModelID < right.UpstreamModelID + }) + } } } } @@ -260,27 +338,18 @@ func validateCompileInput(input CompileInput) error { return fmt.Errorf("duplicate group id %d", group.ID) } groupIDs[group.ID] = struct{}{} - if group.ProviderID != nil { - if _, err := pricing.ProviderScopeKey(*group.ProviderID); err != nil { - return fmt.Errorf("group %d has invalid provider id: %w", group.ID, err) - } + if input.ChannelRegistry == nil { + return fmt.Errorf("group %d channel registry is required", group.ID) + } + if group.ChannelID == "" { + return fmt.Errorf("group %d channel id is required", group.ID) + } + if _, err := input.ChannelRegistry.Resolve(group.ChannelID, group.Params); err != nil { + return fmt.Errorf("group %d channel %q: %w", group.ID, group.ChannelID, err) } if err := validateManualWeight(fmt.Sprintf("group %d", group.ID), group.WeightManual); err != nil { return err } - if len(group.Protocols) == 0 { - return fmt.Errorf("group %d protocols are required", group.ID) - } - seenProtocols := make(map[protocol.Protocol]struct{}, len(group.Protocols)) - for _, value := range group.Protocols { - if !value.Valid() { - return fmt.Errorf("group %d has invalid protocol %q", group.ID, value) - } - if _, duplicate := seenProtocols[value]; duplicate { - return fmt.Errorf("group %d has duplicate protocol %q", group.ID, value) - } - seenProtocols[value] = struct{}{} - } seenModels := make(map[string]struct{}, len(group.Models)) for _, model := range group.Models { if strings.TrimSpace(model.ID) == "" { @@ -294,6 +363,40 @@ func validateCompileInput(input CompileInput) error { } } + credentialIDs := make(map[uint]struct{}, len(input.Credentials)) + for _, credential := range input.Credentials { + if credential.ID == 0 { + return fmt.Errorf("credential id is required") + } + if _, duplicate := credentialIDs[credential.ID]; duplicate { + return fmt.Errorf("duplicate credential id %d", credential.ID) + } + credentialIDs[credential.ID] = struct{}{} + if credential.GroupID == 0 { + return fmt.Errorf("credential %d group id is required", credential.ID) + } + if _, ok := groupIDs[credential.GroupID]; !ok { + return fmt.Errorf("credential %d belongs to unknown group %d", credential.ID, credential.GroupID) + } + switch credential.Status { + case CredentialStatusActive, CredentialStatusDisabled: + default: + return fmt.Errorf("credential %d has invalid status %q", credential.ID, credential.Status) + } + if err := validateManualWeight(fmt.Sprintf("credential %d", credential.ID), credential.WeightManual); err != nil { + return err + } + if credential.Version == 0 { + return fmt.Errorf("credential %d version is required", credential.ID) + } + if credential.IdentityGeneration == 0 { + return fmt.Errorf("credential %d identity generation is required", credential.ID) + } + if strings.TrimSpace(credential.Fingerprint) == "" { + return fmt.Errorf("credential %d fingerprint is required", credential.ID) + } + } + accessKeyIDs := make(map[uint]struct{}, len(input.AccessKeys)) hashes := make(map[string]struct{}, len(input.AccessKeys)) for _, accessKey := range input.AccessKeys { diff --git a/internal/state/snapshot_test.go b/internal/state/snapshot_test.go index ebc500af..4cec0b87 100644 --- a/internal/state/snapshot_test.go +++ b/internal/state/snapshot_test.go @@ -5,83 +5,76 @@ import ( "reflect" "strings" "testing" - "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/platform/config" "gpt-load/internal/protocol" ) func TestCompileIndexesExternalModelsAndPreservesUpstreamIDs(t *testing.T) { - input := CompileInput{Groups: []GroupConfig{ - { - ID: 1, Name: "one", UpstreamURL: "https://one.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{ - {ID: "provider-a", Alias: "public"}, - {ID: "provider-a", Alias: "secondary"}, - {ID: "plain"}, + t.Parallel() + + input := CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []GroupConfig{ + { + ID: 1, Name: "one", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []ModelConfig{ + {ID: "provider-a", Alias: "public"}, + {ID: "provider-a", Alias: "secondary"}, + {ID: "plain"}, + }, + Enabled: true, + }, + { + ID: 2, Name: "two", ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://proxy.example/v1"}`), + Models: []ModelConfig{{ID: "provider-b", Alias: "public"}}, + Enabled: true, }, - Enabled: true, }, - { - ID: 2, Name: "two", UpstreamURL: "https://two.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "provider-b", Alias: "public"}}, Enabled: true, - }, - }} + } snapshot, err := Compile(input) if err != nil { t.Fatalf("Compile() error = %v", err) } - wantPublic := []RouteTarget{ - {GroupID: 1, UpstreamModelID: "provider-a"}, - {GroupID: 2, UpstreamModelID: "provider-b"}, + index := snapshot.ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion] + public := index["public"] + if len(public) != 2 || public[0].UpstreamModelID != "provider-a" || public[1].UpstreamModelID != "provider-b" { + t.Fatalf("public targets = %#v", public) } - if got := snapshot.Candidates[protocol.OpenAICompletions]["public"]; !reflect.DeepEqual(got, wantPublic) { - t.Fatalf("public targets = %#v, want %#v", got, wantPublic) - } - if got := snapshot.Candidates[protocol.OpenAICompletions]["secondary"]; len(got) != 1 || got[0].UpstreamModelID != "provider-a" { + if got := index["secondary"]; len(got) != 1 || got[0].UpstreamModelID != "provider-a" { t.Fatalf("secondary targets = %#v", got) } - if got := snapshot.Candidates[protocol.OpenAICompletions]["plain"]; len(got) != 1 || got[0].UpstreamModelID != "plain" { + if got := index["plain"]; len(got) != 1 || got[0].UpstreamModelID != "plain" { t.Fatalf("plain targets = %#v", got) } - if _, exists := snapshot.Candidates[protocol.OpenAICompletions]["provider-a"]; exists { + if _, exists := index["provider-a"]; exists { t.Fatal("aliased upstream id entered external index") } } func TestCompileBuildsManagementCatalogsWithoutChangingActiveIndexes(t *testing.T) { + t.Parallel() + disabledWeight := 20 input := CompileInput{ + ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{ { - ID: 2, Name: "disabled", UpstreamURL: "https://disabled.example", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "provider-disabled", Alias: "public"}}, - WeightManual: &disabledWeight, Enabled: false, + ID: 2, Name: "disabled", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []ModelConfig{{ID: "provider-disabled", Alias: "public"}}, WeightManual: &disabledWeight, }, { - ID: 1, Name: "active", UpstreamURL: "https://active.example", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "provider-active", Alias: "public"}}, - Enabled: true, + ID: 1, Name: "active", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []ModelConfig{{ID: "provider-active", Alias: "public"}}, Enabled: true, }, }, AccessKeys: []AccessKeyConfig{ - { - ID: 11, Name: "active-client", KeyHash: "active-hash", - Status: AccessKeyStatusActive, - Filters: FilterSet{Groups: map[uint]struct{}{1: {}}}, - RPMLimit: 10, - }, - { - ID: 12, Name: "disabled-client", KeyHash: "disabled-hash", - Status: AccessKeyStatusDisabled, - Filters: FilterSet{Models: map[string]struct{}{"public": {}}}, - RPMLimit: 20, - }, + {ID: 11, Name: "active-client", KeyHash: "active-hash", Status: AccessKeyStatusActive, Filters: FilterSet{Groups: map[uint]struct{}{1: {}}}, RPMLimit: 10}, + {ID: 12, Name: "disabled-client", KeyHash: "disabled-hash", Status: AccessKeyStatusDisabled, Filters: FilterSet{Models: map[string]struct{}{"public": {}}}, RPMLimit: 20}, }, } @@ -92,174 +85,125 @@ func TestCompileBuildsManagementCatalogsWithoutChangingActiveIndexes(t *testing. if len(snapshot.Groups) != 1 || snapshot.Groups[1].Name != "active" { t.Fatalf("active Groups = %#v", snapshot.Groups) } - if got := snapshot.Candidates[protocol.OpenAICompletions]["public"]; !reflect.DeepEqual(got, []RouteTarget{{ - GroupID: 1, UpstreamModelID: "provider-active", - }}) { - t.Fatalf("active Candidates = %#v", got) + active := snapshot.ExecutionCandidates[protocol.OpenAICompletions][execution.OperationChatCompletion]["public"] + if len(active) != 1 || active[0].GroupID != 1 { + t.Fatalf("active candidates = %#v", active) } - wantRoutes := []RouteTarget{ - {GroupID: 1, UpstreamModelID: "provider-active"}, - {GroupID: 2, UpstreamModelID: "provider-disabled"}, + routes := snapshot.ExecutionRouteCatalog[protocol.OpenAICompletions][execution.OperationChatCompletion]["public"] + if len(routes) != 2 || routes[0].GroupID != 1 || routes[1].GroupID != 2 { + t.Fatalf("route catalog = %#v", routes) } - if got := snapshot.RouteCatalog[protocol.OpenAICompletions]["public"]; !reflect.DeepEqual(got, wantRoutes) { - t.Fatalf("RouteCatalog = %#v, want %#v", got, wantRoutes) + if got := snapshot.GroupCatalog[2]; got.Enabled || got.WeightManual == nil || *got.WeightManual != 20 { + t.Fatalf("disabled group catalog = %#v", got) } - if got := snapshot.GroupCatalog[2]; got.ID != 2 || got.Name != "disabled" || - got.Enabled || got.WeightManual == nil || *got.WeightManual != 20 { - t.Fatalf("disabled GroupCatalog entry = %#v", got) + if _, ok := snapshot.AccessKeysByHash["disabled-hash"]; ok { + t.Fatal("disabled access key entered active hash index") } - if len(snapshot.AccessKeysByHash) != 1 || - snapshot.AccessKeysByHash["active-hash"].Status != AccessKeyStatusActive { - t.Fatalf("active AccessKeysByHash = %#v", snapshot.AccessKeysByHash) - } - if _, exists := snapshot.AccessKeysByHash["disabled-hash"]; exists { - t.Fatal("disabled key entered AccessKeysByHash") - } - if got := snapshot.AccessKeysByID[12]; got.ID != 12 || - got.Status != AccessKeyStatusDisabled || got.RPMLimit != 20 { - t.Fatalf("disabled AccessKeysByID entry = %#v", got) + if got := snapshot.AccessKeysByID[12]; got.Status != AccessKeyStatusDisabled || got.RPMLimit != 20 { + t.Fatalf("disabled access key catalog = %#v", got) } } -func TestCompileBuildsProtocolOnlyRuntimeAndCatalogIndexes(t *testing.T) { +func TestCompileCarriesSettingsAndValidationModel(t *testing.T) { t.Parallel() - snapshot, err := Compile(CompileInput{Groups: []GroupConfig{ - { - ID: 3, Name: "disabled responses", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Enabled: false, - }, - { - ID: 2, Name: "active both", - Protocols: []protocol.Protocol{ - protocol.OpenAICompletions, - protocol.OpenAIResponses, - }, - Models: []ModelConfig{{ID: "provider", Alias: "public"}}, - Enabled: true, - }, - { - ID: 1, Name: "active responses without models", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Enabled: true, - }, - }}) + snapshot, err := Compile(CompileInput{ + ChannelRegistry: channel.NewRegistry(), + SystemSettings: config.Settings{"connect_timeout": json.Number("20")}, + Groups: []GroupConfig{{ + ID: 1, Name: "one", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + ValidationModel: " probe-model ", + Models: []ModelConfig{{ID: "real-model", Alias: "public-model"}}, + Settings: config.Settings{"request_timeout": json.Number("30")}, Enabled: true, + }}, + }) if err != nil { t.Fatalf("Compile() error = %v", err) } - - if got, want := snapshot.ProtocolCandidates[protocol.OpenAIResponses], []uint{1, 2}; !reflect.DeepEqual(got, want) { - t.Fatalf("Responses ProtocolCandidates = %#v, want %#v", got, want) - } - if got, want := snapshot.ProtocolRouteCatalog[protocol.OpenAIResponses], []uint{1, 2, 3}; !reflect.DeepEqual(got, want) { - t.Fatalf("Responses ProtocolRouteCatalog = %#v, want %#v", got, want) - } - if got, want := snapshot.ProtocolCandidates[protocol.OpenAICompletions], []uint{2}; !reflect.DeepEqual(got, want) { - t.Fatalf("Chat ProtocolCandidates = %#v, want %#v", got, want) + view := snapshot.Groups[1] + if view.ValidationModel != "probe-model" || view.Timeouts.Connect.Seconds() != 20 || view.Timeouts.Request.Seconds() != 30 { + t.Fatalf("group runtime view = %#v", view) } } -func TestCompileCarriesValidationModelWithoutChangingCandidates(t *testing.T) { - snapshot, err := Compile(CompileInput{Groups: []GroupConfig{{ - ID: 1, - Name: "one", - UpstreamURL: "https://one.example.com", - ValidationModel: "probe-model", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "real-model", Alias: "public-model"}}, - Enabled: true, - }}}) +func TestCompileOwnsInputData(t *testing.T) { + t.Parallel() + + weight := 25 + filters := FilterSet{ + Groups: map[uint]struct{}{1: {}}, + Protocols: map[protocol.Protocol]struct{}{protocol.OpenAICompletions: {}}, + Models: map[string]struct{}{"public": {}}, + } + input := CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []GroupConfig{{ + ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []ModelConfig{{ID: "upstream", Alias: "public"}}, WeightManual: &weight, Enabled: true, + }}, + AccessKeys: []AccessKeyConfig{{ID: 1, KeyHash: "hash", Status: AccessKeyStatusActive, Filters: filters}}, + } + snapshot, err := Compile(input) if err != nil { t.Fatalf("Compile() error = %v", err) } - if got := snapshot.Groups[1].ValidationModel; got != "probe-model" { - t.Fatalf("ValidationModel = %q, want probe-model", got) + input.Groups[0].Models[0] = ModelConfig{ID: "changed"} + weight = 99 + filters.Groups[2] = struct{}{} + filters.Protocols[protocol.Gemini] = struct{}{} + filters.Models["changed"] = struct{}{} + + view := snapshot.Groups[1] + if !reflect.DeepEqual(view.Models, []ModelConfig{{ID: "upstream", Alias: "public"}}) || view.WeightManual == nil || *view.WeightManual != 25 { + t.Fatalf("group view changed with input = %#v", view) } - if got := snapshot.Candidates[protocol.OpenAICompletions]["public-model"][0].UpstreamModelID; got != "real-model" { - t.Fatalf("candidate upstream model = %q, want real-model", got) + gotFilters := snapshot.AccessKeysByID[1].Filters + if _, ok := gotFilters.Groups[2]; ok { + t.Fatal("group filter retained caller mutation") + } + if _, ok := gotFilters.Protocols[protocol.Gemini]; ok { + t.Fatal("protocol filter retained caller mutation") + } + if _, ok := gotFilters.Models["changed"]; ok { + t.Fatal("model filter retained caller mutation") } } -func TestCompileRejectsDuplicateExternalModelWithinGroup(t *testing.T) { - tests := []struct { - name string - models []ModelConfig - }{ - {name: "two ids share alias", models: []ModelConfig{{ID: "a", Alias: "public"}, {ID: "b", Alias: "public"}}}, - {name: "alias collides with plain id", models: []ModelConfig{{ID: "a"}, {ID: "b", Alias: "a"}}}, - {name: "duplicate entry", models: []ModelConfig{{ID: "a"}, {ID: "a"}}}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - _, err := Compile(CompileInput{Groups: []GroupConfig{{ - ID: 1, Name: "group", UpstreamURL: "https://example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, Models: test.models, Enabled: true, - }}}) - if err == nil || !strings.Contains(err.Error(), "duplicate external model") { - t.Fatalf("Compile() error = %v", err) - } - }) - } -} +func TestCompileRejectsInvalidCoreConfiguration(t *testing.T) { + t.Parallel() -func TestCompileValidatesDisabledCatalogEntries(t *testing.T) { - invalidWeight := -1 tests := []struct { name string input CompileInput wantErr string }{ { - name: "disabled group zero id", - input: CompileInput{Groups: []GroupConfig{{ - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "model"}}, Enabled: false, + name: "duplicate external model", + input: CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{{ + ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), Models: []ModelConfig{{ID: "a"}, {ID: "b", Alias: "a"}}, Enabled: true, }}}, - wantErr: "group id is required", + wantErr: "duplicate external model", }, { - name: "disabled group invalid protocol", - input: CompileInput{Groups: []GroupConfig{{ - ID: 1, Protocols: []protocol.Protocol{"invalid"}, - Models: []ModelConfig{{ID: "model"}}, Enabled: false, - }}}, - wantErr: "invalid protocol", + name: "duplicate group id", + input: CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{ + {ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`)}, + {ID: 1, ChannelID: channel.Anthropic, Params: json.RawMessage(`{}`)}, + }}, + wantErr: "duplicate group id", }, { - name: "disabled group invalid weight", - input: CompileInput{Groups: []GroupConfig{{ - ID: 1, Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "model"}}, - WeightManual: &invalidWeight, Enabled: false, - }}}, - wantErr: "manual weight", - }, - { - name: "disabled access key zero id", - input: CompileInput{AccessKeys: []AccessKeyConfig{{ - KeyHash: "disabled-hash", Status: AccessKeyStatusDisabled, - }}}, - wantErr: "access key id is required", - }, - { - name: "disabled access key empty hash", - input: CompileInput{AccessKeys: []AccessKeyConfig{{ - ID: 1, Status: AccessKeyStatusDisabled, - }}}, - wantErr: "key hash is required", - }, - { - name: "disabled access key invalid filter", - input: CompileInput{AccessKeys: []AccessKeyConfig{{ - ID: 1, KeyHash: "disabled-hash", Status: AccessKeyStatusDisabled, - Filters: FilterSet{Protocols: map[protocol.Protocol]struct{}{"invalid": {}}}, - }}}, - wantErr: "invalid protocol", + name: "duplicate access key hash", + input: CompileInput{AccessKeys: []AccessKeyConfig{ + {ID: 1, KeyHash: "same", Status: AccessKeyStatusActive}, + {ID: 2, KeyHash: "same", Status: AccessKeyStatusActive}, + }}, + wantErr: "duplicate access key hash", }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { + t.Parallel() _, err := Compile(test.input) if err == nil || !strings.Contains(err.Error(), test.wantErr) { t.Fatalf("Compile() error = %v, want substring %q", err, test.wantErr) @@ -267,738 +211,3 @@ func TestCompileValidatesDisabledCatalogEntries(t *testing.T) { }) } } - -func TestCompileRejectsLegacyProtocolValuesWithObjectIdentity(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - input CompileInput - wantObject string - wantValue string - }{ - { - name: "group legacy openai", - input: CompileInput{Groups: []GroupConfig{{ - ID: 41, Protocols: []protocol.Protocol{"openai"}, Enabled: true, - }}}, - wantObject: "group 41", - wantValue: `"openai"`, - }, - { - name: "access key legacy openai response", - input: CompileInput{AccessKeys: []AccessKeyConfig{{ - ID: 52, KeyHash: "legacy", Status: AccessKeyStatusActive, - Filters: FilterSet{Protocols: map[protocol.Protocol]struct{}{ - "openai-response": {}, - }}, - }}}, - wantObject: "access key 52", - wantValue: `"openai-response"`, - }, - { - name: "group replaced openai completions", - input: CompileInput{Groups: []GroupConfig{{ - ID: 63, Protocols: []protocol.Protocol{"openai-chat-completions"}, Enabled: true, - }}}, - wantObject: "group 63", - wantValue: `"openai-chat-completions"`, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - t.Parallel() - _, err := Compile(test.input) - if err == nil { - t.Fatal("Compile() error = nil") - } - if !strings.Contains(err.Error(), test.wantObject) || - !strings.Contains(err.Error(), test.wantValue) { - t.Fatalf( - "Compile() error = %q, want object %q and value %q", - err, - test.wantObject, - test.wantValue, - ) - } - }) - } -} - -func TestCompileRejectsDuplicateCatalogIDs(t *testing.T) { - validGroup := func(id uint, enabled bool) GroupConfig { - return GroupConfig{ - ID: id, Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "model"}}, Enabled: enabled, - } - } - if _, err := Compile(CompileInput{Groups: []GroupConfig{ - validGroup(1, true), validGroup(1, false), - }}); err == nil || !strings.Contains(err.Error(), "duplicate group id") { - t.Fatalf("duplicate Group catalog error = %v", err) - } - if _, err := Compile(CompileInput{AccessKeys: []AccessKeyConfig{ - {ID: 1, KeyHash: "one", Status: AccessKeyStatusActive}, - {ID: 1, KeyHash: "two", Status: AccessKeyStatusDisabled}, - }}); err == nil || !strings.Contains(err.Error(), "duplicate access key id") { - t.Fatalf("duplicate AccessKey catalog error = %v", err) - } -} - -func TestCompileRejectsDuplicateCatalogHashes(t *testing.T) { - _, err := Compile(CompileInput{AccessKeys: []AccessKeyConfig{ - {ID: 1, KeyHash: "duplicate", Status: AccessKeyStatusActive}, - {ID: 2, KeyHash: "duplicate", Status: AccessKeyStatusDisabled}, - }}) - if err == nil || !strings.Contains(err.Error(), "duplicate access key hash") { - t.Fatalf("duplicate AccessKey catalog hash error = %v", err) - } -} - -func TestCompileManagementCatalogsOwnCopies(t *testing.T) { - groupWeight := 20 - filterGroups := map[uint]struct{}{1: {}} - input := CompileInput{ - Groups: []GroupConfig{{ - ID: 1, Name: "group", Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "provider", Alias: "public"}}, - WeightManual: &groupWeight, Enabled: false, - }}, - AccessKeys: []AccessKeyConfig{{ - ID: 2, Name: "client", KeyHash: "hash", Status: AccessKeyStatusDisabled, - Filters: FilterSet{Groups: filterGroups}, - }}, - } - snapshot, err := Compile(input) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - groupWeight = 99 - delete(filterGroups, 1) - input.Groups[0].Models[0] = ModelConfig{ID: "changed", Alias: "changed"} - if got := snapshot.GroupCatalog[1].WeightManual; got == nil || *got != 20 { - t.Fatalf("GroupCatalog weight = %v, want 20", got) - } - if got := snapshot.RouteCatalog[protocol.OpenAICompletions]["public"]; len(got) != 1 || - got[0].UpstreamModelID != "provider" { - t.Fatalf("RouteCatalog after caller mutation = %#v", got) - } - if _, ok := snapshot.AccessKeysByID[2].Filters.Groups[1]; !ok { - t.Fatal("AccessKeysByID filters changed with caller map") - } -} - -func TestCompileMergesTimeoutsAndHeaderRules(t *testing.T) { - systemSettings := config.Settings{ - "connect_timeout": 10.0, - "first_byte_timeout": 90.0, - "header_rules": map[string]any{ - "set": map[string]any{"X-System": "system"}, - "remove": []any{"X-System-Remove"}, - }, - } - groupSettings := config.Settings{ - "connect_timeout": 20.0, - "request_timeout": 500.0, - "stream_idle_timeout": 250.0, - "header_rules": map[string]any{ - "set": map[string]any{"X-Group": "${API_KEY}"}, - "remove": []any{"X-Group-Remove"}, - }, - } - - snapshot, err := Compile(runtimeSettingsInput(systemSettings, groupSettings)) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - - wantSettings := RuntimeSettings{ - ConnectTimeout: 10 * time.Second, - FirstByteTimeout: 90 * time.Second, - RequestTimeout: 600 * time.Second, - StreamIdleTimeout: 300 * time.Second, - HeaderRules: HeaderRules{ - Set: map[string]string{"X-System": "system"}, - Remove: []string{"X-System-Remove"}, - }, - InjectUsageOptions: true, - RequestLogRetentionDays: 7, - ModelsDevAutoSyncEnabled: true, - } - if !reflect.DeepEqual(snapshot.Settings, wantSettings) { - t.Errorf("ConfigSnapshot.Settings = %#v, want %#v", snapshot.Settings, wantSettings) - } - - group := snapshot.Groups[1] - wantTimeouts := TimeoutConfig{ - Connect: 20 * time.Second, - FirstByte: 90 * time.Second, - Request: 500 * time.Second, - StreamIdle: 250 * time.Second, - } - if group.Timeouts != wantTimeouts { - t.Errorf("GroupView.Timeouts = %#v, want %#v", group.Timeouts, wantTimeouts) - } - wantRules := HeaderRules{ - Set: map[string]string{"X-Group": "${API_KEY}"}, - Remove: []string{"X-Group-Remove"}, - } - if !reflect.DeepEqual(group.HeaderRules, wantRules) { - t.Errorf("GroupView.HeaderRules = %#v, want %#v", group.HeaderRules, wantRules) - } -} - -func TestCompileUsesDefaultRuntimeSettings(t *testing.T) { - snapshot, err := Compile(runtimeSettingsInput(nil, nil)) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - - wantSettings := RuntimeSettings{ - ConnectTimeout: 15 * time.Second, - FirstByteTimeout: 120 * time.Second, - RequestTimeout: 600 * time.Second, - StreamIdleTimeout: 300 * time.Second, - HeaderRules: HeaderRules{Set: map[string]string{}}, - InjectUsageOptions: true, - RequestLogRetentionDays: 7, - ModelsDevAutoSyncEnabled: true, - } - if !reflect.DeepEqual(snapshot.Settings, wantSettings) { - t.Errorf("ConfigSnapshot.Settings = %#v, want %#v", snapshot.Settings, wantSettings) - } - - group := snapshot.Groups[1] - wantTimeouts := TimeoutConfig{ - Connect: 15 * time.Second, - FirstByte: 120 * time.Second, - Request: 600 * time.Second, - StreamIdle: 300 * time.Second, - } - if group.Timeouts != wantTimeouts { - t.Errorf("GroupView.Timeouts = %#v, want %#v", group.Timeouts, wantTimeouts) - } - wantRules := HeaderRules{Set: map[string]string{}} - if !reflect.DeepEqual(group.HeaderRules, wantRules) { - t.Errorf("GroupView.HeaderRules = %#v, want %#v", group.HeaderRules, wantRules) - } -} - -func TestCompileFreezesInjectUsageOptionsInGroupView(t *testing.T) { - groupSettings := config.Settings{SettingInjectUsageOptions: false} - snapshot, err := Compile(runtimeSettingsInput(nil, groupSettings)) - if err != nil { - t.Fatal(err) - } - if snapshot.Groups[1].InjectUsageOptions { - t.Fatalf("GroupView.InjectUsageOptions = true, want false") - } - groupSettings[SettingInjectUsageOptions] = true - if snapshot.Groups[1].InjectUsageOptions { - t.Fatal("published GroupView changed with source config") - } -} - -func TestCompileDeepClonesGroupProviderID(t *testing.T) { - providerID := "openai" - input := CompileInput{Groups: []GroupConfig{{ - ID: 1, - Name: "provider-group", - ProviderID: &providerID, - UpstreamURL: "https://api.openai.com/v1", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "gpt-4o"}}, - Enabled: true, - }}} - - snapshot, err := Compile(input) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - providerID = "mutated-source" - *input.Groups[0].ProviderID = "mutated-input" - - view := snapshot.Groups[1] - if view.ProviderID == nil || *view.ProviderID != "openai" { - t.Fatalf("GroupView.ProviderID = %v, want independent openai", view.ProviderID) - } -} - -func TestCompileRejectsInvalidProviderIDForZeroModelAndDisabledGroups(t *testing.T) { - invalidProviderID := "OpenAI" - for _, test := range []struct { - name string - enabled bool - }{ - {name: "enabled zero-model Group", enabled: true}, - {name: "disabled zero-model Group", enabled: false}, - } { - t.Run(test.name, func(t *testing.T) { - _, err := Compile(CompileInput{Groups: []GroupConfig{{ - ID: 1, - Name: "corrupt-provider", - ProviderID: &invalidProviderID, - UpstreamURL: "https://corrupt-provider.example/v1", - Protocols: []protocol.Protocol{protocol.OpenAIResponses}, - Models: []ModelConfig{}, - Enabled: test.enabled, - }}}) - if err == nil { - t.Fatal("Compile() error = nil, want invalid provider_id rejection") - } - }) - } -} - -func TestCompileCanonicalizesHeaderRuleNames(t *testing.T) { - snapshot, err := Compile(runtimeSettingsInput(nil, config.Settings{ - "header_rules": map[string]any{ - "set": map[string]any{"x-custom-key": "value"}, - "remove": []any{"x-remove-me"}, - }, - })) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - - want := HeaderRules{ - Set: map[string]string{"X-Custom-Key": "value"}, - Remove: []string{"X-Remove-Me"}, - } - if got := snapshot.Groups[1].HeaderRules; !reflect.DeepEqual(got, want) { - t.Fatalf("GroupView.HeaderRules = %#v, want %#v", got, want) - } -} - -func TestCompileAcceptsJSONNumberWholeSecondTimeouts(t *testing.T) { - tests := []struct { - literal string - want time.Duration - }{ - {literal: "20.0", want: 20 * time.Second}, - {literal: "2e1", want: 20 * time.Second}, - {literal: "9223372036.0", want: 9223372036 * time.Second}, - } - for _, test := range tests { - t.Run(test.literal, func(t *testing.T) { - snapshot, err := Compile(runtimeSettingsInput(nil, config.Settings{ - "connect_timeout": json.Number(test.literal), - })) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - if got := snapshot.Groups[1].Timeouts.Connect; got != test.want { - t.Fatalf("Connect timeout = %s, want %s", got, test.want) - } - }) - } -} - -func TestCompileRejectsMalformedRuntimeSettings(t *testing.T) { - tests := []struct { - name string - groupSettings config.Settings - wantErr string - }{ - { - name: "zero timeout", - groupSettings: config.Settings{"connect_timeout": 0}, - wantErr: "connect_timeout must be a positive whole number", - }, - { - name: "negative timeout", - groupSettings: config.Settings{"first_byte_timeout": int64(-1)}, - wantErr: "first_byte_timeout must be a positive whole number", - }, - { - name: "fractional timeout", - groupSettings: config.Settings{"request_timeout": 1.5}, - wantErr: "request_timeout must be a positive whole number", - }, - { - name: "precise fractional JSON timeout", - groupSettings: config.Settings{"request_timeout": json.Number("20.0000000000000001")}, - wantErr: "request_timeout must be a positive whole number", - }, - { - name: "precise fractional JSON timeout at limit", - groupSettings: config.Settings{"request_timeout": json.Number("9223372036.0000001")}, - wantErr: "request_timeout must be a positive whole number", - }, - { - name: "overflow timeout", - groupSettings: config.Settings{"stream_idle_timeout": json.Number("9223372037")}, - wantErr: "stream_idle_timeout must be a positive whole number", - }, - { - name: "non-object header rules", - groupSettings: config.Settings{"header_rules": []any{}}, - wantErr: "header_rules must be an object", - }, - { - name: "non-object header set", - groupSettings: config.Settings{"header_rules": map[string]any{ - "set": []any{}, - }}, - wantErr: "header_rules.set must be an object", - }, - { - name: "non-string header set value", - groupSettings: config.Settings{"header_rules": map[string]any{ - "set": map[string]any{"X-Invalid": 1.0}, - }}, - wantErr: "header_rules.set.X-Invalid must be a string", - }, - { - name: "empty header set name", - groupSettings: config.Settings{"header_rules": map[string]any{ - "set": map[string]any{"": "value"}, - }}, - wantErr: "header_rules.set contains invalid header name \"\"", - }, - { - name: "header set name with space", - groupSettings: config.Settings{"header_rules": map[string]any{ - "set": map[string]any{"Bad Header": "value"}, - }}, - wantErr: "header_rules.set contains invalid header name \"Bad Header\"", - }, - { - name: "header set value with newline", - groupSettings: config.Settings{"header_rules": map[string]any{ - "set": map[string]any{"X-Credential": "prefix\r\nInjected: value"}, - }}, - wantErr: "header_rules.set.X-Credential contains invalid header value", - }, - { - name: "case-insensitive duplicate header set", - groupSettings: config.Settings{"header_rules": map[string]any{ - "set": map[string]any{ - "Authorization": "Bearer ${API_KEY}", - "authorization": "Token ${API_KEY}", - }, - }}, - wantErr: "header_rules.set contains duplicate header \"Authorization\"", - }, - { - name: "non-array header remove", - groupSettings: config.Settings{"header_rules": map[string]any{ - "remove": "X-Invalid", - }}, - wantErr: "header_rules.remove must be an array", - }, - { - name: "non-string header remove entry", - groupSettings: config.Settings{"header_rules": map[string]any{ - "remove": []any{"X-Valid", 1.0}, - }}, - wantErr: "header_rules.remove[1] must be a string", - }, - { - name: "invalid header remove name", - groupSettings: config.Settings{"header_rules": map[string]any{ - "remove": []any{"Bad Header"}, - }}, - wantErr: "header_rules.remove[0] contains invalid header name \"Bad Header\"", - }, - { - name: "unknown header rules field", - groupSettings: config.Settings{"header_rules": map[string]any{ - "append": map[string]any{"X-Invalid": "value"}, - }}, - wantErr: "unknown header_rules field \"append\"", - }, - { - name: "unknown group setting", - groupSettings: config.Settings{"retry_count": 3.0}, - wantErr: "unknown group setting \"retry_count\"", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - _, err := Compile(runtimeSettingsInput(nil, tt.groupSettings)) - if err == nil || !strings.Contains(err.Error(), tt.wantErr) { - t.Fatalf("Compile() error = %v, want error containing %q", err, tt.wantErr) - } - }) - } -} - -func TestCompileRuntimeSettingsOwnsIndependentCopies(t *testing.T) { - set := map[string]any{"X-Original": "original"} - remove := []any{"X-Original-Remove"} - headerRules := map[string]any{"set": set, "remove": remove} - groupSettings := config.Settings{"header_rules": headerRules} - - snapshot, err := Compile(runtimeSettingsInput(nil, groupSettings)) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - - set["X-Original"] = "mutated" - set["X-Added"] = "added" - remove[0] = "X-Mutated-Remove" - headerRules["set"] = map[string]any{"X-Replaced": "replaced"} - headerRules["remove"] = []any{"X-Replaced-Remove"} - groupSettings["header_rules"] = map[string]any{} - - wantRules := HeaderRules{ - Set: map[string]string{"X-Original": "original"}, - Remove: []string{"X-Original-Remove"}, - } - if got := snapshot.Groups[1].HeaderRules; !reflect.DeepEqual(got, wantRules) { - t.Errorf("GroupView.HeaderRules = %#v after input mutation, want %#v", got, wantRules) - } -} - -func runtimeSettingsInput(systemSettings, groupSettings config.Settings) CompileInput { - return CompileInput{ - SystemSettings: systemSettings, - Groups: []GroupConfig{{ - ID: 1, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "model"}}, - Settings: groupSettings, - Enabled: true, - }}, - } -} - -func TestCompileRejectsInvalidCandidateConfiguration(t *testing.T) { - tests := []struct { - name string - input CompileInput - wantErr string - }{ - { - name: "zero group id", - input: CompileInput{Groups: []GroupConfig{{ - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "model"}}, - Enabled: true, - }}}, - wantErr: "group id is required", - }, - { - name: "duplicate enabled group id", - input: CompileInput{Groups: []GroupConfig{ - { - ID: 1, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "model-one"}}, - Enabled: true, - }, - { - ID: 1, - Protocols: []protocol.Protocol{protocol.Anthropic}, - Models: []ModelConfig{{ID: "model-two"}}, - Enabled: true, - }, - }}, - wantErr: "duplicate group id", - }, - { - name: "empty protocol list", - input: CompileInput{Groups: []GroupConfig{{ - ID: 1, - Models: []ModelConfig{{ID: "model"}}, - Enabled: true, - }}}, - wantErr: "protocols are required", - }, - { - name: "unknown protocol", - input: CompileInput{Groups: []GroupConfig{{ - ID: 1, - Protocols: []protocol.Protocol{"unknown"}, - Models: []ModelConfig{{ID: "model"}}, - Enabled: true, - }}}, - wantErr: "invalid protocol", - }, - { - name: "duplicate protocol", - input: CompileInput{Groups: []GroupConfig{{ - ID: 1, - Protocols: []protocol.Protocol{protocol.OpenAICompletions, protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "model"}}, - Enabled: true, - }}}, - wantErr: "duplicate protocol", - }, - { - name: "whitespace model id", - input: CompileInput{Groups: []GroupConfig{{ - ID: 1, - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: " \t"}}, - Enabled: true, - }}}, - wantErr: "model id is required", - }, - { - name: "active access key without hash", - input: CompileInput{AccessKeys: []AccessKeyConfig{{ - ID: 10, - Status: AccessKeyStatusActive, - }}}, - wantErr: "key hash is required", - }, - { - name: "invalid access key status", - input: CompileInput{AccessKeys: []AccessKeyConfig{{ - ID: 10, - KeyHash: "hash", - Status: AccessKeyStatus("invalid"), - }}}, - wantErr: "invalid status", - }, - { - name: "duplicate active access key hash", - input: CompileInput{AccessKeys: []AccessKeyConfig{ - {ID: 10, KeyHash: "duplicate", Status: AccessKeyStatusActive}, - {ID: 11, KeyHash: "duplicate", Status: AccessKeyStatusActive}, - }}, - wantErr: "duplicate access key hash", - }, - { - name: "invalid filter protocol", - input: CompileInput{AccessKeys: []AccessKeyConfig{{ - ID: 10, - KeyHash: "hash", - Status: AccessKeyStatusActive, - Filters: FilterSet{Protocols: map[protocol.Protocol]struct{}{"invalid": {}}}, - }}}, - wantErr: "invalid protocol", - }, - { - name: "blank filter model", - input: CompileInput{AccessKeys: []AccessKeyConfig{{ - ID: 10, - KeyHash: "hash", - Status: AccessKeyStatusActive, - Filters: FilterSet{Models: map[string]struct{}{" ": {}}}, - }}}, - wantErr: "filter model is required", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - _, err := Compile(tt.input) - if err == nil || !strings.Contains(err.Error(), tt.wantErr) { - t.Fatalf("Compile() error = %v, want error containing %q", err, tt.wantErr) - } - }) - } -} - -func TestCompileAllowsDanglingFilterGroupIDs(t *testing.T) { - input := CompileInput{AccessKeys: []AccessKeyConfig{{ - ID: 10, - KeyHash: "hash", - Status: AccessKeyStatusActive, - Filters: FilterSet{Groups: map[uint]struct{}{999: {}}}, - }}} - - snapshot, err := Compile(input) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - accessKey, ok := snapshot.AccessKeysByHash["hash"] - if !ok { - t.Fatal("active access key missing") - } - if _, ok := accessKey.Filters.Groups[999]; !ok { - t.Fatal("dangling group filter id was not preserved") - } -} - -func TestCompileOwnsIndependentCopiesOfInput(t *testing.T) { - protocols := []protocol.Protocol{protocol.OpenAICompletions} - models := []ModelConfig{{ID: "model-real", Alias: "model-alias"}} - filterGroups := map[uint]struct{}{1: {}} - filterProtocols := map[protocol.Protocol]struct{}{protocol.OpenAICompletions: {}} - filterModels := map[string]struct{}{"model-real": {}} - input := CompileInput{ - Groups: []GroupConfig{{ - ID: 1, - Protocols: protocols, - Models: models, - Enabled: true, - }}, - AccessKeys: []AccessKeyConfig{{ - ID: 10, - KeyHash: "hash", - Status: AccessKeyStatusActive, - Filters: FilterSet{ - Groups: filterGroups, - Protocols: filterProtocols, - Models: filterModels, - }, - }}, - } - - snapshot, err := Compile(input) - if err != nil { - t.Fatalf("Compile() error = %v", err) - } - - protocols[0] = protocol.Anthropic - models[0] = ModelConfig{ID: "changed", Alias: "changed"} - delete(filterGroups, 1) - filterGroups[2] = struct{}{} - delete(filterProtocols, protocol.OpenAICompletions) - filterProtocols[protocol.Gemini] = struct{}{} - delete(filterModels, "model-real") - filterModels["changed"] = struct{}{} - - group := snapshot.Groups[1] - if got := group.Protocols[0]; got != protocol.OpenAICompletions { - t.Errorf("GroupView.Protocols[0] = %q, want %q", got, protocol.OpenAICompletions) - } - if got := group.Models[0]; got != (ModelConfig{ID: "model-real", Alias: "model-alias"}) { - t.Errorf("GroupView.Models[0] = %#v, want original model", got) - } - filters := snapshot.AccessKeysByHash["hash"].Filters - if _, ok := filters.Groups[1]; !ok { - t.Error("AccessKeyView.Filters.Groups lost original group") - } - if _, ok := filters.Groups[2]; ok { - t.Error("AccessKeyView.Filters.Groups contains caller mutation") - } - if _, ok := filters.Protocols[protocol.OpenAICompletions]; !ok { - t.Error("AccessKeyView.Filters.Protocols lost original protocol") - } - if _, ok := filters.Protocols[protocol.Gemini]; ok { - t.Error("AccessKeyView.Filters.Protocols contains caller mutation") - } - if _, ok := filters.Models["model-real"]; !ok { - t.Error("AccessKeyView.Filters.Models lost original model") - } - if _, ok := filters.Models["changed"]; ok { - t.Error("AccessKeyView.Filters.Models contains caller mutation") - } -} - -func TestCompileCarriesRPMLimitAndValidatesDisabledKeys(t *testing.T) { - input := CompileInput{AccessKeys: []AccessKeyConfig{{ - ID: 1, Name: "client", KeyHash: "hash", Status: AccessKeyStatusActive, - RPMLimit: 12, - }}} - snapshot, err := Compile(input) - if err != nil { - t.Fatal(err) - } - if got := snapshot.AccessKeysByHash["hash"].RPMLimit; got != 12 { - t.Fatalf("RPMLimit = %d, want 12", got) - } - - input.AccessKeys[0].Status = AccessKeyStatusDisabled - input.AccessKeys[0].RPMLimit = -1 - if _, err := Compile(input); err == nil { - t.Fatal("Compile() accepted negative rpm_limit on disabled key") - } -} diff --git a/internal/state/weights_test.go b/internal/state/weights_test.go index fbdad028..11ff5c33 100644 --- a/internal/state/weights_test.go +++ b/internal/state/weights_test.go @@ -1,20 +1,20 @@ package state import ( + "encoding/json" "reflect" "sync" "testing" "time" - "gpt-load/internal/protocol" + "gpt-load/internal/channel" ) func TestCompileCopiesAndValidatesGroupManualWeight(t *testing.T) { weight := 25 - input := CompileInput{Groups: []GroupConfig{{ - ID: 1, Name: "weighted", UpstreamURL: "https://weighted.example.com", - Protocols: []protocol.Protocol{protocol.OpenAICompletions}, - Models: []ModelConfig{{ID: "gpt-weighted"}}, WeightManual: &weight, Enabled: true, + input := CompileInput{ChannelRegistry: channel.NewRegistry(), Groups: []GroupConfig{{ + ID: 1, Name: "weighted", ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []ModelConfig{{ID: "gpt-weighted"}}, WeightManual: &weight, Enabled: true, }}} snapshot, err := Compile(input) @@ -36,22 +36,22 @@ func TestCompileCopiesAndValidatesGroupManualWeight(t *testing.T) { func TestKeyRegistryCollectCandidatesExcludesRuntimeUnavailableKeys(t *testing.T) { now := time.Date(2026, time.July, 20, 12, 0, 0, 0, time.UTC) - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "active"}, - {ID: 2, GroupID: 10, Status: KeyStatusActive, CooldownUntil: now.Add(time.Second), EncryptedValue: "cooling"}, - {ID: 3, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, EncryptedValue: "blacklisted"}, - {ID: 4, GroupID: 10, Status: KeyStatusActive, CooldownUntil: now, EncryptedValue: "expired"}, - {ID: 5, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "disabled"}, - {ID: 6, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "excluded"}, + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "active"}, + {ID: 2, GroupID: 10, Status: CredentialStatusActive, CooldownUntil: now.Add(time.Second), Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cooling"}, + {ID: 3, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "blacklisted"}, + {ID: 4, GroupID: 10, Status: CredentialStatusActive, CooldownUntil: now, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "expired"}, + {ID: 5, GroupID: 10, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "disabled"}, + {ID: 6, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "excluded"}, }) - got := registry.CollectCandidates([]uint{10}, func(keyID uint) bool { + got := registry.CollectCredentialCandidates([]uint{10}, func(keyID uint) bool { return keyID == 6 }, now) - want := []KeyMeta{ - {ID: 1, GroupID: 10, WeightAuto: DefaultWeight}, - {ID: 4, GroupID: 10, WeightAuto: DefaultWeight}, + want := []CredentialMeta{ + {ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, WeightAuto: DefaultWeight}, + {ID: 4, GroupID: 10, Version: 1, IdentityGeneration: 1, WeightAuto: DefaultWeight}, } if !reflect.DeepEqual(got, want) { t.Fatalf("CollectCandidates() = %#v, want %#v", got, want) @@ -62,43 +62,43 @@ func TestKeyRuntimeViewClassifiesAvailability(t *testing.T) { now := time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC) tests := []struct { name string - view KeyRuntimeView - want KeyRuntimeState + view CredentialRuntimeView + want CredentialRuntimeState }{ { name: "disabled wins", - view: KeyRuntimeView{ - Status: KeyStatusDisabled, Blacklisted: true, + view: CredentialRuntimeView{ + Status: CredentialStatusDisabled, Blacklisted: true, CooldownUntil: now.Add(time.Minute), }, - want: KeyRuntimeDisabled, + want: CredentialRuntimeDisabled, }, { name: "blacklist wins cooldown", - view: KeyRuntimeView{ - Status: KeyStatusActive, Blacklisted: true, + view: CredentialRuntimeView{ + Status: CredentialStatusActive, Blacklisted: true, CooldownUntil: now.Add(time.Minute), }, - want: KeyRuntimeBlacklisted, + want: CredentialRuntimeBlacklisted, }, { name: "future cooldown", - view: KeyRuntimeView{ - Status: KeyStatusActive, CooldownUntil: now.Add(time.Nanosecond), + view: CredentialRuntimeView{ + Status: CredentialStatusActive, CooldownUntil: now.Add(time.Nanosecond), }, - want: KeyRuntimeCooldown, + want: CredentialRuntimeCooldown, }, { name: "cooldown equality is available", - view: KeyRuntimeView{ - Status: KeyStatusActive, CooldownUntil: now, + view: CredentialRuntimeView{ + Status: CredentialStatusActive, CooldownUntil: now, }, - want: KeyRuntimeAvailable, + want: CredentialRuntimeAvailable, }, { name: "active", - view: KeyRuntimeView{Status: KeyStatusActive}, - want: KeyRuntimeAvailable, + view: CredentialRuntimeView{Status: CredentialStatusActive}, + want: CredentialRuntimeAvailable, }, } for _, test := range tests { @@ -112,25 +112,25 @@ func TestKeyRuntimeViewClassifiesAvailability(t *testing.T) { func TestCollectCandidatesUsesRuntimeViewBoundary(t *testing.T) { now := time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC) - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "active"}, + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "active"}, { - ID: 2, GroupID: 10, Status: KeyStatusActive, - CooldownUntil: now.Add(time.Nanosecond), EncryptedValue: "cooling", + ID: 2, GroupID: 10, Status: CredentialStatusActive, + CooldownUntil: now.Add(time.Nanosecond), Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cooling", }, - {ID: 3, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, EncryptedValue: "blacklisted"}, - {ID: 4, GroupID: 10, Status: KeyStatusActive, CooldownUntil: now, EncryptedValue: "boundary"}, - {ID: 5, GroupID: 10, Status: KeyStatusDisabled, EncryptedValue: "disabled"}, + {ID: 3, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "blacklisted"}, + {ID: 4, GroupID: 10, Status: CredentialStatusActive, CooldownUntil: now, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "boundary"}, + {ID: 5, GroupID: 10, Status: CredentialStatusDisabled, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "disabled"}, }) views := registry.Snapshot() availableIDs := make([]uint, 0) for _, view := range views { - if view.RuntimeState(now) == KeyRuntimeAvailable { + if view.RuntimeState(now) == CredentialRuntimeAvailable { availableIDs = append(availableIDs, view.ID) } } - candidates := registry.CollectCandidates([]uint{10}, nil, now) + candidates := registry.CollectCredentialCandidates([]uint{10}, nil, now) candidateIDs := make([]uint, 0, len(candidates)) for _, candidate := range candidates { candidateIDs = append(candidateIDs, candidate.ID) @@ -142,9 +142,9 @@ func TestCollectCandidatesUsesRuntimeViewBoundary(t *testing.T) { func TestKeyRegistrySetCooldownNeverShortensDeadline(t *testing.T) { now := time.Date(2026, time.July, 22, 12, 0, 0, 0, time.UTC) - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher", }}) if ok := registry.SetCooldown(1, now.Add(time.Hour)); !ok { @@ -153,19 +153,19 @@ func TestKeyRegistrySetCooldownNeverShortensDeadline(t *testing.T) { if ok := registry.SetCooldown(1, now.Add(time.Minute)); !ok { t.Fatal("SetCooldown(short deadline) = false, want true") } - if got := registry.CollectCandidates([]uint{10}, nil, now.Add(2*time.Minute)); len(got) != 0 { + if got := registry.CollectCredentialCandidates([]uint{10}, nil, now.Add(2*time.Minute)); len(got) != 0 { t.Fatalf("CollectCandidates() before longest cooldown expires = %#v, want none", got) } - if got := registry.CollectCandidates([]uint{10}, nil, now.Add(time.Hour)); len(got) != 1 { + if got := registry.CollectCredentialCandidates([]uint{10}, nil, now.Add(time.Hour)); len(got) != 1 { t.Fatalf("CollectCandidates() at longest cooldown boundary = %#v, want key 1", got) } } func TestKeyRegistrySetCooldownConcurrentWritersKeepLatestDeadline(t *testing.T) { now := time.Date(2026, time.July, 22, 12, 0, 0, 0, time.UTC) - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher", }}) const writers = 32 @@ -186,23 +186,23 @@ func TestKeyRegistrySetCooldownConcurrentWritersKeepLatestDeadline(t *testing.T) wait.Wait() latest := now.Add(writers * time.Minute) - if got := registry.CollectCandidates([]uint{10}, nil, latest.Add(-time.Nanosecond)); len(got) != 0 { + if got := registry.CollectCredentialCandidates([]uint{10}, nil, latest.Add(-time.Nanosecond)); len(got) != 0 { t.Fatalf("CollectCandidates() before latest cooldown expires = %#v, want none", got) } - if got := registry.CollectCandidates([]uint{10}, nil, latest); len(got) != 1 { + if got := registry.CollectCredentialCandidates([]uint{10}, nil, latest); len(got) != 1 { t.Fatalf("CollectCandidates() at latest cooldown boundary = %#v, want key 1", got) } } func TestKeyRegistryDefaultsAndSetsAutoWeight(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher", }}) assertAutoWeight := func(want int) { t.Helper() - got := registry.CollectCandidates([]uint{10}, nil, time.Time{}) + got := registry.CollectCredentialCandidates([]uint{10}, nil, time.Time{}) if len(got) != 1 || got[0].WeightAuto != want { t.Fatalf("CollectCandidates() = %#v, want WeightAuto %d", got, want) } @@ -225,10 +225,10 @@ func TestKeyRegistryDefaultsAndSetsAutoWeight(t *testing.T) { } func TestKeyRegistryClearFailureAndRecover(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, - FailureCount: 3, EncryptedValue: "cipher", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, + FailureCount: 3, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher", }}) if ok := registry.ClearFailure(1); !ok { @@ -256,31 +256,31 @@ func TestKeyRegistryClearFailureAndRecover(t *testing.T) { } func TestKeyRegistryRecoverIfMatchRestoresMatchingBlacklistedActiveKey(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, - FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-one", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, + FailureCount: 3, WeightAuto: 17, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }}) - if ok := registry.RecoverIfMatch(KeyRef{ID: 1, GroupID: 10, EncryptedValue: "cipher-one"}, DefaultWeight); !ok { + if ok := registry.RecoverIfMatch(CredentialRef{ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, DefaultWeight); !ok { t.Fatal("RecoverIfMatch() = false, want true") } - if got, want := registryEntry(t, registry, 1), (KeyEntry{ - ID: 1, GroupID: 10, Status: KeyStatusActive, WeightAuto: DefaultWeight, - FailureGeneration: 1, EncryptedValue: "cipher-one", + if got, want := registryEntry(t, registry, 1), (CredentialEntry{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, WeightAuto: DefaultWeight, + FailureGeneration: 1, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }); !reflect.DeepEqual(got, want) { t.Fatalf("entry after RecoverIfMatch() = %#v, want %#v", got, want) } } func TestKeyRegistryRecoverIfMatchRejectsStaleGeneration(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{{ - ID: 1, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, - FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-one", + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, + FailureCount: 3, WeightAuto: 17, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }}) - stale := registry.BlacklistedKeys()[0] + stale := registry.BlacklistedCredentials()[0] if stale.FailureGeneration != 0 { t.Fatalf("captured FailureGeneration = %d, want 0", stale.FailureGeneration) } @@ -295,16 +295,16 @@ func TestKeyRegistryRecoverIfMatchRejectsStaleGeneration(t *testing.T) { t.Fatalf("entry after stale RecoverIfMatch() = %#v, want unchanged %#v", got, before) } - fresh := registry.BlacklistedKeys()[0] + fresh := registry.BlacklistedCredentials()[0] if fresh.FailureGeneration != 1 { t.Fatalf("fresh FailureGeneration = %d, want 1", fresh.FailureGeneration) } if ok := registry.RecoverIfMatch(fresh, DefaultWeight); !ok { t.Fatal("RecoverIfMatch(fresh ref) = false, want true") } - if got, want := registryEntry(t, registry, 1), (KeyEntry{ - ID: 1, GroupID: 10, Status: KeyStatusActive, WeightAuto: DefaultWeight, - FailureGeneration: 2, EncryptedValue: "cipher-one", + if got, want := registryEntry(t, registry, 1), (CredentialEntry{ + ID: 1, GroupID: 10, Status: CredentialStatusActive, WeightAuto: DefaultWeight, + FailureGeneration: 2, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one", }); !reflect.DeepEqual(got, want) { t.Fatalf("entry after fresh RecoverIfMatch() = %#v, want %#v", got, want) } @@ -313,44 +313,44 @@ func TestKeyRegistryRecoverIfMatchRejectsStaleGeneration(t *testing.T) { func TestKeyRegistryRecoverIfMatchRejectsNonMatchingOrInvalidRecoveryWithoutMutation(t *testing.T) { tests := []struct { name string - entry KeyEntry - ref KeyRef + entry CredentialEntry + ref CredentialRef weight int }{ { - name: "disabled", entry: KeyEntry{ID: 1, GroupID: 10, Status: KeyStatusDisabled, Blacklisted: true, FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-one"}, - ref: KeyRef{ID: 1, GroupID: 10, EncryptedValue: "cipher-one"}, weight: DefaultWeight, + name: "disabled", entry: CredentialEntry{ID: 1, GroupID: 10, Status: CredentialStatusDisabled, Blacklisted: true, FailureCount: 3, WeightAuto: 17, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + ref: CredentialRef{ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, weight: DefaultWeight, }, { - name: "not blacklisted", entry: KeyEntry{ID: 1, GroupID: 10, Status: KeyStatusActive, FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-one"}, - ref: KeyRef{ID: 1, GroupID: 10, EncryptedValue: "cipher-one"}, weight: DefaultWeight, + name: "not blacklisted", entry: CredentialEntry{ID: 1, GroupID: 10, Status: CredentialStatusActive, FailureCount: 3, WeightAuto: 17, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + ref: CredentialRef{ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, weight: DefaultWeight, }, { - name: "group mismatch", entry: KeyEntry{ID: 1, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-one"}, - ref: KeyRef{ID: 1, GroupID: 11, EncryptedValue: "cipher-one"}, weight: DefaultWeight, + name: "group mismatch", entry: CredentialEntry{ID: 1, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + ref: CredentialRef{ID: 1, GroupID: 11, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, weight: DefaultWeight, }, { - name: "cipher mismatch", entry: KeyEntry{ID: 1, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-one"}, - ref: KeyRef{ID: 1, GroupID: 10, EncryptedValue: "cipher-replaced"}, weight: DefaultWeight, + name: "cipher mismatch", entry: CredentialEntry{ID: 1, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + ref: CredentialRef{ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-replaced"}, weight: DefaultWeight, }, { - name: "missing", entry: KeyEntry{ID: 2, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-two"}, - ref: KeyRef{ID: 1, GroupID: 10, EncryptedValue: "cipher-one"}, weight: DefaultWeight, + name: "missing", entry: CredentialEntry{ID: 2, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, + ref: CredentialRef{ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, weight: DefaultWeight, }, { - name: "weight too low", entry: KeyEntry{ID: 1, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-one"}, - ref: KeyRef{ID: 1, GroupID: 10, EncryptedValue: "cipher-one"}, weight: 0, + name: "weight too low", entry: CredentialEntry{ID: 1, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + ref: CredentialRef{ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, weight: 0, }, { - name: "weight too high", entry: KeyEntry{ID: 1, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, EncryptedValue: "cipher-one"}, - ref: KeyRef{ID: 1, GroupID: 10, EncryptedValue: "cipher-one"}, weight: MaxWeight + 1, + name: "weight too high", entry: CredentialEntry{ID: 1, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, FailureCount: 3, WeightAuto: 17, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + ref: CredentialRef{ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, weight: MaxWeight + 1, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{test.entry}) + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{test.entry}) before := registryEntry(t, registry, test.entry.ID) if ok := registry.RecoverIfMatch(test.ref, test.weight); ok { @@ -364,22 +364,22 @@ func TestKeyRegistryRecoverIfMatchRejectsNonMatchingOrInvalidRecoveryWithoutMuta } func TestKeyRegistryBlacklistedKeysReturnsActiveSortedRefs(t *testing.T) { - registry := NewKeyRegistry() - mustReplaceKeyEntries(t, registry, []KeyEntry{ - {ID: 3, GroupID: 20, Status: KeyStatusActive, Blacklisted: true, EncryptedValue: "cipher-three"}, - {ID: 2, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, EncryptedValue: "cipher-two"}, - {ID: 1, GroupID: 10, Status: KeyStatusActive, Blacklisted: true, EncryptedValue: "cipher-one"}, - {ID: 4, GroupID: 10, Status: KeyStatusDisabled, Blacklisted: true, EncryptedValue: "cipher-disabled"}, - {ID: 5, GroupID: 10, Status: KeyStatusActive, EncryptedValue: "cipher-healthy"}, + registry := NewCredentialRegistry() + mustReplaceKeyEntries(t, registry, []CredentialEntry{ + {ID: 3, GroupID: 20, Status: CredentialStatusActive, Blacklisted: true, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three"}, + {ID: 2, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, + {ID: 1, GroupID: 10, Status: CredentialStatusActive, Blacklisted: true, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + {ID: 4, GroupID: 10, Status: CredentialStatusDisabled, Blacklisted: true, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-disabled"}, + {ID: 5, GroupID: 10, Status: CredentialStatusActive, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-healthy"}, }) - want := []KeyRef{ - {ID: 1, GroupID: 10, EncryptedValue: "cipher-one"}, - {ID: 2, GroupID: 10, EncryptedValue: "cipher-two"}, - {ID: 3, GroupID: 20, EncryptedValue: "cipher-three"}, + want := []CredentialRef{ + {ID: 1, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-one"}, + {ID: 2, GroupID: 10, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-two"}, + {ID: 3, GroupID: 20, Version: 1, IdentityGeneration: 1, Fingerprint: "test-fingerprint", EncryptedValue: "cipher-three"}, } - if got := registry.BlacklistedKeys(); !reflect.DeepEqual(got, want) { - t.Fatalf("BlacklistedKeys() = %#v, want %#v", got, want) + if got := registry.BlacklistedCredentials(); !reflect.DeepEqual(got, want) { + t.Fatalf("BlacklistedCredentials() = %#v, want %#v", got, want) } } @@ -388,31 +388,31 @@ func TestValidateKeyEntriesRejectsInvalidWeights(t *testing.T) { manualTooHigh := MaxWeight + 1 tests := []struct { name string - entry KeyEntry + entry CredentialEntry }{ - {name: "manual below range", entry: KeyEntry{WeightManual: &manualTooLow}}, - {name: "manual above range", entry: KeyEntry{WeightManual: &manualTooHigh}}, - {name: "auto below range", entry: KeyEntry{WeightAuto: -1}}, - {name: "auto above range", entry: KeyEntry{WeightAuto: MaxWeight + 1}}, + {name: "manual below range", entry: CredentialEntry{WeightManual: &manualTooLow}}, + {name: "manual above range", entry: CredentialEntry{WeightManual: &manualTooHigh}}, + {name: "auto below range", entry: CredentialEntry{WeightAuto: -1}}, + {name: "auto above range", entry: CredentialEntry{WeightAuto: MaxWeight + 1}}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { test.entry.ID = 1 test.entry.GroupID = 10 - test.entry.Status = KeyStatusActive + test.entry.Status = CredentialStatusActive test.entry.EncryptedValue = "cipher" - if err := ValidateKeyEntries([]KeyEntry{test.entry}); err == nil { - t.Fatal("ValidateKeyEntries() error = nil, want error") + if err := ValidateCredentialEntries([]CredentialEntry{test.entry}); err == nil { + t.Fatal("ValidateCredentialEntries() error = nil, want error") } }) } } -func registryEntry(t *testing.T, registry *KeyRegistry, keyID uint) KeyEntry { +func registryEntry(t *testing.T, registry *CredentialRegistry, keyID uint) CredentialEntry { t.Helper() registry.mu.RLock() defer registry.mu.RUnlock() - groupID, ok := registry.keyGroups[keyID] + groupID, ok := registry.credentialGroups[keyID] if !ok { t.Fatalf("key %d missing", keyID) } diff --git a/internal/storage/database.go b/internal/storage/database.go index eba61e80..6ae13f6b 100644 --- a/internal/storage/database.go +++ b/internal/storage/database.go @@ -69,7 +69,13 @@ func newDatabaseDialector(database config.DatabaseConfig) (gorm.Dialector, error } return gormmysql.Open(dsn), nil case config.DatabaseDriverPostgreSQL: - return gormpostgres.Open(database.DSN), nil + // Schema migrations rename/rebuild tables while the process is running. + // pgx's implicit statement cache otherwise can retain a result shape from + // the legacy table and fail the first query against the rebuilt table. + return gormpostgres.New(gormpostgres.Config{ + DSN: database.DSN, + PreferSimpleProtocol: true, + }), nil default: return nil, fmt.Errorf("unsupported database driver") } diff --git a/internal/storage/database_integration_test.go b/internal/storage/database_integration_test.go index d1038781..dd4df4a9 100644 --- a/internal/storage/database_integration_test.go +++ b/internal/storage/database_integration_test.go @@ -86,7 +86,7 @@ func TestExternalDatabaseLifecycle(t *testing.T) { t.Fatalf("second AutoMigrate() error = %v", err) } for _, table := range []string{ - "groups", "upstream_keys", "access_keys", "request_logs", "request_log_attempts", + "groups", "credentials", "access_keys", "request_logs", "request_log_attempts", "usage_aggregation_journal", "usage_stats", "model_prices", "system_settings", "jobs", "control_operations", "schema_migrations", } { @@ -97,23 +97,19 @@ func TestExternalDatabaseLifecycle(t *testing.T) { suffix := strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) group := models.Group{ - Name: fmt.Sprintf("integration-%s-%d", suffix, time.Now().UnixNano()), - UpstreamURL: "https://integration.example.com", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[]`), - Config: models.JSON(`{}`), - Enabled: true, + Name: fmt.Sprintf("integration-%s-%d", suffix, time.Now().UnixNano()), + ChannelID: "openai_compatible", Params: models.JSON(`{"base_url":"https://integration.example.com"}`), + Models: models.JSON(`[]`), Overrides: models.JSON(`{}`), Enabled: true, } if err := db.Create(&group).Error; err != nil { t.Fatalf("create group: %v", err) } - key := models.UpstreamKey{ - GroupID: group.ID, - KeyValue: "ciphertext", - KeyHash: fmt.Sprintf("integration-%d", time.Now().UnixNano()), + credential := models.Credential{ + GroupID: group.ID, Data: "encrypted-data", + Fingerprint: fmt.Sprintf("integration-credential-%d", time.Now().UnixNano()), } - if err := db.Create(&key).Error; err != nil { - t.Fatalf("create upstream key: %v", err) + if err := db.Create(&credential).Error; err != nil { + t.Fatalf("create credential: %v", err) } setting := models.SystemSetting{ @@ -152,7 +148,9 @@ func TestExternalDatabaseLifecycle(t *testing.T) { stat := models.UsageStat{ BucketStartMS: time.Now().UnixMilli(), AccessKeyID: 1, + ChannelID: group.ChannelID, GroupID: group.ID, + CredentialID: credential.ID, Model: "integration-model", RequestCount: 1, SuccessCount: 1, @@ -160,7 +158,8 @@ func TestExternalDatabaseLifecycle(t *testing.T) { upsert := clause.OnConflict{ Columns: []clause.Column{ {Name: "bucket_start_ms"}, {Name: "access_key_id"}, - {Name: "group_id"}, {Name: "model"}, + {Name: "channel_id"}, {Name: "group_id"}, + {Name: "credential_id"}, {Name: "model"}, }, DoUpdates: clause.AssignmentColumns([]string{"request_count", "success_count"}), } @@ -198,8 +197,8 @@ func TestExternalDatabaseModelPriceIdentityUsesExactComparison(t *testing.T) { } base := fmt.Sprintf("external-model-%d", time.Now().UnixNano()) - upper := models.ModelPrice{ModelID: "Model-" + base} - lower := models.ModelPrice{ModelID: "model-" + base} + upper := models.ModelPrice{ChannelID: "openai", ModelID: "Model-" + base} + lower := models.ModelPrice{ChannelID: "openai", ModelID: "model-" + base} if err := db.Create(&upper).Error; err != nil { t.Fatalf("create first case-distinct model price: %v", err) } @@ -248,7 +247,7 @@ func TestExternalDatabaseReadSnapshotUsesRepeatableRead(t *testing.T) { if before != 0 { return fmt.Errorf("initial model price count = %d, want 0", before) } - if err := db.Create(&models.ModelPrice{ModelID: modelID}).Error; err != nil { + if err := db.Create(&models.ModelPrice{ChannelID: "openai", ModelID: modelID}).Error; err != nil { return err } var after int64 diff --git a/internal/storage/db_test.go b/internal/storage/db_test.go index 4f6e226f..d14a6b33 100644 --- a/internal/storage/db_test.go +++ b/internal/storage/db_test.go @@ -5,6 +5,7 @@ import ( "context" "os" "path/filepath" + "reflect" "strings" "testing" "time" @@ -17,70 +18,56 @@ import ( "gpt-load/internal/storage/models" ) -func TestAutoMigrateKeepsUpstreamRuntimeFailuresOutOfDatabase(t *testing.T) { - t.Parallel() - - db := openMigratedDatabase(t) - type columnInfo struct { - Name string - DefaultValue *string `gorm:"column:dflt_value"` - } - var columns []columnInfo - if err := db.Raw("PRAGMA table_info('upstream_keys')").Scan(&columns).Error; err != nil { - t.Fatalf("inspect upstream_keys columns: %v", err) - } - - var statusDefault string - for _, column := range columns { - if column.Name == "failure_count" { - t.Fatal("upstream_keys contains failure_count; runtime failure state must not be persisted") - } - if column.Name == "status" && column.DefaultValue != nil { - statusDefault = strings.Trim(*column.DefaultValue, "'\"") - } - } - if statusDefault != string(models.UpstreamKeyStatusActive) { - t.Fatalf("upstream_keys status default = %q, want %q", statusDefault, models.UpstreamKeyStatusActive) - } -} - -func TestUpstreamKeyStatusAcceptsOnlyDurableOperatorStates(t *testing.T) { +func TestCredentialStatusAcceptsOnlyDurableOperatorStates(t *testing.T) { t.Parallel() db := openMigratedDatabase(t) group := models.Group{ - Name: "status-parent", - UpstreamURL: "https://status.example.com", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[]`), + Name: "credential-status-parent", ChannelID: "openai_compatible", + Params: models.JSON(`{"base_url":"https://credential-status.example.com"}`), Models: models.JSON(`[]`), } if err := db.Create(&group).Error; err != nil { t.Fatalf("create parent group: %v", err) } - - for index, status := range []models.UpstreamKeyStatus{ - models.UpstreamKeyStatusActive, - models.UpstreamKeyStatusDisabled, + for index, status := range []models.CredentialStatus{ + models.CredentialStatusActive, + models.CredentialStatusDisabled, } { - key := models.UpstreamKey{ - GroupID: group.ID, - KeyValue: "ciphertext", - KeyHash: "allowed-status-" + string(rune('a'+index)), - Status: status, + credential := models.Credential{ + GroupID: group.ID, Data: "encrypted-data", + Fingerprint: "allowed-credential-status-" + string(rune('a'+index)), + Status: status, } - if err := db.Create(&key).Error; err != nil { - t.Fatalf("create upstream key with status %q: %v", status, err) + if err := db.Create(&credential).Error; err != nil { + t.Fatalf("create credential with status %q: %v", status, err) } } - - invalid := models.UpstreamKey{ - GroupID: group.ID, - KeyValue: "ciphertext", - KeyHash: "invalid-status", - Status: models.UpstreamKeyStatus("blacklisted"), + invalid := models.Credential{ + GroupID: group.ID, Data: "encrypted-data", Fingerprint: "invalid-credential-status", + Status: models.CredentialStatus("blacklisted"), } if err := db.Create(&invalid).Error; err == nil { - t.Fatal("create upstream key with runtime-only blacklisted status error = nil, want constraint error") + t.Fatal("create credential with runtime-only blacklisted status error = nil, want constraint error") + } +} + +func TestGroupNormalizesMissingChannelParamsToEmptyObject(t *testing.T) { + t.Parallel() + + db := openMigratedDatabase(t) + group := models.Group{ + Name: "normalized-channel-params", ChannelID: "staged-channel", + Models: models.JSON(`[]`), + } + if err := db.Create(&group).Error; err != nil { + t.Fatalf("create group without explicit channel params: %v", err) + } + var params string + if err := db.Table("groups").Select("params").Where("id = ?", group.ID).Scan(¶ms).Error; err != nil { + t.Fatalf("load normalized channel params: %v", err) + } + if params != `{}` { + t.Fatalf("normalized channel params = %q, want {}", params) } } @@ -151,7 +138,10 @@ func TestAutoMigrateCreatesReviewedIndexesAndPrimaryKeys(t *testing.T) { } } - for _, table := range []string{"upstream_keys", "access_keys"} { + if db.Migrator().HasTable("upstream_keys") { + t.Fatal("retired upstream_keys table exists") + } + for _, table := range []string{"credentials", "access_keys"} { var indexes []pragmaIndex if err := db.Raw("PRAGMA index_list('" + table + "')").Scan(&indexes).Error; err != nil { t.Fatalf("inspect %s indexes: %v", table, err) @@ -600,7 +590,7 @@ func TestAutoMigrateCreatesUsageJournalAndMigrationLedger(t *testing.T) { wantTables := []string{ "groups", - "upstream_keys", + "credentials", "access_keys", "request_logs", "usage_aggregation_journal", @@ -621,14 +611,16 @@ func TestAutoMigrateCreatesUsageJournalAndMigrationLedger(t *testing.T) { if err := db.Table("schema_migrations").Order("id ASC").Pluck("id", &migrationIDs).Error; err != nil { t.Fatalf("read schema_migrations: %v", err) } - if len(migrationIDs) != 5 || migrationIDs[0] != "0001_initial_v2" || - migrationIDs[1] != "0002_request_log_reasoning" || migrationIDs[2] != "0003_global_model_prices" || - migrationIDs[3] != "0004_mysql_model_price_identity" || - migrationIDs[4] != "0005_request_log_model_consistency" { - t.Fatalf( - "schema_migrations IDs = %v, want [0001_initial_v2 0002_request_log_reasoning 0003_global_model_prices 0004_mysql_model_price_identity 0005_request_log_model_consistency]", - migrationIDs, - ) + wantMigrationIDs := []string{ + "0001_initial_v2", + "0002_request_log_reasoning", + "0003_global_model_prices", + "0004_mysql_model_price_identity", + "0005_request_log_model_consistency", + "0006_channel_execution", + } + if !reflect.DeepEqual(migrationIDs, wantMigrationIDs) { + t.Fatalf("schema_migrations IDs = %v, want %v", migrationIDs, wantMigrationIDs) } if err := storage.AutoMigrate(db); err != nil { @@ -638,8 +630,8 @@ func TestAutoMigrateCreatesUsageJournalAndMigrationLedger(t *testing.T) { if err := db.Table("schema_migrations").Count(&count).Error; err != nil { t.Fatalf("count schema_migrations: %v", err) } - if count != 5 { - t.Fatalf("schema_migrations row count after a second migration = %d, want 5", count) + if count != int64(len(wantMigrationIDs)) { + t.Fatalf("schema_migrations row count after a second migration = %d, want %d", count, len(wantMigrationIDs)) } } @@ -784,17 +776,17 @@ func TestAutoMigrateOmitsGroupSignature(t *testing.T) { } } -func TestAutoMigrateAllowsDuplicateUpstreamURLs(t *testing.T) { +func TestAutoMigrateAllowsDuplicateChannelTargets(t *testing.T) { t.Parallel() db := openMigratedDatabase(t) first := models.Group{ - Name: "group-one", - UpstreamURL: "https://same.example.com/v1", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[]`), - Config: models.JSON(`{}`), - Enabled: true, + Name: "group-one", + ChannelID: "openai_compatible", + Params: models.JSON(`{"base_url":"https://same.example.com/v1"}`), + Models: models.JSON(`[]`), + Overrides: models.JSON(`{}`), + Enabled: true, } second := first second.Name = "group-two" @@ -802,7 +794,7 @@ func TestAutoMigrateAllowsDuplicateUpstreamURLs(t *testing.T) { t.Fatalf("create first group: %v", err) } if err := db.Create(&second).Error; err != nil { - t.Fatalf("create group with duplicate upstream URL: %v", err) + t.Fatalf("create group with duplicate channel target: %v", err) } } @@ -849,42 +841,54 @@ func TestAutoMigrateCreatesCriticalUniqueConstraints(t *testing.T) { t.Run("group name", func(t *testing.T) { first := models.Group{ - Name: "group-one", - UpstreamURL: "https://one.example.com", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[]`), - Config: models.JSON(`{}`), + Name: "group-one", + ChannelID: "openai_compatible", + Params: models.JSON(`{"base_url":"https://one.example.com"}`), + Models: models.JSON(`[]`), + Overrides: models.JSON(`{}`), } second := first second.ID = 0 - second.UpstreamURL = "https://two.example.com" + second.Params = models.JSON(`{"base_url":"https://two.example.com"}`) assertDuplicateRejected(t, db.Create(&first).Error, db.Create(&second).Error) }) - t.Run("upstream key group and hash", func(t *testing.T) { + t.Run("credential group and fingerprint", func(t *testing.T) { group := models.Group{ - Name: "upstream-key-parent", - UpstreamURL: "https://keys.example.com", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[]`), + Name: "credential-parent", + ChannelID: "openai_compatible", + Params: models.JSON(`{"base_url":"https://credentials.example.com"}`), + Models: models.JSON(`[]`), } if err := db.Create(&group).Error; err != nil { - t.Fatalf("create parent group: %v", err) + t.Fatalf("create credential parent group: %v", err) } - - first := models.UpstreamKey{ - GroupID: group.ID, - KeyValue: "ciphertext-one", - KeyHash: "same-key-hash", + first := models.Credential{ + GroupID: group.ID, Data: "encrypted-one", Fingerprint: "same-fingerprint", } second := first second.ID = 0 - second.KeyValue = "ciphertext-two" - + second.Data = "encrypted-two" assertDuplicateRejected(t, db.Create(&first).Error, db.Create(&second).Error) }) + t.Run("model price channel and model", func(t *testing.T) { + first := models.ModelPrice{ChannelID: "openai", ModelID: "same-model"} + duplicate := first + otherChannel := first + otherChannel.ChannelID = "anthropic" + if err := db.Create(&first).Error; err != nil { + t.Fatalf("create model price: %v", err) + } + if err := db.Create(&otherChannel).Error; err != nil { + t.Fatalf("create same model for another channel: %v", err) + } + if err := db.Create(&duplicate).Error; err == nil { + t.Fatal("create duplicate channel/model price error = nil, want unique constraint error") + } + }) + t.Run("access key hash", func(t *testing.T) { first := models.AccessKey{ Name: "access-one", @@ -916,11 +920,10 @@ func TestAutoMigrateCreatesCriticalUniqueConstraints(t *testing.T) { }) } -func TestAutoMigrateCreatesUpstreamKeyForeignKeyWithCascade(t *testing.T) { +func TestAutoMigrateCreatesCredentialForeignKeyWithCascade(t *testing.T) { t.Parallel() db := openMigratedDatabase(t) - type foreignKey struct { Table string From string @@ -928,47 +931,45 @@ func TestAutoMigrateCreatesUpstreamKeyForeignKeyWithCascade(t *testing.T) { OnDelete string `gorm:"column:on_delete"` } var foreignKeys []foreignKey - if err := db.Raw("PRAGMA foreign_key_list('upstream_keys')").Scan(&foreignKeys).Error; err != nil { - t.Fatalf("inspect upstream_keys foreign keys: %v", err) + if err := db.Raw("PRAGMA foreign_key_list('credentials')").Scan(&foreignKeys).Error; err != nil { + t.Fatalf("inspect credentials foreign keys: %v", err) } - var groupForeignKey *foreignKey - for i := range foreignKeys { - if foreignKeys[i].Table == "groups" && foreignKeys[i].From == "group_id" && foreignKeys[i].To == "id" { - groupForeignKey = &foreignKeys[i] + for index := range foreignKeys { + candidate := &foreignKeys[index] + if candidate.Table == "groups" && candidate.From == "group_id" && candidate.To == "id" { + groupForeignKey = candidate break } } - if groupForeignKey == nil { - t.Fatalf("upstream_keys foreign keys = %+v, want group_id -> groups.id", foreignKeys) - } - if groupForeignKey.OnDelete != "CASCADE" { - t.Fatalf("upstream_keys group foreign key ON DELETE = %q, want CASCADE", groupForeignKey.OnDelete) + if groupForeignKey == nil || groupForeignKey.OnDelete != "CASCADE" { + t.Fatalf("credentials foreign keys = %+v, want cascading group_id -> groups.id", foreignKeys) } group := models.Group{ - Name: "cascade-parent", - UpstreamURL: "https://cascade.example.com", - Protocols: models.JSON(`["openai-completions"]`), - Models: models.JSON(`[]`), + Name: "credential-cascade-parent", + ChannelID: "openai_compatible", + Params: models.JSON(`{"base_url":"https://credential-cascade.example.com"}`), + Models: models.JSON(`[]`), } if err := db.Create(&group).Error; err != nil { - t.Fatalf("create parent group: %v", err) + t.Fatalf("create credential parent group: %v", err) } - key := models.UpstreamKey{GroupID: group.ID, KeyValue: "ciphertext", KeyHash: "cascade-key-hash"} - if err := db.Create(&key).Error; err != nil { - t.Fatalf("create upstream key: %v", err) + credential := models.Credential{ + GroupID: group.ID, Data: "encrypted-data", Fingerprint: "cascade-fingerprint", + } + if err := db.Create(&credential).Error; err != nil { + t.Fatalf("create credential: %v", err) } if err := db.Delete(&group).Error; err != nil { - t.Fatalf("delete parent group: %v", err) + t.Fatalf("delete credential parent group: %v", err) } - - var keyCount int64 - if err := db.Model(&models.UpstreamKey{}).Where("id = ?", key.ID).Count(&keyCount).Error; err != nil { - t.Fatalf("count child key: %v", err) + var count int64 + if err := db.Model(&models.Credential{}).Where("id = ?", credential.ID).Count(&count).Error; err != nil { + t.Fatalf("count credential after deleting group: %v", err) } - if keyCount != 0 { - t.Fatalf("upstream key count after deleting group = %d, want 0", keyCount) + if count != 0 { + t.Fatalf("credential count after deleting group = %d, want 0", count) } } diff --git a/internal/storage/migration_0001_initial_v2.go b/internal/storage/migration_0001_initial_v2.go index 6839b1e7..aea942b1 100644 --- a/internal/storage/migration_0001_initial_v2.go +++ b/internal/storage/migration_0001_initial_v2.go @@ -14,7 +14,7 @@ import ( func createInitialV2Tables(db *gorm.DB) error { if err := db.AutoMigrate( &models.Group{}, - &models.UpstreamKey{}, + &models.Credential{}, &models.AccessKey{}, &models.RequestLog{}, &models.RequestLogAttempt{}, diff --git a/internal/storage/migration_0001_initial_v2_test.go b/internal/storage/migration_0001_initial_v2_test.go index c4690bed..e8395dd1 100644 --- a/internal/storage/migration_0001_initial_v2_test.go +++ b/internal/storage/migration_0001_initial_v2_test.go @@ -29,14 +29,14 @@ func TestAutoMigrateCreatesFinalPricingSchema(t *testing.T) { t.Fatal("schema_migrations is missing") } assertColumns(t, db, "model_prices", []string{ - "model_id", + "channel_id", "model_id", "input_price_nano_usd_per_million_tokens", "output_price_nano_usd_per_million_tokens", "cache_read_price_nano_usd_per_million_tokens", "cache_write_price_nano_usd_per_million_tokens", "context_price_tiers", "is_manual", "created_at_ms", "updated_at_ms", }) - assertUniqueIndex(t, db, "model_prices", "idx_model_prices_model", []string{"model_id"}) + assertUniqueIndex(t, db, "model_prices", "idx_model_prices_channel_model", []string{"channel_id", "model_id"}) columns := initialV2Columns(t, db, "model_prices") for _, name := range []string{ @@ -83,9 +83,11 @@ func TestAutoMigrateCreatesNormalizedRequestLogInitialV2(t *testing.T) { } assertColumns(t, db, "request_log_attempts", []string{ "request_id", "sequence", "completed_at_ms", "group_id", "group_name", - "key_id", "upstream_model", "status_code", "duration_ms", + "channel_id", "credential_id", "operation", "route_mode", "upstream_model", + "upstream_request_id", "dispatch_state", "response_started", + "status_code", "duration_ms", "failure_category", "action", "will_retry", "error_code", "error_summary", - "pricing_receipt", + "committed", "pricing_receipt", }) var foreignKeys []struct { @@ -129,18 +131,23 @@ func TestAutoMigrateCreatesFinalCatalogAndUsageColumns(t *testing.T) { } groups := initialV2Columns(t, db, "groups") - providerID, ok := groups["provider_id"] - if !ok { - t.Fatal("groups.provider_id is missing") + for _, name := range []string{"channel_id", "params", "models", "overrides"} { + if _, ok := groups[name]; !ok { + t.Errorf("groups.%s is missing", name) + } } - if providerID.NotNull != 0 || !strings.EqualFold(providerID.Type, "varchar(255)") { - t.Fatalf("groups.provider_id = type:%q notnull:%d, want varchar(255) nullable", - providerID.Type, providerID.NotNull) + for _, removed := range []string{"provider_id", "upstream_url", "protocols", "convert_enabled"} { + if _, ok := groups[removed]; ok { + t.Errorf("groups retains retired column %q", removed) + } + } + if db.Migrator().HasTable("upstream_keys") { + t.Fatal("retired upstream_keys table exists") } for table, columns := range map[string][]string{ - "request_logs": {"cache_write_unknown_tokens", "pricing_completeness"}, - "usage_stats": {"cache_write_unknown_tokens", "pricing_partial_count"}, + "request_logs": {"channel_id", "credential_id", "cache_write_unknown_tokens", "pricing_completeness"}, + "usage_stats": {"channel_id", "credential_id", "cache_write_unknown_tokens", "pricing_partial_count"}, } { got := initialV2Columns(t, db, table) for _, name := range columns { @@ -157,7 +164,8 @@ func TestAutoMigrateCreatesFinalCatalogAndUsageColumns(t *testing.T) { } journal := initialV2Columns(t, db, "usage_aggregation_journal") for _, name := range []string{ - "request_id", "bucket_start_ms", "access_key_id", "group_id", "model", + "request_id", "bucket_start_ms", "access_key_id", "channel_id", "group_id", + "credential_id", "model", "request_count", "success_count", "failure_count", "uncached_input_tokens", "output_tokens", "cache_read_tokens", "cache_write_5m_tokens", "cache_write_1h_tokens", @@ -197,8 +205,8 @@ func TestFinalModelPriceRejectsNegativeScalarPrices(t *testing.T) { for index, column := range priceColumns { t.Run(column, func(t *testing.T) { query := fmt.Sprintf(`INSERT INTO model_prices ( - model_id, %s, is_manual, created_at_ms, updated_at_ms - ) VALUES (?, -1, false, 1, 1)`, column) + channel_id, model_id, %s, is_manual, created_at_ms, updated_at_ms + ) VALUES ('openai', ?, -1, false, 1, 1)`, column) if err := db.Exec(query, fmt.Sprintf("model-%d", index)).Error; err == nil { t.Fatalf("negative %s was accepted", column) } @@ -210,6 +218,7 @@ func TestModelPriceValidatesAndNormalizesContextTiersBeforePersistence(t *testin db := openMigratedInitialV2TestDatabase(t) empty := models.ModelPrice{ + ChannelID: "openai", ModelID: "empty-tiers", ContextPriceTiers: models.JSON(`[]`), } @@ -227,7 +236,8 @@ func TestModelPriceValidatesAndNormalizesContextTiersBeforePersistence(t *testin } valid := models.ModelPrice{ - ModelID: "valid-tiers", + ChannelID: "openai", + ModelID: "valid-tiers", ContextPriceTiers: models.JSON(`[ {"threshold_tokens":0,"input_price_nano_usd_per_million_tokens":1}, {"threshold_tokens":100,"cache_write_price_nano_usd_per_million_tokens":0} @@ -257,6 +267,7 @@ func TestModelPriceValidatesAndNormalizesContextTiersBeforePersistence(t *testin for index, test := range invalid { t.Run(test.name, func(t *testing.T) { row := models.ModelPrice{ + ChannelID: "openai", ModelID: fmt.Sprintf("invalid-tier-%d", index), ContextPriceTiers: models.JSON(test.raw), } @@ -304,7 +315,7 @@ func TestModelPriceValidatesContextTiersAcrossGORMWritePaths(t *testing.T) { incoming := row incoming.ID = 0 return db.Clauses(clause.OnConflict{ - Columns: []clause.Column{{Name: "model_id"}}, + Columns: []clause.Column{{Name: "channel_id"}, {Name: "model_id"}}, DoUpdates: clause.Assignments(map[string]any{ "context_price_tiers": tiers, }), @@ -368,7 +379,7 @@ func TestFinalSchemaUsesMillisecondIntegersAndEnforcesCounters(t *testing.T) { for table, names := range map[string][]string{ "groups": {"created_at_ms", "updated_at_ms"}, - "upstream_keys": {"created_at_ms", "updated_at_ms"}, + "credentials": {"created_at_ms", "updated_at_ms"}, "access_keys": {"created_at_ms", "updated_at_ms"}, "model_prices": {"created_at_ms", "updated_at_ms"}, "request_logs": {"completed_at_ms"}, @@ -514,7 +525,7 @@ func TestAutoMigrateDoesNotReapplyCompletedMigrationForChangedTable(t *testing.T func TestAutoMigrateDoesNotReapplyCompletedMigrationForMissingIndex(t *testing.T) { db := openMigratedInitialV2TestDatabase(t) - if err := db.Exec("DROP INDEX idx_model_prices_model").Error; err != nil { + if err := db.Exec("DROP INDEX idx_model_prices_channel_model").Error; err != nil { t.Fatal(err) } if err := storage.AutoMigrate(db); err != nil { @@ -549,7 +560,8 @@ func createModelPriceForWritePath(t *testing.T, modelID string) (*gorm.DB, model t.Helper() db := openMigratedInitialV2TestDatabase(t) row := models.ModelPrice{ - ModelID: modelID, + ChannelID: "openai", + ModelID: modelID, ContextPriceTiers: models.JSON( `[{"threshold_tokens":0,"input_price_nano_usd_per_million_tokens":1}]`, ), diff --git a/internal/storage/migration_0002_request_log_reasoning.go b/internal/storage/migration_0002_request_log_reasoning.go deleted file mode 100644 index e799fb6b..00000000 --- a/internal/storage/migration_0002_request_log_reasoning.go +++ /dev/null @@ -1,26 +0,0 @@ -package storage - -import ( - "fmt" - - "gorm.io/gorm" - - "gpt-load/internal/storage/models" -) - -// addRequestLogReasoningColumns upgrades databases that were created by the -// original SQLite-only 0001 migration. Fresh schemas already contain these -// fields because the current model is the single source of schema metadata; -// HasColumn makes this migration idempotent for both paths. -func addRequestLogReasoningColumns(db *gorm.DB) error { - fields := []string{"ReasoningMode", "ReasoningEffort", "ReasoningBudgetTokens"} - for _, field := range fields { - if db.Migrator().HasColumn(&models.RequestLog{}, field) { - continue - } - if err := db.Migrator().AddColumn(&models.RequestLog{}, field); err != nil { - return fmt.Errorf("add request log reasoning column %s: %w", field, err) - } - } - return nil -} diff --git a/internal/storage/migration_0003_global_model_prices.go b/internal/storage/migration_0003_global_model_prices.go deleted file mode 100644 index 42ec88de..00000000 --- a/internal/storage/migration_0003_global_model_prices.go +++ /dev/null @@ -1,166 +0,0 @@ -package storage - -import ( - "bytes" - "fmt" - "sort" - "strings" - - "gorm.io/gorm" - - "gpt-load/internal/storage/models" -) - -// migrateGlobalModelPrices collapses the former (scope, model) rows into one -// row per exact upstream model. A conflicting manual value cannot be selected -// safely, so the migration fails before changing either data or its ledger. -func migrateGlobalModelPrices(db *gorm.DB) error { - if !db.Migrator().HasColumn(&models.ModelPrice{}, "price_scope_key") { - return nil - } - driver := strings.ToLower(db.Dialector.Name()) - if driver != "sqlite" && driver != "mysql" && driver != "postgres" && driver != "postgresql" { - return fmt.Errorf("migrate global model prices: unsupported database driver %q", db.Dialector.Name()) - } - - var rows []models.ModelPrice - if err := db.Order("model_id ASC").Order("id ASC").Find(&rows).Error; err != nil { - return fmt.Errorf("load scoped model prices: %w", err) - } - - loserIDs := make([]uint, 0) - conflicts := make([]string, 0) - for start := 0; start < len(rows); { - end := start + 1 - for end < len(rows) && rows[end].ModelID == rows[start].ModelID { - end++ - } - group := append([]models.ModelPrice(nil), rows[start:end]...) - sort.SliceStable(group, func(left, right int) bool { - if group[left].IsManual != group[right].IsManual { - return group[left].IsManual - } - if group[left].UpdatedAtMS != group[right].UpdatedAtMS { - return group[left].UpdatedAtMS > group[right].UpdatedAtMS - } - return group[left].ID > group[right].ID - }) - - manuals := make([]models.ModelPrice, 0, len(group)) - for _, row := range group { - if row.IsManual { - manuals = append(manuals, row) - } - } - if len(manuals) > 1 { - for _, candidate := range manuals[1:] { - if !globalModelPriceValuesEqual(manuals[0], candidate) { - conflicts = append(conflicts, rows[start].ModelID) - break - } - } - } - for _, row := range group[1:] { - loserIDs = append(loserIDs, row.ID) - } - start = end - } - if len(conflicts) > 0 { - return fmt.Errorf("migrate global model prices: conflicting manual prices for %s", strings.Join(conflicts, ", ")) - } - if len(loserIDs) > 0 { - if err := db.Where("id IN ?", loserIDs).Delete(&models.ModelPrice{}).Error; err != nil { - return fmt.Errorf("remove superseded scoped model prices: %w", err) - } - } - return migrateGlobalModelPriceSchema(db) -} - -func migrateGlobalModelPriceSchema(db *gorm.DB) error { - if !strings.EqualFold(db.Dialector.Name(), "sqlite") { - statements, err := globalModelPriceSchemaStatements(db.Dialector.Name()) - if err != nil { - return err - } - for _, statement := range statements { - if err := db.Exec(statement).Error; err != nil { - return fmt.Errorf("update global model price schema: %w", err) - } - } - return nil - } - if err := db.Exec(`CREATE TABLE model_prices_global_migration ( - id integer PRIMARY KEY AUTOINCREMENT, - model_id varchar(255) NOT NULL, - input_price_nano_usd_per_million_tokens integer NULL CONSTRAINT chk_model_price_input_nano CHECK (input_price_nano_usd_per_million_tokens IS NULL OR input_price_nano_usd_per_million_tokens >= 0), - output_price_nano_usd_per_million_tokens integer NULL CONSTRAINT chk_model_price_output_nano CHECK (output_price_nano_usd_per_million_tokens IS NULL OR output_price_nano_usd_per_million_tokens >= 0), - cache_read_price_nano_usd_per_million_tokens integer NULL CONSTRAINT chk_model_price_cache_read_nano CHECK (cache_read_price_nano_usd_per_million_tokens IS NULL OR cache_read_price_nano_usd_per_million_tokens >= 0), - cache_write_price_nano_usd_per_million_tokens integer NULL CONSTRAINT chk_model_price_cache_write_nano CHECK (cache_write_price_nano_usd_per_million_tokens IS NULL OR cache_write_price_nano_usd_per_million_tokens >= 0), - context_price_tiers json NULL, - is_manual boolean NOT NULL DEFAULT false, - created_at_ms integer NOT NULL CONSTRAINT chk_model_price_created_at CHECK (created_at_ms >= 0), - updated_at_ms integer NOT NULL CONSTRAINT chk_model_price_updated_at CHECK (updated_at_ms >= 0) - )`).Error; err != nil { - return fmt.Errorf("create global model price table: %w", err) - } - if err := db.Exec(`INSERT INTO model_prices_global_migration ( - id, model_id, input_price_nano_usd_per_million_tokens, - output_price_nano_usd_per_million_tokens, cache_read_price_nano_usd_per_million_tokens, - cache_write_price_nano_usd_per_million_tokens, context_price_tiers, is_manual, - created_at_ms, updated_at_ms - ) SELECT id, model_id, input_price_nano_usd_per_million_tokens, - output_price_nano_usd_per_million_tokens, cache_read_price_nano_usd_per_million_tokens, - cache_write_price_nano_usd_per_million_tokens, context_price_tiers, is_manual, - created_at_ms, updated_at_ms FROM model_prices`).Error; err != nil { - return fmt.Errorf("copy global model prices: %w", err) - } - if err := db.Exec("DROP TABLE model_prices").Error; err != nil { - return fmt.Errorf("drop scoped model price table: %w", err) - } - if err := db.Exec("ALTER TABLE model_prices_global_migration RENAME TO model_prices").Error; err != nil { - return fmt.Errorf("rename global model price table: %w", err) - } - if err := db.Exec("CREATE UNIQUE INDEX idx_model_prices_model ON model_prices(model_id)").Error; err != nil { - return fmt.Errorf("create global model price index: %w", err) - } - return nil -} - -// globalModelPriceSchemaStatements contains only fixed, schema-owned -// identifiers. MySQL combines the schema conversion into one ALTER TABLE so a -// failed DDL statement cannot leave the old table shape half-applied. -func globalModelPriceSchemaStatements(driver string) ([]string, error) { - switch strings.ToLower(driver) { - case "postgres", "postgresql": - return []string{ - "DROP INDEX idx_model_prices_scope_model", - "ALTER TABLE model_prices DROP COLUMN price_scope_key", - "CREATE UNIQUE INDEX idx_model_prices_model ON model_prices(model_id)", - }, nil - case "mysql": - return []string{ - "ALTER TABLE model_prices DROP INDEX idx_model_prices_scope_model, DROP COLUMN price_scope_key, MODIFY COLUMN model_id varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, ADD UNIQUE INDEX idx_model_prices_model (model_id)", - }, nil - default: - return nil, fmt.Errorf("migrate global model prices: unsupported database driver %q", driver) - } -} - -func globalModelPriceValuesEqual(left, right models.ModelPrice) bool { - if !globalModelPricePointerEqual(left.InputPriceNanoUSDPerMillionTokens, right.InputPriceNanoUSDPerMillionTokens) || - !globalModelPricePointerEqual(left.OutputPriceNanoUSDPerMillionTokens, right.OutputPriceNanoUSDPerMillionTokens) || - !globalModelPricePointerEqual(left.CacheReadPriceNanoUSDPerMillionTokens, right.CacheReadPriceNanoUSDPerMillionTokens) || - !globalModelPricePointerEqual(left.CacheWritePriceNanoUSDPerMillionTokens, right.CacheWritePriceNanoUSDPerMillionTokens) { - return false - } - leftTiers, leftErr := models.NormalizeContextPriceTiers(left.ContextPriceTiers) - rightTiers, rightErr := models.NormalizeContextPriceTiers(right.ContextPriceTiers) - return leftErr == nil && rightErr == nil && bytes.Equal(leftTiers, rightTiers) -} - -func globalModelPricePointerEqual(left, right *int64) bool { - if left == nil || right == nil { - return left == nil && right == nil - } - return *left == *right -} diff --git a/internal/storage/migration_0003_global_model_prices_sql_test.go b/internal/storage/migration_0003_global_model_prices_sql_test.go deleted file mode 100644 index 7e10def8..00000000 --- a/internal/storage/migration_0003_global_model_prices_sql_test.go +++ /dev/null @@ -1,58 +0,0 @@ -package storage - -import ( - "reflect" - "testing" -) - -func TestGlobalModelPriceSchemaStatementsSupportExternalDrivers(t *testing.T) { - tests := []struct { - name string - driver string - want []string - }{ - { - name: "PostgreSQL", - driver: "postgres", - want: []string{ - "DROP INDEX idx_model_prices_scope_model", - "ALTER TABLE model_prices DROP COLUMN price_scope_key", - "CREATE UNIQUE INDEX idx_model_prices_model ON model_prices(model_id)", - }, - }, - { - name: "PostgreSQL alias", - driver: "postgresql", - want: []string{ - "DROP INDEX idx_model_prices_scope_model", - "ALTER TABLE model_prices DROP COLUMN price_scope_key", - "CREATE UNIQUE INDEX idx_model_prices_model ON model_prices(model_id)", - }, - }, - { - name: "MySQL", - driver: "mysql", - want: []string{ - "ALTER TABLE model_prices DROP INDEX idx_model_prices_scope_model, DROP COLUMN price_scope_key, MODIFY COLUMN model_id varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, ADD UNIQUE INDEX idx_model_prices_model (model_id)", - }, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - got, err := globalModelPriceSchemaStatements(test.driver) - if err != nil { - t.Fatalf("globalModelPriceSchemaStatements() error = %v", err) - } - if !reflect.DeepEqual(got, test.want) { - t.Fatalf("globalModelPriceSchemaStatements() = %#v, want %#v", got, test.want) - } - }) - } -} - -func TestGlobalModelPriceSchemaStatementsRejectUnsupportedDriver(t *testing.T) { - if _, err := globalModelPriceSchemaStatements("sqlserver"); err == nil { - t.Fatal("globalModelPriceSchemaStatements() accepted an unsupported driver") - } -} diff --git a/internal/storage/migration_0003_global_model_prices_test.go b/internal/storage/migration_0003_global_model_prices_test.go deleted file mode 100644 index 244c7d96..00000000 --- a/internal/storage/migration_0003_global_model_prices_test.go +++ /dev/null @@ -1,108 +0,0 @@ -package storage_test - -import ( - "testing" - - "gorm.io/gorm" - - "gpt-load/internal/storage" - "gpt-load/internal/storage/models" -) - -func TestAutoMigrateCollapsesScopedModelPricesIntoGlobalRows(t *testing.T) { - db := openInitialV2TestDatabase(t) - createScopedModelPriceSchema(t, db) - if err := db.Exec(`INSERT INTO model_prices ( - id, price_scope_key, model_id, input_price_nano_usd_per_million_tokens, - is_manual, created_at_ms, updated_at_ms - ) VALUES - (1, 'provider:openai', 'shared', 1, false, 1, 10), - (2, 'group:7', 'shared', 2, true, 2, 20), - (3, 'provider:anthropic', 'automatic', 3, false, 3, 30)`).Error; err != nil { - t.Fatalf("seed scoped model prices: %v", err) - } - - if err := storage.AutoMigrate(db); err != nil { - t.Fatalf("AutoMigrate() error = %v", err) - } - if db.Migrator().HasColumn(&models.ModelPrice{}, "price_scope_key") { - t.Fatal("model_prices retains price_scope_key after global migration") - } - assertUniqueIndex(t, db, "model_prices", "idx_model_prices_model", []string{"model_id"}) - - var rows []models.ModelPrice - if err := db.Order("model_id ASC").Find(&rows).Error; err != nil { - t.Fatalf("load migrated prices: %v", err) - } - if len(rows) != 2 || rows[0].ModelID != "automatic" || rows[1].ModelID != "shared" { - t.Fatalf("migrated rows = %#v, want one row per model", rows) - } - if rows[1].ID != 2 || !rows[1].IsManual || rows[1].InputPriceNanoUSDPerMillionTokens == nil || - *rows[1].InputPriceNanoUSDPerMillionTokens != 2 { - t.Fatalf("shared winner = %#v, want newest manual row", rows[1]) - } -} - -func TestAutoMigrateRejectsConflictingScopedManualPricesWithoutMutation(t *testing.T) { - db := openInitialV2TestDatabase(t) - createScopedModelPriceSchema(t, db) - if err := db.Exec(`INSERT INTO model_prices ( - id, price_scope_key, model_id, input_price_nano_usd_per_million_tokens, - is_manual, created_at_ms, updated_at_ms - ) VALUES - (1, 'provider:openai', 'shared', 1, true, 1, 10), - (2, 'group:7', 'shared', 2, true, 2, 20)`).Error; err != nil { - t.Fatalf("seed conflicting scoped prices: %v", err) - } - - if err := storage.AutoMigrate(db); err == nil { - t.Fatal("AutoMigrate() accepted conflicting manual model prices") - } - if !db.Migrator().HasColumn(&models.ModelPrice{}, "price_scope_key") { - t.Fatal("failed global migration mutated model_prices schema") - } - var count int64 - if err := db.Table("model_prices").Count(&count).Error; err != nil { - t.Fatalf("count scoped prices after failed migration: %v", err) - } - if count != 2 { - t.Fatalf("scoped prices after failed migration = %d, want 2", count) - } - if err := db.Table("schema_migrations").Count(&count).Error; err != nil { - t.Fatalf("count migration ledger after failure: %v", err) - } - if count != 2 { - t.Fatalf("migration ledger after failure = %d, want 2", count) - } -} - -func createScopedModelPriceSchema(t *testing.T, db *gorm.DB) { - t.Helper() - for _, statement := range []string{ - `CREATE TABLE schema_migrations (id varchar(255) PRIMARY KEY NOT NULL)`, - `INSERT INTO schema_migrations(id) VALUES ('0001_initial_v2'), ('0002_request_log_reasoning')`, - `CREATE TABLE request_logs ( - id varchar(36) PRIMARY KEY NOT NULL, - status varchar(32) NOT NULL DEFAULT '', - upstream_model varchar(255) NOT NULL DEFAULT '' - )`, - `CREATE TABLE model_prices ( - id integer PRIMARY KEY AUTOINCREMENT, - price_scope_key varchar(255) NOT NULL, - model_id varchar(255) NOT NULL, - input_price_nano_usd_per_million_tokens integer NULL, - output_price_nano_usd_per_million_tokens integer NULL, - cache_read_price_nano_usd_per_million_tokens integer NULL, - cache_write_price_nano_usd_per_million_tokens integer NULL, - context_price_tiers json NULL, - is_manual boolean NOT NULL DEFAULT false, - created_at_ms integer NOT NULL, - updated_at_ms integer NOT NULL - )`, - `CREATE UNIQUE INDEX idx_model_prices_scope_model ON model_prices(price_scope_key, model_id)`, - } { - if err := db.Exec(statement).Error; err != nil { - t.Fatalf("create scoped model price fixture: %v", err) - } - } -} diff --git a/internal/storage/migration_0004_mysql_model_price_identity.go b/internal/storage/migration_0004_mysql_model_price_identity.go deleted file mode 100644 index 8fcff17a..00000000 --- a/internal/storage/migration_0004_mysql_model_price_identity.go +++ /dev/null @@ -1,38 +0,0 @@ -package storage - -import ( - "fmt" - "strings" - - "gorm.io/gorm" -) - -// migrateMySQLModelPriceIdentity makes MySQL enforce the same exact model-ID -// identity that SQLite and PostgreSQL already provide. A connection collation -// does not change the unique index's column collation, so this is schema work -// rather than a DSN-only setting. -func migrateMySQLModelPriceIdentity(db *gorm.DB) error { - statements, err := mysqlModelPriceIdentityStatements(db.Dialector.Name()) - if err != nil { - return err - } - for _, statement := range statements { - if err := db.Exec(statement).Error; err != nil { - return fmt.Errorf("set MySQL model price identity collation: %w", err) - } - } - return nil -} - -func mysqlModelPriceIdentityStatements(driver string) ([]string, error) { - switch strings.ToLower(driver) { - case "mysql": - return []string{ - "ALTER TABLE model_prices MODIFY COLUMN model_id varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL", - }, nil - case "sqlite", "postgres", "postgresql": - return nil, nil - default: - return nil, fmt.Errorf("migrate MySQL model price identity: unsupported database driver %q", driver) - } -} diff --git a/internal/storage/migration_0004_mysql_model_price_identity_test.go b/internal/storage/migration_0004_mysql_model_price_identity_test.go deleted file mode 100644 index d38e5d3d..00000000 --- a/internal/storage/migration_0004_mysql_model_price_identity_test.go +++ /dev/null @@ -1,39 +0,0 @@ -package storage - -import ( - "reflect" - "testing" -) - -func TestMySQLModelPriceIdentityStatementsPreserveExactModelIDs(t *testing.T) { - got, err := mysqlModelPriceIdentityStatements("mysql") - if err != nil { - t.Fatalf("mysqlModelPriceIdentityStatements() error = %v", err) - } - want := []string{ - "ALTER TABLE model_prices MODIFY COLUMN model_id varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL", - } - if !reflect.DeepEqual(got, want) { - t.Fatalf("mysqlModelPriceIdentityStatements() = %#v, want %#v", got, want) - } -} - -func TestMySQLModelPriceIdentityStatementsLeaveCaseSensitiveDriversUntouched(t *testing.T) { - for _, driver := range []string{"sqlite", "postgres", "postgresql"} { - t.Run(driver, func(t *testing.T) { - got, err := mysqlModelPriceIdentityStatements(driver) - if err != nil { - t.Fatalf("mysqlModelPriceIdentityStatements(%q) error = %v", driver, err) - } - if len(got) != 0 { - t.Fatalf("mysqlModelPriceIdentityStatements(%q) = %#v, want no statements", driver, got) - } - }) - } -} - -func TestMySQLModelPriceIdentityStatementsRejectUnsupportedDriver(t *testing.T) { - if _, err := mysqlModelPriceIdentityStatements("sqlserver"); err == nil { - t.Fatal("mysqlModelPriceIdentityStatements() accepted an unsupported driver") - } -} diff --git a/internal/storage/migration_0005_request_log_model_consistency.go b/internal/storage/migration_0005_request_log_model_consistency.go deleted file mode 100644 index a0ff3906..00000000 --- a/internal/storage/migration_0005_request_log_model_consistency.go +++ /dev/null @@ -1,51 +0,0 @@ -package storage - -import ( - "fmt" - - "gorm.io/gorm" - - "gpt-load/internal/storage/models" - "gpt-load/internal/telemetry" -) - -// addRequestLogModelConsistencyColumns preserves existing logs while enabling -// new requests to persist the upstream-declared response model. Historical -// successful modeled requests become unknown because they were not observed at -// the time; failed and model-less requests remain not applicable. -func addRequestLogModelConsistencyColumns(db *gorm.DB) error { - const constraint = "chk_request_log_model_consistency" - for _, field := range []string{"UpstreamReportedModel", "ModelConsistency"} { - if db.Migrator().HasColumn(&models.RequestLog{}, field) { - continue - } - var err error - if field == "ModelConsistency" && db.Dialector.Name() == "sqlite" { - err = db.Exec(`ALTER TABLE request_logs ADD COLUMN model_consistency varchar(32) NOT NULL DEFAULT 'not_applicable' CONSTRAINT "chk_request_log_model_consistency" CHECK (model_consistency IN ('not_applicable','match','unknown','mismatch'))`).Error - } else { - err = db.Migrator().AddColumn(&models.RequestLog{}, field) - } - if err != nil { - return fmt.Errorf("add request log model consistency column %s: %w", field, err) - } - } - - if err := db.Model(&models.RequestLog{}). - Where( - "status = ? AND upstream_model <> ?", - string(telemetry.RequestStatusSuccess), - "", - ). - Update("model_consistency", string(telemetry.ModelConsistencyUnknown)).Error; err != nil { - return fmt.Errorf("backfill request log model consistency: %w", err) - } - if !db.Migrator().HasConstraint(&models.RequestLog{}, constraint) { - if db.Dialector.Name() == "sqlite" { - return fmt.Errorf("request log model consistency constraint is missing after column migration") - } - if err := db.Migrator().CreateConstraint(&models.RequestLog{}, constraint); err != nil { - return fmt.Errorf("create request log model consistency constraint: %w", err) - } - } - return nil -} diff --git a/internal/storage/migration_0006_channel_execution.go b/internal/storage/migration_0006_channel_execution.go new file mode 100644 index 00000000..9ac4836d --- /dev/null +++ b/internal/storage/migration_0006_channel_execution.go @@ -0,0 +1,731 @@ +package storage + +import ( + "encoding/json" + "fmt" + "net/url" + "sort" + "strings" + + "gorm.io/gorm" + "gorm.io/gorm/clause" + + "gpt-load/internal/channel" + "gpt-load/internal/storage/models" +) + +func retainLegacyMigrationPosition(*gorm.DB) error { return nil } + +type legacyExecutionGroup struct { + ID uint + Name string + ProviderID *string + UpstreamURL string + Protocols models.JSON + Models models.JSON + Config models.JSON + WeightManual *int + ValidationModel *string + Enabled bool + CreatedAtMS int64 + UpdatedAtMS int64 +} + +func (legacyExecutionGroup) TableName() string { return "groups" } + +type legacyExecutionCredential struct { + ID uint + GroupID uint + KeyValue string + KeyHash string + Status string + WeightManual *int + CreatedAtMS int64 + UpdatedAtMS int64 +} + +// migrationCredential mirrors the final credential columns without the Group +// association. Following that association would try to migrate the legacy +// groups table before its channel fields have been backfilled. +type legacyExecutionPrice struct { + ID uint + ModelID string + InputPriceNanoUSDPerMillionTokens *int64 + OutputPriceNanoUSDPerMillionTokens *int64 + CacheReadPriceNanoUSDPerMillionTokens *int64 + CacheWritePriceNanoUSDPerMillionTokens *int64 + ContextPriceTiers models.JSON + IsManual bool + CreatedAtMS int64 + UpdatedAtMS int64 +} + +type legacyRequestLogAttempt struct { + RequestID string + Sequence int + CompletedAtMS int64 + GroupID uint + GroupName string + KeyID uint + UpstreamModel string + StatusCode int + DurationMs int64 + FailureCategory string + Action string + WillRetry bool + ErrorCode string + ErrorSummary string + PricingReceipt models.JSON +} + +func (legacyRequestLogAttempt) TableName() string { return "request_log_attempts" } + +type mappedExecutionGroup struct { + legacyExecutionGroup + channelID channel.ID + params json.RawMessage + overrides json.RawMessage +} + +type migrationGroup struct { + ID uint `gorm:"primaryKey;autoIncrement"` + Name string `gorm:"type:varchar(255);not null;uniqueIndex"` + ChannelID string `gorm:"type:varchar(64);not null"` + Params models.JSON `gorm:"type:json;not null"` + Models models.JSON `gorm:"type:json;not null"` + WeightManual *int + ValidationModel *string `gorm:"type:varchar(255)"` + Overrides models.JSON `gorm:"type:json"` + Enabled bool `gorm:"not null;default:true"` + CreatedAtMS int64 `gorm:"column:created_at_ms;not null;check:chk_group_created_at,created_at_ms >= 0"` + UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null;check:chk_group_updated_at,updated_at_ms >= 0"` +} + +func (migrationGroup) TableName() string { return "groups" } + +func migrateChannelExecution(db *gorm.DB) error { + if db == nil || db.Dialector == nil { + return fmt.Errorf("migrate channel execution: database is unavailable") + } + if finalChannelExecutionSchema(db) { + return nil + } + if !db.Migrator().HasTable("groups") { + return fmt.Errorf("migrate channel execution: groups table is missing") + } + + mappedGroups, err := loadAndMapLegacyExecutionGroups(db) + if err != nil { + return err + } + legacyCredentials, err := loadLegacyExecutionCredentials(db) + if err != nil { + return err + } + legacyPrices, err := loadLegacyExecutionPrices(db) + if err != nil { + return err + } + + if db.Migrator().HasTable("upstream_keys") { + if err := dropChannelMigrationTable(db, "upstream_keys"); err != nil { + return fmt.Errorf("migrate channel execution: drop upstream_keys: %w", err) + } + } + if err := migrateLegacyGroups(db, mappedGroups); err != nil { + return err + } + if err := migrateLegacyCredentials(db, legacyCredentials); err != nil { + return err + } + if err := migrateLegacyModelPrices(db, mappedGroups, legacyPrices); err != nil { + return err + } + if err := prepareLegacyRequestAndUsageSchema(db); err != nil { + return err + } + if err := createInitialV2Tables(db); err != nil { + return fmt.Errorf("migrate channel execution: build final schema: %w", err) + } + if err := backfillChannelExecutionIdentity(db, mappedGroups); err != nil { + return err + } + return nil +} + +func finalChannelExecutionSchema(db *gorm.DB) bool { + return db.Migrator().HasColumn("groups", "channel_id") && + db.Migrator().HasColumn("groups", "params") && + db.Migrator().HasTable("credentials") && + !db.Migrator().HasTable("upstream_keys") && + db.Migrator().HasColumn("model_prices", "channel_id") && + db.Migrator().HasColumn("request_log_attempts", "credential_id") +} + +func loadAndMapLegacyExecutionGroups(db *gorm.DB) ([]mappedExecutionGroup, error) { + if !db.Migrator().HasColumn("groups", "upstream_url") || + !db.Migrator().HasColumn("groups", "protocols") { + return nil, fmt.Errorf("migrate channel execution: groups schema is partial") + } + var rows []legacyExecutionGroup + if err := db.Table("groups").Order("id ASC").Find(&rows).Error; err != nil { + return nil, fmt.Errorf("migrate channel execution: load groups: %w", err) + } + registry := channel.NewRegistry() + result := make([]mappedExecutionGroup, 0, len(rows)) + for _, row := range rows { + channelID, params, err := mapLegacyExecutionChannel(registry, row) + if err != nil { + return nil, fmt.Errorf("migrate channel execution: group %d: %w", row.ID, err) + } + overrides := json.RawMessage(`{}`) + if trimmed := strings.TrimSpace(string(row.Config)); trimmed != "" && trimmed != "null" { + if !json.Valid(row.Config) { + return nil, fmt.Errorf("migrate channel execution: group %d has invalid config", row.ID) + } + overrides = append(json.RawMessage(nil), row.Config...) + } + result = append(result, mappedExecutionGroup{ + legacyExecutionGroup: row, + channelID: channelID, + params: params, + overrides: overrides, + }) + } + return result, nil +} + +func mapLegacyExecutionChannel( + registry *channel.Registry, + row legacyExecutionGroup, +) (channel.ID, json.RawMessage, error) { + normalizedURL, err := normalizeLegacyChannelURL(row.UpstreamURL) + if err != nil { + return "", nil, err + } + if id, ok := fixedLegacyChannelURLs[normalizedURL]; ok { + return id, json.RawMessage(`{}`), nil + } + var protocols []string + if err := json.Unmarshal(row.Protocols, &protocols); err != nil || len(protocols) == 0 { + return "", nil, fmt.Errorf("protocols must be a non-empty JSON array") + } + family := "" + for _, item := range protocols { + candidate := "" + switch strings.TrimSpace(item) { + case "openai-completions", "openai-responses": + candidate = "openai" + case "anthropic": + candidate = "anthropic" + case "gemini": + candidate = "gemini" + default: + return "", nil, fmt.Errorf("unsupported protocol %q", item) + } + if family != "" && family != candidate { + return "", nil, fmt.Errorf("multiple upstream protocol families require an explicit channel") + } + family = candidate + } + var channelID channel.ID + switch family { + case "openai": + channelID = channel.OpenAICompatible + case "anthropic": + channelID = channel.AnthropicCompatible + case "gemini": + channelID = channel.GeminiCompatible + default: + return "", nil, fmt.Errorf("cannot infer channel") + } + encoded, err := json.Marshal(map[string]string{"base_url": normalizedURL}) + if err != nil { + return "", nil, fmt.Errorf("encode channel params: %w", err) + } + resolved, err := registry.Resolve(channelID, encoded) + if err != nil { + return "", nil, fmt.Errorf("validate channel params: %w", err) + } + return channelID, append(json.RawMessage(nil), resolved.TargetConfig...), nil +} + +func normalizeLegacyChannelURL(value string) (string, error) { + parsed, err := url.Parse(strings.TrimSpace(value)) + if err != nil || parsed == nil || parsed.Opaque != "" || parsed.User != nil || parsed.Host == "" || + (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Fragment != "" { + return "", fmt.Errorf("invalid upstream URL") + } + parsed.Scheme = strings.ToLower(parsed.Scheme) + parsed.Host = strings.ToLower(parsed.Host) + parsed.Path = strings.TrimRight(parsed.Path, "/") + parsed.RawPath = strings.TrimRight(parsed.RawPath, "/") + parsed.ForceQuery = false + return parsed.String(), nil +} + +var fixedLegacyChannelURLs = map[string]channel.ID{ + "https://api.openai.com": channel.OpenAI, + "https://api.openai.com/v1": channel.OpenAI, + "https://api.anthropic.com": channel.Anthropic, + "https://api.anthropic.com/v1": channel.Anthropic, + "https://generativelanguage.googleapis.com": channel.Gemini, + "https://generativelanguage.googleapis.com/v1beta": channel.Gemini, + "https://api.deepseek.com/v1": channel.DeepSeek, + "https://api.moonshot.cn/v1": channel.MoonshotAI, + "https://api.siliconflow.cn/v1": channel.SiliconFlow, + "https://open.bigmodel.cn/api/paas/v4": channel.ZhipuAI, + "https://dashscope.aliyuncs.com/compatible-mode/v1": channel.Alibaba, + "https://ark.cn-beijing.volces.com/api/v3": channel.Volcengine, + "https://openrouter.ai/api/v1": channel.OpenRouter, + "https://api.groq.com/openai/v1": channel.Groq, + "https://api.x.ai/v1": channel.XAI, +} + +func loadLegacyExecutionCredentials(db *gorm.DB) ([]legacyExecutionCredential, error) { + if !db.Migrator().HasTable("upstream_keys") { + return nil, nil + } + var rows []legacyExecutionCredential + if err := db.Table("upstream_keys").Order("id ASC").Find(&rows).Error; err != nil { + return nil, fmt.Errorf("migrate channel execution: load upstream keys: %w", err) + } + return rows, nil +} + +func migrateLegacyCredentials(db *gorm.DB, rows []legacyExecutionCredential) error { + if err := db.AutoMigrate(&models.Credential{}); err != nil { + return fmt.Errorf("migrate channel execution: create credentials: %w", err) + } + for _, row := range rows { + credential := models.Credential{ + ID: row.ID, GroupID: row.GroupID, Data: row.KeyValue, Fingerprint: row.KeyHash, + Status: models.CredentialStatus(row.Status), WeightManual: cloneMigrationWeight(row.WeightManual), + CreatedAtMS: row.CreatedAtMS, UpdatedAtMS: row.UpdatedAtMS, + } + if err := db.Clauses(clause.OnConflict{DoNothing: true}).Create(&credential).Error; err != nil { + return fmt.Errorf("migrate channel execution: copy credential %d: %w", row.ID, err) + } + } + if err := resetMigrationIdentitySequence(db, "credentials"); err != nil { + return err + } + return nil +} + +func migrateLegacyGroups(db *gorm.DB, rows []mappedExecutionGroup) error { + migrator := db.Migrator() + const legacyTable = "groups_channel_legacy" + if migrator.HasTable(legacyTable) { + return fmt.Errorf("migrate channel execution: partial group migration") + } + if migrator.HasIndex(&legacyExecutionGroup{}, "idx_groups_name") { + if err := dropChannelMigrationIndex(db, "groups", "idx_groups_name"); err != nil { + return fmt.Errorf("migrate channel execution: drop legacy group name index: %w", err) + } + } + if err := migrator.RenameTable("groups", legacyTable); err != nil { + return fmt.Errorf("migrate channel execution: preserve groups: %w", err) + } + if err := preservePostgresLegacySequence(db, legacyTable); err != nil { + return err + } + if err := db.AutoMigrate(&migrationGroup{}); err != nil { + return fmt.Errorf("migrate channel execution: create channel groups: %w", err) + } + for _, row := range rows { + group := migrationGroup{ + ID: row.ID, Name: row.Name, ChannelID: string(row.channelID), Params: append(models.JSON(nil), row.params...), + Models: append(models.JSON(nil), row.Models...), WeightManual: cloneMigrationWeight(row.WeightManual), + ValidationModel: row.ValidationModel, Overrides: append(models.JSON(nil), row.overrides...), Enabled: row.Enabled, + CreatedAtMS: row.CreatedAtMS, UpdatedAtMS: row.UpdatedAtMS, + } + if err := db.Create(&group).Error; err != nil { + return fmt.Errorf("migrate channel execution: copy group %d: %w", row.ID, err) + } + } + if err := resetMigrationIdentitySequence(db, "groups"); err != nil { + return err + } + if err := dropChannelMigrationTable(db, legacyTable); err != nil { + return fmt.Errorf("migrate channel execution: drop legacy groups: %w", err) + } + return nil +} + +func loadLegacyExecutionPrices(db *gorm.DB) ([]legacyExecutionPrice, error) { + if !db.Migrator().HasTable("model_prices") { + return nil, nil + } + var rows []legacyExecutionPrice + if err := db.Table("model_prices").Order("model_id ASC").Order("id ASC").Find(&rows).Error; err != nil { + return nil, fmt.Errorf("migrate channel execution: load model prices: %w", err) + } + return rows, nil +} + +func migrateLegacyModelPrices( + db *gorm.DB, + groups []mappedExecutionGroup, + prices []legacyExecutionPrice, +) error { + if db.Migrator().HasTable("model_prices_channel_legacy") { + return fmt.Errorf("migrate channel execution: partial model price migration") + } + if db.Migrator().HasTable("model_prices") { + if err := db.Migrator().RenameTable("model_prices", "model_prices_channel_legacy"); err != nil { + return fmt.Errorf("migrate channel execution: preserve model prices: %w", err) + } + if err := preservePostgresLegacySequence(db, "model_prices_channel_legacy"); err != nil { + return err + } + } + if err := db.AutoMigrate(&models.ModelPrice{}); err != nil { + return fmt.Errorf("migrate channel execution: create channel model prices: %w", err) + } + channelsByModel, allChannels, err := legacyPriceChannels(groups) + if err != nil { + return err + } + for _, price := range prices { + channels := channelsByModel[price.ModelID] + if len(channels) == 0 { + channels = allChannels + } + if len(channels) == 0 { + channels = []channel.ID{channel.OpenAI} + } + for index, channelID := range channels { + id := uint(0) + if index == 0 { + id = price.ID + } + row := models.ModelPrice{ + ID: id, ChannelID: string(channelID), ModelID: price.ModelID, + InputPriceNanoUSDPerMillionTokens: cloneMigrationInt64(price.InputPriceNanoUSDPerMillionTokens), + OutputPriceNanoUSDPerMillionTokens: cloneMigrationInt64(price.OutputPriceNanoUSDPerMillionTokens), + CacheReadPriceNanoUSDPerMillionTokens: cloneMigrationInt64(price.CacheReadPriceNanoUSDPerMillionTokens), + CacheWritePriceNanoUSDPerMillionTokens: cloneMigrationInt64(price.CacheWritePriceNanoUSDPerMillionTokens), + ContextPriceTiers: append(models.JSON(nil), price.ContextPriceTiers...), IsManual: price.IsManual, + CreatedAtMS: price.CreatedAtMS, UpdatedAtMS: price.UpdatedAtMS, + } + if err := db.Create(&row).Error; err != nil { + return fmt.Errorf("migrate channel execution: copy price %s/%s: %w", channelID, price.ModelID, err) + } + } + } + if err := resetMigrationIdentitySequence(db, "model_prices"); err != nil { + return err + } + if db.Migrator().HasTable("model_prices_channel_legacy") { + if err := dropChannelMigrationTable(db, "model_prices_channel_legacy"); err != nil { + return fmt.Errorf("migrate channel execution: drop legacy model prices: %w", err) + } + } + return nil +} + +func legacyPriceChannels( + groups []mappedExecutionGroup, +) (map[string][]channel.ID, []channel.ID, error) { + sets := make(map[string]map[channel.ID]struct{}) + allSet := make(map[channel.ID]struct{}) + for _, group := range groups { + allSet[group.channelID] = struct{}{} + var items []struct { + ID string `json:"id"` + } + if err := json.Unmarshal(group.Models, &items); err != nil { + return nil, nil, fmt.Errorf("migrate channel execution: group %d has invalid models", group.ID) + } + for _, item := range items { + modelID := strings.TrimSpace(item.ID) + if modelID == "" { + continue + } + if sets[modelID] == nil { + sets[modelID] = make(map[channel.ID]struct{}) + } + sets[modelID][group.channelID] = struct{}{} + } + } + allChannels := sortedMigrationChannels(allSet) + result := make(map[string][]channel.ID, len(sets)) + for modelID, set := range sets { + result[modelID] = sortedMigrationChannels(set) + } + return result, allChannels, nil +} + +func sortedMigrationChannels(set map[channel.ID]struct{}) []channel.ID { + result := make([]channel.ID, 0, len(set)) + for id := range set { + result = append(result, id) + } + sort.Slice(result, func(i, j int) bool { return result[i] < result[j] }) + return result +} + +func prepareLegacyRequestAndUsageSchema(db *gorm.DB) error { + migrator := db.Migrator() + if migrator.HasTable("request_log_attempts") && migrator.HasColumn("request_log_attempts", "key_id") { + if err := rebuildLegacyRequestLogAttempts(db); err != nil { + return err + } + } + for table, index := range map[string]string{ + "usage_stats": "idx_usage_stats_bucket_access_group_model", + "usage_aggregation_journal": "idx_usage_aggregation_journal_identity", + } { + if migrator.HasTable(table) && migrator.HasIndex(table, index) { + if err := dropChannelMigrationIndex(db, table, index); err != nil { + return fmt.Errorf("migrate channel execution: drop %s: %w", index, err) + } + } + } + return nil +} + +func rebuildLegacyRequestLogAttempts(db *gorm.DB) error { + const legacyTable = "request_log_attempts_channel_legacy" + migrator := db.Migrator() + if migrator.HasTable(legacyTable) { + return fmt.Errorf("migrate channel execution: partial request attempt migration") + } + var rows []legacyRequestLogAttempt + if err := db.Table("request_log_attempts").Order("request_id ASC").Order("sequence ASC").Find(&rows).Error; err != nil { + return fmt.Errorf("migrate channel execution: load request attempts: %w", err) + } + for _, index := range []string{ + "idx_request_log_attempts_error_completed_request", + "idx_request_log_attempts_failure_completed_request", + "idx_request_log_attempts_status_completed_request", + "idx_request_log_attempts_model_completed_request", + "idx_request_log_attempts_key_completed_request", + "idx_request_log_attempts_group_completed_request", + } { + if migrator.HasIndex(&legacyRequestLogAttempt{}, index) { + if err := dropChannelMigrationIndex(db, "request_log_attempts", index); err != nil { + return fmt.Errorf("migrate channel execution: drop legacy attempt index %s: %w", index, err) + } + } + } + if err := dropLegacyAttemptForeignKey(db); err != nil { + return err + } + if err := migrator.RenameTable("request_log_attempts", legacyTable); err != nil { + return fmt.Errorf("migrate channel execution: preserve request attempts: %w", err) + } + if err := db.AutoMigrate(&models.RequestLogAttempt{}); err != nil { + return fmt.Errorf("migrate channel execution: create request attempts: %w", err) + } + for _, row := range rows { + action := row.Action + switch action { + case "cooldown_key": + action = "cooldown_credential" + case "fail_key": + action = "fail_credential" + } + attempt := models.RequestLogAttempt{ + RequestID: row.RequestID, Sequence: row.Sequence, CompletedAtMS: row.CompletedAtMS, + GroupID: row.GroupID, GroupName: row.GroupName, CredentialID: row.KeyID, + UpstreamModel: row.UpstreamModel, StatusCode: row.StatusCode, DurationMs: row.DurationMs, + FailureCategory: row.FailureCategory, Action: action, WillRetry: row.WillRetry, + ErrorCode: row.ErrorCode, ErrorSummary: row.ErrorSummary, + PricingReceipt: append(models.JSON(nil), row.PricingReceipt...), + } + if err := db.Omit("RequestLog").Create(&attempt).Error; err != nil { + return fmt.Errorf("migrate channel execution: copy request attempt %s/%d: %w", row.RequestID, row.Sequence, err) + } + } + if err := dropChannelMigrationTable(db, legacyTable); err != nil { + return fmt.Errorf("migrate channel execution: drop legacy request attempts: %w", err) + } + return nil +} + +func dropChannelMigrationTable(db *gorm.DB, table string) error { + switch table { + case "upstream_keys", "groups_channel_legacy", "model_prices_channel_legacy", "request_log_attempts_channel_legacy": + default: + return fmt.Errorf("unsupported migration table %q", table) + } + if !strings.EqualFold(db.Dialector.Name(), "mysql") { + return db.Migrator().DropTable(table) + } + // The migration already runs on a pinned connection for the advisory lock. + // GORM's MySQL DropTable opens another pinned connection internally, which + // fails with gorm.ErrInvalidDB here. Execute the fixed-table drop on the + // existing connection and always restore foreign-key checks. + if err := db.Exec("SET FOREIGN_KEY_CHECKS = 0").Error; err != nil { + return fmt.Errorf("disable foreign key checks: %w", err) + } + dropErr := db.Exec("DROP TABLE IF EXISTS ?", clause.Table{Name: table}).Error + restoreErr := db.Exec("SET FOREIGN_KEY_CHECKS = 1").Error + if dropErr != nil { + if restoreErr != nil { + return fmt.Errorf("drop table: %v; restore foreign key checks: %w", dropErr, restoreErr) + } + return dropErr + } + if restoreErr != nil { + return fmt.Errorf("restore foreign key checks: %w", restoreErr) + } + return nil +} + +func dropChannelMigrationIndex(db *gorm.DB, table, index string) error { + allowed := map[string]map[string]struct{}{ + "groups": { + "idx_groups_name": {}, + }, + "request_log_attempts": { + "idx_request_log_attempts_error_completed_request": {}, + "idx_request_log_attempts_failure_completed_request": {}, + "idx_request_log_attempts_status_completed_request": {}, + "idx_request_log_attempts_model_completed_request": {}, + "idx_request_log_attempts_key_completed_request": {}, + "idx_request_log_attempts_group_completed_request": {}, + }, + "usage_stats": { + "idx_usage_stats_bucket_access_group_model": {}, + }, + "usage_aggregation_journal": { + "idx_usage_aggregation_journal_identity": {}, + }, + } + if _, ok := allowed[table][index]; !ok { + return fmt.Errorf("unsupported migration index %s.%s", table, index) + } + if strings.EqualFold(db.Dialector.Name(), "mysql") { + return db.Exec("ALTER TABLE ? DROP INDEX ?", clause.Table{Name: table}, clause.Column{Name: index}).Error + } + return db.Exec("DROP INDEX IF EXISTS ?", clause.Column{Name: index}).Error +} + +func dropLegacyAttemptForeignKey(db *gorm.DB) error { + switch strings.ToLower(db.Dialector.Name()) { + case "sqlite": + return nil + case "mysql": + if err := db.Exec("ALTER TABLE request_log_attempts DROP FOREIGN KEY fk_request_log_attempts_request_log").Error; err != nil { + return fmt.Errorf("migrate channel execution: drop legacy attempt foreign key: %w", err) + } + return nil + case "postgres", "postgresql": + if err := db.Exec("ALTER TABLE request_log_attempts DROP CONSTRAINT IF EXISTS fk_request_log_attempts_request_log").Error; err != nil { + return fmt.Errorf("migrate channel execution: drop legacy attempt foreign key: %w", err) + } + return nil + default: + return fmt.Errorf("migrate channel execution: unsupported database driver %q", db.Dialector.Name()) + } +} + +func backfillChannelExecutionIdentity(db *gorm.DB, groups []mappedExecutionGroup) error { + for _, group := range groups { + for _, table := range []string{"request_logs", "request_log_attempts", "usage_stats", "usage_aggregation_journal"} { + if !db.Migrator().HasTable(table) || !db.Migrator().HasColumn(table, "channel_id") { + continue + } + if err := db.Table(table).Where("group_id = ?", group.ID).Update("channel_id", string(group.channelID)).Error; err != nil { + return fmt.Errorf("migrate channel execution: backfill %s channel: %w", table, err) + } + } + } + if db.Migrator().HasTable("request_logs") && db.Migrator().HasTable("request_log_attempts") && + db.Migrator().HasColumn("request_logs", "credential_id") { + type attemptIdentity struct { + RequestID string + CredentialID uint + } + var attempts []attemptIdentity + if err := db.Table("request_log_attempts").Select("request_id", "credential_id").Order("sequence ASC").Find(&attempts).Error; err != nil { + return fmt.Errorf("migrate channel execution: load attempt identities: %w", err) + } + byRequest := make(map[string]uint) + ambiguous := make(map[string]bool) + for _, attempt := range attempts { + if existing, seen := byRequest[attempt.RequestID]; seen && existing != attempt.CredentialID { + ambiguous[attempt.RequestID] = true + continue + } + byRequest[attempt.RequestID] = attempt.CredentialID + } + for requestID, credentialID := range byRequest { + if ambiguous[requestID] || credentialID == 0 { + continue + } + if err := db.Table("request_logs").Where("id = ?", requestID).Update("credential_id", credentialID).Error; err != nil { + return fmt.Errorf("migrate channel execution: backfill request credential: %w", err) + } + } + } + return nil +} + +func cloneMigrationWeight(value *int) *int { + if value == nil { + return nil + } + copy := *value + return © +} + +func preservePostgresLegacySequence(db *gorm.DB, legacyTable string) error { + if db == nil || db.Dialector == nil || + !strings.EqualFold(db.Dialector.Name(), "postgres") && + !strings.EqualFold(db.Dialector.Name(), "postgresql") { + return nil + } + sequenceNames := map[string]string{ + "groups_channel_legacy": "groups_channel_legacy_id_seq", + "model_prices_channel_legacy": "model_prices_channel_legacy_id_seq", + } + legacySequence, ok := sequenceNames[legacyTable] + if !ok { + return fmt.Errorf("migrate channel execution: unsupported legacy sequence table %q", legacyTable) + } + statement := fmt.Sprintf(`DO $migration$ +DECLARE sequence_name regclass; +BEGIN + sequence_name := pg_get_serial_sequence('%s', 'id')::regclass; + IF sequence_name IS NOT NULL THEN + EXECUTE format('ALTER SEQUENCE %%s RENAME TO %%I', sequence_name, '%s'); + END IF; +END +$migration$;`, legacyTable, legacySequence) + if err := db.Exec(statement).Error; err != nil { + return fmt.Errorf("migrate channel execution: preserve %s sequence: %w", legacyTable, err) + } + return nil +} + +func resetMigrationIdentitySequence(db *gorm.DB, table string) error { + if db == nil || db.Dialector == nil || + !strings.EqualFold(db.Dialector.Name(), "postgres") && + !strings.EqualFold(db.Dialector.Name(), "postgresql") { + return nil + } + if table != "groups" && table != "credentials" && table != "model_prices" { + return fmt.Errorf("migrate channel execution: unsupported identity table %q", table) + } + statement := fmt.Sprintf( + "SELECT setval(pg_get_serial_sequence('%s', 'id'), COALESCE(MAX(id), 1), COUNT(*) > 0) FROM %s", + table, + table, + ) + if err := db.Exec(statement).Error; err != nil { + return fmt.Errorf("migrate channel execution: reset %s identity sequence: %w", table, err) + } + return nil +} + +func cloneMigrationInt64(value *int64) *int64 { + if value == nil { + return nil + } + copy := *value + return © +} diff --git a/internal/storage/migration_0006_channel_execution_test.go b/internal/storage/migration_0006_channel_execution_test.go new file mode 100644 index 00000000..d329513a --- /dev/null +++ b/internal/storage/migration_0006_channel_execution_test.go @@ -0,0 +1,263 @@ +package storage_test + +import ( + "encoding/json" + "os" + "reflect" + "strings" + "testing" + + "gorm.io/gorm" + + "gpt-load/internal/platform/config" + "gpt-load/internal/storage" + "gpt-load/internal/storage/models" +) + +type legacyChannelGroup struct { + ID uint `gorm:"primaryKey;autoIncrement"` + Name string `gorm:"type:varchar(255);not null;uniqueIndex"` + ProviderID *string `gorm:"type:varchar(255)"` + UpstreamURL string `gorm:"type:text;not null"` + Protocols models.JSON `gorm:"type:json;not null"` + Models models.JSON `gorm:"type:json;not null"` + ConvertEnabled bool `gorm:"not null;default:false"` + WeightManual *int + ValidationModel *string `gorm:"type:varchar(255)"` + Config models.JSON `gorm:"type:json"` + Enabled bool `gorm:"not null;default:true"` + CreatedAtMS int64 `gorm:"column:created_at_ms;not null"` + UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null"` +} + +func (legacyChannelGroup) TableName() string { return "groups" } + +type legacyChannelKey struct { + ID uint `gorm:"primaryKey;autoIncrement"` + GroupID uint `gorm:"not null;uniqueIndex:idx_upstream_keys_group_hash,priority:1"` + KeyValue string `gorm:"type:text;not null"` + KeyHash string `gorm:"type:varchar(128);not null;uniqueIndex:idx_upstream_keys_group_hash,priority:2"` + Status string `gorm:"type:varchar(32);not null"` + WeightManual *int + CreatedAtMS int64 `gorm:"column:created_at_ms;not null"` + UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null"` +} + +func (legacyChannelKey) TableName() string { return "upstream_keys" } + +type legacyGlobalPrice struct { + ID uint `gorm:"primaryKey;autoIncrement"` + ModelID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_model_prices_model"` + InputPriceNanoUSDPerMillionTokens *int64 + OutputPriceNanoUSDPerMillionTokens *int64 + CacheReadPriceNanoUSDPerMillionTokens *int64 + CacheWritePriceNanoUSDPerMillionTokens *int64 + ContextPriceTiers models.JSON `gorm:"type:json"` + IsManual bool `gorm:"not null;default:false"` + CreatedAtMS int64 `gorm:"column:created_at_ms;not null"` + UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null"` +} + +func (legacyGlobalPrice) TableName() string { return "model_prices" } + +func TestAutoMigrateUpgradesExistingV2ChannelsCredentialsAndPrices(t *testing.T) { + db := openInitialV2TestDatabase(t) + runChannelExecutionMigrationContract(t, db) + assertChannelMigrationIdentitySequences(t, db) +} + +func TestExternalDatabaseChannelExecutionMigration(t *testing.T) { + dsn := strings.TrimSpace(os.Getenv("GPT_LOAD_DATABASE_MIGRATION_TEST_DSN")) + if dsn == "" { + t.Skip("GPT_LOAD_DATABASE_MIGRATION_TEST_DSN is not set") + } + db, err := storage.OpenWithSource(dsn, config.DatabaseSourceExternal) + if err != nil { + t.Fatalf("OpenWithSource() error = %v", err) + } + sqlDB, err := db.DB() + if err != nil { + t.Fatalf("DB() error = %v", err) + } + t.Cleanup(func() { _ = sqlDB.Close() }) + tables, err := db.Migrator().GetTables() + if err != nil { + t.Fatalf("GetTables() error = %v", err) + } + if len(tables) != 0 { + t.Fatalf("migration contract requires an empty dedicated database, found %v", tables) + } + + runChannelExecutionMigrationContract(t, db) + assertChannelMigrationIdentitySequences(t, db) +} + +func runChannelExecutionMigrationContract(t *testing.T, db *gorm.DB) { + t.Helper() + if err := db.AutoMigrate( + &legacyChannelGroup{}, + &legacyChannelKey{}, + &legacyGlobalPrice{}, + ); err != nil { + t.Fatalf("create legacy schema: %v", err) + } + if err := db.Exec(`CREATE TABLE schema_migrations ( + id varchar(255) PRIMARY KEY NOT NULL + )`).Error; err != nil { + t.Fatalf("create migration ledger: %v", err) + } + for _, id := range []string{ + "0001_initial_v2", + "0002_request_log_reasoning", + "0003_global_model_prices", + "0004_mysql_model_price_identity", + "0005_request_log_model_consistency", + } { + if err := db.Exec("INSERT INTO schema_migrations(id) VALUES (?)", id).Error; err != nil { + t.Fatalf("seed migration %s: %v", id, err) + } + } + openAI, deepSeek := "openai", "deepseek" + groups := []legacyChannelGroup{ + { + ID: 1, Name: "official", ProviderID: &openAI, + UpstreamURL: "https://api.openai.com/v1/", + Protocols: models.JSON(`["openai-completions","openai-responses"]`), + Models: models.JSON(`[{"id":"shared-model"}]`), + Config: models.JSON(`{"request_timeout":31}`), Enabled: true, + CreatedAtMS: 1, UpdatedAtMS: 2, + }, + { + ID: 2, Name: "curated", ProviderID: &deepSeek, + UpstreamURL: "https://api.deepseek.com/v1", + Protocols: models.JSON(`["openai-completions","anthropic"]`), + Models: models.JSON(`[{"id":"shared-model"}]`), Enabled: true, + CreatedAtMS: 3, UpdatedAtMS: 4, + }, + } + if err := db.Create(&groups).Error; err != nil { + t.Fatalf("seed legacy groups: %v", err) + } + keys := []legacyChannelKey{ + {ID: 11, GroupID: 1, KeyValue: "cipher-openai", KeyHash: "hash-openai", Status: "active", CreatedAtMS: 1, UpdatedAtMS: 2}, + {ID: 21, GroupID: 2, KeyValue: "cipher-deepseek", KeyHash: "hash-deepseek", Status: "disabled", CreatedAtMS: 3, UpdatedAtMS: 4}, + } + if err := db.Create(&keys).Error; err != nil { + t.Fatalf("seed legacy keys: %v", err) + } + inputPrice, outputPrice := int64(100), int64(200) + if err := db.Create(&legacyGlobalPrice{ + ID: 31, ModelID: "shared-model", + InputPriceNanoUSDPerMillionTokens: &inputPrice, + OutputPriceNanoUSDPerMillionTokens: &outputPrice, + IsManual: true, CreatedAtMS: 5, UpdatedAtMS: 6, + }).Error; err != nil { + t.Fatalf("seed legacy price: %v", err) + } + + if err := storage.AutoMigrate(db); err != nil { + t.Fatalf("AutoMigrate() error = %v", err) + } + + var migratedGroups []models.Group + if err := db.Order("id ASC").Find(&migratedGroups).Error; err != nil { + t.Fatalf("load migrated groups: %v", err) + } + if len(migratedGroups) != 2 || migratedGroups[0].ChannelID != "openai" || + migratedGroups[1].ChannelID != "deepseek" { + t.Fatalf("migrated groups = %#v", migratedGroups) + } + assertMigrationJSON(t, migratedGroups[0].Params, `{}`) + assertMigrationJSON(t, migratedGroups[0].Overrides, `{"request_timeout":31}`) + assertMigrationJSON(t, migratedGroups[1].Params, `{}`) + var credentials []models.Credential + if err := db.Order("id ASC").Find(&credentials).Error; err != nil { + t.Fatalf("load migrated credentials: %v", err) + } + if got, want := credentials, []models.Credential{ + {ID: 11, GroupID: 1, Data: "cipher-openai", Fingerprint: "hash-openai", Status: models.CredentialStatusActive, CreatedAtMS: 1, UpdatedAtMS: 2}, + {ID: 21, GroupID: 2, Data: "cipher-deepseek", Fingerprint: "hash-deepseek", Status: models.CredentialStatusDisabled, CreatedAtMS: 3, UpdatedAtMS: 4}, + }; !reflect.DeepEqual(got, want) { + t.Fatalf("migrated credentials = %#v, want %#v", got, want) + } + var prices []models.ModelPrice + if err := db.Order("channel_id ASC").Find(&prices).Error; err != nil { + t.Fatalf("load migrated prices: %v", err) + } + if len(prices) != 2 || prices[0].ChannelID != "deepseek" || prices[1].ChannelID != "openai" || + prices[0].ModelID != "shared-model" || prices[1].ModelID != "shared-model" { + t.Fatalf("migrated prices = %#v", prices) + } + if db.Migrator().HasTable("upstream_keys") { + t.Fatal("legacy upstream_keys table remains") + } + for _, retired := range []string{"provider_id", "upstream_url", "protocols", "convert_enabled", "config"} { + if db.Migrator().HasColumn("groups", retired) { + t.Errorf("groups retains %s", retired) + } + } + + if err := storage.AutoMigrate(db); err != nil { + t.Fatalf("second AutoMigrate() error = %v", err) + } + var ledger []string + if err := db.Table("schema_migrations").Order("id ASC").Pluck("id", &ledger).Error; err != nil { + t.Fatalf("load migration ledger: %v", err) + } + if want := []string{ + "0001_initial_v2", "0002_request_log_reasoning", "0003_global_model_prices", + "0004_mysql_model_price_identity", "0005_request_log_model_consistency", + "0006_channel_execution", + }; !reflect.DeepEqual(ledger, want) { + t.Fatalf("migration ledger = %#v, want %#v", ledger, want) + } +} + +func assertMigrationJSON(t *testing.T, got models.JSON, want string) { + t.Helper() + var gotValue any + if err := json.Unmarshal(got, &gotValue); err != nil { + t.Fatalf("decode migrated JSON %q: %v", got, err) + } + var wantValue any + if err := json.Unmarshal([]byte(want), &wantValue); err != nil { + t.Fatalf("decode expected JSON %q: %v", want, err) + } + if !reflect.DeepEqual(gotValue, wantValue) { + t.Fatalf("migrated JSON = %s, want %s", got, want) + } +} + +func assertChannelMigrationIdentitySequences(t *testing.T, db *gorm.DB) { + t.Helper() + group := models.Group{ + Name: "post-migration-sequence", ChannelID: "openai", Params: models.JSON(`{}`), + Models: models.JSON(`[]`), Overrides: models.JSON(`{}`), Enabled: true, + CreatedAtMS: 7, UpdatedAtMS: 7, + } + if err := db.Create(&group).Error; err != nil { + t.Fatalf("create post-migration group: %v", err) + } + if group.ID <= 2 { + t.Fatalf("post-migration group id = %d, want > 2", group.ID) + } + credential := models.Credential{ + GroupID: group.ID, Data: "cipher-new", Fingerprint: "hash-new", + Status: models.CredentialStatusActive, CreatedAtMS: 7, UpdatedAtMS: 7, + } + if err := db.Create(&credential).Error; err != nil { + t.Fatalf("create post-migration credential: %v", err) + } + if credential.ID <= 21 { + t.Fatalf("post-migration credential id = %d, want > 21", credential.ID) + } + price := models.ModelPrice{ + ChannelID: "openai", ModelID: "post-migration-model", CreatedAtMS: 7, UpdatedAtMS: 7, + } + if err := db.Create(&price).Error; err != nil { + t.Fatalf("create post-migration price: %v", err) + } + if price.ID <= 31 { + t.Fatalf("post-migration price id = %d, want > 31", price.ID) + } +} diff --git a/internal/storage/migration_compatibility_test.go b/internal/storage/migration_compatibility_test.go deleted file mode 100644 index 4216f23e..00000000 --- a/internal/storage/migration_compatibility_test.go +++ /dev/null @@ -1,99 +0,0 @@ -package storage_test - -import ( - "testing" - - "gpt-load/internal/storage" - "gpt-load/internal/storage/models" -) - -func TestAutoMigrateUpgradesLegacyRequestLogColumns(t *testing.T) { - db, err := storage.Open(":memory:") - if err != nil { - t.Fatalf("Open(:memory:) error = %v", err) - } - sqlDB, err := db.DB() - if err != nil { - t.Fatalf("DB() error = %v", err) - } - t.Cleanup(func() { _ = sqlDB.Close() }) - - if err := db.Exec(`CREATE TABLE request_logs ( - id varchar(36) PRIMARY KEY NOT NULL, - status varchar(32) NOT NULL DEFAULT '', - upstream_model varchar(255) NOT NULL DEFAULT '' - )`).Error; err != nil { - t.Fatalf("create legacy request_logs: %v", err) - } - if err := db.Exec(`INSERT INTO request_logs(id, status, upstream_model) VALUES - ('success-modeled', 'success', 'model-a'), - ('success-model-less', 'success', ''), - ('failed-modeled', 'error', 'model-a')`).Error; err != nil { - t.Fatalf("seed legacy request logs: %v", err) - } - if err := db.Exec(`CREATE TABLE schema_migrations ( - id varchar(255) PRIMARY KEY NOT NULL - )`).Error; err != nil { - t.Fatalf("create schema_migrations: %v", err) - } - if err := db.Exec("INSERT INTO schema_migrations(id) VALUES (?)", "0001_initial_v2").Error; err != nil { - t.Fatalf("seed legacy migration ledger: %v", err) - } - - if err := storage.AutoMigrate(db); err != nil { - t.Fatalf("AutoMigrate() error = %v", err) - } - for _, field := range []string{ - "ReasoningMode", "ReasoningEffort", "ReasoningBudgetTokens", - "UpstreamReportedModel", "ModelConsistency", - } { - if !db.Migrator().HasColumn(&models.RequestLog{}, field) { - t.Errorf("legacy upgrade did not add request_logs.%s", field) - } - } - if !db.Migrator().HasConstraint(&models.RequestLog{}, "chk_request_log_model_consistency") { - t.Fatal("legacy upgrade did not create the request log model consistency constraint") - } - var migrationCount int64 - if err := db.Table("schema_migrations").Count(&migrationCount).Error; err != nil { - t.Fatalf("count migration ledger: %v", err) - } - if migrationCount != 5 { - t.Fatalf("migration ledger count = %d, want 5", migrationCount) - } - var consistencyByID map[string]string - var rows []struct { - ID string - ModelConsistency string - } - if err := db.Table("request_logs"). - Select("id, model_consistency"). - Order("id ASC"). - Scan(&rows).Error; err != nil { - t.Fatalf("read upgraded request logs: %v", err) - } - consistencyByID = make(map[string]string, len(rows)) - for _, row := range rows { - consistencyByID[row.ID] = row.ModelConsistency - } - if consistencyByID["success-modeled"] != "unknown" || - consistencyByID["success-model-less"] != "not_applicable" || - consistencyByID["failed-modeled"] != "not_applicable" { - t.Fatalf("upgraded model consistency = %#v", consistencyByID) - } - if err := db.Model(&models.RequestLog{}). - Where("id = ?", "success-modeled"). - Update("model_consistency", "invalid").Error; err == nil { - t.Fatal("legacy upgrade accepted an invalid request log model consistency value") - } - - if err := storage.AutoMigrate(db); err != nil { - t.Fatalf("second AutoMigrate() error = %v", err) - } - if err := db.Table("schema_migrations").Count(&migrationCount).Error; err != nil { - t.Fatalf("count migration ledger after second run: %v", err) - } - if migrationCount != 5 { - t.Fatalf("migration ledger count after second run = %d, want 5", migrationCount) - } -} diff --git a/internal/storage/migrations.go b/internal/storage/migrations.go index 0ae52b1b..70364e5c 100644 --- a/internal/storage/migrations.go +++ b/internal/storage/migrations.go @@ -38,22 +38,11 @@ var migrations = []migration{ ID: "0001_initial_v2", Up: createInitialV2Tables, }, - { - ID: "0002_request_log_reasoning", - Up: addRequestLogReasoningColumns, - }, - { - ID: "0003_global_model_prices", - Up: migrateGlobalModelPrices, - }, - { - ID: "0004_mysql_model_price_identity", - Up: migrateMySQLModelPriceIdentity, - }, - { - ID: "0005_request_log_model_consistency", - Up: addRequestLogModelConsistencyColumns, - }, + {ID: "0002_request_log_reasoning", Up: retainLegacyMigrationPosition}, + {ID: "0003_global_model_prices", Up: retainLegacyMigrationPosition}, + {ID: "0004_mysql_model_price_identity", Up: retainLegacyMigrationPosition}, + {ID: "0005_request_log_model_consistency", Up: retainLegacyMigrationPosition}, + {ID: "0006_channel_execution", Up: migrateChannelExecution}, } func applyMigrations(db *gorm.DB) error { diff --git a/internal/storage/models/group.go b/internal/storage/models/group.go index 611737ec..c58fb320 100644 --- a/internal/storage/models/group.go +++ b/internal/storage/models/group.go @@ -1,42 +1,56 @@ package models -// UpstreamKeyStatus is the durable operator-controlled state of a credential. -// Runtime cooldown and failure state belongs to state.KeyRegistry. -type UpstreamKeyStatus string +import ( + "bytes" -const ( - UpstreamKeyStatusActive UpstreamKeyStatus = "active" - UpstreamKeyStatusDisabled UpstreamKeyStatus = "disabled" + "gorm.io/gorm" ) // Group is the persisted configuration for an upstream service group. // API DTOs and runtime state views must be defined outside the storage package. type Group struct { - ID uint `gorm:"primaryKey;autoIncrement"` - Name string `gorm:"type:varchar(255);not null;uniqueIndex"` - ProviderID *string `gorm:"type:varchar(255)"` - UpstreamURL string `gorm:"type:text;not null"` - Protocols JSON `gorm:"type:json;not null"` - Models JSON `gorm:"type:json;not null"` - ConvertEnabled bool `gorm:"not null;default:false"` + ID uint `gorm:"primaryKey;autoIncrement"` + Name string `gorm:"type:varchar(255);not null;uniqueIndex"` + ChannelID string `gorm:"type:varchar(64);not null"` + Params JSON `gorm:"type:json;not null"` + Models JSON `gorm:"type:json;not null"` WeightManual *int - ValidationModel *string `gorm:"type:varchar(255)"` - Config JSON `gorm:"type:json"` - Enabled bool `gorm:"not null;default:true"` - UpstreamKeys []UpstreamKey `gorm:"foreignKey:GroupID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` - CreatedAtMS int64 `gorm:"column:created_at_ms;not null;autoCreateTime:milli;check:chk_group_created_at,created_at_ms >= 0"` - UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null;autoUpdateTime:milli;check:chk_group_updated_at,updated_at_ms >= 0"` + ValidationModel *string `gorm:"type:varchar(255)"` + Overrides JSON `gorm:"type:json"` + Enabled bool `gorm:"not null;default:true"` + Credentials []Credential `gorm:"foreignKey:GroupID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` + CreatedAtMS int64 `gorm:"column:created_at_ms;not null;autoCreateTime:milli;check:chk_group_created_at,created_at_ms >= 0"` + UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null;autoUpdateTime:milli;check:chk_group_updated_at,updated_at_ms >= 0"` } -// UpstreamKey is an encrypted provider credential that belongs to one group. -type UpstreamKey struct { - ID uint `gorm:"primaryKey;autoIncrement"` - GroupID uint `gorm:"not null;uniqueIndex:idx_upstream_keys_group_hash,priority:1"` - KeyValue string `gorm:"type:text;not null"` - KeyHash string `gorm:"type:varchar(128);not null;uniqueIndex:idx_upstream_keys_group_hash,priority:2"` - Status UpstreamKeyStatus `gorm:"type:varchar(32);not null;default:'active';check:chk_upstream_key_status,status IN ('active','disabled')"` +// BeforeSave keeps channel parameters representable. Channel-specific shape +// validation belongs to the code-owned channel registry. +func (group *Group) BeforeSave(_ *gorm.DB) error { + trimmed := bytes.TrimSpace(group.Params) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + group.Params = JSON(`{}`) + } + return nil +} + +// CredentialStatus is the durable operator-controlled state of channel data. +// Runtime cooldown and failure state belongs to the runtime credential registry. +type CredentialStatus string + +const ( + CredentialStatusActive CredentialStatus = "active" + CredentialStatusDisabled CredentialStatus = "disabled" +) + +// Credential is encrypted channel credential data that belongs to one group. +type Credential struct { + ID uint `gorm:"primaryKey;autoIncrement"` + GroupID uint `gorm:"not null;uniqueIndex:idx_credentials_group_fingerprint,priority:1"` + Data string `gorm:"type:text;not null"` + Fingerprint string `gorm:"type:varchar(128);not null;uniqueIndex:idx_credentials_group_fingerprint,priority:2"` + Status CredentialStatus `gorm:"type:varchar(32);not null;default:'active';check:chk_credential_status,status IN ('active','disabled')"` WeightManual *int Group *Group `gorm:"foreignKey:GroupID;references:ID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` - CreatedAtMS int64 `gorm:"column:created_at_ms;not null;autoCreateTime:milli;check:chk_upstream_key_created_at,created_at_ms >= 0"` - UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null;autoUpdateTime:milli;check:chk_upstream_key_updated_at,updated_at_ms >= 0"` + CreatedAtMS int64 `gorm:"column:created_at_ms;not null;autoCreateTime:milli;check:chk_credential_created_at,created_at_ms >= 0"` + UpdatedAtMS int64 `gorm:"column:updated_at_ms;not null;autoUpdateTime:milli;check:chk_credential_updated_at,updated_at_ms >= 0"` } diff --git a/internal/storage/models/model_price.go b/internal/storage/models/model_price.go index 473f7f25..89756817 100644 --- a/internal/storage/models/model_price.go +++ b/internal/storage/models/model_price.go @@ -28,11 +28,11 @@ type contextPriceTierJSON struct { CacheWritePriceNanoUSDPerMillionTokens *int64 `json:"cache_write_price_nano_usd_per_million_tokens"` } -// ModelPrice stores one catalog or manual global price for an exact upstream -// model ID. Routing scope never participates in pricing. +// ModelPrice stores one catalog or manual price for an exact channel/model pair. type ModelPrice struct { ID uint `gorm:"primaryKey;autoIncrement"` - ModelID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_model_prices_model"` + ChannelID string `gorm:"type:varchar(64);not null;uniqueIndex:idx_model_prices_channel_model,priority:1"` + ModelID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_model_prices_channel_model,priority:2"` InputPriceNanoUSDPerMillionTokens *int64 `gorm:"column:input_price_nano_usd_per_million_tokens;check:chk_model_price_input_nano,input_price_nano_usd_per_million_tokens IS NULL OR input_price_nano_usd_per_million_tokens >= 0"` OutputPriceNanoUSDPerMillionTokens *int64 `gorm:"column:output_price_nano_usd_per_million_tokens;check:chk_model_price_output_nano,output_price_nano_usd_per_million_tokens IS NULL OR output_price_nano_usd_per_million_tokens >= 0"` CacheReadPriceNanoUSDPerMillionTokens *int64 `gorm:"column:cache_read_price_nano_usd_per_million_tokens;check:chk_model_price_cache_read_nano,cache_read_price_nano_usd_per_million_tokens IS NULL OR cache_read_price_nano_usd_per_million_tokens >= 0"` diff --git a/internal/storage/models/request.go b/internal/storage/models/request.go index d706661a..a4b11d45 100644 --- a/internal/storage/models/request.go +++ b/internal/storage/models/request.go @@ -6,6 +6,8 @@ type RequestLog struct { CompletedAtMS int64 `gorm:"column:completed_at_ms;not null;check:chk_request_log_completed_at,completed_at_ms >= 0;index:idx_request_logs_completed_id,priority:1,sort:desc;index:idx_request_logs_access_completed_id,priority:2,sort:desc;index:idx_request_logs_status_completed_id,priority:2,sort:desc;index:idx_request_logs_model_completed_id,priority:2,sort:desc;index:idx_request_logs_upstream_model_completed_id,priority:2,sort:desc"` AccessKeyID uint `gorm:"not null;index:idx_request_logs_access_completed_id,priority:1"` GroupID uint `gorm:"not null;default:0"` + ChannelID string `gorm:"type:varchar(64);not null;default:''"` + CredentialID uint `gorm:"not null;default:0"` Protocol string `gorm:"type:varchar(32);not null"` ClientModel string `gorm:"type:varchar(255);not null;index:idx_request_logs_model_completed_id,priority:1"` UpstreamModel string `gorm:"type:varchar(255);not null;index:idx_request_logs_upstream_model_completed_id,priority:1"` @@ -38,24 +40,31 @@ type RequestLog struct { // RequestLogAttempt is one durable upstream attempt belonging to a client // request. Identity fields are snapshots and intentionally do not reference -// mutable Group or Key catalog rows. +// mutable Group or Credential catalog rows. type RequestLogAttempt struct { - RequestID string `gorm:"type:varchar(36);primaryKey;not null;index:idx_request_log_attempts_group_completed_request,priority:3;index:idx_request_log_attempts_key_completed_request,priority:3;index:idx_request_log_attempts_model_completed_request,priority:3;index:idx_request_log_attempts_status_completed_request,priority:3;index:idx_request_log_attempts_failure_completed_request,priority:3;index:idx_request_log_attempts_error_completed_request,priority:3"` - Sequence int `gorm:"primaryKey;not null;check:chk_request_log_attempt_sequence,sequence > 0"` - CompletedAtMS int64 `gorm:"column:completed_at_ms;not null;check:chk_request_log_attempt_completed_at,completed_at_ms >= 0;index:idx_request_log_attempts_group_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_key_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_model_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_status_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_failure_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_error_completed_request,priority:2,sort:desc"` - GroupID uint `gorm:"not null;check:chk_request_log_attempt_group,group_id > 0;index:idx_request_log_attempts_group_completed_request,priority:1"` - GroupName string `gorm:"type:varchar(255);not null"` - KeyID uint `gorm:"not null;check:chk_request_log_attempt_key,key_id > 0;index:idx_request_log_attempts_key_completed_request,priority:1"` - UpstreamModel string `gorm:"type:varchar(255);not null;default:''"` - StatusCode int `gorm:"not null;index:idx_request_log_attempts_status_completed_request,priority:1"` - DurationMs int64 `gorm:"not null;check:chk_request_log_attempt_duration,duration_ms >= 0"` - FailureCategory string `gorm:"type:varchar(32);not null;check:chk_request_log_attempt_failure_category,failure_category IN ('ok','rate_limited','model_unavailable','invalid_key','upstream_host_error','client_error','downstream_cancel','ambiguous');index:idx_request_log_attempts_failure_completed_request,priority:1"` - Action string `gorm:"type:varchar(32);not null;check:chk_request_log_attempt_action,action IN ('terminate','retry','cooldown_key','fail_key','skip_group')"` - WillRetry bool `gorm:"not null;default:false"` - ErrorCode string `gorm:"type:varchar(64);not null;default:'';index:idx_request_log_attempts_error_completed_request,priority:1"` - ErrorSummary string `gorm:"type:text;not null"` - PricingReceipt JSON `gorm:"type:json"` - RequestLog *RequestLog `gorm:"foreignKey:RequestID;references:ID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` + RequestID string `gorm:"type:varchar(36);primaryKey;not null;index:idx_request_log_attempts_group_completed_request,priority:3;index:idx_request_log_attempts_channel_completed_request,priority:3;index:idx_request_log_attempts_credential_completed_request,priority:3;index:idx_request_log_attempts_model_completed_request,priority:3;index:idx_request_log_attempts_status_completed_request,priority:3;index:idx_request_log_attempts_failure_completed_request,priority:3;index:idx_request_log_attempts_error_completed_request,priority:3"` + Sequence int `gorm:"primaryKey;not null;check:chk_request_log_attempt_sequence,sequence > 0"` + CompletedAtMS int64 `gorm:"column:completed_at_ms;not null;check:chk_request_log_attempt_completed_at,completed_at_ms >= 0;index:idx_request_log_attempts_group_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_credential_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_channel_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_model_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_status_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_failure_completed_request,priority:2,sort:desc;index:idx_request_log_attempts_error_completed_request,priority:2,sort:desc"` + GroupID uint `gorm:"not null;check:chk_request_log_attempt_group,group_id > 0;index:idx_request_log_attempts_group_completed_request,priority:1"` + GroupName string `gorm:"type:varchar(255);not null"` + ChannelID string `gorm:"type:varchar(64);not null;default:'';index:idx_request_log_attempts_channel_completed_request,priority:1"` + CredentialID uint `gorm:"not null;check:chk_request_log_attempt_credential,credential_id > 0;index:idx_request_log_attempts_credential_completed_request,priority:1"` + Operation string `gorm:"type:varchar(64);not null;default:''"` + RouteMode string `gorm:"type:varchar(32);not null;default:''"` + UpstreamModel string `gorm:"type:varchar(255);not null;default:''"` + UpstreamRequestID string `gorm:"type:varchar(255);not null;default:''"` + DispatchState string `gorm:"type:varchar(32);not null;default:''"` + ResponseStarted bool `gorm:"not null;default:false"` + StatusCode int `gorm:"not null;index:idx_request_log_attempts_status_completed_request,priority:1"` + DurationMs int64 `gorm:"not null;check:chk_request_log_attempt_duration,duration_ms >= 0"` + FailureCategory string `gorm:"type:varchar(32);not null;check:chk_request_log_attempt_failure_category,failure_category IN ('ok','rate_limited','model_unavailable','invalid_key','upstream_host_error','client_error','downstream_cancel','ambiguous');index:idx_request_log_attempts_failure_completed_request,priority:1"` + Action string `gorm:"type:varchar(32);not null;check:chk_request_log_attempt_action,action IN ('terminate','retry','cooldown_credential','fail_credential','skip_group')"` + WillRetry bool `gorm:"not null;default:false"` + ErrorCode string `gorm:"type:varchar(64);not null;default:'';index:idx_request_log_attempts_error_completed_request,priority:1"` + ErrorSummary string `gorm:"type:text;not null"` + Committed bool `gorm:"not null;default:false"` + PricingReceipt JSON `gorm:"type:json"` + RequestLog *RequestLog `gorm:"foreignKey:RequestID;references:ID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE"` } // UsageAggregationJournal is the request-idempotent input for hourly usage @@ -66,6 +75,8 @@ type UsageAggregationJournal struct { BucketStartMS int64 `gorm:"column:bucket_start_ms;not null;check:chk_usage_journal_bucket, bucket_start_ms >= 0;index:idx_usage_aggregation_journal_pending_bucket,priority:2"` AccessKeyID uint `gorm:"not null"` GroupID uint `gorm:"not null"` + ChannelID string `gorm:"type:varchar(64);not null;default:''"` + CredentialID uint `gorm:"not null;default:0"` Model string `gorm:"type:varchar(255);not null"` RequestCount int64 `gorm:"not null;check:chk_usage_journal_request_count,request_count = 1;check:chk_usage_journal_request_outcome,request_count = success_count + failure_count"` SuccessCount int64 `gorm:"not null;check:chk_usage_journal_success_count,success_count >= 0"` @@ -88,13 +99,16 @@ func (UsageAggregationJournal) TableName() string { return "usage_aggregation_journal" } -// UsageStat is an hourly aggregate by access key, upstream group, and upstream model. +// UsageStat is an hourly aggregate by access key, channel, upstream group, +// credential, and upstream model. type UsageStat struct { ID uint `gorm:"primaryKey;autoIncrement"` - BucketStartMS int64 `gorm:"column:bucket_start_ms;not null;check:chk_usage_stat_bucket,bucket_start_ms >= 0;uniqueIndex:idx_usage_stats_bucket_access_group_model,priority:1"` - AccessKeyID uint `gorm:"not null;uniqueIndex:idx_usage_stats_bucket_access_group_model,priority:2"` - GroupID uint `gorm:"not null;uniqueIndex:idx_usage_stats_bucket_access_group_model,priority:3"` - Model string `gorm:"type:varchar(255);not null;uniqueIndex:idx_usage_stats_bucket_access_group_model,priority:4"` + BucketStartMS int64 `gorm:"column:bucket_start_ms;not null;check:chk_usage_stat_bucket,bucket_start_ms >= 0;uniqueIndex:idx_usage_stats_identity,priority:1"` + AccessKeyID uint `gorm:"not null;uniqueIndex:idx_usage_stats_identity,priority:2"` + ChannelID string `gorm:"type:varchar(64);not null;default:'';uniqueIndex:idx_usage_stats_identity,priority:3"` + GroupID uint `gorm:"not null;uniqueIndex:idx_usage_stats_identity,priority:4"` + CredentialID uint `gorm:"not null;default:0;uniqueIndex:idx_usage_stats_identity,priority:5"` + Model string `gorm:"type:varchar(255);not null;uniqueIndex:idx_usage_stats_identity,priority:6"` RequestCount int64 `gorm:"not null;default:0;check:chk_usage_stat_request_count,request_count >= 0;check:chk_usage_stat_request_outcome,request_count = success_count + failure_count"` SuccessCount int64 `gorm:"not null;default:0;check:chk_usage_stat_success_count,success_count >= 0"` FailureCount int64 `gorm:"not null;default:0;check:chk_usage_stat_failure_count,failure_count >= 0"` diff --git a/internal/telemetry/requestlog.go b/internal/telemetry/requestlog.go index 84d35378..79a8f02f 100644 --- a/internal/telemetry/requestlog.go +++ b/internal/telemetry/requestlog.go @@ -3,6 +3,8 @@ package telemetry import ( "time" + "gpt-load/internal/channel" + "gpt-load/internal/execution" "gpt-load/internal/protocol" "gpt-load/internal/reasoning" "gpt-load/internal/usage" @@ -42,27 +44,33 @@ const ( type Action string const ( - ActionTerminate Action = "terminate" - ActionRetry Action = "retry" - ActionCooldownKey Action = "cooldown_key" - ActionFailKey Action = "fail_key" - ActionSkipGroup Action = "skip_group" + ActionTerminate Action = "terminate" + ActionRetry Action = "retry" + ActionCooldownCredential Action = "cooldown_credential" + ActionFailCredential Action = "fail_credential" + ActionSkipGroup Action = "skip_group" ) type Attempt struct { - Sequence int - GroupID uint - GroupName string - KeyID uint - UpstreamModel string - StatusCode int - DurationMs int64 - FailureCategory FailureCategory - Action Action - WillRetry bool - ErrorCode string - ErrorSummary string - Committed bool + Sequence int + GroupID uint + GroupName string + ChannelID channel.ID + CredentialID uint + Operation execution.Operation + RouteMode channel.RouteMode + UpstreamModel string + UpstreamRequestID string + DispatchState execution.DispatchState + ResponseStarted bool + StatusCode int + DurationMs int64 + FailureCategory FailureCategory + Action Action + WillRetry bool + ErrorCode string + ErrorSummary string + Committed bool } // PricingObservation is the frozen, dependency-neutral quote selected by the @@ -78,7 +86,8 @@ type PricingObservation struct { type UsageObservation struct { Result usage.Result GroupID uint - KeyID uint + ChannelID channel.ID + CredentialID uint AttemptSequence int Pricing PricingObservation } diff --git a/internal/telemetry/requestlog_test.go b/internal/telemetry/requestlog_test.go index 717e2f55..45a34bb5 100644 --- a/internal/telemetry/requestlog_test.go +++ b/internal/telemetry/requestlog_test.go @@ -41,8 +41,14 @@ func TestRequestTelemetryContractUsesExactFieldAllowlist(t *testing.T) { "Sequence", "GroupID", "GroupName", - "KeyID", + "ChannelID", + "CredentialID", + "Operation", + "RouteMode", "UpstreamModel", + "UpstreamRequestID", + "DispatchState", + "ResponseStarted", "StatusCode", "DurationMs", "FailureCategory", @@ -62,7 +68,8 @@ func TestRequestTelemetryContractUsesExactFieldAllowlist(t *testing.T) { reflect.TypeOf(UsageObservation{}): { "Result", "GroupID", - "KeyID", + "ChannelID", + "CredentialID", "AttemptSequence", "Pricing", }, diff --git a/internal/webui/workflow_test.go b/internal/webui/workflow_test.go index 62b160e6..9be464d2 100644 --- a/internal/webui/workflow_test.go +++ b/internal/webui/workflow_test.go @@ -58,7 +58,7 @@ func TestWebCICompositeActionRunsCompleteFrontendGate(t *testing.T) { } } -func TestBranchCIContainsStaticGoAndFreshRaceGates(t *testing.T) { +func TestBranchCIContainsStaticGoAndBackendGates(t *testing.T) { content := readRepositoryFile(t, ".github/workflows/ci.yml") testJob := workflowJobBlock(t, content, "test") workflowGoFormattingScript(t, content, "test") @@ -69,7 +69,7 @@ func TestBranchCIContainsStaticGoAndFreshRaceGates(t *testing.T) { {name: "Check module graph", run: "go mod tidy -diff"}, {name: "Audit Go dependencies", run: "go run golang.org/x/vuln/cmd/govulncheck@v1.1.4 ./..."}, {name: "Run Go vet", run: "go vet ./..."}, - {name: "Run race-enabled tests", run: "go test -race -count=1 . ./internal/..."}, + {name: "Run backend tests", run: "go test -count=1 . ./internal/..."}, {name: "Check repository invariants", run: "git diff --check"}, } { assertWorkflowGateStep(t, testJob, test.name, test.run) @@ -83,8 +83,8 @@ func TestBranchCIContainsStaticGoAndFreshRaceGates(t *testing.T) { assertWorkflowGateStep( t, releaseJob, - "Run race-enabled tests", - "go test -race -count=1 . ./internal/...", + "Run backend tests", + "go test -count=1 . ./internal/...", ) } @@ -1195,7 +1195,7 @@ func TestReleaseWorkflowPostPublishDownloadsExactAssetsAndRunsFiveNativeSmokes(t "SHA256SUMS", "sha256sum --check", "asset_count", - `test "${asset_count}" = "6"`, + `test "${asset_count}" = "10"`, } { if !strings.Contains(inventoryJob, required) { t.Fatalf("post-publish asset inventory does not contain %q:\n%s", required, inventoryJob) @@ -1375,7 +1375,7 @@ func TestReleaseDockerSmokeCanUsePublishedSourceWithoutDeletingIt(t *testing.T) `docker image rm "${image}"`, `model_price_list_path="/api/model-prices?usage=in_use&status=all&page=1&page_size=100"`, `api_write PUT "/api/model-prices/${model_price_id}"`, - `item.scope.kind==="group"&&item.scope.id===groupID`, + `item.channel_id==="openai_compatible"&&item.model_id===modelID`, `item.model_id===modelID`, `!Number.isSafeInteger(matches[0].id)||matches[0].id<=0`, `api_get "${model_price_list_path}" >"${task_tmp}/prices-second.json"`, diff --git a/web/src/api/control/types.ts b/web/src/api/control/types.ts index cbcebbf5..d6f2e839 100644 --- a/web/src/api/control/types.ts +++ b/web/src/api/control/types.ts @@ -13,13 +13,13 @@ export type FailureCategory = | 'ambiguous' export type GroupCollectionStatus = 'available' | 'unavailable' | 'disabled' -export type GroupCollectionSort = 'status' | 'name' | 'keys' | 'created' +export type GroupCollectionSort = 'status' | 'name' | 'credentials' | 'created' export type ModelPricingStatus = 'pending' | 'configured' +export type ChannelParamsDto = Record export interface GroupCollectionFilters { q?: string status?: GroupCollectionStatus - protocol?: GroupProtocol sort: GroupCollectionSort page: number page_size: 20 @@ -35,12 +35,11 @@ export interface GroupCollectionSummaryDto { export interface GroupCollectionItemDto { id: number name: string - provider_id: string | null + channel_id: string + params: ChannelParamsDto status: GroupCollectionStatus - upstream_url: string - protocols: GroupProtocol[] model_count: number - key_counts: KeyCounts + credential_counts: CredentialCounts } export interface GroupCollectionPaginationDto { @@ -60,11 +59,10 @@ export interface GroupCollectionResponseDto { export interface GroupSummaryDto { id: number name: string - provider_id: string | null + channel_id: string + params: ChannelParamsDto service_status: GroupCollectionStatus - upstream_url: string - protocols: GroupProtocol[] - key_count: number + credential_count: number model_count: number } @@ -93,9 +91,8 @@ export interface GroupEffectiveConfigDto { export interface GroupSettingsDto { name: string - provider_id: string | null - upstream_url: string - protocols: GroupProtocol[] + channel_id: string + params: ChannelParamsDto validation_model: string | null enabled: boolean weight_manual: number | null @@ -117,23 +114,23 @@ export interface GroupModelsDto { pending: number } -export type GroupKeyStatus = 'available' | 'cooldown' | 'blacklisted' | 'disabled' -export type GroupKeyConfiguredStatus = 'active' | 'disabled' -export type GroupKeyWeightMode = 'auto' | 'manual' -export type GroupKeyRecoveryMode = 'none' | 'cooldown' | 'probe' | 'manual' +export type CredentialStatus = 'available' | 'cooldown' | 'blacklisted' | 'disabled' +export type CredentialConfiguredStatus = 'active' | 'disabled' +export type CredentialWeightMode = 'auto' | 'manual' +export type CredentialRecoveryMode = 'none' | 'cooldown' | 'probe' | 'manual' -export interface GroupKeyRecoveryDto { - mode: GroupKeyRecoveryMode +export interface CredentialRecoveryDto { + mode: CredentialRecoveryMode automatic: boolean at_ms: number | null } -export interface GroupKeyItemDto { - id: number +export interface CredentialItemDto { + credential_id: number mask: string - configured_status: GroupKeyConfiguredStatus - effective_status: GroupKeyStatus - weight_mode: GroupKeyWeightMode + configured_status: CredentialConfiguredStatus + effective_status: CredentialStatus + weight_mode: CredentialWeightMode weight: number | null recent_success_count: number recent_failure_count: number @@ -141,16 +138,16 @@ export interface GroupKeyItemDto { last_failure_category: FailureCategory last_status_code: number | null cooldown_until_ms: number | null - recovery: GroupKeyRecoveryDto + recovery: CredentialRecoveryDto } -export interface GroupKeyRevealDto { - id: number - key: string +export interface CredentialRevealDto { + credential_id: number + credential: Record revealed_at_ms: number } -export interface GroupKeySummaryDto { +export interface CredentialSummaryDto { total: number available: number cooldown: number @@ -158,42 +155,43 @@ export interface GroupKeySummaryDto { disabled: number } -export interface GroupKeyPaginationDto { +export interface CredentialPaginationDto { page: number page_size: 20 | 50 | 100 total_items: number total_pages: number } -export interface GroupKeyCollectionDto { +export interface CredentialCollectionDto { observed_at_ms: number stats_window_seconds: number - summary: GroupKeySummaryDto - items: GroupKeyItemDto[] - pagination: GroupKeyPaginationDto + summary: CredentialSummaryDto + items: CredentialItemDto[] + pagination: CredentialPaginationDto } -export interface GroupKeyCollectionFilters { +export interface CredentialCollectionFilters { q?: string - status?: GroupKeyStatus + status?: CredentialStatus page: number page_size: 20 | 50 | 100 } -export interface GroupKeyBatchResultDto { - affected_ids: number[] - summary: GroupKeySummaryDto +export interface CredentialBatchResultDto { + affected_credential_ids: number[] + summary: CredentialSummaryDto } export interface GroupOptionDto { id: number name: string + channel_id: string + params: ChannelParamsDto enabled: boolean - protocols: GroupProtocol[] models: string[] } -export interface KeyCounts { +export interface CredentialCounts { total: number available: number cooldown: number @@ -201,11 +199,18 @@ export interface KeyCounts { disabled: number } +export interface HealthCredentialCountsDto { + credentials: number + available: number + cooldown: number + blacklisted: number +} + export interface HealthGroupDto { id: number name: string enabled: boolean - counts: KeyCounts + counts: HealthCredentialCountsDto } export interface HealthRecoveryDto { @@ -214,8 +219,8 @@ export interface HealthRecoveryDto { at_ms: number | null } -export interface HealthProblemKeyDto { - key_id: number +export interface HealthProblemCredentialDto { + credential_id: number group_id: number group_name: string cooldown_until_ms: number | null @@ -254,10 +259,10 @@ export interface RuntimeHealthDto { uptime_seconds: number snapshot_revision: number stats_window_seconds: number - counts: KeyCounts + counts: HealthCredentialCountsDto groups: HealthGroupDto[] - cooldown_keys: HealthProblemKeyDto[] - blacklisted_keys: HealthProblemKeyDto[] + cooldown_credentials: HealthProblemCredentialDto[] + blacklisted_credentials: HealthProblemCredentialDto[] request_log: RequestLogHealthDto } diff --git a/web/src/app/mutation-outcome.ts b/web/src/app/mutation-outcome.ts index 14d354db..d7eec502 100644 --- a/web/src/app/mutation-outcome.ts +++ b/web/src/app/mutation-outcome.ts @@ -1,6 +1,6 @@ import { ApiError, InvalidResponseError, NetworkError, RequestCancelledError } from '@/api/errors' -type OperationKind = 'access_key_create' | 'group_create' | 'group_key_import' +type OperationKind = 'access_key_create' | 'group_create' | 'credential_import' interface IncompleteOperation { operation_id: string @@ -36,7 +36,7 @@ function isRecord(value: unknown): value is Record { } function isOperationKind(value: unknown): value is OperationKind { - return ['access_key_create', 'group_create', 'group_key_import'].includes(String(value)) + return ['access_key_create', 'group_create', 'credential_import'].includes(String(value)) } function incompleteOperation(data: unknown): IncompleteOperation | undefined { diff --git a/web/src/app/query-keys.ts b/web/src/app/query-keys.ts index bd419cb6..6e9c02ba 100644 --- a/web/src/app/query-keys.ts +++ b/web/src/app/query-keys.ts @@ -1,14 +1,13 @@ import type { RequestLogFilters } from '@/app/resources/request-logs' import type { HomeRange } from '@/app/resources/home' import type { UsageFilters } from '@/app/resources/usage' -import type { ProviderModelFilters } from '@/app/resources/providers' import type { ModelPriceFilters } from '@/app/resources/model-prices' import type { ModelCollectionFilters } from '@/app/resources/models' import { normalizeRequestLogFilters } from '@/app/resources/request-log-filters' import type { AccessKeyCollectionFilters, + CredentialCollectionFilters, GroupCollectionFilters, - GroupKeyCollectionFilters, } from '@/api/control/types' export function normalizeGroupCollectionFilters( @@ -22,14 +21,13 @@ export function normalizeGroupCollectionFilters( const query = filters.q?.trim() if (query) normalized.q = query if (filters.status !== undefined) normalized.status = filters.status - if (filters.protocol !== undefined) normalized.protocol = filters.protocol return normalized } -export function normalizeGroupKeyCollectionFilters( - filters: GroupKeyCollectionFilters, -): GroupKeyCollectionFilters { - const normalized: GroupKeyCollectionFilters = { +export function normalizeCredentialCollectionFilters( + filters: CredentialCollectionFilters, +): CredentialCollectionFilters { + const normalized: CredentialCollectionFilters = { page: filters.page, page_size: filters.page_size, } @@ -58,7 +56,9 @@ function normalizeUsageFilters(filters: UsageFilters): UsageFilters { breakdown_order: filters.breakdown_order ?? 'requests', } if (filters.group_id !== undefined) result.group_id = filters.group_id - if (filters.model !== undefined) result.model = filters.model + if (filters.channel_id !== undefined) result.channel_id = filters.channel_id + if (filters.credential_id !== undefined) result.credential_id = filters.credential_id + if (filters.upstream_model !== undefined) result.upstream_model = filters.upstream_model return result } @@ -76,17 +76,21 @@ export const controlQueryKeys = { settings: (id: number) => ['control', 'groups', 'settings', id] as const, modelsAll: () => ['control', 'groups', 'models'] as const, models: (id: number) => ['control', 'groups', 'models', id] as const, - keysAll: (id: number) => ['control', 'groups', 'keys', id] as const, - keys: (id: number, filters: GroupKeyCollectionFilters) => + credentialsAll: (id: number) => ['control', 'groups', 'credentials', id] as const, + credentials: (id: number, filters: CredentialCollectionFilters) => [ 'control', 'groups', - 'keys', + 'credentials', id, 'collection', - normalizeGroupKeyCollectionFilters(filters), + normalizeCredentialCollectionFilters(filters), ] as const, }, + channels: { + all: ['control', 'channels'] as const, + list: (search: string) => ['control', 'channels', 'list', search] as const, + }, health: () => ['control', 'health'] as const, logs: { all: ['control', 'logs'] as const, @@ -118,16 +122,6 @@ export const controlQueryKeys = { ] as const, options: () => ['control', 'access-keys', 'options'] as const, }, - providers: { - all: ['control', 'providers'] as const, - suggestionsAll: () => ['control', 'providers', 'suggestions'] as const, - suggestions: (search: string) => ['control', 'providers', 'suggestions', search] as const, - suggestionsByIDs: (providerIDs: readonly string[]) => - ['control', 'providers', 'suggestions-by-ids', [...providerIDs]] as const, - modelsAll: () => ['control', 'providers', 'models'] as const, - models: (providerID: string, filters: ProviderModelFilters) => - ['control', 'providers', 'models', providerID, filters] as const, - }, settingsAll: ['control', 'settings'] as const, settings: (locale: string) => ['control', 'settings', locale] as const, systemInfo: () => ['control', 'system-info'] as const, diff --git a/web/src/app/resources/channels.ts b/web/src/app/resources/channels.ts new file mode 100644 index 00000000..04b4ed4a --- /dev/null +++ b/web/src/app/resources/channels.ts @@ -0,0 +1,163 @@ +import { keepPreviousData, queryOptions } from '@tanstack/vue-query' +import { computed, toValue, type MaybeRefOrGetter } from 'vue' + +import type { ApiClient } from '@/api/client' +import { enabledDataProtocols } from '@/api/control/protocols' +import type { AccessProtocol } from '@/api/control/types' +import { InvalidResponseError } from '@/api/errors' +import { controlQueryKeys } from '@/app/query-keys' + +import { + assertNoSecretLikeFields, + projectArray, + projectBoolean, + projectEnum, + projectRecord, + projectString, +} from './projector' + +export type ChannelFieldInputKind = 'text' | 'url' | 'secret' + +export interface ChannelFieldDto { + key: string + label: string + input_kind: ChannelFieldInputKind + required: boolean + sensitive: boolean +} + +export interface ChannelDto { + channel_id: string + name: string + mark: string + description: string + param_fields: ChannelFieldDto[] + credential_fields: ChannelFieldDto[] + client_protocols: AccessProtocol[] +} + +export interface ChannelListDto { + items: ChannelDto[] + total: number +} + +const channelFields = [ + 'channel_id', + 'name', + 'mark', + 'description', + 'param_fields', + 'credential_fields', + 'client_protocols', +] as const +const fieldFields = ['key', 'label', 'input_kind', 'required', 'sensitive'] as const +const listFields = ['items', 'total'] as const +const inputKinds = ['text', 'url', 'secret'] as const + +function invalidResponse(): never { + throw new InvalidResponseError() +} + +export function projectChannelID(value: unknown): string { + const result = projectString(value) + if (result !== result.trim() || !/^[a-z][a-z0-9_]*$/u.test(result) || result.length > 100) { + invalidResponse() + } + return result +} + +function projectChannelField(value: unknown): ChannelFieldDto { + const record = projectRecord(value) + assertNoSecretLikeFields(record, fieldFields) + const key = projectString(record.key) + const label = projectString(record.label) + const inputKind = projectEnum(record.input_kind, inputKinds) + const sensitive = projectBoolean(record.sensitive) + if ( + key !== key.trim() || + !/^[a-z][a-z0-9_]*$/u.test(key) || + label.trim().length === 0 || + sensitive !== (inputKind === 'secret') + ) { + invalidResponse() + } + return { + key, + label, + input_kind: inputKind, + required: projectBoolean(record.required), + sensitive, + } +} + +function projectChannel(value: unknown): ChannelDto { + const record = projectRecord(value) + assertNoSecretLikeFields(record, channelFields) + const paramFields = projectArray(record.param_fields, projectChannelField) + const credentialFields = projectArray(record.credential_fields, projectChannelField) + const clientProtocols = projectArray(record.client_protocols, (protocol) => + projectEnum(protocol, enabledDataProtocols), + ) + if ( + new Set(paramFields.map(({ key }) => key)).size !== paramFields.length || + new Set(credentialFields.map(({ key }) => key)).size !== credentialFields.length || + credentialFields.some(({ sensitive }) => !sensitive) || + new Set(clientProtocols).size !== clientProtocols.length + ) { + invalidResponse() + } + return { + channel_id: projectChannelID(record.channel_id), + name: projectString(record.name), + mark: projectString(record.mark, { allowEmpty: true }), + description: projectString(record.description, { allowEmpty: true }), + param_fields: paramFields, + credential_fields: credentialFields, + client_protocols: clientProtocols, + } +} + +export function projectChannelList(value: unknown): ChannelListDto { + const record = projectRecord(value) + assertNoSecretLikeFields(record, listFields) + const items = projectArray(record.items, projectChannel) + const total = Number(record.total) + if ( + !Number.isSafeInteger(total) || + total < 0 || + total !== items.length || + new Set(items.map(({ channel_id }) => channel_id)).size !== items.length + ) { + invalidResponse() + } + return { items, total } +} + +export function normalizeChannelSearch(value: string): string { + return [...value.trim()].slice(0, 200).join('') +} + +export async function listChannels( + client: ApiClient, + search: string, + signal?: AbortSignal, +): Promise { + const normalized = normalizeChannelSearch(search) + const params = new URLSearchParams() + if (normalized) params.set('q', normalized) + const suffix = params.size > 0 ? `?${params.toString()}` : '' + return projectChannelList( + await client.request(`/api/channels${suffix}`, { method: 'GET', signal }), + ) +} + +export function channelsQueryOptions(client: ApiClient, search: MaybeRefOrGetter) { + return queryOptions({ + queryKey: computed(() => + controlQueryKeys.channels.list(normalizeChannelSearch(toValue(search))), + ), + queryFn: ({ queryKey, signal }) => listChannels(client, queryKey[3], signal), + staleTime: 5 * 60 * 1_000, + placeholderData: keepPreviousData, + }) +} diff --git a/web/src/app/resources/upstream-keys.ts b/web/src/app/resources/credentials.ts similarity index 59% rename from web/src/app/resources/upstream-keys.ts rename to web/src/app/resources/credentials.ts index 23c05b8f..34b703ee 100644 --- a/web/src/app/resources/upstream-keys.ts +++ b/web/src/app/resources/credentials.ts @@ -8,18 +8,18 @@ import { computed, toValue, type MaybeRefOrGetter } from 'vue' import type { ApiClient } from '@/api/client' import type { - GroupKeyBatchResultDto, - GroupKeyCollectionDto, - GroupKeyCollectionFilters, - GroupKeyConfiguredStatus, - GroupKeyItemDto, - GroupKeyRecoveryDto, - GroupKeyRevealDto, - GroupKeyStatus, - GroupKeySummaryDto, + CredentialBatchResultDto, + CredentialCollectionDto, + CredentialCollectionFilters, + CredentialConfiguredStatus, + CredentialItemDto, + CredentialRecoveryDto, + CredentialRevealDto, + CredentialStatus, + CredentialSummaryDto, } from '@/api/control/types' import { InvalidResponseError } from '@/api/errors' -import { controlQueryKeys, normalizeGroupKeyCollectionFilters } from '@/app/query-keys' +import { controlQueryKeys, normalizeCredentialCollectionFilters } from '@/app/query-keys' import { assertNoSecretLikeFields, @@ -34,41 +34,44 @@ import { } from './projector' export type { - GroupKeyBatchResultDto, - GroupKeyCollectionDto, - GroupKeyCollectionFilters, - GroupKeyConfiguredStatus, - GroupKeyItemDto, - GroupKeyRecoveryDto, - GroupKeyRevealDto, - GroupKeyStatus, - GroupKeySummaryDto, - GroupKeyWeightMode, + CredentialBatchResultDto, + CredentialCollectionDto, + CredentialCollectionFilters, + CredentialConfiguredStatus, + CredentialItemDto, + CredentialRecoveryDto, + CredentialRevealDto, + CredentialStatus, + CredentialSummaryDto, + CredentialWeightMode, } from '@/api/control/types' -export type UpstreamKeyStatus = GroupKeyConfiguredStatus -export type UpstreamKeyEffectiveStatus = GroupKeyStatus - -export interface UpstreamKeyPatch { - status?: UpstreamKeyStatus +export interface CredentialPatch { + status?: CredentialConfiguredStatus weight_manual?: number | null } -export interface GroupKeyBatchRequest { +export interface CredentialBatchRequest { action: 'enable' | 'disable' | 'delete' - key_ids: number[] + credential_ids: number[] } -const groupKeyCollectionFields = [ +const credentialCollectionFields = [ 'observed_at_ms', 'stats_window_seconds', 'summary', 'items', 'pagination', ] as const -const groupKeySummaryFields = ['total', 'available', 'cooldown', 'blacklisted', 'disabled'] as const -const groupKeyItemFields = [ - 'id', +const credentialSummaryFields = [ + 'total', + 'available', + 'cooldown', + 'blacklisted', + 'disabled', +] as const +const credentialItemFields = [ + 'credential_id', 'mask', 'configured_status', 'effective_status', @@ -82,9 +85,9 @@ const groupKeyItemFields = [ 'cooldown_until_ms', 'recovery', ] as const -const groupKeyRecoveryFields = ['mode', 'automatic', 'at_ms'] as const -const groupKeyPaginationFields = ['page', 'page_size', 'total_items', 'total_pages'] as const -const groupKeyBatchFields = ['affected_ids', 'summary'] as const +const credentialRecoveryFields = ['mode', 'automatic', 'at_ms'] as const +const credentialPaginationFields = ['page', 'page_size', 'total_items', 'total_pages'] as const +const credentialBatchFields = ['affected_credential_ids', 'summary'] as const const configuredStatuses = ['active', 'disabled'] as const const effectiveStatuses = ['available', 'cooldown', 'blacklisted', 'disabled'] as const const weightModes = ['auto', 'manual'] as const @@ -111,9 +114,9 @@ function projectMask(value: unknown): string { return mask } -export function projectGroupKeySummary(value: unknown): GroupKeySummaryDto { +export function projectCredentialSummary(value: unknown): CredentialSummaryDto { const record = projectRecord(value) - assertNoSecretLikeFields(record, groupKeySummaryFields) + assertNoSecretLikeFields(record, credentialSummaryFields) const result = { total: projectSafeInteger(record.total, { minimum: 0 }), available: projectSafeInteger(record.available, { minimum: 0 }), @@ -127,9 +130,9 @@ export function projectGroupKeySummary(value: unknown): GroupKeySummaryDto { return result } -function projectRecovery(value: unknown): GroupKeyRecoveryDto { +function projectRecovery(value: unknown): CredentialRecoveryDto { const record = projectRecord(value) - assertNoSecretLikeFields(record, groupKeyRecoveryFields) + assertNoSecretLikeFields(record, credentialRecoveryFields) const result = { mode: projectEnum(record.mode, recoveryModes), automatic: projectBoolean(record.automatic), @@ -146,9 +149,9 @@ function projectRecovery(value: unknown): GroupKeyRecoveryDto { return result } -export function projectGroupKeyItem(value: unknown): GroupKeyItemDto { +export function projectCredentialItem(value: unknown): CredentialItemDto { const record = projectRecord(value) - assertNoSecretLikeFields(record, groupKeyItemFields) + assertNoSecretLikeFields(record, credentialItemFields) const configuredStatus = projectEnum(record.configured_status, configuredStatuses) const effectiveStatus = projectEnum(record.effective_status, effectiveStatuses) const weightMode = projectEnum(record.weight_mode, weightModes) @@ -157,7 +160,7 @@ export function projectGroupKeyItem(value: unknown): GroupKeyItemDto { const cooldownUntil = projectNullableEpochMilliseconds(record.cooldown_until_ms) const recovery = projectRecovery(record.recovery) if ( - // 分组停用或手动权重为 0 时，active 密钥的运行时状态也会是 disabled。 + // 分组停用或手动权重为 0 时，active 凭据的运行时状态也会是 disabled。 (configuredStatus === 'disabled' && effectiveStatus !== 'disabled') || (effectiveStatus === 'available') !== (weight !== null) || (effectiveStatus === 'cooldown') !== (cooldownUntil !== null) || @@ -166,7 +169,7 @@ export function projectGroupKeyItem(value: unknown): GroupKeyItemDto { invalidResponse() } return { - id: projectSafeInteger(record.id, { minimum: 1 }), + credential_id: projectSafeInteger(record.credential_id, { minimum: 1 }), mask: projectMask(record.mask), configured_status: configuredStatus, effective_status: effectiveStatus, @@ -185,9 +188,9 @@ export function projectGroupKeyItem(value: unknown): GroupKeyItemDto { } } -function projectPagination(value: unknown): GroupKeyCollectionDto['pagination'] { +function projectPagination(value: unknown): CredentialCollectionDto['pagination'] { const record = projectRecord(value) - assertNoSecretLikeFields(record, groupKeyPaginationFields) + assertNoSecretLikeFields(record, credentialPaginationFields) const pageSize = projectSafeInteger(record.page_size, { minimum: 20, maximum: 100 }) if (pageSize !== 20 && pageSize !== 50 && pageSize !== 100) invalidResponse() return { @@ -202,24 +205,24 @@ function expectedPageCount(totalItems: number, pageSize: number): number { return totalItems === 0 ? 0 : Math.ceil(totalItems / pageSize) } -function expectedPageItems(pagination: GroupKeyCollectionDto['pagination']): number { +function expectedPageItems(pagination: CredentialCollectionDto['pagination']): number { if (pagination.total_items === 0 || pagination.page > pagination.total_pages) return 0 if (pagination.page < pagination.total_pages) return pagination.page_size const finalPageItems = pagination.total_items % pagination.page_size return finalPageItems === 0 ? pagination.page_size : finalPageItems } -export function projectGroupKeyCollection(value: unknown): GroupKeyCollectionDto { +export function projectCredentialCollection(value: unknown): CredentialCollectionDto { const record = projectRecord(value) - assertNoSecretLikeFields(record, groupKeyCollectionFields) - const items = projectArray(record.items, projectGroupKeyItem) - const summary = projectGroupKeySummary(record.summary) + assertNoSecretLikeFields(record, credentialCollectionFields) + const items = projectArray(record.items, projectCredentialItem) + const summary = projectCredentialSummary(record.summary) const pagination = projectPagination(record.pagination) if ( pagination.total_items > summary.total || pagination.total_pages !== expectedPageCount(pagination.total_items, pagination.page_size) || items.length !== expectedPageItems(pagination) || - new Set(items.map(({ id }) => id)).size !== items.length + new Set(items.map(({ credential_id }) => credential_id)).size !== items.length ) { invalidResponse() } @@ -232,12 +235,12 @@ export function projectGroupKeyCollection(value: unknown): GroupKeyCollectionDto } } -function normalizePatch(patch: UpstreamKeyPatch): UpstreamKeyPatch { +function normalizePatch(patch: CredentialPatch): CredentialPatch { const keys = Object.keys(patch) if (keys.length === 0 || keys.some((key) => key !== 'status' && key !== 'weight_manual')) { - throw new Error('INVALID_UPSTREAM_KEY_PATCH') + throw new Error('INVALID_CREDENTIAL_PATCH') } - const body: UpstreamKeyPatch = {} + const body: CredentialPatch = {} if (Object.prototype.hasOwnProperty.call(patch, 'status')) { body.status = projectEnum(patch.status, configuredStatuses) } @@ -247,36 +250,36 @@ function normalizePatch(patch: UpstreamKeyPatch): UpstreamKeyPatch { weight === undefined || (weight !== null && (!Number.isInteger(weight) || weight < 1 || weight > 100)) ) { - throw new Error('INVALID_UPSTREAM_KEY_WEIGHT') + throw new Error('INVALID_CREDENTIAL_WEIGHT') } body.weight_manual = weight } return body } -function groupKeyCollectionURL( +function credentialCollectionURL( groupId: number, - filters: GroupKeyCollectionFilters, + filters: CredentialCollectionFilters, ): `/api/${string}` { - const normalized = normalizeGroupKeyCollectionFilters(filters) + const normalized = normalizeCredentialCollectionFilters(filters) const params = new URLSearchParams({ page: String(normalized.page), page_size: String(normalized.page_size), }) if (normalized.q !== undefined) params.set('q', normalized.q) if (normalized.status !== undefined) params.set('status', normalized.status) - return `/api/groups/${groupId}/keys?${params.toString()}` + return `/api/groups/${groupId}/credentials?${params.toString()}` } -export async function getGroupKeyCollection( +export async function getCredentialCollection( client: ApiClient, groupId: number, - filters: GroupKeyCollectionFilters, + filters: CredentialCollectionFilters, signal?: AbortSignal, -): Promise { - const normalized = normalizeGroupKeyCollectionFilters(filters) - const result = projectGroupKeyCollection( - await client.request(groupKeyCollectionURL(groupId, normalized), { method: 'GET', signal }), +): Promise { + const normalized = normalizeCredentialCollectionFilters(filters) + const result = projectCredentialCollection( + await client.request(credentialCollectionURL(groupId, normalized), { method: 'GET', signal }), ) if ( result.pagination.page !== normalized.page || @@ -292,29 +295,31 @@ const manualGroupQueryOptions = { refetchOnReconnect: false, } as const -export function groupKeyCollectionQueryOptions( +export function credentialCollectionQueryOptions( client: ApiClient, groupID: MaybeRefOrGetter, - filters: MaybeRefOrGetter, + filters: MaybeRefOrGetter, ) { return queryOptions({ ...manualGroupQueryOptions, - queryKey: computed(() => controlQueryKeys.groups.keys(toValue(groupID), toValue(filters))), + queryKey: computed(() => + controlQueryKeys.groups.credentials(toValue(groupID), toValue(filters)), + ), queryFn: ({ queryKey, signal }) => - getGroupKeyCollection(client, queryKey[3], queryKey[5], signal), + getCredentialCollection(client, queryKey[3], queryKey[5], signal), placeholderData: keepPreviousData, }) } -export async function updateGroupKey( +export async function updateCredential( client: ApiClient, groupId: number, - keyId: number, - patch: UpstreamKeyPatch, + credentialId: number, + patch: CredentialPatch, signal?: AbortSignal, -): Promise { - return projectGroupKeyItem( - await client.request(`/api/groups/${groupId}/keys/${keyId}`, { +): Promise { + return projectCredentialItem( + await client.request(`/api/groups/${groupId}/credentials/${credentialId}`, { method: 'PUT', json: normalizePatch(patch), signal, @@ -322,35 +327,42 @@ export async function updateGroupKey( ) } -export async function revealGroupKey( +export async function revealCredential( client: ApiClient, groupId: number, - keyId: number, + credentialId: number, signal?: AbortSignal, -): Promise { +): Promise { const record = projectRecord( - await client.request(`/api/groups/${groupId}/keys/${keyId}/reveal`, { + await client.request(`/api/groups/${groupId}/credentials/${credentialId}/reveal`, { method: 'POST', signal, }), ) - assertNoSecretLikeFields(record, ['id', 'key', 'revealed_at_ms']) - if (projectSafeInteger(record.id, { minimum: 1 }) !== keyId) invalidResponse() + assertNoSecretLikeFields(record, ['credential_id', 'credential', 'revealed_at_ms']) + if (projectSafeInteger(record.credential_id, { minimum: 1 }) !== credentialId) invalidResponse() + const credentialRecord = projectRecord(record.credential) + const credential: Record = {} + for (const [key, value] of Object.entries(credentialRecord)) { + if (key !== key.trim() || !/^[a-z][a-z0-9_]*$/u.test(key)) invalidResponse() + credential[key] = projectString(value) + } + if (Object.keys(credential).length === 0) invalidResponse() return { - id: keyId, - key: projectString(record.key), + credential_id: credentialId, + credential, revealed_at_ms: projectEpochMilliseconds(record.revealed_at_ms), } } -export async function restoreGroupKey( +export async function restoreCredential( client: ApiClient, groupId: number, - keyId: number, + credentialId: number, signal?: AbortSignal, -): Promise { - return projectGroupKeyItem( - await client.request(`/api/groups/${groupId}/keys/${keyId}/restore`, { +): Promise { + return projectCredentialItem( + await client.request(`/api/groups/${groupId}/credentials/${credentialId}/restore`, { method: 'POST', json: {}, signal, @@ -358,13 +370,13 @@ export async function restoreGroupKey( ) } -export async function batchGroupKeys( +export async function batchCredentials( client: ApiClient, groupId: number, - body: GroupKeyBatchRequest, + body: CredentialBatchRequest, signal?: AbortSignal, -): Promise { - const ids = body.key_ids +): Promise { + const ids = body.credential_ids if ( !['enable', 'disable', 'delete'].includes(body.action) || ids.length < 1 || @@ -372,34 +384,43 @@ export async function batchGroupKeys( ids.some((id) => !Number.isSafeInteger(id) || id < 1) || new Set(ids).size !== ids.length ) { - throw new Error('INVALID_GROUP_KEY_BATCH') + throw new Error('INVALID_CREDENTIAL_BATCH') } const record = projectRecord( - await client.request(`/api/groups/${groupId}/keys/batch`, { + await client.request(`/api/groups/${groupId}/credentials/batch`, { method: 'POST', json: body, signal, }), ) - assertNoSecretLikeFields(record, groupKeyBatchFields) - const affectedIDs = projectArray(record.affected_ids, (id) => + assertNoSecretLikeFields(record, credentialBatchFields) + const affectedCredentialIDs = projectArray(record.affected_credential_ids, (id) => projectSafeInteger(id, { minimum: 1 }), ) - if (new Set(affectedIDs).size !== affectedIDs.length || affectedIDs.length !== ids.length) + if ( + new Set(affectedCredentialIDs).size !== affectedCredentialIDs.length || + affectedCredentialIDs.length !== ids.length + ) invalidResponse() - return { affected_ids: affectedIDs, summary: projectGroupKeySummary(record.summary) } + return { + affected_credential_ids: affectedCredentialIDs, + summary: projectCredentialSummary(record.summary), + } } -export async function deleteGroupKey( +export async function deleteCredential( client: ApiClient, groupId: number, - keyId: number, + credentialId: number, signal?: AbortSignal, ): Promise { - await client.request(`/api/groups/${groupId}/keys/${keyId}`, { method: 'DELETE', signal }) + await client.request(`/api/groups/${groupId}/credentials/${credentialId}`, { + method: 'DELETE', + signal, + }) } -function queryFilters(queryKey: QueryKey): GroupKeyCollectionFilters | undefined { +function queryFilters(queryKey: QueryKey): CredentialCollectionFilters | undefined { const filters = queryKey[5] if (typeof filters !== 'object' || filters === null || Array.isArray(filters)) return undefined const record = filters as Record @@ -407,7 +428,7 @@ function queryFilters(queryKey: QueryKey): GroupKeyCollectionFilters | undefined !Number.isSafeInteger(record.page) || (record.page_size !== 20 && record.page_size !== 50 && record.page_size !== 100) || (record.q !== undefined && typeof record.q !== 'string') || - (record.status !== undefined && !effectiveStatuses.includes(record.status as GroupKeyStatus)) + (record.status !== undefined && !effectiveStatuses.includes(record.status as CredentialStatus)) ) { return undefined } @@ -415,20 +436,20 @@ function queryFilters(queryKey: QueryKey): GroupKeyCollectionFilters | undefined page: record.page as number, page_size: record.page_size as 20 | 50 | 100, ...(record.q === undefined ? {} : { q: record.q }), - ...(record.status === undefined ? {} : { status: record.status as GroupKeyStatus }), + ...(record.status === undefined ? {} : { status: record.status as CredentialStatus }), } } -function matchesFilters(item: GroupKeyItemDto, filters: GroupKeyCollectionFilters): boolean { +function matchesFilters(item: CredentialItemDto, filters: CredentialCollectionFilters): boolean { if (filters.status !== undefined && item.effective_status !== filters.status) return false return filters.q === undefined || item.mask.toLowerCase().includes(filters.q.toLowerCase()) } function withSummaryDelta( - summary: GroupKeySummaryDto, - previous: GroupKeyItemDto, - next: GroupKeyItemDto | undefined, -): GroupKeySummaryDto { + summary: CredentialSummaryDto, + previous: CredentialItemDto, + next: CredentialItemDto | undefined, +): CredentialSummaryDto { const result = { ...summary } result[previous.effective_status]-- if (next !== undefined) result[next.effective_status]++ @@ -436,11 +457,11 @@ function withSummaryDelta( } function withRemovedItem( - collection: GroupKeyCollectionDto, + collection: CredentialCollectionDto, id: number, - summary: GroupKeySummaryDto, -): GroupKeyCollectionDto { - const items = collection.items.filter((item) => item.id !== id) + summary: CredentialSummaryDto, +): CredentialCollectionDto { + const items = collection.items.filter((item) => item.credential_id !== id) const totalItems = collection.pagination.total_items - 1 return { ...collection, @@ -454,17 +475,20 @@ function withRemovedItem( } } -async function invalidateExactKeyPage(queryClient: QueryClient, queryKey: QueryKey): Promise { +async function invalidateExactCredentialPage( + queryClient: QueryClient, + queryKey: QueryKey, +): Promise { await queryClient.invalidateQueries({ queryKey, exact: true, refetchType: 'none' }) } -interface MaterializedKeyPage { +interface MaterializedCredentialPage { queryKey: QueryKey - collection: GroupKeyCollectionDto - filters: GroupKeyCollectionFilters + collection: CredentialCollectionDto + filters: CredentialCollectionFilters } -function keyFilterSetID(filters: GroupKeyCollectionFilters): string { +function credentialFilterSetID(filters: CredentialCollectionFilters): string { return JSON.stringify({ q: filters.q ?? null, status: filters.status ?? null, @@ -473,9 +497,9 @@ function keyFilterSetID(filters: GroupKeyCollectionFilters): string { } function totalItemsAfterBatchDelete( - pages: MaterializedKeyPage[], + pages: MaterializedCredentialPage[], knownDeletedIDs: Set, - summary: GroupKeySummaryDto, + summary: CredentialSummaryDto, ): number { const { q, status } = pages[0].filters if (q === undefined) return status === undefined ? summary.total : summary[status] @@ -490,75 +514,82 @@ function totalItemsAfterBatchDelete( * Reconciles only a cached page containing the previous item. Pages whose * membership or sort position cannot be known are explicitly marked stale. */ -export async function cacheGroupKeyItem( +export async function cacheCredentialItem( queryClient: QueryClient, groupId: number, - item: GroupKeyItemDto, + item: CredentialItemDto, ): Promise { const queries = queryClient .getQueryCache() - .findAll({ queryKey: controlQueryKeys.groups.keysAll(groupId) }) + .findAll({ queryKey: controlQueryKeys.groups.credentialsAll(groupId) }) const previous = queries - .map((query) => (query.state.data as GroupKeyCollectionDto | undefined)?.items) + .map((query) => (query.state.data as CredentialCollectionDto | undefined)?.items) .flatMap((items) => items ?? []) - .find(({ id }) => id === item.id) + .find(({ credential_id }) => credential_id === item.credential_id) for (const query of queries) { const filters = queryFilters(query.queryKey) - const collection = query.state.data as GroupKeyCollectionDto | undefined - const current = collection?.items.find(({ id }) => id === item.id) + const collection = query.state.data as CredentialCollectionDto | undefined + const current = collection?.items.find( + ({ credential_id }) => credential_id === item.credential_id, + ) if (collection === undefined || previous === undefined) { - await invalidateExactKeyPage(queryClient, query.queryKey) + await invalidateExactCredentialPage(queryClient, query.queryKey) continue } const summary = withSummaryDelta(collection.summary, previous, item) if (filters === undefined || current === undefined) { - queryClient.setQueryData(query.queryKey, { ...collection, summary }) - await invalidateExactKeyPage(queryClient, query.queryKey) + queryClient.setQueryData(query.queryKey, { ...collection, summary }) + await invalidateExactCredentialPage(queryClient, query.queryKey) continue } const nextMatches = matchesFilters(item, filters) if (!nextMatches) { - queryClient.setQueryData(query.queryKey, withRemovedItem(collection, item.id, summary)) - await invalidateExactKeyPage(queryClient, query.queryKey) + queryClient.setQueryData( + query.queryKey, + withRemovedItem(collection, item.credential_id, summary), + ) + await invalidateExactCredentialPage(queryClient, query.queryKey) continue } if (current.effective_status !== item.effective_status) { - queryClient.setQueryData(query.queryKey, { ...collection, summary }) - await invalidateExactKeyPage(queryClient, query.queryKey) + queryClient.setQueryData(query.queryKey, { ...collection, summary }) + await invalidateExactCredentialPage(queryClient, query.queryKey) continue } - queryClient.setQueryData(query.queryKey, { + queryClient.setQueryData(query.queryKey, { ...collection, summary, - items: collection.items.map((current) => (current.id === item.id ? item : current)), + items: collection.items.map((current) => + current.credential_id === item.credential_id ? item : current, + ), }) } } -/** Marks only this Group's key pages stale without scheduling an automatic request. */ -export async function invalidateGroupKeyCollections( +/** Marks only this Group's credential pages stale without scheduling an automatic request. */ +export async function invalidateCredentialCollections( queryClient: QueryClient, groupId: number, ): Promise { await queryClient.invalidateQueries({ - queryKey: controlQueryKeys.groups.keysAll(groupId), + queryKey: controlQueryKeys.groups.credentialsAll(groupId), refetchType: 'none', }) } /** A batch result carries the authoritative aggregate, so cached pages can be reconciled deterministically. */ -export async function cacheGroupKeyBatch( +export async function cacheCredentialBatch( queryClient: QueryClient, groupId: number, - action: GroupKeyBatchRequest['action'], - result: GroupKeyBatchResultDto, + action: CredentialBatchRequest['action'], + result: CredentialBatchResultDto, ): Promise { - const affected = new Set(result.affected_ids) + const affected = new Set(result.affected_credential_ids) const queries = queryClient .getQueryCache() - .findAll({ queryKey: controlQueryKeys.groups.keysAll(groupId) }) - const pages = queries.flatMap((query): MaterializedKeyPage[] => { - const collection = query.state.data as GroupKeyCollectionDto | undefined + .findAll({ queryKey: controlQueryKeys.groups.credentialsAll(groupId) }) + const pages = queries.flatMap((query): MaterializedCredentialPage[] => { + const collection = query.state.data as CredentialCollectionDto | undefined const filters = queryFilters(query.queryKey) return collection === undefined || filters === undefined ? [] @@ -566,18 +597,18 @@ export async function cacheGroupKeyBatch( }) if (action !== 'delete') { for (const { queryKey, collection } of pages) { - queryClient.setQueryData(queryKey, { + queryClient.setQueryData(queryKey, { ...collection, summary: result.summary, }) - await invalidateExactKeyPage(queryClient, queryKey) + await invalidateExactCredentialPage(queryClient, queryKey) } return } - const pageSets = new Map() + const pageSets = new Map() for (const page of pages) { - const id = keyFilterSetID(page.filters) + const id = credentialFilterSetID(page.filters) const set = pageSets.get(id) if (set === undefined) pageSets.set(id, [page]) else set.push(page) @@ -586,23 +617,23 @@ export async function cacheGroupKeyBatch( const knownDeletedIDs = new Set( pageSet .flatMap(({ collection }) => collection.items) - .filter(({ id }) => affected.has(id)) - .map(({ id }) => id), + .filter(({ credential_id }) => affected.has(credential_id)) + .map(({ credential_id }) => credential_id), ) const totalItems = totalItemsAfterBatchDelete(pageSet, knownDeletedIDs, result.summary) const totalPages = expectedPageCount(totalItems, pageSet[0].filters.page_size) for (const { queryKey, collection } of pageSet) { - queryClient.setQueryData(queryKey, { + queryClient.setQueryData(queryKey, { ...collection, summary: result.summary, - items: collection.items.filter(({ id }) => !affected.has(id)), + items: collection.items.filter(({ credential_id }) => !affected.has(credential_id)), pagination: { ...collection.pagination, total_items: totalItems, total_pages: totalPages, }, }) - await invalidateExactKeyPage(queryClient, queryKey) + await invalidateExactCredentialPage(queryClient, queryKey) } } } diff --git a/web/src/app/resources/groups.ts b/web/src/app/resources/groups.ts index 30b19a27..e56d0833 100644 --- a/web/src/app/resources/groups.ts +++ b/web/src/app/resources/groups.ts @@ -2,8 +2,9 @@ import { keepPreviousData, queryOptions, type QueryClient } from '@tanstack/vue- import { computed, toValue, type MaybeRefOrGetter } from 'vue' import type { ApiClient } from '@/api/client' -import { enabledDataProtocols } from '@/api/control/protocols' import type { + ChannelParamsDto, + CredentialCounts, GroupCollectionFilters, GroupCollectionItemDto, GroupCollectionPaginationDto, @@ -13,18 +14,13 @@ import type { GroupModelItemDto, GroupModelsDto, GroupOptionDto, - GroupProtocol, GroupSettingsDto, GroupSummaryDto, - KeyCounts, } from '@/api/control/types' import { InvalidResponseError } from '@/api/errors' import { controlQueryKeys, normalizeGroupCollectionFilters } from '@/app/query-keys' -import { - projectModelCandidate, - projectProviderID, - type ModelCandidate, -} from '@/app/resources/providers' +import { projectModelCandidate, type ModelCandidate } from '@/app/resources/providers' +import { projectChannelID } from '@/app/resources/channels' import { assertNoSecretLikeFields, @@ -32,7 +28,6 @@ import { projectBoolean, projectEpochMilliseconds, projectEnum, - projectHTTPURL, projectRecord, projectSafeInteger, projectString, @@ -41,18 +36,16 @@ import { const groupSummaryFields = [ 'id', 'name', - 'provider_id', + 'channel_id', + 'params', 'service_status', - 'upstream_url', - 'protocols', - 'key_count', + 'credential_count', 'model_count', ] as const const groupSettingsFields = [ 'name', - 'provider_id', - 'upstream_url', - 'protocols', + 'channel_id', + 'params', 'validation_model', 'enabled', 'weight_manual', @@ -72,16 +65,15 @@ const groupCollectionSummaryFields = ['total', 'available', 'unavailable', 'disa const groupCollectionItemFields = [ 'id', 'name', - 'provider_id', + 'channel_id', + 'params', 'status', - 'upstream_url', - 'protocols', 'model_count', - 'key_counts', + 'credential_counts', ] as const const groupCollectionPaginationFields = ['page', 'page_size', 'total_items', 'total_pages'] as const -const groupOptionFields = ['id', 'name', 'enabled', 'protocols', 'models'] as const -const keyCountFields = ['total', 'available', 'cooldown', 'blacklisted', 'disabled'] as const +const groupOptionFields = ['id', 'name', 'channel_id', 'params', 'enabled', 'models'] as const +const credentialCountFields = ['total', 'available', 'cooldown', 'blacklisted', 'disabled'] as const const groupCollectionStatuses = ['available', 'unavailable', 'disabled'] as const const runtimeSettingFields = [ 'connect_timeout', @@ -124,9 +116,7 @@ export type { export type GroupSettingsUpdateRequest = Partial<{ name: string - provider_id: string | null - upstream_url: string - protocols: GroupProtocol[] + params: ChannelParamsDto validation_model: string | null enabled: boolean weight_manual: number | null @@ -143,10 +133,9 @@ export interface GroupInUseData { } export interface ModelDiscoveryRequest { - provider_id: string | null - upstream_url: string - protocols: readonly GroupProtocol[] - keys: string + channel_id: string + params: ChannelParamsDto + credentials: string } export interface ModelDiscoveryResult { @@ -165,32 +154,31 @@ export interface GroupModelsReplaceRequest { export interface GroupCreateRequest { name?: string - provider_id: string | null - upstream_url: string - protocols: readonly GroupProtocol[] + channel_id: string + params: ChannelParamsDto models: GroupModelUpdateDto[] - keys: string - confirm_same_upstream_url: boolean + credentials: string + confirm_same_target: boolean } export interface GroupCreateResult { group_id: number group_name: string - keys_added: number - keys_duplicated: number + credentials_added: number + credentials_duplicated: number } -export interface GroupKeyImportRequest { - keys: string +export interface CredentialImportRequest { + credentials: string } -export interface GroupKeyImportResult { +export interface CredentialImportResult { group_id: number - keys_added: number - keys_duplicated: number + credentials_added: number + credentials_duplicated: number } -export interface UpstreamUrlConflictData { +export interface SameTargetConflictData { groups: Array<{ id: number; name: string }> } @@ -200,8 +188,14 @@ function projectNonBlankString(value: unknown): string { return result } -function projectNullableProviderID(value: unknown): string | null { - return value === null ? null : projectProviderID(value) +function projectChannelParams(value: unknown): ChannelParamsDto { + const record = projectRecord(value) + const params: ChannelParamsDto = {} + for (const [key, paramValue] of Object.entries(record)) { + if (key !== key.trim() || !/^[a-z][a-z0-9_]*$/u.test(key)) throw new InvalidResponseError() + params[key] = projectString(paramValue) + } + return params } function projectHeaderRules(value: unknown): HeaderRulesDto { @@ -251,18 +245,13 @@ function projectRuntimeConfig( export function projectGroupSummary(value: unknown): GroupSummaryDto { const record = projectRecord(value) assertNoSecretLikeFields(record, groupSummaryFields) - const protocols = projectArray(record.protocols, (protocol) => - projectEnum(protocol, enabledDataProtocols), - ) - if (new Set(protocols).size !== protocols.length) throw new InvalidResponseError() return { id: projectSafeInteger(record.id, { minimum: 1 }), name: projectNonBlankString(record.name), - provider_id: projectNullableProviderID(record.provider_id), + channel_id: projectChannelID(record.channel_id), + params: projectChannelParams(record.params), service_status: projectEnum(record.service_status, groupCollectionStatuses), - upstream_url: projectHTTPURL(record.upstream_url), - protocols, - key_count: projectSafeInteger(record.key_count, { minimum: 0 }), + credential_count: projectSafeInteger(record.credential_count, { minimum: 0 }), model_count: projectSafeInteger(record.model_count, { minimum: 0 }), } } @@ -270,15 +259,10 @@ export function projectGroupSummary(value: unknown): GroupSummaryDto { export function projectGroupSettings(value: unknown): GroupSettingsDto { const record = projectRecord(value) assertNoSecretLikeFields(record, groupSettingsFields) - const protocols = projectArray(record.protocols, (protocol) => - projectEnum(protocol, enabledDataProtocols), - ) - if (new Set(protocols).size !== protocols.length) throw new InvalidResponseError() return { name: projectNonBlankString(record.name), - provider_id: projectNullableProviderID(record.provider_id), - upstream_url: projectHTTPURL(record.upstream_url), - protocols, + channel_id: projectChannelID(record.channel_id), + params: projectChannelParams(record.params), validation_model: record.validation_model === null ? null : projectNonBlankString(record.validation_model), enabled: projectBoolean(record.enabled), @@ -327,9 +311,9 @@ export function projectGroupModels(value: unknown): GroupModelsDto { return { items, total, pending } } -function projectKeyCounts(value: unknown): KeyCounts { +function projectCredentialCounts(value: unknown): CredentialCounts { const record = projectRecord(value) - assertNoSecretLikeFields(record, keyCountFields) + assertNoSecretLikeFields(record, credentialCountFields) const result = { total: projectSafeInteger(record.total, { minimum: 0 }), available: projectSafeInteger(record.available, { minimum: 0 }), @@ -361,31 +345,20 @@ function projectGroupCollectionSummary(value: unknown): GroupCollectionSummaryDt function projectGroupCollectionItem(value: unknown): GroupCollectionItemDto { const record = projectRecord(value) assertNoSecretLikeFields(record, groupCollectionItemFields) - const protocols = projectArray(record.protocols, (protocol) => - projectEnum(protocol, enabledDataProtocols), - ) - if (new Set(protocols).size !== protocols.length) throw new InvalidResponseError() const status = projectEnum(record.status, groupCollectionStatuses) as GroupCollectionStatus - const keyCounts = projectKeyCounts(record.key_counts) + const credentialCounts = projectCredentialCounts(record.credential_counts) const modelCount = projectSafeInteger(record.model_count, { minimum: 0 }) - const hasRoutableRequest = modelCount > 0 || protocols.includes('openai-responses') - const routeAvailable = keyCounts.available > 0 && hasRoutableRequest - if ( - (status === 'available' && !routeAvailable) || - (status === 'unavailable' && routeAvailable) || - (status === 'disabled' && keyCounts.disabled !== keyCounts.total) - ) { + if (status === 'disabled' && credentialCounts.disabled !== credentialCounts.total) { throw new InvalidResponseError() } return { id: projectSafeInteger(record.id, { minimum: 1 }), name: projectNonBlankString(record.name), - provider_id: projectNullableProviderID(record.provider_id), + channel_id: projectChannelID(record.channel_id), + params: projectChannelParams(record.params), status, - upstream_url: projectHTTPURL(record.upstream_url), - protocols, model_count: modelCount, - key_counts: keyCounts, + credential_counts: credentialCounts, } } @@ -439,19 +412,14 @@ export function projectGroupCollection(value: unknown): GroupCollectionResponseD function projectGroupOption(value: unknown): GroupOptionDto { const record = projectRecord(value) assertNoSecretLikeFields(record, groupOptionFields) - const protocols = projectArray(record.protocols, (protocol) => - projectEnum(protocol, enabledDataProtocols), - ) - if (protocols.length === 0 || new Set(protocols).size !== protocols.length) { - throw new InvalidResponseError() - } const models = projectArray(record.models, projectNonBlankString) if (new Set(models).size !== models.length) throw new InvalidResponseError() return { id: projectSafeInteger(record.id, { minimum: 1 }), name: projectNonBlankString(record.name), + channel_id: projectChannelID(record.channel_id), + params: projectChannelParams(record.params), enabled: projectBoolean(record.enabled), - protocols, models, } } @@ -474,22 +442,27 @@ function projectDiscoveryResult(value: unknown): ModelDiscoveryResult { function projectGroupCreateResult(value: unknown): GroupCreateResult { const record = projectRecord(value) - assertNoSecretLikeFields(record, ['group_id', 'group_name', 'keys_added', 'keys_duplicated']) + assertNoSecretLikeFields(record, [ + 'group_id', + 'group_name', + 'credentials_added', + 'credentials_duplicated', + ]) return { group_id: projectSafeInteger(record.group_id, { minimum: 1 }), group_name: projectNonBlankString(record.group_name), - keys_added: projectSafeInteger(record.keys_added, { minimum: 0 }), - keys_duplicated: projectSafeInteger(record.keys_duplicated, { minimum: 0 }), + credentials_added: projectSafeInteger(record.credentials_added, { minimum: 0 }), + credentials_duplicated: projectSafeInteger(record.credentials_duplicated, { minimum: 0 }), } } -function projectGroupKeyImportResult(value: unknown): GroupKeyImportResult { +function projectCredentialImportResult(value: unknown): CredentialImportResult { const record = projectRecord(value) - assertNoSecretLikeFields(record, ['group_id', 'keys_added', 'keys_duplicated']) + assertNoSecretLikeFields(record, ['group_id', 'credentials_added', 'credentials_duplicated']) return { group_id: projectSafeInteger(record.group_id, { minimum: 1 }), - keys_added: projectSafeInteger(record.keys_added, { minimum: 0 }), - keys_duplicated: projectSafeInteger(record.keys_duplicated, { minimum: 0 }), + credentials_added: projectSafeInteger(record.credentials_added, { minimum: 0 }), + credentials_duplicated: projectSafeInteger(record.credentials_duplicated, { minimum: 0 }), } } @@ -502,7 +475,7 @@ function projectIDName(value: unknown): { id: number; name: string } { } } -export function isUpstreamUrlConflictData(value: unknown): value is UpstreamUrlConflictData { +export function isSameTargetConflictData(value: unknown): value is SameTargetConflictData { try { const record = projectRecord(value) const groups = projectArray(record.groups, projectIDName) @@ -534,7 +507,6 @@ export async function listGroupCollection( params.set('page_size', String(normalized.page_size)) if (normalized.q !== undefined) params.set('q', normalized.q) if (normalized.status !== undefined) params.set('status', normalized.status) - if (normalized.protocol !== undefined) params.set('protocol', normalized.protocol) const result = projectGroupCollection( await client.request(`/api/groups?${params.toString()}`, { method: 'GET', signal }), ) @@ -765,10 +737,6 @@ export async function invalidateGroupSummary( queryKey: controlQueryKeys.models.all, refetchType: 'none', }), - queryClient.invalidateQueries({ - queryKey: controlQueryKeys.providers.modelsAll(), - refetchType: 'none', - }), ]) } @@ -789,7 +757,7 @@ export async function invalidateGroupSettingsDependents( refetchType: 'active', }), queryClient.invalidateQueries({ - queryKey: controlQueryKeys.groups.keysAll(groupID), + queryKey: controlQueryKeys.groups.credentialsAll(groupID), refetchType: 'active', }), queryClient.invalidateQueries({ @@ -809,10 +777,6 @@ export async function invalidateGroupSettingsDependents( queryKey: controlQueryKeys.models.all, refetchType: 'active', }), - queryClient.invalidateQueries({ - queryKey: controlQueryKeys.providers.modelsAll(), - refetchType: 'active', - }), ]) } @@ -832,7 +796,10 @@ export function clearGroupResourceCaches(queryClient: QueryClient, groupID: numb queryClient.removeQueries({ queryKey: controlQueryKeys.groups.summary(groupID), exact: true }) queryClient.removeQueries({ queryKey: controlQueryKeys.groups.models(groupID), exact: true }) queryClient.removeQueries({ queryKey: controlQueryKeys.groups.settings(groupID), exact: true }) - queryClient.removeQueries({ queryKey: controlQueryKeys.groups.keysAll(groupID), exact: false }) + queryClient.removeQueries({ + queryKey: controlQueryKeys.groups.credentialsAll(groupID), + exact: false, + }) } export async function createGroup( @@ -851,15 +818,15 @@ export async function createGroup( ) } -export async function importGroupKeys( +export async function importGroupCredentials( client: ApiClient, groupID: number, - body: GroupKeyImportRequest, + body: CredentialImportRequest, idempotencyKey: string, signal?: AbortSignal, -): Promise { - return projectGroupKeyImportResult( - await client.request(`/api/groups/${groupID}/keys/import`, { +): Promise { + return projectCredentialImportResult( + await client.request(`/api/groups/${groupID}/credentials/import`, { method: 'POST', headers: { 'Idempotency-Key': idempotencyKey }, json: body, diff --git a/web/src/app/resources/health.ts b/web/src/app/resources/health.ts index 9a86814f..276f80f8 100644 --- a/web/src/app/resources/health.ts +++ b/web/src/app/resources/health.ts @@ -3,9 +3,9 @@ import { queryOptions } from '@tanstack/vue-query' import type { ApiClient } from '@/api/client' import type { HealthGroupDto, - HealthProblemKeyDto, + HealthProblemCredentialDto, HealthRecoveryDto, - KeyCounts, + HealthCredentialCountsDto, RequestLogHealthDto, RuntimeHealthDto, } from '@/api/control/types' @@ -26,14 +26,14 @@ import { export type { HealthGroupDto, - HealthProblemKeyDto, + HealthProblemCredentialDto, HealthRecoveryDto, - KeyCounts, + HealthCredentialCountsDto, RequestLogHealthDto, RuntimeHealthDto, } from '@/api/control/types' -const countFields = ['total', 'available', 'cooldown', 'blacklisted', 'disabled'] as const +const countFields = ['credentials', 'available', 'cooldown', 'blacklisted'] as const const healthFields = [ 'observed_at_ms', 'version', @@ -42,12 +42,12 @@ const healthFields = [ 'stats_window_seconds', 'counts', 'groups', - 'cooldown_keys', - 'blacklisted_keys', + 'cooldown_credentials', + 'blacklisted_credentials', 'request_log', ] as const -const problemKeyFields = [ - 'key_id', +const problemCredentialFields = [ + 'credential_id', 'group_id', 'group_name', 'cooldown_until_ms', @@ -101,23 +101,22 @@ function projectNonBlankString(value: unknown): string { return result } -function projectHealthKeyMask(value: unknown): string { +function projectHealthCredentialMask(value: unknown): string { const mask = projectString(value) if (mask !== '****' && !longMaskPattern.test(mask)) invalidResponse() return mask } -export function projectHealthCounts(value: unknown): KeyCounts { +export function projectHealthCounts(value: unknown): HealthCredentialCountsDto { const record = projectRecord(value) assertNoSecretLikeFields(record, countFields) const result = { - total: projectSafeInteger(record.total, { minimum: 0 }), + credentials: projectSafeInteger(record.credentials, { minimum: 0 }), available: projectSafeInteger(record.available, { minimum: 0 }), cooldown: projectSafeInteger(record.cooldown, { minimum: 0 }), blacklisted: projectSafeInteger(record.blacklisted, { minimum: 0 }), - disabled: projectSafeInteger(record.disabled, { minimum: 0 }), } - if (result.total !== result.available + result.cooldown + result.blacklisted + result.disabled) { + if (result.credentials !== result.available + result.cooldown + result.blacklisted) { invalidResponse() } return result @@ -144,9 +143,9 @@ function projectRecovery(value: unknown): HealthRecoveryDto { } } -function projectProblemKey(value: unknown): HealthProblemKeyDto { +function projectProblemCredential(value: unknown): HealthProblemCredentialDto { const record = projectRecord(value) - assertNoSecretLikeFields(record, problemKeyFields) + assertNoSecretLikeFields(record, problemCredentialFields) const recovery = projectRecovery(record.recovery) const cooldownUntilMS = projectNullableEpochMilliseconds(record.cooldown_until_ms) @@ -161,7 +160,7 @@ function projectProblemKey(value: unknown): HealthProblemKeyDto { } return { - key_id: projectSafeInteger(record.key_id, { minimum: 1 }), + credential_id: projectSafeInteger(record.credential_id, { minimum: 1 }), group_id: projectSafeInteger(record.group_id, { minimum: 1 }), group_name: projectNonBlankString(record.group_name), cooldown_until_ms: cooldownUntilMS, @@ -177,7 +176,7 @@ function projectProblemKey(value: unknown): HealthProblemKeyDto { : projectSafeInteger(record.weight_manual, { minimum: 0, maximum: 100 }), weight_auto: projectSafeInteger(record.weight_auto, { minimum: 0, maximum: 100 }), recovery, - mask: projectHealthKeyMask(record.mask), + mask: projectHealthCredentialMask(record.mask), last_failure_category: projectEnum(record.last_failure_category, problemFailureCategories), last_status_code: record.last_status_code === null @@ -218,8 +217,8 @@ export function projectRuntimeHealth(value: unknown): RuntimeHealthDto { stats_window_seconds: projectSafeInteger(record.stats_window_seconds, { minimum: 1 }), counts: projectHealthCounts(record.counts), groups: projectArray(record.groups, projectHealthGroup), - cooldown_keys: projectArray(record.cooldown_keys, projectProblemKey), - blacklisted_keys: projectArray(record.blacklisted_keys, projectProblemKey), + cooldown_credentials: projectArray(record.cooldown_credentials, projectProblemCredential), + blacklisted_credentials: projectArray(record.blacklisted_credentials, projectProblemCredential), request_log: projectRequestLogHealth(record.request_log), } } diff --git a/web/src/app/resources/home.ts b/web/src/app/resources/home.ts index 0786d445..a01b11c3 100644 --- a/web/src/app/resources/home.ts +++ b/web/src/app/resources/home.ts @@ -31,8 +31,8 @@ export interface HomeBaseDto { version: string inventory: { group_count: number - upstream_key_count: number - available_upstream_key_count: number + credential_count: number + available_credential_count: number model_count: number } access_keys: Array<{ @@ -120,8 +120,8 @@ const homeBaseFields = [ ] as const const inventoryFields = [ 'group_count', - 'upstream_key_count', - 'available_upstream_key_count', + 'credential_count', + 'available_credential_count', 'model_count', ] as const const accessKeyFields = ['id', 'name', 'masked_key', 'protocols'] as const @@ -181,13 +181,13 @@ function projectHomeInventory(value: unknown): HomeBaseDto['inventory'] { assertNoSecretLikeFields(record, inventoryFields) const result = { group_count: projectSafeInteger(record.group_count, { minimum: 0 }), - upstream_key_count: projectSafeInteger(record.upstream_key_count, { minimum: 0 }), - available_upstream_key_count: projectSafeInteger(record.available_upstream_key_count, { + credential_count: projectSafeInteger(record.credential_count, { minimum: 0 }), + available_credential_count: projectSafeInteger(record.available_credential_count, { minimum: 0, }), model_count: projectSafeInteger(record.model_count, { minimum: 0 }), } - if (result.available_upstream_key_count > result.upstream_key_count) invalidResponse() + if (result.available_credential_count > result.credential_count) invalidResponse() return result } diff --git a/web/src/app/resources/invalidation.ts b/web/src/app/resources/invalidation.ts index 9ec386e6..256d3a70 100644 --- a/web/src/app/resources/invalidation.ts +++ b/web/src/app/resources/invalidation.ts @@ -14,11 +14,11 @@ function plan( return { exact, prefixes } } -const importedKeyResourcePlan = (groupID: number) => +const importedCredentialResourcePlan = (groupID: number) => plan( [controlQueryKeys.groups.summary(groupID), controlQueryKeys.health()], [ - controlQueryKeys.groups.keysAll(groupID), + controlQueryKeys.groups.credentialsAll(groupID), controlQueryKeys.groups.collectionAll, controlQueryKeys.home.all, ], @@ -36,7 +36,6 @@ export const mutationInvalidationPlans = { controlQueryKeys.home.all, controlQueryKeys.models.all, controlQueryKeys.modelPrices(), - controlQueryKeys.providers.modelsAll(), ], ), delete: plan( @@ -46,10 +45,9 @@ export const mutationInvalidationPlans = { controlQueryKeys.home.all, controlQueryKeys.models.all, controlQueryKeys.modelPrices(), - controlQueryKeys.providers.modelsAll(), ], ), - importKeys: importedKeyResourcePlan, + importCredentials: importedCredentialResourcePlan, }, accessKey: { create: plan( @@ -80,7 +78,6 @@ export const mutationInvalidationPlans = { [ controlQueryKeys.modelPrices(), controlQueryKeys.models.all, - controlQueryKeys.providers.modelsAll(), controlQueryKeys.groups.modelsAll(), ], ), @@ -89,7 +86,6 @@ export const mutationInvalidationPlans = { [ controlQueryKeys.modelPrices(), controlQueryKeys.models.all, - controlQueryKeys.providers.modelsAll(), controlQueryKeys.groups.modelsAll(), ], ), @@ -98,7 +94,6 @@ export const mutationInvalidationPlans = { [ controlQueryKeys.modelPrices(), controlQueryKeys.models.all, - controlQueryKeys.providers.modelsAll(), controlQueryKeys.groups.modelsAll(), ], ), @@ -107,7 +102,6 @@ export const mutationInvalidationPlans = { [ controlQueryKeys.modelPrices(), controlQueryKeys.models.all, - controlQueryKeys.providers.modelsAll(), controlQueryKeys.groups.modelsAll(), ], ), diff --git a/web/src/app/resources/model-prices.ts b/web/src/app/resources/model-prices.ts index 90e49f07..4ba42c81 100644 --- a/web/src/app/resources/model-prices.ts +++ b/web/src/app/resources/model-prices.ts @@ -4,6 +4,7 @@ import { computed, toValue, type MaybeRefOrGetter } from 'vue' import type { ApiClient } from '@/api/client' import { ApiError, InvalidResponseError } from '@/api/errors' import { controlQueryKeys } from '@/app/query-keys' +import { projectChannelID } from '@/app/resources/channels' import { assertNoSecretLikeFields, @@ -21,6 +22,7 @@ export type ModelPriceStatus = 'pending' | 'configured' export type ModelPriceStatusFilter = ModelPriceStatus | 'all' export type ModelPriceUsageFilter = 'in_use' | 'unreferenced' | 'all' export type ModelPriceMethod = 'auto_sync' | 'user_set' | 'user_marked_unpriced' +export type ModelPriceMatchSource = 'channel_catalog_provider' | 'provider_priority_fallback' export interface ModelPriceFilters { usage: ModelPriceUsageFilter @@ -44,11 +46,13 @@ export interface ModelPriceContextTierDto { export interface ModelPriceDto { id: number + channel_id: string model_id: string prices: ModelPriceSlotsDto pricing_status: ModelPriceStatus method: ModelPriceMethod | null matched_provider_id: string | null + match_source: ModelPriceMatchSource | null referenced: boolean reference_count: number reference_group_count: number @@ -102,11 +106,13 @@ const collectionFields = ['items', 'pagination'] as const const paginationFields = ['page', 'page_size', 'total_items', 'total_pages'] as const const itemFields = [ 'id', + 'channel_id', 'model_id', 'prices', 'pricing_status', 'method', 'matched_provider_id', + 'match_source', 'referenced', 'reference_count', 'reference_group_count', @@ -175,6 +181,7 @@ export function projectModelPrice(value: unknown): ModelPriceDto { assertNoSecretLikeFields(record, itemFields) const result: ModelPriceDto = { id: projectSafeInteger(record.id, { minimum: 1 }), + channel_id: projectChannelID(record.channel_id), model_id: projectIdentityString(record.model_id), prices: projectPrices(record.prices), pricing_status: projectEnum(record.pricing_status, ['pending', 'configured'] as const), @@ -183,6 +190,13 @@ export function projectModelPrice(value: unknown): ModelPriceDto { record.matched_provider_id === null ? null : projectIdentityString(record.matched_provider_id), + match_source: + record.match_source === null + ? null + : projectEnum(record.match_source, [ + 'channel_catalog_provider', + 'provider_priority_fallback', + ] as const), referenced: projectBoolean(record.referenced), reference_count: projectSafeInteger(record.reference_count, { minimum: 0 }), reference_group_count: projectSafeInteger(record.reference_group_count, { minimum: 0 }), @@ -199,10 +213,13 @@ export function projectModelPrice(value: unknown): ModelPriceDto { result.referenced !== result.reference_count > 0 || (result.pricing_status === 'pending' && (result.method !== null || result.matched_provider_id !== null)) || - (hasAutomaticReference && result.matched_provider_id === null) || - (hasManualMethod && result.matched_provider_id !== null) || - (result.method === null && result.matched_provider_id !== null) || - (result.matched_provider_id !== null && !hasAutomaticReference) || + (hasAutomaticReference && + (result.matched_provider_id === null || result.match_source === null)) || + (hasManualMethod && (result.matched_provider_id !== null || result.match_source !== null)) || + (result.method === null && + (result.matched_provider_id !== null || result.match_source !== null)) || + ((result.matched_provider_id !== null || result.match_source !== null) && + !hasAutomaticReference) || (result.can_delete && (!result.can_reset || result.referenced)) ) { invalidResponse() diff --git a/web/src/app/resources/models.ts b/web/src/app/resources/models.ts index a326c54e..e894a946 100644 --- a/web/src/app/resources/models.ts +++ b/web/src/app/resources/models.ts @@ -3,10 +3,11 @@ import { computed, toValue, type MaybeRefOrGetter } from 'vue' import type { ApiClient } from '@/api/client' import { enabledDataProtocols } from '@/api/control/protocols' -import type { GroupProtocol } from '@/api/control/types' +import type { ChannelParamsDto, GroupProtocol } from '@/api/control/types' import { InvalidResponseError } from '@/api/errors' import { controlQueryKeys } from '@/app/query-keys' import { projectModelPrice, type ModelPriceDto } from '@/app/resources/model-prices' +import { projectChannelID } from '@/app/resources/channels' import { assertNoSecretLikeFields, @@ -67,9 +68,10 @@ export interface ModelCatalogMetadataDto { export interface ModelRouteGroupDto { id: number name: string - provider_id: string | null + channel_id: string + params: ChannelParamsDto enabled: boolean - protocols: GroupProtocol[] + client_protocols: GroupProtocol[] } export type ModelCatalogReferenceSource = 'actual_provider' | 'reference_provider' @@ -160,7 +162,7 @@ const upstreamDetailFields = [ ] as const const associationFields = ['client_model', 'alias_applied', 'group'] as const const catalogReferenceFields = ['source', 'provider_id', 'provider_name', 'model'] as const -const groupFields = ['id', 'name', 'provider_id', 'enabled', 'protocols'] as const +const groupFields = ['id', 'name', 'channel_id', 'params', 'enabled', 'client_protocols'] as const const catalogModelFields = [ 'id', 'name', @@ -223,6 +225,16 @@ function projectProtocols(value: unknown): GroupProtocol[] { return protocols } +function projectChannelParams(value: unknown): ChannelParamsDto { + const record = projectRecord(value) + const result: ChannelParamsDto = {} + for (const [key, param] of Object.entries(record)) { + if (!/^[a-z][a-z0-9_]*$/u.test(key)) invalidResponse() + result[key] = projectString(param, { allowEmpty: true }) + } + return result +} + function projectNullableBoolean(value: unknown): boolean | null { return value === null ? null : projectBoolean(value) } @@ -286,9 +298,10 @@ function projectRouteGroup(value: unknown): ModelRouteGroupDto { return { id: projectSafeInteger(record.id, { minimum: 1 }), name: projectIdentityString(record.name), - provider_id: record.provider_id === null ? null : projectIdentityString(record.provider_id), + channel_id: projectChannelID(record.channel_id), + params: projectChannelParams(record.params), enabled: projectBoolean(record.enabled), - protocols: projectProtocols(record.protocols), + client_protocols: projectProtocols(record.client_protocols), } } @@ -318,6 +331,8 @@ function projectUpstreamPrice( price.model_id !== upstreamModelID || routeGroups.length === 0 || affectedGroups.length === 0 || + routeGroups.some(({ channel_id }) => channel_id !== price.channel_id) || + affectedGroups.some(({ channel_id }) => channel_id !== price.channel_id) || new Set(routeGroups.map(({ id }) => id)).size !== routeGroups.length || affectedGroupIDs.size !== affectedGroups.length || routeGroups.some(({ id }) => !affectedGroupIDs.has(id)) || @@ -369,6 +384,7 @@ export function projectUpstreamModelDetail(value: unknown): UpstreamModelDetailD const groupCount = projectSafeInteger(record.group_count, { minimum: 0 }) if ( price.model_id !== modelID || + associations.some(({ group }) => group.channel_id !== price.channel_id) || clientModelCount !== new Set(associations.map(({ client_model }) => client_model)).size || groupCount !== new Set(associations.map(({ group }) => group.id)).size || price.reference_count !== associations.length || @@ -397,7 +413,7 @@ function projectClientModel(value: unknown): ClientModelDto { ) if ( upstreamModels.length === 0 || - new Set(upstreamModels.map(({ model_id }) => model_id)).size !== upstreamModels.length + new Set(upstreamModels.map(({ price }) => price.id)).size !== upstreamModels.length ) { invalidResponse() } diff --git a/web/src/app/resources/providers.ts b/web/src/app/resources/providers.ts index 94e92231..f603d810 100644 --- a/web/src/app/resources/providers.ts +++ b/web/src/app/resources/providers.ts @@ -1,11 +1,6 @@ -import { keepPreviousData, queryOptions } from '@tanstack/vue-query' -import { computed, toValue, type MaybeRefOrGetter } from 'vue' - import type { ApiClient } from '@/api/client' -import { enabledDataProtocols } from '@/api/control/protocols' -import type { GroupProtocol, ModelPricingStatus } from '@/api/control/types' +import type { ModelPricingStatus } from '@/api/control/types' import { InvalidResponseError } from '@/api/errors' -import { controlQueryKeys } from '@/app/query-keys' import { assertNoSecretLikeFields, @@ -13,28 +8,10 @@ import { projectBoolean, projectEnum, projectEpochMilliseconds, - projectHTTPURL, projectRecord, - projectSafeInteger, projectString, } from './projector' -export type ProviderSuggestionSource = 'official' | 'curated' | 'catalog' - -export interface ProviderSuggestion { - provider_id: string - name: string - api_url: string | null - protocols: GroupProtocol[] - mark: string - source: ProviderSuggestionSource -} - -export interface ProviderSuggestionList { - items: ProviderSuggestion[] - total: number -} - export type ModelCandidateSource = 'catalog' | 'live' export interface ModelCandidate { @@ -45,16 +22,6 @@ export interface ModelCandidate { pricing_source: string | null } -export interface ProviderModelFilters { - q?: string - status?: ModelPricingStatus -} - -export interface ProviderModelList { - items: ModelCandidate[] - total: number -} - export interface CatalogSyncStatus { trigger: 'startup' | 'periodic' | 'group_change' | 'manual' checked_at_ms: number @@ -64,18 +31,7 @@ export interface CatalogSyncStatus { error_code: string | null } -const providerSuggestionFields = [ - 'provider_id', - 'name', - 'api_url', - 'protocols', - 'mark', - 'source', -] as const -const providerSuggestionSources = ['official', 'curated', 'catalog'] as const -const providerSuggestionListFields = ['items', 'total'] as const const modelCandidateFields = ['id', 'name', 'sources', 'pricing_status', 'pricing_source'] as const -const providerModelListFields = ['items', 'total'] as const const catalogSyncStatusFields = [ 'trigger', 'checked_at_ms', @@ -92,65 +48,6 @@ function invalidResponse(): never { throw new InvalidResponseError() } -export function projectProviderID(value: unknown): string { - const providerID = projectString(value) - if ( - providerID !== providerID.trim() || - providerID.includes(':') || - /[\u0000-\u001f\u007f]/u.test(providerID) - ) { - invalidResponse() - } - return providerID -} - -function projectProtocols(value: unknown): GroupProtocol[] { - const protocols = projectArray(value, (protocol) => projectEnum(protocol, enabledDataProtocols)) - if (new Set(protocols).size !== protocols.length) invalidResponse() - return protocols -} - -function projectProviderSuggestion(value: unknown): ProviderSuggestion { - const record = projectRecord(value) - assertNoSecretLikeFields(record, providerSuggestionFields) - return { - provider_id: projectProviderID(record.provider_id), - name: projectString(record.name), - api_url: record.api_url === undefined ? null : projectHTTPURL(record.api_url), - protocols: projectProtocols(record.protocols), - mark: projectString(record.mark, { allowEmpty: true }), - source: projectEnum(record.source, providerSuggestionSources), - } -} - -export function projectProviderSuggestionList(value: unknown): ProviderSuggestionList { - return projectProviderSuggestionListValue(value, true) -} - -function projectProviderSuggestionListByIDs(value: unknown): ProviderSuggestionList { - const result = projectProviderSuggestionListValue(value, false) - if (result.items.some(({ source }) => source !== 'catalog')) invalidResponse() - return result -} - -function projectProviderSuggestionListValue( - value: unknown, - requireOfficialSuggestions: boolean, -): ProviderSuggestionList { - const record = projectRecord(value) - assertNoSecretLikeFields(record, providerSuggestionListFields) - const items = projectArray(record.items, projectProviderSuggestion) - const total = projectSafeInteger(record.total, { minimum: 0 }) - if ( - total !== items.length || - new Set(items.map(({ provider_id }) => provider_id)).size !== items.length || - (requireOfficialSuggestions && items.filter(({ source }) => source === 'official').length !== 3) - ) { - invalidResponse() - } - return { items, total } -} - export function projectModelCandidate(value: unknown): ModelCandidate { const record = projectRecord(value) assertNoSecretLikeFields(record, modelCandidateFields) @@ -181,17 +78,6 @@ export function projectModelCandidate(value: unknown): ModelCandidate { } } -function projectProviderModelList(value: unknown): ProviderModelList { - const record = projectRecord(value) - assertNoSecretLikeFields(record, providerModelListFields) - const items = projectArray(record.items, projectModelCandidate) - const total = projectSafeInteger(record.total, { minimum: 0 }) - if (items.length > total || new Set(items.map(({ id }) => id)).size !== items.length) { - invalidResponse() - } - return { items, total } -} - function projectCatalogSyncStatus(value: unknown): CatalogSyncStatus { const record = projectRecord(value) assertNoSecretLikeFields(record, catalogSyncStatusFields) @@ -208,81 +94,6 @@ function projectCatalogSyncStatus(value: unknown): CatalogSyncStatus { } } -export function normalizeProviderSearch(value: string): string { - return value.trim().slice(0, 200) -} - -export function normalizeProviderModelFilters(filters: ProviderModelFilters): ProviderModelFilters { - const normalized: ProviderModelFilters = {} - const query = normalizeProviderSearch(filters.q ?? '') - if (query) normalized.q = query - if (filters.status !== undefined) normalized.status = filters.status - return normalized -} - -export async function listProviderSuggestions( - client: ApiClient, - search: string, - signal?: AbortSignal, -): Promise { - const query = normalizeProviderSearch(search) - const params = new URLSearchParams() - if (query) params.set('q', query) - const suffix = params.size ? `?${params.toString()}` : '' - return projectProviderSuggestionList( - await client.request(`/api/provider-suggestions${suffix}`, { method: 'GET', signal }), - ) -} - -export async function listProviderSuggestionsByIDs( - client: ApiClient, - providerIDs: readonly string[], - signal?: AbortSignal, -): Promise { - const normalizedIDs = normalizeProviderIDs(providerIDs) - if (!normalizedIDs.length) return { items: [], total: 0 } - const params = new URLSearchParams() - params.set('provider_ids', normalizedIDs.join(',')) - return projectProviderSuggestionListByIDs( - await client.request(`/api/provider-suggestions?${params.toString()}`, { - method: 'GET', - signal, - }), - ) -} - -function normalizeProviderIDs(providerIDs: readonly string[]): string[] { - const seen = new Set() - const result: string[] = [] - for (const providerID of providerIDs) { - const normalized = providerID.trim() - if (!normalized || seen.has(normalized)) continue - seen.add(normalized) - result.push(normalized) - } - return result -} - -export async function listProviderModels( - client: ApiClient, - providerID: string, - filters: ProviderModelFilters, - signal?: AbortSignal, -): Promise { - const canonicalProviderID = projectProviderID(providerID) - const normalized = normalizeProviderModelFilters(filters) - const params = new URLSearchParams() - if (normalized.q !== undefined) params.set('q', normalized.q) - if (normalized.status !== undefined) params.set('status', normalized.status) - const suffix = params.size ? `?${params.toString()}` : '' - return projectProviderModelList( - await client.request( - `/api/providers/${encodeURIComponent(canonicalProviderID)}/models${suffix}`, - { method: 'GET', signal }, - ), - ) -} - export async function syncModelPrices( client: ApiClient, signal?: AbortSignal, @@ -291,56 +102,3 @@ export async function syncModelPrices( await client.request('/api/model-prices/sync', { method: 'POST', signal }), ) } - -export function providerSuggestionsQueryOptions( - client: ApiClient, - search: MaybeRefOrGetter, -) { - return queryOptions({ - queryKey: computed(() => - controlQueryKeys.providers.suggestions(normalizeProviderSearch(toValue(search))), - ), - queryFn: ({ queryKey, signal }) => listProviderSuggestions(client, queryKey[3], signal), - staleTime: 5 * 60 * 1_000, - placeholderData: keepPreviousData, - }) -} - -export function providerSuggestionsByIDsQueryOptions( - client: ApiClient, - providerIDs: MaybeRefOrGetter, -) { - return queryOptions({ - queryKey: computed(() => { - const normalizedIDs = normalizeProviderIDs(toValue(providerIDs)) - return controlQueryKeys.providers.suggestionsByIDs(normalizedIDs) - }), - queryFn: ({ queryKey, signal }) => listProviderSuggestionsByIDs(client, queryKey[3], signal), - enabled: computed(() => normalizeProviderIDs(toValue(providerIDs)).length > 0), - staleTime: 5 * 60 * 1_000, - placeholderData: keepPreviousData, - }) -} - -export function providerModelsQueryOptions( - client: ApiClient, - providerID: MaybeRefOrGetter, - filters: MaybeRefOrGetter, -) { - return queryOptions({ - queryKey: computed(() => { - const id = toValue(providerID) - return id === null - ? controlQueryKeys.providers.modelsAll() - : controlQueryKeys.providers.models(id, normalizeProviderModelFilters(toValue(filters))) - }), - queryFn: ({ signal }) => { - const id = toValue(providerID) - if (id === null) invalidResponse() - return listProviderModels(client, id, toValue(filters), signal) - }, - enabled: computed(() => toValue(providerID) !== null), - staleTime: 5 * 60 * 1_000, - placeholderData: keepPreviousData, - }) -} diff --git a/web/src/app/resources/request-log-filters.ts b/web/src/app/resources/request-log-filters.ts index 88ea5c58..91d6383a 100644 --- a/web/src/app/resources/request-log-filters.ts +++ b/web/src/app/resources/request-log-filters.ts @@ -3,6 +3,8 @@ export const requestLogFilterFields = [ 'to_ms', 'limit', 'group_id', + 'channel_id', + 'credential_id', 'client_model', 'upstream_model', 'access_key_id', @@ -15,7 +17,6 @@ export const requestLogFilterFields = [ 'cost_state', 'pricing_completeness', 'cache_present', - 'upstream_key_id', 'attempt_status_code', 'failure_category', 'error_code', diff --git a/web/src/app/resources/request-logs.ts b/web/src/app/resources/request-logs.ts index c597e963..218238b1 100644 --- a/web/src/app/resources/request-logs.ts +++ b/web/src/app/resources/request-logs.ts @@ -6,6 +6,7 @@ import { enabledDataProtocols } from '@/api/control/protocols' import type { AccessProtocol, FailureCategory } from '@/api/control/types' import { InvalidResponseError } from '@/api/errors' import { controlQueryKeys } from '@/app/query-keys' +import { projectChannelID } from '@/app/resources/channels' import { normalizeRequestLogFilters, requestLogFilterFields } from './request-log-filters' @@ -24,7 +25,8 @@ import { export type RequestLogStatus = 'success' | 'error' | 'incomplete' | 'canceled' export type RequestLogModelConsistency = 'not_applicable' | 'match' | 'unknown' | 'mismatch' -export type RequestLogAction = 'terminate' | 'retry' | 'cooldown_key' | 'fail_key' | 'skip_group' +export type RequestLogAction = + 'terminate' | 'retry' | 'cooldown_credential' | 'fail_credential' | 'skip_group' export type RequestLogUsageState = 'complete' | 'partial' | 'missing' | 'not_applicable' export type RequestLogCostState = 'priced' | 'unpriced' | 'not_applicable' export type RequestLogPricingCompleteness = @@ -32,6 +34,20 @@ export type RequestLogPricingCompleteness = export type RequestLogRetryState = 'retried' | 'not_retried' export type RequestLogReceiptLineState = 'priced' | 'unpriced' export type RequestLogPageSize = 20 | 50 | 100 +export type RequestLogOperation = + | 'chat_completion' + | 'responses_create' + | 'responses_retrieve' + | 'responses_delete' + | 'responses_cancel' + | 'responses_input_items' + | 'responses_compact' + | 'responses_input_tokens' + | 'responses_passthrough' + | 'list_models' + | 'probe' +export type RequestLogRouteMode = 'native' | 'converted' +export type RequestLogDispatchState = 'not_sent' | 'maybe_sent' export type { FailureCategory } from '@/api/control/types' @@ -40,6 +56,8 @@ export interface RequestLogFilters { to_ms?: number limit?: RequestLogPageSize group_id?: number + channel_id?: string + credential_id?: number client_model?: string upstream_model?: string access_key_id?: number @@ -52,7 +70,6 @@ export interface RequestLogFilters { cost_state?: RequestLogCostState pricing_completeness?: RequestLogPricingCompleteness cache_present?: boolean - upstream_key_id?: number attempt_status_code?: number failure_category?: FailureCategory error_code?: string @@ -81,11 +98,11 @@ export interface RequestLogPricingLineDto { } export interface RequestLogPricingReceiptDto { - schema_version: 1 | 2 + schema_version: 1 | 2 | 3 method: 'unit_rate_sum' method_version: 1 currency: 'USD' - rule: { scope_key?: string; model_id: string } + rule: { scope_key?: string; channel_id?: string; model_id: string } context_threshold_tokens: string | null line_items: RequestLogPricingLineDto[] total_nano_usd: string @@ -95,8 +112,14 @@ export interface RequestLogAttemptDto { sequence: number group_id: number group_name: string - key_id: number + channel_id: string | null + credential_id: number | null + operation: RequestLogOperation | null + route_mode: RequestLogRouteMode | null upstream_model: string | null + upstream_request_id: string | null + dispatch_state: RequestLogDispatchState | null + response_started: boolean status_code: number duration_ms: number failure_category: FailureCategory @@ -104,6 +127,7 @@ export interface RequestLogAttemptDto { will_retry: boolean error_code: string error_summary: string + committed: boolean pricing_receipt: RequestLogPricingReceiptDto | null } @@ -133,6 +157,8 @@ export interface RequestLogItemDto { error_summary: string affinity_hit: boolean group_id: number | null + channel_id: string | null + credential_id: number | null usage_state: RequestLogUsageState cost_state: RequestLogCostState pricing_completeness: RequestLogPricingCompleteness @@ -167,7 +193,28 @@ const failureCategories = [ 'downstream_cancel', 'ambiguous', ] as const -const actions = ['terminate', 'retry', 'cooldown_key', 'fail_key', 'skip_group'] as const +const actions = [ + 'terminate', + 'retry', + 'cooldown_credential', + 'fail_credential', + 'skip_group', +] as const +const operations = [ + 'chat_completion', + 'responses_create', + 'responses_retrieve', + 'responses_delete', + 'responses_cancel', + 'responses_input_items', + 'responses_compact', + 'responses_input_tokens', + 'responses_passthrough', + 'list_models', + 'probe', +] as const +const routeModes = ['native', 'converted'] as const +const dispatchStates = ['not_sent', 'maybe_sent'] as const const usageStates = ['complete', 'partial', 'missing', 'not_applicable'] as const const costStates = ['priced', 'unpriced', 'not_applicable'] as const const pricingCompletenessValues = ['complete', 'partial', 'unavailable', 'not_applicable'] as const @@ -200,6 +247,8 @@ const itemFields = [ 'error_summary', 'affinity_hit', 'group_id', + 'channel_id', + 'credential_id', 'usage_state', 'cost_state', 'pricing_completeness', @@ -260,7 +309,7 @@ function projectPricingReceipt(value: unknown): RequestLogPricingReceiptDto | nu 'total_nano_usd', ]) const rule = projectRecord(record.rule) - assertNoSecretLikeFields(rule, ['scope_key', 'model_id']) + assertNoSecretLikeFields(rule, ['scope_key', 'channel_id', 'model_id']) const lines = projectArray(record.line_items, (lineValue): RequestLogPricingLineDto => { const line = projectRecord(lineValue) assertNoSecretLikeFields(line, [ @@ -292,12 +341,14 @@ function projectPricingReceipt(value: unknown): RequestLogPricingReceiptDto | nu line.amount_nano_usd === null ? null : projectNonNegativeInt64String(line.amount_nano_usd), } }) - const schemaVersion = projectSafeInteger(record.schema_version, { minimum: 1, maximum: 2 }) as - 1 | 2 + const schemaVersion = projectSafeInteger(record.schema_version, { minimum: 1, maximum: 3 }) as + 1 | 2 | 3 const scopeKey = rule.scope_key === undefined ? undefined : projectNonBlankString(rule.scope_key) + const channelID = rule.channel_id === undefined ? undefined : projectChannelID(rule.channel_id) if ( - (schemaVersion === 1 && scopeKey === undefined) || - (schemaVersion === 2 && scopeKey !== undefined) + (schemaVersion === 1 && (scopeKey === undefined || channelID !== undefined)) || + (schemaVersion === 2 && (scopeKey !== undefined || channelID !== undefined)) || + (schemaVersion === 3 && (scopeKey !== undefined || channelID === undefined)) ) { invalidResponse() } @@ -308,6 +359,7 @@ function projectPricingReceipt(value: unknown): RequestLogPricingReceiptDto | nu currency: projectEnum(record.currency, ['USD'] as const), rule: { ...(scopeKey === undefined ? {} : { scope_key: scopeKey }), + ...(channelID === undefined ? {} : { channel_id: channelID }), model_id: projectNonBlankString(rule.model_id), }, context_threshold_tokens: @@ -325,8 +377,14 @@ function projectAttempt(value: unknown): RequestLogAttemptDto { 'sequence', 'group_id', 'group_name', - 'key_id', + 'channel_id', + 'credential_id', + 'operation', + 'route_mode', 'upstream_model', + 'upstream_request_id', + 'dispatch_state', + 'response_started', 'status_code', 'duration_ms', 'failure_category', @@ -334,14 +392,25 @@ function projectAttempt(value: unknown): RequestLogAttemptDto { 'will_retry', 'error_code', 'error_summary', + 'committed', 'pricing_receipt', ]) return { sequence: projectSafeInteger(record.sequence, { minimum: 1 }), group_id: projectSafeInteger(record.group_id, { minimum: 1 }), group_name: projectNonBlankString(record.group_name), - key_id: projectSafeInteger(record.key_id, { minimum: 1 }), + channel_id: record.channel_id === null ? null : projectChannelID(record.channel_id), + credential_id: + record.credential_id === null + ? null + : projectSafeInteger(record.credential_id, { minimum: 1 }), + operation: record.operation === null ? null : projectEnum(record.operation, operations), + route_mode: record.route_mode === null ? null : projectEnum(record.route_mode, routeModes), upstream_model: projectNullableModel(record.upstream_model), + upstream_request_id: projectNullableModel(record.upstream_request_id), + dispatch_state: + record.dispatch_state === null ? null : projectEnum(record.dispatch_state, dispatchStates), + response_started: projectBoolean(record.response_started), status_code: projectStatusCode(record.status_code), duration_ms: projectSafeInteger(record.duration_ms, { minimum: 0 }), failure_category: projectEnum(record.failure_category, failureCategories), @@ -349,6 +418,7 @@ function projectAttempt(value: unknown): RequestLogAttemptDto { will_retry: projectBoolean(record.will_retry), error_code: projectString(record.error_code, { allowEmpty: true }), error_summary: projectString(record.error_summary, { allowEmpty: true }), + committed: projectBoolean(record.committed), pricing_receipt: projectPricingReceipt(record.pricing_receipt), } } @@ -444,6 +514,11 @@ function projectItemRecord(record: Record): RequestLogItemDto { error_summary: projectString(record.error_summary, { allowEmpty: true }), affinity_hit: projectBoolean(record.affinity_hit), group_id: record.group_id === null ? null : projectSafeInteger(record.group_id, { minimum: 1 }), + channel_id: record.channel_id === null ? null : projectChannelID(record.channel_id), + credential_id: + record.credential_id === null + ? null + : projectSafeInteger(record.credential_id, { minimum: 1 }), ...projectUsageCost(record), } } diff --git a/web/src/app/resources/route-inspection.ts b/web/src/app/resources/route-inspection.ts index 56b71680..0574e0f9 100644 --- a/web/src/app/resources/route-inspection.ts +++ b/web/src/app/resources/route-inspection.ts @@ -20,17 +20,19 @@ export type RouteInspectReasonCode = | 'protocol_filtered' | 'model_filtered' | 'model_required_by_filter' + | 'capability_unsupported' | 'no_route_target' | 'group_disabled' | 'group_filtered' | 'no_available_group' - | 'no_keys' + | 'no_credentials' | 'group_weight_zero' - | 'key_disabled' - | 'key_blacklisted' - | 'key_cooldown' - | 'key_weight_zero' - | 'no_available_key' + | 'credential_disabled' + | 'credential_blacklisted' + | 'credential_cooldown' + | 'credential_weight_zero' + | 'credential_not_allowed' + | 'no_available_credential' export interface RouteInspectRequest { protocol: AccessProtocol @@ -38,8 +40,8 @@ export interface RouteInspectRequest { access_key_id: number } -export interface RouteInspectKeyDto { - key_id: number +export interface RouteInspectCredentialDto { + credential_id: number available: boolean reason_code: RouteInspectReasonCode | null weight_manual: number | null @@ -56,7 +58,7 @@ export interface RouteInspectGroupDto { included: boolean routable: boolean reason_code: RouteInspectReasonCode | null - keys: RouteInspectKeyDto[] + credentials: RouteInspectCredentialDto[] } export interface RouteInspectResponseDto { @@ -80,17 +82,19 @@ const reasonCodes = [ 'protocol_filtered', 'model_filtered', 'model_required_by_filter', + 'capability_unsupported', 'no_route_target', 'group_disabled', 'group_filtered', 'no_available_group', - 'no_keys', + 'no_credentials', 'group_weight_zero', - 'key_disabled', - 'key_blacklisted', - 'key_cooldown', - 'key_weight_zero', - 'no_available_key', + 'credential_disabled', + 'credential_blacklisted', + 'credential_cooldown', + 'credential_weight_zero', + 'credential_not_allowed', + 'no_available_credential', ] as const function invalidResponse(): never { @@ -115,10 +119,10 @@ function projectNullableWeight(value: unknown): number | null { return value === null ? null : projectSafeInteger(value, { minimum: 0, maximum: 100 }) } -function projectRouteKey(value: unknown): RouteInspectKeyDto { +function projectRouteCredential(value: unknown): RouteInspectCredentialDto { const record = projectRecord(value) assertNoSecretLikeFields(record, [ - 'key_id', + 'credential_id', 'available', 'reason_code', 'weight_manual', @@ -127,7 +131,7 @@ function projectRouteKey(value: unknown): RouteInspectKeyDto { 'cooldown_until_ms', ]) return { - key_id: projectSafeInteger(record.key_id, { minimum: 1 }), + credential_id: projectSafeInteger(record.credential_id, { minimum: 1 }), available: projectBoolean(record.available), reason_code: projectReason(record.reason_code), weight_manual: projectNullableWeight(record.weight_manual), @@ -147,7 +151,7 @@ function projectRouteGroup(value: unknown): RouteInspectGroupDto { 'included', 'routable', 'reason_code', - 'keys', + 'credentials', ]) return { group_id: projectSafeInteger(record.group_id, { minimum: 1 }), @@ -157,7 +161,7 @@ function projectRouteGroup(value: unknown): RouteInspectGroupDto { included: projectBoolean(record.included), routable: projectBoolean(record.routable), reason_code: projectReason(record.reason_code), - keys: projectArray(record.keys, projectRouteKey), + credentials: projectArray(record.credentials, projectRouteCredential), } } diff --git a/web/src/app/resources/usage.ts b/web/src/app/resources/usage.ts index d9e4d923..74066dab 100644 --- a/web/src/app/resources/usage.ts +++ b/web/src/app/resources/usage.ts @@ -4,6 +4,7 @@ import { computed, toValue, type MaybeRefOrGetter } from 'vue' import type { ApiClient } from '@/api/client' import { InvalidResponseError } from '@/api/errors' import { controlQueryKeys } from '@/app/query-keys' +import { projectChannelID } from '@/app/resources/channels' import { timeRanges, type TimeRange } from '@/lib/time' import { @@ -27,7 +28,9 @@ export interface UsageFilters { range: UsageRange breakdown_order?: UsageBreakdownOrder group_id?: number - model?: string + channel_id?: string + credential_id?: number + upstream_model?: string } type NormalizedUsageFilters = UsageFilters & { @@ -61,7 +64,14 @@ export interface UsageReportDto { observed_at_ms: number summary: UsageAggregateDto series: Array - breakdown: Array + breakdown: Array< + UsageAggregateDto & { + group_id: number + channel_id: string | null + credential_id: number | null + model: string + } + > breakdown_truncated: boolean breakdown_order: UsageBreakdownOrder breakdown_count: number @@ -243,6 +253,8 @@ export function projectUsageReport(value: unknown): UsageReportDto { ...aggregateKeys, 'estimated_cost_nano_usd', 'group_id', + 'channel_id', + 'credential_id', 'model', ]) const model = projectString(item.model, { allowEmpty: true }) @@ -252,6 +264,9 @@ export function projectUsageReport(value: unknown): UsageReportDto { Object.fromEntries(aggregateFields.map((field) => [field, item[field]])), ), group_id: projectSafeInteger(item.group_id, { minimum: 0 }), + channel_id: item.channel_id === null ? null : projectChannelID(item.channel_id), + credential_id: + item.credential_id === null ? null : projectSafeInteger(item.credential_id, { minimum: 1 }), model, } }) @@ -280,7 +295,9 @@ export function normalizeUsageFilters(filters: UsageFilters): NormalizedUsageFil breakdown_order: breakdownOrder, } if (filters.group_id !== undefined) result.group_id = filters.group_id - if (filters.model !== undefined) result.model = filters.model + if (filters.channel_id !== undefined) result.channel_id = filters.channel_id + if (filters.credential_id !== undefined) result.credential_id = filters.credential_id + if (filters.upstream_model !== undefined) result.upstream_model = filters.upstream_model return result } @@ -299,7 +316,13 @@ export async function getUsageReport( params.append('breakdown_order', normalized.breakdown_order) } if (normalized.group_id !== undefined) params.append('group_id', String(normalized.group_id)) - if (normalized.model !== undefined) params.append('upstream_model', normalized.model) + if (normalized.channel_id !== undefined) params.append('channel_id', normalized.channel_id) + if (normalized.credential_id !== undefined) { + params.append('credential_id', String(normalized.credential_id)) + } + if (normalized.upstream_model !== undefined) { + params.append('upstream_model', normalized.upstream_model) + } const report = projectUsageReport( await client.request(`/api/usage?${params.toString()}`, { method: 'GET', signal }), ) diff --git a/web/src/components/ui/CredentialHealthBar.vue b/web/src/components/ui/CredentialHealthBar.vue new file mode 100644 index 00000000..130ba77f --- /dev/null +++ b/web/src/components/ui/CredentialHealthBar.vue @@ -0,0 +1,187 @@ + + + + + diff --git a/web/src/components/ui/KeyHealthBar.vue b/web/src/components/ui/KeyHealthBar.vue deleted file mode 100644 index cc51b111..00000000 --- a/web/src/components/ui/KeyHealthBar.vue +++ /dev/null @@ -1,164 +0,0 @@ - - - - - diff --git a/web/src/features/access-keys/AccessKeyDrawer.vue b/web/src/features/access-keys/AccessKeyDrawer.vue index 5dde7964..1a01b533 100644 --- a/web/src/features/access-keys/AccessKeyDrawer.vue +++ b/web/src/features/access-keys/AccessKeyDrawer.vue @@ -11,6 +11,7 @@ import { type CreateAccessKeyRequest, } from '@/app/resources/access-keys' import type { AccessKeyDto, AccessProtocol, GroupOptionDto } from '@/api/control/types' +import type { ChannelDto } from '@/app/resources/channels' import { RequestCancelledError } from '@/api/errors' import { classifyMutationOutcome } from '@/app/mutation-outcome' import { applyInvalidationPlan, mutationInvalidationPlans } from '@/app/resources/invalidation' @@ -59,6 +60,7 @@ const props = withDefaults( open: boolean accessKey: AccessKeyDto | null groups: GroupOptionDto[] + channels: ChannelDto[] total: number groupCatalogState?: GroupCatalogState createOperation?: PendingAccessKeyCreateOperation | null @@ -107,7 +109,7 @@ const selectedGroupIDs = computed(() => draft.value.scopeModes.groups === 'restricted' ? draft.value.filters.groups : [], ) const supportedProtocolOptions = computed(() => - buildAccessKeyProtocolCandidates(props.groups, selectedGroupIDs.value), + buildAccessKeyProtocolCandidates(props.groups, props.channels, selectedGroupIDs.value), ) const catalogModelOptions = computed(() => buildAccessKeyModelOptions(props.groups, [], selectedGroupIDs.value), diff --git a/web/src/features/access-keys/AccessKeysView.vue b/web/src/features/access-keys/AccessKeysView.vue index fb75722c..4969cfdd 100644 --- a/web/src/features/access-keys/AccessKeysView.vue +++ b/web/src/features/access-keys/AccessKeysView.vue @@ -19,6 +19,7 @@ import { updateAccessKey, } from '@/app/resources/access-keys' import { groupOptionsQueryOptions } from '@/app/resources/groups' +import { channelsQueryOptions } from '@/app/resources/channels' import { applyInvalidationPlan, mutationInvalidationPlans } from '@/app/resources/invalidation' import { accessKeysLocation } from '@/app/route-locations' import { useToast } from '@/app/toast' @@ -70,6 +71,7 @@ const pendingStatusIDs = ref(new Set()) const statusControllers = new Map() const accessKeysQuery = useQuery(accessKeyCollectionQueryOptions(client, filters)) const groupsQuery = useQuery(groupOptionsQueryOptions(client)) +const channelsQuery = useQuery(channelsQueryOptions(client, '')) const data = computed(() => accessKeysQuery.data.value) const collectionBusy = computed(() => data.value !== undefined && accessKeysQuery.isFetching.value) const { @@ -90,7 +92,8 @@ const { const pageRefreshing = computed( () => collectionRefreshing.value || - (groupsQuery.data.value !== undefined && groupsQuery.isFetching.value), + (groupsQuery.data.value !== undefined && groupsQuery.isFetching.value) || + (channelsQuery.data.value !== undefined && channelsQuery.isFetching.value), ) const hasFilterCriteria = computed( () => filters.value.q !== undefined || filters.value.status !== undefined, @@ -120,8 +123,10 @@ const statusSummaryItems = computed(() => { ] }) const groupCatalogState = computed(() => { - if (groupsQuery.isError.value) return groupsQuery.data.value ? 'stale' : 'error' - if (groupsQuery.isPending.value) return 'loading' + if (groupsQuery.isError.value || channelsQuery.isError.value) { + return groupsQuery.data.value && channelsQuery.data.value ? 'stale' : 'error' + } + if (groupsQuery.isPending.value || channelsQuery.isPending.value) return 'loading' return 'ready' }) const operationNoticeKey = computed(() => @@ -256,6 +261,10 @@ function resetConditions(): void { routeWithFilters({ page: 1, page_size: 20 }) } +function retryGroupCatalog(): void { + void Promise.allSettled([groupsQuery.refetch(), channelsQuery.refetch()]) +} + function setPage(page: number): void { routeWithFilters({ ...filters.value, page }) } @@ -403,6 +412,7 @@ async function toggleStatus(accessKey: AccessKeyDto): Promise { :open="drawerOpen" :access-key="selected" :groups="groupsQuery.data.value ?? []" + :channels="channelsQuery.data.value?.items ?? []" :total="data?.summary.total ?? 0" :group-catalog-state="groupCatalogState" :create-operation="createOperation" @@ -483,12 +493,12 @@ async function toggleStatus(accessKey: AccessKeyDto): Promise { @retry="accessKeysQuery.refetch()" /> -