From 22a848b3604a9b182e7ea26053e48b70bd9ff2c5 Mon Sep 17 00:00:00 2001 From: Zyy Date: Wed, 7 Oct 2026 19:18:40 +0800 Subject: [PATCH] refactor(backend): clarify ownership and align xai contracts Separate execution, provider ports, administration, and plugin lifecycle responsibilities. Fix shutdown ordering, account lock reclamation, multipart cancellation, budget windows, and WebSocket close handling. Align xAI catalog, reasoning, tools, replay, quota, and cost handling with current Grok contracts and remove legacy fallbacks. Update API and architecture documentation with regression coverage. Validated backend suites, strict workspace Clippy, isolated PostgreSQL/Redis integration, and real Codex/Grok client flows. Co-authored-by: Codex --- backend/apps/gateway/src/bootstrap/command.rs | 4 +- backend/apps/gateway/src/bootstrap/plugins.rs | 6 +- backend/apps/gateway/src/bootstrap/server.rs | 8 +- backend/apps/gateway/tests/architecture.rs | 171 +- backend/crates/gateway-admin/src/lib.rs | 10 +- .../src/ports/plugin_management.rs | 2 +- .../crates/gateway-admin/src/ports/plugins.rs | 2 +- .../src/{service => public_service}/mod.rs | 6 +- .../{service => public_service}/registry.rs | 2 +- .../{service => public_service}/settings.rs | 0 .../src/use_case/credential_mutation.rs | 371 +++ .../gateway-admin/src/use_case/error.rs | 67 + .../crates/gateway-admin/src/use_case/mod.rs | 458 +--- .../src/use_case/plugins/management.rs | 3 +- .../gateway-admin/src/use_case/publication.rs | 14 + .../gateway-admin/tests/use_case/mod.rs | 7 +- .../tests/use_case/plugins/artifacts.rs | 12 +- .../tests/use_case/plugins/instances.rs | 6 +- .../tests/use_case/plugins/mod.rs | 20 +- .../tests/use_case/service_contract.rs | 4 +- .../gateway-admin/tests/use_case/settings.rs | 2 +- backend/crates/gateway-api/src/lib.rs | 2 +- .../crates/gateway-api/src/middleware/mod.rs | 28 +- .../gateway-api/src/middleware/websocket.rs | 17 +- .../responses/websocket/connection/frame.rs | 171 ++ .../responses/websocket/connection/handle.rs | 362 +++ .../websocket/connection/middleware.rs | 62 +- .../responses/websocket/connection/mod.rs | 811 +------ .../responses/websocket/connection/pump.rs | 230 ++ .../responses/websocket/connection/state.rs | 75 + .../tests/admin/plugins/management.rs | 52 +- .../gateway-api/tests/admin/plugins/mod.rs | 20 +- .../crates/gateway-api/tests/architecture.rs | 160 +- .../gateway-api/tests/middleware/mod.rs | 41 +- .../crates/gateway-api/tests/openai/auth.rs | 6 +- .../tests/openai/middleware/mod.rs | 28 +- .../websocket/connection/middleware.rs | 102 +- .../responses/websocket/connection/mod.rs | 109 +- .../openai/responses/websocket/forward.rs | 2 +- .../gateway-core/src/engine/authentication.rs | 15 +- .../src/engine/execution/admission.rs | 71 + .../src/engine/execution/contract.rs | 480 ++++ .../src/engine/execution/failure.rs | 100 + .../gateway-core/src/engine/execution/mod.rs | 24 + .../src/engine/execution/probe_store.rs | 64 + .../src/engine/execution/scope.rs | 301 +++ .../{execution.rs => execution/service.rs} | 1286 +---------- .../src/engine/execution/session.rs | 213 ++ .../gateway-core/src/engine/extensions.rs | 81 +- .../{middleware.rs => middleware/contract.rs} | 182 +- .../src/{ => engine}/middleware/http/mod.rs | 6 +- .../{ => engine}/middleware/http/upgrade.rs | 2 +- .../gateway-core/src/engine/middleware/mod.rs | 15 + .../src/engine/middleware/plan.rs | 152 ++ .../src/{ => engine}/middleware/service.rs | 2 +- .../src/{ => engine}/middleware/websocket.rs | 2 +- backend/crates/gateway-core/src/engine/mod.rs | 12 +- .../crates/gateway-core/src/engine/nested.rs | 2 +- .../gateway-core/src/engine/observation.rs | 2 +- .../crates/gateway-core/src/engine/policy.rs | 2 +- .../src/engine/upstream_adapter.rs | 2 +- backend/crates/gateway-core/src/error/mod.rs | 4 +- backend/crates/gateway-core/src/lib.rs | 34 +- .../src/{middleware/mod.rs => middleware.rs} | 4 - .../crates/gateway-core/src/provider_ports.rs | 1645 -------------- .../src/provider_ports/affinity.rs | 174 ++ .../gateway-core/src/provider_ports/bundle.rs | 126 ++ .../src/provider_ports/catalog.rs | 167 ++ .../src/provider_ports/cooldown.rs | 219 ++ .../src/provider_ports/credential.rs | 90 + .../gateway-core/src/provider_ports/error.rs | 48 + .../gateway-core/src/provider_ports/lease.rs | 216 ++ .../gateway-core/src/provider_ports/mod.rs | 41 + .../gateway-core/src/provider_ports/oauth.rs | 184 ++ .../src/provider_ports/refresh.rs | 161 ++ .../src/provider_ports/runtime_policy.rs | 315 +++ .../src/{runtime => routing}/extensions.rs | 16 - .../crates/gateway-core/src/routing/mod.rs | 2 + .../src/routing/request_settings.rs | 374 ++++ .../gateway-core/src/routing/snapshot.rs | 12 +- .../crates/gateway-core/src/runtime/mod.rs | 544 +---- .../gateway-core/src/runtime/publications.rs | 120 + .../gateway-core/src/runtime/publisher.rs | 429 ++++ .../crates/gateway-core/src/settings/mod.rs | 374 +--- backend/crates/gateway-core/src/upstream.rs | 14 +- .../engine/{execution.rs => execution/mod.rs} | 20 +- .../gateway-core/tests/engine/extensions.rs | 77 +- .../{middleware.rs => middleware/mod.rs} | 28 +- .../gateway-core/tests/engine/policy.rs | 2 +- backend/crates/gateway-core/tests/main.rs | 1 - .../{middleware/mod.rs => middleware.rs} | 0 .../mod.rs} | 0 .../crates/gateway-core/tests/routing/mod.rs | 1 + .../mod.rs => routing/request_settings.rs} | 3 +- .../gateway-core/tests/routing/snapshot.rs | 12 +- .../crates/gateway-core/tests/runtime/mod.rs | 2 +- .../{extensions.rs => publications.rs} | 6 +- backend/crates/gateway-host/src/lib.rs | 12 +- .../gateway-host/src/system_update/config.rs | 165 ++ .../src/system_update/download.rs | 3 +- .../gateway-host/src/system_update/error.rs | 21 + .../gateway-host/src/system_update/events.rs | 130 ++ .../src/system_update/installation.rs | 6 +- .../gateway-host/src/system_update/mod.rs | 961 +------- .../src/system_update/operations.rs | 605 +++++ .../gateway-host/src/system_update/process.rs | 8 +- .../gateway-host/src/system_update/release.rs | 99 +- .../gateway-host/src/system_update/state.rs | 10 +- backend/crates/gateway-host/src/workers.rs | 24 +- .../crates/gateway-host/tests/command_line.rs | 85 + backend/crates/gateway-host/tests/main.rs | 1 + backend/crates/gateway-host/tests/serve.rs | 364 +++ .../crates/gateway-host/tests/support/host.rs | 65 + .../crates/gateway-host/tests/support/mod.rs | 1 + .../tests/system_update/download.rs | 14 + .../tests/system_update/events.rs | 154 ++ .../gateway-host/tests/system_update/mod.rs | 1815 +-------------- .../tests/system_update/operations.rs | 1083 +++++++++ .../tests/system_update/release.rs | 512 +++++ .../gateway-host/tests/system_update/state.rs | 60 + backend/crates/gateway-host/tests/workers.rs | 34 + .../runtime/src/adapter/command_line/mod.rs | 2 +- .../runtime/src/adapter/observer.rs | 2 +- .../execution.rs} | 30 +- .../adapter/policy/{ => middleware}/http.rs | 4 +- .../src/adapter/policy/middleware/mod.rs | 6 + .../policy/{ => middleware}/service.rs | 4 +- .../policy/{ => middleware}/websocket.rs | 4 +- .../runtime/src/adapter/policy/mod.rs | 3 - .../runtime/src/callback/call_resources.rs | 128 ++ .../runtime/src/callback/http_dispatch.rs | 2 +- .../src/callback/http_middleware/mod.rs | 8 +- .../src/callback/http_middleware/websocket.rs | 2 +- .../runtime/src/callback/middleware/mod.rs | 4 +- .../runtime/src/callback/mod.rs | 111 +- .../runtime/src/callback/services.rs | 18 +- .../src/callback/websocket_middleware.rs | 2 +- .../runtime/src/compatibility/fast.rs | 2 +- .../src/generation/prepare/diagnostics.rs | 255 +++ .../src/generation/prepare/instance.rs | 17 +- .../src/generation/prepare/maintenance.rs | 3 +- .../src/generation/prepare/management.rs | 137 ++ .../runtime/src/generation/prepare/mod.rs | 1271 +---------- .../runtime/src/generation/prepare/runtime.rs | 240 ++ .../runtime/src/generation/prepare/set.rs | 140 ++ .../src/generation/prepare/snapshot.rs | 343 +++ .../runtime/src/generation/prepare/state.rs | 226 ++ .../runtime/tests/generation/data.rs | 16 +- .../runtime/tests/generation/observer.rs | 20 +- .../tests/generation/policy/compatibility.rs | 12 +- .../tests/generation/policy/facts/mod.rs | 5 +- .../tests/generation/policy/http/dispatch.rs | 7 +- .../tests/generation/policy/http/mod.rs | 14 +- .../generation/policy/http/websocket/mod.rs | 6 +- .../runtime/tests/generation/policy/mod.rs | 105 +- .../runtime/tests/generation/policy/mounts.rs | 10 +- .../tests/generation/policy/service/mod.rs | 13 +- .../tests/generation/policy/websocket/mod.rs | 7 +- .../tests/generation/prepare/diagnostics.rs | 191 ++ .../tests/generation/prepare/instance.rs | 511 +++++ .../tests/generation/prepare/management.rs | 50 + .../runtime/tests/generation/prepare/mod.rs | 1024 +-------- .../tests/generation/prepare/runtime.rs | 27 + .../runtime/tests/generation/prepare/set.rs | 34 + .../tests/generation/prepare/snapshot.rs | 248 ++ .../runtime/tests/generation/private_state.rs | 2 +- .../tests/generation/upstream_adapter.rs | 131 +- .../runtime/tests/support/environment.rs | 8 +- backend/crates/gateway-plugin/sdk/README.md | 2 +- .../sdk/src/call/services/settings.rs | 2 +- .../gateway-plugin/sdk/src/client/session.rs | 1479 ------------ .../sdk/src/client/session/call.rs | 66 + .../sdk/src/client/session/config.rs | 54 + .../sdk/src/client/session/credit.rs | 109 + .../sdk/src/client/session/driver.rs | 712 ++++++ .../sdk/src/client/session/error.rs | 47 + .../sdk/src/client/session/lifecycle.rs | 116 + .../sdk/src/client/session/mod.rs | 19 + .../sdk/src/client/session/output.rs | 84 + .../sdk/src/client/session/registry.rs | 251 +++ .../sdk/src/client/session/stream.rs | 142 ++ .../gateway-plugin/sdk/tests/client/frame.rs | 87 + .../gateway-plugin/sdk/tests/client/mod.rs | 1 + .../sdk/tests/client/session.rs | 1760 --------------- .../sdk/tests/client/session/credit.rs | 271 +++ .../sdk/tests/client/session/driver.rs | 545 +++++ .../sdk/tests/client/session/lifecycle.rs | 112 + .../sdk/tests/client/session/mod.rs | 470 ++++ .../sdk/tests/client/session/registry.rs | 299 +++ backend/crates/gateway-store/src/backup/s3.rs | 114 +- .../src/postgres/client_budgets.rs | 9 +- .../gateway-store/src/postgres/connection.rs | 104 + .../src/postgres/control_plane.rs | 286 +++ .../crates/gateway-store/src/postgres/mod.rs | 397 +--- .../src/postgres/query_pattern.rs | 13 + .../crates/gateway-store/tests/backup/s3.rs | 224 +- .../tests/postgres/client_budgets.rs | 35 + .../tests/postgres/connection.rs | 189 +- .../tests/postgres/control_plane.rs | 104 + .../gateway-store/tests/postgres/mod.rs | 188 +- .../tests/postgres/runtime_settings.rs | 101 +- .../openai/src/credential/affinity.rs | 96 +- .../openai/src/credential/api_key.rs | 2 +- .../openai/src/credential/catalog.rs | 2 +- .../providers/openai/src/credential/mod.rs | 10 +- .../openai/src/credential/quota/document.rs | 72 + .../openai/src/credential/quota/mod.rs | 1987 +---------------- .../openai/src/credential/quota/passive.rs | 236 ++ .../src/credential/quota/reset_credits.rs | 315 +++ .../openai/src/credential/quota/scheduling.rs | 414 ++++ .../openai/src/credential/quota/service.rs | 756 +++++++ .../openai/src/credential/quota/warmup.rs | 260 +++ .../openai/src/credential/security.rs | 8 + backend/crates/providers/openai/src/lib.rs | 4 + .../openai/src/provider/execution.rs | 2 +- .../providers/openai/src/provider/live.rs | 2 +- .../providers/openai/src/provider/mod.rs | 14 +- .../providers/openai/src/request_identity.rs | 102 + .../providers/openai/src/transport/client.rs | 11 +- .../providers/openai/src/transport/request.rs | 2 +- .../tests/credential/contract/capacity.rs | 2 +- .../openai/tests/credential/quota/mod.rs | 1 + .../tests/credential/quota/reset_credits.rs | 245 ++ .../openai/tests/provider/contract/mod.rs | 2 +- .../provider/contract/session_binding.rs | 4 +- .../crates/providers/openai/tests/support.rs | 11 + backend/crates/providers/xai/src/admin.rs | 2 - .../providers/xai/src/credential/catalog.rs | 701 +----- .../providers/xai/src/credential/mod.rs | 11 +- .../providers/xai/src/credential/quota.rs | 669 ++++++ .../providers/xai/src/credential/selector.rs | 3 +- backend/crates/providers/xai/src/lib.rs | 2 +- .../xai/src/provider/continuation.rs | 21 +- .../providers/xai/src/provider/failure.rs | 12 +- .../crates/providers/xai/src/provider/mod.rs | 90 +- .../providers/xai/src/reasoning_replay.rs | 12 +- .../providers/xai/src/transport/canonical.rs | 40 +- .../providers/xai/src/transport/catalog.rs | 169 +- .../crates/providers/xai/src/transport/mod.rs | 1 + .../xai/src/transport/request/mod.rs | 68 +- .../xai/src/transport/request/model.rs | 144 -- .../xai/src/transport/request/parameters.rs | 76 + .../xai/src/transport/request/tools.rs | 96 +- backend/crates/providers/xai/tests/admin.rs | 20 +- .../providers/xai/tests/credential/catalog.rs | 612 +---- .../providers/xai/tests/credential/mod.rs | 1 + .../providers/xai/tests/credential/quota.rs | 788 +++++++ .../providers/xai/tests/credential/refresh.rs | 6 +- .../provider/{contract.rs => contract/mod.rs} | 285 ++- .../provider/contract/replay_tool_names.rs | 185 ++ .../xai/tests/transport/canonical.rs | 188 +- .../providers/xai/tests/transport/catalog.rs | 228 +- .../catalog/fixtures/cli_proxy_models.json | 42 + .../official_grok_models_snapshot.json | 32 - .../tests/transport/request/hosted_names.rs | 243 ++ .../xai/tests/transport/request/mod.rs | 42 +- .../tests/transport/request/model_identity.rs | 54 + .../xai/tests/transport/request/parameters.rs | 75 + .../xai/tests/transport/request/web_search.rs | 136 ++ docs/api.md | 28 +- docs/architecture.md | 16 +- 261 files changed, 22602 insertions(+), 19019 deletions(-) rename backend/crates/gateway-admin/src/{service => public_service}/mod.rs (98%) rename backend/crates/gateway-admin/src/{service => public_service}/registry.rs (97%) rename backend/crates/gateway-admin/src/{service => public_service}/settings.rs (100%) create mode 100644 backend/crates/gateway-admin/src/use_case/credential_mutation.rs create mode 100644 backend/crates/gateway-admin/src/use_case/error.rs create mode 100644 backend/crates/gateway-admin/src/use_case/publication.rs create mode 100644 backend/crates/gateway-api/src/openai/responses/websocket/connection/frame.rs create mode 100644 backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs create mode 100644 backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs create mode 100644 backend/crates/gateway-api/src/openai/responses/websocket/connection/state.rs create mode 100644 backend/crates/gateway-core/src/engine/execution/admission.rs create mode 100644 backend/crates/gateway-core/src/engine/execution/contract.rs create mode 100644 backend/crates/gateway-core/src/engine/execution/failure.rs create mode 100644 backend/crates/gateway-core/src/engine/execution/mod.rs create mode 100644 backend/crates/gateway-core/src/engine/execution/probe_store.rs create mode 100644 backend/crates/gateway-core/src/engine/execution/scope.rs rename backend/crates/gateway-core/src/engine/{execution.rs => execution/service.rs} (64%) create mode 100644 backend/crates/gateway-core/src/engine/execution/session.rs rename backend/crates/gateway-core/src/engine/{middleware.rs => middleware/contract.rs} (80%) rename backend/crates/gateway-core/src/{ => engine}/middleware/http/mod.rs (86%) rename backend/crates/gateway-core/src/{ => engine}/middleware/http/upgrade.rs (93%) create mode 100644 backend/crates/gateway-core/src/engine/middleware/mod.rs create mode 100644 backend/crates/gateway-core/src/engine/middleware/plan.rs rename backend/crates/gateway-core/src/{ => engine}/middleware/service.rs (95%) rename backend/crates/gateway-core/src/{ => engine}/middleware/websocket.rs (93%) rename backend/crates/gateway-core/src/{middleware/mod.rs => middleware.rs} (97%) delete mode 100644 backend/crates/gateway-core/src/provider_ports.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/affinity.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/bundle.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/catalog.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/cooldown.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/credential.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/error.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/lease.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/mod.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/oauth.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/refresh.rs create mode 100644 backend/crates/gateway-core/src/provider_ports/runtime_policy.rs rename backend/crates/gateway-core/src/{runtime => routing}/extensions.rs (82%) create mode 100644 backend/crates/gateway-core/src/routing/request_settings.rs create mode 100644 backend/crates/gateway-core/src/runtime/publications.rs create mode 100644 backend/crates/gateway-core/src/runtime/publisher.rs rename backend/crates/gateway-core/tests/engine/{execution.rs => execution/mod.rs} (99%) rename backend/crates/gateway-core/tests/engine/{middleware.rs => middleware/mod.rs} (95%) rename backend/crates/gateway-core/tests/{middleware/mod.rs => middleware.rs} (100%) rename backend/crates/gateway-core/tests/{provider_ports.rs => provider_ports/mod.rs} (100%) rename backend/crates/gateway-core/tests/{settings/mod.rs => routing/request_settings.rs} (99%) rename backend/crates/gateway-core/tests/runtime/{extensions.rs => publications.rs} (96%) create mode 100644 backend/crates/gateway-host/src/system_update/config.rs create mode 100644 backend/crates/gateway-host/src/system_update/error.rs create mode 100644 backend/crates/gateway-host/src/system_update/events.rs create mode 100644 backend/crates/gateway-host/src/system_update/operations.rs create mode 100644 backend/crates/gateway-host/tests/command_line.rs create mode 100644 backend/crates/gateway-host/tests/support/host.rs create mode 100644 backend/crates/gateway-host/tests/system_update/download.rs create mode 100644 backend/crates/gateway-host/tests/system_update/events.rs create mode 100644 backend/crates/gateway-host/tests/system_update/operations.rs create mode 100644 backend/crates/gateway-host/tests/system_update/release.rs create mode 100644 backend/crates/gateway-host/tests/system_update/state.rs rename backend/crates/gateway-plugin/runtime/src/adapter/policy/{middleware.rs => middleware/execution.rs} (96%) rename backend/crates/gateway-plugin/runtime/src/adapter/policy/{ => middleware}/http.rs (97%) create mode 100644 backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/mod.rs rename backend/crates/gateway-plugin/runtime/src/adapter/policy/{ => middleware}/service.rs (96%) rename backend/crates/gateway-plugin/runtime/src/adapter/policy/{ => middleware}/websocket.rs (95%) create mode 100644 backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs create mode 100644 backend/crates/gateway-plugin/runtime/src/generation/prepare/diagnostics.rs create mode 100644 backend/crates/gateway-plugin/runtime/src/generation/prepare/management.rs create mode 100644 backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs create mode 100644 backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs create mode 100644 backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs create mode 100644 backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs create mode 100644 backend/crates/gateway-plugin/runtime/tests/generation/prepare/diagnostics.rs create mode 100644 backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs create mode 100644 backend/crates/gateway-plugin/runtime/tests/generation/prepare/management.rs create mode 100644 backend/crates/gateway-plugin/runtime/tests/generation/prepare/runtime.rs create mode 100644 backend/crates/gateway-plugin/runtime/tests/generation/prepare/set.rs create mode 100644 backend/crates/gateway-plugin/runtime/tests/generation/prepare/snapshot.rs delete mode 100644 backend/crates/gateway-plugin/sdk/src/client/session.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/call.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/config.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/credit.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/driver.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/error.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/lifecycle.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/mod.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/output.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/registry.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/client/session/stream.rs create mode 100644 backend/crates/gateway-plugin/sdk/tests/client/frame.rs delete mode 100644 backend/crates/gateway-plugin/sdk/tests/client/session.rs create mode 100644 backend/crates/gateway-plugin/sdk/tests/client/session/credit.rs create mode 100644 backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs create mode 100644 backend/crates/gateway-plugin/sdk/tests/client/session/lifecycle.rs create mode 100644 backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs create mode 100644 backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs create mode 100644 backend/crates/gateway-store/src/postgres/connection.rs create mode 100644 backend/crates/gateway-store/src/postgres/control_plane.rs create mode 100644 backend/crates/gateway-store/src/postgres/query_pattern.rs create mode 100644 backend/crates/gateway-store/tests/postgres/control_plane.rs create mode 100644 backend/crates/providers/openai/src/credential/quota/passive.rs create mode 100644 backend/crates/providers/openai/src/credential/quota/reset_credits.rs create mode 100644 backend/crates/providers/openai/src/credential/quota/scheduling.rs create mode 100644 backend/crates/providers/openai/src/credential/quota/service.rs create mode 100644 backend/crates/providers/openai/src/credential/quota/warmup.rs create mode 100644 backend/crates/providers/openai/src/request_identity.rs create mode 100644 backend/crates/providers/openai/tests/credential/quota/reset_credits.rs create mode 100644 backend/crates/providers/xai/src/credential/quota.rs delete mode 100644 backend/crates/providers/xai/src/transport/request/model.rs create mode 100644 backend/crates/providers/xai/src/transport/request/parameters.rs create mode 100644 backend/crates/providers/xai/tests/credential/quota.rs rename backend/crates/providers/xai/tests/provider/{contract.rs => contract/mod.rs} (94%) create mode 100644 backend/crates/providers/xai/tests/provider/contract/replay_tool_names.rs create mode 100644 backend/crates/providers/xai/tests/transport/catalog/fixtures/cli_proxy_models.json delete mode 100644 backend/crates/providers/xai/tests/transport/catalog/fixtures/official_grok_models_snapshot.json create mode 100644 backend/crates/providers/xai/tests/transport/request/hosted_names.rs create mode 100644 backend/crates/providers/xai/tests/transport/request/model_identity.rs create mode 100644 backend/crates/providers/xai/tests/transport/request/parameters.rs create mode 100644 backend/crates/providers/xai/tests/transport/request/web_search.rs diff --git a/backend/apps/gateway/src/bootstrap/command.rs b/backend/apps/gateway/src/bootstrap/command.rs index 354db1ba0..058cdc832 100644 --- a/backend/apps/gateway/src/bootstrap/command.rs +++ b/backend/apps/gateway/src/bootstrap/command.rs @@ -93,7 +93,7 @@ async fn execute( Some(plugin_runtime.clone()), Some(plugin_runtime.observer_registry()), Some(plugin_runtime.policy_registry()), - Some(plugin_runtime.middleware_registry()), + Some(plugin_runtime.execution_registry()), ); let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control()); or_shutdown!(plugin_runtime, management.bind(&plugin_runtime)); @@ -110,7 +110,7 @@ async fn execute( providers.admin, std::sync::Arc::new(gateway_host::pricing::ModelsDevPricing), ); - let mut services = gateway_admin::service::Registry::new(service_middleware( + let mut services = gateway_admin::public_service::Registry::new(service_middleware( core.snapshots(), &plugin_runtime, )); diff --git a/backend/apps/gateway/src/bootstrap/plugins.rs b/backend/apps/gateway/src/bootstrap/plugins.rs index d6cc8e6f2..00ee71caa 100644 --- a/backend/apps/gateway/src/bootstrap/plugins.rs +++ b/backend/apps/gateway/src/bootstrap/plugins.rs @@ -72,11 +72,11 @@ impl ManagementPorts { pub(super) fn service_middleware( snapshots: RuntimeSnapshotHandle, runtime: &PluginRuntime, -) -> gateway_admin::service::PlanSource { - let middleware = runtime.middleware_registry(); +) -> gateway_admin::public_service::PlanSource { + let middleware = runtime.execution_registry(); Arc::new(move || { let snapshot = snapshots.snapshot_for_diagnostics()?; - middleware.resolve(snapshot.extensions()?) + middleware.middleware(snapshot.extensions()?) }) } diff --git a/backend/apps/gateway/src/bootstrap/server.rs b/backend/apps/gateway/src/bootstrap/server.rs index 569f032ec..28cae2fcc 100644 --- a/backend/apps/gateway/src/bootstrap/server.rs +++ b/backend/apps/gateway/src/bootstrap/server.rs @@ -28,7 +28,7 @@ pub async fn run() -> Result<(), BootstrapError> { Some(plugin_runtime.clone()), Some(plugin_runtime.observer_registry()), Some(plugin_runtime.policy_registry()), - Some(plugin_runtime.middleware_registry()), + Some(plugin_runtime.execution_registry()), Some(plugin_runtime.frontend_authentication_registry()), ); let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control()); @@ -122,14 +122,14 @@ pub async fn run() -> Result<(), BootstrapError> { ) .with_middleware({ let snapshots = core.snapshots(); - let middleware = plugin_runtime.middleware_registry(); + let middleware = plugin_runtime.execution_registry(); move |snapshot| { if let Some(snapshot) = snapshot { - return middleware.resolve(snapshot.extensions()?); + return middleware.middleware(snapshot.extensions()?); } // 数据面未就绪时，管理与诊断路由仍可使用发布候选中的插件 let diagnostic = snapshots.snapshot_for_diagnostics()?; - middleware.resolve(diagnostic.extensions()?) + middleware.middleware(diagnostic.extensions()?) } }); let http_dispatcher = api.dispatcher(); diff --git a/backend/apps/gateway/tests/architecture.rs b/backend/apps/gateway/tests/architecture.rs index 79e31f542..cfafa8811 100644 --- a/backend/apps/gateway/tests/architecture.rs +++ b/backend/apps/gateway/tests/architecture.rs @@ -79,18 +79,25 @@ fn core_value_owners_do_not_depend_on_execution_or_routing() { path if path.starts_with("policy/") => { Some(&["account", "identity", "policy", "validation"]) } - "settings/values.rs" | "settings/compiled.rs" => { - Some(&["account", "concurrency", "identity", "metering", "policy"]) - } + path if path.starts_with("settings/") => Some(&[ + "account", + "concurrency", + "identity", + "metering", + "policy", + "settings", + ]), path if path.starts_with("account/") => Some(&["account", "identity", "validation"]), _ => None, }; let Some(allowed) = allowed else { continue }; let source = fs::read_to_string(root.join(&relative)).expect("read core source"); let syntax = syn::parse_file(&source).expect("parse core source"); - let mut references = CrateReferences::default(); + let mut references = CrateReferences::in_module(&relative); references.visit_file(&syntax); - for dependency in references.0 { + // 相对路径中的 self 指向当前 owner，不是跨 owner 依赖 + references.owners.remove(&references.module_path[0]); + for dependency in references.owners { assert!( allowed.contains(&dependency.as_str()), "{} must not depend on crate::{dependency}", @@ -100,49 +107,155 @@ fn core_value_owners_do_not_depend_on_execution_or_routing() { } } +#[test] +fn core_owner_paths_form_a_dag() { + let root = backend_root().join("crates/gateway-core/src"); + let mut graph = BTreeMap::>::new(); + for relative in super::rust_files(&root) { + if relative == Path::new("lib.rs") { + continue; + } + let mut references = CrateReferences::in_module(&relative); + let owner = references.module_path[0].clone(); + let source = fs::read_to_string(root.join(&relative)).expect("read core source"); + references.visit_file(&syn::parse_file(&source).expect("parse core source")); + references.owners.remove(&owner); + graph.entry(owner).or_default().extend(references.owners); + } + for owner in graph.keys() { + let mut pending = graph[owner].iter().cloned().collect::>(); + let mut visited = BTreeSet::new(); + while let Some(dependency) = pending.pop() { + assert_ne!( + &dependency, owner, + "Core owner dependency cycle reaches {owner}" + ); + if visited.insert(dependency.clone()) + && let Some(next) = graph.get(&dependency) + { + pending.extend(next.iter().cloned()); + } + } + } +} + +#[test] +fn core_reference_scan_resolves_groups_reexports_and_relative_paths() { + let syntax = syn::parse_file( + r#" + pub use crate::{engine::AttemptContext, settings::SettingsValues as Values}; + use super::super::runtime as publication; + use self::compiled::Value; + // crate::ignored::Comment + const EXAMPLE: &str = "crate::ignored::Literal"; + "#, + ) + .expect("parse reference fixture"); + let mut references = CrateReferences::in_module(Path::new("routing/snapshot.rs")); + references.visit_file(&syntax); + assert_eq!( + references.owners, + BTreeSet::from([ + "engine".to_owned(), + "settings".to_owned(), + "runtime".to_owned(), + "routing".to_owned(), + ]) + ); +} + #[derive(Default)] -struct CrateReferences(BTreeSet); +struct CrateReferences { + owners: BTreeSet, + module_path: Vec, +} impl<'ast> Visit<'ast> for CrateReferences { fn visit_path(&mut self, path: &'ast syn::Path) { - let mut segments = path.segments.iter(); - if segments - .next() - .is_some_and(|segment| segment.ident == "crate") - && let Some(owner) = segments.next() - { - self.0.insert(owner.ident.to_string()); - } + self.record( + &path + .segments + .iter() + .map(|part| part.ident.to_string()) + .collect::>(), + ); syn::visit::visit_path(self, path); } fn visit_item_use(&mut self, item: &'ast syn::ItemUse) { - self.use_tree(&item.tree, false); + self.use_tree(&item.tree, &mut Vec::new()); } } impl CrateReferences { - fn use_tree(&mut self, tree: &syn::UseTree, crate_root: bool) { + fn in_module(relative: &Path) -> Self { + let mut module = relative + .with_extension("") + .components() + .map(|part| { + part.as_os_str() + .to_str() + .expect("UTF-8 module path") + .to_owned() + }) + .collect::>(); + if module.last().is_some_and(|part| part == "mod") { + module.pop(); + } + Self { + owners: BTreeSet::new(), + module_path: module, + } + } + + fn record(&mut self, path: &[String]) { + let Some(first) = path.first() else { return }; + let owner = if first == "crate" { + path.get(1).cloned() + } else if first == "self" || first == "super" { + let mut module = self.module_path.clone(); + let mut rest = path; + while rest + .first() + .is_some_and(|part| part == "super" || part == "self") + { + if rest[0] == "super" { + module.pop(); + } + rest = &rest[1..]; + } + module.first().or_else(|| rest.first()).cloned() + } else { + None + }; + if let Some(owner) = owner { + self.owners.insert(owner); + } + } + + fn use_tree(&mut self, tree: &syn::UseTree, prefix: &mut Vec) { match tree { - syn::UseTree::Path(path) if crate_root => { - self.0.insert(path.ident.to_string()); + syn::UseTree::Path(path) => { + prefix.push(path.ident.to_string()); + self.use_tree(&path.tree, prefix); + prefix.pop(); } - syn::UseTree::Name(name) if crate_root => { - self.0.insert(name.ident.to_string()); + syn::UseTree::Name(name) => { + prefix.push(name.ident.to_string()); + self.record(prefix); + prefix.pop(); } - syn::UseTree::Rename(rename) if crate_root => { - self.0.insert(rename.ident.to_string()); + syn::UseTree::Rename(rename) => { + prefix.push(rename.ident.to_string()); + self.record(prefix); + prefix.pop(); } - syn::UseTree::Path(path) => self.use_tree(&path.tree, path.ident == "crate"), syn::UseTree::Group(group) => { for item in &group.items { - self.use_tree(item, crate_root); + self.use_tree(item, prefix); } } - syn::UseTree::Glob(_) if crate_root => { - self.0.insert("*".to_owned()); - } - _ => {} + syn::UseTree::Glob(_) => self.record(prefix), } } } @@ -221,7 +334,7 @@ const ADAPTER_PUBLIC_MODULES: &[(&str, &[&str])] = &[ "crates/providers/openai", &["config", "credential", "transport"], ), - ("crates/providers/xai", &["credential", "transport"]), + ("crates/providers/xai", &["transport"]), ]; /// 不对应单一生产模块、而是校验 crate/workspace 整体契约的根级测试场景 diff --git a/backend/crates/gateway-admin/src/lib.rs b/backend/crates/gateway-admin/src/lib.rs index 4b5df63bb..9a89db5e4 100644 --- a/backend/crates/gateway-admin/src/lib.rs +++ b/backend/crates/gateway-admin/src/lib.rs @@ -20,7 +20,7 @@ pub mod backup; pub mod freeze_recovery; pub mod model; pub mod ports; -pub mod service; +pub mod public_service; mod use_case; pub use use_case::plugins::{PluginDistributionPorts, PluginManagementService, PluginsService}; @@ -205,7 +205,7 @@ pub enum AdminConfigError { #[derive(Clone)] pub struct AdminServices { timezone: gateway_core::time::DeploymentTimeZone, - public_services: Arc, + public_services: Arc, plugins: Arc, plugin_management: Arc, proxies: Arc, @@ -231,7 +231,7 @@ impl AdminServices { self.timezone } - pub fn public_services(&self) -> Arc { + pub fn public_services(&self) -> Arc { self.public_services.clone() } @@ -331,7 +331,7 @@ impl AdminBundle { /// 组合根提供给控制面的运行能力；与配置和存储端口分别传入 pub struct AdminRuntimePorts { pub timezone: gateway_core::time::DeploymentTimeZone, - pub service_middleware: service::PlanSource, + pub service_middleware: public_service::PlanSource, pub plugin_preparation: Arc, pub plugin_management: Arc, pub published_snapshot: gateway_core::runtime::RuntimeSnapshotHandle, @@ -466,7 +466,7 @@ async fn initialize_inner( registry.clone(), pricing_source, ); - let mut public_services = service::Registry::new(service_middleware); + let mut public_services = public_service::Registry::new(service_middleware); public_services.register_settings(&settings)?; let services = AdminServices { timezone, diff --git a/backend/crates/gateway-admin/src/ports/plugin_management.rs b/backend/crates/gateway-admin/src/ports/plugin_management.rs index c6a898875..9ac9f1e75 100644 --- a/backend/crates/gateway-admin/src/ports/plugin_management.rs +++ b/backend/crates/gateway-admin/src/ports/plugin_management.rs @@ -1,7 +1,7 @@ //! 插件管理页面、资源访问与授权回调的运行时端口 use async_trait::async_trait; -use gateway_core::runtime::extensions::ExtensionSetReference; +use gateway_core::routing::extensions::ExtensionSetReference; use crate::model::{ AdminError, diff --git a/backend/crates/gateway-admin/src/ports/plugins.rs b/backend/crates/gateway-admin/src/ports/plugins.rs index b1a71f7ee..5648a4b72 100644 --- a/backend/crates/gateway-admin/src/ports/plugins.rs +++ b/backend/crates/gateway-admin/src/ports/plugins.rs @@ -26,7 +26,7 @@ use crate::model::{ PluginCompatibilityRequirements, PluginSource, }, }; -use gateway_core::runtime::extensions::ExtensionSetReference; +use gateway_core::routing::extensions::ExtensionSetReference; /// Admin 准备完整候选并保活到提交后的发布结束；进程状态不能代替持久启用状态 #[async_trait] diff --git a/backend/crates/gateway-admin/src/service/mod.rs b/backend/crates/gateway-admin/src/public_service/mod.rs similarity index 98% rename from backend/crates/gateway-admin/src/service/mod.rs rename to backend/crates/gateway-admin/src/public_service/mod.rs index 382d0fccb..a82142cee 100644 --- a/backend/crates/gateway-admin/src/service/mod.rs +++ b/backend/crates/gateway-admin/src/public_service/mod.rs @@ -8,11 +8,11 @@ use std::{future::Future, sync::Arc}; use futures::future::BoxFuture; use gateway_core::{ - engine::middleware::FrozenMiddlewarePlan, - middleware::{ - compose, + engine::middleware::{ + FrozenMiddlewarePlan, service::{Context, Error, Value}, }, + middleware::compose, }; use crate::model::{AdminError, AdminErrorKind}; diff --git a/backend/crates/gateway-admin/src/service/registry.rs b/backend/crates/gateway-admin/src/public_service/registry.rs similarity index 97% rename from backend/crates/gateway-admin/src/service/registry.rs rename to backend/crates/gateway-admin/src/public_service/registry.rs index 760d123d1..e6b1d8cf7 100644 --- a/backend/crates/gateway-admin/src/service/registry.rs +++ b/backend/crates/gateway-admin/src/public_service/registry.rs @@ -2,7 +2,7 @@ use crate::model::AdminError; use futures::future::BoxFuture; -use gateway_core::middleware::service::{Error, Value}; +use gateway_core::engine::middleware::service::{Error, Value}; use serde::{Serialize, de::DeserializeOwned}; use std::{collections::BTreeMap, sync::Arc}; diff --git a/backend/crates/gateway-admin/src/service/settings.rs b/backend/crates/gateway-admin/src/public_service/settings.rs similarity index 100% rename from backend/crates/gateway-admin/src/service/settings.rs rename to backend/crates/gateway-admin/src/public_service/settings.rs diff --git a/backend/crates/gateway-admin/src/use_case/credential_mutation.rs b/backend/crates/gateway-admin/src/use_case/credential_mutation.rs new file mode 100644 index 000000000..d5cc1d816 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/credential_mutation.rs @@ -0,0 +1,371 @@ +//! 凭据变更共享的绑定校验、事务结算、失效与提交后观察 + +use crate::{ + model::{ + AdminError, MutationContext, + accounts::DeleteAccounts, + provider_credentials::{ + AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult, + CredentialDetails, CredentialMutationResult, CredentialRotationCommit, + PendingAuthorizationMutation, PreparedAuthorizationCommit, + PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation, + ProviderQuotaRequest, StartAuthorization, + }, + }, + ports::{provider::ProviderAdmin, store::AccountStore}, +}; +use gateway_core::{account::ProviderAccountId, routing::ProviderKind, runtime::SnapshotControl}; + +use super::{map_store_error, publish_committed}; + +pub(super) async fn publish_credentials_and_observe_quota( + provider: &std::sync::Arc, + snapshot: &dyn SnapshotControl, + revision: crate::model::Revision, + account_ids: &[ProviderAccountId], + request_id: &str, +) -> Result<(), AdminError> { + provider.account_facts_changed(account_ids).await; + publish_committed(snapshot, revision).await?; + + // 凭据已经提交；额度只是可重建的观察，不能拖住管理请求或回滚提交结果 + let provider = provider.clone(); + let account_ids = account_ids.to_vec(); + let request_id = request_id.to_owned(); + tokio::spawn(async move { + for account_id in account_ids { + if let Err(error) = provider + .quota(ProviderQuotaRequest { + account_id: account_id.clone(), + refresh: true, + rolling_usage: None, + }) + .await + { + tracing::warn!( + request_id, + provider = %provider.provider_kind().as_str(), + account_id = %account_id.as_str(), + quota_error = ?error.kind(), + "initial quota observation failed" + ); + } + } + }); + Ok(()) +} + +pub(super) async fn required_credential( + accounts: &dyn AccountStore, + provider_kind: &ProviderKind, + account_id: &ProviderAccountId, + resource: &'static str, +) -> Result { + accounts + .credential_details(provider_kind, account_id) + .await + .map_err(|error| map_store_error(error, resource))? + .ok_or_else(|| AdminError::not_found("Provider 凭据不存在")) +} + +pub(super) async fn required_plugin_credential( + accounts: &dyn AccountStore, + account_id: &ProviderAccountId, + resource: &'static str, +) -> Result { + accounts + .credential_details_by_id(account_id) + .await + .map_err(|error| map_store_error(error, resource))? + .ok_or_else(|| AdminError::not_found("Provider 凭据不存在")) +} + +pub(super) async fn import_proxy_binding( + proxies: &dyn crate::ports::proxy::ProxyStore, + id: Option<&str>, +) -> Result, AdminError> { + let Some(id) = id else { return Ok(None) }; + let reservation = proxies + .reserve_import(id) + .await + .map_err(|error| map_store_error(error, "proxy"))?; + Ok(Some(reservation)) +} + +pub(super) async fn pending_authorization( + accounts: &dyn AccountStore, + proxies: &dyn crate::ports::proxy::ProxyStore, + provider_kind: &ProviderKind, + command: &StartAuthorization, + resource: &'static str, +) -> Result { + let proxy; + let mut proxy_id = None; + let target = match &command.reauthorization { + Some(account_id) => { + proxy = required_credential(accounts, provider_kind, account_id, resource) + .await? + .credential + .outbound_proxy; + AuthorizationMutationTarget::Reauthorize { + account_id: account_id.clone(), + } + } + None => { + use crate::model::proxies::AccountProxySelection; + proxy = match &command.outbound_proxy { + Some(AccountProxySelection::Url(value)) => Some(value.clone()), + Some(AccountProxySelection::Saved(id)) => { + let record = proxies + .get(id) + .await + .map_err(|error| map_store_error(error, "proxy"))?; + proxy_id = Some(record.id); + Some(record.proxy) + } + Some(AccountProxySelection::Direct) | None => None, + }; + AuthorizationMutationTarget::Create { + name: command.name.clone(), + } + } + }; + Ok(PendingAuthorizationMutation::new( + provider_kind.clone(), + target, + crate::model::provider_credentials::AuthorizationOwnerBinding::from_context( + &command.context, + ), + ) + .with_outbound_proxy(proxy) + .with_outbound_proxy_id(proxy_id)) +} + +pub(super) fn validate_prepared_import( + provider_kind: &ProviderKind, + prepared: &PreparedCredentialImport, + _resource: &'static str, +) -> Result<(), AdminError> { + if prepared.provider_kind != *provider_kind + || prepared + .credentials + .iter() + .any(|credential| credential.provider_kind != *provider_kind) + { + return Err(AdminError::conflict("Provider 准备结果与请求范围不一致")); + } + Ok(()) +} + +pub(super) fn validate_prepared_rotation( + account: &crate::model::accounts::AccountRecord, + prepared: &PreparedCredentialRotation, + _resource: &'static str, +) -> Result<(), AdminError> { + validate_prepared_rotation_facts(account, prepared.facts()) +} + +pub(super) fn validate_prepared_rotation_facts( + account: &crate::model::accounts::AccountRecord, + facts: &crate::model::provider_credentials::PreparedCredentialRotationFacts, +) -> Result<(), AdminError> { + if facts.account_id.as_str() != account.id.as_str() + || facts.provider_kind != account.provider_kind + { + return Err(AdminError::conflict("Provider 准备结果与当前凭据不一致")); + } + Ok(()) +} + +pub(super) async fn validate_authorization_commit( + provider_kind: &ProviderKind, + context: &MutationContext, + prepared: PreparedAuthorizationCommit, + resource: &'static str, +) -> Result { + if prepared.pending.provider_kind() != provider_kind + || !prepared.pending.owner_binding().matches_context(context) + { + let validation_error = AdminError::conflict("OAuth 授权绑定已失效，请重新发起授权"); + if let Err(error) = prepared.abort().await { + tracing::warn!( + resource, + settlement_error = %error, + "OAuth authorization claim release failed after preparation validation" + ); + return Err(error); + } + return Err(validation_error); + } + let matches_target = match (prepared.pending.target(), &prepared.credential) { + ( + AuthorizationMutationTarget::Create { .. }, + PreparedAuthorizationCredential::Create(credential), + ) => credential.provider_kind == *provider_kind, + ( + AuthorizationMutationTarget::Reauthorize { account_id }, + PreparedAuthorizationCredential::Reauthorize(credential), + ) => { + let facts = credential.facts(); + facts.provider_kind == *provider_kind && facts.account_id == *account_id + } + _ => false, + }; + if !matches_target { + let validation_error = AdminError::conflict("OAuth 凭据与待提交目标不一致"); + if let Err(error) = prepared.abort().await { + tracing::warn!( + resource, + settlement_error = %error, + "OAuth authorization claim release failed after preparation validation" + ); + return Err(error); + } + return Err(validation_error); + } + Ok(prepared) +} + +pub(super) async fn commit_authorization( + accounts: &dyn AccountStore, + prepared: PreparedAuthorizationCommit, + key: crate::model::provider_credentials::AuthorizationReceiptKey, + settings: Option, + context: &MutationContext, + resource: &'static str, +) -> Result { + if settings.is_some() + && matches!( + &prepared.credential, + PreparedAuthorizationCredential::Reauthorize(_) + ) + { + prepared.abort().await?; + return Err(AdminError::invalid("重新授权不能修改账号设置")); + } + let crate::model::provider_credentials::AuthorizationCommitSettlement { + command, + credential_guard, + authorization_guard, + } = prepared.into_commit(settings, key); + match accounts.commit_authorization(command, context).await { + Ok(result) => { + if let Some(guard) = credential_guard + && result.newly_committed + { + guard.finish(); + } + if let Some(guard) = authorization_guard + && let Err(error) = guard.commit().await + { + // 账号与回执已经原子提交，Redis 清理失败不能把确定的成功改写为失败 + tracing::warn!(resource, settlement_error = %error, "authorization committed but pending cleanup failed"); + } + Ok(result) + } + Err(error) => { + drop(credential_guard); + if let Some(guard) = authorization_guard + && let Err(settlement_error) = guard.abort().await + { + tracing::warn!( + resource, + store_error = %error, + settlement_error = %settlement_error, + "OAuth authorization claim release failed after Store commit failure" + ); + return Err(settlement_error); + } + Err(map_store_error(error, resource)) + } + } +} + +pub(super) async fn commit_credential_rotation( + accounts: &dyn AccountStore, + prepared: PreparedCredentialRotation, + settings: Option, + context: &MutationContext, + resource: &'static str, +) -> Result { + let (facts, guard) = prepared.into_parts(); + match accounts + .commit_credential_rotation( + CredentialRotationCommit { + prepared: facts, + settings, + }, + context, + ) + .await + { + Ok(result) => { + guard.finish(); + Ok(result) + } + Err(error) => { + drop(guard); + Err(map_store_error(error, resource)) + } + } +} + +pub(super) async fn commit_credential_refresh( + accounts: &dyn AccountStore, + prepared: PreparedCredentialRotation, + context: &MutationContext, + resource: &'static str, +) -> Result { + let (facts, guard) = prepared.into_parts(); + match accounts + .commit_credential_refresh( + CredentialRotationCommit { + prepared: facts, + settings: None, + }, + context, + ) + .await + { + Ok(result) => { + guard.finish(); + Ok(result) + } + Err(error) => { + drop(guard); + Err(map_store_error(error, resource)) + } + } +} + +pub(super) async fn delete_credentials( + accounts: &dyn AccountStore, + provider: &dyn ProviderAdmin, + command: CredentialDeletion, + resource: &'static str, +) -> Result { + for account_id in &command.account_ids { + required_credential(accounts, provider.provider_kind(), account_id, resource).await?; + } + let account_ids = command.account_ids; + let revision = accounts + .delete_accounts( + DeleteAccounts { + account_ids: account_ids + .iter() + .map(|account_id| account_id.as_str().to_owned()) + .collect(), + }, + &command.context, + ) + .await + .map_err(|error| map_store_error(error, resource))?; + for account_id in &account_ids { + provider.account_unavailable(account_id).await; + } + provider.account_facts_changed(&account_ids).await; + Ok(CredentialDeletionResult { + config_revision: revision, + account_ids, + }) +} diff --git a/backend/crates/gateway-admin/src/use_case/error.rs b/backend/crates/gateway-admin/src/use_case/error.rs new file mode 100644 index 000000000..08f9065ef --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/error.rs @@ -0,0 +1,67 @@ +//! 管理端口错误到用例错误的统一分类与来源保留 + +use crate::model::{AdminError, AdminErrorKind}; +use crate::ports::store::{AdminStoreError, AdminStoreErrorKind}; + +pub(super) fn map_store_error(error: AdminStoreError, resource: &'static str) -> AdminError { + let kind = match error.kind() { + AdminStoreErrorKind::Invalid => AdminErrorKind::Invalid, + AdminStoreErrorKind::NotFound => AdminErrorKind::NotFound, + AdminStoreErrorKind::StaleRevision + | AdminStoreErrorKind::DuplicateName + | AdminStoreErrorKind::Conflict => AdminErrorKind::Conflict, + AdminStoreErrorKind::Unavailable => AdminErrorKind::Unavailable, + }; + let message = match kind { + AdminErrorKind::Invalid => "请求参数不合法", + AdminErrorKind::NotFound => "请求的资源不存在", + AdminErrorKind::Conflict => "当前资源状态冲突，请刷新后重试", + AdminErrorKind::Unavailable => "依赖服务暂不可用", + _ => "服务内部错误", + }; + AdminError::new(kind, message).with_source(OperationFailure { + operation: resource, + source: error.into(), + }) +} + +pub(super) fn map_provider_error( + error: crate::ports::provider::ProviderAdminError, + resource: &'static str, +) -> AdminError { + use crate::ports::provider::ProviderAdminErrorKind; + + let kind = match error.kind() { + ProviderAdminErrorKind::Invalid => AdminErrorKind::Invalid, + ProviderAdminErrorKind::Unsupported => AdminErrorKind::Invalid, + ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound, + ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict, + ProviderAdminErrorKind::Ambiguous => AdminErrorKind::UpstreamResultUnknown, + ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable, + ProviderAdminErrorKind::CredentialRefreshRequired => AdminErrorKind::Unavailable, + ProviderAdminErrorKind::BadGateway => AdminErrorKind::BadGateway, + ProviderAdminErrorKind::Internal => AdminErrorKind::Internal, + }; + let message = match kind { + AdminErrorKind::Invalid => "Provider 请求不合法", + AdminErrorKind::NotFound => "Provider 资源不存在", + AdminErrorKind::Conflict => "Provider 资源状态冲突，请刷新后重试", + AdminErrorKind::BadGateway => "上游服务请求失败", + AdminErrorKind::UpstreamResultUnknown => "上游执行结果未知，请刷新状态后再决定是否重试", + AdminErrorKind::Unavailable => "Provider 服务暂不可用", + AdminErrorKind::Internal => "服务内部错误", + _ => "Provider 操作失败", + }; + AdminError::new(kind, error.public_message().unwrap_or(message)).with_source(OperationFailure { + operation: resource, + source: error.into(), + }) +} + +// 用例的操作名补充端口资源上下文，底层分类与原始原因仍由来源持有 +#[derive(Debug, thiserror::Error)] +#[error("admin operation failed: {operation}")] +struct OperationFailure { + operation: &'static str, + source: gateway_core::error::ErrorSource, +} diff --git a/backend/crates/gateway-admin/src/use_case/mod.rs b/backend/crates/gateway-admin/src/use_case/mod.rs index 8ed47dccd..877800c7d 100644 --- a/backend/crates/gateway-admin/src/use_case/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/mod.rs @@ -19,450 +19,16 @@ pub mod proxies; pub mod settings; pub mod system; -use crate::{ - model::{ - AdminError, AdminErrorKind, MutationContext, - accounts::DeleteAccounts, - provider_credentials::{ - AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult, - CredentialDetails, CredentialMutationResult, CredentialRotationCommit, - PendingAuthorizationMutation, PreparedAuthorizationCommit, - PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation, - ProviderQuotaRequest, StartAuthorization, - }, - }, - ports::{ - provider::ProviderAdmin, - store::{AccountStore, AdminStoreError, AdminStoreErrorKind}, - }, +mod credential_mutation; +mod error; +mod publication; + +use credential_mutation::{ + commit_authorization, commit_credential_refresh, commit_credential_rotation, + delete_credentials, import_proxy_binding, pending_authorization, + publish_credentials_and_observe_quota, required_credential, required_plugin_credential, + validate_authorization_commit, validate_prepared_import, validate_prepared_rotation, + validate_prepared_rotation_facts, }; -use gateway_core::{ - account::ProviderAccountId, - routing::{ConfigRevision, ProviderKind}, - runtime::SnapshotControl, -}; - -fn map_store_error(error: AdminStoreError, resource: &'static str) -> AdminError { - let kind = match error.kind() { - AdminStoreErrorKind::Invalid => AdminErrorKind::Invalid, - AdminStoreErrorKind::NotFound => AdminErrorKind::NotFound, - AdminStoreErrorKind::StaleRevision - | AdminStoreErrorKind::DuplicateName - | AdminStoreErrorKind::Conflict => AdminErrorKind::Conflict, - AdminStoreErrorKind::Unavailable => AdminErrorKind::Unavailable, - }; - let message = match kind { - AdminErrorKind::Invalid => "请求参数不合法", - AdminErrorKind::NotFound => "请求的资源不存在", - AdminErrorKind::Conflict => "当前资源状态冲突，请刷新后重试", - AdminErrorKind::Unavailable => "依赖服务暂不可用", - _ => "服务内部错误", - }; - AdminError::new(kind, message).with_source(OperationFailure { - operation: resource, - source: error.into(), - }) -} - -fn map_provider_error( - error: crate::ports::provider::ProviderAdminError, - resource: &'static str, -) -> AdminError { - use crate::ports::provider::ProviderAdminErrorKind; - - let kind = match error.kind() { - ProviderAdminErrorKind::Invalid => AdminErrorKind::Invalid, - ProviderAdminErrorKind::Unsupported => AdminErrorKind::Invalid, - ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound, - ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict, - ProviderAdminErrorKind::Ambiguous => AdminErrorKind::UpstreamResultUnknown, - ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable, - ProviderAdminErrorKind::CredentialRefreshRequired => AdminErrorKind::Unavailable, - ProviderAdminErrorKind::BadGateway => AdminErrorKind::BadGateway, - ProviderAdminErrorKind::Internal => AdminErrorKind::Internal, - }; - let message = match kind { - AdminErrorKind::Invalid => "Provider 请求不合法", - AdminErrorKind::NotFound => "Provider 资源不存在", - AdminErrorKind::Conflict => "Provider 资源状态冲突，请刷新后重试", - AdminErrorKind::BadGateway => "上游服务请求失败", - AdminErrorKind::UpstreamResultUnknown => "上游执行结果未知，请刷新状态后再决定是否重试", - AdminErrorKind::Unavailable => "Provider 服务暂不可用", - AdminErrorKind::Internal => "服务内部错误", - _ => "Provider 操作失败", - }; - AdminError::new(kind, error.public_message().unwrap_or(message)).with_source(OperationFailure { - operation: resource, - source: error.into(), - }) -} - -async fn publish_committed( - snapshot: &dyn SnapshotControl, - revision: crate::model::Revision, -) -> Result<(), AdminError> { - let revision = ConfigRevision::new(revision.get()) - .map_err(|_| AdminError::internal("已提交的配置版本不合法"))?; - snapshot.publish_committed(revision).await; - Ok(()) -} - -async fn publish_credentials_and_observe_quota( - provider: &std::sync::Arc, - snapshot: &dyn SnapshotControl, - revision: crate::model::Revision, - account_ids: &[ProviderAccountId], - request_id: &str, -) -> Result<(), AdminError> { - provider.account_facts_changed(account_ids).await; - publish_committed(snapshot, revision).await?; - - // 凭据已经提交；额度只是可重建的观察，不能拖住管理请求或回滚提交结果 - let provider = provider.clone(); - let account_ids = account_ids.to_vec(); - let request_id = request_id.to_owned(); - tokio::spawn(async move { - for account_id in account_ids { - if let Err(error) = provider - .quota(ProviderQuotaRequest { - account_id: account_id.clone(), - refresh: true, - rolling_usage: None, - }) - .await - { - tracing::warn!( - request_id, - provider = %provider.provider_kind().as_str(), - account_id = %account_id.as_str(), - quota_error = ?error.kind(), - "initial quota observation failed" - ); - } - } - }); - Ok(()) -} - -async fn required_credential( - accounts: &dyn AccountStore, - provider_kind: &ProviderKind, - account_id: &ProviderAccountId, - resource: &'static str, -) -> Result { - accounts - .credential_details(provider_kind, account_id) - .await - .map_err(|error| map_store_error(error, resource))? - .ok_or_else(|| AdminError::not_found("Provider 凭据不存在")) -} - -async fn required_plugin_credential( - accounts: &dyn AccountStore, - account_id: &ProviderAccountId, - resource: &'static str, -) -> Result { - accounts - .credential_details_by_id(account_id) - .await - .map_err(|error| map_store_error(error, resource))? - .ok_or_else(|| AdminError::not_found("Provider 凭据不存在")) -} - -async fn import_proxy_binding( - proxies: &dyn crate::ports::proxy::ProxyStore, - id: Option<&str>, -) -> Result, AdminError> { - let Some(id) = id else { return Ok(None) }; - let reservation = proxies - .reserve_import(id) - .await - .map_err(|error| map_store_error(error, "proxy"))?; - Ok(Some(reservation)) -} - -async fn pending_authorization( - accounts: &dyn AccountStore, - proxies: &dyn crate::ports::proxy::ProxyStore, - provider_kind: &ProviderKind, - command: &StartAuthorization, - resource: &'static str, -) -> Result { - let proxy; - let mut proxy_id = None; - let target = match &command.reauthorization { - Some(account_id) => { - proxy = required_credential(accounts, provider_kind, account_id, resource) - .await? - .credential - .outbound_proxy; - AuthorizationMutationTarget::Reauthorize { - account_id: account_id.clone(), - } - } - None => { - use crate::model::proxies::AccountProxySelection; - proxy = match &command.outbound_proxy { - Some(AccountProxySelection::Url(value)) => Some(value.clone()), - Some(AccountProxySelection::Saved(id)) => { - let record = proxies - .get(id) - .await - .map_err(|error| map_store_error(error, "proxy"))?; - proxy_id = Some(record.id); - Some(record.proxy) - } - Some(AccountProxySelection::Direct) | None => None, - }; - AuthorizationMutationTarget::Create { - name: command.name.clone(), - } - } - }; - Ok(PendingAuthorizationMutation::new( - provider_kind.clone(), - target, - crate::model::provider_credentials::AuthorizationOwnerBinding::from_context( - &command.context, - ), - ) - .with_outbound_proxy(proxy) - .with_outbound_proxy_id(proxy_id)) -} - -fn validate_prepared_import( - provider_kind: &ProviderKind, - prepared: &PreparedCredentialImport, - _resource: &'static str, -) -> Result<(), AdminError> { - if prepared.provider_kind != *provider_kind - || prepared - .credentials - .iter() - .any(|credential| credential.provider_kind != *provider_kind) - { - return Err(AdminError::conflict("Provider 准备结果与请求范围不一致")); - } - Ok(()) -} - -fn validate_prepared_rotation( - account: &crate::model::accounts::AccountRecord, - prepared: &PreparedCredentialRotation, - _resource: &'static str, -) -> Result<(), AdminError> { - validate_prepared_rotation_facts(account, prepared.facts()) -} - -fn validate_prepared_rotation_facts( - account: &crate::model::accounts::AccountRecord, - facts: &crate::model::provider_credentials::PreparedCredentialRotationFacts, -) -> Result<(), AdminError> { - if facts.account_id.as_str() != account.id.as_str() - || facts.provider_kind != account.provider_kind - { - return Err(AdminError::conflict("Provider 准备结果与当前凭据不一致")); - } - Ok(()) -} - -async fn validate_authorization_commit( - provider_kind: &ProviderKind, - context: &MutationContext, - prepared: PreparedAuthorizationCommit, - resource: &'static str, -) -> Result { - if prepared.pending.provider_kind() != provider_kind - || !prepared.pending.owner_binding().matches_context(context) - { - let validation_error = AdminError::conflict("OAuth 授权绑定已失效，请重新发起授权"); - if let Err(error) = prepared.abort().await { - tracing::warn!( - resource, - settlement_error = %error, - "OAuth authorization claim release failed after preparation validation" - ); - return Err(error); - } - return Err(validation_error); - } - let matches_target = match (prepared.pending.target(), &prepared.credential) { - ( - AuthorizationMutationTarget::Create { .. }, - PreparedAuthorizationCredential::Create(credential), - ) => credential.provider_kind == *provider_kind, - ( - AuthorizationMutationTarget::Reauthorize { account_id }, - PreparedAuthorizationCredential::Reauthorize(credential), - ) => { - let facts = credential.facts(); - facts.provider_kind == *provider_kind && facts.account_id == *account_id - } - _ => false, - }; - if !matches_target { - let validation_error = AdminError::conflict("OAuth 凭据与待提交目标不一致"); - if let Err(error) = prepared.abort().await { - tracing::warn!( - resource, - settlement_error = %error, - "OAuth authorization claim release failed after preparation validation" - ); - return Err(error); - } - return Err(validation_error); - } - Ok(prepared) -} - -async fn commit_authorization( - accounts: &dyn AccountStore, - prepared: PreparedAuthorizationCommit, - key: crate::model::provider_credentials::AuthorizationReceiptKey, - settings: Option, - context: &MutationContext, - resource: &'static str, -) -> Result { - if settings.is_some() - && matches!( - &prepared.credential, - PreparedAuthorizationCredential::Reauthorize(_) - ) - { - prepared.abort().await?; - return Err(AdminError::invalid("重新授权不能修改账号设置")); - } - let crate::model::provider_credentials::AuthorizationCommitSettlement { - command, - credential_guard, - authorization_guard, - } = prepared.into_commit(settings, key); - match accounts.commit_authorization(command, context).await { - Ok(result) => { - if let Some(guard) = credential_guard - && result.newly_committed - { - guard.finish(); - } - if let Some(guard) = authorization_guard - && let Err(error) = guard.commit().await - { - // 账号与回执已经原子提交，Redis 清理失败不能把确定的成功改写为失败 - tracing::warn!(resource, settlement_error = %error, "authorization committed but pending cleanup failed"); - } - Ok(result) - } - Err(error) => { - drop(credential_guard); - if let Some(guard) = authorization_guard - && let Err(settlement_error) = guard.abort().await - { - tracing::warn!( - resource, - store_error = %error, - settlement_error = %settlement_error, - "OAuth authorization claim release failed after Store commit failure" - ); - return Err(settlement_error); - } - Err(map_store_error(error, resource)) - } - } -} - -async fn commit_credential_rotation( - accounts: &dyn AccountStore, - prepared: PreparedCredentialRotation, - settings: Option, - context: &MutationContext, - resource: &'static str, -) -> Result { - let (facts, guard) = prepared.into_parts(); - match accounts - .commit_credential_rotation( - CredentialRotationCommit { - prepared: facts, - settings, - }, - context, - ) - .await - { - Ok(result) => { - guard.finish(); - Ok(result) - } - Err(error) => { - drop(guard); - Err(map_store_error(error, resource)) - } - } -} - -async fn commit_credential_refresh( - accounts: &dyn AccountStore, - prepared: PreparedCredentialRotation, - context: &MutationContext, - resource: &'static str, -) -> Result { - let (facts, guard) = prepared.into_parts(); - match accounts - .commit_credential_refresh( - CredentialRotationCommit { - prepared: facts, - settings: None, - }, - context, - ) - .await - { - Ok(result) => { - guard.finish(); - Ok(result) - } - Err(error) => { - drop(guard); - Err(map_store_error(error, resource)) - } - } -} - -async fn delete_credentials( - accounts: &dyn AccountStore, - provider: &dyn ProviderAdmin, - command: CredentialDeletion, - resource: &'static str, -) -> Result { - for account_id in &command.account_ids { - required_credential(accounts, provider.provider_kind(), account_id, resource).await?; - } - let account_ids = command.account_ids; - let revision = accounts - .delete_accounts( - DeleteAccounts { - account_ids: account_ids - .iter() - .map(|account_id| account_id.as_str().to_owned()) - .collect(), - }, - &command.context, - ) - .await - .map_err(|error| map_store_error(error, resource))?; - for account_id in &account_ids { - provider.account_unavailable(account_id).await; - } - provider.account_facts_changed(&account_ids).await; - Ok(CredentialDeletionResult { - config_revision: revision, - account_ids, - }) -} - -// 用例的操作名补充端口资源上下文，底层分类与原始原因仍由来源持有 -#[derive(Debug, thiserror::Error)] -#[error("admin operation failed: {operation}")] -struct OperationFailure { - operation: &'static str, - source: gateway_core::error::ErrorSource, -} +use error::{map_provider_error, map_store_error}; +use publication::publish_committed; diff --git a/backend/crates/gateway-admin/src/use_case/plugins/management.rs b/backend/crates/gateway-admin/src/use_case/plugins/management.rs index cb7b8f417..c8f0d7deb 100644 --- a/backend/crates/gateway-admin/src/use_case/plugins/management.rs +++ b/backend/crates/gateway-admin/src/use_case/plugins/management.rs @@ -2,7 +2,8 @@ use std::sync::Arc; -use gateway_core::runtime::{RuntimeSnapshotHandle, extensions::ExtensionSetReference}; +use gateway_core::routing::extensions::ExtensionSetReference; +use gateway_core::runtime::RuntimeSnapshotHandle; use crate::{ model::{ diff --git a/backend/crates/gateway-admin/src/use_case/publication.rs b/backend/crates/gateway-admin/src/use_case/publication.rs new file mode 100644 index 000000000..5bc2807c8 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/publication.rs @@ -0,0 +1,14 @@ +//! 管理事务提交后的运行快照发布入口 + +use crate::model::AdminError; +use gateway_core::{routing::ConfigRevision, runtime::SnapshotControl}; + +pub(super) async fn publish_committed( + snapshot: &dyn SnapshotControl, + revision: crate::model::Revision, +) -> Result<(), AdminError> { + let revision = ConfigRevision::new(revision.get()) + .map_err(|_| AdminError::internal("已提交的配置版本不合法"))?; + snapshot.publish_committed(revision).await; + Ok(()) +} diff --git a/backend/crates/gateway-admin/tests/use_case/mod.rs b/backend/crates/gateway-admin/tests/use_case/mod.rs index 7f8e617cb..12a9a7c25 100644 --- a/backend/crates/gateway-admin/tests/use_case/mod.rs +++ b/backend/crates/gateway-admin/tests/use_case/mod.rs @@ -107,7 +107,7 @@ pub(super) struct AdminHarness { plugin_store: Arc, plugin_inspector: Arc, client_key_verifier: Arc, - service_middleware: gateway_admin::service::PlanSource, + service_middleware: gateway_admin::public_service::PlanSource, } impl AdminHarness { @@ -194,7 +194,10 @@ impl AdminHarness { self } - pub(super) fn service_middleware(mut self, source: gateway_admin::service::PlanSource) -> Self { + pub(super) fn service_middleware( + mut self, + source: gateway_admin::public_service::PlanSource, + ) -> Self { self.service_middleware = source; self } diff --git a/backend/crates/gateway-admin/tests/use_case/plugins/artifacts.rs b/backend/crates/gateway-admin/tests/use_case/plugins/artifacts.rs index 3b1d6c598..859cb75f1 100644 --- a/backend/crates/gateway-admin/tests/use_case/plugins/artifacts.rs +++ b/backend/crates/gateway-admin/tests/use_case/plugins/artifacts.rs @@ -24,12 +24,10 @@ use gateway_admin::{ store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult}, }, }; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::{ routing::ConfigRevision, - runtime::{ - RuntimeSnapshotHandle, SnapshotControl, - extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, - }, + runtime::{RuntimeSnapshotHandle, SnapshotControl}, }; use serde_json::json; @@ -117,7 +115,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for Fixture { &self, _: &gateway_admin::model::plugins::instances::PluginInstanceSnapshot, _: Option, - _: Option<&gateway_core::runtime::extensions::ExtensionSetReference>, + _: Option<&gateway_core::routing::extensions::ExtensionSetReference>, ) -> Option< std::collections::BTreeMap< String, @@ -132,7 +130,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for Fixture { impl gateway_admin::ports::plugins::PluginStateLifecycle for Fixture { async fn activate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::instances::PluginInstance, ) -> Result<(), AdminError> { Ok(()) @@ -142,7 +140,7 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for Fixture { } async fn migrate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: gateway_admin::model::plugins::state::PluginStateTransition, ) -> Result<(), AdminError> { panic!("unexpected state migration") diff --git a/backend/crates/gateway-admin/tests/use_case/plugins/instances.rs b/backend/crates/gateway-admin/tests/use_case/plugins/instances.rs index b1305cf23..796f936fd 100644 --- a/backend/crates/gateway-admin/tests/use_case/plugins/instances.rs +++ b/backend/crates/gateway-admin/tests/use_case/plugins/instances.rs @@ -40,12 +40,10 @@ use gateway_admin::{ store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult}, }, }; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::{ routing::ConfigRevision, - runtime::{ - RuntimeSnapshotHandle, SnapshotControl, - extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, - }, + runtime::{RuntimeSnapshotHandle, SnapshotControl}, }; use serde_json::json; diff --git a/backend/crates/gateway-admin/tests/use_case/plugins/mod.rs b/backend/crates/gateway-admin/tests/use_case/plugins/mod.rs index 072cee754..560654cb4 100644 --- a/backend/crates/gateway-admin/tests/use_case/plugins/mod.rs +++ b/backend/crates/gateway-admin/tests/use_case/plugins/mod.rs @@ -111,7 +111,7 @@ impl gateway_admin::ports::proxy::ProxyStore for TestPluginPorts { impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPorts { async fn validate_target( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, ) -> Result<(), AdminError> { Err(AdminError::not_found("unused plugin fixture")) @@ -119,7 +119,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor async fn start_callback( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: gateway_admin::model::plugins::management::StartPluginManagementCallback, _: &gateway_admin::model::auth::AdminRequestContext, @@ -129,7 +129,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn callback( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: &str, _: gateway_admin::model::plugins::management::PluginManagementRequest, @@ -139,14 +139,14 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn views( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, ) -> Result, AdminError> { Ok(Vec::new()) } async fn resource( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: &str, _: bool, @@ -156,7 +156,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn handle( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: gateway_admin::model::plugins::management::PluginManagementRequest, ) -> Result @@ -171,7 +171,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for TestPluginPorts &self, _: &gateway_admin::model::plugins::instances::PluginInstanceSnapshot, _: Option, - _: Option<&gateway_core::runtime::extensions::ExtensionSetReference>, + _: Option<&gateway_core::routing::extensions::ExtensionSetReference>, ) -> Option< std::collections::BTreeMap< String, @@ -186,7 +186,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for TestPluginPorts impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts { async fn activate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::instances::PluginInstance, ) -> Result<(), AdminError> { panic!("unexpected state activation") @@ -196,7 +196,7 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts { } async fn migrate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: gateway_admin::model::plugins::state::PluginStateTransition, ) -> Result<(), AdminError> { panic!("unexpected state migration") @@ -222,7 +222,7 @@ impl gateway_admin::ports::plugins::PluginPreparation for TestPluginPorts { async fn prepare( &self, _: gateway_admin::model::plugins::instances::PluginInstanceSnapshot, - ) -> Result { + ) -> Result { Err(AdminError::invalid("unused plugin fixture")) } } diff --git a/backend/crates/gateway-admin/tests/use_case/service_contract.rs b/backend/crates/gateway-admin/tests/use_case/service_contract.rs index 3d34cf4a5..35cb471d7 100644 --- a/backend/crates/gateway-admin/tests/use_case/service_contract.rs +++ b/backend/crates/gateway-admin/tests/use_case/service_contract.rs @@ -247,7 +247,7 @@ fn sdk_settings_contract_matches_host_declarations() { _ => None, }) .collect(); - let register = source(root, "src/service/settings.rs") + let register = source(root, "src/public_service/settings.rs") .items .into_iter() .find_map(|item| match item { @@ -293,7 +293,7 @@ fn sdk_settings_contract_matches_host_declarations() { let generated = quote! { //! 宿主设置服务的操作标识与请求、响应数据合同 //! - //! 从宿主设置类型与 service/settings.rs 生成；更新命令见 SDK 维护说明 + //! 从宿主设置类型与 public_service/settings.rs 生成；更新命令见 SDK 维护说明 use super::Operation; use serde::{Deserialize, Serialize}; diff --git a/backend/crates/gateway-admin/tests/use_case/settings.rs b/backend/crates/gateway-admin/tests/use_case/settings.rs index e7dbbdc6c..517ca0448 100644 --- a/backend/crates/gateway-admin/tests/use_case/settings.rs +++ b/backend/crates/gateway-admin/tests/use_case/settings.rs @@ -28,7 +28,7 @@ struct PricingSettingsStore { #[tokio::test] async fn explicit_service_calls_freeze_empty_plan_and_cancel_before_entering_terminal() { - use gateway_admin::service::{Origin, Plan, scope}; + use gateway_admin::public_service::{Origin, Plan, scope}; use std::sync::atomic::{AtomicUsize, Ordering}; let resolutions = Arc::new(AtomicUsize::new(0)); diff --git a/backend/crates/gateway-api/src/lib.rs b/backend/crates/gateway-api/src/lib.rs index 0f0c4559f..2ccb140ac 100644 --- a/backend/crates/gateway-api/src/lib.rs +++ b/backend/crates/gateway-api/src/lib.rs @@ -138,7 +138,7 @@ impl ApiBundle { } /// 内部调用使用同一总路由；组合根持有强引用，Runtime 只保存 Weak - pub fn dispatcher(&self) -> Arc { + pub fn dispatcher(&self) -> Arc { Arc::new(middleware::RouterDispatcher { router: middleware::wrap( self.router.clone(), diff --git a/backend/crates/gateway-api/src/middleware/mod.rs b/backend/crates/gateway-api/src/middleware/mod.rs index 3e05fe103..8319a4e2c 100644 --- a/backend/crates/gateway-api/src/middleware/mod.rs +++ b/backend/crates/gateway-api/src/middleware/mod.rs @@ -17,23 +17,21 @@ use axum::{ http::{HeaderName, HeaderValue, Request, StatusCode}, response::{IntoResponse as _, Response}, }; -use gateway_core::{ - engine::middleware::FrozenMiddlewarePlan, - lifecycle::CancellationToken, - middleware::{ - compose, - http::{self as contract, Settings}, - }, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::http as contract; +use gateway_core::engine::middleware::http::Settings; +use gateway_core::lifecycle::CancellationToken; +use gateway_core::middleware::compose; use http_body::{Body as HttpBody, Frame, SizeHint}; use http_body_util::BodyExt as _; use tokio::time::Instant; use tower::ServiceExt as _; tokio::task_local! { static CURRENT: contract::Context; } -tokio::task_local! { static SETTINGS: Option; } +tokio::task_local! { static SETTINGS: Option; } -pub(crate) fn current_settings() -> Option { +pub(crate) fn current_settings() -> Option +{ SETTINGS.try_with(Clone::clone).ok().flatten() } @@ -52,7 +50,7 @@ pub(crate) async fn scope( } pub(crate) type SettingsSource = - Arc Option + Send + Sync>; + Arc Option + Send + Sync>; pub(crate) struct RouterDispatcher { pub(crate) router: Router, @@ -60,7 +58,7 @@ pub(crate) struct RouterDispatcher { } impl contract::Dispatcher for RouterDispatcher { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { (self.settings)() } @@ -195,12 +193,12 @@ async fn handle( }); request.extensions_mut().insert(guard.clone()); request.extensions_mut().insert(context.clone()); - let service_context = gateway_admin::service::Origin { + let service_context = gateway_admin::public_service::Origin { request_id: context.request_id.clone(), call_id: context.call_id.clone(), cancellation: context.cancellation.clone(), extensions: context.extensions.clone(), - plan: gateway_admin::service::Plan::Frozen(plan.clone()), + plan: gateway_admin::public_service::Plan::Frozen(plan.clone()), }; request .extensions_mut() @@ -220,7 +218,7 @@ async fn handle( terminal_context, SETTINGS.scope( settings, - gateway_admin::service::scope( + gateway_admin::public_service::scope( service_context, dispatch(router, request.map(Body::new), started), ), diff --git a/backend/crates/gateway-api/src/middleware/websocket.rs b/backend/crates/gateway-api/src/middleware/websocket.rs index 487c85ea8..891c187a0 100644 --- a/backend/crates/gateway-api/src/middleware/websocket.rs +++ b/backend/crates/gateway-api/src/middleware/websocket.rs @@ -14,15 +14,14 @@ use futures::{ future::BoxFuture, stream::{SplitSink, SplitStream}, }; -use gateway_core::{ - engine::middleware::{FrozenMiddlewarePlan, MiddlewareError, MiddlewareHeader}, - lifecycle::CancellationToken, - middleware::{ - compose, - http::{self as http, upgrade}, - websocket as core, - }, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareHeader; +use gateway_core::engine::middleware::http; +use gateway_core::engine::middleware::http::upgrade; +use gateway_core::engine::middleware::websocket as core; +use gateway_core::lifecycle::CancellationToken; +use gateway_core::middleware::compose; use http_body_util::BodyExt as _; pub(super) struct Upgrade { diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/frame.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/frame.rs new file mode 100644 index 000000000..4e9bd4a74 --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/frame.rs @@ -0,0 +1,171 @@ +//! 下游连接的帧、写入结果、退出原因与队列预算 + +use std::{sync::Arc, time::Duration}; +use thiserror::Error; +use tokio::sync::OwnedSemaphorePermit; + +pub(super) const OUTBOUND_COMMAND_BUFFER: usize = 32; +pub(super) const INBOUND_EVENT_BUFFER: usize = 32; +pub(super) const DOWNSTREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(5 * 60); +pub(super) const DOWNSTREAM_CLOSE_TIMEOUT: Duration = Duration::from_secs(1); +pub(super) const CONNECTION_MAX_AGE: Duration = Duration::from_secs(60 * 60); +pub(super) const DOWNSTREAM_PING_INTERVAL: Duration = Duration::from_secs(25); + +/// WebSocket pump 的写入与生命周期预算 +#[derive(Clone, Copy)] +pub struct ConnectionConfig { + pub(super) write_timeout: Duration, + pub(super) max_age: Duration, +} + +impl ConnectionConfig { + /// 生产环境固定预算 + pub const PRODUCTION: Self = Self { + write_timeout: DOWNSTREAM_WRITE_TIMEOUT, + max_age: CONNECTION_MAX_AGE, + }; +} + +/// 业务层可观察的客户端输入；Ping/Pong 始终由 pump 消费 +pub enum ConnectionEvent { + Text(String), + Binary, + Expired, + Exited(PumpExitReason), +} + +/// 接收队列与活动期间暂存的请求共用容量；移动帧不释放它占用的名额 +pub(in super::super) struct PendingConnectionEvent { + pub(in super::super) event: ConnectionEvent, + pub(super) _permit: Option, +} + +impl PendingConnectionEvent { + pub(super) fn exited(reason: PumpExitReason) -> Self { + Self { + event: ConnectionEvent::Exited(reason), + _permit: None, + } + } +} + +/// 下游写入阶段；名称刻意使用 write，而不是暗示客户端已消费的 delivery +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FramePhase { + Metadata, + First, + Data, + Terminal, + FirstAndTerminal, + Error, + ConnectionLimit, + Close, +} + +impl FramePhase { + pub(super) const fn as_str(self) -> &'static str { + match self { + Self::Metadata => "metadata", + Self::First => "first", + Self::Data => "data", + Self::Terminal => "terminal", + Self::FirstAndTerminal => "first_and_terminal", + Self::Error => "error", + Self::ConnectionLimit => "connection_limit", + Self::Close => "close", + } + } + + pub(super) const fn is_milestone(self) -> bool { + matches!( + self, + Self::First + | Self::Terminal + | Self::FirstAndTerminal + | Self::Error + | Self::ConnectionLimit + | Self::Close + ) + } +} + +/// 一次下游写入的请求归属与协议阶段 +#[derive(Clone)] +pub struct WriteContext { + pub(super) request_id: Option>, + pub(super) phase: FramePhase, +} + +impl WriteContext { + /// 创建归属于某个请求的写入上下文 + #[must_use] + pub fn request(request_id: &Arc, phase: FramePhase) -> Self { + Self { + request_id: Some(Arc::clone(request_id)), + phase, + } + } + + /// 创建连接级写入上下文 + #[must_use] + pub const fn connection(phase: FramePhase) -> Self { + Self { + request_id: None, + phase, + } + } +} + +/// WebSocket pump 停止的稳定原因 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PumpExitReason { + ClientClose, + PeerEof, + ReadError, + WriteError, + WriteTimeout, + LifecycleShutdown, + ConnectionMaxAge, + CoordinatorDropped, + InboundOverload, + ServerClose, + PumpStopped, + MiddlewareFailed, +} + +impl PumpExitReason { + pub(in super::super) const fn as_str(self) -> &'static str { + match self { + Self::ClientClose => "client_close", + Self::PeerEof => "peer_eof", + Self::ReadError => "read_error", + Self::WriteError => "write_error", + Self::WriteTimeout => "write_timeout", + Self::LifecycleShutdown => "lifecycle_shutdown", + Self::ConnectionMaxAge => "connection_max_age", + Self::CoordinatorDropped => "coordinator_dropped", + Self::InboundOverload => "inbound_overload", + Self::ServerClose => "server_close", + Self::PumpStopped => "pump_stopped", + Self::MiddlewareFailed => "middleware_failed", + } + } +} + +/// 下游 WebSocket 写入失败 +#[derive(Debug, Error)] +pub enum ConnectionWriteError { + #[error("downstream WebSocket pump is closed")] + Closed, + #[error("downstream WebSocket write timed out after {timeout:?}")] + Timeout { timeout: Duration }, + #[error("downstream WebSocket transport write failed: {message}")] + Transport { message: String }, +} + +/// 只有实际 transport 写入才记为 Written；插件丢弃消息不伪造写入成功 +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WriteOutcome { + Written, + Suppressed, +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs new file mode 100644 index 000000000..59d79562c --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs @@ -0,0 +1,362 @@ +//! 协调层持有的连接句柄与唯一 pump 启动入口 + +use super::{ + frame::*, + pump::run_pump, + state::{ConnectionCommand, ConnectionStats}, +}; +use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code}; +use futures::{Sink, Stream}; +use gateway_core::{ + engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader}, + lifecycle::CancellationToken, +}; +use std::{ + collections::VecDeque, + fmt, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; +use tokio::{ + sync::{mpsc, oneshot}, + task::JoinHandle, + time::{Instant, timeout}, +}; + +/// 协调层持有的单 owner WebSocket 连接句柄 +pub struct ResponsesWebSocketConnection { + connection_id: Arc, + opened_at: Instant, + expired: Arc, + commands: Option>, + incoming: mpsc::Receiver, + deferred: VecDeque, + exited: oneshot::Receiver, + pump_task: Option>, + stats: Arc, + config: ConnectionConfig, + exit_reason: Option, +} + +impl ResponsesWebSocketConnection { + pub(in super::super) fn new( + socket: WebSocket, + connection_id: String, + cancellation: CancellationToken, + middleware: Option, + headers: Arc<[MiddlewareHeader]>, + ) -> Self { + spawn_connection( + socket, + Arc::::from(connection_id), + cancellation, + ConnectionConfig::PRODUCTION, + middleware, + headers, + ) + } + + pub(in super::super) fn id(&self) -> &str { + &self.connection_id + } + + /// 返回连接是否已经达到生命周期上限 + #[must_use] + pub fn is_expired(&self) -> bool { + self.expired.load(Ordering::Acquire) + } + + pub(in super::super) fn age(&self) -> Duration { + self.opened_at.elapsed() + } + + /// 等待下一个需要业务层处理的客户端事件 + pub async fn next_event(&mut self) -> Option { + if self.exit_reason.is_none() { + match self.exited.try_recv() { + Ok(reason) => { + self.exit_reason = Some(reason); + } + Err(tokio::sync::oneshot::error::TryRecvError::Closed) => { + self.exit_reason = Some(PumpExitReason::PumpStopped); + } + Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {} + } + } + if let Some(reason) = self.exit_reason { + return Some(ConnectionEvent::Exited(reason)); + } + if let Some(event) = self.deferred.pop_front() { + return Some(event.event); + } + self.next_active_event().await.map(|event| event.event) + } + + /// 活动响应期间只消费新到的帧，已排队的下一轮请求继续保持串行 + pub(in super::super) async fn next_active_event(&mut self) -> Option { + if let Some(reason) = self.exit_reason { + return Some(PendingConnectionEvent::exited(reason)); + } + let event = tokio::select! { + biased; + reason = &mut self.exited => { + Some(PendingConnectionEvent::exited(reason.unwrap_or(PumpExitReason::PumpStopped))) + } + event = self.incoming.recv() => event, + }; + if let Some(PendingConnectionEvent { + event: ConnectionEvent::Exited(reason), + .. + }) = event.as_ref() + { + self.exit_reason.get_or_insert(*reason); + } else if event.is_none() { + self.exit_reason.get_or_insert(PumpExitReason::PumpStopped); + } + event + } + + pub(in super::super) fn defer(&mut self, event: PendingConnectionEvent) { + self.deferred.push_back(event); + } + + /// 等待连接退出，不消费留给后续串行请求的业务帧 + pub async fn wait_for_exit(&mut self) -> PumpExitReason { + if let Some(reason) = self.exit_reason { + return reason; + } + let reason = (&mut self.exited) + .await + .unwrap_or(PumpExitReason::PumpStopped); + self.exit_reason = Some(reason); + reason + } + + /// 串行写入文本帧，并等待 pump 确认 transport 写入结果 + /// + /// # Errors + /// + /// pump 已关闭、写入超时或 transport 失败时返回稳定错误 + pub async fn send_text( + &mut self, + payload: String, + context: WriteContext, + ) -> Result { + self.send(Message::Text(payload.into()), context).await + } + + pub(in super::super) async fn close_policy( + &mut self, + reason: &'static str, + request_id: Option<&Arc>, + ) { + let context = request_id.map_or_else( + || WriteContext::connection(FramePhase::Close), + |request_id| WriteContext::request(request_id, FramePhase::Close), + ); + self.close( + CloseFrame { + code: close_code::POLICY, + reason: reason.into(), + }, + context, + PumpExitReason::ServerClose, + ) + .await; + } + + pub(in super::super) async fn close_for_connection_limit(&mut self, reason: &'static str) { + self.close( + CloseFrame { + code: close_code::NORMAL, + reason: reason.into(), + }, + WriteContext::connection(FramePhase::Close), + PumpExitReason::ConnectionMaxAge, + ) + .await; + } + + async fn close( + &mut self, + frame: CloseFrame, + context: WriteContext, + exit_reason: PumpExitReason, + ) { + if self + .send(Message::Close(Some(frame)), context) + .await + .is_ok() + { + self.exit_reason.get_or_insert(exit_reason); + } + } + + async fn send( + &mut self, + message: Message, + context: WriteContext, + ) -> Result { + let Some(commands) = self.commands.as_ref().cloned() else { + return Err(ConnectionWriteError::Closed); + }; + self.stats.observe_command_queue(&commands); + let (acknowledged, acknowledgement) = oneshot::channel(); + let started_at = Instant::now(); + let request_id = context.request_id.clone(); + let phase = context.phase; + let command = ConnectionCommand { + message, + context, + acknowledged, + }; + let write = async move { + commands + .send(command) + .await + .map_err(|_| ConnectionWriteError::Closed)?; + acknowledgement + .await + .map_err(|_| ConnectionWriteError::Closed)? + }; + let result = match timeout(self.config.write_timeout, write).await { + Ok(result) => result, + Err(_) => Err(ConnectionWriteError::Timeout { + timeout: self.config.write_timeout, + }), + }; + let duration = started_at.elapsed(); + match &result { + Ok(WriteOutcome::Written) if phase.is_milestone() => tracing::info!( + websocket_connection_id = %self.connection_id, + request_id = request_id.as_deref().unwrap_or(""), + frame_phase = phase.as_str(), + write_duration_ms = duration.as_millis(), + "Responses WebSocket frame write succeeded" + ), + Ok(outcome) => tracing::debug!( + websocket_connection_id = %self.connection_id, + request_id = request_id.as_deref().unwrap_or(""), + frame_phase = phase.as_str(), + write_duration_ms = duration.as_millis(), + ?outcome, + "Responses WebSocket frame processed" + ), + Err(error) => tracing::info!( + websocket_connection_id = %self.connection_id, + request_id = request_id.as_deref().unwrap_or(""), + frame_phase = phase.as_str(), + write_duration_ms = duration.as_millis(), + error = %error, + "Responses WebSocket frame write failed" + ), + } + if let Err(error) = &result { + let reason = match error { + ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout, + ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError, + ConnectionWriteError::Closed => PumpExitReason::PumpStopped, + }; + self.terminate(reason); + } + result + } + + fn terminate(&mut self, reason: PumpExitReason) { + self.exit_reason.get_or_insert(reason); + self.commands.take(); + if let Some(task) = self.pump_task.take() { + task.abort(); + } + } + + pub(in super::super) fn log_summary(&self, request_count: u64) { + let connection_age_ms = self.age().as_millis(); + let read_idle_ms = connection_age_ms + .saturating_sub(u128::from(self.stats.last_read_ms.load(Ordering::Relaxed))); + let write_idle_ms = connection_age_ms + .saturating_sub(u128::from(self.stats.last_write_ms.load(Ordering::Relaxed))); + tracing::info!( + websocket_connection_id = %self.connection_id, + request_count, + connection_age_ms, + pump_exit_reason = self + .exit_reason + .map_or("unknown", PumpExitReason::as_str), + command_queue_high_water = self + .stats + .command_queue_high_water + .load(Ordering::Relaxed), + command_backpressure_count = self + .stats + .command_backpressure_count + .load(Ordering::Relaxed), + ping_received_count = self.stats.ping_received_count.load(Ordering::Relaxed), + ping_written_count = self.stats.ping_written_count.load(Ordering::Relaxed), + pong_received_count = self.stats.pong_received_count.load(Ordering::Relaxed), + read_idle_ms, + write_idle_ms, + "Responses WebSocket disconnected" + ); + } +} + +impl Drop for ResponsesWebSocketConnection { + fn drop(&mut self) { + self.commands.take(); + if let Some(task) = self.pump_task.take() { + task.abort(); + } + } +} + +/// 为 WebSocket transport 启动包含冻结中间件计划的单 owner pump +pub fn spawn_connection( + socket: S, + connection_id: Arc, + cancellation: CancellationToken, + config: ConnectionConfig, + middleware: Option, + headers: Arc<[MiddlewareHeader]>, +) -> ResponsesWebSocketConnection +where + S: Stream> + Sink + Unpin + Send + 'static, + E: fmt::Display + Send + 'static, +{ + let opened_at = Instant::now(); + let expired = Arc::new(AtomicBool::new(false)); + let stats = Arc::new(ConnectionStats::default()); + let (command_tx, command_rx) = mpsc::channel(OUTBOUND_COMMAND_BUFFER); + let (incoming_tx, incoming_rx) = mpsc::channel(INBOUND_EVENT_BUFFER); + let (exit_tx, exit_rx) = oneshot::channel(); + let pump_task = tokio::spawn(run_pump( + socket, + command_rx, + incoming_tx, + exit_tx, + Arc::clone(&connection_id), + cancellation, + opened_at, + Arc::clone(&expired), + Arc::clone(&stats), + config, + middleware, + headers, + )); + ResponsesWebSocketConnection { + connection_id, + opened_at, + expired, + commands: Some(command_tx), + incoming: incoming_rx, + deferred: VecDeque::new(), + exited: exit_rx, + pump_task: Some(pump_task), + stats, + config, + exit_reason: None, + } +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/middleware.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/middleware.rs index 4f2f64dae..44267c054 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/connection/middleware.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/middleware.rs @@ -1,11 +1,13 @@ //! 插件消息组合在 transport 锁之外执行，主动发送复用唯一 writer -use super::{ConnectionWriteError, PumpContext, write_message}; +use super::{ConnectionWriteError, frame::DOWNSTREAM_CLOSE_TIMEOUT, state::PumpContext}; pub(super) use crate::middleware::websocket::CancelOnDrop; use crate::middleware::websocket::{from_core, into_core}; use axum::extract::ws::Message; -use futures::{Sink, future::BoxFuture}; -use gateway_core::{engine::middleware::MiddlewareError, middleware::websocket as core}; -use std::sync::Arc; +use futures::{Sink, SinkExt, future::BoxFuture}; +use gateway_core::engine::middleware::{MiddlewareError, websocket as core}; +use gateway_core::lifecycle::CancellationToken; +use std::{fmt, sync::Arc}; +use tokio::time::timeout; pub(super) struct Writer { socket: tokio::sync::Mutex, @@ -37,6 +39,36 @@ where }) => result.unwrap_or(Err(ConnectionWriteError::Timeout { timeout: self.context.config.write_timeout })), } } + + pub(super) async fn finish_peer_close( + &self, + cancellation: &CancellationToken, + ) -> Result<(), ConnectionWriteError> { + // 主动 Sender 也可能在 pump 之外持锁；先独占 transport 再取消业务， + // 有在途写或已取消的半帧时立即跳过，不能等待它释放锁后继续刷新 + let socket = self + .socket + .try_lock() + .ok() + .filter(|_| !self.context.cancellation.is_cancelled()); + self.context.cancellation.cancel(); + let Some(mut socket) = socket else { + return Err(ConnectionWriteError::Closed); + }; + // 收到 Close 后 transport 已排队应答；只刷新它，不能再发送一份 Close + // 业务调用已取消，清理只受宿主取消与独立短预算约束，不沿用业务写入的长超时 + tokio::select! { + biased; + () = cancellation.cancelled() => Err(ConnectionWriteError::Closed), + result = timeout(DOWNSTREAM_CLOSE_TIMEOUT, async { + socket.flush().await.map_err(|error| { + ConnectionWriteError::Transport { message: error.to_string() } + }) + }) => result.unwrap_or(Err(ConnectionWriteError::Timeout { + timeout: DOWNSTREAM_CLOSE_TIMEOUT, + })), + } + } } impl core::Sender for Writer where @@ -81,3 +113,25 @@ pub(super) async fn transform( .map(from_core) .transpose() } + +async fn write_message( + socket: &mut S, + message: Message, + context: &PumpContext, +) -> Result<(), ConnectionWriteError> +where + S: Sink + Unpin, + E: fmt::Display, +{ + tokio::select! { + biased; + () = context.cancellation.cancelled() => Err(ConnectionWriteError::Closed), + result = timeout(context.config.write_timeout, socket.send(message)) => { + match result { + Ok(Ok(())) => { context.stats.record_write(context.opened_at); Ok(()) } + Ok(Err(error)) => Err(ConnectionWriteError::Transport { message: error.to_string() }), + Err(_) => Err(ConnectionWriteError::Timeout { timeout: context.config.write_timeout }), + } + } + } +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/mod.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/mod.rs index 42db1d2fb..981bd2117 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/connection/mod.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/mod.rs @@ -1,806 +1,13 @@ -//! 下游客户端 WebSocket 的单 owner pump 与有界收发边界 +//! 下游客户端 WebSocket 的连接合同、句柄与单 owner pump +mod frame; +mod handle; mod middleware; +mod pump; +mod state; -use std::{ - collections::VecDeque, - fmt, - sync::{ - Arc, - atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, - }, - time::Duration, +pub use frame::{ + ConnectionConfig, ConnectionEvent, ConnectionWriteError, FramePhase, PumpExitReason, + WriteContext, WriteOutcome, }; - -use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code}; -use futures::{Sink, SinkExt, Stream, StreamExt}; -use gateway_core::{ - engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader}, - lifecycle::CancellationToken, - middleware::websocket as contract, -}; -use thiserror::Error; -use tokio::{ - sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot}, - task::JoinHandle, - time::{Instant, timeout}, -}; - -const OUTBOUND_COMMAND_BUFFER: usize = 32; -const INBOUND_EVENT_BUFFER: usize = 32; -const DOWNSTREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(5 * 60); -const CONNECTION_MAX_AGE: Duration = Duration::from_secs(60 * 60); -const DOWNSTREAM_PING_INTERVAL: Duration = Duration::from_secs(25); - -/// WebSocket pump 的写入与生命周期预算 -#[derive(Clone, Copy)] -pub struct ConnectionConfig { - write_timeout: Duration, - max_age: Duration, -} - -impl ConnectionConfig { - /// 生产环境固定预算 - pub const PRODUCTION: Self = Self { - write_timeout: DOWNSTREAM_WRITE_TIMEOUT, - max_age: CONNECTION_MAX_AGE, - }; -} - -/// 业务层可观察的客户端输入；Ping/Pong 始终由 pump 消费 -pub enum ConnectionEvent { - Text(String), - Binary, - Expired, - Exited(PumpExitReason), -} - -/// 接收队列与活动期间暂存的请求共用容量；移动帧不释放它占用的名额 -pub(super) struct PendingConnectionEvent { - pub(super) event: ConnectionEvent, - _permit: Option, -} - -impl PendingConnectionEvent { - fn exited(reason: PumpExitReason) -> Self { - Self { - event: ConnectionEvent::Exited(reason), - _permit: None, - } - } -} - -/// 下游写入阶段；名称刻意使用 write，而不是暗示客户端已消费的 delivery -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum FramePhase { - Metadata, - First, - Data, - Terminal, - FirstAndTerminal, - Error, - ConnectionLimit, - Close, -} - -impl FramePhase { - const fn as_str(self) -> &'static str { - match self { - Self::Metadata => "metadata", - Self::First => "first", - Self::Data => "data", - Self::Terminal => "terminal", - Self::FirstAndTerminal => "first_and_terminal", - Self::Error => "error", - Self::ConnectionLimit => "connection_limit", - Self::Close => "close", - } - } - - const fn is_milestone(self) -> bool { - matches!( - self, - Self::First - | Self::Terminal - | Self::FirstAndTerminal - | Self::Error - | Self::ConnectionLimit - | Self::Close - ) - } -} - -/// 一次下游写入的请求归属与协议阶段 -#[derive(Clone)] -pub struct WriteContext { - request_id: Option>, - phase: FramePhase, -} - -impl WriteContext { - /// 创建归属于某个请求的写入上下文 - #[must_use] - pub fn request(request_id: &Arc, phase: FramePhase) -> Self { - Self { - request_id: Some(Arc::clone(request_id)), - phase, - } - } - - /// 创建连接级写入上下文 - #[must_use] - pub const fn connection(phase: FramePhase) -> Self { - Self { - request_id: None, - phase, - } - } -} - -/// WebSocket pump 停止的稳定原因 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum PumpExitReason { - ClientClose, - PeerEof, - ReadError, - WriteError, - WriteTimeout, - LifecycleShutdown, - ConnectionMaxAge, - CoordinatorDropped, - InboundOverload, - ServerClose, - PumpStopped, - MiddlewareFailed, -} - -impl PumpExitReason { - pub(super) const fn as_str(self) -> &'static str { - match self { - Self::ClientClose => "client_close", - Self::PeerEof => "peer_eof", - Self::ReadError => "read_error", - Self::WriteError => "write_error", - Self::WriteTimeout => "write_timeout", - Self::LifecycleShutdown => "lifecycle_shutdown", - Self::ConnectionMaxAge => "connection_max_age", - Self::CoordinatorDropped => "coordinator_dropped", - Self::InboundOverload => "inbound_overload", - Self::ServerClose => "server_close", - Self::PumpStopped => "pump_stopped", - Self::MiddlewareFailed => "middleware_failed", - } - } -} - -/// 下游 WebSocket 写入失败 -#[derive(Debug, Error)] -pub enum ConnectionWriteError { - #[error("downstream WebSocket pump is closed")] - Closed, - #[error("downstream WebSocket write timed out after {timeout:?}")] - Timeout { timeout: Duration }, - #[error("downstream WebSocket transport write failed: {message}")] - Transport { message: String }, -} - -/// 只有实际 transport 写入才记为 Written；插件丢弃消息不伪造写入成功 -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum WriteOutcome { - Written, - Suppressed, -} - -struct ConnectionCommand { - message: Message, - context: WriteContext, - acknowledged: oneshot::Sender>, -} - -#[derive(Default)] -struct ConnectionStats { - command_queue_high_water: AtomicUsize, - command_backpressure_count: AtomicU64, - ping_received_count: AtomicU64, - ping_written_count: AtomicU64, - pong_received_count: AtomicU64, - last_read_ms: AtomicU64, - last_write_ms: AtomicU64, -} - -impl ConnectionStats { - fn record_read(&self, opened_at: Instant) { - self.last_read_ms.store( - u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX), - Ordering::Relaxed, - ); - } - - fn record_write(&self, opened_at: Instant) { - self.last_write_ms.store( - u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX), - Ordering::Relaxed, - ); - } - - fn observe_command_queue(&self, sender: &mpsc::Sender) { - let queued = OUTBOUND_COMMAND_BUFFER - .saturating_sub(sender.capacity()) - .saturating_add(1) - .min(OUTBOUND_COMMAND_BUFFER); - self.command_queue_high_water - .fetch_max(queued, Ordering::Relaxed); - if sender.capacity() == 0 { - self.command_backpressure_count - .fetch_add(1, Ordering::Relaxed); - } - } -} - -/// 协调层持有的单 owner WebSocket 连接句柄 -pub struct ResponsesWebSocketConnection { - connection_id: Arc, - opened_at: Instant, - expired: Arc, - commands: Option>, - incoming: mpsc::Receiver, - deferred: VecDeque, - exited: oneshot::Receiver, - pump_task: Option>, - stats: Arc, - config: ConnectionConfig, - exit_reason: Option, -} - -impl ResponsesWebSocketConnection { - pub(super) fn new( - socket: WebSocket, - connection_id: String, - cancellation: CancellationToken, - middleware: Option, - headers: Arc<[MiddlewareHeader]>, - ) -> Self { - spawn_connection( - socket, - Arc::::from(connection_id), - cancellation, - ConnectionConfig::PRODUCTION, - middleware, - headers, - ) - } - - pub(super) fn id(&self) -> &str { - &self.connection_id - } - - /// 返回连接是否已经达到生命周期上限 - #[must_use] - pub fn is_expired(&self) -> bool { - self.expired.load(Ordering::Acquire) - } - - pub(super) fn age(&self) -> Duration { - self.opened_at.elapsed() - } - - /// 等待下一个需要业务层处理的客户端事件 - pub async fn next_event(&mut self) -> Option { - if self.exit_reason.is_none() { - match self.exited.try_recv() { - Ok(reason) => { - self.exit_reason = Some(reason); - } - Err(tokio::sync::oneshot::error::TryRecvError::Closed) => { - self.exit_reason = Some(PumpExitReason::PumpStopped); - } - Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {} - } - } - if let Some(reason) = self.exit_reason { - return Some(ConnectionEvent::Exited(reason)); - } - if let Some(event) = self.deferred.pop_front() { - return Some(event.event); - } - self.next_active_event().await.map(|event| event.event) - } - - /// 活动响应期间只消费新到的帧，已排队的下一轮请求继续保持串行 - pub(super) async fn next_active_event(&mut self) -> Option { - if let Some(reason) = self.exit_reason { - return Some(PendingConnectionEvent::exited(reason)); - } - let event = tokio::select! { - biased; - reason = &mut self.exited => { - Some(PendingConnectionEvent::exited(reason.unwrap_or(PumpExitReason::PumpStopped))) - } - event = self.incoming.recv() => event, - }; - if let Some(PendingConnectionEvent { - event: ConnectionEvent::Exited(reason), - .. - }) = event.as_ref() - { - self.exit_reason.get_or_insert(*reason); - } else if event.is_none() { - self.exit_reason.get_or_insert(PumpExitReason::PumpStopped); - } - event - } - - pub(super) fn defer(&mut self, event: PendingConnectionEvent) { - self.deferred.push_back(event); - } - - /// 等待连接退出，不消费留给后续串行请求的业务帧 - pub async fn wait_for_exit(&mut self) -> PumpExitReason { - if let Some(reason) = self.exit_reason { - return reason; - } - let reason = (&mut self.exited) - .await - .unwrap_or(PumpExitReason::PumpStopped); - self.exit_reason = Some(reason); - reason - } - - /// 串行写入文本帧，并等待 pump 确认 transport 写入结果 - /// - /// # Errors - /// - /// pump 已关闭、写入超时或 transport 失败时返回稳定错误 - pub async fn send_text( - &mut self, - payload: String, - context: WriteContext, - ) -> Result { - self.send(Message::Text(payload.into()), context).await - } - - pub(super) async fn close_policy( - &mut self, - reason: &'static str, - request_id: Option<&Arc>, - ) { - let context = request_id.map_or_else( - || WriteContext::connection(FramePhase::Close), - |request_id| WriteContext::request(request_id, FramePhase::Close), - ); - self.close( - CloseFrame { - code: close_code::POLICY, - reason: reason.into(), - }, - context, - PumpExitReason::ServerClose, - ) - .await; - } - - pub(super) async fn close_for_connection_limit(&mut self, reason: &'static str) { - self.close( - CloseFrame { - code: close_code::NORMAL, - reason: reason.into(), - }, - WriteContext::connection(FramePhase::Close), - PumpExitReason::ConnectionMaxAge, - ) - .await; - } - - async fn close( - &mut self, - frame: CloseFrame, - context: WriteContext, - exit_reason: PumpExitReason, - ) { - if self - .send(Message::Close(Some(frame)), context) - .await - .is_ok() - { - self.exit_reason.get_or_insert(exit_reason); - } - } - - async fn send( - &mut self, - message: Message, - context: WriteContext, - ) -> Result { - let Some(commands) = self.commands.as_ref().cloned() else { - return Err(ConnectionWriteError::Closed); - }; - self.stats.observe_command_queue(&commands); - let (acknowledged, acknowledgement) = oneshot::channel(); - let started_at = Instant::now(); - let request_id = context.request_id.clone(); - let phase = context.phase; - let command = ConnectionCommand { - message, - context, - acknowledged, - }; - let write = async move { - commands - .send(command) - .await - .map_err(|_| ConnectionWriteError::Closed)?; - acknowledgement - .await - .map_err(|_| ConnectionWriteError::Closed)? - }; - let result = match timeout(self.config.write_timeout, write).await { - Ok(result) => result, - Err(_) => Err(ConnectionWriteError::Timeout { - timeout: self.config.write_timeout, - }), - }; - let duration = started_at.elapsed(); - match &result { - Ok(WriteOutcome::Written) if phase.is_milestone() => tracing::info!( - websocket_connection_id = %self.connection_id, - request_id = request_id.as_deref().unwrap_or(""), - frame_phase = phase.as_str(), - write_duration_ms = duration.as_millis(), - "Responses WebSocket frame write succeeded" - ), - Ok(outcome) => tracing::debug!( - websocket_connection_id = %self.connection_id, - request_id = request_id.as_deref().unwrap_or(""), - frame_phase = phase.as_str(), - write_duration_ms = duration.as_millis(), - ?outcome, - "Responses WebSocket frame processed" - ), - Err(error) => tracing::info!( - websocket_connection_id = %self.connection_id, - request_id = request_id.as_deref().unwrap_or(""), - frame_phase = phase.as_str(), - write_duration_ms = duration.as_millis(), - error = %error, - "Responses WebSocket frame write failed" - ), - } - if let Err(error) = &result { - let reason = match error { - ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout, - ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError, - ConnectionWriteError::Closed => PumpExitReason::PumpStopped, - }; - self.terminate(reason); - } - result - } - - fn terminate(&mut self, reason: PumpExitReason) { - self.exit_reason.get_or_insert(reason); - self.commands.take(); - if let Some(task) = self.pump_task.take() { - task.abort(); - } - } - - pub(super) fn log_summary(&self, request_count: u64) { - let connection_age_ms = self.age().as_millis(); - let read_idle_ms = connection_age_ms - .saturating_sub(u128::from(self.stats.last_read_ms.load(Ordering::Relaxed))); - let write_idle_ms = connection_age_ms - .saturating_sub(u128::from(self.stats.last_write_ms.load(Ordering::Relaxed))); - tracing::info!( - websocket_connection_id = %self.connection_id, - request_count, - connection_age_ms, - pump_exit_reason = self - .exit_reason - .map_or("unknown", PumpExitReason::as_str), - command_queue_high_water = self - .stats - .command_queue_high_water - .load(Ordering::Relaxed), - command_backpressure_count = self - .stats - .command_backpressure_count - .load(Ordering::Relaxed), - ping_received_count = self.stats.ping_received_count.load(Ordering::Relaxed), - ping_written_count = self.stats.ping_written_count.load(Ordering::Relaxed), - pong_received_count = self.stats.pong_received_count.load(Ordering::Relaxed), - read_idle_ms, - write_idle_ms, - "Responses WebSocket disconnected" - ); - } -} - -impl Drop for ResponsesWebSocketConnection { - fn drop(&mut self) { - self.commands.take(); - if let Some(task) = self.pump_task.take() { - task.abort(); - } - } -} - -/// 为 WebSocket transport 启动包含冻结中间件计划的单 owner pump -pub fn spawn_connection( - socket: S, - connection_id: Arc, - cancellation: CancellationToken, - config: ConnectionConfig, - middleware: Option, - headers: Arc<[MiddlewareHeader]>, -) -> ResponsesWebSocketConnection -where - S: Stream> + Sink + Unpin + Send + 'static, - E: fmt::Display + Send + 'static, -{ - let opened_at = Instant::now(); - let expired = Arc::new(AtomicBool::new(false)); - let stats = Arc::new(ConnectionStats::default()); - let (command_tx, command_rx) = mpsc::channel(OUTBOUND_COMMAND_BUFFER); - let (incoming_tx, incoming_rx) = mpsc::channel(INBOUND_EVENT_BUFFER); - let (exit_tx, exit_rx) = oneshot::channel(); - let pump_task = tokio::spawn(run_pump( - socket, - command_rx, - incoming_tx, - exit_tx, - Arc::clone(&connection_id), - cancellation, - opened_at, - Arc::clone(&expired), - Arc::clone(&stats), - config, - middleware, - headers, - )); - ResponsesWebSocketConnection { - connection_id, - opened_at, - expired, - commands: Some(command_tx), - incoming: incoming_rx, - deferred: VecDeque::new(), - exited: exit_rx, - pump_task: Some(pump_task), - stats, - config, - exit_reason: None, - } -} - -#[expect(clippy::too_many_arguments)] -async fn run_pump( - socket: S, - commands: mpsc::Receiver, - incoming: mpsc::Sender, - exited: oneshot::Sender, - connection_id: Arc, - cancellation: CancellationToken, - opened_at: Instant, - expired: Arc, - stats: Arc, - config: ConnectionConfig, - middleware: Option, - headers: Arc<[MiddlewareHeader]>, -) where - S: Stream> + Sink + Unpin + Send + 'static, - E: fmt::Display + Send + 'static, -{ - let context = Arc::new(PumpContext { - connection_id, - cancellation: cancellation.child_token(), - opened_at, - expired, - stats, - config, - middleware, - headers, - }); - let _lifetime = middleware::CancelOnDrop(Some(context.cancellation.clone())); - let (writer, mut reader) = socket.split(); - let writer = Arc::new(middleware::Writer::new(writer, context.clone())); - let sender: Arc = writer.clone(); - // 同一个受管任务同时驱动两侧；等待写入不能阻塞入站控制、关闭或取消 - let reason = tokio::select! { - biased; - reason = read_connection(&mut reader, &incoming, sender.clone(), &context) => reason, - reason = write_connection(writer.clone(), commands, &context) => reason, - }; - context.cancellation.cancel(); - let _ = exited.send(reason); - tracing::debug!(websocket_connection_id = %context.connection_id, pump_exit_reason = reason.as_str(), "Responses WebSocket pump exited"); -} - -struct PumpContext { - middleware: Option, - headers: Arc<[MiddlewareHeader]>, - connection_id: Arc, - cancellation: CancellationToken, - opened_at: Instant, - expired: Arc, - stats: Arc, - config: ConnectionConfig, -} - -async fn read_connection( - socket: &mut S, - incoming: &mpsc::Sender, - sender: Arc, - context: &PumpContext, -) -> PumpExitReason -where - S: Stream> + Unpin, - E: fmt::Display, -{ - let event_slots = Arc::new(Semaphore::new(INBOUND_EVENT_BUFFER)); - let deadline = tokio::time::sleep_until(context.opened_at + context.config.max_age); - tokio::pin!(deadline); - let mut deadline_elapsed = false; - loop { - tokio::select! { - biased; - () = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown, - () = &mut deadline, if !deadline_elapsed => { - deadline_elapsed = true; - context.expired.store(true, Ordering::Release); - // 到期只阻止下一轮，当前响应仍由协调层按原有合同收尾 - match emit_incoming(incoming, &event_slots, ConnectionEvent::Expired) { - Ok(()) | Err(PumpExitReason::InboundOverload) => {} - Err(reason) => return reason, - } - } - message = socket.next() => { - if matches!(&message, Some(Ok(_))) { context.stats.record_read(context.opened_at); } - let message = match message { - Some(Ok(message)) => match middleware::transform(message, contract::Direction::Incoming, sender.clone(), context).await { - Ok(Some(message)) => Some(Ok(message)), - Ok(None) => continue, - Err(_) => return PumpExitReason::MiddlewareFailed, - }, - other => other, - }; - let event = match message { - Some(Ok(Message::Text(payload))) => Some(ConnectionEvent::Text(payload.to_string())), - Some(Ok(Message::Binary(_))) => Some(ConnectionEvent::Binary), - Some(Ok(Message::Ping(_))) => { - // Axum/tungstenite 在继续读取时自动刷新 Pong，不能再手工发送一份 - context.stats.ping_received_count.fetch_add(1, Ordering::Relaxed); - None - } - Some(Ok(Message::Pong(_))) => { - context.stats.pong_received_count.fetch_add(1, Ordering::Relaxed); - None - } - Some(Ok(Message::Close(_))) => return PumpExitReason::ClientClose, - Some(Err(error)) => { - tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket receive failed"); - return PumpExitReason::ReadError; - } - None => return PumpExitReason::PeerEof, - }; - if let Some(event) = event && let Err(reason) = emit_incoming(incoming, &event_slots, event) { return reason; } - } - } - } -} - -async fn write_connection( - writer: Arc>, - mut commands: mpsc::Receiver, - context: &PumpContext, -) -> PumpExitReason -where - S: Sink + Unpin + Send + 'static, - E: fmt::Display + Send + 'static, -{ - let mut heartbeat = tokio::time::interval_at( - context.opened_at + DOWNSTREAM_PING_INTERVAL, - DOWNSTREAM_PING_INTERVAL, - ); - heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - let mut ping_sequence = 0_u64; - loop { - let message = tokio::select! { - biased; - () = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown, - command = commands.recv() => { - let Some(command) = command else { return PumpExitReason::CoordinatorDropped; }; - let ConnectionCommand { message, context: write_context, acknowledged } = command; - let closing = matches!(message, Message::Close(_)); - let result = match middleware::transform(message, contract::Direction::Outgoing, writer.clone(), context).await { - Ok(Some(message)) => writer.write(message).await.map(|()| WriteOutcome::Written), - Ok(None) => Ok(WriteOutcome::Suppressed), - Err(error) => Err(ConnectionWriteError::Transport { message: error.to_string() }), - }; - let reason = result.as_ref().err().map(|error| write_exit_reason(error, context)); - if let Err(error) = &result { - tracing::debug!(websocket_connection_id = %context.connection_id, request_id = write_context.request_id.as_deref().unwrap_or(""), frame_phase = write_context.phase.as_str(), error = %error, "Responses WebSocket pump transport write failed"); - } - let _ = acknowledged.send(result); - if let Some(reason) = reason { return reason; } - if closing { return PumpExitReason::ServerClose; } - continue; - } - _ = heartbeat.tick() => { - ping_sequence = ping_sequence.wrapping_add(1); - Message::Ping(ping_sequence.to_be_bytes().to_vec().into()) - } - }; - let message = match middleware::transform( - message, - contract::Direction::Outgoing, - writer.clone(), - context, - ) - .await - { - Ok(Some(message)) => message, - Ok(None) => continue, - Err(_) => return PumpExitReason::MiddlewareFailed, - }; - if let Err(error) = writer.write(message).await { - tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket control write failed"); - return write_exit_reason(&error, context); - } - context - .stats - .ping_written_count - .fetch_add(1, Ordering::Relaxed); - } -} - -async fn write_message( - socket: &mut S, - message: Message, - context: &PumpContext, -) -> Result<(), ConnectionWriteError> -where - S: Sink + Unpin, - E: fmt::Display, -{ - tokio::select! { - biased; - () = context.cancellation.cancelled() => Err(ConnectionWriteError::Closed), - result = timeout(context.config.write_timeout, socket.send(message)) => { - match result { - Ok(Ok(())) => { context.stats.record_write(context.opened_at); Ok(()) } - Ok(Err(error)) => Err(ConnectionWriteError::Transport { message: error.to_string() }), - Err(_) => Err(ConnectionWriteError::Timeout { timeout: context.config.write_timeout }), - } - } - } -} - -fn write_exit_reason(error: &ConnectionWriteError, context: &PumpContext) -> PumpExitReason { - match error { - ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout, - ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError, - ConnectionWriteError::Closed if context.cancellation.is_cancelled() => { - PumpExitReason::LifecycleShutdown - } - ConnectionWriteError::Closed => PumpExitReason::PumpStopped, - } -} - -fn emit_incoming( - incoming: &mpsc::Sender, - slots: &Arc, - event: ConnectionEvent, -) -> Result<(), PumpExitReason> { - let permit = Arc::clone(slots) - .try_acquire_owned() - .map_err(|_| PumpExitReason::InboundOverload)?; - incoming - .try_send(PendingConnectionEvent { - event, - _permit: Some(permit), - }) - .map_err(|error| match error { - mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload, - mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped, - }) -} +pub use handle::{ResponsesWebSocketConnection, spawn_connection}; diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs new file mode 100644 index 000000000..14168917a --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs @@ -0,0 +1,230 @@ +//! 单 owner 驱动连接读写、心跳与有界入站队列 + +use super::{ + frame::*, + middleware, + state::{ConnectionCommand, ConnectionStats, PumpContext}, +}; +use axum::extract::ws::Message; +use futures::{Sink, Stream, StreamExt}; +use gateway_core::{ + engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader, websocket as contract}, + lifecycle::CancellationToken, +}; +use std::{ + fmt, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, +}; +use tokio::{ + sync::{Semaphore, mpsc, oneshot}, + time::Instant, +}; + +#[expect(clippy::too_many_arguments)] +pub(super) async fn run_pump( + socket: S, + commands: mpsc::Receiver, + incoming: mpsc::Sender, + exited: oneshot::Sender, + connection_id: Arc, + cancellation: CancellationToken, + opened_at: Instant, + expired: Arc, + stats: Arc, + config: ConnectionConfig, + middleware: Option, + headers: Arc<[MiddlewareHeader]>, +) where + S: Stream> + Sink + Unpin + Send + 'static, + E: fmt::Display + Send + 'static, +{ + let context = Arc::new(PumpContext { + connection_id, + cancellation: cancellation.child_token(), + opened_at, + expired, + stats, + config, + middleware, + headers, + }); + let _lifetime = middleware::CancelOnDrop(Some(context.cancellation.clone())); + let (writer, mut reader) = socket.split(); + let writer = Arc::new(middleware::Writer::new(writer, context.clone())); + let sender: Arc = writer.clone(); + // 同一个受管任务同时驱动两侧；等待写入不能阻塞入站控制、关闭或取消 + let reason = tokio::select! { + biased; + reason = read_connection(&mut reader, &incoming, sender.clone(), &context) => reason, + reason = write_connection(writer.clone(), commands, &context) => reason, + }; + if matches!(reason, PumpExitReason::ClientClose) { + if let Err(error) = writer.finish_peer_close(&cancellation).await { + tracing::debug!(websocket_connection_id = %context.connection_id, %error, "Responses WebSocket close acknowledgement could not be flushed"); + } + } else { + context.cancellation.cancel(); + } + let _ = exited.send(reason); + tracing::debug!(websocket_connection_id = %context.connection_id, pump_exit_reason = reason.as_str(), "Responses WebSocket pump exited"); +} + +async fn read_connection( + socket: &mut S, + incoming: &mpsc::Sender, + sender: Arc, + context: &PumpContext, +) -> PumpExitReason +where + S: Stream> + Unpin, + E: fmt::Display, +{ + let event_slots = Arc::new(Semaphore::new(INBOUND_EVENT_BUFFER)); + let deadline = tokio::time::sleep_until(context.opened_at + context.config.max_age); + tokio::pin!(deadline); + let mut deadline_elapsed = false; + loop { + tokio::select! { + biased; + () = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown, + () = &mut deadline, if !deadline_elapsed => { + deadline_elapsed = true; + context.expired.store(true, Ordering::Release); + // 到期只阻止下一轮，当前响应仍由协调层按原有合同收尾 + match emit_incoming(incoming, &event_slots, ConnectionEvent::Expired) { + Ok(()) | Err(PumpExitReason::InboundOverload) => {} + Err(reason) => return reason, + } + } + message = socket.next() => { + if matches!(&message, Some(Ok(_))) { context.stats.record_read(context.opened_at); } + let message = match message { + Some(Ok(message)) => match middleware::transform(message, contract::Direction::Incoming, sender.clone(), context).await { + Ok(Some(message)) => Some(Ok(message)), + Ok(None) => continue, + Err(_) => return PumpExitReason::MiddlewareFailed, + }, + other => other, + }; + let event = match message { + Some(Ok(Message::Text(payload))) => Some(ConnectionEvent::Text(payload.to_string())), + Some(Ok(Message::Binary(_))) => Some(ConnectionEvent::Binary), + Some(Ok(Message::Ping(_))) => { + // Axum/tungstenite 在继续读取时自动刷新 Pong，不能再手工发送一份 + context.stats.ping_received_count.fetch_add(1, Ordering::Relaxed); + None + } + Some(Ok(Message::Pong(_))) => { + context.stats.pong_received_count.fetch_add(1, Ordering::Relaxed); + None + } + Some(Ok(Message::Close(_))) => return PumpExitReason::ClientClose, + Some(Err(error)) => { + tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket receive failed"); + return PumpExitReason::ReadError; + } + None => return PumpExitReason::PeerEof, + }; + if let Some(event) = event && let Err(reason) = emit_incoming(incoming, &event_slots, event) { return reason; } + } + } + } +} + +async fn write_connection( + writer: Arc>, + mut commands: mpsc::Receiver, + context: &PumpContext, +) -> PumpExitReason +where + S: Sink + Unpin + Send + 'static, + E: fmt::Display + Send + 'static, +{ + let mut heartbeat = tokio::time::interval_at( + context.opened_at + DOWNSTREAM_PING_INTERVAL, + DOWNSTREAM_PING_INTERVAL, + ); + heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut ping_sequence = 0_u64; + loop { + let message = tokio::select! { + biased; + () = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown, + command = commands.recv() => { + let Some(command) = command else { return PumpExitReason::CoordinatorDropped; }; + let ConnectionCommand { message, context: write_context, acknowledged } = command; + let closing = matches!(message, Message::Close(_)); + let result = match middleware::transform(message, contract::Direction::Outgoing, writer.clone(), context).await { + Ok(Some(message)) => writer.write(message).await.map(|()| WriteOutcome::Written), + Ok(None) => Ok(WriteOutcome::Suppressed), + Err(error) => Err(ConnectionWriteError::Transport { message: error.to_string() }), + }; + let reason = result.as_ref().err().map(|error| write_exit_reason(error, context)); + if let Err(error) = &result { + tracing::debug!(websocket_connection_id = %context.connection_id, request_id = write_context.request_id.as_deref().unwrap_or(""), frame_phase = write_context.phase.as_str(), error = %error, "Responses WebSocket pump transport write failed"); + } + let _ = acknowledged.send(result); + if let Some(reason) = reason { return reason; } + if closing { return PumpExitReason::ServerClose; } + continue; + } + _ = heartbeat.tick() => { + ping_sequence = ping_sequence.wrapping_add(1); + Message::Ping(ping_sequence.to_be_bytes().to_vec().into()) + } + }; + let message = match middleware::transform( + message, + contract::Direction::Outgoing, + writer.clone(), + context, + ) + .await + { + Ok(Some(message)) => message, + Ok(None) => continue, + Err(_) => return PumpExitReason::MiddlewareFailed, + }; + if let Err(error) = writer.write(message).await { + tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket control write failed"); + return write_exit_reason(&error, context); + } + context + .stats + .ping_written_count + .fetch_add(1, Ordering::Relaxed); + } +} + +fn write_exit_reason(error: &ConnectionWriteError, context: &PumpContext) -> PumpExitReason { + match error { + ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout, + ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError, + ConnectionWriteError::Closed if context.cancellation.is_cancelled() => { + PumpExitReason::LifecycleShutdown + } + ConnectionWriteError::Closed => PumpExitReason::PumpStopped, + } +} + +fn emit_incoming( + incoming: &mpsc::Sender, + slots: &Arc, + event: ConnectionEvent, +) -> Result<(), PumpExitReason> { + let permit = Arc::clone(slots) + .try_acquire_owned() + .map_err(|_| PumpExitReason::InboundOverload)?; + incoming + .try_send(PendingConnectionEvent { + event, + _permit: Some(permit), + }) + .map_err(|error| match error { + mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload, + mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped, + }) +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/state.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/state.rs new file mode 100644 index 000000000..558ea7073 --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/state.rs @@ -0,0 +1,75 @@ +//! 单连接共享的命令、运行预算与观测状态 + +use super::frame::{ + ConnectionConfig, ConnectionWriteError, OUTBOUND_COMMAND_BUFFER, WriteContext, WriteOutcome, +}; +use axum::extract::ws::Message; +use gateway_core::{ + engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader}, + lifecycle::CancellationToken, +}; +use std::sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, +}; +use tokio::{ + sync::{mpsc, oneshot}, + time::Instant, +}; + +pub(super) struct ConnectionCommand { + pub(super) message: Message, + pub(super) context: WriteContext, + pub(super) acknowledged: oneshot::Sender>, +} + +#[derive(Default)] +pub(super) struct ConnectionStats { + pub(super) command_queue_high_water: AtomicUsize, + pub(super) command_backpressure_count: AtomicU64, + pub(super) ping_received_count: AtomicU64, + pub(super) ping_written_count: AtomicU64, + pub(super) pong_received_count: AtomicU64, + pub(super) last_read_ms: AtomicU64, + pub(super) last_write_ms: AtomicU64, +} + +impl ConnectionStats { + pub(super) fn record_read(&self, opened_at: Instant) { + self.last_read_ms.store( + u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + } + + pub(super) fn record_write(&self, opened_at: Instant) { + self.last_write_ms.store( + u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + } + + pub(super) fn observe_command_queue(&self, sender: &mpsc::Sender) { + let queued = OUTBOUND_COMMAND_BUFFER + .saturating_sub(sender.capacity()) + .saturating_add(1) + .min(OUTBOUND_COMMAND_BUFFER); + self.command_queue_high_water + .fetch_max(queued, Ordering::Relaxed); + if sender.capacity() == 0 { + self.command_backpressure_count + .fetch_add(1, Ordering::Relaxed); + } + } +} + +pub(super) struct PumpContext { + pub(super) middleware: Option, + pub(super) headers: Arc<[MiddlewareHeader]>, + pub(super) connection_id: Arc, + pub(super) cancellation: CancellationToken, + pub(super) opened_at: Instant, + pub(super) expired: Arc, + pub(super) stats: Arc, + pub(super) config: ConnectionConfig, +} diff --git a/backend/crates/gateway-api/tests/admin/plugins/management.rs b/backend/crates/gateway-api/tests/admin/plugins/management.rs index 1954286b3..43bc95851 100644 --- a/backend/crates/gateway-api/tests/admin/plugins/management.rs +++ b/backend/crates/gateway-api/tests/admin/plugins/management.rs @@ -5,24 +5,31 @@ use axum::{ http::{Request, StatusCode}, }; use futures::future::BoxFuture; -use gateway_core::{ - engine::{ - execution::{ - AuthenticatedClient, ClientAuthenticationError, ExecutionService, - PreparedRootExecution, StartExecution, StartProviderExecution, StartedExecution, - }, - middleware::{ - FrozenMiddlewarePlan, MiddlewareBody, MiddlewareContext, MiddlewareError, - MiddlewareFrame, MiddlewareFraming, MiddlewareHeader, MiddlewareNext, MiddlewarePlan, - MiddlewareRequest, MiddlewareResponse, - }, - }, - error::GatewayError, - lifecycle::CancellationToken, - policy::ClientApiKeyId, - routing::PublicModelId, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::engine::execution::AuthenticatedClient; +use gateway_core::engine::execution::ClientAuthenticationError; +use gateway_core::engine::execution::ExecutionService; +use gateway_core::engine::execution::PreparedRootExecution; +use gateway_core::engine::execution::StartExecution; +use gateway_core::engine::execution::StartProviderExecution; +use gateway_core::engine::execution::StartedExecution; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareBody; +use gateway_core::engine::middleware::MiddlewareContext; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareFrame; +use gateway_core::engine::middleware::MiddlewareFraming; +use gateway_core::engine::middleware::MiddlewareHeader; +use gateway_core::engine::middleware::MiddlewareNext; +use gateway_core::engine::middleware::MiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareRequest; +use gateway_core::engine::middleware::MiddlewareResponse; +use gateway_core::error::GatewayError; +use gateway_core::lifecycle::CancellationToken; +use gateway_core::policy::ClientApiKeyId; +use gateway_core::routing::PublicModelId; +use gateway_core::routing::extensions::ExtensionSetId; +use gateway_core::routing::extensions::ExtensionSetLease; +use gateway_core::routing::extensions::ExtensionSetReference; use std::sync::{ Arc, Mutex, atomic::{AtomicBool, Ordering}, @@ -368,10 +375,11 @@ pub(super) fn resource_fixture( } async fn published_fixture() -> AdminTestFixture { - use gateway_core::{ - routing::{ConfigRevision, RuntimeSnapshot}, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, - }; + use gateway_core::routing::ConfigRevision; + use gateway_core::routing::RuntimeSnapshot; + use gateway_core::routing::extensions::ExtensionSetId; + use gateway_core::routing::extensions::ExtensionSetLease; + use gateway_core::routing::extensions::ExtensionSetReference; use std::sync::Arc; struct Lease; impl ExtensionSetLease for Lease { diff --git a/backend/crates/gateway-api/tests/admin/plugins/mod.rs b/backend/crates/gateway-api/tests/admin/plugins/mod.rs index da463be7c..b5b11361e 100644 --- a/backend/crates/gateway-api/tests/admin/plugins/mod.rs +++ b/backend/crates/gateway-api/tests/admin/plugins/mod.rs @@ -546,7 +546,7 @@ pub(super) struct TestPluginPorts { impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPorts { async fn validate_target( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, ) -> Result<(), AdminError> { Ok(()) @@ -554,7 +554,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor async fn start_callback( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: gateway_admin::model::plugins::management::StartPluginManagementCallback, _: &gateway_admin::model::auth::AdminRequestContext, @@ -564,7 +564,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn callback( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: &str, request: gateway_admin::model::plugins::management::PluginManagementRequest, @@ -576,14 +576,14 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn views( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, ) -> Result, AdminError> { Ok(Vec::new()) } async fn resource( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, path: &str, public: bool, @@ -596,7 +596,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn handle( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, request: gateway_admin::model::plugins::management::PluginManagementRequest, ) -> Result @@ -611,7 +611,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts { async fn activate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::instances::PluginInstance, ) -> Result<(), AdminError> { Ok(()) @@ -621,7 +621,7 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts { } async fn migrate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: gateway_admin::model::plugins::state::PluginStateTransition, ) -> Result<(), AdminError> { panic!("unexpected state migration") @@ -647,7 +647,7 @@ impl gateway_admin::ports::plugins::PluginPreparation for TestPluginPorts { async fn prepare( &self, _: gateway_admin::model::plugins::instances::PluginInstanceSnapshot, - ) -> Result { + ) -> Result { Err(AdminError::invalid("unused plugin fixture")) } } @@ -658,7 +658,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for TestPluginPorts &self, snapshot: &gateway_admin::model::plugins::instances::PluginInstanceSnapshot, _: Option, - _: Option<&gateway_core::runtime::extensions::ExtensionSetReference>, + _: Option<&gateway_core::routing::extensions::ExtensionSetReference>, ) -> Option> { Some( diff --git a/backend/crates/gateway-api/tests/architecture.rs b/backend/crates/gateway-api/tests/architecture.rs index 24ca76d0a..32a05ed04 100644 --- a/backend/crates/gateway-api/tests/architecture.rs +++ b/backend/crates/gateway-api/tests/architecture.rs @@ -1,4 +1,4 @@ -//! 检查 API 层依赖、源码与测试模块边界，以及管理路由约定 +//! 检查 API 层依赖与管理路由约定，模块布局和测试镜像由 workspace 架构检查覆盖 use std::{fs, path::Path}; @@ -29,150 +29,6 @@ fn workspace_should_include_gateway_api() { assert!(workspace.contains("\"crates/gateway-api\"")); } -#[test] -fn source_tree_should_match_frozen_machine_manifest() { - let root = Path::new(env!("CARGO_MANIFEST_DIR")); - let mut actual = rust_files(&root.join("src")); - actual.sort(); - let mut expected = vec![ - "src/admin/account_groups.rs", - "src/admin/accounts/credentials.rs", - "src/admin/accounts/handlers.rs", - "src/admin/accounts/import_tasks.rs", - "src/admin/accounts/mod.rs", - "src/admin/accounts/presenter.rs", - "src/admin/accounts/wire.rs", - "src/admin/auth.rs", - "src/admin/backups.rs", - "src/admin/client_keys.rs", - "src/admin/diagnostics.rs", - "src/admin/extract.rs", - "src/admin/mod.rs", - "src/admin/observability/mod.rs", - "src/admin/observability/presenter.rs", - "src/admin/observability/query.rs", - "src/admin/observability/routes.rs", - "src/admin/observability/wire.rs", - "src/admin/plugins/artifacts.rs", - "src/admin/plugins/distribution.rs", - "src/admin/plugins/instances.rs", - "src/admin/plugins/management.rs", - "src/admin/plugins/mod.rs", - "src/admin/presenter.rs", - "src/admin/proxies.rs", - "src/admin/settings.rs", - "src/admin/system.rs", - "src/admin/wire.rs", - "src/auth.rs", - "src/health.rs", - "src/key_usage/mod.rs", - "src/key_usage/query.rs", - "src/key_usage/presenter.rs", - "src/lib.rs", - "src/middleware/headers.rs", - "src/middleware/mod.rs", - "src/middleware/websocket.rs", - "src/openai/middleware/http.rs", - "src/openai/middleware/mod.rs", - "src/openai/middleware/request.rs", - "src/openai/middleware/session.rs", - "src/openai/auth.rs", - "src/openai/error.rs", - "src/openai/images/http.rs", - "src/openai/images/mod.rs", - "src/openai/live/http.rs", - "src/openai/live/mod.rs", - "src/openai/live/websocket.rs", - "src/openai/mod.rs", - "src/openai/models.rs", - "src/openai/endpoint.rs", - "src/openai/responses/error.rs", - "src/openai/responses/http.rs", - "src/openai/responses/mod.rs", - "src/openai/responses/request.rs", - "src/openai/responses/response.rs", - "src/openai/responses/validation.rs", - "src/openai/responses/websocket/connection/mod.rs", - "src/openai/responses/websocket/connection/middleware.rs", - "src/openai/responses/websocket/forward.rs", - "src/openai/responses/websocket/mod.rs", - "src/openai/responses/websocket/protocol.rs", - "src/openai/router.rs", - "src/openai/search/http.rs", - "src/openai/search/mod.rs", - "src/openai/service.rs", - "src/openai/usage.rs", - "src/provider.rs", - "src/session_cookie.rs", - "src/time.rs", - ]; - expected.sort_unstable(); - - assert_eq!(actual, expected); -} - -#[test] -fn test_tree_should_match_frozen_rust_mirror() { - let root = Path::new(env!("CARGO_MANIFEST_DIR")); - let mut actual = rust_files(&root.join("tests")); - actual.sort(); - let mut expected = vec![ - "tests/admin/account_groups.rs", - "tests/admin/accounts/handlers.rs", - "tests/admin/accounts/import_tasks.rs", - "tests/admin/accounts/mod.rs", - "tests/admin/accounts/presenter.rs", - "tests/admin/auth.rs", - "tests/admin/client_keys.rs", - "tests/admin/errors.rs", - "tests/admin/error_diagnostics.rs", - "tests/admin/mod.rs", - "tests/admin/observability/mod.rs", - "tests/admin/observability/query.rs", - "tests/admin/observability/response.rs", - "tests/admin/plugins/management.rs", - "tests/admin/plugins/mod.rs", - "tests/admin/proxies.rs", - "tests/admin/settings.rs", - "tests/admin/system.rs", - "tests/admin/wire.rs", - "tests/architecture.rs", - "tests/auth.rs", - "tests/support.rs", - "tests/health.rs", - "tests/key_usage/mod.rs", - "tests/key_usage/fixtures.rs", - "tests/main.rs", - "tests/middleware/mod.rs", - "tests/middleware/websocket.rs", - "tests/openai/middleware/mod.rs", - "tests/openai/auth.rs", - "tests/openai/error.rs", - "tests/openai/images/mod.rs", - "tests/openai/mod.rs", - "tests/openai/models.rs", - "tests/openai/endpoint.rs", - "tests/openai/responses/http.rs", - "tests/openai/responses/mod.rs", - "tests/openai/responses/request.rs", - "tests/openai/responses/websocket/connection/mod.rs", - "tests/openai/responses/websocket/connection/middleware.rs", - "tests/openai/responses/websocket/forward.rs", - "tests/openai/responses/websocket/interrupt.rs", - "tests/openai/responses/websocket/mod.rs", - "tests/openai/responses/websocket/protocol.rs", - "tests/openai/live/mod.rs", - "tests/openai/live/websocket.rs", - "tests/openai/router.rs", - "tests/openai/search/mod.rs", - "tests/openai/usage.rs", - "tests/provider.rs", - ]; - expected.sort_unstable(); - - assert_eq!(actual, expected); -} - #[test] fn admin_routes_should_use_only_get_post_and_static_paths() { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/admin"); @@ -192,20 +48,6 @@ fn admin_routes_should_use_only_get_post_and_static_paths() { ); } -fn rust_files(root: &Path) -> Vec { - let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); - all_files(root) - .into_iter() - .filter(|path| path.extension().and_then(|value| value.to_str()) == Some("rs")) - .map(|path| { - path.strip_prefix(manifest) - .expect("path below manifest") - .to_string_lossy() - .replace('\\', "/") - }) - .collect() -} - fn all_files(root: &Path) -> Vec { let mut files = Vec::new(); let mut pending = vec![root.to_path_buf()]; diff --git a/backend/crates/gateway-api/tests/middleware/mod.rs b/backend/crates/gateway-api/tests/middleware/mod.rs index 9721385dd..86817603f 100644 --- a/backend/crates/gateway-api/tests/middleware/mod.rs +++ b/backend/crates/gateway-api/tests/middleware/mod.rs @@ -16,15 +16,18 @@ use axum::{ }; use bytes::Bytes; use futures::{future::BoxFuture, stream}; -use gateway_core::{ - engine::middleware::{ - FrozenMiddlewarePlan, MiddlewareContext, MiddlewareError, MiddlewareNext, MiddlewarePlan, - MiddlewareRequest, MiddlewareResponse, - }, - lifecycle::CancellationToken, - middleware::http as contract, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareContext; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareNext; +use gateway_core::engine::middleware::MiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareRequest; +use gateway_core::engine::middleware::MiddlewareResponse; +use gateway_core::engine::middleware::http as contract; +use gateway_core::lifecycle::CancellationToken; +use gateway_core::routing::extensions::ExtensionSetId; +use gateway_core::routing::extensions::ExtensionSetLease; +use gateway_core::routing::extensions::ExtensionSetReference; use http_body::Frame; use http_body_util::{BodyExt as _, StreamBody}; use tower::ServiceExt as _; @@ -70,12 +73,12 @@ impl MiddlewarePlan for Plan { } fn handle_websocket( &self, - context: gateway_core::middleware::websocket::Context, - message: gateway_core::middleware::websocket::Message, - next: gateway_core::middleware::websocket::Next, + context: gateway_core::engine::middleware::websocket::Context, + message: gateway_core::engine::middleware::websocket::Message, + next: gateway_core::engine::middleware::websocket::Next, ) -> BoxFuture< 'static, - Result, MiddlewareError>, + Result, MiddlewareError>, > { assert!(context.plan.is_some()); self.calls.fetch_add(1, Ordering::SeqCst); @@ -103,11 +106,13 @@ impl MiddlewarePlan for Plan { } fn handle_service( &self, - context: gateway_core::middleware::service::Context, + context: gateway_core::engine::middleware::service::Context, input: serde_json::Value, - next: gateway_core::middleware::service::Next, - ) -> BoxFuture<'static, Result> - { + next: gateway_core::engine::middleware::service::Next, + ) -> BoxFuture< + 'static, + Result, + > { self.calls.fetch_add(1, Ordering::SeqCst); assert_eq!(context.operation, "settings.load"); assert_eq!(context.request_id, "service-request"); @@ -146,7 +151,7 @@ impl MiddlewarePlan for Plan { response.extensions_mut().insert(contract::upgrade::Upgraded::new(status, Box::pin(async move { let session = session.await.map_err(|_| MiddlewareError::Fault)?; if idle { return std::future::pending().await; } - use gateway_core::middleware::websocket::{Kind, Message}; + use gateway_core::engine::middleware::websocket::{Kind, Message}; let receive = session.receive(); tokio::pin!(receive); tokio::select! { diff --git a/backend/crates/gateway-api/tests/openai/auth.rs b/backend/crates/gateway-api/tests/openai/auth.rs index 9cc8ffe81..298e0190d 100644 --- a/backend/crates/gateway-api/tests/openai/auth.rs +++ b/backend/crates/gateway-api/tests/openai/auth.rs @@ -374,9 +374,9 @@ fn chatgpt_remote_suffix_created_by_header_truncation_should_not_skip_the_gate() #[tokio::test] async fn http_settings_freeze_before_plan_resolution_and_apply_before_admission() { - use gateway_core::{ - engine::middleware::*, middleware::http as contract, runtime::extensions::*, - }; + use gateway_core::engine::middleware::http as contract; + use gateway_core::engine::middleware::*; + use gateway_core::routing::extensions::*; struct Lease; impl ExtensionSetLease for Lease { fn is_ready(&self) -> bool { diff --git a/backend/crates/gateway-api/tests/openai/middleware/mod.rs b/backend/crates/gateway-api/tests/openai/middleware/mod.rs index 4505a4c71..a61a3a182 100644 --- a/backend/crates/gateway-api/tests/openai/middleware/mod.rs +++ b/backend/crates/gateway-api/tests/openai/middleware/mod.rs @@ -10,15 +10,22 @@ use std::{ use bytes::Bytes; use futures::future::BoxFuture; -use gateway_core::{ - engine::middleware::{ - FrozenMiddlewarePlan, MiddlewareBody, MiddlewareContext, MiddlewareError, MiddlewareFrame, - MiddlewareFraming, MiddlewareHeader, MiddlewareMount, MiddlewareNext, MiddlewarePlan, - MiddlewareRequest, MiddlewareResponse, - }, - operation::OperationKind, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareBody; +use gateway_core::engine::middleware::MiddlewareContext; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareFrame; +use gateway_core::engine::middleware::MiddlewareFraming; +use gateway_core::engine::middleware::MiddlewareHeader; +use gateway_core::engine::middleware::MiddlewareMount; +use gateway_core::engine::middleware::MiddlewareNext; +use gateway_core::engine::middleware::MiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareRequest; +use gateway_core::engine::middleware::MiddlewareResponse; +use gateway_core::operation::OperationKind; +use gateway_core::routing::extensions::ExtensionSetId; +use gateway_core::routing::extensions::ExtensionSetLease; +use gateway_core::routing::extensions::ExtensionSetReference; /// OpenAI API 入口共用的中间件夹具；授权和 RPC 由 Runtime 行为测试覆盖 #[derive(Debug, Default)] @@ -30,7 +37,8 @@ pub(super) struct RequestMiddleware { pub(super) response_actions: Mutex>, pub(super) response_headers: Vec, pub(super) short_circuit_json: Option, - pub(super) settings: Option, + pub(super) settings: + Option, } #[derive(Debug, Clone)] diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs index e54596e90..60ce032ea 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs @@ -2,14 +2,100 @@ use super::*; use bytes::Bytes; -use gateway_core::{ - engine::middleware::{ - FrozenMiddlewarePlan, MiddlewareContext, MiddlewareError, MiddlewareNext, MiddlewarePlan, - MiddlewareRequest, MiddlewareResponse, - }, - middleware::{http as http_contract, websocket as core}, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareContext; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareNext; +use gateway_core::engine::middleware::MiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareRequest; +use gateway_core::engine::middleware::MiddlewareResponse; +use gateway_core::engine::middleware::http as http_contract; +use gateway_core::engine::middleware::websocket as core; +use gateway_core::routing::extensions::ExtensionSetId; +use gateway_core::routing::extensions::ExtensionSetLease; +use gateway_core::routing::extensions::ExtensionSetReference; + +#[tokio::test(start_paused = true)] +async fn client_close_does_not_flush_a_retained_sender_with_an_inflight_write() { + #[derive(Default)] + struct RetainedSenderPlan(Mutex>>); + impl fmt::Debug for RetainedSenderPlan { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("RetainedSenderPlan") + } + } + impl MiddlewarePlan for RetainedSenderPlan { + fn has_websocket(&self) -> bool { + true + } + fn handle_websocket( + &self, + context: core::Context, + message: core::Message, + next: core::Next, + ) -> BoxFuture<'static, Result, MiddlewareError>> { + *self.0.lock().unwrap() = Some(context.sender); + next.run(message) + } + fn handle( + &self, + _: MiddlewareContext, + request: MiddlewareRequest, + next: MiddlewareNext, + ) -> BoxFuture<'static, Result> { + next.run(request) + } + } + + let plan = Arc::new(RetainedSenderPlan::default()); + let frozen = FrozenMiddlewarePlan::new( + plan.clone(), + ExtensionSetReference::new( + ExtensionSetId::new("retained-writer".into()).unwrap(), + Arc::new(Lease), + ), + ); + let (incoming, received) = unbounded_channel(); + let (written, mut output) = unbounded_channel(); + let flush_started = Arc::new(AtomicBool::new(false)); + let socket = TestSocket { + incoming: received, + written, + stall_writes: true, + flush_started: flush_started.clone(), + dropped: Arc::new(AtomicBool::new(false)), + }; + let mut connection = spawn_connection( + socket, + Arc::from("retained-writer"), + CancellationToken::new(), + ConnectionConfig::PRODUCTION, + Some(frozen), + Arc::from([]), + ); + incoming.send(Ok(Message::Text("capture".into()))).unwrap(); + assert!(matches!( + connection.next_event().await, + Some(ConnectionEvent::Text(_)) + )); + let sender = plan.0.lock().unwrap().clone().unwrap(); + let mut write = sender.send(core::Message { + kind: core::Kind::Text, + payload: Bytes::from_static(b"pending"), + }); + assert!(futures::poll!(write.as_mut()).is_pending()); + // 写 future 在 pump 的 select 之外持锁，而且暂不继续 poll 以响应取消 + incoming.send(Ok(Message::Close(None))).unwrap(); + let reason = tokio::time::timeout(Duration::from_millis(50), connection.wait_for_exit()) + .await + .expect("peer close must not wait for a retained sender"); + assert_eq!(reason, PumpExitReason::ClientClose); + assert!(!flush_started.load(Ordering::Acquire)); + assert!(write.await.is_err()); + assert!(matches!(output.try_recv(), Err(TryRecvError::Empty))); + plan.0.lock().unwrap().take(); +} + #[derive(Debug, Default)] struct Plan { block_terminal: bool, diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs index 4ce5dea46..4d6bc7c47 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs @@ -234,6 +234,7 @@ struct TestSocket { incoming: UnboundedReceiver>, written: UnboundedSender, stall_writes: bool, + flush_started: Arc, dropped: Arc, } @@ -268,7 +269,12 @@ impl Sink for TestSocket { } fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) + self.flush_started.store(true, Ordering::Release); + if self.stall_writes { + Poll::Pending + } else { + Poll::Ready(Ok(())) + } } fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { @@ -286,6 +292,7 @@ struct PumpHarness { connection: ResponsesWebSocketConnection, incoming: UnboundedSender>, written: UnboundedReceiver, + flush_started: Arc, dropped: Arc, cancellation: CancellationToken, } @@ -294,11 +301,13 @@ fn test_connection(stall_writes: bool) -> PumpHarness { let (incoming_tx, incoming_rx) = unbounded_channel(); let (written_tx, written_rx) = unbounded_channel(); let dropped = Arc::new(AtomicBool::new(false)); + let flush_started = Arc::new(AtomicBool::new(false)); let cancellation = CancellationToken::new(); let socket = TestSocket { incoming: incoming_rx, written: written_tx, stall_writes, + flush_started: Arc::clone(&flush_started), dropped: Arc::clone(&dropped), }; let connection = spawn_connection( @@ -313,6 +322,7 @@ fn test_connection(stall_writes: bool) -> PumpHarness { connection, incoming: incoming_tx, written: written_rx, + flush_started, dropped, cancellation, } @@ -404,6 +414,62 @@ async fn client_close_should_take_priority_over_queued_business_frames() { )); } +#[tokio::test(start_paused = true)] +async fn client_close_flush_has_a_short_deadline_without_sending_another_frame() { + let PumpHarness { + mut connection, + incoming, + mut written, + flush_started, + dropped, + .. + } = test_connection(true); + incoming.send(Ok(Message::Close(None))).expect("close peer"); + let started = tokio::time::Instant::now(); + + let reason = tokio::time::timeout(Duration::from_secs(2), connection.wait_for_exit()) + .await + .expect("peer close cannot consume the ordinary write timeout"); + assert_eq!(reason, PumpExitReason::ClientClose); + assert_eq!(started.elapsed(), Duration::from_secs(1)); + assert!(flush_started.load(Ordering::Acquire)); + assert!(dropped.load(Ordering::Acquire)); + assert!( + written.recv().await.is_none(), + "do not send a duplicate Close" + ); +} + +#[tokio::test(start_paused = true)] +async fn lifecycle_shutdown_interrupts_a_pending_client_close_flush() { + let PumpHarness { + mut connection, + incoming, + flush_started, + cancellation, + dropped, + .. + } = test_connection(true); + incoming.send(Ok(Message::Close(None))).expect("close peer"); + for _ in 0..100 { + if flush_started.load(Ordering::Acquire) { + break; + } + tokio::task::yield_now().await; + } + assert!( + flush_started.load(Ordering::Acquire), + "close flush must start" + ); + cancellation.cancel(); + + let reason = tokio::time::timeout(Duration::from_millis(50), connection.wait_for_exit()) + .await + .expect("host cancellation must not wait for the close timeout"); + assert_eq!(reason, PumpExitReason::ClientClose); + assert!(dropped.load(Ordering::Acquire)); +} + #[tokio::test] async fn full_business_queue_should_not_block_lifecycle_shutdown() { let PumpHarness { @@ -988,6 +1054,40 @@ async fn requests_deferred_during_active_response_share_the_inbound_queue_limit( server.abort(); } +#[tokio::test] +async fn client_normal_close_after_completed_response_receives_normal_close_reply() { + let (trace, mut socket, server) = start_active_response().await; + trace.release_terminal.notify_one(); + loop { + let message = tokio::time::timeout(Duration::from_secs(2), socket.next()) + .await + .expect("terminal response timeout") + .expect("response remains open") + .expect("valid response frame"); + if let ClientMessage::Text(text) = message { + let event: Value = serde_json::from_str(&text).expect("response JSON"); + if event["type"] == "response.completed" { + break; + } + } + } + socket + .close(Some(tokio_tungstenite::tungstenite::protocol::CloseFrame { + code: CloseCode::Normal, + reason: "client finished".into(), + })) + .await + .expect("send normal close"); + // close() 只发送本端关闭帧；必须读取对端应答，才能证明握手正常完成 + let reply = tokio::time::timeout(Duration::from_secs(2), socket.next()).await; + server.abort(); + assert!( + matches!(&reply, Ok(Some(Ok(ClientMessage::Close(Some(frame))))) + if frame.code == CloseCode::Normal && frame.reason == "client finished"), + "client must receive the close acknowledgement, got {reply:?}" + ); +} + #[tokio::test] async fn client_close_should_cancel_active_execution_without_starting_queued_requests() { let (trace, mut socket, server) = start_active_response().await; @@ -1009,6 +1109,12 @@ async fn client_close_should_cancel_active_execution_without_starting_queued_req .expect("cancel active execution promptly"); assert_eq!(trace.starts.load(Ordering::Acquire), 1); assert!(trace.cancelled.load(Ordering::Acquire)); + let reply = tokio::time::timeout(Duration::from_secs(2), socket.next()) + .await + .expect("close reply timeout") + .expect("close reply") + .expect("clean close handshake"); + assert!(matches!(reply, ClientMessage::Close(None))); server.abort(); } @@ -1073,6 +1179,7 @@ async fn stalled_heartbeat_can_be_cancelled_without_waiting_for_write_timeout() written: _written, dropped, cancellation, + .. } = test_connection(true); tokio::time::advance(Duration::from_secs(25)).await; tokio::task::yield_now().await; diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs index c0c407169..f6c069a2c 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs @@ -19,7 +19,7 @@ use gateway_core::error::{ ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ProviderError, ProviderErrorKind, }; use gateway_core::event::{ProtocolWireEvent, ProviderEvent, ProviderResponseHeader}; -use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::upstream::{OpaqueUpstreamValue, UpstreamSendState}; use serde_json::{Value, json}; use tokio::net::TcpStream; diff --git a/backend/crates/gateway-core/src/engine/authentication.rs b/backend/crates/gateway-core/src/engine/authentication.rs index d6a546ac2..c2659a30e 100644 --- a/backend/crates/gateway-core/src/engine/authentication.rs +++ b/backend/crates/gateway-core/src/engine/authentication.rs @@ -10,7 +10,7 @@ use futures::future::BoxFuture; use crate::{ policy::{ClientApiKeyId, PlaintextClientApiKey}, - runtime::extensions::{ExtensionSetId, ExtensionSetReference}, + routing::extensions::{ExtensionSetId, ExtensionSetReference}, }; pub const MAXIMUM_AUTHORIZATION_BYTES: usize = 8 * 1024; @@ -18,24 +18,29 @@ pub const MAXIMUM_AUTHORIZATION_BYTES: usize = 8 * 1024; /// API 从数据面 Authorization 头构造的有界认证信封 #[derive(Clone, PartialEq, Eq)] pub struct ClientAuthenticationRequest { - settings: Option, + settings: Option, authorization: Arc, native_bearer: Option, } impl ClientAuthenticationRequest { #[must_use] - pub fn with_settings(mut self, settings: crate::settings::RequestSettings) -> Self { + pub fn with_settings( + mut self, + settings: crate::routing::request_settings::RequestSettings, + ) -> Self { self.settings = Some(settings); self } #[must_use] - pub fn settings(&self) -> Option<&crate::settings::RequestSettings> { + pub fn settings(&self) -> Option<&crate::routing::request_settings::RequestSettings> { self.settings.as_ref() } - pub(crate) fn take_settings(&mut self) -> Option { + pub(crate) fn take_settings( + &mut self, + ) -> Option { self.settings.take() } diff --git a/backend/crates/gateway-core/src/engine/execution/admission.rs b/backend/crates/gateway-core/src/engine/execution/admission.rs new file mode 100644 index 000000000..0acb2075c --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/admission.rs @@ -0,0 +1,71 @@ +//! 请求准入租约及费用结算的释放边界 + +use super::failure::record_resource_failure; +use crate::{ + engine::{ + ModelRequestId, + admission::ClientAdmissionPort, + budget::{ClientBudgetCharge, ClientBudgetPort}, + }, + lifecycle::LeaseGuard, + policy::ClientApiKeyId, +}; +use std::sync::Arc; + +pub(super) struct AdmissionLease { + pub(super) renewal: Option>, + pub(super) armed: bool, + pub(super) port: Arc, + pub(super) diagnostics: Arc, + pub(super) client_api_key_id: ClientApiKeyId, + pub(super) model_request_id: ModelRequestId, +} + +pub(super) async fn settle_budget( + port: &dyn ClientBudgetPort, + charge: ClientBudgetCharge, + diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, +) { + let request_id = charge.request_id.clone(); + if let Err(error) = port.settle(charge).await { + record_resource_failure( + diagnostics, + &request_id, + "settle_client_budget", + "Client budget settlement failed; storage will retry on the next request", + error, + ) + .await; + } +} + +impl AdmissionLease { + pub(super) async fn release(mut self) { + self.renewal.take(); + if let Err(error) = self + .port + .release(&self.client_api_key_id, &self.model_request_id) + .await + { + record_resource_failure( + self.diagnostics.as_ref(), + &self.model_request_id, + "release_client_admission", + "Client admission release failed; lease TTL remains active", + error, + ) + .await; + } + self.armed = false; + } +} + +impl Drop for AdmissionLease { + fn drop(&mut self) { + self.renewal.take(); + if self.armed { + self.port + .abandon(&self.client_api_key_id, &self.model_request_id); + } + } +} diff --git a/backend/crates/gateway-core/src/engine/execution/contract.rs b/backend/crates/gateway-core/src/engine/execution/contract.rs new file mode 100644 index 000000000..432821eaa --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/contract.rs @@ -0,0 +1,480 @@ +//! 执行请求、认证事实与对外会话端口 + +use crate::{ + engine::{ + CoordinatedEvent, EngineError, ModelRequestId, + authentication::ClientAuthenticationRequest, + continuation::PreviousResponseId, + extensions::ExtensionCallScope, + middleware::{ + FrozenMiddlewarePlan, MiddlewareAuthority, MiddlewareContext, MiddlewareTarget, + }, + nested::ExecutionEffects, + }, + error::{GatewayError, GatewayErrorKind}, + event::{ProviderEvent, ProviderResponseHeader}, + identity::ProviderKind, + lifecycle::{CancellationToken, Deadline}, + operation::Operation, + policy::{ClientApiKeyId, ClientPolicy}, + routing::{ + ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId, RuntimeSnapshot, + UpstreamModelId, request_settings::RequestSettings, + }, +}; +use futures::future::BoxFuture; +use std::{ + fmt, + net::IpAddr, + sync::Arc, + time::{Duration, SystemTime}, +}; +use uuid::Uuid; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClientTransport { + HttpJson, + HttpSse, + WebSocket, + InternalProbe, + InternalPlugin, +} + +impl ClientTransport { + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::HttpJson => "http_json", + Self::HttpSse => "http_sse", + Self::WebSocket => "websocket", + Self::InternalProbe => "internal", + Self::InternalPlugin => "internal_plugin", + } + } +} + +/// API 解码后交给 Core 的稳定请求元数据 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecutionRequestMetadata { + pub protocol: String, + pub endpoint: String, + pub transport: ClientTransport, + pub stream: bool, + pub client_ip: Option, + pub user_agent: Option, + pub previous_response_id: Option, +} + +#[derive(Clone)] +pub struct AuthenticatedClient { + pub(super) settings: Option, + pub(super) snapshot: Arc, + pub(super) policy: ClientPolicy, + pub(super) authentication: Option, +} + +impl AuthenticatedClient { + pub(super) fn execution_timeout(&self) -> Option { + self.settings + .as_ref() + .and_then(|settings| settings.execution_timeout(self.policy.key_id())) + } + + #[must_use] + pub fn request_settings(&self) -> Option<&RequestSettings> { + self.settings.as_ref() + } + + /// 后续仍须重新认证；这里只更新下一次认证使用的冻结配置 + #[must_use] + pub fn with_request_settings(mut self, settings: RequestSettings) -> Self { + self.settings = Some(settings); + self + } + + #[must_use] + pub const fn snapshot(&self) -> &Arc { + &self.snapshot + } + + #[must_use] + pub const fn policy(&self) -> &ClientPolicy { + &self.policy + } +} + +impl fmt::Debug for AuthenticatedClient { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("AuthenticatedClient") + .field("key_id", &self.policy.key_id()) + .field("revision", &self.snapshot.revision()) + .finish_non_exhaustive() + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum ClientAuthenticationError { + #[error("client API key is invalid")] + InvalidKey, + #[error("runtime snapshot is unavailable")] + SnapshotUnavailable, + #[error("frontend authentication provider is unavailable")] + ProviderUnavailable, +} + +pub struct StartExecution { + pub client: AuthenticatedClient, + pub public_model: PublicModelId, + pub operation: Operation, + pub metadata: ExecutionRequestMetadata, +} + +/// 入口中间件开始前冻结的根请求身份与生命周期 +/// +/// 中间件短路时本值直接释放；调用 `next` 后必须原样交回 +/// [`ExecutionService::start_prepared`]，不能重新认证为另一个 Key +pub struct PreparedRootExecution { + pub(super) extension_scope: ExtensionCallScope, + pub(super) response_control: crate::engine::response_control::ResponseControl, + pub(super) client: AuthenticatedClient, + pub(super) request_id: ModelRequestId, + pub(super) started_at: SystemTime, + pub(super) deadline_at: Deadline, + pub(super) cancellation: CancellationToken, + pub(super) execution_effects: Arc, + pub(super) execution_effects_baseline: usize, +} + +impl PreparedRootExecution { + #[must_use] + pub fn request_settings(&self) -> RequestSettings { + self.client + .settings + .clone() + .unwrap_or_else(|| RequestSettings::new(self.client.snapshot.clone())) + .with_execution( + &self.client.policy, + self.deadline_at + .at() + .map(|at| duration_ms(at.duration_since(self.started_at).unwrap_or_default())), + ) + } + + /// 入口延续发起实例集合；保留已经交给调用方的取消句柄 + #[must_use] + pub fn with_extension_scope(mut self, scope: ExtensionCallScope) -> Self { + self.extension_scope = scope; + self + } + + #[must_use] + pub fn extension_scope(&self) -> ExtensionCallScope { + self.extension_scope.clone() + } + + /// 设置已经在改写时完成编译；此处只将同一份结果应用到冻结身份和生命周期 + pub fn apply_settings(&mut self, settings: &RequestSettings) -> Result<(), GatewayError> { + let deadline = settings + .execution_deadline(self.client.policy.key_id(), self.started_at) + .map_err(|_| { + GatewayError::new( + GatewayErrorKind::InvalidRequest, + "execution settings are invalid", + ) + })?; + self.client.policy = settings.apply_policy(self.client.policy.clone()); + self.client.snapshot = settings.snapshot(); + self.client.settings = Some(settings.clone()); + self.deadline_at = deadline; + Ok(()) + } + + #[must_use] + pub fn response_control(&self) -> crate::engine::response_control::ResponseControl { + self.response_control.clone() + } + + pub(super) fn new(client: AuthenticatedClient) -> Result { + let started_at = SystemTime::now(); + let deadline_at = Deadline::from_timeout(started_at, client.execution_timeout()) + .ok_or_else(|| { + GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid") + })?; + let execution_effects = Arc::new(ExecutionEffects::default()); + let execution_effects_baseline = execution_effects.epoch(); + Ok(Self { + extension_scope: ExtensionCallScope::default(), + client, + response_control: crate::engine::response_control::ResponseControl::default(), + request_id: new_request_id()?, + started_at, + deadline_at, + cancellation: CancellationToken::new(), + execution_effects, + execution_effects_baseline, + }) + } + + #[must_use] + pub const fn request_id(&self) -> &ModelRequestId { + &self.request_id + } + + #[must_use] + pub const fn client(&self) -> &AuthenticatedClient { + &self.client + } + + #[must_use] + pub const fn started_at(&self) -> SystemTime { + self.started_at + } + + #[must_use] + pub const fn deadline_at(&self) -> Deadline { + self.deadline_at + } + + #[must_use] + pub fn cancellation(&self) -> CancellationToken { + self.cancellation.clone() + } + + /// 以冻结的入口身份补齐 Runtime-only binding 事实；调用位置事实由 API/Core 提供 + #[must_use] + pub fn middleware_context( + &self, + mut target: MiddlewareTarget, + extension_scope: ExtensionCallScope, + ) -> MiddlewareContext { + target.request_id = self.request_id.clone(); + let account_group_ids = self + .client + .policy + .account_scope() + .routing_snapshot() + .groups_snapshot() + .iter() + .map(|group| group.id().clone()) + .collect::>(); + MiddlewareContext::new( + target, + MiddlewareAuthority { + client_key_id: self.client.policy.key_id().clone(), + account_group_ids: Arc::from(account_group_ids), + cancellation: self.cancellation.clone(), + deadline: self.deadline_at, + extension_scope, + execution_effects: Some(Arc::clone(&self.execution_effects)), + }, + ) + } +} + +/// 入口中间件调用 `next` 后交给 Core 的已解码执行输入 +pub struct PreparedExecutionRequest { + pub public_model: PublicModelId, + pub operation: Operation, + pub metadata: ExecutionRequestMetadata, +} + +/// 启动一个由协议 adapter 明确绑定到 Provider 自有端点的请求 +pub struct StartProviderExecution { + pub client: AuthenticatedClient, + pub provider: ProviderKind, + /// Provider 自有模型端点可固定上游模型；非模型端点保持 `None` + pub upstream_model: Option, + pub operation: Operation, + pub metadata: ExecutionRequestMetadata, +} + +pub struct StartedExecution { + pub request_id: ModelRequestId, + pub created_at: SystemTime, + pub stream: bool, + pub session: Box, +} + +pub trait ExecutionSession: Send { + fn trace(&self) -> crate::diagnostics::TraceContext { + crate::diagnostics::TraceContext::default() + } + fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>>; + fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>>; + fn response_headers(&self) -> &[ProviderResponseHeader]; + fn response_status_code(&self) -> Option { + None + } + fn discard_pending_delivery(&mut self) -> Result<(), EngineError> { + Err(EngineError::InvalidDeliveryState) + } + fn commit_downstream( + &mut self, + client_status_code: Option, + ) -> BoxFuture<'_, Result<(), EngineError>>; + fn record_client_status( + &mut self, + client_status_code: u16, + ) -> BoxFuture<'_, Result<(), EngineError>>; + /// 执行已终结且结算、准入释放均已返回，协议层才可以放弃清理责任 + /// 结算失败仍由 Store 保留费用重试，释放失败仍按租约 TTL 收敛 + fn is_finalized(&self) -> bool; + fn cancel(&self); + /// 将会话交给宿主持续驱动清理；取消请求事件的等待不会丢弃已开始的结算 + fn detach_finalize(self: Box) -> BoxFuture<'static, ()>; +} + +pub trait ExecutionService: Send + Sync { + /// 未就绪时管理与健康接口仍可运行，模型认证按既有路径报告不可用 + fn request_settings(&self) -> Option { + None + } + + fn authenticate( + &self, + plaintext: &str, + ) -> Result; + fn authenticate_request( + &self, + request: ClientAuthenticationRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let plaintext = request + .native_bearer() + .ok_or(ClientAuthenticationError::InvalidKey)?; + self.authenticate(plaintext.expose_for_auth()) + }) + } + /// 只验证入口认证信封，不记录 Key 使用事实或执行推理准入 + fn verify_request( + &self, + _request: ClientAuthenticationRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Err(ClientAuthenticationError::InvalidKey) }) + } + fn public_models(&self, client: &AuthenticatedClient) -> Vec; + fn client_model_catalog<'a>( + &'a self, + _client: &'a AuthenticatedClient, + _protocol: &'a str, + _client_version: &'a str, + ) -> BoxFuture<'a, Result, ProviderCatalogUnavailable>> { + Box::pin(async { Err(ProviderCatalogUnavailable) }) + } + fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool; + + /// 重新鉴权并冻结本次根请求身份、请求 ID、deadline 与取消域 + fn prepare_execution( + &self, + client: AuthenticatedClient, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { PreparedRootExecution::new(client) }) + } + + /// 由已经验证管理会话与插件 models 域的宿主入口选择 Key + /// 不向页面提供明文，也不跳过普通请求的扩展、准入和计量 + fn prepare_plugin_execution( + &self, + _client_key_id: &ClientApiKeyId, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + Err(GatewayError::new( + GatewayErrorKind::Unauthorized, + "plugin model entry is unavailable", + )) + }) + } + + /// 为已经通过 `verify_request` 的只读入口建立中间件生命周期，不重复写 Key 使用事实 + fn prepare_verified_execution( + &self, + client: AuthenticatedClient, + ) -> Result { + PreparedRootExecution::new(client) + } + + /// 解析与准备阶段冻结的发布代次一致的中间件计划 + fn middleware_plan(&self, _prepared: &PreparedRootExecution) -> Option { + None + } + + /// 消费一次已准备身份并进入原有 Core 路由、准入、attempt 与结算路径 + fn start_prepared( + &self, + prepared: PreparedRootExecution, + request: PreparedExecutionRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + self.start(StartExecution { + client: prepared.client, + public_model: request.public_model, + operation: request.operation, + metadata: request.metadata, + }) + .await + }) + } + + /// 消费一次已准备身份，并进入 Provider 自有端点的原有准入与结算路径 + fn start_prepared_provider_endpoint( + &self, + prepared: PreparedRootExecution, + provider: ProviderKind, + upstream_model: Option, + operation: Operation, + metadata: ExecutionRequestMetadata, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + self.start_provider_endpoint(StartProviderExecution { + client: prepared.client, + provider, + upstream_model, + operation, + metadata, + }) + .await + }) + } + + fn start( + &self, + request: StartExecution, + ) -> BoxFuture<'_, Result>; + fn start_provider_endpoint( + &self, + request: StartProviderExecution, + ) -> BoxFuture<'_, Result>; + + /// 返回 Provider 注册表中首个声明 Live sideband 能力的网关。 + /// `None` 表示当前组合不含语音 sideband；协议 adapter 据此回退到稳定 501。 + fn live_gateway(&self) -> Option> { + None + } +} + +/// 只验证 Client API Key 并返回稳定 Key ID，不产生请求使用事实 +/// +/// 管理侧登录等非数据面场景必须使用该端口，避免把认证本身误记为一次 Key 使用 +pub trait ClientKeyVerifier: Send + Sync { + fn verify_client_key( + &self, + plaintext: &str, + ) -> Result; +} + +/// 成功认证后的 API Key 使用事实接收器 +/// +/// 认证仍是同步快照读取；实现必须自行异步、去重地持久化，不得阻塞客户端请求 +pub trait ClientApiKeyUsageSink: Send + Sync { + fn record_used(&self, key_id: &ClientApiKeyId); +} + +pub(super) fn duration_ms(duration: Duration) -> u64 { + u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) +} + +pub(super) fn new_request_id() -> Result { + ModelRequestId::new(format!("req_{}", Uuid::now_v7().simple())) + .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "failed to allocate request ID")) +} diff --git a/backend/crates/gateway-core/src/engine/execution/failure.rs b/backend/crates/gateway-core/src/engine/execution/failure.rs new file mode 100644 index 000000000..37149ec17 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/failure.rs @@ -0,0 +1,100 @@ +//! 执行边界的错误分类与资源释放诊断 + +use super::contract::ClientAuthenticationError; +use crate::{ + engine::{EngineError, ModelRequestId}, + error::{GatewayError, GatewayErrorKind}, +}; + +pub(super) fn authentication_gateway_error(error: ClientAuthenticationError) -> GatewayError { + match error { + ClientAuthenticationError::InvalidKey => { + GatewayError::new(GatewayErrorKind::Unauthorized, "client API key is invalid") + } + ClientAuthenticationError::SnapshotUnavailable => GatewayError::new( + GatewayErrorKind::Internal, + "runtime snapshot is unavailable", + ), + ClientAuthenticationError::ProviderUnavailable => GatewayError::new( + GatewayErrorKind::Internal, + "frontend authentication provider is unavailable", + ), + } +} + +pub(super) fn map_routing_error(error: crate::validation::RoutingError) -> GatewayError { + match error { + crate::validation::RoutingError::ModelNotFound { + model, + mapped_model, + } => GatewayError::new( + GatewayErrorKind::ModelNotFound, + if model == mapped_model { + "the requested model was not found in the provider catalogs available to this API key; check the model name" + } else { + "the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping" + }, + ), + crate::validation::RoutingError::NoCapableProvider { .. } + | crate::validation::RoutingError::NoCapableProviderEndpoint { .. } + | crate::validation::RoutingError::EmptyAccountScope => GatewayError::new( + GatewayErrorKind::NoAvailableProvider, + "no provider can execute this request", + ), + crate::validation::RoutingError::UnsupportedProviderEndpoint { .. } => GatewayError::new( + GatewayErrorKind::Unsupported, + "the selected provider does not support this operation", + ), + _ => GatewayError::new( + GatewayErrorKind::Internal, + "runtime routing configuration is invalid", + ), + } +} + +pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError { + match error { + EngineError::Cancelled => { + GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled") + } + EngineError::Deadline => { + GatewayError::new(GatewayErrorKind::Timeout, "request deadline elapsed") + } + EngineError::Provider(provider) => GatewayError::from_provider(provider), + EngineError::EmptyRoutingPlan | EngineError::ProviderNotRegistered { .. } => { + GatewayError::new( + GatewayErrorKind::NoAvailableProvider, + "no provider is available", + ) + } + EngineError::Store(source) => { + GatewayError::new(GatewayErrorKind::Internal, "request execution failed") + .with_source(source.clone()) + } + _ => GatewayError::new(GatewayErrorKind::Internal, "request execution failed"), + } +} + +pub(super) async fn record_resource_failure( + diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, + request_id: &ModelRequestId, + operation: &'static str, + message: &'static str, + source: impl Into, +) { + let mut failure = crate::diagnostics::OperationalFailure::new( + "core", + operation, + "resource_unavailable", + message, + ); + failure.correlation_id = Some(request_id.as_str().to_owned()); + failure.details = crate::error::ErrorDetails::capture(Some(&source.into()), None, false); + if diagnostics.record_failure(failure).await.is_err() { + tracing::warn!( + request_id = request_id.as_str(), + operation, + "resource diagnostic could not be recorded" + ); + } +} diff --git a/backend/crates/gateway-core/src/engine/execution/mod.rs b/backend/crates/gateway-core/src/engine/execution/mod.rs new file mode 100644 index 000000000..198340669 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/mod.rs @@ -0,0 +1,24 @@ +//! 数据面执行入口与请求、作用域、会话的单一所有权 + +mod admission; +mod contract; +mod failure; +pub use contract::{ + AuthenticatedClient, ClientApiKeyUsageSink, ClientAuthenticationError, ClientKeyVerifier, + ClientTransport, ExecutionRequestMetadata, ExecutionService, ExecutionSession, + PreparedExecutionRequest, PreparedRootExecution, StartExecution, StartProviderExecution, + StartedExecution, +}; +mod scope; +pub use scope::BoundModelExecutionContext; +mod service; +pub use failure::gateway_error_from_engine; +pub use service::DefaultExecutionService; +mod probe_store; +mod session; + +use std::time::Duration; +const DIAGNOSTIC_TIMEOUT: Duration = Duration::from_secs(10 * 60); +const COORDINATION_TIMEOUT: Duration = Duration::from_millis(100); +const MAX_NESTED_EXECUTIONS: usize = 16; +const MAX_CONCURRENT_NESTED_EXECUTIONS: usize = 4; diff --git a/backend/crates/gateway-core/src/engine/execution/probe_store.rs b/backend/crates/gateway-core/src/engine/execution/probe_store.rs new file mode 100644 index 000000000..fe10c93d1 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/probe_store.rs @@ -0,0 +1,64 @@ +//! 诊断探测使用的非持久执行记录与观测事实 + +use crate::{ + engine::{ + AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization, + ModelRequestId, NewModelRequest, ProviderAccountId, RecoveryReport, UpstreamSendState, + }, + error::StoreError, + identity::ProviderKind, + routing::UpstreamModelId, +}; +use std::time::SystemTime; + +pub(super) struct ProbeObservation { + pub(super) provider_kind: ProviderKind, + pub(super) account_id: ProviderAccountId, + pub(super) upstream_model: UpstreamModelId, +} + +pub(super) struct TransientExecutionStore; + +#[async_trait::async_trait] +impl ExecutionStore for TransientExecutionStore { + async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> { + Ok(()) + } + + async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> { + Ok(()) + } + + async fn mark_send_state( + &self, + _: &ModelRequestId, + _: UpstreamSendState, + ) -> Result<(), StoreError> { + Ok(()) + } + + async fn mark_downstream_committed( + &self, + _: &ModelRequestId, + _: SystemTime, + _: Option, + ) -> Result<(), StoreError> { + Ok(()) + } + + async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> { + Ok(()) + } + + async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> { + Ok(()) + } + + async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> { + Ok(()) + } + + async fn recover_expired(&self, _: SystemTime) -> Result { + Ok(RecoveryReport::default()) + } +} diff --git a/backend/crates/gateway-core/src/engine/execution/scope.rs b/backend/crates/gateway-core/src/engine/execution/scope.rs new file mode 100644 index 000000000..deef09dda --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/scope.rs @@ -0,0 +1,301 @@ +//! 请求嵌套图、活动权限与作用域 guard + +use super::admission::{AdmissionLease, settle_budget}; +use super::contract::{AuthenticatedClient, ExecutionRequestMetadata}; +use super::{MAX_CONCURRENT_NESTED_EXECUTIONS, MAX_NESTED_EXECUTIONS}; +use crate::engine::budget::{ClientBudgetCharge, ClientBudgetPort}; +use crate::{ + concurrency::ConcurrencyWaitBudget, + engine::{ + ModelRequestId, ProviderAccountId, coordinator::CoordinationExtensions, + extensions::ExtensionCallScope, nested::ExecutionEffects, + }, + error::{GatewayError, GatewayErrorKind}, + identity::ProviderKind, + lifecycle::{CancellationToken, Deadline}, + operation::Operation, + policy::ClientApiKeyId, + routing::{ + FrozenAccountScope, PublicModelId, UpstreamModelId, request_settings::RequestSettings, + }, +}; +use std::{ + collections::BTreeMap, + fmt, + sync::{ + Arc, Mutex, Weak, + atomic::{AtomicUsize, Ordering}, + }, + time::SystemTime, +}; + +pub(super) enum ExecutionTarget { + Model(PublicModelId), + ProviderEndpoint { + provider: ProviderKind, + upstream_model: Option, + }, +} + +impl ExecutionTarget { + pub(super) fn public_model(&self) -> Option<&PublicModelId> { + match self { + Self::Model(model) => Some(model), + Self::ProviderEndpoint { .. } => None, + } + } + + pub(super) fn into_public_model(self) -> Option { + match self { + Self::Model(model) => Some(model), + Self::ProviderEndpoint { .. } => None, + } + } +} + +pub(super) struct PendingStartExecution { + pub(super) client: AuthenticatedClient, + pub(super) target: ExecutionTarget, + pub(super) operation: Operation, + pub(super) metadata: ExecutionRequestMetadata, +} + +pub(super) struct AuthorizedExecution { + pub(super) response_control: Option, + pub(super) account_scope: Arc, + pub(super) deadline_at: Deadline, + pub(super) cancellation: CancellationToken, + pub(super) extension_scope: ExtensionCallScope, + pub(super) required_provider: Option, + pub(super) required_account: Option, + pub(super) nested: Option, + pub(super) bound: Option, + pub(super) graph: Option>, + pub(super) execution_effects_baseline: usize, + pub(super) nested_permit: Option, +} + +pub(super) struct ExecutionStartGuard { + pub(super) admission: ExecutionAdmission, + pub(super) active_request: Option, + pub(super) concurrency_wait_budget: ConcurrencyWaitBudget, + pub(super) admission_decision_ms: Option, +} + +pub(super) struct PreparedExecutionStart { + pub(super) request_id: ModelRequestId, + pub(super) started_at: SystemTime, + pub(super) plan: crate::routing::RoutingPlan, + pub(super) extensions: CoordinationExtensions, + pub(super) authorization: AuthorizedExecution, + pub(super) guard: ExecutionStartGuard, +} + +pub(super) struct NestedExecutionFacts { + pub(super) parent_request_id: ModelRequestId, + pub(super) initiating_plugin_instance_id: String, +} + +pub(super) struct BoundExecutionFacts { + pub(super) initiating_plugin_instance_id: String, +} + +pub(super) struct NestedExecutionGraph { + pub(super) total_started: AtomicUsize, + pub(super) active: AtomicUsize, + pub(super) effects: Arc, +} + +impl NestedExecutionGraph { + pub(super) fn new(effects: Arc) -> Self { + Self { + total_started: AtomicUsize::new(0), + active: AtomicUsize::new(0), + effects, + } + } + + pub(super) fn acquire(self: &Arc) -> Result { + if self + .active + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |active| { + (active < MAX_CONCURRENT_NESTED_EXECUTIONS).then_some(active + 1) + }) + .is_err() + { + return Err(GatewayError::new( + GatewayErrorKind::ConcurrencyQueueFull, + "nested model execution concurrency is exhausted", + )); + } + if self + .total_started + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |total| { + (total < MAX_NESTED_EXECUTIONS).then_some(total + 1) + }) + .is_err() + { + self.active.fetch_sub(1, Ordering::AcqRel); + return Err(GatewayError::new( + GatewayErrorKind::PolicyDenied, + "nested model execution budget is exhausted", + )); + } + Ok(NestedExecutionPermit { + graph: Arc::clone(self), + armed: true, + }) + } +} + +pub(super) struct NestedExecutionPermit { + pub(super) graph: Arc, + pub(super) armed: bool, +} + +impl NestedExecutionPermit { + pub(super) fn release(mut self) { + if self.armed { + self.graph.active.fetch_sub(1, Ordering::AcqRel); + self.armed = false; + } + } +} + +impl Drop for NestedExecutionPermit { + fn drop(&mut self) { + if self.armed { + self.graph.active.fetch_sub(1, Ordering::AcqRel); + } + } +} + +pub(super) struct ActiveRequestAuthority { + pub(super) client: AuthenticatedClient, + pub(super) account_scope: Arc, + pub(super) deadline_at: Deadline, + pub(super) cancellation: CancellationToken, + pub(super) extension_scope: ExtensionCallScope, + pub(super) graph: Arc, +} + +pub(super) type ActiveRequestRegistry = + Arc>>>; + +pub(super) struct ActiveRequestLease { + pub(super) registry: ActiveRequestRegistry, + pub(super) request_id: ModelRequestId, + pub(super) authority: Arc, +} + +/// 管理页或 CLI 一次调用持有的显式模型执行身份 +/// +/// 结构不暴露 Key 明文；Runtime 只能把它原样交回 Core 发起模型请求 +#[derive(Clone)] +pub struct BoundModelExecutionContext { + pub(super) authority: Arc, +} + +impl BoundModelExecutionContext { + #[must_use] + pub fn request_settings(&self) -> Option<&RequestSettings> { + self.authority.client.request_settings() + } +} + +pub(super) struct BoundModelExecutionAuthority { + pub(super) client: AuthenticatedClient, + pub(super) account_scope: Arc, + pub(super) cancellation: CancellationToken, + pub(super) extension_scope: ExtensionCallScope, + pub(super) initiating_plugin_instance_id: String, +} + +impl fmt::Debug for BoundModelExecutionContext { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("BoundModelExecutionContext") + .field("key_id", &self.authority.client.policy.key_id()) + .field( + "initiating_plugin_instance_id", + &self.authority.initiating_plugin_instance_id, + ) + .finish_non_exhaustive() + } +} + +impl Drop for BoundModelExecutionAuthority { + fn drop(&mut self) { + self.cancellation.cancel(); + } +} + +impl ActiveRequestLease { + pub(super) fn release(self) { + drop(self); + } +} + +impl Drop for ActiveRequestLease { + fn drop(&mut self) { + self.authority.cancellation.cancel(); + let mut active = self + .registry + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if active + .get(&self.request_id) + .and_then(Weak::upgrade) + .is_some_and(|authority| Arc::ptr_eq(&authority, &self.authority)) + { + active.remove(&self.request_id); + } + } +} + +impl ExecutionStartGuard { + pub(super) async fn release_failed( + self, + budget: Option<&dyn ClientBudgetPort>, + request_id: ModelRequestId, + key_id: ClientApiKeyId, + diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, + ) { + if let Some(active_request) = self.active_request { + active_request.release(); + } + if let Some(budget) = budget { + settle_budget( + budget, + ClientBudgetCharge { + key_id, + request_id, + amount_usd: crate::metering::Decimal::ZERO, + completed_at: SystemTime::now(), + }, + diagnostics, + ) + .await; + } + self.admission.release().await; + } +} + +pub(super) enum ExecutionAdmission { + Client(AdmissionLease), + Nested(NestedExecutionPermit), + Bound(AdmissionLease, NestedExecutionPermit), +} + +impl ExecutionAdmission { + pub(super) async fn release(self) { + match self { + Self::Client(admission) => admission.release().await, + Self::Nested(permit) => permit.release(), + Self::Bound(admission, permit) => { + admission.release().await; + permit.release(); + } + } + } +} diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution/service.rs similarity index 64% rename from backend/crates/gateway-core/src/engine/execution.rs rename to backend/crates/gateway-core/src/engine/execution/service.rs index 29975b7f9..8cacc9b36 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution/service.rs @@ -1,739 +1,80 @@ -//! 数据面执行用例：认证、准入、路由、continuation 与会话生命周期 +//! 唯一执行服务的认证、准入、路由与调用编排 -use std::collections::{BTreeMap, BTreeSet}; -use std::fmt; -use std::net::IpAddr; -use std::sync::{ - Arc, Mutex, Weak, - atomic::{AtomicUsize, Ordering}, +use super::failure::{authentication_gateway_error, gateway_error_from_engine, map_routing_error}; + +use super::admission::{AdmissionLease, settle_budget}; +use super::contract::{ + AuthenticatedClient, ClientApiKeyUsageSink, ClientAuthenticationError, ClientKeyVerifier, + ClientTransport, ExecutionRequestMetadata, ExecutionService, PreparedExecutionRequest, + PreparedRootExecution, StartExecution, StartProviderExecution, StartedExecution, duration_ms, + new_request_id, }; -use std::time::{Duration, Instant, SystemTime}; - -use futures::{FutureExt, future::BoxFuture, pin_mut, select_biased}; -use futures_timer::Delay; -use uuid::Uuid; - -use crate::concurrency::{CapacityWait, ConcurrencyWaitBudget, ConcurrencyWaitQueue}; -use crate::engine::admission::{ - ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest, +use super::probe_store::{ProbeObservation, TransientExecutionStore}; +use super::scope::{ + ActiveRequestAuthority, ActiveRequestLease, ActiveRequestRegistry, AuthorizedExecution, + BoundExecutionFacts, BoundModelExecutionAuthority, BoundModelExecutionContext, + ExecutionStartGuard, ExecutionTarget, NestedExecutionFacts, NestedExecutionGraph, + PendingStartExecution, PreparedExecutionStart, }; -use crate::engine::authentication::{ - ClientAuthenticationRequest, FrontendAuthenticationDecision, - FrontendAuthenticationExtensionIndex, -}; -use crate::engine::budget::{ClientBudgetCharge, ClientBudgetPort}; -use crate::engine::continuation::{ - ContinuationBinding, NativeContinuationPin, NativeContinuationPort, - NativeContinuationStoreErrorKind, PreviousResponseId, -}; -use crate::engine::coordinator::{CoordinationExtensions, ResponseExecutionSession}; -use crate::engine::extensions::ExtensionCallScope; -use crate::engine::middleware::{ - FrozenMiddlewarePlan, MiddlewareAuthority, MiddlewareContext, MiddlewareExtensionIndex, - MiddlewareTarget, -}; -use crate::engine::nested::{ - AffinityLookupPort, AffinityLookupRequest, AffinityLookupResult, BoundModelExecutionBinding, - BoundModelExecutionRequest, ExecutionEffects, NestedModelExecutionPort, - NestedModelExecutionRequest, -}; -use crate::engine::observation::{ - FrozenRequestObservationContext, RequestObservationDispatch, RequestObservationScope, - RequestObserverExtensionIndex, -}; -use crate::engine::policy::{ - ModelRouteDecision, RequestPolicyContext, RequestPolicyExtensionIndex, -}; -use crate::engine::probe::{ - AccountProbe, AccountProbeError, AccountProbeErrorSource, AccountProbeRequest, - AccountProbeResult, AccountProbeUpstreamResponse, -}; -use crate::engine::provider::ProviderRegistry; -use crate::engine::{ - AttemptCoordinator, AttemptRecord, CoordinatedEvent, EngineError, ExecutionStore, - GatewayEngine, IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest, - ProbeFailure, ProviderAccountId, RecoveryReport, UpstreamSendState, -}; -use crate::error::{GatewayError, GatewayErrorKind, StoreError}; -use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader}; -use crate::identity::ProviderKind; -use crate::lifecycle::{CancellationToken, Deadline, LeaseGuard, REQUEST_LEASE_TTL}; -use crate::operation::{Operation, ProviderSessionState}; -use crate::policy::{ClientApiKeyId, ClientPolicy}; -use crate::provider_ports::{ProviderSessionAffinityPort, ProviderStoreErrorKind}; -use crate::routing::{ - FrozenAccountScope, ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId, - RoutingContext, RuntimeSnapshot, UpstreamModelId, -}; -use crate::runtime::{RuntimeSnapshotHandle, RuntimeSnapshotPublisher}; -use crate::settings::RequestSettings; - -const DIAGNOSTIC_TIMEOUT: Duration = Duration::from_secs(10 * 60); -const COORDINATION_TIMEOUT: Duration = Duration::from_millis(100); -const MAX_NESTED_EXECUTIONS: usize = 16; -const MAX_CONCURRENT_NESTED_EXECUTIONS: usize = 4; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ClientTransport { - HttpJson, - HttpSse, - WebSocket, - InternalProbe, - InternalPlugin, -} - -impl ClientTransport { - #[must_use] - pub const fn as_str(self) -> &'static str { - match self { - Self::HttpJson => "http_json", - Self::HttpSse => "http_sse", - Self::WebSocket => "websocket", - Self::InternalProbe => "internal", - Self::InternalPlugin => "internal_plugin", - } - } -} - -/// API 解码后交给 Core 的稳定请求元数据 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ExecutionRequestMetadata { - pub protocol: String, - pub endpoint: String, - pub transport: ClientTransport, - pub stream: bool, - pub client_ip: Option, - pub user_agent: Option, - pub previous_response_id: Option, -} - -#[derive(Clone)] -pub struct AuthenticatedClient { - settings: Option, - snapshot: Arc, - policy: ClientPolicy, - authentication: Option, -} - -impl AuthenticatedClient { - fn execution_timeout(&self) -> Option { - self.settings - .as_ref() - .and_then(|settings| settings.execution_timeout(self.policy.key_id())) - } - - #[must_use] - pub fn request_settings(&self) -> Option<&RequestSettings> { - self.settings.as_ref() - } - - /// 后续仍须重新认证；这里只更新下一次认证使用的冻结配置 - #[must_use] - pub fn with_request_settings(mut self, settings: RequestSettings) -> Self { - self.settings = Some(settings); - self - } - - #[must_use] - pub const fn snapshot(&self) -> &Arc { - &self.snapshot - } - - #[must_use] - pub const fn policy(&self) -> &ClientPolicy { - &self.policy - } -} - -impl fmt::Debug for AuthenticatedClient { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("AuthenticatedClient") - .field("key_id", &self.policy.key_id()) - .field("revision", &self.snapshot.revision()) - .finish_non_exhaustive() - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] -pub enum ClientAuthenticationError { - #[error("client API key is invalid")] - InvalidKey, - #[error("runtime snapshot is unavailable")] - SnapshotUnavailable, - #[error("frontend authentication provider is unavailable")] - ProviderUnavailable, -} - -pub struct StartExecution { - pub client: AuthenticatedClient, - pub public_model: PublicModelId, - pub operation: Operation, - pub metadata: ExecutionRequestMetadata, -} - -/// 入口中间件开始前冻结的根请求身份与生命周期 -/// -/// 中间件短路时本值直接释放；调用 `next` 后必须原样交回 -/// [`ExecutionService::start_prepared`]，不能重新认证为另一个 Key -pub struct PreparedRootExecution { - extension_scope: ExtensionCallScope, - response_control: super::response_control::ResponseControl, - client: AuthenticatedClient, - request_id: ModelRequestId, - started_at: SystemTime, - deadline_at: Deadline, - cancellation: CancellationToken, - execution_effects: Arc, - execution_effects_baseline: usize, -} - -impl PreparedRootExecution { - #[must_use] - pub fn request_settings(&self) -> RequestSettings { - self.client - .settings - .clone() - .unwrap_or_else(|| RequestSettings::new(self.client.snapshot.clone())) - .with_execution( - &self.client.policy, - self.deadline_at - .at() - .map(|at| duration_ms(at.duration_since(self.started_at).unwrap_or_default())), - ) - } - - /// 入口延续发起实例集合；保留已经交给调用方的取消句柄 - #[must_use] - pub fn with_extension_scope(mut self, scope: ExtensionCallScope) -> Self { - self.extension_scope = scope; - self - } - - #[must_use] - pub fn extension_scope(&self) -> ExtensionCallScope { - self.extension_scope.clone() - } - - /// 设置已经在改写时完成编译；此处只将同一份结果应用到冻结身份和生命周期 - pub fn apply_settings(&mut self, settings: &RequestSettings) -> Result<(), GatewayError> { - let deadline = settings - .execution_deadline(self.client.policy.key_id(), self.started_at) - .map_err(|_| { - GatewayError::new( - GatewayErrorKind::InvalidRequest, - "execution settings are invalid", - ) - })?; - self.client.policy = settings.apply_policy(self.client.policy.clone()); - self.client.snapshot = settings.snapshot(); - self.client.settings = Some(settings.clone()); - self.deadline_at = deadline; - Ok(()) - } - - #[must_use] - pub fn response_control(&self) -> super::response_control::ResponseControl { - self.response_control.clone() - } - - fn new(client: AuthenticatedClient) -> Result { - let started_at = SystemTime::now(); - let deadline_at = Deadline::from_timeout(started_at, client.execution_timeout()) - .ok_or_else(|| { - GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid") - })?; - let execution_effects = Arc::new(ExecutionEffects::default()); - let execution_effects_baseline = execution_effects.epoch(); - Ok(Self { - extension_scope: ExtensionCallScope::default(), - client, - response_control: super::response_control::ResponseControl::default(), - request_id: new_request_id()?, - started_at, - deadline_at, - cancellation: CancellationToken::new(), - execution_effects, - execution_effects_baseline, - }) - } - - #[must_use] - pub const fn request_id(&self) -> &ModelRequestId { - &self.request_id - } - - #[must_use] - pub const fn client(&self) -> &AuthenticatedClient { - &self.client - } - - #[must_use] - pub const fn started_at(&self) -> SystemTime { - self.started_at - } - - #[must_use] - pub const fn deadline_at(&self) -> Deadline { - self.deadline_at - } - - #[must_use] - pub fn cancellation(&self) -> CancellationToken { - self.cancellation.clone() - } - - /// 以冻结的入口身份补齐 Runtime-only binding 事实；调用位置事实由 API/Core 提供 - #[must_use] - pub fn middleware_context( - &self, - mut target: MiddlewareTarget, - extension_scope: ExtensionCallScope, - ) -> MiddlewareContext { - target.request_id = self.request_id.clone(); - let account_group_ids = self - .client - .policy - .account_scope() - .routing_snapshot() - .groups_snapshot() - .iter() - .map(|group| group.id().clone()) - .collect::>(); - MiddlewareContext::new( - target, - MiddlewareAuthority { - client_key_id: self.client.policy.key_id().clone(), - account_group_ids: Arc::from(account_group_ids), - cancellation: self.cancellation.clone(), - deadline: self.deadline_at, - extension_scope, - execution_effects: Some(Arc::clone(&self.execution_effects)), - }, - ) - } -} - -/// 入口中间件调用 `next` 后交给 Core 的已解码执行输入 -pub struct PreparedExecutionRequest { - pub public_model: PublicModelId, - pub operation: Operation, - pub metadata: ExecutionRequestMetadata, -} - -/// 启动一个由协议 adapter 明确绑定到 Provider 自有端点的请求 -pub struct StartProviderExecution { - pub client: AuthenticatedClient, - pub provider: ProviderKind, - /// Provider 自有模型端点可固定上游模型；非模型端点保持 `None` - pub upstream_model: Option, - pub operation: Operation, - pub metadata: ExecutionRequestMetadata, -} - -enum ExecutionTarget { - Model(PublicModelId), - ProviderEndpoint { - provider: ProviderKind, - upstream_model: Option, +use super::{COORDINATION_TIMEOUT, DIAGNOSTIC_TIMEOUT}; +use super::{scope::ExecutionAdmission, session::DefaultExecutionSession}; +use crate::{ + concurrency::{CapacityWait, ConcurrencyWaitBudget, ConcurrencyWaitQueue}, + engine::{ + AttemptCoordinator, EngineError, ExecutionStore, GatewayEngine, ModelRequestId, + NewModelRequest, ProbeFailure, ProviderAccountId, UpstreamSendState, + admission::{ + ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, + ClientAdmissionRequest, + }, + authentication::{ + ClientAuthenticationRequest, FrontendAuthenticationDecision, + FrontendAuthenticationExtensionIndex, + }, + budget::{ClientBudgetCharge, ClientBudgetPort}, + continuation::{ + ContinuationBinding, NativeContinuationPin, NativeContinuationPort, + NativeContinuationStoreErrorKind, + }, + coordinator::CoordinationExtensions, + extensions::{ExecutionExtensionIndex, ExtensionCallScope}, + middleware::FrozenMiddlewarePlan, + nested::{ + AffinityLookupPort, AffinityLookupRequest, AffinityLookupResult, + BoundModelExecutionBinding, BoundModelExecutionRequest, ExecutionEffects, + NestedModelExecutionPort, NestedModelExecutionRequest, + }, + observation::{ + FrozenRequestObservationContext, RequestObservationDispatch, RequestObservationScope, + RequestObserverExtensionIndex, + }, + policy::{ModelRouteDecision, RequestPolicyContext, RequestPolicyExtensionIndex}, + probe::{ + AccountProbe, AccountProbeError, AccountProbeErrorSource, AccountProbeRequest, + AccountProbeResult, AccountProbeUpstreamResponse, + }, + provider::ProviderRegistry, }, -} - -impl ExecutionTarget { - fn public_model(&self) -> Option<&PublicModelId> { - match self { - Self::Model(model) => Some(model), - Self::ProviderEndpoint { .. } => None, - } - } - - fn into_public_model(self) -> Option { - match self { - Self::Model(model) => Some(model), - Self::ProviderEndpoint { .. } => None, - } - } -} - -struct PendingStartExecution { - client: AuthenticatedClient, - target: ExecutionTarget, - operation: Operation, - metadata: ExecutionRequestMetadata, -} - -struct AuthorizedExecution { - response_control: Option, - account_scope: Arc, - deadline_at: Deadline, - cancellation: CancellationToken, - extension_scope: ExtensionCallScope, - required_provider: Option, - required_account: Option, - nested: Option, - bound: Option, - graph: Option>, - execution_effects_baseline: usize, - nested_permit: Option, -} - -struct ExecutionStartGuard { - admission: ExecutionAdmission, - active_request: Option, - concurrency_wait_budget: ConcurrencyWaitBudget, - admission_decision_ms: Option, -} - -struct PreparedExecutionStart { - request_id: ModelRequestId, - started_at: SystemTime, - plan: crate::routing::RoutingPlan, - extensions: CoordinationExtensions, - authorization: AuthorizedExecution, - guard: ExecutionStartGuard, -} - -struct NestedExecutionFacts { - parent_request_id: ModelRequestId, - initiating_plugin_instance_id: String, -} - -struct BoundExecutionFacts { - initiating_plugin_instance_id: String, -} - -struct NestedExecutionGraph { - total_started: AtomicUsize, - active: AtomicUsize, - effects: Arc, -} - -impl NestedExecutionGraph { - fn new(effects: Arc) -> Self { - Self { - total_started: AtomicUsize::new(0), - active: AtomicUsize::new(0), - effects, - } - } - - fn acquire(self: &Arc) -> Result { - if self - .active - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |active| { - (active < MAX_CONCURRENT_NESTED_EXECUTIONS).then_some(active + 1) - }) - .is_err() - { - return Err(GatewayError::new( - GatewayErrorKind::ConcurrencyQueueFull, - "nested model execution concurrency is exhausted", - )); - } - if self - .total_started - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |total| { - (total < MAX_NESTED_EXECUTIONS).then_some(total + 1) - }) - .is_err() - { - self.active.fetch_sub(1, Ordering::AcqRel); - return Err(GatewayError::new( - GatewayErrorKind::PolicyDenied, - "nested model execution budget is exhausted", - )); - } - Ok(NestedExecutionPermit { - graph: Arc::clone(self), - armed: true, - }) - } -} - -struct NestedExecutionPermit { - graph: Arc, - armed: bool, -} - -impl NestedExecutionPermit { - fn release(mut self) { - if self.armed { - self.graph.active.fetch_sub(1, Ordering::AcqRel); - self.armed = false; - } - } -} - -impl Drop for NestedExecutionPermit { - fn drop(&mut self) { - if self.armed { - self.graph.active.fetch_sub(1, Ordering::AcqRel); - } - } -} - -struct ActiveRequestAuthority { - client: AuthenticatedClient, - account_scope: Arc, - deadline_at: Deadline, - cancellation: CancellationToken, - extension_scope: ExtensionCallScope, - graph: Arc, -} - -type ActiveRequestRegistry = Arc>>>; - -struct ActiveRequestLease { - registry: ActiveRequestRegistry, - request_id: ModelRequestId, - authority: Arc, -} - -/// 管理页或 CLI 一次调用持有的显式模型执行身份 -/// -/// 结构不暴露 Key 明文；Runtime 只能把它原样交回 Core 发起模型请求 -#[derive(Clone)] -pub struct BoundModelExecutionContext { - authority: Arc, -} - -impl BoundModelExecutionContext { - #[must_use] - pub fn request_settings(&self) -> Option<&RequestSettings> { - self.authority.client.request_settings() - } -} - -struct BoundModelExecutionAuthority { - client: AuthenticatedClient, - account_scope: Arc, - cancellation: CancellationToken, - extension_scope: ExtensionCallScope, - initiating_plugin_instance_id: String, -} - -impl fmt::Debug for BoundModelExecutionContext { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("BoundModelExecutionContext") - .field("key_id", &self.authority.client.policy.key_id()) - .field( - "initiating_plugin_instance_id", - &self.authority.initiating_plugin_instance_id, - ) - .finish_non_exhaustive() - } -} - -impl Drop for BoundModelExecutionAuthority { - fn drop(&mut self) { - self.cancellation.cancel(); - } -} - -impl ActiveRequestLease { - fn release(self) { - drop(self); - } -} - -impl Drop for ActiveRequestLease { - fn drop(&mut self) { - self.authority.cancellation.cancel(); - let mut active = self - .registry - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if active - .get(&self.request_id) - .and_then(Weak::upgrade) - .is_some_and(|authority| Arc::ptr_eq(&authority, &self.authority)) - { - active.remove(&self.request_id); - } - } -} - -pub struct StartedExecution { - pub request_id: ModelRequestId, - pub created_at: SystemTime, - pub stream: bool, - pub session: Box, -} - -pub trait ExecutionSession: Send { - fn trace(&self) -> crate::diagnostics::TraceContext { - crate::diagnostics::TraceContext::default() - } - fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>>; - fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>>; - fn response_headers(&self) -> &[ProviderResponseHeader]; - fn response_status_code(&self) -> Option { - None - } - fn discard_pending_delivery(&mut self) -> Result<(), EngineError> { - Err(EngineError::InvalidDeliveryState) - } - fn commit_downstream( - &mut self, - client_status_code: Option, - ) -> BoxFuture<'_, Result<(), EngineError>>; - fn record_client_status( - &mut self, - client_status_code: u16, - ) -> BoxFuture<'_, Result<(), EngineError>>; - /// 执行已终结且结算、准入释放均已返回，协议层才可以放弃清理责任 - /// 结算失败仍由 Store 保留费用重试，释放失败仍按租约 TTL 收敛 - fn is_finalized(&self) -> bool; - fn cancel(&self); - /// 将会话交给宿主持续驱动清理；取消请求事件的等待不会丢弃已开始的结算 - fn detach_finalize(self: Box) -> BoxFuture<'static, ()>; -} - -pub trait ExecutionService: Send + Sync { - /// 未就绪时管理与健康接口仍可运行，模型认证按既有路径报告不可用 - fn request_settings(&self) -> Option { - None - } - - fn authenticate( - &self, - plaintext: &str, - ) -> Result; - fn authenticate_request( - &self, - request: ClientAuthenticationRequest, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { - let plaintext = request - .native_bearer() - .ok_or(ClientAuthenticationError::InvalidKey)?; - self.authenticate(plaintext.expose_for_auth()) - }) - } - /// 只验证入口认证信封，不记录 Key 使用事实或执行推理准入 - fn verify_request( - &self, - _request: ClientAuthenticationRequest, - ) -> BoxFuture<'_, Result> { - Box::pin(async { Err(ClientAuthenticationError::InvalidKey) }) - } - fn public_models(&self, client: &AuthenticatedClient) -> Vec; - fn client_model_catalog<'a>( - &'a self, - _client: &'a AuthenticatedClient, - _protocol: &'a str, - _client_version: &'a str, - ) -> BoxFuture<'a, Result, ProviderCatalogUnavailable>> { - Box::pin(async { Err(ProviderCatalogUnavailable) }) - } - fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool; - - /// 重新鉴权并冻结本次根请求身份、请求 ID、deadline 与取消域 - fn prepare_execution( - &self, - client: AuthenticatedClient, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { PreparedRootExecution::new(client) }) - } - - /// 由已经验证管理会话与插件 models 域的宿主入口选择 Key - /// 不向页面提供明文，也不跳过普通请求的扩展、准入和计量 - fn prepare_plugin_execution( - &self, - _client_key_id: &ClientApiKeyId, - ) -> BoxFuture<'_, Result> { - Box::pin(async { - Err(GatewayError::new( - GatewayErrorKind::Unauthorized, - "plugin model entry is unavailable", - )) - }) - } - - /// 为已经通过 `verify_request` 的只读入口建立中间件生命周期，不重复写 Key 使用事实 - fn prepare_verified_execution( - &self, - client: AuthenticatedClient, - ) -> Result { - PreparedRootExecution::new(client) - } - - /// 解析与准备阶段冻结的发布代次一致的中间件计划 - fn middleware_plan(&self, _prepared: &PreparedRootExecution) -> Option { - None - } - - /// 消费一次已准备身份并进入原有 Core 路由、准入、attempt 与结算路径 - fn start_prepared( - &self, - prepared: PreparedRootExecution, - request: PreparedExecutionRequest, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { - self.start(StartExecution { - client: prepared.client, - public_model: request.public_model, - operation: request.operation, - metadata: request.metadata, - }) - .await - }) - } - - /// 消费一次已准备身份，并进入 Provider 自有端点的原有准入与结算路径 - fn start_prepared_provider_endpoint( - &self, - prepared: PreparedRootExecution, - provider: ProviderKind, - upstream_model: Option, - operation: Operation, - metadata: ExecutionRequestMetadata, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { - self.start_provider_endpoint(StartProviderExecution { - client: prepared.client, - provider, - upstream_model, - operation, - metadata, - }) - .await - }) - } - - fn start( - &self, - request: StartExecution, - ) -> BoxFuture<'_, Result>; - fn start_provider_endpoint( - &self, - request: StartProviderExecution, - ) -> BoxFuture<'_, Result>; - - /// 返回 Provider 注册表中首个声明 Live sideband 能力的网关。 - /// `None` 表示当前组合不含语音 sideband；协议 adapter 据此回退到稳定 501。 - fn live_gateway(&self) -> Option> { - None - } -} - -/// 只验证 Client API Key 并返回稳定 Key ID，不产生请求使用事实 -/// -/// 管理侧登录等非数据面场景必须使用该端口，避免把认证本身误记为一次 Key 使用 -pub trait ClientKeyVerifier: Send + Sync { - fn verify_client_key( - &self, - plaintext: &str, - ) -> Result; -} - -/// 成功认证后的 API Key 使用事实接收器 -/// -/// 认证仍是同步快照读取；实现必须自行异步、去重地持久化，不得阻塞客户端请求 -pub trait ClientApiKeyUsageSink: Send + Sync { - fn record_used(&self, key_id: &ClientApiKeyId); -} + error::{GatewayError, GatewayErrorKind}, + event::GatewayEvent, + identity::ProviderKind, + lifecycle::{CancellationToken, Deadline, REQUEST_LEASE_TTL}, + operation::Operation, + policy::{ClientApiKeyId, ClientPolicy}, + provider_ports::{ProviderSessionAffinityPort, ProviderStoreErrorKind}, + routing::{ + FrozenAccountScope, ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId, + RoutingContext, RuntimeSnapshot, UpstreamModelId, request_settings::RequestSettings, + }, + runtime::{RuntimeSnapshotHandle, RuntimeSnapshotPublisher}, +}; +use futures::{FutureExt as _, future::BoxFuture, pin_mut, select_biased}; +use futures_timer::Delay; +use std::{ + collections::BTreeSet, + sync::{Arc, Weak}, + time::{Instant, SystemTime}, +}; pub struct DefaultExecutionService { snapshots: RuntimeSnapshotHandle, @@ -748,7 +89,7 @@ pub struct DefaultExecutionService { budget: Option>, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, session_affinity: Option>, snapshot_refresh: Option>, @@ -778,7 +119,7 @@ impl DefaultExecutionService { budget: None, request_observers: None, request_policies: None, - middlewares: None, + execution_extensions: None, frontend_authentication: None, session_affinity: None, snapshot_refresh: None, @@ -805,8 +146,11 @@ impl DefaultExecutionService { } #[must_use] - pub fn with_middlewares(mut self, middlewares: MiddlewareExtensionIndex) -> Self { - self.middlewares = Some(middlewares); + pub fn with_execution_extensions( + mut self, + execution_extensions: ExecutionExtensionIndex, + ) -> Self { + self.execution_extensions = Some(execution_extensions); self } @@ -1069,11 +413,14 @@ impl DefaultExecutionService { .client .snapshot .extensions() - .and_then(|set| set.upstream_adapters()); - let middleware = self.middlewares.as_ref().and_then(|middlewares| { - let generation = request.client.snapshot.extensions()?.clone(); - middlewares.resolve(&generation) - }); + .and_then(|set| self.execution_extensions.as_ref()?.upstream_adapters(set)); + let middleware = self + .execution_extensions + .as_ref() + .and_then(|execution_extensions| { + let generation = request.client.snapshot.extensions()?.clone(); + execution_extensions.middleware(&generation) + }); let request_policy = self.request_policies.as_ref().and_then(|policies| { let generation = request.client.snapshot.extensions()?.clone(); let plan = policies.resolve(&generation)?; @@ -1303,7 +650,7 @@ impl DefaultExecutionService { failure_kind = error.kind().as_str(), "请求在路由或准入阶段被拒绝" ); - let rejection = super::EntryRejection { + let rejection = crate::engine::EntryRejection { request_id: request_id.clone(), client_key_id: budget_key_id.clone(), error: error.clone(), @@ -2140,58 +1487,6 @@ impl DefaultExecutionService { } } -struct ProbeObservation { - provider_kind: ProviderKind, - account_id: ProviderAccountId, - upstream_model: UpstreamModelId, -} - -struct TransientExecutionStore; - -#[async_trait::async_trait] -impl ExecutionStore for TransientExecutionStore { - async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> { - Ok(()) - } - - async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> { - Ok(()) - } - - async fn mark_send_state( - &self, - _: &ModelRequestId, - _: UpstreamSendState, - ) -> Result<(), StoreError> { - Ok(()) - } - - async fn mark_downstream_committed( - &self, - _: &ModelRequestId, - _: SystemTime, - _: Option, - ) -> Result<(), StoreError> { - Ok(()) - } - - async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> { - Ok(()) - } - - async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> { - Ok(()) - } - - async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> { - Ok(()) - } - - async fn recover_expired(&self, _: SystemTime) -> Result { - Ok(RecoveryReport::default()) - } -} - impl ExecutionService for DefaultExecutionService { fn request_settings(&self) -> Option { self.snapshots.acquire().ok().map(RequestSettings::new) @@ -2328,7 +1623,7 @@ impl ExecutionService for DefaultExecutionService { fn middleware_plan(&self, prepared: &PreparedRootExecution) -> Option { let generation = prepared.client.snapshot.extensions()?.clone(); - self.middlewares.as_ref()?.resolve(&generation) + self.execution_extensions.as_ref()?.middleware(&generation) } fn prepare_plugin_execution( @@ -2484,305 +1779,6 @@ impl AccountProbe for DefaultExecutionService { } } -struct AdmissionLease { - renewal: Option>, - armed: bool, - port: Arc, - diagnostics: Arc, - client_api_key_id: ClientApiKeyId, - model_request_id: ModelRequestId, -} - -async fn settle_budget( - port: &dyn ClientBudgetPort, - charge: ClientBudgetCharge, - diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, -) { - let request_id = charge.request_id.clone(); - if let Err(error) = port.settle(charge).await { - record_resource_failure( - diagnostics, - &request_id, - "settle_client_budget", - "Client budget settlement failed; storage will retry on the next request", - error, - ) - .await; - } -} - -impl AdmissionLease { - async fn release(mut self) { - self.renewal.take(); - if let Err(error) = self - .port - .release(&self.client_api_key_id, &self.model_request_id) - .await - { - record_resource_failure( - self.diagnostics.as_ref(), - &self.model_request_id, - "release_client_admission", - "Client admission release failed; lease TTL remains active", - error, - ) - .await; - } - self.armed = false; - } -} - -impl Drop for AdmissionLease { - fn drop(&mut self) { - self.renewal.take(); - if self.armed { - self.port - .abandon(&self.client_api_key_id, &self.model_request_id); - } - } -} - -enum ExecutionAdmission { - Client(AdmissionLease), - Nested(NestedExecutionPermit), - Bound(AdmissionLease, NestedExecutionPermit), -} - -impl ExecutionAdmission { - async fn release(self) { - match self { - Self::Client(admission) => admission.release().await, - Self::Nested(permit) => permit.release(), - Self::Bound(admission, permit) => { - admission.release().await; - permit.release(); - } - } - } -} - -impl ExecutionStartGuard { - async fn release_failed( - self, - budget: Option<&dyn ClientBudgetPort>, - request_id: ModelRequestId, - key_id: ClientApiKeyId, - diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, - ) { - if let Some(active_request) = self.active_request { - active_request.release(); - } - if let Some(budget) = budget { - settle_budget( - budget, - ClientBudgetCharge { - key_id, - request_id, - amount_usd: crate::metering::Decimal::ZERO, - completed_at: SystemTime::now(), - }, - diagnostics, - ) - .await; - } - self.admission.release().await; - } -} - -struct DefaultExecutionSession { - core: ResponseExecutionSession, - admission: Option, - active_request: Option, - cleanup: Option>, - continuation: Arc, - continuation_recorded: bool, - budget: Option>, - diagnostics: Arc, -} - -impl DefaultExecutionSession { - fn new( - core: ResponseExecutionSession, - admission: ExecutionAdmission, - active_request: Option, - continuation: Arc, - budget: Option>, - diagnostics: Arc, - ) -> Self { - Self { - core, - admission: Some(admission), - active_request, - cleanup: None, - continuation, - continuation_recorded: false, - budget, - diagnostics, - } - } - - async fn settle_if_finalized(&mut self) { - if self.core.is_finalized() - && let Some(admission) = self.admission.take() - { - if let Some(active_request) = self.active_request.take() { - active_request.release(); - } - let budget = self.budget.take(); - let charge = self.core.budget_charge(); - let diagnostics = self.diagnostics.clone(); - // 在首次 await 前把完整清理责任留在会话内 - // 事件等待被取消后，后续 poll - // 或 detach 继续同一个 future，既不丢失费用，也不重启已完成的结算 - self.cleanup = Some(Box::pin(async move { - if let Some(budget) = budget { - settle_budget(budget.as_ref(), charge, diagnostics.as_ref()).await; - } - admission.release().await; - })); - } - if let Some(cleanup) = self.cleanup.as_mut() { - cleanup.await; - self.cleanup = None; - } - } - - async fn record_continuation(&mut self, state: Option<&ProviderSessionState>) { - if self.continuation_recorded { - return; - } - let Some(state) = state else { - return; - }; - let Some(pin) = self.core.native_continuation_pin(state) else { - return; - }; - self.continuation_recorded = true; - record_native_continuation(self.continuation.as_ref(), pin).await; - } - - async fn finalize_detached(&mut self) { - if let Err(error) = self.core.cancel_and_finalize().await { - tracing::warn!(request_id = self.core.request_id().as_str(), operation = "finalize_detached_execution", %error, "Detached execution 终态收敛失败"); - } - self.settle_if_finalized().await; - } -} - -impl Drop for DefaultExecutionSession { - fn drop(&mut self) { - self.core.cancel(); - drop(self.active_request.take()); - } -} - -impl ExecutionSession for DefaultExecutionSession { - fn trace(&self) -> crate::diagnostics::TraceContext { - self.core.trace() - } - fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>> { - Box::pin(async move { - let result = self.core.next_event().await; - if let Ok(Some(event)) = result.as_ref() { - self.record_continuation(event.session_update()).await; - } - self.settle_if_finalized().await; - result - }) - } - - fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>> { - Box::pin(async move { - let result = self.core.collect_uncommitted().await; - if let Ok(events) = result.as_ref() { - let state = events.iter().find_map(ProviderEvent::session_update); - self.record_continuation(state).await; - } - self.settle_if_finalized().await; - result - }) - } - - fn response_headers(&self) -> &[ProviderResponseHeader] { - self.core.response_headers() - } - - fn response_status_code(&self) -> Option { - self.core.response_status_code() - } - - fn discard_pending_delivery(&mut self) -> Result<(), EngineError> { - self.core.discard_pending_delivery() - } - - fn commit_downstream( - &mut self, - client_status_code: Option, - ) -> BoxFuture<'_, Result<(), EngineError>> { - Box::pin(async move { - let result = self.core.commit_downstream(client_status_code).await; - self.settle_if_finalized().await; - result - }) - } - - fn record_client_status( - &mut self, - client_status_code: u16, - ) -> BoxFuture<'_, Result<(), EngineError>> { - Box::pin(async move { - let result = self.core.record_client_status(client_status_code).await; - self.settle_if_finalized().await; - result - }) - } - - fn is_finalized(&self) -> bool { - self.core.is_finalized() - && self.admission.is_none() - && self.active_request.is_none() - && self.cleanup.is_none() - } - - fn cancel(&self) { - self.core.cancel(); - } - - fn detach_finalize(mut self: Box) -> BoxFuture<'static, ()> { - Box::pin(async move { self.finalize_detached().await }) - } -} - -async fn record_native_continuation( - continuation: &dyn NativeContinuationPort, - pin: NativeContinuationPin, -) { - let provider = pin.provider().as_str().to_owned(); - let account = pin.account().as_str().to_owned(); - let record = continuation.record(pin).fuse(); - let timeout = Delay::new(COORDINATION_TIMEOUT).fuse(); - pin_mut!(record, timeout); - select_biased! { - result = record => { - if let Err(error) = result { - tracing::warn!( - provider = %provider, - account = %account, - %error, - "Continuation affinity 写入失败，后续请求将退化为外部续接" - ); - } - }, - _ = timeout => { - tracing::warn!( - provider = %provider, - account = %account, - "Continuation affinity 后台写入超时，已丢弃本次亲和记录" - ); - }, - } -} - fn constant_time_equal(left: &str, right: &str) -> bool { if left.len() != right.len() { return false; @@ -2877,61 +1873,6 @@ fn native_policy( .ok_or(ClientAuthenticationError::InvalidKey) } -fn duration_ms(duration: Duration) -> u64 { - u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) -} - -fn new_request_id() -> Result { - ModelRequestId::new(format!("req_{}", Uuid::now_v7().simple())) - .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "failed to allocate request ID")) -} - -fn authentication_gateway_error(error: ClientAuthenticationError) -> GatewayError { - match error { - ClientAuthenticationError::InvalidKey => { - GatewayError::new(GatewayErrorKind::Unauthorized, "client API key is invalid") - } - ClientAuthenticationError::SnapshotUnavailable => GatewayError::new( - GatewayErrorKind::Internal, - "runtime snapshot is unavailable", - ), - ClientAuthenticationError::ProviderUnavailable => GatewayError::new( - GatewayErrorKind::Internal, - "frontend authentication provider is unavailable", - ), - } -} - -fn map_routing_error(error: crate::validation::RoutingError) -> GatewayError { - match error { - crate::validation::RoutingError::ModelNotFound { - model, - mapped_model, - } => GatewayError::new( - GatewayErrorKind::ModelNotFound, - if model == mapped_model { - "the requested model was not found in the provider catalogs available to this API key; check the model name" - } else { - "the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping" - }, - ), - crate::validation::RoutingError::NoCapableProvider { .. } - | crate::validation::RoutingError::NoCapableProviderEndpoint { .. } - | crate::validation::RoutingError::EmptyAccountScope => GatewayError::new( - GatewayErrorKind::NoAvailableProvider, - "no provider can execute this request", - ), - crate::validation::RoutingError::UnsupportedProviderEndpoint { .. } => GatewayError::new( - GatewayErrorKind::Unsupported, - "the selected provider does not support this operation", - ), - _ => GatewayError::new( - GatewayErrorKind::Internal, - "runtime routing configuration is invalid", - ), - } -} - fn attach_continuation_session_state(operation: &mut Operation, pin: &NativeContinuationPin) { let Some(state) = pin.session_state() else { return; @@ -2943,50 +1884,3 @@ fn attach_continuation_session_state(operation: &mut Operation, pin: &NativeCont } operation.set_provider_session_state(state.clone()); } - -pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError { - match error { - EngineError::Cancelled => { - GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled") - } - EngineError::Deadline => { - GatewayError::new(GatewayErrorKind::Timeout, "request deadline elapsed") - } - EngineError::Provider(provider) => GatewayError::from_provider(provider), - EngineError::EmptyRoutingPlan | EngineError::ProviderNotRegistered { .. } => { - GatewayError::new( - GatewayErrorKind::NoAvailableProvider, - "no provider is available", - ) - } - EngineError::Store(source) => { - GatewayError::new(GatewayErrorKind::Internal, "request execution failed") - .with_source(source.clone()) - } - _ => GatewayError::new(GatewayErrorKind::Internal, "request execution failed"), - } -} - -async fn record_resource_failure( - diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, - request_id: &ModelRequestId, - operation: &'static str, - message: &'static str, - source: impl Into, -) { - let mut failure = crate::diagnostics::OperationalFailure::new( - "core", - operation, - "resource_unavailable", - message, - ); - failure.correlation_id = Some(request_id.as_str().to_owned()); - failure.details = crate::error::ErrorDetails::capture(Some(&source.into()), None, false); - if diagnostics.record_failure(failure).await.is_err() { - tracing::warn!( - request_id = request_id.as_str(), - operation, - "resource diagnostic could not be recorded" - ); - } -} diff --git a/backend/crates/gateway-core/src/engine/execution/session.rs b/backend/crates/gateway-core/src/engine/execution/session.rs new file mode 100644 index 000000000..cd5570413 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/session.rs @@ -0,0 +1,213 @@ +//! 执行会话的准入释放、费用结算与终态收敛 + +use super::COORDINATION_TIMEOUT; +use super::admission::settle_budget; +use super::contract::ExecutionSession; +use super::scope::{ActiveRequestLease, ExecutionAdmission}; +use crate::{ + engine::{ + CoordinatedEvent, EngineError, ExecutionStore, + budget::ClientBudgetPort, + continuation::{NativeContinuationPin, NativeContinuationPort}, + coordinator::ResponseExecutionSession, + }, + event::{ProviderEvent, ProviderResponseHeader}, + operation::ProviderSessionState, +}; +use futures::{FutureExt as _, future::BoxFuture, pin_mut, select_biased}; +use futures_timer::Delay; +use std::sync::Arc; + +pub(super) struct DefaultExecutionSession { + pub(super) core: ResponseExecutionSession, + pub(super) admission: Option, + pub(super) active_request: Option, + pub(super) cleanup: Option>, + pub(super) continuation: Arc, + pub(super) continuation_recorded: bool, + pub(super) budget: Option>, + pub(super) diagnostics: Arc, +} + +impl DefaultExecutionSession { + pub(super) fn new( + core: ResponseExecutionSession, + admission: ExecutionAdmission, + active_request: Option, + continuation: Arc, + budget: Option>, + diagnostics: Arc, + ) -> Self { + Self { + core, + admission: Some(admission), + active_request, + cleanup: None, + continuation, + continuation_recorded: false, + budget, + diagnostics, + } + } + + pub(super) async fn settle_if_finalized(&mut self) { + if self.core.is_finalized() + && let Some(admission) = self.admission.take() + { + if let Some(active_request) = self.active_request.take() { + active_request.release(); + } + let budget = self.budget.take(); + let charge = self.core.budget_charge(); + let diagnostics = self.diagnostics.clone(); + // 在首次 await 前把完整清理责任留在会话内 + // 事件等待被取消后，后续 poll + // 或 detach 继续同一个 future，既不丢失费用，也不重启已完成的结算 + self.cleanup = Some(Box::pin(async move { + if let Some(budget) = budget { + settle_budget(budget.as_ref(), charge, diagnostics.as_ref()).await; + } + admission.release().await; + })); + } + if let Some(cleanup) = self.cleanup.as_mut() { + cleanup.await; + self.cleanup = None; + } + } + + pub(super) async fn record_continuation(&mut self, state: Option<&ProviderSessionState>) { + if self.continuation_recorded { + return; + } + let Some(state) = state else { + return; + }; + let Some(pin) = self.core.native_continuation_pin(state) else { + return; + }; + self.continuation_recorded = true; + record_native_continuation(self.continuation.as_ref(), pin).await; + } + + pub(super) async fn finalize_detached(&mut self) { + if let Err(error) = self.core.cancel_and_finalize().await { + tracing::warn!(request_id = self.core.request_id().as_str(), operation = "finalize_detached_execution", %error, "Detached execution 终态收敛失败"); + } + self.settle_if_finalized().await; + } +} + +impl Drop for DefaultExecutionSession { + fn drop(&mut self) { + self.core.cancel(); + drop(self.active_request.take()); + } +} + +impl ExecutionSession for DefaultExecutionSession { + fn trace(&self) -> crate::diagnostics::TraceContext { + self.core.trace() + } + fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>> { + Box::pin(async move { + let result = self.core.next_event().await; + if let Ok(Some(event)) = result.as_ref() { + self.record_continuation(event.session_update()).await; + } + self.settle_if_finalized().await; + result + }) + } + + fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>> { + Box::pin(async move { + let result = self.core.collect_uncommitted().await; + if let Ok(events) = result.as_ref() { + let state = events.iter().find_map(ProviderEvent::session_update); + self.record_continuation(state).await; + } + self.settle_if_finalized().await; + result + }) + } + + fn response_headers(&self) -> &[ProviderResponseHeader] { + self.core.response_headers() + } + + fn response_status_code(&self) -> Option { + self.core.response_status_code() + } + + fn discard_pending_delivery(&mut self) -> Result<(), EngineError> { + self.core.discard_pending_delivery() + } + + fn commit_downstream( + &mut self, + client_status_code: Option, + ) -> BoxFuture<'_, Result<(), EngineError>> { + Box::pin(async move { + let result = self.core.commit_downstream(client_status_code).await; + self.settle_if_finalized().await; + result + }) + } + + fn record_client_status( + &mut self, + client_status_code: u16, + ) -> BoxFuture<'_, Result<(), EngineError>> { + Box::pin(async move { + let result = self.core.record_client_status(client_status_code).await; + self.settle_if_finalized().await; + result + }) + } + + fn is_finalized(&self) -> bool { + self.core.is_finalized() + && self.admission.is_none() + && self.active_request.is_none() + && self.cleanup.is_none() + } + + fn cancel(&self) { + self.core.cancel(); + } + + fn detach_finalize(mut self: Box) -> BoxFuture<'static, ()> { + Box::pin(async move { self.finalize_detached().await }) + } +} + +pub(super) async fn record_native_continuation( + continuation: &dyn NativeContinuationPort, + pin: NativeContinuationPin, +) { + let provider = pin.provider().as_str().to_owned(); + let account = pin.account().as_str().to_owned(); + let record = continuation.record(pin).fuse(); + let timeout = Delay::new(COORDINATION_TIMEOUT).fuse(); + pin_mut!(record, timeout); + select_biased! { + result = record => { + if let Err(error) = result { + tracing::warn!( + provider = %provider, + account = %account, + %error, + "Continuation affinity 写入失败，后续请求将退化为外部续接" + ); + } + }, + _ = timeout => { + tracing::warn!( + provider = %provider, + account = %account, + "Continuation affinity 后台写入超时，已丢弃本次亲和记录" + ); + }, + } +} diff --git a/backend/crates/gateway-core/src/engine/extensions.rs b/backend/crates/gateway-core/src/engine/extensions.rs index 1d058c3d3..e0726de4a 100644 --- a/backend/crates/gateway-core/src/engine/extensions.rs +++ b/backend/crates/gateway-core/src/engine/extensions.rs @@ -1,4 +1,4 @@ -//! 插件嵌套调用的递归作用域 +//! 执行扩展计划的弱引用索引、冻结解析与插件递归作用域 use std::{collections::BTreeSet, sync::Arc}; @@ -43,3 +43,82 @@ impl ExtensionCallScope { }) } } + +/// 同一代次的具体执行计划；不拥有发布状态，运行资源由对应发布集合保活 +#[derive(Debug, Default)] +pub struct ExecutionExtensionPlans { + pub middleware: Option>, + pub upstream_adapters: Option>, +} + +/// 同一发布代次不能被静默替换 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("execution extension generation is already registered")] +pub struct ExtensionRegistrationError; + +/// 仅以弱引用按冻结身份解析执行计划，不延长旧代次寿命或决定何时发布 +#[derive(Clone, Default)] +pub struct ExecutionExtensionIndex { + sets: Arc< + std::sync::RwLock< + std::collections::BTreeMap< + crate::routing::extensions::ExtensionSetId, + std::sync::Weak, + >, + >, + >, +} + +impl ExecutionExtensionIndex { + pub fn register( + &self, + id: crate::routing::extensions::ExtensionSetId, + plans: Arc, + ) -> Result, ExtensionRegistrationError> { + let mut sets = self + .sets + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + sets.retain(|_, plans| plans.strong_count() > 0); + if sets.contains_key(&id) { + return Err(ExtensionRegistrationError); + } + sets.insert(id, Arc::downgrade(&plans)); + Ok(plans) + } + + fn resolve( + &self, + generation: &crate::routing::extensions::ExtensionSetReference, + ) -> Option> { + self.sets + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(generation.id()) + .and_then(std::sync::Weak::upgrade) + } + + #[must_use] + pub fn middleware( + &self, + generation: &crate::routing::extensions::ExtensionSetReference, + ) -> Option { + let plan = self.resolve(generation)?.middleware.clone()?; + Some(super::middleware::FrozenMiddlewarePlan::new( + plan, + generation.clone(), + )) + } + + #[must_use] + pub fn upstream_adapters( + &self, + generation: &crate::routing::extensions::ExtensionSetReference, + ) -> Option { + let plan = self.resolve(generation)?.upstream_adapters.clone()?; + Some(super::upstream_adapter::FrozenUpstreamAdapterPlan::new( + plan, + generation.clone(), + )) + } +} diff --git a/backend/crates/gateway-core/src/engine/middleware.rs b/backend/crates/gateway-core/src/engine/middleware/contract.rs similarity index 80% rename from backend/crates/gateway-core/src/engine/middleware.rs rename to backend/crates/gateway-core/src/engine/middleware/contract.rs index 5cf19e0f8..386c338af 100644 --- a/backend/crates/gateway-core/src/engine/middleware.rs +++ b/backend/crates/gateway-core/src/engine/middleware/contract.rs @@ -5,10 +5,11 @@ //! Runtime/SDK 负责；身份、路由、账号租约、发送事实、重试、计量与结算仍由 Core //! 持有，不能通过本端口交给插件解释 -use std::collections::{BTreeMap, BTreeSet}; +use super::FrozenMiddlewarePlan; +use std::collections::BTreeSet; use std::fmt; use std::num::NonZeroU32; -use std::sync::{Arc, RwLock, Weak}; +use std::sync::Arc; use bytes::Bytes; use futures::future::BoxFuture; @@ -27,8 +28,7 @@ use crate::operation::{ CapabilityRequirements, Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload, }; use crate::policy::ClientApiKeyId; -use crate::runtime::extensions::{ExtensionSetId, ExtensionSetReference}; -use crate::settings::RequestSettings; +use crate::routing::request_settings::RequestSettings; /// 中间件在一次逻辑请求中的挂载位置 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] @@ -543,7 +543,7 @@ impl fmt::Debug for MiddlewareResponseEnvelope { /// Core 冻结的调用事实；完整业务身份可由插件读取 #[derive(Clone)] pub struct MiddlewareContext { - plan: Option, + pub(super) plan: Option, request_id: ModelRequestId, mount: MiddlewareMount, attempt_index: Option, @@ -799,175 +799,3 @@ impl From for MiddlewareError { /// 绑定当前调用且只能消费一次的宿主续体 pub type MiddlewareNext = crate::middleware::Next; - -/// 一个发布代次冻结的统一数据面中间件计划 -pub trait MiddlewarePlan: Send + Sync + fmt::Debug { - fn has_service(&self) -> bool { - false - } - fn handle_service( - &self, - _context: crate::middleware::service::Context, - input: crate::middleware::service::Value, - next: crate::middleware::service::Next, - ) -> BoxFuture< - 'static, - Result, - > { - next.run(input) - } - - fn has_websocket(&self) -> bool { - false - } - fn handle_websocket( - &self, - _context: crate::middleware::websocket::Context, - message: crate::middleware::websocket::Message, - next: crate::middleware::websocket::Next, - ) -> BoxFuture<'static, Result, MiddlewareError>> - { - next.run(message) - } - - /// HTTP 入口在认证、路由与正文解析之前组合，无匹配时保留直接路径 - fn has_http(&self) -> bool { - false - } - - fn handle_http( - &self, - _context: crate::middleware::http::Context, - request: crate::middleware::http::Request, - next: crate::middleware::http::Next, - ) -> BoxFuture<'static, Result> { - next.run(request) - } - - fn handle( - &self, - context: MiddlewareContext, - request: MiddlewareRequest, - next: MiddlewareNext, - ) -> BoxFuture<'static, Result>; -} - -/// 调用期间同时保活发布代次与中间件计划 -#[derive(Clone)] -pub struct FrozenMiddlewarePlan { - plan: Arc, - _generation: ExtensionSetReference, -} - -impl FrozenMiddlewarePlan { - #[must_use] - pub fn has_service(&self) -> bool { - self.plan.has_service() - } - pub fn handle_service( - &self, - context: crate::middleware::service::Context, - input: crate::middleware::service::Value, - next: crate::middleware::service::Next, - ) -> BoxFuture< - 'static, - Result, - > { - self.plan.handle_service(context, input, next) - } - - #[must_use] - pub fn has_websocket(&self) -> bool { - self.plan.has_websocket() - } - pub fn handle_websocket( - &self, - context: crate::middleware::websocket::Context, - message: crate::middleware::websocket::Message, - next: crate::middleware::websocket::Next, - ) -> BoxFuture<'static, Result, MiddlewareError>> - { - self.plan.handle_websocket(context, message, next) - } - - #[must_use] - pub fn has_http(&self) -> bool { - self.plan.has_http() - } - - pub fn handle_http( - &self, - mut context: crate::middleware::http::Context, - request: crate::middleware::http::Request, - next: crate::middleware::http::Next, - ) -> BoxFuture<'static, Result> { - context.plan = Some(self.clone()); - self.plan.handle_http(context, request, next) - } - - #[must_use] - pub fn new(plan: Arc, generation: ExtensionSetReference) -> Self { - Self { - plan, - _generation: generation, - } - } - - pub fn handle( - &self, - mut context: MiddlewareContext, - request: MiddlewareRequest, - next: MiddlewareNext, - ) -> BoxFuture<'static, Result> { - context.plan = Some(self.clone()); - self.plan.handle(context, request, next) - } -} - -impl fmt::Debug for FrozenMiddlewarePlan { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("FrozenMiddlewarePlan") - .finish_non_exhaustive() - } -} - -/// 同一发布代次不能被静默替换 -#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] -#[error("middleware generation is already registered")] -pub struct MiddlewareRegistrationError; - -/// 按请求冻结的发布集合解析中间件计划；索引不延长旧代次寿命 -#[derive(Clone, Default)] -pub struct MiddlewareExtensionIndex { - sets: Arc>>>, -} - -impl MiddlewareExtensionIndex { - pub fn register( - &self, - id: ExtensionSetId, - plan: Arc, - ) -> Result, MiddlewareRegistrationError> { - let mut sets = self - .sets - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - sets.retain(|_, plan| plan.strong_count() > 0); - if sets.contains_key(&id) { - return Err(MiddlewareRegistrationError); - } - sets.insert(id, Arc::downgrade(&plan)); - Ok(plan) - } - - #[must_use] - pub fn resolve(&self, generation: &ExtensionSetReference) -> Option { - self.sets - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .get(generation.id()) - .and_then(Weak::upgrade) - .map(|plan| FrozenMiddlewarePlan::new(plan, generation.clone())) - } -} diff --git a/backend/crates/gateway-core/src/middleware/http/mod.rs b/backend/crates/gateway-core/src/engine/middleware/http/mod.rs similarity index 86% rename from backend/crates/gateway-core/src/middleware/http/mod.rs rename to backend/crates/gateway-core/src/engine/middleware/http/mod.rs index 7743556a9..096c6c0e2 100644 --- a/backend/crates/gateway-core/src/middleware/http/mod.rs +++ b/backend/crates/gateway-core/src/engine/middleware/http/mod.rs @@ -13,7 +13,7 @@ use crate::{engine::middleware::MiddlewareError, lifecycle::CancellationToken}; pub type Body = UnsyncBoxBody>; pub type Request = http::Request; pub type Response = http::Response; -pub type Next = super::Next; +pub type Next = crate::middleware::Next; /// 路由前调用没有 Client Key；身份在默认认证流程中产生 #[derive(Clone, Debug)] @@ -29,7 +29,7 @@ pub struct Context { /// 内部子请求复用 API 总路由；端口不解释业务路径，也不经网络回环 pub trait Dispatcher: Send + Sync { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { None } @@ -44,7 +44,7 @@ pub trait Dispatcher: Send + Sync { #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct Settings { pub timeout: Option, - pub runtime: Option, + pub runtime: Option, } #[must_use] diff --git a/backend/crates/gateway-core/src/middleware/http/upgrade.rs b/backend/crates/gateway-core/src/engine/middleware/http/upgrade.rs similarity index 93% rename from backend/crates/gateway-core/src/middleware/http/upgrade.rs rename to backend/crates/gateway-core/src/engine/middleware/http/upgrade.rs index e29829170..25b062181 100644 --- a/backend/crates/gateway-core/src/middleware/http/upgrade.rs +++ b/backend/crates/gateway-core/src/engine/middleware/http/upgrade.rs @@ -4,7 +4,7 @@ use std::sync::{Arc, Mutex}; use futures::future::{BoxFuture, Shared}; -use crate::{engine::middleware::MiddlewareError, middleware::websocket::Session}; +use crate::{engine::middleware::MiddlewareError, engine::middleware::websocket::Session}; pub type PendingSession = Shared, Arc>>>; diff --git a/backend/crates/gateway-core/src/engine/middleware/mod.rs b/backend/crates/gateway-core/src/engine/middleware/mod.rs new file mode 100644 index 000000000..c193773fa --- /dev/null +++ b/backend/crates/gateway-core/src/engine/middleware/mod.rs @@ -0,0 +1,15 @@ +//! 执行层的请求、HTTP、WebSocket 与公开服务中间件合同 + +mod contract; +pub mod http; +mod plan; +pub mod service; +pub mod websocket; + +pub use contract::{ + MiddlewareAuthority, MiddlewareBody, MiddlewareCapabilityDeclaration, MiddlewareContext, + MiddlewareError, MiddlewareFrame, MiddlewareFrameEnvelope, MiddlewareFraming, MiddlewareHeader, + MiddlewareMount, MiddlewareNext, MiddlewareRequest, MiddlewareResponse, + MiddlewareResponseEnvelope, MiddlewareTarget, +}; +pub use plan::{FrozenMiddlewarePlan, MiddlewarePlan}; diff --git a/backend/crates/gateway-core/src/engine/middleware/plan.rs b/backend/crates/gateway-core/src/engine/middleware/plan.rs new file mode 100644 index 000000000..d96ab8f86 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/middleware/plan.rs @@ -0,0 +1,152 @@ +//! 与发布代次一起冻结的协议及请求中间件计划 + +use super::{ + MiddlewareContext, MiddlewareError, MiddlewareNext, MiddlewareRequest, MiddlewareResponse, +}; +use crate::routing::extensions::ExtensionSetReference; +use futures::future::BoxFuture; +use std::{fmt, sync::Arc}; + +/// 一个发布代次冻结的统一数据面中间件计划 +pub trait MiddlewarePlan: Send + Sync + fmt::Debug { + fn has_service(&self) -> bool { + false + } + fn handle_service( + &self, + _context: crate::engine::middleware::service::Context, + input: crate::engine::middleware::service::Value, + next: crate::engine::middleware::service::Next, + ) -> BoxFuture< + 'static, + Result< + crate::engine::middleware::service::Value, + crate::engine::middleware::service::Error, + >, + > { + next.run(input) + } + + fn has_websocket(&self) -> bool { + false + } + fn handle_websocket( + &self, + _context: crate::engine::middleware::websocket::Context, + message: crate::engine::middleware::websocket::Message, + next: crate::engine::middleware::websocket::Next, + ) -> BoxFuture< + 'static, + Result, MiddlewareError>, + > { + next.run(message) + } + + /// HTTP 入口在认证、路由与正文解析之前组合，无匹配时保留直接路径 + fn has_http(&self) -> bool { + false + } + + fn handle_http( + &self, + _context: crate::engine::middleware::http::Context, + request: crate::engine::middleware::http::Request, + next: crate::engine::middleware::http::Next, + ) -> BoxFuture<'static, Result> + { + next.run(request) + } + + fn handle( + &self, + context: MiddlewareContext, + request: MiddlewareRequest, + next: MiddlewareNext, + ) -> BoxFuture<'static, Result>; +} + +/// 调用期间同时保活发布代次与中间件计划 +#[derive(Clone)] +pub struct FrozenMiddlewarePlan { + plan: Arc, + _generation: ExtensionSetReference, +} + +impl FrozenMiddlewarePlan { + #[must_use] + pub fn has_service(&self) -> bool { + self.plan.has_service() + } + pub fn handle_service( + &self, + context: crate::engine::middleware::service::Context, + input: crate::engine::middleware::service::Value, + next: crate::engine::middleware::service::Next, + ) -> BoxFuture< + 'static, + Result< + crate::engine::middleware::service::Value, + crate::engine::middleware::service::Error, + >, + > { + self.plan.handle_service(context, input, next) + } + + #[must_use] + pub fn has_websocket(&self) -> bool { + self.plan.has_websocket() + } + pub fn handle_websocket( + &self, + context: crate::engine::middleware::websocket::Context, + message: crate::engine::middleware::websocket::Message, + next: crate::engine::middleware::websocket::Next, + ) -> BoxFuture< + 'static, + Result, MiddlewareError>, + > { + self.plan.handle_websocket(context, message, next) + } + + #[must_use] + pub fn has_http(&self) -> bool { + self.plan.has_http() + } + + pub fn handle_http( + &self, + mut context: crate::engine::middleware::http::Context, + request: crate::engine::middleware::http::Request, + next: crate::engine::middleware::http::Next, + ) -> BoxFuture<'static, Result> + { + context.plan = Some(self.clone()); + self.plan.handle_http(context, request, next) + } + + #[must_use] + pub fn new(plan: Arc, generation: ExtensionSetReference) -> Self { + Self { + plan, + _generation: generation, + } + } + + pub fn handle( + &self, + mut context: MiddlewareContext, + request: MiddlewareRequest, + next: MiddlewareNext, + ) -> BoxFuture<'static, Result> { + context.plan = Some(self.clone()); + self.plan.handle(context, request, next) + } +} + +impl fmt::Debug for FrozenMiddlewarePlan { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("FrozenMiddlewarePlan") + .finish_non_exhaustive() + } +} diff --git a/backend/crates/gateway-core/src/middleware/service.rs b/backend/crates/gateway-core/src/engine/middleware/service.rs similarity index 95% rename from backend/crates/gateway-core/src/middleware/service.rs rename to backend/crates/gateway-core/src/engine/middleware/service.rs index 8e4d55500..eeffba953 100644 --- a/backend/crates/gateway-core/src/middleware/service.rs +++ b/backend/crates/gateway-core/src/engine/middleware/service.rs @@ -7,7 +7,7 @@ use crate::{ /// 值只在进入插件链时编码；操作的领域校验与副作用由终端所属服务执行 pub type Value = serde_json::Value; -pub type Next = super::Next; +pub type Next = crate::middleware::Next; #[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, thiserror::Error)] #[error("{message}")] diff --git a/backend/crates/gateway-core/src/middleware/websocket.rs b/backend/crates/gateway-core/src/engine/middleware/websocket.rs similarity index 93% rename from backend/crates/gateway-core/src/middleware/websocket.rs rename to backend/crates/gateway-core/src/engine/middleware/websocket.rs index 16fc1cb38..2caba924c 100644 --- a/backend/crates/gateway-core/src/middleware/websocket.rs +++ b/backend/crates/gateway-core/src/engine/middleware/websocket.rs @@ -44,4 +44,4 @@ pub struct Context { pub cancellation: CancellationToken, pub sender: Arc, } -pub type Next = super::Next, MiddlewareError>; +pub type Next = crate::middleware::Next, MiddlewareError>; diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index 6fd4b841b..34893182b 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -101,17 +101,7 @@ pub enum AttemptTrigger { AccountRetry, } -/// Provider 可为当前 attempt 选择的请求局部传输档位 -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] -pub enum AttemptTransport { - /// 使用 Provider 的默认传输策略 - #[default] - Default, - /// 固定账号重试 Provider 的首选传输；序号由 Provider 的独立预算驱动 - Retry(NonZeroU32), - /// 使用 Provider 定义的备用传输 - Fallback, -} +pub use crate::upstream::AttemptTransport; impl AttemptTrigger { #[must_use] diff --git a/backend/crates/gateway-core/src/engine/nested.rs b/backend/crates/gateway-core/src/engine/nested.rs index 3ef1f35fb..82dede6a3 100644 --- a/backend/crates/gateway-core/src/engine/nested.rs +++ b/backend/crates/gateway-core/src/engine/nested.rs @@ -65,7 +65,7 @@ pub struct NestedModelExecutionRequest { /// Core 在绑定时冻结当前 /// Key 策略；每个实际模型请求仍独立经过准入、预算、账本与计费 pub struct BoundModelExecutionBinding { - pub settings: Option, + pub settings: Option, pub client_key_id: ClientApiKeyId, pub initiating_plugin_instance_id: String, /// 身份随父调用或连接回收；每个模型请求使用独立的执行期限 diff --git a/backend/crates/gateway-core/src/engine/observation.rs b/backend/crates/gateway-core/src/engine/observation.rs index 3dddea682..700b81e4d 100644 --- a/backend/crates/gateway-core/src/engine/observation.rs +++ b/backend/crates/gateway-core/src/engine/observation.rs @@ -22,8 +22,8 @@ use crate::{ metering::{CostEstimate, CostSource, Usage}, operation::OperationKind, policy::ClientApiKeyId, + routing::extensions::{ExtensionSetId, ExtensionSetReference}, routing::{AccountGroupId, ConfigRevision, PublicModelId}, - runtime::extensions::{ExtensionSetId, ExtensionSetReference}, upstream::UpstreamSendState, }; diff --git a/backend/crates/gateway-core/src/engine/policy.rs b/backend/crates/gateway-core/src/engine/policy.rs index 73d9935bd..6ee86cef8 100644 --- a/backend/crates/gateway-core/src/engine/policy.rs +++ b/backend/crates/gateway-core/src/engine/policy.rs @@ -19,8 +19,8 @@ use crate::{ identity::ProviderKind, operation::Operation, policy::ClientApiKeyId, + routing::extensions::{ExtensionSetId, ExtensionSetReference}, routing::{AccountGroupId, PublicModelId}, - runtime::extensions::{ExtensionSetId, ExtensionSetReference}, }; /// 模型路由插件的一次输入；正文仍保存在 `operation` 中，由 Runtime 按权限投影 diff --git a/backend/crates/gateway-core/src/engine/upstream_adapter.rs b/backend/crates/gateway-core/src/engine/upstream_adapter.rs index 1d1041aba..836dc1f02 100644 --- a/backend/crates/gateway-core/src/engine/upstream_adapter.rs +++ b/backend/crates/gateway-core/src/engine/upstream_adapter.rs @@ -15,7 +15,7 @@ use crate::{ metering::{CalculatedCost, Usage}, operation::Operation, routing::UpstreamModelId, - runtime::extensions::ExtensionSetReference, + routing::extensions::ExtensionSetReference, }; use super::{ diff --git a/backend/crates/gateway-core/src/error/mod.rs b/backend/crates/gateway-core/src/error/mod.rs index cc97f3ffa..2455a26e7 100644 --- a/backend/crates/gateway-core/src/error/mod.rs +++ b/backend/crates/gateway-core/src/error/mod.rs @@ -144,7 +144,7 @@ pub enum PreDeliveryRetry { /// 可靠 NotSent 的建连恢复；次数与时间由 Core 的请求级预算裁决 SameAccountConnectionRetry { /// Provider 明确指定实际失败的传输，避免 HTTP 恢复重新进入 WS - transport: crate::engine::AttemptTransport, + transport: crate::upstream::AttemptTransport, }, /// 固定本次账号，并要求 Provider 使用备用传输 SameAccountTransportFallback, @@ -666,7 +666,7 @@ impl ProviderError { } #[must_use] - pub fn with_connection_retry(mut self, transport: crate::engine::AttemptTransport) -> Self { + pub fn with_connection_retry(mut self, transport: crate::upstream::AttemptTransport) -> Self { self.pre_delivery_retry = Some(Box::new(PreDeliveryRetry::SameAccountConnectionRetry { transport, })); diff --git a/backend/crates/gateway-core/src/lib.rs b/backend/crates/gateway-core/src/lib.rs index ff130bf22..7b206691f 100644 --- a/backend/crates/gateway-core/src/lib.rs +++ b/backend/crates/gateway-core/src/lib.rs @@ -128,7 +128,7 @@ pub struct CoreStartup { providers: ProviderRegistry, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, control: CoreControlPlaneStartup, } @@ -164,7 +164,7 @@ impl CoreStartup { self.providers, self.request_observers, self.request_policies, - self.middlewares, + self.execution_extensions, self.frontend_authentication, )); let execution: Arc = service.clone(); @@ -194,7 +194,7 @@ fn build_execution_service( providers: ProviderRegistry, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, ) -> DefaultExecutionService { let mut service = DefaultExecutionService::new( @@ -215,8 +215,8 @@ fn build_execution_service( if let Some(request_policies) = request_policies { service = service.with_request_policies(request_policies); } - if let Some(middlewares) = middlewares { - service = service.with_middlewares(middlewares); + if let Some(execution_extensions) = execution_extensions { + service = service.with_execution_extensions(execution_extensions); } if let Some(frontend_authentication) = frontend_authentication { service = service.with_frontend_authentication(frontend_authentication); @@ -276,10 +276,10 @@ impl CoreBundle { pub async fn initialize( ports: CoreStorePorts, providers: ProviderRegistry, - extensions: Option>, + extensions: Option>, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, ) -> Result { prepare( @@ -288,7 +288,7 @@ pub async fn initialize( extensions, request_observers, request_policies, - middlewares, + execution_extensions, frontend_authentication, ) .activate() @@ -300,10 +300,10 @@ pub async fn initialize( pub fn prepare( ports: CoreStorePorts, providers: ProviderRegistry, - extensions: Option>, + extensions: Option>, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, ) -> CoreStartup { let control = prepare_control_plane(ports.clone(), providers.clone(), extensions); @@ -312,7 +312,7 @@ pub fn prepare( providers, request_observers, request_policies, - middlewares, + execution_extensions, frontend_authentication, control, } @@ -377,7 +377,7 @@ pub struct CoreCommandPlaneStartup { providers: ProviderRegistry, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, control: CoreControlPlaneStartup, } @@ -407,7 +407,7 @@ impl CoreCommandPlaneStartup { self.providers, self.request_observers, self.request_policies, - self.middlewares, + self.execution_extensions, None, ) .with_snapshot_refresh(Arc::clone(&control.publisher)), @@ -449,10 +449,10 @@ impl CoreCommandPlaneBundle { pub fn prepare_command_plane( ports: CoreStorePorts, providers: ProviderRegistry, - extensions: Option>, + extensions: Option>, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, ) -> CoreCommandPlaneStartup { let control = prepare_control_plane(ports.clone(), providers.clone(), extensions); CoreCommandPlaneStartup { @@ -460,7 +460,7 @@ pub fn prepare_command_plane( providers, request_observers, request_policies, - middlewares, + execution_extensions, control, } } @@ -470,7 +470,7 @@ pub fn prepare_command_plane( pub fn prepare_control_plane( ports: CoreStorePorts, providers: ProviderRegistry, - extensions: Option>, + extensions: Option>, ) -> CoreControlPlaneStartup { let mut compiler = RuntimeSnapshotCompiler::new(ports.snapshots, Arc::new(providers)); if let Some(extensions) = extensions { diff --git a/backend/crates/gateway-core/src/middleware/mod.rs b/backend/crates/gateway-core/src/middleware.rs similarity index 97% rename from backend/crates/gateway-core/src/middleware/mod.rs rename to backend/crates/gateway-core/src/middleware.rs index 7e8b3b818..c851004bd 100644 --- a/backend/crates/gateway-core/src/middleware/mod.rs +++ b/backend/crates/gateway-core/src/middleware.rs @@ -1,9 +1,5 @@ //! 与协议无关的洋葱组合器；输入、结果与错误保留调用方的领域类型 -pub mod http; -pub mod service; -pub mod websocket; - use futures::future::BoxFuture; /// 单次调用中的一层；上下文放在类型化输入中或由闭包捕获 diff --git a/backend/crates/gateway-core/src/provider_ports.rs b/backend/crates/gateway-core/src/provider_ports.rs deleted file mode 100644 index b9df9ffcb..000000000 --- a/backend/crates/gateway-core/src/provider_ports.rs +++ /dev/null @@ -1,1645 +0,0 @@ -//! Provider 运行时所需的中立存储能力 - -use std::collections::{BTreeMap, BTreeSet}; -use std::fmt; -use std::num::NonZeroU32; -use std::sync::Arc; -use std::time::{Duration, SystemTime}; - -use futures::future::BoxFuture; - -use crate::account::{ - AccountConcurrency, AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision, - CredentialState, OpaqueProviderData, ProviderAccountId, ProviderAccountStore, -}; -use crate::identity::ProviderKind; -use crate::policy::ClientApiKeyId; -use crate::routing::{ConfigRevision, UpstreamModelId}; -use crate::validation::{IdentifierError, validate_text}; - -const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60); - -/// Provider 可据此决定是否重试，但看不到 SQL、Redis 或秘密原文 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ProviderStoreErrorKind { - Unavailable, - InvalidData, - Conflict, -} - -/// Provider 存储端口的脱敏错误 -#[derive(Debug, Clone, thiserror::Error)] -#[error("provider store {operation} failed: {kind:?}")] -pub struct ProviderStoreError { - kind: ProviderStoreErrorKind, - operation: &'static str, - source: Option, -} - -impl ProviderStoreError { - #[must_use] - pub const fn new(kind: ProviderStoreErrorKind, operation: &'static str) -> Self { - Self { - kind, - operation, - source: None, - } - } - - /// 包装基础设施失败，分类与原始来源分别保留 - #[must_use] - pub fn caused_by( - kind: ProviderStoreErrorKind, - operation: &'static str, - source: impl Into, - ) -> Self { - Self { - kind, - operation, - source: Some(source.into()), - } - } - - #[must_use] - pub const fn kind(&self) -> ProviderStoreErrorKind { - self.kind - } -} - -/// 一个 Provider 的完整可重建调度状态 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderSchedulingState { - signals: BTreeMap, - round_robin_cursor: u64, -} - -impl ProviderSchedulingState { - #[must_use] - pub const fn new( - signals: BTreeMap, - round_robin_cursor: u64, - ) -> Self { - Self { - signals, - round_robin_cursor, - } - } - - #[must_use] - pub const fn signals(&self) -> &BTreeMap { - &self.signals - } - - #[must_use] - pub const fn round_robin_cursor(&self) -> u64 { - self.round_robin_cursor - } -} - -/// 请求级账号 lease 的全部中立事实 -#[derive(Debug, Clone)] -pub struct ProviderSchedulingLeaseRequest { - provider_kind: ProviderKind, - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - max_concurrent: AccountConcurrency, - request_interval: Duration, - deadline: crate::lifecycle::Deadline, - cancellation: crate::lifecycle::CancellationToken, -} - -impl ProviderSchedulingLeaseRequest { - #[must_use] - pub fn new( - provider_kind: ProviderKind, - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - max_concurrent: impl Into, - request_interval: Duration, - deadline: impl Into, - ) -> Self { - Self { - provider_kind, - account_id, - credential_revision, - max_concurrent: max_concurrent.into(), - request_interval, - deadline: deadline.into(), - cancellation: crate::lifecycle::CancellationToken::new(), - } - } - - #[must_use] - pub fn with_cancellation(mut self, cancellation: crate::lifecycle::CancellationToken) -> Self { - self.cancellation = cancellation; - self - } - - #[must_use] - pub fn cancellation(&self) -> crate::lifecycle::CancellationToken { - self.cancellation.clone() - } - - #[must_use] - pub const fn provider_kind(&self) -> &ProviderKind { - &self.provider_kind - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn max_concurrent(&self) -> AccountConcurrency { - self.max_concurrent - } - - #[must_use] - pub const fn request_interval(&self) -> Duration { - self.request_interval - } - - #[must_use] - pub const fn deadline(&self) -> crate::lifecycle::Deadline { - self.deadline - } -} - -/// Lease 生命周期由具体 Store guard 管理，Provider 只能持有 -pub trait ProviderLeaseGuard: Send + Sync + 'static {} - -impl ProviderLeaseGuard for T where T: Send + Sync + 'static {} - -pub enum ProviderLeaseAcquisition { - Acquired(Box), - Busy { retry_after: Option }, -} - -impl fmt::Debug for ProviderLeaseAcquisition { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Acquired(_) => formatter.write_str("Acquired([LEASE])"), - Self::Busy { retry_after } => formatter - .debug_struct("Busy") - .field("retry_after", retry_after) - .finish(), - } - } -} - -/// Provider 运行时会持有的三类 lease；刷新必须同时持有全局容量与账号互斥 lease -#[derive(Debug, Clone)] -pub enum ProviderLeaseRequest { - Scheduling(ProviderSchedulingLeaseRequest), - RefreshCapacity(ProviderRefreshCapacityRequest), - Refresh(ProviderRefreshLeaseRequest), -} - -pub trait ProviderLeasePort: Send + Sync { - fn load_state<'a>( - &'a self, - client_api_key_id: &'a ClientApiKeyId, - provider_kind: &'a ProviderKind, - accounts: &'a [ProviderAccountId], - ) -> BoxFuture<'a, Result>; - - fn try_acquire( - &self, - request: ProviderLeaseRequest, - ) -> BoxFuture<'_, Result>; - - /// 读取指定账号当前的在途请求数；只用于容量熔断的峰值证据， - /// 支持租约信号的存储实现覆盖，否则视为不可观测（空映射） - fn account_in_flight<'a>( - &'a self, - account_ids: &'a [ProviderAccountId], - ) -> BoxFuture<'a, Result, ProviderStoreError>> { - let _ = account_ids; - Box::pin(async move { Ok(BTreeMap::new()) }) - } -} - -/// Provider 从原始会话锚点派生的不可逆亲和键 -#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ProviderSessionAffinityKey(String); - -impl ProviderSessionAffinityKey { - pub fn try_new(value: impl Into) -> Result { - let value = value.into(); - if value.is_empty() - || value.len() > 128 - || !value.bytes().all(|byte| { - byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'-' | b'_') - }) - { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate provider session affinity key", - )); - } - Ok(Self(value)) - } - - #[must_use] - pub fn expose_to_store(&self) -> &str { - &self.0 - } -} - -impl fmt::Debug for ProviderSessionAffinityKey { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("ProviderSessionAffinityKey([OPAQUE])") - } -} - -/// 会话绑定快照；版本区分同一账号的不同认领，防止过期和 A → B → A 后的旧写入 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderSessionBinding { - account_id: ProviderAccountId, - revision: String, -} - -impl ProviderSessionBinding { - pub fn new( - account_id: ProviderAccountId, - revision: String, - ) -> Result { - if revision.len() != 32 || !revision.bytes().all(|byte| byte.is_ascii_hexdigit()) { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "decode provider session binding revision", - )); - } - Ok(Self { - account_id, - revision, - }) - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub fn revision(&self) -> &str { - &self.revision - } -} - -/// 已观测请求关联及其账号迁移权限，由 Provider 解释协议后写入 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderSessionAlias { - pub session_key: ProviderSessionAffinityKey, - pub root_session_key: Option, - pub follow_only: bool, -} - -/// 客户端作用域内的会话账号绑定，原始会话身份由 Provider 哈希后传入 -pub trait ProviderSessionAffinityPort: Send + Sync { - fn load<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - /// 在发送前原子认领、续期或迁移；None 表示快照已变化，调用方必须释放租约重新选择 - /// expected 为 None 只允许首次认领，已有快照必须连同版本匹配 - fn compare_and_bind<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - expected: Option<&'a ProviderSessionBinding>, - account_id: &'a ProviderAccountId, - ttl: Duration, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - /// 显式观测到的请求关联只指向会话键，不缓存账号，迁移后仍读取当前绑定 - fn load_alias<'a>( - &'a self, - provider: &'a ProviderKind, - alias: &'a ProviderSessionAffinityKey, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - /// 只允许首次关联或同目标续期，冲突时禁止把同一轮次改指其他会话 - fn bind_alias<'a>( - &'a self, - provider: &'a ProviderKind, - alias: &'a ProviderSessionAffinityKey, - session: &'a ProviderSessionAlias, - ttl: Duration, - ) -> BoxFuture<'a, Result>; -} - -/// Provider 会话内已失败账号的可丢失排除集 -/// -/// Provider 自行派生会话键并决定何时写入或清理；Core 只承载调度所需的账号 ID -/// 与 compare-and-swap revision，不解释任一 Provider 协议字段 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderSessionExclusions { - excluded_accounts: BTreeSet, - revision: String, -} - -impl ProviderSessionExclusions { - #[must_use] - pub const fn new(excluded_accounts: BTreeSet, revision: String) -> Self { - Self { - excluded_accounts, - revision, - } - } - - #[must_use] - pub const fn excluded_accounts(&self) -> &BTreeSet { - &self.excluded_accounts - } - - #[must_use] - pub fn revision(&self) -> &str { - &self.revision - } -} - -/// 可丢失的 Provider 会话级账号排除状态 -/// -/// 该端口不接收协议正文；Provider 只能以不可逆会话键、账号 ID 和固定 TTL 操作 -pub trait ProviderSessionExclusionPort: Send + Sync { - fn load<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - fn record_failure<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - account_id: &'a ProviderAccountId, - ttl: Duration, - ) -> BoxFuture<'a, Result>; - - fn clear<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - expected_revision: &'a str, - ) -> BoxFuture<'a, Result>; -} - -/// 所有 Provider 共享的 OAuth refresh 并发容量 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ProviderRefreshCapacityRequest { - max_concurrent: NonZeroU32, -} - -impl ProviderRefreshCapacityRequest { - #[must_use] - pub const fn new(max_concurrent: NonZeroU32) -> Self { - Self { max_concurrent } - } - - #[must_use] - pub const fn max_concurrent(self) -> NonZeroU32 { - self.max_concurrent - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderRefreshLeaseRequest { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, -} - -impl ProviderRefreshLeaseRequest { - #[must_use] - pub const fn new( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - ) -> Self { - Self { - account_id, - credential_revision, - } - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } -} - -/// Provider 定义的 catalog cache 作用域 -/// -/// Core 不解释其值；例如 Provider 可以使用套餐、区域或产品线作为共享目录边界 -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ProviderCatalogScope(String); - -impl ProviderCatalogScope { - /// 创建一个稳定且可用于 Redis 隔离键的 Provider-owned 作用域 - /// - /// # Errors - /// - /// 空值、过长文本或控制字符会被拒绝 - pub fn new(value: impl Into) -> Result { - let value = value.into(); - validate_text(&value, 128, false, None)?; - Ok(Self(value)) - } - - /// 返回 Provider-owned 作用域文本 - #[must_use] - pub fn as_str(&self) -> &str { - &self.0 - } -} - -/// Opaque catalog cache 的 Provider 与 Provider-owned 作用域 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderCatalogCacheKey { - provider_kind: ProviderKind, - scope: ProviderCatalogScope, -} - -impl ProviderCatalogCacheKey { - #[must_use] - pub const fn new(provider_kind: ProviderKind, scope: ProviderCatalogScope) -> Self { - Self { - provider_kind, - scope, - } - } - - #[must_use] - pub const fn provider_kind(&self) -> &ProviderKind { - &self.provider_kind - } - - #[must_use] - pub const fn scope(&self) -> &ProviderCatalogScope { - &self.scope - } -} - -pub trait ProviderCatalogCachePort: Send + Sync { - fn replace<'a>( - &'a self, - key: &'a ProviderCatalogCacheKey, - catalog: &'a OpaqueProviderData, - ttl: Duration, - ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; - - fn read<'a>( - &'a self, - key: &'a ProviderCatalogCacheKey, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; -} - -/// Provider 从官方制品核验出的可重建请求画像 -/// -/// Core 只用单调制品序号约束覆盖顺序；具体版本字段由对应 Provider 放在 -/// `profile` 中解释 -/// 每个 Provider 的各制品分别保留一份最新画像 -#[derive(Clone, PartialEq)] -pub struct ProviderArtifactProfile { - provider_kind: ProviderKind, - artifact_key: String, - artifact_sequence: u64, - verified_at: SystemTime, - profile: OpaqueProviderData, -} - -impl ProviderArtifactProfile { - #[must_use] - pub const fn new( - provider_kind: ProviderKind, - artifact_key: String, - artifact_sequence: u64, - verified_at: SystemTime, - profile: OpaqueProviderData, - ) -> Self { - Self { - provider_kind, - artifact_key, - artifact_sequence, - verified_at, - profile, - } - } - - #[must_use] - pub const fn provider_kind(&self) -> &ProviderKind { - &self.provider_kind - } - - #[must_use] - pub fn artifact_key(&self) -> &str { - &self.artifact_key - } - - #[must_use] - pub const fn artifact_sequence(&self) -> u64 { - self.artifact_sequence - } - - #[must_use] - pub const fn verified_at(&self) -> SystemTime { - self.verified_at - } - - #[must_use] - pub const fn profile(&self) -> &OpaqueProviderData { - &self.profile - } -} - -impl fmt::Debug for ProviderArtifactProfile { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("ProviderArtifactProfile") - .field("provider_kind", &self.provider_kind) - .field("artifact_sequence", &self.artifact_sequence) - .field("verified_at", &self.verified_at) - .field("profile", &"[PROVIDER-OWNED]") - .finish() - } -} - -pub trait ProviderArtifactProfileCachePort: Send + Sync { - /// 覆盖同一 Provider、同一制品的固定 cache key - /// - /// 返回 `false` 表示 Store 已持有更高的制品序号；相同序号但内容不同必须返回 - /// [`ProviderStoreErrorKind::Conflict`]，不能静默改写已核验画像 - fn replace_if_newer( - &self, - profile: ProviderArtifactProfile, - ttl: Duration, - ) -> BoxFuture<'_, Result>; - - fn read<'a>( - &'a self, - provider_kind: &'a ProviderKind, - artifact_key: &'a str, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; -} - -/// Redis 中可重建的账号状态投影 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderCredentialState { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - enabled: bool, - credential_state: CredentialState, - observed_at: SystemTime, -} - -impl ProviderCredentialState { - #[must_use] - pub const fn new( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - enabled: bool, - credential_state: CredentialState, - observed_at: SystemTime, - ) -> Self { - Self { - account_id, - credential_revision, - enabled, - credential_state, - observed_at, - } - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn enabled(&self) -> bool { - self.enabled - } - - #[must_use] - pub const fn credential_state(&self) -> CredentialState { - self.credential_state - } - - #[must_use] - pub const fn observed_at(&self) -> SystemTime { - self.observed_at - } -} - -pub trait ProviderCredentialStatePort: Send + Sync { - fn replace( - &self, - state: ProviderCredentialState, - ) -> BoxFuture<'_, Result<(), ProviderStoreError>>; - - fn read<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - fn clear<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result>; - - /// 记录一次瞬态刷新失败并返回窗口内累计失败次数；每次失败刷新 TTL - fn record_refresh_backoff<'a>( - &'a self, - account_id: &'a ProviderAccountId, - window: Duration, - ) -> BoxFuture<'a, Result>; - - /// 凭据完整成功轮换后清零失败计数，退避窗口重新从 base 起步 - fn clear_refresh_backoff<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; -} - -/// 账号级 cooldown 的来源类别；调度状态统一按 `rate_limited` 处理 -pub use crate::account::AccountCooldownKind as ProviderCooldownKind; - -/// 可丢失的账号冷却事实，不进入持久状态；探测冻结到期后仍需确认恢复 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderCooldown { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - until: SystemTime, - kind: ProviderCooldownKind, -} - -impl ProviderCooldown { - #[must_use] - pub const fn new( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - until: SystemTime, - ) -> Self { - Self::new_with_kind( - account_id, - credential_revision, - until, - ProviderCooldownKind::RateLimit, - ) - } - - #[must_use] - pub const fn new_with_kind( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - until: SystemTime, - kind: ProviderCooldownKind, - ) -> Self { - Self { - account_id, - credential_revision, - until, - kind, - } - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn until(&self) -> SystemTime { - self.until - } - - #[must_use] - pub const fn scheduling_state(&self) -> crate::account::AccountCooldown { - crate::account::AccountCooldown { - until: self.until, - kind: self.kind, - } - } - - #[must_use] - pub const fn kind(&self) -> ProviderCooldownKind { - self.kind - } -} - -/// 可丢失 cooldown 的细粒度作用域 -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub enum ProviderCooldownScope { - /// 只阻止同一账号调用指定上游模型 - UpstreamModel(UpstreamModelId), -} - -impl ProviderCooldownScope { - #[must_use] - pub const fn upstream_model(model: UpstreamModelId) -> Self { - Self::UpstreamModel(model) - } - - #[must_use] - pub const fn kind(&self) -> &'static str { - match self { - Self::UpstreamModel(_) => "model", - } - } - - #[must_use] - pub fn value(&self) -> &str { - match self { - Self::UpstreamModel(model) => model.as_str(), - } - } -} - -/// 不进入账号持久状态的账号+作用域 cooldown -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderScopedCooldown { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - scope: ProviderCooldownScope, - until: SystemTime, -} - -impl ProviderScopedCooldown { - #[must_use] - pub const fn new( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - scope: ProviderCooldownScope, - until: SystemTime, - ) -> Self { - Self { - account_id, - credential_revision, - scope, - until, - } - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn scope(&self) -> &ProviderCooldownScope { - &self.scope - } - - #[must_use] - pub const fn until(&self) -> SystemTime { - self.until - } -} - -pub trait ProviderCooldownPort: Send + Sync { - fn put_if_later( - &self, - cooldown: ProviderCooldown, - ) -> BoxFuture<'_, Result>; - - fn read<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - fn clear<'a>( - &'a self, - account_id: &'a ProviderAccountId, - through_revision: CredentialRevision, - ) -> BoxFuture<'a, Result>; - - fn put_scoped_if_later( - &self, - cooldown: ProviderScopedCooldown, - ) -> BoxFuture<'_, Result>; - - fn read_scoped<'a>( - &'a self, - account_id: &'a ProviderAccountId, - scope: &'a ProviderCooldownScope, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - fn clear_scoped<'a>( - &'a self, - account_id: &'a ProviderAccountId, - scope: &'a ProviderCooldownScope, - through_revision: CredentialRevision, - ) -> BoxFuture<'a, Result>; - - /// 删除账号时清除该账号全部 account/model scope cooldown key - /// 生命周期与 credential revision 无关；不要求调用方持有 revision - fn clear_all<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result>; - - /// 记录一次容量类失败并返回滑动窗口内的累计次数，同时把观测到的账号 - /// 在途并发并入窗口峰值（`in_flight` 为 0 表示本次未观测，跳过峰值更新） - /// 只服务容量熔断触发器；调用频率受失败频率约束，不需要批量接口 - fn record_capacity_failure<'a>( - &'a self, - account_id: &'a ProviderAccountId, - window: Duration, - in_flight: u32, - ) -> BoxFuture<'a, Result>; - - /// 普通请求成功后原子清除临时限流及失败证据；必须保留任何容量冻结 - fn clear_after_success<'a>( - &'a self, - account_id: &'a ProviderAccountId, - through_revision: CredentialRevision, - ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; - - /// 读取窗口内观测到的在途并发峰值；无证据时返回 `None` - fn capacity_peak_in_flight<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ProviderRefreshPolicy { - margin: Duration, - concurrency: NonZeroU32, -} - -impl ProviderRefreshPolicy { - pub fn try_new(margin: Duration, concurrency: NonZeroU32) -> Result { - if margin.is_zero() { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate refresh policy", - )); - } - Ok(Self { - margin, - concurrency, - }) - } - - #[must_use] - pub const fn margin(self) -> Duration { - self.margin - } - - #[must_use] - pub const fn concurrency(self) -> NonZeroU32 { - self.concurrency - } - - /// 判断 AT 是否已经进入当前配置的提前刷新窗口 - /// - /// `margin` 是运行时策略，不投影为账号的持久化时间字段；已过期 AT 也需要 - /// 尝试用 RT 恢复，因此视为到期 - #[must_use] - pub fn is_refresh_due( - self, - access_token_expires_at: SystemTime, - observed_at: SystemTime, - ) -> bool { - match access_token_expires_at.duration_since(observed_at) { - Ok(remaining) => remaining <= self.margin, - Err(_) => true, - } - } - - /// 按账号派生 `[0, margin]` 内的整秒错峰偏移。 - /// - /// 网关以固定周期扫描到期账号，到期时刻相同的账号会在同一轮齐刷， - /// 在 auth.openai.com 侧形成单 IP 批量刷新特征。偏移从账号 ID 的 - /// 稳定哈希派生（与退避扰动同源、独立 salt），跨扫描与重启保持不变； - /// 只会增大有效提前量，不会比配置的 margin 更贴近过期时刻。 - #[must_use] - pub fn refresh_stagger(self, account_id: &ProviderAccountId) -> Duration { - // margin 以整秒配置；秒级粒度在默认 300s 下对应约 10 个扫描槽位。 - let bound = u32::try_from(self.margin.as_secs()).unwrap_or(u32::MAX); - Duration::from_secs(u64::from(stable_factor( - account_id.as_str(), - "refresh-stagger", - 0, - bound, - ))) - } - - /// 含账号错峰偏移的到期判定；扫描路径先按 [`Self::staggered_refresh_bound`] - /// 取回候选超集，再用它在内存中收窄到本轮真正到期的账号。 - #[must_use] - pub fn is_refresh_due_staggered( - self, - account_id: &ProviderAccountId, - access_token_expires_at: SystemTime, - observed_at: SystemTime, - ) -> bool { - let staggered_margin = self.margin.saturating_add(self.refresh_stagger(account_id)); - match access_token_expires_at.duration_since(observed_at) { - Ok(remaining) => remaining <= staggered_margin, - Err(_) => true, - } - } - - /// 错峰候选窗口上界；覆盖 margin 与最大偏移之和。 - /// - /// 与 [`Self::refresh_stagger`] 的值域同址维护：偏移上限为 margin， - /// 因此 `2 × margin` 必然覆盖最大有效提前量。 - #[must_use] - pub fn staggered_refresh_bound(self) -> Duration { - self.margin.saturating_mul(2) - } -} - -/// 指数退避基准延迟；attempt=1 即为该值 -const REFRESH_BACKOFF_BASE_DELAY: Duration = Duration::from_secs(5); -/// 每多一次连续失败，基准延迟乘以该因子 -const REFRESH_BACKOFF_FACTOR: u32 = 3; -/// 退避延迟上限，避免连续失败时无限增长 -const REFRESH_BACKOFF_CAP: Duration = Duration::from_secs(300); - -/// 基于连续失败计数的指数退避重试时刻；复用 `provider_refresh_retry_at` 的稳定扰动 -/// -/// `attempt` 为窗口内累计失败次数（0 与 1 等价，均取基准延迟） -/// 延迟按 -/// `base * factor^(attempt-1)` 增长并封顶到 `REFRESH_BACKOFF_CAP` -pub fn provider_refresh_backoff_at( - account_id: &ProviderAccountId, - observed_at: SystemTime, - attempt: u32, - reason: &'static str, -) -> Result { - let exponent = attempt.saturating_sub(1); - let multiplier = REFRESH_BACKOFF_FACTOR.saturating_pow(exponent); - let scaled_seconds = REFRESH_BACKOFF_BASE_DELAY - .as_secs() - .saturating_mul(u64::from(multiplier)) - .min(REFRESH_BACKOFF_CAP.as_secs()); - let base_delay = Duration::from_secs(scaled_seconds); - provider_refresh_retry_at(account_id, observed_at, base_delay, reason) -} - -/// 临时失败后的持久重试时刻；稳定扰动避免多实例同频重试 -pub fn provider_refresh_retry_at( - account_id: &ProviderAccountId, - observed_at: SystemTime, - base_delay: Duration, - reason: &'static str, -) -> Result { - if base_delay.is_zero() { - return Err(invalid_refresh_policy("schedule refresh retry")); - } - let factor = stable_factor(account_id.as_str(), reason, 800, 1_200); - let millis = u64::try_from(base_delay.as_millis()) - .unwrap_or(u64::MAX) - .saturating_mul(u64::from(factor)) - .saturating_add(500) - / 1_000; - observed_at - .checked_add(Duration::from_millis(millis.max(1))) - .ok_or_else(|| invalid_refresh_policy("schedule refresh retry")) -} - -fn stable_factor(value: &str, salt: &str, minimum: u32, maximum: u32) -> u32 { - let mut hash = 0xcbf2_9ce4_8422_2325_u64; - for byte in value.bytes().chain([0]).chain(salt.bytes()) { - hash ^= u64::from(byte); - hash = hash.wrapping_mul(0x0000_0100_0000_01b3); - } - let width = u64::from(maximum - minimum) + 1; - minimum + u32::try_from(hash % width).unwrap_or_default() -} - -fn invalid_refresh_policy(operation: &'static str) -> ProviderStoreError { - ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) -} - -pub trait ProviderRuntimePolicyPort: Send + Sync { - /// 原子推进预热执行游标，重启或时钟回拨后不重复领取已消费的时刻 - fn claim_warmup_slot<'a>( - &'a self, - _timezone: crate::time::DeploymentTimeZone, - _slot: chrono::NaiveDateTime, - ) -> BoxFuture<'a, Result> { - Box::pin(async { - Err(ProviderStoreError::new( - ProviderStoreErrorKind::Unavailable, - "claim warmup slot", - )) - }) - } - /// 仅首次启动写入该 Provider 的默认选择；已保存的管理配置始终优先 - fn initialize_request_profile<'a>( - &'a self, - _provider: &'a ProviderKind, - initial: OpaqueProviderData, - ) -> BoxFuture<'a, Result> { - Box::pin(async move { Ok(initial) }) - } - - /// 读取候选配置版本实际引用的全局与 Client Key 画像配置 - /// - /// 实现必须在同一数据库快照内核对 revision，且只返回画像投影，不能读取 Key - /// 明文 - /// Provider 代次据此在发布前拒绝已失效的选择 - fn load_request_profile_configurations<'a>( - &'a self, - _revision: ConfigRevision, - _provider: &'a ProviderKind, - ) -> BoxFuture<'a, Result, ProviderStoreError>> { - Box::pin(async move { Ok(Vec::new()) }) - } - - fn load_refresh_policy( - &self, - ) -> BoxFuture<'_, Result>; - - /// 读取账号容量熔断策略；默认关闭，只有实现运行时设置的存储需要覆盖 - fn load_freeze_policy( - &self, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { Ok(ProviderFreezePolicy::disabled()) }) - } - - /// 读取账号模型预激活策略；默认关闭，只有实现运行时设置的存储需要覆盖 - fn load_warmup_policy( - &self, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { Ok(ProviderWarmupPolicy::disabled()) }) - } -} - -/// 账号容量熔断（自动冻结）策略；来源于 `runtime_settings`， -/// 由 Provider 触发路径与恢复 worker 共享同一份配置事实 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderFreezePolicy { - enabled: bool, - threshold: u32, - window: Duration, - freeze_duration: Duration, - probe_enabled: bool, - probe_model: Option, - adaptive_concurrency: bool, -} - -impl ProviderFreezePolicy { - /// 管理写入与运行策略构造共用的冻结约束 - pub fn validate_values( - threshold: u32, - window_seconds: u64, - duration_seconds: u64, - probe_model: Option<&str>, - ) -> Result<(), &'static str> { - for (valid, field) in [ - ( - (2..=1_000).contains(&threshold), - "account_auto_freeze_threshold", - ), - ( - (60..=3_600).contains(&window_seconds), - "account_auto_freeze_window_seconds", - ), - ( - (300..=604_800).contains(&duration_seconds), - "account_auto_freeze_duration_seconds", - ), - ( - valid_optional_probe_model(probe_model), - "account_auto_freeze_probe_model", - ), - ] { - if !valid { - return Err(field); - } - } - Ok(()) - } - - /// 边界与迁移 `0010_account_auto_freeze.sql` 的 check 约束一致； - /// store 层写入前已校验，这里兜底防御越界配置 - pub fn try_new( - enabled: bool, - threshold: u32, - window_seconds: u64, - freeze_duration_seconds: u64, - probe_enabled: bool, - probe_model: Option, - adaptive_concurrency: bool, - ) -> Result { - if Self::validate_values( - threshold, - window_seconds, - freeze_duration_seconds, - probe_model.as_deref(), - ) - .is_err() - { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate freeze policy", - )); - } - Ok(Self { - enabled, - threshold, - window: Duration::from_secs(window_seconds), - freeze_duration: Duration::from_secs(freeze_duration_seconds), - probe_enabled, - probe_model, - adaptive_concurrency, - }) - } - - /// 功能关闭时的全零策略；触发路径与 worker 都以此短路 - #[must_use] - pub const fn disabled() -> Self { - Self { - enabled: false, - threshold: 2, - window: Duration::from_secs(60), - freeze_duration: Duration::from_secs(300), - probe_enabled: false, - probe_model: None, - adaptive_concurrency: false, - } - } - - #[must_use] - pub const fn enabled(&self) -> bool { - self.enabled - } - - /// 窗口内触发冻结的请求级失败次数阈值 - #[must_use] - pub const fn threshold(&self) -> u32 { - self.threshold - } - - /// 失败计数滑动窗口；每次失败都会顺延窗口 - #[must_use] - pub const fn window(&self) -> Duration { - self.window - } - - /// 冻结时长；探测失败后的顺延也使用该值 - #[must_use] - pub const fn freeze_duration(&self) -> Duration { - self.freeze_duration - } - - #[must_use] - pub const fn probe_enabled(&self) -> bool { - self.probe_enabled - } - - /// 探测模型；`None` 表示由 worker 选择账号可用的第一个模型 - #[must_use] - pub fn probe_model(&self) -> Option<&str> { - self.probe_model.as_deref() - } - - #[must_use] - pub const fn adaptive_concurrency(&self) -> bool { - self.adaptive_concurrency - } -} - -/// 校验每日预激活时间格式，如 "08:00" 或 "08:00,13:00" -#[must_use] -pub fn valid_warmup_schedule_time(value: &str) -> bool { - if value.is_empty() - || value.len() > 255 - || value != value.trim() - || value.chars().any(char::is_control) - { - return false; - } - for part in value.split(',') { - let bytes = part.as_bytes(); - if bytes.len() != 5 || bytes[2] != b':' { - return false; - } - let Ok(hour) = part[0..2].parse::() else { - return false; - }; - let Ok(minute) = part[3..5].parse::() else { - return false; - }; - if hour > 23 || minute > 59 { - return false; - } - } - true -} - -/// 账号模型预激活（预热）策略；来源于 `runtime_settings` -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderWarmupPolicy { - enabled: bool, - schedule_time: String, - model: Option, -} - -impl ProviderWarmupPolicy { - /// 保存设置与创建预热策略使用相同的时间表和模型约束 - pub fn validate_values( - enabled: bool, - schedule_time: &str, - model: Option<&str>, - ) -> Result<(), &'static str> { - if !valid_warmup_schedule_time(schedule_time) { - return Err("account_warmup_schedule_time"); - } - if (enabled && model.is_none()) || !valid_optional_probe_model(model) { - return Err("account_warmup_model"); - } - Ok(()) - } - - pub fn try_new( - enabled: bool, - schedule_time: String, - model: Option, - ) -> Result { - if Self::validate_values(enabled, &schedule_time, model.as_deref()).is_err() { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate warmup policy", - )); - } - Ok(Self { - enabled, - schedule_time, - model, - }) - } - - #[must_use] - pub const fn disabled() -> Self { - Self { - enabled: false, - schedule_time: String::new(), - model: None, - } - } - - #[must_use] - pub const fn enabled(&self) -> bool { - self.enabled - } - - #[must_use] - pub fn schedule_time(&self) -> &str { - &self.schedule_time - } - - #[must_use] - pub fn model(&self) -> Option<&str> { - self.model.as_deref() - } - - /// 解析每日时间列表，返回如 `vec![(8, 0)]` - #[must_use] - pub fn scheduled_times(&self) -> Vec<(u32, u32)> { - self.schedule_time - .split(',') - .filter_map(|part| { - let part = part.trim(); - let mut iter = part.split(':'); - let hour = iter.next()?.parse::().ok()?; - let minute = iter.next()?.parse::().ok()?; - Some((hour, minute)) - }) - .collect() - } -} - -/// OAuth pending flow 的原始绑定只在 Provider 与 Store 边界内短暂存在 -#[derive(Clone, PartialEq, Eq)] -pub struct OAuthPendingBinding(String); - -impl OAuthPendingBinding { - pub fn try_new(value: impl Into) -> Result { - let value = value.into(); - if value.is_empty() || value.len() > 512 { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate OAuth pending binding", - )); - } - Ok(Self(value)) - } - - #[must_use] - pub fn expose_to_store(&self) -> &str { - &self.0 - } -} - -impl fmt::Debug for OAuthPendingBinding { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("OAuthPendingBinding([REDACTED])") - } -} - -#[derive(Clone, PartialEq)] -pub struct NewOAuthPendingFlow { - provider_kind: ProviderKind, - flow: OAuthPendingBinding, - owner: OAuthPendingBinding, - ttl: Duration, - payload: OpaqueProviderData, -} - -impl NewOAuthPendingFlow { - pub fn try_new( - provider_kind: ProviderKind, - flow: OAuthPendingBinding, - owner: OAuthPendingBinding, - ttl: Duration, - payload: OpaqueProviderData, - ) -> Result { - if ttl.is_zero() || ttl > MAX_PENDING_FLOW_TTL { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate OAuth pending TTL", - )); - } - Ok(Self { - provider_kind, - flow, - owner, - ttl, - payload, - }) - } - - #[must_use] - pub const fn provider_kind(&self) -> &ProviderKind { - &self.provider_kind - } - - #[must_use] - pub const fn flow(&self) -> &OAuthPendingBinding { - &self.flow - } - - #[must_use] - pub const fn owner(&self) -> &OAuthPendingBinding { - &self.owner - } - - #[must_use] - pub const fn ttl(&self) -> Duration { - self.ttl - } - - #[must_use] - pub const fn payload(&self) -> &OpaqueProviderData { - &self.payload - } -} - -impl fmt::Debug for NewOAuthPendingFlow { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("NewOAuthPendingFlow") - .field("provider_kind", &self.provider_kind) - .field("flow", &self.flow) - .field("owner", &self.owner) - .field("ttl", &self.ttl) - .field("payload", &"[PROVIDER-OWNED]") - .finish() - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum OAuthPendingPutOutcome { - Stored, - AlreadyExists, -} - -/// OAuth 回调处理取得临时 flow 的独占权结果 -/// -/// 与一次性消费不同，Provider 在上游换取令牌失败时可以释放 claim，让同一 flow -/// 使用新的回调地址重试；只有完整校验成功后才会消费 flow -#[derive(Clone, PartialEq)] -pub enum OAuthPendingClaimOutcome { - Claimed(OpaqueProviderData), - NotFound, - OwnerMismatch, - InProgress, -} - -impl fmt::Debug for OAuthPendingClaimOutcome { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Claimed(_) => formatter.write_str("Claimed([PROVIDER-OWNED])"), - Self::NotFound => formatter.write_str("NotFound"), - Self::OwnerMismatch => formatter.write_str("OwnerMismatch"), - Self::InProgress => formatter.write_str("InProgress"), - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum OAuthPendingReleaseOutcome { - Released, - NotFound, - OwnerMismatch, - ClaimMismatch, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum OAuthPendingConsumeOutcome { - Consumed, - NotFound, - OwnerMismatch, - ClaimMismatch, -} - -pub trait OAuthPendingFlowPort: Send + Sync { - fn put_if_absent( - &self, - flow: NewOAuthPendingFlow, - ) -> BoxFuture<'_, Result>; - - fn claim_if_owner<'a>( - &'a self, - provider_kind: &'a ProviderKind, - flow: &'a OAuthPendingBinding, - owner: &'a OAuthPendingBinding, - claim: &'a OAuthPendingBinding, - claim_ttl: Duration, - ) -> BoxFuture<'a, Result>; - - fn release_claim<'a>( - &'a self, - provider_kind: &'a ProviderKind, - flow: &'a OAuthPendingBinding, - owner: &'a OAuthPendingBinding, - claim: &'a OAuthPendingBinding, - ) -> BoxFuture<'a, Result>; - - fn consume_claim<'a>( - &'a self, - provider_kind: &'a ProviderKind, - flow: &'a OAuthPendingBinding, - owner: &'a OAuthPendingBinding, - claim: &'a OAuthPendingBinding, - ) -> BoxFuture<'a, Result>; -} - -/// Provider 只能按能力取用端口，无法取得 Redis client 或 repository 集合 -#[derive(Clone)] -pub struct ProviderStorePorts { - accounts: Arc, - leases: Arc, - session_affinity: Arc, - session_exclusions: Arc, - account_feedback: Arc, - catalog_cache: Arc, - artifact_profiles: Arc, - credential_state: Arc, - cooldowns: Arc, - runtime_policy: Arc, - oauth_pending: Arc, - diagnostics: Arc, -} - -impl ProviderStorePorts { - #[must_use] - // 每个参数代表独立能力端口；合并为单一配置对象会隐藏 Provider 能力边界 - #[expect(clippy::too_many_arguments)] - pub fn new( - accounts: Arc, - leases: Arc, - session_affinity: Arc, - session_exclusions: Arc, - catalog_cache: Arc, - artifact_profiles: Arc, - credential_state: Arc, - cooldowns: Arc, - runtime_policy: Arc, - oauth_pending: Arc, - diagnostics: Arc, - ) -> Self { - Self { - accounts, - leases, - session_affinity, - session_exclusions, - account_feedback: Arc::new(AccountFeedbackStats::default()), - catalog_cache, - artifact_profiles, - credential_state, - cooldowns, - runtime_policy, - oauth_pending, - diagnostics, - } - } - - #[must_use] - pub fn diagnostics(&self) -> Arc { - self.diagnostics.clone() - } - - #[must_use] - pub fn accounts(&self) -> Arc { - Arc::clone(&self.accounts) - } - - #[must_use] - pub fn leases(&self) -> Arc { - Arc::clone(&self.leases) - } - - #[must_use] - pub fn session_affinity(&self) -> Arc { - Arc::clone(&self.session_affinity) - } - - #[must_use] - pub fn session_exclusions(&self) -> Arc { - Arc::clone(&self.session_exclusions) - } - - #[must_use] - pub fn account_feedback(&self) -> Arc { - Arc::clone(&self.account_feedback) - } - - #[must_use] - pub fn catalog_cache(&self) -> Arc { - Arc::clone(&self.catalog_cache) - } - - #[must_use] - pub fn artifact_profiles(&self) -> Arc { - Arc::clone(&self.artifact_profiles) - } - - #[must_use] - pub fn credential_state(&self) -> Arc { - Arc::clone(&self.credential_state) - } - - #[must_use] - pub fn cooldowns(&self) -> Arc { - Arc::clone(&self.cooldowns) - } - - #[must_use] - pub fn runtime_policy(&self) -> Arc { - Arc::clone(&self.runtime_policy) - } - - #[must_use] - pub fn oauth_pending(&self) -> Arc { - Arc::clone(&self.oauth_pending) - } -} - -impl fmt::Debug for ProviderStorePorts { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("ProviderStorePorts([CAPABILITIES])") - } -} - -fn valid_optional_probe_model(model: Option<&str>) -> bool { - model.is_none_or(|model| { - !model.is_empty() - && model.len() <= 128 - && model == model.trim() - && !model.bytes().any(|byte| byte.is_ascii_control()) - }) -} diff --git a/backend/crates/gateway-core/src/provider_ports/affinity.rs b/backend/crates/gateway-core/src/provider_ports/affinity.rs new file mode 100644 index 000000000..b1fafa772 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/affinity.rs @@ -0,0 +1,174 @@ +//! Provider 会话亲和、绑定和排除事实 + +use super::{ProviderStoreError, ProviderStoreErrorKind}; +use crate::{account::ProviderAccountId, identity::ProviderKind}; +use futures::future::BoxFuture; +use std::{collections::BTreeSet, fmt, time::Duration}; + +/// Provider 从原始会话锚点派生的不可逆亲和键 +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ProviderSessionAffinityKey(String); + +impl ProviderSessionAffinityKey { + pub fn try_new(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() + || value.len() > 128 + || !value.bytes().all(|byte| { + byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'-' | b'_') + }) + { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate provider session affinity key", + )); + } + Ok(Self(value)) + } + + #[must_use] + pub fn expose_to_store(&self) -> &str { + &self.0 + } +} + +impl fmt::Debug for ProviderSessionAffinityKey { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ProviderSessionAffinityKey([OPAQUE])") + } +} + +/// 会话绑定快照；版本区分同一账号的不同认领，防止过期和 A → B → A 后的旧写入 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSessionBinding { + account_id: ProviderAccountId, + revision: String, +} + +impl ProviderSessionBinding { + pub fn new( + account_id: ProviderAccountId, + revision: String, + ) -> Result { + if revision.len() != 32 || !revision.bytes().all(|byte| byte.is_ascii_hexdigit()) { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "decode provider session binding revision", + )); + } + Ok(Self { + account_id, + revision, + }) + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub fn revision(&self) -> &str { + &self.revision + } +} + +/// 已观测请求关联及其账号迁移权限，由 Provider 解释协议后写入 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSessionAlias { + pub session_key: ProviderSessionAffinityKey, + pub root_session_key: Option, + pub follow_only: bool, +} + +/// 客户端作用域内的会话账号绑定，原始会话身份由 Provider 哈希后传入 +pub trait ProviderSessionAffinityPort: Send + Sync { + fn load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + /// 在发送前原子认领、续期或迁移；None 表示快照已变化，调用方必须释放租约重新选择 + /// expected 为 None 只允许首次认领，已有快照必须连同版本匹配 + fn compare_and_bind<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected: Option<&'a ProviderSessionBinding>, + account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + /// 显式观测到的请求关联只指向会话键，不缓存账号，迁移后仍读取当前绑定 + fn load_alias<'a>( + &'a self, + provider: &'a ProviderKind, + alias: &'a ProviderSessionAffinityKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + /// 只允许首次关联或同目标续期，冲突时禁止把同一轮次改指其他会话 + fn bind_alias<'a>( + &'a self, + provider: &'a ProviderKind, + alias: &'a ProviderSessionAffinityKey, + session: &'a ProviderSessionAlias, + ttl: Duration, + ) -> BoxFuture<'a, Result>; +} + +/// Provider 会话内已失败账号的可丢失排除集 +/// +/// Provider 自行派生会话键并决定何时写入或清理；Core 只承载调度所需的账号 ID +/// 与 compare-and-swap revision，不解释任一 Provider 协议字段 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSessionExclusions { + excluded_accounts: BTreeSet, + revision: String, +} + +impl ProviderSessionExclusions { + #[must_use] + pub const fn new(excluded_accounts: BTreeSet, revision: String) -> Self { + Self { + excluded_accounts, + revision, + } + } + + #[must_use] + pub const fn excluded_accounts(&self) -> &BTreeSet { + &self.excluded_accounts + } + + #[must_use] + pub fn revision(&self) -> &str { + &self.revision + } +} + +/// 可丢失的 Provider 会话级账号排除状态 +/// +/// 该端口不接收协议正文；Provider 只能以不可逆会话键、账号 ID 和固定 TTL 操作 +pub trait ProviderSessionExclusionPort: Send + Sync { + fn load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn record_failure<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> BoxFuture<'a, Result>; + + fn clear<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected_revision: &'a str, + ) -> BoxFuture<'a, Result>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/bundle.rs b/backend/crates/gateway-core/src/provider_ports/bundle.rs new file mode 100644 index 000000000..13ec981d0 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/bundle.rs @@ -0,0 +1,126 @@ +//! 按能力提供 Provider 所需的存储端口 + +use super::{ + OAuthPendingFlowPort, ProviderArtifactProfileCachePort, ProviderCatalogCachePort, + ProviderCooldownPort, ProviderCredentialStatePort, ProviderLeasePort, + ProviderRuntimePolicyPort, ProviderSessionAffinityPort, ProviderSessionExclusionPort, +}; +use crate::account::{AccountFeedbackStats, ProviderAccountStore}; +use std::{fmt, sync::Arc}; + +/// Provider 只能按能力取用端口，无法取得 Redis client 或 repository 集合 +#[derive(Clone)] +pub struct ProviderStorePorts { + accounts: Arc, + leases: Arc, + session_affinity: Arc, + session_exclusions: Arc, + account_feedback: Arc, + catalog_cache: Arc, + artifact_profiles: Arc, + credential_state: Arc, + cooldowns: Arc, + runtime_policy: Arc, + oauth_pending: Arc, + diagnostics: Arc, +} + +impl ProviderStorePorts { + #[must_use] + // 每个参数代表独立能力端口；合并为单一配置对象会隐藏 Provider 能力边界 + #[expect(clippy::too_many_arguments)] + pub fn new( + accounts: Arc, + leases: Arc, + session_affinity: Arc, + session_exclusions: Arc, + catalog_cache: Arc, + artifact_profiles: Arc, + credential_state: Arc, + cooldowns: Arc, + runtime_policy: Arc, + oauth_pending: Arc, + diagnostics: Arc, + ) -> Self { + Self { + accounts, + leases, + session_affinity, + session_exclusions, + account_feedback: Arc::new(AccountFeedbackStats::default()), + catalog_cache, + artifact_profiles, + credential_state, + cooldowns, + runtime_policy, + oauth_pending, + diagnostics, + } + } + + #[must_use] + pub fn diagnostics(&self) -> Arc { + self.diagnostics.clone() + } + + #[must_use] + pub fn accounts(&self) -> Arc { + Arc::clone(&self.accounts) + } + + #[must_use] + pub fn leases(&self) -> Arc { + Arc::clone(&self.leases) + } + + #[must_use] + pub fn session_affinity(&self) -> Arc { + Arc::clone(&self.session_affinity) + } + + #[must_use] + pub fn session_exclusions(&self) -> Arc { + Arc::clone(&self.session_exclusions) + } + + #[must_use] + pub fn account_feedback(&self) -> Arc { + Arc::clone(&self.account_feedback) + } + + #[must_use] + pub fn catalog_cache(&self) -> Arc { + Arc::clone(&self.catalog_cache) + } + + #[must_use] + pub fn artifact_profiles(&self) -> Arc { + Arc::clone(&self.artifact_profiles) + } + + #[must_use] + pub fn credential_state(&self) -> Arc { + Arc::clone(&self.credential_state) + } + + #[must_use] + pub fn cooldowns(&self) -> Arc { + Arc::clone(&self.cooldowns) + } + + #[must_use] + pub fn runtime_policy(&self) -> Arc { + Arc::clone(&self.runtime_policy) + } + + #[must_use] + pub fn oauth_pending(&self) -> Arc { + Arc::clone(&self.oauth_pending) + } +} + +impl fmt::Debug for ProviderStorePorts { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ProviderStorePorts([CAPABILITIES])") + } +} diff --git a/backend/crates/gateway-core/src/provider_ports/catalog.rs b/backend/crates/gateway-core/src/provider_ports/catalog.rs new file mode 100644 index 000000000..8fc28a2be --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/catalog.rs @@ -0,0 +1,167 @@ +//! Provider 目录与身份画像缓存端口 + +use super::ProviderStoreError; +use crate::{ + account::OpaqueProviderData, + identity::ProviderKind, + validation::{IdentifierError, validate_text}, +}; +use futures::future::BoxFuture; +use std::{ + fmt, + time::{Duration, SystemTime}, +}; + +/// Provider 定义的 catalog cache 作用域 +/// +/// Core 不解释其值；例如 Provider 可以使用套餐、区域或产品线作为共享目录边界 +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ProviderCatalogScope(String); + +impl ProviderCatalogScope { + /// 创建一个稳定且可用于 Redis 隔离键的 Provider-owned 作用域 + /// + /// # Errors + /// + /// 空值、过长文本或控制字符会被拒绝 + pub fn new(value: impl Into) -> Result { + let value = value.into(); + validate_text(&value, 128, false, None)?; + Ok(Self(value)) + } + + /// 返回 Provider-owned 作用域文本 + #[must_use] + pub fn as_str(&self) -> &str { + &self.0 + } +} + +/// Opaque catalog cache 的 Provider 与 Provider-owned 作用域 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderCatalogCacheKey { + provider_kind: ProviderKind, + scope: ProviderCatalogScope, +} + +impl ProviderCatalogCacheKey { + #[must_use] + pub const fn new(provider_kind: ProviderKind, scope: ProviderCatalogScope) -> Self { + Self { + provider_kind, + scope, + } + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub const fn scope(&self) -> &ProviderCatalogScope { + &self.scope + } +} + +pub trait ProviderCatalogCachePort: Send + Sync { + fn replace<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + catalog: &'a OpaqueProviderData, + ttl: Duration, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; + + fn read<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; +} + +/// Provider 从官方制品核验出的可重建请求画像 +/// +/// Core 只用单调制品序号约束覆盖顺序；具体版本字段由对应 Provider 放在 +/// `profile` 中解释 +/// 每个 Provider 的各制品分别保留一份最新画像 +#[derive(Clone, PartialEq)] +pub struct ProviderArtifactProfile { + provider_kind: ProviderKind, + artifact_key: String, + artifact_sequence: u64, + verified_at: SystemTime, + profile: OpaqueProviderData, +} + +impl ProviderArtifactProfile { + #[must_use] + pub const fn new( + provider_kind: ProviderKind, + artifact_key: String, + artifact_sequence: u64, + verified_at: SystemTime, + profile: OpaqueProviderData, + ) -> Self { + Self { + provider_kind, + artifact_key, + artifact_sequence, + verified_at, + profile, + } + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub fn artifact_key(&self) -> &str { + &self.artifact_key + } + + #[must_use] + pub const fn artifact_sequence(&self) -> u64 { + self.artifact_sequence + } + + #[must_use] + pub const fn verified_at(&self) -> SystemTime { + self.verified_at + } + + #[must_use] + pub const fn profile(&self) -> &OpaqueProviderData { + &self.profile + } +} + +impl fmt::Debug for ProviderArtifactProfile { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ProviderArtifactProfile") + .field("provider_kind", &self.provider_kind) + .field("artifact_sequence", &self.artifact_sequence) + .field("verified_at", &self.verified_at) + .field("profile", &"[PROVIDER-OWNED]") + .finish() + } +} + +pub trait ProviderArtifactProfileCachePort: Send + Sync { + /// 覆盖同一 Provider、同一制品的固定 cache key + /// + /// 返回 `false` 表示 Store 已持有更高的制品序号；相同序号但内容不同必须返回 + /// [`super::ProviderStoreErrorKind::Conflict`]，不能静默改写已核验画像 + fn replace_if_newer( + &self, + profile: ProviderArtifactProfile, + ttl: Duration, + ) -> BoxFuture<'_, Result>; + + fn read<'a>( + &'a self, + provider_kind: &'a ProviderKind, + artifact_key: &'a str, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/cooldown.rs b/backend/crates/gateway-core/src/provider_ports/cooldown.rs new file mode 100644 index 000000000..99df144d5 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/cooldown.rs @@ -0,0 +1,219 @@ +//! Provider 冷却、容量窗口和运行状态端口 + +use super::ProviderStoreError; +/// 账号级 cooldown 的来源类别；调度状态统一按 `rate_limited` 处理 +pub use crate::account::AccountCooldownKind as ProviderCooldownKind; +use crate::{ + account::{CredentialRevision, ProviderAccountId}, + routing::UpstreamModelId, +}; +use futures::future::BoxFuture; +use std::time::{Duration, SystemTime}; + +/// 可丢失的账号冷却事实，不进入持久状态；探测冻结到期后仍需确认恢复 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderCooldown { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + until: SystemTime, + kind: ProviderCooldownKind, +} + +impl ProviderCooldown { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + until: SystemTime, + ) -> Self { + Self::new_with_kind( + account_id, + credential_revision, + until, + ProviderCooldownKind::RateLimit, + ) + } + + #[must_use] + pub const fn new_with_kind( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + until: SystemTime, + kind: ProviderCooldownKind, + ) -> Self { + Self { + account_id, + credential_revision, + until, + kind, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn until(&self) -> SystemTime { + self.until + } + + #[must_use] + pub const fn scheduling_state(&self) -> crate::account::AccountCooldown { + crate::account::AccountCooldown { + until: self.until, + kind: self.kind, + } + } + + #[must_use] + pub const fn kind(&self) -> ProviderCooldownKind { + self.kind + } +} + +/// 可丢失 cooldown 的细粒度作用域 +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum ProviderCooldownScope { + /// 只阻止同一账号调用指定上游模型 + UpstreamModel(UpstreamModelId), +} + +impl ProviderCooldownScope { + #[must_use] + pub const fn upstream_model(model: UpstreamModelId) -> Self { + Self::UpstreamModel(model) + } + + #[must_use] + pub const fn kind(&self) -> &'static str { + match self { + Self::UpstreamModel(_) => "model", + } + } + + #[must_use] + pub fn value(&self) -> &str { + match self { + Self::UpstreamModel(model) => model.as_str(), + } + } +} + +/// 不进入账号持久状态的账号+作用域 cooldown +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderScopedCooldown { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + scope: ProviderCooldownScope, + until: SystemTime, +} + +impl ProviderScopedCooldown { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + scope: ProviderCooldownScope, + until: SystemTime, + ) -> Self { + Self { + account_id, + credential_revision, + scope, + until, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn scope(&self) -> &ProviderCooldownScope { + &self.scope + } + + #[must_use] + pub const fn until(&self) -> SystemTime { + self.until + } +} + +pub trait ProviderCooldownPort: Send + Sync { + fn put_if_later( + &self, + cooldown: ProviderCooldown, + ) -> BoxFuture<'_, Result>; + + fn read<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn clear<'a>( + &'a self, + account_id: &'a ProviderAccountId, + through_revision: CredentialRevision, + ) -> BoxFuture<'a, Result>; + + fn put_scoped_if_later( + &self, + cooldown: ProviderScopedCooldown, + ) -> BoxFuture<'_, Result>; + + fn read_scoped<'a>( + &'a self, + account_id: &'a ProviderAccountId, + scope: &'a ProviderCooldownScope, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn clear_scoped<'a>( + &'a self, + account_id: &'a ProviderAccountId, + scope: &'a ProviderCooldownScope, + through_revision: CredentialRevision, + ) -> BoxFuture<'a, Result>; + + /// 删除账号时清除该账号全部 account/model scope cooldown key + /// 生命周期与 credential revision 无关；不要求调用方持有 revision + fn clear_all<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result>; + + /// 记录一次容量类失败并返回滑动窗口内的累计次数，同时把观测到的账号 + /// 在途并发并入窗口峰值（`in_flight` 为 0 表示本次未观测，跳过峰值更新） + /// 只服务容量熔断触发器；调用频率受失败频率约束，不需要批量接口 + fn record_capacity_failure<'a>( + &'a self, + account_id: &'a ProviderAccountId, + window: Duration, + in_flight: u32, + ) -> BoxFuture<'a, Result>; + + /// 普通请求成功后原子清除临时限流及失败证据；必须保留任何容量冻结 + fn clear_after_success<'a>( + &'a self, + account_id: &'a ProviderAccountId, + through_revision: CredentialRevision, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; + + /// 读取窗口内观测到的在途并发峰值；无证据时返回 `None` + fn capacity_peak_in_flight<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/credential.rs b/backend/crates/gateway-core/src/provider_ports/credential.rs new file mode 100644 index 000000000..abbd1c75a --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/credential.rs @@ -0,0 +1,90 @@ +//! Provider 凭据运行状态读写端口 + +use super::ProviderStoreError; +use crate::account::{CredentialRevision, CredentialState, ProviderAccountId}; +use futures::future::BoxFuture; +use std::time::{Duration, SystemTime}; + +/// Redis 中可重建的账号状态投影 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderCredentialState { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + enabled: bool, + credential_state: CredentialState, + observed_at: SystemTime, +} + +impl ProviderCredentialState { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + enabled: bool, + credential_state: CredentialState, + observed_at: SystemTime, + ) -> Self { + Self { + account_id, + credential_revision, + enabled, + credential_state, + observed_at, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn enabled(&self) -> bool { + self.enabled + } + + #[must_use] + pub const fn credential_state(&self) -> CredentialState { + self.credential_state + } + + #[must_use] + pub const fn observed_at(&self) -> SystemTime { + self.observed_at + } +} + +pub trait ProviderCredentialStatePort: Send + Sync { + fn replace( + &self, + state: ProviderCredentialState, + ) -> BoxFuture<'_, Result<(), ProviderStoreError>>; + + fn read<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn clear<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result>; + + /// 记录一次瞬态刷新失败并返回窗口内累计失败次数；每次失败刷新 TTL + fn record_refresh_backoff<'a>( + &'a self, + account_id: &'a ProviderAccountId, + window: Duration, + ) -> BoxFuture<'a, Result>; + + /// 凭据完整成功轮换后清零失败计数，退避窗口重新从 base 起步 + fn clear_refresh_backoff<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/error.rs b/backend/crates/gateway-core/src/provider_ports/error.rs new file mode 100644 index 000000000..4d4917be4 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/error.rs @@ -0,0 +1,48 @@ +//! Provider 存储端口的错误分类与原始来源 + +/// Provider 可据此决定是否重试，但看不到 SQL、Redis 或秘密原文 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProviderStoreErrorKind { + Unavailable, + InvalidData, + Conflict, +} + +/// Provider 存储端口的脱敏错误 +#[derive(Debug, Clone, thiserror::Error)] +#[error("provider store {operation} failed: {kind:?}")] +pub struct ProviderStoreError { + kind: ProviderStoreErrorKind, + operation: &'static str, + source: Option, +} + +impl ProviderStoreError { + #[must_use] + pub const fn new(kind: ProviderStoreErrorKind, operation: &'static str) -> Self { + Self { + kind, + operation, + source: None, + } + } + + /// 包装基础设施失败，分类与原始来源分别保留 + #[must_use] + pub fn caused_by( + kind: ProviderStoreErrorKind, + operation: &'static str, + source: impl Into, + ) -> Self { + Self { + kind, + operation, + source: Some(source.into()), + } + } + + #[must_use] + pub const fn kind(&self) -> ProviderStoreErrorKind { + self.kind + } +} diff --git a/backend/crates/gateway-core/src/provider_ports/lease.rs b/backend/crates/gateway-core/src/provider_ports/lease.rs new file mode 100644 index 000000000..6b3202085 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/lease.rs @@ -0,0 +1,216 @@ +//! Provider 调度与刷新租约的中立端口 + +use super::ProviderStoreError; +use crate::{ + account::{AccountConcurrency, AccountRuntimeSignals, CredentialRevision, ProviderAccountId}, + identity::ProviderKind, + policy::ClientApiKeyId, +}; +use futures::future::BoxFuture; +use std::{collections::BTreeMap, fmt, num::NonZeroU32, time::Duration}; + +/// 一个 Provider 的完整可重建调度状态 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSchedulingState { + signals: BTreeMap, + round_robin_cursor: u64, +} + +impl ProviderSchedulingState { + #[must_use] + pub const fn new( + signals: BTreeMap, + round_robin_cursor: u64, + ) -> Self { + Self { + signals, + round_robin_cursor, + } + } + + #[must_use] + pub const fn signals(&self) -> &BTreeMap { + &self.signals + } + + #[must_use] + pub const fn round_robin_cursor(&self) -> u64 { + self.round_robin_cursor + } +} + +/// 请求级账号 lease 的全部中立事实 +#[derive(Debug, Clone)] +pub struct ProviderSchedulingLeaseRequest { + provider_kind: ProviderKind, + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + max_concurrent: AccountConcurrency, + request_interval: Duration, + deadline: crate::lifecycle::Deadline, + cancellation: crate::lifecycle::CancellationToken, +} + +impl ProviderSchedulingLeaseRequest { + #[must_use] + pub fn new( + provider_kind: ProviderKind, + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + max_concurrent: impl Into, + request_interval: Duration, + deadline: impl Into, + ) -> Self { + Self { + provider_kind, + account_id, + credential_revision, + max_concurrent: max_concurrent.into(), + request_interval, + deadline: deadline.into(), + cancellation: crate::lifecycle::CancellationToken::new(), + } + } + + #[must_use] + pub fn with_cancellation(mut self, cancellation: crate::lifecycle::CancellationToken) -> Self { + self.cancellation = cancellation; + self + } + + #[must_use] + pub fn cancellation(&self) -> crate::lifecycle::CancellationToken { + self.cancellation.clone() + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn max_concurrent(&self) -> AccountConcurrency { + self.max_concurrent + } + + #[must_use] + pub const fn request_interval(&self) -> Duration { + self.request_interval + } + + #[must_use] + pub const fn deadline(&self) -> crate::lifecycle::Deadline { + self.deadline + } +} + +/// Lease 生命周期由具体 Store guard 管理，Provider 只能持有 +pub trait ProviderLeaseGuard: Send + Sync + 'static {} + +impl ProviderLeaseGuard for T where T: Send + Sync + 'static {} + +pub enum ProviderLeaseAcquisition { + Acquired(Box), + Busy { retry_after: Option }, +} + +impl fmt::Debug for ProviderLeaseAcquisition { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Acquired(_) => formatter.write_str("Acquired([LEASE])"), + Self::Busy { retry_after } => formatter + .debug_struct("Busy") + .field("retry_after", retry_after) + .finish(), + } + } +} + +/// Provider 运行时会持有的三类 lease；刷新必须同时持有全局容量与账号互斥 lease +#[derive(Debug, Clone)] +pub enum ProviderLeaseRequest { + Scheduling(ProviderSchedulingLeaseRequest), + RefreshCapacity(ProviderRefreshCapacityRequest), + Refresh(ProviderRefreshLeaseRequest), +} + +pub trait ProviderLeasePort: Send + Sync { + fn load_state<'a>( + &'a self, + client_api_key_id: &'a ClientApiKeyId, + provider_kind: &'a ProviderKind, + accounts: &'a [ProviderAccountId], + ) -> BoxFuture<'a, Result>; + + fn try_acquire( + &self, + request: ProviderLeaseRequest, + ) -> BoxFuture<'_, Result>; + + /// 读取指定账号当前的在途请求数；只用于容量熔断的峰值证据， + /// 支持租约信号的存储实现覆盖，否则视为不可观测（空映射） + fn account_in_flight<'a>( + &'a self, + account_ids: &'a [ProviderAccountId], + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + let _ = account_ids; + Box::pin(async move { Ok(BTreeMap::new()) }) + } +} + +/// 所有 Provider 共享的 OAuth refresh 并发容量 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ProviderRefreshCapacityRequest { + max_concurrent: NonZeroU32, +} + +impl ProviderRefreshCapacityRequest { + #[must_use] + pub const fn new(max_concurrent: NonZeroU32) -> Self { + Self { max_concurrent } + } + + #[must_use] + pub const fn max_concurrent(self) -> NonZeroU32 { + self.max_concurrent + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderRefreshLeaseRequest { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, +} + +impl ProviderRefreshLeaseRequest { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + ) -> Self { + Self { + account_id, + credential_revision, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } +} diff --git a/backend/crates/gateway-core/src/provider_ports/mod.rs b/backend/crates/gateway-core/src/provider_ports/mod.rs new file mode 100644 index 000000000..858c3e1b2 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/mod.rs @@ -0,0 +1,41 @@ +//! Provider 运行时所需的存储合同，按能力组织端口与共享策略 + +mod error; +pub use error::{ProviderStoreError, ProviderStoreErrorKind}; +mod lease; +pub use lease::{ + ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest, + ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderSchedulingLeaseRequest, + ProviderSchedulingState, +}; +mod affinity; +pub use affinity::{ + ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionAlias, + ProviderSessionBinding, ProviderSessionExclusionPort, ProviderSessionExclusions, +}; +mod catalog; +pub use catalog::{ + ProviderArtifactProfile, ProviderArtifactProfileCachePort, ProviderCatalogCacheKey, + ProviderCatalogCachePort, ProviderCatalogScope, +}; +mod credential; +pub use credential::{ProviderCredentialState, ProviderCredentialStatePort}; +mod cooldown; +pub use cooldown::{ + ProviderCooldown, ProviderCooldownKind, ProviderCooldownPort, ProviderCooldownScope, + ProviderScopedCooldown, +}; +mod refresh; +pub use refresh::{ProviderRefreshPolicy, provider_refresh_backoff_at, provider_refresh_retry_at}; +mod runtime_policy; +pub use runtime_policy::{ + ProviderFreezePolicy, ProviderRuntimePolicyPort, ProviderWarmupPolicy, + valid_warmup_schedule_time, +}; +mod oauth; +pub use oauth::{ + NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome, + OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome, +}; +mod bundle; +pub use bundle::ProviderStorePorts; diff --git a/backend/crates/gateway-core/src/provider_ports/oauth.rs b/backend/crates/gateway-core/src/provider_ports/oauth.rs new file mode 100644 index 000000000..a84a965ac --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/oauth.rs @@ -0,0 +1,184 @@ +//! OAuth 授权 pending flow 的领取与消费合同 + +use super::{ProviderStoreError, ProviderStoreErrorKind}; +use crate::{account::OpaqueProviderData, identity::ProviderKind}; +use futures::future::BoxFuture; +use std::{fmt, time::Duration}; + +const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60); + +/// OAuth pending flow 的原始绑定只在 Provider 与 Store 边界内短暂存在 +#[derive(Clone, PartialEq, Eq)] +pub struct OAuthPendingBinding(String); + +impl OAuthPendingBinding { + pub fn try_new(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() || value.len() > 512 { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate OAuth pending binding", + )); + } + Ok(Self(value)) + } + + #[must_use] + pub fn expose_to_store(&self) -> &str { + &self.0 + } +} + +impl fmt::Debug for OAuthPendingBinding { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("OAuthPendingBinding([REDACTED])") + } +} + +#[derive(Clone, PartialEq)] +pub struct NewOAuthPendingFlow { + provider_kind: ProviderKind, + flow: OAuthPendingBinding, + owner: OAuthPendingBinding, + ttl: Duration, + payload: OpaqueProviderData, +} + +impl NewOAuthPendingFlow { + pub fn try_new( + provider_kind: ProviderKind, + flow: OAuthPendingBinding, + owner: OAuthPendingBinding, + ttl: Duration, + payload: OpaqueProviderData, + ) -> Result { + if ttl.is_zero() || ttl > MAX_PENDING_FLOW_TTL { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate OAuth pending TTL", + )); + } + Ok(Self { + provider_kind, + flow, + owner, + ttl, + payload, + }) + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub const fn flow(&self) -> &OAuthPendingBinding { + &self.flow + } + + #[must_use] + pub const fn owner(&self) -> &OAuthPendingBinding { + &self.owner + } + + #[must_use] + pub const fn ttl(&self) -> Duration { + self.ttl + } + + #[must_use] + pub const fn payload(&self) -> &OpaqueProviderData { + &self.payload + } +} + +impl fmt::Debug for NewOAuthPendingFlow { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("NewOAuthPendingFlow") + .field("provider_kind", &self.provider_kind) + .field("flow", &self.flow) + .field("owner", &self.owner) + .field("ttl", &self.ttl) + .field("payload", &"[PROVIDER-OWNED]") + .finish() + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthPendingPutOutcome { + Stored, + AlreadyExists, +} + +/// OAuth 回调处理取得临时 flow 的独占权结果 +/// +/// 与一次性消费不同，Provider 在上游换取令牌失败时可以释放 claim，让同一 flow +/// 使用新的回调地址重试；只有完整校验成功后才会消费 flow +#[derive(Clone, PartialEq)] +pub enum OAuthPendingClaimOutcome { + Claimed(OpaqueProviderData), + NotFound, + OwnerMismatch, + InProgress, +} + +impl fmt::Debug for OAuthPendingClaimOutcome { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Claimed(_) => formatter.write_str("Claimed([PROVIDER-OWNED])"), + Self::NotFound => formatter.write_str("NotFound"), + Self::OwnerMismatch => formatter.write_str("OwnerMismatch"), + Self::InProgress => formatter.write_str("InProgress"), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthPendingReleaseOutcome { + Released, + NotFound, + OwnerMismatch, + ClaimMismatch, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthPendingConsumeOutcome { + Consumed, + NotFound, + OwnerMismatch, + ClaimMismatch, +} + +pub trait OAuthPendingFlowPort: Send + Sync { + fn put_if_absent( + &self, + flow: NewOAuthPendingFlow, + ) -> BoxFuture<'_, Result>; + + fn claim_if_owner<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a OAuthPendingBinding, + owner: &'a OAuthPendingBinding, + claim: &'a OAuthPendingBinding, + claim_ttl: Duration, + ) -> BoxFuture<'a, Result>; + + fn release_claim<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a OAuthPendingBinding, + owner: &'a OAuthPendingBinding, + claim: &'a OAuthPendingBinding, + ) -> BoxFuture<'a, Result>; + + fn consume_claim<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a OAuthPendingBinding, + owner: &'a OAuthPendingBinding, + claim: &'a OAuthPendingBinding, + ) -> BoxFuture<'a, Result>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/refresh.rs b/backend/crates/gateway-core/src/provider_ports/refresh.rs new file mode 100644 index 000000000..620288f47 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/refresh.rs @@ -0,0 +1,161 @@ +//! Provider 刷新时机、退避与重试策略 + +use super::{ProviderStoreError, ProviderStoreErrorKind}; +use crate::account::ProviderAccountId; +use std::{ + num::NonZeroU32, + time::{Duration, SystemTime}, +}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ProviderRefreshPolicy { + margin: Duration, + concurrency: NonZeroU32, +} + +impl ProviderRefreshPolicy { + pub fn try_new(margin: Duration, concurrency: NonZeroU32) -> Result { + if margin.is_zero() { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate refresh policy", + )); + } + Ok(Self { + margin, + concurrency, + }) + } + + #[must_use] + pub const fn margin(self) -> Duration { + self.margin + } + + #[must_use] + pub const fn concurrency(self) -> NonZeroU32 { + self.concurrency + } + + /// 判断 AT 是否已经进入当前配置的提前刷新窗口 + /// + /// `margin` 是运行时策略，不投影为账号的持久化时间字段；已过期 AT 也需要 + /// 尝试用 RT 恢复，因此视为到期 + #[must_use] + pub fn is_refresh_due( + self, + access_token_expires_at: SystemTime, + observed_at: SystemTime, + ) -> bool { + match access_token_expires_at.duration_since(observed_at) { + Ok(remaining) => remaining <= self.margin, + Err(_) => true, + } + } + + /// 按账号派生 `[0, margin]` 内的整秒错峰偏移。 + /// + /// 网关以固定周期扫描到期账号，到期时刻相同的账号会在同一轮齐刷， + /// 在 auth.openai.com 侧形成单 IP 批量刷新特征。偏移从账号 ID 的 + /// 稳定哈希派生（与退避扰动同源、独立 salt），跨扫描与重启保持不变； + /// 只会增大有效提前量，不会比配置的 margin 更贴近过期时刻。 + #[must_use] + pub fn refresh_stagger(self, account_id: &ProviderAccountId) -> Duration { + // margin 以整秒配置；秒级粒度在默认 300s 下对应约 10 个扫描槽位。 + let bound = u32::try_from(self.margin.as_secs()).unwrap_or(u32::MAX); + Duration::from_secs(u64::from(stable_factor( + account_id.as_str(), + "refresh-stagger", + 0, + bound, + ))) + } + + /// 含账号错峰偏移的到期判定；扫描路径先按 [`Self::staggered_refresh_bound`] + /// 取回候选超集，再用它在内存中收窄到本轮真正到期的账号。 + #[must_use] + pub fn is_refresh_due_staggered( + self, + account_id: &ProviderAccountId, + access_token_expires_at: SystemTime, + observed_at: SystemTime, + ) -> bool { + let staggered_margin = self.margin.saturating_add(self.refresh_stagger(account_id)); + match access_token_expires_at.duration_since(observed_at) { + Ok(remaining) => remaining <= staggered_margin, + Err(_) => true, + } + } + + /// 错峰候选窗口上界；覆盖 margin 与最大偏移之和。 + /// + /// 与 [`Self::refresh_stagger`] 的值域同址维护：偏移上限为 margin， + /// 因此 `2 × margin` 必然覆盖最大有效提前量。 + #[must_use] + pub fn staggered_refresh_bound(self) -> Duration { + self.margin.saturating_mul(2) + } +} + +/// 指数退避基准延迟；attempt=1 即为该值 +const REFRESH_BACKOFF_BASE_DELAY: Duration = Duration::from_secs(5); +/// 每多一次连续失败，基准延迟乘以该因子 +const REFRESH_BACKOFF_FACTOR: u32 = 3; +/// 退避延迟上限，避免连续失败时无限增长 +const REFRESH_BACKOFF_CAP: Duration = Duration::from_secs(300); + +/// 基于连续失败计数的指数退避重试时刻；复用 `provider_refresh_retry_at` 的稳定扰动 +/// +/// `attempt` 为窗口内累计失败次数（0 与 1 等价，均取基准延迟） +/// 延迟按 +/// `base * factor^(attempt-1)` 增长并封顶到 `REFRESH_BACKOFF_CAP` +pub fn provider_refresh_backoff_at( + account_id: &ProviderAccountId, + observed_at: SystemTime, + attempt: u32, + reason: &'static str, +) -> Result { + let exponent = attempt.saturating_sub(1); + let multiplier = REFRESH_BACKOFF_FACTOR.saturating_pow(exponent); + let scaled_seconds = REFRESH_BACKOFF_BASE_DELAY + .as_secs() + .saturating_mul(u64::from(multiplier)) + .min(REFRESH_BACKOFF_CAP.as_secs()); + let base_delay = Duration::from_secs(scaled_seconds); + provider_refresh_retry_at(account_id, observed_at, base_delay, reason) +} + +/// 临时失败后的持久重试时刻；稳定扰动避免多实例同频重试 +pub fn provider_refresh_retry_at( + account_id: &ProviderAccountId, + observed_at: SystemTime, + base_delay: Duration, + reason: &'static str, +) -> Result { + if base_delay.is_zero() { + return Err(invalid_refresh_policy("schedule refresh retry")); + } + let factor = stable_factor(account_id.as_str(), reason, 800, 1_200); + let millis = u64::try_from(base_delay.as_millis()) + .unwrap_or(u64::MAX) + .saturating_mul(u64::from(factor)) + .saturating_add(500) + / 1_000; + observed_at + .checked_add(Duration::from_millis(millis.max(1))) + .ok_or_else(|| invalid_refresh_policy("schedule refresh retry")) +} + +fn stable_factor(value: &str, salt: &str, minimum: u32, maximum: u32) -> u32 { + let mut hash = 0xcbf2_9ce4_8422_2325_u64; + for byte in value.bytes().chain([0]).chain(salt.bytes()) { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + let width = u64::from(maximum - minimum) + 1; + minimum + u32::try_from(hash % width).unwrap_or_default() +} + +fn invalid_refresh_policy(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) +} diff --git a/backend/crates/gateway-core/src/provider_ports/runtime_policy.rs b/backend/crates/gateway-core/src/provider_ports/runtime_policy.rs new file mode 100644 index 000000000..f4d8255e5 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/runtime_policy.rs @@ -0,0 +1,315 @@ +//! Provider 运行策略端口及冻结、预热配置 + +use super::{ProviderRefreshPolicy, ProviderStoreError, ProviderStoreErrorKind}; +use crate::{account::OpaqueProviderData, identity::ProviderKind, routing::ConfigRevision}; +use futures::future::BoxFuture; +use std::time::Duration; + +pub trait ProviderRuntimePolicyPort: Send + Sync { + /// 原子推进预热执行游标，重启或时钟回拨后不重复领取已消费的时刻 + fn claim_warmup_slot<'a>( + &'a self, + _timezone: crate::time::DeploymentTimeZone, + _slot: chrono::NaiveDateTime, + ) -> BoxFuture<'a, Result> { + Box::pin(async { + Err(ProviderStoreError::new( + ProviderStoreErrorKind::Unavailable, + "claim warmup slot", + )) + }) + } + /// 仅首次启动写入该 Provider 的默认选择；已保存的管理配置始终优先 + fn initialize_request_profile<'a>( + &'a self, + _provider: &'a ProviderKind, + initial: OpaqueProviderData, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { Ok(initial) }) + } + + /// 读取候选配置版本实际引用的全局与 Client Key 画像配置 + /// + /// 实现必须在同一数据库快照内核对 revision，且只返回画像投影，不能读取 Key + /// 明文 + /// Provider 代次据此在发布前拒绝已失效的选择 + fn load_request_profile_configurations<'a>( + &'a self, + _revision: ConfigRevision, + _provider: &'a ProviderKind, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { Ok(Vec::new()) }) + } + + fn load_refresh_policy( + &self, + ) -> BoxFuture<'_, Result>; + + /// 读取账号容量熔断策略；默认关闭，只有实现运行时设置的存储需要覆盖 + fn load_freeze_policy( + &self, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(ProviderFreezePolicy::disabled()) }) + } + + /// 读取账号模型预激活策略；默认关闭，只有实现运行时设置的存储需要覆盖 + fn load_warmup_policy( + &self, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(ProviderWarmupPolicy::disabled()) }) + } +} + +/// 账号容量熔断（自动冻结）策略；来源于 `runtime_settings`， +/// 由 Provider 触发路径与恢复 worker 共享同一份配置事实 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderFreezePolicy { + enabled: bool, + threshold: u32, + window: Duration, + freeze_duration: Duration, + probe_enabled: bool, + probe_model: Option, + adaptive_concurrency: bool, +} + +impl ProviderFreezePolicy { + /// 管理写入与运行策略构造共用的冻结约束 + pub fn validate_values( + threshold: u32, + window_seconds: u64, + duration_seconds: u64, + probe_model: Option<&str>, + ) -> Result<(), &'static str> { + for (valid, field) in [ + ( + (2..=1_000).contains(&threshold), + "account_auto_freeze_threshold", + ), + ( + (60..=3_600).contains(&window_seconds), + "account_auto_freeze_window_seconds", + ), + ( + (300..=604_800).contains(&duration_seconds), + "account_auto_freeze_duration_seconds", + ), + ( + valid_optional_probe_model(probe_model), + "account_auto_freeze_probe_model", + ), + ] { + if !valid { + return Err(field); + } + } + Ok(()) + } + + /// 边界与迁移 `0010_account_auto_freeze.sql` 的 check 约束一致； + /// store 层写入前已校验，这里兜底防御越界配置 + pub fn try_new( + enabled: bool, + threshold: u32, + window_seconds: u64, + freeze_duration_seconds: u64, + probe_enabled: bool, + probe_model: Option, + adaptive_concurrency: bool, + ) -> Result { + if Self::validate_values( + threshold, + window_seconds, + freeze_duration_seconds, + probe_model.as_deref(), + ) + .is_err() + { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate freeze policy", + )); + } + Ok(Self { + enabled, + threshold, + window: Duration::from_secs(window_seconds), + freeze_duration: Duration::from_secs(freeze_duration_seconds), + probe_enabled, + probe_model, + adaptive_concurrency, + }) + } + + /// 功能关闭时的全零策略；触发路径与 worker 都以此短路 + #[must_use] + pub const fn disabled() -> Self { + Self { + enabled: false, + threshold: 2, + window: Duration::from_secs(60), + freeze_duration: Duration::from_secs(300), + probe_enabled: false, + probe_model: None, + adaptive_concurrency: false, + } + } + + #[must_use] + pub const fn enabled(&self) -> bool { + self.enabled + } + + /// 窗口内触发冻结的请求级失败次数阈值 + #[must_use] + pub const fn threshold(&self) -> u32 { + self.threshold + } + + /// 失败计数滑动窗口；每次失败都会顺延窗口 + #[must_use] + pub const fn window(&self) -> Duration { + self.window + } + + /// 冻结时长；探测失败后的顺延也使用该值 + #[must_use] + pub const fn freeze_duration(&self) -> Duration { + self.freeze_duration + } + + #[must_use] + pub const fn probe_enabled(&self) -> bool { + self.probe_enabled + } + + /// 探测模型；`None` 表示由 worker 选择账号可用的第一个模型 + #[must_use] + pub fn probe_model(&self) -> Option<&str> { + self.probe_model.as_deref() + } + + #[must_use] + pub const fn adaptive_concurrency(&self) -> bool { + self.adaptive_concurrency + } +} + +/// 校验每日预激活时间格式，如 "08:00" 或 "08:00,13:00" +#[must_use] +pub fn valid_warmup_schedule_time(value: &str) -> bool { + if value.is_empty() + || value.len() > 255 + || value != value.trim() + || value.chars().any(char::is_control) + { + return false; + } + for part in value.split(',') { + let bytes = part.as_bytes(); + if bytes.len() != 5 || bytes[2] != b':' { + return false; + } + let Ok(hour) = part[0..2].parse::() else { + return false; + }; + let Ok(minute) = part[3..5].parse::() else { + return false; + }; + if hour > 23 || minute > 59 { + return false; + } + } + true +} + +/// 账号模型预激活（预热）策略；来源于 `runtime_settings` +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderWarmupPolicy { + enabled: bool, + schedule_time: String, + model: Option, +} + +impl ProviderWarmupPolicy { + /// 保存设置与创建预热策略使用相同的时间表和模型约束 + pub fn validate_values( + enabled: bool, + schedule_time: &str, + model: Option<&str>, + ) -> Result<(), &'static str> { + if !valid_warmup_schedule_time(schedule_time) { + return Err("account_warmup_schedule_time"); + } + if (enabled && model.is_none()) || !valid_optional_probe_model(model) { + return Err("account_warmup_model"); + } + Ok(()) + } + + pub fn try_new( + enabled: bool, + schedule_time: String, + model: Option, + ) -> Result { + if Self::validate_values(enabled, &schedule_time, model.as_deref()).is_err() { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate warmup policy", + )); + } + Ok(Self { + enabled, + schedule_time, + model, + }) + } + + #[must_use] + pub const fn disabled() -> Self { + Self { + enabled: false, + schedule_time: String::new(), + model: None, + } + } + + #[must_use] + pub const fn enabled(&self) -> bool { + self.enabled + } + + #[must_use] + pub fn schedule_time(&self) -> &str { + &self.schedule_time + } + + #[must_use] + pub fn model(&self) -> Option<&str> { + self.model.as_deref() + } + + /// 解析每日时间列表，返回如 `vec![(8, 0)]` + #[must_use] + pub fn scheduled_times(&self) -> Vec<(u32, u32)> { + self.schedule_time + .split(',') + .filter_map(|part| { + let part = part.trim(); + let mut iter = part.split(':'); + let hour = iter.next()?.parse::().ok()?; + let minute = iter.next()?.parse::().ok()?; + Some((hour, minute)) + }) + .collect() + } +} + +fn valid_optional_probe_model(model: Option<&str>) -> bool { + model.is_none_or(|model| { + !model.is_empty() + && model.len() <= 128 + && model == model.trim() + && !model.bytes().any(|byte| byte.is_ascii_control()) + }) +} diff --git a/backend/crates/gateway-core/src/runtime/extensions.rs b/backend/crates/gateway-core/src/routing/extensions.rs similarity index 82% rename from backend/crates/gateway-core/src/runtime/extensions.rs rename to backend/crates/gateway-core/src/routing/extensions.rs index 059cd26cd..e44169104 100644 --- a/backend/crates/gateway-core/src/runtime/extensions.rs +++ b/backend/crates/gateway-core/src/routing/extensions.rs @@ -32,13 +32,6 @@ impl ExtensionSetId { pub trait ExtensionSetLease: Send + Sync { fn is_ready(&self) -> bool; - /// 与发布集合一起冻结的上游适配计划，不能在请求中解析为新代次 - fn upstream_adapters( - &self, - ) -> Option> { - None - } - /// 与集合一起冻结的目录事实，读取不调用插件进程 fn model_aliases(&self) -> &[ContributedModelAlias] { &[] @@ -57,15 +50,6 @@ pub struct ExtensionSetReference { } impl ExtensionSetReference { - #[must_use] - pub fn upstream_adapters( - &self, - ) -> Option { - self.lease.upstream_adapters().map(|plan| { - crate::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new(plan, self.clone()) - }) - } - #[must_use] pub fn new(id: ExtensionSetId, lease: Arc) -> Self { Self { id, lease } diff --git a/backend/crates/gateway-core/src/routing/mod.rs b/backend/crates/gateway-core/src/routing/mod.rs index 2118b334d..532b212b5 100644 --- a/backend/crates/gateway-core/src/routing/mod.rs +++ b/backend/crates/gateway-core/src/routing/mod.rs @@ -1,6 +1,8 @@ //! Provider、模型目录、精确模型映射与请求级候选计划 mod catalog; +pub mod extensions; +pub mod request_settings; pub mod snapshot; pub use crate::account::scope::{ diff --git a/backend/crates/gateway-core/src/routing/request_settings.rs b/backend/crates/gateway-core/src/routing/request_settings.rs new file mode 100644 index 000000000..bb4837b57 --- /dev/null +++ b/backend/crates/gateway-core/src/routing/request_settings.rs @@ -0,0 +1,374 @@ +//! 冻结宿主基线与请求显式改写；派生快照只属于当前调用，不发布全局配置 + +use std::{ + collections::BTreeMap, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::account::FastMode; +use crate::policy::{ClientApiKeyId, ClientPolicy, ClientSettings, RateLimits}; +use crate::routing::RuntimeSnapshot; +use crate::settings::{InvalidSettings, SettingsValues}; + +/// 一次模型调用的有效设置；改写只影响当前请求，不发布配置或修改持久化 revision +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ExecutionSettings { + pub runtime: SettingsValues, + pub fast_mode: FastMode, + pub client_limits: RateLimits, + /// null 表示不限制总时长；显式时限从请求开始计时，改写不重置计时原点 + pub timeout_ms: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SettingOverride { + pub instance_id: String, + pub order: u64, + pub value: T, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +struct ExecutionOverrides { + #[serde(skip)] + defaults: Arc, + client_key_id: String, + input: Arc, + fast_mode: Option>, + client_limits: Option>, + timeout_ms: Option>>, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +struct HttpTimeout { + input_ms: Option, + change: Option>>, +} + +/// 共享值均不可变；同级子调用的修改不会回写父调用 +#[derive(Clone, Debug)] +pub struct RequestSettings { + baseline: Arc, + effective: Arc, + overrides: Arc>, + execution: Option>, + http_timeout: Option, + order: u64, +} + +impl PartialEq for RequestSettings { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.baseline, &other.baseline) + && self.overrides == other.overrides + && self.execution == other.execution + && self.http_timeout == other.http_timeout + } +} +impl Eq for RequestSettings {} + +impl RequestSettings { + #[must_use] + pub fn new(snapshot: Arc) -> Self { + Self { + baseline: snapshot.clone(), + effective: snapshot, + overrides: Arc::default(), + execution: None, + http_timeout: None, + order: 0, + } + } + + #[must_use] + pub fn snapshot(&self) -> Arc { + self.effective.clone() + } + + #[must_use] + pub fn values(&self) -> &SettingsValues { + self.effective.settings() + } + + /// 来源与有效值分开查询，不允许插件伪造其他实例的写入来源 + pub fn inspect(&self) -> Value { + serde_json::json!({ + "config_revision": self.baseline.revision().get(), + "host": self.baseline.settings(), + "overrides": self.overrides, + "execution": self.execution, + "http_timeout": self.http_timeout, + }) + } + + pub fn replace( + &self, + values: SettingsValues, + instance_id: &str, + ) -> Result { + self.replace_scoped(self.values(), values, instance_id) + } + + /// Key 画像等隐式作用域结果不记录成插件改写，避免随后切换 Key 时泄漏 + fn replace_scoped( + &self, + previous: &SettingsValues, + values: SettingsValues, + instance_id: &str, + ) -> Result { + if previous == &values { + return Ok(self.clone()); + } + let current = serde_json::to_value(previous).map_err(|_| InvalidSettings)?; + let next = serde_json::to_value(&values).map_err(|_| InvalidSettings)?; + let order = self.order.checked_add(1).ok_or(InvalidSettings)?; + let mut overrides = self.overrides.as_ref().clone(); + let effective_values = self.values(); + // 每个已声明设置项整体替换，不猜测画像对象、数组或 null 的深合并含义 + for (name, value) in next.as_object().ok_or(InvalidSettings)? { + if current.get(name) != Some(value) { + overrides.insert( + name.clone(), + SettingOverride { + instance_id: instance_id.to_owned(), + order, + value: value.clone(), + }, + ); + } + } + // 只把本层实际修改的字段写回宿主层，作用域默认值不参与继承 + let values = if previous == effective_values { + values + } else { + apply_overrides(self.baseline.settings(), &overrides)? + }; + let effective = self + .effective + .resolve_settings(&values) + .map_err(|_| InvalidSettings)?; + Ok(Self { + baseline: self.baseline.clone(), + effective, + overrides: Arc::new(overrides), + execution: self.execution.clone(), + http_timeout: self.http_timeout.clone(), + order, + }) + } + + /// 持续会话的新请求读取当前宿主快照，只继承明确改写的字段 + pub fn rebase(&self, snapshot: Arc) -> Result { + if Arc::ptr_eq(&self.baseline, &snapshot) { + return Ok(self.clone()); + } + let effective = if self.overrides.is_empty() { + snapshot.clone() + } else { + snapshot + .resolve_settings(&apply_overrides(snapshot.settings(), &self.overrides)?) + .map_err(|_| InvalidSettings)? + }; + Ok(Self { + baseline: snapshot, + effective, + overrides: self.overrides.clone(), + execution: self.execution.clone(), + http_timeout: self.http_timeout.clone(), + order: self.order, + }) + } + + /// 认证、模型入口和插件视图共用解析规则，派生策略不改变 Key 默认值 + #[must_use] + pub fn apply_policy(&self, policy: ClientPolicy) -> ClientPolicy { + let values = self.resolve_execution(policy.key_id().as_str(), policy.defaults(), None); + policy.with_settings( + values.runtime.request_profiles(), + values.fast_mode, + values.client_limits, + ) + } + + #[must_use] + pub fn with_execution(mut self, policy: &ClientPolicy, timeout_ms: Option) -> Self { + let input = self.resolve_execution(policy.key_id().as_str(), policy.defaults(), timeout_ms); + let previous = self + .execution + .as_ref() + .filter(|scope| scope.client_key_id == policy.key_id().as_str()); + self.execution = Some(Arc::new(ExecutionOverrides { + defaults: policy.defaults().clone(), + client_key_id: policy.key_id().as_str().to_owned(), + input: Arc::new(input), + fast_mode: previous.and_then(|scope| scope.fast_mode.clone()), + client_limits: previous.and_then(|scope| scope.client_limits.clone()), + timeout_ms: previous.and_then(|scope| scope.timeout_ms.clone()), + })); + self + } + + #[must_use] + pub fn execution_values(&self) -> Option { + let scope = self.execution.as_ref()?; + Some(self.resolve_execution( + &scope.client_key_id, + &scope.defaults, + scope.input.timeout_ms, + )) + } + + fn resolve_execution( + &self, + key: &str, + defaults: &ClientSettings, + timeout_ms: Option, + ) -> ExecutionSettings { + let overrides = self + .execution + .as_ref() + .filter(|scope| scope.client_key_id == key); + let mut runtime = self.values().clone(); + // Key 画像覆盖宿主默认，插件显式覆盖整个设置项；此优先级只在这里解释 + if !self.overrides.contains_key("request_profiles") && !defaults.request_profiles.is_empty() + { + let mut profiles = runtime.request_profiles().clone(); + profiles.extend(defaults.request_profiles.clone()); + runtime = runtime.with_request_profiles(profiles); + } + ExecutionSettings { + runtime, + fast_mode: overrides + .and_then(|scope| scope.fast_mode.as_ref()) + .map_or(defaults.fast_mode, |change| change.value), + client_limits: overrides + .and_then(|scope| scope.client_limits.as_ref()) + .map_or(defaults.limits, |change| change.value), + timeout_ms: overrides + .and_then(|scope| scope.timeout_ms.as_ref()) + .map_or(timeout_ms, |change| change.value), + } + } + + pub(crate) fn execution_deadline( + &self, + key: &ClientApiKeyId, + started_at: SystemTime, + ) -> Result { + let scope = self + .execution + .as_ref() + .filter(|scope| scope.client_key_id == key.as_str()) + .ok_or(InvalidSettings)?; + let timeout_ms = scope + .timeout_ms + .as_ref() + .map_or(scope.input.timeout_ms, |change| change.value); + crate::lifecycle::Deadline::from_timeout(started_at, timeout_ms.map(Duration::from_millis)) + .ok_or(InvalidSettings) + } + + pub fn replace_execution( + &self, + values: &ExecutionSettings, + instance_id: &str, + ) -> Result { + let previous = self.execution_values().ok_or(InvalidSettings)?; + let mut updated = + self.replace_scoped(&previous.runtime, values.runtime.clone(), instance_id)?; + if previous.fast_mode == values.fast_mode + && previous.client_limits == values.client_limits + && previous.timeout_ms == values.timeout_ms + { + return Ok(updated); + } + let mut scope = self.execution.as_deref().cloned().ok_or(InvalidSettings)?; + let order = if updated.order == self.order { + self.order.checked_add(1).ok_or(InvalidSettings)? + } else { + updated.order + }; + if previous.fast_mode != values.fast_mode { + scope.fast_mode = Some(SettingOverride { + instance_id: instance_id.to_owned(), + order, + value: values.fast_mode, + }); + } + if previous.client_limits != values.client_limits { + scope.client_limits = Some(SettingOverride { + instance_id: instance_id.to_owned(), + order, + value: values.client_limits, + }); + } + if previous.timeout_ms != values.timeout_ms { + scope.timeout_ms = Some(SettingOverride { + instance_id: instance_id.to_owned(), + order, + value: values.timeout_ms, + }); + } + updated.execution = Some(Arc::new(scope)); + updated.order = order; + Ok(updated) + } + + #[must_use] + pub fn execution_timeout(&self, key: &ClientApiKeyId) -> Option { + self.execution + .as_ref() + .filter(|scope| scope.client_key_id == key.as_str()) + .and_then(|scope| scope.timeout_ms.as_ref()) + .and_then(|change| change.value.map(Duration::from_millis)) + } + + #[must_use] + pub fn with_http_timeout(mut self, timeout_ms: Option) -> Self { + self.http_timeout.get_or_insert(HttpTimeout { + input_ms: timeout_ms, + change: None, + }); + self + } + + pub fn replace_http_timeout( + mut self, + timeout_ms: Option, + instance_id: &str, + ) -> Result { + let timeout = self.http_timeout.get_or_insert(HttpTimeout { + input_ms: None, + change: None, + }); + if timeout + .change + .as_ref() + .map_or(timeout.input_ms, |change| change.value) + != timeout_ms + { + self.order = self.order.checked_add(1).ok_or(InvalidSettings)?; + timeout.change = Some(SettingOverride { + instance_id: instance_id.to_owned(), + order: self.order, + value: timeout_ms, + }); + } + Ok(self) + } +} + +fn apply_overrides( + values: &SettingsValues, + overrides: &BTreeMap, +) -> Result { + let mut values = serde_json::to_value(values).map_err(|_| InvalidSettings)?; + let object = values.as_object_mut().ok_or(InvalidSettings)?; + for (name, change) in overrides { + object.insert(name.clone(), change.value.clone()); + } + serde_json::from_value(values).map_err(|_| InvalidSettings) +} diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs index d059f3915..c5dd3fe50 100644 --- a/backend/crates/gateway-core/src/routing/snapshot.rs +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -220,7 +220,7 @@ pub enum RuntimeSnapshotCompileError { pub struct RuntimeSnapshotCompiler { store: Arc, catalogs: Arc, - extensions: Option>, + extensions: Option>, } impl RuntimeSnapshotCompiler { @@ -239,7 +239,7 @@ impl RuntimeSnapshotCompiler { #[must_use] pub fn with_extensions( mut self, - extensions: Arc, + extensions: Arc, ) -> Self { self.extensions = Some(extensions); self @@ -275,7 +275,7 @@ impl RuntimeSnapshotCompiler { // 同一配置 revision 的重试复用扩展候选，保留插件策略的发布一致性 let mut prepared_extensions: Option<( ConfigRevision, - crate::runtime::extensions::ExtensionSetReference, + crate::routing::extensions::ExtensionSetReference, )> = None; for _ in 0..MAXIMUM_CATALOG_STABILITY_ATTEMPTS { let facts = self @@ -524,7 +524,7 @@ async fn compile_runtime_snapshot( #[derive(Debug, Clone)] pub struct RuntimeSnapshot { settings: Arc, - extensions: Option, + extensions: Option, revision: ConfigRevision, providers: Arc>, provider_models: Arc>>, @@ -570,14 +570,14 @@ impl RuntimeSnapshot { #[must_use] pub fn with_extensions( mut self, - extensions: Option, + extensions: Option, ) -> Self { self.extensions = extensions; self } #[must_use] - pub const fn extensions(&self) -> Option<&crate::runtime::extensions::ExtensionSetReference> { + pub const fn extensions(&self) -> Option<&crate::routing::extensions::ExtensionSetReference> { self.extensions.as_ref() } diff --git a/backend/crates/gateway-core/src/runtime/mod.rs b/backend/crates/gateway-core/src/runtime/mod.rs index 5391cb6f7..d43baa5ce 100644 --- a/backend/crates/gateway-core/src/runtime/mod.rs +++ b/backend/crates/gateway-core/src/runtime/mod.rs @@ -1,542 +1,10 @@ //! 运行时快照的原子发布、健康状态与跨进程版本收敛 -pub mod extensions; +mod publications; +mod publisher; -use std::collections::BTreeMap; -use std::future::Future; -use std::pin::Pin; -use std::sync::{Arc, Mutex as SyncMutex, RwLock}; -use std::time::Duration; - -use futures::future::BoxFuture; -use futures::lock::Mutex; -use futures::{FutureExt as _, Stream, StreamExt as _, pin_mut, select_biased}; -use futures_timer::Delay; - -use crate::health::{HealthProbe, HealthState}; -use crate::identity::ProviderKind; -use crate::lifecycle::CancellationToken; -use crate::routing::snapshot::{ - RuntimeSnapshot, RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, +pub use publications::{RuntimeSnapshotHandle, RuntimeSnapshotUnavailable}; +pub use publisher::{ + RuntimeSnapshotPublisher, SnapshotControl, SnapshotRevisionStream, SnapshotSubscriptionError, + SnapshotSubscriptionPort, runtime_revision_needs_refresh, }; -use crate::routing::{ConfigRevision, ProviderCatalogGeneration}; -use crate::task::{ - DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, - WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable, - WorkerSchedule, WorkerTaskError, -}; - -const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5); -const INITIAL_BACKOFF: Duration = Duration::from_secs(1); -const MAXIMUM_BACKOFF: Duration = Duration::from_secs(30); -const UNUSED_LEASE_TTL: Duration = Duration::from_secs(30); -const UNUSED_LEASE_RENEWAL: Duration = Duration::from_secs(10); - -/// 不泄漏订阅基础设施细节的通知错误 -#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] -#[error("runtime snapshot notification is unavailable")] -pub struct SnapshotSubscriptionError; - -impl SnapshotSubscriptionError { - #[must_use] - pub const fn unavailable() -> Self { - Self - } -} - -/// 可丢失的配置 revision 通知流；权威 revision 始终由 Store 端口读取 -pub type SnapshotRevisionStream = - Pin> + Send + 'static>>; - -/// 跨进程 revision 通知的基础设施中立端口 -pub trait SnapshotSubscriptionPort: Send + Sync { - fn publish_snapshot_revision( - &self, - revision: ConfigRevision, - ) -> BoxFuture<'_, Result<(), SnapshotSubscriptionError>>; - - fn subscribe_snapshot_revisions( - &self, - ) -> BoxFuture<'_, Result>; -} - -/// 请求级冻结失败；此状态必须 fail closed -#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] -#[error("runtime snapshot is unavailable")] -pub struct RuntimeSnapshotUnavailable; - -/// RuntimeSnapshot 原子发布和请求级冻结句柄 -#[derive(Clone, Default)] -pub struct RuntimeSnapshotHandle { - current: Arc>>>, - publications: Arc>>>, -} - -impl RuntimeSnapshotHandle { - #[must_use] - pub fn new(initial: RuntimeSnapshot) -> Self { - Self { - current: Arc::new(RwLock::new(Some(Arc::new(initial)))), - publications: Default::default(), - } - } - - pub fn publish(&self, snapshot: RuntimeSnapshot) { - *write_unpoisoned(&self.current) = Some(Arc::new(snapshot)); - self.notify_publication(); - } - - pub fn suspend(&self) { - *write_unpoisoned(&self.current) = None; - self.notify_publication(); - } - - /// 合并发布与暂停通知；订阅者收到后读取当前快照，不依赖逐条投递 - pub fn subscribe_publications(&self) -> futures::channel::mpsc::Receiver<()> { - let (sender, receiver) = futures::channel::mpsc::channel(0); - self.publications - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(sender); - receiver - } - - fn notify_publication(&self) { - self.publications - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .retain_mut(|sender| match sender.try_send(()) { - Ok(()) => true, - Err(error) => error.is_full(), - }); - } - - #[must_use] - pub fn revision(&self) -> Option { - read_unpoisoned(&self.current) - .as_ref() - .map(|snapshot| snapshot.revision()) - } - - #[must_use] - pub fn provider_catalog_generations( - &self, - ) -> Option> { - read_unpoisoned(&self.current) - .as_ref() - .map(|snapshot| snapshot.provider_catalog_generations().clone()) - } - - /// 控制面只读诊断可观察尚未就绪的发布候选；数据面仍必须通过 [`Self::acquire`] - #[must_use] - pub fn snapshot_for_diagnostics(&self) -> Option> { - read_unpoisoned(&self.current).clone() - } - - /// 冻结当前 Arc；后续发布不改变已经开始的请求 - pub fn acquire(&self) -> Result, RuntimeSnapshotUnavailable> { - read_unpoisoned(&self.current) - .clone() - .filter(|snapshot| snapshot.extensions().is_none_or(|set| set.can_serve())) - .ok_or(RuntimeSnapshotUnavailable) - } -} - -impl HealthProbe for RuntimeSnapshotHandle { - fn name(&self) -> &'static str { - "runtime_snapshot" - } - - fn check(&self) -> BoxFuture<'_, HealthState> { - Box::pin(async move { - if self.acquire().is_ok() { - HealthState::Healthy - } else { - HealthState::Unhealthy("Runtime snapshot is unavailable".to_owned()) - } - }) - } -} - -/// Admin 提交配置后触发本进程刷新与跨进程通知的对象安全端口 -pub trait SnapshotControl: Send + Sync { - fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()>; -} - -#[derive(Default)] -struct RefreshPriorityState { - pending_commits: usize, - active_background: Option, -} - -struct CommittedRefreshGuard { - priority: Arc>, -} - -impl CommittedRefreshGuard { - fn begin(priority: Arc>) -> Self { - let active = { - let mut state = lock_unpoisoned(&priority); - state.pending_commits += 1; - state.active_background.clone() - }; - // 先登记待发布提交，再取消正在等待目录的后台刷新，避免后者抢先重入 - if let Some(active) = active { - active.cancel(); - } - Self { priority } - } -} - -impl Drop for CommittedRefreshGuard { - fn drop(&mut self) { - lock_unpoisoned(&self.priority).pending_commits -= 1; - } -} - -struct BackgroundRefreshGuard { - priority: Arc>, - cancellation: CancellationToken, -} - -impl BackgroundRefreshGuard { - fn begin(priority: Arc>) -> Option { - let mut state = lock_unpoisoned(&priority); - if state.pending_commits != 0 { - return None; - } - let cancellation = CancellationToken::new(); - state.active_background = Some(cancellation.clone()); - drop(state); - Some(Self { - priority, - cancellation, - }) - } -} - -impl Drop for BackgroundRefreshGuard { - fn drop(&mut self) { - lock_unpoisoned(&self.priority).active_background = None; - } -} - -// 仅中断未完成的后台读取或编译；配置提交负责发布最新事实并发送通知 -async fn preempt_for_committed_refresh( - cancellation: Option<&CancellationToken>, - operation: impl Future, -) -> Result { - let Some(cancellation) = cancellation else { - return Ok(operation.await); - }; - let cancelled = cancellation.cancelled().fuse(); - let operation = operation.fuse(); - pin_mut!(cancelled, operation); - select_biased! { - _ = cancelled => Err(RuntimeSnapshotCompileError::RevisionChanged), - result = operation => Ok(result), - } -} - -/// 配置提交后的本进程快照发布与跨进程失效通知 -#[derive(Clone)] -pub struct RuntimeSnapshotPublisher { - compiler: Arc, - snapshots: RuntimeSnapshotHandle, - subscriptions: Arc, - refresh_lock: Arc>, - refresh_priority: Arc>, -} - -enum RefreshMode { - Required, - Reconcile, - Committed, -} - -impl RuntimeSnapshotPublisher { - #[must_use] - pub fn new( - compiler: Arc, - snapshots: RuntimeSnapshotHandle, - subscriptions: Arc, - ) -> Self { - Self { - compiler, - snapshots, - subscriptions, - refresh_lock: Arc::new(Mutex::new(())), - refresh_priority: Arc::new(SyncMutex::new(RefreshPriorityState::default())), - } - } - - /// 串行重编译并替换本进程快照；无法确认配置时暂停新请求 - pub async fn refresh(&self) -> Result { - self.refresh_with_mode(RefreshMode::Required).await - } - - async fn refresh_with_mode( - &self, - mode: RefreshMode, - ) -> Result { - // 提交先登记并中断慢后台读取；发布/暂停仍在完整临界区内串行执行 - let _committed = matches!(mode, RefreshMode::Committed) - .then(|| CommittedRefreshGuard::begin(Arc::clone(&self.refresh_priority))); - let _refresh = self.refresh_lock.lock().await; - self.refresh_locked_with_priority(mode).await - } - - async fn refresh_locked_with_priority( - &self, - mode: RefreshMode, - ) -> Result { - if matches!(mode, RefreshMode::Committed) { - return self.refresh_locked(mode, None).await; - } - let background = BackgroundRefreshGuard::begin(Arc::clone(&self.refresh_priority)) - .ok_or(RuntimeSnapshotCompileError::RevisionChanged)?; - self.refresh_locked(mode, Some(&background.cancellation)) - .await - } - - async fn refresh_locked( - &self, - mode: RefreshMode, - cancellation: Option<&CancellationToken>, - ) -> Result { - let configuration_changed = if matches!(mode, RefreshMode::Reconcile) { - let persisted_revision = preempt_for_committed_refresh( - cancellation, - self.compiler.store().current_config_revision(), - ) - .await? - .map_err(|_| { - self.snapshots.suspend(); - RuntimeSnapshotCompileError::StoreUnavailable - })?; - let changed = runtime_revision_needs_refresh( - self.published_revision().map(ConfigRevision::get), - persisted_revision.get(), - ); - if !changed - && !self.provider_catalogs_need_refresh() - && self - .snapshots - .acquire() - .is_ok_and(|snapshot| snapshot.extensions().is_none_or(|set| set.is_ready())) - { - return Ok(persisted_revision); - } - changed - } else { - true - }; - let compiled = if matches!(mode, RefreshMode::Committed) { - match self.snapshots.snapshot_for_diagnostics() { - Some(previous) => self.compiler.compile_with_cached_catalog(&previous).await, - None => self.compiler.compile().await, - } - } else { - preempt_for_committed_refresh(cancellation, self.compiler.compile()).await? - }; - let snapshot = match compiled { - Ok(snapshot) => snapshot, - Err(error) => { - // 目录对账失败可继续服务旧快照；配置或权限变化无法确认时须暂停 - if configuration_changed { - self.snapshots.suspend(); - } - return Err(error); - } - }; - let revision = snapshot.revision(); - self.snapshots.publish(snapshot); - Ok(revision) - } - - #[must_use] - pub fn published_revision(&self) -> Option { - self.snapshots.revision() - } - - #[must_use] - fn provider_catalogs_need_refresh(&self) -> bool { - let Ok(snapshot) = self.snapshots.acquire() else { - return true; - }; - self.compiler.provider_catalog_generations() != *snapshot.provider_catalog_generations() - } - - /// 数据库提交不能被目录或通知基础设施的暂时故障伪装成回滚 - async fn publish_committed_inner(&self, committed_revision: ConfigRevision) { - if let Err(error) = self.refresh_with_mode(RefreshMode::Committed).await { - tracing::warn!(?error, "committed runtime snapshot refresh failed"); - } - self.notify_committed_revision(committed_revision).await; - } - - async fn notify_committed_revision(&self, committed_revision: ConfigRevision) { - if let Err(error) = self - .subscriptions - .publish_snapshot_revision(committed_revision) - .await - { - tracing::warn!( - ?error, - revision = committed_revision.get(), - "runtime snapshot revision notification failed" - ); - } - } - - /// 交给 Host 的周期对账与长驻订阅任务 - pub fn worker_contributions(&self) -> Result, WorkerDefinitionError> { - let reconciliation_id = WorkerId::try_new( - WorkerKind::RuntimeSnapshotReconciliation, - "runtime_snapshot", - )?; - let schedule = WorkerSchedule::try_new( - RECONCILIATION_INTERVAL, - INITIAL_BACKOFF, - MAXIMUM_BACKOFF, - UNUSED_LEASE_TTL, - UNUSED_LEASE_RENEWAL, - )?; - let reconciliation = WorkerRegistration::try_new( - reconciliation_id, - WorkerRunnable::Scheduled { - schedule, - lease: None, - task: Box::new(RuntimeSnapshotReconciliationTask { - publisher: self.clone(), - }), - }, - )?; - let subscription_id = - WorkerId::try_new(WorkerKind::RuntimeChangeSubscription, "runtime_snapshot")?; - let restart = DaemonRestartPolicy::try_new(INITIAL_BACKOFF, MAXIMUM_BACKOFF)?; - let subscription = WorkerRegistration::try_new( - subscription_id, - WorkerRunnable::Daemon { - restart, - task: Box::new(RuntimeSnapshotSubscriptionTask { - subscriptions: Arc::clone(&self.subscriptions), - publisher: self.clone(), - }), - }, - )?; - Ok(vec![ - WorkerContribution::Registration(reconciliation), - WorkerContribution::Registration(subscription), - ]) - } -} - -impl SnapshotControl for RuntimeSnapshotPublisher { - fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()> { - Box::pin(async move { - self.publish_committed_inner(committed_revision).await; - }) - } -} - -struct RuntimeSnapshotReconciliationTask { - publisher: RuntimeSnapshotPublisher, -} - -impl ScheduledTask for RuntimeSnapshotReconciliationTask { - fn run_cycle( - &self, - _context: WorkerCycleContext, - ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { - Box::pin(async move { - self.publisher - .refresh_with_mode(RefreshMode::Reconcile) - .await - .map(|_| ()) - .map_err(|source| { - WorkerTaskError::safe("runtime snapshot reconciliation failed") - .with_source(source) - }) - }) - } -} - -struct RuntimeSnapshotSubscriptionTask { - subscriptions: Arc, - publisher: RuntimeSnapshotPublisher, -} - -impl DaemonTask for RuntimeSnapshotSubscriptionTask { - fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { - Box::pin(async move { - let mut retry_delay = INITIAL_BACKOFF; - loop { - if cancellation.is_cancelled() { - return Ok(()); - } - let subscription = self.subscriptions.subscribe_snapshot_revisions().await; - let mut subscription = match subscription { - Ok(subscription) => { - retry_delay = INITIAL_BACKOFF; - subscription - } - Err(_) => { - wait_or_cancel(&cancellation, retry_delay).await; - retry_delay = (retry_delay * 2).min(MAXIMUM_BACKOFF); - continue; - } - }; - loop { - let cancelled = cancellation.cancelled().fuse(); - let next = subscription.next().fuse(); - pin_mut!(cancelled, next); - let notified = select_biased! { - _ = cancelled => return Ok(()), - next = next => next, - }; - match notified { - Some(Ok(_)) => { - let _ = self.publisher.refresh().await; - } - Some(Err(_)) | None => break, - } - } - } - }) - } -} - -async fn wait_or_cancel(cancellation: &CancellationToken, duration: Duration) { - let cancelled = cancellation.cancelled().fuse(); - let delay = Delay::new(duration).fuse(); - pin_mut!(cancelled, delay); - select_biased! { - _ = cancelled => {}, - _ = delay => {}, - } -} - -/// 当前发布版本与持久版本不一致时必须重载；缺失和回退同样 fail closed -#[must_use] -pub fn runtime_revision_needs_refresh( - published_revision: Option, - persisted_revision: u64, -) -> bool { - published_revision != Some(persisted_revision) -} - -fn read_unpoisoned(lock: &RwLock) -> std::sync::RwLockReadGuard<'_, T> { - lock.read() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} - -fn write_unpoisoned(lock: &RwLock) -> std::sync::RwLockWriteGuard<'_, T> { - lock.write() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} - -fn lock_unpoisoned(lock: &SyncMutex) -> std::sync::MutexGuard<'_, T> { - lock.lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} diff --git a/backend/crates/gateway-core/src/runtime/publications.rs b/backend/crates/gateway-core/src/runtime/publications.rs new file mode 100644 index 000000000..693c68776 --- /dev/null +++ b/backend/crates/gateway-core/src/runtime/publications.rs @@ -0,0 +1,120 @@ +//! 当前路由快照的原子发布、冻结与健康投影 + +use crate::{ + health::{HealthProbe, HealthState}, + identity::ProviderKind, + routing::{ConfigRevision, ProviderCatalogGeneration, RuntimeSnapshot}, +}; +use futures::future::BoxFuture; +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex as SyncMutex, RwLock}, +}; + +/// 请求级冻结失败；此状态必须 fail closed +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("runtime snapshot is unavailable")] +pub struct RuntimeSnapshotUnavailable; + +/// RuntimeSnapshot 原子发布和请求级冻结句柄 +#[derive(Clone, Default)] +pub struct RuntimeSnapshotHandle { + current: Arc>>>, + publications: Arc>>>, +} + +impl RuntimeSnapshotHandle { + #[must_use] + pub fn new(initial: RuntimeSnapshot) -> Self { + Self { + current: Arc::new(RwLock::new(Some(Arc::new(initial)))), + publications: Default::default(), + } + } + + pub fn publish(&self, snapshot: RuntimeSnapshot) { + *write_unpoisoned(&self.current) = Some(Arc::new(snapshot)); + self.notify_publication(); + } + + pub fn suspend(&self) { + *write_unpoisoned(&self.current) = None; + self.notify_publication(); + } + + /// 合并发布与暂停通知；订阅者收到后读取当前快照，不依赖逐条投递 + pub fn subscribe_publications(&self) -> futures::channel::mpsc::Receiver<()> { + let (sender, receiver) = futures::channel::mpsc::channel(0); + self.publications + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(sender); + receiver + } + + fn notify_publication(&self) { + self.publications + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .retain_mut(|sender| match sender.try_send(()) { + Ok(()) => true, + Err(error) => error.is_full(), + }); + } + + #[must_use] + pub fn revision(&self) -> Option { + read_unpoisoned(&self.current) + .as_ref() + .map(|snapshot| snapshot.revision()) + } + + #[must_use] + pub fn provider_catalog_generations( + &self, + ) -> Option> { + read_unpoisoned(&self.current) + .as_ref() + .map(|snapshot| snapshot.provider_catalog_generations().clone()) + } + + /// 控制面只读诊断可观察尚未就绪的发布候选；数据面仍必须通过 [`Self::acquire`] + #[must_use] + pub fn snapshot_for_diagnostics(&self) -> Option> { + read_unpoisoned(&self.current).clone() + } + + /// 冻结当前 Arc；后续发布不改变已经开始的请求 + pub fn acquire(&self) -> Result, RuntimeSnapshotUnavailable> { + read_unpoisoned(&self.current) + .clone() + .filter(|snapshot| snapshot.extensions().is_none_or(|set| set.can_serve())) + .ok_or(RuntimeSnapshotUnavailable) + } +} + +impl HealthProbe for RuntimeSnapshotHandle { + fn name(&self) -> &'static str { + "runtime_snapshot" + } + + fn check(&self) -> BoxFuture<'_, HealthState> { + Box::pin(async move { + if self.acquire().is_ok() { + HealthState::Healthy + } else { + HealthState::Unhealthy("Runtime snapshot is unavailable".to_owned()) + } + }) + } +} + +fn read_unpoisoned(lock: &RwLock) -> std::sync::RwLockReadGuard<'_, T> { + lock.read() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +fn write_unpoisoned(lock: &RwLock) -> std::sync::RwLockWriteGuard<'_, T> { + lock.write() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} diff --git a/backend/crates/gateway-core/src/runtime/publisher.rs b/backend/crates/gateway-core/src/runtime/publisher.rs new file mode 100644 index 000000000..6944e96a9 --- /dev/null +++ b/backend/crates/gateway-core/src/runtime/publisher.rs @@ -0,0 +1,429 @@ +//! 运行时快照的原子发布、健康状态与跨进程版本收敛 + +use super::publications::RuntimeSnapshotHandle; + +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, Mutex as SyncMutex}; +use std::time::Duration; + +use futures::future::BoxFuture; +use futures::lock::Mutex; +use futures::{FutureExt as _, Stream, StreamExt as _, pin_mut, select_biased}; +use futures_timer::Delay; + +use crate::lifecycle::CancellationToken; +use crate::routing::ConfigRevision; +use crate::routing::snapshot::{RuntimeSnapshotCompileError, RuntimeSnapshotCompiler}; +use crate::task::{ + DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, + WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable, + WorkerSchedule, WorkerTaskError, +}; + +const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5); +const INITIAL_BACKOFF: Duration = Duration::from_secs(1); +const MAXIMUM_BACKOFF: Duration = Duration::from_secs(30); +const UNUSED_LEASE_TTL: Duration = Duration::from_secs(30); +const UNUSED_LEASE_RENEWAL: Duration = Duration::from_secs(10); + +/// 不泄漏订阅基础设施细节的通知错误 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("runtime snapshot notification is unavailable")] +pub struct SnapshotSubscriptionError; + +impl SnapshotSubscriptionError { + #[must_use] + pub const fn unavailable() -> Self { + Self + } +} + +/// 可丢失的配置 revision 通知流；权威 revision 始终由 Store 端口读取 +pub type SnapshotRevisionStream = + Pin> + Send + 'static>>; + +/// 跨进程 revision 通知的基础设施中立端口 +pub trait SnapshotSubscriptionPort: Send + Sync { + fn publish_snapshot_revision( + &self, + revision: ConfigRevision, + ) -> BoxFuture<'_, Result<(), SnapshotSubscriptionError>>; + + fn subscribe_snapshot_revisions( + &self, + ) -> BoxFuture<'_, Result>; +} + +/// Admin 提交配置后触发本进程刷新与跨进程通知的对象安全端口 +pub trait SnapshotControl: Send + Sync { + fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()>; +} + +#[derive(Default)] +struct RefreshPriorityState { + pending_commits: usize, + active_background: Option, +} + +struct CommittedRefreshGuard { + priority: Arc>, +} + +impl CommittedRefreshGuard { + fn begin(priority: Arc>) -> Self { + let active = { + let mut state = lock_unpoisoned(&priority); + state.pending_commits += 1; + state.active_background.clone() + }; + // 先登记待发布提交，再取消正在等待目录的后台刷新，避免后者抢先重入 + if let Some(active) = active { + active.cancel(); + } + Self { priority } + } +} + +impl Drop for CommittedRefreshGuard { + fn drop(&mut self) { + lock_unpoisoned(&self.priority).pending_commits -= 1; + } +} + +struct BackgroundRefreshGuard { + priority: Arc>, + cancellation: CancellationToken, +} + +impl BackgroundRefreshGuard { + fn begin(priority: Arc>) -> Option { + let mut state = lock_unpoisoned(&priority); + if state.pending_commits != 0 { + return None; + } + let cancellation = CancellationToken::new(); + state.active_background = Some(cancellation.clone()); + drop(state); + Some(Self { + priority, + cancellation, + }) + } +} + +impl Drop for BackgroundRefreshGuard { + fn drop(&mut self) { + lock_unpoisoned(&self.priority).active_background = None; + } +} + +// 仅中断未完成的后台读取或编译；配置提交负责发布最新事实并发送通知 +async fn preempt_for_committed_refresh( + cancellation: Option<&CancellationToken>, + operation: impl Future, +) -> Result { + let Some(cancellation) = cancellation else { + return Ok(operation.await); + }; + let cancelled = cancellation.cancelled().fuse(); + let operation = operation.fuse(); + pin_mut!(cancelled, operation); + select_biased! { + _ = cancelled => Err(RuntimeSnapshotCompileError::RevisionChanged), + result = operation => Ok(result), + } +} + +/// 配置提交后的本进程快照发布与跨进程失效通知 +#[derive(Clone)] +pub struct RuntimeSnapshotPublisher { + compiler: Arc, + snapshots: RuntimeSnapshotHandle, + subscriptions: Arc, + refresh_lock: Arc>, + refresh_priority: Arc>, +} + +enum RefreshMode { + Required, + Reconcile, + Committed, +} + +impl RuntimeSnapshotPublisher { + #[must_use] + pub fn new( + compiler: Arc, + snapshots: RuntimeSnapshotHandle, + subscriptions: Arc, + ) -> Self { + Self { + compiler, + snapshots, + subscriptions, + refresh_lock: Arc::new(Mutex::new(())), + refresh_priority: Arc::new(SyncMutex::new(RefreshPriorityState::default())), + } + } + + /// 串行重编译并替换本进程快照；无法确认配置时暂停新请求 + pub async fn refresh(&self) -> Result { + self.refresh_with_mode(RefreshMode::Required).await + } + + async fn refresh_with_mode( + &self, + mode: RefreshMode, + ) -> Result { + // 提交先登记并中断慢后台读取；发布/暂停仍在完整临界区内串行执行 + let _committed = matches!(mode, RefreshMode::Committed) + .then(|| CommittedRefreshGuard::begin(Arc::clone(&self.refresh_priority))); + let _refresh = self.refresh_lock.lock().await; + self.refresh_locked_with_priority(mode).await + } + + async fn refresh_locked_with_priority( + &self, + mode: RefreshMode, + ) -> Result { + if matches!(mode, RefreshMode::Committed) { + return self.refresh_locked(mode, None).await; + } + let background = BackgroundRefreshGuard::begin(Arc::clone(&self.refresh_priority)) + .ok_or(RuntimeSnapshotCompileError::RevisionChanged)?; + self.refresh_locked(mode, Some(&background.cancellation)) + .await + } + + async fn refresh_locked( + &self, + mode: RefreshMode, + cancellation: Option<&CancellationToken>, + ) -> Result { + let configuration_changed = if matches!(mode, RefreshMode::Reconcile) { + let persisted_revision = preempt_for_committed_refresh( + cancellation, + self.compiler.store().current_config_revision(), + ) + .await? + .map_err(|_| { + self.snapshots.suspend(); + RuntimeSnapshotCompileError::StoreUnavailable + })?; + let changed = runtime_revision_needs_refresh( + self.published_revision().map(ConfigRevision::get), + persisted_revision.get(), + ); + if !changed + && !self.provider_catalogs_need_refresh() + && self + .snapshots + .acquire() + .is_ok_and(|snapshot| snapshot.extensions().is_none_or(|set| set.is_ready())) + { + return Ok(persisted_revision); + } + changed + } else { + true + }; + let compiled = if matches!(mode, RefreshMode::Committed) { + match self.snapshots.snapshot_for_diagnostics() { + Some(previous) => self.compiler.compile_with_cached_catalog(&previous).await, + None => self.compiler.compile().await, + } + } else { + preempt_for_committed_refresh(cancellation, self.compiler.compile()).await? + }; + let snapshot = match compiled { + Ok(snapshot) => snapshot, + Err(error) => { + // 目录对账失败可继续服务旧快照；配置或权限变化无法确认时须暂停 + if configuration_changed { + self.snapshots.suspend(); + } + return Err(error); + } + }; + let revision = snapshot.revision(); + self.snapshots.publish(snapshot); + Ok(revision) + } + + #[must_use] + pub fn published_revision(&self) -> Option { + self.snapshots.revision() + } + + #[must_use] + fn provider_catalogs_need_refresh(&self) -> bool { + let Ok(snapshot) = self.snapshots.acquire() else { + return true; + }; + self.compiler.provider_catalog_generations() != *snapshot.provider_catalog_generations() + } + + /// 数据库提交不能被目录或通知基础设施的暂时故障伪装成回滚 + async fn publish_committed_inner(&self, committed_revision: ConfigRevision) { + if let Err(error) = self.refresh_with_mode(RefreshMode::Committed).await { + tracing::warn!(?error, "committed runtime snapshot refresh failed"); + } + self.notify_committed_revision(committed_revision).await; + } + + async fn notify_committed_revision(&self, committed_revision: ConfigRevision) { + if let Err(error) = self + .subscriptions + .publish_snapshot_revision(committed_revision) + .await + { + tracing::warn!( + ?error, + revision = committed_revision.get(), + "runtime snapshot revision notification failed" + ); + } + } + + /// 交给 Host 的周期对账与长驻订阅任务 + pub fn worker_contributions(&self) -> Result, WorkerDefinitionError> { + let reconciliation_id = WorkerId::try_new( + WorkerKind::RuntimeSnapshotReconciliation, + "runtime_snapshot", + )?; + let schedule = WorkerSchedule::try_new( + RECONCILIATION_INTERVAL, + INITIAL_BACKOFF, + MAXIMUM_BACKOFF, + UNUSED_LEASE_TTL, + UNUSED_LEASE_RENEWAL, + )?; + let reconciliation = WorkerRegistration::try_new( + reconciliation_id, + WorkerRunnable::Scheduled { + schedule, + lease: None, + task: Box::new(RuntimeSnapshotReconciliationTask { + publisher: self.clone(), + }), + }, + )?; + let subscription_id = + WorkerId::try_new(WorkerKind::RuntimeChangeSubscription, "runtime_snapshot")?; + let restart = DaemonRestartPolicy::try_new(INITIAL_BACKOFF, MAXIMUM_BACKOFF)?; + let subscription = WorkerRegistration::try_new( + subscription_id, + WorkerRunnable::Daemon { + restart, + task: Box::new(RuntimeSnapshotSubscriptionTask { + subscriptions: Arc::clone(&self.subscriptions), + publisher: self.clone(), + }), + }, + )?; + Ok(vec![ + WorkerContribution::Registration(reconciliation), + WorkerContribution::Registration(subscription), + ]) + } +} + +impl SnapshotControl for RuntimeSnapshotPublisher { + fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()> { + Box::pin(async move { + self.publish_committed_inner(committed_revision).await; + }) + } +} + +struct RuntimeSnapshotReconciliationTask { + publisher: RuntimeSnapshotPublisher, +} + +impl ScheduledTask for RuntimeSnapshotReconciliationTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + self.publisher + .refresh_with_mode(RefreshMode::Reconcile) + .await + .map(|_| ()) + .map_err(|source| { + WorkerTaskError::safe("runtime snapshot reconciliation failed") + .with_source(source) + }) + }) + } +} + +struct RuntimeSnapshotSubscriptionTask { + subscriptions: Arc, + publisher: RuntimeSnapshotPublisher, +} + +impl DaemonTask for RuntimeSnapshotSubscriptionTask { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let mut retry_delay = INITIAL_BACKOFF; + loop { + if cancellation.is_cancelled() { + return Ok(()); + } + let subscription = self.subscriptions.subscribe_snapshot_revisions().await; + let mut subscription = match subscription { + Ok(subscription) => { + retry_delay = INITIAL_BACKOFF; + subscription + } + Err(_) => { + wait_or_cancel(&cancellation, retry_delay).await; + retry_delay = (retry_delay * 2).min(MAXIMUM_BACKOFF); + continue; + } + }; + loop { + let cancelled = cancellation.cancelled().fuse(); + let next = subscription.next().fuse(); + pin_mut!(cancelled, next); + let notified = select_biased! { + _ = cancelled => return Ok(()), + next = next => next, + }; + match notified { + Some(Ok(_)) => { + let _ = self.publisher.refresh().await; + } + Some(Err(_)) | None => break, + } + } + } + }) + } +} + +async fn wait_or_cancel(cancellation: &CancellationToken, duration: Duration) { + let cancelled = cancellation.cancelled().fuse(); + let delay = Delay::new(duration).fuse(); + pin_mut!(cancelled, delay); + select_biased! { + _ = cancelled => {}, + _ = delay => {}, + } +} + +/// 当前发布版本与持久版本不一致时必须重载；缺失和回退同样 fail closed +#[must_use] +pub fn runtime_revision_needs_refresh( + published_revision: Option, + persisted_revision: u64, +) -> bool { + published_revision != Some(persisted_revision) +} + +fn lock_unpoisoned(lock: &SyncMutex) -> std::sync::MutexGuard<'_, T> { + lock.lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} diff --git a/backend/crates/gateway-core/src/settings/mod.rs b/backend/crates/gateway-core/src/settings/mod.rs index b521c09c0..10adbfb08 100644 --- a/backend/crates/gateway-core/src/settings/mod.rs +++ b/backend/crates/gateway-core/src/settings/mod.rs @@ -1,17 +1,4 @@ -//! 冻结宿主基线与请求显式改写；派生快照只属于当前调用，不发布全局配置 - -use std::{ - collections::BTreeMap, - sync::Arc, - time::{Duration, SystemTime}, -}; - -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -use crate::account::FastMode; -use crate::policy::{ClientApiKeyId, ClientPolicy, ClientSettings, RateLimits}; -use crate::routing::RuntimeSnapshot; +//! 运行设置的纯值、校验与编译结果 pub(crate) mod compiled; mod values; @@ -19,365 +6,6 @@ pub use values::{ SettingsValues, client_min_versions, response_body_limit, validate_request_limits, }; -/// 一次模型调用的有效设置；改写只影响当前请求，不发布配置或修改持久化 revision -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ExecutionSettings { - pub runtime: SettingsValues, - pub fast_mode: FastMode, - pub client_limits: RateLimits, - /// null 表示不限制总时长；显式时限从请求开始计时，改写不重置计时原点 - pub timeout_ms: Option, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct SettingOverride { - pub instance_id: String, - pub order: u64, - pub value: T, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] -struct ExecutionOverrides { - #[serde(skip)] - defaults: Arc, - client_key_id: String, - input: Arc, - fast_mode: Option>, - client_limits: Option>, - timeout_ms: Option>>, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] -struct HttpTimeout { - input_ms: Option, - change: Option>>, -} - -/// 共享值均不可变；同级子调用的修改不会回写父调用 -#[derive(Clone, Debug)] -pub struct RequestSettings { - baseline: Arc, - effective: Arc, - overrides: Arc>, - execution: Option>, - http_timeout: Option, - order: u64, -} - -impl PartialEq for RequestSettings { - fn eq(&self, other: &Self) -> bool { - Arc::ptr_eq(&self.baseline, &other.baseline) - && self.overrides == other.overrides - && self.execution == other.execution - && self.http_timeout == other.http_timeout - } -} -impl Eq for RequestSettings {} - #[derive(Debug, thiserror::Error)] #[error("request settings are invalid")] pub struct InvalidSettings; - -impl RequestSettings { - #[must_use] - pub fn new(snapshot: Arc) -> Self { - Self { - baseline: snapshot.clone(), - effective: snapshot, - overrides: Arc::default(), - execution: None, - http_timeout: None, - order: 0, - } - } - - #[must_use] - pub fn snapshot(&self) -> Arc { - self.effective.clone() - } - - #[must_use] - pub fn values(&self) -> &SettingsValues { - self.effective.settings() - } - - /// 来源与有效值分开查询，不允许插件伪造其他实例的写入来源 - pub fn inspect(&self) -> Value { - serde_json::json!({ - "config_revision": self.baseline.revision().get(), - "host": self.baseline.settings(), - "overrides": self.overrides, - "execution": self.execution, - "http_timeout": self.http_timeout, - }) - } - - pub fn replace( - &self, - values: SettingsValues, - instance_id: &str, - ) -> Result { - self.replace_scoped(self.values(), values, instance_id) - } - - /// Key 画像等隐式作用域结果不记录成插件改写，避免随后切换 Key 时泄漏 - fn replace_scoped( - &self, - previous: &SettingsValues, - values: SettingsValues, - instance_id: &str, - ) -> Result { - if previous == &values { - return Ok(self.clone()); - } - let current = serde_json::to_value(previous).map_err(|_| InvalidSettings)?; - let next = serde_json::to_value(&values).map_err(|_| InvalidSettings)?; - let order = self.order.checked_add(1).ok_or(InvalidSettings)?; - let mut overrides = self.overrides.as_ref().clone(); - let effective_values = self.values(); - // 每个已声明设置项整体替换，不猜测画像对象、数组或 null 的深合并含义 - for (name, value) in next.as_object().ok_or(InvalidSettings)? { - if current.get(name) != Some(value) { - overrides.insert( - name.clone(), - SettingOverride { - instance_id: instance_id.to_owned(), - order, - value: value.clone(), - }, - ); - } - } - // 只把本层实际修改的字段写回宿主层，作用域默认值不参与继承 - let values = if previous == effective_values { - values - } else { - apply_overrides(self.baseline.settings(), &overrides)? - }; - let effective = self - .effective - .resolve_settings(&values) - .map_err(|_| InvalidSettings)?; - Ok(Self { - baseline: self.baseline.clone(), - effective, - overrides: Arc::new(overrides), - execution: self.execution.clone(), - http_timeout: self.http_timeout.clone(), - order, - }) - } - - /// 持续会话的新请求读取当前宿主快照，只继承明确改写的字段 - pub fn rebase(&self, snapshot: Arc) -> Result { - if Arc::ptr_eq(&self.baseline, &snapshot) { - return Ok(self.clone()); - } - let effective = if self.overrides.is_empty() { - snapshot.clone() - } else { - snapshot - .resolve_settings(&apply_overrides(snapshot.settings(), &self.overrides)?) - .map_err(|_| InvalidSettings)? - }; - Ok(Self { - baseline: snapshot, - effective, - overrides: self.overrides.clone(), - execution: self.execution.clone(), - http_timeout: self.http_timeout.clone(), - order: self.order, - }) - } - - /// 认证、模型入口和插件视图共用解析规则，派生策略不改变 Key 默认值 - #[must_use] - pub fn apply_policy(&self, policy: ClientPolicy) -> ClientPolicy { - let values = self.resolve_execution(policy.key_id().as_str(), policy.defaults(), None); - policy.with_settings( - values.runtime.request_profiles(), - values.fast_mode, - values.client_limits, - ) - } - - #[must_use] - pub fn with_execution(mut self, policy: &ClientPolicy, timeout_ms: Option) -> Self { - let input = self.resolve_execution(policy.key_id().as_str(), policy.defaults(), timeout_ms); - let previous = self - .execution - .as_ref() - .filter(|scope| scope.client_key_id == policy.key_id().as_str()); - self.execution = Some(Arc::new(ExecutionOverrides { - defaults: policy.defaults().clone(), - client_key_id: policy.key_id().as_str().to_owned(), - input: Arc::new(input), - fast_mode: previous.and_then(|scope| scope.fast_mode.clone()), - client_limits: previous.and_then(|scope| scope.client_limits.clone()), - timeout_ms: previous.and_then(|scope| scope.timeout_ms.clone()), - })); - self - } - - #[must_use] - pub fn execution_values(&self) -> Option { - let scope = self.execution.as_ref()?; - Some(self.resolve_execution( - &scope.client_key_id, - &scope.defaults, - scope.input.timeout_ms, - )) - } - - fn resolve_execution( - &self, - key: &str, - defaults: &ClientSettings, - timeout_ms: Option, - ) -> ExecutionSettings { - let overrides = self - .execution - .as_ref() - .filter(|scope| scope.client_key_id == key); - let mut runtime = self.values().clone(); - // Key 画像覆盖宿主默认，插件显式覆盖整个设置项；此优先级只在这里解释 - if !self.overrides.contains_key("request_profiles") && !defaults.request_profiles.is_empty() - { - let mut profiles = runtime.request_profiles().clone(); - profiles.extend(defaults.request_profiles.clone()); - runtime = runtime.with_request_profiles(profiles); - } - ExecutionSettings { - runtime, - fast_mode: overrides - .and_then(|scope| scope.fast_mode.as_ref()) - .map_or(defaults.fast_mode, |change| change.value), - client_limits: overrides - .and_then(|scope| scope.client_limits.as_ref()) - .map_or(defaults.limits, |change| change.value), - timeout_ms: overrides - .and_then(|scope| scope.timeout_ms.as_ref()) - .map_or(timeout_ms, |change| change.value), - } - } - - pub(crate) fn execution_deadline( - &self, - key: &ClientApiKeyId, - started_at: SystemTime, - ) -> Result { - let scope = self - .execution - .as_ref() - .filter(|scope| scope.client_key_id == key.as_str()) - .ok_or(InvalidSettings)?; - let timeout_ms = scope - .timeout_ms - .as_ref() - .map_or(scope.input.timeout_ms, |change| change.value); - crate::lifecycle::Deadline::from_timeout(started_at, timeout_ms.map(Duration::from_millis)) - .ok_or(InvalidSettings) - } - - pub fn replace_execution( - &self, - values: &ExecutionSettings, - instance_id: &str, - ) -> Result { - let previous = self.execution_values().ok_or(InvalidSettings)?; - let mut updated = - self.replace_scoped(&previous.runtime, values.runtime.clone(), instance_id)?; - if previous.fast_mode == values.fast_mode - && previous.client_limits == values.client_limits - && previous.timeout_ms == values.timeout_ms - { - return Ok(updated); - } - let mut scope = self.execution.as_deref().cloned().ok_or(InvalidSettings)?; - let order = if updated.order == self.order { - self.order.checked_add(1).ok_or(InvalidSettings)? - } else { - updated.order - }; - if previous.fast_mode != values.fast_mode { - scope.fast_mode = Some(SettingOverride { - instance_id: instance_id.to_owned(), - order, - value: values.fast_mode, - }); - } - if previous.client_limits != values.client_limits { - scope.client_limits = Some(SettingOverride { - instance_id: instance_id.to_owned(), - order, - value: values.client_limits, - }); - } - if previous.timeout_ms != values.timeout_ms { - scope.timeout_ms = Some(SettingOverride { - instance_id: instance_id.to_owned(), - order, - value: values.timeout_ms, - }); - } - updated.execution = Some(Arc::new(scope)); - updated.order = order; - Ok(updated) - } - - #[must_use] - pub fn execution_timeout(&self, key: &ClientApiKeyId) -> Option { - self.execution - .as_ref() - .filter(|scope| scope.client_key_id == key.as_str()) - .and_then(|scope| scope.timeout_ms.as_ref()) - .and_then(|change| change.value.map(Duration::from_millis)) - } - - #[must_use] - pub fn with_http_timeout(mut self, timeout_ms: Option) -> Self { - self.http_timeout.get_or_insert(HttpTimeout { - input_ms: timeout_ms, - change: None, - }); - self - } - - pub fn replace_http_timeout( - mut self, - timeout_ms: Option, - instance_id: &str, - ) -> Result { - let timeout = self.http_timeout.get_or_insert(HttpTimeout { - input_ms: None, - change: None, - }); - if timeout - .change - .as_ref() - .map_or(timeout.input_ms, |change| change.value) - != timeout_ms - { - self.order = self.order.checked_add(1).ok_or(InvalidSettings)?; - timeout.change = Some(SettingOverride { - instance_id: instance_id.to_owned(), - order: self.order, - value: timeout_ms, - }); - } - Ok(self) - } -} - -fn apply_overrides( - values: &SettingsValues, - overrides: &BTreeMap, -) -> Result { - let mut values = serde_json::to_value(values).map_err(|_| InvalidSettings)?; - let object = values.as_object_mut().ok_or(InvalidSettings)?; - for (name, change) in overrides { - object.insert(name.clone(), change.value.clone()); - } - serde_json::from_value(values).map_err(|_| InvalidSettings) -} diff --git a/backend/crates/gateway-core/src/upstream.rs b/backend/crates/gateway-core/src/upstream.rs index ec4db950c..9c17c6123 100644 --- a/backend/crates/gateway-core/src/upstream.rs +++ b/backend/crates/gateway-core/src/upstream.rs @@ -1,6 +1,6 @@ //! 上游调用跨 Engine、Event、Error 与 Provider 共享的中立边界事实 -use std::fmt; +use std::{fmt, num::NonZeroU32}; use crate::validation::{IdentifierError, validate_text}; @@ -73,3 +73,15 @@ impl fmt::Debug for OpaqueUpstreamValue { formatter.write_str("OpaqueUpstreamValue()") } } + +/// Provider 可为当前 attempt 选择的请求局部传输档位 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum AttemptTransport { + /// 使用 Provider 的默认传输策略 + #[default] + Default, + /// 固定账号重试 Provider 的首选传输；序号由 Provider 的独立预算驱动 + Retry(NonZeroU32), + /// 使用 Provider 定义的备用传输 + Fallback, +} diff --git a/backend/crates/gateway-core/tests/engine/execution.rs b/backend/crates/gateway-core/tests/engine/execution/mod.rs similarity index 99% rename from backend/crates/gateway-core/tests/engine/execution.rs rename to backend/crates/gateway-core/tests/engine/execution/mod.rs index 3ae6cce71..790f7b151 100644 --- a/backend/crates/gateway-core/tests/engine/execution.rs +++ b/backend/crates/gateway-core/tests/engine/execution/mod.rs @@ -1,4 +1,5 @@ //! 执行服务的准入、预算、请求记录与终结行为测试 +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use std::{ collections::{BTreeMap, BTreeSet, VecDeque}, @@ -2964,10 +2965,7 @@ use gateway_core::routing::{ PublicModelId, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId, }; -use gateway_core::runtime::{ - RuntimeSnapshotHandle, - extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::runtime::RuntimeSnapshotHandle; use gateway_core::upstream::{UpstreamSendState, UpstreamTransport}; use serde_json::{Value, json}; @@ -3218,7 +3216,7 @@ impl Provider for NativeCatalogProvider { fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models() { #[derive(Debug)] struct Aliases(Vec); - impl gateway_core::runtime::extensions::ExtensionSetLease for Aliases { + impl gateway_core::routing::extensions::ExtensionSetLease for Aliases { fn is_ready(&self) -> bool { true } @@ -3239,7 +3237,7 @@ fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models( .with_settings(&settings) .unwrap() .with_extensions(Some(ExtensionSetReference::new( - gateway_core::runtime::extensions::ExtensionSetId::new("catalog-aliases".into()) + gateway_core::routing::extensions::ExtensionSetId::new("catalog-aliases".into()) .unwrap(), Arc::new(Aliases(vec![ gateway_core::routing::ContributedModelAlias { @@ -4500,7 +4498,7 @@ fn invalid_request_settings_leave_the_prepared_execution_unchanged() { #[test] fn execution_settings_must_match_the_prepared_key() { - use gateway_core::settings::RequestSettings; + use gateway_core::routing::request_settings::RequestSettings; let service = service(Arc::new(Admissions::default()), Arc::new(Budget::default())); let client = service.authenticate("sk_start_test").unwrap(); @@ -4595,7 +4593,7 @@ fn repeated_connection_failures_never_block_later_requests_for_the_provider() { #[test] fn entry_settings_freeze_authentication_and_rebase_only_explicit_overrides() { block_on(async { - use gateway_core::settings::RequestSettings; + use gateway_core::routing::request_settings::RequestSettings; let snapshots = RuntimeSnapshotHandle::new(start_snapshot()); let service = DefaultExecutionService::new( snapshots.clone(), @@ -4707,7 +4705,7 @@ fn entry_settings_freeze_authentication_and_rebase_only_explicit_overrides() { #[test] fn unchanged_and_precompiled_request_settings_reuse_the_frozen_snapshot() { - use gateway_core::settings::RequestSettings; + use gateway_core::routing::request_settings::RequestSettings; let snapshot = Arc::new(start_snapshot()); let baseline = RequestSettings::new(snapshot.clone()); @@ -4785,7 +4783,9 @@ fn unchanged_and_precompiled_request_settings_reuse_the_frozen_snapshot() { #[test] fn child_settings_recompute_key_scope_without_inheriting_parent_defaults() { block_on(async { - use gateway_core::{account::OpaqueProviderData, settings::RequestSettings}; + use gateway_core::{ + account::OpaqueProviderData, routing::request_settings::RequestSettings, + }; let provider = ProviderKind::new("openai").unwrap(); let profiles = |name: &str| { BTreeMap::from([( diff --git a/backend/crates/gateway-core/tests/engine/extensions.rs b/backend/crates/gateway-core/tests/engine/extensions.rs index 739fbb0da..4fead6150 100644 --- a/backend/crates/gateway-core/tests/engine/extensions.rs +++ b/backend/crates/gateway-core/tests/engine/extensions.rs @@ -1,8 +1,13 @@ -//! 执行测试共用的扩展代次引用与就绪租约替身 +//! 执行扩展索引的代次隔离、计划保活与租约替身 use std::sync::Arc; -use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; +use gateway_core::engine::extensions::{ExecutionExtensionIndex, ExecutionExtensionPlans}; +use gateway_core::engine::upstream_adapter::{UpstreamAdapter, UpstreamAdapterPlan}; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; +use gateway_core::{ + engine::AttemptContext, error::ProviderError, identity::ProviderKind, routing::UpstreamModelId, +}; struct ReadyLease; @@ -16,3 +21,71 @@ pub(super) fn reference(id: &str) -> ExtensionSetReference { let id = ExtensionSetId::new(id.into()).unwrap(); ExtensionSetReference::new(id, Arc::new(ReadyLease)) } + +#[derive(Debug)] +struct EmptyAdapterPlan; + +impl UpstreamAdapterPlan for EmptyAdapterPlan { + fn select( + &self, + _: &AttemptContext, + _: &ProviderKind, + _: &UpstreamModelId, + ) -> Result>, ProviderError> { + Ok(None) + } +} + +struct PublishedPlans(Arc); + +impl ExtensionSetLease for PublishedPlans { + fn is_ready(&self) -> bool { + self.0.upstream_adapters.is_some() + } +} + +#[test] +fn adapter_only_generation_stays_alive_until_the_last_frozen_plan_is_released() { + let index = ExecutionExtensionIndex::default(); + let id = ExtensionSetId::new("adapter-only".into()).unwrap(); + let plans = Arc::new(ExecutionExtensionPlans { + middleware: None, + upstream_adapters: Some(Arc::new(EmptyAdapterPlan)), + }); + let weak_plans = Arc::downgrade(&plans); + let published = Arc::new(PublishedPlans(index.register(id.clone(), plans).unwrap())); + let weak_publication = Arc::downgrade(&published); + let reference = ExtensionSetReference::new(id.clone(), published); + assert!(index.middleware(&reference).is_none()); + let frozen = index.upstream_adapters(&reference).unwrap(); + drop(reference); + + // 当前发布可以退场；在途计划仍保活整个代次，不能覆盖相同身份 + assert!(weak_publication.upgrade().is_some()); + assert!(weak_plans.upgrade().is_some()); + assert!(index.register(id.clone(), Arc::default()).is_err()); + let in_flight = frozen.clone(); + drop(frozen); + assert!(weak_publication.upgrade().is_some()); + drop(in_flight); + + assert!(weak_publication.upgrade().is_none()); + assert!(weak_plans.upgrade().is_none()); + let released = ExtensionSetReference::new(id.clone(), Arc::new(ReadyLease)); + assert!(index.upstream_adapters(&released).is_none()); + let _replacement = index.register(id, Arc::default()).unwrap(); +} + +#[test] +fn execution_index_does_not_own_registered_generations() { + let index = ExecutionExtensionIndex::default(); + let reference = reference("unpublished"); + let plans = Arc::new(ExecutionExtensionPlans::default()); + let weak = Arc::downgrade(&plans); + let owner = index.register(reference.id().clone(), plans).unwrap(); + assert!(weak.upgrade().is_some()); + drop(owner); + assert!(weak.upgrade().is_none()); + assert!(index.middleware(&reference).is_none()); + assert!(index.upstream_adapters(&reference).is_none()); +} diff --git a/backend/crates/gateway-core/tests/engine/middleware.rs b/backend/crates/gateway-core/tests/engine/middleware/mod.rs similarity index 95% rename from backend/crates/gateway-core/tests/engine/middleware.rs rename to backend/crates/gateway-core/tests/engine/middleware/mod.rs index 24bb2b665..8d9ddc034 100644 --- a/backend/crates/gateway-core/tests/engine/middleware.rs +++ b/backend/crates/gateway-core/tests/engine/middleware/mod.rs @@ -12,11 +12,13 @@ use gateway_core::account::ProviderAccountId; use gateway_core::account::scope::AccountGroupId; use gateway_core::engine::ModelRequestId; use gateway_core::engine::execution::ClientTransport; -use gateway_core::engine::extensions::ExtensionCallScope; +use gateway_core::engine::extensions::{ + ExecutionExtensionIndex, ExecutionExtensionPlans, ExtensionCallScope, +}; use gateway_core::engine::middleware::{ FrozenMiddlewarePlan, MiddlewareAuthority, MiddlewareBody, MiddlewareContext, MiddlewareError, - MiddlewareExtensionIndex, MiddlewareFrame, MiddlewareHeader, MiddlewareMount, MiddlewareNext, - MiddlewarePlan, MiddlewareRequest, MiddlewareResponse, MiddlewareTarget, + MiddlewareFrame, MiddlewareHeader, MiddlewareMount, MiddlewareNext, MiddlewarePlan, + MiddlewareRequest, MiddlewareResponse, MiddlewareTarget, }; use gateway_core::engine::provider::{ ProviderCallMetadata, ProviderMiddlewareTerminal, ProviderStream, execute_attempt_middleware, @@ -27,7 +29,7 @@ use gateway_core::lifecycle::CancellationToken; use gateway_core::operation::{GenerateRequest, Operation, OperationKind, ProtocolPayload}; use gateway_core::policy::ClientApiKeyId; use gateway_core::routing::UpstreamModelId; -use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::upstream::UpstreamTransport; use serde_json::{Value, json}; @@ -303,14 +305,17 @@ fn attempt_middleware_consumes_owned_next_once_and_preserves_host_envelopes() { futures::executor::block_on(async { let plan_calls = Arc::new(AtomicUsize::new(0)); let terminal_calls = Arc::new(AtomicUsize::new(0)); - let index = MiddlewareExtensionIndex::default(); + let index = ExecutionExtensionIndex::default(); let id = ExtensionSetId::new("middleware-generation".to_owned()).unwrap(); let reference = ExtensionSetReference::new(id.clone(), Arc::new(ReadyGeneration)); let owner = index .register( id.clone(), - Arc::new(RewritingPlan { - calls: Arc::clone(&plan_calls), + Arc::new(ExecutionExtensionPlans { + middleware: Some(Arc::new(RewritingPlan { + calls: Arc::clone(&plan_calls), + })), + upstream_adapters: None, }), ) .unwrap(); @@ -318,14 +323,17 @@ fn attempt_middleware_consumes_owned_next_once_and_preserves_host_envelopes() { index .register( id, - Arc::new(RewritingPlan { - calls: Arc::new(AtomicUsize::new(0)), + Arc::new(ExecutionExtensionPlans { + middleware: Some(Arc::new(RewritingPlan { + calls: Arc::new(AtomicUsize::new(0)), + })), + upstream_adapters: None }), ) .is_err(), "同一发布代次不能静默替换中间件计划", ); - let plan = index.resolve(&reference).unwrap(); + let plan = index.middleware(&reference).unwrap(); let terminal: ProviderMiddlewareTerminal = { let terminal_calls = Arc::clone(&terminal_calls); diff --git a/backend/crates/gateway-core/tests/engine/policy.rs b/backend/crates/gateway-core/tests/engine/policy.rs index dfd6a6792..e876a1e97 100644 --- a/backend/crates/gateway-core/tests/engine/policy.rs +++ b/backend/crates/gateway-core/tests/engine/policy.rs @@ -24,7 +24,7 @@ use gateway_core::{ }, policy::ClientApiKeyId, routing::ProviderKind, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, + routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, }; #[derive(Debug)] diff --git a/backend/crates/gateway-core/tests/main.rs b/backend/crates/gateway-core/tests/main.rs index 48d0a3cde..15edffcc2 100644 --- a/backend/crates/gateway-core/tests/main.rs +++ b/backend/crates/gateway-core/tests/main.rs @@ -16,7 +16,6 @@ mod policy; mod provider_ports; mod routing; mod runtime; -mod settings; mod task; mod upstream; diff --git a/backend/crates/gateway-core/tests/middleware/mod.rs b/backend/crates/gateway-core/tests/middleware.rs similarity index 100% rename from backend/crates/gateway-core/tests/middleware/mod.rs rename to backend/crates/gateway-core/tests/middleware.rs diff --git a/backend/crates/gateway-core/tests/provider_ports.rs b/backend/crates/gateway-core/tests/provider_ports/mod.rs similarity index 100% rename from backend/crates/gateway-core/tests/provider_ports.rs rename to backend/crates/gateway-core/tests/provider_ports/mod.rs diff --git a/backend/crates/gateway-core/tests/routing/mod.rs b/backend/crates/gateway-core/tests/routing/mod.rs index e0ad53344..92d21416c 100644 --- a/backend/crates/gateway-core/tests/routing/mod.rs +++ b/backend/crates/gateway-core/tests/routing/mod.rs @@ -18,6 +18,7 @@ use gateway_core::routing::{ UpstreamModelId, }; +mod request_settings; mod snapshot; fn settings() -> gateway_core::settings::SettingsValues { diff --git a/backend/crates/gateway-core/tests/settings/mod.rs b/backend/crates/gateway-core/tests/routing/request_settings.rs similarity index 99% rename from backend/crates/gateway-core/tests/settings/mod.rs rename to backend/crates/gateway-core/tests/routing/request_settings.rs index c3b6f0291..b574e985e 100644 --- a/backend/crates/gateway-core/tests/settings/mod.rs +++ b/backend/crates/gateway-core/tests/routing/request_settings.rs @@ -7,8 +7,9 @@ use gateway_core::{ account::OpaqueProviderData, identity::ProviderKind, policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits}, + routing::request_settings::RequestSettings, routing::{ClientRoutingScope, ConfigRevision, FrozenAccountScope, RuntimeSnapshot}, - settings::{RequestSettings, SettingsValues}, + settings::SettingsValues, }; use serde_json::json; diff --git a/backend/crates/gateway-core/tests/routing/snapshot.rs b/backend/crates/gateway-core/tests/routing/snapshot.rs index f39bf6cb5..d76a16531 100644 --- a/backend/crates/gateway-core/tests/routing/snapshot.rs +++ b/backend/crates/gateway-core/tests/routing/snapshot.rs @@ -10,6 +10,10 @@ use futures::future::BoxFuture; use gateway_core::account::FastMode; use gateway_core::operation::OperationKind; use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits}; +use gateway_core::routing::extensions::{ + ExtensionPreparationError, ExtensionPreparationPort, ExtensionSetId, ExtensionSetLease, + ExtensionSetReference, +}; use gateway_core::routing::snapshot::{ RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts, SnapshotClientPolicyFacts, SnapshotFacts, @@ -20,10 +24,6 @@ use gateway_core::routing::{ ProviderCatalogGeneration, ProviderCatalogPort, ProviderCatalogUnavailable, ProviderKind, ProviderModelCapabilities, PublicModelId, UpstreamModelId, }; -use gateway_core::runtime::extensions::{ - ExtensionPreparationError, ExtensionPreparationPort, ExtensionSetId, ExtensionSetLease, - ExtensionSetReference, -}; use gateway_core::settings::SettingsValues; #[derive(Clone)] @@ -1084,7 +1084,9 @@ fn key_profiles_replace_whole_global_choice_and_previous_snapshot_stays_frozen() .unwrap() }; let values = |snapshot: &gateway_core::routing::RuntimeSnapshot| { - let settings = gateway_core::settings::RequestSettings::new(Arc::new(snapshot.clone())); + let settings = gateway_core::routing::request_settings::RequestSettings::new(Arc::new( + snapshot.clone(), + )); let mut values: Vec<_> = snapshot .client_policies() .map(|policy| { diff --git a/backend/crates/gateway-core/tests/runtime/mod.rs b/backend/crates/gateway-core/tests/runtime/mod.rs index 7915b1dbc..fc8652468 100644 --- a/backend/crates/gateway-core/tests/runtime/mod.rs +++ b/backend/crates/gateway-core/tests/runtime/mod.rs @@ -1,7 +1,7 @@ //! 验证运行时快照发布、版本对账与并发更新行为 use std::collections::BTreeMap; -mod extensions; +mod publications; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; diff --git a/backend/crates/gateway-core/tests/runtime/extensions.rs b/backend/crates/gateway-core/tests/runtime/publications.rs similarity index 96% rename from backend/crates/gateway-core/tests/runtime/extensions.rs rename to backend/crates/gateway-core/tests/runtime/publications.rs index b8488ebc1..43e832952 100644 --- a/backend/crates/gateway-core/tests/runtime/extensions.rs +++ b/backend/crates/gateway-core/tests/runtime/publications.rs @@ -5,9 +5,9 @@ use std::sync::{ atomic::{AtomicBool, AtomicUsize, Ordering}, }; -use gateway_core::runtime::{ - RuntimeSnapshotHandle, - extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, +use gateway_core::{ + routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, + runtime::RuntimeSnapshotHandle, }; struct Lease { diff --git a/backend/crates/gateway-host/src/lib.rs b/backend/crates/gateway-host/src/lib.rs index 91efed550..7ef96b4b7 100644 --- a/backend/crates/gateway-host/src/lib.rs +++ b/backend/crates/gateway-host/src/lib.rs @@ -59,7 +59,9 @@ pub async fn initialize(config: HostConfig) -> Result { let log_guard = initialize_logging(&config.logging, config.timezone)?; let cancellation = CancellationToken::new(); let connections = Arc::new(ConnectionTracker::new(cancellation.clone())); - let workers = WorkerSupervisor::new(cancellation.clone()); + // 关闭请求先排空 HTTP/WS，写泵仍需接收在途请求的最后写入 + // Worker 使用独立信号，只由 drain 结束后的 shutdown 或析构触发 + let workers = WorkerSupervisor::new(CancellationToken::new()); let system = Arc::new(ProcessSystemOperations::new( cancellation.clone(), config.system_update.clone(), @@ -190,6 +192,14 @@ impl HostBundle { } } +impl Drop for HostBundle { + fn drop(&mut self) { + // CLI、启动失败和 serve future 被丢弃也要通知公共能力停止 + // WorkerSupervisor 自己负责取消并终止剩余任务 + self.cancellation.cancel(); + } +} + #[derive(Debug, thiserror::Error)] pub enum HostError { #[error(transparent)] diff --git a/backend/crates/gateway-host/src/system_update/config.rs b/backend/crates/gateway-host/src/system_update/config.rs new file mode 100644 index 000000000..65fcb6cef --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/config.rs @@ -0,0 +1,165 @@ +//! 系统更新的部署配置、编译元数据与本地路径解析 + +use std::{ + env, fs, + path::{Path, PathBuf}, +}; + +use serde::Deserialize; + +use super::{OperationError, conflict, internal}; +use crate::config::ConfigError; + +const APP_BINARY_NAME: &str = "codex-proxy-rs"; +const DEFAULT_GITHUB_API_BASE: &str = "https://api.github.com/repos"; +const DEFAULT_UPDATE_REPOSITORY: &str = "zyycn/codex-proxy-rs"; + +/// 系统更新与重启配置；所有字段只由 Host 解释 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(default)] +pub struct SystemUpdateConfig { + pub version: String, + pub git_sha: String, + pub build_time: String, + pub deployment_mode: String, + pub build_type: String, + pub update_repository: Option, + pub github_api_base: String, + pub executable_path: Option, + /// 未显式指定时，由组合根传入 API 实际使用的静态资源目录 + pub web_dist_dir: Option, + pub update_state_file: PathBuf, + pub update_lock_file: PathBuf, + pub update_temp_dir: PathBuf, + pub self_restart_enabled: bool, +} + +impl Default for SystemUpdateConfig { + fn default() -> Self { + let deployment_mode = + environment_value("CPR_DEPLOYMENT_MODE").unwrap_or_else(|| "source".to_owned()); + let executable_path = environment_value("CPR_UPDATE_EXE_PATH") + .map(PathBuf::from) + .or_else(|| { + (deployment_mode == "docker") + .then(|| PathBuf::from("/app/bin").join(APP_BINARY_NAME)) + }); + let update_state_file = environment_value("CPR_UPDATE_STATE_FILE") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("update-state.json")); + let update_lock_file = environment_value("CPR_UPDATE_LOCK_FILE") + .map(PathBuf::from) + .unwrap_or_else(|| update_state_file.with_file_name("update.lock")); + let update_temp_dir = environment_value("CPR_UPDATE_TEMP_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| default_temp_dir(&update_state_file)); + Self { + version: option_env!("CPR_VERSION") + .unwrap_or(env!("CARGO_PKG_VERSION")) + .to_owned(), + git_sha: option_env!("CPR_GIT_SHA").unwrap_or("unknown").to_owned(), + build_time: option_env!("CPR_BUILD_TIME") + .unwrap_or("unknown") + .to_owned(), + deployment_mode, + build_type: option_env!("CPR_BUILD_TYPE").unwrap_or("source").to_owned(), + update_repository: Some( + environment_value("CPR_UPDATE_REPOSITORY") + .unwrap_or_else(|| DEFAULT_UPDATE_REPOSITORY.to_owned()), + ), + github_api_base: environment_value("CPR_GITHUB_API_BASE") + .unwrap_or_else(|| DEFAULT_GITHUB_API_BASE.to_owned()), + executable_path, + web_dist_dir: None, + update_state_file, + update_lock_file, + update_temp_dir, + self_restart_enabled: environment_value("CPR_ENABLE_SELF_RESTART").as_deref() + == Some("true"), + } + } +} + +impl SystemUpdateConfig { + pub(crate) fn resolve_and_validate( + &mut self, + source_dir: &Path, + runtime_data_dir: &Path, + asset_directory: &Path, + ) -> Result<(), ConfigError> { + let web_dist_dir = self + .web_dist_dir + .get_or_insert_with(|| asset_directory.to_path_buf()); + if web_dist_dir.as_os_str().is_empty() { + return Err(ConfigError::InvalidField("host.system_update.web_dist_dir")); + } + for path in [self.executable_path.as_mut(), self.web_dist_dir.as_mut()] + .into_iter() + .flatten() + { + if path.is_relative() { + *path = source_dir.join(&*path); + } + } + for path in [ + &mut self.update_state_file, + &mut self.update_lock_file, + &mut self.update_temp_dir, + ] { + if path.is_relative() { + *path = runtime_data_dir.join(&*path); + } + } + if self.version.trim().is_empty() { + return Err(ConfigError::InvalidField("host.system_update.version")); + } + if !matches!( + self.deployment_mode.as_str(), + "source" | "binary" | "docker" + ) { + return Err(ConfigError::InvalidField( + "host.system_update.deployment_mode", + )); + } + Ok(()) + } + + pub(crate) fn executable_path(&self) -> Result { + if let Some(path) = &self.executable_path { + return Ok(path.clone()); + } + env::current_exe() + .and_then(fs::canonicalize) + .map_err(|error| internal(format!("failed to resolve executable: {error}"))) + } + + pub(crate) fn official_plugins_dir(&self) -> Result { + let executable = self.executable_path()?; + let parent = executable + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .ok_or_else(|| internal("failed to resolve official plugin release directory"))?; + Ok(parent.join("plugins").join("official")) + } + + pub(crate) fn web_dist_dir(&self) -> Result<&Path, OperationError> { + self.web_dist_dir + .as_deref() + .filter(|path| !path.as_os_str().is_empty()) + .ok_or_else(|| conflict("web assets directory is not configured")) + } +} + +pub(super) fn environment_value(key: &str) -> Option { + env::var(key) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +fn default_temp_dir(state_file: &Path) -> PathBuf { + state_file + .parent() + .map(|parent| parent.join("update-tmp")) + .unwrap_or_else(|| std::env::temp_dir().join("codex-proxy-rs-update")) +} diff --git a/backend/crates/gateway-host/src/system_update/download.rs b/backend/crates/gateway-host/src/system_update/download.rs index c99ed258d..bd7685304 100644 --- a/backend/crates/gateway-host/src/system_update/download.rs +++ b/backend/crates/gateway-host/src/system_update/download.rs @@ -8,8 +8,9 @@ use std::time::Duration; use futures::StreamExt as _; use sha2::Digest as _; +use super::events::UpdateEvents; use super::release::{download_client, validate_download_url}; -use super::{OperationError, UpdateEvents, invalid, upstream}; +use super::{OperationError, invalid, upstream}; pub(crate) const MAX_DOWNLOAD_SIZE: u64 = 500 * 1024 * 1024; pub(crate) const MAX_CHECKSUM_SIZE: u64 = 1024 * 1024; diff --git a/backend/crates/gateway-host/src/system_update/error.rs b/backend/crates/gateway-host/src/system_update/error.rs new file mode 100644 index 000000000..bdf7d85bb --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/error.rs @@ -0,0 +1,21 @@ +//! 系统更新各基础设施步骤共用的管理端错误分类 + +use gateway_admin::ports::system::{SystemOperationError, SystemOperationErrorKind}; + +pub(super) type OperationError = SystemOperationError; + +pub(super) fn invalid(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Invalid, message) +} + +pub(super) fn conflict(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Conflict, message) +} + +pub(super) fn upstream(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Upstream, message) +} + +pub(super) fn internal(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Internal, message) +} diff --git a/backend/crates/gateway-host/src/system_update/events.rs b/backend/crates/gateway-host/src/system_update/events.rs new file mode 100644 index 000000000..f58fcbd8c --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/events.rs @@ -0,0 +1,130 @@ +//! 更新日志事件的有界广播、递增标识与终态标记 + +use chrono::Utc; +use gateway_admin::model::system::{SystemUpdateEvent, SystemUpdateEventLevel}; +use std::sync::atomic::{AtomicU64, Ordering}; +use tokio::sync::broadcast; + +pub(super) struct UpdateEvents { + sender: broadcast::Sender<(SystemUpdateEvent, bool)>, + sequence: AtomicU64, +} + +impl Default for UpdateEvents { + fn default() -> Self { + let (sender, _) = broadcast::channel(256); + Self { + sender, + sequence: AtomicU64::new(0), + } + } +} + +impl UpdateEvents { + pub(super) fn subscribe(&self) -> broadcast::Receiver<(SystemUpdateEvent, bool)> { + self.sender.subscribe() + } + + pub(super) fn info( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Info, + operation_id, + step, + message, + None, + false, + ); + } + + pub(super) fn success( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Success, + operation_id, + step, + message, + None, + false, + ); + } + + pub(super) fn success_terminal( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Success, + operation_id, + step, + message, + Some(100), + true, + ); + } + + pub(super) fn error_terminal( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Error, + operation_id, + step, + message, + None, + true, + ); + } + + pub(super) fn emit_progress(&self, operation_id: &str, message: &str, progress: u8) { + self.emit( + SystemUpdateEventLevel::Info, + Some(operation_id), + Some("download"), + message, + Some(progress), + false, + ); + } + + fn emit( + &self, + level: SystemUpdateEventLevel, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + progress_percent: Option, + terminal: bool, + ) { + let now = Utc::now(); + let sequence = self.sequence.fetch_add(1, Ordering::Relaxed); + let event = SystemUpdateEvent { + id: format!( + "update-log-{}-{sequence}", + now.timestamp_nanos_opt() + .unwrap_or_else(|| now.timestamp_millis()) + ), + operation_id: operation_id.map(str::to_owned), + level, + step: step.map(str::to_owned), + message: message.into(), + terminal, + progress_percent, + occurred_at: now, + }; + let _ = self.sender.send((event, terminal)); + } +} diff --git a/backend/crates/gateway-host/src/system_update/installation.rs b/backend/crates/gateway-host/src/system_update/installation.rs index d4d827b81..8b107548c 100644 --- a/backend/crates/gateway-host/src/system_update/installation.rs +++ b/backend/crates/gateway-host/src/system_update/installation.rs @@ -10,6 +10,8 @@ use sha2::{Digest as _, Sha256}; use super::swap::backup_path_for; use super::{OperationError, SystemUpdateConfig, conflict, internal}; +pub(super) const OFFICIAL_PLUGIN_MANIFEST: &str = "plugin-release-manifest.json"; + #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] pub(crate) struct ReleaseFiles { @@ -36,9 +38,7 @@ impl ReleaseFiles { } fn read(binary: &Path, web: &Path, plugins: &Path) -> Result { - if !web.join("index.html").is_file() - || !plugins.join(super::OFFICIAL_PLUGIN_MANIFEST).is_file() - { + if !web.join("index.html").is_file() || !plugins.join(OFFICIAL_PLUGIN_MANIFEST).is_file() { return Err(conflict("安装文件不完整，请核对部署后重试")); } let mut binary_hash = Sha256::new(); diff --git a/backend/crates/gateway-host/src/system_update/mod.rs b/backend/crates/gateway-host/src/system_update/mod.rs index d26757702..38b9e17e4 100644 --- a/backend/crates/gateway-host/src/system_update/mod.rs +++ b/backend/crates/gateway-host/src/system_update/mod.rs @@ -1,962 +1,19 @@ -//! 版本、安全自更新、回滚与进程重启的 Host-owned 实现 +//! 版本、安全自更新、回滚与进程重启的 Host-owned 实现入口 mod archive; +mod config; mod download; +mod error; +mod events; mod installation; +mod operations; mod process; mod release; mod state; mod swap; -use std::env; -use std::fs; -use std::path::{Path, PathBuf}; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; +pub use config::SystemUpdateConfig; +pub use operations::ProcessSystemOperations; +pub use release::validate_download_url; -use async_trait::async_trait; -use chrono::Utc; -use gateway_admin::model::system::{ - SystemOperationAccepted, SystemOperationKind, SystemUpdateChannel as UpdateChannel, - SystemUpdateDetail, SystemUpdateEvent, SystemUpdateEventLevel, SystemUpdatePolicy, - SystemUpdateStatus, SystemVersion, -}; -use gateway_admin::ports::system::{ - SystemOperationError, SystemOperationErrorKind, SystemOperations, SystemRestartPreflight, - SystemUpdateCandidate, SystemUpdateEventStream, SystemUpdatePreflight, -}; -use gateway_core::lifecycle::CancellationToken; -use serde::Deserialize; -use tokio::sync::{Mutex as AsyncMutex, broadcast}; - -use crate::config::ConfigError; - -use self::archive::extract_release; -use self::download::{MAX_CHECKSUM_SIZE, MAX_DOWNLOAD_SIZE, download_file, verify_checksum}; -use self::installation::ReleaseFiles; -use self::process::{environment_value, spawn_replacement}; -use self::release::{ - ReleaseCache, confirmed_target, detail_from_release, fetch_latest, select_archive, - validate_update_target, version_channel, -}; -use self::state::{ - OperationFileLock, UpdateOperation, UpdateTempDir, finish, operation_id, read_status, - reconcile_installation, recover_interrupted, set_running, -}; -use self::swap::{replace_release_files, rollback_official_plugins_dir, rollback_release}; - -pub use self::release::validate_download_url; - -const APP_BINARY_NAME: &str = "codex-proxy-rs"; -const DEFAULT_GITHUB_API_BASE: &str = "https://api.github.com/repos"; -const DEFAULT_UPDATE_REPOSITORY: &str = "zyycn/codex-proxy-rs"; -const OFFICIAL_PLUGIN_MANIFEST: &str = "plugin-release-manifest.json"; -const MAX_OFFICIAL_PLUGIN_MANIFEST_SIZE: u64 = 256 * 1024; - -type OperationError = SystemOperationError; - -/// 系统更新与重启配置；所有字段只由 Host 解释 -#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] -#[serde(default)] -pub struct SystemUpdateConfig { - pub version: String, - pub git_sha: String, - pub build_time: String, - pub deployment_mode: String, - pub build_type: String, - pub update_repository: Option, - pub github_api_base: String, - pub executable_path: Option, - /// 未显式指定时，由组合根传入 API 实际使用的静态资源目录 - pub web_dist_dir: Option, - pub update_state_file: PathBuf, - pub update_lock_file: PathBuf, - pub update_temp_dir: PathBuf, - pub self_restart_enabled: bool, -} - -impl Default for SystemUpdateConfig { - fn default() -> Self { - let deployment_mode = - environment_value("CPR_DEPLOYMENT_MODE").unwrap_or_else(|| "source".to_owned()); - let executable_path = environment_value("CPR_UPDATE_EXE_PATH") - .map(PathBuf::from) - .or_else(|| { - (deployment_mode == "docker") - .then(|| PathBuf::from("/app/bin").join(APP_BINARY_NAME)) - }); - let update_state_file = environment_value("CPR_UPDATE_STATE_FILE") - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from("update-state.json")); - let update_lock_file = environment_value("CPR_UPDATE_LOCK_FILE") - .map(PathBuf::from) - .unwrap_or_else(|| update_state_file.with_file_name("update.lock")); - let update_temp_dir = environment_value("CPR_UPDATE_TEMP_DIR") - .map(PathBuf::from) - .unwrap_or_else(|| state::default_temp_dir(&update_state_file)); - Self { - version: option_env!("CPR_VERSION") - .unwrap_or(env!("CARGO_PKG_VERSION")) - .to_owned(), - git_sha: option_env!("CPR_GIT_SHA").unwrap_or("unknown").to_owned(), - build_time: option_env!("CPR_BUILD_TIME") - .unwrap_or("unknown") - .to_owned(), - deployment_mode, - build_type: option_env!("CPR_BUILD_TYPE").unwrap_or("source").to_owned(), - update_repository: Some( - environment_value("CPR_UPDATE_REPOSITORY") - .unwrap_or_else(|| DEFAULT_UPDATE_REPOSITORY.to_owned()), - ), - github_api_base: environment_value("CPR_GITHUB_API_BASE") - .unwrap_or_else(|| DEFAULT_GITHUB_API_BASE.to_owned()), - executable_path, - web_dist_dir: None, - update_state_file, - update_lock_file, - update_temp_dir, - self_restart_enabled: environment_value("CPR_ENABLE_SELF_RESTART").as_deref() - == Some("true"), - } - } -} - -impl SystemUpdateConfig { - pub(crate) fn resolve_and_validate( - &mut self, - source_dir: &Path, - runtime_data_dir: &Path, - asset_directory: &Path, - ) -> Result<(), ConfigError> { - let web_dist_dir = self - .web_dist_dir - .get_or_insert_with(|| asset_directory.to_path_buf()); - if web_dist_dir.as_os_str().is_empty() { - return Err(ConfigError::InvalidField("host.system_update.web_dist_dir")); - } - for path in [self.executable_path.as_mut(), self.web_dist_dir.as_mut()] - .into_iter() - .flatten() - { - if path.is_relative() { - *path = source_dir.join(&*path); - } - } - for path in [ - &mut self.update_state_file, - &mut self.update_lock_file, - &mut self.update_temp_dir, - ] { - if path.is_relative() { - *path = runtime_data_dir.join(&*path); - } - } - if self.version.trim().is_empty() { - return Err(ConfigError::InvalidField("host.system_update.version")); - } - if !matches!( - self.deployment_mode.as_str(), - "source" | "binary" | "docker" - ) { - return Err(ConfigError::InvalidField( - "host.system_update.deployment_mode", - )); - } - Ok(()) - } - - fn update_support_error(&self) -> Option { - if !matches!(self.build_type.as_str(), "release" | "experimental") { - return Some("在线更新需要官方发布构建".to_owned()); - } - let Some(channel) = version_channel(&self.version) else { - return Some("当前版本不符合发行命名规范，无法确定更新通道".to_owned()); - }; - if self.build_type == "experimental" && channel != UpdateChannel::Experimental { - return Some("实验构建必须使用 exp.N 版本，无法在线更新".to_owned()); - } - let Some(repository) = self.update_repository.as_deref() else { - return Some("检查更新需要配置 CPR_UPDATE_REPOSITORY".to_owned()); - }; - if let Err(error) = release::validate_repository(repository) { - return Some(error.to_string()); - } - if let Err(error) = release::validate_api_base(&self.github_api_base) { - return Some(error); - } - if let Err(error) = self.web_dist_dir() { - return Some(error.to_string()); - } - None - } - - fn release_cache_key(&self) -> String { - format!( - "{}|{}|{}|{}|{}", - self.update_repository.as_deref().unwrap_or_default(), - self.github_api_base, - self.version, - self.deployment_mode, - self.build_type, - ) - } - - pub(crate) fn executable_path(&self) -> Result { - if let Some(path) = &self.executable_path { - return Ok(path.clone()); - } - env::current_exe() - .and_then(fs::canonicalize) - .map_err(|error| internal(format!("failed to resolve executable: {error}"))) - } - - pub(crate) fn official_plugins_dir(&self) -> Result { - let executable = self.executable_path()?; - let parent = executable - .parent() - .filter(|path| !path.as_os_str().is_empty()) - .ok_or_else(|| internal("failed to resolve official plugin release directory"))?; - Ok(parent.join("plugins").join("official")) - } - - pub(crate) fn web_dist_dir(&self) -> Result<&Path, OperationError> { - self.web_dist_dir - .as_deref() - .filter(|path| !path.as_os_str().is_empty()) - .ok_or_else(|| conflict("web assets directory is not configured")) - } -} - -/// gateway-admin 消费的真实进程/文件系统操作实现 -#[derive(Clone)] -pub struct ProcessSystemOperations { - cancellation: CancellationToken, - events: Arc, - config: Arc, - operation_lock: Arc>, - release_cache: Arc, - running_files: Arc>, -} - -impl ProcessSystemOperations { - #[must_use] - pub fn new(cancellation: CancellationToken, mut config: SystemUpdateConfig) -> Self { - // 组合根在接收请求前固定安装路径和发行文件；旧程序重命名后， - // current_exe 可能指向备份，后续更新、回滚和重启不能再反查运行文件 - let running_files = Arc::new(config.executable_path().and_then(|executable| { - config.executable_path = Some(executable); - ReleaseFiles::installed(&config) - })); - Self { - cancellation, - events: Arc::new(UpdateEvents::default()), - config: Arc::new(config), - operation_lock: Arc::new(AsyncMutex::new(())), - release_cache: Arc::new(ReleaseCache::default()), - running_files, - } - } - - fn start_update( - &self, - target_version: Option, - channel: Option, - preflight: Arc, - ) -> Result { - let operation_lock = Arc::clone(&self.operation_lock) - .try_lock_owned() - .map_err(|_| conflict("system operation is already running"))?; - if self.cancellation.is_cancelled() { - return Err(conflict("服务正在关闭")); - } - if let Some(reason) = self.config.update_support_error() { - self.events - .error_terminal(None, Some("preflight"), reason.clone()); - return Err(conflict(reason)); - } - let target = confirmed_target(target_version)?; - let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; - let selected = self.resolve_channel(channel)?; - if let Err(error) = validate_update_target(&self.config.version, &target, selected) { - self.events - .error_terminal(None, Some("preflight"), error.to_string()); - return Err(error); - } - let status = self.reconcile_installation()?; - if status.need_restart { - return Err(conflict("更新已完成，请先重启服务")); - } - let operation_id = operation_id("update"); - let mut operation = UpdateOperation::start( - &self.config.update_state_file, - &operation_id, - &target, - &self.config.version, - file_lock, - )?; - let accepted = SystemOperationAccepted::Update { - operation_id: operation_id.clone(), - deployment_mode: self.config.deployment_mode.clone(), - message: "更新已开始".to_owned(), - target_version: target.clone(), - }; - let service = self.clone(); - // 任务持有互斥锁和落盘守卫，HTTP 断开不会取消更新；Host 关闭仍会收敛终态 - drop(tokio::spawn(async move { - let _operation_lock = operation_lock; - let result = tokio::select! { - biased; - () = service.cancellation.cancelled() => Err(conflict("服务关闭，更新已中断")), - result = service.perform_update_inner(&target, selected, &operation_id, preflight.as_ref()) => result, - }; - if let Err(error) = operation.complete(&result) { - tracing::warn!(error = %error, "系统更新终态落盘失败"); - return; - } - match result { - Ok(_) => service.events.success_terminal( - Some(&operation_id), - Some("done"), - "更新文件已替换，等待服务重启生效", - ), - Err(error) => service.events.error_terminal( - Some(&operation_id), - Some("failed"), - error.to_string(), - ), - } - })); - Ok(accepted) - } - - async fn perform_update_inner( - &self, - target: &str, - channel: UpdateChannel, - operation_id: &str, - preflight: &dyn SystemUpdatePreflight, - ) -> Result { - let repository = self - .config - .update_repository - .as_deref() - .ok_or_else(|| conflict("update repository is not configured"))?; - self.events.info( - Some(operation_id), - Some("release"), - "正在获取最新 Release 信息", - ); - let release = fetch_latest( - &self.config.github_api_base, - repository, - &self.config.version, - channel, - ) - .await? - .ok_or_else(|| conflict("当前发行通道没有可用更新"))?; - let detail = detail_from_release(&self.config, &release, channel); - if detail.latest_version != target { - return Err(conflict("远端最新版本已变化，请重新确认")); - } - if !detail.has_update { - return Err(conflict("当前没有可用更新")); - } - self.events.info( - Some(operation_id), - Some("prepare"), - format!("准备更新到 v{target}"), - ); - self.install_release(&release, target, operation_id, preflight) - .await - } - - async fn install_release( - &self, - release: &release::GitHubRelease, - version: &str, - operation_id: &str, - preflight: &dyn SystemUpdatePreflight, - ) -> Result { - self.events.info( - Some(operation_id), - Some("asset"), - "正在选择匹配当前平台的更新包", - ); - let archive = select_archive(release, version)?; - self.events.info( - Some(operation_id), - Some("asset"), - format!( - "已选择更新包 {} ({})", - archive.name, - download::format_bytes(archive.size) - ), - ); - self.events - .info(Some(operation_id), Some("verify"), "正在校验更新资源"); - if archive.size == 0 || archive.size > MAX_DOWNLOAD_SIZE { - return Err(invalid("release archive size is invalid")); - } - let checksum = release - .assets - .iter() - .find(|asset| asset.name == "checksums.txt") - .ok_or_else(|| upstream("release checksums.txt is required"))?; - if checksum.size == 0 || checksum.size > MAX_CHECKSUM_SIZE { - return Err(invalid("release checksum size is invalid")); - } - self.events - .info(Some(operation_id), Some("prepare"), "正在创建临时更新目录"); - fs::create_dir_all(&self.config.update_temp_dir) - .map_err(|error| internal(format!("failed to prepare update temp dir: {error}")))?; - let temp_root = fs::canonicalize(&self.config.update_temp_dir) - .map_err(|error| internal(format!("failed to resolve update temp dir: {error}")))?; - let temp = UpdateTempDir::create(&temp_root)?; - let archive_path = temp.path().join(&archive.name); - self.events - .info(Some(operation_id), Some("download"), "开始下载更新包"); - download_file( - &archive.browser_download_url, - &archive_path, - archive.size, - &self.config.github_api_base, - operation_id, - &self.events, - ) - .await?; - self.events - .success(Some(operation_id), Some("download"), "更新包下载完成"); - self.events - .info(Some(operation_id), Some("checksum"), "正在校验 checksum"); - verify_checksum( - &archive_path, - &archive.name, - &checksum.browser_download_url, - checksum.size, - &self.config.github_api_base, - ) - .await?; - self.events - .success(Some(operation_id), Some("checksum"), "checksum 校验通过"); - self.events - .info(Some(operation_id), Some("extract"), "正在解压更新包"); - let extracted = extract_release(&archive_path, temp.path())?; - self.events - .success(Some(operation_id), Some("extract"), "更新包解压完成"); - let release_manifest = read_release_manifest(&extracted.official_plugins_dir)?; - let candidate = SystemUpdateCandidate { - target_version: version.to_owned(), - release_manifest: release_manifest.clone(), - }; - self.events.info( - Some(operation_id), - Some("preflight"), - "正在校验目标发行信息", - ); - let plugin_revision = preflight.validate(candidate).await?; - preflight.confirm_revision(plugin_revision).await?; - self.events.success( - Some(operation_id), - Some("preflight"), - "目标发行信息校验通过，插件兼容性将在重启前检查", - ); - self.events - .info(Some(operation_id), Some("replace"), "正在替换应用文件"); - replace_release_files( - &self.config.executable_path()?, - self.config.web_dist_dir()?, - extracted, - )?; - let mut applied = AppliedReleaseGuard::new(&self.config); - let applied_manifest = read_release_manifest(&self.config.official_plugins_dir()?)?; - if applied_manifest.as_ref() != release_manifest.as_ref() { - return Err(applied.rollback(conflict("已应用的更新候选与预检候选不一致"))); - } - if let Err(error) = preflight.confirm_revision(plugin_revision).await { - return Err(applied.rollback(error)); - } - let files = ReleaseFiles::installed(&self.config)?; - applied.commit(); - self.events - .success(Some(operation_id), Some("replace"), "应用文件替换完成"); - Ok(files) - } - - fn resolve_channel( - &self, - requested: Option, - ) -> Result { - let channel = requested.unwrap_or_else(|| { - version_channel(&self.config.version).unwrap_or(UpdateChannel::Stable) - }); - if !update_policy(&self.config, channel) - .available_channels - .contains(&channel) - { - return Err(conflict("当前发行版本不支持此更新通道")); - } - Ok(channel) - } - - fn installed_restart_candidate(&self) -> Result, OperationError> { - if self.config.build_type == "source" { - return Ok(None); - } - let release_manifest = read_release_manifest(&self.config.official_plugins_dir()?)?; - #[derive(Deserialize)] - struct ReleaseIdentity { - gateway_version: String, - } - let identity: ReleaseIdentity = serde_json::from_slice(&release_manifest) - .map_err(|_| invalid("已安装发行清单缺少目标版本"))?; - Ok(Some(SystemUpdateCandidate { - target_version: identity.gateway_version, - release_manifest, - })) - } - - fn reconcile_installation(&self) -> Result { - let running_files = self.running_files.as_ref().as_ref().map_err(Clone::clone)?; - reconcile_installation(&self.config, running_files) - } -} - -#[async_trait] -impl SystemOperations for ProcessSystemOperations { - async fn version(&self) -> Result { - let detail = self.update_detail(false, None).await?; - Ok(SystemVersion { - version: self.config.version.clone(), - git_sha: self.config.git_sha.clone(), - build_time: self.config.build_time.clone(), - deployment_mode: self.config.deployment_mode.clone(), - update_channel: version_channel(&self.config.version) - .map(|channel| channel.as_str().to_owned()) - .unwrap_or_else(|| "unknown".to_owned()), - latest_version: detail.latest_version, - has_update: detail.has_update, - update_cached: detail.cached, - update_warning: detail.warning, - }) - } - - async fn update_detail( - &self, - refresh: bool, - channel: Option, - ) -> Result { - let channel = self.resolve_channel(channel)?; - match self - .release_cache - .detail(&self.config, refresh, channel) - .await - { - Ok(detail) => Ok(detail), - Err(error) => Ok(base_update_detail( - &self.config, - channel, - self.config.update_support_error(), - Some(error.to_string()), - )), - } - } - - fn update_events(&self) -> SystemUpdateEventStream { - let receiver = self.events.subscribe(); - Box::pin(futures::stream::unfold( - (receiver, false), - |(mut receiver, close)| async move { - if close { - return None; - } - loop { - match receiver.recv().await { - Ok((event, terminal)) => return Some((event, (receiver, terminal))), - Err(broadcast::error::RecvError::Lagged(_)) => continue, - Err(broadcast::error::RecvError::Closed) => return None, - } - } - }, - )) - } - - async fn perform_update( - &self, - target_version: Option, - channel: Option, - preflight: Arc, - ) -> Result { - self.start_update(target_version, channel, preflight) - } - - async fn update_status(&self) -> Result { - if let Ok(_operation) = self.operation_lock.try_lock() { - recover_interrupted( - &self.config.update_state_file, - &self.config.update_lock_file, - )?; - if self.config.update_support_error().is_none() { - match OperationFileLock::acquire(&self.config.update_lock_file) { - Ok(_lock) => return self.reconcile_installation(), - Err(error) if error.kind() == SystemOperationErrorKind::Conflict => {} - Err(error) => return Err(error), - } - } - } - read_status(&self.config.update_state_file, false) - } - - async fn rollback( - &self, - preflight: Arc, - ) -> Result { - let _operation = self - .operation_lock - .try_lock() - .map_err(|_| conflict("system operation is already running"))?; - if let Some(reason) = self.config.update_support_error() { - return Err(conflict(reason)); - } - let operation_id = operation_id("rollback"); - let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; - let target_version = self - .reconcile_installation()? - .previous_version - .ok_or_else(|| conflict("没有可用于回滚的上一版本"))?; - let release_manifest = - read_release_manifest(&rollback_official_plugins_dir(&self.config)?)?; - let plugin_revision = preflight - .validate_rollback(SystemUpdateCandidate { - target_version, - release_manifest: Arc::clone(&release_manifest), - }) - .await?; - preflight.confirm_revision(plugin_revision).await?; - set_running( - &self.config.update_state_file, - &operation_id, - SystemOperationKind::Rollback, - None, - &self.config.version, - )?; - let result = match rollback_release(&self.config) { - Ok(()) => { - let mut applied = AppliedReleaseGuard::new(&self.config); - let result = match read_release_manifest(&self.config.official_plugins_dir()?) { - Ok(manifest) if manifest.as_ref() == release_manifest.as_ref() => { - preflight.confirm_revision(plugin_revision).await - } - Ok(_) => Err(conflict("已应用的回滚候选与预检候选不一致")), - Err(error) => Err(error), - }; - match result { - Ok(()) => { - applied.commit(); - Ok(()) - } - Err(error) => Err(applied.rollback(error)), - } - } - Err(error) => Err(error), - }; - finish( - &self.config.update_state_file, - &operation_id, - SystemOperationKind::Rollback, - None, - None, - result.as_ref().err().map(ToString::to_string), - )?; - match &result { - Ok(()) => self.events.success_terminal( - Some(&operation_id), - Some("done"), - "previous release restored", - ), - Err(error) => { - self.events - .error_terminal(Some(&operation_id), Some("failed"), error.to_string()) - } - } - drop(file_lock); - result?; - Ok(SystemOperationAccepted::Rollback { - operation_id, - message: "回滚完成，请重启服务。".to_owned(), - need_restart: true, - }) - } - - async fn restart_candidate(&self) -> Result, OperationError> { - let _operation = self - .operation_lock - .try_lock() - .map_err(|_| conflict("system operation is already running"))?; - self.installed_restart_candidate() - } - - async fn restart( - &self, - preflight: Arc, - ) -> Result { - // 与 update/rollback 互斥：更新替换文件期间触发自重启会让新进程 - // 载入半成品产物 - let _operation = self - .operation_lock - .try_lock() - .map_err(|_| conflict("system operation is already running"))?; - if !self.config.self_restart_enabled { - return Err(conflict("self restart is disabled")); - } - let _file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; - let candidate = self.installed_restart_candidate()?; - let files = candidate - .as_ref() - .map(|_| ReleaseFiles::installed(&self.config)) - .transpose()?; - preflight.prepare(candidate).await?; - if let Some(files) = files - && files != ReleaseFiles::installed(&self.config)? - { - return Err(conflict("重启检查期间安装文件已变化，请重新确认")); - } - let message = if self.config.deployment_mode == "docker" { - "已安排进程内重启" - } else { - spawn_replacement(&self.config)?; - "已安排自重启" - }; - let operation_id = operation_id("restart"); - let cancellation = self.cancellation.clone(); - drop(tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(500)).await; - cancellation.cancel(); - })); - Ok(SystemOperationAccepted::Restart { - operation_id, - message: message.to_owned(), - }) - } -} - -fn read_release_manifest(directory: &Path) -> Result, OperationError> { - let path = directory.join(OFFICIAL_PLUGIN_MANIFEST); - let metadata = fs::symlink_metadata(&path) - .map_err(|error| invalid(format!("official plugin manifest is unavailable: {error}")))?; - if !metadata.file_type().is_file() || metadata.file_type().is_symlink() { - return Err(invalid("official plugin manifest is not a regular file")); - } - if metadata.len() == 0 || metadata.len() > MAX_OFFICIAL_PLUGIN_MANIFEST_SIZE { - return Err(invalid("official plugin manifest size is invalid")); - } - fs::read(path) - .map(Arc::<[u8]>::from) - .map_err(|error| invalid(format!("official plugin manifest cannot be read: {error}"))) -} - -fn restore_applied_release(config: &SystemUpdateConfig, cause: OperationError) -> OperationError { - match rollback_release(config) { - Ok(()) => cause, - Err(error) => internal(format!( - "release verification failed and restoring the previous files failed: {error}" - )), - } -} - -struct AppliedReleaseGuard<'a> { - config: &'a SystemUpdateConfig, - armed: bool, -} - -impl<'a> AppliedReleaseGuard<'a> { - const fn new(config: &'a SystemUpdateConfig) -> Self { - Self { - config, - armed: true, - } - } - - const fn commit(&mut self) { - self.armed = false; - } - - fn rollback(mut self, cause: OperationError) -> OperationError { - self.armed = false; - restore_applied_release(self.config, cause) - } -} - -impl Drop for AppliedReleaseGuard<'_> { - fn drop(&mut self) { - if self.armed - && let Err(error) = rollback_release(self.config) - { - tracing::error!(error = %error, "系统发行物临界区取消后的恢复失败"); - } - } -} - -fn update_policy(config: &SystemUpdateConfig, channel: UpdateChannel) -> SystemUpdatePolicy { - use UpdateChannel::{Alpha, Beta, Experimental, Rc, Stable}; - SystemUpdatePolicy { - channel, - available_channels: if version_channel(&config.version) == Some(Experimental) { - vec![Experimental] - } else { - vec![Stable, Rc, Beta, Alpha] - }, - } -} - -fn base_update_detail( - config: &SystemUpdateConfig, - channel: UpdateChannel, - unsupported_reason: Option, - warning: Option, -) -> SystemUpdateDetail { - SystemUpdateDetail { - policy: update_policy(config, channel), - current_version: config.version.clone(), - latest_version: config.version.clone(), - has_update: false, - deployment_mode: config.deployment_mode.clone(), - build_type: config.build_type.clone(), - release_url: None, - notes: None, - cached: false, - update_supported: unsupported_reason.is_none() && warning.is_none(), - unsupported_reason, - warning, - } -} - -pub(crate) struct UpdateEvents { - sender: broadcast::Sender<(SystemUpdateEvent, bool)>, - sequence: AtomicU64, -} - -impl Default for UpdateEvents { - fn default() -> Self { - let (sender, _) = broadcast::channel(256); - Self { - sender, - sequence: AtomicU64::new(0), - } - } -} - -impl UpdateEvents { - fn subscribe(&self) -> broadcast::Receiver<(SystemUpdateEvent, bool)> { - self.sender.subscribe() - } - - fn info(&self, operation_id: Option<&str>, step: Option<&str>, message: impl Into) { - self.emit( - SystemUpdateEventLevel::Info, - operation_id, - step, - message, - None, - false, - ); - } - - fn success(&self, operation_id: Option<&str>, step: Option<&str>, message: impl Into) { - self.emit( - SystemUpdateEventLevel::Success, - operation_id, - step, - message, - None, - false, - ); - } - - fn success_terminal( - &self, - operation_id: Option<&str>, - step: Option<&str>, - message: impl Into, - ) { - self.emit( - SystemUpdateEventLevel::Success, - operation_id, - step, - message, - Some(100), - true, - ); - } - - fn error_terminal( - &self, - operation_id: Option<&str>, - step: Option<&str>, - message: impl Into, - ) { - self.emit( - SystemUpdateEventLevel::Error, - operation_id, - step, - message, - None, - true, - ); - } - - pub(crate) fn emit_progress(&self, operation_id: &str, message: &str, progress: u8) { - self.emit( - SystemUpdateEventLevel::Info, - Some(operation_id), - Some("download"), - message, - Some(progress), - false, - ); - } - - fn emit( - &self, - level: SystemUpdateEventLevel, - operation_id: Option<&str>, - step: Option<&str>, - message: impl Into, - progress_percent: Option, - terminal: bool, - ) { - let now = Utc::now(); - let sequence = self.sequence.fetch_add(1, Ordering::Relaxed); - let event = SystemUpdateEvent { - id: format!( - "update-log-{}-{sequence}", - now.timestamp_nanos_opt() - .unwrap_or_else(|| now.timestamp_millis()) - ), - operation_id: operation_id.map(str::to_owned), - level, - step: step.map(str::to_owned), - message: message.into(), - terminal, - progress_percent, - occurred_at: now, - }; - let _ = self.sender.send((event, terminal)); - } -} - -fn invalid(message: impl Into) -> OperationError { - SystemOperationError::new(SystemOperationErrorKind::Invalid, message) -} - -fn conflict(message: impl Into) -> OperationError { - SystemOperationError::new(SystemOperationErrorKind::Conflict, message) -} - -fn upstream(message: impl Into) -> OperationError { - SystemOperationError::new(SystemOperationErrorKind::Upstream, message) -} - -fn internal(message: impl Into) -> OperationError { - SystemOperationError::new(SystemOperationErrorKind::Internal, message) -} +use error::{OperationError, conflict, internal, invalid, upstream}; diff --git a/backend/crates/gateway-host/src/system_update/operations.rs b/backend/crates/gateway-host/src/system_update/operations.rs new file mode 100644 index 000000000..1e485419c --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/operations.rs @@ -0,0 +1,605 @@ +//! 系统更新、回滚与重启的受理、互斥和文件交换生命周期 + +use std::{fs, path::Path, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use gateway_admin::model::system::{ + SystemOperationAccepted, SystemOperationKind, SystemUpdateChannel as UpdateChannel, + SystemUpdateDetail, SystemUpdateStatus, SystemVersion, +}; +use gateway_admin::ports::system::{ + SystemOperationErrorKind, SystemOperations, SystemRestartPreflight, SystemUpdateCandidate, + SystemUpdateEventStream, SystemUpdatePreflight, +}; +use gateway_core::lifecycle::CancellationToken; +use serde::Deserialize; +use tokio::sync::{Mutex as AsyncMutex, broadcast}; + +use super::archive::extract_release; +use super::download::{ + MAX_CHECKSUM_SIZE, MAX_DOWNLOAD_SIZE, download_file, format_bytes, verify_checksum, +}; +use super::events::UpdateEvents; +use super::installation::{OFFICIAL_PLUGIN_MANIFEST, ReleaseFiles}; +use super::process::spawn_replacement; +use super::release::{ + GitHubRelease, ReleaseCache, base_update_detail, confirmed_target, detail_from_release, + fetch_latest, select_archive, update_policy, validate_update_target, version_channel, +}; +use super::state::{ + OperationFileLock, UpdateOperation, UpdateTempDir, finish, operation_id, read_status, + reconcile_installation, recover_interrupted, set_running, +}; +use super::swap::{replace_release_files, rollback_official_plugins_dir, rollback_release}; +use super::{OperationError, SystemUpdateConfig, conflict, internal, invalid, upstream}; + +const MAX_OFFICIAL_PLUGIN_MANIFEST_SIZE: u64 = 256 * 1024; + +/// gateway-admin 消费的真实进程/文件系统操作实现 +#[derive(Clone)] +pub struct ProcessSystemOperations { + cancellation: CancellationToken, + events: Arc, + config: Arc, + operation_lock: Arc>, + release_cache: Arc, + running_files: Arc>, +} + +impl ProcessSystemOperations { + #[must_use] + pub fn new(cancellation: CancellationToken, mut config: SystemUpdateConfig) -> Self { + // 组合根在接收请求前固定安装路径和发行文件；旧程序重命名后， + // current_exe 可能指向备份，后续更新、回滚和重启不能再反查运行文件 + let running_files = Arc::new(config.executable_path().and_then(|executable| { + config.executable_path = Some(executable); + ReleaseFiles::installed(&config) + })); + Self { + cancellation, + events: Arc::new(UpdateEvents::default()), + config: Arc::new(config), + operation_lock: Arc::new(AsyncMutex::new(())), + release_cache: Arc::new(ReleaseCache::default()), + running_files, + } + } + + fn start_update( + &self, + target_version: Option, + channel: Option, + preflight: Arc, + ) -> Result { + let operation_lock = Arc::clone(&self.operation_lock) + .try_lock_owned() + .map_err(|_| conflict("system operation is already running"))?; + if self.cancellation.is_cancelled() { + return Err(conflict("服务正在关闭")); + } + if let Some(reason) = self.config.update_support_error() { + self.events + .error_terminal(None, Some("preflight"), reason.clone()); + return Err(conflict(reason)); + } + let target = confirmed_target(target_version)?; + let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; + let selected = self.resolve_channel(channel)?; + if let Err(error) = validate_update_target(&self.config.version, &target, selected) { + self.events + .error_terminal(None, Some("preflight"), error.to_string()); + return Err(error); + } + let status = self.reconcile_installation()?; + if status.need_restart { + return Err(conflict("更新已完成，请先重启服务")); + } + let operation_id = operation_id("update"); + let mut operation = UpdateOperation::start( + &self.config.update_state_file, + &operation_id, + &target, + &self.config.version, + file_lock, + )?; + let accepted = SystemOperationAccepted::Update { + operation_id: operation_id.clone(), + deployment_mode: self.config.deployment_mode.clone(), + message: "更新已开始".to_owned(), + target_version: target.clone(), + }; + let service = self.clone(); + // 任务持有互斥锁和落盘守卫，HTTP 断开不会取消更新；Host 关闭仍会收敛终态 + drop(tokio::spawn(async move { + let _operation_lock = operation_lock; + let result = tokio::select! { + biased; + () = service.cancellation.cancelled() => Err(conflict("服务关闭，更新已中断")), + result = service.perform_update_inner(&target, selected, &operation_id, preflight.as_ref()) => result, + }; + if let Err(error) = operation.complete(&result) { + tracing::warn!(error = %error, "系统更新终态落盘失败"); + return; + } + match result { + Ok(_) => service.events.success_terminal( + Some(&operation_id), + Some("done"), + "更新文件已替换，等待服务重启生效", + ), + Err(error) => service.events.error_terminal( + Some(&operation_id), + Some("failed"), + error.to_string(), + ), + } + })); + Ok(accepted) + } + + async fn perform_update_inner( + &self, + target: &str, + channel: UpdateChannel, + operation_id: &str, + preflight: &dyn SystemUpdatePreflight, + ) -> Result { + let repository = self + .config + .update_repository + .as_deref() + .ok_or_else(|| conflict("update repository is not configured"))?; + self.events.info( + Some(operation_id), + Some("release"), + "正在获取最新 Release 信息", + ); + let release = fetch_latest( + &self.config.github_api_base, + repository, + &self.config.version, + channel, + ) + .await? + .ok_or_else(|| conflict("当前发行通道没有可用更新"))?; + let detail = detail_from_release(&self.config, &release, channel); + if detail.latest_version != target { + return Err(conflict("远端最新版本已变化，请重新确认")); + } + if !detail.has_update { + return Err(conflict("当前没有可用更新")); + } + self.events.info( + Some(operation_id), + Some("prepare"), + format!("准备更新到 v{target}"), + ); + self.install_release(&release, target, operation_id, preflight) + .await + } + + async fn install_release( + &self, + release: &GitHubRelease, + version: &str, + operation_id: &str, + preflight: &dyn SystemUpdatePreflight, + ) -> Result { + self.events.info( + Some(operation_id), + Some("asset"), + "正在选择匹配当前平台的更新包", + ); + let archive = select_archive(release, version)?; + self.events.info( + Some(operation_id), + Some("asset"), + format!( + "已选择更新包 {} ({})", + archive.name, + format_bytes(archive.size) + ), + ); + self.events + .info(Some(operation_id), Some("verify"), "正在校验更新资源"); + if archive.size == 0 || archive.size > MAX_DOWNLOAD_SIZE { + return Err(invalid("release archive size is invalid")); + } + let checksum = release + .assets + .iter() + .find(|asset| asset.name == "checksums.txt") + .ok_or_else(|| upstream("release checksums.txt is required"))?; + if checksum.size == 0 || checksum.size > MAX_CHECKSUM_SIZE { + return Err(invalid("release checksum size is invalid")); + } + self.events + .info(Some(operation_id), Some("prepare"), "正在创建临时更新目录"); + fs::create_dir_all(&self.config.update_temp_dir) + .map_err(|error| internal(format!("failed to prepare update temp dir: {error}")))?; + let temp_root = fs::canonicalize(&self.config.update_temp_dir) + .map_err(|error| internal(format!("failed to resolve update temp dir: {error}")))?; + let temp = UpdateTempDir::create(&temp_root)?; + let archive_path = temp.path().join(&archive.name); + self.events + .info(Some(operation_id), Some("download"), "开始下载更新包"); + download_file( + &archive.browser_download_url, + &archive_path, + archive.size, + &self.config.github_api_base, + operation_id, + &self.events, + ) + .await?; + self.events + .success(Some(operation_id), Some("download"), "更新包下载完成"); + self.events + .info(Some(operation_id), Some("checksum"), "正在校验 checksum"); + verify_checksum( + &archive_path, + &archive.name, + &checksum.browser_download_url, + checksum.size, + &self.config.github_api_base, + ) + .await?; + self.events + .success(Some(operation_id), Some("checksum"), "checksum 校验通过"); + self.events + .info(Some(operation_id), Some("extract"), "正在解压更新包"); + let extracted = extract_release(&archive_path, temp.path())?; + self.events + .success(Some(operation_id), Some("extract"), "更新包解压完成"); + let release_manifest = read_release_manifest(&extracted.official_plugins_dir)?; + let candidate = SystemUpdateCandidate { + target_version: version.to_owned(), + release_manifest: release_manifest.clone(), + }; + self.events.info( + Some(operation_id), + Some("preflight"), + "正在校验目标发行信息", + ); + let plugin_revision = preflight.validate(candidate).await?; + preflight.confirm_revision(plugin_revision).await?; + self.events.success( + Some(operation_id), + Some("preflight"), + "目标发行信息校验通过，插件兼容性将在重启前检查", + ); + self.events + .info(Some(operation_id), Some("replace"), "正在替换应用文件"); + replace_release_files( + &self.config.executable_path()?, + self.config.web_dist_dir()?, + extracted, + )?; + let mut applied = AppliedReleaseGuard::new(&self.config); + let applied_manifest = read_release_manifest(&self.config.official_plugins_dir()?)?; + if applied_manifest.as_ref() != release_manifest.as_ref() { + return Err(applied.rollback(conflict("已应用的更新候选与预检候选不一致"))); + } + if let Err(error) = preflight.confirm_revision(plugin_revision).await { + return Err(applied.rollback(error)); + } + let files = ReleaseFiles::installed(&self.config)?; + applied.commit(); + self.events + .success(Some(operation_id), Some("replace"), "应用文件替换完成"); + Ok(files) + } + + fn resolve_channel( + &self, + requested: Option, + ) -> Result { + let channel = requested.unwrap_or_else(|| { + version_channel(&self.config.version).unwrap_or(UpdateChannel::Stable) + }); + if !update_policy(&self.config, channel) + .available_channels + .contains(&channel) + { + return Err(conflict("当前发行版本不支持此更新通道")); + } + Ok(channel) + } + + fn installed_restart_candidate(&self) -> Result, OperationError> { + if self.config.build_type == "source" { + return Ok(None); + } + let release_manifest = read_release_manifest(&self.config.official_plugins_dir()?)?; + #[derive(Deserialize)] + struct ReleaseIdentity { + gateway_version: String, + } + let identity: ReleaseIdentity = serde_json::from_slice(&release_manifest) + .map_err(|_| invalid("已安装发行清单缺少目标版本"))?; + Ok(Some(SystemUpdateCandidate { + target_version: identity.gateway_version, + release_manifest, + })) + } + + fn reconcile_installation(&self) -> Result { + let running_files = self.running_files.as_ref().as_ref().map_err(Clone::clone)?; + reconcile_installation(&self.config, running_files) + } +} + +#[async_trait] +impl SystemOperations for ProcessSystemOperations { + async fn version(&self) -> Result { + let detail = self.update_detail(false, None).await?; + Ok(SystemVersion { + version: self.config.version.clone(), + git_sha: self.config.git_sha.clone(), + build_time: self.config.build_time.clone(), + deployment_mode: self.config.deployment_mode.clone(), + update_channel: version_channel(&self.config.version) + .map(|channel| channel.as_str().to_owned()) + .unwrap_or_else(|| "unknown".to_owned()), + latest_version: detail.latest_version, + has_update: detail.has_update, + update_cached: detail.cached, + update_warning: detail.warning, + }) + } + + async fn update_detail( + &self, + refresh: bool, + channel: Option, + ) -> Result { + let channel = self.resolve_channel(channel)?; + match self + .release_cache + .detail(&self.config, refresh, channel) + .await + { + Ok(detail) => Ok(detail), + Err(error) => Ok(base_update_detail( + &self.config, + channel, + self.config.update_support_error(), + Some(error.to_string()), + )), + } + } + + fn update_events(&self) -> SystemUpdateEventStream { + let receiver = self.events.subscribe(); + Box::pin(futures::stream::unfold( + (receiver, false), + |(mut receiver, close)| async move { + if close { + return None; + } + loop { + match receiver.recv().await { + Ok((event, terminal)) => return Some((event, (receiver, terminal))), + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => return None, + } + } + }, + )) + } + + async fn perform_update( + &self, + target_version: Option, + channel: Option, + preflight: Arc, + ) -> Result { + self.start_update(target_version, channel, preflight) + } + + async fn update_status(&self) -> Result { + if let Ok(_operation) = self.operation_lock.try_lock() { + recover_interrupted( + &self.config.update_state_file, + &self.config.update_lock_file, + )?; + if self.config.update_support_error().is_none() { + match OperationFileLock::acquire(&self.config.update_lock_file) { + Ok(_lock) => return self.reconcile_installation(), + Err(error) if error.kind() == SystemOperationErrorKind::Conflict => {} + Err(error) => return Err(error), + } + } + } + read_status(&self.config.update_state_file, false) + } + + async fn rollback( + &self, + preflight: Arc, + ) -> Result { + let _operation = self + .operation_lock + .try_lock() + .map_err(|_| conflict("system operation is already running"))?; + if let Some(reason) = self.config.update_support_error() { + return Err(conflict(reason)); + } + let operation_id = operation_id("rollback"); + let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; + let target_version = self + .reconcile_installation()? + .previous_version + .ok_or_else(|| conflict("没有可用于回滚的上一版本"))?; + let release_manifest = + read_release_manifest(&rollback_official_plugins_dir(&self.config)?)?; + let plugin_revision = preflight + .validate_rollback(SystemUpdateCandidate { + target_version, + release_manifest: Arc::clone(&release_manifest), + }) + .await?; + preflight.confirm_revision(plugin_revision).await?; + set_running( + &self.config.update_state_file, + &operation_id, + SystemOperationKind::Rollback, + None, + &self.config.version, + )?; + let result = match rollback_release(&self.config) { + Ok(()) => { + let mut applied = AppliedReleaseGuard::new(&self.config); + let result = match read_release_manifest(&self.config.official_plugins_dir()?) { + Ok(manifest) if manifest.as_ref() == release_manifest.as_ref() => { + preflight.confirm_revision(plugin_revision).await + } + Ok(_) => Err(conflict("已应用的回滚候选与预检候选不一致")), + Err(error) => Err(error), + }; + match result { + Ok(()) => { + applied.commit(); + Ok(()) + } + Err(error) => Err(applied.rollback(error)), + } + } + Err(error) => Err(error), + }; + finish( + &self.config.update_state_file, + &operation_id, + SystemOperationKind::Rollback, + None, + None, + result.as_ref().err().map(ToString::to_string), + )?; + match &result { + Ok(()) => self.events.success_terminal( + Some(&operation_id), + Some("done"), + "previous release restored", + ), + Err(error) => { + self.events + .error_terminal(Some(&operation_id), Some("failed"), error.to_string()) + } + } + drop(file_lock); + result?; + Ok(SystemOperationAccepted::Rollback { + operation_id, + message: "回滚完成，请重启服务。".to_owned(), + need_restart: true, + }) + } + + async fn restart_candidate(&self) -> Result, OperationError> { + let _operation = self + .operation_lock + .try_lock() + .map_err(|_| conflict("system operation is already running"))?; + self.installed_restart_candidate() + } + + async fn restart( + &self, + preflight: Arc, + ) -> Result { + // 与 update/rollback 互斥：更新替换文件期间触发自重启会让新进程 + // 载入半成品产物 + let _operation = self + .operation_lock + .try_lock() + .map_err(|_| conflict("system operation is already running"))?; + if !self.config.self_restart_enabled { + return Err(conflict("self restart is disabled")); + } + let _file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; + let candidate = self.installed_restart_candidate()?; + let files = candidate + .as_ref() + .map(|_| ReleaseFiles::installed(&self.config)) + .transpose()?; + preflight.prepare(candidate).await?; + if let Some(files) = files + && files != ReleaseFiles::installed(&self.config)? + { + return Err(conflict("重启检查期间安装文件已变化，请重新确认")); + } + let message = if self.config.deployment_mode == "docker" { + "已安排进程内重启" + } else { + spawn_replacement(&self.config)?; + "已安排自重启" + }; + let operation_id = operation_id("restart"); + let cancellation = self.cancellation.clone(); + drop(tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(500)).await; + cancellation.cancel(); + })); + Ok(SystemOperationAccepted::Restart { + operation_id, + message: message.to_owned(), + }) + } +} + +fn read_release_manifest(directory: &Path) -> Result, OperationError> { + let path = directory.join(OFFICIAL_PLUGIN_MANIFEST); + let metadata = fs::symlink_metadata(&path) + .map_err(|error| invalid(format!("official plugin manifest is unavailable: {error}")))?; + if !metadata.file_type().is_file() || metadata.file_type().is_symlink() { + return Err(invalid("official plugin manifest is not a regular file")); + } + if metadata.len() == 0 || metadata.len() > MAX_OFFICIAL_PLUGIN_MANIFEST_SIZE { + return Err(invalid("official plugin manifest size is invalid")); + } + fs::read(path) + .map(Arc::<[u8]>::from) + .map_err(|error| invalid(format!("official plugin manifest cannot be read: {error}"))) +} + +fn restore_applied_release(config: &SystemUpdateConfig, cause: OperationError) -> OperationError { + match rollback_release(config) { + Ok(()) => cause, + Err(error) => internal(format!( + "release verification failed and restoring the previous files failed: {error}" + )), + } +} + +struct AppliedReleaseGuard<'a> { + config: &'a SystemUpdateConfig, + armed: bool, +} + +impl<'a> AppliedReleaseGuard<'a> { + const fn new(config: &'a SystemUpdateConfig) -> Self { + Self { + config, + armed: true, + } + } + + const fn commit(&mut self) { + self.armed = false; + } + + fn rollback(mut self, cause: OperationError) -> OperationError { + self.armed = false; + restore_applied_release(self.config, cause) + } +} + +impl Drop for AppliedReleaseGuard<'_> { + fn drop(&mut self) { + if self.armed + && let Err(error) = rollback_release(self.config) + { + tracing::error!(error = %error, "系统发行物临界区取消后的恢复失败"); + } + } +} diff --git a/backend/crates/gateway-host/src/system_update/process.rs b/backend/crates/gateway-host/src/system_update/process.rs index 4f6d8f6c5..d95fe6fd9 100644 --- a/backend/crates/gateway-host/src/system_update/process.rs +++ b/backend/crates/gateway-host/src/system_update/process.rs @@ -4,6 +4,7 @@ use std::env; use std::fs; use std::process::{Command, Stdio}; +use super::config::environment_value; use super::{OperationError, SystemUpdateConfig, internal}; const RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS"; @@ -59,10 +60,3 @@ pub(crate) fn spawn_replacement(config: &SystemUpdateConfig) -> Result<(), Opera .map(|_| ()) .map_err(|error| internal(format!("failed to schedule replacement process: {error}"))) } - -pub(crate) fn environment_value(key: &str) -> Option { - env::var(key) - .ok() - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()) -} diff --git a/backend/crates/gateway-host/src/system_update/release.rs b/backend/crates/gateway-host/src/system_update/release.rs index 69f9114e5..0544d4aef 100644 --- a/backend/crates/gateway-host/src/system_update/release.rs +++ b/backend/crates/gateway-host/src/system_update/release.rs @@ -4,7 +4,9 @@ use std::env; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; -use gateway_admin::model::system::{SystemUpdateChannel as UpdateChannel, SystemUpdateDetail}; +use gateway_admin::model::system::{ + SystemUpdateChannel as UpdateChannel, SystemUpdateDetail, SystemUpdatePolicy, +}; use serde::Deserialize; use tokio::sync::Mutex; @@ -53,12 +55,7 @@ impl ReleaseCache { channel: UpdateChannel, ) -> Result { if let Some(reason) = config.update_support_error() { - return Ok(super::base_update_detail( - config, - channel, - Some(reason), - None, - )); + return Ok(base_update_detail(config, channel, Some(reason), None)); } let repository = config .update_repository @@ -81,14 +78,7 @@ impl ReleaseCache { { Ok(release) => { let detail = release.as_ref().map_or_else( - || { - super::base_update_detail( - config, - channel, - config.update_support_error(), - None, - ) - }, + || base_update_detail(config, channel, config.update_support_error(), None), |release| detail_from_release(config, release, channel), ); let mut entry = self.entry.lock().await; @@ -218,10 +208,10 @@ pub(crate) fn detail_from_release( if unsupported_reason.is_some() || (!has_update && latest_version != normalize_version(&config.version)) { - return super::base_update_detail(config, channel, unsupported_reason, None); + return base_update_detail(config, channel, unsupported_reason, None); } SystemUpdateDetail { - policy: super::update_policy(config, channel), + policy: update_policy(config, channel), current_version: config.version.clone(), latest_version, has_update, @@ -488,3 +478,78 @@ fn url_host_is_loopback(url: &reqwest::Url) -> bool { .is_ok_and(|address| address.is_loopback()) }) } + +impl SystemUpdateConfig { + pub(super) fn update_support_error(&self) -> Option { + if !matches!(self.build_type.as_str(), "release" | "experimental") { + return Some("在线更新需要官方发布构建".to_owned()); + } + let Some(channel) = version_channel(&self.version) else { + return Some("当前版本不符合发行命名规范，无法确定更新通道".to_owned()); + }; + if self.build_type == "experimental" && channel != UpdateChannel::Experimental { + return Some("实验构建必须使用 exp.N 版本，无法在线更新".to_owned()); + } + let Some(repository) = self.update_repository.as_deref() else { + return Some("检查更新需要配置 CPR_UPDATE_REPOSITORY".to_owned()); + }; + if let Err(error) = validate_repository(repository) { + return Some(error.to_string()); + } + if let Err(error) = validate_api_base(&self.github_api_base) { + return Some(error); + } + if let Err(error) = self.web_dist_dir() { + return Some(error.to_string()); + } + None + } + + fn release_cache_key(&self) -> String { + format!( + "{}|{}|{}|{}|{}", + self.update_repository.as_deref().unwrap_or_default(), + self.github_api_base, + self.version, + self.deployment_mode, + self.build_type, + ) + } +} + +pub(super) fn update_policy( + config: &SystemUpdateConfig, + channel: UpdateChannel, +) -> SystemUpdatePolicy { + use UpdateChannel::{Alpha, Beta, Experimental, Rc, Stable}; + SystemUpdatePolicy { + channel, + available_channels: if version_channel(&config.version) == Some(Experimental) { + vec![Experimental] + } else { + vec![Stable, Rc, Beta, Alpha] + }, + } +} + +pub(super) fn base_update_detail( + config: &SystemUpdateConfig, + channel: UpdateChannel, + unsupported_reason: Option, + warning: Option, +) -> SystemUpdateDetail { + SystemUpdateDetail { + policy: update_policy(config, channel), + current_version: config.version.clone(), + latest_version: config.version.clone(), + has_update: false, + deployment_mode: config.deployment_mode.clone(), + build_type: config.build_type.clone(), + release_url: None, + notes: None, + cached: false, + update_supported: unsupported_reason.is_none() && warning.is_none(), + unsupported_reason, + warning, + } +} diff --git a/backend/crates/gateway-host/src/system_update/state.rs b/backend/crates/gateway-host/src/system_update/state.rs index 2daa64721..274208a78 100644 --- a/backend/crates/gateway-host/src/system_update/state.rs +++ b/backend/crates/gateway-host/src/system_update/state.rs @@ -10,6 +10,7 @@ use chrono::{DateTime, Utc}; use gateway_admin::model::system::{ SystemOperationKind, SystemOperationState, SystemOperationStatus, SystemUpdateStatus, }; +use gateway_admin::ports::system::SystemOperationErrorKind; use serde::{Deserialize, Serialize}; use super::installation::ReleaseFiles; @@ -235,7 +236,7 @@ pub(crate) fn recover_interrupted(path: &Path, lock_path: &Path) -> Result<(), O } let _lock = match OperationFileLock::acquire(lock_path) { Ok(lock) => lock, - Err(error) if error.kind() == super::SystemOperationErrorKind::Conflict => return Ok(()), + Err(error) if error.kind() == SystemOperationErrorKind::Conflict => return Ok(()), Err(error) => return Err(error), }; if let (Some(operation_id), Some(kind)) = (state.operation.operation_id, state.operation.kind) { @@ -394,13 +395,6 @@ pub(crate) fn reconcile_installation( read_status(path, need_restart) } -pub(crate) fn default_temp_dir(state_file: &Path) -> PathBuf { - state_file - .parent() - .map(|parent| parent.join("update-tmp")) - .unwrap_or_else(|| std::env::temp_dir().join("codex-proxy-rs-update")) -} - fn stale_lock(path: &Path) -> Result { let modified = fs::metadata(path) .and_then(|metadata| metadata.modified()) diff --git a/backend/crates/gateway-host/src/workers.rs b/backend/crates/gateway-host/src/workers.rs index 989f5f785..889bb5045 100644 --- a/backend/crates/gateway-host/src/workers.rs +++ b/backend/crates/gateway-host/src/workers.rs @@ -117,10 +117,13 @@ impl WorkerSupervisor { /// 不会被二次 poll，逾期的在 abort 后等待其真正终止 pub async fn shutdown(&self, timeout: Duration) { self.cancellation.cancel(); - let handles = std::mem::take(&mut *lock_unpoisoned(&self.handles)); + let mut shutdown = WorkerShutdown { + handles: std::mem::take(&mut *lock_unpoisoned(&self.handles)), + health: &self.health, + }; let deadline = tokio::time::Instant::now() + timeout; - for (_, mut handle) in handles { - if tokio::time::timeout_at(deadline, &mut handle) + for (_, handle) in &mut shutdown.handles { + if tokio::time::timeout_at(deadline, &mut *handle) .await .is_err() { @@ -128,6 +131,21 @@ impl WorkerSupervisor { let _ = handle.await; } } + } +} + +struct WorkerShutdown<'a> { + handles: Vec<(WorkerId, JoinHandle<()>)>, + health: &'a WorkerHealthRegistry, +} + +impl Drop for WorkerShutdown<'_> { + fn drop(&mut self) { + // shutdown future 可在任一 join 点被丢弃；句柄始终留在同一 owner 内， + // 防止当前等待或尚未轮到的任务因 JoinHandle 析构而脱离监督 + for (_, handle) in &self.handles { + handle.abort(); + } self.health.stop_all(); } } diff --git a/backend/crates/gateway-host/tests/command_line.rs b/backend/crates/gateway-host/tests/command_line.rs new file mode 100644 index 000000000..ce776141e --- /dev/null +++ b/backend/crates/gateway-host/tests/command_line.rs @@ -0,0 +1,85 @@ +//! 验证 CLI 不进入 HTTP serve 时的信号取消与 Host 析构 + +use crate::support::host::{configuration, run_in_child}; + +#[test] +fn dropping_command_line_host_cancels_issued_capabilities() { + run_in_child( + "command_line::dropping_command_line_host_cancels_issued_capabilities", + async { + let directory = tempfile::tempdir().unwrap(); + let host = gateway_host::initialize_command_line(configuration(directory.path())) + .await + .unwrap(); + let command = host.cancellation().child_token(); + let connection = host.connection_lifecycle().cancellation(); + + drop(host); + + assert!(command.is_cancelled()); + assert!(connection.is_cancelled()); + }, + ); +} + +#[cfg(unix)] +#[test] +fn sigterm_cancels_command_line_without_http_serve() { + use std::io::{BufRead as _, Write as _}; + use std::process::{Command, Stdio}; + + use crate::support::host::{CHILD_PROCESS_ENV, child_command}; + + const NAME: &str = "command_line::sigterm_cancels_command_line_without_http_serve"; + const READY: &str = "host command signal listener ready"; + + if std::env::var(CHILD_PROCESS_ENV).as_deref() == Ok(NAME) { + run_in_child(NAME, async { + let directory = tempfile::tempdir().unwrap(); + let host = gateway_host::initialize_command_line(configuration(directory.path())) + .await + .unwrap(); + let command = host.cancellation().child_token(); + // SignalGuard 必须先 poll，父进程才可发送信号，避免测试命中默认 SIGTERM 行为 + tokio::task::yield_now().await; + println!("{READY}"); + std::io::stdout().flush().unwrap(); + command.cancelled().await; + assert!(host.connection_lifecycle().cancellation().is_cancelled()); + assert!(host.worker_health().snapshot().is_empty()); + }); + return; + } + + let mut child = child_command(NAME) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let mut output = std::io::BufReader::new(child.stdout.take().unwrap()); + let mut line = String::new(); + loop { + assert_ne!( + output.read_line(&mut line).unwrap(), + 0, + "child was not ready" + ); + if line.contains(READY) { + break; + } + line.clear(); + } + assert!( + Command::new("kill") + .args(["-TERM", &child.id().to_string()]) + .status() + .unwrap() + .success() + ); + let output = child.wait_with_output().unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); +} diff --git a/backend/crates/gateway-host/tests/main.rs b/backend/crates/gateway-host/tests/main.rs index 712529fd7..c71ba8584 100644 --- a/backend/crates/gateway-host/tests/main.rs +++ b/backend/crates/gateway-host/tests/main.rs @@ -1,6 +1,7 @@ //! 宿主配置、网络、进程、日志与后台任务的测试入口 mod client_distribution; +mod command_line; mod config; mod logging; mod official_plugins; diff --git a/backend/crates/gateway-host/tests/serve.rs b/backend/crates/gateway-host/tests/serve.rs index 170d63b30..e6fda0b2a 100644 --- a/backend/crates/gateway-host/tests/serve.rs +++ b/backend/crates/gateway-host/tests/serve.rs @@ -1,12 +1,25 @@ //! 验证宿主服务关闭时的连接排空、超时与并发唤醒 use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; +use axum::{Router, routing::get}; +use futures::future::BoxFuture; +use gateway_core::diagnostics::{OperationalDiagnostics, OperationalFailure}; use gateway_core::lifecycle::CancellationToken; use gateway_core::lifecycle::ConnectionLifecycle as _; +use gateway_core::task::{ + DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerId, WorkerKind, + WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseError, WorkerLeaseRequest, + WorkerRegistration, WorkerRunnable, WorkerTaskError, +}; use gateway_host::HostBundle; use gateway_host::serve::{ConnectionTracker, bind_listener}; +use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; +use tokio::sync::{Semaphore, mpsc}; + +use crate::support::host::{configuration, run_in_child}; #[test] fn host_bundle_serve_is_a_consuming_process_entrypoint() { @@ -95,3 +108,354 @@ async fn bind_listener_should_retry_until_previous_listener_releases_port() { address ); } + +#[test] +fn http_drain_keeps_writer_running_until_the_request_finishes() { + run_in_child( + "serve::http_drain_keeps_writer_running_until_the_request_finishes", + async { + let directory = tempfile::tempdir().unwrap(); + let (host, address) = listening_host(directory.path()).await; + let cancellation = host.cancellation(); + let connections = host.connection_lifecycle(); + let probe = start_workers(&host, false); + probe.started.cancelled().await; + let (router, entered, finish_request) = request_router(&host, &probe); + let server = tokio::spawn(host.serve(router)); + let mut client = connect_request(address).await; + entered.cancelled().await; + + cancellation.cancel(); + wait_for_draining(connections.as_ref()).await; + assert!(connections.try_register().is_err()); + assert!( + tokio::time::timeout(Duration::from_millis(50), probe.cancelled.cancelled()) + .await + .is_err(), + "writer cancelled before the in-flight request finished" + ); + finish_request.add_permits(1); + let mut response = String::new(); + client.read_to_string(&mut response).await.unwrap(); + server.await.unwrap().unwrap(); + + assert!(response.starts_with("HTTP/1.1 200 OK")); + assert_eq!(*probe.writes.lock().unwrap(), [1]); + assert!(probe.finished.is_cancelled()); + }, + ); +} + +#[test] +fn drain_timeout_starts_a_separate_worker_shutdown_phase() { + run_in_child( + "serve::drain_timeout_starts_a_separate_worker_shutdown_phase", + async { + let directory = tempfile::tempdir().unwrap(); + let (host, address) = listening_host(directory.path()).await; + let cancellation = host.cancellation(); + let connections = host.connection_lifecycle(); + let probe = start_workers(&host, true); + probe.started.cancelled().await; + let (router, entered, finish_request) = request_router(&host, &probe); + let server = tokio::spawn(host.serve(router)); + let client = connect_request(address).await; + entered.cancelled().await; + + let started = tokio::time::Instant::now(); + cancellation.cancel(); + wait_for_draining(connections.as_ref()).await; + probe.cancelled.cancelled().await; + assert!(started.elapsed() >= Duration::from_secs(1)); + assert!(!server.is_finished(), "worker cleanup was not joined"); + assert!(!probe.finished.is_cancelled()); + probe.finish_shutdown.add_permits(1); + server.await.unwrap().unwrap(); + assert!(probe.finished.is_cancelled()); + + // 超时后真实连接可以仍在 axum 任务中；释放测试端，避免遗留挂起 handler + drop(client); + finish_request.add_permits(1); + }, + ); +} + +#[test] +fn bind_failure_cancels_and_joins_workers_before_returning() { + run_in_child( + "serve::bind_failure_cancels_and_joins_workers_before_returning", + async { + let directory = tempfile::tempdir().unwrap(); + let mut config = configuration(directory.path()); + config.listen.host = "invalid address".to_owned(); + let host = gateway_host::initialize(config).await.unwrap(); + let cancellation = host.cancellation(); + let probe = start_workers(&host, false); + probe.started.cancelled().await; + + let result = host.serve(Router::new()).await; + + assert!(matches!( + result, + Err(gateway_host::HostError::Serve( + gateway_host::serve::ServeError::Bind(_) + )) + )); + assert!(cancellation.is_cancelled()); + assert!(probe.finished.is_cancelled()); + }, + ); +} + +#[test] +fn dropping_serve_cancels_public_operations_and_aborts_workers() { + run_in_child( + "serve::dropping_serve_cancels_public_operations_and_aborts_workers", + async { + let directory = tempfile::tempdir().unwrap(); + let (host, address) = listening_host(directory.path()).await; + let cancellation = host.cancellation(); + let probe = start_workers(&host, true); + probe.started.cancelled().await; + let server = tokio::spawn(host.serve(Router::new())); + let client = connect_request(address).await; + + server.abort(); + assert!(server.await.unwrap_err().is_cancelled()); + probe.dropped.cancelled().await; + + assert!(cancellation.is_cancelled()); + assert!(!probe.finished.is_cancelled()); + drop(client); + }, + ); +} + +#[test] +fn dropping_serve_during_worker_shutdown_aborts_the_joined_worker() { + run_in_child( + "serve::dropping_serve_during_worker_shutdown_aborts_the_joined_worker", + async { + let directory = tempfile::tempdir().unwrap(); + let (host, address) = listening_host(directory.path()).await; + let cancellation = host.cancellation(); + let probe = start_workers(&host, true); + probe.started.cancelled().await; + let server = tokio::spawn(host.serve(Router::new())); + let client = connect_request(address).await; + + cancellation.cancel(); + probe.cancelled.cancelled().await; + assert!(!probe.finished.is_cancelled()); + assert!( + !server.is_finished(), + "worker shutdown must still be waiting" + ); + // Worker 已观察独立的取消信号，serve 正持有 shutdown 的 join future + server.abort(); + assert!(server.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_secs(1), probe.dropped.cancelled()) + .await + .expect("aborting shutdown must not detach the worker"); + + assert!(!probe.finished.is_cancelled()); + drop(client); + }, + ); +} + +async fn listening_host(directory: &std::path::Path) -> (HostBundle, std::net::SocketAddr) { + let reservation = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let address = reservation.local_addr().unwrap(); + let mut config = configuration(directory); + config.listen.port = address.port(); + let host = gateway_host::initialize(config).await.unwrap(); + drop(reservation); + (host, address) +} + +async fn connect_request(address: std::net::SocketAddr) -> tokio::net::TcpStream { + let mut client = loop { + match tokio::net::TcpStream::connect(address).await { + Ok(client) => break client, + Err(error) if error.kind() == std::io::ErrorKind::ConnectionRefused => { + tokio::time::sleep(Duration::from_millis(5)).await; + } + Err(error) => panic!("connect host: {error}"), + } + }; + client + .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n") + .await + .unwrap(); + client +} + +async fn wait_for_draining(connections: &dyn gateway_core::lifecycle::ConnectionLifecycle) { + while !connections.is_draining() { + tokio::task::yield_now().await; + } +} + +fn request_router( + host: &HostBundle, + probe: &WriterProbe, +) -> (Router, CancellationToken, Arc) { + let entered = CancellationToken::new(); + let finish = Arc::new(Semaphore::new(0)); + let connections = host.connection_lifecycle(); + let sender = probe.sender.clone(); + let router = Router::new().route( + "/", + get({ + let entered = entered.clone(); + let finish = Arc::clone(&finish); + move || { + let entered = entered.clone(); + let finish = Arc::clone(&finish); + let connections = Arc::clone(&connections); + let sender = sender.clone(); + async move { + let _connection = connections.try_register().unwrap(); + entered.cancel(); + finish.acquire().await.unwrap().forget(); + // 对应请求结束后写入终态的时机；drain 超时后的写入允许被拒绝 + if sender.send(1).is_ok() { + "recorded" + } else { + "writer stopped" + } + } + } + }), + ); + (router, entered, finish) +} + +struct WriterProbe { + sender: mpsc::UnboundedSender, + writes: Arc>>, + started: CancellationToken, + cancelled: CancellationToken, + finish_shutdown: Arc, + finished: CancellationToken, + dropped: CancellationToken, +} + +struct WriterTask { + received: tokio::sync::Mutex>, + probe: Arc, +} + +struct TaskLifetime(CancellationToken); + +impl Drop for TaskLifetime { + fn drop(&mut self) { + self.0.cancel(); + } +} + +impl DaemonTask for WriterTask { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let _lifetime = TaskLifetime(self.probe.dropped.clone()); + let mut received = self.received.lock().await; + self.probe.started.cancel(); + loop { + tokio::select! { + biased; + () = cancellation.cancelled() => break, + value = received.recv() => { + self.probe.writes.lock().unwrap().push(value.unwrap()); + } + } + } + self.probe.cancelled.cancel(); + received.close(); + while let Some(value) = received.recv().await { + self.probe.writes.lock().unwrap().push(value); + } + self.probe.finish_shutdown.acquire().await.unwrap().forget(); + self.probe.finished.cancel(); + Ok(()) + }) + } +} + +struct IdleTask; + +impl DaemonTask for IdleTask { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + cancellation.cancelled().await; + Ok(()) + }) + } +} + +struct UnusedLease; + +impl WorkerLeaderLeasePort for UnusedLease { + fn try_acquire( + &self, + _: WorkerLeaseRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Err(WorkerLeaseError::safe("daemon must not acquire a lease")) }) + } +} + +struct Diagnostics; + +#[async_trait::async_trait] +impl OperationalDiagnostics for Diagnostics { + async fn record_failure( + &self, + _: OperationalFailure, + ) -> Result<(), gateway_core::error::StoreError> { + Ok(()) + } +} + +fn start_workers(host: &HostBundle, hold_shutdown: bool) -> Arc { + let (sender, received) = mpsc::unbounded_channel(); + let probe = Arc::new(WriterProbe { + sender, + writes: Arc::new(Mutex::new(Vec::new())), + started: CancellationToken::new(), + cancelled: CancellationToken::new(), + finish_shutdown: Arc::new(Semaphore::new(usize::from(!hold_shutdown))), + finished: CancellationToken::new(), + dropped: CancellationToken::new(), + }); + let mut writer = Some(WriterTask { + received: tokio::sync::Mutex::new(received), + probe: Arc::clone(&probe), + }); + let plan = WorkerKind::ALL + .into_iter() + .map(|kind| { + let task: Box = if kind == WorkerKind::OpsFlush { + Box::new(writer.take().unwrap()) + } else { + Box::new(IdleTask) + }; + WorkerContribution::Registration( + WorkerRegistration::try_new( + WorkerId::try_new(kind, "host_test").unwrap(), + WorkerRunnable::Daemon { + restart: DaemonRestartPolicy::try_new( + Duration::from_secs(1), + Duration::from_secs(2), + ) + .unwrap(), + task, + }, + ) + .unwrap(), + ) + }) + .collect(); + host.start_workers(plan, Arc::new(UnusedLease), Arc::new(Diagnostics)) + .unwrap(); + probe +} diff --git a/backend/crates/gateway-host/tests/support/host.rs b/backend/crates/gateway-host/tests/support/host.rs new file mode 100644 index 000000000..7bf4ae22c --- /dev/null +++ b/backend/crates/gateway-host/tests/support/host.rs @@ -0,0 +1,65 @@ +//! 真实 Host 生命周期测试的配置与独立进程，隔离全局日志订阅器和信号监听 + +use std::{future::Future, path::Path, process::Command, time::Duration}; + +use gateway_host::config::{FileLoggingConfig, HostConfig, ListenConfig, LoggingConfig}; + +pub const CHILD_PROCESS_ENV: &str = "CPR_HOST_LIFECYCLE_TEST_CHILD"; + +pub fn run_in_child(name: &str, scenario: impl Future) { + if std::env::var(CHILD_PROCESS_ENV).as_deref() == Ok(name) { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + tokio::time::timeout(Duration::from_secs(10), scenario) + .await + .expect("host lifecycle scenario timed out"); + }); + return; + } + let output = child_command(name).output().unwrap(); + assert!( + output.status.success(), + "child failed: {}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); +} + +pub fn child_command(name: &str) -> Command { + let mut command = Command::new(std::env::current_exe().unwrap()); + command + .args(["--exact", name, "--nocapture"]) + .env(CHILD_PROCESS_ENV, name) + .env("RUST_LOG", "off"); + command +} + +pub fn configuration(directory: &Path) -> HostConfig { + HostConfig { + timezone: Default::default(), + listen: ListenConfig { + host: "127.0.0.1".to_owned(), + port: 0, + }, + runtime_data_dir: directory.to_path_buf(), + logging: LoggingConfig { + level: "off".to_owned(), + stdout: false, + file: FileLoggingConfig { + enabled: true, + directory: directory.join("logs"), + retention_days: 1, + max_file_size_mb: 1, + }, + oauth_recovery: false, + request_dump: false, + request_dump_retention_days: 1, + }, + system_update: Default::default(), + drain_timeout_seconds: 1, + worker_shutdown_timeout_seconds: 1, + } +} diff --git a/backend/crates/gateway-host/tests/support/mod.rs b/backend/crates/gateway-host/tests/support/mod.rs index 205a7c466..dd1f2b0b0 100644 --- a/backend/crates/gateway-host/tests/support/mod.rs +++ b/backend/crates/gateway-host/tests/support/mod.rs @@ -1,3 +1,4 @@ //! 宿主测试共用的网络与服务辅助入口 +pub mod host; pub mod network; diff --git a/backend/crates/gateway-host/tests/system_update/download.rs b/backend/crates/gateway-host/tests/system_update/download.rs new file mode 100644 index 000000000..16c74c30b --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/download.rs @@ -0,0 +1,14 @@ +//! 更新下载来源信任边界测试 + +use super::*; + +#[test] +fn update_should_trust_github_release_asset_redirect_host() { + assert!( + validate_download_url( + "https://release-assets.githubusercontent.com/github-production-release-asset/archive", + "https://api.github.com/repos", + ) + .is_ok() + ); +} diff --git a/backend/crates/gateway-host/tests/system_update/events.rs b/backend/crates/gateway-host/tests/system_update/events.rs new file mode 100644 index 000000000..f232a2b4f --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/events.rs @@ -0,0 +1,154 @@ +//! 系统更新事件的订阅、阶段进度与终态可见性测试 + +use super::*; + +#[tokio::test] +async fn update_events_should_close_after_terminal_update_log() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.build_type = "source".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let mut events = service.update_events(); + let _ = service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await; + let first = events.next().await.expect("terminal event"); + + assert_eq!(first.level, SystemUpdateEventLevel::Error); + assert!(events.next().await.is_none()); +} + +#[tokio::test] +async fn update_events_should_open_authenticated_sse_stream() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.build_type = "source".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let mut stream = service.update_events(); + let _ = service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await; + + assert!(stream.next().await.is_some()); +} + +#[tokio::test] +async fn update_events_should_preserve_complete_release_stage_sequence() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + let events = service.update_events(); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + let steps = events + .map(|event| event.step.expect("release update step")) + .collect::>() + .await; + + assert_eq!( + steps, + [ + "release", + "prepare", + "asset", + "asset", + "verify", + "prepare", + "download", + "download", + "download", + "checksum", + "checksum", + "extract", + "extract", + "preflight", + "preflight", + "replace", + "replace", + "done", + ] + ); +} + +#[tokio::test] +async fn update_events_should_report_download_bytes_and_percent() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + let events = service.update_events(); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + let progress = events + .filter_map( + |event| async move { event.progress_percent.map(|value| (value, event.message)) }, + ) + .collect::>() + .await; + + assert!(progress.iter().any(|(percent, message)| { + *percent == 100 && message.starts_with("已下载 ") && message.ends_with(" (100%)") + })); +} + +#[tokio::test] +async fn terminal_event_should_only_be_visible_after_status_is_persisted() { + let upstream = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &upstream, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Mismatch, + ) + .await; + let service = fixture.service(&upstream); + let mut events = service.update_events(); + let SystemOperationAccepted::Update { operation_id, .. } = service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect("accepted") + else { + panic!("update") + }; + tokio::time::timeout(Duration::from_secs(5), async { + while let Some(event) = events.next().await { + assert_eq!(event.operation_id.as_deref(), Some(operation_id.as_str())); + if event.terminal { + let status = service.update_status().await.expect("terminal status"); + assert_eq!(status.operation.status, SystemOperationStatus::Failed); + assert!(status.operation.finished_at.is_some()); + return; + } + } + panic!("missing terminal event"); + }) + .await + .expect("terminal event"); +} diff --git a/backend/crates/gateway-host/tests/system_update/mod.rs b/backend/crates/gateway-host/tests/system_update/mod.rs index b3cd7142c..197117bb1 100644 --- a/backend/crates/gateway-host/tests/system_update/mod.rs +++ b/backend/crates/gateway-host/tests/system_update/mod.rs @@ -1,7 +1,12 @@ -//! 系统更新测试入口，以及重启准备和操作互斥测试 +//! 系统更新测试入口与发行包、前置校验和本地安装夹具 mod channels; +mod download; +mod events; mod installation; +mod operations; +mod release; +mod state; use std::fs; use std::path::PathBuf; @@ -164,1789 +169,6 @@ impl SystemUpdatePreflight for BlockingRollbackPreflight { } } -#[tokio::test] -async fn restart_should_not_shutdown_when_replacement_spawn_fails() { - let fixture = Fixture::new(); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.self_restart_enabled = true; - config.executable_path = Some(fixture.root.path().join("missing")); - let shutdown = CancellationToken::new(); - let service = ProcessSystemOperations::new(shutdown.clone(), config); - - assert!( - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .is_err() - ); - assert!(!shutdown.is_cancelled()); -} - -#[tokio::test] -async fn restart_should_request_process_restart_inside_docker() { - let fixture = Fixture::new(); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.self_restart_enabled = true; - config.deployment_mode = "docker".to_owned(); - let shutdown = CancellationToken::new(); - let service = ProcessSystemOperations::new(shutdown.clone(), config); - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect("restart accepted"); - - tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) - .await - .expect("shutdown requested"); -} - -#[tokio::test] -async fn restart_should_spawn_replacement_before_shutdown_outside_docker() { - let fixture = Fixture::new(); - let replacement_marker = fixture.root.path().join("replacement-ran"); - fixture.write_executable("#!/bin/sh\n: > \"${0%/*}/replacement-ran\"\n"); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.self_restart_enabled = true; - let shutdown = CancellationToken::new(); - let service = ProcessSystemOperations::new(shutdown.clone(), config); - - assert_eq!( - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect("replacement scheduled") - .kind(), - SystemOperationKind::Restart - ); - tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) - .await - .expect("shutdown requested after spawn"); - tokio::time::timeout(Duration::from_secs(2), async { - while !replacement_marker.is_file() { - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("replacement executed before fixture cleanup"); -} - -#[cfg(target_os = "linux")] -#[tokio::test] -async fn restart_should_use_startup_path_after_running_executable_is_renamed() { - const CHILD_ENV: &str = "CPR_TEST_RESTART_RENAMED_EXECUTABLE_CHILD"; - let executable = std::env::current_exe().expect("test executable"); - if std::env::var_os(CHILD_ENV).is_none() { - // 用例只移动目录项，不修改旧程序内容；同盘硬链接避免复制整个调试二进制 - let directory = - tempfile::tempdir_in(executable.parent().unwrap()).expect("isolated test directory"); - let child = directory.path().join("restart-test"); - fs::hard_link(&executable, &child).expect("link isolated test executable"); - let output = std::process::Command::new(child) - .args([ - "--exact", - "system_update::restart_should_use_startup_path_after_running_executable_is_renamed", - ]) - .env(CHILD_ENV, "1") - .env("CPR_RESTART_DELAY_MS", "1200") - .output() - .expect("isolated restart test"); - assert!( - output.status.success(), - "stdout: {}\nstderr: {}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); - return; - } - if executable.file_name() == Some(std::ffi::OsStr::new("codex-proxy-rs.backup")) { - // 误启动旧程序时留下可观察结果并退出，避免旧副本再次发起重启 - fs::write( - executable - .parent() - .expect("installation directory") - .join("replacement-version"), - "old", - ) - .expect("old replacement marker"); - return; - } - - let fixture = Fixture::new_in(executable.parent().unwrap()); - fs::rename(executable, fixture.executable()).expect("move isolated running executable"); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.executable_path = None; - config.self_restart_enabled = true; - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let backup = fixture.root.path().join("codex-proxy-rs.backup"); - fs::rename(fixture.executable(), &backup).expect("back up running executable"); - fixture.write_executable("#!/bin/sh\nprintf new > \"${0%/*}/replacement-version\"\n"); - assert_eq!(std::env::current_exe().expect("renamed executable"), backup); - - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect("restart accepted"); - let marker = fixture.root.path().join("replacement-version"); - tokio::time::timeout(Duration::from_secs(5), async { - while !marker.is_file() { - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("replacement executed"); - assert_eq!( - fs::read_to_string(marker).expect("replacement version"), - "new" - ); -} - -#[tokio::test] -async fn restart_should_conflict_while_another_system_operation_is_running() { - let fixture = Fixture::new(); - // 裸 TCP 监听不回包：update 持有操作锁后停在 release 拉取阶段，accept - // 信号保证断言时锁一定已被占用 - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("release listener"); - let api_base = format!("http://{}/repos", listener.local_addr().expect("addr")); - let (connected_tx, connected_rx) = tokio::sync::oneshot::channel(); - let upstream = tokio::spawn(async move { - let connection = listener.accept().await; - let _ = connected_tx.send(()); - tokio::time::sleep(Duration::from_secs(30)).await; - drop(connection); - }); - let mut config = fixture.config(&api_base); - config.self_restart_enabled = true; - config.deployment_mode = "docker".to_owned(); - let shutdown = CancellationToken::new(); - let service = Arc::new(ProcessSystemOperations::new(shutdown.clone(), config)); - let update = tokio::spawn({ - let service = Arc::clone(&service); - async move { - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - } - }); - tokio::time::timeout(Duration::from_secs(5), connected_rx) - .await - .expect("update reaches release fetch while holding the operation lock") - .expect("connected signal"); - - let error = service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect_err("restart during update"); - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert!(!shutdown.is_cancelled()); - - update.await.expect("request").expect("accepted"); - assert!(fixture.lock().exists()); - upstream.abort(); - assert_eq!( - wait_for_update(&service).await.operation.status, - SystemOperationStatus::Failed - ); - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect("restart after lock release"); -} - -#[tokio::test] -async fn rollback_should_restore_binary_web_and_version_state() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - let updated = complete_update(&service, TARGET_VERSION).await; - assert_eq!( - updated.operation.status, - SystemOperationStatus::Succeeded, - "{updated:?}" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - b"new-manifest" - ); - assert!(fixture.official().join("new-plugin.tar.gz").is_file()); - assert!(!fixture.official().join("old-plugin.tar.gz").exists()); - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt as _; - - assert_eq!( - fs::metadata(fixture.official()) - .expect("official directory") - .permissions() - .mode() - & 0o777, - 0o555 - ); - } - service - .rollback(Arc::new(AllowingUpdatePreflight)) - .await - .expect("rollback"); - - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - OLD_RELEASE_MANIFEST.as_bytes() - ); - assert_eq!( - fs::read(fixture.official().join("old-plugin.tar.gz")).expect("old plugin"), - b"old-plugin" - ); - assert!(!fixture.official().join("new-plugin.tar.gz").exists()); - assert_eq!( - service - .update_status() - .await - .expect("status") - .operation - .status, - SystemOperationStatus::Succeeded - ); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded, - "第二次更新必须能清理只读的官方目录备份" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - b"new-manifest" - ); -} - -#[tokio::test] -async fn update_should_restore_all_files_when_plugin_revision_changes_during_swap() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - SystemOperations::perform_update( - &service, - Some(TARGET_VERSION.to_owned()), - None, - Arc::new(ChangingRevisionPreflight { - confirmations: AtomicUsize::new(0), - }), - ) - .await - .expect("update accepted"); - - let status = wait_for_update(&service).await; - assert_eq!(status.operation.status, SystemOperationStatus::Failed); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - OLD_RELEASE_MANIFEST.as_bytes() - ); - assert!(fixture.official().join("old-plugin.tar.gz").is_file()); - assert!(!fixture.official().join("new-plugin.tar.gz").exists()); -} - -#[tokio::test] -async fn rollback_should_restore_current_release_when_plugin_revision_changes_during_swap() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - - let error = service - .rollback(Arc::new(ChangingRollbackRevisionPreflight { - confirmations: AtomicUsize::new(0), - })) - .await - .expect_err("revision changed after rollback swap"); - - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - b"new-manifest" - ); - assert_eq!( - service - .update_status() - .await - .expect("status") - .operation - .status, - SystemOperationStatus::Failed - ); -} - -#[tokio::test] -async fn cancelled_rollback_should_restore_current_release_after_the_swap() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - let reached_after_swap = Arc::new(tokio::sync::Notify::new()); - let task_service = service.clone(); - let task_reached = Arc::clone(&reached_after_swap); - let task = tokio::spawn(async move { - task_service - .rollback(Arc::new(BlockingRollbackPreflight { - confirmations: AtomicUsize::new(0), - reached_after_swap: task_reached, - })) - .await - }); - tokio::time::timeout(Duration::from_secs(5), reached_after_swap.notified()) - .await - .expect("post-swap confirmation reached"); - task.abort(); - assert!( - task.await - .expect_err("rollback task cancelled") - .is_cancelled() - ); - - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - b"new-manifest" - ); - assert_eq!( - service - .update_status() - .await - .expect("recovered status") - .operation - .status, - SystemOperationStatus::Failed - ); -} - -#[tokio::test] -async fn update_detail_should_report_refresh_failure_and_clear_previous_update() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture.mount_release_once(&server, TARGET_VERSION).await; - let service = fixture.service(&server); - service - .update_detail(true, None) - .await - .expect("prime cache"); - - let detail = service - .update_detail(true, None) - .await - .expect("failure detail"); - assert!(!detail.has_update); - assert!(detail.warning.is_some()); - assert!(!detail.cached); - assert_eq!(detail.latest_version, "1.0.0"); - assert!(detail.notes.is_none()); - let version = service - .version() - .await - .expect("version after failed refresh"); - assert!(!version.has_update); - assert!(version.update_warning.is_some()); -} - -#[tokio::test] -async fn update_detail_should_offer_latest_non_draft_release_in_current_channel() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { - "tag_name": "v1.1.0-beta.3", - "name": "Draft 1.1.0-beta.3", - "prerelease": true, - "draft": true, - "assets": [], - }, - { - "tag_name": "v1.1.0-beta.2", - "name": "Release 1.1.0-beta.2", - "body": "notes", - "html_url": "https://github.com/owner/repository/releases", - "prerelease": true, - "assets": [], - }, - ]))) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "1.1.0-beta.1".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.latest_version, "1.1.0-beta.2"); - assert!(detail.has_update); -} - -#[tokio::test] -async fn update_detail_should_reject_untrusted_github_api_base() { - let fixture = Fixture::new(); - let service = ProcessSystemOperations::new( - CancellationToken::new(), - fixture.config("https://api.github.example/repos"), - ); - - let detail = service - .update_detail(true, None) - .await - .expect("safe rejection"); - assert!(!detail.update_supported); - assert!(detail.warning.is_some() || detail.unsupported_reason.is_some()); -} - -#[tokio::test] -async fn update_detail_should_withhold_updates_for_source_builds() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.build_type = "source".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.latest_version, "1.0.0"); - assert!(!detail.update_supported); - assert!(detail.unsupported_reason.is_some()); - assert!(!detail.has_update); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); - assert!(!service.version().await.expect("version").has_update); - assert!( - server - .received_requests() - .await - .expect("requests") - .is_empty() - ); -} - -#[tokio::test] -async fn experimental_build_should_ignore_newer_stable_and_other_experimental_releases() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { "tag_name": "v4.0.0", "prerelease": false }, - { "tag_name": "v3.12.0", "prerelease": false }, - { "tag_name": "v3.10.0", "prerelease": false }, - { "tag_name": "v3.11.0-exp.3", "prerelease": true }, - { "tag_name": "v3.10.1-exp.3", "prerelease": true }, - { "tag_name": "v3.11.0-beta.1", "prerelease": true }, - { "tag_name": "v3.10.0-exp.other.3", "prerelease": true }, - { "tag_name": "v3.10.0-exp.3", "prerelease": true, "draft": true }, - { "tag_name": "v3.10.0-exp.4", "prerelease": false }, - { "tag_name": "invalid", "prerelease": true } - ]))) - .expect(1) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "3.10.0-exp.2".to_owned(); - config.build_type = "experimental".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.latest_version, "3.10.0-exp.2"); - assert!(!detail.has_update); - assert!(detail.update_supported); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); - assert!(detail.warning.is_none()); - let version = service.version().await.expect("version"); - assert!(!version.has_update); - assert_eq!(version.update_channel, "exp"); - let cached = service - .update_detail(false, None) - .await - .expect("cached detail"); - assert!(cached.cached); - assert!(!cached.has_update); -} - -#[tokio::test] -async fn experimental_update_should_find_highest_same_channel_version_across_pages() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut first_page = vec![ - serde_json::json!({ - "tag_name": "v3.12.0", "prerelease": false - }); - 100 - ]; - first_page[1] = serde_json::json!({ "tag_name": "v3.10.0-exp.3", "prerelease": true }); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .and(query_param("page", "1")) - .and(query_param("per_page", "100")) - .respond_with(ResponseTemplate::new(200).set_body_json(first_page)) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .and(query_param("page", "2")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { "tag_name": "v3.10.0-exp.10", "prerelease": true, "body": "exp.10 notes" }, - { "tag_name": "v3.10.0-exp.2", "prerelease": true } - ]))) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "3.10.0-exp.2".to_owned(); - config.build_type = "experimental".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.latest_version, "3.10.0-exp.10"); - assert_eq!(detail.notes.as_deref(), Some("exp.10 notes")); - assert!(detail.has_update); - assert!(detail.update_supported); - assert!(service.version().await.expect("version").has_update); -} - -#[tokio::test] -async fn experimental_update_should_install_same_channel_release() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - "3.10.0-exp.3", - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "3.10.0-exp.2".to_owned(); - config.build_type = "experimental".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - assert_eq!( - complete_update(&service, "3.10.0-exp.3") - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - assert_eq!( - service - .update_status() - .await - .expect("status") - .operation - .status, - SystemOperationStatus::Succeeded - ); -} - -#[tokio::test] -async fn update_should_reject_cross_channel_targets_before_fetching_or_replacing_files() { - for (build_type, current, target) in [ - ("experimental", "3.10.0-exp.2", "3.12.0"), - ("experimental", "3.10.0-exp.2", "3.10.0"), - ("experimental", "3.10.0-exp.2", "3.10.1-exp.3"), - ("experimental", "3.10.0-exp.2", "3.10.0-exp.1"), - ("release", "3.10.0-beta.2", "3.10.0-alpha.3"), - ("release", "3.10.0", "3.10.0+new"), - ("release", "3.10.0", "3.9.0"), - ("experimental", "3.10.0-exp.2", "3.11.0-beta.1"), - ("experimental", "3.10.0-exp.2", "3.10.0-exp.other.3"), - ("release", "3.10.0", "3.11.0-exp.1"), - ("release", "3.10.0", "3.11.0-beta.1"), - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - config.build_type = build_type.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let error = service - .perform_test_update(Some(target.to_owned())) - .await - .expect_err("cross-channel update"); - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert!(error.to_string().contains("所选通道")); - assert!( - server - .received_requests() - .await - .expect("requests") - .is_empty() - ); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); - assert!(!fixture.state().exists()); - } -} - -#[tokio::test] -async fn update_checks_should_follow_stage_promotion_and_experimental_isolation() { - let cases = [ - ("3.12.0-alpha.1", "alpha", [true, true, true, true, false]), - ("3.12.0-beta.1", "beta", [false, true, true, true, false]), - ("3.12.0-rc.1", "rc", [false, false, true, true, false]), - ("3.12.0-exp.1", "exp", [false, false, false, false, true]), - ("3.11.0", "stable", [false, false, false, true, false]), - ]; - for (current, channel, expected) in cases { - for (target, allowed) in [ - "3.12.0-alpha.2", - "3.12.0-beta.2", - "3.12.0-rc.2", - "3.12.0", - "3.12.0-exp.2", - ] - .into_iter() - .zip(expected) - { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture.mount_release_once(&server, target).await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.has_update, allowed, "{current} -> {target}"); - assert_eq!( - detail.latest_version, - if allowed { target } else { current } - ); - assert_eq!(detail.notes.is_some(), allowed); - assert_eq!(detail.release_url.is_some(), allowed); - assert!(detail.warning.is_none()); - let version = service.version().await.expect("version"); - assert_eq!(version.has_update, allowed); - assert_eq!(version.update_channel, channel); - } - } -} - -#[tokio::test] -async fn update_detail_should_keep_current_release_notes_without_allowing_reinstallation() { - for current in [ - "3.12.1", - "3.12.1-alpha.1", - "3.12.1-beta.1", - "3.12.1-rc.1", - "3.12.1-exp.1", - "3.12.1+build.1", - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let release_url = format!("https://github.com/owner/repository/releases/tag/v{current}"); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { - "tag_name": format!("v{current}"), - "prerelease": !semver::Version::parse(current).expect("version").pre.is_empty(), - "body": "## 当前版本\n\n- 修复更新状态", - "html_url": release_url, - } - ]))) - .expect(1) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - if current.contains("-exp.") { - config.build_type = "experimental".to_owned(); - } - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service - .update_detail(true, None) - .await - .expect("current release"); - assert_eq!(detail.latest_version, current); - assert!(!detail.has_update); - assert!(detail.update_supported); - assert_eq!( - detail.notes.as_deref(), - Some("## 当前版本\n\n- 修复更新状态") - ); - assert_eq!(detail.release_url.as_deref(), Some(release_url.as_str())); - let cached = service - .update_detail(false, None) - .await - .expect("cached release"); - assert!(cached.cached); - assert_eq!(cached.notes, detail.notes); - assert_eq!(cached.release_url, detail.release_url); - assert!(!service.version().await.expect("version").has_update); - - let error = service - .perform_test_update(Some(current.to_owned())) - .await - .expect_err("current release must not be installed again"); - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert!(!fixture.state().exists()); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - } -} - -#[tokio::test] -async fn update_detail_should_prefer_available_update_notes_over_current_release() { - for current_first in [true, false] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut releases = vec![ - serde_json::json!({ - "tag_name": "v1.0.0", "prerelease": false, "body": "current notes" - }), - serde_json::json!({ - "tag_name": "v1.9.9", "prerelease": false, "body": "update notes", - "html_url": "https://github.com/owner/repository/releases/tag/v1.9.9" - }), - ]; - if !current_first { - releases.reverse(); - } - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(releases)) - .mount(&server) - .await; - - let detail = fixture - .service(&server) - .update_detail(true, None) - .await - .expect("detail"); - assert!(detail.has_update); - assert_eq!(detail.latest_version, TARGET_VERSION); - assert_eq!(detail.notes.as_deref(), Some("update notes")); - assert_eq!( - detail.release_url.as_deref(), - Some("https://github.com/owner/repository/releases/tag/v1.9.9") - ); - } -} - -#[tokio::test] -async fn update_detail_should_find_current_release_notes_across_pages() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .and(query_param("page", "1")) - .respond_with(ResponseTemplate::new(200).set_body_json(vec![ - serde_json::json!({ "tag_name": "v2.0.0", "prerelease": false }); - 100 - ])) - .expect(1) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .and(query_param("page", "2")) - .respond_with(release_response("1.0.0", Vec::new())) - .expect(1) - .mount(&server) - .await; - - let detail = fixture - .service(&server) - .update_detail(true, None) - .await - .expect("detail"); - assert!(!detail.has_update); - assert_eq!(detail.latest_version, "1.0.0"); - assert_eq!(detail.notes.as_deref(), Some("notes")); - assert!(detail.release_url.is_some()); -} - -#[tokio::test] -async fn update_detail_should_not_use_draft_or_mislabeled_current_release_notes() { - for (current, prerelease) in [("1.0.0", false), ("1.0.0-beta.1", true)] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { "tag_name": format!("v{current}"), "prerelease": prerelease, "draft": true, "body": "draft notes" }, - { "tag_name": format!("v{current}"), "prerelease": !prerelease, "body": "mislabeled notes" }, - { "tag_name": "v0.9.0", "prerelease": false, "body": "older notes" }, - ]))) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert!(!detail.has_update); - assert_eq!(detail.latest_version, current); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); - } -} - -#[tokio::test] -async fn update_checks_should_ignore_downgrades_unknown_stages_and_build_metadata() { - for (current, target) in [ - ("3.12.0-beta.2", "3.12.0-beta.1"), - ("3.12.0", "3.12.0+new-build"), - ("3.12.0-beta.1+aaa", "3.12.0-beta.1+zzz"), - ("3.12.0", "3.11.0"), - ("3.12.0-alpha.1", "3.12.0-preview.2"), - ("3.12.0-alpha.1", "3.12.0-beta"), - ("3.12.0-alpha.1", "3.12.0-beta.0"), - ("3.12.0-alpha.1", "3.12.0-beta.2.extra"), - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture.mount_release_once(&server, target).await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let detail = service.update_detail(true, None).await.expect("detail"); - assert!(!detail.has_update, "{current} -> {target}"); - assert_eq!(detail.latest_version, current); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); - } -} - -#[tokio::test] -async fn update_checks_should_fail_closed_for_unknown_current_channels() { - for (current, build_type) in [ - ("3.12.0-preview.1", "release"), - ("3.12.0-alpha", "release"), - ("3.12.0-exp.0", "experimental"), - ("invalid", "release"), - ("3.12.0", "experimental"), - ("3.12.0-beta.1", "experimental"), - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - config.build_type = build_type.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let detail = service.update_detail(true, None).await.expect("detail"); - assert!(!detail.has_update); - assert!(!detail.update_supported); - assert!(detail.unsupported_reason.is_some()); - assert!( - server - .received_requests() - .await - .expect("requests") - .is_empty() - ); - } -} - -#[tokio::test] -async fn stable_update_should_select_allowed_version_below_newer_major_releases() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { "tag_name": "v4.0.0", "prerelease": false }, - { "tag_name": "v3.13.0-beta.1", "prerelease": true }, - { "tag_name": "v3.12.1", "prerelease": true }, - { "tag_name": "v3.12.0", "prerelease": false, "draft": true }, - { "tag_name": "v3.11.2", "prerelease": false }, - { "tag_name": "v3.11.10", "prerelease": false, "body": "stable notes" } - ]))) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "3.11.0".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let detail = service.update_detail(true, None).await.expect("detail"); - assert!(detail.has_update); - assert_eq!(detail.latest_version, "3.11.10"); - assert_eq!(detail.notes.as_deref(), Some("stable notes")); -} - -#[tokio::test] -async fn prerelease_updates_should_install_later_stages_and_stable_release() { - for (current, target) in [ - ("3.12.0-alpha.1", "3.12.0-beta.1"), - ("3.12.0-beta.1", "3.12.0-rc.1"), - ("3.12.0-rc.1", "3.12.0"), - ("3.12.0-alpha.1", "3.12.0"), - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release(&server, target, ArchiveKind::Safe, ChecksumKind::Valid) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - assert!( - service - .update_detail(true, None) - .await - .expect("detail") - .has_update - ); - assert_eq!( - complete_update(&service, target).await.operation.status, - SystemOperationStatus::Succeeded - ); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - assert_eq!( - service - .update_status() - .await - .expect("status") - .operation - .status, - SystemOperationStatus::Succeeded - ); - } -} - -#[tokio::test] -async fn update_detail_should_use_cached_release_when_not_refreshed() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture.mount_release_once(&server, TARGET_VERSION).await; - let service = fixture.service(&server); - service - .update_detail(true, None) - .await - .expect("prime cache"); - - assert!(service.update_detail(false, None).await.is_ok()); -} - -#[tokio::test] -async fn update_detail_should_withhold_cross_major_release() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release_once(&server, CROSS_MAJOR_VERSION) - .await; - - let detail = fixture - .service(&server) - .update_detail(true, None) - .await - .expect("detail"); - assert_eq!(detail.latest_version, "1.0.0"); - assert!(!detail.has_update); - assert!(detail.warning.is_none()); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); -} - -#[tokio::test] -async fn update_should_reject_cross_major_target_before_fetching_release() { - let fixture = Fixture::new(); - let service = ProcessSystemOperations::new( - CancellationToken::new(), - fixture.config("http://127.0.0.1:1/repos"), - ); - - let error = service - .perform_test_update(Some(CROSS_MAJOR_VERSION.to_owned())) - .await - .expect_err("cross-major target must be rejected"); - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); -} - -#[tokio::test] -async fn update_events_should_close_after_terminal_update_log() { - let fixture = Fixture::new(); - let mut config = fixture.config("https://api.github.com/repos"); - config.build_type = "source".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let mut events = service.update_events(); - let _ = service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await; - let first = events.next().await.expect("terminal event"); - - assert_eq!(first.level, SystemUpdateEventLevel::Error); - assert!(events.next().await.is_none()); -} - -#[tokio::test] -async fn update_events_should_open_authenticated_sse_stream() { - let fixture = Fixture::new(); - let mut config = fixture.config("https://api.github.com/repos"); - config.build_type = "source".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let mut stream = service.update_events(); - let _ = service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await; - - assert!(stream.next().await.is_some()); -} - -#[tokio::test] -async fn update_events_should_preserve_complete_release_stage_sequence() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - let events = service.update_events(); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - let steps = events - .map(|event| event.step.expect("release update step")) - .collect::>() - .await; - - assert_eq!( - steps, - [ - "release", - "prepare", - "asset", - "asset", - "verify", - "prepare", - "download", - "download", - "download", - "checksum", - "checksum", - "extract", - "extract", - "preflight", - "preflight", - "replace", - "replace", - "done", - ] - ); -} - -#[tokio::test] -async fn update_events_should_report_download_bytes_and_percent() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - let events = service.update_events(); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - let progress = events - .filter_map( - |event| async move { event.progress_percent.map(|value| (value, event.message)) }, - ) - .collect::>() - .await; - - assert!(progress.iter().any(|(percent, message)| { - *percent == 100 && message.starts_with("已下载 ") && message.ends_with(" (100%)") - })); -} - -#[tokio::test] -async fn update_should_fail_when_release_checksum_is_missing() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Missing, - ) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_fail_when_release_checksum_mismatches() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Mismatch, - ) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_reject_insecure_release_archive() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_custom_release(&server, "http://github.com/archive", None) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_reject_release_archive_from_untrusted_host() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_custom_release(&server, "https://evil.example/archive", None) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[test] -fn update_should_trust_github_release_asset_redirect_host() { - assert!( - validate_download_url( - "https://release-assets.githubusercontent.com/github-production-release-asset/archive", - "https://api.github.com/repos", - ) - .is_ok() - ); -} - -#[tokio::test] -async fn update_should_reject_release_archive_with_unsafe_path() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::UnsafePath, - ChecksumKind::Valid, - ) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_require_a_flat_regular_official_plugin_bundle() { - for archive_kind in [ - ArchiveKind::MissingOfficialManifest, - ArchiveKind::NestedOfficialAsset, - ArchiveKind::OfficialSymlink, - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release(&server, TARGET_VERSION, archive_kind, ChecksumKind::Valid) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); - } -} - -#[tokio::test] -async fn update_should_reject_untrusted_github_api_base() { - let fixture = Fixture::new(); - let service = ProcessSystemOperations::new( - CancellationToken::new(), - fixture.config("https://api.github.example/repos"), - ); - let error = service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect_err("untrusted API rejected"); - - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); -} - -#[tokio::test] -async fn update_should_reject_when_confirmed_target_differs_from_remote_latest() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release(&server, "1.9.8", ArchiveKind::Safe, ChecksumKind::Valid) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_remove_stale_file_lock_and_continue() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - fs::write(fixture.lock(), "stale").expect("lock"); - filetime::set_file_mtime(fixture.lock(), FileTime::from_unix_time(1, 0)).expect("old mtime"); - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded, - ); -} - -#[tokio::test] -async fn update_and_rollback_should_use_the_default_api_asset_directory() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let mut host: gateway_host::HostConfig = serde_json::from_value(serde_json::json!({ - "listen": { "host": "127.0.0.1", "port": 8080 }, - "runtime_data_dir": "../.runtime/data", - "logging": { - "level": "info", "stdout": true, - "file": { "enabled": false, "directory": "../.runtime/logs", "max_file_size_mb": 20 } - } - })) - .expect("binary host config"); - host.system_update = fixture.config(&format!("{}/repos", server.uri())); - host.system_update.web_dist_dir = None; - host.resolve_and_validate(&fixture.root.path().join("deploy"), &fixture.web()) - .expect("inherit API asset directory"); - let service = ProcessSystemOperations::new(CancellationToken::new(), host.system_update); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - - service - .rollback(Arc::new(AllowingUpdatePreflight)) - .await - .expect("rollback"); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); -} - -#[tokio::test] -async fn update_should_replace_web_assets_across_filesystems() { - let Ok(external) = tempfile::tempdir_in("/dev/shm") else { - return; - }; - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - let web_dist = external.path().join("dist"); - config.web_dist_dir = Some(web_dist.clone()); - fs::create_dir_all(&web_dist).expect("web dir"); - fs::write(web_dist.join("index.html"), "old-web").expect("web"); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - - assert_eq!( - fs::read(web_dist.join("index.html")).expect("web"), - b"new-web" - ); -} - -#[tokio::test] -async fn update_should_restore_web_assets_when_binary_backup_fails() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - fs::create_dir(fixture.executable().with_extension("backup")).expect("blocking backup dir"); - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - OLD_RELEASE_MANIFEST.as_bytes() - ); - assert!(fixture.official().join("old-plugin.tar.gz").is_file()); - assert!(!fixture.official().join("new-plugin.tar.gz").exists()); -} - -#[tokio::test] -async fn update_status_should_reconcile_manual_deployment_without_reusing_legacy_success() { - let fixture = Fixture::new(); - fs::write(fixture.state(), r#"{"previousVersion":"3.14.0","currentVersion":"3.14.1","operation":{"operationId":"legacy","kind":"update","status":"succeeded","targetVersion":"3.14.1"}}"#).expect("legacy state"); - let mut config = fixture.config("https://api.github.com/repos"); - config.version = "3.15.0".to_owned(); - config.deployment_mode = "docker".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let status = service.update_status().await.expect("status"); - assert!(!status.need_restart); - assert_eq!(status.current_version.as_deref(), Some("3.15.0")); - assert!(status.previous_version.is_none()); - assert_eq!(status.operation.target_version.as_deref(), Some("3.14.1")); - assert_eq!(status.operation.status, SystemOperationStatus::Succeeded); -} - -#[tokio::test] -async fn version_should_return_backend_build_metadata() { - let fixture = Fixture::new(); - let mut config = fixture.config("https://api.github.com/repos"); - config.version = "2.3.4".to_owned(); - config.git_sha = "abc123".to_owned(); - config.update_repository = None; - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let version = service.version().await.expect("version"); - - assert_eq!( - (version.version.as_str(), version.git_sha.as_str()), - ("2.3.4", "abc123") - ); -} - -#[tokio::test] -async fn accepted_update_should_survive_a_lost_http_response() { - let upstream = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &upstream, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - Mock::given(method("GET")) - .and(path("/archive")) - .respond_with( - ResponseTemplate::new(200) - .set_body_bytes(release_archive(ArchiveKind::Safe)) - .set_delay(Duration::from_secs(2)), - ) - .with_priority(1) - .mount(&upstream) - .await; - let service = Arc::new(fixture.service(&upstream)); - let handler_service = Arc::clone(&service); - let accepted = Arc::new(tokio::sync::Notify::new()); - let handler_accepted = Arc::clone(&accepted); - let router = axum::Router::new().route( - "/update", - axum::routing::post(move || { - let service = Arc::clone(&handler_service); - let accepted = Arc::clone(&handler_accepted); - async move { - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect("accepted"); - accepted.notify_one(); - // 模拟受理响应丢失，客户端断开会取消此 HTTP handler - std::future::pending::().await - } - }), - ); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("listener"); - let address = listener.local_addr().expect("address"); - let server = tokio::spawn(async move { - axum::serve(listener, router).await.expect("serve"); - }); - let client = tokio::spawn(async move { - reqwest::Client::new() - .post(format!("http://{address}/update")) - .send() - .await - }); - tokio::time::timeout(Duration::from_secs(1), accepted.notified()) - .await - .expect("accepted before slow download finishes"); - client.abort(); - let _ = client.await; - assert_eq!( - service - .update_status() - .await - .expect("running status") - .operation - .status, - SystemOperationStatus::Running - ); - assert_eq!( - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect_err("duplicate update") - .kind(), - SystemOperationErrorKind::Conflict - ); - assert_eq!( - service - .rollback(Arc::new(AllowingUpdatePreflight)) - .await - .expect_err("rollback during update") - .kind(), - SystemOperationErrorKind::Conflict - ); - let status = wait_for_update(&service).await; - assert_eq!(status.operation.status, SystemOperationStatus::Succeeded); - assert!(status.need_restart); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert!(!fixture.lock().exists()); - assert_eq!( - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect_err("restart required") - .kind(), - SystemOperationErrorKind::Conflict - ); - let mut restarted_config = fixture.config(&format!("{}/repos", upstream.uri())); - restarted_config.version = TARGET_VERSION.to_owned(); - let restarted = ProcessSystemOperations::new(CancellationToken::new(), restarted_config); - assert!( - !restarted - .update_status() - .await - .expect("restarted status") - .need_restart - ); - server.abort(); -} - -#[tokio::test] -async fn host_shutdown_should_finish_an_accepted_update_and_release_its_lock() { - let upstream = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &upstream, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let cancellation = CancellationToken::new(); - let service = ProcessSystemOperations::new( - cancellation.clone(), - fixture.config(&format!("{}/repos", upstream.uri())), - ); - // 在后台任务第一次 poll 之前触发关闭，覆盖受理后的生命周期交接 - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect("accepted"); - cancellation.cancel(); - let status = wait_for_update(&service).await; - assert_eq!(status.operation.status, SystemOperationStatus::Failed); - assert!(status.operation.error.expect("reason").contains("中断")); - assert!(!fixture.lock().exists()); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); -} - -#[test] -fn dropped_update_task_should_persist_failure_even_before_its_first_poll() { - let fixture = Fixture::new(); - let service = fixture.service_for_url("http://127.0.0.1:1/repos"); - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("runtime"); - runtime - .block_on(service.perform_test_update(Some(TARGET_VERSION.to_owned()))) - .expect("accepted"); - drop(runtime); - let status: serde_json::Value = - serde_json::from_slice(&fs::read(fixture.state()).expect("state")).expect("JSON"); - assert_eq!(status["operation"]["status"], "failed"); - assert!(status["operation"]["finishedAt"].is_string()); - assert!(!fixture.lock().exists()); -} - -#[tokio::test] -async fn status_should_recover_abandoned_running_state_without_stealing_a_live_lock() { - let fixture = Fixture::new(); - fs::write(fixture.state(), r#"{"currentVersion":"1.0.0","operation":{"operationId":"abandoned","kind":"update","status":"running","targetVersion":"1.9.9"}}"#).expect("state"); - fs::write(fixture.lock(), "another process").expect("live lock"); - let service = fixture.service_for_url("http://127.0.0.1:1/repos"); - assert_eq!( - service - .update_status() - .await - .expect("locked status") - .operation - .status, - SystemOperationStatus::Running - ); - fs::remove_file(fixture.lock()).expect("release abandoned lock"); - let status = service.update_status().await.expect("recovered status"); - assert_eq!(status.operation.status, SystemOperationStatus::Failed); - assert!(status.operation.finished_at.is_some()); - assert_eq!(status.operation.operation_id.as_deref(), Some("abandoned")); -} - -#[tokio::test] -async fn terminal_event_should_only_be_visible_after_status_is_persisted() { - let upstream = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &upstream, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Mismatch, - ) - .await; - let service = fixture.service(&upstream); - let mut events = service.update_events(); - let SystemOperationAccepted::Update { operation_id, .. } = service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect("accepted") - else { - panic!("update") - }; - tokio::time::timeout(Duration::from_secs(5), async { - while let Some(event) = events.next().await { - assert_eq!(event.operation_id.as_deref(), Some(operation_id.as_str())); - if event.terminal { - let status = service.update_status().await.expect("terminal status"); - assert_eq!(status.operation.status, SystemOperationStatus::Failed); - assert!(status.operation.finished_at.is_some()); - return; - } - } - panic!("missing terminal event"); - }) - .await - .expect("terminal event"); -} - async fn wait_for_update(service: &ProcessSystemOperations) -> SystemUpdateStatus { tokio::time::timeout(Duration::from_secs(5), async { loop { @@ -2278,28 +500,3 @@ impl gateway_admin::ports::system::SystemRestartPreflight for RejectingRestartPr )) } } - -#[tokio::test] -async fn restart_checks_plugins_before_shutdown_and_releases_locks_on_rejection() { - let fixture = Fixture::new(); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.self_restart_enabled = true; - config.deployment_mode = "docker".into(); - let shutdown = CancellationToken::new(); - let service = ProcessSystemOperations::new(shutdown.clone(), config); - assert!( - service - .restart(Arc::new(RejectingRestartPreflight)) - .await - .is_err() - ); - assert!(!shutdown.is_cancelled()); - assert!(service.restart_candidate().await.unwrap().is_some()); - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .unwrap(); - tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) - .await - .unwrap(); -} diff --git a/backend/crates/gateway-host/tests/system_update/operations.rs b/backend/crates/gateway-host/tests/system_update/operations.rs new file mode 100644 index 000000000..5387b7ddd --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/operations.rs @@ -0,0 +1,1083 @@ +//! 系统更新、回滚与重启的受理、互斥和文件交换测试 + +use super::*; + +#[tokio::test] +async fn restart_should_not_shutdown_when_replacement_spawn_fails() { + let fixture = Fixture::new(); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + config.executable_path = Some(fixture.root.path().join("missing")); + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + + assert!( + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .is_err() + ); + assert!(!shutdown.is_cancelled()); +} + +#[tokio::test] +async fn restart_should_request_process_restart_inside_docker() { + let fixture = Fixture::new(); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + config.deployment_mode = "docker".to_owned(); + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect("restart accepted"); + + tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) + .await + .expect("shutdown requested"); +} + +#[tokio::test] +async fn restart_should_spawn_replacement_before_shutdown_outside_docker() { + let fixture = Fixture::new(); + let replacement_marker = fixture.root.path().join("replacement-ran"); + fixture.write_executable("#!/bin/sh\n: > \"${0%/*}/replacement-ran\"\n"); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + + assert_eq!( + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect("replacement scheduled") + .kind(), + SystemOperationKind::Restart + ); + tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) + .await + .expect("shutdown requested after spawn"); + tokio::time::timeout(Duration::from_secs(2), async { + while !replacement_marker.is_file() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("replacement executed before fixture cleanup"); +} + +#[cfg(target_os = "linux")] +#[tokio::test] +async fn restart_should_use_startup_path_after_running_executable_is_renamed() { + const CHILD_ENV: &str = "CPR_TEST_RESTART_RENAMED_EXECUTABLE_CHILD"; + let executable = std::env::current_exe().expect("test executable"); + if std::env::var_os(CHILD_ENV).is_none() { + // 用例只移动目录项，不修改旧程序内容；同盘硬链接避免复制整个调试二进制 + let directory = + tempfile::tempdir_in(executable.parent().unwrap()).expect("isolated test directory"); + let child = directory.path().join("restart-test"); + fs::hard_link(&executable, &child).expect("link isolated test executable"); + let output = std::process::Command::new(child) + .args([ + "--exact", + "system_update::restart_should_use_startup_path_after_running_executable_is_renamed", + ]) + .env(CHILD_ENV, "1") + .env("CPR_RESTART_DELAY_MS", "1200") + .output() + .expect("isolated restart test"); + assert!( + output.status.success(), + "stdout: {}\nstderr: {}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + return; + } + if executable.file_name() == Some(std::ffi::OsStr::new("codex-proxy-rs.backup")) { + // 误启动旧程序时留下可观察结果并退出，避免旧副本再次发起重启 + fs::write( + executable + .parent() + .expect("installation directory") + .join("replacement-version"), + "old", + ) + .expect("old replacement marker"); + return; + } + + let fixture = Fixture::new_in(executable.parent().unwrap()); + fs::rename(executable, fixture.executable()).expect("move isolated running executable"); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.executable_path = None; + config.self_restart_enabled = true; + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let backup = fixture.root.path().join("codex-proxy-rs.backup"); + fs::rename(fixture.executable(), &backup).expect("back up running executable"); + fixture.write_executable("#!/bin/sh\nprintf new > \"${0%/*}/replacement-version\"\n"); + assert_eq!(std::env::current_exe().expect("renamed executable"), backup); + + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect("restart accepted"); + let marker = fixture.root.path().join("replacement-version"); + tokio::time::timeout(Duration::from_secs(5), async { + while !marker.is_file() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("replacement executed"); + assert_eq!( + fs::read_to_string(marker).expect("replacement version"), + "new" + ); +} + +#[tokio::test] +async fn restart_should_conflict_while_another_system_operation_is_running() { + let fixture = Fixture::new(); + // 裸 TCP 监听不回包：update 持有操作锁后停在 release 拉取阶段，accept + // 信号保证断言时锁一定已被占用 + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("release listener"); + let api_base = format!("http://{}/repos", listener.local_addr().expect("addr")); + let (connected_tx, connected_rx) = tokio::sync::oneshot::channel(); + let upstream = tokio::spawn(async move { + let connection = listener.accept().await; + let _ = connected_tx.send(()); + tokio::time::sleep(Duration::from_secs(30)).await; + drop(connection); + }); + let mut config = fixture.config(&api_base); + config.self_restart_enabled = true; + config.deployment_mode = "docker".to_owned(); + let shutdown = CancellationToken::new(); + let service = Arc::new(ProcessSystemOperations::new(shutdown.clone(), config)); + let update = tokio::spawn({ + let service = Arc::clone(&service); + async move { + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + } + }); + tokio::time::timeout(Duration::from_secs(5), connected_rx) + .await + .expect("update reaches release fetch while holding the operation lock") + .expect("connected signal"); + + let error = service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect_err("restart during update"); + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert!(!shutdown.is_cancelled()); + + update.await.expect("request").expect("accepted"); + assert!(fixture.lock().exists()); + upstream.abort(); + assert_eq!( + wait_for_update(&service).await.operation.status, + SystemOperationStatus::Failed + ); + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect("restart after lock release"); +} + +#[tokio::test] +async fn rollback_should_restore_binary_web_and_version_state() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + let updated = complete_update(&service, TARGET_VERSION).await; + assert_eq!( + updated.operation.status, + SystemOperationStatus::Succeeded, + "{updated:?}" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + b"new-manifest" + ); + assert!(fixture.official().join("new-plugin.tar.gz").is_file()); + assert!(!fixture.official().join("old-plugin.tar.gz").exists()); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt as _; + + assert_eq!( + fs::metadata(fixture.official()) + .expect("official directory") + .permissions() + .mode() + & 0o777, + 0o555 + ); + } + service + .rollback(Arc::new(AllowingUpdatePreflight)) + .await + .expect("rollback"); + + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + OLD_RELEASE_MANIFEST.as_bytes() + ); + assert_eq!( + fs::read(fixture.official().join("old-plugin.tar.gz")).expect("old plugin"), + b"old-plugin" + ); + assert!(!fixture.official().join("new-plugin.tar.gz").exists()); + assert_eq!( + service + .update_status() + .await + .expect("status") + .operation + .status, + SystemOperationStatus::Succeeded + ); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded, + "第二次更新必须能清理只读的官方目录备份" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + b"new-manifest" + ); +} + +#[tokio::test] +async fn update_should_restore_all_files_when_plugin_revision_changes_during_swap() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + SystemOperations::perform_update( + &service, + Some(TARGET_VERSION.to_owned()), + None, + Arc::new(ChangingRevisionPreflight { + confirmations: AtomicUsize::new(0), + }), + ) + .await + .expect("update accepted"); + + let status = wait_for_update(&service).await; + assert_eq!(status.operation.status, SystemOperationStatus::Failed); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + OLD_RELEASE_MANIFEST.as_bytes() + ); + assert!(fixture.official().join("old-plugin.tar.gz").is_file()); + assert!(!fixture.official().join("new-plugin.tar.gz").exists()); +} + +#[tokio::test] +async fn rollback_should_restore_current_release_when_plugin_revision_changes_during_swap() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + + let error = service + .rollback(Arc::new(ChangingRollbackRevisionPreflight { + confirmations: AtomicUsize::new(0), + })) + .await + .expect_err("revision changed after rollback swap"); + + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + b"new-manifest" + ); + assert_eq!( + service + .update_status() + .await + .expect("status") + .operation + .status, + SystemOperationStatus::Failed + ); +} + +#[tokio::test] +async fn cancelled_rollback_should_restore_current_release_after_the_swap() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + let reached_after_swap = Arc::new(tokio::sync::Notify::new()); + let task_service = service.clone(); + let task_reached = Arc::clone(&reached_after_swap); + let task = tokio::spawn(async move { + task_service + .rollback(Arc::new(BlockingRollbackPreflight { + confirmations: AtomicUsize::new(0), + reached_after_swap: task_reached, + })) + .await + }); + tokio::time::timeout(Duration::from_secs(5), reached_after_swap.notified()) + .await + .expect("post-swap confirmation reached"); + task.abort(); + assert!( + task.await + .expect_err("rollback task cancelled") + .is_cancelled() + ); + + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + b"new-manifest" + ); + assert_eq!( + service + .update_status() + .await + .expect("recovered status") + .operation + .status, + SystemOperationStatus::Failed + ); +} + +#[tokio::test] +async fn experimental_update_should_install_same_channel_release() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + "3.10.0-exp.3", + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "3.10.0-exp.2".to_owned(); + config.build_type = "experimental".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + assert_eq!( + complete_update(&service, "3.10.0-exp.3") + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + assert_eq!( + service + .update_status() + .await + .expect("status") + .operation + .status, + SystemOperationStatus::Succeeded + ); +} + +#[tokio::test] +async fn update_should_reject_cross_channel_targets_before_fetching_or_replacing_files() { + for (build_type, current, target) in [ + ("experimental", "3.10.0-exp.2", "3.12.0"), + ("experimental", "3.10.0-exp.2", "3.10.0"), + ("experimental", "3.10.0-exp.2", "3.10.1-exp.3"), + ("experimental", "3.10.0-exp.2", "3.10.0-exp.1"), + ("release", "3.10.0-beta.2", "3.10.0-alpha.3"), + ("release", "3.10.0", "3.10.0+new"), + ("release", "3.10.0", "3.9.0"), + ("experimental", "3.10.0-exp.2", "3.11.0-beta.1"), + ("experimental", "3.10.0-exp.2", "3.10.0-exp.other.3"), + ("release", "3.10.0", "3.11.0-exp.1"), + ("release", "3.10.0", "3.11.0-beta.1"), + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + config.build_type = build_type.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let error = service + .perform_test_update(Some(target.to_owned())) + .await + .expect_err("cross-channel update"); + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert!(error.to_string().contains("所选通道")); + assert!( + server + .received_requests() + .await + .expect("requests") + .is_empty() + ); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); + assert!(!fixture.state().exists()); + } +} + +#[tokio::test] +async fn prerelease_updates_should_install_later_stages_and_stable_release() { + for (current, target) in [ + ("3.12.0-alpha.1", "3.12.0-beta.1"), + ("3.12.0-beta.1", "3.12.0-rc.1"), + ("3.12.0-rc.1", "3.12.0"), + ("3.12.0-alpha.1", "3.12.0"), + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release(&server, target, ArchiveKind::Safe, ChecksumKind::Valid) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + assert!( + service + .update_detail(true, None) + .await + .expect("detail") + .has_update + ); + assert_eq!( + complete_update(&service, target).await.operation.status, + SystemOperationStatus::Succeeded + ); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + assert_eq!( + service + .update_status() + .await + .expect("status") + .operation + .status, + SystemOperationStatus::Succeeded + ); + } +} + +#[tokio::test] +async fn update_should_reject_cross_major_target_before_fetching_release() { + let fixture = Fixture::new(); + let service = ProcessSystemOperations::new( + CancellationToken::new(), + fixture.config("http://127.0.0.1:1/repos"), + ); + + let error = service + .perform_test_update(Some(CROSS_MAJOR_VERSION.to_owned())) + .await + .expect_err("cross-major target must be rejected"); + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); +} + +#[tokio::test] +async fn update_should_fail_when_release_checksum_is_missing() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Missing, + ) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_fail_when_release_checksum_mismatches() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Mismatch, + ) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_reject_insecure_release_archive() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_custom_release(&server, "http://github.com/archive", None) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_reject_release_archive_from_untrusted_host() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_custom_release(&server, "https://evil.example/archive", None) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_reject_release_archive_with_unsafe_path() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::UnsafePath, + ChecksumKind::Valid, + ) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_require_a_flat_regular_official_plugin_bundle() { + for archive_kind in [ + ArchiveKind::MissingOfficialManifest, + ArchiveKind::NestedOfficialAsset, + ArchiveKind::OfficialSymlink, + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release(&server, TARGET_VERSION, archive_kind, ChecksumKind::Valid) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); + } +} + +#[tokio::test] +async fn update_should_reject_untrusted_github_api_base() { + let fixture = Fixture::new(); + let service = ProcessSystemOperations::new( + CancellationToken::new(), + fixture.config("https://api.github.example/repos"), + ); + let error = service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect_err("untrusted API rejected"); + + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); +} + +#[tokio::test] +async fn update_should_reject_when_confirmed_target_differs_from_remote_latest() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release(&server, "1.9.8", ArchiveKind::Safe, ChecksumKind::Valid) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_remove_stale_file_lock_and_continue() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + fs::write(fixture.lock(), "stale").expect("lock"); + filetime::set_file_mtime(fixture.lock(), FileTime::from_unix_time(1, 0)).expect("old mtime"); + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded, + ); +} + +#[tokio::test] +async fn update_and_rollback_should_use_the_default_api_asset_directory() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let mut host: gateway_host::HostConfig = serde_json::from_value(serde_json::json!({ + "listen": { "host": "127.0.0.1", "port": 8080 }, + "runtime_data_dir": "../.runtime/data", + "logging": { + "level": "info", "stdout": true, + "file": { "enabled": false, "directory": "../.runtime/logs", "max_file_size_mb": 20 } + } + })) + .expect("binary host config"); + host.system_update = fixture.config(&format!("{}/repos", server.uri())); + host.system_update.web_dist_dir = None; + host.resolve_and_validate(&fixture.root.path().join("deploy"), &fixture.web()) + .expect("inherit API asset directory"); + let service = ProcessSystemOperations::new(CancellationToken::new(), host.system_update); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + + service + .rollback(Arc::new(AllowingUpdatePreflight)) + .await + .expect("rollback"); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); +} + +#[tokio::test] +async fn update_should_replace_web_assets_across_filesystems() { + let Ok(external) = tempfile::tempdir_in("/dev/shm") else { + return; + }; + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + let web_dist = external.path().join("dist"); + config.web_dist_dir = Some(web_dist.clone()); + fs::create_dir_all(&web_dist).expect("web dir"); + fs::write(web_dist.join("index.html"), "old-web").expect("web"); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + + assert_eq!( + fs::read(web_dist.join("index.html")).expect("web"), + b"new-web" + ); +} + +#[tokio::test] +async fn update_should_restore_web_assets_when_binary_backup_fails() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + fs::create_dir(fixture.executable().with_extension("backup")).expect("blocking backup dir"); + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + OLD_RELEASE_MANIFEST.as_bytes() + ); + assert!(fixture.official().join("old-plugin.tar.gz").is_file()); + assert!(!fixture.official().join("new-plugin.tar.gz").exists()); +} + +#[tokio::test] +async fn version_should_return_backend_build_metadata() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.version = "2.3.4".to_owned(); + config.git_sha = "abc123".to_owned(); + config.update_repository = None; + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let version = service.version().await.expect("version"); + + assert_eq!( + (version.version.as_str(), version.git_sha.as_str()), + ("2.3.4", "abc123") + ); +} + +#[tokio::test] +async fn accepted_update_should_survive_a_lost_http_response() { + let upstream = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &upstream, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + Mock::given(method("GET")) + .and(path("/archive")) + .respond_with( + ResponseTemplate::new(200) + .set_body_bytes(release_archive(ArchiveKind::Safe)) + .set_delay(Duration::from_secs(2)), + ) + .with_priority(1) + .mount(&upstream) + .await; + let service = Arc::new(fixture.service(&upstream)); + let handler_service = Arc::clone(&service); + let accepted = Arc::new(tokio::sync::Notify::new()); + let handler_accepted = Arc::clone(&accepted); + let router = axum::Router::new().route( + "/update", + axum::routing::post(move || { + let service = Arc::clone(&handler_service); + let accepted = Arc::clone(&handler_accepted); + async move { + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect("accepted"); + accepted.notify_one(); + // 模拟受理响应丢失，客户端断开会取消此 HTTP handler + std::future::pending::().await + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener"); + let address = listener.local_addr().expect("address"); + let server = tokio::spawn(async move { + axum::serve(listener, router).await.expect("serve"); + }); + let client = tokio::spawn(async move { + reqwest::Client::new() + .post(format!("http://{address}/update")) + .send() + .await + }); + tokio::time::timeout(Duration::from_secs(1), accepted.notified()) + .await + .expect("accepted before slow download finishes"); + client.abort(); + let _ = client.await; + assert_eq!( + service + .update_status() + .await + .expect("running status") + .operation + .status, + SystemOperationStatus::Running + ); + assert_eq!( + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect_err("duplicate update") + .kind(), + SystemOperationErrorKind::Conflict + ); + assert_eq!( + service + .rollback(Arc::new(AllowingUpdatePreflight)) + .await + .expect_err("rollback during update") + .kind(), + SystemOperationErrorKind::Conflict + ); + let status = wait_for_update(&service).await; + assert_eq!(status.operation.status, SystemOperationStatus::Succeeded); + assert!(status.need_restart); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert!(!fixture.lock().exists()); + assert_eq!( + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect_err("restart required") + .kind(), + SystemOperationErrorKind::Conflict + ); + let mut restarted_config = fixture.config(&format!("{}/repos", upstream.uri())); + restarted_config.version = TARGET_VERSION.to_owned(); + let restarted = ProcessSystemOperations::new(CancellationToken::new(), restarted_config); + assert!( + !restarted + .update_status() + .await + .expect("restarted status") + .need_restart + ); + server.abort(); +} + +#[tokio::test] +async fn host_shutdown_should_finish_an_accepted_update_and_release_its_lock() { + let upstream = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &upstream, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let cancellation = CancellationToken::new(); + let service = ProcessSystemOperations::new( + cancellation.clone(), + fixture.config(&format!("{}/repos", upstream.uri())), + ); + // 在后台任务第一次 poll 之前触发关闭，覆盖受理后的生命周期交接 + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect("accepted"); + cancellation.cancel(); + let status = wait_for_update(&service).await; + assert_eq!(status.operation.status, SystemOperationStatus::Failed); + assert!(status.operation.error.expect("reason").contains("中断")); + assert!(!fixture.lock().exists()); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); +} + +#[tokio::test] +async fn restart_checks_plugins_before_shutdown_and_releases_locks_on_rejection() { + let fixture = Fixture::new(); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + config.deployment_mode = "docker".into(); + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + assert!( + service + .restart(Arc::new(RejectingRestartPreflight)) + .await + .is_err() + ); + assert!(!shutdown.is_cancelled()); + assert!(service.restart_candidate().await.unwrap().is_some()); + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) + .await + .unwrap(); +} diff --git a/backend/crates/gateway-host/tests/system_update/release.rs b/backend/crates/gateway-host/tests/system_update/release.rs new file mode 100644 index 000000000..b54675d06 --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/release.rs @@ -0,0 +1,512 @@ +//! 发行查询、缓存、版本通道与目标策略测试 + +use super::*; + +#[tokio::test] +async fn update_detail_should_report_refresh_failure_and_clear_previous_update() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, TARGET_VERSION).await; + let service = fixture.service(&server); + service + .update_detail(true, None) + .await + .expect("prime cache"); + + let detail = service + .update_detail(true, None) + .await + .expect("failure detail"); + assert!(!detail.has_update); + assert!(detail.warning.is_some()); + assert!(!detail.cached); + assert_eq!(detail.latest_version, "1.0.0"); + assert!(detail.notes.is_none()); + let version = service + .version() + .await + .expect("version after failed refresh"); + assert!(!version.has_update); + assert!(version.update_warning.is_some()); +} + +#[tokio::test] +async fn update_detail_should_offer_latest_non_draft_release_in_current_channel() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { + "tag_name": "v1.1.0-beta.3", + "name": "Draft 1.1.0-beta.3", + "prerelease": true, + "draft": true, + "assets": [], + }, + { + "tag_name": "v1.1.0-beta.2", + "name": "Release 1.1.0-beta.2", + "body": "notes", + "html_url": "https://github.com/owner/repository/releases", + "prerelease": true, + "assets": [], + }, + ]))) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "1.1.0-beta.1".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.latest_version, "1.1.0-beta.2"); + assert!(detail.has_update); +} + +#[tokio::test] +async fn update_detail_should_reject_untrusted_github_api_base() { + let fixture = Fixture::new(); + let service = ProcessSystemOperations::new( + CancellationToken::new(), + fixture.config("https://api.github.example/repos"), + ); + + let detail = service + .update_detail(true, None) + .await + .expect("safe rejection"); + assert!(!detail.update_supported); + assert!(detail.warning.is_some() || detail.unsupported_reason.is_some()); +} + +#[tokio::test] +async fn update_detail_should_withhold_updates_for_source_builds() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.build_type = "source".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.latest_version, "1.0.0"); + assert!(!detail.update_supported); + assert!(detail.unsupported_reason.is_some()); + assert!(!detail.has_update); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); + assert!(!service.version().await.expect("version").has_update); + assert!( + server + .received_requests() + .await + .expect("requests") + .is_empty() + ); +} + +#[tokio::test] +async fn experimental_build_should_ignore_newer_stable_and_other_experimental_releases() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { "tag_name": "v4.0.0", "prerelease": false }, + { "tag_name": "v3.12.0", "prerelease": false }, + { "tag_name": "v3.10.0", "prerelease": false }, + { "tag_name": "v3.11.0-exp.3", "prerelease": true }, + { "tag_name": "v3.10.1-exp.3", "prerelease": true }, + { "tag_name": "v3.11.0-beta.1", "prerelease": true }, + { "tag_name": "v3.10.0-exp.other.3", "prerelease": true }, + { "tag_name": "v3.10.0-exp.3", "prerelease": true, "draft": true }, + { "tag_name": "v3.10.0-exp.4", "prerelease": false }, + { "tag_name": "invalid", "prerelease": true } + ]))) + .expect(1) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "3.10.0-exp.2".to_owned(); + config.build_type = "experimental".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.latest_version, "3.10.0-exp.2"); + assert!(!detail.has_update); + assert!(detail.update_supported); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); + assert!(detail.warning.is_none()); + let version = service.version().await.expect("version"); + assert!(!version.has_update); + assert_eq!(version.update_channel, "exp"); + let cached = service + .update_detail(false, None) + .await + .expect("cached detail"); + assert!(cached.cached); + assert!(!cached.has_update); +} + +#[tokio::test] +async fn experimental_update_should_find_highest_same_channel_version_across_pages() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut first_page = vec![ + serde_json::json!({ + "tag_name": "v3.12.0", "prerelease": false + }); + 100 + ]; + first_page[1] = serde_json::json!({ "tag_name": "v3.10.0-exp.3", "prerelease": true }); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .and(query_param("page", "1")) + .and(query_param("per_page", "100")) + .respond_with(ResponseTemplate::new(200).set_body_json(first_page)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .and(query_param("page", "2")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { "tag_name": "v3.10.0-exp.10", "prerelease": true, "body": "exp.10 notes" }, + { "tag_name": "v3.10.0-exp.2", "prerelease": true } + ]))) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "3.10.0-exp.2".to_owned(); + config.build_type = "experimental".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.latest_version, "3.10.0-exp.10"); + assert_eq!(detail.notes.as_deref(), Some("exp.10 notes")); + assert!(detail.has_update); + assert!(detail.update_supported); + assert!(service.version().await.expect("version").has_update); +} + +#[tokio::test] +async fn update_checks_should_follow_stage_promotion_and_experimental_isolation() { + let cases = [ + ("3.12.0-alpha.1", "alpha", [true, true, true, true, false]), + ("3.12.0-beta.1", "beta", [false, true, true, true, false]), + ("3.12.0-rc.1", "rc", [false, false, true, true, false]), + ("3.12.0-exp.1", "exp", [false, false, false, false, true]), + ("3.11.0", "stable", [false, false, false, true, false]), + ]; + for (current, channel, expected) in cases { + for (target, allowed) in [ + "3.12.0-alpha.2", + "3.12.0-beta.2", + "3.12.0-rc.2", + "3.12.0", + "3.12.0-exp.2", + ] + .into_iter() + .zip(expected) + { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, target).await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.has_update, allowed, "{current} -> {target}"); + assert_eq!( + detail.latest_version, + if allowed { target } else { current } + ); + assert_eq!(detail.notes.is_some(), allowed); + assert_eq!(detail.release_url.is_some(), allowed); + assert!(detail.warning.is_none()); + let version = service.version().await.expect("version"); + assert_eq!(version.has_update, allowed); + assert_eq!(version.update_channel, channel); + } + } +} + +#[tokio::test] +async fn update_detail_should_keep_current_release_notes_without_allowing_reinstallation() { + for current in [ + "3.12.1", + "3.12.1-alpha.1", + "3.12.1-beta.1", + "3.12.1-rc.1", + "3.12.1-exp.1", + "3.12.1+build.1", + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let release_url = format!("https://github.com/owner/repository/releases/tag/v{current}"); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { + "tag_name": format!("v{current}"), + "prerelease": !semver::Version::parse(current).expect("version").pre.is_empty(), + "body": "## 当前版本\n\n- 修复更新状态", + "html_url": release_url, + } + ]))) + .expect(1) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + if current.contains("-exp.") { + config.build_type = "experimental".to_owned(); + } + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service + .update_detail(true, None) + .await + .expect("current release"); + assert_eq!(detail.latest_version, current); + assert!(!detail.has_update); + assert!(detail.update_supported); + assert_eq!( + detail.notes.as_deref(), + Some("## 当前版本\n\n- 修复更新状态") + ); + assert_eq!(detail.release_url.as_deref(), Some(release_url.as_str())); + let cached = service + .update_detail(false, None) + .await + .expect("cached release"); + assert!(cached.cached); + assert_eq!(cached.notes, detail.notes); + assert_eq!(cached.release_url, detail.release_url); + assert!(!service.version().await.expect("version").has_update); + + let error = service + .perform_test_update(Some(current.to_owned())) + .await + .expect_err("current release must not be installed again"); + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert!(!fixture.state().exists()); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + } +} + +#[tokio::test] +async fn update_detail_should_prefer_available_update_notes_over_current_release() { + for current_first in [true, false] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut releases = vec![ + serde_json::json!({ + "tag_name": "v1.0.0", "prerelease": false, "body": "current notes" + }), + serde_json::json!({ + "tag_name": "v1.9.9", "prerelease": false, "body": "update notes", + "html_url": "https://github.com/owner/repository/releases/tag/v1.9.9" + }), + ]; + if !current_first { + releases.reverse(); + } + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(releases)) + .mount(&server) + .await; + + let detail = fixture + .service(&server) + .update_detail(true, None) + .await + .expect("detail"); + assert!(detail.has_update); + assert_eq!(detail.latest_version, TARGET_VERSION); + assert_eq!(detail.notes.as_deref(), Some("update notes")); + assert_eq!( + detail.release_url.as_deref(), + Some("https://github.com/owner/repository/releases/tag/v1.9.9") + ); + } +} + +#[tokio::test] +async fn update_detail_should_find_current_release_notes_across_pages() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .and(query_param("page", "1")) + .respond_with(ResponseTemplate::new(200).set_body_json(vec![ + serde_json::json!({ "tag_name": "v2.0.0", "prerelease": false }); + 100 + ])) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .and(query_param("page", "2")) + .respond_with(release_response("1.0.0", Vec::new())) + .expect(1) + .mount(&server) + .await; + + let detail = fixture + .service(&server) + .update_detail(true, None) + .await + .expect("detail"); + assert!(!detail.has_update); + assert_eq!(detail.latest_version, "1.0.0"); + assert_eq!(detail.notes.as_deref(), Some("notes")); + assert!(detail.release_url.is_some()); +} + +#[tokio::test] +async fn update_detail_should_not_use_draft_or_mislabeled_current_release_notes() { + for (current, prerelease) in [("1.0.0", false), ("1.0.0-beta.1", true)] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { "tag_name": format!("v{current}"), "prerelease": prerelease, "draft": true, "body": "draft notes" }, + { "tag_name": format!("v{current}"), "prerelease": !prerelease, "body": "mislabeled notes" }, + { "tag_name": "v0.9.0", "prerelease": false, "body": "older notes" }, + ]))) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert!(!detail.has_update); + assert_eq!(detail.latest_version, current); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); + } +} + +#[tokio::test] +async fn update_checks_should_ignore_downgrades_unknown_stages_and_build_metadata() { + for (current, target) in [ + ("3.12.0-beta.2", "3.12.0-beta.1"), + ("3.12.0", "3.12.0+new-build"), + ("3.12.0-beta.1+aaa", "3.12.0-beta.1+zzz"), + ("3.12.0", "3.11.0"), + ("3.12.0-alpha.1", "3.12.0-preview.2"), + ("3.12.0-alpha.1", "3.12.0-beta"), + ("3.12.0-alpha.1", "3.12.0-beta.0"), + ("3.12.0-alpha.1", "3.12.0-beta.2.extra"), + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, target).await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let detail = service.update_detail(true, None).await.expect("detail"); + assert!(!detail.has_update, "{current} -> {target}"); + assert_eq!(detail.latest_version, current); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); + } +} + +#[tokio::test] +async fn update_checks_should_fail_closed_for_unknown_current_channels() { + for (current, build_type) in [ + ("3.12.0-preview.1", "release"), + ("3.12.0-alpha", "release"), + ("3.12.0-exp.0", "experimental"), + ("invalid", "release"), + ("3.12.0", "experimental"), + ("3.12.0-beta.1", "experimental"), + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + config.build_type = build_type.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let detail = service.update_detail(true, None).await.expect("detail"); + assert!(!detail.has_update); + assert!(!detail.update_supported); + assert!(detail.unsupported_reason.is_some()); + assert!( + server + .received_requests() + .await + .expect("requests") + .is_empty() + ); + } +} + +#[tokio::test] +async fn stable_update_should_select_allowed_version_below_newer_major_releases() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { "tag_name": "v4.0.0", "prerelease": false }, + { "tag_name": "v3.13.0-beta.1", "prerelease": true }, + { "tag_name": "v3.12.1", "prerelease": true }, + { "tag_name": "v3.12.0", "prerelease": false, "draft": true }, + { "tag_name": "v3.11.2", "prerelease": false }, + { "tag_name": "v3.11.10", "prerelease": false, "body": "stable notes" } + ]))) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "3.11.0".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let detail = service.update_detail(true, None).await.expect("detail"); + assert!(detail.has_update); + assert_eq!(detail.latest_version, "3.11.10"); + assert_eq!(detail.notes.as_deref(), Some("stable notes")); +} + +#[tokio::test] +async fn update_detail_should_use_cached_release_when_not_refreshed() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, TARGET_VERSION).await; + let service = fixture.service(&server); + service + .update_detail(true, None) + .await + .expect("prime cache"); + + assert!(service.update_detail(false, None).await.is_ok()); +} + +#[tokio::test] +async fn update_detail_should_withhold_cross_major_release() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release_once(&server, CROSS_MAJOR_VERSION) + .await; + + let detail = fixture + .service(&server) + .update_detail(true, None) + .await + .expect("detail"); + assert_eq!(detail.latest_version, "1.0.0"); + assert!(!detail.has_update); + assert!(detail.warning.is_none()); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); +} diff --git a/backend/crates/gateway-host/tests/system_update/state.rs b/backend/crates/gateway-host/tests/system_update/state.rs new file mode 100644 index 000000000..04cf95e83 --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/state.rs @@ -0,0 +1,60 @@ +//! 系统更新状态恢复与任务中止持久化测试 + +use super::*; + +#[tokio::test] +async fn update_status_should_reconcile_manual_deployment_without_reusing_legacy_success() { + let fixture = Fixture::new(); + fs::write(fixture.state(), r#"{"previousVersion":"3.14.0","currentVersion":"3.14.1","operation":{"operationId":"legacy","kind":"update","status":"succeeded","targetVersion":"3.14.1"}}"#).expect("legacy state"); + let mut config = fixture.config("https://api.github.com/repos"); + config.version = "3.15.0".to_owned(); + config.deployment_mode = "docker".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let status = service.update_status().await.expect("status"); + assert!(!status.need_restart); + assert_eq!(status.current_version.as_deref(), Some("3.15.0")); + assert!(status.previous_version.is_none()); + assert_eq!(status.operation.target_version.as_deref(), Some("3.14.1")); + assert_eq!(status.operation.status, SystemOperationStatus::Succeeded); +} + +#[test] +fn dropped_update_task_should_persist_failure_even_before_its_first_poll() { + let fixture = Fixture::new(); + let service = fixture.service_for_url("http://127.0.0.1:1/repos"); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime"); + runtime + .block_on(service.perform_test_update(Some(TARGET_VERSION.to_owned()))) + .expect("accepted"); + drop(runtime); + let status: serde_json::Value = + serde_json::from_slice(&fs::read(fixture.state()).expect("state")).expect("JSON"); + assert_eq!(status["operation"]["status"], "failed"); + assert!(status["operation"]["finishedAt"].is_string()); + assert!(!fixture.lock().exists()); +} + +#[tokio::test] +async fn status_should_recover_abandoned_running_state_without_stealing_a_live_lock() { + let fixture = Fixture::new(); + fs::write(fixture.state(), r#"{"currentVersion":"1.0.0","operation":{"operationId":"abandoned","kind":"update","status":"running","targetVersion":"1.9.9"}}"#).expect("state"); + fs::write(fixture.lock(), "another process").expect("live lock"); + let service = fixture.service_for_url("http://127.0.0.1:1/repos"); + assert_eq!( + service + .update_status() + .await + .expect("locked status") + .operation + .status, + SystemOperationStatus::Running + ); + fs::remove_file(fixture.lock()).expect("release abandoned lock"); + let status = service.update_status().await.expect("recovered status"); + assert_eq!(status.operation.status, SystemOperationStatus::Failed); + assert!(status.operation.finished_at.is_some()); + assert_eq!(status.operation.operation_id.as_deref(), Some("abandoned")); +} diff --git a/backend/crates/gateway-host/tests/workers.rs b/backend/crates/gateway-host/tests/workers.rs index 9ff6ecb3b..f86ff63f1 100644 --- a/backend/crates/gateway-host/tests/workers.rs +++ b/backend/crates/gateway-host/tests/workers.rs @@ -668,6 +668,40 @@ async fn shutdown_timeout_aborts_uncooperative_task_and_drops_guard() { ); } +#[tokio::test] +async fn dropping_shutdown_aborts_current_and_unjoined_workers() { + let task = HungTask { + entered: Arc::new(AtomicBool::new(false)), + notification: Arc::new(Notify::new()), + }; + let leases = Arc::new(FakeLeasePort::default()); + let supervisor = supervisor(); + supervisor + .start( + all_active_plan(task), + leases.clone(), + Arc::new(RecordingDiagnostics::default()), + ) + .expect("complete plan"); + wait_until(|| lock_unpoisoned(&leases.active_guards).len() == ACTIVE_KINDS.len()).await; + + let mut shutdown = Box::pin(supervisor.shutdown(Duration::from_secs(60))); + assert!(futures::poll!(&mut shutdown).is_pending()); + drop(shutdown); + // Supervisor 仍存活，shutdown 自己必须回收当前等待和尚未轮到的所有任务 + wait_until(|| lock_unpoisoned(&leases.active_guards).is_empty()).await; + assert!( + supervisor + .health_source() + .snapshot() + .into_iter() + .all(|entry| matches!( + entry.state, + WorkerRuntimeState::Stopped | WorkerRuntimeState::Disabled + )) + ); +} + #[tokio::test] async fn shutdown_timeout_must_not_rejoin_workers_that_already_exited() { // 挂死任务放在 join 顺序最末的 kind（Retention），保证超时触发时 diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/command_line/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/command_line/mod.rs index aae868186..a3b5aadaa 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/command_line/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/command_line/mod.rs @@ -13,7 +13,7 @@ use gateway_admin::{ model::{AdminError, Revision}, ports::plugins::PluginStore, }; -use gateway_core::runtime::extensions::ExtensionSetReference; +use gateway_core::routing::extensions::ExtensionSetReference; use gateway_plugin_sdk::call::management::CommandDescriptor; use crate::{RpcLimits, RpcSession, callback::PluginCallbacks}; diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs b/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs index 03cbd504d..782006158 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs @@ -16,7 +16,7 @@ use gateway_core::{ WebSocketResponseObservation, }, metering::{CostEstimateStatus, CostSource}, - runtime::extensions::ExtensionSetReference, + routing::extensions::ExtensionSetReference, upstream::UpstreamSendState, }; use gateway_plugin_sdk::{ diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/execution.rs similarity index 96% rename from backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware.rs rename to backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/execution.rs index 121f16d97..30a11da26 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/execution.rs @@ -28,7 +28,7 @@ use crate::{ callback::{MiddlewareBodyAuthority, MiddlewareCompletionBody, MiddlewareInvocation}, }; -use super::{MiddlewareEntry, PluginRequestPolicyPlan}; +use super::super::{MiddlewareEntry, PluginRequestPolicyPlan}; const MAXIMUM_MAPPED_FRAMES_PER_SOURCE: usize = 64; const MAXIMUM_MAPPED_BYTES_PER_SOURCE: usize = 8 * 1024 * 1024; @@ -80,12 +80,15 @@ impl MiddlewarePlan for PluginRequestPolicyPlan { } fn handle_service( &self, - context: gateway_core::middleware::service::Context, - input: gateway_core::middleware::service::Value, - next: gateway_core::middleware::service::Next, + context: gateway_core::engine::middleware::service::Context, + input: gateway_core::engine::middleware::service::Value, + next: gateway_core::engine::middleware::service::Next, ) -> BoxFuture< 'static, - Result, + Result< + gateway_core::engine::middleware::service::Value, + gateway_core::engine::middleware::service::Error, + >, > { self.chain( MiddlewareMount::Service, @@ -104,12 +107,12 @@ impl MiddlewarePlan for PluginRequestPolicyPlan { } fn handle_websocket( &self, - context: gateway_core::middleware::websocket::Context, - message: gateway_core::middleware::websocket::Message, - next: gateway_core::middleware::websocket::Next, + context: gateway_core::engine::middleware::websocket::Context, + message: gateway_core::engine::middleware::websocket::Message, + next: gateway_core::engine::middleware::websocket::Next, ) -> BoxFuture< 'static, - Result, MiddlewareError>, + Result, MiddlewareError>, > { self.chain( MiddlewareMount::WebSocket, @@ -130,10 +133,11 @@ impl MiddlewarePlan for PluginRequestPolicyPlan { fn handle_http( &self, - context: gateway_core::middleware::http::Context, - request: gateway_core::middleware::http::Request, - next: gateway_core::middleware::http::Next, - ) -> BoxFuture<'static, Result> { + context: gateway_core::engine::middleware::http::Context, + request: gateway_core::engine::middleware::http::Request, + next: gateway_core::engine::middleware::http::Next, + ) -> BoxFuture<'static, Result> + { self.chain( MiddlewareMount::Http, context, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/http.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/http.rs similarity index 97% rename from backend/crates/gateway-plugin/runtime/src/adapter/policy/http.rs rename to backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/http.rs index 99586684e..57766a92f 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/http.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/http.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use futures::stream; use gateway_admin::model::plugins::instances::PluginFailurePolicy; -use gateway_core::{engine::middleware::MiddlewareError, middleware::http as core}; +use gateway_core::{engine::middleware::MiddlewareError, engine::middleware::http as core}; use gateway_plugin_sdk::{ Stage, call::middleware::{HANDLE_METHOD, http as wire}, @@ -12,7 +12,7 @@ use gateway_plugin_sdk::{ use http_body::Frame; use http_body_util::{BodyExt as _, StreamBody}; -use super::MiddlewareEntry; +use super::super::MiddlewareEntry; use crate::callback::http_middleware::{Invocation, headers}; pub(super) async fn invoke( diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/mod.rs new file mode 100644 index 000000000..6b4e3cdd6 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/mod.rs @@ -0,0 +1,6 @@ +//! HTTP、WebSocket、逻辑请求与服务中间件的协议适配入口 + +mod execution; +mod http; +mod service; +mod websocket; diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/service.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/service.rs similarity index 96% rename from backend/crates/gateway-plugin/runtime/src/adapter/policy/service.rs rename to backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/service.rs index ef930e6c9..b836a7fd3 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/service.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/service.rs @@ -2,10 +2,10 @@ use std::sync::Arc; -use super::MiddlewareEntry; +use super::super::MiddlewareEntry; use crate::callback::services::Invocation; use gateway_admin::model::plugins::instances::PluginFailurePolicy; -use gateway_core::middleware::service as core; +use gateway_core::engine::middleware::service as core; use gateway_plugin_sdk::{ Stage, call::{middleware::HANDLE_METHOD, services as wire}, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/websocket.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/websocket.rs similarity index 95% rename from backend/crates/gateway-plugin/runtime/src/adapter/policy/websocket.rs rename to backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/websocket.rs index 2daa33bcc..749420dfc 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/websocket.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/websocket.rs @@ -2,10 +2,10 @@ use std::sync::Arc; -use super::MiddlewareEntry; +use super::super::MiddlewareEntry; use crate::callback::websocket_middleware::Invocation; use gateway_admin::model::plugins::instances::PluginFailurePolicy; -use gateway_core::{engine::middleware::MiddlewareError, middleware::websocket as core}; +use gateway_core::{engine::middleware::MiddlewareError, engine::middleware::websocket as core}; use gateway_plugin_sdk::{ Stage, call::middleware::{HANDLE_METHOD, MiddlewareHeader, websocket as wire}, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs index 3a6ba63cc..b212343f0 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs @@ -1,11 +1,8 @@ //! 编译插件请求策略与中间件绑定，组合为请求级执行计划 -mod http; mod middleware; mod retry; mod route_schedule; -mod service; -mod websocket; use std::{collections::BTreeMap, fmt, sync::Arc, time::Duration}; diff --git a/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs b/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs new file mode 100644 index 000000000..308023a36 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs @@ -0,0 +1,128 @@ +//! 单次宿主回调的资源状态、流式预算与统一终结 + +use super::{CallbackScope, MiddlewareCallback, denied, http_middleware, model}; +use gateway_core::lifecycle::CancellationToken; +use gateway_host::outbound::HttpBody; +use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault}; +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, + time::Duration, +}; +use tokio::time::Instant; + +pub(super) struct CallResources { + pub(super) deadline: Instant, + operation_timeout: Duration, + pub(super) cancellation: CancellationToken, + pub(super) scope: Arc, + pub(super) http_resources: Arc, + pub(super) model_bindings: tokio::sync::Mutex< + BTreeMap, + >, + pub(super) state: Mutex, +} + +impl CallResources { + pub(super) fn new( + context: &CallContext, + scope: Arc, + middleware: Option>, + ) -> Self { + Self { + deadline: Instant::now() + Duration::from_millis(context.timeout_ms), + operation_timeout: Duration::from_millis(context.timeout_ms), + cancellation: CancellationToken::new(), + scope, + http_resources: middleware + .as_ref() + .and_then(|middleware| middleware.http_resources()) + .unwrap_or_default(), + model_bindings: tokio::sync::Mutex::new(BTreeMap::new()), + state: Mutex::new(CallState { + middleware, + ..CallState::default() + }), + } + } + + pub(super) fn set_streaming(&self, streaming: bool) { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .resource_stream = streaming; + } + + pub(super) fn close(&self) { + self.cancellation.cancel(); + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.closed = true; + // RPC End 可以先于消费者取完已入队正文到达，这里只撤销新 callback + // 正文包装器持有 invocation，最后一个 owner 释放时再关闭，保留计量和终态信封 + state.middleware.take(); + for stream in state.streams.values() { + stream.close(); + } + state.streams.clear(); + for stream in state.model_streams.values() { + stream.close(); + } + state.model_streams.clear(); + } + + pub(super) fn request_settings( + &self, + ) -> Option { + let middleware = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .middleware + .clone(); + middleware.and_then(|middleware| middleware.request_settings()) + } + + /// 建立受管流后，网络操作各自计时；连接空闲时间不消耗后续操作的预算 + pub(super) fn timeout(&self) -> Result { + let state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.closed { + return Err(denied()); + } + if state.resource_stream { + return Ok(self.operation_timeout); + } + self.deadline + .checked_duration_since(Instant::now()) + .filter(|timeout| !timeout.is_zero()) + .ok_or_else(|| PluginFault::new(ErrorCode::Timeout, "callback deadline elapsed")) + } +} + +#[derive(Default)] +pub(super) struct CallState { + pub(super) closed: bool, + pub(super) resource_stream: bool, + pub(super) middleware: Option>, + pub(super) streams: BTreeMap>, + pub(super) model_streams: BTreeMap>, +} + +pub(super) struct HttpStream { + pub(super) body: tokio::sync::Mutex>, + pub(super) closed: tokio::sync::watch::Sender, +} + +impl HttpStream { + pub(super) fn close(&self) { + self.closed.send_replace(true); + if let Ok(mut body) = self.body.try_lock() { + body.take(); + } + } +} diff --git a/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs b/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs index 83cb05274..1c88e9870 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs @@ -1,7 +1,7 @@ //! 主动 HTTP 子调用只进入 API 分派端口，正文复用当前 RPC 的资源池 use super::{CallResources, http_middleware, invalid, services::ServicePorts}; use crate::RpcReply; -use gateway_core::middleware::http as core; +use gateway_core::engine::middleware::http as core; use gateway_plugin_sdk::{CallContext, PluginFault, call::middleware::http as wire}; pub(super) async fn call( diff --git a/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/mod.rs b/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/mod.rs index 30d2f5ba6..e06fc05dd 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/mod.rs @@ -9,7 +9,7 @@ use std::{ }; use futures::future::BoxFuture; -use gateway_core::{engine::middleware::MiddlewareError, middleware::http as core}; +use gateway_core::{engine::middleware::MiddlewareError, engine::middleware::http as core}; use gateway_plugin_sdk::{ ErrorCode, PluginFault, call::middleware::{MiddlewareHeader, NEXT_METHOD, http as wire}, @@ -27,13 +27,13 @@ pub(crate) struct Invocation { } struct State { - settings: Option, + settings: Option, request: Option, next: Option, } pub(super) fn resolve_settings( - settings: Option, + settings: Option, value: serde_json::Value, timeout_ms: Option, instance_id: &str, @@ -171,7 +171,7 @@ impl Invocation { } impl MiddlewareCallback for Invocation { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) diff --git a/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/websocket.rs b/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/websocket.rs index d908bfbb5..85af02e17 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/websocket.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/websocket.rs @@ -2,7 +2,7 @@ use super::*; use crate::callback::websocket_middleware::Invocation as Messages; -use gateway_core::middleware::{http::upgrade, websocket as ws}; +use gateway_core::engine::middleware::{http::upgrade, websocket as ws}; use gateway_plugin_sdk::call::middleware::websocket as message; use gateway_plugin_sdk::call::middleware::{BODY_CLOSE_METHOD, BODY_READ_METHOD}; diff --git a/backend/crates/gateway-plugin/runtime/src/callback/middleware/mod.rs b/backend/crates/gateway-plugin/runtime/src/callback/middleware/mod.rs index cbb3a675a..23fae55dd 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/middleware/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/middleware/mod.rs @@ -7,7 +7,7 @@ use futures::future::BoxFuture; use std::sync::{Arc, Mutex}; pub(crate) trait MiddlewareCallback: Send + Sync { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { None } @@ -25,7 +25,7 @@ pub(crate) trait MiddlewareCallback: Send + Sync { } impl MiddlewareCallback for MiddlewareInvocation { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) diff --git a/backend/crates/gateway-plugin/runtime/src/callback/mod.rs b/backend/crates/gateway-plugin/runtime/src/callback/mod.rs index 9923bf480..82a654ba1 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/mod.rs @@ -3,6 +3,7 @@ mod accounts; mod admin; mod affinity; +mod call_resources; mod data; pub(crate) mod error; mod facts; @@ -23,19 +24,19 @@ pub(crate) mod websocket_middleware; use std::{ collections::BTreeMap, sync::{Arc, Mutex, Weak}, - time::Duration, }; use futures::future::BoxFuture; use gateway_admin::model::AdminError; use gateway_core::{account::OutboundProxy, lifecycle::CancellationToken}; -use gateway_host::outbound::{HttpBody, HttpClient, NetworkPolicy}; +use gateway_host::outbound::{HttpClient, NetworkPolicy}; use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault}; use tokio::time::Instant; use crate::{CallbackHandler, RpcReply}; pub(crate) use accounts::PluginAccountPortSlot; pub(crate) use affinity::PluginAffinityPortSlot; +use call_resources::{CallResources, HttpStream}; pub(crate) use keys::PluginClientKeyPortSlot; pub(crate) use middleware::{ MiddlewareBinding, MiddlewareBodyAuthority, MiddlewareCallback, MiddlewareCompletionBody, @@ -97,71 +98,6 @@ pub(crate) struct PluginCallbacks { maximum_payload: usize, } -struct CallResources { - deadline: Instant, - operation_timeout: Duration, - cancellation: CancellationToken, - scope: Arc, - http_resources: Arc, - model_bindings: tokio::sync::Mutex< - BTreeMap, - >, - state: Mutex, -} - -impl CallResources { - fn request_settings(&self) -> Option { - let middleware = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .middleware - .clone(); - middleware.and_then(|middleware| middleware.request_settings()) - } - - /// 建立受管流后，网络操作各自计时；连接空闲时间不消耗后续操作的预算 - fn timeout(&self) -> Result { - let state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state.closed { - return Err(denied()); - } - if state.resource_stream { - return Ok(self.operation_timeout); - } - self.deadline - .checked_duration_since(Instant::now()) - .filter(|timeout| !timeout.is_zero()) - .ok_or_else(|| PluginFault::new(ErrorCode::Timeout, "callback deadline elapsed")) - } -} - -#[derive(Default)] -struct CallState { - closed: bool, - resource_stream: bool, - middleware: Option>, - streams: BTreeMap>, - model_streams: BTreeMap>, -} - -struct HttpStream { - body: tokio::sync::Mutex>, - closed: tokio::sync::watch::Sender, -} - -impl HttpStream { - fn close(&self) { - self.closed.send_replace(true); - if let Ok(mut body) = self.body.try_lock() { - body.take(); - } - } -} - impl PluginCallbacks { pub(crate) fn new( instance: &gateway_admin::model::plugins::instances::PluginInstance, @@ -411,21 +347,7 @@ impl CallbackHandler for PluginCallbacks { .unwrap_or_else(std::sync::PoisonError::into_inner) .insert( context.call_id, - Arc::new(CallResources { - deadline: Instant::now() + Duration::from_millis(context.timeout_ms), - operation_timeout: Duration::from_millis(context.timeout_ms), - cancellation: CancellationToken::new(), - scope, - http_resources: middleware - .as_ref() - .and_then(|middleware| middleware.http_resources()) - .unwrap_or_default(), - model_bindings: tokio::sync::Mutex::new(BTreeMap::new()), - state: Mutex::new(CallState { - middleware, - ..CallState::default() - }), - }), + Arc::new(CallResources::new(context, scope, middleware)), ); } @@ -436,10 +358,7 @@ impl CallbackHandler for PluginCallbacks { .unwrap_or_else(std::sync::PoisonError::into_inner) .get(&context.call_id) { - call.state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .resource_stream = context.resource_stream; + call.set_streaming(context.resource_stream); } } @@ -450,25 +369,7 @@ impl CallbackHandler for PluginCallbacks { .unwrap_or_else(std::sync::PoisonError::into_inner) .remove(&context.call_id) { - call.cancellation.cancel(); - let mut state = call - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - state.closed = true; - // RPC End 可以先于消费者取完已入队的中间件正文到达 - // 这里只撤销新的 - // callback 入口；正文包装器持有 invocation，最后一个 owner 释放时再关闭 - // 下游，避免提前丢失尚未搬运的计量/终态信封 - state.middleware.take(); - for stream in state.streams.values() { - stream.close(); - } - state.streams.clear(); - for stream in state.model_streams.values() { - stream.close(); - } - state.model_streams.clear(); + call.close(); } } diff --git a/backend/crates/gateway-plugin/runtime/src/callback/services.rs b/backend/crates/gateway-plugin/runtime/src/callback/services.rs index 9496c779d..01814dd15 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/services.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/services.rs @@ -3,8 +3,8 @@ use super::{CallbackScope, MiddlewareCallback, invalid}; use crate::RpcReply; use futures::future::BoxFuture; -use gateway_admin::{model::AdminError, service::Registry}; -use gateway_core::middleware::service as core; +use gateway_admin::{model::AdminError, public_service::Registry}; +use gateway_core::engine::middleware::service as core; use gateway_plugin_sdk::{ CallContext, PluginFault, call::{middleware::NEXT_METHOD, services as wire}, @@ -14,13 +14,13 @@ use std::sync::{Arc, Mutex, OnceLock, Weak}; #[derive(Default)] pub(crate) struct ServicePorts { registry: OnceLock>, - http: OnceLock>, + http: OnceLock>, } impl ServicePorts { pub(crate) fn bind_http( &self, - dispatcher: &Arc, + dispatcher: &Arc, ) -> Result<(), AdminError> { self.http .set(Arc::downgrade(dispatcher)) @@ -29,7 +29,7 @@ impl ServicePorts { pub(super) fn http( &self, - ) -> Result, PluginFault> { + ) -> Result, PluginFault> { self.http.get().and_then(Weak::upgrade).ok_or_else(|| { PluginFault::new( gateway_plugin_sdk::ErrorCode::Fault, @@ -64,7 +64,7 @@ impl ServicePorts { })?; let call = registry.call(&request.operation, request.input); let origin = scope.origin.as_ref(); - let parent = gateway_admin::service::Origin { + let parent = gateway_admin::public_service::Origin { request_id: origin.map_or_else( || { context @@ -80,11 +80,11 @@ impl ServicePorts { ), cancellation, extensions: scope.child_extensions(&context.instance_id)?, - plan: origin.map_or(gateway_admin::service::Plan::Current, |origin| { - gateway_admin::service::Plan::Frozen(Some(origin.plan.clone())) + plan: origin.map_or(gateway_admin::public_service::Plan::Current, |origin| { + gateway_admin::public_service::Plan::Frozen(Some(origin.plan.clone())) }), }; - let result = gateway_admin::service::scope(parent, call).await; + let result = gateway_admin::public_service::scope(parent, call).await; encode(result) } } diff --git a/backend/crates/gateway-plugin/runtime/src/callback/websocket_middleware.rs b/backend/crates/gateway-plugin/runtime/src/callback/websocket_middleware.rs index 4fa7da028..190a42474 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/websocket_middleware.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/websocket_middleware.rs @@ -3,7 +3,7 @@ use super::MiddlewareCallback; use crate::RpcReply; use bytes::Bytes; use futures::future::BoxFuture; -use gateway_core::middleware::websocket as core; +use gateway_core::engine::middleware::websocket as core; use gateway_plugin_sdk::{ ErrorCode, PluginFault, call::middleware::{ diff --git a/backend/crates/gateway-plugin/runtime/src/compatibility/fast.rs b/backend/crates/gateway-plugin/runtime/src/compatibility/fast.rs index c9d3d187a..c030b4c89 100644 --- a/backend/crates/gateway-plugin/runtime/src/compatibility/fast.rs +++ b/backend/crates/gateway-plugin/runtime/src/compatibility/fast.rs @@ -2,7 +2,7 @@ use gateway_core::{ account::FastMode, - settings::{ExecutionSettings, RequestSettings}, + routing::request_settings::{ExecutionSettings, RequestSettings}, }; use gateway_plugin_sdk::{ Capability, ErrorCode, PluginFault, diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/diagnostics.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/diagnostics.rs new file mode 100644 index 000000000..27a1a63ed --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/diagnostics.rs @@ -0,0 +1,255 @@ +//! 将已发布与准备中的实例事实投影为运行诊断，不参与候选提交 + +use super::super::restart_circuit::RestartIdentity; +use super::{ + PluginRuntime, + set::{PreparedSet, instance_fingerprint}, +}; +use async_trait::async_trait; +use gateway_admin::model::{ + AdminError, AdminErrorKind, Revision, + plugins::instances::{ + PluginInstance, PluginInstanceRuntime, PluginInstanceRuntimeFailure, + PluginInstanceRuntimeStatus, PluginInstanceSnapshot, + }, +}; +use gateway_core::routing::extensions::ExtensionSetReference; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::{Arc, Weak}, +}; + +const MAXIMUM_PREPARATION_DIAGNOSTICS: usize = 8; +#[derive(Clone)] +pub(super) enum PreparationDiagnostic { + Preparing, + Prepared { set_id: String }, + Failed(PluginInstanceRuntimeFailure), +} +impl PluginRuntime { + pub(super) fn record_preparation(&self, revision: u64, diagnostic: PreparationDiagnostic) { + let mut diagnostics = self + .preparation_diagnostics + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + diagnostics.insert(revision, diagnostic); + while diagnostics.len() > MAXIMUM_PREPARATION_DIAGNOSTICS { + diagnostics.pop_first(); + } + } + + fn preparation_diagnostic(&self, revision: u64) -> Option { + self.preparation_diagnostics + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(&revision) + .cloned() + } +} +#[async_trait] +impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for PluginRuntime { + async fn runtime_diagnostics( + &self, + snapshot: &PluginInstanceSnapshot, + published_revision: Option, + published: Option<&ExtensionSetReference>, + ) -> Option> { + let sets = { + let mut prepared = self.prepared.lock().await; + prepared.retain(|_, set| set.strong_count() > 0); + prepared + .values() + .filter_map(Weak::upgrade) + .collect::>() + }; + let open_restart_circuits = snapshot + .instances + .iter() + .filter(|instance| instance.enabled) + .filter_map(|instance| { + let identity = + RestartIdentity::new(instance.id.clone(), instance_fingerprint(instance).ok()?); + self.restart_circuit_is_open(&identity) + .then(|| instance.id.clone()) + }) + .collect::>(); + let published_set = published + .and_then(|reference| sets.iter().find(|set| set.id == *reference.id()).cloned()); + let published_ready = published.is_some_and(ExtensionSetReference::can_serve); + let preparation = self.preparation_diagnostic(snapshot.config_revision.get()); + let projection = InstanceRuntimeProjection { + target_revision: snapshot.config_revision, + published_revision, + published_ready, + published_set: published_set.as_deref(), + sets: &sets, + preparation: preparation.as_ref(), + open_restart_circuits: &open_restart_circuits, + }; + Some( + snapshot + .instances + .iter() + .map(|instance| { + let runtime = projection.project(instance); + (instance.id.clone(), runtime) + }) + .collect(), + ) + } +} +struct InstanceRuntimeProjection<'a> { + target_revision: Revision, + published_revision: Option, + published_ready: bool, + published_set: Option<&'a PreparedSet>, + sets: &'a [Arc], + preparation: Option<&'a PreparationDiagnostic>, + open_restart_circuits: &'a BTreeSet, +} + +impl InstanceRuntimeProjection<'_> { + fn project(&self, expected: &PluginInstance) -> PluginInstanceRuntime { + let active = self.published_set.and_then(|set| { + set.sessions + .iter() + .find(|instance| instance.instance_id == expected.id) + }); + let actual_matches = active.is_some_and(|instance| { + instance.revision == expected.revision + && instance.artifact_sha256 == expected.artifact_sha256 + }); + let target_is_published = self.published_revision == Some(self.target_revision.get()); + let candidate_set = match self.preparation { + Some(PreparationDiagnostic::Prepared { set_id }) => Some(set_id.as_str()), + _ => None, + }; + let published_set_id = self.published_set.map(|set| set.id.as_str()); + let mut draining_revisions = BTreeSet::new(); + for set in self.sets { + if Some(set.id.as_str()) == published_set_id || Some(set.id.as_str()) == candidate_set { + continue; + } + if let Some(instance) = set + .sessions + .iter() + .find(|instance| instance.instance_id == expected.id) + && instance.revision.get() < expected.revision.get() + { + draining_revisions.insert(instance.revision.get()); + } + } + if !expected.enabled + && let Some(instance) = active + && instance.revision != expected.revision + { + draining_revisions.insert(instance.revision.get()); + } + + let diagnostic = active.map(|instance| instance.session.diagnostic()); + let mut failure = None; + let status = if !expected.enabled { + if active.is_some() || !draining_revisions.is_empty() { + PluginInstanceRuntimeStatus::Draining + } else { + PluginInstanceRuntimeStatus::Disabled + } + } else if target_is_published + && let Some(reason) = self + .published_set + .and_then(|set| set.failures.get(&expected.id)) + { + failure = Some(reason.clone()); + PluginInstanceRuntimeStatus::PreparationFailed + } else if self.open_restart_circuits.contains(&expected.id) { + let message = match diagnostic { + Some(crate::rpc::RpcSessionDiagnostic::Failed { message, .. }) => { + format!("{message}，已暂停自动重启") + } + _ => "插件实例连续异常退出，已暂停自动重启".to_owned(), + }; + failure = Some(runtime_failure("restart_circuit_open", &message)); + PluginInstanceRuntimeStatus::Faulted + } else if actual_matches + && let Some(crate::rpc::RpcSessionDiagnostic::Failed { code, message }) = diagnostic + { + failure = Some(runtime_failure(code, message)); + PluginInstanceRuntimeStatus::Faulted + } else if !target_is_published { + match self.preparation { + Some(PreparationDiagnostic::Preparing) => PluginInstanceRuntimeStatus::Preparing, + Some(PreparationDiagnostic::Failed(reason)) => { + failure = Some(reason.clone()); + PluginInstanceRuntimeStatus::PreparationFailed + } + Some(PreparationDiagnostic::Prepared { .. }) | None => { + PluginInstanceRuntimeStatus::AwaitingPublication + } + } + } else if !actual_matches { + failure = Some(runtime_failure( + "published_instance_mismatch", + "已发布插件实例与期望版本不一致", + )); + PluginInstanceRuntimeStatus::Blocked + } else { + match diagnostic { + Some(crate::rpc::RpcSessionDiagnostic::Ready) if self.published_ready => { + PluginInstanceRuntimeStatus::Running + } + Some(crate::rpc::RpcSessionDiagnostic::Ready) => { + failure = Some(runtime_failure( + "published_set_unready", + "同一发布集合中存在未就绪插件实例", + )); + PluginInstanceRuntimeStatus::Blocked + } + Some(crate::rpc::RpcSessionDiagnostic::Quiescing) => { + PluginInstanceRuntimeStatus::Draining + } + Some(crate::rpc::RpcSessionDiagnostic::Failed { code, message }) => { + failure = Some(runtime_failure(code, message)); + PluginInstanceRuntimeStatus::Faulted + } + None => { + failure = Some(runtime_failure( + "instance_not_published", + "已发布集合缺少该插件实例", + )); + PluginInstanceRuntimeStatus::Blocked + } + } + }; + + PluginInstanceRuntime { + status, + actual_revision: active.map(|instance| instance.revision.get()), + actual_artifact_sha256: active.map(|instance| instance.artifact_sha256.clone()), + failure, + draining_revisions: draining_revisions.into_iter().collect(), + } + } +} + +fn runtime_failure(code: &str, message: &str) -> PluginInstanceRuntimeFailure { + PluginInstanceRuntimeFailure { + code: code.to_owned(), + message: message.to_owned(), + } +} + +pub(super) fn runtime_failure_from_admin(error: &AdminError) -> PluginInstanceRuntimeFailure { + let code = match error.kind() { + AdminErrorKind::Invalid => "invalid", + AdminErrorKind::Unauthorized => "unauthorized", + AdminErrorKind::Forbidden => "forbidden", + AdminErrorKind::NotFound => "not_found", + AdminErrorKind::Conflict => "conflict", + AdminErrorKind::RateLimited => "rate_limited", + AdminErrorKind::BadGateway => "bad_gateway", + AdminErrorKind::UpstreamResultUnknown => "upstream_result_unknown", + AdminErrorKind::Unavailable => "unavailable", + AdminErrorKind::Internal => "internal", + }; + runtime_failure(code, error.message()) +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs index 1dd82987d..5331e1369 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs @@ -1,11 +1,20 @@ //! 准备单个插件实例的进程、RPC 会话与能力贡献 +use super::super::restart_circuit::RestartIdentity; use super::{ - AdminError, Arc, Duration, Handshake, PluginCallbackPorts, PluginCallbacks, PluginInstance, - PluginPrivateState, PluginRuntime, PluginStateStoreErrorKind, PreparedContributions, - PreparedInstance, Registration, RestartIdentity, Revision, RpcSession, Stage, ValidatedPackage, - instance_fingerprint, + PluginRuntime, + set::{PreparedContributions, PreparedInstance, instance_fingerprint}, }; +use crate::{ + RpcSession, ValidatedPackage, + callback::{PluginCallbackPorts, PluginCallbacks, private_state::PluginPrivateState}, +}; +use gateway_admin::{ + model::{AdminError, Revision, plugins::instances::PluginInstance}, + ports::plugins::PluginStateStoreErrorKind, +}; +use gateway_plugin_sdk::{Handshake, Stage, call::registration::Registration}; +use std::{sync::Arc, time::Duration}; impl PluginRuntime { pub(super) async fn prepare_instance( diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs index a91e1012c..edb32f731 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs @@ -1,6 +1,7 @@ //! 对账只从已发布集合开始；通知合并，实例串行，失败和停用不阻塞其他实例 -use super::{PluginRuntime, PreparedSet, RpcSession}; +use super::{PluginRuntime, set::PreparedSet}; +use crate::RpcSession; use futures::{StreamExt as _, future::BoxFuture}; use gateway_core::{ lifecycle::CancellationToken, diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/management.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/management.rs new file mode 100644 index 000000000..2a18fddec --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/management.rs @@ -0,0 +1,137 @@ +//! 按发布集合适配插件管理接口与命令行会话 + +use super::PluginRuntime; +use async_trait::async_trait; +use gateway_admin::model::AdminError; +use gateway_core::routing::extensions::ExtensionSetReference; + +impl PluginRuntime { + /// CLI 复用同一候选准备与弱索引；读取帮助不会发布 Core 快照或执行命令 + pub async fn prepare_command_line(&self) -> Result { + let snapshot = self + .store + .load_instances() + .await + .map_err(|_| AdminError::unavailable("无法读取已安装插件"))?; + let reference = self.prepare_snapshot(snapshot, None).await?; + let set = self.prepared_set(&reference).await?; + Ok(crate::PluginCommandSession::new( + reference, + set.commands.clone(), + set.sessions + .iter() + .map(|instance| instance.session.clone()) + .collect(), + self.store.clone(), + self.config.rpc_limits, + )) + } +} +#[async_trait] +impl gateway_admin::ports::plugin_management::PluginManagement for PluginRuntime { + async fn validate_target( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + ) -> Result<(), AdminError> { + let set = self.prepared_set(published).await?; + set.management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))?; + + Ok(()) + } + + async fn start_callback( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + command: gateway_admin::model::plugins::management::StartPluginManagementCallback, + context: &gateway_admin::model::auth::AdminRequestContext, + ) -> Result + { + let store = self + .oauth_pending + .as_ref() + .ok_or_else(|| AdminError::unavailable("插件登录状态端口不可用"))?; + self.prepared_set(published) + .await? + .management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? + .start_callback(store.as_ref(), command, context) + .await + } + + async fn callback( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + state: &str, + request: gateway_admin::model::plugins::management::PluginManagementRequest, + ) -> Result + { + let store = self + .oauth_pending + .as_ref() + .ok_or_else(|| AdminError::unavailable("插件登录状态端口不可用"))?; + self.prepared_set(published) + .await? + .management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? + .callback(store.as_ref(), state, request, self.config.rpc_limits) + .await + } + + async fn views( + &self, + published: &ExtensionSetReference, + ) -> Result, AdminError> + { + Ok(self + .prepared_set(published) + .await? + .management + .iter() + .map(|entry| entry.view.clone()) + .collect()) + } + + async fn resource( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + path: &str, + public: bool, + ) -> Result + { + self.prepared_set(published) + .await? + .management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? + .resource(path, public) + } + + async fn handle( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + request: gateway_admin::model::plugins::management::PluginManagementRequest, + ) -> Result + { + self.prepared_set(published) + .await? + .management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? + .handle(request, self.config.rpc_limits) + .await + } +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/mod.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/mod.rs index f6cf6cfdc..c22a17afd 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/mod.rs @@ -1,1267 +1,12 @@ -//! 组装插件运行时端口，准备并持有可发布的实例与能力集合 +//! 插件运行时端口与代次准备的公共入口 +mod diagnostics; mod instance; mod maintenance; +mod management; +mod runtime; +mod set; +mod snapshot; +mod state; -use std::{ - collections::{BTreeMap, BTreeSet}, - path::PathBuf, - sync::{ - Arc, Mutex as SyncMutex, Weak, - atomic::{AtomicBool, Ordering}, - }, - time::Duration, -}; - -use async_trait::async_trait; -use futures::future::BoxFuture; -use gateway_admin::{ - model::{ - AdminError, AdminErrorKind, Revision, - plugins::{ - instances::{ - PluginInstance, PluginInstanceRuntime, PluginInstanceRuntimeFailure, - PluginInstanceRuntimeStatus, PluginInstanceSnapshot, - }, - state::{ - ApplyPluginStateMigration, PluginStateConfiguration, PluginStateMigrationAction, - PluginStateMigrationChange, PluginStateTransition, - }, - }, - }, - ports::{ - plugin_accounts::PluginAccountAccess, - plugin_client_keys::PluginClientKeyAccess, - plugins::{PluginPreparation, PluginStateStore, PluginStateStoreErrorKind, PluginStore}, - }, -}; -use gateway_core::{ - engine::{ - authentication::FrontendAuthenticationExtensionIndex, - middleware::{MiddlewareExtensionIndex, MiddlewarePlan}, - nested::{AffinityLookupPort, NestedModelExecutionPort}, - observation::{RequestObserverExtensionIndex, RequestObserverPlan}, - policy::{RequestPolicyExtensionIndex, RequestPolicyPlan}, - }, - provider_ports::OAuthPendingFlowPort, - routing::ConfigRevision, - runtime::extensions::{ - ExtensionPreparationError, ExtensionPreparationPort, ExtensionSetId, ExtensionSetLease, - ExtensionSetReference, - }, -}; -use gateway_plugin_sdk::call::registration::Registration; -use gateway_plugin_sdk::{ - Handshake, Stage, - call::host::{ - StateMigrationChange, StateMigrationRecord, StateMigrationRequest, StateMigrationResult, - }, -}; -use secrecy::ExposeSecret as _; -use sha2::{Digest as _, Sha256}; -use tokio::sync::{Mutex, Semaphore}; - -use crate::{ - PackageLimits, RpcLimits, RpcSession, ValidatedPackage, - callback::{ - PluginAccountPortSlot, PluginAffinityPortSlot, PluginCallbackPorts, PluginCallbacks, - PluginClientKeyPortSlot, PluginModelPortSlot, private_state::PluginPrivateState, - }, -}; - -use super::restart_circuit::{PluginRestartCircuitConfig, RestartCircuits, RestartIdentity}; - -const RUNTIME_SHUTDOWN_GRACE: Duration = Duration::from_secs(2); -const MAXIMUM_PREPARATION_DIAGNOSTICS: usize = 8; - -pub struct PluginRuntimeConfig { - pub cache_directory: PathBuf, - pub host_version: semver::Version, - pub package_limits: PackageLimits, - pub rpc_limits: RpcLimits, - pub restart_circuit: PluginRestartCircuitConfig, -} - -/// 只索引有主人的候选；当前代次仅由 Core 的发布视图持有 -pub struct PluginRuntime { - service_ports: Arc, - pub(super) store: Arc, - state: Arc, - config: PluginRuntimeConfig, - prepare_lock: Mutex<()>, - prepared: Mutex>>, - preparation_diagnostics: SyncMutex>, - restart_circuits: Arc, - shutting_down: Arc, - shutdown_lock: Mutex<()>, - processes: Arc, - validators: Arc, - log_slots: Arc, - pub(super) account_ports: Arc, - client_key_ports: Arc, - resource_ports: Arc, - model_ports: Arc, - affinity_ports: Arc, - observers: RequestObserverExtensionIndex, - policies: RequestPolicyExtensionIndex, - middleware: MiddlewareExtensionIndex, - authentications: FrontendAuthenticationExtensionIndex, - oauth_pending: Option>, - http: Arc, -} - -#[derive(Clone)] -enum PreparationDiagnostic { - Preparing, - Prepared { set_id: String }, - Failed(PluginInstanceRuntimeFailure), -} - -pub(super) struct PreparedSet { - id: ExtensionSetId, - _observers: Option>, - _policies: Option>, - _middleware: Option>, - upstream_adapters: Option>, - _authentication: - Option>, - sessions: Vec, - failures: BTreeMap, - shutting_down: Arc, - commands: Vec>, - management: Vec, - model_aliases: Vec, -} - -#[derive(Default)] -struct PreparedContributions { - sessions: Vec, - observer_entries: Vec, - commands: Vec>, - management: Vec, - policy_entries: Vec, - upstream_entries: Vec, - authentication_entries: - Vec, - model_aliases: Vec, -} - -impl PreparedContributions { - fn append(&mut self, other: Self) { - self.sessions.extend(other.sessions); - self.observer_entries.extend(other.observer_entries); - self.commands.extend(other.commands); - self.management.extend(other.management); - self.model_aliases.extend(other.model_aliases); - self.policy_entries.extend(other.policy_entries); - self.upstream_entries.extend(other.upstream_entries); - self.authentication_entries - .extend(other.authentication_entries); - } -} - -struct PreparedInstance { - instance_id: String, - artifact_sha256: String, - revision: Revision, - session: Arc, - private_state: Arc, - maintenance: bool, -} - -impl PreparedSet { - fn sessions_ready(&self) -> bool { - self.sessions - .iter() - .all(|instance| instance.session.is_ready()) - } -} - -impl ExtensionSetLease for PreparedSet { - fn upstream_adapters( - &self, - ) -> Option> { - self.upstream_adapters.clone() - } - - fn is_ready(&self) -> bool { - self.sessions_ready() && self.can_serve() - } - - fn model_aliases(&self) -> &[gateway_core::routing::ContributedModelAlias] { - &self.model_aliases - } - - fn can_serve(&self) -> bool { - // 发布的是可用能力和故障绑定组成的完整计划；单个进程退出不能使原生转发失效 - !self.shutting_down.load(Ordering::Acquire) - } -} - -impl Drop for PreparedSet { - fn drop(&mut self) { - for instance in self.sessions.drain(..) { - if let Ok(handle) = tokio::runtime::Handle::try_current() { - handle.spawn(async move { - let session = instance.session; - session.quiesce(); - session.shutdown(Duration::from_secs(1)).await; - }); - } - } - } -} - -impl PluginRuntime { - #[must_use] - pub fn new( - store: Arc, - state: Arc, - config: PluginRuntimeConfig, - http: Arc, - processes: Arc, - ) -> Self { - let restart_circuit = config.restart_circuit; - Self { - store, - state, - log_slots: Arc::new(Semaphore::new( - config - .rpc_limits - .maximum_callbacks - .min(Semaphore::MAX_PERMITS), - )), - config, - prepare_lock: Mutex::new(()), - prepared: Mutex::new(BTreeMap::new()), - preparation_diagnostics: SyncMutex::new(BTreeMap::new()), - restart_circuits: RestartCircuits::new(restart_circuit), - shutting_down: Arc::new(AtomicBool::new(false)), - shutdown_lock: Mutex::new(()), - processes, - validators: Arc::new(Semaphore::new(2)), - service_ports: Arc::new(crate::callback::services::ServicePorts::default()), - account_ports: Arc::new(PluginAccountPortSlot::new()), - client_key_ports: Arc::new(PluginClientKeyPortSlot::new()), - resource_ports: Arc::new(crate::callback::PluginResourcePorts::new()), - model_ports: Arc::new(PluginModelPortSlot::new()), - affinity_ports: Arc::new(PluginAffinityPortSlot::new()), - observers: RequestObserverExtensionIndex::default(), - policies: RequestPolicyExtensionIndex::default(), - middleware: MiddlewareExtensionIndex::default(), - authentications: FrontendAuthenticationExtensionIndex::default(), - oauth_pending: None, - http, - } - } - - #[must_use] - pub fn with_oauth_pending(mut self, port: Arc) -> Self { - self.oauth_pending = Some(port); - self - } - - /// API 完成组装后绑定唯一内部 HTTP 分派端口；不持有服务端强引用 - pub fn bind_http( - &self, - dispatcher: &Arc, - ) -> Result<(), AdminError> { - self.service_ports.bind_http(dispatcher) - } - - pub fn bind_services( - &self, - registry: &Arc, - ) -> Result<(), AdminError> { - self.service_ports.bind(registry) - } - - pub fn bind_account_ports( - &self, - access: &Arc, - ) -> Result<(), AdminError> { - self.account_ports.bind(access) - } - - /// Key 目录与预算端口由 Admin 组合并保活，完整字段由对应公开合同返回 - pub fn bind_client_key_ports( - &self, - access: &Arc, - ) -> Result<(), AdminError> { - self.client_key_ports.bind(access) - } - - /// Core 激活后一次性绑定嵌套执行与亲和查询端口；Runtime 仅保存 Weak，避免组合根强环 - pub fn bind_model_ports( - &self, - models: &Arc, - affinity: &Arc, - ) -> Result<(), AdminError> { - self.model_ports.bind(models)?; - self.affinity_ports.bind(affinity) - } - - /// 调用方必须先停止接收 HTTP 与 Worker；这里拒绝新候选并等待全部插件 I/O 退出 - pub async fn shutdown(&self) { - self.shutting_down.store(true, Ordering::Release); - let _shutdown = self.shutdown_lock.lock().await; - let _prepare = self.prepare_lock.lock().await; - let sets = { - let mut prepared = self.prepared.lock().await; - prepared.retain(|_, set| set.strong_count() > 0); - let sets = prepared - .values() - .filter_map(Weak::upgrade) - .collect::>(); - prepared.clear(); - sets - }; - let mut sessions = Vec::new(); - for set in &sets { - for instance in &set.sessions { - if !sessions - .iter() - .any(|session| Arc::ptr_eq(session, &instance.session)) - { - sessions.push(Arc::clone(&instance.session)); - } - } - } - shutdown_sessions(sessions, RUNTIME_SHUTDOWN_GRACE).await; - } - - #[must_use] - pub fn observer_registry(&self) -> RequestObserverExtensionIndex { - self.observers.clone() - } - - #[must_use] - pub fn policy_registry(&self) -> RequestPolicyExtensionIndex { - self.policies.clone() - } - - #[must_use] - pub fn middleware_registry(&self) -> MiddlewareExtensionIndex { - self.middleware.clone() - } - - #[must_use] - pub fn frontend_authentication_registry(&self) -> FrontendAuthenticationExtensionIndex { - self.authentications.clone() - } - - fn record_preparation(&self, revision: u64, diagnostic: PreparationDiagnostic) { - let mut diagnostics = self - .preparation_diagnostics - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - diagnostics.insert(revision, diagnostic); - while diagnostics.len() > MAXIMUM_PREPARATION_DIAGNOSTICS { - diagnostics.pop_first(); - } - } - - fn preparation_diagnostic(&self, revision: u64) -> Option { - self.preparation_diagnostics - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .get(&revision) - .cloned() - } - - fn restart_circuit_is_open(&self, identity: &RestartIdentity) -> bool { - self.restart_circuits.is_open(identity) - } - - async fn prepare_snapshot( - &self, - snapshot: PluginInstanceSnapshot, - required_revision: Option, - ) -> Result { - let target_revision = snapshot.config_revision.get(); - self.record_preparation(target_revision, PreparationDiagnostic::Preparing); - let result = self.prepare_inner(snapshot, required_revision).await; - match &result { - Ok(reference) => self.record_preparation( - target_revision, - PreparationDiagnostic::Prepared { - set_id: reference.id().as_str().to_owned(), - }, - ), - Err(error) => self.record_preparation( - target_revision, - PreparationDiagnostic::Failed(runtime_failure_from_admin(error)), - ), - } - result - } - - async fn prepare_inner( - &self, - snapshot: PluginInstanceSnapshot, - required_revision: Option, - ) -> Result { - if self.shutting_down.load(Ordering::Acquire) { - return Err(AdminError::unavailable("插件运行时正在关闭")); - } - if snapshot.instances.len() > 64 { - return Err(AdminError::invalid("最多配置 64 个插件实例")); - } - let fingerprint = fingerprint(&snapshot)?; - // 准备串行化与弱索引分锁；诊断和已发布请求不会等待外部 I/O - let _prepare = self.prepare_lock.lock().await; - if self.shutting_down.load(Ordering::Acquire) { - return Err(AdminError::unavailable("插件运行时正在关闭")); - } - let existing = { - let mut prepared = self.prepared.lock().await; - prepared.retain(|_, set| set.strong_count() > 0); - prepared - .get(&fingerprint) - .and_then(Weak::upgrade) - .filter(|set| set.sessions_ready()) - }; - if let Some(set) = existing { - if snapshot.instances.iter().any(|instance| { - required_revision == Some(instance.revision) - && set.failures.contains_key(&instance.id) - }) { - return Err(AdminError::unavailable("插件启动失败，请重新启用后重试")); - } - return Ok(ExtensionSetReference::new(set.id.clone(), set)); - } - let mut contributions = PreparedContributions::default(); - let mut failures = BTreeMap::new(); - let mut identities = BTreeSet::new(); - for instance in snapshot - .instances - .iter() - .filter(|instance| instance.enabled) - { - if !identities.insert(instance.id.clone()) { - return Err(AdminError::invalid("插件实例 ID 重复")); - } - let result = tokio::time::timeout( - // 沿用准备阶段原有的 30 秒预算；打包校验与能力恢复也在预算内 - Duration::from_secs(30), - self.prepare_instance(instance.clone(), snapshot.config_revision), - ) - .await - .unwrap_or_else(|_| Err(AdminError::unavailable("插件启动超时"))); - match result { - Ok(candidate) => contributions.append(candidate), - Err(error) if required_revision == Some(instance.revision) => return Err(error), - Err(error) => { - // 目录没有拒绝请求用的 binding，不能在恢复失败时悄悄撤销仍启用的别名 - // 不复制其他候选的注册结果；让发布事务保留调用方持有的旧有效快照 - let has_retained_catalog = self - .prepared - .lock() - .await - .values() - .filter_map(Weak::upgrade) - .any(|set| { - set.model_aliases - .iter() - .any(|alias| alias.owner == instance.id) - }); - if has_retained_catalog { - return Err(AdminError::unavailable( - "已启用的模型目录插件未就绪，请修复或停用后重试", - )); - } - // 恢复失败只隔离该实例；绑定的拒绝策略仍保留，不能绕过认证或必需处理 - contributions - .policy_entries - .extend(crate::adapter::policy::unavailable_entries(instance)?); - contributions.upstream_entries.extend( - crate::adapter::upstream_adapter::unavailable_entries(instance)?, - ); - if let Some(entry) = - crate::adapter::frontend_authentication::unavailable_entry(instance) - { - contributions.authentication_entries.push(entry); - } - failures.insert(instance.id.clone(), runtime_failure_from_admin(&error)); - } - } - } - let PreparedContributions { - sessions, - observer_entries, - commands, - management, - policy_entries, - upstream_entries, - authentication_entries, - model_aliases, - } = contributions; - let mut alias_owners = BTreeMap::new(); - for alias in &model_aliases { - if let Some(previous) = alias_owners.insert(&alias.id, &alias.owner) { - return Err(AdminError::invalid(format!( - "模型别名 {} 在插件实例 {} 与 {} 之间冲突", - alias.id, previous, alias.owner - ))); - } - } - let id = ExtensionSetId::new(uuid::Uuid::new_v4().to_string()) - .map_err(|_| AdminError::internal("插件集合 ID 无效"))?; - let observers = crate::adapter::observer::PluginObserverPlan::compile( - observer_entries, - self.config.rpc_limits.maximum_calls, - self.config.rpc_limits.maximum_call_timeout, - self.config.rpc_limits.maximum_buffered_body_bytes, - ) - .map(|plan| { - self.observers - .register(id.clone(), plan) - .map_err(|_| AdminError::invalid("插件观察计划注册冲突")) - }) - .transpose()?; - let policy_plan = crate::adapter::policy::PluginRequestPolicyPlan::compile( - policy_entries, - self.config.rpc_limits.maximum_call_timeout, - )?; - let policies = policy_plan - .as_ref() - .filter(|plan| plan.has_request_policy()) - .cloned() - .map(|plan| { - self.policies - .register(id.clone(), plan) - .map_err(|_| AdminError::invalid("插件请求策略计划注册冲突")) - }) - .transpose()?; - let middleware = policy_plan - .filter(|plan| plan.has_middleware()) - .map(|plan| { - self.middleware - .register(id.clone(), plan) - .map_err(|_| AdminError::invalid("插件中间件计划注册冲突")) - }) - .transpose()?; - let authentication = - crate::adapter::frontend_authentication::PluginFrontendAuthenticationPlan::compile( - authentication_entries, - self.config.rpc_limits.maximum_call_timeout, - )? - .map(|plan| { - self.authentications - .register(id.clone(), plan) - .map_err(|_| AdminError::invalid("插件入口认证计划注册冲突")) - }) - .transpose()?; - let set = Arc::new(PreparedSet { - id, - _observers: observers, - _policies: policies, - _middleware: middleware, - upstream_adapters: - crate::adapter::upstream_adapter::PluginUpstreamAdapterPlan::compile( - upstream_entries, - )?, - _authentication: authentication, - sessions, - commands, - management, - failures, - model_aliases, - shutting_down: self.shutting_down.clone(), - }); - self.prepared - .lock() - .await - .insert(fingerprint, Arc::downgrade(&set)); - Ok(ExtensionSetReference::new(set.id.clone(), set)) - } - - /// CLI 复用同一候选准备与弱索引；读取帮助不会发布 Core 快照或执行命令 - pub async fn prepare_command_line(&self) -> Result { - let snapshot = self - .store - .load_instances() - .await - .map_err(|_| AdminError::unavailable("无法读取已安装插件"))?; - let reference = self.prepare_snapshot(snapshot, None).await?; - let set = self.prepared_set(&reference).await?; - Ok(crate::PluginCommandSession::new( - reference, - set.commands.clone(), - set.sessions - .iter() - .map(|instance| instance.session.clone()) - .collect(), - self.store.clone(), - self.config.rpc_limits, - )) - } - - pub(super) async fn prepared_set( - &self, - reference: &ExtensionSetReference, - ) -> Result, AdminError> { - self.prepared - .lock() - .await - .values() - .filter_map(Weak::upgrade) - .find(|set| set.id == *reference.id()) - .ok_or_else(|| AdminError::conflict("插件候选已过期，请重试")) - } - - async fn run_state_migration( - &self, - prepared: &ExtensionSetReference, - transition: PluginStateTransition, - ) -> Result<(), AdminError> { - let set = self.prepared_set(prepared).await?; - let instance = set - .sessions - .iter() - .find(|instance| { - instance.instance_id == transition.instance_id - && instance.artifact_sha256 == transition.artifact_sha256 - }) - .ok_or_else(|| AdminError::conflict("状态迁移候选插件已失效"))?; - for namespace in &transition.namespaces { - let mut batches = 0usize; - loop { - batches += 1; - if batches > 200 { - return Err(AdminError::invalid("插件状态迁移超过有界批次数")); - } - let batch = self - .state - .migration_batch(&transition.id, &namespace.namespace, 64) - .await - .map_err(map_state_admin_error)?; - let expected_keys = batch - .records - .iter() - .map(|record| record.key.clone()) - .collect::>(); - if batch.records.is_empty() { - self.state - .apply_migration_batch(ApplyPluginStateMigration { - transition_id: transition.id.clone(), - namespace: namespace.namespace.clone(), - cursor: batch.cursor, - expected_keys, - changes: Vec::new(), - }) - .await - .map_err(map_state_admin_error)?; - break; - } - let request = StateMigrationRequest { - namespace: namespace.namespace.clone(), - from_schema_version: namespace.from_schema_version, - to_schema_version: namespace.to_schema_version, - records: batch - .records - .iter() - .map(|record| StateMigrationRecord { - key: record.key.clone(), - value: record.value.clone(), - version: record.version, - }) - .collect(), - }; - let reply = instance - .session - .call( - "plugin.state.migrate", - instance - .session - .context(Stage::Configuration, Duration::from_secs(30)), - serde_json::json!({}), - serde_json::to_vec(&request) - .map_err(|_| AdminError::internal("状态迁移批次无法编码"))?, - ) - .await - .map_err(|_| AdminError::invalid("插件状态迁移调用失败"))?; - if reply.result != serde_json::json!({}) { - return Err(AdminError::invalid("插件状态迁移结果信封无效")); - } - let result: StateMigrationResult = serde_json::from_slice(&reply.payload) - .map_err(|_| AdminError::invalid("插件状态迁移结果无效"))?; - if result.changes.len() != batch.records.len() - || result - .changes - .iter() - .zip(&expected_keys) - .any(|(change, expected)| change.key() != expected) - { - return Err(AdminError::invalid("插件状态迁移结果未逐项对应源批次")); - } - let mut changes = Vec::with_capacity(result.changes.len()); - for (change, source) in result.changes.into_iter().zip(&batch.records) { - let (key, action) = match change { - StateMigrationChange::Keep { key } => { - instance - .private_state - .validate_value(&namespace.namespace, &source.value)?; - (key, PluginStateMigrationAction::Keep) - } - StateMigrationChange::Replace { key, value } => { - instance - .private_state - .validate_value(&namespace.namespace, &value)?; - (key, PluginStateMigrationAction::Replace(value)) - } - StateMigrationChange::Delete { key } => { - (key, PluginStateMigrationAction::Delete) - } - }; - changes.push(PluginStateMigrationChange { key, action }); - } - self.state - .apply_migration_batch(ApplyPluginStateMigration { - transition_id: transition.id.clone(), - namespace: namespace.namespace.clone(), - cursor: batch.cursor, - expected_keys, - changes, - }) - .await - .map_err(map_state_admin_error)?; - } - } - Ok(()) - } -} - -#[async_trait] -impl gateway_admin::ports::plugin_management::PluginManagement for PluginRuntime { - async fn validate_target( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - ) -> Result<(), AdminError> { - let set = self.prepared_set(published).await?; - set.management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))?; - - Ok(()) - } - - async fn start_callback( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - command: gateway_admin::model::plugins::management::StartPluginManagementCallback, - context: &gateway_admin::model::auth::AdminRequestContext, - ) -> Result - { - let store = self - .oauth_pending - .as_ref() - .ok_or_else(|| AdminError::unavailable("插件登录状态端口不可用"))?; - self.prepared_set(published) - .await? - .management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? - .start_callback(store.as_ref(), command, context) - .await - } - - async fn callback( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - state: &str, - request: gateway_admin::model::plugins::management::PluginManagementRequest, - ) -> Result - { - let store = self - .oauth_pending - .as_ref() - .ok_or_else(|| AdminError::unavailable("插件登录状态端口不可用"))?; - self.prepared_set(published) - .await? - .management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? - .callback(store.as_ref(), state, request, self.config.rpc_limits) - .await - } - - async fn views( - &self, - published: &ExtensionSetReference, - ) -> Result, AdminError> - { - Ok(self - .prepared_set(published) - .await? - .management - .iter() - .map(|entry| entry.view.clone()) - .collect()) - } - - async fn resource( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - path: &str, - public: bool, - ) -> Result - { - self.prepared_set(published) - .await? - .management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? - .resource(path, public) - } - - async fn handle( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - request: gateway_admin::model::plugins::management::PluginManagementRequest, - ) -> Result - { - self.prepared_set(published) - .await? - .management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? - .handle(request, self.config.rpc_limits) - .await - } -} - -#[async_trait] -impl PluginPreparation for PluginRuntime { - async fn configuration_ready( - &self, - instance: PluginInstance, - metadata: &gateway_admin::model::plugins::PluginArtifactMetadata, - ) -> Result { - if instance.artifact_sha256 != metadata.sha256 { - return Err(AdminError::invalid("插件配置与制品不匹配")); - } - // 安装时已验证声明；只读列表不重新解包，也不占用安装与启用的校验配额 - let schema = metadata.configuration_schema.clone(); - let secret_fields = metadata.secret_fields.iter().cloned().collect(); - tokio::task::spawn_blocking(move || { - super::configuration::configuration_ready(&instance, &schema, &secret_fields) - }) - .await - .map_err(|_| AdminError::internal("插件校验任务失败"))? - } - - async fn validate( - &self, - instance: PluginInstance, - ) -> Result { - let artifact = self - .store - .load_artifact(&instance.artifact_sha256) - .await - .map_err(|_| AdminError::not_found("插件制品不存在"))?; - let limits = self.config.package_limits; - let slot = self - .validators - .clone() - .try_acquire_owned() - .map_err(|_| AdminError::unavailable("插件校验繁忙"))?; - tokio::task::spawn_blocking(move || { - let _slot = slot; - let package = - ValidatedPackage::read(artifact.archive, Some(&instance.artifact_sha256), limits) - .map_err(|_| AdminError::invalid("插件制品校验失败"))?; - let (_, state) = super::configuration::validate(&instance, package.manifest())?; - crate::adapter::observer::validate_bindings(package.manifest(), &instance.bindings)?; - crate::adapter::policy::validate_bindings(package.manifest(), &instance.bindings)?; - crate::adapter::catalog::validate_bindings(package.manifest(), &instance.bindings)?; - crate::adapter::upstream_adapter::validate_bindings( - package.manifest(), - &instance.bindings, - )?; - crate::adapter::frontend_authentication::validate_bindings( - package.manifest(), - &instance.bindings, - )?; - Ok(state) - }) - .await - .map_err(|_| AdminError::internal("插件校验任务失败"))? - } - async fn prepare( - &self, - snapshot: PluginInstanceSnapshot, - ) -> Result { - let required_revision = snapshot.config_revision; - self.prepare_snapshot(snapshot, Some(required_revision)) - .await - } -} - -#[async_trait] -impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for PluginRuntime { - async fn runtime_diagnostics( - &self, - snapshot: &PluginInstanceSnapshot, - published_revision: Option, - published: Option<&ExtensionSetReference>, - ) -> Option> { - let sets = { - let mut prepared = self.prepared.lock().await; - prepared.retain(|_, set| set.strong_count() > 0); - prepared - .values() - .filter_map(Weak::upgrade) - .collect::>() - }; - let open_restart_circuits = snapshot - .instances - .iter() - .filter(|instance| instance.enabled) - .filter_map(|instance| { - let identity = - RestartIdentity::new(instance.id.clone(), instance_fingerprint(instance).ok()?); - self.restart_circuit_is_open(&identity) - .then(|| instance.id.clone()) - }) - .collect::>(); - let published_set = published - .and_then(|reference| sets.iter().find(|set| set.id == *reference.id()).cloned()); - let published_ready = published.is_some_and(ExtensionSetReference::can_serve); - let preparation = self.preparation_diagnostic(snapshot.config_revision.get()); - let projection = InstanceRuntimeProjection { - target_revision: snapshot.config_revision, - published_revision, - published_ready, - published_set: published_set.as_deref(), - sets: &sets, - preparation: preparation.as_ref(), - open_restart_circuits: &open_restart_circuits, - }; - Some( - snapshot - .instances - .iter() - .map(|instance| { - let runtime = projection.project(instance); - (instance.id.clone(), runtime) - }) - .collect(), - ) - } -} - -#[async_trait] -impl gateway_admin::ports::plugins::PluginStateLifecycle for PluginRuntime { - async fn activate_state( - &self, - prepared: &ExtensionSetReference, - instance: &PluginInstance, - ) -> Result<(), AdminError> { - if !instance.enabled { - return Ok(()); - } - let set = self.prepared_set(prepared).await?; - let prepared = set - .sessions - .iter() - .find(|prepared| { - prepared.instance_id == instance.id - && prepared.artifact_sha256 == instance.artifact_sha256 - }) - .ok_or_else(|| AdminError::conflict("插件候选实例已失效"))?; - let owner = self - .state - .load_owner(prepared.private_state.owner_request(instance)) - .await - .map_err(map_state_admin_error)? - .ok_or_else(|| AdminError::conflict("插件状态 fence 尚未提交"))?; - prepared.private_state.activate(owner); - Ok(()) - } - - async fn quiesce_instance(&self, instance_id: &str, artifact_sha256: &str, revision: Revision) { - let mut sessions = Vec::new(); - for set in self - .prepared - .lock() - .await - .values() - .filter_map(Weak::upgrade) - { - for prepared in &set.sessions { - if prepared.instance_id == instance_id - && prepared.artifact_sha256 == artifact_sha256 - && prepared.revision == revision - && !sessions - .iter() - .any(|session| Arc::ptr_eq(session, &prepared.session)) - { - sessions.push(prepared.session.clone()); - } - } - } - shutdown_sessions(sessions, Duration::from_secs(2)).await; - } - - async fn migrate_state( - &self, - prepared: &ExtensionSetReference, - transition: PluginStateTransition, - ) -> Result<(), AdminError> { - self.run_state_migration(prepared, transition).await - } -} - -struct InstanceRuntimeProjection<'a> { - target_revision: Revision, - published_revision: Option, - published_ready: bool, - published_set: Option<&'a PreparedSet>, - sets: &'a [Arc], - preparation: Option<&'a PreparationDiagnostic>, - open_restart_circuits: &'a BTreeSet, -} - -impl InstanceRuntimeProjection<'_> { - fn project(&self, expected: &PluginInstance) -> PluginInstanceRuntime { - let active = self.published_set.and_then(|set| { - set.sessions - .iter() - .find(|instance| instance.instance_id == expected.id) - }); - let actual_matches = active.is_some_and(|instance| { - instance.revision == expected.revision - && instance.artifact_sha256 == expected.artifact_sha256 - }); - let target_is_published = self.published_revision == Some(self.target_revision.get()); - let candidate_set = match self.preparation { - Some(PreparationDiagnostic::Prepared { set_id }) => Some(set_id.as_str()), - _ => None, - }; - let published_set_id = self.published_set.map(|set| set.id.as_str()); - let mut draining_revisions = BTreeSet::new(); - for set in self.sets { - if Some(set.id.as_str()) == published_set_id || Some(set.id.as_str()) == candidate_set { - continue; - } - if let Some(instance) = set - .sessions - .iter() - .find(|instance| instance.instance_id == expected.id) - && instance.revision.get() < expected.revision.get() - { - draining_revisions.insert(instance.revision.get()); - } - } - if !expected.enabled - && let Some(instance) = active - && instance.revision != expected.revision - { - draining_revisions.insert(instance.revision.get()); - } - - let diagnostic = active.map(|instance| instance.session.diagnostic()); - let mut failure = None; - let status = if !expected.enabled { - if active.is_some() || !draining_revisions.is_empty() { - PluginInstanceRuntimeStatus::Draining - } else { - PluginInstanceRuntimeStatus::Disabled - } - } else if target_is_published - && let Some(reason) = self - .published_set - .and_then(|set| set.failures.get(&expected.id)) - { - failure = Some(reason.clone()); - PluginInstanceRuntimeStatus::PreparationFailed - } else if self.open_restart_circuits.contains(&expected.id) { - let message = match diagnostic { - Some(crate::rpc::RpcSessionDiagnostic::Failed { message, .. }) => { - format!("{message}，已暂停自动重启") - } - _ => "插件实例连续异常退出，已暂停自动重启".to_owned(), - }; - failure = Some(runtime_failure("restart_circuit_open", &message)); - PluginInstanceRuntimeStatus::Faulted - } else if actual_matches - && let Some(crate::rpc::RpcSessionDiagnostic::Failed { code, message }) = diagnostic - { - failure = Some(runtime_failure(code, message)); - PluginInstanceRuntimeStatus::Faulted - } else if !target_is_published { - match self.preparation { - Some(PreparationDiagnostic::Preparing) => PluginInstanceRuntimeStatus::Preparing, - Some(PreparationDiagnostic::Failed(reason)) => { - failure = Some(reason.clone()); - PluginInstanceRuntimeStatus::PreparationFailed - } - Some(PreparationDiagnostic::Prepared { .. }) | None => { - PluginInstanceRuntimeStatus::AwaitingPublication - } - } - } else if !actual_matches { - failure = Some(runtime_failure( - "published_instance_mismatch", - "已发布插件实例与期望版本不一致", - )); - PluginInstanceRuntimeStatus::Blocked - } else { - match diagnostic { - Some(crate::rpc::RpcSessionDiagnostic::Ready) if self.published_ready => { - PluginInstanceRuntimeStatus::Running - } - Some(crate::rpc::RpcSessionDiagnostic::Ready) => { - failure = Some(runtime_failure( - "published_set_unready", - "同一发布集合中存在未就绪插件实例", - )); - PluginInstanceRuntimeStatus::Blocked - } - Some(crate::rpc::RpcSessionDiagnostic::Quiescing) => { - PluginInstanceRuntimeStatus::Draining - } - Some(crate::rpc::RpcSessionDiagnostic::Failed { code, message }) => { - failure = Some(runtime_failure(code, message)); - PluginInstanceRuntimeStatus::Faulted - } - None => { - failure = Some(runtime_failure( - "instance_not_published", - "已发布集合缺少该插件实例", - )); - PluginInstanceRuntimeStatus::Blocked - } - } - }; - - PluginInstanceRuntime { - status, - actual_revision: active.map(|instance| instance.revision.get()), - actual_artifact_sha256: active.map(|instance| instance.artifact_sha256.clone()), - failure, - draining_revisions: draining_revisions.into_iter().collect(), - } - } -} - -fn runtime_failure(code: &str, message: &str) -> PluginInstanceRuntimeFailure { - PluginInstanceRuntimeFailure { - code: code.to_owned(), - message: message.to_owned(), - } -} - -fn runtime_failure_from_admin(error: &AdminError) -> PluginInstanceRuntimeFailure { - let code = match error.kind() { - AdminErrorKind::Invalid => "invalid", - AdminErrorKind::Unauthorized => "unauthorized", - AdminErrorKind::Forbidden => "forbidden", - AdminErrorKind::NotFound => "not_found", - AdminErrorKind::Conflict => "conflict", - AdminErrorKind::RateLimited => "rate_limited", - AdminErrorKind::BadGateway => "bad_gateway", - AdminErrorKind::UpstreamResultUnknown => "upstream_result_unknown", - AdminErrorKind::Unavailable => "unavailable", - AdminErrorKind::Internal => "internal", - }; - runtime_failure(code, error.message()) -} - -async fn shutdown_sessions(sessions: Vec>, grace: Duration) { - for session in &sessions { - session.quiesce(); - } - let deadline = tokio::time::Instant::now() + grace; - futures::future::join_all( - sessions - .iter() - .map(|session| session.shutdown_until(deadline)), - ) - .await; -} - -fn map_state_admin_error( - error: gateway_admin::ports::plugins::PluginStateStoreError, -) -> AdminError { - match error.kind() { - PluginStateStoreErrorKind::Invalid => AdminError::invalid("插件状态请求不合法"), - PluginStateStoreErrorKind::NotFound => AdminError::not_found("插件状态迁移不存在"), - PluginStateStoreErrorKind::Conflict => AdminError::conflict("插件状态版本或迁移游标冲突"), - PluginStateStoreErrorKind::Quota => AdminError::conflict("插件状态超过声明配额"), - PluginStateStoreErrorKind::PermissionDenied => { - AdminError::conflict("插件状态 fence 已失效") - } - PluginStateStoreErrorKind::Unavailable => AdminError::unavailable("插件状态存储暂不可用"), - } -} - -impl ExtensionPreparationPort for PluginRuntime { - fn prepare( - &self, - revision: ConfigRevision, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { - let snapshot = self - .store - .load_instances() - .await - .map_err(|_| ExtensionPreparationError)?; - if snapshot.config_revision.get() != revision.get() { - return Err(ExtensionPreparationError); - } - self.prepare_snapshot(snapshot, None) - .await - .map_err(|_| ExtensionPreparationError) - }) - } -} - -fn fingerprint(snapshot: &PluginInstanceSnapshot) -> Result { - let mut instances: Vec<_> = snapshot - .instances - .iter() - .filter(|instance| instance.enabled) - .collect(); - instances.sort_by(|left, right| left.id.cmp(&right.id)); - let mut hash = Sha256::new(); - for instance in instances { - hash.update(instance_fingerprint(instance)?.as_bytes()); - } - Ok(hex::encode(hash.finalize())) -} - -fn instance_fingerprint( - instance: &gateway_admin::model::plugins::instances::PluginInstance, -) -> Result { - let secrets: BTreeMap<_, _> = instance - .secrets - .iter() - .map(|(name, value)| (name, value.expose_secret())) - .collect(); - let mut value = serde_json::json!({"id":instance.id,"artifact":instance.artifact_sha256,"revision":instance.revision.get(),"trusted":instance.trusted_process,"configuration":instance.configuration,"secrets":secrets,"bindings":instance.bindings}); - // JSONB 恢复可以改变对象键顺序，不能因此替换已经准备好的同一配置 - value.sort_all_objects(); - let bytes = serde_json::to_vec(&value).map_err(|_| AdminError::invalid("插件配置无法编码"))?; - Ok(hex::encode(Sha256::digest(bytes))) -} - -impl PluginRuntime { - pub fn bind_resource_ports( - &self, - access: &Arc, - ) -> Result<(), AdminError> { - self.resource_ports.bind(access) - } -} +pub use runtime::{PluginRuntime, PluginRuntimeConfig}; diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs new file mode 100644 index 000000000..be6244605 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs @@ -0,0 +1,240 @@ +//! 插件运行时配置、端口装配与整体关闭；候选和调用均复用同一状态 + +use super::super::restart_circuit::{PluginRestartCircuitConfig, RestartCircuits, RestartIdentity}; +use super::{diagnostics::PreparationDiagnostic, set::PreparedSet}; +use crate::{ + PackageLimits, RpcLimits, RpcSession, + callback::{ + PluginAccountPortSlot, PluginAffinityPortSlot, PluginClientKeyPortSlot, PluginModelPortSlot, + }, +}; +use gateway_admin::{ + model::AdminError, + ports::{ + plugin_accounts::PluginAccountAccess, + plugin_client_keys::PluginClientKeyAccess, + plugins::{PluginStateStore, PluginStore}, + }, +}; +use gateway_core::{ + engine::{ + authentication::FrontendAuthenticationExtensionIndex, + extensions::ExecutionExtensionIndex, + nested::{AffinityLookupPort, NestedModelExecutionPort}, + observation::RequestObserverExtensionIndex, + policy::RequestPolicyExtensionIndex, + }, + provider_ports::OAuthPendingFlowPort, +}; +use std::{ + collections::BTreeMap, + path::PathBuf, + sync::{ + Arc, Mutex as SyncMutex, Weak, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; +use tokio::sync::{Mutex, Semaphore}; + +const RUNTIME_SHUTDOWN_GRACE: Duration = Duration::from_secs(2); +pub struct PluginRuntimeConfig { + pub cache_directory: PathBuf, + pub host_version: semver::Version, + pub package_limits: PackageLimits, + pub rpc_limits: RpcLimits, + pub restart_circuit: PluginRestartCircuitConfig, +} + +/// 只索引有主人的候选；当前代次仅由 Core 的发布视图持有 +pub struct PluginRuntime { + pub(super) service_ports: Arc, + pub(super) store: Arc, + pub(super) state: Arc, + pub(super) config: PluginRuntimeConfig, + pub(super) prepare_lock: Mutex<()>, + pub(super) prepared: Mutex>>, + pub(super) preparation_diagnostics: SyncMutex>, + pub(super) restart_circuits: Arc, + pub(super) shutting_down: Arc, + shutdown_lock: Mutex<()>, + pub(super) processes: Arc, + pub(super) validators: Arc, + pub(super) log_slots: Arc, + pub(super) account_ports: Arc, + pub(super) client_key_ports: Arc, + pub(super) resource_ports: Arc, + pub(super) model_ports: Arc, + pub(super) affinity_ports: Arc, + pub(super) observers: RequestObserverExtensionIndex, + pub(super) policies: RequestPolicyExtensionIndex, + pub(super) execution: ExecutionExtensionIndex, + pub(super) authentications: FrontendAuthenticationExtensionIndex, + pub(super) oauth_pending: Option>, + pub(super) http: Arc, +} +impl PluginRuntime { + #[must_use] + pub fn new( + store: Arc, + state: Arc, + config: PluginRuntimeConfig, + http: Arc, + processes: Arc, + ) -> Self { + let restart_circuit = config.restart_circuit; + Self { + store, + state, + log_slots: Arc::new(Semaphore::new( + config + .rpc_limits + .maximum_callbacks + .min(Semaphore::MAX_PERMITS), + )), + config, + prepare_lock: Mutex::new(()), + prepared: Mutex::new(BTreeMap::new()), + preparation_diagnostics: SyncMutex::new(BTreeMap::new()), + restart_circuits: RestartCircuits::new(restart_circuit), + shutting_down: Arc::new(AtomicBool::new(false)), + shutdown_lock: Mutex::new(()), + processes, + validators: Arc::new(Semaphore::new(2)), + service_ports: Arc::new(crate::callback::services::ServicePorts::default()), + account_ports: Arc::new(PluginAccountPortSlot::new()), + client_key_ports: Arc::new(PluginClientKeyPortSlot::new()), + resource_ports: Arc::new(crate::callback::PluginResourcePorts::new()), + model_ports: Arc::new(PluginModelPortSlot::new()), + affinity_ports: Arc::new(PluginAffinityPortSlot::new()), + observers: RequestObserverExtensionIndex::default(), + policies: RequestPolicyExtensionIndex::default(), + execution: ExecutionExtensionIndex::default(), + authentications: FrontendAuthenticationExtensionIndex::default(), + oauth_pending: None, + http, + } + } + + #[must_use] + pub fn with_oauth_pending(mut self, port: Arc) -> Self { + self.oauth_pending = Some(port); + self + } + + /// API 完成组装后绑定唯一内部 HTTP 分派端口；不持有服务端强引用 + pub fn bind_http( + &self, + dispatcher: &Arc, + ) -> Result<(), AdminError> { + self.service_ports.bind_http(dispatcher) + } + + pub fn bind_services( + &self, + registry: &Arc, + ) -> Result<(), AdminError> { + self.service_ports.bind(registry) + } + + pub fn bind_account_ports( + &self, + access: &Arc, + ) -> Result<(), AdminError> { + self.account_ports.bind(access) + } + + /// Key 目录与预算端口由 Admin 组合并保活，完整字段由对应公开合同返回 + pub fn bind_client_key_ports( + &self, + access: &Arc, + ) -> Result<(), AdminError> { + self.client_key_ports.bind(access) + } + + /// Core 激活后一次性绑定嵌套执行与亲和查询端口；Runtime 仅保存 Weak，避免组合根强环 + pub fn bind_model_ports( + &self, + models: &Arc, + affinity: &Arc, + ) -> Result<(), AdminError> { + self.model_ports.bind(models)?; + self.affinity_ports.bind(affinity) + } + + /// 调用方必须先停止接收 HTTP 与 Worker；这里拒绝新候选并等待全部插件 I/O 退出 + pub async fn shutdown(&self) { + self.shutting_down.store(true, Ordering::Release); + let _shutdown = self.shutdown_lock.lock().await; + let _prepare = self.prepare_lock.lock().await; + let sets = { + let mut prepared = self.prepared.lock().await; + prepared.retain(|_, set| set.strong_count() > 0); + let sets = prepared + .values() + .filter_map(Weak::upgrade) + .collect::>(); + prepared.clear(); + sets + }; + let mut sessions = Vec::new(); + for set in &sets { + for instance in &set.sessions { + if !sessions + .iter() + .any(|session| Arc::ptr_eq(session, &instance.session)) + { + sessions.push(Arc::clone(&instance.session)); + } + } + } + shutdown_sessions(sessions, RUNTIME_SHUTDOWN_GRACE).await; + } + + #[must_use] + pub fn observer_registry(&self) -> RequestObserverExtensionIndex { + self.observers.clone() + } + + #[must_use] + pub fn policy_registry(&self) -> RequestPolicyExtensionIndex { + self.policies.clone() + } + + /// 执行计划按同一发布身份解析，索引不延长代次寿命 + #[must_use] + pub fn execution_registry(&self) -> ExecutionExtensionIndex { + self.execution.clone() + } + + #[must_use] + pub fn frontend_authentication_registry(&self) -> FrontendAuthenticationExtensionIndex { + self.authentications.clone() + } +} +pub(super) async fn shutdown_sessions(sessions: Vec>, grace: Duration) { + for session in &sessions { + session.quiesce(); + } + let deadline = tokio::time::Instant::now() + grace; + futures::future::join_all( + sessions + .iter() + .map(|session| session.shutdown_until(deadline)), + ) + .await; +} +impl PluginRuntime { + pub fn bind_resource_ports( + &self, + access: &Arc, + ) -> Result<(), AdminError> { + self.resource_ports.bind(access) + } +} + +impl PluginRuntime { + pub(super) fn restart_circuit_is_open(&self, identity: &RestartIdentity) -> bool { + self.restart_circuits.is_open(identity) + } +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs new file mode 100644 index 000000000..aae6eb632 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs @@ -0,0 +1,140 @@ +//! 发布集合及其保活、候选贡献与配置身份指纹 + +use crate::{RpcSession, callback::private_state::PluginPrivateState}; +use gateway_admin::model::{ + AdminError, Revision, + plugins::instances::{PluginInstanceRuntimeFailure, PluginInstanceSnapshot}, +}; +use gateway_core::{ + engine::{ + extensions::ExecutionExtensionPlans, observation::RequestObserverPlan, + policy::RequestPolicyPlan, + }, + routing::extensions::{ExtensionSetId, ExtensionSetLease}, +}; +use secrecy::ExposeSecret as _; +use sha2::{Digest as _, Sha256}; +use std::{ + collections::BTreeMap, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; + +pub(super) struct PreparedSet { + pub(super) id: ExtensionSetId, + pub(super) _observers: Option>, + pub(super) _policies: Option>, + pub(super) _execution: Option>, + pub(super) _authentication: + Option>, + pub(super) sessions: Vec, + pub(super) failures: BTreeMap, + pub(super) shutting_down: Arc, + pub(super) commands: Vec>, + pub(super) management: Vec, + pub(super) model_aliases: Vec, +} + +#[derive(Default)] +pub(super) struct PreparedContributions { + pub(super) sessions: Vec, + pub(super) observer_entries: Vec, + pub(super) commands: Vec>, + pub(super) management: Vec, + pub(super) policy_entries: Vec, + pub(super) upstream_entries: Vec, + pub(super) authentication_entries: + Vec, + pub(super) model_aliases: Vec, +} + +impl PreparedContributions { + pub(super) fn append(&mut self, other: Self) { + self.sessions.extend(other.sessions); + self.observer_entries.extend(other.observer_entries); + self.commands.extend(other.commands); + self.management.extend(other.management); + self.model_aliases.extend(other.model_aliases); + self.policy_entries.extend(other.policy_entries); + self.upstream_entries.extend(other.upstream_entries); + self.authentication_entries + .extend(other.authentication_entries); + } +} + +pub(super) struct PreparedInstance { + pub(super) instance_id: String, + pub(super) artifact_sha256: String, + pub(super) revision: Revision, + pub(super) session: Arc, + pub(super) private_state: Arc, + pub(super) maintenance: bool, +} + +impl PreparedSet { + pub(super) fn sessions_ready(&self) -> bool { + self.sessions + .iter() + .all(|instance| instance.session.is_ready()) + } +} + +impl ExtensionSetLease for PreparedSet { + fn is_ready(&self) -> bool { + self.sessions_ready() && self.can_serve() + } + + fn model_aliases(&self) -> &[gateway_core::routing::ContributedModelAlias] { + &self.model_aliases + } + + fn can_serve(&self) -> bool { + // 发布的是可用能力和故障绑定组成的完整计划；单个进程退出不能使原生转发失效 + !self.shutting_down.load(Ordering::Acquire) + } +} + +impl Drop for PreparedSet { + fn drop(&mut self) { + for instance in self.sessions.drain(..) { + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + let session = instance.session; + session.quiesce(); + session.shutdown(Duration::from_secs(1)).await; + }); + } + } + } +} +pub(super) fn fingerprint(snapshot: &PluginInstanceSnapshot) -> Result { + let mut instances: Vec<_> = snapshot + .instances + .iter() + .filter(|instance| instance.enabled) + .collect(); + instances.sort_by(|left, right| left.id.cmp(&right.id)); + let mut hash = Sha256::new(); + for instance in instances { + hash.update(instance_fingerprint(instance)?.as_bytes()); + } + Ok(hex::encode(hash.finalize())) +} + +pub(super) fn instance_fingerprint( + instance: &gateway_admin::model::plugins::instances::PluginInstance, +) -> Result { + let secrets: BTreeMap<_, _> = instance + .secrets + .iter() + .map(|(name, value)| (name, value.expose_secret())) + .collect(); + let mut value = serde_json::json!({"id":instance.id,"artifact":instance.artifact_sha256,"revision":instance.revision.get(),"trusted":instance.trusted_process,"configuration":instance.configuration,"secrets":secrets,"bindings":instance.bindings}); + // JSONB 恢复可以改变对象键顺序，不能因此替换已经准备好的同一配置 + value.sort_all_objects(); + let bytes = serde_json::to_vec(&value).map_err(|_| AdminError::invalid("插件配置无法编码"))?; + Ok(hex::encode(Sha256::digest(bytes))) +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs new file mode 100644 index 000000000..e55ee2e13 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs @@ -0,0 +1,343 @@ +//! 校验实例配置并准备可发布的完整能力集合 + +use super::{ + PluginRuntime, + diagnostics::{PreparationDiagnostic, runtime_failure_from_admin}, + set::{PreparedContributions, PreparedSet, fingerprint}, +}; +use crate::ValidatedPackage; +use async_trait::async_trait; +use futures::future::BoxFuture; +use gateway_admin::{ + model::{ + AdminError, Revision, + plugins::{ + instances::{PluginInstance, PluginInstanceSnapshot}, + state::PluginStateConfiguration, + }, + }, + ports::plugins::PluginPreparation, +}; +use gateway_core::{ + routing::ConfigRevision, + routing::extensions::{ + ExtensionPreparationError, ExtensionPreparationPort, ExtensionSetId, ExtensionSetReference, + }, +}; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::{Arc, Weak, atomic::Ordering}, + time::Duration, +}; + +impl PluginRuntime { + pub(super) async fn prepare_snapshot( + &self, + snapshot: PluginInstanceSnapshot, + required_revision: Option, + ) -> Result { + let target_revision = snapshot.config_revision.get(); + self.record_preparation(target_revision, PreparationDiagnostic::Preparing); + let result = self.prepare_inner(snapshot, required_revision).await; + match &result { + Ok(reference) => self.record_preparation( + target_revision, + PreparationDiagnostic::Prepared { + set_id: reference.id().as_str().to_owned(), + }, + ), + Err(error) => self.record_preparation( + target_revision, + PreparationDiagnostic::Failed(runtime_failure_from_admin(error)), + ), + } + result + } + + async fn prepare_inner( + &self, + snapshot: PluginInstanceSnapshot, + required_revision: Option, + ) -> Result { + if self.shutting_down.load(Ordering::Acquire) { + return Err(AdminError::unavailable("插件运行时正在关闭")); + } + if snapshot.instances.len() > 64 { + return Err(AdminError::invalid("最多配置 64 个插件实例")); + } + let fingerprint = fingerprint(&snapshot)?; + // 准备串行化与弱索引分锁；诊断和已发布请求不会等待外部 I/O + let _prepare = self.prepare_lock.lock().await; + if self.shutting_down.load(Ordering::Acquire) { + return Err(AdminError::unavailable("插件运行时正在关闭")); + } + let existing = { + let mut prepared = self.prepared.lock().await; + prepared.retain(|_, set| set.strong_count() > 0); + prepared + .get(&fingerprint) + .and_then(Weak::upgrade) + .filter(|set| set.sessions_ready()) + }; + if let Some(set) = existing { + if snapshot.instances.iter().any(|instance| { + required_revision == Some(instance.revision) + && set.failures.contains_key(&instance.id) + }) { + return Err(AdminError::unavailable("插件启动失败，请重新启用后重试")); + } + return Ok(ExtensionSetReference::new(set.id.clone(), set)); + } + let mut contributions = PreparedContributions::default(); + let mut failures = BTreeMap::new(); + let mut identities = BTreeSet::new(); + for instance in snapshot + .instances + .iter() + .filter(|instance| instance.enabled) + { + if !identities.insert(instance.id.clone()) { + return Err(AdminError::invalid("插件实例 ID 重复")); + } + let result = tokio::time::timeout( + // 沿用准备阶段原有的 30 秒预算；打包校验与能力恢复也在预算内 + Duration::from_secs(30), + self.prepare_instance(instance.clone(), snapshot.config_revision), + ) + .await + .unwrap_or_else(|_| Err(AdminError::unavailable("插件启动超时"))); + match result { + Ok(candidate) => contributions.append(candidate), + Err(error) if required_revision == Some(instance.revision) => return Err(error), + Err(error) => { + // 目录没有拒绝请求用的 binding，不能在恢复失败时悄悄撤销仍启用的别名 + // 不复制其他候选的注册结果；让发布事务保留调用方持有的旧有效快照 + let has_retained_catalog = self + .prepared + .lock() + .await + .values() + .filter_map(Weak::upgrade) + .any(|set| { + set.model_aliases + .iter() + .any(|alias| alias.owner == instance.id) + }); + if has_retained_catalog { + return Err(AdminError::unavailable( + "已启用的模型目录插件未就绪，请修复或停用后重试", + )); + } + // 恢复失败只隔离该实例；绑定的拒绝策略仍保留，不能绕过认证或必需处理 + contributions + .policy_entries + .extend(crate::adapter::policy::unavailable_entries(instance)?); + contributions.upstream_entries.extend( + crate::adapter::upstream_adapter::unavailable_entries(instance)?, + ); + if let Some(entry) = + crate::adapter::frontend_authentication::unavailable_entry(instance) + { + contributions.authentication_entries.push(entry); + } + failures.insert(instance.id.clone(), runtime_failure_from_admin(&error)); + } + } + } + let PreparedContributions { + sessions, + observer_entries, + commands, + management, + policy_entries, + upstream_entries, + authentication_entries, + model_aliases, + } = contributions; + let mut alias_owners = BTreeMap::new(); + for alias in &model_aliases { + if let Some(previous) = alias_owners.insert(&alias.id, &alias.owner) { + return Err(AdminError::invalid(format!( + "模型别名 {} 在插件实例 {} 与 {} 之间冲突", + alias.id, previous, alias.owner + ))); + } + } + let id = ExtensionSetId::new(uuid::Uuid::new_v4().to_string()) + .map_err(|_| AdminError::internal("插件集合 ID 无效"))?; + let observers = crate::adapter::observer::PluginObserverPlan::compile( + observer_entries, + self.config.rpc_limits.maximum_calls, + self.config.rpc_limits.maximum_call_timeout, + self.config.rpc_limits.maximum_buffered_body_bytes, + ) + .map(|plan| { + self.observers + .register(id.clone(), plan) + .map_err(|_| AdminError::invalid("插件观察计划注册冲突")) + }) + .transpose()?; + let policy_plan = crate::adapter::policy::PluginRequestPolicyPlan::compile( + policy_entries, + self.config.rpc_limits.maximum_call_timeout, + )?; + let policies = policy_plan + .as_ref() + .filter(|plan| plan.has_request_policy()) + .cloned() + .map(|plan| { + self.policies + .register(id.clone(), plan) + .map_err(|_| AdminError::invalid("插件请求策略计划注册冲突")) + }) + .transpose()?; + let middleware = policy_plan + .filter(|plan| plan.has_middleware()) + .map(|plan| plan as Arc); + let upstream_adapters = + crate::adapter::upstream_adapter::PluginUpstreamAdapterPlan::compile(upstream_entries)?; + let execution = if middleware.is_some() || upstream_adapters.is_some() { + Some( + self.execution + .register( + id.clone(), + Arc::new(gateway_core::engine::extensions::ExecutionExtensionPlans { + middleware, + upstream_adapters, + }), + ) + .map_err(|_| AdminError::invalid("插件执行计划注册冲突"))?, + ) + } else { + None + }; + let authentication = + crate::adapter::frontend_authentication::PluginFrontendAuthenticationPlan::compile( + authentication_entries, + self.config.rpc_limits.maximum_call_timeout, + )? + .map(|plan| { + self.authentications + .register(id.clone(), plan) + .map_err(|_| AdminError::invalid("插件入口认证计划注册冲突")) + }) + .transpose()?; + let set = Arc::new(PreparedSet { + id, + _observers: observers, + _policies: policies, + _execution: execution, + _authentication: authentication, + sessions, + commands, + management, + failures, + model_aliases, + shutting_down: self.shutting_down.clone(), + }); + self.prepared + .lock() + .await + .insert(fingerprint, Arc::downgrade(&set)); + Ok(ExtensionSetReference::new(set.id.clone(), set)) + } + pub(super) async fn prepared_set( + &self, + reference: &ExtensionSetReference, + ) -> Result, AdminError> { + self.prepared + .lock() + .await + .values() + .filter_map(Weak::upgrade) + .find(|set| set.id == *reference.id()) + .ok_or_else(|| AdminError::conflict("插件候选已过期，请重试")) + } +} +#[async_trait] +impl PluginPreparation for PluginRuntime { + async fn configuration_ready( + &self, + instance: PluginInstance, + metadata: &gateway_admin::model::plugins::PluginArtifactMetadata, + ) -> Result { + if instance.artifact_sha256 != metadata.sha256 { + return Err(AdminError::invalid("插件配置与制品不匹配")); + } + // 安装时已验证声明；只读列表不重新解包，也不占用安装与启用的校验配额 + let schema = metadata.configuration_schema.clone(); + let secret_fields = metadata.secret_fields.iter().cloned().collect(); + tokio::task::spawn_blocking(move || { + super::super::configuration::configuration_ready(&instance, &schema, &secret_fields) + }) + .await + .map_err(|_| AdminError::internal("插件校验任务失败"))? + } + + async fn validate( + &self, + instance: PluginInstance, + ) -> Result { + let artifact = self + .store + .load_artifact(&instance.artifact_sha256) + .await + .map_err(|_| AdminError::not_found("插件制品不存在"))?; + let limits = self.config.package_limits; + let slot = self + .validators + .clone() + .try_acquire_owned() + .map_err(|_| AdminError::unavailable("插件校验繁忙"))?; + tokio::task::spawn_blocking(move || { + let _slot = slot; + let package = + ValidatedPackage::read(artifact.archive, Some(&instance.artifact_sha256), limits) + .map_err(|_| AdminError::invalid("插件制品校验失败"))?; + let (_, state) = super::super::configuration::validate(&instance, package.manifest())?; + crate::adapter::observer::validate_bindings(package.manifest(), &instance.bindings)?; + crate::adapter::policy::validate_bindings(package.manifest(), &instance.bindings)?; + crate::adapter::catalog::validate_bindings(package.manifest(), &instance.bindings)?; + crate::adapter::upstream_adapter::validate_bindings( + package.manifest(), + &instance.bindings, + )?; + crate::adapter::frontend_authentication::validate_bindings( + package.manifest(), + &instance.bindings, + )?; + Ok(state) + }) + .await + .map_err(|_| AdminError::internal("插件校验任务失败"))? + } + async fn prepare( + &self, + snapshot: PluginInstanceSnapshot, + ) -> Result { + let required_revision = snapshot.config_revision; + self.prepare_snapshot(snapshot, Some(required_revision)) + .await + } +} + +impl ExtensionPreparationPort for PluginRuntime { + fn prepare( + &self, + revision: ConfigRevision, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let snapshot = self + .store + .load_instances() + .await + .map_err(|_| ExtensionPreparationError)?; + if snapshot.config_revision.get() != revision.get() { + return Err(ExtensionPreparationError); + } + self.prepare_snapshot(snapshot, None) + .await + .map_err(|_| ExtensionPreparationError) + }) + } +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs new file mode 100644 index 000000000..2b83fa96a --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs @@ -0,0 +1,226 @@ +//! 插件状态的激活、迁移与实例排空，事务事实由状态端口提交 + +use super::{PluginRuntime, runtime::shutdown_sessions}; +use async_trait::async_trait; +use gateway_admin::{ + model::{ + AdminError, Revision, + plugins::{ + instances::PluginInstance, + state::{ + ApplyPluginStateMigration, PluginStateMigrationAction, PluginStateMigrationChange, + PluginStateTransition, + }, + }, + }, + ports::plugins::PluginStateStoreErrorKind, +}; +use gateway_core::routing::extensions::ExtensionSetReference; +use gateway_plugin_sdk::{ + Stage, + call::host::{ + StateMigrationChange, StateMigrationRecord, StateMigrationRequest, StateMigrationResult, + }, +}; +use std::{ + sync::{Arc, Weak}, + time::Duration, +}; + +impl PluginRuntime { + async fn run_state_migration( + &self, + prepared: &ExtensionSetReference, + transition: PluginStateTransition, + ) -> Result<(), AdminError> { + let set = self.prepared_set(prepared).await?; + let instance = set + .sessions + .iter() + .find(|instance| { + instance.instance_id == transition.instance_id + && instance.artifact_sha256 == transition.artifact_sha256 + }) + .ok_or_else(|| AdminError::conflict("状态迁移候选插件已失效"))?; + for namespace in &transition.namespaces { + let mut batches = 0usize; + loop { + batches += 1; + if batches > 200 { + return Err(AdminError::invalid("插件状态迁移超过有界批次数")); + } + let batch = self + .state + .migration_batch(&transition.id, &namespace.namespace, 64) + .await + .map_err(map_state_admin_error)?; + let expected_keys = batch + .records + .iter() + .map(|record| record.key.clone()) + .collect::>(); + if batch.records.is_empty() { + self.state + .apply_migration_batch(ApplyPluginStateMigration { + transition_id: transition.id.clone(), + namespace: namespace.namespace.clone(), + cursor: batch.cursor, + expected_keys, + changes: Vec::new(), + }) + .await + .map_err(map_state_admin_error)?; + break; + } + let request = StateMigrationRequest { + namespace: namespace.namespace.clone(), + from_schema_version: namespace.from_schema_version, + to_schema_version: namespace.to_schema_version, + records: batch + .records + .iter() + .map(|record| StateMigrationRecord { + key: record.key.clone(), + value: record.value.clone(), + version: record.version, + }) + .collect(), + }; + let reply = instance + .session + .call( + "plugin.state.migrate", + instance + .session + .context(Stage::Configuration, Duration::from_secs(30)), + serde_json::json!({}), + serde_json::to_vec(&request) + .map_err(|_| AdminError::internal("状态迁移批次无法编码"))?, + ) + .await + .map_err(|_| AdminError::invalid("插件状态迁移调用失败"))?; + if reply.result != serde_json::json!({}) { + return Err(AdminError::invalid("插件状态迁移结果信封无效")); + } + let result: StateMigrationResult = serde_json::from_slice(&reply.payload) + .map_err(|_| AdminError::invalid("插件状态迁移结果无效"))?; + if result.changes.len() != batch.records.len() + || result + .changes + .iter() + .zip(&expected_keys) + .any(|(change, expected)| change.key() != expected) + { + return Err(AdminError::invalid("插件状态迁移结果未逐项对应源批次")); + } + let mut changes = Vec::with_capacity(result.changes.len()); + for (change, source) in result.changes.into_iter().zip(&batch.records) { + let (key, action) = match change { + StateMigrationChange::Keep { key } => { + instance + .private_state + .validate_value(&namespace.namespace, &source.value)?; + (key, PluginStateMigrationAction::Keep) + } + StateMigrationChange::Replace { key, value } => { + instance + .private_state + .validate_value(&namespace.namespace, &value)?; + (key, PluginStateMigrationAction::Replace(value)) + } + StateMigrationChange::Delete { key } => { + (key, PluginStateMigrationAction::Delete) + } + }; + changes.push(PluginStateMigrationChange { key, action }); + } + self.state + .apply_migration_batch(ApplyPluginStateMigration { + transition_id: transition.id.clone(), + namespace: namespace.namespace.clone(), + cursor: batch.cursor, + expected_keys, + changes, + }) + .await + .map_err(map_state_admin_error)?; + } + } + Ok(()) + } +} +#[async_trait] +impl gateway_admin::ports::plugins::PluginStateLifecycle for PluginRuntime { + async fn activate_state( + &self, + prepared: &ExtensionSetReference, + instance: &PluginInstance, + ) -> Result<(), AdminError> { + if !instance.enabled { + return Ok(()); + } + let set = self.prepared_set(prepared).await?; + let prepared = set + .sessions + .iter() + .find(|prepared| { + prepared.instance_id == instance.id + && prepared.artifact_sha256 == instance.artifact_sha256 + }) + .ok_or_else(|| AdminError::conflict("插件候选实例已失效"))?; + let owner = self + .state + .load_owner(prepared.private_state.owner_request(instance)) + .await + .map_err(map_state_admin_error)? + .ok_or_else(|| AdminError::conflict("插件状态 fence 尚未提交"))?; + prepared.private_state.activate(owner); + Ok(()) + } + + async fn quiesce_instance(&self, instance_id: &str, artifact_sha256: &str, revision: Revision) { + let mut sessions = Vec::new(); + for set in self + .prepared + .lock() + .await + .values() + .filter_map(Weak::upgrade) + { + for prepared in &set.sessions { + if prepared.instance_id == instance_id + && prepared.artifact_sha256 == artifact_sha256 + && prepared.revision == revision + && !sessions + .iter() + .any(|session| Arc::ptr_eq(session, &prepared.session)) + { + sessions.push(prepared.session.clone()); + } + } + } + shutdown_sessions(sessions, Duration::from_secs(2)).await; + } + + async fn migrate_state( + &self, + prepared: &ExtensionSetReference, + transition: PluginStateTransition, + ) -> Result<(), AdminError> { + self.run_state_migration(prepared, transition).await + } +} +pub(super) fn map_state_admin_error( + error: gateway_admin::ports::plugins::PluginStateStoreError, +) -> AdminError { + match error.kind() { + PluginStateStoreErrorKind::Invalid => AdminError::invalid("插件状态请求不合法"), + PluginStateStoreErrorKind::NotFound => AdminError::not_found("插件状态迁移不存在"), + PluginStateStoreErrorKind::Conflict => AdminError::conflict("插件状态版本或迁移游标冲突"), + PluginStateStoreErrorKind::Quota => AdminError::conflict("插件状态超过声明配额"), + PluginStateStoreErrorKind::PermissionDenied => { + AdminError::conflict("插件状态 fence 已失效") + } + PluginStateStoreErrorKind::Unavailable => AdminError::unavailable("插件状态存储暂不可用"), + } +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/data.rs b/backend/crates/gateway-plugin/runtime/tests/generation/data.rs index a3dfa1baa..95fc020e3 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/data.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/data.rs @@ -15,8 +15,8 @@ use gateway_admin::{ }; use gateway_core::{ account::{AccountWeight, CredentialState, ProviderAccountId, QuotaState}, + routing::extensions::ExtensionPreparationPort, routing::{ConfigRevision, ProviderKind}, - runtime::extensions::ExtensionPreparationPort, }; use gateway_plugin_sdk::{Capability, Contributions, Stage}; use serde_json::{Value, json}; @@ -180,7 +180,12 @@ async fn management_facts_validate_queries_without_permission_declarations() { ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - assert!(runtime.middleware_registry().resolve(&generation).is_none()); + assert!( + runtime + .execution_registry() + .middleware(&generation) + .is_none() + ); assert!(runtime.policy_registry().resolve(&generation).is_none()); let view = runtime.views(&generation).await.unwrap().remove(0); let response = runtime @@ -266,7 +271,12 @@ async fn quota_refresh_is_available_without_permission_declarations() { ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - assert!(runtime.middleware_registry().resolve(&generation).is_none()); + assert!( + runtime + .execution_registry() + .middleware(&generation) + .is_none() + ); assert!(runtime.policy_registry().resolve(&generation).is_none()); let view = runtime.views(&generation).await.unwrap().remove(0); let response = runtime diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/observer.rs b/backend/crates/gateway-plugin/runtime/tests/generation/observer.rs index d1bcc37b2..1537fa9f4 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/observer.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/observer.rs @@ -251,7 +251,7 @@ async fn frozen_plan_applies_scope_order_and_instance_deduplication() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -322,7 +322,7 @@ async fn websocket_observer_preserves_order_generation_and_complete_payload() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -416,7 +416,7 @@ async fn websocket_observer_bounds_event_count_and_payload_bytes_without_blockin limits, ) .await; - let count_generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let count_generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &count_runtime, ConfigRevision::new(1).unwrap(), ) @@ -487,7 +487,7 @@ async fn websocket_observer_bounds_event_count_and_payload_bytes_without_blockin RpcLimits::default(), ) .await; - let byte_generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let byte_generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &byte_runtime, ConfigRevision::new(1).unwrap(), ) @@ -547,7 +547,7 @@ async fn malformed_observer_response_stops_only_the_observer() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -598,7 +598,7 @@ async fn observer_failure_continues_the_plan_and_timeout_is_bounded() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -690,7 +690,7 @@ async fn client_scope_filters_observations_and_exposes_only_the_key_identifier() RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -767,7 +767,7 @@ async fn completed_observation_contains_terminal_cost_timings_and_failure() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -874,7 +874,7 @@ async fn observer_backpressure_drops_excess_and_inflight_dispatch_keeps_generati limits, ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -924,7 +924,7 @@ async fn one_observer_matches_completed_and_websocket_scopes_independently() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/compatibility.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/compatibility.rs index a9334a57a..157ee2e2b 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/compatibility.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/compatibility.rs @@ -1,7 +1,7 @@ //! 验证旧中间件的 Fast 设置适配与不兼容插件的实例故障隔离 use super::*; -use gateway_core::{policy::ClientPolicy, settings::RequestSettings}; +use gateway_core::{policy::ClientPolicy, routing::request_settings::RequestSettings}; use serde_json::{Value, json}; #[tokio::test] @@ -75,7 +75,10 @@ async fn legacy_middleware_preserves_three_state_fast_and_projects_sources() { bindings: vec![binding(MIDDLEWARE_CONTRIBUTION,"request",0,PluginFailurePolicy::Reject)], }], package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let next = (Downstream { calls: Arc::default(), closes: Arc::default(), @@ -158,7 +161,10 @@ async fn unsupported_contract_starts_and_runtime_faults_only_stop_the_failed_ins } let generation = prepare(&runtime).await; assert!(generation.is_ready(), "版本警告允许随宿主启动"); - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); for _ in 0..2 { let response = plan diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/facts/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/facts/mod.rs index 6d60184cf..31d74a8da 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/facts/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/facts/mod.rs @@ -41,7 +41,10 @@ async fn facts_are_readable_without_replacing_or_duplicating_the_original_execut bindings: vec![binding(MIDDLEWARE_CONTRIBUTION, "attempt", 0, PluginFailurePolicy::Reject)], }], package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let metadata = ProviderCallMetadata::new( ProviderKind::new("openai").unwrap(), UpstreamModelId::new("native-model").unwrap(), diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/dispatch.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/dispatch.rs index 93055beca..561776160 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/dispatch.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/dispatch.rs @@ -116,7 +116,10 @@ async fn setup( ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); (cache, runtime, generation, plan) } @@ -378,7 +381,7 @@ async fn cancelling_a_parent_drops_its_active_http_child() { #[tokio::test] async fn active_http_children_inherit_settings_and_preserve_override_sources() { - use gateway_core::{routing::RuntimeSnapshot, settings::RequestSettings}; + use gateway_core::{routing::RuntimeSnapshot, routing::request_settings::RequestSettings}; let (_cache, runtime, generation, plan) = setup("settings").await; let settings = Arc::new(Mutex::new(Vec::new())); let dispatcher: Arc = Arc::new(Dispatcher { diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs index c093e6af1..161b69492 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs @@ -4,7 +4,7 @@ mod dispatch; mod websocket; use super::*; -use gateway_core::middleware::{compose, http as core}; +use gateway_core::{engine::middleware::http as core, middleware::compose}; use http_body::Frame; use http_body_util::{BodyExt as _, StreamBody}; @@ -54,7 +54,10 @@ async fn http_plugin_with_limits( ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); assert!(plan.has_http()); (cache, runtime, generation, plan) } @@ -521,7 +524,7 @@ async fn http_onion_performance_baseline() { .collect(); let (_cache, runtime) = setup_package(instances, package.clone()).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation); + let plan = runtime.execution_registry().middleware(&generation); let snapshot = gateway_core::routing::RuntimeSnapshot::new( ConfigRevision::new(1).unwrap(), gateway_core::settings::SettingsValues::new( @@ -537,7 +540,8 @@ async fn http_onion_performance_baseline() { vec![], ) .unwrap(); - let settings = gateway_core::settings::RequestSettings::new(Arc::new(snapshot)); + let settings = + gateway_core::routing::request_settings::RequestSettings::new(Arc::new(snapshot)); for size in [64 * 1024, 2 * 1024 * 1024] { let payload = Bytes::from(vec![b'a'; size]); for _ in 0..5 { @@ -585,7 +589,7 @@ async fn http_onion_performance_baseline() { #[cfg(not(debug_assertions))] async fn benchmark_http_request( plan: Option<&gateway_core::engine::middleware::FrozenMiddlewarePlan>, - settings: &gateway_core::settings::RequestSettings, + settings: &gateway_core::routing::request_settings::RequestSettings, payload: Bytes, uppercase: bool, ) -> (Duration, Duration) { diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs index b724bd3e8..b7fa01445 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs @@ -2,7 +2,7 @@ use super::*; use futures::FutureExt as _; -use gateway_core::middleware::{http::upgrade, websocket as ws}; +use gateway_core::engine::middleware::{http::upgrade, websocket as ws}; use std::sync::atomic::AtomicBool; struct Transport { @@ -75,8 +75,8 @@ async fn session_model_and_network_callbacks_outlive_handshake_and_each_executio environment.store.start_command_line_writes().unwrap(); let snapshot = bundle.snapshots().snapshot_for_diagnostics().unwrap(); let plan = runtime - .middleware_registry() - .resolve(snapshot.extensions().unwrap()) + .execution_registry() + .middleware(snapshot.extensions().unwrap()) .unwrap(); let (send, incoming) = tokio::sync::mpsc::channel(1); let (outgoing, mut receive) = tokio::sync::mpsc::channel(1); diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs index cbd7ea501..4d30c7702 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs @@ -51,8 +51,8 @@ use gateway_core::{ lifecycle::CancellationToken, operation::{GenerateRequest, Operation, OperationKind, ProtocolPayload}, policy::ClientApiKeyId, + routing::extensions::{ExtensionPreparationPort, ExtensionSetReference}, routing::{AccountGroupId, ConfigRevision, PublicModelId}, - runtime::extensions::{ExtensionPreparationPort, ExtensionSetReference}, }; use gateway_plugin_runtime::{ PackageInspector, PackageLimits, PluginRuntime, PluginRuntimeConfig, RpcLimits, @@ -766,8 +766,8 @@ async fn attempt_middleware_can_delegate_after_a_managed_http_side_effect() { .await; let generation = prepare(&runtime).await; let plan = runtime - .middleware_registry() - .resolve(&generation) + .execution_registry() + .middleware(&generation) .expect("middleware plan"); let calls = Arc::new(AtomicUsize::new(0)); let closes = Arc::new(AtomicUsize::new(0)); @@ -821,8 +821,8 @@ async fn middleware_next_and_lazy_body_mapping_use_one_real_rpc_call() { .await; let generation = prepare(&runtime).await; let plan = runtime - .middleware_registry() - .resolve(&generation) + .execution_registry() + .middleware(&generation) .expect("middleware plan"); let calls = Arc::new(AtomicUsize::new(0)); let closes = Arc::new(AtomicUsize::new(0)); @@ -974,7 +974,10 @@ async fn middleware_classifies_real_rpc_output_after_a_terminal_frame() { }]) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let closes = Arc::new(AtomicUsize::new(0)); let response = plan @@ -1035,7 +1038,10 @@ async fn outer_plugin_receives_inner_rejection_details_and_preserves_rejection_s InstanceFixture {id:"inner-error", configuration:serde_json::json!({"mode":"rejected"}), bindings:vec![binding(MIDDLEWARE_CONTRIBUTION,"request",1,PluginFailurePolicy::Reject)]}, ],package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let result = plan .handle( middleware_context(ClientTransport::HttpJson), @@ -1090,7 +1096,10 @@ async fn real_request_plugin_receives_downstream_error_details_without_consuming bindings:vec![binding(MIDDLEWARE_CONTRIBUTION,"request",0,PluginFailurePolicy::Reject)], }],package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let result = plan .handle( middleware_context(ClientTransport::HttpJson), @@ -1140,7 +1149,9 @@ async fn real_request_plugin_receives_downstream_error_details_without_consuming #[tokio::test] async fn real_request_plugins_rewrite_settings_in_onion_order() { - use gateway_core::{policy::RateLimits, settings::ExecutionSettings, settings::SettingsValues}; + use gateway_core::{ + policy::RateLimits, routing::request_settings::ExecutionSettings, settings::SettingsValues, + }; let profiles = |identity| { BTreeMap::from([( gateway_core::identity::ProviderKind::new("openai").unwrap(), @@ -1178,26 +1189,27 @@ async fn real_request_plugins_rewrite_settings_in_onion_order() { .with_request_profiles(profiles("host")), ) .unwrap(); - let configuration = gateway_core::settings::RequestSettings::new(Arc::new(host_snapshot)) - .with_execution( - &gateway_core::policy::ClientPolicy::new( - middleware_context(ClientTransport::HttpJson) - .client_key_id() - .clone(), - gateway_core::policy::PlaintextClientApiKey::new("fixture-key").unwrap(), - Arc::new( - gateway_core::routing::FrozenAccountScope::new( - Arc::default(), - gateway_core::routing::ClientRoutingScope::all_accounts(), - ) - .with_request_profiles(profiles("key-default")) - .with_fast_mode(baseline.fast_mode), + let configuration = + gateway_core::routing::request_settings::RequestSettings::new(Arc::new(host_snapshot)) + .with_execution( + &gateway_core::policy::ClientPolicy::new( + middleware_context(ClientTransport::HttpJson) + .client_key_id() + .clone(), + gateway_core::policy::PlaintextClientApiKey::new("fixture-key").unwrap(), + Arc::new( + gateway_core::routing::FrozenAccountScope::new( + Arc::default(), + gateway_core::routing::ClientRoutingScope::all_accounts(), + ) + .with_request_profiles(profiles("key-default")) + .with_fast_mode(baseline.fast_mode), + ), + true, + baseline.client_limits, ), - true, - baseline.client_limits, - ), - baseline.timeout_ms, - ); + baseline.timeout_ms, + ); let mut first = baseline_json.clone(); first["fast_mode"] = serde_json::json!("default"); first["runtime"]["request_interval_ms"] = serde_json::json!(0); @@ -1223,7 +1235,10 @@ async fn real_request_plugins_rewrite_settings_in_onion_order() { InstanceFixture { id: "settings-pass", configuration: serde_json::json!({"mode":"settings","expected_settings":last}), bindings: vec![binding(MIDDLEWARE_CONTRIBUTION,"request",2,PluginFailurePolicy::Reject)] }, ], package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); // 连续调用使用同一宿主基线，前一次插件改写不能泄漏到后一次 for _ in 0..2 { let expected: ExecutionSettings = serde_json::from_value(last.clone()).unwrap(); @@ -1310,7 +1325,10 @@ async fn sdk_middleware_entry_registers_and_maps_a_real_process_response() { ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let closes = Arc::new(AtomicUsize::new(0)); let reads = Arc::new(AtomicUsize::new(0)); @@ -1448,7 +1466,12 @@ async fn instances_without_data_plane_bindings_publish_neither_plan() { .await; let generation = prepare(&runtime).await; assert!(runtime.policy_registry().resolve(&generation).is_none()); - assert!(runtime.middleware_registry().resolve(&generation).is_none()); + assert!( + runtime + .execution_registry() + .middleware(&generation) + .is_none() + ); drop(generation); super::wait_until_empty(cache.path()).await; } @@ -1522,7 +1545,10 @@ async fn sdk_capability_declarations_require_request_stage() { bindings: vec![binding(MIDDLEWARE_CONTRIBUTION, if attempt { "attempt" } else { "request" }, 0, PluginFailurePolicy::Reject)], }], package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let context = if attempt { attempt_middleware_context(Arc::default()) @@ -1604,7 +1630,10 @@ async fn assert_sdk_transparency(mode: &'static str) { ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let mut headers = middleware_headers(); headers.extend([ MiddlewareHeader::new("X-Repeated", Bytes::from_static(b"first")), @@ -1777,7 +1806,10 @@ async fn a_failed_middleware_preserves_scope_and_the_configured_failure_policy() }]) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let response = plan .handle( @@ -1853,7 +1885,10 @@ async fn middleware_preserves_large_body_and_following_requests() { }]) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); for length in [4 * 1024 * 1024 + 17, 32] { let payload = Bytes::from( (0..length) diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mounts.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mounts.rs index d44f76fea..aeefa995a 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mounts.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mounts.rs @@ -1,7 +1,10 @@ //! 验证混合中间件挂载保持独立顺序与各次调用的续接状态 use super::*; -use gateway_core::middleware::{compose, http, service, websocket}; +use gateway_core::{ + engine::middleware::{http, service, websocket}, + middleware::compose, +}; use http_body_util::BodyExt as _; struct Sender; @@ -71,7 +74,10 @@ async fn mixed_mounts_keep_independent_order_and_fresh_continuations() { ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); assert!(plan.has_http() && plan.has_websocket() && plan.has_service()); let mut expected = Vec::new(); for direction in [ diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/service/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/service/mod.rs index a39bb461f..f46f0337c 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/service/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/service/mod.rs @@ -294,10 +294,10 @@ async fn cli_service_callback_uses_the_same_public_settings_and_plugin_chain() { crate::support::native::admin_registry(), Arc::new(Pricing), ); - let mut registry = gateway_admin::service::Registry::new({ + let mut registry = gateway_admin::public_service::Registry::new({ let snapshots = core.snapshots(); - let middleware = runtime.middleware_registry(); - Arc::new(move || middleware.resolve(snapshots.snapshot_for_diagnostics()?.extensions()?)) + let middleware = runtime.execution_registry(); + Arc::new(move || middleware.middleware(snapshots.snapshot_for_diagnostics()?.extensions()?)) }); registry.register_settings(&settings).unwrap(); let registry = Arc::new(registry); @@ -345,7 +345,7 @@ async fn cli_service_callback_uses_the_same_public_settings_and_plugin_chain() { #[tokio::test] async fn service_chain_preserves_onion_order_parent_scope_and_native_results_without_a_database() { use super::*; - use gateway_core::middleware::{compose, service as core}; + use gateway_core::{engine::middleware::service as core, middleware::compose}; let directory = tempfile::tempdir().unwrap(); let records = directory.path().join("onion.jsonl"); @@ -377,7 +377,10 @@ async fn service_chain_preserves_onion_order_parent_scope_and_native_results_wit ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let input = json!({"opaque":[null,false,0,""]}); for reject in [false, true] { diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/websocket/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/websocket/mod.rs index d6df33241..db14cf597 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/websocket/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/websocket/mod.rs @@ -1,7 +1,7 @@ //! 验证真实插件 WebSocket 消息转发、改写与控制消息处理 use super::*; -use gateway_core::middleware::{compose, websocket as core}; +use gateway_core::{engine::middleware::websocket as core, middleware::compose}; #[derive(Default)] struct Sender(std::sync::Mutex>); @@ -46,7 +46,10 @@ async fn websocket_plugin( ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); assert!(plan.has_websocket()); (cache, runtime, generation, plan) } diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/diagnostics.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/diagnostics.rs new file mode 100644 index 000000000..fcf3bbfa6 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/diagnostics.rs @@ -0,0 +1,191 @@ +//! 验证准备、发布进程故障与排空状态的诊断投影 + +use super::wait_until_unready; +use gateway_admin::model::Revision; +use gateway_admin::model::plugins::instances::PluginInstanceRuntimeStatus; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; + +#[tokio::test] +async fn diagnostics_distinguish_running_failed_preparation_and_draining_generations() { + let (cache, store, runtime) = super::super::setup().await; + let snapshot = store.snapshot.lock().unwrap().clone(); + let active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .unwrap(); + + let running = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("runtime diagnostics"); + let running = &running["instance-one"]; + assert_eq!(running.status, PluginInstanceRuntimeStatus::Running); + assert_eq!(running.actual_revision, Some(1)); + assert_eq!( + running.actual_artifact_sha256.as_deref(), + Some(snapshot.instances[0].artifact_sha256.as_str()) + ); + assert!(running.failure.is_none()); + + let mut failed = snapshot.clone(); + failed.config_revision = Revision::new(2).unwrap(); + failed.instances[0].revision = Revision::new(2).unwrap(); + failed.instances[0].configuration = serde_json::json!({"startup":"fail"}); + assert!( + PluginPreparation::prepare(&runtime, failed.clone()) + .await + .is_err() + ); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &failed, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("failed preparation diagnostics"); + let failed_runtime = &diagnostics["instance-one"]; + assert_eq!( + failed_runtime.status, + PluginInstanceRuntimeStatus::PreparationFailed + ); + assert!(failed_runtime.failure.is_some()); + assert_eq!(failed_runtime.actual_revision, Some(1)); + + let mut disabled = failed; + disabled.config_revision = Revision::new(3).unwrap(); + disabled.instances[0].revision = Revision::new(3).unwrap(); + disabled.instances[0].enabled = false; + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &disabled, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("draining diagnostics"); + let draining = &diagnostics["instance-one"]; + assert_eq!(draining.status, PluginInstanceRuntimeStatus::Draining); + assert_eq!(draining.draining_revisions, [1]); + + drop(active); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn preparing_diagnostics_never_wait_for_the_serialized_prepare_io() { + let (cache, store, runtime) = super::super::setup().await; + let source = store.snapshot.lock().unwrap().clone(); + let mut candidate = source.clone(); + candidate.config_revision = Revision::new(2).unwrap(); + candidate.instances[0].revision = Revision::new(2).unwrap(); + candidate.instances[0].configuration = serde_json::json!({"startup_delay_ms":400}); + let diagnostic_snapshot = candidate.clone(); + let mut preparation = Box::pin(PluginPreparation::prepare(&runtime, candidate)); + assert!(futures::poll!(preparation.as_mut()).is_pending()); + + let diagnostics = tokio::time::timeout( + std::time::Duration::from_millis(100), + gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &diagnostic_snapshot, + None, + None, + ), + ) + .await + .expect("diagnostics must not wait for plugin startup") + .expect("runtime diagnostics"); + assert_eq!( + diagnostics["instance-one"].status, + PluginInstanceRuntimeStatus::Preparing + ); + + let candidate = preparation.await.expect("delayed candidate"); + drop(candidate); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn diagnostics_report_a_published_process_fault_without_plugin_details() { + let (cache, store, runtime) = super::super::setup().await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({"exit_after_ready_ms":20}); + let active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("candidate starts before the controlled exit"); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while active.is_ready() { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("worker exit observed"); + + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("runtime diagnostics"); + let fault = &diagnostics["instance-one"]; + assert_eq!(fault.status, PluginInstanceRuntimeStatus::Faulted); + assert_eq!( + fault.failure.as_ref().map(|failure| failure.code.as_str()), + Some("closed") + ); + assert_eq!( + fault + .failure + .as_ref() + .map(|failure| failure.message.as_str()), + Some("插件进程或传输已停止") + ); + + drop(active); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn a_published_process_crash_preserves_the_serving_snapshot_and_other_plugin_status() { + let (cache, store, runtime) = super::super::setup().await; + let snapshot = { + let mut snapshot = store.snapshot.lock().unwrap(); + let mut failed = snapshot.instances[0].clone(); + failed.id = "crashing-plugin".into(); + failed.configuration = serde_json::json!({"exit_after_ready_ms":100}); + snapshot.instances.push(failed); + snapshot.clone() + }; + let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + wait_until_unready(&generation).await; + assert!(generation.can_serve()); + assert!(!generation.is_ready()); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(1), + Some(&generation), + ) + .await + .unwrap(); + assert_eq!( + diagnostics["instance-one"].status, + PluginInstanceRuntimeStatus::Running + ); + assert_eq!( + diagnostics["crashing-plugin"].status, + PluginInstanceRuntimeStatus::Faulted + ); + drop(generation); + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs new file mode 100644 index 000000000..d711f0584 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs @@ -0,0 +1,511 @@ +//! 验证实例握手、声明与重启熔断的生命周期 + +use super::wait_until_unready; +use gateway_admin::model::Revision; +use gateway_admin::model::plugins::instances::PluginInstanceRuntimeStatus; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; +use gateway_plugin_sdk::Capability; +use gateway_plugin_sdk::Contributions; +use gateway_plugin_sdk::Stage; +use std::num::NonZeroUsize; +use std::path::Path; +use std::time::Duration; + +fn startup_count(marker: &Path) -> usize { + std::fs::read_to_string(marker) + .map(|content| content.lines().count()) + .unwrap_or_default() +} + +#[tokio::test] +async fn registration_with_a_changed_contribution_id_is_rejected() { + let (cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::Observer, + vec![Stage::Observation], + Vec::new(), + Vec::new(), + ), + ])) + .await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("the declared contribution prepares before its registration id changes"); + drop(candidate); + super::super::wait_until_empty(cache.path()).await; + + snapshot.instances[0].configuration = serde_json::json!({"registration_mismatch":true}); + let error = PluginPreparation::prepare(&runtime, snapshot) + .await + .expect_err("a changed registration contribution id must be rejected"); + assert!( + error.to_string().contains("插件注册结果与清单不符"), + "unexpected preparation failure: {error}" + ); +} + +#[tokio::test] +async fn restart_circuit_caps_short_lived_crashes_and_a_new_revision_recovers() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(3).unwrap(), + stability_window: Duration::from_secs(1), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("starts.jsonl"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker": marker, + "exit_after_ready_ms": 20, + }); + let mut active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("first incarnation"); + for expected_start in 2..=3 { + wait_until_unready(&active).await; + let replacement = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("failure budget still permits a restart"); + drop(active); + active = replacement; + assert_eq!(startup_count(&marker), expected_start); + } + wait_until_unready(&active).await; + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err(), + "the fourth incarnation must not start" + ); + assert_eq!(startup_count(&marker), 3); + + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("runtime diagnostics"); + assert_eq!( + diagnostics["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("restart_circuit_open") + ); + + let mut disabled = snapshot; + disabled.config_revision = Revision::new(2).unwrap(); + disabled.instances[0].revision = Revision::new(2).unwrap(); + disabled.instances[0].enabled = false; + let disabled_generation = PluginPreparation::prepare(&runtime, disabled.clone()) + .await + .expect("disabling the instance does not reuse its open circuit"); + assert_eq!(startup_count(&marker), 3); + + let mut repaired = disabled; + repaired.config_revision = Revision::new(3).unwrap(); + repaired.instances[0].revision = Revision::new(3).unwrap(); + repaired.instances[0].enabled = true; + repaired.instances[0].configuration = serde_json::json!({"startup_marker":marker}); + let replacement = PluginPreparation::prepare(&runtime, repaired) + .await + .expect("a new instance revision has an independent failure budget"); + assert!(replacement.is_ready()); + assert_eq!(startup_count(&marker), 4); + + drop(active); + drop(disabled_generation); + drop(replacement); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_resets_after_a_stable_incarnation() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(3).unwrap(), + stability_window: Duration::from_millis(80), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("stable-reset-starts.jsonl"); + let exit = controls.path().join("stable-exit"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "startup_failures":[true,true,false,true,true], + "exit_after_ready_signals":[null,null,exit], + }); + // 握手前失败的有效运行时间固定为零；不能用短 sleep 假定宿主一定及时观察到退出 + for _ in 0..2 { + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + } + let active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("the third incarnation starts before the circuit opens"); + assert!(active.is_ready()); + // 明确等到稳定窗口之后才允许第三个进程退出；调度变慢不会把短命进程误变为稳定进程 + tokio::time::sleep(restart_circuit.stability_window).await; + tokio::fs::write(exit, b"exit").await.unwrap(); + wait_until_unready(&active).await; + for expected_start in 4..=5 { + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + assert_eq!(startup_count(&marker), expected_start); + } + assert!( + PluginPreparation::prepare(&runtime, snapshot) + .await + .is_err(), + "three new failures after the stable incarnation reopen the circuit" + ); + assert_eq!(startup_count(&marker), 5); + + drop(active); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_ignores_planned_generation_shutdown() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::MIN, + stability_window: Duration::from_secs(60), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("planned-shutdown-starts.jsonl"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({"startup_marker":marker}); + let first = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("first generation"); + let instance = &snapshot.instances[0]; + gateway_admin::ports::plugins::PluginStateLifecycle::quiesce_instance( + &runtime, + &instance.id, + &instance.artifact_sha256, + instance.revision, + ) + .await; + assert!(!first.is_ready()); + + let second = PluginPreparation::prepare(&runtime, snapshot) + .await + .expect("planned shutdown must not open the one-failure circuit"); + assert!(second.is_ready()); + assert_eq!(startup_count(&marker), 2); + + drop(first); + drop(second); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_counts_exit_before_the_candidate_is_indexed() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(3).unwrap(), + stability_window: Duration::from_secs(1), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("registration-exits.jsonl"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "exit_during_registration":true, + }); + for _ in 0..3 { + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + } + assert_eq!(startup_count(&marker), 3); + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + assert_eq!( + startup_count(&marker), + 3, + "open circuit must gate process spawn" + ); + + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, &snapshot, None, None, + ) + .await + .expect("runtime diagnostics"); + assert_eq!( + diagnostics["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("restart_circuit_open") + ); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_counts_a_started_process_that_exits_during_handshake() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(2).unwrap(), + stability_window: Duration::from_millis(10), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("handshake-exits.jsonl"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "startup":"fail", + "startup_fail_delay_ms":50, + }); + for _ in 0..2 { + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + } + assert_eq!(startup_count(&marker), 2); + assert!( + PluginPreparation::prepare(&runtime, snapshot) + .await + .is_err() + ); + assert_eq!( + startup_count(&marker), + 2, + "host-side gating must happen before another process spawn" + ); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_does_not_let_an_older_generation_clear_newer_failures() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(3).unwrap(), + stability_window: Duration::from_millis(50), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("overlapping-generations.jsonl"); + let old_exit = controls.path().join("old-generation-exit"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "startup_failures":[false,true,true,true], + "exit_after_ready_signals":[old_exit], + }); + let mut other = snapshot.instances[0].clone(); + other.id = "other-instance".into(); + other.name = "Other".into(); + other.configuration = serde_json::json!({}); + snapshot.instances.push(other); + let published = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("published generation"); + assert!(published.is_ready()); + // 旧代次明确越过稳定窗口；新代次在握手前失败，运行时间固定为零 + // 不依赖 15 ms 退出与注册完成的竞速，也不让调度延迟重置新失败预算 + tokio::time::sleep(restart_circuit.stability_window).await; + + let mut failed_candidate = None; + for revision in 2..=4 { + snapshot.config_revision = Revision::new(revision).unwrap(); + snapshot.instances[1].revision = Revision::new(revision).unwrap(); + snapshot.instances[1].configuration = serde_json::json!({"revision":revision}); + let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("the unrelated instance starts while the newer incarnation fails"); + assert_eq!(startup_count(&marker), revision as usize); + drop(failed_candidate.replace(candidate)); + assert!(published.is_ready(), "the older generation remains healthy"); + } + drop(failed_candidate.take()); + std::fs::write(&old_exit, b"exit").unwrap(); + wait_until_unready(&published).await; + snapshot.config_revision = Revision::new(5).unwrap(); + snapshot.instances[1].revision = Revision::new(5).unwrap(); + snapshot.instances[1].configuration = serde_json::json!({"revision":5}); + let circuit = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, &snapshot, None, None, + ) + .await + .unwrap(); + assert_eq!( + circuit["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("restart_circuit_open"), + ); + let recovered = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("an older failed instance is isolated during unrelated recovery"); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(5), + Some(&recovered), + ) + .await + .unwrap(); + assert_eq!( + diagnostics["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("unavailable"), + "the older healthy session must not erase newer failures for the same identity" + ); + assert_eq!(startup_count(&marker), 4); + + drop(recovered); + drop(published); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_restores_an_older_failure_after_candidate_shutdown() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::MIN, + stability_window: Duration::from_secs(1), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("abandoned-candidate.jsonl"); + let old_exit = controls.path().join("abandoned-candidate-old-exit"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "exit_after_ready_delays_ms":[null,null], + "exit_after_ready_signals":[old_exit,null], + }); + let mut other = snapshot.instances[0].clone(); + other.id = "other-instance".into(); + other.name = "Other".into(); + other.configuration = serde_json::json!({}); + snapshot.instances.push(other); + let published = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("published generation"); + + snapshot.config_revision = Revision::new(2).unwrap(); + snapshot.instances[1].revision = Revision::new(2).unwrap(); + snapshot.instances[1].configuration = serde_json::json!({"revision":2}); + let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("healthy replacement candidate"); + std::fs::write(&old_exit, b"exit").unwrap(); + wait_until_unready(&published).await; + assert!(candidate.is_ready()); + drop(candidate); + + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let diagnostics = + gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(1), + Some(&published), + ) + .await + .expect("runtime diagnostics"); + if diagnostics["instance-one"] + .failure + .as_ref() + .is_some_and(|failure| failure.code == "restart_circuit_open") + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("older failure is restored after candidate shutdown"); + + snapshot.config_revision = Revision::new(3).unwrap(); + snapshot.instances[1].revision = Revision::new(3).unwrap(); + snapshot.instances[1].configuration = serde_json::json!({"revision":3}); + let recovered = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("an older failed instance is isolated during unrelated recovery"); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(3), + Some(&recovered), + ) + .await + .unwrap(); + assert_eq!( + diagnostics["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("unavailable"), + ); + assert_eq!(startup_count(&marker), 2); + + drop(recovered); + drop(published); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn a_host_version_mismatch_still_restores_the_enabled_plugin() { + use std::sync::Arc; + let (cache, store, _) = super::super::setup().await; + let runtime = gateway_plugin_runtime::PluginRuntime::new( + store.clone(), + store.clone(), + gateway_plugin_runtime::PluginRuntimeConfig { + cache_directory: cache.path().to_owned(), + host_version: "3.0.0".parse().unwrap(), + package_limits: Default::default(), + rpc_limits: Default::default(), + restart_circuit: Default::default(), + }, + Arc::new(gateway_host::outbound::HttpClient::new().unwrap()), + Arc::new(gateway_host::process::ProcessSupervisor::default()), + ); + let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert!(restored.can_serve()); + let snapshot = store.snapshot.lock().unwrap().clone(); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(1), + Some(&restored), + ) + .await + .unwrap(); + assert!(restored.is_ready()); + assert!(snapshot.instances[0].enabled); + assert_eq!( + diagnostics["instance-one"].status, + PluginInstanceRuntimeStatus::Running + ); + assert!(diagnostics["instance-one"].failure.is_none()); + drop(restored); + runtime.shutdown().await; + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/management.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/management.rs new file mode 100644 index 000000000..e674a5fd1 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/management.rs @@ -0,0 +1,50 @@ +//! 验证管理调用绑定、执行身份与候选配置 + +use gateway_admin::model::plugins::instances::PluginCapabilityBinding; +use gateway_admin::model::plugins::instances::PluginFailurePolicy; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_plugin_sdk::Capability; +use gateway_plugin_sdk::Contributions; +use gateway_plugin_sdk::Stage; + +#[tokio::test] +async fn management_needs_no_binding_and_rejects_stale_execution_identity_configuration() { + let (_cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::Management, + vec![Stage::Management], + Vec::new(), + Vec::new(), + ), + ])) + .await; + let mut instance = store.snapshot.lock().unwrap().instances[0].clone(); + assert!( + PluginPreparation::validate(&runtime, instance.clone()) + .await + .is_ok() + ); + instance.bindings = vec![PluginCapabilityBinding { + contribution: "test.example.management".into(), + stage: "management".into(), + order: 0, + failure_policy: PluginFailurePolicy::Reject, + client_key_ids: vec!["key-one".into(), "key-two".into()], + account_group_ids: Vec::new(), + provider_ids: Vec::new(), + models: Vec::new(), + event: None, + identity_bindings: Vec::new(), + }]; + assert!( + PluginPreparation::validate(&runtime, instance.clone()) + .await + .is_err() + ); + instance.enabled = false; + assert!( + PluginPreparation::validate(&runtime, instance) + .await + .is_err() + ); +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs index 101d49c44..e7531c5cc 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs @@ -1,19 +1,15 @@ -//! 验证插件发布准备的目录冻结、声明校验与运行时关闭 +//! 插件发布准备各职责的集成测试入口 -use gateway_admin::{ - model::{ - Revision, - plugins::instances::{ - PluginCapabilityBinding, PluginFailurePolicy, PluginInstanceRuntimeStatus, - }, - }, - ports::plugins::PluginPreparation, -}; -use gateway_core::{routing::ConfigRevision, runtime::extensions::ExtensionPreparationPort}; -use gateway_plugin_sdk::{Capability, Contributions, Stage}; -use std::{num::NonZeroUsize, path::Path, time::Duration}; +mod diagnostics; +mod instance; +mod management; +mod runtime; +mod set; +mod snapshot; -async fn wait_until_unready(reference: &gateway_core::runtime::extensions::ExtensionSetReference) { +use std::time::Duration; + +async fn wait_until_unready(reference: &gateway_core::routing::extensions::ExtensionSetReference) { tokio::time::timeout(Duration::from_secs(2), async { while reference.is_ready() { tokio::time::sleep(Duration::from_millis(10)).await; @@ -22,1003 +18,3 @@ async fn wait_until_unready(reference: &gateway_core::runtime::extensions::Exten .await .expect("worker exit observed"); } - -fn startup_count(marker: &Path) -> usize { - std::fs::read_to_string(marker) - .map(|content| content.lines().count()) - .unwrap_or_default() -} - -#[tokio::test] -async fn model_catalog_is_frozen_without_entering_the_request_chain() { - let (cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::ModelCatalog, - vec![Stage::Registration], - vec![], - vec![], - )])) - .await; - store.snapshot.lock().unwrap().instances[0].configuration = serde_json::json!({ - "model_catalog": {"models": [{"id": "my-model", "provider": "openai", "model": "upstream-model"}]} - }); - let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert_eq!(generation.model_aliases().len(), 1); - assert_eq!(generation.model_aliases()[0].id.as_str(), "my-model"); - assert_eq!( - generation.model_aliases()[0].target.as_str(), - "upstream-model" - ); - assert!(runtime.policy_registry().resolve(&generation).is_none()); - assert!(runtime.middleware_registry().resolve(&generation).is_none()); - let reused = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert_eq!(reused.id(), generation.id()); - assert_eq!(reused.model_aliases(), generation.model_aliases()); - { - let mut snapshot = store.snapshot.lock().unwrap(); - snapshot.config_revision = Revision::new(2).unwrap(); - snapshot.instances[0].revision = Revision::new(2).unwrap(); - snapshot.instances[0].configuration = serde_json::json!({"startup":"fail"}); - } - assert!( - ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) - .await - .is_err(), - "failed recovery must not publish a catalog with silently removed aliases" - ); - assert!(generation.is_ready()); - assert_eq!(generation.model_aliases()[0].id.as_str(), "my-model"); - { - let mut snapshot = store.snapshot.lock().unwrap(); - snapshot.config_revision = Revision::new(2).unwrap(); - snapshot.instances[0].enabled = false; - } - let disabled = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) - .await - .unwrap(); - assert!(disabled.model_aliases().is_empty()); - assert_eq!(generation.model_aliases().len(), 1); - drop(disabled); - drop(reused); - drop(generation); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn model_catalog_rejects_invalid_registrations_and_cross_instance_conflicts() { - for models in [ - serde_json::json!([{"id":"alias","provider":"custom","model":"target"}]), - serde_json::json!([{"id":"alias","provider":"openai","model":"alias"}]), - serde_json::json!([ - {"id":"alias","provider":"openai","model":"target"}, - {"id":"alias","provider":"xai","model":"target"} - ]), - serde_json::json!([{"id":"alias","provider":"openai","model":"target","capabilities":[]}]), - ] { - let (cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::ModelCatalog, - vec![Stage::Registration], - vec![], - vec![], - )])) - .await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = - serde_json::json!({"model_catalog":{"models":models}}); - assert!( - PluginPreparation::prepare(&runtime, snapshot) - .await - .is_err() - ); - runtime.shutdown().await; - super::wait_until_empty(cache.path()).await; - } - - let (cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::ModelCatalog, - vec![Stage::Registration], - vec![], - vec![], - )])) - .await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "model_catalog":{"models":[{"id":"alias","provider":"openai","model":"target"}]} - }); - let mut other = snapshot.instances[0].clone(); - other.id = "instance-two".into(); - snapshot.instances.push(other); - let error = PluginPreparation::prepare(&runtime, snapshot) - .await - .unwrap_err(); - assert!(error.to_string().contains("instance-one")); - assert!(error.to_string().contains("instance-two")); - runtime.shutdown().await; - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn runtime_shutdown_reaps_live_generation_and_rejects_late_preparation() { - let (cache, store, runtime) = super::setup().await; - let published = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert!(published.is_ready()); - assert_eq!(std::fs::read_dir(cache.path()).unwrap().count(), 1); - - runtime.shutdown().await; - assert!(!published.is_ready()); - super::wait_until_empty(cache.path()).await; - assert!( - ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .is_err() - ); - runtime.shutdown().await; - - drop(published); - drop(store); -} - -#[tokio::test] -async fn management_needs_no_binding_and_rejects_stale_execution_identity_configuration() { - let (_cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::Management, - vec![Stage::Management], - Vec::new(), - Vec::new(), - )])) - .await; - let mut instance = store.snapshot.lock().unwrap().instances[0].clone(); - assert!( - PluginPreparation::validate(&runtime, instance.clone()) - .await - .is_ok() - ); - instance.bindings = vec![PluginCapabilityBinding { - contribution: "test.example.management".into(), - stage: "management".into(), - order: 0, - failure_policy: PluginFailurePolicy::Reject, - client_key_ids: vec!["key-one".into(), "key-two".into()], - account_group_ids: Vec::new(), - provider_ids: Vec::new(), - models: Vec::new(), - event: None, - identity_bindings: Vec::new(), - }]; - assert!( - PluginPreparation::validate(&runtime, instance.clone()) - .await - .is_err() - ); - instance.enabled = false; - assert!( - PluginPreparation::validate(&runtime, instance) - .await - .is_err() - ); -} - -#[tokio::test] -async fn registration_with_a_changed_contribution_id_is_rejected() { - let (cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::Observer, - vec![Stage::Observation], - Vec::new(), - Vec::new(), - )])) - .await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("the declared contribution prepares before its registration id changes"); - drop(candidate); - super::wait_until_empty(cache.path()).await; - - snapshot.instances[0].configuration = serde_json::json!({"registration_mismatch":true}); - let error = PluginPreparation::prepare(&runtime, snapshot) - .await - .expect_err("a changed registration contribution id must be rejected"); - assert!( - error.to_string().contains("插件注册结果与清单不符"), - "unexpected preparation failure: {error}" - ); -} - -#[tokio::test] -async fn candidate_is_reused_after_commit_and_released_when_no_generation_owner_remains() { - let (cache, store, runtime) = super::setup().await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = - serde_json::from_str(r#"{"z":1,"a":{"z":2,"a":3}}"#).unwrap(); - let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .unwrap(); - assert!( - runtime.observer_registry().resolve(&candidate).is_none(), - "没有观察绑定时不应编译派发计划" - ); - snapshot.instances[0].configuration = - serde_json::from_str(r#"{"a":{"a":3,"z":2},"z":1}"#).unwrap(); - snapshot.config_revision = Revision::new(2).unwrap(); - *store.snapshot.lock().unwrap() = snapshot; - let published = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) - .await - .unwrap(); - assert_eq!(candidate.id(), published.id()); - assert!(published.is_ready()); - drop(candidate); - assert_eq!(std::fs::read_dir(cache.path()).unwrap().count(), 1); - drop(published); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn diagnostics_distinguish_running_failed_preparation_and_draining_generations() { - let (cache, store, runtime) = super::setup().await; - let snapshot = store.snapshot.lock().unwrap().clone(); - let active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .unwrap(); - - let running = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("runtime diagnostics"); - let running = &running["instance-one"]; - assert_eq!(running.status, PluginInstanceRuntimeStatus::Running); - assert_eq!(running.actual_revision, Some(1)); - assert_eq!( - running.actual_artifact_sha256.as_deref(), - Some(snapshot.instances[0].artifact_sha256.as_str()) - ); - assert!(running.failure.is_none()); - - let mut failed = snapshot.clone(); - failed.config_revision = Revision::new(2).unwrap(); - failed.instances[0].revision = Revision::new(2).unwrap(); - failed.instances[0].configuration = serde_json::json!({"startup":"fail"}); - assert!( - PluginPreparation::prepare(&runtime, failed.clone()) - .await - .is_err() - ); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &failed, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("failed preparation diagnostics"); - let failed_runtime = &diagnostics["instance-one"]; - assert_eq!( - failed_runtime.status, - PluginInstanceRuntimeStatus::PreparationFailed - ); - assert!(failed_runtime.failure.is_some()); - assert_eq!(failed_runtime.actual_revision, Some(1)); - - let mut disabled = failed; - disabled.config_revision = Revision::new(3).unwrap(); - disabled.instances[0].revision = Revision::new(3).unwrap(); - disabled.instances[0].enabled = false; - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &disabled, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("draining diagnostics"); - let draining = &diagnostics["instance-one"]; - assert_eq!(draining.status, PluginInstanceRuntimeStatus::Draining); - assert_eq!(draining.draining_revisions, [1]); - - drop(active); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn preparing_diagnostics_never_wait_for_the_serialized_prepare_io() { - let (cache, store, runtime) = super::setup().await; - let source = store.snapshot.lock().unwrap().clone(); - let mut candidate = source.clone(); - candidate.config_revision = Revision::new(2).unwrap(); - candidate.instances[0].revision = Revision::new(2).unwrap(); - candidate.instances[0].configuration = serde_json::json!({"startup_delay_ms":400}); - let diagnostic_snapshot = candidate.clone(); - let mut preparation = Box::pin(PluginPreparation::prepare(&runtime, candidate)); - assert!(futures::poll!(preparation.as_mut()).is_pending()); - - let diagnostics = tokio::time::timeout( - std::time::Duration::from_millis(100), - gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &diagnostic_snapshot, - None, - None, - ), - ) - .await - .expect("diagnostics must not wait for plugin startup") - .expect("runtime diagnostics"); - assert_eq!( - diagnostics["instance-one"].status, - PluginInstanceRuntimeStatus::Preparing - ); - - let candidate = preparation.await.expect("delayed candidate"); - drop(candidate); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn diagnostics_report_a_published_process_fault_without_plugin_details() { - let (cache, store, runtime) = super::setup().await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({"exit_after_ready_ms":20}); - let active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("candidate starts before the controlled exit"); - tokio::time::timeout(std::time::Duration::from_secs(2), async { - while active.is_ready() { - tokio::time::sleep(std::time::Duration::from_millis(10)).await; - } - }) - .await - .expect("worker exit observed"); - - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("runtime diagnostics"); - let fault = &diagnostics["instance-one"]; - assert_eq!(fault.status, PluginInstanceRuntimeStatus::Faulted); - assert_eq!( - fault.failure.as_ref().map(|failure| failure.code.as_str()), - Some("closed") - ); - assert_eq!( - fault - .failure - .as_ref() - .map(|failure| failure.message.as_str()), - Some("插件进程或传输已停止") - ); - - drop(active); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_caps_short_lived_crashes_and_a_new_revision_recovers() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(3).unwrap(), - stability_window: Duration::from_secs(1), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("starts.jsonl"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker": marker, - "exit_after_ready_ms": 20, - }); - let mut active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("first incarnation"); - for expected_start in 2..=3 { - wait_until_unready(&active).await; - let replacement = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("failure budget still permits a restart"); - drop(active); - active = replacement; - assert_eq!(startup_count(&marker), expected_start); - } - wait_until_unready(&active).await; - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err(), - "the fourth incarnation must not start" - ); - assert_eq!(startup_count(&marker), 3); - - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("runtime diagnostics"); - assert_eq!( - diagnostics["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("restart_circuit_open") - ); - - let mut disabled = snapshot; - disabled.config_revision = Revision::new(2).unwrap(); - disabled.instances[0].revision = Revision::new(2).unwrap(); - disabled.instances[0].enabled = false; - let disabled_generation = PluginPreparation::prepare(&runtime, disabled.clone()) - .await - .expect("disabling the instance does not reuse its open circuit"); - assert_eq!(startup_count(&marker), 3); - - let mut repaired = disabled; - repaired.config_revision = Revision::new(3).unwrap(); - repaired.instances[0].revision = Revision::new(3).unwrap(); - repaired.instances[0].enabled = true; - repaired.instances[0].configuration = serde_json::json!({"startup_marker":marker}); - let replacement = PluginPreparation::prepare(&runtime, repaired) - .await - .expect("a new instance revision has an independent failure budget"); - assert!(replacement.is_ready()); - assert_eq!(startup_count(&marker), 4); - - drop(active); - drop(disabled_generation); - drop(replacement); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_resets_after_a_stable_incarnation() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(3).unwrap(), - stability_window: Duration::from_millis(80), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("stable-reset-starts.jsonl"); - let exit = controls.path().join("stable-exit"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "startup_failures":[true,true,false,true,true], - "exit_after_ready_signals":[null,null,exit], - }); - // 握手前失败的有效运行时间固定为零；不能用短 sleep 假定宿主一定及时观察到退出 - for _ in 0..2 { - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - } - let active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("the third incarnation starts before the circuit opens"); - assert!(active.is_ready()); - // 明确等到稳定窗口之后才允许第三个进程退出；调度变慢不会把短命进程误变为稳定进程 - tokio::time::sleep(restart_circuit.stability_window).await; - tokio::fs::write(exit, b"exit").await.unwrap(); - wait_until_unready(&active).await; - for expected_start in 4..=5 { - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - assert_eq!(startup_count(&marker), expected_start); - } - assert!( - PluginPreparation::prepare(&runtime, snapshot) - .await - .is_err(), - "three new failures after the stable incarnation reopen the circuit" - ); - assert_eq!(startup_count(&marker), 5); - - drop(active); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_ignores_planned_generation_shutdown() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::MIN, - stability_window: Duration::from_secs(60), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("planned-shutdown-starts.jsonl"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({"startup_marker":marker}); - let first = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("first generation"); - let instance = &snapshot.instances[0]; - gateway_admin::ports::plugins::PluginStateLifecycle::quiesce_instance( - &runtime, - &instance.id, - &instance.artifact_sha256, - instance.revision, - ) - .await; - assert!(!first.is_ready()); - - let second = PluginPreparation::prepare(&runtime, snapshot) - .await - .expect("planned shutdown must not open the one-failure circuit"); - assert!(second.is_ready()); - assert_eq!(startup_count(&marker), 2); - - drop(first); - drop(second); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_counts_exit_before_the_candidate_is_indexed() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(3).unwrap(), - stability_window: Duration::from_secs(1), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("registration-exits.jsonl"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "exit_during_registration":true, - }); - for _ in 0..3 { - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - } - assert_eq!(startup_count(&marker), 3); - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - assert_eq!( - startup_count(&marker), - 3, - "open circuit must gate process spawn" - ); - - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, &snapshot, None, None, - ) - .await - .expect("runtime diagnostics"); - assert_eq!( - diagnostics["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("restart_circuit_open") - ); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_counts_a_started_process_that_exits_during_handshake() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(2).unwrap(), - stability_window: Duration::from_millis(10), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("handshake-exits.jsonl"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "startup":"fail", - "startup_fail_delay_ms":50, - }); - for _ in 0..2 { - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - } - assert_eq!(startup_count(&marker), 2); - assert!( - PluginPreparation::prepare(&runtime, snapshot) - .await - .is_err() - ); - assert_eq!( - startup_count(&marker), - 2, - "host-side gating must happen before another process spawn" - ); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_does_not_let_an_older_generation_clear_newer_failures() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(3).unwrap(), - stability_window: Duration::from_millis(50), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("overlapping-generations.jsonl"); - let old_exit = controls.path().join("old-generation-exit"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "startup_failures":[false,true,true,true], - "exit_after_ready_signals":[old_exit], - }); - let mut other = snapshot.instances[0].clone(); - other.id = "other-instance".into(); - other.name = "Other".into(); - other.configuration = serde_json::json!({}); - snapshot.instances.push(other); - let published = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("published generation"); - assert!(published.is_ready()); - // 旧代次明确越过稳定窗口；新代次在握手前失败，运行时间固定为零 - // 不依赖 15 ms 退出与注册完成的竞速，也不让调度延迟重置新失败预算 - tokio::time::sleep(restart_circuit.stability_window).await; - - let mut failed_candidate = None; - for revision in 2..=4 { - snapshot.config_revision = Revision::new(revision).unwrap(); - snapshot.instances[1].revision = Revision::new(revision).unwrap(); - snapshot.instances[1].configuration = serde_json::json!({"revision":revision}); - let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("the unrelated instance starts while the newer incarnation fails"); - assert_eq!(startup_count(&marker), revision as usize); - drop(failed_candidate.replace(candidate)); - assert!(published.is_ready(), "the older generation remains healthy"); - } - drop(failed_candidate.take()); - std::fs::write(&old_exit, b"exit").unwrap(); - wait_until_unready(&published).await; - snapshot.config_revision = Revision::new(5).unwrap(); - snapshot.instances[1].revision = Revision::new(5).unwrap(); - snapshot.instances[1].configuration = serde_json::json!({"revision":5}); - let circuit = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, &snapshot, None, None, - ) - .await - .unwrap(); - assert_eq!( - circuit["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("restart_circuit_open"), - ); - let recovered = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("an older failed instance is isolated during unrelated recovery"); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(5), - Some(&recovered), - ) - .await - .unwrap(); - assert_eq!( - diagnostics["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("unavailable"), - "the older healthy session must not erase newer failures for the same identity" - ); - assert_eq!(startup_count(&marker), 4); - - drop(recovered); - drop(published); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_restores_an_older_failure_after_candidate_shutdown() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::MIN, - stability_window: Duration::from_secs(1), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("abandoned-candidate.jsonl"); - let old_exit = controls.path().join("abandoned-candidate-old-exit"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "exit_after_ready_delays_ms":[null,null], - "exit_after_ready_signals":[old_exit,null], - }); - let mut other = snapshot.instances[0].clone(); - other.id = "other-instance".into(); - other.name = "Other".into(); - other.configuration = serde_json::json!({}); - snapshot.instances.push(other); - let published = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("published generation"); - - snapshot.config_revision = Revision::new(2).unwrap(); - snapshot.instances[1].revision = Revision::new(2).unwrap(); - snapshot.instances[1].configuration = serde_json::json!({"revision":2}); - let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("healthy replacement candidate"); - std::fs::write(&old_exit, b"exit").unwrap(); - wait_until_unready(&published).await; - assert!(candidate.is_ready()); - drop(candidate); - - tokio::time::timeout(Duration::from_secs(2), async { - loop { - let diagnostics = - gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(1), - Some(&published), - ) - .await - .expect("runtime diagnostics"); - if diagnostics["instance-one"] - .failure - .as_ref() - .is_some_and(|failure| failure.code == "restart_circuit_open") - { - break; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("older failure is restored after candidate shutdown"); - - snapshot.config_revision = Revision::new(3).unwrap(); - snapshot.instances[1].revision = Revision::new(3).unwrap(); - snapshot.instances[1].configuration = serde_json::json!({"revision":3}); - let recovered = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("an older failed instance is isolated during unrelated recovery"); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(3), - Some(&recovered), - ) - .await - .unwrap(); - assert_eq!( - diagnostics["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("unavailable"), - ); - assert_eq!(startup_count(&marker), 2); - - drop(recovered); - drop(published); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn failed_candidate_leaves_the_previous_generation_ready() { - let (cache, store, runtime) = super::setup().await; - let snapshot = store.snapshot.lock().unwrap().clone(); - let active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .unwrap(); - let mut candidate = snapshot; - candidate.instances[0].configuration = serde_json::json!({"startup":"fail"}); - assert!( - PluginPreparation::prepare(&runtime, candidate) - .await - .is_err() - ); - assert!(active.is_ready()); - drop(active); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn disabled_missing_artifacts_do_not_block_the_remaining_generation_or_offline_recovery() { - let (cache, store, runtime) = super::setup().await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - let mut missing = snapshot.instances[0].clone(); - missing.id = "missing".into(); - missing.enabled = false; - missing.artifact_sha256 = "0".repeat(64); - snapshot.instances.push(missing); - let active = PluginPreparation::prepare(&runtime, snapshot) - .await - .unwrap(); - assert!(active.is_ready()); - drop(active); - super::wait_until_empty(cache.path()).await; - let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert!(restored.is_ready()); - drop(restored); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn preparation_cannot_mix_a_newer_persisted_revision_with_an_older_snapshot() { - let (_cache, _store, runtime) = super::setup().await; - assert!( - ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) - .await - .is_err() - ); -} - -#[tokio::test] -async fn restoring_a_failed_plugin_keeps_other_instances_ready_and_reports_only_its_failure() { - let (cache, store, runtime) = super::setup().await; - let snapshot = { - let mut snapshot = store.snapshot.lock().unwrap(); - let mut failed = snapshot.instances[0].clone(); - failed.id = "failed-plugin".into(); - failed.configuration = serde_json::json!({"startup":"fail"}); - snapshot.instances.insert(0, failed); - snapshot.clone() - }; - let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .expect("one failed plugin must not stop restoration"); - assert!(generation.can_serve()); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(1), - Some(&generation), - ) - .await - .unwrap(); - assert_eq!( - diagnostics["failed-plugin"].status, - PluginInstanceRuntimeStatus::PreparationFailed - ); - assert!(diagnostics["failed-plugin"].failure.is_some()); - assert_eq!( - diagnostics["instance-one"].status, - PluginInstanceRuntimeStatus::Running - ); - let reused = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert_eq!( - generation.id(), - reused.id(), - "quarantined failures do not restart on every request" - ); - drop(reused); - drop(generation); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn a_published_process_crash_preserves_the_serving_snapshot_and_other_plugin_status() { - let (cache, store, runtime) = super::setup().await; - let snapshot = { - let mut snapshot = store.snapshot.lock().unwrap(); - let mut failed = snapshot.instances[0].clone(); - failed.id = "crashing-plugin".into(); - failed.configuration = serde_json::json!({"exit_after_ready_ms":100}); - snapshot.instances.push(failed); - snapshot.clone() - }; - let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - wait_until_unready(&generation).await; - assert!(generation.can_serve()); - assert!(!generation.is_ready()); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(1), - Some(&generation), - ) - .await - .unwrap(); - assert_eq!( - diagnostics["instance-one"].status, - PluginInstanceRuntimeStatus::Running - ); - assert_eq!( - diagnostics["crashing-plugin"].status, - PluginInstanceRuntimeStatus::Faulted - ); - drop(generation); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn a_host_version_mismatch_still_restores_the_enabled_plugin() { - use std::sync::Arc; - let (cache, store, _) = super::setup().await; - let runtime = gateway_plugin_runtime::PluginRuntime::new( - store.clone(), - store.clone(), - gateway_plugin_runtime::PluginRuntimeConfig { - cache_directory: cache.path().to_owned(), - host_version: "3.0.0".parse().unwrap(), - package_limits: Default::default(), - rpc_limits: Default::default(), - restart_circuit: Default::default(), - }, - Arc::new(gateway_host::outbound::HttpClient::new().unwrap()), - Arc::new(gateway_host::process::ProcessSupervisor::default()), - ); - let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert!(restored.can_serve()); - let snapshot = store.snapshot.lock().unwrap().clone(); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(1), - Some(&restored), - ) - .await - .unwrap(); - assert!(restored.is_ready()); - assert!(snapshot.instances[0].enabled); - assert_eq!( - diagnostics["instance-one"].status, - PluginInstanceRuntimeStatus::Running - ); - assert!(diagnostics["instance-one"].failure.is_none()); - drop(restored); - runtime.shutdown().await; - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn an_existing_failed_plugin_does_not_prevent_editing_an_unrelated_instance() { - let (cache, store, runtime) = super::setup().await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - let mut failed = snapshot.instances[0].clone(); - failed.id = "failed-plugin".into(); - failed.configuration = serde_json::json!({"startup":"fail"}); - snapshot.instances.push(failed); - snapshot.config_revision = Revision::new(2).unwrap(); - snapshot.instances[0].revision = snapshot.config_revision; - let candidate = PluginPreparation::prepare(&runtime, snapshot) - .await - .expect("only the instance being changed must successfully prepare"); - assert!(candidate.can_serve()); - drop(candidate); - super::wait_until_empty(cache.path()).await; -} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/runtime.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/runtime.rs new file mode 100644 index 000000000..ee3e633dd --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/runtime.rs @@ -0,0 +1,27 @@ +//! 验证运行时整体关闭回收当前代次并拒绝新候选 + +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; + +#[tokio::test] +async fn runtime_shutdown_reaps_live_generation_and_rejects_late_preparation() { + let (cache, store, runtime) = super::super::setup().await; + let published = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert!(published.is_ready()); + assert_eq!(std::fs::read_dir(cache.path()).unwrap().count(), 1); + + runtime.shutdown().await; + assert!(!published.is_ready()); + super::super::wait_until_empty(cache.path()).await; + assert!( + ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .is_err() + ); + runtime.shutdown().await; + + drop(published); + drop(store); +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/set.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/set.rs new file mode 100644 index 000000000..0fd200820 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/set.rs @@ -0,0 +1,34 @@ +//! 验证候选复用与最后一个代次持有者释放后的回收 + +use gateway_admin::model::Revision; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; + +#[tokio::test] +async fn candidate_is_reused_after_commit_and_released_when_no_generation_owner_remains() { + let (cache, store, runtime) = super::super::setup().await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = + serde_json::from_str(r#"{"z":1,"a":{"z":2,"a":3}}"#).unwrap(); + let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .unwrap(); + assert!( + runtime.observer_registry().resolve(&candidate).is_none(), + "没有观察绑定时不应编译派发计划" + ); + snapshot.instances[0].configuration = + serde_json::from_str(r#"{"a":{"a":3,"z":2},"z":1}"#).unwrap(); + snapshot.config_revision = Revision::new(2).unwrap(); + *store.snapshot.lock().unwrap() = snapshot; + let published = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) + .await + .unwrap(); + assert_eq!(candidate.id(), published.id()); + assert!(published.is_ready()); + drop(candidate); + assert_eq!(std::fs::read_dir(cache.path()).unwrap().count(), 1); + drop(published); + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/snapshot.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/snapshot.rs new file mode 100644 index 000000000..fe01dcb11 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/snapshot.rs @@ -0,0 +1,248 @@ +//! 验证完整快照准备、能力目录与失败实例恢复 + +use gateway_admin::model::Revision; +use gateway_admin::model::plugins::instances::PluginInstanceRuntimeStatus; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; +use gateway_plugin_sdk::Capability; +use gateway_plugin_sdk::Contributions; +use gateway_plugin_sdk::Stage; + +#[tokio::test] +async fn model_catalog_is_frozen_without_entering_the_request_chain() { + let (cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::ModelCatalog, + vec![Stage::Registration], + vec![], + vec![], + ), + ])) + .await; + store.snapshot.lock().unwrap().instances[0].configuration = serde_json::json!({ + "model_catalog": {"models": [{"id": "my-model", "provider": "openai", "model": "upstream-model"}]} + }); + let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert_eq!(generation.model_aliases().len(), 1); + assert_eq!(generation.model_aliases()[0].id.as_str(), "my-model"); + assert_eq!( + generation.model_aliases()[0].target.as_str(), + "upstream-model" + ); + assert!(runtime.policy_registry().resolve(&generation).is_none()); + assert!( + runtime + .execution_registry() + .middleware(&generation) + .is_none() + ); + let reused = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert_eq!(reused.id(), generation.id()); + assert_eq!(reused.model_aliases(), generation.model_aliases()); + { + let mut snapshot = store.snapshot.lock().unwrap(); + snapshot.config_revision = Revision::new(2).unwrap(); + snapshot.instances[0].revision = Revision::new(2).unwrap(); + snapshot.instances[0].configuration = serde_json::json!({"startup":"fail"}); + } + assert!( + ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) + .await + .is_err(), + "failed recovery must not publish a catalog with silently removed aliases" + ); + assert!(generation.is_ready()); + assert_eq!(generation.model_aliases()[0].id.as_str(), "my-model"); + { + let mut snapshot = store.snapshot.lock().unwrap(); + snapshot.config_revision = Revision::new(2).unwrap(); + snapshot.instances[0].enabled = false; + } + let disabled = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) + .await + .unwrap(); + assert!(disabled.model_aliases().is_empty()); + assert_eq!(generation.model_aliases().len(), 1); + drop(disabled); + drop(reused); + drop(generation); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn model_catalog_rejects_invalid_registrations_and_cross_instance_conflicts() { + for models in [ + serde_json::json!([{"id":"alias","provider":"custom","model":"target"}]), + serde_json::json!([{"id":"alias","provider":"openai","model":"alias"}]), + serde_json::json!([ + {"id":"alias","provider":"openai","model":"target"}, + {"id":"alias","provider":"xai","model":"target"} + ]), + serde_json::json!([{"id":"alias","provider":"openai","model":"target","capabilities":[]}]), + ] { + let (cache, store, runtime) = + super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::ModelCatalog, + vec![Stage::Registration], + vec![], + vec![], + ), + ])) + .await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = + serde_json::json!({"model_catalog":{"models":models}}); + assert!( + PluginPreparation::prepare(&runtime, snapshot) + .await + .is_err() + ); + runtime.shutdown().await; + super::super::wait_until_empty(cache.path()).await; + } + + let (cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::ModelCatalog, + vec![Stage::Registration], + vec![], + vec![], + ), + ])) + .await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "model_catalog":{"models":[{"id":"alias","provider":"openai","model":"target"}]} + }); + let mut other = snapshot.instances[0].clone(); + other.id = "instance-two".into(); + snapshot.instances.push(other); + let error = PluginPreparation::prepare(&runtime, snapshot) + .await + .unwrap_err(); + assert!(error.to_string().contains("instance-one")); + assert!(error.to_string().contains("instance-two")); + runtime.shutdown().await; + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn failed_candidate_leaves_the_previous_generation_ready() { + let (cache, store, runtime) = super::super::setup().await; + let snapshot = store.snapshot.lock().unwrap().clone(); + let active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .unwrap(); + let mut candidate = snapshot; + candidate.instances[0].configuration = serde_json::json!({"startup":"fail"}); + assert!( + PluginPreparation::prepare(&runtime, candidate) + .await + .is_err() + ); + assert!(active.is_ready()); + drop(active); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn disabled_missing_artifacts_do_not_block_the_remaining_generation_or_offline_recovery() { + let (cache, store, runtime) = super::super::setup().await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + let mut missing = snapshot.instances[0].clone(); + missing.id = "missing".into(); + missing.enabled = false; + missing.artifact_sha256 = "0".repeat(64); + snapshot.instances.push(missing); + let active = PluginPreparation::prepare(&runtime, snapshot) + .await + .unwrap(); + assert!(active.is_ready()); + drop(active); + super::super::wait_until_empty(cache.path()).await; + let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert!(restored.is_ready()); + drop(restored); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn preparation_cannot_mix_a_newer_persisted_revision_with_an_older_snapshot() { + let (_cache, _store, runtime) = super::super::setup().await; + assert!( + ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) + .await + .is_err() + ); +} + +#[tokio::test] +async fn restoring_a_failed_plugin_keeps_other_instances_ready_and_reports_only_its_failure() { + let (cache, store, runtime) = super::super::setup().await; + let snapshot = { + let mut snapshot = store.snapshot.lock().unwrap(); + let mut failed = snapshot.instances[0].clone(); + failed.id = "failed-plugin".into(); + failed.configuration = serde_json::json!({"startup":"fail"}); + snapshot.instances.insert(0, failed); + snapshot.clone() + }; + let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .expect("one failed plugin must not stop restoration"); + assert!(generation.can_serve()); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(1), + Some(&generation), + ) + .await + .unwrap(); + assert_eq!( + diagnostics["failed-plugin"].status, + PluginInstanceRuntimeStatus::PreparationFailed + ); + assert!(diagnostics["failed-plugin"].failure.is_some()); + assert_eq!( + diagnostics["instance-one"].status, + PluginInstanceRuntimeStatus::Running + ); + let reused = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert_eq!( + generation.id(), + reused.id(), + "quarantined failures do not restart on every request" + ); + drop(reused); + drop(generation); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn an_existing_failed_plugin_does_not_prevent_editing_an_unrelated_instance() { + let (cache, store, runtime) = super::super::setup().await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + let mut failed = snapshot.instances[0].clone(); + failed.id = "failed-plugin".into(); + failed.configuration = serde_json::json!({"startup":"fail"}); + snapshot.instances.push(failed); + snapshot.config_revision = Revision::new(2).unwrap(); + snapshot.instances[0].revision = snapshot.config_revision; + let candidate = PluginPreparation::prepare(&runtime, snapshot) + .await + .expect("only the instance being changed must successfully prepare"); + assert!(candidate.can_serve()); + drop(candidate); + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/private_state.rs b/backend/crates/gateway-plugin/runtime/tests/generation/private_state.rs index ccf90a50b..a5421e9e1 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/private_state.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/private_state.rs @@ -38,8 +38,8 @@ use gateway_core::{ }, identity::ProviderKind, operation::OperationKind, + routing::extensions::ExtensionPreparationPort, routing::{ConfigRevision, PublicModelId}, - runtime::extensions::ExtensionPreparationPort, upstream::UpstreamSendState, }; use gateway_plugin_runtime::{ diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs b/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs index d287a56aa..7589dd7a8 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs @@ -27,8 +27,8 @@ use gateway_core::{ metering::{CalculatedCost, Usage}, operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState}, policy::ClientApiKeyId, + routing::extensions::{ExtensionPreparationPort, ExtensionSetReference}, routing::{ConfigRevision, UpstreamModelId}, - runtime::extensions::{ExtensionPreparationPort, ExtensionSetReference}, upstream::{UpstreamSendState, UpstreamTransport}, }; use gateway_plugin_sdk::{Capability, Contributions, Stage}; @@ -136,11 +136,21 @@ async fn setup( (cache, store, runtime) } -fn context(generation: &ExtensionSetReference, key: &str) -> AttemptContext { - context_with_fast_mode(generation, key, gateway_core::account::FastMode::Default) +fn context( + runtime: &gateway_plugin_runtime::PluginRuntime, + generation: &ExtensionSetReference, + key: &str, +) -> AttemptContext { + context_with_fast_mode( + runtime, + generation, + key, + gateway_core::account::FastMode::Default, + ) } fn context_with_fast_mode( + runtime: &gateway_plugin_runtime::PluginRuntime, generation: &ExtensionSetReference, key: &str, mode: gateway_core::account::FastMode, @@ -150,7 +160,7 @@ fn context_with_fast_mode( ModelRequestId::new("req_one").unwrap(), ClientApiKeyId::new(key).unwrap(), ) - .with_upstream_adapters(generation.upstream_adapters()) + .with_upstream_adapters(runtime.execution_registry().upstream_adapters(generation)) .with_fast_mode(mode), NonZeroU32::new(1).unwrap(), SystemTime::now() + Duration::from_secs(5), @@ -166,12 +176,13 @@ fn context_with_fast_mode( } fn execute( + runtime: &gateway_plugin_runtime::PluginRuntime, generation: &ExtensionSetReference, account: Arc, key: &str, previous: Option, ) -> EventStream { - let context = context(generation, key); + let context = context(runtime, generation, key); execute_with_context(account, context, previous) } @@ -235,7 +246,7 @@ async fn legacy_adapter_decodes_exact_v1_metadata_for_all_fast_modes() { .unwrap(); let mut stream = execute_with_context( Account::new(), - context_with_fast_mode(&generation, "key-one", mode), + context_with_fast_mode(&runtime, &generation, "key-one", mode), None, ); let mut completed = 0; @@ -277,7 +288,7 @@ async fn adapter_http_and_sse_are_cold_authenticated_and_settled_once() { .await .unwrap(); let account = Account::new(); - let mut stream = execute(&generation, Arc::clone(&account), "key-one", None); + let mut stream = execute(&runtime, &generation, Arc::clone(&account), "key-one", None); assert!( server.received_requests().await.unwrap().is_empty(), "构造冷流不能先发送业务请求" @@ -321,7 +332,7 @@ async fn adapter_accepts_service_tier_resolved_by_the_terminal_response() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let mut tiers = Vec::new(); let mut completed = false; while let Some(event) = stream.next().await { @@ -368,7 +379,7 @@ async fn adapter_failure_preserves_upstream_diagnostics_and_native_feedback() { .await .unwrap(); let account = Account::new(); - let mut stream = execute(&generation, Arc::clone(&account), "key-one", None); + let mut stream = execute(&runtime, &generation, Arc::clone(&account), "key-one", None); let error = stream.next().await.unwrap().unwrap_err(); assert_eq!(error.kind(), ProviderErrorKind::RateLimited); assert_eq!(error.send_state(), UpstreamSendState::Sent); @@ -409,7 +420,7 @@ async fn adapter_overrides_auth_headers_and_rejects_stale_continuation() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let mut state = None; while let Some(event) = stream.next().await { state = event.unwrap().take_session_update().or(state); @@ -420,7 +431,7 @@ async fn adapter_overrides_auth_headers_and_rejects_stale_continuation() { ("key-one", Account::with_identity("acct_other", 1)), ("key-one", Account::with_identity("acct_one", 2)), ] { - let error = execute(&generation, account, key, state.clone()) + let error = execute(&runtime, &generation, account, key, state.clone()) .next() .await .unwrap() @@ -438,7 +449,7 @@ async fn adapter_overrides_auth_headers_and_rejects_stale_continuation() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); while let Some(event) = stream.next().await { event.unwrap(); } @@ -467,7 +478,7 @@ async fn adapter_does_not_publish_completed_when_rpc_fails_after_terminal_frame( let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); stream.next().await.unwrap().unwrap(); let error = stream.next().await.unwrap().err().unwrap(); assert_eq!(error.kind(), ProviderErrorKind::Unavailable); @@ -576,7 +587,7 @@ async fn websocket_continuation_reuses_exact_connection_and_consumes_each_handle let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut first = execute(&generation, Account::new(), "key-one", None); + let mut first = execute(&runtime, &generation, Account::new(), "key-one", None); let mut state = None; while let Some(event) = first.next().await { state = event.unwrap().take_session_update().or(state); @@ -584,24 +595,42 @@ async fn websocket_continuation_reuses_exact_connection_and_consumes_each_handle drop(first); let state = state.unwrap(); assert!(state.extension_owner().unwrap().connection_local); - let error = execute(&generation, Account::new(), "key-two", Some(state.clone())) - .next() - .await - .unwrap() - .err() - .unwrap(); + let error = execute( + &runtime, + &generation, + Account::new(), + "key-two", + Some(state.clone()), + ) + .next() + .await + .unwrap() + .err() + .unwrap(); assert_eq!(error.send_state(), UpstreamSendState::NotSent); - let mut second = execute(&generation, Account::new(), "key-one", Some(state.clone())); + let mut second = execute( + &runtime, + &generation, + Account::new(), + "key-one", + Some(state.clone()), + ); while let Some(event) = second.next().await { event.unwrap(); } drop(second); - let error = execute(&generation, Account::new(), "key-one", Some(state)) - .next() - .await - .unwrap() - .err() - .unwrap(); + let error = execute( + &runtime, + &generation, + Account::new(), + "key-one", + Some(state), + ) + .next() + .await + .unwrap() + .err() + .unwrap(); assert_eq!(error.send_state(), UpstreamSendState::NotSent); drop(generation); runtime.shutdown().await; @@ -642,7 +671,7 @@ async fn websocket_adapter_can_send_while_a_read_callback_is_pending() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let mut continuation = None; tokio::time::timeout(Duration::from_secs(3), async { while let Some(event) = stream.next().await { @@ -696,7 +725,7 @@ async fn cancelling_adapter_body_closes_pending_websocket_and_releases_selected_ .await .unwrap(); let account = Account::new(); - let mut stream = execute(&generation, account.clone(), "key-one", None); + let mut stream = execute(&runtime, &generation, account.clone(), "key-one", None); let mut reading = Box::pin(stream.next()); tokio::select! { result = &mut reading => panic!("上游等待期间不应提前产生结果：{}", result.is_some()), @@ -734,7 +763,7 @@ async fn usage_without_response_model_still_uses_native_account_pricing() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let mut costs = 0; while let Some(event) = stream.next().await { costs += event @@ -834,7 +863,7 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a .await .unwrap(); let account = Account::new(); - let mut old_stream = execute(&old, account.clone(), "key-one", None); + let mut old_stream = execute(&runtime, &old, account.clone(), "key-one", None); old_stream.next().await.unwrap().unwrap(); { let mut snapshot = store.snapshot.lock().unwrap(); @@ -846,7 +875,7 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a .await .unwrap(); assert_ne!(old.id(), new.id()); - let mut new_stream = execute(&new, account.clone(), "key-one", None); + let mut new_stream = execute(&runtime, &new, account.clone(), "key-one", None); new_stream.next().await.unwrap().unwrap(); drop(old); let old_state = old_stream @@ -868,7 +897,7 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a .await .unwrap(); assert!( - context(&disabled, "key-one") + context(&runtime, &disabled, "key-one") .upstream_adapter( &ProviderKind::new("openai").unwrap(), &UpstreamModelId::new("native-model").unwrap() @@ -903,13 +932,19 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a .await .unwrap(); // 恢复同一配置仍是新代次，不能复活旧进程中的续接身份 - let error = execute(&restored, account.clone(), "key-one", Some(old_state)) - .next() - .await - .unwrap() - .unwrap_err(); + let error = execute( + &runtime, + &restored, + account.clone(), + "key-one", + Some(old_state), + ) + .next() + .await + .unwrap() + .unwrap_err(); assert_eq!(error.send_state(), UpstreamSendState::NotSent); - let mut fresh = execute(&restored, account.clone(), "key-one", None); + let mut fresh = execute(&runtime, &restored, account.clone(), "key-one", None); while let Some(event) = fresh.next().await { event.unwrap(); } @@ -923,15 +958,15 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a #[tokio::test] async fn upstream_adapter_can_dispatch_http_after_its_instance_entered_the_scope() { struct Http(Arc); - impl gateway_core::middleware::http::Dispatcher for Http { + impl gateway_core::engine::middleware::http::Dispatcher for Http { fn dispatch( &self, - context: gateway_core::middleware::http::Context, - request: gateway_core::middleware::http::Request, + context: gateway_core::engine::middleware::http::Context, + request: gateway_core::engine::middleware::http::Request, ) -> BoxFuture< 'static, Result< - gateway_core::middleware::http::Response, + gateway_core::engine::middleware::http::Response, gateway_core::engine::middleware::MiddlewareError, >, > { @@ -944,8 +979,8 @@ async fn upstream_adapter_can_dispatch_http_after_its_instance_entered_the_scope assert_eq!(request.uri(), "/api/admin/settings"); self.0.fetch_add(1, Ordering::SeqCst); Box::pin(async { - Ok(gateway_core::middleware::http::Response::new( - gateway_core::middleware::http::empty_body(), + Ok(gateway_core::engine::middleware::http::Response::new( + gateway_core::engine::middleware::http::empty_body(), )) }) } @@ -961,13 +996,13 @@ async fn upstream_adapter_can_dispatch_http_after_its_instance_entered_the_scope })); let (_cache, _, runtime) = setup(config).await; let calls = Arc::new(AtomicUsize::new(0)); - let dispatcher: Arc = + let dispatcher: Arc = Arc::new(Http(calls.clone())); runtime.bind_http(&dispatcher).unwrap(); let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); while let Some(event) = stream.next().await { event.unwrap(); } @@ -1004,7 +1039,7 @@ async fn malformed_adapter_events_stop_only_the_plugin_session() { ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let error = stream.next().await.unwrap().unwrap_err(); assert_eq!(error.kind(), ProviderErrorKind::Protocol); if has_sequence_error { diff --git a/backend/crates/gateway-plugin/runtime/tests/support/environment.rs b/backend/crates/gateway-plugin/runtime/tests/support/environment.rs index 5b6ce314b..c61b1b789 100644 --- a/backend/crates/gateway-plugin/runtime/tests/support/environment.rs +++ b/backend/crates/gateway-plugin/runtime/tests/support/environment.rs @@ -439,7 +439,7 @@ impl Environment { Some(runtime.clone()), Some(runtime.observer_registry()), Some(runtime.policy_registry()), - Some(runtime.middleware_registry()), + Some(runtime.execution_registry()), Some(runtime.frontend_authentication_registry()), ); let access = gateway_admin::initialize_plugin_accounts( @@ -472,7 +472,7 @@ impl Environment { Some(runtime.clone()), Some(runtime.observer_registry()), Some(runtime.policy_registry()), - Some(runtime.middleware_registry()), + Some(runtime.execution_registry()), ) .activate() .await @@ -536,10 +536,10 @@ impl Environment { timezone: Default::default(), service_middleware: { let snapshots = core.snapshots(); - let middleware = runtime.middleware_registry(); + let middleware = runtime.execution_registry(); std::sync::Arc::new(move || { let snapshot = snapshots.snapshot_for_diagnostics()?; - middleware.resolve(snapshot.extensions()?) + middleware.middleware(snapshot.extensions()?) }) }, plugin_preparation: runtime.clone(), diff --git a/backend/crates/gateway-plugin/sdk/README.md b/backend/crates/gateway-plugin/sdk/README.md index 485af3fcb..160964c82 100644 --- a/backend/crates/gateway-plugin/sdk/README.md +++ b/backend/crates/gateway-plugin/sdk/README.md @@ -101,7 +101,7 @@ sequenceDiagram ## 本地验证 -设置服务的 SDK 类型和操作由宿主领域类型及[操作声明](../../gateway-admin/src/service/settings.rs)生成，SDK 构建不依赖宿主源码。 +设置服务的 SDK 类型和操作由宿主领域类型及[操作声明](../../gateway-admin/src/public_service/settings.rs)生成，SDK 构建不依赖宿主源码。 修改宿主合同后，在仓库根目录更新生成文件；Admin 合同测试会检查生成结果是否同步： ```bash diff --git a/backend/crates/gateway-plugin/sdk/src/call/services/settings.rs b/backend/crates/gateway-plugin/sdk/src/call/services/settings.rs index 42f123220..bab9cdc92 100644 --- a/backend/crates/gateway-plugin/sdk/src/call/services/settings.rs +++ b/backend/crates/gateway-plugin/sdk/src/call/services/settings.rs @@ -1,6 +1,6 @@ //! 宿主设置服务的操作标识与请求、响应数据合同 //! -//! 从宿主设置类型与 service/settings.rs 生成；更新命令见 SDK 维护说明 +//! 从宿主设置类型与 public_service/settings.rs 生成；更新命令见 SDK 维护说明 use super::Operation; use serde::{Deserialize, Serialize}; use std::{ diff --git a/backend/crates/gateway-plugin/sdk/src/client/session.rs b/backend/crates/gateway-plugin/sdk/src/client/session.rs deleted file mode 100644 index 6f7e22b44..000000000 --- a/backend/crates/gateway-plugin/sdk/src/client/session.rs +++ /dev/null @@ -1,1479 +0,0 @@ -//! 插件侧双向 RPC 会话的握手、调用分派、流控与取消生命周期 - -use std::{ - collections::{BTreeMap, BTreeSet, VecDeque}, - future::Future, - num::NonZeroUsize, - pin::Pin, - sync::{Arc, Mutex}, - time::Duration, -}; - -use serde_json::Value; -use tokio::{ - io::{AsyncRead, AsyncWrite}, - sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot, watch}, - task::JoinHandle, - time::Instant, -}; - -use crate::{ - CallContext, ErrorCode, Frame, FrameError, Handshake, Message, PROTOCOL_VERSION, PluginFault, -}; - -use super::frame::{read_frame, validate_frame, write_frame}; - -const MAXIMUM_STREAM_CHUNK_BYTES: usize = 16 * 1024 * 1024; -const MAXIMUM_CONCURRENCY: usize = 256; -const MAXIMUM_BUFFERED_STREAM_CHUNKS: usize = 65_536; - -/// 插件侧会话的内存、并发与期限边界 -#[derive(Debug, Clone, Copy)] -pub struct SessionConfig { - /// 单个业务流分块的预算，不限制普通调用正文的总长度 - pub maximum_stream_chunk_bytes: usize, - pub maximum_calls: usize, - pub maximum_callbacks: usize, - pub maximum_buffered_stream_chunks: usize, - pub handshake_timeout: Duration, - pub maximum_call_timeout: Duration, -} - -impl Default for SessionConfig { - fn default() -> Self { - Self { - // 流分块还需匹配宿主授予的信用窗口 - maximum_stream_chunk_bytes: 1024 * 1024, - maximum_calls: 32, - maximum_callbacks: 32, - maximum_buffered_stream_chunks: 1_024, - handshake_timeout: Duration::from_secs(5), - maximum_call_timeout: Duration::from_secs(120), - } - } -} - -impl SessionConfig { - fn validate(self) -> Result { - if self.maximum_stream_chunk_bytes < 1_024 - || self.maximum_stream_chunk_bytes > MAXIMUM_STREAM_CHUNK_BYTES - || self.maximum_calls == 0 - || self.maximum_calls > MAXIMUM_CONCURRENCY - || self.maximum_callbacks == 0 - || self.maximum_callbacks > MAXIMUM_CONCURRENCY - || self.maximum_buffered_stream_chunks == 0 - || self.maximum_buffered_stream_chunks > MAXIMUM_BUFFERED_STREAM_CHUNKS - || self.handshake_timeout.is_zero() - || self.maximum_call_timeout.is_zero() - { - return Err(SessionError::Configuration); - } - Ok(self) - } -} - -/// 插件侧传输和会话错误 -#[derive(Debug, thiserror::Error)] -pub enum SessionError { - #[error("plugin session configuration is invalid")] - Configuration, - #[error("plugin handshake is invalid")] - Handshake, - #[error("plugin protocol is invalid")] - Protocol, - #[error("plugin transport is closed")] - Closed, - #[error("plugin operation exceeded its deadline")] - Timeout, - #[error("plugin call was cancelled")] - Cancelled, - #[error("plugin session capacity is exhausted")] - Capacity, - #[error("plugin handler stopped unexpectedly")] - HandlerStopped, - #[error("host callback returned an error: {0:?}")] - Remote(PluginFault), - #[error(transparent)] - Frame(#[from] FrameError), -} - -impl SessionError { - /// 将本地会话失败收敛为可返回宿主的稳定插件错误 - #[must_use] - pub fn into_plugin_fault(self) -> PluginFault { - match self { - Self::Remote(fault) => fault, - Self::Timeout => PluginFault::new(ErrorCode::Timeout, "host callback timed out"), - Self::Cancelled => PluginFault::new(ErrorCode::Cancelled, "parent call was cancelled"), - Self::Capacity => PluginFault::new(ErrorCode::Capacity, "plugin capacity is exhausted"), - Self::Configuration | Self::Protocol | Self::Handshake | Self::Frame(_) => { - PluginFault::new(ErrorCode::InvalidInput, "plugin session input is invalid") - } - Self::Closed | Self::HandlerStopped => { - PluginFault::new(ErrorCode::Fault, "plugin session is unavailable") - } - } - } -} - -/// 一次宿主调用的可克隆取消信号 -#[derive(Clone)] -pub struct CallCancellation { - receiver: watch::Receiver, -} - -impl CallCancellation { - #[must_use] - pub fn is_cancelled(&self) -> bool { - *self.receiver.borrow() - } - - pub async fn cancelled(&self) { - let mut receiver = self.receiver.clone(); - if *receiver.borrow() { - return; - } - while receiver.changed().await.is_ok() { - if *receiver.borrow() { - return; - } - } - } -} - -struct CancellationSource { - sender: watch::Sender, -} - -impl CancellationSource { - fn new() -> (Self, CallCancellation) { - let (sender, receiver) = watch::channel(false); - (Self { sender }, CallCancellation { receiver }) - } - - fn cancel(&self) { - self.sender.send_replace(true); - } -} - -/// 宿主回调的成功结果；不实现 `Debug`，避免载荷进入普通诊断 -pub struct HostReply { - pub result: Value, - pub payload: Vec, -} - -/// 绑定父调用、期限和取消信号的宿主回调客户端 -#[derive(Clone)] -pub struct HostClient { - parent_id: u64, - deadline: tokio::sync::watch::Receiver>, - cancellation: CallCancellation, - callbacks: Arc, - output: Output, - maximum_stream_chunk_bytes: usize, - credits: Arc, -} - -impl HostClient { - pub(crate) fn maximum_stream_chunk_bytes(&self) -> usize { - let window = lock_unpoisoned(&self.credits.state).maximum_bytes; - if window == 0 { - self.maximum_stream_chunk_bytes - } else { - self.maximum_stream_chunk_bytes - .min(usize::try_from(window).unwrap_or(usize::MAX)) - } - } - - /// 发起一次与父调用关联的宿主回调 - /// - /// # Errors - /// - /// 方法或帧无效、容量耗尽、父调用取消、期限到达、连接关闭或宿主返回错误时失败 - pub async fn call( - &self, - method: impl Into, - params: Value, - payload: Vec, - ) -> Result { - let method = method.into(); - if method.is_empty() || method.len() > 128 { - return Err(SessionError::Protocol); - } - let ordered = self.callbacks.order.lock().await; - let (id, received) = self.callbacks.begin(self.parent_id)?; - let frame = Frame { - message: Message::Callback { - id, - parent_id: self.parent_id, - method, - params, - }, - payload, - }; - if validate_frame(&frame).is_err() { - self.callbacks.retire(id); - return Err(SessionError::Protocol); - } - let sent = tokio::select! { - biased; - () = self.cancellation.cancelled() => Err(SessionError::Cancelled), - () = callback_deadline(self.deadline.clone()) => Err(SessionError::Timeout), - result = self.output.send_data(frame) => result, - }; - if let Err(error) = sent { - self.callbacks.retire(id); - return Err(error); - } - drop(ordered); - tokio::select! { - biased; - () = self.cancellation.cancelled() => { - self.callbacks.retire(id); - Err(SessionError::Cancelled) - } - () = callback_deadline(self.deadline.clone()) => { - self.callbacks.retire(id); - Err(SessionError::Timeout) - } - result = received => result.map_err(|_| SessionError::Closed)?, - } - } -} - -/// 交给插件业务处理器的一次宿主调用；不实现 `Debug`，避免敏感载荷泄露 -pub struct PluginCall { - pub method: String, - pub context: CallContext, - pub params: Value, - pub payload: Vec, - pub host: HostClient, - pub cancellation: CallCancellation, -} - -/// 插件业务调用 future -pub type CallFuture<'a> = Pin> + Send + 'a>>; - -/// 插件只实现业务调用；会话负责握手、关联、流控、取消和关闭 -pub trait PluginHandler: Send + Sync + 'static { - fn call(&self, call: PluginCall) -> CallFuture<'_>; - - /// 取消时释放只属于该调用的临时业务状态 - /// 实现必须幂等、快速且不得执行阻塞 I/O - fn cancel(&self, _context: &CallContext) {} - - /// 会话停止接收新调用时触发 - /// 实现必须幂等、快速且不得执行阻塞 I/O - fn quiesce(&self) {} - - /// 会话关闭时释放插件自有临时状态 - /// 实现必须幂等、快速且不得执行阻塞 I/O - fn shutdown(&self) {} -} - -/// 一次业务成功结果；`stream` 存在时 SDK 自动产生单一流终态 -pub struct CallReply { - result: Value, - payload: Vec, - stream: Option, -} - -impl CallReply { - #[must_use] - pub fn unary(result: Value, payload: Vec) -> Self { - Self { - result, - payload, - stream: None, - } - } - - #[must_use] - pub fn stream(result: Value, payload: Vec, stream: ResponseStream) -> Self { - Self { - result, - payload, - stream: Some(stream), - } - } -} - -enum StreamSource { - Buffered(VecDeque>), - Channel(mpsc::Receiver, PluginFault>>), - Pull(Box), -} - -pub type PullResponseFuture<'a> = - Pin, PluginFault>>> + Send + 'a>>; - -pub trait PullResponseStream: Send { - fn next(&mut self) -> PullResponseFuture<'_>; -} - -/// SDK 管理 sequence、Credit 和终态的响应流 -pub struct ResponseStream { - source: StreamSource, - declared_capacity: usize, -} - -impl ResponseStream { - /// 从已完成业务校验的有限分块构造流；SDK 会在发送 `Result` 前预检全部分块 - #[must_use] - pub fn from_chunks(chunks: Vec>) -> Self { - let declared_capacity = chunks.len().max(1); - Self { - source: StreamSource::Buffered(chunks.into()), - declared_capacity, - } - } - - /// 创建由生产者驱动的有界流队列 - #[must_use] - pub fn channel(capacity: NonZeroUsize) -> (StreamSender, Self) { - let (sender, receiver) = mpsc::channel(capacity.get()); - ( - StreamSender { sender }, - Self { - source: StreamSource::Channel(receiver), - declared_capacity: capacity.get(), - }, - ) - } - - /// 按消费进度拉取下一帧；返回 None 表示完成，丢弃流会丢弃生产者 - pub fn pull(source: Box) -> Self { - Self { - source: StreamSource::Pull(source), - declared_capacity: 1, - } - } - - fn validate_buffered( - &self, - maximum_chunks: usize, - maximum_stream_chunk_bytes: usize, - window_bytes: u64, - ) -> Result<(), PluginFault> { - if self.declared_capacity > maximum_chunks { - return Err(capacity_fault( - "response stream exceeds the local queue limit", - )); - } - if let StreamSource::Buffered(chunks) = &self.source { - for chunk in chunks { - validate_stream_chunk(chunk, maximum_stream_chunk_bytes, window_bytes)?; - } - } - Ok(()) - } - - pub(crate) async fn next(&mut self) -> Option, PluginFault>> { - match &mut self.source { - StreamSource::Buffered(chunks) => chunks.pop_front().map(Ok), - StreamSource::Channel(receiver) => receiver.recv().await, - StreamSource::Pull(source) => source.next().await, - } - } -} - -/// 动态响应流的有界生产端；最后一个 sender 被释放表示成功终态 -#[derive(Clone)] -pub struct StreamSender { - sender: mpsc::Sender, PluginFault>>, -} - -impl StreamSender { - /// 排队一个业务分块；实际线协议流控由会话完成 - pub async fn send(&self, payload: Vec) -> Result<(), SessionError> { - self.sender - .send(Ok(payload)) - .await - .map_err(|_| SessionError::Closed) - } - - /// 以业务错误结束动态流 - pub async fn fail(&self, fault: PluginFault) -> Result<(), SessionError> { - self.sender - .send(Err(fault)) - .await - .map_err(|_| SessionError::Closed) - } -} - -/// 已读取并校验 Hello、尚未向宿主发送 Ready 的插件会话 -pub struct PluginSession { - reader: R, - writer: W, - handshake: Handshake, - config: SessionConfig, -} - -impl PluginSession -where - R: AsyncRead + Unpin + Send + 'static, - W: AsyncWrite + Unpin + Send + 'static, -{ - /// 读取并校验宿主握手 - /// 业务处理器可以在 `run` 前读取不可变 Handshake 构造注册描述 - /// - /// # Errors - /// - /// 配置无效、握手超时、传输关闭或 Hello 不符合协议时失败 - pub async fn accept( - mut reader: R, - writer: W, - config: SessionConfig, - ) -> Result { - let config = config.validate()?; - let frame = tokio::time::timeout(config.handshake_timeout, read_frame(&mut reader)) - .await - .map_err(|_| SessionError::Timeout)??; - let Message::Hello { handshake } = frame.message else { - return Err(SessionError::Handshake); - }; - if !frame.payload.is_empty() - || handshake.protocol_version != PROTOCOL_VERSION - || handshake.plugin_id.is_empty() - || handshake.instance_id.is_empty() - || handshake.generation == 0 - || handshake.incarnation.is_empty() - { - return Err(SessionError::Handshake); - } - Ok(Self { - reader, - writer, - handshake, - config, - }) - } - - #[must_use] - pub const fn handshake(&self) -> &Handshake { - &self.handshake - } - - /// 发送 Ready 并持续处理调用、回调、流控、取消与关闭 - /// - /// # Errors - /// - /// 帧损坏、协议违规、处理器异常停止或传输关闭时失败 - pub async fn run(mut self, handler: H) -> Result<(), SessionError> - where - H: PluginHandler, - { - let data_capacity = self - .config - .maximum_calls - .checked_add(self.config.maximum_callbacks) - .ok_or(SessionError::Configuration)?; - let control_reserve = self - .config - .maximum_calls - .checked_add(8) - .ok_or(SessionError::Configuration)?; - let queue_capacity = data_capacity - .checked_add(control_reserve) - .ok_or(SessionError::Configuration)?; - let (outbound, received) = mpsc::channel(queue_capacity); - let output = Output { - sender: outbound, - data_slots: Arc::new(Semaphore::new(data_capacity)), - }; - let callbacks = Arc::new(CallbackRegistry::new(self.config.maximum_callbacks)); - let handler: Arc = Arc::new(handler); - let lifecycle = LifecycleHooks::new(Arc::clone(&handler), control_reserve)?; - write_frame( - &mut self.writer, - &Frame::control(Message::Ready { - protocol_version: PROTOCOL_VERSION, - incarnation: self.handshake.incarnation.clone(), - }), - ) - .await?; - let writer = tokio::spawn(writer_loop(self.writer, received)); - let (completed, completions) = mpsc::channel(self.config.maximum_calls); - let mut driver = SessionDriver { - config: self.config, - instance_id: self.handshake.instance_id, - generation: self.handshake.generation, - incarnation: self.handshake.incarnation, - output, - callbacks, - handler, - lifecycle, - writer, - completed, - completions, - active: BTreeMap::new(), - last_call: 0, - quiescing: false, - }; - let result = driver.drive(self.reader).await; - driver.close().await; - result - } -} - -struct Output { - sender: mpsc::Sender, - data_slots: Arc, -} - -impl Clone for Output { - fn clone(&self) -> Self { - Self { - sender: self.sender.clone(), - data_slots: Arc::clone(&self.data_slots), - } - } -} - -impl Output { - async fn send_data(&self, frame: Frame) -> Result<(), SessionError> { - validate_frame(&frame)?; - let permit = Arc::clone(&self.data_slots) - .acquire_owned() - .await - .map_err(|_| SessionError::Closed)?; - self.sender - .send(Outbound { - frame, - _data_slot: Some(permit), - }) - .await - .map_err(|_| SessionError::Closed) - } - - async fn send_data_until( - &self, - frame: Frame, - deadline: Option, - cancellation: &CallCancellation, - ) -> Result<(), SessionError> { - tokio::select! { - biased; - () = cancellation.cancelled() => Err(SessionError::Cancelled), - () = wait_deadline(deadline) => Err(SessionError::Timeout), - result = self.send_data(frame) => result, - } - } - - fn send_control(&self, frame: Frame) -> Result<(), SessionError> { - validate_frame(&frame)?; - self.sender - .try_send(Outbound { - frame, - _data_slot: None, - }) - .map_err(|_| SessionError::Capacity) - } -} - -struct Outbound { - frame: Frame, - _data_slot: Option, -} - -async fn writer_loop( - mut writer: W, - mut received: mpsc::Receiver, -) -> Result<(), SessionError> { - while let Some(outbound) = received.recv().await { - write_frame(&mut writer, &outbound.frame).await?; - } - Ok(()) -} - -struct CallbackRegistry { - state: Mutex, - order: tokio::sync::Mutex<()>, - maximum: usize, - maximum_retired: usize, -} - -struct CallbackState { - next_id: u64, - pending: BTreeMap, - retired: BTreeSet, - retired_order: VecDeque, -} - -struct PendingCallback { - parent_id: u64, - result: oneshot::Sender>, -} - -impl CallbackRegistry { - fn new(maximum: usize) -> Self { - Self { - state: Mutex::new(CallbackState { - next_id: 2, - pending: BTreeMap::new(), - retired: BTreeSet::new(), - retired_order: VecDeque::new(), - }), - order: tokio::sync::Mutex::new(()), - maximum, - maximum_retired: maximum.saturating_mul(2), - } - } - - fn begin( - &self, - parent_id: u64, - ) -> Result<(u64, oneshot::Receiver>), SessionError> { - let mut state = lock_unpoisoned(&self.state); - if state.pending.len() >= self.maximum { - return Err(SessionError::Capacity); - } - let id = state.next_id; - state.next_id = state.next_id.checked_add(2).ok_or(SessionError::Capacity)?; - let (result, received) = oneshot::channel(); - state - .pending - .insert(id, PendingCallback { parent_id, result }); - Ok((id, received)) - } - - fn resolve(&self, frame: Frame) -> Result<(), SessionError> { - let id = match &frame.message { - Message::Result { id, .. } | Message::Error { id, .. } => *id, - _ => return Err(SessionError::Protocol), - }; - let pending = { - let mut state = lock_unpoisoned(&self.state); - if let Some(pending) = state.pending.remove(&id) { - Some(pending) - } else if state.retired.remove(&id) { - state.retired_order.retain(|retired| *retired != id); - None - } else { - return Err(SessionError::Protocol); - } - }; - let Some(pending) = pending else { - return Ok(()); - }; - let result = match frame.message { - Message::Result { result, .. } => Ok(HostReply { - result, - payload: frame.payload, - }), - Message::Error { error, .. } if frame.payload.is_empty() => { - Err(SessionError::Remote(error)) - } - _ => return Err(SessionError::Protocol), - }; - let _ = pending.result.send(result); - Ok(()) - } - - fn retire(&self, id: u64) { - let mut state = lock_unpoisoned(&self.state); - if state.pending.remove(&id).is_some() { - retire_callback(&mut state, id, self.maximum_retired); - } - } - - fn finish_parent(&self, parent_id: u64) { - let pending = { - let mut state = lock_unpoisoned(&self.state); - let mut remaining = BTreeMap::new(); - let mut finished = Vec::new(); - for (id, pending) in std::mem::take(&mut state.pending) { - if pending.parent_id == parent_id { - retire_callback(&mut state, id, self.maximum_retired); - finished.push(pending); - } else { - remaining.insert(id, pending); - } - } - state.pending = remaining; - finished - }; - for pending in pending { - let _ = pending.result.send(Err(SessionError::Cancelled)); - } - } - - fn close(&self) { - let pending = { - let mut state = lock_unpoisoned(&self.state); - std::mem::take(&mut state.pending) - }; - for pending in pending.into_values() { - let _ = pending.result.send(Err(SessionError::Closed)); - } - } -} - -fn retire_callback(state: &mut CallbackState, id: u64, maximum: usize) { - if state.retired.insert(id) { - state.retired_order.push_back(id); - } - while state.retired_order.len() > maximum { - if let Some(expired) = state.retired_order.pop_front() { - state.retired.remove(&expired); - } - } -} - -enum LifecycleEvent { - Cancel(CallContext), - Quiesce, - Shutdown, -} - -#[derive(Clone)] -struct LifecycleHooks { - events: std::sync::mpsc::SyncSender, -} - -impl LifecycleHooks { - fn new(handler: Arc, capacity: usize) -> Result { - let (events, received) = std::sync::mpsc::sync_channel(capacity); - std::thread::Builder::new() - .name("gateway-plugin-sdk-lifecycle".into()) - .spawn(move || { - while let Ok(event) = received.recv() { - match event { - LifecycleEvent::Cancel(context) => handler.cancel(&context), - LifecycleEvent::Quiesce => handler.quiesce(), - LifecycleEvent::Shutdown => { - handler.shutdown(); - break; - } - } - } - }) - .map_err(|_| SessionError::Closed)?; - Ok(Self { events }) - } - - fn cancel(&self, context: CallContext) { - let _ = self.events.try_send(LifecycleEvent::Cancel(context)); - } - - fn quiesce(&self) { - let _ = self.events.try_send(LifecycleEvent::Quiesce); - } - - fn shutdown(&self) { - let _ = self.events.try_send(LifecycleEvent::Shutdown); - } -} - -struct ActiveCall { - context: CallContext, - credits: Arc, - cancellation: CancellationSource, - task: JoinHandle<()>, -} - -#[derive(Clone, Copy)] -enum CallTaskOutcome { - Complete, - TransportFailed, - HandlerStopped, -} - -struct CallFinished { - id: u64, - outcome: CallTaskOutcome, -} - -struct CompletionGuard { - id: u64, - completed: mpsc::Sender, - armed: bool, -} - -impl CompletionGuard { - fn new(id: u64, completed: mpsc::Sender) -> Self { - Self { - id, - completed, - armed: true, - } - } - - fn finish(mut self, outcome: CallTaskOutcome) { - self.armed = false; - let _ = self.completed.try_send(CallFinished { - id: self.id, - outcome, - }); - } -} - -impl Drop for CompletionGuard { - fn drop(&mut self) { - if self.armed { - let _ = self.completed.try_send(CallFinished { - id: self.id, - outcome: CallTaskOutcome::HandlerStopped, - }); - } - } -} - -struct SessionDriver { - config: SessionConfig, - instance_id: String, - generation: u64, - incarnation: String, - output: Output, - callbacks: Arc, - handler: Arc, - lifecycle: LifecycleHooks, - writer: JoinHandle>, - completed: mpsc::Sender, - completions: mpsc::Receiver, - active: BTreeMap, - last_call: u64, - quiescing: bool, -} - -impl SessionDriver { - async fn drive(&mut self, mut reader: R) -> Result<(), SessionError> { - loop { - // read_exact 不是取消安全的；完成通知可以穿插处理，但必须保留同一个半帧读取 future - let incoming = read_frame(&mut reader); - tokio::pin!(incoming); - let frame = loop { - tokio::select! { - biased; - writer = &mut self.writer => { - return writer.map_err(|_| SessionError::Closed)?; - } - Some(finished) = self.completions.recv() => { - self.finish_call(finished)?; - } - frame = &mut incoming => { - break frame.map_err(|error| match error { - FrameError::Io(io_error) - if io_error.kind() == std::io::ErrorKind::UnexpectedEof => - { - SessionError::Closed - } - error => SessionError::Frame(error), - })?; - } - } - }; - if self.handle_frame(frame)? { - return Ok(()); - } - } - } - - fn handle_frame(&mut self, frame: Frame) -> Result { - match frame.message { - Message::Call { - id, - method, - context, - params, - } => { - self.start_call(id, method, context, params, frame.payload)?; - Ok(false) - } - Message::Credit { id, bytes, frames } if frame.payload.is_empty() => { - if id == 0 - || id.is_multiple_of(2) - || id > self.last_call - || bytes == 0 - || frames == 0 - { - return Err(SessionError::Protocol); - } - // End/Cancelled 已进入同一 FIFO 后，宿主先前为已消费分块排队的 - // Credit 仍可能稍晚抵达 - // 此时调用资源已经回收，忽略迟到信用即可； - // 未来或回调 ID 仍按协议错误处理 - if let Some(active) = self.active.get(&id) { - active.credits.grant(bytes, frames)?; - } - Ok(false) - } - Message::Cancel { id } if frame.payload.is_empty() => { - self.cancel_call(id)?; - Ok(false) - } - Message::Result { id, result } => { - self.callbacks.resolve(Frame { - message: Message::Result { id, result }, - payload: frame.payload, - })?; - Ok(false) - } - Message::Error { id, error } => { - self.callbacks.resolve(Frame { - message: Message::Error { id, error }, - payload: frame.payload, - })?; - Ok(false) - } - Message::Quiesce if frame.payload.is_empty() => { - if !self.quiescing { - self.quiescing = true; - self.lifecycle.quiesce(); - } - Ok(false) - } - Message::Shutdown if frame.payload.is_empty() => Ok(true), - _ => Err(SessionError::Protocol), - } - } - - fn start_call( - &mut self, - id: u64, - method: String, - context: CallContext, - params: Value, - payload: Vec, - ) -> Result<(), SessionError> { - if id == 0 - || id.is_multiple_of(2) - || id <= self.last_call - || context.call_id != id - || context.instance_id != self.instance_id - || context.generation != self.generation - || context.incarnation != self.incarnation - || context.timeout_ms == 0 - || Duration::from_millis(context.timeout_ms) > self.config.maximum_call_timeout - || method.is_empty() - || method.len() > 128 - { - return Err(SessionError::Protocol); - } - self.last_call = id; - if self.quiescing { - return self.output.send_control(fault_frame( - id, - PluginFault::new(ErrorCode::Capacity, "plugin is quiescing"), - )); - } - if self.active.len() >= self.config.maximum_calls { - return self.output.send_control(fault_frame( - id, - PluginFault::new(ErrorCode::Capacity, "plugin call capacity is exhausted"), - )); - } - let deadline = Instant::now() - .checked_add(Duration::from_millis(context.timeout_ms)) - .ok_or(SessionError::Protocol)?; - let deadline = Some(deadline); - let (deadline_source, callback_deadline) = tokio::sync::watch::channel(deadline); - let credits = Arc::new(CreditWindow::default()); - let (cancellation, cancellation_signal) = CancellationSource::new(); - let host = HostClient { - parent_id: id, - deadline: callback_deadline, - cancellation: cancellation_signal.clone(), - callbacks: Arc::clone(&self.callbacks), - output: self.output.clone(), - maximum_stream_chunk_bytes: self.config.maximum_stream_chunk_bytes, - credits: credits.clone(), - }; - let call = PluginCall { - method, - context: context.clone(), - params, - payload, - host, - cancellation: cancellation_signal, - }; - let handler = Arc::clone(&self.handler); - let output = self.output.clone(); - let completed = self.completed.clone(); - let task_credits = Arc::clone(&credits); - let maximum_stream_chunk_bytes = self.config.maximum_stream_chunk_bytes; - let maximum_chunks = self.config.maximum_buffered_stream_chunks; - let task_context = context.clone(); - let lifecycle = self.lifecycle.clone(); - let task = tokio::spawn(async move { - let guard = CompletionGuard::new(id, completed); - let result = run_call( - handler, - call, - task_context, - task_credits, - output, - lifecycle, - deadline, - deadline_source, - maximum_stream_chunk_bytes, - maximum_chunks, - ) - .await; - guard.finish(if result.is_ok() { - CallTaskOutcome::Complete - } else { - CallTaskOutcome::TransportFailed - }); - }); - self.active.insert( - id, - ActiveCall { - context, - credits, - cancellation, - task, - }, - ); - Ok(()) - } - - fn cancel_call(&mut self, id: u64) -> Result<(), SessionError> { - if id == 0 || id.is_multiple_of(2) || id > self.last_call { - return Err(SessionError::Protocol); - } - if let Some(active) = self.active.remove(&id) { - active.cancellation.cancel(); - self.callbacks.finish_parent(id); - self.lifecycle.cancel(active.context); - active.task.abort(); - } - self.output - .send_control(Frame::control(Message::Cancelled { id })) - } - - fn finish_call(&mut self, finished: CallFinished) -> Result<(), SessionError> { - let Some(_active) = self.active.remove(&finished.id) else { - return Ok(()); - }; - self.callbacks.finish_parent(finished.id); - match finished.outcome { - CallTaskOutcome::Complete => Ok(()), - CallTaskOutcome::TransportFailed => Err(SessionError::Closed), - CallTaskOutcome::HandlerStopped => Err(SessionError::HandlerStopped), - } - } - - async fn close(&mut self) { - for (id, active) in std::mem::take(&mut self.active) { - active.cancellation.cancel(); - self.callbacks.finish_parent(id); - self.lifecycle.cancel(active.context); - active.task.abort(); - } - self.callbacks.close(); - self.lifecycle.shutdown(); - self.writer.abort(); - let _ = (&mut self.writer).await; - } -} - -#[expect( - clippy::too_many_arguments, - reason = "调用任务显式接收其完整资源边界，避免隐藏第二份会话状态" -)] -async fn run_call( - handler: Arc, - call: PluginCall, - context: CallContext, - credits: Arc, - output: Output, - lifecycle: LifecycleHooks, - deadline: Option, - deadline_source: tokio::sync::watch::Sender>, - maximum_stream_chunk_bytes: usize, - maximum_chunks: usize, -) -> Result<(), SessionError> { - let id = context.call_id; - let cancellation = call.cancellation.clone(); - let response = tokio::select! { - biased; - () = cancellation.cancelled() => return Ok(()), - () = wait_deadline(deadline) => { - lifecycle.cancel(context.clone()); - output.send_control(fault_frame( - id, - PluginFault::new(ErrorCode::Timeout, "plugin call timed out"), - ))?; - return Ok(()); - } - response = handler.call(call) => response, - }; - let mut reply = match response { - Ok(reply) => reply, - Err(error) => { - let _ = send_initial( - &output, - bounded_fault_frame(id, error), - &context, - &lifecycle, - deadline, - &cancellation, - ) - .await?; - return Ok(()); - } - }; - let initial = Frame { - message: Message::Result { - id, - result: reply.result, - }, - payload: reply.payload, - }; - if validate_frame(&initial).is_err() { - output.send_control(fault_frame( - id, - PluginFault::new( - ErrorCode::InvalidInput, - "plugin response metadata exceeds its limit", - ), - ))?; - return Ok(()); - } - let Some(mut stream) = reply.stream.take() else { - let _ = send_initial( - &output, - initial, - &context, - &lifecycle, - deadline, - &cancellation, - ) - .await?; - return Ok(()); - }; - let window_bytes = match credits.initial_window(deadline, &cancellation).await { - Ok(window) => window, - Err(SessionError::Timeout) => { - lifecycle.cancel(context.clone()); - output.send_control(fault_frame( - id, - PluginFault::new(ErrorCode::Timeout, "stream credit timed out"), - ))?; - return Ok(()); - } - Err(SessionError::Cancelled) => return Ok(()), - Err(error) => return Err(error), - }; - if let Err(fault) = - stream.validate_buffered(maximum_chunks, maximum_stream_chunk_bytes, window_bytes) - { - output.send_control(fault_frame(id, fault))?; - return Ok(()); - } - if !send_initial( - &output, - initial, - &context, - &lifecycle, - deadline, - &cancellation, - ) - .await? - { - return Ok(()); - } - let deadline = if context.resource_stream { - None - } else { - deadline - }; - deadline_source.send_replace(deadline); - let mut sequence = 0_u64; - loop { - let item = tokio::select! { - biased; - () = cancellation.cancelled() => return Ok(()), - () = wait_deadline(deadline) => { - lifecycle.cancel(context.clone()); - output.send_control(end_frame( - id, - Some(PluginFault::new(ErrorCode::Timeout, "plugin stream timed out")), - ))?; - return Ok(()); - } - item = stream.next() => item, - }; - let Some(item) = item else { - output.send_control(end_frame(id, None))?; - return Ok(()); - }; - let payload = match item { - Ok(payload) => payload, - Err(fault) => { - output.send_control(bounded_end_frame(id, fault))?; - return Ok(()); - } - }; - if let Err(fault) = - validate_stream_chunk(&payload, maximum_stream_chunk_bytes, window_bytes) - { - output.send_control(end_frame(id, Some(fault)))?; - return Ok(()); - } - match credits - .take(payload.len() as u64, deadline, &cancellation) - .await - { - Ok(()) => {} - Err(SessionError::Timeout) => { - lifecycle.cancel(context.clone()); - output.send_control(end_frame( - id, - Some(PluginFault::new( - ErrorCode::Timeout, - "stream credit timed out", - )), - ))?; - return Ok(()); - } - Err(SessionError::Cancelled) => return Ok(()), - Err(error) => return Err(error), - } - let frame = Frame { - message: Message::Stream { id, sequence }, - payload, - }; - if !send_stream_data( - &output, - frame, - id, - &context, - &lifecycle, - deadline, - &cancellation, - ) - .await? - { - return Ok(()); - } - sequence = sequence.checked_add(1).ok_or(SessionError::Protocol)?; - } -} - -async fn send_initial( - output: &Output, - frame: Frame, - context: &CallContext, - lifecycle: &LifecycleHooks, - deadline: Option, - cancellation: &CallCancellation, -) -> Result { - match output.send_data_until(frame, deadline, cancellation).await { - Ok(()) => Ok(true), - Err(SessionError::Cancelled) => Ok(false), - Err(SessionError::Timeout) => { - lifecycle.cancel(context.clone()); - output.send_control(fault_frame( - context.call_id, - PluginFault::new(ErrorCode::Timeout, "plugin call timed out"), - ))?; - Ok(false) - } - Err(error) => Err(error), - } -} - -async fn send_stream_data( - output: &Output, - frame: Frame, - id: u64, - context: &CallContext, - lifecycle: &LifecycleHooks, - deadline: Option, - cancellation: &CallCancellation, -) -> Result { - match output.send_data_until(frame, deadline, cancellation).await { - Ok(()) => Ok(true), - Err(SessionError::Cancelled) => Ok(false), - Err(SessionError::Timeout) => { - lifecycle.cancel(context.clone()); - output.send_control(end_frame( - id, - Some(PluginFault::new( - ErrorCode::Timeout, - "plugin stream timed out", - )), - ))?; - Ok(false) - } - Err(error) => Err(error), - } -} - -#[derive(Default)] -struct CreditWindow { - state: Mutex, - changed: tokio::sync::Notify, -} - -#[derive(Default)] -struct CreditState { - bytes: u64, - frames: u64, - maximum_bytes: u64, -} - -impl CreditWindow { - fn grant(&self, bytes: u32, frames: u32) -> Result<(), SessionError> { - if bytes == 0 || frames == 0 { - return Err(SessionError::Protocol); - } - let mut state = lock_unpoisoned(&self.state); - state.bytes = state - .bytes - .checked_add(u64::from(bytes)) - .ok_or(SessionError::Protocol)?; - state.frames = state - .frames - .checked_add(u64::from(frames)) - .ok_or(SessionError::Protocol)?; - state.maximum_bytes = state.maximum_bytes.max(state.bytes); - drop(state); - // 每个调用只有一个顺序消费 Credit 的任务；notify_one 会保留许可，避免 - // grant 恰好发生在状态检查和等待注册之间时丢失唤醒 - self.changed.notify_one(); - Ok(()) - } - - async fn initial_window( - &self, - deadline: Option, - cancellation: &CallCancellation, - ) -> Result { - loop { - let changed = self.changed.notified(); - { - let state = lock_unpoisoned(&self.state); - if state.bytes > 0 && state.frames > 0 { - return Ok(state.maximum_bytes); - } - } - tokio::select! { - biased; - () = cancellation.cancelled() => return Err(SessionError::Cancelled), - () = wait_deadline(deadline) => return Err(SessionError::Timeout), - () = changed => {} - } - } - } - - async fn take( - &self, - bytes: u64, - deadline: Option, - cancellation: &CallCancellation, - ) -> Result<(), SessionError> { - loop { - let changed = self.changed.notified(); - { - let mut state = lock_unpoisoned(&self.state); - if bytes > state.maximum_bytes { - return Err(SessionError::Protocol); - } - if state.bytes >= bytes && state.frames > 0 { - state.bytes -= bytes; - state.frames -= 1; - return Ok(()); - } - } - tokio::select! { - biased; - () = cancellation.cancelled() => return Err(SessionError::Cancelled), - () = wait_deadline(deadline) => return Err(SessionError::Timeout), - () = changed => {} - } - } - } -} - -fn validate_stream_chunk( - payload: &[u8], - maximum_stream_chunk_bytes: usize, - window_bytes: u64, -) -> Result<(), PluginFault> { - if payload.is_empty() { - return Err(PluginFault::new( - ErrorCode::InvalidInput, - "stream chunks must not be empty", - )); - } - if payload.len() as u64 > window_bytes { - return Err(PluginFault::new( - ErrorCode::InvalidInput, - "response event exceeds the host stream credit window", - )); - } - if payload.len() > maximum_stream_chunk_bytes { - return Err(PluginFault::new( - ErrorCode::InvalidInput, - "response event exceeds the stream chunk budget", - )); - } - Ok(()) -} - -fn capacity_fault(message: &'static str) -> PluginFault { - PluginFault::new(ErrorCode::Capacity, message) -} - -fn fault_frame(id: u64, error: PluginFault) -> Frame { - Frame::control(Message::Error { id, error }) -} - -fn end_frame(id: u64, error: Option) -> Frame { - Frame::control(Message::End { id, error }) -} - -fn bounded_fault_frame(id: u64, error: PluginFault) -> Frame { - let frame = fault_frame(id, error); - if validate_frame(&frame).is_ok() { - frame - } else { - fault_frame( - id, - PluginFault::new(ErrorCode::Fault, "plugin error metadata exceeds its limit"), - ) - } -} - -fn bounded_end_frame(id: u64, error: PluginFault) -> Frame { - let frame = end_frame(id, Some(error)); - if validate_frame(&frame).is_ok() { - frame - } else { - end_frame( - id, - Some(PluginFault::new( - ErrorCode::Fault, - "plugin stream error metadata exceeds its limit", - )), - ) - } -} - -fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { - mutex - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} - -async fn wait_deadline(deadline: Option) { - match deadline { - Some(deadline) => tokio::time::sleep_until(deadline).await, - None => std::future::pending().await, - } -} - -async fn callback_deadline(mut deadline: tokio::sync::watch::Receiver>) { - loop { - let current = *deadline.borrow_and_update(); - tokio::select! { - () = wait_deadline(current) => return, - changed = deadline.changed() => { - if changed.is_err() { return wait_deadline(current).await; } - } - } - } -} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/call.rs b/backend/crates/gateway-plugin/sdk/src/client/session/call.rs new file mode 100644 index 000000000..ebc21bd69 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/call.rs @@ -0,0 +1,66 @@ +//! 插件业务调用、处理器与响应合同 + +use std::{future::Future, pin::Pin}; + +use serde_json::Value; + +use crate::{CallContext, PluginFault}; + +use super::{CallCancellation, HostClient, ResponseStream}; + +/// 交给插件业务处理器的一次宿主调用；不实现 `Debug`，避免敏感载荷泄露 +pub struct PluginCall { + pub method: String, + pub context: CallContext, + pub params: Value, + pub payload: Vec, + pub host: HostClient, + pub cancellation: CallCancellation, +} + +/// 插件业务调用 future +pub type CallFuture<'a> = Pin> + Send + 'a>>; + +/// 插件只实现业务调用；会话负责握手、关联、流控、取消和关闭 +pub trait PluginHandler: Send + Sync + 'static { + fn call(&self, call: PluginCall) -> CallFuture<'_>; + + /// 取消时释放只属于该调用的临时业务状态 + /// 实现必须幂等、快速且不得执行阻塞 I/O + fn cancel(&self, _context: &CallContext) {} + + /// 会话停止接收新调用时触发 + /// 实现必须幂等、快速且不得执行阻塞 I/O + fn quiesce(&self) {} + + /// 会话关闭时释放插件自有临时状态 + /// 实现必须幂等、快速且不得执行阻塞 I/O + fn shutdown(&self) {} +} + +/// 一次业务成功结果；`stream` 存在时 SDK 自动产生单一流终态 +pub struct CallReply { + pub(super) result: Value, + pub(super) payload: Vec, + pub(super) stream: Option, +} + +impl CallReply { + #[must_use] + pub fn unary(result: Value, payload: Vec) -> Self { + Self { + result, + payload, + stream: None, + } + } + + #[must_use] + pub fn stream(result: Value, payload: Vec, stream: ResponseStream) -> Self { + Self { + result, + payload, + stream: Some(stream), + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/config.rs b/backend/crates/gateway-plugin/sdk/src/client/session/config.rs new file mode 100644 index 000000000..6c880ad38 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/config.rs @@ -0,0 +1,54 @@ +//! 插件侧会话的内存、并发与期限配置 + +use std::time::Duration; + +use super::SessionError; + +const MAXIMUM_STREAM_CHUNK_BYTES: usize = 16 * 1024 * 1024; +const MAXIMUM_CONCURRENCY: usize = 256; +const MAXIMUM_BUFFERED_STREAM_CHUNKS: usize = 65_536; + +/// 插件侧会话的内存、并发与期限边界 +#[derive(Debug, Clone, Copy)] +pub struct SessionConfig { + /// 单个业务流分块的预算，不限制普通调用正文的总长度 + pub maximum_stream_chunk_bytes: usize, + pub maximum_calls: usize, + pub maximum_callbacks: usize, + pub maximum_buffered_stream_chunks: usize, + pub handshake_timeout: Duration, + pub maximum_call_timeout: Duration, +} + +impl Default for SessionConfig { + fn default() -> Self { + Self { + // 流分块还需匹配宿主授予的信用窗口 + maximum_stream_chunk_bytes: 1024 * 1024, + maximum_calls: 32, + maximum_callbacks: 32, + maximum_buffered_stream_chunks: 1_024, + handshake_timeout: Duration::from_secs(5), + maximum_call_timeout: Duration::from_secs(120), + } + } +} + +impl SessionConfig { + pub(super) fn validate(self) -> Result { + if self.maximum_stream_chunk_bytes < 1_024 + || self.maximum_stream_chunk_bytes > MAXIMUM_STREAM_CHUNK_BYTES + || self.maximum_calls == 0 + || self.maximum_calls > MAXIMUM_CONCURRENCY + || self.maximum_callbacks == 0 + || self.maximum_callbacks > MAXIMUM_CONCURRENCY + || self.maximum_buffered_stream_chunks == 0 + || self.maximum_buffered_stream_chunks > MAXIMUM_BUFFERED_STREAM_CHUNKS + || self.handshake_timeout.is_zero() + || self.maximum_call_timeout.is_zero() + { + return Err(SessionError::Configuration); + } + Ok(self) + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/credit.rs b/backend/crates/gateway-plugin/sdk/src/client/session/credit.rs new file mode 100644 index 000000000..8e6810364 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/credit.rs @@ -0,0 +1,109 @@ +//! 响应流字节与帧信用的唯一状态和等待唤醒 + +use std::sync::Mutex; + +use tokio::time::Instant; + +use super::{CallCancellation, SessionError, lifecycle::wait_deadline}; + +#[derive(Default)] +pub(super) struct CreditWindow { + state: Mutex, + changed: tokio::sync::Notify, +} + +#[derive(Default)] +struct CreditState { + bytes: u64, + frames: u64, + maximum_bytes: u64, +} + +impl CreditWindow { + pub(super) fn maximum_bytes(&self) -> u64 { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .maximum_bytes + } + + pub(super) fn grant(&self, bytes: u32, frames: u32) -> Result<(), SessionError> { + if bytes == 0 || frames == 0 { + return Err(SessionError::Protocol); + } + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.bytes = state + .bytes + .checked_add(u64::from(bytes)) + .ok_or(SessionError::Protocol)?; + state.frames = state + .frames + .checked_add(u64::from(frames)) + .ok_or(SessionError::Protocol)?; + state.maximum_bytes = state.maximum_bytes.max(state.bytes); + drop(state); + // 每个调用只有一个顺序消费 Credit 的任务；notify_one 会保留许可，避免 + // grant 恰好发生在状态检查和等待注册之间时丢失唤醒 + self.changed.notify_one(); + Ok(()) + } + + pub(super) async fn initial_window( + &self, + deadline: Option, + cancellation: &CallCancellation, + ) -> Result { + loop { + let changed = self.changed.notified(); + { + let state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.bytes > 0 && state.frames > 0 { + return Ok(state.maximum_bytes); + } + } + tokio::select! { + biased; + () = cancellation.cancelled() => return Err(SessionError::Cancelled), + () = wait_deadline(deadline) => return Err(SessionError::Timeout), + () = changed => {} + } + } + } + + pub(super) async fn take( + &self, + bytes: u64, + deadline: Option, + cancellation: &CallCancellation, + ) -> Result<(), SessionError> { + loop { + let changed = self.changed.notified(); + { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if bytes > state.maximum_bytes { + return Err(SessionError::Protocol); + } + if state.bytes >= bytes && state.frames > 0 { + state.bytes -= bytes; + state.frames -= 1; + return Ok(()); + } + } + tokio::select! { + biased; + () = cancellation.cancelled() => return Err(SessionError::Cancelled), + () = wait_deadline(deadline) => return Err(SessionError::Timeout), + () = changed => {} + } + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs b/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs new file mode 100644 index 000000000..0c850465b --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs @@ -0,0 +1,712 @@ +//! 插件会话握手、在途调用分派与单一终结责任 + +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use serde_json::Value; +use tokio::{ + io::{AsyncRead, AsyncWrite}, + sync::{Semaphore, mpsc}, + task::JoinHandle, + time::Instant, +}; + +use crate::{ + CallContext, ErrorCode, Frame, FrameError, Handshake, Message, PROTOCOL_VERSION, PluginFault, +}; + +use super::super::frame::{read_frame, validate_frame, write_frame}; +use super::{ + CallCancellation, HostClient, PluginCall, PluginHandler, SessionConfig, SessionError, + credit::CreditWindow, + lifecycle::{CancellationSource, LifecycleHooks, wait_deadline}, + output::{Output, writer_loop}, + registry::CallbackRegistry, + stream::validate_stream_chunk, +}; + +/// 已读取并校验 Hello、尚未向宿主发送 Ready 的插件会话 +pub struct PluginSession { + reader: R, + writer: W, + handshake: Handshake, + config: SessionConfig, +} + +impl PluginSession +where + R: AsyncRead + Unpin + Send + 'static, + W: AsyncWrite + Unpin + Send + 'static, +{ + /// 读取并校验宿主握手 + /// 业务处理器可以在 `run` 前读取不可变 Handshake 构造注册描述 + /// + /// # Errors + /// + /// 配置无效、握手超时、传输关闭或 Hello 不符合协议时失败 + pub async fn accept( + mut reader: R, + writer: W, + config: SessionConfig, + ) -> Result { + let config = config.validate()?; + let frame = tokio::time::timeout(config.handshake_timeout, read_frame(&mut reader)) + .await + .map_err(|_| SessionError::Timeout)??; + let Message::Hello { handshake } = frame.message else { + return Err(SessionError::Handshake); + }; + if !frame.payload.is_empty() + || handshake.protocol_version != PROTOCOL_VERSION + || handshake.plugin_id.is_empty() + || handshake.instance_id.is_empty() + || handshake.generation == 0 + || handshake.incarnation.is_empty() + { + return Err(SessionError::Handshake); + } + Ok(Self { + reader, + writer, + handshake, + config, + }) + } + + #[must_use] + pub const fn handshake(&self) -> &Handshake { + &self.handshake + } + + /// 发送 Ready 并持续处理调用、回调、流控、取消与关闭 + /// + /// # Errors + /// + /// 帧损坏、协议违规、处理器异常停止或传输关闭时失败 + pub async fn run(mut self, handler: H) -> Result<(), SessionError> + where + H: PluginHandler, + { + let data_capacity = self + .config + .maximum_calls + .checked_add(self.config.maximum_callbacks) + .ok_or(SessionError::Configuration)?; + let control_reserve = self + .config + .maximum_calls + .checked_add(8) + .ok_or(SessionError::Configuration)?; + let queue_capacity = data_capacity + .checked_add(control_reserve) + .ok_or(SessionError::Configuration)?; + let (outbound, received) = mpsc::channel(queue_capacity); + let output = Output { + sender: outbound, + data_slots: Arc::new(Semaphore::new(data_capacity)), + }; + let callbacks = Arc::new(CallbackRegistry::new(self.config.maximum_callbacks)); + let handler: Arc = Arc::new(handler); + let lifecycle = LifecycleHooks::new(Arc::clone(&handler), control_reserve)?; + write_frame( + &mut self.writer, + &Frame::control(Message::Ready { + protocol_version: PROTOCOL_VERSION, + incarnation: self.handshake.incarnation.clone(), + }), + ) + .await?; + let writer = tokio::spawn(writer_loop(self.writer, received)); + let (completed, completions) = mpsc::channel(self.config.maximum_calls); + let mut driver = SessionDriver { + config: self.config, + instance_id: self.handshake.instance_id, + generation: self.handshake.generation, + incarnation: self.handshake.incarnation, + output, + callbacks, + handler, + lifecycle, + writer, + completed, + completions, + active: BTreeMap::new(), + last_call: 0, + quiescing: false, + }; + let result = driver.drive(self.reader).await; + driver.close().await; + result + } +} + +struct ActiveCall { + context: CallContext, + credits: Arc, + cancellation: CancellationSource, + task: JoinHandle<()>, +} + +#[derive(Clone, Copy)] +enum CallTaskOutcome { + Complete, + TransportFailed, + HandlerStopped, +} + +struct CallFinished { + id: u64, + outcome: CallTaskOutcome, +} + +struct CompletionGuard { + id: u64, + completed: mpsc::Sender, + armed: bool, +} + +impl CompletionGuard { + fn new(id: u64, completed: mpsc::Sender) -> Self { + Self { + id, + completed, + armed: true, + } + } + + fn finish(mut self, outcome: CallTaskOutcome) { + self.armed = false; + let _ = self.completed.try_send(CallFinished { + id: self.id, + outcome, + }); + } +} + +impl Drop for CompletionGuard { + fn drop(&mut self) { + if self.armed { + let _ = self.completed.try_send(CallFinished { + id: self.id, + outcome: CallTaskOutcome::HandlerStopped, + }); + } + } +} + +struct SessionDriver { + config: SessionConfig, + instance_id: String, + generation: u64, + incarnation: String, + output: Output, + callbacks: Arc, + handler: Arc, + lifecycle: LifecycleHooks, + writer: JoinHandle>, + completed: mpsc::Sender, + completions: mpsc::Receiver, + active: BTreeMap, + last_call: u64, + quiescing: bool, +} + +impl SessionDriver { + async fn drive(&mut self, mut reader: R) -> Result<(), SessionError> { + loop { + // read_exact 不是取消安全的；完成通知可以穿插处理，但必须保留同一个半帧读取 future + let incoming = read_frame(&mut reader); + tokio::pin!(incoming); + let frame = loop { + tokio::select! { + biased; + writer = &mut self.writer => { + return writer.map_err(|_| SessionError::Closed)?; + } + Some(finished) = self.completions.recv() => { + self.finish_call(finished)?; + } + frame = &mut incoming => { + break frame.map_err(|error| match error { + FrameError::Io(io_error) + if io_error.kind() == std::io::ErrorKind::UnexpectedEof => + { + SessionError::Closed + } + error => SessionError::Frame(error), + })?; + } + } + }; + if self.handle_frame(frame)? { + return Ok(()); + } + } + } + + fn handle_frame(&mut self, frame: Frame) -> Result { + match frame.message { + Message::Call { + id, + method, + context, + params, + } => { + self.start_call(id, method, context, params, frame.payload)?; + Ok(false) + } + Message::Credit { id, bytes, frames } if frame.payload.is_empty() => { + if id == 0 + || id.is_multiple_of(2) + || id > self.last_call + || bytes == 0 + || frames == 0 + { + return Err(SessionError::Protocol); + } + // End/Cancelled 已进入同一 FIFO 后，宿主先前为已消费分块排队的 + // Credit 仍可能稍晚抵达 + // 此时调用资源已经回收，忽略迟到信用即可； + // 未来或回调 ID 仍按协议错误处理 + if let Some(active) = self.active.get(&id) { + active.credits.grant(bytes, frames)?; + } + Ok(false) + } + Message::Cancel { id } if frame.payload.is_empty() => { + self.cancel_call(id)?; + Ok(false) + } + Message::Result { id, result } => { + self.callbacks.resolve(Frame { + message: Message::Result { id, result }, + payload: frame.payload, + })?; + Ok(false) + } + Message::Error { id, error } => { + self.callbacks.resolve(Frame { + message: Message::Error { id, error }, + payload: frame.payload, + })?; + Ok(false) + } + Message::Quiesce if frame.payload.is_empty() => { + if !self.quiescing { + self.quiescing = true; + self.lifecycle.quiesce(); + } + Ok(false) + } + Message::Shutdown if frame.payload.is_empty() => Ok(true), + _ => Err(SessionError::Protocol), + } + } + + fn start_call( + &mut self, + id: u64, + method: String, + context: CallContext, + params: Value, + payload: Vec, + ) -> Result<(), SessionError> { + if id == 0 + || id.is_multiple_of(2) + || id <= self.last_call + || context.call_id != id + || context.instance_id != self.instance_id + || context.generation != self.generation + || context.incarnation != self.incarnation + || context.timeout_ms == 0 + || Duration::from_millis(context.timeout_ms) > self.config.maximum_call_timeout + || method.is_empty() + || method.len() > 128 + { + return Err(SessionError::Protocol); + } + self.last_call = id; + if self.quiescing { + return self.output.send_control(fault_frame( + id, + PluginFault::new(ErrorCode::Capacity, "plugin is quiescing"), + )); + } + if self.active.len() >= self.config.maximum_calls { + return self.output.send_control(fault_frame( + id, + PluginFault::new(ErrorCode::Capacity, "plugin call capacity is exhausted"), + )); + } + let deadline = Instant::now() + .checked_add(Duration::from_millis(context.timeout_ms)) + .ok_or(SessionError::Protocol)?; + let deadline = Some(deadline); + let (deadline_source, callback_deadline) = tokio::sync::watch::channel(deadline); + let credits = Arc::new(CreditWindow::default()); + let (cancellation, cancellation_signal) = CancellationSource::new(); + let host = HostClient { + parent_id: id, + deadline: callback_deadline, + cancellation: cancellation_signal.clone(), + callbacks: Arc::clone(&self.callbacks), + output: self.output.clone(), + maximum_stream_chunk_bytes: self.config.maximum_stream_chunk_bytes, + credits: credits.clone(), + }; + let call = PluginCall { + method, + context: context.clone(), + params, + payload, + host, + cancellation: cancellation_signal, + }; + let handler = Arc::clone(&self.handler); + let output = self.output.clone(); + let completed = self.completed.clone(); + let task_credits = Arc::clone(&credits); + let maximum_stream_chunk_bytes = self.config.maximum_stream_chunk_bytes; + let maximum_chunks = self.config.maximum_buffered_stream_chunks; + let task_context = context.clone(); + let lifecycle = self.lifecycle.clone(); + let task = tokio::spawn(async move { + let guard = CompletionGuard::new(id, completed); + let result = run_call( + handler, + call, + task_context, + task_credits, + output, + lifecycle, + deadline, + deadline_source, + maximum_stream_chunk_bytes, + maximum_chunks, + ) + .await; + guard.finish(if result.is_ok() { + CallTaskOutcome::Complete + } else { + CallTaskOutcome::TransportFailed + }); + }); + self.active.insert( + id, + ActiveCall { + context, + credits, + cancellation, + task, + }, + ); + Ok(()) + } + + fn cancel_call(&mut self, id: u64) -> Result<(), SessionError> { + if id == 0 || id.is_multiple_of(2) || id > self.last_call { + return Err(SessionError::Protocol); + } + if let Some(active) = self.active.remove(&id) { + active.cancellation.cancel(); + self.callbacks.finish_parent(id); + self.lifecycle.cancel(active.context); + active.task.abort(); + } + self.output + .send_control(Frame::control(Message::Cancelled { id })) + } + + fn finish_call(&mut self, finished: CallFinished) -> Result<(), SessionError> { + let Some(_active) = self.active.remove(&finished.id) else { + return Ok(()); + }; + self.callbacks.finish_parent(finished.id); + match finished.outcome { + CallTaskOutcome::Complete => Ok(()), + CallTaskOutcome::TransportFailed => Err(SessionError::Closed), + CallTaskOutcome::HandlerStopped => Err(SessionError::HandlerStopped), + } + } + + async fn close(&mut self) { + for (id, active) in std::mem::take(&mut self.active) { + active.cancellation.cancel(); + self.callbacks.finish_parent(id); + self.lifecycle.cancel(active.context); + active.task.abort(); + } + self.callbacks.close(); + self.lifecycle.shutdown(); + self.writer.abort(); + let _ = (&mut self.writer).await; + } +} + +#[expect( + clippy::too_many_arguments, + reason = "调用任务显式接收其完整资源边界，避免隐藏第二份会话状态" +)] +async fn run_call( + handler: Arc, + call: PluginCall, + context: CallContext, + credits: Arc, + output: Output, + lifecycle: LifecycleHooks, + deadline: Option, + deadline_source: tokio::sync::watch::Sender>, + maximum_stream_chunk_bytes: usize, + maximum_chunks: usize, +) -> Result<(), SessionError> { + let id = context.call_id; + let cancellation = call.cancellation.clone(); + let response = tokio::select! { + biased; + () = cancellation.cancelled() => return Ok(()), + () = wait_deadline(deadline) => { + lifecycle.cancel(context.clone()); + output.send_control(fault_frame( + id, + PluginFault::new(ErrorCode::Timeout, "plugin call timed out"), + ))?; + return Ok(()); + } + response = handler.call(call) => response, + }; + let mut reply = match response { + Ok(reply) => reply, + Err(error) => { + let _ = send_initial( + &output, + bounded_fault_frame(id, error), + &context, + &lifecycle, + deadline, + &cancellation, + ) + .await?; + return Ok(()); + } + }; + let initial = Frame { + message: Message::Result { + id, + result: reply.result, + }, + payload: reply.payload, + }; + if validate_frame(&initial).is_err() { + output.send_control(fault_frame( + id, + PluginFault::new( + ErrorCode::InvalidInput, + "plugin response metadata exceeds its limit", + ), + ))?; + return Ok(()); + } + let Some(mut stream) = reply.stream.take() else { + let _ = send_initial( + &output, + initial, + &context, + &lifecycle, + deadline, + &cancellation, + ) + .await?; + return Ok(()); + }; + let window_bytes = match credits.initial_window(deadline, &cancellation).await { + Ok(window) => window, + Err(SessionError::Timeout) => { + lifecycle.cancel(context.clone()); + output.send_control(fault_frame( + id, + PluginFault::new(ErrorCode::Timeout, "stream credit timed out"), + ))?; + return Ok(()); + } + Err(SessionError::Cancelled) => return Ok(()), + Err(error) => return Err(error), + }; + if let Err(fault) = + stream.validate_buffered(maximum_chunks, maximum_stream_chunk_bytes, window_bytes) + { + output.send_control(fault_frame(id, fault))?; + return Ok(()); + } + if !send_initial( + &output, + initial, + &context, + &lifecycle, + deadline, + &cancellation, + ) + .await? + { + return Ok(()); + } + let deadline = if context.resource_stream { + None + } else { + deadline + }; + deadline_source.send_replace(deadline); + let mut sequence = 0_u64; + loop { + let item = tokio::select! { + biased; + () = cancellation.cancelled() => return Ok(()), + () = wait_deadline(deadline) => { + lifecycle.cancel(context.clone()); + output.send_control(end_frame( + id, + Some(PluginFault::new(ErrorCode::Timeout, "plugin stream timed out")), + ))?; + return Ok(()); + } + item = stream.next() => item, + }; + let Some(item) = item else { + output.send_control(end_frame(id, None))?; + return Ok(()); + }; + let payload = match item { + Ok(payload) => payload, + Err(fault) => { + output.send_control(bounded_end_frame(id, fault))?; + return Ok(()); + } + }; + if let Err(fault) = + validate_stream_chunk(&payload, maximum_stream_chunk_bytes, window_bytes) + { + output.send_control(end_frame(id, Some(fault)))?; + return Ok(()); + } + match credits + .take(payload.len() as u64, deadline, &cancellation) + .await + { + Ok(()) => {} + Err(SessionError::Timeout) => { + lifecycle.cancel(context.clone()); + output.send_control(end_frame( + id, + Some(PluginFault::new( + ErrorCode::Timeout, + "stream credit timed out", + )), + ))?; + return Ok(()); + } + Err(SessionError::Cancelled) => return Ok(()), + Err(error) => return Err(error), + } + let frame = Frame { + message: Message::Stream { id, sequence }, + payload, + }; + if !send_stream_data( + &output, + frame, + id, + &context, + &lifecycle, + deadline, + &cancellation, + ) + .await? + { + return Ok(()); + } + sequence = sequence.checked_add(1).ok_or(SessionError::Protocol)?; + } +} + +async fn send_initial( + output: &Output, + frame: Frame, + context: &CallContext, + lifecycle: &LifecycleHooks, + deadline: Option, + cancellation: &CallCancellation, +) -> Result { + match output.send_data_until(frame, deadline, cancellation).await { + Ok(()) => Ok(true), + Err(SessionError::Cancelled) => Ok(false), + Err(SessionError::Timeout) => { + lifecycle.cancel(context.clone()); + output.send_control(fault_frame( + context.call_id, + PluginFault::new(ErrorCode::Timeout, "plugin call timed out"), + ))?; + Ok(false) + } + Err(error) => Err(error), + } +} + +async fn send_stream_data( + output: &Output, + frame: Frame, + id: u64, + context: &CallContext, + lifecycle: &LifecycleHooks, + deadline: Option, + cancellation: &CallCancellation, +) -> Result { + match output.send_data_until(frame, deadline, cancellation).await { + Ok(()) => Ok(true), + Err(SessionError::Cancelled) => Ok(false), + Err(SessionError::Timeout) => { + lifecycle.cancel(context.clone()); + output.send_control(end_frame( + id, + Some(PluginFault::new( + ErrorCode::Timeout, + "plugin stream timed out", + )), + ))?; + Ok(false) + } + Err(error) => Err(error), + } +} + +fn fault_frame(id: u64, error: PluginFault) -> Frame { + Frame::control(Message::Error { id, error }) +} + +fn end_frame(id: u64, error: Option) -> Frame { + Frame::control(Message::End { id, error }) +} + +fn bounded_fault_frame(id: u64, error: PluginFault) -> Frame { + let frame = fault_frame(id, error); + if validate_frame(&frame).is_ok() { + frame + } else { + fault_frame( + id, + PluginFault::new(ErrorCode::Fault, "plugin error metadata exceeds its limit"), + ) + } +} + +fn bounded_end_frame(id: u64, error: PluginFault) -> Frame { + let frame = end_frame(id, Some(error)); + if validate_frame(&frame).is_ok() { + frame + } else { + end_frame( + id, + Some(PluginFault::new( + ErrorCode::Fault, + "plugin stream error metadata exceeds its limit", + )), + ) + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/error.rs b/backend/crates/gateway-plugin/sdk/src/client/session/error.rs new file mode 100644 index 000000000..834027fc8 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/error.rs @@ -0,0 +1,47 @@ +//! 插件侧传输和会话错误及稳定业务错误投影 + +use crate::{ErrorCode, FrameError, PluginFault}; + +/// 插件侧传输和会话错误 +#[derive(Debug, thiserror::Error)] +pub enum SessionError { + #[error("plugin session configuration is invalid")] + Configuration, + #[error("plugin handshake is invalid")] + Handshake, + #[error("plugin protocol is invalid")] + Protocol, + #[error("plugin transport is closed")] + Closed, + #[error("plugin operation exceeded its deadline")] + Timeout, + #[error("plugin call was cancelled")] + Cancelled, + #[error("plugin session capacity is exhausted")] + Capacity, + #[error("plugin handler stopped unexpectedly")] + HandlerStopped, + #[error("host callback returned an error: {0:?}")] + Remote(PluginFault), + #[error(transparent)] + Frame(#[from] FrameError), +} + +impl SessionError { + /// 将本地会话失败收敛为可返回宿主的稳定插件错误 + #[must_use] + pub fn into_plugin_fault(self) -> PluginFault { + match self { + Self::Remote(fault) => fault, + Self::Timeout => PluginFault::new(ErrorCode::Timeout, "host callback timed out"), + Self::Cancelled => PluginFault::new(ErrorCode::Cancelled, "parent call was cancelled"), + Self::Capacity => PluginFault::new(ErrorCode::Capacity, "plugin capacity is exhausted"), + Self::Configuration | Self::Protocol | Self::Handshake | Self::Frame(_) => { + PluginFault::new(ErrorCode::InvalidInput, "plugin session input is invalid") + } + Self::Closed | Self::HandlerStopped => { + PluginFault::new(ErrorCode::Fault, "plugin session is unavailable") + } + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/lifecycle.rs b/backend/crates/gateway-plugin/sdk/src/client/session/lifecycle.rs new file mode 100644 index 000000000..fefc65cdb --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/lifecycle.rs @@ -0,0 +1,116 @@ +//! 调用取消、期限等待与插件生命周期钩子的分派 + +use std::sync::Arc; + +use tokio::{sync::watch, time::Instant}; + +use crate::CallContext; + +use super::{PluginHandler, SessionError}; + +/// 一次宿主调用的可克隆取消信号 +#[derive(Clone)] +pub struct CallCancellation { + receiver: watch::Receiver, +} + +impl CallCancellation { + #[must_use] + pub fn is_cancelled(&self) -> bool { + *self.receiver.borrow() + } + + pub async fn cancelled(&self) { + let mut receiver = self.receiver.clone(); + if *receiver.borrow() { + return; + } + while receiver.changed().await.is_ok() { + if *receiver.borrow() { + return; + } + } + } +} + +pub(super) struct CancellationSource { + sender: watch::Sender, +} + +impl CancellationSource { + pub(super) fn new() -> (Self, CallCancellation) { + let (sender, receiver) = watch::channel(false); + (Self { sender }, CallCancellation { receiver }) + } + + pub(super) fn cancel(&self) { + self.sender.send_replace(true); + } +} + +enum LifecycleEvent { + Cancel(CallContext), + Quiesce, + Shutdown, +} + +#[derive(Clone)] +pub(super) struct LifecycleHooks { + events: std::sync::mpsc::SyncSender, +} + +impl LifecycleHooks { + pub(super) fn new( + handler: Arc, + capacity: usize, + ) -> Result { + let (events, received) = std::sync::mpsc::sync_channel(capacity); + std::thread::Builder::new() + .name("gateway-plugin-sdk-lifecycle".into()) + .spawn(move || { + while let Ok(event) = received.recv() { + match event { + LifecycleEvent::Cancel(context) => handler.cancel(&context), + LifecycleEvent::Quiesce => handler.quiesce(), + LifecycleEvent::Shutdown => { + handler.shutdown(); + break; + } + } + } + }) + .map_err(|_| SessionError::Closed)?; + Ok(Self { events }) + } + + pub(super) fn cancel(&self, context: CallContext) { + let _ = self.events.try_send(LifecycleEvent::Cancel(context)); + } + + pub(super) fn quiesce(&self) { + let _ = self.events.try_send(LifecycleEvent::Quiesce); + } + + pub(super) fn shutdown(&self) { + let _ = self.events.try_send(LifecycleEvent::Shutdown); + } +} + +pub(super) async fn wait_deadline(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending().await, + } +} + +pub(super) async fn callback_deadline(mut deadline: tokio::sync::watch::Receiver>) { + loop { + let current = *deadline.borrow_and_update(); + tokio::select! { + () = wait_deadline(current) => return, + changed = deadline.changed() => { + if changed.is_err() { return wait_deadline(current).await; } + } + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs b/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs new file mode 100644 index 000000000..7e7b5f62b --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs @@ -0,0 +1,19 @@ +//! 插件侧双向 RPC 会话的公开合同与内部职责入口 + +mod call; +mod config; +mod credit; +mod driver; +mod error; +mod lifecycle; +mod output; +mod registry; +mod stream; + +pub use call::{CallFuture, CallReply, PluginCall, PluginHandler}; +pub use config::SessionConfig; +pub use driver::PluginSession; +pub use error::SessionError; +pub use lifecycle::CallCancellation; +pub use registry::{HostClient, HostReply}; +pub use stream::{PullResponseFuture, PullResponseStream, ResponseStream, StreamSender}; diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/output.rs b/backend/crates/gateway-plugin/sdk/src/client/session/output.rs new file mode 100644 index 000000000..aa2a7cffb --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/output.rs @@ -0,0 +1,84 @@ +//! 会话输出队列、控制帧预留与唯一串行写泵 + +use std::sync::Arc; + +use tokio::{ + io::AsyncWrite, + sync::{OwnedSemaphorePermit, Semaphore, mpsc}, + time::Instant, +}; + +use crate::Frame; + +use super::super::frame::{validate_frame, write_frame}; +use super::{CallCancellation, SessionError, lifecycle::wait_deadline}; + +pub(super) struct Output { + pub(super) sender: mpsc::Sender, + pub(super) data_slots: Arc, +} + +impl Clone for Output { + fn clone(&self) -> Self { + Self { + sender: self.sender.clone(), + data_slots: Arc::clone(&self.data_slots), + } + } +} + +impl Output { + pub(super) async fn send_data(&self, frame: Frame) -> Result<(), SessionError> { + validate_frame(&frame)?; + let permit = Arc::clone(&self.data_slots) + .acquire_owned() + .await + .map_err(|_| SessionError::Closed)?; + self.sender + .send(Outbound { + frame, + _data_slot: Some(permit), + }) + .await + .map_err(|_| SessionError::Closed) + } + + pub(super) async fn send_data_until( + &self, + frame: Frame, + deadline: Option, + cancellation: &CallCancellation, + ) -> Result<(), SessionError> { + tokio::select! { + biased; + () = cancellation.cancelled() => Err(SessionError::Cancelled), + () = wait_deadline(deadline) => Err(SessionError::Timeout), + result = self.send_data(frame) => result, + } + } + + pub(super) fn send_control(&self, frame: Frame) -> Result<(), SessionError> { + validate_frame(&frame)?; + self.sender + .try_send(Outbound { + frame, + _data_slot: None, + }) + .map_err(|_| SessionError::Capacity) + } +} + +pub(super) struct Outbound { + frame: Frame, + _data_slot: Option, +} + +pub(super) async fn writer_loop( + mut writer: W, + mut received: mpsc::Receiver, +) -> Result<(), SessionError> { + while let Some(outbound) = received.recv().await { + write_frame(&mut writer, &outbound.frame).await?; + } + Ok(()) +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs b/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs new file mode 100644 index 000000000..cfe11e86e --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs @@ -0,0 +1,251 @@ +//! 宿主回调客户端、关联编号与有界退休回调注册表 + +use std::{ + collections::{BTreeMap, BTreeSet, VecDeque}, + sync::{Arc, Mutex}, +}; + +use serde_json::Value; +use tokio::{sync::oneshot, time::Instant}; + +use crate::{Frame, Message}; + +use super::super::frame::validate_frame; +use super::{ + CallCancellation, SessionError, credit::CreditWindow, lifecycle::callback_deadline, + output::Output, +}; + +/// 宿主回调的成功结果；不实现 `Debug`，避免载荷进入普通诊断 +pub struct HostReply { + pub result: Value, + pub payload: Vec, +} + +/// 绑定父调用、期限和取消信号的宿主回调客户端 +#[derive(Clone)] +pub struct HostClient { + pub(super) parent_id: u64, + pub(super) deadline: tokio::sync::watch::Receiver>, + pub(super) cancellation: CallCancellation, + pub(super) callbacks: Arc, + pub(super) output: Output, + pub(super) maximum_stream_chunk_bytes: usize, + pub(super) credits: Arc, +} + +impl HostClient { + pub(crate) fn maximum_stream_chunk_bytes(&self) -> usize { + let window = self.credits.maximum_bytes(); + if window == 0 { + self.maximum_stream_chunk_bytes + } else { + self.maximum_stream_chunk_bytes + .min(usize::try_from(window).unwrap_or(usize::MAX)) + } + } + + /// 发起一次与父调用关联的宿主回调 + /// + /// # Errors + /// + /// 方法或帧无效、容量耗尽、父调用取消、期限到达、连接关闭或宿主返回错误时失败 + pub async fn call( + &self, + method: impl Into, + params: Value, + payload: Vec, + ) -> Result { + let method = method.into(); + if method.is_empty() || method.len() > 128 { + return Err(SessionError::Protocol); + } + let ordered = self.callbacks.order.lock().await; + let (id, received) = self.callbacks.begin(self.parent_id)?; + let frame = Frame { + message: Message::Callback { + id, + parent_id: self.parent_id, + method, + params, + }, + payload, + }; + if validate_frame(&frame).is_err() { + self.callbacks.retire(id); + return Err(SessionError::Protocol); + } + let sent = tokio::select! { + biased; + () = self.cancellation.cancelled() => Err(SessionError::Cancelled), + () = callback_deadline(self.deadline.clone()) => Err(SessionError::Timeout), + result = self.output.send_data(frame) => result, + }; + if let Err(error) = sent { + self.callbacks.retire(id); + return Err(error); + } + drop(ordered); + tokio::select! { + biased; + () = self.cancellation.cancelled() => { + self.callbacks.retire(id); + Err(SessionError::Cancelled) + } + () = callback_deadline(self.deadline.clone()) => { + self.callbacks.retire(id); + Err(SessionError::Timeout) + } + result = received => result.map_err(|_| SessionError::Closed)?, + } + } +} + +pub(super) struct CallbackRegistry { + state: Mutex, + order: tokio::sync::Mutex<()>, + maximum: usize, + maximum_retired: usize, +} + +struct CallbackState { + next_id: u64, + pending: BTreeMap, + retired: BTreeSet, + retired_order: VecDeque, +} + +struct PendingCallback { + parent_id: u64, + result: oneshot::Sender>, +} + +impl CallbackRegistry { + pub(super) fn new(maximum: usize) -> Self { + Self { + state: Mutex::new(CallbackState { + next_id: 2, + pending: BTreeMap::new(), + retired: BTreeSet::new(), + retired_order: VecDeque::new(), + }), + order: tokio::sync::Mutex::new(()), + maximum, + maximum_retired: maximum.saturating_mul(2), + } + } + + fn begin( + &self, + parent_id: u64, + ) -> Result<(u64, oneshot::Receiver>), SessionError> { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.pending.len() >= self.maximum { + return Err(SessionError::Capacity); + } + let id = state.next_id; + state.next_id = state.next_id.checked_add(2).ok_or(SessionError::Capacity)?; + let (result, received) = oneshot::channel(); + state + .pending + .insert(id, PendingCallback { parent_id, result }); + Ok((id, received)) + } + + pub(super) fn resolve(&self, frame: Frame) -> Result<(), SessionError> { + let id = match &frame.message { + Message::Result { id, .. } | Message::Error { id, .. } => *id, + _ => return Err(SessionError::Protocol), + }; + let pending = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(pending) = state.pending.remove(&id) { + Some(pending) + } else if state.retired.remove(&id) { + state.retired_order.retain(|retired| *retired != id); + None + } else { + return Err(SessionError::Protocol); + } + }; + let Some(pending) = pending else { + return Ok(()); + }; + let result = match frame.message { + Message::Result { result, .. } => Ok(HostReply { + result, + payload: frame.payload, + }), + Message::Error { error, .. } if frame.payload.is_empty() => { + Err(SessionError::Remote(error)) + } + _ => return Err(SessionError::Protocol), + }; + let _ = pending.result.send(result); + Ok(()) + } + + fn retire(&self, id: u64) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.pending.remove(&id).is_some() { + retire_callback(&mut state, id, self.maximum_retired); + } + } + + pub(super) fn finish_parent(&self, parent_id: u64) { + let pending = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut remaining = BTreeMap::new(); + let mut finished = Vec::new(); + for (id, pending) in std::mem::take(&mut state.pending) { + if pending.parent_id == parent_id { + retire_callback(&mut state, id, self.maximum_retired); + finished.push(pending); + } else { + remaining.insert(id, pending); + } + } + state.pending = remaining; + finished + }; + for pending in pending { + let _ = pending.result.send(Err(SessionError::Cancelled)); + } + } + + pub(super) fn close(&self) { + let pending = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + std::mem::take(&mut state.pending) + }; + for pending in pending.into_values() { + let _ = pending.result.send(Err(SessionError::Closed)); + } + } +} + +fn retire_callback(state: &mut CallbackState, id: u64, maximum: usize) { + if state.retired.insert(id) { + state.retired_order.push_back(id); + } + while state.retired_order.len() > maximum { + if let Some(expired) = state.retired_order.pop_front() { + state.retired.remove(&expired); + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/stream.rs b/backend/crates/gateway-plugin/sdk/src/client/session/stream.rs new file mode 100644 index 000000000..c10690122 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/stream.rs @@ -0,0 +1,142 @@ +//! 插件响应流的有界生产、拉取与分块校验 + +use std::{collections::VecDeque, future::Future, num::NonZeroUsize, pin::Pin}; + +use tokio::sync::mpsc; + +use crate::{ErrorCode, PluginFault}; + +use super::SessionError; + +enum StreamSource { + Buffered(VecDeque>), + Channel(mpsc::Receiver, PluginFault>>), + Pull(Box), +} + +pub type PullResponseFuture<'a> = + Pin, PluginFault>>> + Send + 'a>>; + +pub trait PullResponseStream: Send { + fn next(&mut self) -> PullResponseFuture<'_>; +} + +/// SDK 管理 sequence、Credit 和终态的响应流 +pub struct ResponseStream { + source: StreamSource, + declared_capacity: usize, +} + +impl ResponseStream { + /// 从已完成业务校验的有限分块构造流；SDK 会在发送 `Result` 前预检全部分块 + #[must_use] + pub fn from_chunks(chunks: Vec>) -> Self { + let declared_capacity = chunks.len().max(1); + Self { + source: StreamSource::Buffered(chunks.into()), + declared_capacity, + } + } + + /// 创建由生产者驱动的有界流队列 + #[must_use] + pub fn channel(capacity: NonZeroUsize) -> (StreamSender, Self) { + let (sender, receiver) = mpsc::channel(capacity.get()); + ( + StreamSender { sender }, + Self { + source: StreamSource::Channel(receiver), + declared_capacity: capacity.get(), + }, + ) + } + + /// 按消费进度拉取下一帧；返回 None 表示完成，丢弃流会丢弃生产者 + pub fn pull(source: Box) -> Self { + Self { + source: StreamSource::Pull(source), + declared_capacity: 1, + } + } + + pub(super) fn validate_buffered( + &self, + maximum_chunks: usize, + maximum_stream_chunk_bytes: usize, + window_bytes: u64, + ) -> Result<(), PluginFault> { + if self.declared_capacity > maximum_chunks { + return Err(capacity_fault( + "response stream exceeds the local queue limit", + )); + } + if let StreamSource::Buffered(chunks) = &self.source { + for chunk in chunks { + validate_stream_chunk(chunk, maximum_stream_chunk_bytes, window_bytes)?; + } + } + Ok(()) + } + + pub(crate) async fn next(&mut self) -> Option, PluginFault>> { + match &mut self.source { + StreamSource::Buffered(chunks) => chunks.pop_front().map(Ok), + StreamSource::Channel(receiver) => receiver.recv().await, + StreamSource::Pull(source) => source.next().await, + } + } +} + +/// 动态响应流的有界生产端；最后一个 sender 被释放表示成功终态 +#[derive(Clone)] +pub struct StreamSender { + sender: mpsc::Sender, PluginFault>>, +} + +impl StreamSender { + /// 排队一个业务分块；实际线协议流控由会话完成 + pub async fn send(&self, payload: Vec) -> Result<(), SessionError> { + self.sender + .send(Ok(payload)) + .await + .map_err(|_| SessionError::Closed) + } + + /// 以业务错误结束动态流 + pub async fn fail(&self, fault: PluginFault) -> Result<(), SessionError> { + self.sender + .send(Err(fault)) + .await + .map_err(|_| SessionError::Closed) + } +} + +pub(super) fn validate_stream_chunk( + payload: &[u8], + maximum_stream_chunk_bytes: usize, + window_bytes: u64, +) -> Result<(), PluginFault> { + if payload.is_empty() { + return Err(PluginFault::new( + ErrorCode::InvalidInput, + "stream chunks must not be empty", + )); + } + if payload.len() as u64 > window_bytes { + return Err(PluginFault::new( + ErrorCode::InvalidInput, + "response event exceeds the host stream credit window", + )); + } + if payload.len() > maximum_stream_chunk_bytes { + return Err(PluginFault::new( + ErrorCode::InvalidInput, + "response event exceeds the stream chunk budget", + )); + } + Ok(()) +} + +fn capacity_fault(message: &'static str) -> PluginFault { + PluginFault::new(ErrorCode::Capacity, message) +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/frame.rs b/backend/crates/gateway-plugin/sdk/tests/client/frame.rs new file mode 100644 index 000000000..a7ed59402 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/frame.rs @@ -0,0 +1,87 @@ +//! 二进制帧往返、长度预算与不完整帧拒绝测试 + +use gateway_plugin_sdk::{ + Frame, FrameError, Message, + client::{read_frame, write_frame}, +}; +use serde_json::json; + +#[tokio::test] +async fn binary_payload_round_trips_without_json_encoding() { + let frame = Frame { + message: Message::Stream { + id: 17, + sequence: 2, + }, + payload: vec![0, 255, 13, 10, 128], + }; + let mut bytes = Vec::new(); + write_frame(&mut bytes, &frame).await.unwrap(); + assert!(bytes.ends_with(&frame.payload)); + assert_eq!(read_frame(&mut bytes.as_slice()).await.unwrap(), frame); +} + +#[tokio::test] +async fn malicious_lengths_fail_before_reading_or_allocating_payload() { + let bytes = [0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0]; + assert!(matches!( + read_frame(&mut bytes.as_slice()).await, + Err(FrameError::Length) + )); +} + +#[tokio::test] +async fn large_payload_is_exact_across_multiple_io_buffers_and_following_frames() { + let frame = Frame { + message: Message::Result { + id: 1, + result: json!({}), + }, + payload: (0..101 * 1024 * 1024 + 13) + .map(|index| (index % 251) as u8) + .collect(), + }; + let next = Frame::control(Message::Cancel { id: 3 }); + let (mut writer, mut reader) = tokio::io::duplex(16 * 1024); + let writing = tokio::spawn(async move { + write_frame(&mut writer, &frame).await.unwrap(); + write_frame(&mut writer, &next).await.unwrap(); + (frame, next) + }); + let received = read_frame(&mut reader).await.unwrap(); + let following = read_frame(&mut reader).await.unwrap(); + let (original, next) = writing.await.unwrap(); + assert_eq!(received, original); + assert_eq!(following, next); +} + +#[tokio::test] +async fn truncated_large_payload_never_becomes_a_successful_message() { + let message = serde_json::to_vec(&Message::Result { + id: 1, + result: json!({}), + }) + .unwrap(); + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(message.len() as u32).to_be_bytes()); + bytes.extend_from_slice(&(101_u64 * 1024 * 1024).to_be_bytes()); + bytes.extend_from_slice(&message); + bytes.extend_from_slice(&[1, 2, 3]); + assert!(matches!( + read_frame(&mut bytes.as_slice()).await, + Err(FrameError::Io(_)) + )); +} + +#[tokio::test] +async fn truncated_frame_never_becomes_a_successful_message() { + let mut bytes = Vec::new(); + write_frame(&mut bytes, &Frame::control(Message::Cancel { id: 3 })) + .await + .unwrap(); + bytes.pop(); + assert!(matches!( + read_frame(&mut bytes.as_slice()).await, + Err(FrameError::Io(_)) + )); +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/mod.rs b/backend/crates/gateway-plugin/sdk/tests/client/mod.rs index e239f0e45..551836fae 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/mod.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/mod.rs @@ -1,5 +1,6 @@ //! 插件客户端、处理器注册与会话的测试入口 +mod frame; mod http; mod plugin; mod read; diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session.rs b/backend/crates/gateway-plugin/sdk/tests/client/session.rs deleted file mode 100644 index e3ec1c119..000000000 --- a/backend/crates/gateway-plugin/sdk/tests/client/session.rs +++ /dev/null @@ -1,1760 +0,0 @@ -//! 验证插件会话的二进制传输、帧长度限制与不完整消息拒绝 - -use std::{ - future::pending, - num::NonZeroUsize, - sync::{ - Arc, - atomic::{AtomicBool, AtomicUsize, Ordering}, - }, - time::Duration, -}; - -use gateway_plugin_sdk::{ - CallContext, Capability, ContributionDeclaration, Contributions, ErrorCode, Frame, FrameError, - Handshake, Message, PROTOCOL_VERSION, PluginFault, Stage, - call::middleware::{ - BODY_READ_METHOD, HANDLE_METHOD, MiddlewareBodyDisposition, MiddlewareBodyFrame, - MiddlewareBodyFraming, MiddlewareBodyHandle, MiddlewareBodyReadResult, MiddlewareHeader, - MiddlewareHeaderMutation, MiddlewareMount, MiddlewareNextRequest, MiddlewareNextResponse, - MiddlewareRequestBody, MiddlewareRequestHead, MiddlewareResponseBody, - MiddlewareResponseHead, MiddlewareTransport, NEXT_METHOD, - }, - client::{ - CallFuture, CallReply, MiddlewarePlugin, PluginCall, PluginHandler, PluginSession, - RequestCall, ResponseStream, SessionConfig, SessionError, read_frame, write_frame, - }, -}; -use serde_json::{Value, json}; -use tokio::{ - io::{AsyncWriteExt, DuplexStream, ReadHalf, WriteHalf}, - task::JoinHandle, -}; - -const MAXIMUM_STREAM_CHUNK_BYTES: usize = 64 * 1024; - -#[tokio::test] -async fn binary_payload_round_trips_without_json_encoding() { - let frame = Frame { - message: Message::Stream { - id: 17, - sequence: 2, - }, - payload: vec![0, 255, 13, 10, 128], - }; - let mut bytes = Vec::new(); - write_frame(&mut bytes, &frame).await.unwrap(); - assert!(bytes.ends_with(&frame.payload)); - assert_eq!(read_frame(&mut bytes.as_slice()).await.unwrap(), frame); -} - -#[tokio::test] -async fn malicious_lengths_fail_before_reading_or_allocating_payload() { - let bytes = [0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0]; - assert!(matches!( - read_frame(&mut bytes.as_slice()).await, - Err(FrameError::Length) - )); -} - -#[tokio::test] -async fn large_payload_is_exact_across_multiple_io_buffers_and_following_frames() { - let frame = Frame { - message: Message::Result { - id: 1, - result: json!({}), - }, - payload: (0..101 * 1024 * 1024 + 13) - .map(|index| (index % 251) as u8) - .collect(), - }; - let next = Frame::control(Message::Cancel { id: 3 }); - let (mut writer, mut reader) = tokio::io::duplex(16 * 1024); - let writing = tokio::spawn(async move { - write_frame(&mut writer, &frame).await.unwrap(); - write_frame(&mut writer, &next).await.unwrap(); - (frame, next) - }); - let received = read_frame(&mut reader).await.unwrap(); - let following = read_frame(&mut reader).await.unwrap(); - let (original, next) = writing.await.unwrap(); - assert_eq!(received, original); - assert_eq!(following, next); -} - -#[tokio::test] -async fn truncated_large_payload_never_becomes_a_successful_message() { - let message = serde_json::to_vec(&Message::Result { - id: 1, - result: json!({}), - }) - .unwrap(); - let mut bytes = Vec::new(); - bytes.extend_from_slice(&(message.len() as u32).to_be_bytes()); - bytes.extend_from_slice(&(101_u64 * 1024 * 1024).to_be_bytes()); - bytes.extend_from_slice(&message); - bytes.extend_from_slice(&[1, 2, 3]); - assert!(matches!( - read_frame(&mut bytes.as_slice()).await, - Err(FrameError::Io(_)) - )); -} - -#[tokio::test] -async fn truncated_frame_never_becomes_a_successful_message() { - let mut bytes = Vec::new(); - write_frame(&mut bytes, &Frame::control(Message::Cancel { id: 3 })) - .await - .unwrap(); - bytes.pop(); - assert!(matches!( - read_frame(&mut bytes.as_slice()).await, - Err(FrameError::Io(_)) - )); -} - -#[derive(Clone, Default)] -struct TestHandler { - lifecycle: Arc, - cancel_delay: Option, -} - -#[derive(Default)] -struct LifecycleState { - cancel_started: AtomicBool, - cancelled: AtomicUsize, - quiesced: AtomicBool, - shutdown: AtomicBool, -} - -impl PluginHandler for TestHandler { - fn call(&self, call: PluginCall) -> CallFuture<'_> { - Box::pin(async move { - match call.method.as_str() { - "echo" => Ok(CallReply::unary(call.params, call.payload)), - "key_facts" => { - let result = call - .host - .key_facts(gateway_plugin_sdk::call::data::ClientKeyFactsQuery { - client_key_id: "key_1".into(), - }) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "get_budget" => { - let result = call - .host - .get_key_budget( - gateway_plugin_sdk::call::key_budgets::GetKeyBudgetRequest { - client_key_id: "key_1".into(), - }, - ) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "update_budget_limits" => { - let result = call - .host - .update_key_budget_limits( - gateway_plugin_sdk::call::key_budgets::UpdateKeyBudgetLimitsRequest { - client_key_id: "key_1".into(), - daily_limit_usd: None, - weekly_limit_usd: Some("12.5".into()), - }, - ) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "refresh_quota" => { - let result = call - .host - .refresh_account_quota(gateway_plugin_sdk::call::data::QuotaFactsQuery { - account_id: "acct_1".into(), - }) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "reset_budget" => { - use gateway_plugin_sdk::call::{ - host::KeyListRequest, - key_budgets::{BudgetPeriod, ResetKeyBudgetRequest}, - }; - let keys = call - .host - .list_keys(KeyListRequest { - cursor: None, - limit: 10, - }) - .await?; - let result = call - .host - .reset_key_budget(ResetKeyBudgetRequest { - client_key_id: keys.keys[0].id.clone(), - period: BudgetPeriod::Weekly, - }) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "slow" => { - tokio::time::sleep(Duration::from_millis(80)).await; - Ok(CallReply::unary(call.params, call.payload)) - } - "callback" => { - let reply = call - .host - .call("host.log", call.params, call.payload) - .await - .map_err(SessionError::into_plugin_fault)?; - Ok(CallReply::unary(reply.result, reply.payload)) - } - "resource_stream" => Ok(CallReply::stream( - json!({}), - Vec::new(), - ResponseStream::pull(Box::new(ResourceStream(Some(call.host)))), - )), - "stream" => Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - ResponseStream::from_chunks(vec![b"abcd".to_vec(), b"efgh".to_vec()]), - )), - "long_stream" => Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - ResponseStream::from_chunks(vec![b"x".to_vec(); 16]), - )), - "oversized_stream" => Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - ResponseStream::from_chunks(vec![b"abcde".to_vec()]), - )), - "too_many_chunks" => Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - ResponseStream::from_chunks(vec![b"x".to_vec(); 17]), - )), - "dynamic_oversized_stream" => { - let (sender, stream) = ResponseStream::channel(NonZeroUsize::new(1).unwrap()); - tokio::spawn(async move { - let _ = sender.send(b"abcde".to_vec()).await; - }); - Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - stream, - )) - } - "oversized_fault" => Err(PluginFault::new( - ErrorCode::Upstream, - "x".repeat(MAXIMUM_STREAM_CHUNK_BYTES), - )), - "oversized_result" => Ok(CallReply::unary( - Value::String("x".repeat(MAXIMUM_STREAM_CHUNK_BYTES)), - Vec::new(), - )), - "dynamic_oversized_fault" => { - let (sender, stream) = ResponseStream::channel(NonZeroUsize::new(1).unwrap()); - tokio::spawn(async move { - let _ = sender - .fail(PluginFault::new( - ErrorCode::Upstream, - "x".repeat(MAXIMUM_STREAM_CHUNK_BYTES), - )) - .await; - }); - Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - stream, - )) - } - "pending" => pending::>().await, - _ => Err(PluginFault::new( - ErrorCode::Unsupported, - "unsupported test method", - )), - } - }) - } - - fn cancel(&self, _context: &CallContext) { - self.lifecycle.cancel_started.store(true, Ordering::Release); - if let Some(delay) = self.cancel_delay { - std::thread::sleep(delay); - } - self.lifecycle.cancelled.fetch_add(1, Ordering::Relaxed); - } - - fn quiesce(&self) { - self.lifecycle.quiesced.store(true, Ordering::Release); - } - - fn shutdown(&self) { - self.lifecycle.shutdown.store(true, Ordering::Release); - } -} - -pub(super) struct HostPeer { - pub(super) reader: ReadHalf, - pub(super) writer: WriteHalf, -} - -pub(super) async fn start_session( - handler: H, -) -> (HostPeer, JoinHandle>) { - start_session_with_capacity(handler, MAXIMUM_STREAM_CHUNK_BYTES * 2).await -} - -async fn start_session_with_capacity( - handler: H, - capacity: usize, -) -> (HostPeer, JoinHandle>) { - let (host, plugin) = tokio::io::duplex(capacity); - let (plugin_reader, plugin_writer) = tokio::io::split(plugin); - let task = tokio::spawn(async move { - let session = PluginSession::accept( - plugin_reader, - plugin_writer, - SessionConfig { - maximum_stream_chunk_bytes: MAXIMUM_STREAM_CHUNK_BYTES, - maximum_calls: 4, - maximum_callbacks: 4, - maximum_buffered_stream_chunks: 16, - handshake_timeout: Duration::from_secs(1), - maximum_call_timeout: Duration::from_secs(2), - }, - ) - .await?; - session.run(handler).await - }); - let (mut reader, mut writer) = tokio::io::split(host); - write_frame( - &mut writer, - &Frame::control(Message::Hello { - handshake: handshake(), - }), - ) - .await - .unwrap(); - let ready = read_frame(&mut reader).await.unwrap(); - assert!(matches!( - ready, - Frame { - message: Message::Ready { - protocol_version: PROTOCOL_VERSION, - ref incarnation, - }, - ref payload, - } if incarnation == "test-incarnation" && payload.is_empty() - )); - (HostPeer { reader, writer }, task) -} - -fn middleware_contributions() -> Contributions { - Contributions::from([( - Capability::Middleware, - ContributionDeclaration { - id: "test.example.middleware".into(), - version: 4, - stages: vec![Stage::Request], - input_formats: vec!["openai".into()], - output_formats: vec!["openai".into()], - }, - )]) -} - -fn middleware_plugin(map_response: bool) -> impl PluginHandler { - MiddlewarePlugin::new( - &middleware_contributions(), - move |call: RequestCall| async move { - let RequestCall { - mut request, next, .. - } = call; - if map_response { - request.body.push(b'!'); - request.remove_header("x-old"); - request.head.headers.push(MiddlewareHeader { - name: "x-direct".into(), - value: b"request".to_vec(), - }); - } - let mut response = next.run(request).await?; - if map_response { - response.remove_header("x-upstream"); - response.headers.push(MiddlewareHeader { - name: "x-direct-response".into(), - value: b"response".to_vec(), - }); - response.body = response.body.map_frames(|mut frame| { - frame.payload.make_ascii_uppercase(); - Ok(vec![frame]) - })?; - } - Ok(response) - }, - ) - .unwrap() -} - -#[test] -fn middleware_plugin_rejects_declarations_without_matching_handlers() { - let declaration = middleware_contributions() - .remove(&Capability::Middleware) - .unwrap(); - let mut invalid = vec![Contributions::new()]; - invalid.push(Contributions::from([( - Capability::Scheduler, - declaration.clone(), - )])); - invalid.push(Contributions::from([ - (Capability::Middleware, declaration.clone()), - (Capability::Scheduler, declaration.clone()), - ])); - let mut unsupported = declaration.clone(); - unsupported.version = 2; - invalid.push(Contributions::from([(Capability::Middleware, unsupported)])); - for stages in [ - vec![], - vec![Stage::Registration], - vec![Stage::Request, Stage::Request], - vec![Stage::Request, Stage::Attempt, Stage::Request], - ] { - let mut invalid_stage = declaration.clone(); - invalid_stage.stages = stages; - invalid.push(Contributions::from([( - Capability::Middleware, - invalid_stage, - )])); - } - for contributes in invalid { - let result = MiddlewarePlugin::new(&contributes, |call: RequestCall| async move { - call.next.run(call.request).await - }); - assert!(matches!(result, Err(SessionError::Configuration))); - } - let mut both_mounts = declaration; - both_mounts.stages.push(Stage::Attempt); - let contributes = Contributions::from([(Capability::Middleware, both_mounts)]); - assert!( - MiddlewarePlugin::new(&contributes, |call: RequestCall| async move { - call.next.run(call.request).await - }) - .is_ok() - ); -} - -#[tokio::test] -async fn middleware_plugin_registers_without_invoking_business_handler() { - let contributes = middleware_contributions(); - let plugin = MiddlewarePlugin::new(&contributes, |_: RequestCall| async { - panic!("registration must not invoke middleware") - }) - .unwrap(); - let (mut host, task) = start_session(plugin).await; - let mut registration_context = context(1, Duration::from_secs(1)); - registration_context.stage = Stage::Registration; - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id: 1, - method: "plugin.register".into(), - context: registration_context, - params: json!({}), - }, - payload: vec![], - }, - ) - .await - .unwrap(); - let reply = receive(&mut host).await; - let Message::Result { id: 1, result } = reply.message else { - panic!("expected registration result") - }; - let registration: gateway_plugin_sdk::call::registration::Registration = - serde_json::from_value(result).unwrap(); - assert_eq!(registration.contributes, contributes); - assert!(reply.payload.is_empty()); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn middleware_plugin_rejects_invalid_calls_before_business_dispatch() { - let plugin = MiddlewarePlugin::new(&middleware_contributions(), |_: RequestCall| async { - panic!("invalid call must not invoke middleware") - }) - .unwrap(); - let (mut host, task) = start_session(plugin).await; - let cases = [ - ("plugin.register", Stage::Request, json!({}), vec![]), - ("plugin.register", Stage::Registration, json!(null), vec![]), - ( - "plugin.register", - Stage::Registration, - json!({"extra":true}), - vec![], - ), - ("plugin.register", Stage::Registration, json!({}), vec![1]), - (HANDLE_METHOD, Stage::Attempt, json!({}), vec![]), - (HANDLE_METHOD, Stage::Request, json!({}), vec![]), - ("unknown.execute", Stage::Request, json!({}), vec![]), - ]; - for (index, (method, stage, params, payload)) in cases.into_iter().enumerate() { - let id = u64::try_from(index).unwrap() * 2 + 1; - let mut call_context = middleware_context(id); - call_context.stage = stage; - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id, - method: method.into(), - context: call_context, - params, - }, - payload, - }, - ) - .await - .unwrap(); - let reply = receive(&mut host).await; - let Message::Error { - id: reply_id, - error, - } = reply.message - else { - panic!("expected invalid call error") - }; - assert_eq!(reply_id, id); - assert_eq!( - error.code, - if method == "unknown.execute" { - ErrorCode::Unsupported - } else { - ErrorCode::InvalidInput - } - ); - } - shutdown(&mut host, task).await; -} - -fn handshake() -> Handshake { - Handshake { - protocol_version: PROTOCOL_VERSION, - artifact_sha256: "a".repeat(64), - plugin_id: "test.example".into(), - instance_id: "test-instance".into(), - generation: 7, - incarnation: "test-incarnation".into(), - configuration: json!({}), - - contributes: Contributions::new(), - } -} - -pub(super) fn context(id: u64, timeout: Duration) -> CallContext { - CallContext { - call_id: id, - instance_id: "test-instance".into(), - generation: 7, - incarnation: "test-incarnation".into(), - stage: Stage::Request, - timeout_ms: timeout.as_millis() as u64, - resource_stream: false, - resource_scope_id: format!("scope-{id}"), - request_id: None, - attempt_id: None, - account_id: None, - credential_revision: None, - } -} - -fn middleware_context(id: u64) -> CallContext { - CallContext { - stage: Stage::Request, - request_id: Some("request-middleware".into()), - ..context(id, Duration::from_secs(1)) - } -} - -pub(super) fn middleware_request() -> MiddlewareRequestHead { - MiddlewareRequestHead { - settings_sources: serde_json::Value::Null, - settings: serde_json::Value::Null, - client_key_id: "fixture-key".into(), - account_group_ids: vec!["fixture-group".into()], - request_id: "request-middleware".into(), - mount: MiddlewareMount::Request, - attempt_index: None, - operation: "generate".into(), - protocol: "openai".into(), - endpoint: "responses".into(), - transport: MiddlewareTransport::HttpSse, - provider: None, - model: Some("public-model".into()), - account_id: None, - headers: vec![MiddlewareHeader { - name: "x-old".into(), - value: b"old".to_vec(), - }], - } -} - -pub(super) async fn send_call( - host: &mut HostPeer, - id: u64, - method: &str, - params: Value, - payload: Vec, -) { - send_call_with_timeout(host, id, method, params, payload, Duration::from_secs(1)).await; -} - -async fn send_call_with_timeout( - host: &mut HostPeer, - id: u64, - method: &str, - params: Value, - payload: Vec, - timeout: Duration, -) { - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id, - method: method.into(), - context: context(id, timeout), - params, - }, - payload, - }, - ) - .await - .unwrap(); -} - -async fn send_control(host: &mut HostPeer, message: Message) { - write_frame(&mut host.writer, &Frame::control(message)) - .await - .unwrap(); -} - -pub(super) async fn receive(host: &mut HostPeer) -> Frame { - tokio::time::timeout(Duration::from_secs(1), read_frame(&mut host.reader)) - .await - .expect("plugin response timed out") - .expect("plugin response frame must be valid") -} - -pub(super) async fn shutdown(host: &mut HostPeer, task: JoinHandle>) { - send_control(host, Message::Shutdown).await; - tokio::time::timeout(Duration::from_secs(1), task) - .await - .expect("plugin session did not shut down") - .expect("plugin session task panicked") - .expect("plugin session shutdown failed"); -} - -#[tokio::test] -async fn typed_budget_helpers_preserve_wire_contract_and_propagate_failure_without_retry() { - for failed in [false, true] { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "reset_budget", json!({}), vec![]).await; - let list = receive(&mut host).await; - let Message::Callback { - id, - parent_id: 1, - method, - params, - } = list.message - else { - panic!("expected Key list callback") - }; - assert_eq!(method, "host.keys.list"); - assert_eq!(params, json!({"cursor":null,"limit":10})); - assert!(list.payload.is_empty()); - send_control(&mut host, Message::Result { id, result: json!({"keys":[{"id":"key_1","name":"test","enabled":true}],"next_cursor":null}) }).await; - let reset = receive(&mut host).await; - let Message::Callback { - id, - parent_id: 1, - method, - params, - } = reset.message - else { - panic!("expected budget reset callback") - }; - assert_eq!(method, "host.keys.reset_budget"); - assert_eq!(params, json!({})); - assert_eq!( - serde_json::from_slice::(&reset.payload).unwrap(), - json!({"client_key_id":"key_1","period":"weekly"}) - ); - if failed { - send_control( - &mut host, - Message::Error { - id, - error: PluginFault::new(ErrorCode::Conflict, "stale instance"), - }, - ) - .await; - let Message::Error { id: 1, error } = receive(&mut host).await.message else { - panic!("expected error propagated without another callback") - }; - assert_eq!(error.code, ErrorCode::Conflict); - } else { - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id, - result: json!({}), - }, - payload: serde_json::to_vec(&json!({"client_key_id":"key_1"})).unwrap(), - }, - ) - .await - .unwrap(); - let Message::Result { id: 1, result } = receive(&mut host).await.message else { - panic!("expected successful reset") - }; - assert_eq!(result, json!({"client_key_id":"key_1"})); - } - shutdown(&mut host, task).await; - } -} - -#[tokio::test] -async fn independent_calls_complete_concurrently() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "slow", json!("slow"), vec![1]).await; - send_call(&mut host, 3, "echo", json!("fast"), vec![3]).await; - - let fast = receive(&mut host).await; - let slow = receive(&mut host).await; - assert!(matches!( - fast, - Frame { - message: Message::Result { id: 3, result }, - payload, - } if result == json!("fast") && payload == [3] - )); - assert!(matches!( - slow, - Frame { - message: Message::Result { id: 1, result }, - payload, - } if result == json!("slow") && payload == [1] - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn completed_call_does_not_discard_a_partially_read_next_frame() { - let next = Frame { - message: Message::Call { - id: 3, - method: "echo".into(), - context: context(3, Duration::from_secs(1)), - params: json!({"fragmented": true}), - }, - payload: vec![0, 255, 13, 10, 128], - }; - let mut encoded = Vec::new(); - write_frame(&mut encoded, &next).await.unwrap(); - let payload_start = encoded.len() - next.payload.len(); - for split in [2, 6, 10, payload_start + 2] { - tokio::time::timeout(Duration::from_secs(2), async { - // 单字节缓冲保证前缀已经进入读取 future，而不是仍滞留在传输队列 - let (mut host, task) = start_session_with_capacity(TestHandler::default(), 1).await; - send_call(&mut host, 1, "slow", json!("first"), Vec::new()).await; - host.writer.write_all(&encoded[..split]).await.unwrap(); - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - host.writer.write_all(&encoded[split..]).await.unwrap(); - assert_eq!( - receive(&mut host).await, - Frame { - message: Message::Result { - id: 3, - result: json!({"fragmented": true}) - }, - payload: next.payload.clone(), - } - ); - shutdown(&mut host, task).await; - }) - .await - .expect("fragmented frame must survive concurrent call completion"); - } -} - -#[tokio::test] -async fn host_callback_round_trips_while_parent_call_is_waiting() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "callback", json!({"event": "ready"}), vec![7]).await; - - let callback = receive(&mut host).await; - let callback_id = match callback.message { - Message::Callback { - id, - parent_id: 1, - ref method, - ref params, - } if id.is_multiple_of(2) - && method == "host.log" - && params == &json!({"event": "ready"}) => - { - id - } - message => panic!("unexpected callback frame: {message:?}"), - }; - assert_eq!(callback.payload, [7]); - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: callback_id, - result: json!({"recorded": true}), - }, - payload: vec![9], - }, - ) - .await - .unwrap(); - - let result = receive(&mut host).await; - assert!(matches!( - result, - Frame { - message: Message::Result { id: 1, result }, - payload, - } if result == json!({"recorded": true}) && payload == [9] - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn callbacks_from_concurrent_calls_remain_correlated() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "callback", json!({"parent": 1}), vec![1]).await; - send_call(&mut host, 3, "callback", json!({"parent": 3}), vec![3]).await; - - let mut callbacks = Vec::new(); - for _ in 0..2 { - let frame = receive(&mut host).await; - let Message::Callback { - id, - parent_id, - method, - params, - } = frame.message - else { - panic!("expected callback") - }; - assert_eq!(method, "host.log"); - assert_eq!(params, json!({"parent": parent_id})); - assert_eq!(frame.payload, [parent_id as u8]); - callbacks.push((id, parent_id)); - } - assert!(callbacks[0].0 < callbacks[1].0); - assert!(callbacks.iter().all(|(id, _)| id.is_multiple_of(2))); - - for (callback_id, parent_id) in callbacks.into_iter().rev() { - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: callback_id, - result: json!({"parent": parent_id}), - }, - payload: vec![parent_id as u8], - }, - ) - .await - .unwrap(); - } - let mut completed = Vec::new(); - for _ in 0..2 { - let frame = receive(&mut host).await; - let Message::Result { id, result } = frame.message else { - panic!("expected parent result") - }; - assert_eq!(result, json!({"parent": id})); - assert_eq!(frame.payload, [id as u8]); - completed.push(id); - } - completed.sort_unstable(); - assert_eq!(completed, [1, 3]); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn middleware_next_and_body_mapping_reuse_callback_and_credit_flow() { - let (mut host, task) = start_session(middleware_plugin(true)).await; - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id: 1, - method: HANDLE_METHOD.into(), - context: middleware_context(1), - params: serde_json::to_value(middleware_request()).unwrap(), - }, - payload: b"request".to_vec(), - }, - ) - .await - .unwrap(); - - let next = receive(&mut host).await; - let next_id = match next.message { - Message::Callback { - id, - parent_id: 1, - ref method, - ref params, - } if method == NEXT_METHOD => { - let request: MiddlewareNextRequest = serde_json::from_value(params.clone()).unwrap(); - assert_eq!(request.body, MiddlewareRequestBody::Replace); - assert!(request.protocol.is_none()); - assert_eq!(request.header_mutations.len(), 2); - assert!(matches!( - &request.header_mutations[0], - MiddlewareHeaderMutation::Remove { name } if name == "x-old" - )); - assert!(matches!( - &request.header_mutations[1], - MiddlewareHeaderMutation::Append { name, value } - if name == "x-direct" && value == b"request" - )); - id - } - message => panic!("unexpected middleware next frame: {message:?}"), - }; - assert_eq!(next.payload, b"request!"); - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: next_id, - result: serde_json::to_value(MiddlewareNextResponse { - metadata: None, - response: "response-1".into(), - protocol: "openai".into(), - status: 200, - headers: vec![MiddlewareHeader { - name: "x-upstream".into(), - value: b"old".to_vec(), - }], - body: Some(MiddlewareBodyHandle { - handle: "body-1".into(), - framing: MiddlewareBodyFraming::SseEvent, - }), - }) - .unwrap(), - }, - payload: Vec::new(), - }, - ) - .await - .unwrap(); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1_024, - frames: 2, - }, - ) - .await; - - let initial = receive(&mut host).await; - let Message::Result { id: 1, result } = initial.message else { - panic!("expected middleware result") - }; - let head: MiddlewareResponseHead = serde_json::from_value(result).unwrap(); - assert_eq!(head.header_mutations.len(), 2); - assert!(matches!( - &head.header_mutations[0], - MiddlewareHeaderMutation::Remove { name } if name == "x-upstream" - )); - assert!(matches!( - &head.header_mutations[1], - MiddlewareHeaderMutation::Append { name, value } - if name == "x-direct-response" && value == b"response" - )); - assert!(matches!( - head.body, - MiddlewareResponseBody::Stream { - framing: MiddlewareBodyFraming::SseEvent - } - )); - - let read = receive(&mut host).await; - let read_id = match read.message { - Message::Callback { - id, - parent_id: 1, - ref method, - .. - } if method == BODY_READ_METHOD => id, - message => panic!("unexpected middleware body read: {message:?}"), - }; - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: read_id, - result: serde_json::to_value(MiddlewareBodyReadResult { - framing: MiddlewareBodyFraming::SseEvent, - source_id: 1, - eof: false, - terminal: true, - }) - .unwrap(), - }, - payload: b"data: done\n\n".to_vec(), - }, - ) - .await - .unwrap(); - let mapped = receive(&mut host).await; - let Frame { - message: Message::Stream { id: 1, sequence: 0 }, - payload, - } = mapped - else { - panic!("expected mapped middleware frame") - }; - let mapped = MiddlewareBodyFrame::decode(&payload).unwrap(); - assert_eq!(mapped.payload, b"DATA: DONE\n\n"); - assert!(!mapped.terminal); - assert_eq!(mapped.source_id(), 1); - assert_eq!(mapped.disposition(), MiddlewareBodyDisposition::Only); - - let eof = receive(&mut host).await; - let eof_id = match eof.message { - Message::Callback { - id, - parent_id: 1, - ref method, - .. - } if method == BODY_READ_METHOD => id, - message => panic!("unexpected middleware eof read: {message:?}"), - }; - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: eof_id, - result: serde_json::to_value(MiddlewareBodyReadResult { - framing: MiddlewareBodyFraming::SseEvent, - source_id: 0, - eof: true, - terminal: false, - }) - .unwrap(), - }, - payload: Vec::new(), - }, - ) - .await - .unwrap(); - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn untouched_middleware_response_transfers_opaque_body_without_reading_it() { - let (mut host, task) = start_session(middleware_plugin(false)).await; - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id: 1, - method: HANDLE_METHOD.into(), - context: middleware_context(1), - params: serde_json::to_value(middleware_request()).unwrap(), - }, - payload: b"hidden-by-preserve".to_vec(), - }, - ) - .await - .unwrap(); - - let next = receive(&mut host).await; - let next_id = match next.message { - Message::Callback { - id, - parent_id: 1, - ref method, - ref params, - } if method == NEXT_METHOD => { - let request: MiddlewareNextRequest = serde_json::from_value(params.clone()).unwrap(); - assert_eq!(request.body, MiddlewareRequestBody::Preserve); - id - } - message => panic!("unexpected middleware next frame: {message:?}"), - }; - assert!(next.payload.is_empty()); - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: next_id, - result: serde_json::to_value(MiddlewareNextResponse { - metadata: None, - response: "response-1".into(), - protocol: "openai".into(), - status: 200, - headers: Vec::new(), - body: Some(MiddlewareBodyHandle { - handle: "body-pass-through".into(), - framing: MiddlewareBodyFraming::RawBytes, - }), - }) - .unwrap(), - }, - payload: Vec::new(), - }, - ) - .await - .unwrap(); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1, - frames: 1, - }, - ) - .await; - - let initial = receive(&mut host).await; - let Message::Result { id: 1, result } = initial.message else { - panic!("expected middleware result") - }; - let head: MiddlewareResponseHead = serde_json::from_value(result).unwrap(); - assert!(matches!( - head.body, - MiddlewareResponseBody::PassThrough { - body: MiddlewareBodyHandle { ref handle, .. } - } if handle == "body-pass-through" - )); - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn stream_waits_for_credit_and_rejects_uncreditable_buffered_chunk_before_result() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "stream", json!({}), Vec::new()).await; - assert!( - tokio::time::timeout(Duration::from_millis(25), read_frame(&mut host.reader)) - .await - .is_err() - ); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 4, - frames: 1, - }, - ) - .await; - - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - assert!(matches!( - receive(&mut host).await, - Frame { - message: Message::Stream { id: 1, sequence: 0 }, - payload, - } if payload == b"abcd" - )); - assert!( - tokio::time::timeout(Duration::from_millis(25), read_frame(&mut host.reader)) - .await - .is_err() - ); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 4, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await, - Frame { - message: Message::Stream { id: 1, sequence: 1 }, - payload, - } if payload == b"efgh" - )); - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - - send_call(&mut host, 3, "oversized_stream", json!({}), Vec::new()).await; - send_control( - &mut host, - Message::Credit { - id: 3, - bytes: 4, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 3, - error: PluginFault { - code: ErrorCode::InvalidInput, - .. - }, - } - )); - - send_call(&mut host, 5, "too_many_chunks", json!({}), Vec::new()).await; - send_control( - &mut host, - Message::Credit { - id: 5, - bytes: 4, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 5, - error: PluginFault { - code: ErrorCode::Capacity, - .. - }, - } - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn late_credit_after_a_long_stream_terminal_does_not_close_the_session() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "long_stream", json!({}), Vec::new()).await; - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - for sequence in 0..16 { - assert!(matches!( - receive(&mut host).await, - Frame { - message: Message::Stream { id: 1, sequence: observed }, - payload, - } if observed == sequence && payload == b"x" - )); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1, - frames: 1, - }, - ) - .await; - } - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - - send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 3, result } if result == json!("still-ready") - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn malformed_late_credit_remains_a_protocol_error() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "echo", json!({}), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 0, - frames: 1, - }, - ) - .await; - assert!(matches!( - tokio::time::timeout(Duration::from_secs(1), task) - .await - .expect("protocol error did not close the session") - .expect("plugin session task panicked"), - Err(SessionError::Protocol) - )); -} - -#[tokio::test] -async fn dynamic_stream_ends_with_one_error_when_a_later_chunk_cannot_fit_credit() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call( - &mut host, - 1, - "dynamic_oversized_stream", - json!({}), - Vec::new(), - ) - .await; - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 4, - frames: 1, - }, - ) - .await; - - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - assert!(matches!( - receive(&mut host).await.message, - Message::End { - id: 1, - error: Some(PluginFault { - code: ErrorCode::InvalidInput, - .. - }), - } - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn oversized_handler_faults_are_replaced_without_closing_the_session() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "oversized_result", json!({}), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 1, - error: PluginFault { - code: ErrorCode::InvalidInput, - .. - }, - } - )); - - send_call(&mut host, 3, "oversized_fault", json!({}), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 3, - error: PluginFault { - code: ErrorCode::Fault, - ref message, - .. - }, - } if message == "plugin error metadata exceeds its limit" - )); - - send_call( - &mut host, - 5, - "dynamic_oversized_fault", - json!({}), - Vec::new(), - ) - .await; - send_control( - &mut host, - Message::Credit { - id: 5, - bytes: 1_024, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 5, .. } - )); - assert!(matches!( - receive(&mut host).await.message, - Message::End { - id: 5, - error: Some(PluginFault { - code: ErrorCode::Fault, - ref message, - .. - }), - } if message == "plugin stream error metadata exceeds its limit" - )); - - send_call(&mut host, 7, "echo", json!("still-ready"), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 7, result } if result == json!("still-ready") - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn call_deadline_cancels_handler_and_keeps_the_session_available() { - let handler = TestHandler::default(); - let state = Arc::clone(&handler.lifecycle); - let (mut host, task) = start_session(handler).await; - send_call_with_timeout( - &mut host, - 1, - "pending", - json!({}), - Vec::new(), - Duration::from_millis(20), - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 1, - error: PluginFault { - code: ErrorCode::Timeout, - .. - }, - } - )); - tokio::time::timeout(Duration::from_millis(100), async { - while !state.cancel_started.load(Ordering::Acquire) { - tokio::task::yield_now().await; - } - }) - .await - .expect("deadline did not notify the lifecycle hook"); - - send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 3, result } if result == json!("still-ready") - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn cancelling_parent_retires_callback_and_late_reply_does_not_close_session() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "callback", json!({}), Vec::new()).await; - let callback = receive(&mut host).await; - let Message::Callback { - id: callback_id, .. - } = callback.message - else { - panic!("expected callback") - }; - - send_control(&mut host, Message::Cancel { id: 1 }).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Cancelled { id: 1 } - )); - send_control( - &mut host, - Message::Result { - id: callback_id, - result: json!({}), - }, - ) - .await; - send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 3, result } if result == json!("still-ready") - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn blocking_cancel_hook_cannot_delay_cancelled_or_eof_settlement() { - let handler = TestHandler { - cancel_delay: Some(Duration::from_millis(200)), - ..TestHandler::default() - }; - let state = Arc::clone(&handler.lifecycle); - let (mut host, task) = start_session(handler).await; - send_call(&mut host, 1, "pending", json!({}), Vec::new()).await; - send_control(&mut host, Message::Cancel { id: 1 }).await; - - let cancelled = tokio::time::timeout(Duration::from_millis(50), receive(&mut host)) - .await - .expect("blocking lifecycle hook delayed Cancelled"); - assert!(matches!(cancelled.message, Message::Cancelled { id: 1 })); - tokio::time::timeout(Duration::from_millis(50), async { - while !state.cancel_started.load(Ordering::Acquire) { - tokio::task::yield_now().await; - } - }) - .await - .expect("cancel hook never started"); - - host.writer.shutdown().await.unwrap(); - let result = tokio::time::timeout(Duration::from_millis(50), task) - .await - .expect("blocking lifecycle hook delayed EOF settlement") - .expect("plugin session task panicked"); - assert!(matches!(result, Err(SessionError::Closed))); -} - -#[tokio::test] -async fn quiesce_rejects_new_calls_and_shutdown_closes_the_session() { - let handler = TestHandler::default(); - let state = Arc::clone(&handler.lifecycle); - let (mut host, task) = start_session(handler).await; - send_control(&mut host, Message::Quiesce).await; - send_call(&mut host, 1, "echo", json!({}), Vec::new()).await; - - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 1, - error: PluginFault { - code: ErrorCode::Capacity, - .. - }, - } - )); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1, - frames: 1, - }, - ) - .await; - tokio::time::timeout(Duration::from_millis(50), async { - while !state.quiesced.load(Ordering::Acquire) { - tokio::task::yield_now().await; - } - }) - .await - .expect("quiesce hook was not signalled"); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn typed_key_and_quota_calls_keep_payloads_and_do_not_retry_failures() { - for (entry, method, request, response) in [ - ( - "key_facts", - "host.data.keys.get", - json!({"client_key_id":"key_1"}), - json!({"schema_version":1,"client_key_id":"key_1","enabled":false,"group_ids":["grp_1"]}), - ), - ( - "get_budget", - "host.keys.get_budget", - json!({"client_key_id":"key_1"}), - json!({ - "client_key_id":"key_1", "daily_limit_usd":"10", "weekly_limit_usd":"20", - "daily_used_usd":"1.25", "weekly_used_usd":"2.5", "daily_resets_at_ms":null, "weekly_resets_at_ms":null, - }), - ), - ( - "update_budget_limits", - "host.keys.update_budget_limits", - json!({"client_key_id":"key_1","weekly_limit_usd":"12.5"}), - json!({"client_key_id":"key_1"}), - ), - ( - "refresh_quota", - "host.quota_observations.refresh", - json!({"account_id":"acct_1"}), - json!({"schema_version":1,"account_id":"acct_1","observed_at_ms":null,"windows":[]}), - ), - ] { - for failed in [false, true] { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, entry, json!({}), vec![]).await; - let callback = receive(&mut host).await; - let Message::Callback { - id, - parent_id: 1, - method: actual, - params, - } = callback.message - else { - panic!("expected callback"); - }; - assert_eq!(actual, method); - assert_eq!(params, json!({})); - assert_eq!( - serde_json::from_slice::(&callback.payload).unwrap(), - request - ); - if failed { - send_control( - &mut host, - Message::Error { - id, - error: PluginFault::new(ErrorCode::Conflict, "changed"), - }, - ) - .await; - let Message::Error { id: 1, error } = receive(&mut host).await.message else { - panic!("expected propagated error without retry"); - }; - assert_eq!(error.code, ErrorCode::Conflict); - } else { - let mut response_payload = response.clone(); - if matches!(entry, "key_facts" | "refresh_quota") { - response_payload["future_fact"] = json!({"value":1}); - } - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id, - result: json!({}), - }, - payload: serde_json::to_vec(&response_payload).unwrap(), - }, - ) - .await - .unwrap(); - let Message::Result { id: 1, result } = receive(&mut host).await.message else { - panic!("expected projected response"); - }; - assert_eq!(result, response); - } - shutdown(&mut host, task).await; - } - } -} - -struct ResourceStream(Option); -impl gateway_plugin_sdk::client::PullResponseStream for ResourceStream { - fn next(&mut self) -> gateway_plugin_sdk::client::PullResponseFuture<'_> { - Box::pin(async move { - let host = self.0.take()?; - tokio::time::sleep(Duration::from_millis(700)).await; - Some( - host.call("host.log", json!({}), Vec::new()) - .await - .map(|_| b"live".to_vec()) - .map_err(SessionError::into_plugin_fault), - ) - }) - } -} - -#[tokio::test] -async fn resource_stream_keeps_callbacks_and_credit_alive_after_the_initial_deadline() { - let (mut host, task) = start_session(TestHandler::default()).await; - let mut context = context(1, Duration::from_millis(500)); - context.resource_stream = true; - write_frame( - &mut host.writer, - &Frame::control(Message::Call { - id: 1, - method: "resource_stream".into(), - context, - params: json!({}), - }), - ) - .await - .unwrap(); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 16, - frames: 2, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - let Message::Callback { - id, parent_id: 1, .. - } = receive(&mut host).await.message - else { - panic!("stream callback must survive initial timeout"); - }; - send_control( - &mut host, - Message::Result { - id, - result: json!({}), - }, - ) - .await; - assert!( - matches!(receive(&mut host).await, Frame { message: Message::Stream { id: 1, sequence: 0 }, payload } if payload == b"live") - ); - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - shutdown(&mut host, task).await; -} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/credit.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/credit.rs new file mode 100644 index 000000000..72af4f63e --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/credit.rs @@ -0,0 +1,271 @@ +//! 流信用等待、迟到信用、分块上限与资源流期限测试 + +use super::*; + +#[tokio::test] +async fn stream_waits_for_credit_and_rejects_uncreditable_buffered_chunk_before_result() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "stream", json!({}), Vec::new()).await; + assert!( + tokio::time::timeout(Duration::from_millis(25), read_frame(&mut host.reader)) + .await + .is_err() + ); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 4, + frames: 1, + }, + ) + .await; + + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + assert!(matches!( + receive(&mut host).await, + Frame { + message: Message::Stream { id: 1, sequence: 0 }, + payload, + } if payload == b"abcd" + )); + assert!( + tokio::time::timeout(Duration::from_millis(25), read_frame(&mut host.reader)) + .await + .is_err() + ); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 4, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await, + Frame { + message: Message::Stream { id: 1, sequence: 1 }, + payload, + } if payload == b"efgh" + )); + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + + send_call(&mut host, 3, "oversized_stream", json!({}), Vec::new()).await; + send_control( + &mut host, + Message::Credit { + id: 3, + bytes: 4, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 3, + error: PluginFault { + code: ErrorCode::InvalidInput, + .. + }, + } + )); + + send_call(&mut host, 5, "too_many_chunks", json!({}), Vec::new()).await; + send_control( + &mut host, + Message::Credit { + id: 5, + bytes: 4, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 5, + error: PluginFault { + code: ErrorCode::Capacity, + .. + }, + } + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn late_credit_after_a_long_stream_terminal_does_not_close_the_session() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "long_stream", json!({}), Vec::new()).await; + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + for sequence in 0..16 { + assert!(matches!( + receive(&mut host).await, + Frame { + message: Message::Stream { id: 1, sequence: observed }, + payload, + } if observed == sequence && payload == b"x" + )); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1, + frames: 1, + }, + ) + .await; + } + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + + send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 3, result } if result == json!("still-ready") + )); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn malformed_late_credit_remains_a_protocol_error() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "echo", json!({}), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 0, + frames: 1, + }, + ) + .await; + assert!(matches!( + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("protocol error did not close the session") + .expect("plugin session task panicked"), + Err(SessionError::Protocol) + )); +} + +#[tokio::test] +async fn dynamic_stream_ends_with_one_error_when_a_later_chunk_cannot_fit_credit() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call( + &mut host, + 1, + "dynamic_oversized_stream", + json!({}), + Vec::new(), + ) + .await; + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 4, + frames: 1, + }, + ) + .await; + + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + assert!(matches!( + receive(&mut host).await.message, + Message::End { + id: 1, + error: Some(PluginFault { + code: ErrorCode::InvalidInput, + .. + }), + } + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn resource_stream_keeps_callbacks_and_credit_alive_after_the_initial_deadline() { + let (mut host, task) = start_session(TestHandler::default()).await; + let mut context = context(1, Duration::from_millis(500)); + context.resource_stream = true; + write_frame( + &mut host.writer, + &Frame::control(Message::Call { + id: 1, + method: "resource_stream".into(), + context, + params: json!({}), + }), + ) + .await + .unwrap(); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 16, + frames: 2, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + let Message::Callback { + id, parent_id: 1, .. + } = receive(&mut host).await.message + else { + panic!("stream callback must survive initial timeout"); + }; + send_control( + &mut host, + Message::Result { + id, + result: json!({}), + }, + ) + .await; + assert!( + matches!(receive(&mut host).await, Frame { message: Message::Stream { id: 1, sequence: 0 }, payload } if payload == b"live") + ); + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + shutdown(&mut host, task).await; +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs new file mode 100644 index 000000000..06a3acff6 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs @@ -0,0 +1,545 @@ +//! 会话调用分派、并发帧读取与处理器响应边界测试 + +use super::*; + +#[test] +fn middleware_plugin_rejects_declarations_without_matching_handlers() { + let declaration = middleware_contributions() + .remove(&Capability::Middleware) + .unwrap(); + let mut invalid = vec![Contributions::new()]; + invalid.push(Contributions::from([( + Capability::Scheduler, + declaration.clone(), + )])); + invalid.push(Contributions::from([ + (Capability::Middleware, declaration.clone()), + (Capability::Scheduler, declaration.clone()), + ])); + let mut unsupported = declaration.clone(); + unsupported.version = 2; + invalid.push(Contributions::from([(Capability::Middleware, unsupported)])); + for stages in [ + vec![], + vec![Stage::Registration], + vec![Stage::Request, Stage::Request], + vec![Stage::Request, Stage::Attempt, Stage::Request], + ] { + let mut invalid_stage = declaration.clone(); + invalid_stage.stages = stages; + invalid.push(Contributions::from([( + Capability::Middleware, + invalid_stage, + )])); + } + for contributes in invalid { + let result = MiddlewarePlugin::new(&contributes, |call: RequestCall| async move { + call.next.run(call.request).await + }); + assert!(matches!(result, Err(SessionError::Configuration))); + } + let mut both_mounts = declaration; + both_mounts.stages.push(Stage::Attempt); + let contributes = Contributions::from([(Capability::Middleware, both_mounts)]); + assert!( + MiddlewarePlugin::new(&contributes, |call: RequestCall| async move { + call.next.run(call.request).await + }) + .is_ok() + ); +} + +#[tokio::test] +async fn middleware_plugin_registers_without_invoking_business_handler() { + let contributes = middleware_contributions(); + let plugin = MiddlewarePlugin::new(&contributes, |_: RequestCall| async { + panic!("registration must not invoke middleware") + }) + .unwrap(); + let (mut host, task) = start_session(plugin).await; + let mut registration_context = context(1, Duration::from_secs(1)); + registration_context.stage = Stage::Registration; + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id: 1, + method: "plugin.register".into(), + context: registration_context, + params: json!({}), + }, + payload: vec![], + }, + ) + .await + .unwrap(); + let reply = receive(&mut host).await; + let Message::Result { id: 1, result } = reply.message else { + panic!("expected registration result") + }; + let registration: gateway_plugin_sdk::call::registration::Registration = + serde_json::from_value(result).unwrap(); + assert_eq!(registration.contributes, contributes); + assert!(reply.payload.is_empty()); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn middleware_plugin_rejects_invalid_calls_before_business_dispatch() { + let plugin = MiddlewarePlugin::new(&middleware_contributions(), |_: RequestCall| async { + panic!("invalid call must not invoke middleware") + }) + .unwrap(); + let (mut host, task) = start_session(plugin).await; + let cases = [ + ("plugin.register", Stage::Request, json!({}), vec![]), + ("plugin.register", Stage::Registration, json!(null), vec![]), + ( + "plugin.register", + Stage::Registration, + json!({"extra":true}), + vec![], + ), + ("plugin.register", Stage::Registration, json!({}), vec![1]), + (HANDLE_METHOD, Stage::Attempt, json!({}), vec![]), + (HANDLE_METHOD, Stage::Request, json!({}), vec![]), + ("unknown.execute", Stage::Request, json!({}), vec![]), + ]; + for (index, (method, stage, params, payload)) in cases.into_iter().enumerate() { + let id = u64::try_from(index).unwrap() * 2 + 1; + let mut call_context = middleware_context(id); + call_context.stage = stage; + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id, + method: method.into(), + context: call_context, + params, + }, + payload, + }, + ) + .await + .unwrap(); + let reply = receive(&mut host).await; + let Message::Error { + id: reply_id, + error, + } = reply.message + else { + panic!("expected invalid call error") + }; + assert_eq!(reply_id, id); + assert_eq!( + error.code, + if method == "unknown.execute" { + ErrorCode::Unsupported + } else { + ErrorCode::InvalidInput + } + ); + } + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn independent_calls_complete_concurrently() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "slow", json!("slow"), vec![1]).await; + send_call(&mut host, 3, "echo", json!("fast"), vec![3]).await; + + let fast = receive(&mut host).await; + let slow = receive(&mut host).await; + assert!(matches!( + fast, + Frame { + message: Message::Result { id: 3, result }, + payload, + } if result == json!("fast") && payload == [3] + )); + assert!(matches!( + slow, + Frame { + message: Message::Result { id: 1, result }, + payload, + } if result == json!("slow") && payload == [1] + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn completed_call_does_not_discard_a_partially_read_next_frame() { + let next = Frame { + message: Message::Call { + id: 3, + method: "echo".into(), + context: context(3, Duration::from_secs(1)), + params: json!({"fragmented": true}), + }, + payload: vec![0, 255, 13, 10, 128], + }; + let mut encoded = Vec::new(); + write_frame(&mut encoded, &next).await.unwrap(); + let payload_start = encoded.len() - next.payload.len(); + for split in [2, 6, 10, payload_start + 2] { + tokio::time::timeout(Duration::from_secs(2), async { + // 单字节缓冲保证前缀已经进入读取 future，而不是仍滞留在传输队列 + let (mut host, task) = start_session_with_capacity(TestHandler::default(), 1).await; + send_call(&mut host, 1, "slow", json!("first"), Vec::new()).await; + host.writer.write_all(&encoded[..split]).await.unwrap(); + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + host.writer.write_all(&encoded[split..]).await.unwrap(); + assert_eq!( + receive(&mut host).await, + Frame { + message: Message::Result { + id: 3, + result: json!({"fragmented": true}) + }, + payload: next.payload.clone(), + } + ); + shutdown(&mut host, task).await; + }) + .await + .expect("fragmented frame must survive concurrent call completion"); + } +} + +#[tokio::test] +async fn middleware_next_and_body_mapping_reuse_callback_and_credit_flow() { + let (mut host, task) = start_session(middleware_plugin(true)).await; + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id: 1, + method: HANDLE_METHOD.into(), + context: middleware_context(1), + params: serde_json::to_value(middleware_request()).unwrap(), + }, + payload: b"request".to_vec(), + }, + ) + .await + .unwrap(); + + let next = receive(&mut host).await; + let next_id = match next.message { + Message::Callback { + id, + parent_id: 1, + ref method, + ref params, + } if method == NEXT_METHOD => { + let request: MiddlewareNextRequest = serde_json::from_value(params.clone()).unwrap(); + assert_eq!(request.body, MiddlewareRequestBody::Replace); + assert!(request.protocol.is_none()); + assert_eq!(request.header_mutations.len(), 2); + assert!(matches!( + &request.header_mutations[0], + MiddlewareHeaderMutation::Remove { name } if name == "x-old" + )); + assert!(matches!( + &request.header_mutations[1], + MiddlewareHeaderMutation::Append { name, value } + if name == "x-direct" && value == b"request" + )); + id + } + message => panic!("unexpected middleware next frame: {message:?}"), + }; + assert_eq!(next.payload, b"request!"); + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: next_id, + result: serde_json::to_value(MiddlewareNextResponse { + metadata: None, + response: "response-1".into(), + protocol: "openai".into(), + status: 200, + headers: vec![MiddlewareHeader { + name: "x-upstream".into(), + value: b"old".to_vec(), + }], + body: Some(MiddlewareBodyHandle { + handle: "body-1".into(), + framing: MiddlewareBodyFraming::SseEvent, + }), + }) + .unwrap(), + }, + payload: Vec::new(), + }, + ) + .await + .unwrap(); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1_024, + frames: 2, + }, + ) + .await; + + let initial = receive(&mut host).await; + let Message::Result { id: 1, result } = initial.message else { + panic!("expected middleware result") + }; + let head: MiddlewareResponseHead = serde_json::from_value(result).unwrap(); + assert_eq!(head.header_mutations.len(), 2); + assert!(matches!( + &head.header_mutations[0], + MiddlewareHeaderMutation::Remove { name } if name == "x-upstream" + )); + assert!(matches!( + &head.header_mutations[1], + MiddlewareHeaderMutation::Append { name, value } + if name == "x-direct-response" && value == b"response" + )); + assert!(matches!( + head.body, + MiddlewareResponseBody::Stream { + framing: MiddlewareBodyFraming::SseEvent + } + )); + + let read = receive(&mut host).await; + let read_id = match read.message { + Message::Callback { + id, + parent_id: 1, + ref method, + .. + } if method == BODY_READ_METHOD => id, + message => panic!("unexpected middleware body read: {message:?}"), + }; + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: read_id, + result: serde_json::to_value(MiddlewareBodyReadResult { + framing: MiddlewareBodyFraming::SseEvent, + source_id: 1, + eof: false, + terminal: true, + }) + .unwrap(), + }, + payload: b"data: done\n\n".to_vec(), + }, + ) + .await + .unwrap(); + let mapped = receive(&mut host).await; + let Frame { + message: Message::Stream { id: 1, sequence: 0 }, + payload, + } = mapped + else { + panic!("expected mapped middleware frame") + }; + let mapped = MiddlewareBodyFrame::decode(&payload).unwrap(); + assert_eq!(mapped.payload, b"DATA: DONE\n\n"); + assert!(!mapped.terminal); + assert_eq!(mapped.source_id(), 1); + assert_eq!(mapped.disposition(), MiddlewareBodyDisposition::Only); + + let eof = receive(&mut host).await; + let eof_id = match eof.message { + Message::Callback { + id, + parent_id: 1, + ref method, + .. + } if method == BODY_READ_METHOD => id, + message => panic!("unexpected middleware eof read: {message:?}"), + }; + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: eof_id, + result: serde_json::to_value(MiddlewareBodyReadResult { + framing: MiddlewareBodyFraming::SseEvent, + source_id: 0, + eof: true, + terminal: false, + }) + .unwrap(), + }, + payload: Vec::new(), + }, + ) + .await + .unwrap(); + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn untouched_middleware_response_transfers_opaque_body_without_reading_it() { + let (mut host, task) = start_session(middleware_plugin(false)).await; + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id: 1, + method: HANDLE_METHOD.into(), + context: middleware_context(1), + params: serde_json::to_value(middleware_request()).unwrap(), + }, + payload: b"hidden-by-preserve".to_vec(), + }, + ) + .await + .unwrap(); + + let next = receive(&mut host).await; + let next_id = match next.message { + Message::Callback { + id, + parent_id: 1, + ref method, + ref params, + } if method == NEXT_METHOD => { + let request: MiddlewareNextRequest = serde_json::from_value(params.clone()).unwrap(); + assert_eq!(request.body, MiddlewareRequestBody::Preserve); + id + } + message => panic!("unexpected middleware next frame: {message:?}"), + }; + assert!(next.payload.is_empty()); + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: next_id, + result: serde_json::to_value(MiddlewareNextResponse { + metadata: None, + response: "response-1".into(), + protocol: "openai".into(), + status: 200, + headers: Vec::new(), + body: Some(MiddlewareBodyHandle { + handle: "body-pass-through".into(), + framing: MiddlewareBodyFraming::RawBytes, + }), + }) + .unwrap(), + }, + payload: Vec::new(), + }, + ) + .await + .unwrap(); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1, + frames: 1, + }, + ) + .await; + + let initial = receive(&mut host).await; + let Message::Result { id: 1, result } = initial.message else { + panic!("expected middleware result") + }; + let head: MiddlewareResponseHead = serde_json::from_value(result).unwrap(); + assert!(matches!( + head.body, + MiddlewareResponseBody::PassThrough { + body: MiddlewareBodyHandle { ref handle, .. } + } if handle == "body-pass-through" + )); + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn oversized_handler_faults_are_replaced_without_closing_the_session() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "oversized_result", json!({}), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 1, + error: PluginFault { + code: ErrorCode::InvalidInput, + .. + }, + } + )); + + send_call(&mut host, 3, "oversized_fault", json!({}), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 3, + error: PluginFault { + code: ErrorCode::Fault, + ref message, + .. + }, + } if message == "plugin error metadata exceeds its limit" + )); + + send_call( + &mut host, + 5, + "dynamic_oversized_fault", + json!({}), + Vec::new(), + ) + .await; + send_control( + &mut host, + Message::Credit { + id: 5, + bytes: 1_024, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 5, .. } + )); + assert!(matches!( + receive(&mut host).await.message, + Message::End { + id: 5, + error: Some(PluginFault { + code: ErrorCode::Fault, + ref message, + .. + }), + } if message == "plugin stream error metadata exceeds its limit" + )); + + send_call(&mut host, 7, "echo", json!("still-ready"), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 7, result } if result == json!("still-ready") + )); + shutdown(&mut host, task).await; +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/lifecycle.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/lifecycle.rs new file mode 100644 index 000000000..edfd2e860 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/lifecycle.rs @@ -0,0 +1,112 @@ +//! 会话期限、取消钩子、排空与关闭测试 + +use super::*; + +#[tokio::test] +async fn call_deadline_cancels_handler_and_keeps_the_session_available() { + let handler = TestHandler::default(); + let state = Arc::clone(&handler.lifecycle); + let (mut host, task) = start_session(handler).await; + send_call_with_timeout( + &mut host, + 1, + "pending", + json!({}), + Vec::new(), + Duration::from_millis(20), + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 1, + error: PluginFault { + code: ErrorCode::Timeout, + .. + }, + } + )); + tokio::time::timeout(Duration::from_millis(100), async { + while !state.cancel_started.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("deadline did not notify the lifecycle hook"); + + send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 3, result } if result == json!("still-ready") + )); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn blocking_cancel_hook_cannot_delay_cancelled_or_eof_settlement() { + let handler = TestHandler { + cancel_delay: Some(Duration::from_millis(200)), + ..TestHandler::default() + }; + let state = Arc::clone(&handler.lifecycle); + let (mut host, task) = start_session(handler).await; + send_call(&mut host, 1, "pending", json!({}), Vec::new()).await; + send_control(&mut host, Message::Cancel { id: 1 }).await; + + let cancelled = tokio::time::timeout(Duration::from_millis(50), receive(&mut host)) + .await + .expect("blocking lifecycle hook delayed Cancelled"); + assert!(matches!(cancelled.message, Message::Cancelled { id: 1 })); + tokio::time::timeout(Duration::from_millis(50), async { + while !state.cancel_started.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancel hook never started"); + + host.writer.shutdown().await.unwrap(); + let result = tokio::time::timeout(Duration::from_millis(50), task) + .await + .expect("blocking lifecycle hook delayed EOF settlement") + .expect("plugin session task panicked"); + assert!(matches!(result, Err(SessionError::Closed))); +} + +#[tokio::test] +async fn quiesce_rejects_new_calls_and_shutdown_closes_the_session() { + let handler = TestHandler::default(); + let state = Arc::clone(&handler.lifecycle); + let (mut host, task) = start_session(handler).await; + send_control(&mut host, Message::Quiesce).await; + send_call(&mut host, 1, "echo", json!({}), Vec::new()).await; + + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 1, + error: PluginFault { + code: ErrorCode::Capacity, + .. + }, + } + )); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1, + frames: 1, + }, + ) + .await; + tokio::time::timeout(Duration::from_millis(50), async { + while !state.quiesced.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("quiesce hook was not signalled"); + + shutdown(&mut host, task).await; +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs new file mode 100644 index 000000000..5b86d9bec --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs @@ -0,0 +1,470 @@ +//! 插件会话测试入口与双工传输、处理器和宿主回调夹具 + +mod credit; +mod driver; +mod lifecycle; +mod registry; + +use std::{ + future::pending, + num::NonZeroUsize, + sync::{ + Arc, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use gateway_plugin_sdk::{ + CallContext, Capability, ContributionDeclaration, Contributions, ErrorCode, Frame, Handshake, + Message, PROTOCOL_VERSION, PluginFault, Stage, + call::middleware::{ + BODY_READ_METHOD, HANDLE_METHOD, MiddlewareBodyDisposition, MiddlewareBodyFrame, + MiddlewareBodyFraming, MiddlewareBodyHandle, MiddlewareBodyReadResult, MiddlewareHeader, + MiddlewareHeaderMutation, MiddlewareMount, MiddlewareNextRequest, MiddlewareNextResponse, + MiddlewareRequestBody, MiddlewareRequestHead, MiddlewareResponseBody, + MiddlewareResponseHead, MiddlewareTransport, NEXT_METHOD, + }, + client::{ + CallFuture, CallReply, MiddlewarePlugin, PluginCall, PluginHandler, PluginSession, + RequestCall, ResponseStream, SessionConfig, SessionError, read_frame, write_frame, + }, +}; +use serde_json::{Value, json}; +use tokio::{ + io::{AsyncWriteExt, DuplexStream, ReadHalf, WriteHalf}, + task::JoinHandle, +}; + +const MAXIMUM_STREAM_CHUNK_BYTES: usize = 64 * 1024; + +#[derive(Clone, Default)] +struct TestHandler { + lifecycle: Arc, + cancel_delay: Option, +} + +#[derive(Default)] +struct LifecycleState { + cancel_started: AtomicBool, + cancelled: AtomicUsize, + quiesced: AtomicBool, + shutdown: AtomicBool, +} + +impl PluginHandler for TestHandler { + fn call(&self, call: PluginCall) -> CallFuture<'_> { + Box::pin(async move { + match call.method.as_str() { + "echo" => Ok(CallReply::unary(call.params, call.payload)), + "key_facts" => { + let result = call + .host + .key_facts(gateway_plugin_sdk::call::data::ClientKeyFactsQuery { + client_key_id: "key_1".into(), + }) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "get_budget" => { + let result = call + .host + .get_key_budget( + gateway_plugin_sdk::call::key_budgets::GetKeyBudgetRequest { + client_key_id: "key_1".into(), + }, + ) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "update_budget_limits" => { + let result = call + .host + .update_key_budget_limits( + gateway_plugin_sdk::call::key_budgets::UpdateKeyBudgetLimitsRequest { + client_key_id: "key_1".into(), + daily_limit_usd: None, + weekly_limit_usd: Some("12.5".into()), + }, + ) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "refresh_quota" => { + let result = call + .host + .refresh_account_quota(gateway_plugin_sdk::call::data::QuotaFactsQuery { + account_id: "acct_1".into(), + }) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "reset_budget" => { + use gateway_plugin_sdk::call::{ + host::KeyListRequest, + key_budgets::{BudgetPeriod, ResetKeyBudgetRequest}, + }; + let keys = call + .host + .list_keys(KeyListRequest { + cursor: None, + limit: 10, + }) + .await?; + let result = call + .host + .reset_key_budget(ResetKeyBudgetRequest { + client_key_id: keys.keys[0].id.clone(), + period: BudgetPeriod::Weekly, + }) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "slow" => { + tokio::time::sleep(Duration::from_millis(80)).await; + Ok(CallReply::unary(call.params, call.payload)) + } + "callback" => { + let reply = call + .host + .call("host.log", call.params, call.payload) + .await + .map_err(SessionError::into_plugin_fault)?; + Ok(CallReply::unary(reply.result, reply.payload)) + } + "resource_stream" => Ok(CallReply::stream( + json!({}), + Vec::new(), + ResponseStream::pull(Box::new(ResourceStream(Some(call.host)))), + )), + "stream" => Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + ResponseStream::from_chunks(vec![b"abcd".to_vec(), b"efgh".to_vec()]), + )), + "long_stream" => Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + ResponseStream::from_chunks(vec![b"x".to_vec(); 16]), + )), + "oversized_stream" => Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + ResponseStream::from_chunks(vec![b"abcde".to_vec()]), + )), + "too_many_chunks" => Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + ResponseStream::from_chunks(vec![b"x".to_vec(); 17]), + )), + "dynamic_oversized_stream" => { + let (sender, stream) = ResponseStream::channel(NonZeroUsize::new(1).unwrap()); + tokio::spawn(async move { + let _ = sender.send(b"abcde".to_vec()).await; + }); + Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + stream, + )) + } + "oversized_fault" => Err(PluginFault::new( + ErrorCode::Upstream, + "x".repeat(MAXIMUM_STREAM_CHUNK_BYTES), + )), + "oversized_result" => Ok(CallReply::unary( + Value::String("x".repeat(MAXIMUM_STREAM_CHUNK_BYTES)), + Vec::new(), + )), + "dynamic_oversized_fault" => { + let (sender, stream) = ResponseStream::channel(NonZeroUsize::new(1).unwrap()); + tokio::spawn(async move { + let _ = sender + .fail(PluginFault::new( + ErrorCode::Upstream, + "x".repeat(MAXIMUM_STREAM_CHUNK_BYTES), + )) + .await; + }); + Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + stream, + )) + } + "pending" => pending::>().await, + _ => Err(PluginFault::new( + ErrorCode::Unsupported, + "unsupported test method", + )), + } + }) + } + + fn cancel(&self, _context: &CallContext) { + self.lifecycle.cancel_started.store(true, Ordering::Release); + if let Some(delay) = self.cancel_delay { + std::thread::sleep(delay); + } + self.lifecycle.cancelled.fetch_add(1, Ordering::Relaxed); + } + + fn quiesce(&self) { + self.lifecycle.quiesced.store(true, Ordering::Release); + } + + fn shutdown(&self) { + self.lifecycle.shutdown.store(true, Ordering::Release); + } +} + +pub(super) struct HostPeer { + pub(super) reader: ReadHalf, + pub(super) writer: WriteHalf, +} + +pub(super) async fn start_session( + handler: H, +) -> (HostPeer, JoinHandle>) { + start_session_with_capacity(handler, MAXIMUM_STREAM_CHUNK_BYTES * 2).await +} + +async fn start_session_with_capacity( + handler: H, + capacity: usize, +) -> (HostPeer, JoinHandle>) { + let (host, plugin) = tokio::io::duplex(capacity); + let (plugin_reader, plugin_writer) = tokio::io::split(plugin); + let task = tokio::spawn(async move { + let session = PluginSession::accept( + plugin_reader, + plugin_writer, + SessionConfig { + maximum_stream_chunk_bytes: MAXIMUM_STREAM_CHUNK_BYTES, + maximum_calls: 4, + maximum_callbacks: 4, + maximum_buffered_stream_chunks: 16, + handshake_timeout: Duration::from_secs(1), + maximum_call_timeout: Duration::from_secs(2), + }, + ) + .await?; + session.run(handler).await + }); + let (mut reader, mut writer) = tokio::io::split(host); + write_frame( + &mut writer, + &Frame::control(Message::Hello { + handshake: handshake(), + }), + ) + .await + .unwrap(); + let ready = read_frame(&mut reader).await.unwrap(); + assert!(matches!( + ready, + Frame { + message: Message::Ready { + protocol_version: PROTOCOL_VERSION, + ref incarnation, + }, + ref payload, + } if incarnation == "test-incarnation" && payload.is_empty() + )); + (HostPeer { reader, writer }, task) +} + +fn middleware_contributions() -> Contributions { + Contributions::from([( + Capability::Middleware, + ContributionDeclaration { + id: "test.example.middleware".into(), + version: 4, + stages: vec![Stage::Request], + input_formats: vec!["openai".into()], + output_formats: vec!["openai".into()], + }, + )]) +} + +fn middleware_plugin(map_response: bool) -> impl PluginHandler { + MiddlewarePlugin::new( + &middleware_contributions(), + move |call: RequestCall| async move { + let RequestCall { + mut request, next, .. + } = call; + if map_response { + request.body.push(b'!'); + request.remove_header("x-old"); + request.head.headers.push(MiddlewareHeader { + name: "x-direct".into(), + value: b"request".to_vec(), + }); + } + let mut response = next.run(request).await?; + if map_response { + response.remove_header("x-upstream"); + response.headers.push(MiddlewareHeader { + name: "x-direct-response".into(), + value: b"response".to_vec(), + }); + response.body = response.body.map_frames(|mut frame| { + frame.payload.make_ascii_uppercase(); + Ok(vec![frame]) + })?; + } + Ok(response) + }, + ) + .unwrap() +} + +fn handshake() -> Handshake { + Handshake { + protocol_version: PROTOCOL_VERSION, + artifact_sha256: "a".repeat(64), + plugin_id: "test.example".into(), + instance_id: "test-instance".into(), + generation: 7, + incarnation: "test-incarnation".into(), + configuration: json!({}), + + contributes: Contributions::new(), + } +} + +pub(super) fn context(id: u64, timeout: Duration) -> CallContext { + CallContext { + call_id: id, + instance_id: "test-instance".into(), + generation: 7, + incarnation: "test-incarnation".into(), + stage: Stage::Request, + timeout_ms: timeout.as_millis() as u64, + resource_stream: false, + resource_scope_id: format!("scope-{id}"), + request_id: None, + attempt_id: None, + account_id: None, + credential_revision: None, + } +} + +fn middleware_context(id: u64) -> CallContext { + CallContext { + stage: Stage::Request, + request_id: Some("request-middleware".into()), + ..context(id, Duration::from_secs(1)) + } +} + +pub(super) fn middleware_request() -> MiddlewareRequestHead { + MiddlewareRequestHead { + settings_sources: serde_json::Value::Null, + settings: serde_json::Value::Null, + client_key_id: "fixture-key".into(), + account_group_ids: vec!["fixture-group".into()], + request_id: "request-middleware".into(), + mount: MiddlewareMount::Request, + attempt_index: None, + operation: "generate".into(), + protocol: "openai".into(), + endpoint: "responses".into(), + transport: MiddlewareTransport::HttpSse, + provider: None, + model: Some("public-model".into()), + account_id: None, + headers: vec![MiddlewareHeader { + name: "x-old".into(), + value: b"old".to_vec(), + }], + } +} + +pub(super) async fn send_call( + host: &mut HostPeer, + id: u64, + method: &str, + params: Value, + payload: Vec, +) { + send_call_with_timeout(host, id, method, params, payload, Duration::from_secs(1)).await; +} + +async fn send_call_with_timeout( + host: &mut HostPeer, + id: u64, + method: &str, + params: Value, + payload: Vec, + timeout: Duration, +) { + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id, + method: method.into(), + context: context(id, timeout), + params, + }, + payload, + }, + ) + .await + .unwrap(); +} + +async fn send_control(host: &mut HostPeer, message: Message) { + write_frame(&mut host.writer, &Frame::control(message)) + .await + .unwrap(); +} + +pub(super) async fn receive(host: &mut HostPeer) -> Frame { + tokio::time::timeout(Duration::from_secs(1), read_frame(&mut host.reader)) + .await + .expect("plugin response timed out") + .expect("plugin response frame must be valid") +} + +pub(super) async fn shutdown(host: &mut HostPeer, task: JoinHandle>) { + send_control(host, Message::Shutdown).await; + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("plugin session did not shut down") + .expect("plugin session task panicked") + .expect("plugin session shutdown failed"); +} + +struct ResourceStream(Option); +impl gateway_plugin_sdk::client::PullResponseStream for ResourceStream { + fn next(&mut self) -> gateway_plugin_sdk::client::PullResponseFuture<'_> { + Box::pin(async move { + let host = self.0.take()?; + tokio::time::sleep(Duration::from_millis(700)).await; + Some( + host.call("host.log", json!({}), Vec::new()) + .await + .map(|_| b"live".to_vec()) + .map_err(SessionError::into_plugin_fault), + ) + }) + } +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs new file mode 100644 index 000000000..cda12ef26 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs @@ -0,0 +1,299 @@ +//! 宿主回调关联、退休回调与强类型调用合同测试 + +use super::*; + +#[tokio::test] +async fn typed_budget_helpers_preserve_wire_contract_and_propagate_failure_without_retry() { + for failed in [false, true] { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "reset_budget", json!({}), vec![]).await; + let list = receive(&mut host).await; + let Message::Callback { + id, + parent_id: 1, + method, + params, + } = list.message + else { + panic!("expected Key list callback") + }; + assert_eq!(method, "host.keys.list"); + assert_eq!(params, json!({"cursor":null,"limit":10})); + assert!(list.payload.is_empty()); + send_control(&mut host, Message::Result { id, result: json!({"keys":[{"id":"key_1","name":"test","enabled":true}],"next_cursor":null}) }).await; + let reset = receive(&mut host).await; + let Message::Callback { + id, + parent_id: 1, + method, + params, + } = reset.message + else { + panic!("expected budget reset callback") + }; + assert_eq!(method, "host.keys.reset_budget"); + assert_eq!(params, json!({})); + assert_eq!( + serde_json::from_slice::(&reset.payload).unwrap(), + json!({"client_key_id":"key_1","period":"weekly"}) + ); + if failed { + send_control( + &mut host, + Message::Error { + id, + error: PluginFault::new(ErrorCode::Conflict, "stale instance"), + }, + ) + .await; + let Message::Error { id: 1, error } = receive(&mut host).await.message else { + panic!("expected error propagated without another callback") + }; + assert_eq!(error.code, ErrorCode::Conflict); + } else { + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id, + result: json!({}), + }, + payload: serde_json::to_vec(&json!({"client_key_id":"key_1"})).unwrap(), + }, + ) + .await + .unwrap(); + let Message::Result { id: 1, result } = receive(&mut host).await.message else { + panic!("expected successful reset") + }; + assert_eq!(result, json!({"client_key_id":"key_1"})); + } + shutdown(&mut host, task).await; + } +} + +#[tokio::test] +async fn host_callback_round_trips_while_parent_call_is_waiting() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "callback", json!({"event": "ready"}), vec![7]).await; + + let callback = receive(&mut host).await; + let callback_id = match callback.message { + Message::Callback { + id, + parent_id: 1, + ref method, + ref params, + } if id.is_multiple_of(2) + && method == "host.log" + && params == &json!({"event": "ready"}) => + { + id + } + message => panic!("unexpected callback frame: {message:?}"), + }; + assert_eq!(callback.payload, [7]); + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: callback_id, + result: json!({"recorded": true}), + }, + payload: vec![9], + }, + ) + .await + .unwrap(); + + let result = receive(&mut host).await; + assert!(matches!( + result, + Frame { + message: Message::Result { id: 1, result }, + payload, + } if result == json!({"recorded": true}) && payload == [9] + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn callbacks_from_concurrent_calls_remain_correlated() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "callback", json!({"parent": 1}), vec![1]).await; + send_call(&mut host, 3, "callback", json!({"parent": 3}), vec![3]).await; + + let mut callbacks = Vec::new(); + for _ in 0..2 { + let frame = receive(&mut host).await; + let Message::Callback { + id, + parent_id, + method, + params, + } = frame.message + else { + panic!("expected callback") + }; + assert_eq!(method, "host.log"); + assert_eq!(params, json!({"parent": parent_id})); + assert_eq!(frame.payload, [parent_id as u8]); + callbacks.push((id, parent_id)); + } + assert!(callbacks[0].0 < callbacks[1].0); + assert!(callbacks.iter().all(|(id, _)| id.is_multiple_of(2))); + + for (callback_id, parent_id) in callbacks.into_iter().rev() { + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: callback_id, + result: json!({"parent": parent_id}), + }, + payload: vec![parent_id as u8], + }, + ) + .await + .unwrap(); + } + let mut completed = Vec::new(); + for _ in 0..2 { + let frame = receive(&mut host).await; + let Message::Result { id, result } = frame.message else { + panic!("expected parent result") + }; + assert_eq!(result, json!({"parent": id})); + assert_eq!(frame.payload, [id as u8]); + completed.push(id); + } + completed.sort_unstable(); + assert_eq!(completed, [1, 3]); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn cancelling_parent_retires_callback_and_late_reply_does_not_close_session() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "callback", json!({}), Vec::new()).await; + let callback = receive(&mut host).await; + let Message::Callback { + id: callback_id, .. + } = callback.message + else { + panic!("expected callback") + }; + + send_control(&mut host, Message::Cancel { id: 1 }).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Cancelled { id: 1 } + )); + send_control( + &mut host, + Message::Result { + id: callback_id, + result: json!({}), + }, + ) + .await; + send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 3, result } if result == json!("still-ready") + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn typed_key_and_quota_calls_keep_payloads_and_do_not_retry_failures() { + for (entry, method, request, response) in [ + ( + "key_facts", + "host.data.keys.get", + json!({"client_key_id":"key_1"}), + json!({"schema_version":1,"client_key_id":"key_1","enabled":false,"group_ids":["grp_1"]}), + ), + ( + "get_budget", + "host.keys.get_budget", + json!({"client_key_id":"key_1"}), + json!({ + "client_key_id":"key_1", "daily_limit_usd":"10", "weekly_limit_usd":"20", + "daily_used_usd":"1.25", "weekly_used_usd":"2.5", "daily_resets_at_ms":null, "weekly_resets_at_ms":null, + }), + ), + ( + "update_budget_limits", + "host.keys.update_budget_limits", + json!({"client_key_id":"key_1","weekly_limit_usd":"12.5"}), + json!({"client_key_id":"key_1"}), + ), + ( + "refresh_quota", + "host.quota_observations.refresh", + json!({"account_id":"acct_1"}), + json!({"schema_version":1,"account_id":"acct_1","observed_at_ms":null,"windows":[]}), + ), + ] { + for failed in [false, true] { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, entry, json!({}), vec![]).await; + let callback = receive(&mut host).await; + let Message::Callback { + id, + parent_id: 1, + method: actual, + params, + } = callback.message + else { + panic!("expected callback"); + }; + assert_eq!(actual, method); + assert_eq!(params, json!({})); + assert_eq!( + serde_json::from_slice::(&callback.payload).unwrap(), + request + ); + if failed { + send_control( + &mut host, + Message::Error { + id, + error: PluginFault::new(ErrorCode::Conflict, "changed"), + }, + ) + .await; + let Message::Error { id: 1, error } = receive(&mut host).await.message else { + panic!("expected propagated error without retry"); + }; + assert_eq!(error.code, ErrorCode::Conflict); + } else { + let mut response_payload = response.clone(); + if matches!(entry, "key_facts" | "refresh_quota") { + response_payload["future_fact"] = json!({"value":1}); + } + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id, + result: json!({}), + }, + payload: serde_json::to_vec(&response_payload).unwrap(), + }, + ) + .await + .unwrap(); + let Message::Result { id: 1, result } = receive(&mut host).await.message else { + panic!("expected projected response"); + }; + assert_eq!(result, response); + } + shutdown(&mut host, task).await; + } + } +} diff --git a/backend/crates/gateway-store/src/backup/s3.rs b/backend/crates/gateway-store/src/backup/s3.rs index 0cbba7772..2e78fa03a 100644 --- a/backend/crates/gateway-store/src/backup/s3.rs +++ b/backend/crates/gateway-store/src/backup/s3.rs @@ -7,10 +7,7 @@ use std::collections::HashMap; use std::path::Path; -use std::sync::{ - Arc, Mutex, - atomic::{AtomicU64, Ordering}, -}; +use std::sync::Mutex; use std::time::Duration; use aws_credential_types::Credentials; @@ -24,9 +21,9 @@ use aws_smithy_types::error::metadata::ProvideErrorMetadata; use aws_smithy_types::retry::RetryConfig; use aws_smithy_types::timeout::TimeoutConfig; use chrono::{DateTime, Utc}; +use futures::{StreamExt as _, TryStreamExt as _, stream}; use secrecy::ExposeSecret as _; use tokio::io::{AsyncReadExt as _, AsyncSeekExt as _, SeekFrom}; -use tokio::sync::Semaphore; use uuid::Uuid; use gateway_admin::model::backup::{ @@ -242,7 +239,7 @@ impl BackupObjectStorePort for S3ObjectStoreAdapter { })? .to_owned(); - let upload = upload_parts(&client, config, &request, upload_id.clone(), file_size); + let upload = upload_parts(&client, config, &request, &upload_id, file_size); let outcome = tokio::select! { result = upload => result, _ = request.cancellation.cancelled() => { @@ -358,78 +355,43 @@ async fn upload_parts( client: &S3Client, config: &BackupStorageConfig, request: &UploadObjectRequest, - upload_id: String, + upload_id: &str, file_size: u64, ) -> Result<(), BackupError> { let total_parts = file_size.div_ceil(u64::try_from(PART_SIZE).unwrap_or(u64::MAX)); - let next_part = Arc::new(AtomicU64::new(1)); - let semaphore = Arc::new(Semaphore::new(CONCURRENCY)); - let (sender, mut receiver) = tokio::sync::mpsc::channel::<(i32, String)>(CONCURRENCY * 2); - - let mut workers = Vec::with_capacity(CONCURRENCY); - for _ in 0..CONCURRENCY { - let client = client.clone(); - let semaphore = Arc::clone(&semaphore); - let next_part = Arc::clone(&next_part); - let sender = sender.clone(); - let source = request.source.clone(); - let bucket = config.bucket.clone(); - let key = request.object_key.clone(); - let upload_id = upload_id.clone(); - workers.push(tokio::spawn(async move { - loop { - let part = next_part.fetch_add(1, Ordering::Relaxed); - if part > total_parts { - break; - } - let _permit = semaphore.acquire().await.map_err(|_| { - BackupError::new(code::S3_UPLOAD_FAILED, "上传并发闸门关闭".to_owned()) - })?; - let offset = (part - 1) * (PART_SIZE as u64); - let remaining = file_size.checked_sub(offset).ok_or_else(|| { - BackupError::new(code::S3_UPLOAD_FAILED, "分片偏移超出归档大小".to_owned()) - })?; - let part_len = usize::try_from(remaining.min(PART_SIZE as u64)).map_err(|_| { - BackupError::new(code::S3_UPLOAD_FAILED, "分片大小溢出".to_owned()) - })?; - let bytes = read_part(&source, offset, part_len).await?; - let part_number = i32::try_from(part).map_err(|_| { - BackupError::new(code::S3_UPLOAD_FAILED, "分片号溢出".to_owned()) - })?; - let response = client - .upload_part() - .bucket(&bucket) - .key(&key) - .upload_id(&upload_id) - .part_number(part_number) - .body(ByteStream::from(bytes)) - .send() - .await - .map_err(map_s3_error)?; - let etag = response - .e_tag() - .ok_or_else(|| { - BackupError::new(code::S3_UPLOAD_FAILED, "S3 未返回分片 ETag".to_owned()) - })? - .to_owned(); - if sender.send((part_number, etag)).await.is_err() { - break; - } - } - Ok::<(), BackupError>(()) - })); - } - drop(sender); - - let mut parts = Vec::new(); - while let Some((part, etag)) = receiver.recv().await { - parts.push((part, etag)); - } - for worker in workers { - worker.await.map_err(|_| { - BackupError::new(code::S3_UPLOAD_FAILED, "上传分片任务崩溃".to_owned()) - })??; - } + // 分片 future 由父上传持有；取消或首次失败时一并释放，SDK 重试不能逃逸到清理之后 + let mut parts: Vec<_> = stream::iter(1..=total_parts) + .map(|part| async move { + let offset = (part - 1) * (PART_SIZE as u64); + let remaining = file_size.checked_sub(offset).ok_or_else(|| { + BackupError::new(code::S3_UPLOAD_FAILED, "分片偏移超出归档大小".to_owned()) + })?; + let part_len = usize::try_from(remaining.min(PART_SIZE as u64)) + .map_err(|_| BackupError::new(code::S3_UPLOAD_FAILED, "分片大小溢出".to_owned()))?; + let bytes = read_part(&request.source, offset, part_len).await?; + let part_number = i32::try_from(part) + .map_err(|_| BackupError::new(code::S3_UPLOAD_FAILED, "分片号溢出".to_owned()))?; + let response = client + .upload_part() + .bucket(&config.bucket) + .key(&request.object_key) + .upload_id(upload_id) + .part_number(part_number) + .body(ByteStream::from(bytes)) + .send() + .await + .map_err(map_s3_error)?; + let etag = response + .e_tag() + .ok_or_else(|| { + BackupError::new(code::S3_UPLOAD_FAILED, "S3 未返回分片 ETag".to_owned()) + })? + .to_owned(); + Ok::<_, BackupError>((part_number, etag)) + }) + .buffer_unordered(CONCURRENCY) + .try_collect() + .await?; parts.sort_by_key(|(part, _)| *part); let completed = CompletedMultipartUpload::builder() @@ -449,7 +411,7 @@ async fn upload_parts( .complete_multipart_upload() .bucket(&config.bucket) .key(&request.object_key) - .upload_id(&upload_id) + .upload_id(upload_id) .multipart_upload(completed) .send() .await diff --git a/backend/crates/gateway-store/src/postgres/client_budgets.rs b/backend/crates/gateway-store/src/postgres/client_budgets.rs index 190880430..377eb1daa 100644 --- a/backend/crates/gateway-store/src/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/src/postgres/client_budgets.rs @@ -371,12 +371,13 @@ pub(super) async fn load_client_key_budgets( .iter() .map(|record| record.id.as_str()) .collect::>(); + // 起止相等是重置后的未开启窗口，不能因应用时钟领先数据库而投影为活跃预算 let rows = sqlx::query( "select k.id, k.daily_limit_usd::text, k.weekly_limit_usd::text, - (case when w.daily_end > now() then w.daily_used_usd else 0 end)::text as daily_used, - (case when w.weekly_end > now() then w.weekly_used_usd else 0 end)::text as weekly_used, - case when w.daily_end > now() then w.daily_end end as daily_end, - case when w.weekly_end > now() then w.weekly_end end as weekly_end + (case when w.daily_end > w.daily_start and w.daily_end > now() then w.daily_used_usd else 0 end)::text as daily_used, + (case when w.weekly_end > w.weekly_start and w.weekly_end > now() then w.weekly_used_usd else 0 end)::text as weekly_used, + case when w.daily_end > w.daily_start and w.daily_end > now() then w.daily_end end as daily_end, + case when w.weekly_end > w.weekly_start and w.weekly_end > now() then w.weekly_end end as weekly_end from client_api_keys k left join client_key_budget_windows w on w.client_api_key_id = k.id where k.id = any($1)", ) diff --git a/backend/crates/gateway-store/src/postgres/connection.rs b/backend/crates/gateway-store/src/postgres/connection.rs new file mode 100644 index 000000000..6cb104040 --- /dev/null +++ b/backend/crates/gateway-store/src/postgres/connection.rs @@ -0,0 +1,104 @@ +//! PostgreSQL 迁移屏障、连接池与会话级资源预算 + +use sqlx::{ + PgPool, + postgres::{PgConnectOptions, PgPoolOptions}, +}; + +use crate::{ + POSTGRES_IDLE_TRANSACTION_TIMEOUT, POSTGRES_LOCK_TIMEOUT, POSTGRES_STATEMENT_TIMEOUT, + StoreError, StorePoolConfig, StoreResult, postgres_unavailable, +}; + +static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../../migrations"); + +/// 建立 PostgreSQL pool 并只执行冻结的 migration 集 +pub async fn connect_and_migrate( + database_url: &str, + pool_config: StorePoolConfig, +) -> StoreResult { + if database_url.trim().is_empty() { + return Err(StoreError::InvalidData { + source: None, + entity: "PostgreSQL configuration", + message: "database URL is empty".to_owned(), + }); + } + pool_config.validate()?; + let connect_options = database_url + .parse::() + .map_err(|source| postgres_unavailable("parse PostgreSQL connection options", source))?; + let migration_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + connect_options + .clone() + .application_name("codex-proxy-rs:migration"), + ) + .await + .map_err(|source| postgres_unavailable("connect PostgreSQL for migrations", source))?; + if let Err(error) = MIGRATOR.run(&migration_pool).await { + migration_pool.close().await; + return Err(postgres_unavailable("apply PostgreSQL migrations", error)); + } + migration_pool.close().await; + + connect_pool(connect_options, pool_config, false).await +} + +/// 帮助查询不执行迁移，并用连接默认只读事务阻止意外业务写入 +pub(crate) async fn connect_read_only( + database_url: &str, + pool_config: StorePoolConfig, +) -> StoreResult { + pool_config.validate()?; + let options = database_url + .parse::() + .map_err(|source| postgres_unavailable("parse PostgreSQL connection options", source))?; + connect_pool(options, pool_config, true).await +} + +async fn connect_pool( + connect_options: PgConnectOptions, + pool_config: StorePoolConfig, + read_only: bool, +) -> StoreResult { + let statement_timeout = postgres_duration_setting(POSTGRES_STATEMENT_TIMEOUT); + let lock_timeout = postgres_duration_setting(POSTGRES_LOCK_TIMEOUT); + let idle_in_transaction_session_timeout = + postgres_duration_setting(POSTGRES_IDLE_TRANSACTION_TIMEOUT); + let pool = PgPoolOptions::new() + .max_connections(pool_config.max_connections) + .acquire_timeout(std::time::Duration::from_secs( + pool_config.acquire_timeout_seconds, + )) + .after_connect(move |connection, _metadata| { + let statement_timeout = statement_timeout.clone(); + let lock_timeout = lock_timeout.clone(); + let idle_in_transaction_session_timeout = idle_in_transaction_session_timeout.clone(); + Box::pin(async move { + sqlx::query( + "select set_config('statement_timeout', $1, false), + set_config('lock_timeout', $2, false), + set_config('idle_in_transaction_session_timeout', $3, false), + set_config('default_transaction_read_only', $4, false), + set_config('TimeZone', 'UTC', false)", + ) + .bind(statement_timeout) + .bind(lock_timeout) + .bind(idle_in_transaction_session_timeout) + .bind(if read_only { "on" } else { "off" }) + .execute(connection) + .await?; + Ok(()) + }) + }) + .connect_with(connect_options.application_name("codex-proxy-rs")) + .await + .map_err(|source| postgres_unavailable("connect PostgreSQL", source))?; + Ok(pool) +} + +fn postgres_duration_setting(duration: std::time::Duration) -> String { + format!("{}ms", duration.as_millis()) +} diff --git a/backend/crates/gateway-store/src/postgres/control_plane.rs b/backend/crates/gateway-store/src/postgres/control_plane.rs new file mode 100644 index 000000000..1dcf96173 --- /dev/null +++ b/backend/crates/gateway-store/src/postgres/control_plane.rs @@ -0,0 +1,286 @@ +//! 控制面设置与 Key 修改的 revision 锁、条件写入和同事务审计 + +use async_trait::async_trait; +use sqlx::PgPool; + +use crate::{Revision, StoreError, StoreResult, postgres_unavailable}; + +use super::{ + AdminAuditEvent, NewClientApiKey, RuntimeSettings, RuntimeSettingsUpdate, + UpdateClientApiKeyDetails, append_admin_audit_event_in_transaction, + bump_config_revision_in_transaction, delete_client_api_key_in_transaction, + insert_client_api_key_in_transaction, load_runtime_settings_from_pool, + load_runtime_settings_in_transaction, set_client_api_key_enabled_in_transaction, + update_admin_api_key_in_transaction, update_client_api_key_in_transaction, + update_runtime_settings_in_transaction, +}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ControlPlaneSnapshot { + pub settings: RuntimeSettings, +} + +#[derive(Debug, Clone)] +pub struct ControlPlaneReplacement { + pub expected_revision: Revision, + pub settings: RuntimeSettingsUpdate, + pub audit: AdminAuditEvent, +} + +#[async_trait] +pub trait ControlPlaneRepository: Send + Sync { + async fn load_control_plane(&self) -> StoreResult; + + async fn replace_control_plane( + &self, + replacement: ControlPlaneReplacement, + ) -> StoreResult; + + /// 更新 admin_api_key 字段并推进 config revision + async fn replace_admin_api_key( + &self, + admin_api_key: Option, + audit: AdminAuditEvent, + ) -> StoreResult; + + async fn create_client_api_key( + &self, + key: NewClientApiKey, + audit: AdminAuditEvent, + ) -> StoreResult; + + async fn update_client_api_key( + &self, + key: UpdateClientApiKeyDetails, + audit: AdminAuditEvent, + ) -> StoreResult; + + async fn set_client_api_key_enabled( + &self, + id: &str, + enabled: bool, + audit: AdminAuditEvent, + ) -> StoreResult; + + async fn delete_client_api_key( + &self, + id: &str, + audit: AdminAuditEvent, + ) -> StoreResult; +} + +#[derive(Clone)] +pub struct PgControlPlaneRepository { + pub(super) pool: PgPool, +} + +impl PgControlPlaneRepository { + #[must_use] + pub const fn new(pool: PgPool) -> Self { + Self { pool } + } +} + +#[async_trait] +impl ControlPlaneRepository for PgControlPlaneRepository { + async fn load_control_plane(&self) -> StoreResult { + let settings = load_runtime_settings_from_pool(&self.pool).await?; + Ok(ControlPlaneSnapshot { settings }) + } + + async fn replace_control_plane( + &self, + replacement: ControlPlaneReplacement, + ) -> StoreResult { + replacement.settings.validate()?; + let mut transaction = + self.pool.begin().await.map_err(|source| { + postgres_unavailable("begin control plane replacement", source) + })?; + let result: StoreResult<_> = async { + // 锁住所有配置写入共同使用的行，版本检查与修改不能被其他事务穿插 + let current = sqlx::query_scalar::<_, i64>( + "select config_revision from runtime_settings where id = 1 for update", + ) + .fetch_one(&mut *transaction) + .await + .map_err(|source| postgres_unavailable("lock control plane revision", source))?; + if u64::try_from(current).ok() != Some(replacement.expected_revision.get()) { + return Err(StoreError::Conflict { + source: None, + entity: "runtime settings", + id: "1".to_owned(), + kind: crate::ConflictKind::StaleRevision, + }); + } + let revision = + update_runtime_settings_in_transaction(&mut transaction, &replacement.settings) + .await?; + append_admin_audit_event_in_transaction(&mut transaction, replacement.audit, revision) + .await?; + load_control_plane_in_transaction(&mut transaction).await + } + .await; + match result { + Ok(snapshot) => { + transaction.commit().await.map_err(|source| { + postgres_unavailable("commit control plane replacement", source) + })?; + Ok(snapshot) + } + Err(error) => { + let error = match transaction.rollback().await { + Ok(()) => error, + Err(cleanup) => error.with_cleanup(cleanup), + }; + Err(error) + } + } + } + + async fn replace_admin_api_key( + &self, + admin_api_key: Option, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation(ControlPlaneMutation::SetAdminApiKey(admin_api_key), audit) + .await + } + + async fn create_client_api_key( + &self, + key: NewClientApiKey, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation(ControlPlaneMutation::CreateClientApiKey(key), audit) + .await + } + + async fn update_client_api_key( + &self, + key: UpdateClientApiKeyDetails, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation(ControlPlaneMutation::UpdateClientApiKey(key), audit) + .await + } + + async fn set_client_api_key_enabled( + &self, + id: &str, + enabled: bool, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation( + ControlPlaneMutation::SetClientApiKeyEnabled { + id: id.to_owned(), + enabled, + }, + audit, + ) + .await + } + + async fn delete_client_api_key( + &self, + id: &str, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation( + ControlPlaneMutation::DeleteClientApiKey(id.to_owned()), + audit, + ) + .await + } +} + +enum ControlPlaneMutation { + CreateClientApiKey(NewClientApiKey), + UpdateClientApiKey(UpdateClientApiKeyDetails), + SetClientApiKeyEnabled { id: String, enabled: bool }, + DeleteClientApiKey(String), + SetAdminApiKey(Option), +} + +impl PgControlPlaneRepository { + async fn apply_targeted_mutation( + &self, + mutation: ControlPlaneMutation, + mut audit: AdminAuditEvent, + ) -> StoreResult { + let mut transaction = self.pool.begin().await.map_err(|source| { + postgres_unavailable("begin targeted control plane mutation", source) + })?; + let result: StoreResult<_> = async { + let revision = bump_config_revision_in_transaction(&mut transaction).await?; + match mutation { + ControlPlaneMutation::CreateClientApiKey(key) => { + audit.changed_fields.push(if key.group_ids.is_empty() { + "routing_scope:all".to_owned() + } else { + "routing_scope:groups".to_owned() + }); + insert_client_api_key_in_transaction(&mut transaction, &key).await?; + } + ControlPlaneMutation::UpdateClientApiKey(key) => { + let previously_restricted = sqlx::query_scalar::<_, bool>( + "select exists( + select 1 from client_api_key_groups where client_api_key_id = $1 + )", + ) + .bind(&key.id) + .fetch_one(&mut *transaction) + .await + .map_err(|source| { + postgres_unavailable("load client API key routing scope for audit", source) + })?; + if previously_restricted && key.group_ids.is_empty() { + audit + .changed_fields + .push("routing_scope:groups->all".to_owned()); + } else if !previously_restricted && !key.group_ids.is_empty() { + audit + .changed_fields + .push("routing_scope:all->groups".to_owned()); + } + update_client_api_key_in_transaction(&mut transaction, &key).await?; + } + ControlPlaneMutation::SetClientApiKeyEnabled { id, enabled } => { + set_client_api_key_enabled_in_transaction(&mut transaction, &id, enabled) + .await?; + } + ControlPlaneMutation::DeleteClientApiKey(id) => { + delete_client_api_key_in_transaction(&mut transaction, &id).await?; + } + ControlPlaneMutation::SetAdminApiKey(key) => { + update_admin_api_key_in_transaction(&mut transaction, key).await?; + } + } + append_admin_audit_event_in_transaction(&mut transaction, audit, revision).await?; + Ok(revision) + } + .await; + match result { + Ok(revision) => { + transaction.commit().await.map_err(|source| { + postgres_unavailable("commit targeted control plane mutation", source) + })?; + Ok(revision) + } + Err(error) => { + let error = match transaction.rollback().await { + Ok(()) => error, + Err(cleanup) => error.with_cleanup(cleanup), + }; + Err(error) + } + } + } +} + +async fn load_control_plane_in_transaction( + transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>, +) -> StoreResult { + let settings = load_runtime_settings_in_transaction(transaction).await?; + Ok(ControlPlaneSnapshot { settings }) +} diff --git a/backend/crates/gateway-store/src/postgres/mod.rs b/backend/crates/gateway-store/src/postgres/mod.rs index 72c086924..aaff5b330 100644 --- a/backend/crates/gateway-store/src/postgres/mod.rs +++ b/backend/crates/gateway-store/src/postgres/mod.rs @@ -1,15 +1,4 @@ -//! PostgreSQL 业务表的 adapters - -use async_trait::async_trait; -use sqlx::{ - PgPool, - postgres::{PgConnectOptions, PgPoolOptions}, -}; - -use crate::{ - POSTGRES_IDLE_TRANSACTION_TIMEOUT, POSTGRES_LOCK_TIMEOUT, POSTGRES_STATEMENT_TIMEOUT, Revision, - StoreError, StorePoolConfig, StoreResult, postgres_unavailable, -}; +//! PostgreSQL 业务表、连接与控制面事务的适配入口 mod account_groups; mod admin_security_audit; @@ -17,6 +6,8 @@ mod admission_recovery; mod backup; mod client_budgets; mod client_keys; +mod connection; +mod control_plane; mod execution; mod execution_buffer; mod observability; @@ -25,6 +16,7 @@ mod plugins; mod pricing; mod provider_accounts; mod proxies; +mod query_pattern; mod retention; mod runtime_settings; mod snapshot; @@ -36,6 +28,11 @@ pub use admission_recovery::*; pub use backup::*; pub use client_budgets::PgClientBudgetStore; pub use client_keys::*; +pub use connection::connect_and_migrate; +pub(crate) use connection::connect_read_only; +pub use control_plane::{ + ControlPlaneReplacement, ControlPlaneRepository, ControlPlaneSnapshot, PgControlPlaneRepository, +}; pub use execution::*; pub use execution_buffer::*; pub use observability::*; @@ -43,6 +40,7 @@ pub use ops_events::*; pub use plugins::PgPluginStore; pub use provider_accounts::*; pub use proxies::PgProxyRepository; +use query_pattern::literal_prefix_pattern; pub use retention::*; pub use runtime_settings::*; pub use snapshot::*; @@ -50,378 +48,3 @@ pub(crate) use usage_facts::{ completed_usage_fact_predicate, push_completed_usage_fact_filter, push_unrecovered_request_filter, }; - -static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../../migrations"); - -/// 建立 PostgreSQL pool 并只执行冻结的 migration 集 -pub async fn connect_and_migrate( - database_url: &str, - pool_config: StorePoolConfig, -) -> StoreResult { - if database_url.trim().is_empty() { - return Err(StoreError::InvalidData { - source: None, - entity: "PostgreSQL configuration", - message: "database URL is empty".to_owned(), - }); - } - pool_config.validate()?; - let connect_options = database_url - .parse::() - .map_err(|source| postgres_unavailable("parse PostgreSQL connection options", source))?; - let migration_pool = PgPoolOptions::new() - .max_connections(1) - .connect_with( - connect_options - .clone() - .application_name("codex-proxy-rs:migration"), - ) - .await - .map_err(|source| postgres_unavailable("connect PostgreSQL for migrations", source))?; - if let Err(error) = MIGRATOR.run(&migration_pool).await { - migration_pool.close().await; - return Err(postgres_unavailable("apply PostgreSQL migrations", error)); - } - migration_pool.close().await; - - connect_pool(connect_options, pool_config, false).await -} - -/// 帮助查询不执行迁移，并用连接默认只读事务阻止意外业务写入 -pub(crate) async fn connect_read_only( - database_url: &str, - pool_config: StorePoolConfig, -) -> StoreResult { - pool_config.validate()?; - let options = database_url - .parse::() - .map_err(|source| postgres_unavailable("parse PostgreSQL connection options", source))?; - connect_pool(options, pool_config, true).await -} - -async fn connect_pool( - connect_options: PgConnectOptions, - pool_config: StorePoolConfig, - read_only: bool, -) -> StoreResult { - let statement_timeout = postgres_duration_setting(POSTGRES_STATEMENT_TIMEOUT); - let lock_timeout = postgres_duration_setting(POSTGRES_LOCK_TIMEOUT); - let idle_in_transaction_session_timeout = - postgres_duration_setting(POSTGRES_IDLE_TRANSACTION_TIMEOUT); - let pool = PgPoolOptions::new() - .max_connections(pool_config.max_connections) - .acquire_timeout(std::time::Duration::from_secs( - pool_config.acquire_timeout_seconds, - )) - .after_connect(move |connection, _metadata| { - let statement_timeout = statement_timeout.clone(); - let lock_timeout = lock_timeout.clone(); - let idle_in_transaction_session_timeout = idle_in_transaction_session_timeout.clone(); - Box::pin(async move { - sqlx::query( - "select set_config('statement_timeout', $1, false), - set_config('lock_timeout', $2, false), - set_config('idle_in_transaction_session_timeout', $3, false), - set_config('default_transaction_read_only', $4, false), - set_config('TimeZone', 'UTC', false)", - ) - .bind(statement_timeout) - .bind(lock_timeout) - .bind(idle_in_transaction_session_timeout) - .bind(if read_only { "on" } else { "off" }) - .execute(connection) - .await?; - Ok(()) - }) - }) - .connect_with(connect_options.application_name("codex-proxy-rs")) - .await - .map_err(|source| postgres_unavailable("connect PostgreSQL", source))?; - Ok(pool) -} - -fn postgres_duration_setting(duration: std::time::Duration) -> String { - format!("{}ms", duration.as_millis()) -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ControlPlaneSnapshot { - pub settings: RuntimeSettings, -} - -#[derive(Debug, Clone)] -pub struct ControlPlaneReplacement { - pub expected_revision: Revision, - pub settings: RuntimeSettingsUpdate, - pub audit: AdminAuditEvent, -} - -#[async_trait] -pub trait ControlPlaneRepository: Send + Sync { - async fn load_control_plane(&self) -> StoreResult; - - async fn replace_control_plane( - &self, - replacement: ControlPlaneReplacement, - ) -> StoreResult; - - /// 更新 admin_api_key 字段并推进 config revision - async fn replace_admin_api_key( - &self, - admin_api_key: Option, - audit: AdminAuditEvent, - ) -> StoreResult; - - async fn create_client_api_key( - &self, - key: NewClientApiKey, - audit: AdminAuditEvent, - ) -> StoreResult; - - async fn update_client_api_key( - &self, - key: UpdateClientApiKeyDetails, - audit: AdminAuditEvent, - ) -> StoreResult; - - async fn set_client_api_key_enabled( - &self, - id: &str, - enabled: bool, - audit: AdminAuditEvent, - ) -> StoreResult; - - async fn delete_client_api_key( - &self, - id: &str, - audit: AdminAuditEvent, - ) -> StoreResult; -} - -#[derive(Clone)] -pub struct PgControlPlaneRepository { - pool: PgPool, -} - -impl PgControlPlaneRepository { - #[must_use] - pub const fn new(pool: PgPool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl ControlPlaneRepository for PgControlPlaneRepository { - async fn load_control_plane(&self) -> StoreResult { - let settings = load_runtime_settings_from_pool(&self.pool).await?; - Ok(ControlPlaneSnapshot { settings }) - } - - async fn replace_control_plane( - &self, - replacement: ControlPlaneReplacement, - ) -> StoreResult { - replacement.settings.validate()?; - let mut transaction = - self.pool.begin().await.map_err(|source| { - postgres_unavailable("begin control plane replacement", source) - })?; - let result: StoreResult<_> = async { - // 锁住所有配置写入共同使用的行，版本检查与修改不能被其他事务穿插 - let current = sqlx::query_scalar::<_, i64>( - "select config_revision from runtime_settings where id = 1 for update", - ) - .fetch_one(&mut *transaction) - .await - .map_err(|source| postgres_unavailable("lock control plane revision", source))?; - if u64::try_from(current).ok() != Some(replacement.expected_revision.get()) { - return Err(StoreError::Conflict { - source: None, - entity: "runtime settings", - id: "1".to_owned(), - kind: crate::ConflictKind::StaleRevision, - }); - } - let revision = - update_runtime_settings_in_transaction(&mut transaction, &replacement.settings) - .await?; - append_admin_audit_event_in_transaction(&mut transaction, replacement.audit, revision) - .await?; - load_control_plane_in_transaction(&mut transaction).await - } - .await; - match result { - Ok(snapshot) => { - transaction.commit().await.map_err(|source| { - postgres_unavailable("commit control plane replacement", source) - })?; - Ok(snapshot) - } - Err(error) => { - let error = match transaction.rollback().await { - Ok(()) => error, - Err(cleanup) => error.with_cleanup(cleanup), - }; - Err(error) - } - } - } - - async fn replace_admin_api_key( - &self, - admin_api_key: Option, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation(ControlPlaneMutation::SetAdminApiKey(admin_api_key), audit) - .await - } - - async fn create_client_api_key( - &self, - key: NewClientApiKey, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation(ControlPlaneMutation::CreateClientApiKey(key), audit) - .await - } - - async fn update_client_api_key( - &self, - key: UpdateClientApiKeyDetails, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation(ControlPlaneMutation::UpdateClientApiKey(key), audit) - .await - } - - async fn set_client_api_key_enabled( - &self, - id: &str, - enabled: bool, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation( - ControlPlaneMutation::SetClientApiKeyEnabled { - id: id.to_owned(), - enabled, - }, - audit, - ) - .await - } - - async fn delete_client_api_key( - &self, - id: &str, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation( - ControlPlaneMutation::DeleteClientApiKey(id.to_owned()), - audit, - ) - .await - } -} - -enum ControlPlaneMutation { - CreateClientApiKey(NewClientApiKey), - UpdateClientApiKey(UpdateClientApiKeyDetails), - SetClientApiKeyEnabled { id: String, enabled: bool }, - DeleteClientApiKey(String), - SetAdminApiKey(Option), -} - -impl PgControlPlaneRepository { - async fn apply_targeted_mutation( - &self, - mutation: ControlPlaneMutation, - mut audit: AdminAuditEvent, - ) -> StoreResult { - let mut transaction = self.pool.begin().await.map_err(|source| { - postgres_unavailable("begin targeted control plane mutation", source) - })?; - let result: StoreResult<_> = async { - let revision = bump_config_revision_in_transaction(&mut transaction).await?; - match mutation { - ControlPlaneMutation::CreateClientApiKey(key) => { - audit.changed_fields.push(if key.group_ids.is_empty() { - "routing_scope:all".to_owned() - } else { - "routing_scope:groups".to_owned() - }); - insert_client_api_key_in_transaction(&mut transaction, &key).await?; - } - ControlPlaneMutation::UpdateClientApiKey(key) => { - let previously_restricted = sqlx::query_scalar::<_, bool>( - "select exists( - select 1 from client_api_key_groups where client_api_key_id = $1 - )", - ) - .bind(&key.id) - .fetch_one(&mut *transaction) - .await - .map_err(|source| { - postgres_unavailable("load client API key routing scope for audit", source) - })?; - if previously_restricted && key.group_ids.is_empty() { - audit - .changed_fields - .push("routing_scope:groups->all".to_owned()); - } else if !previously_restricted && !key.group_ids.is_empty() { - audit - .changed_fields - .push("routing_scope:all->groups".to_owned()); - } - update_client_api_key_in_transaction(&mut transaction, &key).await?; - } - ControlPlaneMutation::SetClientApiKeyEnabled { id, enabled } => { - set_client_api_key_enabled_in_transaction(&mut transaction, &id, enabled) - .await?; - } - ControlPlaneMutation::DeleteClientApiKey(id) => { - delete_client_api_key_in_transaction(&mut transaction, &id).await?; - } - ControlPlaneMutation::SetAdminApiKey(key) => { - update_admin_api_key_in_transaction(&mut transaction, key).await?; - } - } - append_admin_audit_event_in_transaction(&mut transaction, audit, revision).await?; - Ok(revision) - } - .await; - match result { - Ok(revision) => { - transaction.commit().await.map_err(|source| { - postgres_unavailable("commit targeted control plane mutation", source) - })?; - Ok(revision) - } - Err(error) => { - let error = match transaction.rollback().await { - Ok(()) => error, - Err(cleanup) => error.with_cleanup(cleanup), - }; - Err(error) - } - } - } -} - -async fn load_control_plane_in_transaction( - transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>, -) -> StoreResult { - let settings = load_runtime_settings_in_transaction(transaction).await?; - Ok(ControlPlaneSnapshot { settings }) -} - -fn literal_prefix_pattern(value: &str) -> String { - let mut escaped = String::with_capacity(value.len().saturating_add(1)); - for character in value.to_lowercase().chars() { - if matches!(character, '\\' | '%' | '_') { - escaped.push('\\'); - } - escaped.push(character); - } - escaped.push('%'); - escaped -} diff --git a/backend/crates/gateway-store/src/postgres/query_pattern.rs b/backend/crates/gateway-store/src/postgres/query_pattern.rs new file mode 100644 index 000000000..e44a664be --- /dev/null +++ b/backend/crates/gateway-store/src/postgres/query_pattern.rs @@ -0,0 +1,13 @@ +//! 管理实体搜索共用的 SQL 字面前缀转义 + +pub(super) fn literal_prefix_pattern(value: &str) -> String { + let mut escaped = String::with_capacity(value.len().saturating_add(1)); + for character in value.to_lowercase().chars() { + if matches!(character, '\\' | '%' | '_') { + escaped.push('\\'); + } + escaped.push(character); + } + escaped.push('%'); + escaped +} diff --git a/backend/crates/gateway-store/tests/backup/s3.rs b/backend/crates/gateway-store/tests/backup/s3.rs index cf7efee07..d7a4c2461 100644 --- a/backend/crates/gateway-store/tests/backup/s3.rs +++ b/backend/crates/gateway-store/tests/backup/s3.rs @@ -1,9 +1,9 @@ -//! 验证 S3 上传完整发送超过异步文件默认缓冲区的分块 +//! 验证 S3 分片完整读取、并发上限与取消和失败时的请求回收 -use std::io; +use std::{io, time::Duration}; use chrono::Utc; -use gateway_admin::model::backup::{BackupObjectMetadata, BackupStorageConfig}; +use gateway_admin::model::backup::{BackupError, BackupObjectMetadata, BackupStorageConfig, code}; use gateway_admin::ports::backup::{BackupObjectStorePort, UploadObjectRequest}; use gateway_core::lifecycle::CancellationToken; use gateway_store::backup::s3::S3ObjectStoreAdapter; @@ -12,6 +12,7 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; use tokio::net::{TcpListener, TcpStream}; const PART_SIZE_ABOVE_TOKIO_DEFAULT: usize = 2 * 1024 * 1024 + 1; +const MULTIPART_SIZE: usize = 16 * 1024 * 1024; #[tokio::test] async fn upload_file_sends_full_part_beyond_tokio_default_buffer_limit() { @@ -64,6 +65,223 @@ async fn upload_file_sends_full_part_beyond_tokio_default_buffer_limit() { assert_eq!(uploaded_len, PART_SIZE_ABOVE_TOKIO_DEFAULT); } +#[tokio::test] +async fn cancelled_upload_does_not_retry_a_pending_part_after_returning() { + let probe = start_upload(MULTIPART_SIZE).await; + let (mut part, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + + probe.cancellation.cancel(); + let (mut abort, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("DELETE ") && request.contains("uploadId=")); + respond(&mut abort, &xml_response("")).await; + let error = tokio::time::timeout(Duration::from_secs(3), probe.upload) + .await + .expect("cancelled upload returns without waiting for the part response") + .expect("upload task did not panic") + .expect_err("upload was cancelled"); + assert_eq!(error.code(), code::CANCELLED); + + assert_no_part_retry(&probe.listener, &mut part).await; +} + +#[tokio::test] +async fn a_failed_part_stops_other_pending_parts_before_upload_returns() { + let probe = start_upload(MULTIPART_SIZE * 2).await; + let (mut failed, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + let (mut pending, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + respond( + &mut failed, + &error_response("403 Forbidden", "AccessDenied"), + ) + .await; + + // 另一分片仍未收到响应，首次失败必须立即撤销它并进入 multipart 清理 + let (mut abort, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("DELETE ") && request.contains("uploadId=")); + respond(&mut abort, &xml_response("")).await; + let (mut delete, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("DELETE ") && !request.contains("uploadId=")); + respond(&mut delete, &xml_response("")).await; + let error = tokio::time::timeout(Duration::from_secs(3), probe.upload) + .await + .expect("failed upload returns without waiting for the other part") + .expect("upload task did not panic") + .expect_err("part upload failed"); + assert_eq!(error.code(), code::S3_PERMISSION_DENIED); + + assert_no_part_retry(&probe.listener, &mut pending).await; +} + +#[tokio::test] +async fn dropping_the_upload_future_stops_pending_part_retries() { + let probe = start_upload(MULTIPART_SIZE).await; + let (mut part, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + + probe.upload.abort(); + assert!( + probe + .upload + .await + .expect_err("parent task was aborted") + .is_cancelled() + ); + + assert_no_part_retry(&probe.listener, &mut part).await; +} + +#[tokio::test] +async fn multipart_upload_bounds_concurrency_and_completes_parts_in_order() { + let probe = start_upload(MULTIPART_SIZE * 4 + 1).await; + let mut pending = Vec::new(); + for _ in 0..4 { + let (part, request, body) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + assert_eq!(body.len(), MULTIPART_SIZE); + pending.push(part); + } + assert!( + tokio::time::timeout(Duration::from_millis(150), probe.listener.accept()) + .await + .is_err(), + "a fifth part must wait for one of the four active uploads" + ); + respond(&mut pending.pop().unwrap(), part_response()).await; + let (mut last, request, body) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ") && request.contains("partNumber=5")); + assert_eq!(body.len(), 1); + respond(&mut last, part_response()).await; + for mut part in pending.into_iter().rev() { + respond(&mut part, part_response()).await; + } + let (mut complete, request, body) = accept_request(&probe.listener).await; + assert!(request.starts_with("POST ") && request.contains("uploadId=")); + let body = String::from_utf8(body).unwrap(); + let positions: Vec<_> = (1..=5) + .map(|part| { + body.find(&format!("{part}")) + .unwrap() + }) + .collect(); + assert!(positions.windows(2).all(|pair| pair[0] < pair[1])); + respond( + &mut complete, + &xml_response( + "\"probe-part\"", + ), + ) + .await; + tokio::time::timeout(Duration::from_secs(3), probe.upload) + .await + .expect("multipart upload completes") + .expect("upload task did not panic") + .expect("all parts uploaded"); +} + +struct UploadProbe { + listener: TcpListener, + cancellation: CancellationToken, + upload: tokio::task::JoinHandle>, + _directory: tempfile::TempDir, +} + +async fn start_upload(size: usize) -> UploadProbe { + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("archive.dump"); + let file = tokio::fs::File::create(&source).await.unwrap(); + file.set_len(size as u64).await.unwrap(); + drop(file); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let storage = BackupStorageConfig { + storage_revision: 1, + endpoint: format!("http://{}", listener.local_addr().unwrap()), + region: "auto".into(), + bucket: "backup-bucket".into(), + access_key_id: "test-access-key".into(), + secret_access_key: SecretString::from("test-secret-key"), + prefix: "backups".into(), + force_path_style: true, + }; + let cancellation = CancellationToken::new(); + let request = UploadObjectRequest { + object_key: "backups/archive.dump".into(), + source, + metadata: BackupObjectMetadata::new( + "backup_test".into(), + "a".repeat(64), + Utc::now(), + size as u64, + ), + cancellation: cancellation.clone(), + }; + let upload = tokio::spawn(async move { + S3ObjectStoreAdapter::new() + .upload_file(&storage, request) + .await + }); + let (mut create, request, _) = accept_request(&listener).await; + assert!(request.starts_with("POST ") && request.contains("uploads")); + respond( + &mut create, + &xml_response( + "probe-upload", + ), + ) + .await; + UploadProbe { + listener, + cancellation, + upload, + _directory: directory, + } +} + +async fn accept_request(listener: &TcpListener) -> (TcpStream, String, Vec) { + tokio::time::timeout(Duration::from_secs(3), async { + let (mut stream, _) = listener.accept().await.expect("accept S3 request"); + let (request, body) = read_http_request(&mut stream) + .await + .expect("read S3 request"); + (stream, request, body) + }) + .await + .expect("expected S3 request arrives without waiting for another part") +} + +async fn respond(stream: &mut TcpStream, response: &str) { + stream.write_all(response.as_bytes()).await.unwrap(); + stream.shutdown().await.unwrap(); +} + +async fn assert_no_part_retry(listener: &TcpListener, pending: &mut TcpStream) { + // 旧请求可能已被关闭，写入失败是正常撤销结果；仍活动的 SDK 调用收到 503 后会重试 + let _ = pending + .write_all(error_response("503 Service Unavailable", "SlowDown").as_bytes()) + .await; + let _ = pending.shutdown().await; + assert!( + tokio::time::timeout(Duration::from_secs(2), listener.accept()) + .await + .is_err(), + "a retired part must not issue another UploadPart request" + ); +} + +fn part_response() -> &'static str { + "HTTP/1.1 200 OK\r\ncontent-length: 0\r\nconnection: close\r\netag: \"probe-part\"\r\n\r\n" +} + +fn error_response(status: &str, code: &str) -> String { + let body = format!("{code}probe error"); + format!( + "HTTP/1.1 {status}\r\ncontent-type: application/xml\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ) +} + async fn serve_multipart_probe(listener: TcpListener) -> io::Result { let mut uploaded_len = None; for _ in 0..3 { diff --git a/backend/crates/gateway-store/tests/postgres/client_budgets.rs b/backend/crates/gateway-store/tests/postgres/client_budgets.rs index 2704c998a..6fcbba9e2 100644 --- a/backend/crates/gateway-store/tests/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/tests/postgres/client_budgets.rs @@ -76,6 +76,41 @@ fn context() -> MutationContext { } } +#[tokio::test] +async fn unopened_future_budget_windows_do_not_report_usage_or_reset_times() { + let Some(database) = TestDatabase::create("budget_unopened_future").await else { + return; + }; + seed(&database, "key", "10", "20").await; + PgClientBudgetStore::new(database.pool.clone()) + .settle(charge("key", "before", "3")) + .await + .unwrap(); + // 起止相等表示窗口未开启；未来边界确定性覆盖应用时钟领先数据库的场景 + sqlx::query( + "update client_key_budget_windows set + daily_start = now() + interval '1 hour', daily_end = now() + interval '1 hour', + weekly_start = now() + interval '1 hour', weekly_end = now() + interval '1 hour' + where client_api_key_id = $1", + ) + .bind("key") + .execute(&database.pool) + .await + .unwrap(); + + let budget = status(&database, "key").await; + assert_eq!( + ( + budget.daily_used_usd.canonical(), + budget.weekly_used_usd.canonical(), + budget.daily_resets_at, + budget.weekly_resets_at, + ), + ("0".to_owned(), "0".to_owned(), None, None) + ); + database.close().await; +} + #[tokio::test] async fn manual_reset_clears_selected_windows_and_preserves_policy_and_history() { let Some(database) = TestDatabase::create("budget_manual_reset").await else { diff --git a/backend/crates/gateway-store/tests/postgres/connection.rs b/backend/crates/gateway-store/tests/postgres/connection.rs index 9fd564814..3ef4a780c 100644 --- a/backend/crates/gateway-store/tests/postgres/connection.rs +++ b/backend/crates/gateway-store/tests/postgres/connection.rs @@ -1,7 +1,9 @@ -//! 验证实际基础设施错误在 Store 边界保留来源而不进入普通格式化 +//! PostgreSQL 连接错误、迁移应用与事务归属测试 use std::error::Error as _; +use super::*; + #[tokio::test] async fn invalid_postgres_options_preserve_the_sqlx_source() { let error = gateway_store::postgres::connect_and_migrate( @@ -21,3 +23,188 @@ async fn invalid_postgres_options_preserve_the_sqlx_source() { ); assert!(!format!("{error:?} {error}").contains("invalid-port")); } + +#[tokio::test] +async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly() { + let Some(database_url) = crate::support::test_env("CPR_TEST_DATABASE_URL") else { + return; + }; + let database = format!("cpr_store_migrator_{}", Uuid::new_v4().simple()); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("connect migration test PostgreSQL"); + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + "create database \"{database}\"" + ))) + .execute(&admin) + .await + .expect("create migration test database"); + + let isolated_url = PgConnectOptions::from_str(&database_url) + .expect("parse migration test PostgreSQL URL") + .database(&database) + .to_url_lossy() + .to_string(); + let pool_config = gateway_store::StorePoolConfig::default(); + let first = connect_and_migrate(&isolated_url, pool_config) + .await + .expect("apply migrations through production migrator"); + let session_settings = sqlx::query_as::<_, (String, i64, i64, i64)>( + "select current_setting('application_name'), + extract(epoch from current_setting('statement_timeout')::interval)::bigint, + extract(epoch from current_setting('lock_timeout')::interval)::bigint, + extract(epoch from current_setting( + 'idle_in_transaction_session_timeout' + )::interval)::bigint", + ) + .fetch_one(&first) + .await + .expect("load runtime PostgreSQL session settings"); + let timezone: String = sqlx::query_scalar("show time zone") + .fetch_one(&first) + .await + .unwrap(); + assert_eq!( + timezone, "UTC", + "runtime sessions use an explicit technical time basis" + ); + let first_tables = sqlx::query_scalar::<_, String>( + "select table_name + from information_schema.tables + where table_schema = 'public' + order by table_name", + ) + .fetch_all(&first) + .await + .expect("load migrated tables"); + first.close().await; + + let second = connect_and_migrate(&isolated_url, pool_config) + .await + .expect("reopen database through production migrator"); + let migration_count = + sqlx::query_scalar::<_, i64>("select count(*) from _sqlx_migrations where success") + .fetch_one(&second) + .await + .expect("count successful migrations"); + let response_id_types = sqlx::query_scalar::<_, String>( + "select data_type + from information_schema.columns + where table_schema = 'public' + and table_name = 'model_requests' + and column_name in ('client_response_id', 'upstream_response_id') + order by column_name", + ) + .fetch_all(&second) + .await + .expect("load opaque response ID column types"); + let raw_response_id_index_exists = sqlx::query_scalar::<_, bool>( + "select exists ( + select 1 + from pg_indexes + where schemaname = 'public' + and indexname = 'model_requests_client_response_uq' + )", + ) + .fetch_one(&second) + .await + .expect("check removed raw response ID index"); + let legacy_key_provider_column_exists = sqlx::query_scalar::<_, bool>( + "select exists ( + select 1 from information_schema.columns + where table_schema = 'public' + and table_name = 'client_api_keys' + and column_name = 'provider_kind' + )", + ) + .fetch_one(&second) + .await + .expect("check removed client key provider column"); + let routing_history_columns = sqlx::query_scalar::<_, String>( + "select column_name from information_schema.columns + where table_schema = 'public' + and table_name = 'model_requests' + and column_name in ( + 'routing_scope', 'routing_group_refs', 'routing_group_names_snapshot' + ) + order by column_name", + ) + .fetch_all(&second) + .await + .expect("load routing history columns"); + second.close().await; + + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + "drop database \"{database}\" with (force)" + ))) + .execute(&admin) + .await + .expect("drop migration test database"); + admin.close().await; + + assert_eq!( + first_tables, + [ + "_sqlx_migrations", + "account_group_accounts", + "account_groups", + "admin_audit_events", + "admin_users", + "authorization_receipts", + "backup_records", + "backup_settings", + "client_api_key_groups", + "client_api_keys", + "client_key_budget_windows", + "client_key_charge_events", + "model_requests", + "ops_events", + "outbound_proxies", + "plugin_artifact_credentials", + "plugin_artifact_platforms", + "plugin_artifacts", + "plugin_group_resources", + "plugin_instance_secrets", + "plugin_instances", + "plugin_key_resources", + "plugin_source_credentials", + "plugin_state_generations", + "plugin_state_records", + "plugin_update_sources", + "plugin_version_configurations", + "provider_accounts", + "runtime_settings", + ] + ); + assert_eq!(session_settings, ("codex-proxy-rs".to_owned(), 30, 5, 30)); + assert_eq!( + migration_count, + i64::try_from(TEST_MIGRATOR.iter().count()) + .expect("migration count fits PostgreSQL bigint") + ); + assert_eq!(response_id_types, ["bytea", "bytea"]); + assert!(!raw_response_id_index_exists); + assert!(!legacy_key_provider_column_exists); + assert_eq!( + routing_history_columns, + [ + "routing_group_names_snapshot", + "routing_group_refs", + "routing_scope", + ] + ); +} + +#[test] +fn migrations_should_leave_transaction_ownership_to_sqlx() { + let transaction_statements = TEST_MIGRATOR + .iter() + .flat_map(|migration| migration.sql.as_str().lines()) + .map(str::trim) + .filter(|line| matches!(*line, "begin;" | "commit;")) + .count(); + + assert_eq!(transaction_statements, 0); +} diff --git a/backend/crates/gateway-store/tests/postgres/control_plane.rs b/backend/crates/gateway-store/tests/postgres/control_plane.rs new file mode 100644 index 000000000..53dd10cd4 --- /dev/null +++ b/backend/crates/gateway-store/tests/postgres/control_plane.rs @@ -0,0 +1,104 @@ +//! 控制面替换事务的修订冲突与设置保留测试 + +use chrono::Utc; + +use super::{TestDatabase, runtime_settings::settings_with_margin}; + +#[tokio::test] +async fn control_plane_replacement_commits_one_writer_per_revision_and_preserves_newer_settings() { + use gateway_store::postgres::{ + AdminAuditActorKind, AdminAuditEvent, ControlPlaneReplacement, ControlPlaneRepository, + PgControlPlaneRepository, + }; + use gateway_store::{ConflictKind, StoreError}; + let Some(database) = TestDatabase::create("settings_compare_replace").await else { + return; + }; + let repository = PgControlPlaneRepository::new(database.pool.clone()); + let revision = repository + .load_control_plane() + .await + .unwrap() + .settings + .config_revision; + let replacement = |id: &str, margin| ControlPlaneReplacement { + expected_revision: revision, + settings: settings_with_margin(margin), + audit: AdminAuditEvent { + id: id.into(), + actor_kind: AdminAuditActorKind::System, + actor_admin_user_id: None, + actor_ref: "system".into(), + admin_request_id: Some(id.into()), + action: "settings.replace".into(), + entity_kind: "runtime_settings".into(), + entity_ref: "1".into(), + config_revision: None, + changed_fields: vec!["refresh_margin_seconds".into()], + created_at: Utc::now(), + }, + }; + let (first, second) = tokio::join!( + repository.replace_control_plane(replacement("first", 1800)), + repository.replace_control_plane(replacement("second", 7200)), + ); + let (saved, conflict) = match (first, second) { + (Ok(saved), Err(error)) | (Err(error), Ok(saved)) => (saved.settings, error), + other => panic!("expected exactly one successful transaction: {other:?}"), + }; + assert!(matches!( + conflict, + StoreError::Conflict { + kind: ConflictKind::StaleRevision, + .. + } + )); + let current = repository.load_control_plane().await.unwrap().settings; + assert_eq!(current.config_revision.get(), revision.get() + 1); + assert_eq!( + current.values.refresh_margin_seconds, + saved.values.refresh_margin_seconds + ); + let audit_count: i64 = sqlx::query_scalar( + "select count(*) from admin_audit_events where action = 'settings.replace'", + ) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(audit_count, 1); + // API Key 更新也推进相同版本，旧设置快照不能复活已经替换的 Key + let mut key_audit = replacement("key", 3600).audit; + key_audit.action = "settings.admin_key".into(); + repository + .replace_admin_api_key(Some("new-test-key".into()), key_audit) + .await + .unwrap(); + let mut stale = replacement("stale", 3600); + stale.expected_revision = saved.config_revision; + assert!(matches!( + repository.replace_control_plane(stale).await, + Err(StoreError::Conflict { + kind: ConflictKind::StaleRevision, + .. + }) + )); + let mut fresh = replacement("fresh", 3600); + fresh.expected_revision = repository + .load_control_plane() + .await + .unwrap() + .settings + .config_revision; + repository.replace_control_plane(fresh).await.unwrap(); + assert_eq!( + repository + .load_control_plane() + .await + .unwrap() + .settings + .admin_api_key + .as_deref(), + Some("new-test-key") + ); + database.close().await; +} diff --git a/backend/crates/gateway-store/tests/postgres/mod.rs b/backend/crates/gateway-store/tests/postgres/mod.rs index 1c26a49ef..2bb898ece 100644 --- a/backend/crates/gateway-store/tests/postgres/mod.rs +++ b/backend/crates/gateway-store/tests/postgres/mod.rs @@ -1,4 +1,4 @@ -//! PostgreSQL 适配测试入口，以及迁移应用和事务归属测试 +//! PostgreSQL 适配测试入口与隔离数据库夹具 use std::{str::FromStr, time::Duration}; @@ -19,6 +19,7 @@ mod backup; mod client_budgets; mod client_keys; mod connection; +mod control_plane; mod execution; mod execution_buffer; mod health; @@ -110,188 +111,3 @@ impl TestDatabase { self.admin.close().await; } } - -#[tokio::test] -async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly() { - let Some(database_url) = crate::support::test_env("CPR_TEST_DATABASE_URL") else { - return; - }; - let database = format!("cpr_store_migrator_{}", Uuid::new_v4().simple()); - let admin = PgPoolOptions::new() - .max_connections(1) - .connect(&database_url) - .await - .expect("connect migration test PostgreSQL"); - sqlx::raw_sql(sqlx::AssertSqlSafe(format!( - "create database \"{database}\"" - ))) - .execute(&admin) - .await - .expect("create migration test database"); - - let isolated_url = PgConnectOptions::from_str(&database_url) - .expect("parse migration test PostgreSQL URL") - .database(&database) - .to_url_lossy() - .to_string(); - let pool_config = gateway_store::StorePoolConfig::default(); - let first = connect_and_migrate(&isolated_url, pool_config) - .await - .expect("apply migrations through production migrator"); - let session_settings = sqlx::query_as::<_, (String, i64, i64, i64)>( - "select current_setting('application_name'), - extract(epoch from current_setting('statement_timeout')::interval)::bigint, - extract(epoch from current_setting('lock_timeout')::interval)::bigint, - extract(epoch from current_setting( - 'idle_in_transaction_session_timeout' - )::interval)::bigint", - ) - .fetch_one(&first) - .await - .expect("load runtime PostgreSQL session settings"); - let timezone: String = sqlx::query_scalar("show time zone") - .fetch_one(&first) - .await - .unwrap(); - assert_eq!( - timezone, "UTC", - "runtime sessions use an explicit technical time basis" - ); - let first_tables = sqlx::query_scalar::<_, String>( - "select table_name - from information_schema.tables - where table_schema = 'public' - order by table_name", - ) - .fetch_all(&first) - .await - .expect("load migrated tables"); - first.close().await; - - let second = connect_and_migrate(&isolated_url, pool_config) - .await - .expect("reopen database through production migrator"); - let migration_count = - sqlx::query_scalar::<_, i64>("select count(*) from _sqlx_migrations where success") - .fetch_one(&second) - .await - .expect("count successful migrations"); - let response_id_types = sqlx::query_scalar::<_, String>( - "select data_type - from information_schema.columns - where table_schema = 'public' - and table_name = 'model_requests' - and column_name in ('client_response_id', 'upstream_response_id') - order by column_name", - ) - .fetch_all(&second) - .await - .expect("load opaque response ID column types"); - let raw_response_id_index_exists = sqlx::query_scalar::<_, bool>( - "select exists ( - select 1 - from pg_indexes - where schemaname = 'public' - and indexname = 'model_requests_client_response_uq' - )", - ) - .fetch_one(&second) - .await - .expect("check removed raw response ID index"); - let legacy_key_provider_column_exists = sqlx::query_scalar::<_, bool>( - "select exists ( - select 1 from information_schema.columns - where table_schema = 'public' - and table_name = 'client_api_keys' - and column_name = 'provider_kind' - )", - ) - .fetch_one(&second) - .await - .expect("check removed client key provider column"); - let routing_history_columns = sqlx::query_scalar::<_, String>( - "select column_name from information_schema.columns - where table_schema = 'public' - and table_name = 'model_requests' - and column_name in ( - 'routing_scope', 'routing_group_refs', 'routing_group_names_snapshot' - ) - order by column_name", - ) - .fetch_all(&second) - .await - .expect("load routing history columns"); - second.close().await; - - sqlx::raw_sql(sqlx::AssertSqlSafe(format!( - "drop database \"{database}\" with (force)" - ))) - .execute(&admin) - .await - .expect("drop migration test database"); - admin.close().await; - - assert_eq!( - first_tables, - [ - "_sqlx_migrations", - "account_group_accounts", - "account_groups", - "admin_audit_events", - "admin_users", - "authorization_receipts", - "backup_records", - "backup_settings", - "client_api_key_groups", - "client_api_keys", - "client_key_budget_windows", - "client_key_charge_events", - "model_requests", - "ops_events", - "outbound_proxies", - "plugin_artifact_credentials", - "plugin_artifact_platforms", - "plugin_artifacts", - "plugin_group_resources", - "plugin_instance_secrets", - "plugin_instances", - "plugin_key_resources", - "plugin_source_credentials", - "plugin_state_generations", - "plugin_state_records", - "plugin_update_sources", - "plugin_version_configurations", - "provider_accounts", - "runtime_settings", - ] - ); - assert_eq!(session_settings, ("codex-proxy-rs".to_owned(), 30, 5, 30)); - assert_eq!( - migration_count, - i64::try_from(TEST_MIGRATOR.iter().count()) - .expect("migration count fits PostgreSQL bigint") - ); - assert_eq!(response_id_types, ["bytea", "bytea"]); - assert!(!raw_response_id_index_exists); - assert!(!legacy_key_provider_column_exists); - assert_eq!( - routing_history_columns, - [ - "routing_group_names_snapshot", - "routing_group_refs", - "routing_scope", - ] - ); -} - -#[test] -fn migrations_should_leave_transaction_ownership_to_sqlx() { - let transaction_statements = TEST_MIGRATOR - .iter() - .flat_map(|migration| migration.sql.as_str().lines()) - .map(str::trim) - .filter(|line| matches!(*line, "begin;" | "commit;")) - .count(); - - assert_eq!(transaction_statements, 0); -} diff --git a/backend/crates/gateway-store/tests/postgres/runtime_settings.rs b/backend/crates/gateway-store/tests/postgres/runtime_settings.rs index 2c331d64f..8f2adee0b 100644 --- a/backend/crates/gateway-store/tests/postgres/runtime_settings.rs +++ b/backend/crates/gateway-store/tests/postgres/runtime_settings.rs @@ -9,7 +9,7 @@ use gateway_store::postgres::{ use super::TestDatabase; -fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate { +pub(super) fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate { RuntimeSettingsUpdate { request_profile_updates: BTreeMap::new(), rotation_strategy: "smart".to_owned(), @@ -942,105 +942,6 @@ async fn request_profile_projection_is_revision_consistent_and_includes_key_over database.close().await; } -#[tokio::test] -async fn control_plane_replacement_commits_one_writer_per_revision_and_preserves_newer_settings() { - use gateway_store::postgres::{ - AdminAuditActorKind, AdminAuditEvent, ControlPlaneReplacement, ControlPlaneRepository, - PgControlPlaneRepository, - }; - use gateway_store::{ConflictKind, StoreError}; - let Some(database) = TestDatabase::create("settings_compare_replace").await else { - return; - }; - let repository = PgControlPlaneRepository::new(database.pool.clone()); - let revision = repository - .load_control_plane() - .await - .unwrap() - .settings - .config_revision; - let replacement = |id: &str, margin| ControlPlaneReplacement { - expected_revision: revision, - settings: settings_with_margin(margin), - audit: AdminAuditEvent { - id: id.into(), - actor_kind: AdminAuditActorKind::System, - actor_admin_user_id: None, - actor_ref: "system".into(), - admin_request_id: Some(id.into()), - action: "settings.replace".into(), - entity_kind: "runtime_settings".into(), - entity_ref: "1".into(), - config_revision: None, - changed_fields: vec!["refresh_margin_seconds".into()], - created_at: Utc::now(), - }, - }; - let (first, second) = tokio::join!( - repository.replace_control_plane(replacement("first", 1800)), - repository.replace_control_plane(replacement("second", 7200)), - ); - let (saved, conflict) = match (first, second) { - (Ok(saved), Err(error)) | (Err(error), Ok(saved)) => (saved.settings, error), - other => panic!("expected exactly one successful transaction: {other:?}"), - }; - assert!(matches!( - conflict, - StoreError::Conflict { - kind: ConflictKind::StaleRevision, - .. - } - )); - let current = repository.load_control_plane().await.unwrap().settings; - assert_eq!(current.config_revision.get(), revision.get() + 1); - assert_eq!( - current.values.refresh_margin_seconds, - saved.values.refresh_margin_seconds - ); - let audit_count: i64 = sqlx::query_scalar( - "select count(*) from admin_audit_events where action = 'settings.replace'", - ) - .fetch_one(&database.pool) - .await - .unwrap(); - assert_eq!(audit_count, 1); - // API Key 更新也推进相同版本，旧设置快照不能复活已经替换的 Key - let mut key_audit = replacement("key", 3600).audit; - key_audit.action = "settings.admin_key".into(); - repository - .replace_admin_api_key(Some("new-test-key".into()), key_audit) - .await - .unwrap(); - let mut stale = replacement("stale", 3600); - stale.expected_revision = saved.config_revision; - assert!(matches!( - repository.replace_control_plane(stale).await, - Err(StoreError::Conflict { - kind: ConflictKind::StaleRevision, - .. - }) - )); - let mut fresh = replacement("fresh", 3600); - fresh.expected_revision = repository - .load_control_plane() - .await - .unwrap() - .settings - .config_revision; - repository.replace_control_plane(fresh).await.unwrap(); - assert_eq!( - repository - .load_control_plane() - .await - .unwrap() - .settings - .admin_api_key - .as_deref(), - Some("new-test-key") - ); - database.close().await; -} - #[tokio::test] async fn warmup_cursor_survives_restart_and_settings_updates_without_rewinding() { use gateway_core::{provider_ports::ProviderRuntimePolicyPort, time::DeploymentTimeZone}; diff --git a/backend/crates/providers/openai/src/credential/affinity.rs b/backend/crates/providers/openai/src/credential/affinity.rs index 739ec6077..4095487e7 100644 --- a/backend/crates/providers/openai/src/credential/affinity.rs +++ b/backend/crates/providers/openai/src/credential/affinity.rs @@ -6,6 +6,7 @@ use gateway_core::provider_ports::{ProviderSessionAffinityKey, ProviderSessionAl use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; +use crate::request_identity::{account_session_with_headers, follows_session_with_headers}; use crate::transport::protocol::responses::CodexResponsesRequest; use crate::transport::request::derive_conversation_anchor; @@ -171,80 +172,6 @@ pub(crate) fn derive_endpoint_affinity_with_headers( }) } -/// 只有能识别出的后代线程禁止迁移会话；没有身份的普通客户端沿用现有调度 -pub(crate) fn follows_session_with_headers( - body: &Map, - context: &Map, - headers: &[gateway_core::engine::middleware::MiddlewareHeader], -) -> bool { - let metadata = headers - .iter() - .find(|h| h.name().eq_ignore_ascii_case("x-codex-turn-metadata")) - .and_then(|h| std::str::from_utf8(h.value()).ok()); - let thread = account_thread_with_headers(body, context, headers); - if let Some((session, thread)) = - account_session_with_headers(body, context, headers).zip(thread) - { - return session != thread; - } - gateway_protocol::openai::codex_responses_request_semantics_with_turn_metadata( - body, - metadata.or_else(|| context.get("turn_metadata").and_then(Value::as_str)), - ) - .subagent_kind - .is_some() -} - -fn account_thread_with_headers( - body: &Map, - context: &Map, - headers: &[gateway_core::engine::middleware::MiddlewareHeader], -) -> Option { - let header = |name: &str| { - headers - .iter() - .find(|h| h.name().eq_ignore_ascii_case(name)) - .and_then(|h| std::str::from_utf8(h.value()).ok()) - }; - header("x-codex-turn-metadata") - .and_then(gateway_protocol::openai::turn_metadata_thread_id) - .or_else(|| { - header("thread-id") - .and_then(|value| non_empty(Some(value))) - .map(str::to_owned) - }) - .or_else(|| gateway_protocol::openai::codex_account_thread_id(body, context)) -} - -/// 中间件最终请求头按实际传输的覆盖语义参与身份复验 -pub(crate) fn account_session_with_headers( - body: &Map, - context: &Map, - headers: &[gateway_core::engine::middleware::MiddlewareHeader], -) -> Option { - let header = |name: &str| { - headers - .iter() - .find(|h| h.name().eq_ignore_ascii_case(name)) - .and_then(|h| std::str::from_utf8(h.value()).ok()) - }; - if let Some(session) = - header("x-codex-turn-metadata").and_then(gateway_protocol::openai::turn_metadata_session_id) - { - return Some(session); - } - let mut context = context.clone(); - for (name, field) in [ - ("session-id", "session_id"), - ("x-codex-turn-metadata", "turn_metadata"), - ] { - if let Some(value) = header(name) { - context.insert(field.to_owned(), Value::String(value.to_owned())); - } - } - gateway_protocol::openai::codex_account_session_id(body, &context) -} - pub(crate) fn derive_live_session_affinity( request: &gateway_core::operation::ProviderHttpRequest, headers: &[gateway_core::engine::middleware::MiddlewareHeader], @@ -376,24 +303,3 @@ pub(crate) fn derive_turn_alias( let turn_id = non_empty(Some(turn_id))?; opaque_affinity_key("client-turn", &format!("{}\0{turn_id}", client.as_str())) } - -pub(crate) fn turn_id_with_headers( - body: &Map, - context: &Map, - headers: &[gateway_core::engine::middleware::MiddlewareHeader], -) -> Option { - headers - .iter() - .find(|header| header.name().eq_ignore_ascii_case("x-codex-turn-metadata")) - .and_then(|header| std::str::from_utf8(header.value()).ok()) - .and_then(gateway_protocol::openai::turn_metadata_turn_id) - .or_else(|| { - headers - .iter() - .find(|header| header.name().eq_ignore_ascii_case("x-client-turn-id")) - .and_then(|header| std::str::from_utf8(header.value()).ok()) - .and_then(|value| non_empty(Some(value))) - .map(str::to_owned) - }) - .or_else(|| gateway_protocol::openai::codex_turn_id(body, context)) -} diff --git a/backend/crates/providers/openai/src/credential/api_key.rs b/backend/crates/providers/openai/src/credential/api_key.rs index e2e4e3eca..7678cc2d5 100644 --- a/backend/crates/providers/openai/src/credential/api_key.rs +++ b/backend/crates/providers/openai/src/credential/api_key.rs @@ -5,7 +5,7 @@ use std::fmt; use secrecy::SecretString; use serde::{Deserialize, Serialize}; -pub const CODEX_AUTHENTICATION_KIND_API_KEY: &str = "api_key"; +pub(super) use crate::CODEX_AUTHENTICATION_KIND_API_KEY; use super::types::ResponsesTransport; diff --git a/backend/crates/providers/openai/src/credential/catalog.rs b/backend/crates/providers/openai/src/credential/catalog.rs index eb9f03a27..aad85e27d 100644 --- a/backend/crates/providers/openai/src/credential/catalog.rs +++ b/backend/crates/providers/openai/src/credential/catalog.rs @@ -745,7 +745,7 @@ impl CodexCredentialCatalogService { .map_err(|error| CodexCredentialCatalogError::Upstream { detail: error.to_string(), })? - .with_authentication(&credential.authentication) + .with_responses_api_base_url(credential.authentication.responses_api_base_url()) .fetch_models_with_context( CodexRequestContext::auxiliary( authorization.expose_secret(), diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index 435044416..5f85bd933 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -17,19 +17,17 @@ pub mod token_client; mod types; pub(crate) use affinity::{ - CodexSessionAffinity, account_session_with_headers, derive_codex_cyber_policy_session_key, + CodexSessionAffinity, derive_codex_cyber_policy_session_key, derive_codex_endpoint_session_affinity, derive_codex_session_affinity, derive_codex_transport_key, derive_endpoint_affinity_with_headers, derive_live_session_affinity, derive_previous_response_id_hash, derive_turn_alias, - follows_session_with_headers, turn_id_with_headers, }; pub(crate) use oauth::oauth_owner_ref; pub(crate) use types::parse_access_token_expiration; -pub use api_key::{ - ApiKeyAuthentication, ApiKeyConfiguration, ApiKeyCredentialData, - CODEX_AUTHENTICATION_KIND_API_KEY, -}; +pub(crate) use crate::CODEX_AUTHENTICATION_KIND_API_KEY; + +pub use api_key::{ApiKeyAuthentication, ApiKeyConfiguration, ApiKeyCredentialData}; pub use admin::{ CodexCprExportDocument, CodexCredentialAdmin, CodexCredentialAdminError, diff --git a/backend/crates/providers/openai/src/credential/quota/document.rs b/backend/crates/providers/openai/src/credential/quota/document.rs index 8301a07e8..1ab272832 100644 --- a/backend/crates/providers/openai/src/credential/quota/document.rs +++ b/backend/crates/providers/openai/src/credential/quota/document.rs @@ -172,3 +172,75 @@ fn normalize_limit_id(value: &str) -> Option { (!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)) .then(|| value.to_ascii_lowercase().replace('-', "_")) } + +/// 丢弃 core 限额中上游给出的无事实 `secondary_window` 占位 +/// +/// 生产响应头可能携带 `used_percent=0`、零时长和空 reset 的占位，`/usage` 也会 +/// 返回 `secondary_window: null` +/// 正常被动同步保留该响应事实，Admin 展示会将其 +/// 隐藏；但主动 `/usage` 刷新或 402 确认投影会把一个存在但无事实的字段写成 +/// 100%，从而显示并不存在的“次级额度” +/// 因此这些非被动写入口在落库前移除 +/// 无事实值 +/// 带 reset、时长、正用量、触顶、未知或非法字段的次级窗口均完全按 +/// 原有额度逻辑保留 +pub(super) fn normalize_quota_window_placeholders( + mut quota: Map, +) -> Map { + quota = canonicalize_rate_limit_document(quota); + if let Some(rate_limit) = quota + .get_mut(RATE_LIMITS_BY_LIMIT_ID) + .and_then(Value::as_object_mut) + .and_then(|limits| limits.get_mut(DEFAULT_CODEX_LIMIT_ID)) + .and_then(Value::as_object_mut) + { + drop_secondary_window_placeholder(rate_limit); + } + quota +} + +fn drop_secondary_window_placeholder(rate_limit: &mut Map) { + let placeholder = rate_limit.get("secondary_window").is_some_and(|window| { + window.is_null() + || window + .as_object() + .is_some_and(secondary_window_is_placeholder) + }); + if placeholder { + rate_limit.remove("secondary_window"); + } +} + +fn secondary_window_is_placeholder(window: &Map) -> bool { + window.iter().all(|(field, value)| match field.as_str() { + "used_percent" => value + .as_f64() + .is_some_and(|used_percent| used_percent.is_finite() && used_percent == 0.0), + "limit_reached" => value.as_bool() == Some(false), + _ => false, + }) +} + +/// 上游额度可确认套餐变更；同族泛化值不能丢弃 JWT 已给出的具体 SKU +pub(super) fn observed_account_plan( + current: Option<&str>, + observed: Option<&str>, +) -> Option { + let plan = observed?.trim().to_ascii_lowercase(); + if plan.is_empty() || plan == "unknown" { + return None; + } + // 套餐族沿用官方 codex_protocol::account::PlanType 的分类 + let current = current.unwrap_or_default().trim().to_ascii_lowercase(); + let generalized = matches!( + (plan.as_str(), current.as_str()), + ( + "team", + "self_serve_business_prolite" | "self_serve_business_usage_based" + ) | ( + "business", + "ent26" | "enterprise_cbp_automation" | "enterprise_cbp_usage_based" + ) | ("edu" | "education", "edu_plus" | "edu_pro") + ); + (!generalized).then_some(plan) +} diff --git a/backend/crates/providers/openai/src/credential/quota/mod.rs b/backend/crates/providers/openai/src/credential/quota/mod.rs index a0bd5ad76..387c7119f 100644 --- a/backend/crates/providers/openai/src/credential/quota/mod.rs +++ b/backend/crates/providers/openai/src/credential/quota/mod.rs @@ -1,4 +1,4 @@ -//! Codex quota 服务编排：权威额度事实、展示投影、主动/被动同步与 429 冷却 +//! Codex quota 门面与唯一服务状态；实现按读取、排程、预热和重置消费组织 //! //! - [`document`]：上游 `/usage` 输入到多桶 map 的单向规范化 //! - [`snapshot`]：快照/窗口解析、聚合、O(1) 滚动与调度信号 @@ -7,14 +7,23 @@ mod document; pub(crate) mod evidence; +mod passive; mod recovery; +mod reset_credits; +mod scheduling; +mod service; pub(crate) mod snapshot; +mod warmup; + +pub use reset_credits::CodexResetCreditsError; +use reset_credits::ResetCreditLocks; +use scheduling::CodexQuotaSchedulingProjection; pub use snapshot::{ CodexAccountQuotaSnapshot, CodexQuotaFact, CodexQuotaWindow, CodexQuotaWindowKind, CodexQuotaWindowRole, parse_codex_quota_usage, }; -use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::collections::{BTreeMap, BTreeSet}; use std::sync::{Arc, RwLock}; use std::time::{Duration, Instant, SystemTime}; @@ -50,8 +59,8 @@ use crate::transport::{ use super::repository::{CodexCredentialRepository, CredentialRepositoryError}; use super::security::CodexRuntimeCredential; use document::{ - DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, RateLimitSnapshotsByLimitId, - canonicalize_rate_limit_document, + DEFAULT_CODEX_LIMIT_ID, RateLimitSnapshotsByLimitId, normalize_quota_window_placeholders, + observed_account_plan, }; use evidence::{QuotaEndpointFailure, classify_quota_endpoint_failure}; use recovery::{RECOVERY_FIELD, reconcile_refresh}; @@ -123,74 +132,6 @@ pub struct CodexWarmupSummary { pub failed: u64, } -#[derive(Clone, Copy, PartialEq, Eq)] -enum WarmupTerminal { - Missing, - Completed, - Failed, -} - -fn observe_warmup_events(events: Vec, terminal: &mut WarmupTerminal) { - for event in events { - let parsed = serde_json::from_str::(&event.data).ok(); - let kind = parsed - .as_ref() - .and_then(|data| data.get("type")) - .and_then(Value::as_str) - .or(event.event.as_deref()); - match kind { - Some("response.completed") => { - // 只有明确完成的响应才算预热成功；失败事件即使随后出现完成帧也优先 - if *terminal != WarmupTerminal::Failed { - *terminal = if parsed.is_some() - && parsed - .as_ref() - .and_then(|data| data.pointer("/response/status")) - .and_then(Value::as_str) - .is_none_or(|status| status == "completed") - { - WarmupTerminal::Completed - } else { - WarmupTerminal::Failed - }; - } - } - Some("response.failed" | "response.incomplete" | "error") => { - *terminal = WarmupTerminal::Failed; - } - _ => {} - } - } -} - -async fn consume_warmup_sse( - response: &mut CodexBackendStreamingResponse, -) -> Result<(), &'static str> { - let mut decoder = SseEventDecoder::default(); - let mut terminal = WarmupTerminal::Missing; - let mut received_bytes = 0usize; - while let Some(chunk) = response.body.next().await { - let chunk = chunk.map_err(|_| "stream_read_failed")?; - received_bytes = received_bytes.saturating_add(chunk.len()); - if received_bytes > WARMUP_STREAM_MAX_BYTES { - return Err("stream_size_limit_exceeded"); - } - let events = decoder.push(&chunk).map_err(|_| "invalid_sse")?; - observe_warmup_events(events, &mut terminal); - if terminal != WarmupTerminal::Missing { - break; - } - } - if terminal == WarmupTerminal::Missing { - observe_warmup_events(decoder.finish().map_err(|_| "invalid_sse")?, &mut terminal); - } - match terminal { - WarmupTerminal::Completed => Ok(()), - WarmupTerminal::Missing => Err("completion_event_missing"), - WarmupTerminal::Failed => Err("response_failed"), - } -} - #[derive(Debug, Error)] pub enum CodexCredentialQuotaError { #[error("Codex quota response is invalid")] @@ -215,54 +156,6 @@ pub enum CodexCredentialQuotaError { }, } -/// 主动额度重置卡查询/消费失败 -#[derive(Error)] -pub enum CodexResetCreditsError { - #[error("Codex reset-credit credential data is invalid")] - InvalidCredentialData, - #[error("Codex OAuth access token must be refreshed before using reset credits")] - CredentialRefreshRequired { upstream_body: Option }, - #[error("Codex reset-credit account was not found")] - NotFound, - #[error("provider account store is unavailable: {detail}")] - Store { detail: String }, - #[error("Codex reset-credit upstream returned HTTP {status}")] - Upstream { - status: u16, - body: String, - retry_after_seconds: Option, - }, - #[error("Codex reset-credit query transport is unavailable")] - TransportUnavailable, - #[error("Codex reset-credit consume result is unknown")] - ConsumeResultUnknown, -} - -impl std::fmt::Debug for CodexResetCreditsError { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::InvalidCredentialData => formatter.write_str("InvalidCredentialData"), - Self::CredentialRefreshRequired { .. } => { - formatter.write_str("CredentialRefreshRequired { upstream_body: }") - } - Self::NotFound => formatter.write_str("NotFound"), - Self::Store { .. } => formatter.write_str("Store { detail: }"), - Self::Upstream { - status, - retry_after_seconds, - .. - } => formatter - .debug_struct("Upstream") - .field("status", status) - .field("body", &"") - .field("retry_after_seconds", retry_after_seconds) - .finish(), - Self::TransportUnavailable => formatter.write_str("TransportUnavailable"), - Self::ConsumeResultUnknown => formatter.write_str("ConsumeResultUnknown"), - } - } -} - impl From for CodexCredentialQuotaError { fn from(error: gateway_core::error::StoreError) -> Self { Self::Store { @@ -271,38 +164,6 @@ impl From for CodexCredentialQuotaError { } } -/// 客户端错误携带的 HTTP 状态码；传输失败等为 `None` -fn upstream_error_status(error: &CodexClientError) -> Option { - match error { - CodexClientError::Upstream { status, .. } => Some(status.as_u16()), - _ => None, - } -} - -/// 从上游错误体提取稳定错误码：优先 `/error/code`，其次 `/code` -/// -/// 只接受有界的 ASCII 标识，避免把自由文本作为错误码写入诊断 -fn upstream_error_code(error: &CodexClientError) -> Option { - let CodexClientError::Upstream { body, .. } = error else { - return None; - }; - let value = serde_json::from_str::(body).ok()?; - let code = value - .pointer("/error/code") - .or_else(|| value.pointer("/code")) - .and_then(Value::as_str)? - .trim(); - if code.is_empty() - || code.len() > 64 - || !code - .bytes() - .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.')) - { - return None; - } - Some(code.to_owned()) -} - pub struct CodexCredentialQuotaService { repository: CodexCredentialRepository, store: Arc, @@ -315,7 +176,7 @@ pub struct CodexCredentialQuotaService { /// 冻结策略的短 TTL 缓存：失败路径热读，避免每个容量错误都查询设置 freeze_policy_cache: Mutex>, scheduling: CodexQuotaSchedulingProjection, - reset_consume_locks: Mutex>>>, + reset_consume_locks: ResetCreditLocks, /// 首查随机延迟采样器；测试注入固定值获得确定性节奏。 initial_sync_delays: CodexInitialSyncDelays, } @@ -327,1833 +188,13 @@ pub type CodexInitialSyncDelays = Arc Duration + Send + Sync>; /// 冻结策略缓存活跃期；过期后下一次容量错误重新读取运行时设置 const FREEZE_POLICY_CACHE_TTL: Duration = Duration::from_secs(30); -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum QuotaRefreshAuthority { - ObserveAccess, - PreserveAccess, -} - -#[derive(Clone, Default)] -struct CodexQuotaSchedulingProjection { - state: Arc>, - hydration: Arc>, -} - -#[derive(Default)] -struct CodexQuotaProjectionState { - next_version: u64, - entries: BTreeMap, - last_periodic_refresh_at: BTreeMap, - initial_refresh_not_before: BTreeMap, -} - -struct CodexQuotaRefreshAttempt { - monotonic_at: Instant, - wall_at: SystemTime, -} - -#[derive(Debug, Clone, Copy)] -struct CodexQuotaSchedulingEntry { - version: u64, - revision: CredentialRevision, - expires_at: Instant, - signals: Option, -} - -#[derive(Clone)] -struct CodexQuotaHydrationTarget { - account: ProviderAccount, - expected_version: Option, -} - -struct FetchedCodexQuota { - account: ProviderAccount, - value: Value, -} - struct PreparedCodexRuntimeCredential { account: ProviderAccount, credential: CodexRuntimeCredential, } -enum CodexQuotaFetchAttemptError { - InvalidCredential, - Upstream(CodexClientError), -} - -enum CodexQuotaFetchError { - InvalidCredential, - Upstream { - account: Box, - error: CodexClientError, - }, -} - -enum ResetCreditAttemptError { - InvalidCredential, - Upstream(CodexClientError), -} - -impl CodexQuotaSchedulingProjection { - fn invalidate(&self, account_ids: &[ProviderAccountId]) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - for account_id in account_ids { - state.entries.remove(account_id); - state.last_periodic_refresh_at.remove(account_id); - state.initial_refresh_not_before.remove(account_id); - } - } - - fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { - let state = self - .state - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let now = Instant::now(); - accounts - .iter() - .filter_map(|account| { - let current = state.entries.get(account.id()); - current - .is_none_or(|entry| { - entry.revision != account.revision() || now >= entry.expires_at - }) - .then(|| CodexQuotaHydrationTarget { - account: account.clone(), - expected_version: current.map(|entry| entry.version), - }) - }) - .collect() - } - - fn observe(&self, snapshot: &CodexAccountQuotaSnapshot) -> bool { - let Some(remaining_ttl) = quota_projection_ttl(snapshot) else { - return false; - }; - self.replace( - snapshot.account_id().clone(), - snapshot.credential_revision(), - remaining_ttl, - scheduling_signals_from_snapshot(snapshot), - ); - true - } - - fn mark_unknown_if_unchanged(&self, target: &CodexQuotaHydrationTarget, ttl: Duration) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state - .entries - .get(target.account.id()) - .map(|entry| entry.version) - != target.expected_version - { - return; - } - insert_projection_entry( - &mut state, - target.account.id().clone(), - target.account.revision(), - ttl, - None, - ); - } - - fn observe_if_unchanged( - &self, - target: &CodexQuotaHydrationTarget, - snapshot: &CodexAccountQuotaSnapshot, - ) -> bool { - let Some(remaining_ttl) = quota_projection_ttl(snapshot) else { - return false; - }; - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state - .entries - .get(target.account.id()) - .map(|entry| entry.version) - != target.expected_version - { - return true; - } - insert_projection_entry( - &mut state, - snapshot.account_id().clone(), - snapshot.credential_revision(), - remaining_ttl, - scheduling_signals_from_snapshot(snapshot), - ); - true - } - - fn replace( - &self, - account_id: ProviderAccountId, - revision: CredentialRevision, - ttl: Duration, - signals: Option, - ) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - insert_projection_entry(&mut state, account_id, revision, ttl, signals); - } - - fn signals(&self, account: &ProviderAccount) -> Option { - let state = self - .state - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - state - .entries - .get(account.id()) - .filter(|entry| { - entry.revision == account.revision() && Instant::now() < entry.expires_at - }) - .and_then(|entry| entry.signals) - } - - fn reserve_periodic_refreshes( - &self, - accounts: Vec, - observed_snapshots: &BTreeMap, - now: SystemTime, - ) -> Vec { - let candidates = accounts - .into_iter() - .filter_map(|account| { - let snapshot = observed_snapshots.get(account.id()); - quota_refresh_candidate(account, snapshot, now) - }) - .collect::>(); - let candidate_ids = candidates - .iter() - .map(|(account, _)| account.id().clone()) - .collect::>(); - let refreshed_at = Instant::now(); - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - state - .last_periodic_refresh_at - .retain(|account_id, _| candidate_ids.contains(account_id)); - - // 已耗尽账号首次立即复核，之后每 30 分钟复核，以发现官方提前重置 - // reset + 2 分钟额外触发一次复核，给上游重置留出传播时间 - // 正常账号仅在非零用量窗口经过宽限期后参与，未更新时复用周期节流 - let mut reserved = Vec::new(); - for (account, target_reset) in candidates { - if !periodic_quota_refresh_due(&state, account.id(), target_reset, now, refreshed_at) { - continue; - } - state.last_periodic_refresh_at.insert( - account.id().clone(), - CodexQuotaRefreshAttempt { - monotonic_at: refreshed_at, - wall_at: now, - }, - ); - reserved.push(account); - } - reserved - } - - /// 为尚无 quota 快照的账号分配随机首查时点，并返回本轮已到期的候选。 - /// - /// `quota_observed_at` 为空代表首次异步观察尚未成功；不另建同步状态表。 - /// 首次进入候选的账号抽一次随机启动延迟记为到期时刻，避免新导入账号池 - /// 在 worker 唤醒边界形成批量 `/usage` 突发；到期时刻跨轮保留，零延迟 - /// （含随机源退化）时保持首轮即查的旧行为。单轮数量仍受 - /// `INITIAL_QUOTA_SYNC_BATCH` 上限约束。 - fn reserve_initial_refreshes( - &self, - accounts: &[ProviderAccount], - observed_ids: &BTreeSet, - now: SystemTime, - draw_start_delay: &dyn Fn() -> Duration, - ) -> Vec { - let candidates = accounts - .iter() - .filter(|account| { - !observed_ids.contains(account.id()) && eligible_initial_quota_sync(account, now) - }) - .collect::>(); - let candidate_ids = candidates - .iter() - .map(|account| account.id().clone()) - .collect::>(); - let monotonic_now = Instant::now(); - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - // 已产生快照或离开候选集的账号不再保留首查状态。 - state - .initial_refresh_not_before - .retain(|account_id, _| candidate_ids.contains(account_id)); - candidates - .into_iter() - .filter(|account| { - let due_at = state - .initial_refresh_not_before - .entry(account.id().clone()) - .or_insert_with(|| monotonic_now + draw_start_delay()); - *due_at <= monotonic_now - }) - .take(INITIAL_QUOTA_SYNC_BATCH) - .cloned() - .collect() - } - - /// 首查尝试后仍无快照的账号重抽随机延迟，避免失败账号回到固定 30s 重试节奏。 - /// - /// 重排同时随机化每轮的候选集合，避免持续失败时按账号顺序只重试同一批。 - fn defer_failed_initial_refreshes( - &self, - account_ids: &BTreeSet, - draw_start_delay: &dyn Fn() -> Duration, - ) { - if account_ids.is_empty() { - return; - } - let now = Instant::now(); - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - for account_id in account_ids { - // 账号是否仍在首查路径由下一轮 reserve 按候选集剪枝，这里无需区分。 - state - .initial_refresh_not_before - .insert(account_id.clone(), now + draw_start_delay()); - } - } -} - -fn quota_refresh_candidate( - account: ProviderAccount, - snapshot: Option<&CodexAccountQuotaSnapshot>, - now: SystemTime, -) -> Option<(ProviderAccount, Option)> { - if !eligible_periodic_quota_refresh(&account, now) { - return None; - } - if account.quota().is_exhausted() { - let reset_at = account.quota().reset_at(); - return Some((account, reset_at)); - } - // 正常账号首次进入周期候选也要等待宽限期，不能由“无刷新历史”绕过 - let snapshot = snapshot?; - let expired_window_reset = snapshot - .windows() - .iter() - .filter(|window| { - window.used_percent().is_some_and(|used| used > 0.0) || window.limit_reached() - }) - .filter_map(CodexQuotaWindow::reset_at) - .map(SystemTime::from) - .filter(|reset| { - reset - .checked_add(QUOTA_RESET_GRACE) - .is_some_and(|due_at| due_at <= now) - }) - .min(); - expired_window_reset.map(|reset_at| (account, Some(reset_at))) -} - -fn periodic_quota_refresh_due( - state: &CodexQuotaProjectionState, - account_id: &ProviderAccountId, - reset_at: Option, - now: SystemTime, - monotonic_now: Instant, -) -> bool { - state - .last_periodic_refresh_at - .get(account_id) - .is_none_or(|last| { - monotonic_now.saturating_duration_since(last.monotonic_at) - >= PERIODIC_QUOTA_REFRESH_RETRY_INTERVAL - || reset_at - .and_then(|reset| reset.checked_add(QUOTA_RESET_GRACE)) - // 已在该边界之后复核过时回到周期重试，避免过期 reset 每轮触发 - .is_some_and(|due_at| last.wall_at < due_at && due_at <= now) - }) -} - -fn insert_projection_entry( - state: &mut CodexQuotaProjectionState, - account_id: ProviderAccountId, - revision: CredentialRevision, - ttl: Duration, - signals: Option, -) { - state.next_version = state.next_version.saturating_add(1); - state.entries.insert( - account_id, - CodexQuotaSchedulingEntry { - version: state.next_version, - revision, - expires_at: Instant::now() + ttl, - signals, - }, - ); -} - -impl CodexCredentialQuotaService { - pub fn new( - repository: CodexCredentialRepository, - profile: CodexWireProfileState, - http: Client, - base_url: String, - cooldowns: Arc, - leases: Arc, - runtime_policy: Arc, - ) -> Self { - Self { - store: Arc::clone(repository.store()), - repository, - profile, - http, - base_url, - cooldowns, - leases, - runtime_policy, - freeze_policy_cache: Mutex::new(None), - scheduling: CodexQuotaSchedulingProjection::default(), - reset_consume_locks: Mutex::new(HashMap::new()), - initial_sync_delays: Arc::new(|| { - crate::jitter::uniform_delay( - crate::jitter::random_u64(), - INITIAL_QUOTA_SYNC_MAX_START_DELAY, - ) - }), - } - } - - /// 注入首查随机延迟采样器（测试用确定性节奏）；生产默认均匀 - /// `[0, INITIAL_QUOTA_SYNC_MAX_START_DELAY)`，零延迟等价旧行为 - #[doc(hidden)] - pub fn with_initial_sync_delays(mut self, delays: CodexInitialSyncDelays) -> Self { - self.initial_sync_delays = delays; - self - } - - /// 读取容量熔断策略（带短 TTL 缓存）；读取失败退化为关闭，熔断不得放大失败 - async fn freeze_policy(&self) -> ProviderFreezePolicy { - { - let cache = self.freeze_policy_cache.lock().await; - if let Some((policy, loaded_at)) = cache.as_ref() - && loaded_at.elapsed() < FREEZE_POLICY_CACHE_TTL - { - return policy.clone(); - } - } - let loaded = self.runtime_policy.load_freeze_policy().await.ok(); - let mut cache = self.freeze_policy_cache.lock().await; - if let Some(policy) = loaded { - *cache = Some((policy.clone(), Instant::now())); - return policy; - } - cache - .as_ref() - .filter(|(_, loaded_at)| loaded_at.elapsed() < FREEZE_POLICY_CACHE_TTL) - .map_or_else(ProviderFreezePolicy::disabled, |(policy, _)| policy.clone()) - } - - /// 容量熔断入口：滑动窗口内累计容量类失败，达到阈值即写入带 - /// `CapacityFreeze` 类别的账号级冷却 - /// 调度侧立即屏蔽该账号，恢复由 - /// freeze-recovery worker 处理；本路径只依赖 Redis 可丢失事实 - pub async fn apply_capacity_failure(&self, account: &ProviderAccount, observed_at: SystemTime) { - let policy = self.freeze_policy().await; - if !policy.enabled() { - return; - } - let in_flight = self.current_in_flight(account.id()).await; - let Ok(count) = self - .cooldowns - .record_capacity_failure(account.id(), policy.window(), in_flight) - .await - else { - return; - }; - if count < policy.threshold() { - return; - } - let Some(until) = observed_at.checked_add(policy.freeze_duration()) else { - return; - }; - let cooldown = ProviderCooldown::new_with_kind( - account.id().clone(), - account.revision(), - until, - if policy.probe_enabled() { - ProviderCooldownKind::CapacityFreezeProbe - } else { - ProviderCooldownKind::CapacityFreeze - }, - ); - if self.cooldowns.put_if_later(cooldown).await.is_ok() { - tracing::warn!( - account_id = account.id().as_str(), - threshold = policy.threshold(), - window_seconds = policy.window().as_secs(), - freeze_seconds = policy.freeze_duration().as_secs(), - peak_in_flight = in_flight, - "账号容量熔断触发：冻结该账号一段时间", - ); - } - } - - async fn current_in_flight(&self, account_id: &ProviderAccountId) -> u32 { - self.leases - .account_in_flight(std::slice::from_ref(account_id)) - .await - .ok() - .and_then(|signals| signals.get(account_id).copied()) - .unwrap_or(0) - } - - /// 查询当前账号由 Codex Desktop 暴露的主动额度重置卡 - pub async fn list_reset_credits( - &self, - account_id: &ProviderAccountId, - ) -> Result { - let account = self.reset_credit_account(account_id).await?; - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - let credential = self - .repository - .load_runtime_credential(&account) - .await - .map_err(|_| CodexResetCreditsError::InvalidCredentialData)?; - let prepared = PreparedCodexRuntimeCredential { - account, - credential, - }; - let request_id = format!("reset_credits_{}", Uuid::now_v7().simple()); - list_reset_credits_once(&client, &prepared, &request_id) - .await - .map_err(|error| map_reset_credit_attempt_error(error, false)) - } - - /// 消费一张主动额度重置卡 - /// 相同账号在本进程内串行，且不做传输重试 - pub async fn consume_reset_credit( - &self, - account_id: &ProviderAccountId, - credit_id: Option<&str>, - redeem_request_id: Uuid, - ) -> Result { - let lock = { - let mut locks = self.reset_consume_locks.lock().await; - Arc::clone( - locks - .entry(account_id.clone()) - .or_insert_with(|| Arc::new(Mutex::new(()))), - ) - }; - let _guard = lock.lock().await; - let account = self.reset_credit_account(account_id).await?; - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - let credential = self - .repository - .load_runtime_credential(&account) - .await - .map_err(|_| CodexResetCreditsError::InvalidCredentialData)?; - let prepared = PreparedCodexRuntimeCredential { - account, - credential, - }; - let request_id = format!("reset_credit_consume_{}", Uuid::now_v7().simple()); - consume_reset_credit_once( - &client, - &prepared, - &request_id, - credit_id, - redeem_request_id, - ) - .await - .map_err(|error| map_reset_credit_attempt_error(error, true)) - } - - async fn reset_credit_account( - &self, - account_id: &ProviderAccountId, - ) -> Result { - let account = self - .store - .get_account(account_id) - .await - .map_err(|error| CodexResetCreditsError::Store { - detail: error.to_string(), - })? - .filter(|account| account.provider().as_str() == "openai") - .ok_or(CodexResetCreditsError::NotFound)?; - if !access_token_is_current(&account, SystemTime::now()) { - return Err(CodexResetCreditsError::CredentialRefreshRequired { - upstream_body: None, - }); - } - Ok(account) - } - - /// 真实推理错误只更新额度访问事实，不伪造 Provider JSON 或展示百分比 - pub(crate) async fn record_confirmed_exhaustion( - &self, - account: &ProviderAccount, - evidence: QuotaEvidence, - reset_at: Option, - observed_at: SystemTime, - ) -> Result<(), CodexCredentialQuotaError> { - let outcome = self - .store - .apply_quota_access(QuotaAccessChange { - account_id: account.id().clone(), - expected_revision: account.revision(), - state: QuotaState::exhausted(evidence, observed_at, reset_at), - }) - .await?; - if outcome == QuotaWriteOutcome::Conflict { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - Ok(()) - } - - /// 成功推理是额度可访问的权威证据，同时解除账号级 429 冷却 - pub async fn record_successful_inference( - &self, - account: &ProviderAccount, - observed_at: SystemTime, - ) -> Result<(), CodexCredentialQuotaError> { - // 已经发出的并发请求可能晚于耗尽事实成功返回；真实成功仍是独立的 - // Allowed 权威证据，但调度不会为了试探恢复而放行耗尽账号 - if account.quota().access() == QuotaAccessState::Exhausted { - let outcome = self - .store - .apply_quota_access(QuotaAccessChange { - account_id: account.id().clone(), - expected_revision: account.revision(), - state: QuotaState::allowed(observed_at), - }) - .await?; - if outcome == QuotaWriteOutcome::Conflict { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - } - self.cooldowns - .clear_after_success(account.id(), account.revision()) - .await - .map_err(|error| CodexCredentialQuotaError::Store { - detail: error.to_string(), - })?; - Ok(()) - } - - /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度 - pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { - if self.scheduling.hydration_targets(accounts).is_empty() { - return; - } - let _hydration = self.scheduling.hydration.lock().await; - let pending = self.scheduling.hydration_targets(accounts); - if pending.is_empty() { - return; - } - let account_ids = pending - .iter() - .map(|target| target.account.id().clone()) - .collect::>(); - let Ok(observations) = self.store.get_quotas(&account_ids).await else { - for target in &pending { - self.scheduling - .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); - } - return; - }; - let observations = observations - .into_iter() - .map(|observation| (observation.account_id.clone(), observation)) - .collect::>(); - for target in pending { - let snapshot = observations - .get(target.account.id()) - .filter(|observation| observation.expected_revision == target.account.revision()) - .and_then(quota_snapshot_from_observation); - if !snapshot - .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) - { - self.scheduling - .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); - } - } - } - - #[must_use] - pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { - self.scheduling.signals(account) - } - - pub(crate) fn invalidate_scheduling(&self, account_ids: &[ProviderAccountId]) { - self.scheduling.invalidate(account_ids); - } - - pub async fn synchronize(&self) -> Result { - self.synchronize_at(SystemTime::now()).await - } - - /// 按本轮调度时刻选择到期账号；实际 HTTP 观察与落库仍使用发生时刻 - pub async fn synchronize_at( - &self, - now: SystemTime, - ) -> Result { - let mut accounts = self.repository.list_for_provider().await?; - accounts.retain(|account| { - account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH - }); - let mut summary = CodexQuotaSyncSummary::default(); - let account_ids = accounts - .iter() - .map(|account| account.id().clone()) - .collect::>(); - let observed = self.store.get_quotas(&account_ids).await?; - let observed_snapshots = observed - .iter() - .filter_map(|obs| { - quota_snapshot_from_observation(obs) - .map(|snapshot| (obs.account_id.clone(), snapshot)) - }) - .collect::>(); - let observed_ids = observed - .into_iter() - .map(|observation| observation.account_id) - .collect::>(); - let initial = self.scheduling.reserve_initial_refreshes( - &accounts, - &observed_ids, - now, - self.initial_sync_delays.as_ref(), - ); - // 本轮首查后仍无快照的账号需要重排随机延迟。 - let mut pending_initial: BTreeSet = - initial.iter().map(|account| account.id().clone()).collect(); - let periodic = - self.scheduling - .reserve_periodic_refreshes(accounts, &observed_snapshots, now); - let accounts = initial - .into_iter() - .chain(periodic) - .fold(BTreeMap::new(), |mut unique, account| { - unique.insert(account.id().clone(), account); - unique - }) - .into_values() - .collect::>(); - if accounts.is_empty() { - return Ok(summary); - } - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - for account in accounts { - let observed_at = SystemTime::now(); - match self.fetch_usage(&client, &account).await { - Ok(FetchedCodexQuota { account, value }) => { - // 单账号解析或落库失败只影响该账号；其余账号继续同步 - match self - .apply_fetched_quota(&account, &value, observed_at, &mut summary) - .await - { - Ok(()) => { - pending_initial.remove(account.id()); - } - Err(error) => { - summary.transient += 1; - tracing::warn!( - account_id = %account.id(), - error = %error, - "OpenAI quota synchronization skipped one account" - ); - } - } - } - Err(CodexQuotaFetchError::InvalidCredential) => { - summary.stale += 1; - } - Err(CodexQuotaFetchError::Upstream { account, error }) => { - match classify_quota_endpoint_failure(&error) { - Some(QuotaEndpointFailure::Exhausted(evidence)) => { - summary.exhausted += 1; - if let Err(write_error) = self - .record_confirmed_exhaustion(&account, evidence, None, observed_at) - .await - { - summary.transient += 1; - tracing::warn!( - account_id = %account.id(), - error = %write_error, - "OpenAI quota exhaustion fact write failed" - ); - } - } - Some(QuotaEndpointFailure::Credential { state, reason }) => { - summary.banned += 1; - self.persist_credential_failure(&account, state, reason, observed_at) - .await; - } - None => { - summary.transient += 1; - tracing::warn!( - account_id = %account.id(), - error = %error, - upstream_status = upstream_error_status(&error), - upstream_code = ?upstream_error_code(&error), - "OpenAI quota upstream rejection; refresh cycle will retry later" - ); - } - } - } - } - } - // 首查失败或未能落库的账号重排随机延迟；已成功落库的账号由快照事实自然退出首查路径。 - self.scheduling - .defer_failed_initial_refreshes(&pending_initial, self.initial_sync_delays.as_ref()); - Ok(summary) - } - - #[must_use] - pub fn runtime_policy(&self) -> &Arc { - &self.runtime_policy - } - - /// 批量预激活 OAuth 账号的 5h 配额滑动窗口 - pub async fn execute_warmup( - &self, - model: &str, - ) -> Result { - let mut accounts = self.repository.list_for_provider().await?; - accounts.retain(|account| { - account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH - && account.enabled() - && matches!( - account.credential_state(), - CredentialState::Unknown | CredentialState::Ready - ) - }); - let mut summary = CodexWarmupSummary::default(); - if accounts.is_empty() { - return Ok(summary); - } - let account_ids = accounts - .iter() - .map(|account| account.id().clone()) - .collect::>(); - let observed = self.store.get_quotas(&account_ids).await?; - let observed_snapshots = observed - .iter() - .filter_map(|obs| { - quota_snapshot_from_observation(obs) - .map(|snapshot| (obs.account_id.clone(), snapshot)) - }) - .collect::>(); - - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - - let now_utc = chrono::Utc::now(); - for account in accounts { - if let Some(snapshot) = observed_snapshots.get(account.id()) { - // 1. 周线已触顶或耗尽跳过 - let weekly_exhausted = snapshot - .windows() - .iter() - .any(|w| w.kind() == CodexQuotaWindowKind::Weekly && w.limit_reached()); - if weekly_exhausted { - summary.skipped_exhausted += 1; - continue; - } - // 2. 5h 窗口当前活跃且距重置时间 > 30 分钟跳过 - let has_active_5h = snapshot.windows().iter().any(|w| { - w.kind() == CodexQuotaWindowKind::ShortTerm - && w.reset_at().is_some_and(|reset_at| { - reset_at > now_utc + chrono::Duration::minutes(30) - }) - }); - if has_active_5h { - summary.skipped_active += 1; - continue; - } - } - - let credential = match self.repository.load_runtime_credential(&account).await { - Ok(cred) => cred, - Err(_) => { - summary.failed += 1; - continue; - } - }; - let authorization = match credential.authentication.authorization_header() { - Ok(auth) => auth, - Err(_) => { - summary.failed += 1; - continue; - } - }; - - let mut body = Map::new(); - body.insert("model".to_owned(), Value::String(model.to_owned())); - body.insert( - "input".to_owned(), - serde_json::json!([{ - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "hello"}] - }]), - ); - body.insert("stream".to_owned(), Value::Bool(true)); - body.insert("store".to_owned(), Value::Bool(false)); - body.insert( - "service_tier".to_owned(), - Value::String("default".to_owned()), - ); - body.insert( - "reasoning".to_owned(), - serde_json::json!({"effort": "none"}), - ); - body.insert("text".to_owned(), serde_json::json!({"verbosity": "low"})); - - let upstream_request = CodexResponsesRequest::from_body(body); - let request_id = format!("warmup_{}", Uuid::now_v7().simple()); - let client_for_account = match client.for_account(&account) { - Ok(c) => c, - Err(error) => { - tracing::warn!(account_id = %account.id(), error = %error, "warmup client for account failed"); - summary.failed += 1; - continue; - } - }; - let context = crate::transport::CodexRequestContext::auxiliary( - authorization.expose_secret(), - account.upstream_account_id(), - &request_id, - None, - ); - - // HTTP 200 只确认响应头；必须消费 SSE 直到终态才能确认预热结果和额度事件 - let attempt = tokio::time::timeout(WARMUP_REQUEST_TIMEOUT, async { - let mut response = client_for_account - .create_response_stream_http_sse(&upstream_request, context) - .await - .map_err(|_| "request_rejected")?; - let terminal = consume_warmup_sse(&mut response).await; - let rate_limit_updates = match response.rate_limit_updates.as_ref() { - Some(updates) => std::mem::take(&mut *updates.lock().await), - None => Vec::new(), - }; - Ok::<_, &'static str>((terminal, response.rate_limit_headers, rate_limit_updates)) - }) - .await; - match attempt { - Ok(Ok((terminal, headers, rate_limit_updates))) => { - match terminal { - Ok(()) => { - // 被动额度同步会把成功请求的窗口事实标为可用，失败流不能复用该结论 - if !headers.is_empty() - && let Err(error) = - self.synchronize_passive_headers(&account, &headers).await - { - tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit header sync failed"); - } - if !rate_limit_updates.is_empty() - && let Err(error) = self - .synchronize_passive_rate_limits(&account, &rate_limit_updates) - .await - { - tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit event sync failed"); - } - summary.warmed_up += 1; - tracing::info!( - account_id = %account.id(), - model, - "OpenAI account warmed up successfully" - ); - } - Err(reason) => { - summary.failed += 1; - tracing::warn!( - account_id = %account.id(), - reason, - "OpenAI account warmup stream failed" - ); - } - } - } - Ok(Err(reason)) => { - summary.failed += 1; - tracing::warn!( - account_id = %account.id(), - reason, - "OpenAI account warmup request rejected by upstream" - ); - } - Err(_) => { - summary.failed += 1; - tracing::warn!( - account_id = %account.id(), - "OpenAI account warmup request timed out" - ); - } - } - tokio::time::sleep(Duration::from_secs(2)).await; - } - - Ok(summary) - } - - /// 解析并 revision-fenced 落库单账号的 Provider quota JSON - async fn apply_fetched_quota( - &self, - account: &ProviderAccount, - value: &Value, - observed_at: SystemTime, - summary: &mut CodexQuotaSyncSummary, - ) -> Result<(), CodexCredentialQuotaError> { - let mut object = normalize_quota_window_placeholders( - value - .as_object() - .cloned() - .ok_or(CodexCredentialQuotaError::InvalidCredentialData)?, - ); - object.remove(RECOVERY_FIELD); - let mut snapshot = parse_account_quota_snapshot( - account.id().clone(), - account.revision(), - observed_at, - &Value::Object(object.clone()), - )?; - let previous = if account.quota().is_exhausted() { - self.read_snapshot_for(account).await? - } else { - None - }; - let state = reconcile_refresh( - account.quota(), - &mut snapshot, - previous.as_ref(), - &mut object, - )?; - let snapshot = snapshot.with_quota_state(state); - let outcome = self - .store - .compare_and_swap_quota(QuotaObservation { - plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()), - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: OpaqueProviderData::new(object), - observed_at, - state, - }) - .await?; - if outcome == QuotaWriteOutcome::Conflict { - summary.stale += 1; - return Ok(()); - } - self.scheduling.observe(&snapshot); - if snapshot.quota().is_exhausted() { - summary.exhausted += 1; - } else { - summary.updated += 1; - } - Ok(()) - } - - /// 把正常推理响应携带的限流事实合并进 Provider 原始 quota JSON - pub async fn synchronize_passive_headers( - &self, - account: &ProviderAccount, - headers: &[(String, String)], - ) -> Result { - let Some(rate_limits) = parse_rate_limit_headers(headers) else { - return Ok(false); - }; - self.synchronize_passive_rate_limits(account, std::slice::from_ref(&rate_limits)) - .await - } - - /// 把一次推理响应中采集的结构化限流观察合并后单次落库 - pub async fn synchronize_passive_rate_limits( - &self, - account: &ProviderAccount, - rate_limits: &[ParsedRateLimits], - ) -> Result { - if account.authentication_kind() != crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH - || rate_limits.is_empty() - { - return Ok(false); - } - let has_quota_facts = rate_limits.iter().any(|observation| { - observation - .limits - .values() - .any(|details| passive_rate_limit_snapshot(details).is_some()) - }); - let existing = self - .store - .get_quotas(std::slice::from_ref(account.id())) - .await? - .into_iter() - .find(|observation| { - observation.account_id == *account.id() - && observation.expected_revision == account.revision() - }); - let existing_state = existing.as_ref().map(|observation| observation.state); - let existing = existing - .map(|observation| observation.quota.into_inner()) - .unwrap_or_default(); - // 只用本次响应明确携带的套餐更新账号，不能把合并前的旧快照重新当作新证据 - let observed_plan = rate_limits - .iter() - .rev() - .filter_map(|observation| { - observation.plan_type.as_deref().filter(|plan| { - !plan.trim().is_empty() && !plan.trim().eq_ignore_ascii_case("unknown") - }) - }) - .next(); - let plan_type = observed_account_plan(account.plan_type(), observed_plan); - // 套餐、credits 等元数据可以更新，但没有额度窗口事实时必须保留旧观察时刻， - // 也不能借旧快照重新推导 quota state - if !has_quota_facts { - let Some(state) = existing_state else { - return Ok(false); - }; - let outcome = self - .store - .compare_and_swap_quota(QuotaObservation { - plan_type, - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: OpaqueProviderData::new(merge_passive_quota(existing, rate_limits)), - observed_at: SystemTime::now(), - state, - }) - .await?; - return Ok(outcome != QuotaWriteOutcome::Conflict); - } - let quota = merge_passive_quota(existing, rate_limits); - let observed_at = SystemTime::now(); - let snapshot = parse_account_quota_snapshot( - account.id().clone(), - account.revision(), - observed_at, - &Value::Object(quota.clone()), - )?; - // 这些 headers 来自一次成功推理，访问结论优先于可能滞后的百分比 - let state = QuotaState::allowed(observed_at); - let outcome = self - .store - .compare_and_swap_quota(QuotaObservation { - plan_type, - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: OpaqueProviderData::new(quota), - observed_at, - state, - }) - .await?; - if outcome == QuotaWriteOutcome::Conflict { - return Ok(false); - } - self.scheduling.observe(&snapshot); - Ok(true) - } - - /// 真实 429 的单一事实入口：写入 Redis 临时限流冷却（`until = now + retry_after`） - /// 凭据与额度主窗口（额度重置时间）都不改变——临时限流是独立维度， - /// 到期由 Redis key 过期自动解除，不污染配额耗尽状态 - /// 已有更晚的冷却不会被缩短（put_if_later） - pub async fn apply_rate_limit_429( - &self, - account: &ProviderAccount, - retry_after: Option, - observed_at: SystemTime, - ) -> Result<(), CodexCredentialQuotaError> { - let until = observed_at - .checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)) - .unwrap_or(observed_at); - self.cooldowns - .put_if_later(ProviderCooldown::new( - account.id().clone(), - account.revision(), - until, - )) - .await - .map_err(|error| CodexCredentialQuotaError::Store { - detail: error.to_string(), - })?; - Ok(()) - } - - /// 读取有效的账号冷却事实；等待恢复探测的冻结到期后仍有效 - pub async fn cooldown( - &self, - account_id: &ProviderAccountId, - ) -> Result, CodexCredentialQuotaError> { - let Some(cooldown) = self.cooldowns.read(account_id).await.map_err(|error| { - CodexCredentialQuotaError::Store { - detail: error.to_string(), - } - })? - else { - return Ok(None); - }; - let state = cooldown.scheduling_state(); - Ok(state.is_active(SystemTime::now()).then_some(state)) - } - - /// 读取单账号最后一次落库的 Provider quota，并由 Codex 域解析展示窗口 - pub async fn read_account( - &self, - account_id: &ProviderAccountId, - ) -> Result, CodexCredentialQuotaError> { - let account = self - .store - .get_account(account_id) - .await? - .filter(|account| account.provider().as_str() == "openai") - .ok_or(CodexCredentialQuotaError::NotFound)?; - self.read_snapshot_for(&account).await - } - - async fn read_snapshot_for( - &self, - account: &ProviderAccount, - ) -> Result, CodexCredentialQuotaError> { - if account.provider().as_str() != "openai" { - return Err(CodexCredentialQuotaError::NotFound); - } - let account_id = account.id(); - let Some(observation) = self - .store - .get_quotas(std::slice::from_ref(account_id)) - .await? - .into_iter() - .next() - else { - return Ok(None); - }; - if observation.account_id != *account_id - || observation.expected_revision != account.revision() - { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - if observation.quota.expose_to_provider().is_empty() { - return Ok(None); - } - let observed_at = observation.observed_at; - let snapshot = parse_account_quota_snapshot( - account_id.clone(), - account.revision(), - observed_at, - &Value::Object(observation.quota.expose_to_provider().clone()), - )? - .with_quota_state(observation.state); - self.scheduling.observe(&snapshot); - Ok(Some(snapshot)) - } - - /// 只刷新指定账号，revision-fenced 写入动态 Provider JSON 后返回解析快照 - pub async fn refresh_account( - &self, - account_id: &ProviderAccountId, - ) -> Result { - self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::ObserveAccess) - .await - } - - /// 真实限额失败后的异步刷新只补齐展示快照，不允许 usage 快照撤销已确认的失败状态 - pub(crate) async fn refresh_account_after_failure( - &self, - account_id: &ProviderAccountId, - ) -> Result { - self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::PreserveAccess) - .await - } - - async fn refresh_account_snapshot( - &self, - account_id: &ProviderAccountId, - authority: QuotaRefreshAuthority, - ) -> Result { - let account = self - .store - .get_account(account_id) - .await? - .filter(|account| account.provider().as_str() == "openai") - .ok_or(CodexCredentialQuotaError::NotFound)?; - if account.authentication_kind() != crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH { - return Err(CodexCredentialQuotaError::NotFound); - } - let observed_at = SystemTime::now(); - if !access_token_is_current(&account, observed_at) { - return Err(CodexCredentialQuotaError::CredentialRefreshRequired); - } - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - let FetchedCodexQuota { account, value } = match self.fetch_usage(&client, &account).await { - Ok(fetched) => fetched, - Err(CodexQuotaFetchError::InvalidCredential) => { - return Err(CodexCredentialQuotaError::InvalidCredentialData); - } - Err(CodexQuotaFetchError::Upstream { account, error }) => { - match classify_quota_endpoint_failure(&error) { - Some(QuotaEndpointFailure::Exhausted(evidence)) => { - self.record_confirmed_exhaustion(&account, evidence, None, observed_at) - .await?; - } - Some(QuotaEndpointFailure::Credential { state, reason }) => { - self.persist_credential_failure(&account, state, reason, observed_at) - .await; - } - None => {} - } - return Err(CodexCredentialQuotaError::Upstream { - detail: error.to_string(), - status: upstream_error_status(&error), - code: upstream_error_code(&error), - }); - } - }; - let mut object = normalize_quota_window_placeholders( - value - .as_object() - .cloned() - .ok_or(CodexCredentialQuotaError::InvalidCredentialData)?, - ); - object.remove(RECOVERY_FIELD); - let mut snapshot = parse_account_quota_snapshot( - account.id().clone(), - account.revision(), - observed_at, - &Value::Object(object.clone()), - )?; - let previous = if account.quota().is_exhausted() { - self.read_snapshot_for(&account).await? - } else { - None - }; - if account.quota().is_exhausted() - && matches!(authority, QuotaRefreshAuthority::PreserveAccess) - { - if self - .store - .touch_quota_observation(QuotaObservationTouch { - account_id: account.id().clone(), - expected_revision: account.revision(), - observed_at, - }) - .await? - == QuotaWriteOutcome::Conflict - { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - return Ok(match previous { - Some(previous) => { - let previous = previous.with_observed_at(observed_at); - self.scheduling.observe(&previous); - previous - } - None => snapshot - .with_quota_state(account.quota()) - .with_observed_at(observed_at), - }); - } - let state = match authority { - QuotaRefreshAuthority::ObserveAccess => reconcile_refresh( - account.quota(), - &mut snapshot, - previous.as_ref(), - &mut object, - )?, - QuotaRefreshAuthority::PreserveAccess => account.quota(), - }; - let snapshot = snapshot.with_quota_state(state); - if self - .store - .compare_and_swap_quota(QuotaObservation { - plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()), - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: OpaqueProviderData::new(object), - observed_at, - state, - }) - .await? - == QuotaWriteOutcome::Conflict - { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - self.scheduling.observe(&snapshot); - Ok(snapshot) - } - - async fn persist_credential_failure( - &self, - account: &ProviderAccount, - credential_state: CredentialState, - reason: AccountErrorReason, - observed_at: SystemTime, - ) { - if let Err(error) = self - .repository - .apply_state_with_reason(account, credential_state, observed_at, Some(reason), None) - .await - { - tracing::warn!( - account_id = %account.id(), - ?credential_state, - reason = reason.as_str(), - error = %error, - "OpenAI quota credential fact write failed" - ); - } - } - - async fn fetch_usage( - &self, - client: &CodexBackendClient, - account: &ProviderAccount, - ) -> Result { - let credential = self - .repository - .load_runtime_credential(account) - .await - .map_err(|_| CodexQuotaFetchError::InvalidCredential)?; - let prepared = PreparedCodexRuntimeCredential { - account: account.clone(), - credential, - }; - let result = fetch_usage_with_5xx_retry(client, &prepared).await; - match result { - Ok(value) => Ok(FetchedCodexQuota { - account: prepared.account, - value, - }), - Err(CodexQuotaFetchAttemptError::InvalidCredential) => { - Err(CodexQuotaFetchError::InvalidCredential) - } - Err(CodexQuotaFetchAttemptError::Upstream(error)) => { - Err(CodexQuotaFetchError::Upstream { - account: Box::new(prepared.account), - error, - }) - } - } - } -} - -async fn list_reset_credits_once( - client: &CodexBackendClient, - prepared: &PreparedCodexRuntimeCredential, - request_id: &str, -) -> Result { - let authorization = prepared - .credential - .authentication - .authorization_header() - .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; - client - .for_account(&prepared.account) - .map_err(ResetCreditAttemptError::Upstream)? - .list_rate_limit_reset_credits(CodexRequestContext::auxiliary( - authorization.expose_secret(), - prepared.account.upstream_account_id(), - request_id, - None, - )) - .await - .map_err(ResetCreditAttemptError::Upstream) -} - -async fn consume_reset_credit_once( - client: &CodexBackendClient, - prepared: &PreparedCodexRuntimeCredential, - request_id: &str, - credit_id: Option<&str>, - redeem_request_id: Uuid, -) -> Result { - let authorization = prepared - .credential - .authentication - .authorization_header() - .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; - client - .for_account(&prepared.account) - .map_err(ResetCreditAttemptError::Upstream)? - .consume_rate_limit_reset_credit( - CodexRequestContext::auxiliary( - authorization.expose_secret(), - prepared.account.upstream_account_id(), - request_id, - None, - ), - credit_id, - redeem_request_id, - ) - .await - .map_err(ResetCreditAttemptError::Upstream) -} - -fn map_reset_credit_attempt_error( - error: ResetCreditAttemptError, - consume: bool, -) -> CodexResetCreditsError { - match error { - ResetCreditAttemptError::InvalidCredential => CodexResetCreditsError::InvalidCredentialData, - ResetCreditAttemptError::Upstream(error) => map_reset_credit_client_error(error, consume), - } -} - -fn map_reset_credit_client_error(error: CodexClientError, consume: bool) -> CodexResetCreditsError { - match error { - CodexClientError::Upstream { - status, - body, - diagnostics, - .. - } if status == reqwest::StatusCode::UNAUTHORIZED - && reset_credit_response_was_explicit_rejection(status, &diagnostics) => - { - CodexResetCreditsError::CredentialRefreshRequired { - upstream_body: Some(body), - } - } - CodexClientError::Upstream { - status, - body, - retry_after_seconds, - diagnostics, - .. - } => { - // 2xx 后发生的解码/响应体上限错误使用 synthetic 502 表示，但额度卡 - // 可能已经消费 - // 只有 transport 记录的真实非成功状态与错误状态一致时， - // 才能把它当作确定的上游拒绝并允许前端清除 pending 幂等键 - if consume && !reset_credit_response_was_explicit_rejection(status, &diagnostics) { - CodexResetCreditsError::ConsumeResultUnknown - } else { - CodexResetCreditsError::Upstream { - status: status.as_u16(), - body, - retry_after_seconds, - } - } - } - _ if consume => CodexResetCreditsError::ConsumeResultUnknown, - _ => CodexResetCreditsError::TransportUnavailable, - } -} - -fn reset_credit_response_was_explicit_rejection( - status: reqwest::StatusCode, - diagnostics: &crate::transport::CodexUpstreamDiagnostics, -) -> bool { - !status.is_success() && diagnostics.status_code == Some(status.as_u16()) -} - -async fn fetch_usage_once( - client: &CodexBackendClient, - prepared: &PreparedCodexRuntimeCredential, -) -> Result { - let authorization = prepared - .credential - .authentication - .authorization_header() - .map_err(|_| CodexQuotaFetchAttemptError::InvalidCredential)?; - let request_id = format!("quota_{}", Uuid::now_v7().simple()); - client - .for_account(&prepared.account) - .map_err(CodexQuotaFetchAttemptError::Upstream)? - .fetch_usage(CodexRequestContext::auxiliary( - authorization.expose_secret(), - prepared.account.upstream_account_id(), - &request_id, - None, - )) - .await - .map_err(CodexQuotaFetchAttemptError::Upstream) -} - -/// 对 5xx 上游拒绝做有限次指数退避重试（1s/2s），吞掉瞬时抖动 -/// -/// 4xx（含 402/429）不重试：它们已经走额度状态转换，重试只会放大上游负载 -async fn fetch_usage_with_5xx_retry( - client: &CodexBackendClient, - prepared: &PreparedCodexRuntimeCredential, -) -> Result { - let mut attempt = 0_u32; - loop { - let result = fetch_usage_once(client, prepared).await; - let retryable = match &result { - Ok(_) => false, - Err(CodexQuotaFetchAttemptError::Upstream(CodexClientError::Upstream { - status, - .. - })) => status.is_server_error(), - Err(_) => false, - }; - if !retryable || attempt >= QUOTA_FETCH_5XX_MAX_RETRIES { - return result; - } - attempt += 1; - let delay = QUOTA_FETCH_5XX_BASE_DELAY.saturating_mul(attempt); - tracing::warn!( - account_id = %prepared.account.id(), - retry_attempt = attempt, - retry_delay_ms = delay.as_millis(), - "OpenAI quota usage 5xx upstream rejection; retrying with backoff" - ); - tokio::time::sleep(delay).await; - } -} - -/// 上游额度可确认套餐变更；同族泛化值不能丢弃 JWT 已给出的具体 SKU -fn observed_account_plan(current: Option<&str>, observed: Option<&str>) -> Option { - let plan = observed?.trim().to_ascii_lowercase(); - if plan.is_empty() || plan == "unknown" { - return None; - } - // 套餐族沿用官方 codex_protocol::account::PlanType 的分类 - let current = current.unwrap_or_default().trim().to_ascii_lowercase(); - let generalized = matches!( - (plan.as_str(), current.as_str()), - ( - "team", - "self_serve_business_prolite" | "self_serve_business_usage_based" - ) | ( - "business", - "ent26" | "enterprise_cbp_automation" | "enterprise_cbp_usage_based" - ) | ("edu" | "education", "edu_plus" | "edu_pro") - ); - (!generalized).then_some(plan) -} - -fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool { - account.enabled() && access_token_is_current(account, now) -} - -fn eligible_initial_quota_sync(account: &ProviderAccount, now: SystemTime) -> bool { - // 首次观察只兜底刚入库、尚无 quota 快照的账号 - // 已耗尽等运行时状态必须由 - // periodic 路径处理，才能保留同一账号的最小复核间隔 - account.credential_state() == CredentialState::Ready - && eligible_periodic_quota_refresh(account, now) -} - fn access_token_is_current(account: &ProviderAccount, now: SystemTime) -> bool { account .access_token_expires_at() .is_none_or(|expires_at| expires_at > now) } - -fn merge_passive_quota( - mut quota: Map, - observations: &[ParsedRateLimits], -) -> Map { - let mut snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota); - for rate_limits in observations { - let default_is_named_alias = default_limit_is_named_alias(rate_limits); - let mut resolved_limit_ids = BTreeMap::new(); - for (wire_limit_id, details) in &rate_limits.limits { - // HTTP 与 WebSocket 都可能把活动具名桶镜像为默认 `codex` 窗口 - // 同一 wire 观察里存在相同具名事实时丢弃镜像，不触碰 core 桶 - if wire_limit_id == DEFAULT_CODEX_LIMIT_ID && default_is_named_alias { - continue; - } - let Some(limit_id) = snapshots.resolve_limit_id(details) else { - continue; - }; - let Some(rate_limit) = passive_rate_limit_snapshot(details) else { - continue; - }; - snapshots.upsert(&limit_id, details, rate_limit); - resolved_limit_ids.insert(wire_limit_id.as_str(), limit_id); - } - - let active_limit = rate_limits - .active_limit - .as_deref() - .and_then(|wire_limit_id| { - resolved_limit_ids.get(wire_limit_id).cloned().or_else(|| { - (!rate_limits.limits.contains_key(wire_limit_id)) - .then(|| wire_limit_id.to_owned()) - }) - }) - .or_else(|| resolved_limit_ids.get(DEFAULT_CODEX_LIMIT_ID).cloned()); - if let Some(active_limit) = active_limit { - quota.insert("active_limit".to_owned(), Value::String(active_limit)); - } - merge_passive_metadata(&mut quota, rate_limits); - } - snapshots.write_to_document(&mut quota); - quota -} - -fn merge_passive_metadata(quota: &mut Map, rate_limits: &ParsedRateLimits) { - if let Some(plan_type) = rate_limits.plan_type.as_ref() { - quota.insert("plan_type".to_owned(), Value::String(plan_type.clone())); - } - if let Some(promo_message) = rate_limits.promo_message.as_ref() { - quota.insert( - "promo_message".to_owned(), - Value::String(promo_message.clone()), - ); - } - if let Some(reached_type) = rate_limits.rate_limit_reached_type.as_ref() { - quota.insert( - "rate_limit_reached_type".to_owned(), - Value::String(reached_type.clone()), - ); - } - if let Some(credits) = rate_limits.credits.as_ref() { - let mut value = quota - .remove("credits") - .and_then(|value| value.as_object().cloned()) - .unwrap_or_default(); - value.insert("has_credits".to_owned(), Value::Bool(credits.has_credits)); - value.insert("unlimited".to_owned(), Value::Bool(credits.unlimited)); - if let Some(balance) = credits.balance.as_ref() { - value.insert("balance".to_owned(), Value::String(balance.clone())); - } else { - // 点数对象明确更新但未提供余额时，不能把旧余额继续作为当前余额展示 - value.remove("balance"); - } - quota.insert("credits".to_owned(), Value::Object(value)); - } -} - -fn default_limit_is_named_alias(rate_limits: &ParsedRateLimits) -> bool { - let Some(default_primary) = rate_limits - .limits - .get(DEFAULT_CODEX_LIMIT_ID) - .and_then(|details| details.primary) - else { - return false; - }; - rate_limits.limits.iter().any(|(limit_id, details)| { - limit_id != DEFAULT_CODEX_LIMIT_ID && details.primary == Some(default_primary) - }) -} - -/// 丢弃 core 限额中上游给出的无事实 `secondary_window` 占位 -/// -/// 生产响应头可能携带 `used_percent=0`、零时长和空 reset 的占位，`/usage` 也会 -/// 返回 `secondary_window: null` -/// 正常被动同步保留该响应事实，Admin 展示会将其 -/// 隐藏；但主动 `/usage` 刷新或 402 确认投影会把一个存在但无事实的字段写成 -/// 100%，从而显示并不存在的“次级额度” -/// 因此这些非被动写入口在落库前移除 -/// 无事实值 -/// 带 reset、时长、正用量、触顶、未知或非法字段的次级窗口均完全按 -/// 原有额度逻辑保留 -fn normalize_quota_window_placeholders(mut quota: Map) -> Map { - quota = canonicalize_rate_limit_document(quota); - if let Some(rate_limit) = quota - .get_mut(RATE_LIMITS_BY_LIMIT_ID) - .and_then(Value::as_object_mut) - .and_then(|limits| limits.get_mut(DEFAULT_CODEX_LIMIT_ID)) - .and_then(Value::as_object_mut) - { - drop_secondary_window_placeholder(rate_limit); - } - quota -} - -fn drop_secondary_window_placeholder(rate_limit: &mut Map) { - let placeholder = rate_limit.get("secondary_window").is_some_and(|window| { - window.is_null() - || window - .as_object() - .is_some_and(secondary_window_is_placeholder) - }); - if placeholder { - rate_limit.remove("secondary_window"); - } -} - -fn secondary_window_is_placeholder(window: &Map) -> bool { - window.iter().all(|(field, value)| match field.as_str() { - "used_percent" => value - .as_f64() - .is_some_and(|used_percent| used_percent.is_finite() && used_percent == 0.0), - "limit_reached" => value.as_bool() == Some(false), - _ => false, - }) -} - -fn passive_rate_limit_snapshot(details: &RateLimitDetails) -> Option> { - // 响应头的窗口属于同一次上游观测；跨响应合并会制造不存在的额度窗口 - let mut snapshot = Map::new(); - if let Some(allowed) = details.allowed { - snapshot.insert("allowed".to_owned(), Value::Bool(allowed)); - } - if let Some(limit_reached) = details.limit_reached { - snapshot.insert("limit_reached".to_owned(), Value::Bool(limit_reached)); - } - for (field, window) in [ - ("primary_window", details.primary), - ("secondary_window", details.secondary), - ] { - let Some(window) = window else { - continue; - }; - snapshot.insert( - field.to_owned(), - Value::Object(passive_rate_limit_window(window)), - ); - } - (!snapshot.is_empty()).then_some(snapshot) -} - -fn passive_rate_limit_window(window: RateLimitWindow) -> Map { - let mut snapshot = Map::new(); - if let Some(number) = serde_json::Number::from_f64(window.used_percent) { - snapshot.insert("used_percent".to_owned(), Value::Number(number)); - } - if let Some(seconds) = window - .window_minutes - .and_then(|minutes| minutes.checked_mul(60)) - { - snapshot.insert("limit_window_seconds".to_owned(), Value::from(seconds)); - } - if let Some(reset_at) = window.reset_at { - snapshot.insert("reset_at".to_owned(), Value::from(reset_at)); - } - snapshot -} diff --git a/backend/crates/providers/openai/src/credential/quota/passive.rs b/backend/crates/providers/openai/src/credential/quota/passive.rs new file mode 100644 index 000000000..b91553adf --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/passive.rs @@ -0,0 +1,236 @@ +//! 推理响应携带的额度观察合并与 Provider 快照规范化 + +use super::*; + +impl CodexCredentialQuotaService { + /// 把正常推理响应携带的限流事实合并进 Provider 原始 quota JSON + pub async fn synchronize_passive_headers( + &self, + account: &ProviderAccount, + headers: &[(String, String)], + ) -> Result { + let Some(rate_limits) = parse_rate_limit_headers(headers) else { + return Ok(false); + }; + self.synchronize_passive_rate_limits(account, std::slice::from_ref(&rate_limits)) + .await + } + + /// 把一次推理响应中采集的结构化限流观察合并后单次落库 + pub async fn synchronize_passive_rate_limits( + &self, + account: &ProviderAccount, + rate_limits: &[ParsedRateLimits], + ) -> Result { + if account.authentication_kind() != crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH + || rate_limits.is_empty() + { + return Ok(false); + } + let has_quota_facts = rate_limits.iter().any(|observation| { + observation + .limits + .values() + .any(|details| passive_rate_limit_snapshot(details).is_some()) + }); + let existing = self + .store + .get_quotas(std::slice::from_ref(account.id())) + .await? + .into_iter() + .find(|observation| { + observation.account_id == *account.id() + && observation.expected_revision == account.revision() + }); + let existing_state = existing.as_ref().map(|observation| observation.state); + let existing = existing + .map(|observation| observation.quota.into_inner()) + .unwrap_or_default(); + // 只用本次响应明确携带的套餐更新账号，不能把合并前的旧快照重新当作新证据 + let observed_plan = rate_limits + .iter() + .rev() + .filter_map(|observation| { + observation.plan_type.as_deref().filter(|plan| { + !plan.trim().is_empty() && !plan.trim().eq_ignore_ascii_case("unknown") + }) + }) + .next(); + let plan_type = observed_account_plan(account.plan_type(), observed_plan); + // 套餐、credits 等元数据可以更新，但没有额度窗口事实时必须保留旧观察时刻， + // 也不能借旧快照重新推导 quota state + if !has_quota_facts { + let Some(state) = existing_state else { + return Ok(false); + }; + let outcome = self + .store + .compare_and_swap_quota(QuotaObservation { + plan_type, + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: OpaqueProviderData::new(merge_passive_quota(existing, rate_limits)), + observed_at: SystemTime::now(), + state, + }) + .await?; + return Ok(outcome != QuotaWriteOutcome::Conflict); + } + let quota = merge_passive_quota(existing, rate_limits); + let observed_at = SystemTime::now(); + let snapshot = parse_account_quota_snapshot( + account.id().clone(), + account.revision(), + observed_at, + &Value::Object(quota.clone()), + )?; + // 这些 headers 来自一次成功推理，访问结论优先于可能滞后的百分比 + let state = QuotaState::allowed(observed_at); + let outcome = self + .store + .compare_and_swap_quota(QuotaObservation { + plan_type, + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: OpaqueProviderData::new(quota), + observed_at, + state, + }) + .await?; + if outcome == QuotaWriteOutcome::Conflict { + return Ok(false); + } + self.scheduling.observe(&snapshot); + Ok(true) + } +} + +fn merge_passive_quota( + mut quota: Map, + observations: &[ParsedRateLimits], +) -> Map { + let mut snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota); + for rate_limits in observations { + let default_is_named_alias = default_limit_is_named_alias(rate_limits); + let mut resolved_limit_ids = BTreeMap::new(); + for (wire_limit_id, details) in &rate_limits.limits { + // HTTP 与 WebSocket 都可能把活动具名桶镜像为默认 `codex` 窗口 + // 同一 wire 观察里存在相同具名事实时丢弃镜像，不触碰 core 桶 + if wire_limit_id == DEFAULT_CODEX_LIMIT_ID && default_is_named_alias { + continue; + } + let Some(limit_id) = snapshots.resolve_limit_id(details) else { + continue; + }; + let Some(rate_limit) = passive_rate_limit_snapshot(details) else { + continue; + }; + snapshots.upsert(&limit_id, details, rate_limit); + resolved_limit_ids.insert(wire_limit_id.as_str(), limit_id); + } + + let active_limit = rate_limits + .active_limit + .as_deref() + .and_then(|wire_limit_id| { + resolved_limit_ids.get(wire_limit_id).cloned().or_else(|| { + (!rate_limits.limits.contains_key(wire_limit_id)) + .then(|| wire_limit_id.to_owned()) + }) + }) + .or_else(|| resolved_limit_ids.get(DEFAULT_CODEX_LIMIT_ID).cloned()); + if let Some(active_limit) = active_limit { + quota.insert("active_limit".to_owned(), Value::String(active_limit)); + } + merge_passive_metadata(&mut quota, rate_limits); + } + snapshots.write_to_document(&mut quota); + quota +} + +fn merge_passive_metadata(quota: &mut Map, rate_limits: &ParsedRateLimits) { + if let Some(plan_type) = rate_limits.plan_type.as_ref() { + quota.insert("plan_type".to_owned(), Value::String(plan_type.clone())); + } + if let Some(promo_message) = rate_limits.promo_message.as_ref() { + quota.insert( + "promo_message".to_owned(), + Value::String(promo_message.clone()), + ); + } + if let Some(reached_type) = rate_limits.rate_limit_reached_type.as_ref() { + quota.insert( + "rate_limit_reached_type".to_owned(), + Value::String(reached_type.clone()), + ); + } + if let Some(credits) = rate_limits.credits.as_ref() { + let mut value = quota + .remove("credits") + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + value.insert("has_credits".to_owned(), Value::Bool(credits.has_credits)); + value.insert("unlimited".to_owned(), Value::Bool(credits.unlimited)); + if let Some(balance) = credits.balance.as_ref() { + value.insert("balance".to_owned(), Value::String(balance.clone())); + } else { + // 点数对象明确更新但未提供余额时，不能把旧余额继续作为当前余额展示 + value.remove("balance"); + } + quota.insert("credits".to_owned(), Value::Object(value)); + } +} + +fn default_limit_is_named_alias(rate_limits: &ParsedRateLimits) -> bool { + let Some(default_primary) = rate_limits + .limits + .get(DEFAULT_CODEX_LIMIT_ID) + .and_then(|details| details.primary) + else { + return false; + }; + rate_limits.limits.iter().any(|(limit_id, details)| { + limit_id != DEFAULT_CODEX_LIMIT_ID && details.primary == Some(default_primary) + }) +} + +fn passive_rate_limit_snapshot(details: &RateLimitDetails) -> Option> { + // 响应头的窗口属于同一次上游观测；跨响应合并会制造不存在的额度窗口 + let mut snapshot = Map::new(); + if let Some(allowed) = details.allowed { + snapshot.insert("allowed".to_owned(), Value::Bool(allowed)); + } + if let Some(limit_reached) = details.limit_reached { + snapshot.insert("limit_reached".to_owned(), Value::Bool(limit_reached)); + } + for (field, window) in [ + ("primary_window", details.primary), + ("secondary_window", details.secondary), + ] { + let Some(window) = window else { + continue; + }; + snapshot.insert( + field.to_owned(), + Value::Object(passive_rate_limit_window(window)), + ); + } + (!snapshot.is_empty()).then_some(snapshot) +} + +fn passive_rate_limit_window(window: RateLimitWindow) -> Map { + let mut snapshot = Map::new(); + if let Some(number) = serde_json::Number::from_f64(window.used_percent) { + snapshot.insert("used_percent".to_owned(), Value::Number(number)); + } + if let Some(seconds) = window + .window_minutes + .and_then(|minutes| minutes.checked_mul(60)) + { + snapshot.insert("limit_window_seconds".to_owned(), Value::from(seconds)); + } + if let Some(reset_at) = window.reset_at { + snapshot.insert("reset_at".to_owned(), Value::from(reset_at)); + } + snapshot +} diff --git a/backend/crates/providers/openai/src/credential/quota/reset_credits.rs b/backend/crates/providers/openai/src/credential/quota/reset_credits.rs new file mode 100644 index 000000000..c8d0f49b6 --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/reset_credits.rs @@ -0,0 +1,315 @@ +//! 主动额度重置卡查询、按账号串行消费与不确定结果分类 + +use super::*; + +/// 仅索引活跃的消费或排队者，锁资源由请求持有 +#[derive(Default)] +pub(super) struct ResetCreditLocks { + entries: + std::sync::Mutex>>>, +} + +impl ResetCreditLocks { + fn entry(&self, account_id: &ProviderAccountId) -> ResetCreditEntry<'_> { + let mut entries = self + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let lock = entries + .get(account_id) + .and_then(std::sync::Weak::upgrade) + .unwrap_or_else(|| { + let lock = Arc::new(Mutex::new(())); + entries.insert(account_id.clone(), Arc::downgrade(&lock)); + lock + }); + ResetCreditEntry { + owner: self, + account_id: account_id.clone(), + lock: Some(lock), + } + } +} + +struct ResetCreditEntry<'a> { + owner: &'a ResetCreditLocks, + account_id: ProviderAccountId, + lock: Option>>, +} + +impl ResetCreditEntry<'_> { + async fn lock(&self) -> tokio::sync::MutexGuard<'_, ()> { + self.lock + .as_ref() + .expect("账号锁只在 entry 销毁时释放") + .lock() + .await + } +} + +impl Drop for ResetCreditEntry<'_> { + fn drop(&mut self) { + let mut entries = self + .owner + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // 引用释放也须在表锁内完成，避免并发 Drop 都看见彼此而留下死 Weak + // 有持锁或排队者时仍保留同一锁身份；等待 future 取消也走此清理路径 + drop(self.lock.take()); + if entries + .get(&self.account_id) + .is_some_and(|lock| lock.strong_count() == 0) + { + entries.remove(&self.account_id); + } + } +} + +/// 主动额度重置卡查询/消费失败 +#[derive(Error)] +pub enum CodexResetCreditsError { + #[error("Codex reset-credit credential data is invalid")] + InvalidCredentialData, + #[error("Codex OAuth access token must be refreshed before using reset credits")] + CredentialRefreshRequired { upstream_body: Option }, + #[error("Codex reset-credit account was not found")] + NotFound, + #[error("provider account store is unavailable: {detail}")] + Store { detail: String }, + #[error("Codex reset-credit upstream returned HTTP {status}")] + Upstream { + status: u16, + body: String, + retry_after_seconds: Option, + }, + #[error("Codex reset-credit query transport is unavailable")] + TransportUnavailable, + #[error("Codex reset-credit consume result is unknown")] + ConsumeResultUnknown, +} + +impl std::fmt::Debug for CodexResetCreditsError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InvalidCredentialData => formatter.write_str("InvalidCredentialData"), + Self::CredentialRefreshRequired { .. } => { + formatter.write_str("CredentialRefreshRequired { upstream_body: }") + } + Self::NotFound => formatter.write_str("NotFound"), + Self::Store { .. } => formatter.write_str("Store { detail: }"), + Self::Upstream { + status, + retry_after_seconds, + .. + } => formatter + .debug_struct("Upstream") + .field("status", status) + .field("body", &"") + .field("retry_after_seconds", retry_after_seconds) + .finish(), + Self::TransportUnavailable => formatter.write_str("TransportUnavailable"), + Self::ConsumeResultUnknown => formatter.write_str("ConsumeResultUnknown"), + } + } +} + +enum ResetCreditAttemptError { + InvalidCredential, + Upstream(CodexClientError), +} + +impl CodexCredentialQuotaService { + /// 查询当前账号由 Codex Desktop 暴露的主动额度重置卡 + pub async fn list_reset_credits( + &self, + account_id: &ProviderAccountId, + ) -> Result { + let account = self.reset_credit_account(account_id).await?; + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + let credential = self + .repository + .load_runtime_credential(&account) + .await + .map_err(|_| CodexResetCreditsError::InvalidCredentialData)?; + let prepared = PreparedCodexRuntimeCredential { + account, + credential, + }; + let request_id = format!("reset_credits_{}", Uuid::now_v7().simple()); + list_reset_credits_once(&client, &prepared, &request_id) + .await + .map_err(|error| map_reset_credit_attempt_error(error, false)) + } + + /// 消费一张主动额度重置卡 + /// 相同账号在本进程内串行，且不做传输重试 + pub async fn consume_reset_credit( + &self, + account_id: &ProviderAccountId, + credit_id: Option<&str>, + redeem_request_id: Uuid, + ) -> Result { + let entry = self.reset_consume_locks.entry(account_id); + let _guard = entry.lock().await; + let account = self.reset_credit_account(account_id).await?; + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + let credential = self + .repository + .load_runtime_credential(&account) + .await + .map_err(|_| CodexResetCreditsError::InvalidCredentialData)?; + let prepared = PreparedCodexRuntimeCredential { + account, + credential, + }; + let request_id = format!("reset_credit_consume_{}", Uuid::now_v7().simple()); + consume_reset_credit_once( + &client, + &prepared, + &request_id, + credit_id, + redeem_request_id, + ) + .await + .map_err(|error| map_reset_credit_attempt_error(error, true)) + } + + async fn reset_credit_account( + &self, + account_id: &ProviderAccountId, + ) -> Result { + let account = self + .store + .get_account(account_id) + .await + .map_err(|error| CodexResetCreditsError::Store { + detail: error.to_string(), + })? + .filter(|account| account.provider().as_str() == "openai") + .ok_or(CodexResetCreditsError::NotFound)?; + if !access_token_is_current(&account, SystemTime::now()) { + return Err(CodexResetCreditsError::CredentialRefreshRequired { + upstream_body: None, + }); + } + Ok(account) + } +} + +async fn list_reset_credits_once( + client: &CodexBackendClient, + prepared: &PreparedCodexRuntimeCredential, + request_id: &str, +) -> Result { + let authorization = prepared + .credential + .authentication + .authorization_header() + .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; + client + .for_account(&prepared.account) + .map_err(ResetCreditAttemptError::Upstream)? + .list_rate_limit_reset_credits(CodexRequestContext::auxiliary( + authorization.expose_secret(), + prepared.account.upstream_account_id(), + request_id, + None, + )) + .await + .map_err(ResetCreditAttemptError::Upstream) +} + +async fn consume_reset_credit_once( + client: &CodexBackendClient, + prepared: &PreparedCodexRuntimeCredential, + request_id: &str, + credit_id: Option<&str>, + redeem_request_id: Uuid, +) -> Result { + let authorization = prepared + .credential + .authentication + .authorization_header() + .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; + client + .for_account(&prepared.account) + .map_err(ResetCreditAttemptError::Upstream)? + .consume_rate_limit_reset_credit( + CodexRequestContext::auxiliary( + authorization.expose_secret(), + prepared.account.upstream_account_id(), + request_id, + None, + ), + credit_id, + redeem_request_id, + ) + .await + .map_err(ResetCreditAttemptError::Upstream) +} + +fn map_reset_credit_attempt_error( + error: ResetCreditAttemptError, + consume: bool, +) -> CodexResetCreditsError { + match error { + ResetCreditAttemptError::InvalidCredential => CodexResetCreditsError::InvalidCredentialData, + ResetCreditAttemptError::Upstream(error) => map_reset_credit_client_error(error, consume), + } +} + +fn map_reset_credit_client_error(error: CodexClientError, consume: bool) -> CodexResetCreditsError { + match error { + CodexClientError::Upstream { + status, + body, + diagnostics, + .. + } if status == reqwest::StatusCode::UNAUTHORIZED + && reset_credit_response_was_explicit_rejection(status, &diagnostics) => + { + CodexResetCreditsError::CredentialRefreshRequired { + upstream_body: Some(body), + } + } + CodexClientError::Upstream { + status, + body, + retry_after_seconds, + diagnostics, + .. + } => { + // 2xx 后发生的解码/响应体上限错误使用 synthetic 502 表示，但额度卡 + // 可能已经消费 + // 只有 transport 记录的真实非成功状态与错误状态一致时， + // 才能把它当作确定的上游拒绝并允许前端清除 pending 幂等键 + if consume && !reset_credit_response_was_explicit_rejection(status, &diagnostics) { + CodexResetCreditsError::ConsumeResultUnknown + } else { + CodexResetCreditsError::Upstream { + status: status.as_u16(), + body, + retry_after_seconds, + } + } + } + _ if consume => CodexResetCreditsError::ConsumeResultUnknown, + _ => CodexResetCreditsError::TransportUnavailable, + } +} + +fn reset_credit_response_was_explicit_rejection( + status: reqwest::StatusCode, + diagnostics: &crate::transport::CodexUpstreamDiagnostics, +) -> bool { + !status.is_success() && diagnostics.status_code == Some(status.as_u16()) +} diff --git a/backend/crates/providers/openai/src/credential/quota/scheduling.rs b/backend/crates/providers/openai/src/credential/quota/scheduling.rs new file mode 100644 index 000000000..4b3226ea8 --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/scheduling.rs @@ -0,0 +1,414 @@ +//! 请求级额度投影、首查与周期刷新排程 + +use super::*; + +impl CodexCredentialQuotaService { + /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度 + pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { + if self.scheduling.hydration_targets(accounts).is_empty() { + return; + } + let _hydration = self.scheduling.hydration.lock().await; + let pending = self.scheduling.hydration_targets(accounts); + if pending.is_empty() { + return; + } + let account_ids = pending + .iter() + .map(|target| target.account.id().clone()) + .collect::>(); + let Ok(observations) = self.store.get_quotas(&account_ids).await else { + for target in &pending { + self.scheduling + .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); + } + return; + }; + let observations = observations + .into_iter() + .map(|observation| (observation.account_id.clone(), observation)) + .collect::>(); + for target in pending { + let snapshot = observations + .get(target.account.id()) + .filter(|observation| observation.expected_revision == target.account.revision()) + .and_then(quota_snapshot_from_observation); + if !snapshot + .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) + { + self.scheduling + .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); + } + } + } + + #[must_use] + pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { + self.scheduling.signals(account) + } + + pub(crate) fn invalidate_scheduling(&self, account_ids: &[ProviderAccountId]) { + self.scheduling.invalidate(account_ids); + } +} + +#[derive(Clone, Default)] +pub(super) struct CodexQuotaSchedulingProjection { + state: Arc>, + hydration: Arc>, +} + +#[derive(Default)] +struct CodexQuotaProjectionState { + next_version: u64, + entries: BTreeMap, + last_periodic_refresh_at: BTreeMap, + initial_refresh_not_before: BTreeMap, +} + +struct CodexQuotaRefreshAttempt { + monotonic_at: Instant, + wall_at: SystemTime, +} + +#[derive(Debug, Clone, Copy)] +struct CodexQuotaSchedulingEntry { + version: u64, + revision: CredentialRevision, + expires_at: Instant, + signals: Option, +} + +#[derive(Clone)] +struct CodexQuotaHydrationTarget { + account: ProviderAccount, + expected_version: Option, +} + +impl CodexQuotaSchedulingProjection { + fn invalidate(&self, account_ids: &[ProviderAccountId]) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for account_id in account_ids { + state.entries.remove(account_id); + state.last_periodic_refresh_at.remove(account_id); + state.initial_refresh_not_before.remove(account_id); + } + } + + fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let now = Instant::now(); + accounts + .iter() + .filter_map(|account| { + let current = state.entries.get(account.id()); + current + .is_none_or(|entry| { + entry.revision != account.revision() || now >= entry.expires_at + }) + .then(|| CodexQuotaHydrationTarget { + account: account.clone(), + expected_version: current.map(|entry| entry.version), + }) + }) + .collect() + } + + pub(super) fn observe(&self, snapshot: &CodexAccountQuotaSnapshot) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot) else { + return false; + }; + self.replace( + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + scheduling_signals_from_snapshot(snapshot), + ); + true + } + + fn mark_unknown_if_unchanged(&self, target: &CodexQuotaHydrationTarget, ttl: Duration) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return; + } + insert_projection_entry( + &mut state, + target.account.id().clone(), + target.account.revision(), + ttl, + None, + ); + } + + fn observe_if_unchanged( + &self, + target: &CodexQuotaHydrationTarget, + snapshot: &CodexAccountQuotaSnapshot, + ) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot) else { + return false; + }; + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return true; + } + insert_projection_entry( + &mut state, + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + scheduling_signals_from_snapshot(snapshot), + ); + true + } + + fn replace( + &self, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, + ) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + insert_projection_entry(&mut state, account_id, revision, ttl, signals); + } + + fn signals(&self, account: &ProviderAccount) -> Option { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .entries + .get(account.id()) + .filter(|entry| { + entry.revision == account.revision() && Instant::now() < entry.expires_at + }) + .and_then(|entry| entry.signals) + } + + pub(super) fn reserve_periodic_refreshes( + &self, + accounts: Vec, + observed_snapshots: &BTreeMap, + now: SystemTime, + ) -> Vec { + let candidates = accounts + .into_iter() + .filter_map(|account| { + let snapshot = observed_snapshots.get(account.id()); + quota_refresh_candidate(account, snapshot, now) + }) + .collect::>(); + let candidate_ids = candidates + .iter() + .map(|(account, _)| account.id().clone()) + .collect::>(); + let refreshed_at = Instant::now(); + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .last_periodic_refresh_at + .retain(|account_id, _| candidate_ids.contains(account_id)); + + // 已耗尽账号首次立即复核，之后每 30 分钟复核，以发现官方提前重置 + // reset + 2 分钟额外触发一次复核，给上游重置留出传播时间 + // 正常账号仅在非零用量窗口经过宽限期后参与，未更新时复用周期节流 + let mut reserved = Vec::new(); + for (account, target_reset) in candidates { + if !periodic_quota_refresh_due(&state, account.id(), target_reset, now, refreshed_at) { + continue; + } + state.last_periodic_refresh_at.insert( + account.id().clone(), + CodexQuotaRefreshAttempt { + monotonic_at: refreshed_at, + wall_at: now, + }, + ); + reserved.push(account); + } + reserved + } + + /// 为尚无 quota 快照的账号分配随机首查时点，并返回本轮已到期的候选。 + /// + /// `quota_observed_at` 为空代表首次异步观察尚未成功；不另建同步状态表。 + /// 首次进入候选的账号抽一次随机启动延迟记为到期时刻，避免新导入账号池 + /// 在 worker 唤醒边界形成批量 `/usage` 突发；到期时刻跨轮保留，零延迟 + /// （含随机源退化）时保持首轮即查的旧行为。单轮数量仍受 + /// `INITIAL_QUOTA_SYNC_BATCH` 上限约束。 + pub(super) fn reserve_initial_refreshes( + &self, + accounts: &[ProviderAccount], + observed_ids: &BTreeSet, + now: SystemTime, + draw_start_delay: &dyn Fn() -> Duration, + ) -> Vec { + let candidates = accounts + .iter() + .filter(|account| { + !observed_ids.contains(account.id()) && eligible_initial_quota_sync(account, now) + }) + .collect::>(); + let candidate_ids = candidates + .iter() + .map(|account| account.id().clone()) + .collect::>(); + let monotonic_now = Instant::now(); + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // 已产生快照或离开候选集的账号不再保留首查状态。 + state + .initial_refresh_not_before + .retain(|account_id, _| candidate_ids.contains(account_id)); + candidates + .into_iter() + .filter(|account| { + let due_at = state + .initial_refresh_not_before + .entry(account.id().clone()) + .or_insert_with(|| monotonic_now + draw_start_delay()); + *due_at <= monotonic_now + }) + .take(INITIAL_QUOTA_SYNC_BATCH) + .cloned() + .collect() + } + + /// 首查尝试后仍无快照的账号重抽随机延迟，避免失败账号回到固定 30s 重试节奏。 + /// + /// 重排同时随机化每轮的候选集合，避免持续失败时按账号顺序只重试同一批。 + pub(super) fn defer_failed_initial_refreshes( + &self, + account_ids: &BTreeSet, + draw_start_delay: &dyn Fn() -> Duration, + ) { + if account_ids.is_empty() { + return; + } + let now = Instant::now(); + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for account_id in account_ids { + // 账号是否仍在首查路径由下一轮 reserve 按候选集剪枝，这里无需区分。 + state + .initial_refresh_not_before + .insert(account_id.clone(), now + draw_start_delay()); + } + } +} + +fn quota_refresh_candidate( + account: ProviderAccount, + snapshot: Option<&CodexAccountQuotaSnapshot>, + now: SystemTime, +) -> Option<(ProviderAccount, Option)> { + if !eligible_periodic_quota_refresh(&account, now) { + return None; + } + if account.quota().is_exhausted() { + let reset_at = account.quota().reset_at(); + return Some((account, reset_at)); + } + // 正常账号首次进入周期候选也要等待宽限期，不能由“无刷新历史”绕过 + let snapshot = snapshot?; + let expired_window_reset = snapshot + .windows() + .iter() + .filter(|window| { + window.used_percent().is_some_and(|used| used > 0.0) || window.limit_reached() + }) + .filter_map(CodexQuotaWindow::reset_at) + .map(SystemTime::from) + .filter(|reset| { + reset + .checked_add(QUOTA_RESET_GRACE) + .is_some_and(|due_at| due_at <= now) + }) + .min(); + expired_window_reset.map(|reset_at| (account, Some(reset_at))) +} + +fn periodic_quota_refresh_due( + state: &CodexQuotaProjectionState, + account_id: &ProviderAccountId, + reset_at: Option, + now: SystemTime, + monotonic_now: Instant, +) -> bool { + state + .last_periodic_refresh_at + .get(account_id) + .is_none_or(|last| { + monotonic_now.saturating_duration_since(last.monotonic_at) + >= PERIODIC_QUOTA_REFRESH_RETRY_INTERVAL + || reset_at + .and_then(|reset| reset.checked_add(QUOTA_RESET_GRACE)) + // 已在该边界之后复核过时回到周期重试，避免过期 reset 每轮触发 + .is_some_and(|due_at| last.wall_at < due_at && due_at <= now) + }) +} + +fn insert_projection_entry( + state: &mut CodexQuotaProjectionState, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, +) { + state.next_version = state.next_version.saturating_add(1); + state.entries.insert( + account_id, + CodexQuotaSchedulingEntry { + version: state.next_version, + revision, + expires_at: Instant::now() + ttl, + signals, + }, + ); +} + +fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool { + account.enabled() && access_token_is_current(account, now) +} + +fn eligible_initial_quota_sync(account: &ProviderAccount, now: SystemTime) -> bool { + // 首次观察只兜底刚入库、尚无 quota 快照的账号 + // 已耗尽等运行时状态必须由 + // periodic 路径处理，才能保留同一账号的最小复核间隔 + account.credential_state() == CredentialState::Ready + && eligible_periodic_quota_refresh(account, now) +} diff --git a/backend/crates/providers/openai/src/credential/quota/service.rs b/backend/crates/providers/openai/src/credential/quota/service.rs new file mode 100644 index 000000000..b5b824f21 --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/service.rs @@ -0,0 +1,756 @@ +//! 额度权威读取、刷新、错误反馈与服务组装 + +use super::*; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum QuotaRefreshAuthority { + ObserveAccess, + PreserveAccess, +} + +struct FetchedCodexQuota { + account: ProviderAccount, + value: Value, +} + +enum CodexQuotaFetchAttemptError { + InvalidCredential, + Upstream(CodexClientError), +} + +enum CodexQuotaFetchError { + InvalidCredential, + Upstream { + account: Box, + error: CodexClientError, + }, +} + +/// 客户端错误携带的 HTTP 状态码；传输失败等为 `None` +fn upstream_error_status(error: &CodexClientError) -> Option { + match error { + CodexClientError::Upstream { status, .. } => Some(status.as_u16()), + _ => None, + } +} + +/// 从上游错误体提取稳定错误码：优先 `/error/code`，其次 `/code` +/// +/// 只接受有界的 ASCII 标识，避免把自由文本作为错误码写入诊断 +fn upstream_error_code(error: &CodexClientError) -> Option { + let CodexClientError::Upstream { body, .. } = error else { + return None; + }; + let value = serde_json::from_str::(body).ok()?; + let code = value + .pointer("/error/code") + .or_else(|| value.pointer("/code")) + .and_then(Value::as_str)? + .trim(); + if code.is_empty() + || code.len() > 64 + || !code + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.')) + { + return None; + } + Some(code.to_owned()) +} + +impl CodexCredentialQuotaService { + pub fn new( + repository: CodexCredentialRepository, + profile: CodexWireProfileState, + http: Client, + base_url: String, + cooldowns: Arc, + leases: Arc, + runtime_policy: Arc, + ) -> Self { + Self { + store: Arc::clone(repository.store()), + repository, + profile, + http, + base_url, + cooldowns, + leases, + runtime_policy, + freeze_policy_cache: Mutex::new(None), + scheduling: CodexQuotaSchedulingProjection::default(), + reset_consume_locks: ResetCreditLocks::default(), + initial_sync_delays: Arc::new(|| { + crate::jitter::uniform_delay( + crate::jitter::random_u64(), + INITIAL_QUOTA_SYNC_MAX_START_DELAY, + ) + }), + } + } + + /// 注入首查随机延迟采样器（测试用确定性节奏）；生产默认均匀 + /// `[0, INITIAL_QUOTA_SYNC_MAX_START_DELAY)`，零延迟等价旧行为 + #[doc(hidden)] + pub fn with_initial_sync_delays(mut self, delays: CodexInitialSyncDelays) -> Self { + self.initial_sync_delays = delays; + self + } + + /// 读取容量熔断策略（带短 TTL 缓存）；读取失败退化为关闭，熔断不得放大失败 + async fn freeze_policy(&self) -> ProviderFreezePolicy { + { + let cache = self.freeze_policy_cache.lock().await; + if let Some((policy, loaded_at)) = cache.as_ref() + && loaded_at.elapsed() < FREEZE_POLICY_CACHE_TTL + { + return policy.clone(); + } + } + let loaded = self.runtime_policy.load_freeze_policy().await.ok(); + let mut cache = self.freeze_policy_cache.lock().await; + if let Some(policy) = loaded { + *cache = Some((policy.clone(), Instant::now())); + return policy; + } + cache + .as_ref() + .filter(|(_, loaded_at)| loaded_at.elapsed() < FREEZE_POLICY_CACHE_TTL) + .map_or_else(ProviderFreezePolicy::disabled, |(policy, _)| policy.clone()) + } + + /// 容量熔断入口：滑动窗口内累计容量类失败，达到阈值即写入带 + /// `CapacityFreeze` 类别的账号级冷却 + /// 调度侧立即屏蔽该账号，恢复由 + /// freeze-recovery worker 处理；本路径只依赖 Redis 可丢失事实 + pub async fn apply_capacity_failure(&self, account: &ProviderAccount, observed_at: SystemTime) { + let policy = self.freeze_policy().await; + if !policy.enabled() { + return; + } + let in_flight = self.current_in_flight(account.id()).await; + let Ok(count) = self + .cooldowns + .record_capacity_failure(account.id(), policy.window(), in_flight) + .await + else { + return; + }; + if count < policy.threshold() { + return; + } + let Some(until) = observed_at.checked_add(policy.freeze_duration()) else { + return; + }; + let cooldown = ProviderCooldown::new_with_kind( + account.id().clone(), + account.revision(), + until, + if policy.probe_enabled() { + ProviderCooldownKind::CapacityFreezeProbe + } else { + ProviderCooldownKind::CapacityFreeze + }, + ); + if self.cooldowns.put_if_later(cooldown).await.is_ok() { + tracing::warn!( + account_id = account.id().as_str(), + threshold = policy.threshold(), + window_seconds = policy.window().as_secs(), + freeze_seconds = policy.freeze_duration().as_secs(), + peak_in_flight = in_flight, + "账号容量熔断触发：冻结该账号一段时间", + ); + } + } + + async fn current_in_flight(&self, account_id: &ProviderAccountId) -> u32 { + self.leases + .account_in_flight(std::slice::from_ref(account_id)) + .await + .ok() + .and_then(|signals| signals.get(account_id).copied()) + .unwrap_or(0) + } + + /// 真实推理错误只更新额度访问事实，不伪造 Provider JSON 或展示百分比 + pub(crate) async fn record_confirmed_exhaustion( + &self, + account: &ProviderAccount, + evidence: QuotaEvidence, + reset_at: Option, + observed_at: SystemTime, + ) -> Result<(), CodexCredentialQuotaError> { + let outcome = self + .store + .apply_quota_access(QuotaAccessChange { + account_id: account.id().clone(), + expected_revision: account.revision(), + state: QuotaState::exhausted(evidence, observed_at, reset_at), + }) + .await?; + if outcome == QuotaWriteOutcome::Conflict { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + Ok(()) + } + + /// 成功推理是额度可访问的权威证据，同时解除账号级 429 冷却 + pub async fn record_successful_inference( + &self, + account: &ProviderAccount, + observed_at: SystemTime, + ) -> Result<(), CodexCredentialQuotaError> { + // 已经发出的并发请求可能晚于耗尽事实成功返回；真实成功仍是独立的 + // Allowed 权威证据，但调度不会为了试探恢复而放行耗尽账号 + if account.quota().access() == QuotaAccessState::Exhausted { + let outcome = self + .store + .apply_quota_access(QuotaAccessChange { + account_id: account.id().clone(), + expected_revision: account.revision(), + state: QuotaState::allowed(observed_at), + }) + .await?; + if outcome == QuotaWriteOutcome::Conflict { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + } + self.cooldowns + .clear_after_success(account.id(), account.revision()) + .await + .map_err(|error| CodexCredentialQuotaError::Store { + detail: error.to_string(), + })?; + Ok(()) + } + + pub async fn synchronize(&self) -> Result { + self.synchronize_at(SystemTime::now()).await + } + + /// 按本轮调度时刻选择到期账号；实际 HTTP 观察与落库仍使用发生时刻 + pub async fn synchronize_at( + &self, + now: SystemTime, + ) -> Result { + let mut accounts = self.repository.list_for_provider().await?; + accounts.retain(|account| { + account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH + }); + let mut summary = CodexQuotaSyncSummary::default(); + let account_ids = accounts + .iter() + .map(|account| account.id().clone()) + .collect::>(); + let observed = self.store.get_quotas(&account_ids).await?; + let observed_snapshots = observed + .iter() + .filter_map(|obs| { + quota_snapshot_from_observation(obs) + .map(|snapshot| (obs.account_id.clone(), snapshot)) + }) + .collect::>(); + let observed_ids = observed + .into_iter() + .map(|observation| observation.account_id) + .collect::>(); + let initial = self.scheduling.reserve_initial_refreshes( + &accounts, + &observed_ids, + now, + self.initial_sync_delays.as_ref(), + ); + // 本轮首查后仍无快照的账号需要重排随机延迟。 + let mut pending_initial: BTreeSet = + initial.iter().map(|account| account.id().clone()).collect(); + let periodic = + self.scheduling + .reserve_periodic_refreshes(accounts, &observed_snapshots, now); + let accounts = initial + .into_iter() + .chain(periodic) + .fold(BTreeMap::new(), |mut unique, account| { + unique.insert(account.id().clone(), account); + unique + }) + .into_values() + .collect::>(); + if accounts.is_empty() { + return Ok(summary); + } + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + for account in accounts { + let observed_at = SystemTime::now(); + match self.fetch_usage(&client, &account).await { + Ok(FetchedCodexQuota { account, value }) => { + // 单账号解析或落库失败只影响该账号；其余账号继续同步 + match self + .apply_fetched_quota(&account, &value, observed_at, &mut summary) + .await + { + Ok(()) => { + pending_initial.remove(account.id()); + } + Err(error) => { + summary.transient += 1; + tracing::warn!( + account_id = %account.id(), + error = %error, + "OpenAI quota synchronization skipped one account" + ); + } + } + } + Err(CodexQuotaFetchError::InvalidCredential) => { + summary.stale += 1; + } + Err(CodexQuotaFetchError::Upstream { account, error }) => { + match classify_quota_endpoint_failure(&error) { + Some(QuotaEndpointFailure::Exhausted(evidence)) => { + summary.exhausted += 1; + if let Err(write_error) = self + .record_confirmed_exhaustion(&account, evidence, None, observed_at) + .await + { + summary.transient += 1; + tracing::warn!( + account_id = %account.id(), + error = %write_error, + "OpenAI quota exhaustion fact write failed" + ); + } + } + Some(QuotaEndpointFailure::Credential { state, reason }) => { + summary.banned += 1; + self.persist_credential_failure(&account, state, reason, observed_at) + .await; + } + None => { + summary.transient += 1; + tracing::warn!( + account_id = %account.id(), + error = %error, + upstream_status = upstream_error_status(&error), + upstream_code = ?upstream_error_code(&error), + "OpenAI quota upstream rejection; refresh cycle will retry later" + ); + } + } + } + } + } + // 首查失败或未能落库的账号重排随机延迟；已成功落库的账号由快照事实自然退出首查路径。 + self.scheduling + .defer_failed_initial_refreshes(&pending_initial, self.initial_sync_delays.as_ref()); + Ok(summary) + } + + #[must_use] + pub fn runtime_policy(&self) -> &Arc { + &self.runtime_policy + } + + /// 解析并 revision-fenced 落库单账号的 Provider quota JSON + async fn apply_fetched_quota( + &self, + account: &ProviderAccount, + value: &Value, + observed_at: SystemTime, + summary: &mut CodexQuotaSyncSummary, + ) -> Result<(), CodexCredentialQuotaError> { + let mut object = normalize_quota_window_placeholders( + value + .as_object() + .cloned() + .ok_or(CodexCredentialQuotaError::InvalidCredentialData)?, + ); + object.remove(RECOVERY_FIELD); + let mut snapshot = parse_account_quota_snapshot( + account.id().clone(), + account.revision(), + observed_at, + &Value::Object(object.clone()), + )?; + let previous = if account.quota().is_exhausted() { + self.read_snapshot_for(account).await? + } else { + None + }; + let state = reconcile_refresh( + account.quota(), + &mut snapshot, + previous.as_ref(), + &mut object, + )?; + let snapshot = snapshot.with_quota_state(state); + let outcome = self + .store + .compare_and_swap_quota(QuotaObservation { + plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()), + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: OpaqueProviderData::new(object), + observed_at, + state, + }) + .await?; + if outcome == QuotaWriteOutcome::Conflict { + summary.stale += 1; + return Ok(()); + } + self.scheduling.observe(&snapshot); + if snapshot.quota().is_exhausted() { + summary.exhausted += 1; + } else { + summary.updated += 1; + } + Ok(()) + } + + /// 真实 429 的单一事实入口：写入 Redis 临时限流冷却（`until = now + retry_after`） + /// 凭据与额度主窗口（额度重置时间）都不改变——临时限流是独立维度， + /// 到期由 Redis key 过期自动解除，不污染配额耗尽状态 + /// 已有更晚的冷却不会被缩短（put_if_later） + pub async fn apply_rate_limit_429( + &self, + account: &ProviderAccount, + retry_after: Option, + observed_at: SystemTime, + ) -> Result<(), CodexCredentialQuotaError> { + let until = observed_at + .checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)) + .unwrap_or(observed_at); + self.cooldowns + .put_if_later(ProviderCooldown::new( + account.id().clone(), + account.revision(), + until, + )) + .await + .map_err(|error| CodexCredentialQuotaError::Store { + detail: error.to_string(), + })?; + Ok(()) + } + + /// 读取有效的账号冷却事实；等待恢复探测的冻结到期后仍有效 + pub async fn cooldown( + &self, + account_id: &ProviderAccountId, + ) -> Result, CodexCredentialQuotaError> { + let Some(cooldown) = self.cooldowns.read(account_id).await.map_err(|error| { + CodexCredentialQuotaError::Store { + detail: error.to_string(), + } + })? + else { + return Ok(None); + }; + let state = cooldown.scheduling_state(); + Ok(state.is_active(SystemTime::now()).then_some(state)) + } + + /// 读取单账号最后一次落库的 Provider quota，并由 Codex 域解析展示窗口 + pub async fn read_account( + &self, + account_id: &ProviderAccountId, + ) -> Result, CodexCredentialQuotaError> { + let account = self + .store + .get_account(account_id) + .await? + .filter(|account| account.provider().as_str() == "openai") + .ok_or(CodexCredentialQuotaError::NotFound)?; + self.read_snapshot_for(&account).await + } + + async fn read_snapshot_for( + &self, + account: &ProviderAccount, + ) -> Result, CodexCredentialQuotaError> { + if account.provider().as_str() != "openai" { + return Err(CodexCredentialQuotaError::NotFound); + } + let account_id = account.id(); + let Some(observation) = self + .store + .get_quotas(std::slice::from_ref(account_id)) + .await? + .into_iter() + .next() + else { + return Ok(None); + }; + if observation.account_id != *account_id + || observation.expected_revision != account.revision() + { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + if observation.quota.expose_to_provider().is_empty() { + return Ok(None); + } + let observed_at = observation.observed_at; + let snapshot = parse_account_quota_snapshot( + account_id.clone(), + account.revision(), + observed_at, + &Value::Object(observation.quota.expose_to_provider().clone()), + )? + .with_quota_state(observation.state); + self.scheduling.observe(&snapshot); + Ok(Some(snapshot)) + } + + /// 只刷新指定账号，revision-fenced 写入动态 Provider JSON 后返回解析快照 + pub async fn refresh_account( + &self, + account_id: &ProviderAccountId, + ) -> Result { + self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::ObserveAccess) + .await + } + + /// 真实限额失败后的异步刷新只补齐展示快照，不允许 usage 快照撤销已确认的失败状态 + pub(crate) async fn refresh_account_after_failure( + &self, + account_id: &ProviderAccountId, + ) -> Result { + self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::PreserveAccess) + .await + } + + async fn refresh_account_snapshot( + &self, + account_id: &ProviderAccountId, + authority: QuotaRefreshAuthority, + ) -> Result { + let account = self + .store + .get_account(account_id) + .await? + .filter(|account| account.provider().as_str() == "openai") + .ok_or(CodexCredentialQuotaError::NotFound)?; + if account.authentication_kind() != crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH { + return Err(CodexCredentialQuotaError::NotFound); + } + let observed_at = SystemTime::now(); + if !access_token_is_current(&account, observed_at) { + return Err(CodexCredentialQuotaError::CredentialRefreshRequired); + } + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + let FetchedCodexQuota { account, value } = match self.fetch_usage(&client, &account).await { + Ok(fetched) => fetched, + Err(CodexQuotaFetchError::InvalidCredential) => { + return Err(CodexCredentialQuotaError::InvalidCredentialData); + } + Err(CodexQuotaFetchError::Upstream { account, error }) => { + match classify_quota_endpoint_failure(&error) { + Some(QuotaEndpointFailure::Exhausted(evidence)) => { + self.record_confirmed_exhaustion(&account, evidence, None, observed_at) + .await?; + } + Some(QuotaEndpointFailure::Credential { state, reason }) => { + self.persist_credential_failure(&account, state, reason, observed_at) + .await; + } + None => {} + } + return Err(CodexCredentialQuotaError::Upstream { + detail: error.to_string(), + status: upstream_error_status(&error), + code: upstream_error_code(&error), + }); + } + }; + let mut object = normalize_quota_window_placeholders( + value + .as_object() + .cloned() + .ok_or(CodexCredentialQuotaError::InvalidCredentialData)?, + ); + object.remove(RECOVERY_FIELD); + let mut snapshot = parse_account_quota_snapshot( + account.id().clone(), + account.revision(), + observed_at, + &Value::Object(object.clone()), + )?; + let previous = if account.quota().is_exhausted() { + self.read_snapshot_for(&account).await? + } else { + None + }; + if account.quota().is_exhausted() + && matches!(authority, QuotaRefreshAuthority::PreserveAccess) + { + if self + .store + .touch_quota_observation(QuotaObservationTouch { + account_id: account.id().clone(), + expected_revision: account.revision(), + observed_at, + }) + .await? + == QuotaWriteOutcome::Conflict + { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + return Ok(match previous { + Some(previous) => { + let previous = previous.with_observed_at(observed_at); + self.scheduling.observe(&previous); + previous + } + None => snapshot + .with_quota_state(account.quota()) + .with_observed_at(observed_at), + }); + } + let state = match authority { + QuotaRefreshAuthority::ObserveAccess => reconcile_refresh( + account.quota(), + &mut snapshot, + previous.as_ref(), + &mut object, + )?, + QuotaRefreshAuthority::PreserveAccess => account.quota(), + }; + let snapshot = snapshot.with_quota_state(state); + if self + .store + .compare_and_swap_quota(QuotaObservation { + plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()), + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: OpaqueProviderData::new(object), + observed_at, + state, + }) + .await? + == QuotaWriteOutcome::Conflict + { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + self.scheduling.observe(&snapshot); + Ok(snapshot) + } + + async fn persist_credential_failure( + &self, + account: &ProviderAccount, + credential_state: CredentialState, + reason: AccountErrorReason, + observed_at: SystemTime, + ) { + if let Err(error) = self + .repository + .apply_state_with_reason(account, credential_state, observed_at, Some(reason), None) + .await + { + tracing::warn!( + account_id = %account.id(), + ?credential_state, + reason = reason.as_str(), + error = %error, + "OpenAI quota credential fact write failed" + ); + } + } + + async fn fetch_usage( + &self, + client: &CodexBackendClient, + account: &ProviderAccount, + ) -> Result { + let credential = self + .repository + .load_runtime_credential(account) + .await + .map_err(|_| CodexQuotaFetchError::InvalidCredential)?; + let prepared = PreparedCodexRuntimeCredential { + account: account.clone(), + credential, + }; + let result = fetch_usage_with_5xx_retry(client, &prepared).await; + match result { + Ok(value) => Ok(FetchedCodexQuota { + account: prepared.account, + value, + }), + Err(CodexQuotaFetchAttemptError::InvalidCredential) => { + Err(CodexQuotaFetchError::InvalidCredential) + } + Err(CodexQuotaFetchAttemptError::Upstream(error)) => { + Err(CodexQuotaFetchError::Upstream { + account: Box::new(prepared.account), + error, + }) + } + } + } +} + +async fn fetch_usage_once( + client: &CodexBackendClient, + prepared: &PreparedCodexRuntimeCredential, +) -> Result { + let authorization = prepared + .credential + .authentication + .authorization_header() + .map_err(|_| CodexQuotaFetchAttemptError::InvalidCredential)?; + let request_id = format!("quota_{}", Uuid::now_v7().simple()); + client + .for_account(&prepared.account) + .map_err(CodexQuotaFetchAttemptError::Upstream)? + .fetch_usage(CodexRequestContext::auxiliary( + authorization.expose_secret(), + prepared.account.upstream_account_id(), + &request_id, + None, + )) + .await + .map_err(CodexQuotaFetchAttemptError::Upstream) +} + +/// 对 5xx 上游拒绝做有限次指数退避重试（1s/2s），吞掉瞬时抖动 +/// +/// 4xx（含 402/429）不重试：它们已经走额度状态转换，重试只会放大上游负载 +async fn fetch_usage_with_5xx_retry( + client: &CodexBackendClient, + prepared: &PreparedCodexRuntimeCredential, +) -> Result { + let mut attempt = 0_u32; + loop { + let result = fetch_usage_once(client, prepared).await; + let retryable = match &result { + Ok(_) => false, + Err(CodexQuotaFetchAttemptError::Upstream(CodexClientError::Upstream { + status, + .. + })) => status.is_server_error(), + Err(_) => false, + }; + if !retryable || attempt >= QUOTA_FETCH_5XX_MAX_RETRIES { + return result; + } + attempt += 1; + let delay = QUOTA_FETCH_5XX_BASE_DELAY.saturating_mul(attempt); + tracing::warn!( + account_id = %prepared.account.id(), + retry_attempt = attempt, + retry_delay_ms = delay.as_millis(), + "OpenAI quota usage 5xx upstream rejection; retrying with backoff" + ); + tokio::time::sleep(delay).await; + } +} diff --git a/backend/crates/providers/openai/src/credential/quota/warmup.rs b/backend/crates/providers/openai/src/credential/quota/warmup.rs new file mode 100644 index 000000000..9bf45da8c --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/warmup.rs @@ -0,0 +1,260 @@ +//! 账号额度窗口预热与完成事件校验 + +use super::*; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum WarmupTerminal { + Missing, + Completed, + Failed, +} + +fn observe_warmup_events(events: Vec, terminal: &mut WarmupTerminal) { + for event in events { + let parsed = serde_json::from_str::(&event.data).ok(); + let kind = parsed + .as_ref() + .and_then(|data| data.get("type")) + .and_then(Value::as_str) + .or(event.event.as_deref()); + match kind { + Some("response.completed") => { + // 只有明确完成的响应才算预热成功；失败事件即使随后出现完成帧也优先 + if *terminal != WarmupTerminal::Failed { + *terminal = if parsed.is_some() + && parsed + .as_ref() + .and_then(|data| data.pointer("/response/status")) + .and_then(Value::as_str) + .is_none_or(|status| status == "completed") + { + WarmupTerminal::Completed + } else { + WarmupTerminal::Failed + }; + } + } + Some("response.failed" | "response.incomplete" | "error") => { + *terminal = WarmupTerminal::Failed; + } + _ => {} + } + } +} + +async fn consume_warmup_sse( + response: &mut CodexBackendStreamingResponse, +) -> Result<(), &'static str> { + let mut decoder = SseEventDecoder::default(); + let mut terminal = WarmupTerminal::Missing; + let mut received_bytes = 0usize; + while let Some(chunk) = response.body.next().await { + let chunk = chunk.map_err(|_| "stream_read_failed")?; + received_bytes = received_bytes.saturating_add(chunk.len()); + if received_bytes > WARMUP_STREAM_MAX_BYTES { + return Err("stream_size_limit_exceeded"); + } + let events = decoder.push(&chunk).map_err(|_| "invalid_sse")?; + observe_warmup_events(events, &mut terminal); + if terminal != WarmupTerminal::Missing { + break; + } + } + if terminal == WarmupTerminal::Missing { + observe_warmup_events(decoder.finish().map_err(|_| "invalid_sse")?, &mut terminal); + } + match terminal { + WarmupTerminal::Completed => Ok(()), + WarmupTerminal::Missing => Err("completion_event_missing"), + WarmupTerminal::Failed => Err("response_failed"), + } +} + +impl CodexCredentialQuotaService { + /// 批量预激活 OAuth 账号的 5h 配额滑动窗口 + pub async fn execute_warmup( + &self, + model: &str, + ) -> Result { + let mut accounts = self.repository.list_for_provider().await?; + accounts.retain(|account| { + account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH + && account.enabled() + && matches!( + account.credential_state(), + CredentialState::Unknown | CredentialState::Ready + ) + }); + let mut summary = CodexWarmupSummary::default(); + if accounts.is_empty() { + return Ok(summary); + } + let account_ids = accounts + .iter() + .map(|account| account.id().clone()) + .collect::>(); + let observed = self.store.get_quotas(&account_ids).await?; + let observed_snapshots = observed + .iter() + .filter_map(|obs| { + quota_snapshot_from_observation(obs) + .map(|snapshot| (obs.account_id.clone(), snapshot)) + }) + .collect::>(); + + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + + let now_utc = chrono::Utc::now(); + for account in accounts { + if let Some(snapshot) = observed_snapshots.get(account.id()) { + // 1. 周线已触顶或耗尽跳过 + let weekly_exhausted = snapshot + .windows() + .iter() + .any(|w| w.kind() == CodexQuotaWindowKind::Weekly && w.limit_reached()); + if weekly_exhausted { + summary.skipped_exhausted += 1; + continue; + } + // 2. 5h 窗口当前活跃且距重置时间 > 30 分钟跳过 + let has_active_5h = snapshot.windows().iter().any(|w| { + w.kind() == CodexQuotaWindowKind::ShortTerm + && w.reset_at().is_some_and(|reset_at| { + reset_at > now_utc + chrono::Duration::minutes(30) + }) + }); + if has_active_5h { + summary.skipped_active += 1; + continue; + } + } + + let credential = match self.repository.load_runtime_credential(&account).await { + Ok(cred) => cred, + Err(_) => { + summary.failed += 1; + continue; + } + }; + let authorization = match credential.authentication.authorization_header() { + Ok(auth) => auth, + Err(_) => { + summary.failed += 1; + continue; + } + }; + + let mut body = Map::new(); + body.insert("model".to_owned(), Value::String(model.to_owned())); + body.insert( + "input".to_owned(), + serde_json::json!([{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }]), + ); + body.insert("stream".to_owned(), Value::Bool(true)); + body.insert("store".to_owned(), Value::Bool(false)); + body.insert( + "service_tier".to_owned(), + Value::String("default".to_owned()), + ); + body.insert( + "reasoning".to_owned(), + serde_json::json!({"effort": "none"}), + ); + body.insert("text".to_owned(), serde_json::json!({"verbosity": "low"})); + + let upstream_request = CodexResponsesRequest::from_body(body); + let request_id = format!("warmup_{}", Uuid::now_v7().simple()); + let client_for_account = match client.for_account(&account) { + Ok(c) => c, + Err(error) => { + tracing::warn!(account_id = %account.id(), error = %error, "warmup client for account failed"); + summary.failed += 1; + continue; + } + }; + let context = crate::transport::CodexRequestContext::auxiliary( + authorization.expose_secret(), + account.upstream_account_id(), + &request_id, + None, + ); + + // HTTP 200 只确认响应头；必须消费 SSE 直到终态才能确认预热结果和额度事件 + let attempt = tokio::time::timeout(WARMUP_REQUEST_TIMEOUT, async { + let mut response = client_for_account + .create_response_stream_http_sse(&upstream_request, context) + .await + .map_err(|_| "request_rejected")?; + let terminal = consume_warmup_sse(&mut response).await; + let rate_limit_updates = match response.rate_limit_updates.as_ref() { + Some(updates) => std::mem::take(&mut *updates.lock().await), + None => Vec::new(), + }; + Ok::<_, &'static str>((terminal, response.rate_limit_headers, rate_limit_updates)) + }) + .await; + match attempt { + Ok(Ok((terminal, headers, rate_limit_updates))) => { + match terminal { + Ok(()) => { + // 被动额度同步会把成功请求的窗口事实标为可用，失败流不能复用该结论 + if !headers.is_empty() + && let Err(error) = + self.synchronize_passive_headers(&account, &headers).await + { + tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit header sync failed"); + } + if !rate_limit_updates.is_empty() + && let Err(error) = self + .synchronize_passive_rate_limits(&account, &rate_limit_updates) + .await + { + tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit event sync failed"); + } + summary.warmed_up += 1; + tracing::info!( + account_id = %account.id(), + model, + "OpenAI account warmed up successfully" + ); + } + Err(reason) => { + summary.failed += 1; + tracing::warn!( + account_id = %account.id(), + reason, + "OpenAI account warmup stream failed" + ); + } + } + } + Ok(Err(reason)) => { + summary.failed += 1; + tracing::warn!( + account_id = %account.id(), + reason, + "OpenAI account warmup request rejected by upstream" + ); + } + Err(_) => { + summary.failed += 1; + tracing::warn!( + account_id = %account.id(), + "OpenAI account warmup request timed out" + ); + } + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + + Ok(summary) + } +} diff --git a/backend/crates/providers/openai/src/credential/security.rs b/backend/crates/providers/openai/src/credential/security.rs index bd440045e..96a6c4789 100644 --- a/backend/crates/providers/openai/src/credential/security.rs +++ b/backend/crates/providers/openai/src/credential/security.rs @@ -46,6 +46,14 @@ pub enum CodexRuntimeAuthentication { } impl CodexRuntimeAuthentication { + /// 传输只消费公开连接设置，不读取凭据枚举或秘密 + pub(crate) fn responses_api_base_url(&self) -> Option<&str> { + match self { + Self::ApiKey(auth) => Some(&auth.configuration.base_url), + Self::OAuth(_) => None, + } + } + pub fn authorization_header(&self) -> Result { match self { Self::ApiKey(auth) => Ok(SecretString::from(format!( diff --git a/backend/crates/providers/openai/src/lib.rs b/backend/crates/providers/openai/src/lib.rs index d1506bbf0..17cda4d01 100644 --- a/backend/crates/providers/openai/src/lib.rs +++ b/backend/crates/providers/openai/src/lib.rs @@ -4,8 +4,12 @@ mod admin; pub mod config; mod jitter; mod provider; +mod request_identity; mod session_transport; +/// 持久账号使用的 API Key 认证种类，凭据与账号传输隔离共用 +pub(crate) const CODEX_AUTHENTICATION_KIND_API_KEY: &str = "api_key"; + use std::sync::Arc; use gateway_admin::ports::provider::ProviderAdmin; diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index 145ec52eb..aa05fa48d 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -314,7 +314,7 @@ impl CodexProvider { .client_for_request(&context)? .for_account(lease.account()) .map_err(|error| map_client_error(error, UpstreamSendState::NotSent, false).error)? - .with_authentication(lease.authentication()) + .with_responses_api_base_url(lease.authentication().responses_api_base_url()) .with_middleware_headers(middleware_headers), response_origin: request.response_origin, endpoint_path: request.endpoint_path, diff --git a/backend/crates/providers/openai/src/provider/live.rs b/backend/crates/providers/openai/src/provider/live.rs index e49c37e38..f8b2b57ea 100644 --- a/backend/crates/providers/openai/src/provider/live.rs +++ b/backend/crates/providers/openai/src/provider/live.rs @@ -367,7 +367,7 @@ impl CodexProvider { .client_for_request(&context)? .for_account(lease.account()) .map_err(|error| map_client_error(error, UpstreamSendState::NotSent, false).error)? - .with_authentication(lease.authentication()) + .with_responses_api_base_url(lease.authentication().responses_api_base_url()) .with_middleware_headers(middleware_headers), registry: Arc::clone(&self.live_registry), response_origin: self.live_calls_url.clone(), diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 4066a523f..7121a80ad 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -59,12 +59,12 @@ use crate::credential::{ CodexCredentialRefreshService, CodexCredentialRepository, CodexCredentialSelector, CodexCyberPolicyScope, CodexQuotaRefreshPolicy, CodexSessionAffinity, CredentialSelectionError, RuntimeCodexCookie, SelectCodexCredential, SelectCodexProviderEndpointCredential, - account_session_with_headers, derive_codex_cyber_policy_session_key, - derive_codex_endpoint_session_affinity, derive_codex_session_affinity, - derive_codex_transport_key, derive_endpoint_affinity_with_headers, - derive_live_session_affinity, derive_previous_response_id_hash, derive_turn_alias, - turn_id_with_headers, + derive_codex_cyber_policy_session_key, derive_codex_endpoint_session_affinity, + derive_codex_session_affinity, derive_codex_transport_key, + derive_endpoint_affinity_with_headers, derive_live_session_affinity, + derive_previous_response_id_hash, derive_turn_alias, }; +use crate::request_identity::{account_session_with_headers, turn_id_with_headers}; use crate::session_transport::CodexSessionTransportRecovery; use crate::transport::canonical::{ CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, @@ -704,7 +704,7 @@ impl CodexProvider { validate_openai_reasoning(generate.protocol_payload().body())?; let mut upstream = encode_generate_request(&generate, upstream_model.as_str(), None) .map_err(map_request_error)?; - upstream.client_account_follow_only = crate::credential::follows_session_with_headers( + upstream.client_account_follow_only = crate::request_identity::follows_session_with_headers( generate.protocol_payload().body(), generate.protocol_payload().context(), &middleware_headers, @@ -1002,7 +1002,7 @@ impl CodexProvider { .client_for_request(&context)? .for_account(lease.account()) .map_err(|error| map_client_error(error, UpstreamSendState::NotSent, false).error)? - .with_authentication(lease.authentication()) + .with_responses_api_base_url(lease.authentication().responses_api_base_url()) .with_connection_budget(context.connection_budget().clone()) .with_response_control(context.response_control().cloned()) .with_middleware_headers(middleware_headers), diff --git a/backend/crates/providers/openai/src/request_identity.rs b/backend/crates/providers/openai/src/request_identity.rs new file mode 100644 index 000000000..81e456d4e --- /dev/null +++ b/backend/crates/providers/openai/src/request_identity.rs @@ -0,0 +1,102 @@ +//! 请求正文与最终头部的会话身份解释，供账号亲和与传输编码共用 + +use serde_json::{Map, Value}; + +/// 只有能识别出的后代线程禁止迁移会话；没有身份的普通客户端沿用现有调度 +pub(crate) fn follows_session_with_headers( + body: &Map, + context: &Map, + headers: &[gateway_core::engine::middleware::MiddlewareHeader], +) -> bool { + let metadata = headers + .iter() + .find(|h| h.name().eq_ignore_ascii_case("x-codex-turn-metadata")) + .and_then(|h| std::str::from_utf8(h.value()).ok()); + let thread = account_thread_with_headers(body, context, headers); + if let Some((session, thread)) = + account_session_with_headers(body, context, headers).zip(thread) + { + return session != thread; + } + gateway_protocol::openai::codex_responses_request_semantics_with_turn_metadata( + body, + metadata.or_else(|| context.get("turn_metadata").and_then(Value::as_str)), + ) + .subagent_kind + .is_some() +} + +fn account_thread_with_headers( + body: &Map, + context: &Map, + headers: &[gateway_core::engine::middleware::MiddlewareHeader], +) -> Option { + let header = |name: &str| { + headers + .iter() + .find(|h| h.name().eq_ignore_ascii_case(name)) + .and_then(|h| std::str::from_utf8(h.value()).ok()) + }; + header("x-codex-turn-metadata") + .and_then(gateway_protocol::openai::turn_metadata_thread_id) + .or_else(|| { + header("thread-id") + .and_then(|value| non_empty(Some(value))) + .map(str::to_owned) + }) + .or_else(|| gateway_protocol::openai::codex_account_thread_id(body, context)) +} + +/// 中间件最终请求头按实际传输的覆盖语义参与身份复验 +pub(crate) fn account_session_with_headers( + body: &Map, + context: &Map, + headers: &[gateway_core::engine::middleware::MiddlewareHeader], +) -> Option { + let header = |name: &str| { + headers + .iter() + .find(|h| h.name().eq_ignore_ascii_case(name)) + .and_then(|h| std::str::from_utf8(h.value()).ok()) + }; + if let Some(session) = + header("x-codex-turn-metadata").and_then(gateway_protocol::openai::turn_metadata_session_id) + { + return Some(session); + } + let mut context = context.clone(); + for (name, field) in [ + ("session-id", "session_id"), + ("x-codex-turn-metadata", "turn_metadata"), + ] { + if let Some(value) = header(name) { + context.insert(field.to_owned(), Value::String(value.to_owned())); + } + } + gateway_protocol::openai::codex_account_session_id(body, &context) +} + +pub(crate) fn turn_id_with_headers( + body: &Map, + context: &Map, + headers: &[gateway_core::engine::middleware::MiddlewareHeader], +) -> Option { + headers + .iter() + .find(|header| header.name().eq_ignore_ascii_case("x-codex-turn-metadata")) + .and_then(|header| std::str::from_utf8(header.value()).ok()) + .and_then(gateway_protocol::openai::turn_metadata_turn_id) + .or_else(|| { + headers + .iter() + .find(|header| header.name().eq_ignore_ascii_case("x-client-turn-id")) + .and_then(|header| std::str::from_utf8(header.value()).ok()) + .and_then(|value| non_empty(Some(value))) + .map(str::to_owned) + }) + .or_else(|| gateway_protocol::openai::codex_turn_id(body, context)) +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} diff --git a/backend/crates/providers/openai/src/transport/client.rs b/backend/crates/providers/openai/src/transport/client.rs index bcd130343..adc6c14ea 100644 --- a/backend/crates/providers/openai/src/transport/client.rs +++ b/backend/crates/providers/openai/src/transport/client.rs @@ -775,12 +775,9 @@ impl CodexBackendClient { fallback.send().await.map_err(CodexClientError::HttpJson) } - pub(crate) fn with_authentication( - mut self, - authentication: &crate::credential::CodexRuntimeAuthentication, - ) -> Self { - if let crate::credential::CodexRuntimeAuthentication::ApiKey(auth) = authentication { - self.base_url = auth.configuration.base_url.trim_end_matches('/').to_owned(); + pub(crate) fn with_responses_api_base_url(mut self, base_url: Option<&str>) -> Self { + if let Some(base_url) = base_url { + self.base_url = base_url.trim_end_matches('/').to_owned(); self.protocol = OpenAiUpstreamProtocol::ResponsesApi; self.websocket_origin_key = format!( "{}:{}", @@ -798,7 +795,7 @@ impl CodexBackendClient { let mut client = self.clone(); client.outbound_proxy = account.outbound_proxy().cloned(); client.egress_key = egress_key(account.id().as_str(), account.outbound_proxy()); - if account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_API_KEY { + if account.authentication_kind() == crate::CODEX_AUTHENTICATION_KIND_API_KEY { client .egress_key .push_str(&format!(":revision:{}", account.revision().get())); diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index 6d0cad009..af338662a 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -730,7 +730,7 @@ fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map bool { true } @@ -424,7 +424,7 @@ impl gateway_core::runtime::extensions::ExtensionSetLease for SessionExtensionLe #[tokio::test] async fn child_binding_only_constrains_builtin_scheduling_and_preserves_plugin_choices() { use gateway_core::engine::policy::RequestPolicyContext; - use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetReference}; + use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetReference}; for explicit in [false, true] { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_subagent_a").await; diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index 67c46dba2..8380a77c7 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -54,11 +54,19 @@ pub(crate) struct MemoryAccountStore { accounts: Mutex>, quota_reads: AtomicUsize, fail_provider_listing: AtomicBool, + pending_account_reads: bool, credential_load_hook: Mutex>>, credential_loads: AtomicUsize, } impl MemoryAccountStore { + pub(crate) fn with_pending_account_reads() -> Self { + Self { + pending_account_reads: true, + ..Self::default() + } + } + pub(crate) fn repository(self: &Arc) -> CodexCredentialRepository { CodexCredentialRepository::new(self.clone()) } @@ -236,6 +244,9 @@ impl ProviderAccountStore for MemoryAccountStore { &self, account: &ProviderAccountId, ) -> Result, StoreError> { + if self.pending_account_reads { + futures::future::pending::<()>().await; + } Ok(self .accounts .lock() diff --git a/backend/crates/providers/xai/src/admin.rs b/backend/crates/providers/xai/src/admin.rs index 448d3bc66..fe6ee2cc8 100644 --- a/backend/crates/providers/xai/src/admin.rs +++ b/backend/crates/providers/xai/src/admin.rs @@ -1181,8 +1181,6 @@ fn project_quota( .map_or(Value::Null, |value| Value::String(value.to_owned())), ); for (key, value) in [ - ("monthlyLimitCents", billing.monthly_limit_cents()), - ("includedUsedCents", billing.included_used_cents()), ("onDemandCapCents", billing.on_demand_cap_cents()), ("onDemandUsedCents", billing.on_demand_used_cents()), ("prepaidBalanceCents", billing.prepaid_balance_cents()), diff --git a/backend/crates/providers/xai/src/credential/catalog.rs b/backend/crates/providers/xai/src/credential/catalog.rs index 35dfb578f..54969426d 100644 --- a/backend/crates/providers/xai/src/credential/catalog.rs +++ b/backend/crates/providers/xai/src/credential/catalog.rs @@ -3,536 +3,28 @@ use std::collections::BTreeMap; use std::fmt; use std::sync::{Arc, RwLock}; -use std::time::{Duration, Instant, SystemTime}; +use std::time::SystemTime; use async_trait::async_trait; use chrono::{DateTime, Utc}; use gateway_core::account::{ - AccountQuotaSignals, CredentialRevision, CredentialState, OpaqueProviderData, ProviderAccount, - ProviderAccountId, QuotaObservation, QuotaState, + CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId, }; use gateway_core::provider_ports::{ ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogScope, }; use gateway_core::routing::{ProviderCatalogGeneration, ProviderKind}; -use tokio::sync::Mutex; use super::repository::{GrokCredentialRepository, LoadedGrokCredential}; use crate::XaiWireProfileState; use crate::transport::catalog::{MAX_CATALOG_MODELS, valid_model_slug, validate_etag}; use crate::{ - GrokBillingClient, GrokBillingTransport, GrokCatalogModel, GrokModelCatalogClient, - GrokModelCatalogSession, GrokModelCatalogSnapshot, GrokModelCatalogTransport, SecretValue, - parse_grok_billing, + GrokCatalogModel, GrokModelCatalogClient, GrokModelCatalogSession, GrokModelCatalogSnapshot, + GrokModelCatalogTransport, SecretValue, }; const CATALOG_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(5 * 60); const MAX_CATALOG_FETCH_ATTEMPTS: usize = 3; -const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); -const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); - -/// Grok Build Free 额度的滚动观察窗口 -pub const GROK_FREE_ROLLING_WINDOW_SECONDS: u64 = 86_400; - -/// xAI Provider 从动态 billing JSON 解析出的旧账号页安全投影 -#[derive(Debug, Clone, PartialEq)] -pub struct GrokBillingPresentation { - plan_type: Option, - used_percent: Option, - period_type: Option, - period_start: Option, - period_end: Option, - monthly_limit_cents: Option, - included_used_cents: Option, - on_demand_cap_cents: Option, - on_demand_used_cents: Option, - prepaid_balance_cents: Option, -} - -/// xAI credits 当前周期的官方语义 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum GrokQuotaPeriodKind { - Weekly, - Monthly, - Other, -} - -impl GrokBillingPresentation { - #[must_use] - pub fn plan_type(&self) -> Option<&str> { - self.plan_type.as_deref() - } - - #[must_use] - pub const fn used_percent(&self) -> Option { - self.used_percent - } - - /// 返回 billing 是否包含可直接用于账号额度展示的付费事实 - #[must_use] - pub fn has_authoritative_quota(&self) -> bool { - self.used_percent.is_some() - || [ - self.monthly_limit_cents, - self.included_used_cents, - self.on_demand_cap_cents, - self.on_demand_used_cents, - self.prepaid_balance_cents, - ] - .into_iter() - .flatten() - .any(|value| value > 0) - } - - #[must_use] - pub fn period_type(&self) -> Option<&str> { - self.period_type.as_deref() - } - - #[must_use] - pub fn period_kind(&self) -> GrokQuotaPeriodKind { - match self.period_type.as_deref().map(str::trim) { - Some(value) - if value.eq_ignore_ascii_case("USAGE_PERIOD_TYPE_WEEKLY") - || value.eq_ignore_ascii_case("WEEKLY") => - { - GrokQuotaPeriodKind::Weekly - } - Some(value) - if value.eq_ignore_ascii_case("USAGE_PERIOD_TYPE_MONTHLY") - || value.eq_ignore_ascii_case("MONTHLY") => - { - GrokQuotaPeriodKind::Monthly - } - None if self.monthly_limit_cents.is_some() || self.included_used_cents.is_some() => { - GrokQuotaPeriodKind::Monthly - } - _ => GrokQuotaPeriodKind::Other, - } - } - - #[must_use] - pub fn period_start(&self) -> Option<&str> { - self.period_start.as_deref() - } - - #[must_use] - pub fn period_end(&self) -> Option<&str> { - self.period_end.as_deref() - } - - #[must_use] - pub const fn monthly_limit_cents(&self) -> Option { - self.monthly_limit_cents - } - - #[must_use] - pub const fn included_used_cents(&self) -> Option { - self.included_used_cents - } - - #[must_use] - pub const fn on_demand_cap_cents(&self) -> Option { - self.on_demand_cap_cents - } - - #[must_use] - pub const fn on_demand_used_cents(&self) -> Option { - self.on_demand_used_cents - } - - #[must_use] - pub const fn prepaid_balance_cents(&self) -> Option { - self.prepaid_balance_cents - } -} - -/// 一个账号最近一次 xAI billing 观察结果 -#[derive(Debug, Clone, PartialEq)] -pub struct GrokQuotaSnapshot { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - observed_at: DateTime, - billing: GrokBillingPresentation, -} - -impl GrokQuotaSnapshot { - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn observed_at(&self) -> DateTime { - self.observed_at - } - - #[must_use] - pub const fn billing(&self) -> &GrokBillingPresentation { - &self.billing - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] -pub enum GrokQuotaError { - #[error("xAI quota account is unavailable")] - AccountUnavailable, - #[error("xAI quota credential snapshot is stale")] - StaleCredentialSnapshot, - #[error("xAI quota credential or billing data is invalid")] - InvalidData, - #[error("xAI quota upstream request failed")] - Upstream, - #[error("xAI quota store is unavailable")] - Store, -} - -/// 官方 Grok Build billing 同步与 Provider-owned quota JSON 解析服务 -#[derive(Clone)] -pub struct GrokCredentialQuotaService { - repository: GrokCredentialRepository, - client: Arc, - scheduling: GrokQuotaSchedulingProjection, - wire_profile: XaiWireProfileState, -} - -#[derive(Clone, Default)] -struct GrokQuotaSchedulingProjection { - state: Arc>, - hydration: Arc>, -} - -#[derive(Default)] -struct GrokQuotaProjectionState { - next_version: u64, - entries: BTreeMap, -} - -#[derive(Debug, Clone, Copy)] -struct GrokQuotaSchedulingEntry { - version: u64, - revision: CredentialRevision, - expires_at: Instant, - signals: Option, -} - -#[derive(Clone)] -struct GrokQuotaHydrationTarget { - account: ProviderAccount, - expected_version: Option, -} - -impl GrokQuotaSchedulingProjection { - fn invalidate(&self, account_ids: &[ProviderAccountId]) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - for account_id in account_ids { - state.entries.remove(account_id); - } - } - - fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { - let state = self - .state - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let now = Instant::now(); - accounts - .iter() - .filter_map(|account| { - let current = state.entries.get(account.id()); - current - .is_none_or(|entry| { - entry.revision != account.revision() || now >= entry.expires_at - }) - .then(|| GrokQuotaHydrationTarget { - account: account.clone(), - expected_version: current.map(|entry| entry.version), - }) - }) - .collect() - } - - fn observe(&self, snapshot: &GrokQuotaSnapshot) -> bool { - let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { - return false; - }; - self.replace( - snapshot.account_id().clone(), - snapshot.credential_revision(), - remaining_ttl, - quota_scheduling_signals(snapshot.billing()), - ); - true - } - - fn observe_if_unchanged( - &self, - target: &GrokQuotaHydrationTarget, - snapshot: &GrokQuotaSnapshot, - ) -> bool { - let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { - return false; - }; - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state - .entries - .get(target.account.id()) - .map(|entry| entry.version) - != target.expected_version - { - return true; - } - insert_grok_projection_entry( - &mut state, - snapshot.account_id().clone(), - snapshot.credential_revision(), - remaining_ttl, - quota_scheduling_signals(snapshot.billing()), - ); - true - } - - fn mark_unknown_if_unchanged(&self, target: &GrokQuotaHydrationTarget, ttl: Duration) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state - .entries - .get(target.account.id()) - .map(|entry| entry.version) - != target.expected_version - { - return; - } - insert_grok_projection_entry( - &mut state, - target.account.id().clone(), - target.account.revision(), - ttl, - None, - ); - } - - fn replace( - &self, - account_id: ProviderAccountId, - revision: CredentialRevision, - ttl: Duration, - signals: Option, - ) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - insert_grok_projection_entry(&mut state, account_id, revision, ttl, signals); - } - - fn signals(&self, account: &ProviderAccount) -> Option { - let state = self - .state - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - state - .entries - .get(account.id()) - .filter(|entry| { - entry.revision == account.revision() && Instant::now() < entry.expires_at - }) - .and_then(|entry| entry.signals) - } -} - -fn insert_grok_projection_entry( - state: &mut GrokQuotaProjectionState, - account_id: ProviderAccountId, - revision: CredentialRevision, - ttl: Duration, - signals: Option, -) { - state.next_version = state.next_version.saturating_add(1); - state.entries.insert( - account_id, - GrokQuotaSchedulingEntry { - version: state.next_version, - revision, - expires_at: Instant::now() + ttl, - signals, - }, - ); -} - -impl GrokCredentialQuotaService { - #[must_use] - pub fn new( - repository: GrokCredentialRepository, - transport: Arc, - wire_profile: XaiWireProfileState, - ) -> Self { - Self { - repository, - client: Arc::new(GrokBillingClient::new(transport)), - scheduling: GrokQuotaSchedulingProjection::default(), - wire_profile, - } - } - - /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度 - pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { - if self.scheduling.hydration_targets(accounts).is_empty() { - return; - } - let _hydration = self.scheduling.hydration.lock().await; - let pending = self.scheduling.hydration_targets(accounts); - if pending.is_empty() { - return; - } - let pending_accounts = pending - .iter() - .map(|target| target.account.clone()) - .collect::>(); - let Ok(observations) = self.repository.quota_observations(&pending_accounts).await else { - for target in &pending { - self.scheduling - .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); - } - return; - }; - let observations = observations - .into_iter() - .map(|observation| (observation.account_id.clone(), observation)) - .collect::>(); - for target in pending { - let snapshot = observations - .get(target.account.id()) - .filter(|observation| observation.expected_revision == target.account.revision()) - .and_then(quota_snapshot_from_observation); - if !snapshot - .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) - { - self.scheduling - .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); - } - } - } - - #[must_use] - pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { - self.scheduling.signals(account) - } - - pub(crate) fn invalidate_scheduling(&self, account_ids: &[ProviderAccountId]) { - self.scheduling.invalidate(account_ids); - } - - /// 立即刷新一个账号的动态 billing document，并以 credential revision CAS 写回 - pub async fn refresh_account( - &self, - account_id: &ProviderAccountId, - ) -> Result { - let loaded = self - .repository - .load_current(account_id) - .await - .map_err(map_quota_repository_error)?; - if !loaded.account.enabled() - || loaded - .account - .access_token_expires_at() - .is_none_or(|expires_at| expires_at <= SystemTime::now()) - { - return Err(GrokQuotaError::AccountUnavailable); - } - let session = billing_session(&loaded, &self.wire_profile)?; - let billing = self - .client - .fetch(&session) - .await - .map_err(|_| GrokQuotaError::Upstream)?; - let mut document = billing.into_document(); - // 订阅查询失败不影响已取得的额度，也不能据此写入免费套餐 - match self.client.fetch_subscription(&session).await { - Ok(Some(plan_type)) => { - document.insert( - "subscriptionTier".to_owned(), - serde_json::Value::String(plan_type), - ); - } - Ok(None) => {} - Err(error) => tracing::warn!(%error, "xAI subscription query failed"), - } - let observed_at = Utc::now(); - let presentation = billing_presentation(&document)?; - let quota = loaded - .account - .quota() - .merge_observation(quota_access_observation(&presentation, observed_at.into())); - self.repository - .replace_quota( - loaded.account.id().clone(), - loaded.account.revision(), - document, - observed_at.into(), - quota, - ) - .await - .map_err(map_quota_repository_error)?; - let snapshot = GrokQuotaSnapshot { - account_id: loaded.account.id().clone(), - credential_revision: loaded.account.revision(), - observed_at, - billing: presentation, - }; - self.scheduling.observe(&snapshot); - Ok(snapshot) - } - - /// 读取并重新验证 Store 中的 Provider-owned quota JSON - pub async fn read_account( - &self, - account_id: &ProviderAccountId, - ) -> Result, GrokQuotaError> { - let Some(observation) = self - .repository - .quota(account_id) - .await - .map_err(map_quota_repository_error)? - else { - return Ok(None); - }; - let observed_at = observation.observed_at; - let document = observation.quota.into_inner(); - let snapshot = GrokQuotaSnapshot { - account_id: observation.account_id, - credential_revision: observation.expected_revision, - observed_at: observed_at.into(), - billing: billing_presentation(&document)?, - }; - self.scheduling.observe(&snapshot); - Ok(Some(snapshot)) - } -} - -impl fmt::Debug for GrokCredentialQuotaService { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("GrokCredentialQuotaService") - .field("repository", &self.repository) - .field("client", &self.client) - .finish() - } -} /// 官方 `/v1/models` 形成的一个 account 完整模型集合 #[derive(Clone, PartialEq, Eq)] @@ -1219,188 +711,3 @@ fn strict_model_union( } Ok(union.into_values().collect()) } - -fn billing_session( - loaded: &LoadedGrokCredential, - wire_profile: &XaiWireProfileState, -) -> Result { - let upstream_user_id = loaded - .account - .upstream_user_id() - .ok_or(GrokQuotaError::InvalidData)?; - GrokModelCatalogSession::new( - loaded.access_token.clone(), - SecretValue::new(upstream_user_id), - loaded - .account - .email() - .map(|value| SecretValue::new(value.to_owned())), - wire_profile.clone(), - ) - .map(|session| session.with_outbound_proxy(loaded.account.outbound_proxy().cloned())) - .map_err(|_| GrokQuotaError::InvalidData) -} - -fn billing_presentation( - document: &serde_json::Map, -) -> Result { - let body = serde_json::to_vec(document).map_err(|_| GrokQuotaError::InvalidData)?; - let snapshot = parse_grok_billing(&body).map_err(|_| GrokQuotaError::InvalidData)?; - let config = snapshot - .document() - .get("config") - .and_then(|value| value.as_object()); - let current_period = config - .and_then(|config| config.get("currentPeriod")) - .and_then(|value| value.as_object()); - let monthly_limit_cents = cent_value(config, "monthlyLimit"); - let included_used_cents = cent_value(config, "used"); - let used_percent = config - .and_then(|config| config.get("creditUsagePercent")) - .and_then(serde_json::Value::as_f64) - .or_else(|| match (included_used_cents, monthly_limit_cents) { - (Some(used), Some(limit)) if limit > 0 => { - Some((used.min(limit) as f64 / limit as f64) * 100.0) - } - _ => None, - }); - Ok(GrokBillingPresentation { - plan_type: dynamic_string(Some(snapshot.document()), "subscriptionTier"), - used_percent, - period_type: dynamic_string(current_period, "type"), - period_start: dynamic_string(current_period, "start") - .or_else(|| dynamic_string(config, "billingPeriodStart")), - period_end: dynamic_string(current_period, "end") - .or_else(|| dynamic_string(config, "billingPeriodEnd")), - monthly_limit_cents, - included_used_cents, - on_demand_cap_cents: cent_value(config, "onDemandCap"), - on_demand_used_cents: cent_value(config, "onDemandUsed"), - prepaid_balance_cents: cent_value(config, "prepaidBalance"), - }) -} - -fn quota_snapshot_from_observation(observation: &QuotaObservation) -> Option { - let observed_at = DateTime::::from(observation.observed_at); - let document = observation.quota.expose_to_provider(); - Some(GrokQuotaSnapshot { - account_id: observation.account_id.clone(), - credential_revision: observation.expected_revision, - observed_at, - billing: billing_presentation(document).ok()?, - }) -} - -fn quota_projection_ttl(observed_at: DateTime) -> Option { - let age = SystemTime::now() - .duration_since(SystemTime::from(observed_at)) - .unwrap_or(Duration::ZERO); - QUOTA_SCHEDULING_TTL - .checked_sub(age) - .filter(|remaining| !remaining.is_zero()) -} - -fn quota_scheduling_signals(billing: &GrokBillingPresentation) -> Option { - let now = SystemTime::now(); - let period_end = billing - .period_end() - .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) - .map(|value| SystemTime::from(value.to_utc())); - let observation_expired = period_end.is_some_and(|period_end| period_end <= now); - let remaining_rank = (!observation_expired) - .then(|| { - billing - .used_percent() - .filter(|used| used.is_finite() && (0.0..=100.0).contains(used)) - .map(|used| ((100.0 - used) * 100.0).round() as u64) - }) - .flatten(); - let reset_at = period_end.filter(|reset_at| *reset_at > now); - (remaining_rank.is_some() || reset_at.is_some()) - .then(|| AccountQuotaSignals::new(reset_at, remaining_rank)) -} - -/// 把 xAI billing 文档归一化为访问结论 -/// 缺少正反证据时返回 Unknown， -/// 由 `QuotaState::merge_observation` 保留既有已确认结论 -fn quota_access_observation( - billing: &GrokBillingPresentation, - observed_at: SystemTime, -) -> QuotaState { - if !billing.has_authoritative_quota() { - return QuotaState::observed_unknown(observed_at); - } - if billing_proves_access_allowed(billing, observed_at) { - QuotaState::allowed(observed_at) - } else { - QuotaState::observed_unknown(observed_at) - } -} - -fn billing_proves_access_allowed( - billing: &GrokBillingPresentation, - observed_at: SystemTime, -) -> bool { - if billing - .period_end() - .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) - .map(|value| SystemTime::from(value.to_utc())) - .is_some_and(|period_end| period_end <= observed_at) - { - return false; - } - billing - .prepaid_balance_cents() - .is_some_and(|balance| balance > 0) - || billing - .used_percent() - .is_some_and(|used| used.is_finite() && used < 100.0) - || matches!( - (billing.on_demand_used_cents(), billing.on_demand_cap_cents()), - (Some(used), Some(cap)) if cap > 0 && used < cap - ) - || matches!( - (billing.included_used_cents(), billing.monthly_limit_cents()), - (Some(used), Some(limit)) if limit > 0 && used < limit - ) -} - -fn dynamic_string( - object: Option<&serde_json::Map>, - field: &str, -) -> Option { - object? - .get(field) - .and_then(serde_json::Value::as_str) - .map(str::to_owned) -} - -fn cent_value( - object: Option<&serde_json::Map>, - field: &str, -) -> Option { - let cent = object?.get(field)?.as_object()?; - match cent.get("val") { - Some(value) => value.as_i64(), - None => Some(0), - } -} - -fn map_quota_repository_error( - error: super::repository::GrokCredentialRepositoryError, -) -> GrokQuotaError { - use super::repository::GrokCredentialRepositoryError as RepositoryError; - match error { - RepositoryError::CredentialNotFound | RepositoryError::WrongProviderKind => { - GrokQuotaError::AccountUnavailable - } - RepositoryError::StaleCredentialRevision(_) | RepositoryError::Conflict(_) => { - GrokQuotaError::StaleCredentialSnapshot - } - RepositoryError::Store(_) => GrokQuotaError::Store, - RepositoryError::InvalidInput(_) - | RepositoryError::IdentityRebind - | RepositoryError::InvalidCredentialData(_) - | RepositoryError::RevisionOverflow => GrokQuotaError::InvalidData, - } -} diff --git a/backend/crates/providers/xai/src/credential/mod.rs b/backend/crates/providers/xai/src/credential/mod.rs index f30be95c2..7c1580334 100644 --- a/backend/crates/providers/xai/src/credential/mod.rs +++ b/backend/crates/providers/xai/src/credential/mod.rs @@ -10,6 +10,7 @@ mod http; mod import; mod oidc_verifier; mod pkce; +mod quota; mod refresh; mod repository; mod secret; @@ -26,11 +27,9 @@ pub use config::{ }; pub use catalog::{ - GROK_FREE_ROLLING_WINDOW_SECONDS, GrokBillingPresentation, GrokCatalogCache, - GrokCatalogCacheError, GrokCatalogScope, GrokCredentialCatalogCache, + GrokCatalogCache, GrokCatalogCacheError, GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialCatalogError, GrokCredentialCatalogSeed, GrokCredentialCatalogService, - GrokCredentialQuotaService, GrokPlanCatalog, GrokQuotaError, GrokQuotaPeriodKind, - GrokQuotaSnapshot, + GrokPlanCatalog, }; pub use discovery::DiscoveryDocument; pub use error::{ @@ -47,6 +46,10 @@ pub use import::{ }; pub use oidc_verifier::ReqwestOidcTokenVerifier; pub use pkce::Pkce; +pub use quota::{ + GROK_FREE_ROLLING_WINDOW_SECONDS, GrokBillingPresentation, GrokCredentialQuotaService, + GrokQuotaError, GrokQuotaPeriodKind, GrokQuotaSnapshot, +}; pub use refresh::{ DueGrokCredential, GrokCredentialRecovery, GrokCredentialRecoveryOutcome, GrokCredentialRefreshError, GrokCredentialRefreshOutcome, GrokCredentialRefreshService, diff --git a/backend/crates/providers/xai/src/credential/quota.rs b/backend/crates/providers/xai/src/credential/quota.rs new file mode 100644 index 000000000..7cdd18273 --- /dev/null +++ b/backend/crates/providers/xai/src/credential/quota.rs @@ -0,0 +1,669 @@ +//! Grok billing 额度事实、展示投影与请求调度缓存 + +use std::collections::BTreeMap; +use std::fmt; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant, SystemTime}; + +use chrono::{DateTime, Utc}; +use gateway_core::account::{ + AccountQuotaSignals, CredentialRevision, ProviderAccount, ProviderAccountId, QuotaObservation, + QuotaState, +}; +use tokio::sync::Mutex; + +use super::repository::{GrokCredentialRepository, LoadedGrokCredential}; +use crate::XaiWireProfileState; +use crate::{ + GrokBillingClient, GrokBillingTransport, GrokModelCatalogSession, SecretValue, + parse_grok_billing, +}; + +const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); +const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); + +/// Grok Build Free 额度的滚动观察窗口 +pub const GROK_FREE_ROLLING_WINDOW_SECONDS: u64 = 86_400; + +/// xAI Provider 从当前 credits 文档解析出的账号页安全投影 +#[derive(Debug, Clone, PartialEq)] +pub struct GrokBillingPresentation { + plan_type: Option, + used_percent: Option, + period_type: Option, + period_start: Option, + period_end: Option, + on_demand_cap_cents: Option, + on_demand_used_cents: Option, + prepaid_balance_cents: Option, +} + +/// xAI credits 当前周期的官方语义 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GrokQuotaPeriodKind { + Weekly, + Monthly, + Other, +} + +impl GrokBillingPresentation { + #[must_use] + pub fn plan_type(&self) -> Option<&str> { + self.plan_type.as_deref() + } + + #[must_use] + pub const fn used_percent(&self) -> Option { + self.used_percent + } + + /// 返回 billing 是否包含可直接用于账号额度展示的付费事实 + #[must_use] + pub fn has_authoritative_quota(&self) -> bool { + self.used_percent.is_some() + || self.has_prepaid_credits() + || [self.on_demand_cap_cents, self.on_demand_used_cents] + .into_iter() + .flatten() + .any(|value| value > 0) + } + + fn has_prepaid_credits(&self) -> bool { + // 官方账本可用负数表示充值余额，非零即有预付额度，保留原始符号 + self.prepaid_balance_cents + .is_some_and(|balance| balance != 0) + } + + #[must_use] + pub fn period_type(&self) -> Option<&str> { + self.period_type.as_deref() + } + + #[must_use] + pub fn period_kind(&self) -> GrokQuotaPeriodKind { + match self.period_type.as_deref() { + Some("USAGE_PERIOD_TYPE_WEEKLY") => GrokQuotaPeriodKind::Weekly, + Some("USAGE_PERIOD_TYPE_MONTHLY") => GrokQuotaPeriodKind::Monthly, + _ => GrokQuotaPeriodKind::Other, + } + } + + #[must_use] + pub fn period_start(&self) -> Option<&str> { + self.period_start.as_deref() + } + + #[must_use] + pub fn period_end(&self) -> Option<&str> { + self.period_end.as_deref() + } + + #[must_use] + pub const fn on_demand_cap_cents(&self) -> Option { + self.on_demand_cap_cents + } + + #[must_use] + pub const fn on_demand_used_cents(&self) -> Option { + self.on_demand_used_cents + } + + #[must_use] + pub const fn prepaid_balance_cents(&self) -> Option { + self.prepaid_balance_cents + } +} + +/// 一个账号最近一次 xAI billing 观察结果 +#[derive(Debug, Clone, PartialEq)] +pub struct GrokQuotaSnapshot { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + observed_at: DateTime, + billing: GrokBillingPresentation, +} + +impl GrokQuotaSnapshot { + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn observed_at(&self) -> DateTime { + self.observed_at + } + + #[must_use] + pub const fn billing(&self) -> &GrokBillingPresentation { + &self.billing + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum GrokQuotaError { + #[error("xAI quota account is unavailable")] + AccountUnavailable, + #[error("xAI quota credential snapshot is stale")] + StaleCredentialSnapshot, + #[error("xAI quota credential or billing data is invalid")] + InvalidData, + #[error("xAI quota upstream request failed")] + Upstream, + #[error("xAI quota store is unavailable")] + Store, +} + +/// 官方 Grok Build billing 同步与 Provider-owned quota JSON 解析服务 +#[derive(Clone)] +pub struct GrokCredentialQuotaService { + repository: GrokCredentialRepository, + client: Arc, + scheduling: GrokQuotaSchedulingProjection, + wire_profile: XaiWireProfileState, +} + +#[derive(Clone, Default)] +struct GrokQuotaSchedulingProjection { + state: Arc>, + hydration: Arc>, +} + +#[derive(Default)] +struct GrokQuotaProjectionState { + next_version: u64, + entries: BTreeMap, +} + +#[derive(Debug, Clone, Copy)] +struct GrokQuotaSchedulingEntry { + version: u64, + revision: CredentialRevision, + expires_at: Instant, + signals: Option, +} + +#[derive(Clone)] +struct GrokQuotaHydrationTarget { + account: ProviderAccount, + expected_version: Option, +} + +impl GrokQuotaSchedulingProjection { + fn invalidate(&self, account_ids: &[ProviderAccountId]) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for account_id in account_ids { + state.entries.remove(account_id); + } + } + + fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let now = Instant::now(); + accounts + .iter() + .filter_map(|account| { + let current = state.entries.get(account.id()); + current + .is_none_or(|entry| { + entry.revision != account.revision() || now >= entry.expires_at + }) + .then(|| GrokQuotaHydrationTarget { + account: account.clone(), + expected_version: current.map(|entry| entry.version), + }) + }) + .collect() + } + + fn observe(&self, snapshot: &GrokQuotaSnapshot) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { + return false; + }; + self.replace( + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + quota_scheduling_signals(snapshot.billing()), + ); + true + } + + fn observe_if_unchanged( + &self, + target: &GrokQuotaHydrationTarget, + snapshot: &GrokQuotaSnapshot, + ) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { + return false; + }; + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return true; + } + insert_grok_projection_entry( + &mut state, + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + quota_scheduling_signals(snapshot.billing()), + ); + true + } + + fn mark_unknown_if_unchanged(&self, target: &GrokQuotaHydrationTarget, ttl: Duration) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return; + } + insert_grok_projection_entry( + &mut state, + target.account.id().clone(), + target.account.revision(), + ttl, + None, + ); + } + + fn replace( + &self, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, + ) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + insert_grok_projection_entry(&mut state, account_id, revision, ttl, signals); + } + + fn signals(&self, account: &ProviderAccount) -> Option { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .entries + .get(account.id()) + .filter(|entry| { + entry.revision == account.revision() && Instant::now() < entry.expires_at + }) + .and_then(|entry| entry.signals) + } +} + +fn insert_grok_projection_entry( + state: &mut GrokQuotaProjectionState, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, +) { + state.next_version = state.next_version.saturating_add(1); + state.entries.insert( + account_id, + GrokQuotaSchedulingEntry { + version: state.next_version, + revision, + expires_at: Instant::now() + ttl, + signals, + }, + ); +} + +impl GrokCredentialQuotaService { + #[must_use] + pub fn new( + repository: GrokCredentialRepository, + transport: Arc, + wire_profile: XaiWireProfileState, + ) -> Self { + Self { + repository, + client: Arc::new(GrokBillingClient::new(transport)), + scheduling: GrokQuotaSchedulingProjection::default(), + wire_profile, + } + } + + /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度 + pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { + if self.scheduling.hydration_targets(accounts).is_empty() { + return; + } + let _hydration = self.scheduling.hydration.lock().await; + let pending = self.scheduling.hydration_targets(accounts); + if pending.is_empty() { + return; + } + let pending_accounts = pending + .iter() + .map(|target| target.account.clone()) + .collect::>(); + let Ok(observations) = self.repository.quota_observations(&pending_accounts).await else { + for target in &pending { + self.scheduling + .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); + } + return; + }; + let observations = observations + .into_iter() + .map(|observation| (observation.account_id.clone(), observation)) + .collect::>(); + for target in pending { + let snapshot = observations + .get(target.account.id()) + .filter(|observation| observation.expected_revision == target.account.revision()) + .and_then(quota_snapshot_from_observation); + if !snapshot + .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) + { + self.scheduling + .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); + } + } + } + + #[must_use] + pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { + self.scheduling.signals(account) + } + + pub(crate) fn invalidate_scheduling(&self, account_ids: &[ProviderAccountId]) { + self.scheduling.invalidate(account_ids); + } + + /// 立即刷新一个账号的动态 billing document，并以 credential revision CAS 写回 + pub async fn refresh_account( + &self, + account_id: &ProviderAccountId, + ) -> Result { + let loaded = self + .repository + .load_current(account_id) + .await + .map_err(map_quota_repository_error)?; + if !loaded.account.enabled() + || loaded + .account + .access_token_expires_at() + .is_none_or(|expires_at| expires_at <= SystemTime::now()) + { + return Err(GrokQuotaError::AccountUnavailable); + } + let session = billing_session(&loaded, &self.wire_profile)?; + let billing = self + .client + .fetch(&session) + .await + .map_err(|_| GrokQuotaError::Upstream)?; + let mut document = billing.into_document(); + // 订阅查询失败不影响已取得的额度，也不能据此写入免费套餐 + match self.client.fetch_subscription(&session).await { + Ok(Some(plan_type)) => { + document.insert( + "subscriptionTier".to_owned(), + serde_json::Value::String(plan_type), + ); + } + Ok(None) => {} + Err(error) => tracing::warn!(%error, "xAI subscription query failed"), + } + let observed_at = Utc::now(); + let presentation = billing_presentation(&document)?; + let quota = loaded + .account + .quota() + .merge_observation(quota_access_observation(&presentation, observed_at.into())); + self.repository + .replace_quota( + loaded.account.id().clone(), + loaded.account.revision(), + document, + observed_at.into(), + quota, + ) + .await + .map_err(map_quota_repository_error)?; + let snapshot = GrokQuotaSnapshot { + account_id: loaded.account.id().clone(), + credential_revision: loaded.account.revision(), + observed_at, + billing: presentation, + }; + self.scheduling.observe(&snapshot); + Ok(snapshot) + } + + /// 读取并重新验证 Store 中的 Provider-owned quota JSON + pub async fn read_account( + &self, + account_id: &ProviderAccountId, + ) -> Result, GrokQuotaError> { + let Some(observation) = self + .repository + .quota(account_id) + .await + .map_err(map_quota_repository_error)? + else { + return Ok(None); + }; + let observed_at = observation.observed_at; + let document = observation.quota.into_inner(); + let snapshot = GrokQuotaSnapshot { + account_id: observation.account_id, + credential_revision: observation.expected_revision, + observed_at: observed_at.into(), + billing: billing_presentation(&document)?, + }; + self.scheduling.observe(&snapshot); + Ok(Some(snapshot)) + } +} + +impl fmt::Debug for GrokCredentialQuotaService { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("GrokCredentialQuotaService") + .field("repository", &self.repository) + .field("client", &self.client) + .finish() + } +} + +fn billing_session( + loaded: &LoadedGrokCredential, + wire_profile: &XaiWireProfileState, +) -> Result { + let upstream_user_id = loaded + .account + .upstream_user_id() + .ok_or(GrokQuotaError::InvalidData)?; + GrokModelCatalogSession::new( + loaded.access_token.clone(), + SecretValue::new(upstream_user_id), + loaded + .account + .email() + .map(|value| SecretValue::new(value.to_owned())), + wire_profile.clone(), + ) + .map(|session| session.with_outbound_proxy(loaded.account.outbound_proxy().cloned())) + .map_err(|_| GrokQuotaError::InvalidData) +} + +fn billing_presentation( + document: &serde_json::Map, +) -> Result { + let body = serde_json::to_vec(document).map_err(|_| GrokQuotaError::InvalidData)?; + let snapshot = parse_grok_billing(&body).map_err(|_| GrokQuotaError::InvalidData)?; + let config = snapshot + .document() + .get("config") + .and_then(|value| value.as_object()); + let current_period = config + .and_then(|config| config.get("currentPeriod")) + .and_then(|value| value.as_object()); + let used_percent = config + .and_then(|config| config.get("creditUsagePercent")) + .and_then(serde_json::Value::as_f64); + Ok(GrokBillingPresentation { + plan_type: dynamic_string(Some(snapshot.document()), "subscriptionTier"), + used_percent, + period_type: dynamic_string(current_period, "type"), + period_start: dynamic_string(current_period, "start"), + period_end: dynamic_string(current_period, "end"), + on_demand_cap_cents: cent_value(config, "onDemandCap"), + on_demand_used_cents: cent_value(config, "onDemandUsed"), + prepaid_balance_cents: cent_value(config, "prepaidBalance"), + }) +} + +fn quota_snapshot_from_observation(observation: &QuotaObservation) -> Option { + let observed_at = DateTime::::from(observation.observed_at); + let document = observation.quota.expose_to_provider(); + Some(GrokQuotaSnapshot { + account_id: observation.account_id.clone(), + credential_revision: observation.expected_revision, + observed_at, + billing: billing_presentation(document).ok()?, + }) +} + +fn quota_projection_ttl(observed_at: DateTime) -> Option { + let age = SystemTime::now() + .duration_since(SystemTime::from(observed_at)) + .unwrap_or(Duration::ZERO); + QUOTA_SCHEDULING_TTL + .checked_sub(age) + .filter(|remaining| !remaining.is_zero()) +} + +fn quota_scheduling_signals(billing: &GrokBillingPresentation) -> Option { + let now = SystemTime::now(); + let period_end = billing + .period_end() + .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) + .map(|value| SystemTime::from(value.to_utc())); + let observation_expired = period_end.is_some_and(|period_end| period_end <= now); + let remaining_rank = (!observation_expired) + .then(|| { + billing + .used_percent() + .filter(|used| used.is_finite() && (0.0..=100.0).contains(used)) + .map(|used| ((100.0 - used) * 100.0).round() as u64) + }) + .flatten(); + let reset_at = period_end.filter(|reset_at| *reset_at > now); + (remaining_rank.is_some() || reset_at.is_some()) + .then(|| AccountQuotaSignals::new(reset_at, remaining_rank)) +} + +/// 把 xAI billing 文档归一化为访问结论 +/// 缺少正反证据时返回 Unknown， +/// 由 `QuotaState::merge_observation` 保留既有已确认结论 +fn quota_access_observation( + billing: &GrokBillingPresentation, + observed_at: SystemTime, +) -> QuotaState { + if !billing.has_authoritative_quota() { + return QuotaState::observed_unknown(observed_at); + } + if billing_proves_access_allowed(billing, observed_at) { + QuotaState::allowed(observed_at) + } else { + QuotaState::observed_unknown(observed_at) + } +} + +fn billing_proves_access_allowed( + billing: &GrokBillingPresentation, + observed_at: SystemTime, +) -> bool { + if billing + .period_end() + .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) + .map(|value| SystemTime::from(value.to_utc())) + .is_some_and(|period_end| period_end <= observed_at) + { + return false; + } + billing.has_prepaid_credits() + || billing + .used_percent() + .is_some_and(|used| used.is_finite() && used < 100.0) + || matches!( + (billing.on_demand_used_cents(), billing.on_demand_cap_cents()), + (Some(used), Some(cap)) if cap > 0 && used < cap + ) +} + +fn dynamic_string( + object: Option<&serde_json::Map>, + field: &str, +) -> Option { + object? + .get(field) + .and_then(serde_json::Value::as_str) + .map(str::to_owned) +} + +fn cent_value( + object: Option<&serde_json::Map>, + field: &str, +) -> Option { + let cent = object?.get(field)?.as_object()?; + match cent.get("val") { + Some(value) => value.as_i64(), + None => Some(0), + } +} + +fn map_quota_repository_error( + error: super::repository::GrokCredentialRepositoryError, +) -> GrokQuotaError { + use super::repository::GrokCredentialRepositoryError as RepositoryError; + match error { + RepositoryError::CredentialNotFound | RepositoryError::WrongProviderKind => { + GrokQuotaError::AccountUnavailable + } + RepositoryError::StaleCredentialRevision(_) | RepositoryError::Conflict(_) => { + GrokQuotaError::StaleCredentialSnapshot + } + RepositoryError::Store(_) => GrokQuotaError::Store, + RepositoryError::InvalidInput(_) + | RepositoryError::IdentityRebind + | RepositoryError::InvalidCredentialData(_) + | RepositoryError::RevisionOverflow => GrokQuotaError::InvalidData, + } +} diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index f089575a0..e99c156e6 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -19,7 +19,8 @@ use gateway_core::provider_ports::{ }; use gateway_core::routing::{ProviderKind, UpstreamModelId}; -use super::catalog::{GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialQuotaService}; +use super::catalog::{GrokCatalogScope, GrokCredentialCatalogCache}; +use super::quota::GrokCredentialQuotaService; use super::repository::{GrokCredentialRepository, GrokCredentialRepositoryError}; use super::types::UpdateGrokCredentialState; use crate::{ diff --git a/backend/crates/providers/xai/src/lib.rs b/backend/crates/providers/xai/src/lib.rs index 2ced85de2..96afa76fe 100644 --- a/backend/crates/providers/xai/src/lib.rs +++ b/backend/crates/providers/xai/src/lib.rs @@ -1,7 +1,7 @@ //! 官方 Grok Build OAuth Provider 边界 mod admin; -pub mod credential; +mod credential; mod provider; mod reasoning_replay; pub mod transport; diff --git a/backend/crates/providers/xai/src/provider/continuation.rs b/backend/crates/providers/xai/src/provider/continuation.rs index c9bb83b27..fe8d926fa 100644 --- a/backend/crates/providers/xai/src/provider/continuation.rs +++ b/backend/crates/providers/xai/src/provider/continuation.rs @@ -1,6 +1,7 @@ //! xAI continuation 与 reasoning replay 状态处理 use super::*; +use crate::transport::GrokReplayItem; #[derive(Clone, Serialize, Deserialize)] pub(super) struct XaiSessionState { @@ -147,7 +148,12 @@ fn apply_replay( ) -> Result<(), ProviderError> { let mut input = replay_input_for_account(previous, account.as_str(), true); input.reserve(current_input.len()); - input.extend(current_input.iter().cloned()); + input.extend( + current_input + .iter() + .cloned() + .map(GrokReplayItem::UpstreamInput), + ); request.set_replay_input(input).map_err(map_request_error)?; request.set_previous_response_id(None); request.inherit_session(None); @@ -158,21 +164,24 @@ pub(super) fn replay_input_for_account( state: &XaiSessionState, account_id: &str, force_portable: bool, -) -> Vec { +) -> Vec { state .transcript .iter() .filter_map(|item| match item { - XaiReplayItem::ClientInput(value) | XaiReplayItem::SanitizedOutput(value) => { - Some(value.clone()) + XaiReplayItem::ClientInput(value) => Some(GrokReplayItem::UpstreamInput(value.clone())), + XaiReplayItem::SanitizedOutput(value) => { + Some(GrokReplayItem::ClientOutput(value.clone())) } XaiReplayItem::AccountOutput { account_id: owner, item, } if owner == account_id && !force_portable => { - portable_output_item(item.clone(), false) + portable_output_item(item.clone(), false).map(GrokReplayItem::ClientOutput) + } + XaiReplayItem::AccountOutput { item, .. } => { + portable_output_item(item.clone(), true).map(GrokReplayItem::ClientOutput) } - XaiReplayItem::AccountOutput { item, .. } => portable_output_item(item.clone(), true), }) .collect() } diff --git a/backend/crates/providers/xai/src/provider/failure.rs b/backend/crates/providers/xai/src/provider/failure.rs index 8a038c0c0..3e6b493b1 100644 --- a/backend/crates/providers/xai/src/provider/failure.rs +++ b/backend/crates/providers/xai/src/provider/failure.rs @@ -383,15 +383,23 @@ pub(super) fn map_request_error(error: GrokRequestEncodeError) -> ProviderError GrokRequestEncodeError::InvalidProtocolPayload | GrokRequestEncodeError::InvalidRequestNormalization | GrokRequestEncodeError::InvalidRequestField { .. } => ProviderErrorKind::InvalidRequest, + GrokRequestEncodeError::UnsupportedPrewarm => ProviderErrorKind::Unsupported, GrokRequestEncodeError::Serialization => ProviderErrorKind::Protocol, }; let provider_error = provider_error(kind, UpstreamSendState::NotSent); - if kind != ProviderErrorKind::InvalidRequest { + if !matches!( + kind, + ProviderErrorKind::InvalidRequest | ProviderErrorKind::Unsupported + ) { return provider_error; } + let code = match error { + GrokRequestEncodeError::UnsupportedPrewarm => "unsupported_prewarm", + _ => "invalid_request_normalization", + }; provider_error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new( error.to_string(), - Some("invalid_request_normalization".to_owned()), + Some(code.to_owned()), Some("invalid_request_error".to_owned()), )) } diff --git a/backend/crates/providers/xai/src/provider/mod.rs b/backend/crates/providers/xai/src/provider/mod.rs index 1c7053afe..f7b246741 100644 --- a/backend/crates/providers/xai/src/provider/mod.rs +++ b/backend/crates/providers/xai/src/provider/mod.rs @@ -81,7 +81,6 @@ use stream::*; pub(crate) use workers::worker_contributions; const HTTP_SSE_TRANSPORT: &str = "http_sse"; -const DEFAULT_GROK_MODEL: &str = "grok-4.5"; const XAI_SESSION_STATE_MAX_BYTES: usize = 8 * 1024 * 1024; const XAI_SESSION_OUTPUT_LIMIT: usize = 4_096; const REASONING_DECODE_FAILED_CODE: &str = "reasoning_decode_failed"; @@ -205,12 +204,9 @@ impl Provider for GrokBuildProvider { async fn query_model_capabilities( &self, ) -> Result, ProviderError> { - let Ok(models) = self.catalog.query_models().await else { - return Ok(vec![default_grok_model_capabilities()?]); - }; - if models.is_empty() { - return Ok(vec![default_grok_model_capabilities()?]); - } + let models = self.catalog.query_models().await.map_err(|_| { + provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent) + })?; Ok(models .into_iter() .map(compile_grok_model_capabilities) @@ -842,97 +838,41 @@ fn compile_grok_model_capabilities(model: GrokCatalogModel) -> ProviderModelCapa .with_presentation(grok_model_presentation(&model)) } -fn default_grok_model_capabilities() -> Result { - let model = UpstreamModelId::new(DEFAULT_GROK_MODEL.to_owned()) - .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; - let capabilities = ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None) - .with_upstream_feature_validation(); - Ok(ProviderModelCapabilities::new(model, capabilities) - .with_presentation(default_grok_model_presentation())) -} - -fn default_grok_model_presentation() -> ModelPresentation { - ModelPresentation::new( - Some("Grok 4.5".to_owned()), - Some("xAI Grok 4.5 frontier model with reasoning and vision.".to_owned()), - ) - .with_context_window_tokens(Some(500_000)) - .with_max_context_window_tokens(Some(500_000)) - .with_image_input(true) - .with_agent_tools(true, true) -} - fn grok_model_presentation(model: &GrokCatalogModel) -> ModelPresentation { - let slug = model.request_model().as_str(); - let known_grok_4_5 = is_known_grok_4_5_model(slug); - let reasoning_evidence = model.capabilities().reasoning_effort(); - let catalog_reasoning_efforts = model + let reasoning_efforts = model .capabilities() .reasoning_efforts() .iter() .map(|effort| effort.as_str().to_owned()) .collect::>(); - let catalog_default_reasoning = model + let default_reasoning = model .capabilities() .default_reasoning_effort() .map(|effort| effort.as_str().to_owned()); - let (default_reasoning, reasoning_efforts) = match reasoning_evidence { - GrokCatalogCapabilityEvidence::DeclaredNative if !catalog_reasoning_efforts.is_empty() => { - let default_reasoning = catalog_default_reasoning - .filter(|default| { - catalog_reasoning_efforts - .iter() - .any(|effort| effort == default) - }) - .or_else(|| catalog_reasoning_efforts.first().cloned()); - (default_reasoning, catalog_reasoning_efforts) - } - GrokCatalogCapabilityEvidence::DeclaredNative => (None, Vec::new()), - GrokCatalogCapabilityEvidence::DeclaredUnsupported => { - (Some("none".to_owned()), vec!["none".to_owned()]) - } - GrokCatalogCapabilityEvidence::Unknown => (None, Vec::new()), - }; let context_window_tokens = model .limits() .context_window_tokens() - .map(std::num::NonZeroU64::get) - .or(known_grok_4_5.then_some(500_000)); + .map(std::num::NonZeroU64::get); let tool_evidence = model.capabilities().streaming_tool_calls(); ModelPresentation::new( - model - .display_name() - .map(str::to_owned) - .or_else(|| known_grok_4_5.then(|| "Grok 4.5".to_owned())), - model - .metadata() - .description() - .map(str::to_owned) - .or_else(|| { - known_grok_4_5 - .then(|| "xAI Grok 4.5 frontier model with reasoning and vision.".to_owned()) - }), + model.display_name().map(str::to_owned), + model.metadata().description().map(str::to_owned), ) .with_reasoning(default_reasoning, reasoning_efforts) .with_context_window_tokens(context_window_tokens) - // Grok 目录只声明一个窗口，继续将它作为客户端可覆盖上限 - .with_max_context_window_tokens(context_window_tokens) - .with_image_input(known_grok_4_5) + .with_max_context_window_tokens( + model + .limits() + .max_context_window_tokens() + .map(std::num::NonZeroU64::get), + ) .with_agent_tools( tool_evidence != GrokCatalogCapabilityEvidence::DeclaredUnsupported, - tool_evidence == GrokCatalogCapabilityEvidence::DeclaredNative - || (known_grok_4_5 && tool_evidence == GrokCatalogCapabilityEvidence::Unknown), + tool_evidence == GrokCatalogCapabilityEvidence::DeclaredNative, ) .with_search_tool( model.capabilities().backend_search() == GrokCatalogCapabilityEvidence::DeclaredNative, ) .with_hidden(model.metadata().hidden().unwrap_or(false)) } - -fn is_known_grok_4_5_model(slug: &str) -> bool { - matches!( - slug, - DEFAULT_GROK_MODEL | "grok-4.5-latest" | "grok-4.5-build-free" | "grok-build-latest" - ) -} diff --git a/backend/crates/providers/xai/src/reasoning_replay.rs b/backend/crates/providers/xai/src/reasoning_replay.rs index 289bd4a7c..3a99b270f 100644 --- a/backend/crates/providers/xai/src/reasoning_replay.rs +++ b/backend/crates/providers/xai/src/reasoning_replay.rs @@ -8,6 +8,8 @@ use base64::Engine as _; use gateway_core::event::ProviderEvent; use serde_json::{Map, Value}; +use crate::transport::GrokReplayItem; + const REASONING_REPLAY_TTL: Duration = Duration::from_secs(60 * 60); const REASONING_REPLAY_MAX_ENTRIES: usize = 1_024; const REASONING_REPLAY_MAX_BYTES: usize = 8 * 1024 * 1024; @@ -82,7 +84,7 @@ impl GrokReasoningReplay { &self, key: &GrokReasoningReplayKey, input: &[Value], - ) -> Option> { + ) -> Option> { if input.is_empty() { return None; } @@ -560,17 +562,17 @@ fn replay_call_keys(item_type: &str, call_id: &str) -> Vec { .collect() } -fn insert_replay_items(input: &[Value], replay: &[Value]) -> Vec { +fn insert_replay_items(input: &[Value], replay: &[Value]) -> Vec { let insert_at = replay_insert_index(input, replay); let mut result = Vec::new(); for (index, item) in input.iter().enumerate() { if index == insert_at { - result.extend(replay.iter().cloned()); + result.extend(replay.iter().cloned().map(GrokReplayItem::ClientOutput)); } - result.push(item.clone()); + result.push(GrokReplayItem::UpstreamInput(item.clone())); } if insert_at == input.len() { - result.extend(replay.iter().cloned()); + result.extend(replay.iter().cloned().map(GrokReplayItem::ClientOutput)); } result } diff --git a/backend/crates/providers/xai/src/transport/canonical.rs b/backend/crates/providers/xai/src/transport/canonical.rs index 2fa90dc26..a4b7856a4 100644 --- a/backend/crates/providers/xai/src/transport/canonical.rs +++ b/backend/crates/providers/xai/src/transport/canonical.rs @@ -1041,7 +1041,10 @@ fn provider_reported_cost(response: &Value) -> Result FinishReason { .pointer("/incomplete_details/reason") .and_then(Value::as_str) { - Some("max_output_tokens" | "max_tokens") => FinishReason::Length, + Some("max_output_tokens" | "max_tokens" | "max_prompt_tokens" | "max_time_limit") => { + FinishReason::Length + } Some("content_filter") => FinishReason::ContentFilter, _ => FinishReason::Other, } diff --git a/backend/crates/providers/xai/src/transport/catalog.rs b/backend/crates/providers/xai/src/transport/catalog.rs index e457de139..8e270aa0c 100644 --- a/backend/crates/providers/xai/src/transport/catalog.rs +++ b/backend/crates/providers/xai/src/transport/catalog.rs @@ -570,17 +570,12 @@ fn validate_billing_config(config: &Map) -> Result<(), GrokBillin return Err(GrokBillingError::InvalidWire); } } - for field in [ - "monthlyLimit", - "used", - "onDemandCap", - "onDemandUsed", - "prepaidBalance", - ] { + for field in ["onDemandCap", "onDemandUsed", "prepaidBalance"] { if let Some(value) = config.get(field) && !value.is_null() { - validate_cent(value)?; + // 预付余额使用带符号的账本金额，限额和已使用量仍须非负 + validate_cent(value, field == "prepaidBalance")?; } } if let Some(period) = config.get("currentPeriod") @@ -610,12 +605,14 @@ fn validate_billing_config(config: &Map) -> Result<(), GrokBillin Ok(()) } -fn validate_cent(value: &Value) -> Result<(), GrokBillingError> { +fn validate_cent(value: &Value, allow_negative: bool) -> Result<(), GrokBillingError> { let Value::Object(cent) = value else { return Err(GrokBillingError::InvalidWire); }; if let Some(value) = cent.get("val") - && value.as_i64().is_none_or(|value| value < 0) + && value + .as_i64() + .is_none_or(|value| !allow_negative && value < 0) { return Err(GrokBillingError::InvalidWire); } @@ -761,6 +758,7 @@ impl GrokCatalogCapabilities { #[derive(Debug, Clone, PartialEq, Eq)] pub struct GrokCatalogLimits { context_window_tokens: Option, + max_context_window_tokens: Option, max_output_tokens: Option, } @@ -771,6 +769,12 @@ impl GrokCatalogLimits { self.context_window_tokens } + /// 返回目录默认窗口与可选窗口中的最大值；缺失表示未知 + #[must_use] + pub const fn max_context_window_tokens(&self) -> Option { + self.max_context_window_tokens + } + /// 返回明确最大输出 token；缺失表示未知 #[must_use] pub const fn max_output_tokens(&self) -> Option { @@ -958,69 +962,27 @@ enum GrokModelsObject { #[derive(Debug, Deserialize)] struct GrokModelWire { id: Option, - model: Option, - #[serde(rename = "modelId")] - model_id: Option, + model: String, name: Option, description: Option, - #[serde(rename = "contextWindow", alias = "context_window")] context_window: Option, - #[serde(rename = "maxCompletionTokens", alias = "max_completion_tokens")] + #[serde(default)] + context_windows: Vec, max_completion_tokens: Option, - #[serde(rename = "apiBackend", alias = "api_backend")] api_backend: Option, - #[serde(rename = "supportedInApi", alias = "supported_in_api")] supported_in_api: Option, - #[serde( - rename = "supportsReasoningEffort", - alias = "supports_reasoning_effort" - )] - supports_reasoning_effort: Option, - #[serde(rename = "reasoningEffort", alias = "reasoning_effort")] - reasoning_effort: Option, - #[serde(default, rename = "reasoningEfforts", alias = "reasoning_efforts")] + #[serde(default)] reasoning_efforts: Vec, - features: Option, - #[serde(rename = "supportsBackendSearch", alias = "supports_backend_search")] supports_backend_search: Option, - #[serde(rename = "streamToolCalls", alias = "stream_tool_calls")] stream_tool_calls: Option, hidden: Option, } #[derive(Debug, Deserialize)] -struct GrokModelFeaturesWire { - reasoning: Option, - #[serde(rename = "reasoningEffortOptions", alias = "reasoning_effort_options")] - reasoning_effort_options: Option, -} - -#[derive(Debug, Deserialize)] -struct GrokReasoningEffortOptionsWire { - #[serde(default, rename = "supportedEfforts", alias = "supported_efforts")] - supported_efforts: Vec, - #[serde(rename = "defaultEffort", alias = "default_effort")] - default_effort: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(untagged)] -enum GrokReasoningEffortOptionWire { - Bare(GrokCatalogReasoningEffort), - Detailed { - value: GrokCatalogReasoningEffort, - #[serde(default)] - default: bool, - }, -} - -impl GrokReasoningEffortOptionWire { - fn into_parts(self) -> (GrokCatalogReasoningEffort, bool) { - match self { - Self::Bare(value) => (value, false), - Self::Detailed { value, default } => (value, default), - } - } +struct GrokReasoningEffortOptionWire { + value: GrokCatalogReasoningEffort, + #[serde(default)] + default: bool, } /// 解析唯一受支持的 Grok 官方完整快照 @@ -1062,8 +1024,8 @@ pub fn parse_grok_model_catalog( } fn normalize_model(wire: GrokModelWire) -> Result { - let actual_model = wire.model.or(wire.model_id).or_else(|| wire.id.clone()); - let actual_model = actual_model.ok_or(GrokModelCatalogError::InvalidModelSlug)?; + let (context_window_tokens, max_context_window_tokens) = context_window_limits(&wire)?; + let actual_model = wire.model; if !valid_model_slug(&actual_model) { return Err(GrokModelCatalogError::InvalidModelSlug); } @@ -1082,55 +1044,28 @@ fn normalize_model(wire: GrokModelWire) -> Result>(), - options - .default_effort - .filter(|effort| *effort != GrokCatalogReasoningEffort::Unknown), - ) - }, - ); - let reasoning_efforts = if feature_reasoning_efforts.is_empty() { - top_level_reasoning_efforts - } else { - feature_reasoning_efforts - }; - let default_reasoning_effort = feature_default_reasoning_effort - .or_else(|| { - wire.reasoning_effort - .filter(|effort| *effort != GrokCatalogReasoningEffort::Unknown) + // 固定 CLI proxy 的对象菜单拥有档位与默认值，缺失声明时不替用户选择首项 + let mut default_reasoning_effort = None; + let reasoning_efforts = wire + .reasoning_efforts + .into_iter() + .filter_map(|option| { + if option.value == GrokCatalogReasoningEffort::Unknown { + return None; + } + if option.default && default_reasoning_effort.is_none() { + default_reasoning_effort = Some(option.value); + } + Some(option.value) }) - .or(top_level_menu_default) - .or_else(|| reasoning_efforts.first().copied()); - let default_reasoning_effort = default_reasoning_effort - .filter(|default| reasoning_efforts.iter().any(|effort| effort == default)) - .or_else(|| reasoning_efforts.first().copied()); + .collect::>(); let reasoning_effort = if reasoning_efforts.is_empty() { - GrokCatalogCapabilityEvidence::from_wire( - wire.supports_reasoning_effort.or(feature_reasoning), - ) + GrokCatalogCapabilityEvidence::Unknown } else { GrokCatalogCapabilityEvidence::DeclaredNative }; @@ -1149,6 +1084,7 @@ fn normalize_model(wire: GrokModelWire) -> Result Result, - menu_default: &mut Option, -) -> Vec { - values +fn context_window_limits( + wire: &GrokModelWire, +) -> Result<(Option, Option), GrokModelCatalogError> { + let default = + optional_positive(wire.context_window)?.or_else(|| wire.context_windows.first().copied()); + let maximum = default .into_iter() - .filter_map(|option| { - let (effort, is_default) = option.into_parts(); - if effort == GrokCatalogReasoningEffort::Unknown { - return None; - } - if is_default && menu_default.is_none() { - *menu_default = Some(effort); - } - Some(effort) - }) - .collect() + .chain(wire.context_windows.iter().copied()) + .max(); + Ok((default, maximum)) } fn responses_evidence( diff --git a/backend/crates/providers/xai/src/transport/mod.rs b/backend/crates/providers/xai/src/transport/mod.rs index d7fe899a2..6c7549b50 100644 --- a/backend/crates/providers/xai/src/transport/mod.rs +++ b/backend/crates/providers/xai/src/transport/mod.rs @@ -49,6 +49,7 @@ pub use network::{ ReqwestGrokModelCatalogTransport, ReqwestOAuthTransport, }; pub use profile::{XaiWireProfile, XaiWireProfileState}; +pub(crate) use request::GrokReplayItem; pub use request::{GrokRequestEncodeError, GrokResponsesRequest}; pub use session::{ GrokCredentialFailure, GrokCredentialFeedbackFuture, GrokSessionAffinityKey, diff --git a/backend/crates/providers/xai/src/transport/request/mod.rs b/backend/crates/providers/xai/src/transport/request/mod.rs index bcfa08aa2..c96434c99 100644 --- a/backend/crates/providers/xai/src/transport/request/mod.rs +++ b/backend/crates/providers/xai/src/transport/request/mod.rs @@ -12,7 +12,7 @@ use super::{GrokSessionAffinityKey, XAI_PROVIDER_NAME}; mod history; mod identity; -mod model; +mod parameters; mod response; mod schema; mod tools; @@ -22,7 +22,7 @@ pub(crate) use response::GrokResponseTransform; use history::*; use identity::*; -use model::*; +use parameters::*; use schema::*; use tools::*; @@ -40,6 +40,14 @@ pub struct GrokResponsesRequest { response_transform: GrokResponseTransform, } +/// 重放条目在进入 wire 编码前的来源，避免将已编码别名再次当作客户端名称解释 +pub(crate) enum GrokReplayItem { + /// 已由请求转换器编码的输入，保留上游别名 + UpstreamInput(Value), + /// 已向客户端回译的输出，需要按本轮工具声明重新编码 + ClientOutput(Value), +} + impl GrokResponsesRequest { /// 返回发送到 `/v1/responses` 的 JSON object #[must_use] @@ -99,18 +107,29 @@ impl GrokResponsesRequest { pub(crate) fn set_replay_input( &mut self, - input: Vec, + input: Vec, ) -> Result<(), GrokRequestEncodeError> { let mut normalizer = ToolNormalizer::for_replay(self.response_transform.clone()); let mut normalized = Vec::with_capacity(input.len()); for item in input { match item { - Value::Object(item) if string_field(&item, "type") == "custom_tool_call" => { + GrokReplayItem::ClientOutput(Value::Object(item)) + if string_field(&item, "type") == "custom_tool_call" => + { normalized.push(Value::Object( normalizer.normalize_custom_tool_call_input(&item)?, )); } - item => normalized.push(item), + GrokReplayItem::ClientOutput(Value::Object(item)) + if string_field(&item, "type") == "function_call" => + { + normalized.push(Value::Object( + normalizer.normalize_function_call_input(&item)?, + )); + } + GrokReplayItem::UpstreamInput(item) | GrokReplayItem::ClientOutput(item) => { + normalized.push(item); + } } } self.response_transform = normalizer.response; @@ -201,11 +220,16 @@ impl GrokResponsesRequest { if payload.protocol() != "openai" { return Err(GrokRequestEncodeError::InvalidProtocolPayload); } + // Grok Build 会把 Codex 的预热请求当作真实生成；必须在发送前拒绝， + // 不能删掉控制字段后继续调用，或在已经计费后隐藏响应与费用 + if payload.body().get("generate").and_then(Value::as_bool) == Some(false) { + return Err(GrokRequestEncodeError::UnsupportedPrewarm); + } let mut body = payload.body().clone(); if should_consume_terminal_compaction_trigger { consume_terminal_compaction_trigger(&mut body)?; } - let upstream_model = resolve_grok_text_responses_model_id(upstream_model); + let upstream_model = upstream_model.to_owned(); // 这些字段属于 Codex/OpenAI 侧请求控制，不是 xAI 上游协议字段 // OpenAI 透明路径会保留未知字段；这里只在 xAI adapter 内做最小剥离 body.remove("provider_options"); @@ -220,7 +244,7 @@ impl GrokResponsesRequest { ); sanitize_account_identity(&mut body); sanitize_client_metadata(&mut body); - normalize_build_request(&mut body, &upstream_model)?; + normalize_build_request(&mut body)?; let mut response_transform = normalize_responses_request(&mut body)?; response_transform.observe_client_cache_tools(); if enable_cache_route { @@ -258,33 +282,6 @@ impl GrokResponsesRequest { } } -fn resolve_grok_text_responses_model_id(model: &str) -> String { - let stripped = strip_grok_provider_prefix(model); - let normalized = stripped.to_ascii_lowercase(); - match normalized.as_str() { - "grok" | "grok-latest" | "grok-4.5" | "grok-4.5-latest" | "grok-build-latest" => { - "grok-4.5".to_owned() - } - "grok-4.6" | "grok-4.6-latest" => "grok-4.6".to_owned(), - "grok-4.3" | "grok-4.3-latest" => "grok-4.3".to_owned(), - "grok-3-mini" - | "grok-3-mini-fast" - | "grok-build-0.1" - | "grok-composer-2.5-fast" - | "grok-4.20-0309-reasoning" - | "grok-4.20-0309-non-reasoning" - | "grok-4.20-multi-agent-0309" => normalized, - "grok-build" => "grok-build-0.1".to_owned(), - "grok-composer" | "composer-2.5" => "grok-composer-2.5-fast".to_owned(), - "grok-4.20-reasoning" => "grok-4.20-0309-reasoning".to_owned(), - "grok-4.20-non-reasoning" => "grok-4.20-0309-non-reasoning".to_owned(), - "grok-4.20-multi-agent" | "grok-4.20-multi-agent-latest" => { - "grok-4.20-multi-agent-0309".to_owned() - } - _ => stripped.to_owned(), - } -} - fn consume_terminal_compaction_trigger( body: &mut Map, ) -> Result<(), GrokRequestEncodeError> { @@ -322,6 +319,9 @@ pub enum GrokRequestEncodeError { /// 数据面只接受 OpenAI adapter 保留的原始 Responses object #[error("Grok Build request is missing its OpenAI protocol payload")] InvalidProtocolPayload, + /// Grok Build 不支持 Codex 的非生成预热语义 + #[error("Grok Build does not support non-generating prewarm requests (generate=false)")] + UnsupportedPrewarm, /// JSON 序列化意外失败 #[error("Grok Build request serialization failed")] Serialization, diff --git a/backend/crates/providers/xai/src/transport/request/model.rs b/backend/crates/providers/xai/src/transport/request/model.rs deleted file mode 100644 index 28f805f02..000000000 --- a/backend/crates/providers/xai/src/transport/request/model.rs +++ /dev/null @@ -1,144 +0,0 @@ -//! 当前 Responses 参数到 Grok 模型能力的转换 - -use super::*; - -pub(super) fn normalize_build_request( - body: &mut Map, - upstream_model: &str, -) -> Result<(), GrokRequestEncodeError> { - apply_build_response_defaults(body)?; - normalize_build_reasoning_effort(body, upstream_model); - sanitize_build_model_fields(body, upstream_model); - Ok(()) -} - -fn apply_build_response_defaults( - body: &mut Map, -) -> Result<(), GrokRequestEncodeError> { - if body.get("store").is_none_or(Value::is_null) { - body.insert("store".to_owned(), Value::Bool(false)); - } - - let include = body - .entry("include".to_owned()) - .or_insert_with(|| Value::Array(Vec::new())); - if include.is_null() { - *include = Value::Array(Vec::new()); - } - let include = include - .as_array_mut() - .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "include" })?; - if include.iter().any(|value| !value.is_string()) { - return Err(GrokRequestEncodeError::InvalidRequestField { field: "include" }); - } - if !include - .iter() - .any(|value| value.as_str() == Some("reasoning.encrypted_content")) - { - include.push(Value::String("reasoning.encrypted_content".to_owned())); - } - Ok(()) -} - -fn normalize_build_reasoning_effort(body: &mut Map, upstream_model: &str) { - let supports_effort = grok_model_supports_reasoning_effort(upstream_model); - if let Some(reasoning) = body.get_mut("reasoning").and_then(Value::as_object_mut) { - if let Some(effort) = reasoning.remove("effort") - && supports_effort - && let Some(effort) = normalize_grok_reasoning_effort_value(&effort, upstream_model) - { - reasoning.insert("effort".to_owned(), Value::String(effort.to_owned())); - } - if reasoning.is_empty() { - body.remove("reasoning"); - } - } - - if is_grok_composer_model(upstream_model) { - body.remove("reasoning"); - } -} - -fn normalize_grok_reasoning_effort_value(value: &Value, model: &str) -> Option<&'static str> { - let normalized = value.as_str()?.trim().to_ascii_lowercase(); - match normalized.as_str() { - "none" => Some("none"), - "minimal" | "low" => Some("low"), - "medium" => Some("medium"), - "xhigh" if matches!(model, "grok-4.6" | "grok-4.6-latest") => Some("xhigh"), - "high" | "xhigh" => Some("high"), - _ => None, - } -} - -pub(super) fn strip_grok_provider_prefix(model: &str) -> &str { - let model = model.trim(); - let lower = model.to_ascii_lowercase(); - ["xai/", "x-ai/", "grok/"] - .into_iter() - .find_map(|prefix| { - lower - .starts_with(prefix) - .then(|| model[prefix.len()..].trim()) - }) - .unwrap_or(model) -} - -pub(super) fn grok_model_last_slug(model: &str) -> String { - model - .trim() - .rsplit_once('/') - .map_or_else(|| model.trim(), |(_, slug)| slug.trim()) - .to_ascii_lowercase() -} - -pub(super) fn is_grok_composer_model(model: &str) -> bool { - matches!( - grok_model_last_slug(model).as_str(), - "grok-composer" | "grok-composer-2.5-fast" | "composer-2.5" - ) -} - -fn grok_model_supports_reasoning_effort(model: &str) -> bool { - matches!( - strip_grok_provider_prefix(model) - .to_ascii_lowercase() - .as_str(), - "grok-4.5" - | "grok-4.5-latest" - | "grok-4.6" - | "grok-4.6-latest" - | "grok-4.3" - | "grok-4.3-latest" - | "grok-3-mini" - | "grok-3-mini-fast" - | "grok-4.20-0309-reasoning" - | "grok-4.20-reasoning" - | "grok-4.20-multi-agent-0309" - ) -} - -fn sanitize_build_model_fields(body: &mut Map, upstream_model: &str) { - for field in ["prompt_cache_retention", "safety_identifier"] { - body.remove(field); - } - if upstream_model.trim().eq_ignore_ascii_case("grok-4.5") { - for field in [ - "presence_penalty", - "presencePenalty", - "frequency_penalty", - "frequencyPenalty", - "stop", - ] { - body.remove(field); - } - } - if ["grok-4.20", "grok-4.3", "grok-4.5", "grok-4.6"] - .iter() - .any(|prefix| grok_model_last_slug(upstream_model).starts_with(prefix)) - { - body.remove("logprobs"); - body.remove("top_logprobs"); - } - body.remove("external_web_access"); -} diff --git a/backend/crates/providers/xai/src/transport/request/parameters.rs b/backend/crates/providers/xai/src/transport/request/parameters.rs new file mode 100644 index 000000000..99c33ec88 --- /dev/null +++ b/backend/crates/providers/xai/src/transport/request/parameters.rs @@ -0,0 +1,76 @@ +//! 当前 Responses 参数默认值与 Grok wire 字段校验 + +use super::*; + +pub(super) fn normalize_build_request( + body: &mut Map, +) -> Result<(), GrokRequestEncodeError> { + apply_build_response_defaults(body)?; + normalize_build_reasoning_effort(body)?; + sanitize_build_fields(body); + Ok(()) +} + +fn apply_build_response_defaults( + body: &mut Map, +) -> Result<(), GrokRequestEncodeError> { + if body.get("store").is_none_or(Value::is_null) { + body.insert("store".to_owned(), Value::Bool(false)); + } + + let include = body + .entry("include".to_owned()) + .or_insert_with(|| Value::Array(Vec::new())); + if include.is_null() { + *include = Value::Array(Vec::new()); + } + let include = include + .as_array_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "include" })?; + if include.iter().any(|value| !value.is_string()) { + return Err(GrokRequestEncodeError::InvalidRequestField { field: "include" }); + } + if !include + .iter() + .any(|value| value.as_str() == Some("reasoning.encrypted_content")) + { + include.push(Value::String("reasoning.encrypted_content".to_owned())); + } + Ok(()) +} + +fn normalize_build_reasoning_effort( + body: &mut Map, +) -> Result<(), GrokRequestEncodeError> { + let Some(reasoning) = body.get_mut("reasoning").filter(|value| !value.is_null()) else { + return Ok(()); + }; + let reasoning = reasoning + .as_object_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "reasoning" })?; + let Some(effort) = reasoning.get_mut("effort").filter(|value| !value.is_null()) else { + return Ok(()); + }; + let normalized = effort + .as_str() + .map(|value| value.trim().to_ascii_lowercase()) + .filter(|value| { + matches!( + value.as_str(), + "none" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max" + ) + }) + .ok_or(GrokRequestEncodeError::InvalidRequestField { + field: "reasoning.effort", + })?; + // 菜单属于目录事实；请求编码只校验 wire，不根据 slug 改写用户选择的档位 + *effort = Value::String(normalized); + Ok(()) +} + +fn sanitize_build_fields(body: &mut Map) { + for field in ["prompt_cache_retention", "safety_identifier"] { + body.remove(field); + } + body.remove("external_web_access"); +} diff --git a/backend/crates/providers/xai/src/transport/request/tools.rs b/backend/crates/providers/xai/src/transport/request/tools.rs index e30e2e3ab..fdac10f73 100644 --- a/backend/crates/providers/xai/src/transport/request/tools.rs +++ b/backend/crates/providers/xai/src/transport/request/tools.rs @@ -121,13 +121,26 @@ impl ToolNormalizer { identity.kind as u8, identity.namespace, identity.name ); let mut alias = truncate_tool_alias(&base, &key); - if self + // Hosted tool 和 function 共用上游名称空间；缓存路由还会在归一化后注入搜索工具， + // 因此这两个官方 wire name 始终保留，避免同一函数随会话工具集变化而改名 + if is_hosted_search_wire_name(&alias) + || self + .response + .aliases + .get(&alias) + .is_some_and(|existing| existing != &identity) + { + alias = hashed_tool_alias(&base, &key); + } + let mut collision_index = 0; + while self .response .aliases .get(&alias) .is_some_and(|existing| existing != &identity) { - alias = hashed_tool_alias(&base, &key); + collision_index += 1; + alias = hashed_tool_alias(&base, &format!("{key}\0{collision_index}")); } self.response .aliases @@ -396,6 +409,10 @@ impl ToolNormalizer { } } +fn is_hosted_search_wire_name(name: &str) -> bool { + matches!(name, "web_search" | "x_search") +} + pub(super) fn optional_array( value: Option<&Value>, ) -> Result<(Vec, bool), GrokRequestEncodeError> { @@ -765,7 +782,7 @@ impl ToolNormalizer { normalized_tools: &[Value], ) -> Result<(), GrokRequestEncodeError> { if let Some(Value::Object(function)) = object.get_mut("function") { - rewrite_namespace_choice(function, &self.identity_aliases)?; + rewrite_function_choice(function, &self.identity_aliases)?; let name = string_field(function, "name").trim(); if !name.is_empty() && !has_named_tool(normalized_tools, "function", name) { payload.remove("tool_choice"); @@ -774,7 +791,7 @@ impl ToolNormalizer { payload.insert("tool_choice".to_owned(), Value::Object(object)); return Ok(()); } - rewrite_namespace_choice(&mut object, &self.identity_aliases)?; + rewrite_function_choice(&mut object, &self.identity_aliases)?; let name = string_field(&object, "name").trim(); if !name.is_empty() && !has_named_tool(normalized_tools, "function", name) { payload.remove("tool_choice"); @@ -788,27 +805,40 @@ impl ToolNormalizer { pub(super) fn normalize_web_search_filters( tool: &Map, ) -> Result>, GrokRequestEncodeError> { - let nested = match tool.get("filters") { - None | Some(Value::Null) => None, - Some(Value::Object(filters)) => filters - .get("allowed_domains") - .map(normalize_allowed_domains) - .transpose()?, - Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - }; - let top_level = tool - .get("allowed_domains") - .map(normalize_allowed_domains) - .transpose()?; - if nested.is_some() && top_level.is_some() && nested != top_level { + // 域名限制只接受当前 filters 合同；丢弃已知旧字段会把受限搜索变成无限制搜索 + if tool.contains_key("allowed_domains") { return Err(GrokRequestEncodeError::InvalidRequestNormalization); } - Ok(nested - .or(top_level) + let nested = match tool.get("filters") { + None | Some(Value::Null) => None, + Some(Value::Object(filters)) => Some(filters), + Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + }; + let allowed = nested + .and_then(|filters| filters.get("allowed_domains")) + .map(normalize_web_search_domains) + .transpose()? + .filter(|domains| !domains.is_empty()); + let excluded = nested + .and_then(|filters| filters.get("excluded_domains")) + .map(normalize_web_search_domains) + .transpose()? + .filter(|domains| !domains.is_empty()); + if let Some(excluded) = excluded { + // Grok Build 的非空允许表和排除表互斥；空表不限制搜索，不能覆盖另一张表 + if allowed.is_some() { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + return Ok(Some(Map::from_iter([( + "excluded_domains".to_owned(), + Value::Array(excluded), + )]))); + } + Ok(allowed .map(|domains| Map::from_iter([("allowed_domains".to_owned(), Value::Array(domains))]))) } -pub(super) fn normalize_allowed_domains( +pub(super) fn normalize_web_search_domains( value: &Value, ) -> Result, GrokRequestEncodeError> { let domains = value @@ -862,21 +892,24 @@ pub(super) fn tools_of_type(tools: &[Value], kind: &str) -> Vec { .collect() } -pub(super) fn rewrite_namespace_choice( +pub(super) fn rewrite_function_choice( object: &mut Map, aliases: &BTreeMap, ) -> Result<(), GrokRequestEncodeError> { let name = string_field(object, "name").trim(); let namespace = string_field(object, "namespace").trim(); - if name.is_empty() || namespace.is_empty() { + if name.is_empty() { return Ok(()); } let identity = ToolIdentity::new(ToolKind::Function, namespace, name); - let alias = aliases - .get(&identity) - .cloned() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - object.insert("name".to_owned(), Value::String(alias)); + let Some(alias) = aliases.get(&identity) else { + return if namespace.is_empty() { + Ok(()) + } else { + Err(GrokRequestEncodeError::InvalidRequestNormalization) + }; + }; + object.insert("name".to_owned(), Value::String(alias.clone())); object.remove("namespace"); Ok(()) } @@ -1085,7 +1118,7 @@ impl ToolNormalizer { Ok(converted) } - fn normalize_function_call_input( + pub(super) fn normalize_function_call_input( &mut self, item: &Map, ) -> Result, GrokRequestEncodeError> { @@ -1093,8 +1126,11 @@ impl ToolNormalizer { let call_id = required_trimmed_string(item, "call_id")?; let arguments = encode_function_arguments(item.get("arguments"))?; let namespace = string_field(item, "namespace").trim(); - if !namespace.is_empty() { - name = self.alias(ToolIdentity::new(ToolKind::Function, namespace, &name)); + let identity = ToolIdentity::new(ToolKind::Function, namespace, &name); + if !namespace.is_empty() || is_hosted_search_wire_name(&name) { + name = self.alias(identity); + } else if let Some(alias) = self.identity_aliases.get(&identity) { + name.clone_from(alias); } let mut converted = item.clone(); strip_grok_internal_keys(&mut converted); diff --git a/backend/crates/providers/xai/tests/admin.rs b/backend/crates/providers/xai/tests/admin.rs index c0620025a..73031744b 100644 --- a/backend/crates/providers/xai/tests/admin.rs +++ b/backend/crates/providers/xai/tests/admin.rs @@ -485,7 +485,12 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() quota: OpaqueProviderData::new( serde_json::from_value(serde_json::json!({ "subscriptionTier": "Free", - "config": {} + "config": { + "creditUsagePercent": 25, + "onDemandCap": {"val": 500}, + "onDemandUsed": {"val": 100}, + "prepaidBalance": {"val": -500} + } })) .expect("quota document"), ), @@ -503,6 +508,19 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() .await .expect("cached subscription"); assert_eq!(quota.plan_type.as_deref(), Some("Free")); + assert_eq!(quota.windows.len(), 1); + assert_eq!(quota.windows[0].used_percent, Some(25.0)); + let window_data = quota.windows[0] + .provider_data + .as_ref() + .expect("quota details") + .expose_to_provider() + .expose_to_provider(); + assert_eq!(window_data["onDemandCapCents"], 500); + assert_eq!(window_data["onDemandUsedCents"], 100); + assert_eq!(window_data["prepaidBalanceCents"], -500); + assert!(!window_data.contains_key("monthlyLimitCents")); + assert!(!window_data.contains_key("includedUsedCents")); let models = admin .models(account.id(), false) .await diff --git a/backend/crates/providers/xai/tests/credential/catalog.rs b/backend/crates/providers/xai/tests/credential/catalog.rs index 4835c4ce1..ff8b9d4ff 100644 --- a/backend/crates/providers/xai/tests/credential/catalog.rs +++ b/backend/crates/providers/xai/tests/credential/catalog.rs @@ -5,31 +5,27 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; -use futures::future::{BoxFuture, join_all}; +use futures::future::BoxFuture; use gateway_core::account::{ - AccountStateChange, CredentialCasUpdate, CredentialRevision, CredentialState, - OpaqueProviderData, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, - QuotaAccessChange, QuotaEvidence, QuotaObservation, QuotaState, + CredentialRevision, OpaqueProviderData, ProviderAccountStore, ProviderAccountUpdate, + QuotaAccessChange, QuotaEvidence, QuotaState, }; use gateway_core::provider_ports::{ ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderStoreError, }; use provider_xai::{ - GROK_SUBSCRIPTION_URL, GrokBillingRequest, GrokBillingTransport, GrokBillingTransportError, - GrokBillingTransportErrorKind, GrokBillingTransportFuture, GrokBillingTransportResponse, GrokCatalogCache, GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialCatalogError, GrokCredentialCatalogSeed, GrokCredentialRepository, GrokModelCatalogRequest, GrokModelCatalogTransport, GrokModelCatalogTransportError, GrokModelCatalogTransportErrorKind, GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, GrokPlanCatalog, - GrokQuotaError, }; use crate::support::{ MemoryGrokCatalogCache, MemoryProviderAccountStore, account_id, create_input, seed_input, }; -const OFFICIAL_FIXTURE: &[u8] = - include_bytes!("../transport/catalog/fixtures/official_grok_models_snapshot.json"); +const CLI_PROXY_FIXTURE: &[u8] = + include_bytes!("../transport/catalog/fixtures/cli_proxy_models.json"); struct QueueCatalogTransport { calls: AtomicUsize, @@ -88,58 +84,6 @@ impl GrokModelCatalogTransport for QueueCatalogTransport { } } -struct QueueBillingTransport { - calls: AtomicUsize, - subscription: Option>, - responses: Mutex>>, -} - -impl QueueBillingTransport { - fn success(body: &[u8]) -> Arc { - Arc::new(Self { - calls: AtomicUsize::new(0), - subscription: None, - responses: Mutex::new(VecDeque::from([Ok(GrokBillingTransportResponse::new( - body, - ))])), - }) - } - - fn failure() -> Arc { - Arc::new(Self { - calls: AtomicUsize::new(0), - subscription: None, - responses: Mutex::new(VecDeque::from([Err(GrokBillingTransportError::new( - GrokBillingTransportErrorKind::Unavailable, - ))])), - }) - } -} - -impl GrokBillingTransport for QueueBillingTransport { - fn execute(&self, request: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { - if request.endpoint().as_str() == GROK_SUBSCRIPTION_URL { - let response = self - .subscription - .clone() - .map(GrokBillingTransportResponse::new) - .ok_or_else(|| { - GrokBillingTransportError::new(GrokBillingTransportErrorKind::Unavailable) - }); - return Box::pin(async move { response }); - } - self.calls.fetch_add(1, Ordering::SeqCst); - let response = self - .responses - .lock() - .expect("billing response queue") - .pop_front() - .expect("one billing response"); - Box::pin(async move { response }) - } -} - -/// 每次读取都返回同一份无法解码的 catalog 文档的存储端口 struct CorruptCatalogCachePort; impl ProviderCatalogCachePort for CorruptCatalogCachePort { @@ -164,47 +108,6 @@ impl ProviderCatalogCachePort for CorruptCatalogCachePort { } } -enum BillingMutation { - State(AccountStateChange), - Credential(CredentialCasUpdate), -} - -struct MutatingBillingTransport { - store: Arc, - mutation: Mutex>, - body: Vec, -} - -impl GrokBillingTransport for MutatingBillingTransport { - fn execute(&self, request: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { - if request.endpoint().as_str() == GROK_SUBSCRIPTION_URL { - return Box::pin(async { - Err(GrokBillingTransportError::new( - GrokBillingTransportErrorKind::Unavailable, - )) - }); - } - let store = Arc::clone(&self.store); - let mutation = self.mutation.lock().expect("mutation").take(); - let body = self.body.clone(); - Box::pin(async move { - match mutation.expect("one mutation per request") { - BillingMutation::State(change) => store - .apply_state_change(change) - .await - .expect("apply concurrent state"), - BillingMutation::Credential(update) => { - store - .compare_and_swap_credential(update) - .await - .expect("apply concurrent credential rotation"); - } - } - Ok(GrokBillingTransportResponse::new(body)) - }) - } -} - async fn repository_with_accounts( suffixes: &[(&str, &str)], ) -> (Arc, GrokCredentialRepository) { @@ -219,38 +122,6 @@ async fn repository_with_accounts( (store, repository) } -async fn set_account_state( - store: &MemoryProviderAccountStore, - id: &ProviderAccountId, - credential_state: CredentialState, -) { - store - .apply_state_change(AccountStateChange { - message: None, - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - credential_state, - error_reason: credential_state.error_reason(), - observed_at: SystemTime::now(), - }) - .await - .expect("set account state"); -} - -#[tokio::test] -async fn concurrent_cold_scheduling_hydration_reads_quota_once() { - let (store, repository) = - repository_with_accounts(&[("quota-hydration", "subject-hydration")]).await; - let account = store - .account(&account_id("quota-hydration")) - .expect("created account"); - let service = crate::support::grok_quota_service(repository, QueueBillingTransport::failure()); - - join_all((0..32).map(|_| service.prepare_scheduling(std::slice::from_ref(&account)))).await; - - assert_eq!(store.quota_reads(), 1); -} - #[tokio::test] async fn catalog_query_caches_each_plan_and_returns_strict_union() { let (store, repository) = @@ -258,10 +129,10 @@ async fn catalog_query_caches_each_plan_and_returns_strict_union() { let cache = MemoryGrokCatalogCache::shared(); let cache_port: Arc = cache.clone(); let transport = QueueCatalogTransport::from_bodies([ - OFFICIAL_FIXTURE.to_vec(), - OFFICIAL_FIXTURE.to_vec(), - OFFICIAL_FIXTURE.to_vec(), - OFFICIAL_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), ]); let service = crate::support::grok_catalog_service(repository, transport, cache_port); assert_eq!(service.catalog_generation().get(), 0); @@ -293,7 +164,7 @@ async fn single_account_catalog_refresh_and_read_use_provider_cache_boundary() { let cache_port: Arc = cache; let service = crate::support::grok_catalog_service( repository, - QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]), + QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]), cache_port, ); let refreshed = service @@ -324,7 +195,7 @@ async fn disabled_account_refresh_discovers_models_with_the_pinned_account() { .expect("disable account"); let service = crate::support::grok_catalog_service( repository, - QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]), + QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]), MemoryGrokCatalogCache::shared(), ); @@ -374,7 +245,7 @@ async fn corrupt_catalog_cache_entry_reads_as_miss_and_refetches() { .is_none() ); - let transport = QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]); + let transport = QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]); let service = crate::support::grok_catalog_service(repository, transport.clone(), Arc::new(cache)); let catalog = service @@ -432,7 +303,7 @@ async fn manual_catalog_refresh_replaces_the_shared_plan_cache() { )) .await .expect("cache catalog"); - let transport = QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]); + let transport = QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]); let service = crate::support::grok_catalog_service(repository, transport.clone(), cache); let refreshed = service @@ -456,7 +327,7 @@ async fn catalog_refresh_falls_back_within_the_same_plan() { GrokModelCatalogTransportErrorKind::Unavailable, )), Ok(GrokModelCatalogTransportResponse::new( - OFFICIAL_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), None, )), ]); @@ -538,7 +409,7 @@ async fn quota_exhausted_account_remains_eligible_for_catalog_discovery() { .expect("mark quota exhausted"); let service = crate::support::grok_catalog_service( repository, - QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]), + QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]), MemoryGrokCatalogCache::shared(), ); @@ -582,7 +453,7 @@ async fn failed_plan_scope_is_skipped_and_surviving_plans_still_cache() { GrokModelCatalogTransportErrorKind::Unavailable, )), Ok(GrokModelCatalogTransportResponse::new( - OFFICIAL_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), None, )), ]); @@ -638,12 +509,12 @@ async fn conflicting_facts_for_same_slug_fail_closed() { .await .expect("separate plan catalog"); let mut conflicting: serde_json::Value = - serde_json::from_slice(OFFICIAL_FIXTURE).expect("fixture JSON"); + serde_json::from_slice(CLI_PROXY_FIXTURE).expect("fixture JSON"); conflicting["data"][0]["name"] = serde_json::json!("Different name"); let service = crate::support::grok_catalog_service( repository, QueueCatalogTransport::from_bodies([ - OFFICIAL_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), serde_json::to_vec(&conflicting).expect("conflicting JSON"), ]), MemoryGrokCatalogCache::shared(), @@ -665,450 +536,3 @@ fn seed_rejects_duplicates_and_supports_exact_membership() { assert!(seed.permits("grok-4.5")); assert!(!seed.permits("grok-4")); } - -#[tokio::test] -async fn quota_refresh_persists_dynamic_provider_document_and_projects_known_fields() { - let (store, repository) = repository_with_accounts(&[("quota", "subject-quota")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":37.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-13T00:00:00Z","end":"2026-07-20T00:00:00Z"},"prepaidBalance":{"val":2500},"futureWindow":{"kind":"rolling"}}}"#, - ); - let service = crate::support::grok_quota_service(repository, transport.clone()); - - let snapshot = service - .refresh_account(&account_id("quota")) - .await - .expect("refresh quota"); - let persisted = store - .get_quotas(&[account_id("quota")]) - .await - .expect("read persisted quota") - .pop() - .expect("quota exists"); - let document = persisted.quota; - - assert_eq!(transport.calls.load(Ordering::SeqCst), 1); - assert_eq!(snapshot.billing().used_percent(), Some(37.5)); - assert_eq!(snapshot.billing().plan_type(), None); - assert_eq!( - snapshot.billing().period_kind(), - provider_xai::GrokQuotaPeriodKind::Weekly - ); - assert_eq!( - snapshot.billing().period_type(), - Some("USAGE_PERIOD_TYPE_WEEKLY") - ); - assert_eq!(snapshot.billing().prepaid_balance_cents(), Some(2500)); - assert!( - document.expose_to_provider()["config"] - .get("futureWindow") - .is_some() - ); -} - -#[tokio::test] -async fn quota_refresh_persists_subscription_without_user_profile_and_reads_it_from_cache() { - let (store, repository) = - repository_with_accounts(&[("subscription", "subject-subscription")]).await; - let mut transport = QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":25}}"#); - Arc::get_mut(&mut transport).expect("unique transport").subscription = Some( - br#"{"userId":"verified-user","email":"private@example.com","subscriptionTier":"SuperGrokPro"}"#.to_vec(), - ); - let service = crate::support::grok_quota_service(repository, transport.clone()); - let snapshot = service - .refresh_account(&account_id("subscription")) - .await - .expect("refresh quota"); - assert_eq!(snapshot.billing().plan_type(), Some("SuperGrokPro")); - let cached = service - .read_account(&account_id("subscription")) - .await - .expect("read quota") - .expect("cached quota"); - assert_eq!(cached.billing().plan_type(), Some("SuperGrokPro")); - assert_eq!(transport.calls.load(Ordering::SeqCst), 1); - let stored = store - .get_quotas(&[account_id("subscription")]) - .await - .expect("persisted quota"); - let document = stored[0].quota.expose_to_provider(); - assert_eq!(document["subscriptionTier"], "SuperGrokPro"); - assert!(document.get("email").is_none()); - assert!(document.get("userId").is_none()); -} - -#[tokio::test] -async fn billing_percent_does_not_rewrite_quota_access_fact() { - let (store, repository) = - repository_with_accounts(&[("still-exhausted", "still-exhausted")]).await; - let id = account_id("still-exhausted"); - store - .apply_quota_access(QuotaAccessChange { - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), - }) - .await - .expect("mark quota exhausted"); - - crate::support::grok_quota_service( - repository, - QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":100}}"#), - ) - .refresh_account(&id) - .await - .expect("refresh exhausted quota"); - - assert_eq!( - store.account(&id).expect("account").quota().access(), - gateway_core::account::QuotaAccessState::Exhausted - ); -} - -#[tokio::test] -async fn authoritative_billing_refresh_clears_existing_quota_exhaustion() { - let (store, repository) = - repository_with_accounts(&[("recovered-quota", "recovered-quota")]).await; - let id = account_id("recovered-quota"); - store - .apply_quota_access(QuotaAccessChange { - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), - }) - .await - .expect("mark quota exhausted"); - - crate::support::grok_quota_service( - repository, - QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2099-07-22T00:00:00Z"}}}"#, - ), - ) - .refresh_account(&id) - .await - .expect("refresh recovered quota"); - - assert_eq!( - store.account(&id).expect("account").quota().access(), - gateway_core::account::QuotaAccessState::Allowed - ); -} - -#[tokio::test] -async fn full_usage_display_does_not_invent_quota_exhaustion() { - let (store, repository) = - repository_with_accounts(&[("exhausted-quota", "exhausted-quota")]).await; - let id = account_id("exhausted-quota"); - - crate::support::grok_quota_service( - repository, - QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":100,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2099-07-22T00:00:00Z"}}}"#, - ), - ) - .refresh_account(&id) - .await - .expect("refresh exhausted quota"); - - let quota = store.account(&id).expect("account").quota(); - assert_eq!( - quota.access(), - gateway_core::account::QuotaAccessState::Unknown - ); - assert!(quota.reset_at().is_none()); -} - -#[tokio::test] -async fn recovered_quota_does_not_clear_terminal_account_states() { - for (suffix, credential_state) in [ - ("keep-banned", CredentialState::Banned), - ("keep-expired", CredentialState::Expired), - ("keep-invalid", CredentialState::Invalid), - ] { - let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await; - let id = account_id(suffix); - set_account_state(&store, &id, credential_state).await; - crate::support::grok_quota_service( - repository, - QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#), - ) - .refresh_account(&id) - .await - .expect("refresh terminal account quota"); - - assert_eq!( - store.account(&id).expect("account").credential_state(), - credential_state - ); - } -} - -#[tokio::test] -async fn quota_refresh_preserves_ready_credential_state_across_concurrent_write() { - let (store, repository) = repository_with_accounts(&[("new-cooldown", "new-cooldown")]).await; - let id = account_id("new-cooldown"); - let transport = Arc::new(MutatingBillingTransport { - store: Arc::clone(&store), - mutation: Mutex::new(Some(BillingMutation::State(AccountStateChange { - message: None, - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - credential_state: CredentialState::Ready, - error_reason: None, - observed_at: SystemTime::now(), - }))), - body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), - }); - - crate::support::grok_quota_service(repository, transport) - .refresh_account(&id) - .await - .expect("refresh around concurrent state write"); - - let account = store.account(&id).expect("account"); - assert_eq!(account.credential_state(), CredentialState::Ready); -} - -#[tokio::test] -async fn quota_refresh_rejects_a_concurrent_credential_revision() { - let (store, repository) = repository_with_accounts(&[("new-revision", "new-revision")]).await; - let id = account_id("new-revision"); - let account = store.account(&id).expect("account"); - let update = CredentialCasUpdate::new( - id.clone(), - account.revision(), - ProviderAccountUpdate { - account_id: id.clone(), - name: account.name().to_owned(), - email: account.email().map(str::to_owned), - plan_type: account.plan_type().map(str::to_owned), - }, - store.credential(&id).expect("credential"), - account.has_refresh_token(), - account.access_token_expires_at(), - account.next_refresh_at(), - ) - .expect("credential update"); - let transport = Arc::new(MutatingBillingTransport { - store: Arc::clone(&store), - mutation: Mutex::new(Some(BillingMutation::Credential(update))), - body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), - }); - - assert!(matches!( - crate::support::grok_quota_service(repository, transport) - .refresh_account(&id) - .await, - Err(GrokQuotaError::StaleCredentialSnapshot) - )); - assert_eq!(store.account(&id).expect("account").revision().get(), 2); -} - -#[tokio::test] -async fn quota_projection_falls_back_to_legacy_monthly_usage() { - let (_, repository) = - repository_with_accounts(&[("monthly-quota", "subject-monthly-quota")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"monthlyLimit":{"val":10000},"used":{"val":2500},"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&account_id("monthly-quota")) - .await - .expect("refresh monthly quota"); - - assert_eq!(snapshot.billing().used_percent(), Some(25.0)); - assert_eq!( - snapshot.billing().period_kind(), - provider_xai::GrokQuotaPeriodKind::Monthly - ); - assert_eq!( - snapshot.billing().period_end(), - Some("2026-08-01T00:00:00Z") - ); -} - -#[tokio::test] -async fn quota_projection_preserves_unknown_period_for_dynamic_duration_fallback() { - let (_, repository) = - repository_with_accounts(&[("dynamic-quota", "subject-dynamic-quota")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_FORTNIGHT","start":"2026-07-01T00:00:00Z","end":"2026-07-15T00:00:00Z"}}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&account_id("dynamic-quota")) - .await - .expect("refresh dynamic quota"); - - assert_eq!(snapshot.billing().used_percent(), Some(12.5)); - assert_eq!( - snapshot.billing().period_kind(), - provider_xai::GrokQuotaPeriodKind::Other - ); - assert_eq!( - snapshot.billing().period_start(), - Some("2026-07-01T00:00:00Z") - ); - assert_eq!( - snapshot.billing().period_end(), - Some("2026-07-15T00:00:00Z") - ); -} - -#[tokio::test] -async fn expired_billing_window_does_not_participate_in_scheduling_rank() { - let (store, repository) = - repository_with_accounts(&[("expired-billing", "expired-billing")]).await; - let id = account_id("expired-billing"); - let account = store.account(&id).expect("account"); - let service = crate::support::grok_quota_service( - repository, - QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":25,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2000-01-01T00:00:00Z","end":"2000-01-08T00:00:00Z"}}}"#, - ), - ); - - service - .refresh_account(&id) - .await - .expect("refresh expired billing window"); - - assert_eq!(service.scheduling_signals(&account), None); - assert_eq!( - store - .account(&id) - .expect("refreshed account") - .quota() - .access(), - gateway_core::account::QuotaAccessState::Unknown - ); -} - -#[tokio::test] -async fn non_authoritative_quota_refresh_does_not_clear_existing_quota_exhaustion() { - let (store, repository) = - repository_with_accounts(&[("free-quota", "subject-free-quota")]).await; - let id = account_id("free-quota"); - store - .apply_quota_access(QuotaAccessChange { - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), - }) - .await - .expect("mark quota exhausted"); - let transport = QueueBillingTransport::success( - br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"onDemandCap":{"val":0},"onDemandUsed":{"val":0},"prepaidBalance":{"val":0}}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&id) - .await - .expect("refresh Free quota"); - - assert!(!snapshot.billing().has_authoritative_quota()); - assert_eq!( - store.account(&id).expect("account").quota().access(), - gateway_core::account::QuotaAccessState::Exhausted - ); -} - -#[tokio::test] -async fn reported_zero_percent_is_authoritative_quota() { - let (_, repository) = - repository_with_accounts(&[("zero-percent", "subject-zero-percent")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":0,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"}}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&account_id("zero-percent")) - .await - .expect("refresh reported quota"); - - assert!(snapshot.billing().has_authoritative_quota()); -} - -#[tokio::test] -async fn positive_prepaid_balance_is_authoritative_quota() { - let (_, repository) = - repository_with_accounts(&[("prepaid-quota", "subject-prepaid-quota")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"prepaidBalance":{"val":500}}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&account_id("prepaid-quota")) - .await - .expect("refresh prepaid quota"); - - assert!(snapshot.billing().has_authoritative_quota()); -} - -#[tokio::test] -async fn quota_read_rejects_corrupt_provider_document() { - let (store, repository) = repository_with_accounts(&[("corrupt", "subject-corrupt")]).await; - let mut document = serde_json::Map::new(); - document.insert("config".to_owned(), serde_json::json!([])); - store - .compare_and_swap_quota(QuotaObservation { - plan_type: None, - account_id: account_id("corrupt"), - expected_revision: CredentialRevision::new(1).expect("revision"), - quota: OpaqueProviderData::new(document), - observed_at: SystemTime::now(), - state: QuotaState::unknown(), - }) - .await - .expect("seed corrupt quota"); - let service = crate::support::grok_quota_service( - repository, - QueueBillingTransport::success(br#"{"config":null}"#), - ); - - assert!(matches!( - service.read_account(&account_id("corrupt")).await, - Err(GrokQuotaError::InvalidData) - )); -} - -#[tokio::test] -async fn disabled_account_quota_refresh_never_calls_upstream() { - let (store, repository) = - repository_with_accounts(&[("disabled-quota", "subject-disabled")]).await; - store - .set_enabled(&account_id("disabled-quota"), false) - .await - .expect("disable account"); - let transport = QueueBillingTransport::success(br#"{"config":null}"#); - let service = crate::support::grok_quota_service(repository, transport.clone()); - - assert!(matches!( - service.refresh_account(&account_id("disabled-quota")).await, - Err(GrokQuotaError::AccountUnavailable) - )); - assert_eq!(transport.calls.load(Ordering::SeqCst), 0); -} - -#[tokio::test] -async fn failed_quota_refresh_does_not_replace_last_good_observation() { - let (store, repository) = repository_with_accounts(&[("stable", "subject-stable")]).await; - let good = QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#); - crate::support::grok_quota_service(repository.clone(), good) - .refresh_account(&account_id("stable")) - .await - .expect("seed good observation"); - let service = crate::support::grok_quota_service(repository, QueueBillingTransport::failure()); - - assert!(matches!( - service.refresh_account(&account_id("stable")).await, - Err(GrokQuotaError::Upstream) - )); - let persisted = store - .get_quotas(&[account_id("stable")]) - .await - .expect("read quota") - .pop() - .expect("quota remains") - .quota; - assert_eq!( - persisted.expose_to_provider()["config"]["creditUsagePercent"].as_f64(), - Some(10.0), - ); -} diff --git a/backend/crates/providers/xai/tests/credential/mod.rs b/backend/crates/providers/xai/tests/credential/mod.rs index f48673bf5..df4f292fd 100644 --- a/backend/crates/providers/xai/tests/credential/mod.rs +++ b/backend/crates/providers/xai/tests/credential/mod.rs @@ -10,6 +10,7 @@ mod http; mod import; mod oidc_verifier; mod pkce; +mod quota; mod refresh; mod repository; mod secret; diff --git a/backend/crates/providers/xai/tests/credential/quota.rs b/backend/crates/providers/xai/tests/credential/quota.rs new file mode 100644 index 000000000..0efd180eb --- /dev/null +++ b/backend/crates/providers/xai/tests/credential/quota.rs @@ -0,0 +1,788 @@ +//! 验证 xAI 额度观察、revision隔离、权限状态与调度投影 + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::SystemTime; + +use futures::future::join_all; +use gateway_core::account::{ + AccountStateChange, CredentialCasUpdate, CredentialRevision, CredentialState, + OpaqueProviderData, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, + QuotaAccessChange, QuotaEvidence, QuotaObservation, QuotaState, +}; +use provider_xai::{ + GROK_SUBSCRIPTION_URL, GrokBillingRequest, GrokBillingTransport, GrokBillingTransportError, + GrokBillingTransportErrorKind, GrokBillingTransportFuture, GrokBillingTransportResponse, + GrokCredentialRepository, GrokQuotaError, +}; + +use crate::support::{MemoryProviderAccountStore, account_id, create_input, seed_input}; + +struct QueueBillingTransport { + calls: AtomicUsize, + subscription: Option>, + responses: Mutex>>, +} + +impl QueueBillingTransport { + fn success(body: &[u8]) -> Arc { + Arc::new(Self { + calls: AtomicUsize::new(0), + subscription: None, + responses: Mutex::new(VecDeque::from([Ok(GrokBillingTransportResponse::new( + body, + ))])), + }) + } + + fn failure() -> Arc { + Arc::new(Self { + calls: AtomicUsize::new(0), + subscription: None, + responses: Mutex::new(VecDeque::from([Err(GrokBillingTransportError::new( + GrokBillingTransportErrorKind::Unavailable, + ))])), + }) + } +} + +impl GrokBillingTransport for QueueBillingTransport { + fn execute(&self, request: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { + if request.endpoint().as_str() == GROK_SUBSCRIPTION_URL { + let response = self + .subscription + .clone() + .map(GrokBillingTransportResponse::new) + .ok_or_else(|| { + GrokBillingTransportError::new(GrokBillingTransportErrorKind::Unavailable) + }); + return Box::pin(async move { response }); + } + self.calls.fetch_add(1, Ordering::SeqCst); + let response = self + .responses + .lock() + .expect("billing response queue") + .pop_front() + .expect("one billing response"); + Box::pin(async move { response }) + } +} + +/// 每次读取都返回同一份无法解码的 catalog 文档的存储端口 +enum BillingMutation { + State(AccountStateChange), + Credential(CredentialCasUpdate), +} + +struct MutatingBillingTransport { + store: Arc, + mutation: Mutex>, + body: Vec, +} + +impl GrokBillingTransport for MutatingBillingTransport { + fn execute(&self, request: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { + if request.endpoint().as_str() == GROK_SUBSCRIPTION_URL { + return Box::pin(async { + Err(GrokBillingTransportError::new( + GrokBillingTransportErrorKind::Unavailable, + )) + }); + } + let store = Arc::clone(&self.store); + let mutation = self.mutation.lock().expect("mutation").take(); + let body = self.body.clone(); + Box::pin(async move { + match mutation.expect("one mutation per request") { + BillingMutation::State(change) => store + .apply_state_change(change) + .await + .expect("apply concurrent state"), + BillingMutation::Credential(update) => { + store + .compare_and_swap_credential(update) + .await + .expect("apply concurrent credential rotation"); + } + } + Ok(GrokBillingTransportResponse::new(body)) + }) + } +} + +async fn repository_with_accounts( + suffixes: &[(&str, &str)], +) -> (Arc, GrokCredentialRepository) { + let store = MemoryProviderAccountStore::shared(); + let account_store: Arc = store.clone(); + let repository = GrokCredentialRepository::new(account_store); + for (suffix, subject) in suffixes { + seed_input(&store, &create_input(suffix, subject)) + .await + .expect("create account"); + } + (store, repository) +} + +async fn set_account_state( + store: &MemoryProviderAccountStore, + id: &ProviderAccountId, + credential_state: CredentialState, +) { + store + .apply_state_change(AccountStateChange { + message: None, + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + credential_state, + error_reason: credential_state.error_reason(), + observed_at: SystemTime::now(), + }) + .await + .expect("set account state"); +} + +#[tokio::test] +async fn concurrent_cold_scheduling_hydration_reads_quota_once() { + let (store, repository) = + repository_with_accounts(&[("quota-hydration", "subject-hydration")]).await; + let account = store + .account(&account_id("quota-hydration")) + .expect("created account"); + let service = crate::support::grok_quota_service(repository, QueueBillingTransport::failure()); + + join_all((0..32).map(|_| service.prepare_scheduling(std::slice::from_ref(&account)))).await; + + assert_eq!(store.quota_reads(), 1); +} + +#[tokio::test] +async fn quota_refresh_persists_dynamic_provider_document_and_projects_known_fields() { + let (store, repository) = repository_with_accounts(&[("quota", "subject-quota")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":37.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-13T00:00:00Z","end":"2026-07-20T00:00:00Z"},"prepaidBalance":{"val":2500},"futureWindow":{"kind":"rolling"}}}"#, + ); + let service = crate::support::grok_quota_service(repository, transport.clone()); + + let snapshot = service + .refresh_account(&account_id("quota")) + .await + .expect("refresh quota"); + let persisted = store + .get_quotas(&[account_id("quota")]) + .await + .expect("read persisted quota") + .pop() + .expect("quota exists"); + let document = persisted.quota; + + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + assert_eq!(snapshot.billing().used_percent(), Some(37.5)); + assert_eq!(snapshot.billing().plan_type(), None); + assert_eq!( + snapshot.billing().period_kind(), + provider_xai::GrokQuotaPeriodKind::Weekly + ); + assert_eq!( + snapshot.billing().period_type(), + Some("USAGE_PERIOD_TYPE_WEEKLY") + ); + assert_eq!(snapshot.billing().prepaid_balance_cents(), Some(2500)); + assert!( + document.expose_to_provider()["config"] + .get("futureWindow") + .is_some() + ); +} + +#[tokio::test] +async fn quota_refresh_persists_subscription_without_user_profile_and_reads_it_from_cache() { + let (store, repository) = + repository_with_accounts(&[("subscription", "subject-subscription")]).await; + let mut transport = QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":25}}"#); + Arc::get_mut(&mut transport).expect("unique transport").subscription = Some( + br#"{"userId":"verified-user","email":"private@example.com","subscriptionTier":"SuperGrokPro"}"#.to_vec(), + ); + let service = crate::support::grok_quota_service(repository, transport.clone()); + let snapshot = service + .refresh_account(&account_id("subscription")) + .await + .expect("refresh quota"); + assert_eq!(snapshot.billing().plan_type(), Some("SuperGrokPro")); + let cached = service + .read_account(&account_id("subscription")) + .await + .expect("read quota") + .expect("cached quota"); + assert_eq!(cached.billing().plan_type(), Some("SuperGrokPro")); + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + let stored = store + .get_quotas(&[account_id("subscription")]) + .await + .expect("persisted quota"); + let document = stored[0].quota.expose_to_provider(); + assert_eq!(document["subscriptionTier"], "SuperGrokPro"); + assert!(document.get("email").is_none()); + assert!(document.get("userId").is_none()); +} + +#[tokio::test] +async fn billing_percent_does_not_rewrite_quota_access_fact() { + let (store, repository) = + repository_with_accounts(&[("still-exhausted", "still-exhausted")]).await; + let id = account_id("still-exhausted"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), + }) + .await + .expect("mark quota exhausted"); + + crate::support::grok_quota_service( + repository, + QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":100}}"#), + ) + .refresh_account(&id) + .await + .expect("refresh exhausted quota"); + + assert_eq!( + store.account(&id).expect("account").quota().access(), + gateway_core::account::QuotaAccessState::Exhausted + ); +} + +#[tokio::test] +async fn billing_access_and_scheduling_use_current_quota_fields() { + use gateway_core::account::QuotaAccessState::{Allowed, Exhausted}; + + for (config, expected_percent, expected_access) in [ + ( + serde_json::json!({"creditUsagePercent": 100}), + Some(100.0), + Exhausted, + ), + ( + serde_json::json!({"creditUsagePercent": 25}), + Some(25.0), + Allowed, + ), + ( + serde_json::json!({"onDemandCap": {"val": 500}, "onDemandUsed": {"val": 100}}), + None, + Allowed, + ), + ( + serde_json::json!({ + "creditUsagePercent": 100, + "onDemandCap": {"val": 500}, "onDemandUsed": {"val": 100} + }), + Some(100.0), + Allowed, + ), + ] { + let (store, repository) = repository_with_accounts(&[("percent", "percent")]).await; + let id = account_id("percent"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted( + QuotaEvidence::ProviderDenied, + SystemTime::now(), + None, + ), + }) + .await + .expect("mark quota exhausted"); + let body = serde_json::to_vec(&serde_json::json!({"config": config})).expect("billing"); + let service = + crate::support::grok_quota_service(repository, QueueBillingTransport::success(&body)); + let snapshot = service.refresh_account(&id).await.expect("refresh quota"); + let account = store.account(&id).expect("refreshed account"); + + assert_eq!(snapshot.billing().used_percent(), expected_percent); + assert_eq!( + account.quota().access(), + expected_access, + "config: {config}" + ); + assert_eq!( + service + .scheduling_signals(&account) + .and_then(|signals| signals.remaining_rank()), + expected_percent.map(|percent| ((100.0 - percent) * 100.0).round() as u64) + ); + } +} + +#[tokio::test] +async fn signed_prepaid_balance_preserves_wire_and_recovers_access_when_current() { + use gateway_core::account::QuotaAccessState::{Allowed, Exhausted}; + + for (balance, period_end, expected_access) in [ + (-500, None, Allowed), + (i64::MIN, None, Allowed), + (0, None, Exhausted), + (500, None, Allowed), + (-500, Some("2000-01-01T00:00:00Z"), Exhausted), + (-500, Some("2099-01-01T00:00:00Z"), Allowed), + ] { + let (store, repository) = repository_with_accounts(&[("signed", "signed")]).await; + let id = account_id("signed"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted( + QuotaEvidence::ProviderDenied, + SystemTime::now(), + None, + ), + }) + .await + .expect("mark quota exhausted"); + let body = serde_json::to_vec(&serde_json::json!({ + "config": { + "prepaidBalance": {"val": balance}, "currentPeriod": {"end": period_end} + } + })) + .expect("billing"); + let transport = QueueBillingTransport::success(&body); + let service = crate::support::grok_quota_service(repository, transport.clone()); + + let snapshot = service + .refresh_account(&id) + .await + .expect("refresh signed balance"); + let cached = service + .read_account(&id) + .await + .expect("read quota") + .expect("cached quota"); + let persisted = store + .get_quotas(std::slice::from_ref(&id)) + .await + .expect("stored quota"); + + assert_eq!(snapshot.billing().prepaid_balance_cents(), Some(balance)); + assert_eq!(snapshot.billing().has_authoritative_quota(), balance != 0); + assert_eq!(cached.billing().prepaid_balance_cents(), Some(balance)); + assert_eq!( + persisted[0].quota.expose_to_provider()["config"]["prepaidBalance"]["val"], + balance + ); + assert_eq!( + store.account(&id).expect("account").quota().access(), + expected_access, + "balance: {balance}, period_end: {period_end:?}" + ); + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test] +async fn authoritative_billing_refresh_clears_existing_quota_exhaustion() { + let (store, repository) = + repository_with_accounts(&[("recovered-quota", "recovered-quota")]).await; + let id = account_id("recovered-quota"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), + }) + .await + .expect("mark quota exhausted"); + + crate::support::grok_quota_service( + repository, + QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2099-07-22T00:00:00Z"}}}"#, + ), + ) + .refresh_account(&id) + .await + .expect("refresh recovered quota"); + + assert_eq!( + store.account(&id).expect("account").quota().access(), + gateway_core::account::QuotaAccessState::Allowed + ); +} + +#[tokio::test] +async fn full_usage_display_does_not_invent_quota_exhaustion() { + let (store, repository) = + repository_with_accounts(&[("exhausted-quota", "exhausted-quota")]).await; + let id = account_id("exhausted-quota"); + + crate::support::grok_quota_service( + repository, + QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":100,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2099-07-22T00:00:00Z"}}}"#, + ), + ) + .refresh_account(&id) + .await + .expect("refresh exhausted quota"); + + let quota = store.account(&id).expect("account").quota(); + assert_eq!( + quota.access(), + gateway_core::account::QuotaAccessState::Unknown + ); + assert!(quota.reset_at().is_none()); +} + +#[tokio::test] +async fn recovered_quota_does_not_clear_terminal_account_states() { + for (suffix, credential_state) in [ + ("keep-banned", CredentialState::Banned), + ("keep-expired", CredentialState::Expired), + ("keep-invalid", CredentialState::Invalid), + ] { + let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await; + let id = account_id(suffix); + set_account_state(&store, &id, credential_state).await; + crate::support::grok_quota_service( + repository, + QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#), + ) + .refresh_account(&id) + .await + .expect("refresh terminal account quota"); + + assert_eq!( + store.account(&id).expect("account").credential_state(), + credential_state + ); + } +} + +#[tokio::test] +async fn quota_refresh_preserves_ready_credential_state_across_concurrent_write() { + let (store, repository) = repository_with_accounts(&[("new-cooldown", "new-cooldown")]).await; + let id = account_id("new-cooldown"); + let transport = Arc::new(MutatingBillingTransport { + store: Arc::clone(&store), + mutation: Mutex::new(Some(BillingMutation::State(AccountStateChange { + message: None, + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + credential_state: CredentialState::Ready, + error_reason: None, + observed_at: SystemTime::now(), + }))), + body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), + }); + + crate::support::grok_quota_service(repository, transport) + .refresh_account(&id) + .await + .expect("refresh around concurrent state write"); + + let account = store.account(&id).expect("account"); + assert_eq!(account.credential_state(), CredentialState::Ready); +} + +#[tokio::test] +async fn quota_refresh_rejects_a_concurrent_credential_revision() { + let (store, repository) = repository_with_accounts(&[("new-revision", "new-revision")]).await; + let id = account_id("new-revision"); + let account = store.account(&id).expect("account"); + let update = CredentialCasUpdate::new( + id.clone(), + account.revision(), + ProviderAccountUpdate { + account_id: id.clone(), + name: account.name().to_owned(), + email: account.email().map(str::to_owned), + plan_type: account.plan_type().map(str::to_owned), + }, + store.credential(&id).expect("credential"), + account.has_refresh_token(), + account.access_token_expires_at(), + account.next_refresh_at(), + ) + .expect("credential update"); + let transport = Arc::new(MutatingBillingTransport { + store: Arc::clone(&store), + mutation: Mutex::new(Some(BillingMutation::Credential(update))), + body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), + }); + + assert!(matches!( + crate::support::grok_quota_service(repository, transport) + .refresh_account(&id) + .await, + Err(GrokQuotaError::StaleCredentialSnapshot) + )); + assert_eq!(store.account(&id).expect("account").revision().get(), 2); +} + +#[tokio::test] +async fn legacy_billing_fields_do_not_supply_quota_evidence_or_period() { + for (limit, used) in [ + ( + serde_json::json!({"val": 10000}), + serde_json::json!({"val": 2500}), + ), + ( + serde_json::json!("uninterpreted"), + serde_json::json!({"val": -1}), + ), + ] { + let (store, repository) = repository_with_accounts(&[("legacy", "legacy")]).await; + let id = account_id("legacy"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted( + QuotaEvidence::ProviderDenied, + SystemTime::now(), + None, + ), + }) + .await + .expect("mark quota exhausted"); + let body = serde_json::to_vec(&serde_json::json!({"config": { + "monthlyLimit": limit, "used": used, + "billingPeriodStart": "2026-07-01T00:00:00Z", + "billingPeriodEnd": "2099-08-01T00:00:00Z" + }})) + .expect("billing document"); + let service = + crate::support::grok_quota_service(repository, QueueBillingTransport::success(&body)); + let snapshot = service + .refresh_account(&id) + .await + .expect("refresh opaque billing"); + let account = store.account(&id).expect("account"); + + assert_eq!(snapshot.billing().used_percent(), None); + assert!(!snapshot.billing().has_authoritative_quota()); + assert_eq!( + snapshot.billing().period_kind(), + provider_xai::GrokQuotaPeriodKind::Other + ); + assert_eq!(snapshot.billing().period_start(), None); + assert_eq!(snapshot.billing().period_end(), None); + assert_eq!(service.scheduling_signals(&account), None); + assert_eq!( + account.quota().access(), + gateway_core::account::QuotaAccessState::Exhausted + ); + } +} + +#[tokio::test] +async fn current_quota_period_does_not_fill_missing_dates_from_legacy_fields() { + let (_, repository) = repository_with_accounts(&[("current-period", "current-period")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":25,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY"},"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2099-08-01T00:00:00Z"}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&account_id("current-period")) + .await + .expect("current quota"); + + assert_eq!(snapshot.billing().used_percent(), Some(25.0)); + assert_eq!( + snapshot.billing().period_kind(), + provider_xai::GrokQuotaPeriodKind::Weekly + ); + assert_eq!(snapshot.billing().period_start(), None); + assert_eq!(snapshot.billing().period_end(), None); +} + +#[tokio::test] +async fn quota_projection_preserves_unknown_period_for_dynamic_duration_fallback() { + let (_, repository) = + repository_with_accounts(&[("dynamic-quota", "subject-dynamic-quota")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_FORTNIGHT","start":"2026-07-01T00:00:00Z","end":"2026-07-15T00:00:00Z"}}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&account_id("dynamic-quota")) + .await + .expect("refresh dynamic quota"); + + assert_eq!(snapshot.billing().used_percent(), Some(12.5)); + assert_eq!( + snapshot.billing().period_kind(), + provider_xai::GrokQuotaPeriodKind::Other + ); + assert_eq!( + snapshot.billing().period_start(), + Some("2026-07-01T00:00:00Z") + ); + assert_eq!( + snapshot.billing().period_end(), + Some("2026-07-15T00:00:00Z") + ); +} + +#[tokio::test] +async fn expired_billing_window_does_not_participate_in_scheduling_rank() { + let (store, repository) = + repository_with_accounts(&[("expired-billing", "expired-billing")]).await; + let id = account_id("expired-billing"); + let account = store.account(&id).expect("account"); + let service = crate::support::grok_quota_service( + repository, + QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":25,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2000-01-01T00:00:00Z","end":"2000-01-08T00:00:00Z"}}}"#, + ), + ); + + service + .refresh_account(&id) + .await + .expect("refresh expired billing window"); + + assert_eq!(service.scheduling_signals(&account), None); + assert_eq!( + store + .account(&id) + .expect("refreshed account") + .quota() + .access(), + gateway_core::account::QuotaAccessState::Unknown + ); +} + +#[tokio::test] +async fn non_authoritative_quota_refresh_does_not_clear_existing_quota_exhaustion() { + let (store, repository) = + repository_with_accounts(&[("free-quota", "subject-free-quota")]).await; + let id = account_id("free-quota"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), + }) + .await + .expect("mark quota exhausted"); + let transport = QueueBillingTransport::success( + br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"onDemandCap":{"val":0},"onDemandUsed":{"val":0},"prepaidBalance":{"val":0}}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&id) + .await + .expect("refresh Free quota"); + + assert!(!snapshot.billing().has_authoritative_quota()); + assert_eq!( + store.account(&id).expect("account").quota().access(), + gateway_core::account::QuotaAccessState::Exhausted + ); +} + +#[tokio::test] +async fn reported_zero_percent_is_authoritative_quota() { + let (_, repository) = + repository_with_accounts(&[("zero-percent", "subject-zero-percent")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":0,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"}}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&account_id("zero-percent")) + .await + .expect("refresh reported quota"); + + assert!(snapshot.billing().has_authoritative_quota()); +} + +#[tokio::test] +async fn positive_prepaid_balance_is_authoritative_quota() { + let (_, repository) = + repository_with_accounts(&[("prepaid-quota", "subject-prepaid-quota")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"prepaidBalance":{"val":500}}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&account_id("prepaid-quota")) + .await + .expect("refresh prepaid quota"); + + assert!(snapshot.billing().has_authoritative_quota()); +} + +#[tokio::test] +async fn quota_read_rejects_corrupt_provider_document() { + let (store, repository) = repository_with_accounts(&[("corrupt", "subject-corrupt")]).await; + let mut document = serde_json::Map::new(); + document.insert("config".to_owned(), serde_json::json!([])); + store + .compare_and_swap_quota(QuotaObservation { + plan_type: None, + account_id: account_id("corrupt"), + expected_revision: CredentialRevision::new(1).expect("revision"), + quota: OpaqueProviderData::new(document), + observed_at: SystemTime::now(), + state: QuotaState::unknown(), + }) + .await + .expect("seed corrupt quota"); + let service = crate::support::grok_quota_service( + repository, + QueueBillingTransport::success(br#"{"config":null}"#), + ); + + assert!(matches!( + service.read_account(&account_id("corrupt")).await, + Err(GrokQuotaError::InvalidData) + )); +} + +#[tokio::test] +async fn disabled_account_quota_refresh_never_calls_upstream() { + let (store, repository) = + repository_with_accounts(&[("disabled-quota", "subject-disabled")]).await; + store + .set_enabled(&account_id("disabled-quota"), false) + .await + .expect("disable account"); + let transport = QueueBillingTransport::success(br#"{"config":null}"#); + let service = crate::support::grok_quota_service(repository, transport.clone()); + + assert!(matches!( + service.refresh_account(&account_id("disabled-quota")).await, + Err(GrokQuotaError::AccountUnavailable) + )); + assert_eq!(transport.calls.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn failed_quota_refresh_does_not_replace_last_good_observation() { + let (store, repository) = repository_with_accounts(&[("stable", "subject-stable")]).await; + let good = QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#); + crate::support::grok_quota_service(repository.clone(), good) + .refresh_account(&account_id("stable")) + .await + .expect("seed good observation"); + let service = crate::support::grok_quota_service(repository, QueueBillingTransport::failure()); + + assert!(matches!( + service.refresh_account(&account_id("stable")).await, + Err(GrokQuotaError::Upstream) + )); + let persisted = store + .get_quotas(&[account_id("stable")]) + .await + .expect("read quota") + .pop() + .expect("quota remains") + .quota; + assert_eq!( + persisted.expose_to_provider()["config"]["creditUsagePercent"].as_f64(), + Some(10.0), + ); +} diff --git a/backend/crates/providers/xai/tests/credential/refresh.rs b/backend/crates/providers/xai/tests/credential/refresh.rs index 54afb2912..53f9407db 100644 --- a/backend/crates/providers/xai/tests/credential/refresh.rs +++ b/backend/crates/providers/xai/tests/credential/refresh.rs @@ -31,8 +31,8 @@ use crate::support::{ rejecting_account_proxy, runtime_policy, seed_input, }; -const OFFICIAL_FIXTURE: &[u8] = - include_bytes!("../transport/catalog/fixtures/official_grok_models_snapshot.json"); +const CLI_PROXY_FIXTURE: &[u8] = + include_bytes!("../transport/catalog/fixtures/cli_proxy_models.json"); const OAUTH_BACKOFF_ATTEMPTS: u32 = 5; #[tokio::test] @@ -190,7 +190,7 @@ impl GrokModelCatalogTransport for StaticCatalogTransport { fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> { Box::pin(async { Ok(GrokModelCatalogTransportResponse::new( - OFFICIAL_FIXTURE, + CLI_PROXY_FIXTURE, None, )) }) diff --git a/backend/crates/providers/xai/tests/provider/contract.rs b/backend/crates/providers/xai/tests/provider/contract/mod.rs similarity index 94% rename from backend/crates/providers/xai/tests/provider/contract.rs rename to backend/crates/providers/xai/tests/provider/contract/mod.rs index 90fe1b11e..742cbb997 100644 --- a/backend/crates/providers/xai/tests/provider/contract.rs +++ b/backend/crates/providers/xai/tests/provider/contract/mod.rs @@ -1,5 +1,7 @@ //! 验证 xAI 原生执行的协议转换、账号选择与交付状态隔离 +mod replay_tool_names; + use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::num::NonZeroU32; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -38,12 +40,12 @@ use gateway_core::provider_ports::{ ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderSchedulingState, ProviderStoreError, }; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::routing::{ ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, SupportLevel, UpstreamModelId, }; -use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::upstream::UpstreamSendState; use provider_xai::{ GrokAccountSessionSelector, GrokBillingRequest, GrokBillingTransport, @@ -68,7 +70,116 @@ use crate::support::{ const MODEL: &str = "grok-4.5"; const CATALOG_FIXTURE: &[u8] = - include_bytes!("../transport/catalog/fixtures/official_grok_models_snapshot.json"); + include_bytes!("../../transport/catalog/fixtures/cli_proxy_models.json"); + +#[tokio::test] +async fn native_xai_rejects_generate_false_before_upstream_io_or_usage() { + for store in [false, true] { + let transport = StubInferenceTransport::success(); + let selector = StubSelector::success(); + let provider = provider(Arc::clone(&selector), transport.clone()).await; + let request = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + "openai", + object(json!({ + "model": MODEL, + "input": "hello", + "generate": false, + "store": store + })), + ) + .unwrap(), + )); + let error = match provider + .execute( + provider_request_with_operation("xai", request), + context(CancellationToken::new(), None), + ) + .await + { + Ok(_) => panic!("non-generating requests must fail before a usage-producing stream"), + Err(error) => error, + }; + assert_eq!(error.kind(), ProviderErrorKind::Unsupported); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + let detail = error + .client_visible_upstream_error() + .expect("prewarm rejection must explain the unsupported request"); + assert_eq!(detail.code(), Some("unsupported_prewarm")); + assert_eq!(detail.error_type(), Some("invalid_request_error")); + assert!(detail.message().contains("generate=false")); + assert_eq!(selector.calls.load(Ordering::SeqCst), 0); + assert_eq!(transport.calls.load(Ordering::SeqCst), 0); + assert!(transport.requests.lock().unwrap().is_empty()); + } +} + +#[tokio::test] +async fn native_xai_rejects_translated_generate_false_before_upstream_io() { + let transport = StubInferenceTransport::success(); + let selector = StubSelector::success(); + let provider = provider(Arc::clone(&selector), transport.clone()).await; + let source = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object("example-source", object(json!({"opaque": "source"}))) + .unwrap(), + )); + let translations = Arc::new(Mutex::new(0)); + let error = match provider + .execute( + provider_request_with_operation("xai", source), + context_with_middleware(Arc::new(TranslationMiddleware { + translations: Arc::clone(&translations), + target: object(json!({"model": MODEL, "input": "hello", "generate": false})), + })), + ) + .await + { + Ok(_) => panic!("translated prewarm must not start generation"), + Err(error) => error, + }; + assert_eq!(error.kind(), ProviderErrorKind::Unsupported); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + assert_eq!(*translations.lock().unwrap(), 1); + assert_eq!(selector.calls.load(Ordering::SeqCst), 1); + assert_eq!(transport.calls.load(Ordering::SeqCst), 0); + assert!(transport.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn native_xai_preserves_generate_true_and_omitted_requests() { + for generate in [None, Some(true)] { + let transport = StubInferenceTransport::success(); + let provider = provider(StubSelector::success(), transport.clone()).await; + let mut body = object(json!({"model": MODEL, "input": "hello"})); + if let Some(generate) = generate { + body.insert("generate".to_owned(), json!(generate)); + } + let operation = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object("openai", body).unwrap(), + )); + let mut stream = provider + .execute( + provider_request_with_operation("xai", operation), + context(CancellationToken::new(), None), + ) + .await + .expect("ordinary generation remains supported"); + let mut completed = false; + while let Some(event) = stream.next().await { + completed |= event + .expect("normal upstream response") + .into_parts() + .0 + .iter() + .any(|event| matches!(event, GatewayEvent::Completed(_))); + } + assert!(completed); + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + let requests = transport.requests.lock().unwrap(); + let sent: Value = serde_json::from_slice(requests[0].body()).unwrap(); + assert_eq!(sent.get("generate").and_then(Value::as_bool), generate); + } +} #[tokio::test] async fn native_xai_translates_a_non_native_source_before_encoding() { @@ -81,6 +192,7 @@ async fn native_xai_translates_a_non_native_source_before_encoding() { "example-source", Map::from_iter([ ("opaque".to_owned(), json!("source-only")), + ("generate".to_owned(), json!(false)), ("input".to_owned(), json!([{"type":"compaction_trigger"}])), ("prompt_cache_key".to_owned(), json!("must-not-be-read")), ]), @@ -539,7 +651,7 @@ async fn attempt_middleware_can_change_reasoning_before_native_validation() { let requests = transport.requests.lock().unwrap(); let body: Value = serde_json::from_slice(requests[0].body()).unwrap(); - assert_eq!(body["reasoning"]["effort"], "high"); + assert_eq!(body["reasoning"]["effort"], "xhigh"); assert_eq!(body["model"], MODEL); assert_eq!( observed.lock().unwrap()[0].1["reasoning"]["effort"], @@ -583,36 +695,45 @@ fn observed_transport_metrics() -> GrokInferenceTransportMetrics { const CATALOG_WITHOUT_FEATURE_METADATA: &[u8] = br#"{ "object": "list", - "data": [{ - "id": "grok-catalog-entry", - "model": "grok-catalog-entry", - "contextWindow": 1000000, - "maxCompletionTokens": 131072, - "apiBackend": "responses", - "supportedInApi": true, - "supportsReasoningEffort": true - }] + "data": [ + { + "id": "grok-catalog-entry", + "model": "grok-catalog-entry", + "context_window": 1000000, + "max_completion_tokens": 131072, + "api_backend": "responses", + "supported_in_api": true, + "supports_reasoning_effort": true + } + ] }"#; -const CATALOG_WITH_FEATURE_REASONING_OPTIONS: &[u8] = br#"{ +const CATALOG_WITH_REASONING_MENU: &[u8] = br#"{ "object": "list", - "data": [{ - "id": "grok-4.5", - "model": "grok-4.5", - "contextWindow": 500000, - "maxCompletionTokens": 131072, - "apiBackend": "responses", - "supportedInApi": true, - "reasoningEfforts": ["high", {"value": "medium", "default": true}, "low"], - "features": { - "reasoning": true, - "reasoningEffortOptions": { - "supportedEfforts": ["low", "medium", "high", "xhigh"], - "defaultEffort": "high" - } - }, - "supportsReasoningEffort": true, - "streamToolCalls": true - }] + "data": [ + { + "id": "grok-4.5", + "model": "grok-4.5", + "context_window": 500000, + "max_completion_tokens": 131072, + "api_backend": "responses", + "reasoning_efforts": [ + { + "value": "low" + }, + { + "value": "medium" + }, + { + "value": "high", + "default": true + }, + { + "value": "xhigh" + } + ], + "stream_tool_calls": true + } + ] }"#; const SUCCESS_SSE: &[u8] = concat!( "event: response.created\n", @@ -1076,13 +1197,13 @@ impl GrokModelCatalogTransport for CatalogWithoutFeatureMetadataTransport { } } -struct CatalogWithFeatureReasoningOptionsTransport; +struct CatalogWithReasoningMenuTransport; -impl GrokModelCatalogTransport for CatalogWithFeatureReasoningOptionsTransport { +impl GrokModelCatalogTransport for CatalogWithReasoningMenuTransport { fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> { Box::pin(async { Ok(GrokModelCatalogTransportResponse::new( - CATALOG_WITH_FEATURE_REASONING_OPTIONS, + CATALOG_WITH_REASONING_MENU, None, )) }) @@ -1091,6 +1212,22 @@ impl GrokModelCatalogTransport for CatalogWithFeatureReasoningOptionsTransport { struct UnavailableCatalogTransport; +struct CatalogWithServerReasoningDefaultsTransport; + +impl GrokModelCatalogTransport for CatalogWithServerReasoningDefaultsTransport { + fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> { + Box::pin(async { + Ok(GrokModelCatalogTransportResponse::new( + br#"{"object":"list","data":[ + {"model":"grok-4.7","api_backend":"responses","context_window":256000,"context_windows":[256000,500000],"reasoning_efforts":[{"value":"low"},{"value":"high"}]}, + {"model":"grok-future","api_backend":"responses","reasoning_efforts":[{"value":"low"},{"value":"high"},{"value":"unknown","default":true}]} + ]}"#, + None, + )) + }) + } +} + impl GrokModelCatalogTransport for UnavailableCatalogTransport { fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> { Box::pin(async { @@ -1846,7 +1983,7 @@ fn diagnostic_context(request_id: &str, account_suffix: &str) -> AttemptContext } #[tokio::test] -async fn model_alias_should_be_canonical_before_account_selection() { +async fn account_selection_uses_the_exact_routed_model() { let selector = StubSelector::success(); let provider = provider(selector.clone(), StubInferenceTransport::success()).await; let mut stream = provider @@ -1866,7 +2003,7 @@ async fn model_alias_should_be_canonical_before_account_selection() { .lock() .expect("upstream models") .as_slice(), - ["grok-4.6"] + ["grok-4.6-latest"] ); } @@ -4064,7 +4201,7 @@ async fn provider_compiles_realtime_catalog_capabilities() { } #[tokio::test] -async fn provider_publishes_default_codex_profile_when_catalog_is_unavailable() { +async fn provider_propagates_catalog_failure_without_inventing_model_facts() { let provider = provider_with_catalog_transport( StubSelector::success(), StubInferenceTransport::success(), @@ -4072,32 +4209,21 @@ async fn provider_publishes_default_codex_profile_when_catalog_is_unavailable() Arc::new(UnavailableCatalogTransport), ) .await; - let capabilities = provider + let error = provider .query_model_capabilities() .await - .expect("fallback capabilities"); - - assert_eq!(capabilities.len(), 1); - assert_eq!(capabilities[0].upstream_model().as_str(), MODEL); - let presentation = capabilities[0] - .presentation() - .expect("fallback Grok model presentation"); - assert_eq!(presentation.display_name(), Some("Grok 4.5")); - assert_eq!(presentation.default_reasoning_effort(), None); - assert!(presentation.supported_reasoning_efforts().is_empty()); - assert_eq!(presentation.context_window_tokens(), Some(500_000)); - assert_eq!(presentation.max_context_window_tokens(), Some(500_000)); - assert!(presentation.agent_tools()); - assert!(presentation.parallel_tool_calls()); + .expect_err("catalog is unavailable"); + assert_eq!(error.kind(), ProviderErrorKind::Unavailable); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); } #[tokio::test] -async fn provider_publishes_reasoning_efforts_from_feature_options() { +async fn provider_publishes_reasoning_efforts_from_current_menu() { let provider = provider_with_catalog_transport( StubSelector::success(), StubInferenceTransport::success(), StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable), - Arc::new(CatalogWithFeatureReasoningOptionsTransport), + Arc::new(CatalogWithReasoningMenuTransport), ) .await; let capabilities = provider @@ -4115,6 +4241,41 @@ async fn provider_publishes_reasoning_efforts_from_feature_options() { ); } +#[tokio::test] +async fn provider_keeps_server_reasoning_default_when_catalog_does_not_declare_one() { + let provider = provider_with_catalog_transport( + StubSelector::success(), + StubInferenceTransport::success(), + StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable), + Arc::new(CatalogWithServerReasoningDefaultsTransport), + ) + .await; + let capabilities = provider + .query_model_capabilities() + .await + .expect("capabilities"); + assert_eq!(capabilities.len(), 2); + for model in capabilities { + let presentation = model.presentation().expect("model presentation"); + assert_eq!( + presentation.supported_reasoning_efforts(), + ["low", "high"], + "{}", + model.upstream_model() + ); + assert_eq!( + presentation.default_reasoning_effort(), + None, + "{}", + model.upstream_model() + ); + if model.upstream_model().as_str() == "grok-4.7" { + assert_eq!(presentation.context_window_tokens(), Some(256_000)); + assert_eq!(presentation.max_context_window_tokens(), Some(500_000)); + } + } +} + #[tokio::test] async fn missing_catalog_tool_metadata_keeps_build_tools_routable() { let provider = provider_with_catalog_transport( @@ -4301,6 +4462,20 @@ impl gateway_core::engine::upstream_adapter::UpstreamAdapter for AdapterProbe { #[tokio::test] async fn upstream_adapter_reuses_selected_native_account_inside_attempt_onion_and_stays_cold() { + assert_upstream_adapter_stays_cold(operation()).await; + for protocol in ["openai", "example-source"] { + let operation = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + protocol, + object(json!({"model": MODEL, "input": "hello", "generate": false})), + ) + .unwrap(), + )); + assert_upstream_adapter_stays_cold(operation).await; + } +} + +async fn assert_upstream_adapter_stays_cold(operation: Operation) { let selector = StubSelector::success(); let transport = StubInferenceTransport::success(); let provider = provider(selector.clone(), transport.clone()).await; @@ -4344,7 +4519,7 @@ async fn upstream_adapter_reuses_selected_native_account_inside_attempt_onion_an CancellationToken::new(), ); let mut stream = provider - .execute(provider_request("xai"), context) + .execute(provider_request_with_operation("xai", operation), context) .await .unwrap(); assert_eq!(polls.load(std::sync::atomic::Ordering::SeqCst), 0); diff --git a/backend/crates/providers/xai/tests/provider/contract/replay_tool_names.rs b/backend/crates/providers/xai/tests/provider/contract/replay_tool_names.rs new file mode 100644 index 000000000..574b9bfd1 --- /dev/null +++ b/backend/crates/providers/xai/tests/provider/contract/replay_tool_names.rs @@ -0,0 +1,185 @@ +//! 跨轮重放区分已编码输入与客户端可见输出，避免工具别名缺失或二次映射 + +use provider_xai::GrokResponsesRequest; + +use super::*; + +fn generate(body: Value) -> GenerateRequest { + GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object("openai", body.as_object().expect("body").clone()) + .expect("payload"), + ) +} + +fn function(name: &str, label: &str) -> Value { + json!({"type": "function", "name": name, "description": label, + "parameters": {"type": "object"}}) +} + +fn history(names: [&str; 2], prefix: &str) -> Vec { + names.into_iter().enumerate().flat_map(|(index, name)| { + let call_id = format!("{prefix}_{index}"); + [ + json!({"type": "function_call", "name": name, "call_id": call_id, "arguments": "{}"}), + json!({"type": "function_call_output", "call_id": call_id, "output": "done"}), + ] + }).collect() +} + +fn declared_alias(body: &Value, label: &str) -> String { + body["tools"] + .as_array() + .expect("tools") + .iter() + .find(|tool| tool["description"] == label) + .expect("function")["name"] + .as_str() + .expect("alias") + .to_owned() +} + +async fn assert_replay_tool_names(attempt: Option) { + for hosted in ["web_search", "x_search"] { + for literal_first in [false, true] { + let key = ClientApiKeyId::new("key_xai_contract").expect("key"); + let isolated = GrokResponsesRequest::encode( + &generate(json!({"tools": [function(hosted, "reserved")]})), + MODEL, + &key, + ) + .expect("isolated function"); + let literal = isolated.body()["tools"][0]["name"] + .as_str() + .expect("generated alias"); + let mut tools = vec![function(hosted, "reserved"), function(literal, "literal")]; + if literal_first { + tools.reverse(); + } + if hosted == "web_search" { + tools.push(json!({"type": "web_search"})); + } + let first_body = json!({"model": "client-model", "store": false, + "prompt_cache_key": "replay_tool_names", "tools": tools, + "input": history([hosted, literal], "prior")}); + let encoded = GrokResponsesRequest::encode(&generate(first_body.clone()), MODEL, &key) + .expect("first request"); + let encoded_body = Value::Object(encoded.body().clone()); + let aliases = [ + declared_alias(&encoded_body, "reserved"), + declared_alias(&encoded_body, "literal"), + ]; + assert_ne!(aliases[0], aliases[1]); + let output: Vec<_> = aliases + .iter() + .enumerate() + .map(|(index, name)| { + json!({"type": "function_call", "id": format!("fc_returned_{index}"), + "name": name, "call_id": format!("returned_{index}"), "arguments": "{}"}) + }) + .collect(); + let response = [ + json!({"type": "response.created", "response": {"id": "resp_tool_names"}}), + json!({"type": "response.output_item.done", "output_index": 0, "item": output[0]}), + json!({"type": "response.output_item.done", "output_index": 1, "item": output[1]}), + json!({"type": "response.completed", "response": {"id": "resp_tool_names", + "status": "completed", "output": output}}), + ] + .into_iter() + .map(|event| format!("data: {event}\n\n")) + .collect::(); + let transport = StubInferenceTransport::sequence([ + InferenceMode::SuccessBody(response.into_bytes()), + InferenceMode::Success, + ]); + let provider = provider(StubSelector::success(), transport.clone()).await; + let mut first = Arc::clone(&provider) + .execute( + provider_request_with_operation( + "xai", + Operation::Generate(generate(first_body)), + ), + context(CancellationToken::new(), None), + ) + .await + .expect("first stream"); + let state = collect_provider_state(&mut first) + .await + .expect("session state"); + + let mut second_input = history([hosted, literal], "current"); + for index in 0..2 { + second_input.push(json!({"type": "function_call_output", + "call_id": format!("returned_{index}"), "output": "returned result"})); + } + let mut second_body = json!({"model": "client-model", "store": false, + "prompt_cache_key": "replay_tool_names", "tools": tools, "input": second_input}); + let (operation, context) = if let Some(attempt) = attempt { + second_body["previous_response_id"] = json!("resp_gateway_tool_names"); + let pin = NativeContinuationPin::new( + PreviousResponseId::new("resp_gateway_tool_names"), + PreviousResponseId::new("resp_tool_names"), + key, + ProviderKind::new("xai").expect("provider"), + account_id("provider"), + ); + ( + operation_with_state(second_body, state), + context_with_continuation_attempt(ContinuationBinding::Pinned(pin), attempt), + ) + } else { + ( + Operation::Generate(generate(second_body)), + context(CancellationToken::new(), None), + ) + }; + let mut second = provider + .execute(provider_request_with_operation("xai", operation), context) + .await + .expect("second stream"); + while let Some(event) = second.next().await { + event.expect("second event"); + } + let requests = transport.requests.lock().expect("requests"); + let replay: Value = serde_json::from_slice(requests[1].body()).expect("replay body"); + assert!(replay.get("previous_response_id").is_none()); + for (index, expected) in aliases.iter().enumerate() { + assert_eq!( + declared_alias(&replay, if index == 0 { "reserved" } else { "literal" }), + *expected + ); + for prefix in ["prior", "returned", "current"] { + if attempt.is_none() && prefix == "prior" { + continue; + } + let call_id = format!("{prefix}_{index}"); + let call = replay["input"] + .as_array() + .expect("input") + .iter() + .find(|item| item["type"] == "function_call" && item["call_id"] == call_id) + .expect("replayed call"); + assert_eq!( + call["name"], *expected, + "{hosted}, literal_first={literal_first}, {call_id}" + ); + assert_eq!(call["arguments"], "{}"); + } + } + } + } +} + +#[tokio::test] +async fn replay_tool_names_reencode_outputs_on_store_false_native_continuation() { + assert_replay_tool_names(Some(ContinuationAttempt::Native)).await; +} + +#[tokio::test] +async fn replay_tool_names_reencode_outputs_on_replay_owner() { + assert_replay_tool_names(Some(ContinuationAttempt::ReplayOwner)).await; +} + +#[tokio::test] +async fn replay_tool_names_reencode_cached_reasoning_calls_without_reencoding_input() { + assert_replay_tool_names(None).await; +} diff --git a/backend/crates/providers/xai/tests/transport/canonical.rs b/backend/crates/providers/xai/tests/transport/canonical.rs index 5726beef1..51dac36ca 100644 --- a/backend/crates/providers/xai/tests/transport/canonical.rs +++ b/backend/crates/providers/xai/tests/transport/canonical.rs @@ -19,8 +19,9 @@ fn decoder_should_bill_the_sent_model_independently_of_the_response_model() { ("grok-4.5", Some("grok-4.3"), None, Some(2_175_000)), ("grok-4.5", None, None, Some(2_175_000)), ("grok-future", Some("grok-4.5"), None, None), + ("grok-latest", Some("grok-4.3"), Some(0), None), ("grok-4.5", Some("grok-future"), Some(123), None), - ("grok-4.5", Some("grok-future"), Some(0), None), + ("grok-4.5", Some("grok-future"), Some(0), Some(2_175_000)), ] { let created = serde_json::json!({"type":"response.created","response":{"id":"resp_model_cost","model":returned}}); let mut completed = serde_json::json!({ @@ -45,7 +46,10 @@ fn decoder_should_bill_the_sent_model_independently_of_the_response_model() { expected_cost, "{sent} -> {returned:?}" ); - assert_eq!(provider_cost_ticks(&facts), reported_cost); + assert_eq!( + provider_cost_ticks(&facts), + reported_cost.filter(|ticks| *ticks > 0) + ); assert_eq!(decoder.response_model(), returned); assert!(facts.iter().any(|event| matches!( event, @@ -72,6 +76,24 @@ fn terminal_cost_events( cached_tokens: u64, output_tokens: u64, provider_cost_ticks: Option, +) -> Vec { + terminal_cost_events_with_pricing( + model, + input_tokens, + cached_tokens, + output_tokens, + provider_cost_ticks, + None, + ) +} + +fn terminal_cost_events_with_pricing( + model: &str, + input_tokens: u64, + cached_tokens: u64, + output_tokens: u64, + provider_cost_ticks: Option, + pricing: Option, ) -> Vec { let provider_cost = provider_cost_ticks .map(|ticks| format!(",\"cost_in_usd_ticks\":{ticks}")) @@ -89,7 +111,13 @@ fn terminal_cost_events( cached_tokens = cached_tokens, provider_cost = provider_cost, ); - decode_canonical(body.as_bytes(), model).expect("canonical cost response") + GrokCanonicalDecoder::new(model) + .with_pricing(pricing) + .push(body.as_bytes()) + .expect("canonical cost response") + .into_iter() + .flat_map(|event| event.into_parts().0) + .collect() } fn decode_canonical(body: &[u8], upstream_model: &str) -> Result, ProviderError> { @@ -271,11 +299,16 @@ fn decoder_should_price_official_grok_45_build_free_variant() { } #[test] -fn zero_provider_cost_should_remain_an_authoritative_zero() { - let events = terminal_cost_events("grok-4.5", 1, 0, 1, Some(0)); +fn zero_provider_cost_should_use_available_token_pricing() { + for (model, calculated) in [("grok-4.5", Some(80_000)), ("grok-future", None)] { + let events = terminal_cost_events(model, 1, 0, 1, Some(0)); - assert_eq!(calculated_cost_ticks(&events), None); - assert_eq!(provider_cost_ticks(&events), Some(0)); + assert_eq!( + (calculated_cost_ticks(&events), provider_cost_ticks(&events)), + (calculated, None), + "{model}" + ); + } } #[test] @@ -401,7 +434,7 @@ fn billing_should_use_the_final_response_tier_and_preserve_provider_cost() { ("priority", None, Some(400_000_000), None), ("default", None, Some(200_000_000), None), ("future", None, None, None), - ("priority", Some(0), None, Some(0)), + ("priority", Some(0), Some(400_000_000), None), ("priority", Some(123), None, Some(123)), ] { let mut response = serde_json::json!({"service_tier": tier, "output": [], "usage": {"input_tokens": 10_000, "output_tokens": 0}}); @@ -427,12 +460,18 @@ fn billing_should_require_provider_totals_for_server_tools() { "image_generation_call", "future_tool_call", ] { - let response = serde_json::json!({"output": [{"type": kind, "id": "tool_1", "status": "completed"}], "usage": {"input_tokens": 100, "output_tokens": 1}}); - assert_eq!( - calculated_cost_ticks(&pricing_events(response, None)), - None, - "{kind}" - ); + for ticks in [None, Some(0)] { + let mut response = serde_json::json!({"output": [{"type": kind, "id": "tool_1", "status": "completed"}], "usage": {"input_tokens": 100, "output_tokens": 1}}); + if let Some(ticks) = ticks { + response["usage"]["cost_in_usd_ticks"] = serde_json::json!(ticks); + } + let events = pricing_events(response, None); + assert_eq!( + (calculated_cost_ticks(&events), provider_cost_ticks(&events)), + (None, None), + "{kind}: {ticks:?}" + ); + } } let request = tool_request(serde_json::json!({"input": "查资料", "tools": [{"type": "web_search"}]})); @@ -444,9 +483,10 @@ fn billing_should_require_provider_totals_for_server_tools() { None ); response["usage"]["cost_in_usd_ticks"] = serde_json::json!(0); + let events = pricing_events(response, Some(&request)); assert_eq!( - provider_cost_ticks(&pricing_events(response, Some(&request))), - Some(0) + (calculated_cost_ticks(&events), provider_cost_ticks(&events)), + (None, None) ); let response = serde_json::json!({"output": [], "usage": {"input_tokens": 100, "output_tokens": 1, "num_server_side_tool_calls": 1}}); assert_eq!(calculated_cost_ticks(&pricing_events(response, None)), None); @@ -495,6 +535,78 @@ fn decoder_should_leave_unpublished_model_pricing_unavailable() { assert_eq!(provider_cost_ticks(&events), None); } +#[test] +fn decoder_should_not_price_dynamic_or_compatibility_aliases_as_stable_models() { + for model in [ + "grok", + "grok-latest", + "grok-build-latest", + "grok-4.6-latest", + "grok-4.5-latest", + "grok-4.3-latest", + "grok-code-fast", + "grok-4.20-multi-agent-latest", + "grok-4.20-beta", + "grok-4.20-beta-reasoning", + "grok-4.20-beta-non-reasoning", + "grok-4.20-reasoning-gv2", + "grok-4.20-non-reasoning-gv2", + "xai/grok-4.5", + "grok-4.7", + ] { + for reported in [None, Some(0), Some(123)] { + let events = terminal_cost_events(model, 100, 25, 10, reported); + assert_eq!( + calculated_cost_ticks(&events), + None, + "{model}, reported: {reported:?}" + ); + assert_eq!( + provider_cost_ticks(&events), + reported.filter(|ticks| *ticks > 0).map(u128::from) + ); + assert!( + events + .iter() + .any(|event| matches!(event, GatewayEvent::Completed(_))) + ); + } + } +} + +#[test] +fn explicit_alias_pricing_still_applies_when_provider_cost_is_unreported() { + let pricing: gateway_core::metering::ModelPriceOverride = + serde_json::from_value(serde_json::json!({ + "multiplierBps":15000,"bands":{"standard":{ + "input":"3", "output":"12", "cacheRead":"0", "cacheWrite":"0" + }} + })) + .expect("manual price"); + for reported in [None, Some(0), Some(123)] { + let events = terminal_cost_events_with_pricing( + "grok-latest", + 100, + 20, + 10, + reported, + Some(pricing.clone()), + ); + assert_eq!( + calculated_cost_ticks(&events), + if reported == Some(123) { + None + } else { + Some(5_400_000) + } + ); + assert_eq!( + provider_cost_ticks(&events), + reported.filter(|ticks| *ticks > 0).map(u128::from) + ); + } +} + #[test] fn decoder_should_leave_incomplete_usage_pricing_unavailable() { let body = concat!( @@ -1410,20 +1522,38 @@ fn decoder_should_require_terminal_response() { } #[test] -fn decoder_should_preserve_incomplete_length_reason() { - let body = concat!( - "event: response.created\n", - "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_short\"}}\n\n", - "event: response.incomplete\n", - "data: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"resp_short\",\"status\":\"incomplete\",\"incomplete_details\":{\"reason\":\"max_output_tokens\"}}}\n\n", - ); - let events = decode_canonical(body.as_bytes(), "fallback").expect("incomplete response"); +fn decoder_should_classify_incomplete_reasons_and_preserve_native_wire() { + for (reason, expected) in [ + ("max_output_tokens", FinishReason::Length), + ("max_tokens", FinishReason::Length), + ("max_prompt_tokens", FinishReason::Length), + ("max_time_limit", FinishReason::Length), + ("content_filter", FinishReason::ContentFilter), + ("future_reason", FinishReason::Other), + ] { + let created = serde_json::json!({"type":"response.created","response":{"id":"resp_short"}}); + let incomplete = serde_json::json!({ + "type":"response.incomplete", + "response":{ + "id":"resp_short","status":"incomplete", + "incomplete_details":{"reason":reason} + } + }); + let body = format!("data: {created}\n\ndata: {incomplete}\n\n"); + let events = GrokCanonicalDecoder::new("fallback") + .push(body.as_bytes()) + .expect("incomplete response"); + let completed = events + .iter() + .flat_map(ProviderEvent::canonical_facts) + .find_map(|fact| match fact { + GatewayEvent::Completed(meta) => meta.finish_reason(), + _ => None, + }); - assert!(matches!( - events.last(), - Some(GatewayEvent::Completed(meta)) - if meta.finish_reason() == Some(FinishReason::Length) - )); + assert_eq!(completed, Some(expected), "{reason}"); + assert_eq!(wire_events(&events).last().unwrap().data(), &incomplete); + } } #[test] diff --git a/backend/crates/providers/xai/tests/transport/catalog.rs b/backend/crates/providers/xai/tests/transport/catalog.rs index 4940717c8..bdebe5bb4 100644 --- a/backend/crates/providers/xai/tests/transport/catalog.rs +++ b/backend/crates/providers/xai/tests/transport/catalog.rs @@ -13,8 +13,7 @@ use provider_xai::{ MAX_GROK_MODEL_CATALOG_BYTES, SecretValue, parse_grok_billing, parse_grok_model_catalog, }; -const OFFICIAL_FIXTURE: &[u8] = - include_bytes!("catalog/fixtures/official_grok_models_snapshot.json"); +const CLI_PROXY_FIXTURE: &[u8] = include_bytes!("catalog/fixtures/cli_proxy_models.json"); struct CapturingTransport { calls: AtomicUsize, @@ -80,7 +79,7 @@ impl GrokBillingTransport for CapturingBillingTransport { #[tokio::test] async fn client_should_send_exact_oauth_headers_without_api_key() { let transport = Arc::new(CapturingTransport::success( - OFFICIAL_FIXTURE, + CLI_PROXY_FIXTURE, Some("\"grok-v1\""), )); let client = GrokModelCatalogClient::new(transport.clone()); @@ -128,7 +127,7 @@ async fn client_should_send_exact_oauth_headers_without_api_key() { #[tokio::test] async fn client_should_omit_email_when_verified_profile_has_none() { - let transport = Arc::new(CapturingTransport::success(OFFICIAL_FIXTURE, None)); + let transport = Arc::new(CapturingTransport::success(CLI_PROXY_FIXTURE, None)); let client = GrokModelCatalogClient::new(transport.clone()); client .fetch(&session(None)) @@ -262,10 +261,13 @@ async fn subscription_query_rejects_malformed_or_oversized_user_responses() { } #[test] -fn billing_parser_should_accept_credits_and_legacy_shapes() { +fn billing_parser_should_accept_current_credits_fields() { for body in [ br#"{"config":{"creditUsagePercent":31.25,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-13T00:00:00Z","end":"2026-07-20T00:00:00Z"},"prepaidBalance":{"val":2500}}}"#.as_slice(), - br#"{"config":{"monthlyLimit":{"val":2000},"used":{"val":500},"onDemandCap":{"val":1000},"onDemandUsed":{"val":100}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":-500}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":-9223372036854775808}}}"#.as_slice(), + br#"{"config":{"onDemandCap":{},"onDemandUsed":{},"prepaidBalance":{}}}"#.as_slice(), + br#"{"config":{"onDemandCap":{"val":1000},"onDemandUsed":{"val":100}}}"#.as_slice(), br#"{"config":null}"#.as_slice(), ] { parse_grok_billing(body).expect("supported official billing shape"); @@ -278,7 +280,13 @@ fn billing_parser_should_reject_invalid_known_fields() { br#"[]"#.as_slice(), br#"{"config":[]}"#.as_slice(), br#"{"config":{"creditUsagePercent":101}}"#.as_slice(), - br#"{"config":{"used":{"val":-1}}}"#.as_slice(), + br#"{"config":{"onDemandCap":{"val":-1}}}"#.as_slice(), + br#"{"config":{"onDemandUsed":{"val":-1}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":null}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":"-500"}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":-0.5}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":9223372036854775808}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":-9223372036854775809}}}"#.as_slice(), br#"{"config":{"currentPeriod":"weekly"}}"#.as_slice(), br#"{"onDemandEnabled":"yes"}"#.as_slice(), ] { @@ -309,9 +317,9 @@ fn billing_snapshot_debug_should_not_print_values() { } #[test] -fn official_fixture_should_use_actual_model_and_whitelisted_metadata() { - let snapshot = parse_grok_model_catalog(OFFICIAL_FIXTURE, Some("W/\"grok-v1\"")) - .expect("official fixture should parse"); +fn cli_proxy_fixture_should_use_actual_model_and_whitelisted_metadata() { + let snapshot = parse_grok_model_catalog(CLI_PROXY_FIXTURE, Some("W/\"grok-v1\"")) + .expect("CLI proxy fixture should parse"); let model = &snapshot.models()[0]; assert_eq!( @@ -370,7 +378,7 @@ fn official_fixture_should_use_actual_model_and_whitelisted_metadata() { #[test] fn non_whitelisted_wire_fields_should_not_survive_normalization() { let snapshot = - parse_grok_model_catalog(OFFICIAL_FIXTURE, None).expect("official fixture should parse"); + parse_grok_model_catalog(CLI_PROXY_FIXTURE, None).expect("CLI proxy fixture should parse"); let debug = format!("{snapshot:?}"); assert!( @@ -382,16 +390,18 @@ fn non_whitelisted_wire_fields_should_not_survive_normalization() { #[test] fn invalid_etag_should_fail_the_entire_snapshot() { - let result = parse_grok_model_catalog(OFFICIAL_FIXTURE, Some("raw-unquoted-etag")); + let result = parse_grok_model_catalog(CLI_PROXY_FIXTURE, Some("raw-unquoted-etag")); assert!(matches!(result, Err(GrokModelCatalogError::InvalidEtag))); } #[test] fn missing_capability_fields_should_remain_unknown() { - let snapshot = - parse_grok_model_catalog(br#"{"object":"list","data":[{"id":"grok-unknown"}]}"#, None) - .expect("identity-only official entry should parse"); + let snapshot = parse_grok_model_catalog( + br#"{"object":"list","data":[{"id":"grok-unknown","model":"grok-unknown"}]}"#, + None, + ) + .expect("identity-only official entry should parse"); let model = &snapshot.models()[0]; assert_eq!( @@ -417,7 +427,7 @@ fn missing_capability_fields_should_remain_unknown() { #[test] fn reasoning_effort_menu_should_be_capability_evidence_without_legacy_flag() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-reasoning","reasoningEfforts":[{"value":"xhigh","default":true},"low"]}]}"#, + br#"{"object":"list","data":[{"id":"grok-reasoning","reasoning_efforts":[{"value":"xhigh","default":true},{"value":"low"}],"model":"grok-reasoning"}]}"#, None, ) .expect("reasoning effort menu"); @@ -444,38 +454,30 @@ fn reasoning_effort_menu_should_be_capability_evidence_without_legacy_flag() { } #[test] -fn reasoning_effort_options_should_read_official_features_shape() { +fn historical_and_other_endpoint_fields_do_not_supply_model_facts() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-4.5","reasoningEfforts":["high"],"features":{"reasoning":true,"reasoningEffortOptions":{"supportedEfforts":["low","medium","high","xhigh"],"defaultEffort":"high"}}}]}"#, + br#"{"object":"list","data":[{"model":"grok-future","contextWindow":500000,"contextWindows":[1000000],"apiBackend":"responses","reasoningEfforts":[{"value":"high","default":true}],"features":{"reasoning":true,"reasoningEffortOptions":{"supportedEfforts":["high"],"defaultEffort":"high"}},"_meta":{"contextWindow":500000,"reasoningEfforts":[{"value":"high","default":true}]},"capabilities":{"reasoning_effort":["high"],"default_reasoning_effort":"high"},"reasoning_effort":"high"}]}"#, None, - ) - .expect("official feature reasoning options should parse"); - let capabilities = snapshot.models()[0].capabilities(); - + ).expect("unknown fields do not define CLI proxy facts"); + let model = &snapshot.models()[0]; assert_eq!( - ( - capabilities.reasoning_effort(), - capabilities - .reasoning_efforts() - .iter() - .map(|effort| effort.as_str()) - .collect::>(), - capabilities - .default_reasoning_effort() - .map(|effort| effort.as_str()), - ), - ( - GrokCatalogCapabilityEvidence::DeclaredNative, - vec!["low", "medium", "high", "xhigh"], - Some("high"), - ) + model.capabilities().reasoning_effort(), + GrokCatalogCapabilityEvidence::Unknown + ); + assert!(model.capabilities().reasoning_efforts().is_empty()); + assert_eq!(model.capabilities().default_reasoning_effort(), None); + assert_eq!(model.limits().context_window_tokens(), None); + assert_eq!(model.limits().max_context_window_tokens(), None); + assert_eq!( + model.capabilities().responses_api(), + GrokCatalogCapabilityEvidence::Unknown ); } #[test] fn unknown_reasoning_effort_values_should_not_fail_the_snapshot() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-reasoning","reasoningEffort":"ultra","reasoningEfforts":[{"value":"ultra","default":true},"low","hyper"]}]}"#, + br#"{"object":"list","data":[{"id":"grok-reasoning","reasoning_effort":"ultra","reasoning_efforts":[{"value":"ultra","default":true},{"value":"low"},{"value":"hyper"}],"model":"grok-reasoning"}]}"#, None, ) .expect("unknown effort values must not break catalog parsing"); @@ -492,25 +494,143 @@ fn unknown_reasoning_effort_values_should_not_fail_the_snapshot() { .default_reasoning_effort() .map(|effort| effort.as_str()), ), - (vec!["low"], Some("low")) + (vec!["low"], None) ); } #[test] -fn model_id_should_take_priority_over_catalog_id() { +fn reasoning_menu_uses_only_known_explicit_default() { + for (menu, expected) in [ + ( + serde_json::json!([{"value":"low"},{"value":"high","default":true}]), + Some("high"), + ), + (serde_json::json!([{"value":"low"},{"value":"high"}]), None), + ( + serde_json::json!([{"value":"low"},{"value":"quantum","default":true}]), + None, + ), + ] { + let body = serde_json::to_vec(&serde_json::json!({ + "object":"list", "data":[{"model":"grok-4.7","reasoning_efforts":menu}] + })) + .expect("fixture"); + let snapshot = parse_grok_model_catalog(&body, None).expect("current menu"); + assert_eq!( + snapshot.models()[0] + .capabilities() + .default_reasoning_effort() + .map(|effort| effort.as_str()), + expected + ); + } +} + +#[test] +fn current_reasoning_menu_does_not_borrow_defaults_from_other_sources() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"catalog-entry","modelId":"grok-4-fast"}]}"#, + br#"{"object":"list","data":[{"model":"grok-future","reasoning_efforts":[{"value":"low"}],"reasoning_effort":"low","capabilities":{"reasoning_effort":["high"],"default_reasoning_effort":"high"}}]}"#, + None, + ).expect("current menu"); + let capabilities = snapshot.models()[0].capabilities(); + assert_eq!( + capabilities + .reasoning_efforts() + .iter() + .map(|effort| effort.as_str()) + .collect::>(), + ["low"] + ); + assert_eq!(capabilities.default_reasoning_effort(), None); +} + +#[test] +fn unknown_reasoning_menu_does_not_fall_back_to_other_sources() { + let snapshot = parse_grok_model_catalog( + br#"{"object":"list","data":[{"model":"grok-future","reasoning_efforts":[{"value":"quantum","default":true}],"capabilities":{"reasoning_effort":["high"],"default_reasoning_effort":"high"}}]}"#, + None, + ).expect("unknown effort"); + let capabilities = snapshot.models()[0].capabilities(); + assert!(capabilities.reasoning_efforts().is_empty()); + assert_eq!(capabilities.default_reasoning_effort(), None); +} + +#[test] +fn context_window_choices_supply_default_when_scalar_is_missing() { + let snapshot = parse_grok_model_catalog( + br#"{"object":"list","data":[{"model":"grok-4.7","context_windows":[500000,256000]}]}"#, None, ) - .expect("modelId is official fallback"); + .expect("window choices"); + assert_eq!( + snapshot.models()[0] + .limits() + .context_window_tokens() + .map(std::num::NonZeroU64::get), + Some(500000) + ); +} - assert_eq!(snapshot.models()[0].request_model().as_str(), "grok-4-fast"); +#[test] +fn context_window_choices_keep_scalar_default_and_publish_largest_supported_window() { + for (scalar, windows, maximum) in [ + (256000, vec![500000, 256000], 500000), + (1000000, vec![256000, 500000], 1000000), + ] { + let body = serde_json::to_vec(&serde_json::json!({"object":"list","data":[{"model":"grok-4.7","context_window":scalar,"context_windows":windows}]})).expect("fixture"); + let snapshot = parse_grok_model_catalog(&body, None).expect("window choices"); + let limits = snapshot.models()[0].limits(); + assert_eq!( + limits + .context_window_tokens() + .map(std::num::NonZeroU64::get), + Some(scalar) + ); + assert_eq!( + limits + .max_context_window_tokens() + .map(std::num::NonZeroU64::get), + Some(maximum) + ); + } +} + +#[test] +fn malformed_context_window_choices_are_rejected() { + for windows in [ + serde_json::json!([256000, "big"]), + serde_json::json!([0]), + serde_json::json!(500000), + serde_json::json!([-1, 500000]), + serde_json::Value::Null, + ] { + let body = serde_json::to_vec(&serde_json::json!({"object":"list","data":[{"model":"grok-future","context_window":300000,"context_windows":windows}]})).expect("fixture"); + assert!(matches!( + parse_grok_model_catalog(&body, None), + Err(GrokModelCatalogError::InvalidWire | GrokModelCatalogError::InvalidLimits) + )); + } +} + +#[test] +fn current_proxy_requires_the_explicit_request_model() { + for entry in [ + serde_json::json!({"id":"grok-4.7"}), + serde_json::json!({"modelId":"grok-4.7"}), + ] { + let body = serde_json::to_vec(&serde_json::json!({"object":"list","data":[entry]})) + .expect("fixture"); + assert!(matches!( + parse_grok_model_catalog(&body, None), + Err(GrokModelCatalogError::InvalidWire) + )); + } } #[test] fn official_responses_backend_should_be_native_without_redundant_supported_flag() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-responses","api_backend":"responses"}]}"#, + br#"{"object":"list","data":[{"id":"grok-responses","api_backend":"responses","model":"grok-responses"}]}"#, None, ) .expect("Responses backend is explicit capability evidence"); @@ -524,8 +644,8 @@ fn official_responses_backend_should_be_native_without_redundant_supported_flag( #[test] fn explicit_api_disable_and_non_responses_backend_should_remain_unsupported() { for body in [ - br#"{"object":"list","data":[{"id":"grok-disabled","api_backend":"responses","supported_in_api":false}]}"#.as_slice(), - br#"{"object":"list","data":[{"id":"grok-chat","api_backend":"chat_completions","supported_in_api":true}]}"#.as_slice(), + br#"{"object":"list","data":[{"id":"grok-disabled","api_backend":"responses","supported_in_api":false,"model":"grok-disabled"}]}"#.as_slice(), + br#"{"object":"list","data":[{"id":"grok-chat","api_backend":"chat_completions","supported_in_api":true,"model":"grok-chat"}]}"#.as_slice(), ] { let snapshot = parse_grok_model_catalog(body, None).expect("valid unsupported entry"); assert_eq!( @@ -538,8 +658,8 @@ fn explicit_api_disable_and_non_responses_backend_should_remain_unsupported() { #[test] fn list_discriminator_should_be_required_and_exact() { for body in [ - br#"{"data":[{"id":"grok-4"}]}"#.as_slice(), - br#"{"object":"collection","data":[{"id":"grok-4"}]}"#.as_slice(), + br#"{"data":[{"id":"grok-4","model":"grok-4"}]}"#.as_slice(), + br#"{"object":"collection","data":[{"id":"grok-4","model":"grok-4"}]}"#.as_slice(), ] { assert!(matches!( parse_grok_model_catalog(body, None), @@ -565,7 +685,7 @@ fn empty_data_should_fail_the_entire_snapshot() { #[test] fn duplicate_actual_models_should_fail_the_entire_snapshot() { let result = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"entry-a","model":"grok-4"},{"id":"entry-b","modelId":"grok-4"}]}"#, + br#"{"object":"list","data":[{"id":"entry-a","model":"grok-4"},{"id":"entry-b","model":"grok-4"}]}"#, None, ); @@ -578,7 +698,7 @@ fn duplicate_actual_models_should_fail_the_entire_snapshot() { #[test] fn unknown_top_level_pagination_fields_should_not_reject_the_snapshot() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-4"}],"has_more":true,"cursor":"next"}"#, + br#"{"object":"list","data":[{"id":"grok-4","model":"grok-4"}],"has_more":true,"cursor":"next"}"#, None, ) .expect("unknown pagination fields are not part of the model contract"); @@ -589,7 +709,7 @@ fn unknown_top_level_pagination_fields_should_not_reject_the_snapshot() { #[test] fn unknown_api_backend_should_degrade_to_unknown_without_rejecting_the_snapshot() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","future_top_level":true,"data":[{"id":"grok-future","apiBackend":"future_backend","future_field":{"keep":true}}]}"#, + br#"{"object":"list","future_top_level":true,"data":[{"id":"grok-future","future_field":{"keep":true},"api_backend":"future_backend","model":"grok-future"}]}"#, None, ) .expect("future backend should not reject the snapshot"); @@ -605,7 +725,7 @@ fn unknown_api_backend_should_degrade_to_unknown_without_rejecting_the_snapshot( #[test] fn invalid_preferred_model_should_fail_without_falling_back() { let result = parse_grok_model_catalog( - br#"{"object":"list","data":[{"model":"https://evil.invalid/model","modelId":"grok-4","id":"entry"}]}"#, + br#"{"object":"list","data":[{"model":"https://evil.invalid/model","id":"entry"}]}"#, None, ); @@ -686,7 +806,7 @@ fn header_value<'a>(headers: &'a [(String, &str)], name: &str) -> Option<&'a str #[tokio::test] async fn captured_sensitive_headers_should_remain_typed_as_sensitive() { let session = session(Some("person@example.com")); - let transport = Arc::new(CapturingTransport::success(OFFICIAL_FIXTURE, None)); + let transport = Arc::new(CapturingTransport::success(CLI_PROXY_FIXTURE, None)); let client = GrokModelCatalogClient::new(transport.clone()); client.fetch(&session).await.expect("fetch fixture"); let request = transport.request.lock().expect("captured request"); diff --git a/backend/crates/providers/xai/tests/transport/catalog/fixtures/cli_proxy_models.json b/backend/crates/providers/xai/tests/transport/catalog/fixtures/cli_proxy_models.json new file mode 100644 index 000000000..92955a506 --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/catalog/fixtures/cli_proxy_models.json @@ -0,0 +1,42 @@ +{ + "object": "list", + "data": [ + { + "id": "grok-4.5-catalog-entry", + "model": "grok-4.5", + "name": "Grok 4.5", + "description": "Official Grok Build coding model.", + "hidden": false, + "apiKey": "provider-only field must not enter catalog metadata", + "extraHeaders": { + "x-api-key": "provider-only field" + }, + "baseUrl": "https://provider-only.invalid/v1", + "context_window": 1000000, + "max_completion_tokens": 131072, + "api_backend": "responses", + "supported_in_api": true, + "supports_reasoning_effort": true, + "reasoning_efforts": [ + { + "value": "low" + }, + { + "value": "medium", + "label": "Balanced", + "default": true + }, + { + "value": "high" + }, + { + "id": "extended", + "value": "xhigh", + "label": "Extended" + } + ], + "supports_backend_search": true, + "stream_tool_calls": true + } + ] +} diff --git a/backend/crates/providers/xai/tests/transport/catalog/fixtures/official_grok_models_snapshot.json b/backend/crates/providers/xai/tests/transport/catalog/fixtures/official_grok_models_snapshot.json deleted file mode 100644 index fd8111699..000000000 --- a/backend/crates/providers/xai/tests/transport/catalog/fixtures/official_grok_models_snapshot.json +++ /dev/null @@ -1,32 +0,0 @@ -{ - "object": "list", - "data": [ - { - "id": "grok-4.5-catalog-entry", - "model": "grok-4.5", - "modelId": "grok-4.5-fallback", - "name": "Grok 4.5", - "description": "Official Grok Build coding model.", - "contextWindow": 1000000, - "maxCompletionTokens": 131072, - "apiBackend": "responses", - "supportedInApi": true, - "supportsReasoningEffort": true, - "reasoningEffort": "medium", - "reasoningEfforts": [ - "low", - { "value": "medium", "label": "Balanced", "default": true }, - "high", - { "id": "extended", "value": "xhigh", "label": "Extended" } - ], - "supportsBackendSearch": true, - "streamToolCalls": true, - "hidden": false, - "apiKey": "provider-only field must not enter catalog metadata", - "extraHeaders": { - "x-api-key": "provider-only field" - }, - "baseUrl": "https://provider-only.invalid/v1" - } - ] -} diff --git a/backend/crates/providers/xai/tests/transport/request/hosted_names.rs b/backend/crates/providers/xai/tests/transport/request/hosted_names.rs new file mode 100644 index 000000000..415bdda4a --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/request/hosted_names.rs @@ -0,0 +1,243 @@ +//! Hosted tool 与客户端函数同名时，验证声明、历史和响应使用同一别名合同 + +use gateway_core::event::ProviderEvent; +use provider_xai::GrokCanonicalDecoder; + +use super::*; + +fn function_tool(name: &str) -> Value { + json!({"type": "function", "name": name, "parameters": { + "type": "object", "properties": {"query": {"type": "string"}} + }}) +} + +fn function_history(name: &str) -> Value { + json!([ + {"type": "function_call", "id": "fc_previous", "call_id": "call_previous", + "name": name, "arguments": "{\"query\":\"previous\"}"}, + {"type": "function_call_output", "call_id": "call_previous", "output": "result"} + ]) +} + +fn assert_function_response_restored(request: &GrokResponsesRequest, alias: &str, name: &str) { + let item = json!({"type": "function_call", "id": "fc_next", "call_id": "call_next", + "name": alias, "arguments": "{\"query\":\"next\"}"}); + let frames = [ + json!({"type": "response.created", "response": {"id": "resp_names"}}), + json!({"type": "response.output_item.added", "output_index": 0, "item": item}), + json!({"type": "response.function_call_arguments.delta", "output_index": 0, + "item_id": "fc_next", "call_id": "call_next", "delta": "{\"query\":\"next\"}"}), + json!({"type": "response.output_item.done", "output_index": 0, "item": item}), + json!({"type": "response.completed", "response": {"id": "resp_names", "status": "completed", + "tools": request.body()["tools"], "output": [item]}}), + ] + .into_iter() + .map(|frame| format!("data: {frame}\n\n")) + .collect::(); + let events = GrokCanonicalDecoder::for_request("grok-4.7", request) + .push(frames.as_bytes()) + .expect("function response"); + let wire: Vec<_> = events + .iter() + .filter_map(ProviderEvent::wire_event) + .collect(); + let calls: Vec<_> = wire + .iter() + .filter_map(|event| event.data().get("item")) + .collect(); + assert_eq!( + calls.len(), + 2, + "client functions must not be filtered as cache calls" + ); + for call in calls { + assert_eq!(call["name"], name); + assert_eq!(call["type"], "function_call"); + assert_eq!(call["call_id"], "call_next"); + assert_eq!(call["arguments"], "{\"query\":\"next\"}"); + } + let completed = wire.last().expect("completed response").data(); + assert_eq!(completed["response"]["output"][0]["name"], name); + assert!( + completed["response"]["tools"] + .as_array() + .expect("visible tools") + .iter() + .any(|tool| tool["type"] == "function" && tool["name"] == name) + ); +} + +#[test] +fn hosted_names_preserve_explicit_function_declarations_choices_history_and_output() { + for hosted in ["web_search", "x_search"] { + for hosted_first in [false, true] { + for nested_choice in [false, true] { + let mut tools = vec![function_tool(hosted), json!({"type": hosted})]; + if hosted_first { + tools.reverse(); + } + let choice = if nested_choice { + json!({"type": "function", "function": {"name": hosted}}) + } else { + json!({"type": "function", "name": hosted}) + }; + let request = raw_request(json!({ + "input": function_history(hosted), "tools": tools, "tool_choice": choice + })); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("hosted and client function"); + let body = encoded.body(); + let tools = body["tools"].as_array().expect("tools"); + let function = tools + .iter() + .find(|tool| tool["type"] == "function") + .expect("function retained"); + let alias = function["name"].as_str().expect("function alias"); + assert_ne!(alias, hosted, "hosted wire name must be unique"); + assert_eq!(tools.len(), 2); + assert!(tools.iter().any(|tool| tool["type"] == hosted)); + let choice_name = if nested_choice { + &body["tool_choice"]["function"]["name"] + } else { + &body["tool_choice"]["name"] + }; + assert_eq!(choice_name, alias); + assert_eq!(body["input"][0]["name"], alias); + assert_eq!(body["input"][0]["call_id"], "call_previous"); + assert_eq!(body["input"][1]["output"], "result"); + assert_function_response_restored(&encoded, alias, hosted); + } + } + } +} + +#[test] +fn hosted_names_avoid_session_injected_x_search_and_keep_aliases_stable() { + for name in ["web_search", "x_search"] { + let base = json!({"input": function_history(name), "tools": [function_tool(name)], + "tool_choice": {"type": "function", "name": name}}); + let encoded = + GrokResponsesRequest::encode(&raw_request(base.clone()), "grok-4.7", &client_key()) + .expect("without session"); + let alias = encoded.body()["tools"][0]["name"].as_str().expect("alias"); + let mut with_session = base; + with_session["prompt_cache_key"] = json!("hosted_name_session"); + let session = + GrokResponsesRequest::encode(&raw_request(with_session), "grok-4.7", &client_key()) + .expect("with session"); + assert_eq!(session.body()["tools"][0]["name"], alias); + assert_eq!(session.body()["tools"][1], json!({"type": "x_search"})); + assert_ne!(alias, "x_search"); + assert_eq!(session.body()["tool_choice"]["name"], alias); + assert_eq!(session.body()["input"][0]["name"], alias); + assert_function_response_restored(&session, alias, name); + } +} + +#[test] +fn hosted_names_rewrite_history_when_empty_tools_trigger_both_cache_tools() { + for name in ["web_search", "x_search"] { + let request = raw_request(json!({"prompt_cache_key": "history_only_session", + "input": function_history(name)})); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("history only"); + assert_eq!( + encoded.body()["tools"], + json!([{"type": "web_search"}, {"type": "x_search"}]) + ); + assert_eq!(encoded.body()["tool_choice"], "none"); + assert_ne!(encoded.body()["input"][0]["name"], name); + assert_eq!(encoded.body()["input"][0]["call_id"], "call_previous"); + } +} + +#[test] +fn hosted_names_do_not_rename_similar_or_namespaced_client_functions() { + let request = raw_request(json!({"tools": [function_tool("web_search_local"), + {"type": "namespace", "name": "local", "tools": [function_tool("x_search")]}, + {"type": "web_search"}, {"type": "x_search"}], "input": "search"})); + let encoded = + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()).expect("ordinary names"); + assert_eq!(encoded.body()["tools"][0]["name"], "web_search_local"); + assert_eq!(encoded.body()["tools"][1]["name"], "local__x_search"); +} + +#[test] +fn hosted_names_preserve_custom_tools_lowered_to_functions() { + let request = raw_request(json!({ + "tools": [{"type": "custom", "name": "web_search"}, {"type": "web_search"}], + "tool_choice": {"type": "custom", "name": "web_search"}, + "input": [{"type": "custom_tool_call", "call_id": "call_previous", + "name": "web_search", "input": "previous"}] + })); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("custom and hosted search"); + let alias = encoded.body()["tools"][0]["name"] + .as_str() + .expect("custom alias"); + assert_ne!(alias, "web_search"); + assert_eq!(encoded.body()["tool_choice"]["name"], alias); + assert_eq!(encoded.body()["input"][0]["name"], alias); + let item = json!({"type": "function_call", "id": "fc_custom", "call_id": "call_custom", + "name": alias, "arguments": "{\"input\":\"next\"}"}); + let frames = [ + json!({"type": "response.created", "response": {"id": "resp_custom"}}), + json!({"type": "response.output_item.added", "output_index": 0, "item": item}), + json!({"type": "response.output_item.done", "output_index": 0, "item": item}), + json!({"type": "response.completed", "response": {"id": "resp_custom", + "status": "completed", "output": [item]}}), + ] + .into_iter() + .map(|frame| format!("data: {frame}\n\n")) + .collect::(); + let events = GrokCanonicalDecoder::for_request("grok-4.7", &encoded) + .push(frames.as_bytes()) + .expect("custom response"); + let completed = events + .iter() + .filter_map(ProviderEvent::wire_event) + .next_back() + .expect("completed response") + .data(); + let restored = &completed["response"]["output"][0]; + assert_eq!(restored["type"], "custom_tool_call"); + assert_eq!(restored["name"], "web_search"); + assert_eq!(restored["call_id"], "call_custom"); + assert_eq!(restored["input"], "next"); +} + +#[test] +fn hosted_names_keep_existing_function_names_distinct_from_generated_aliases() { + let first = raw_request(json!({"tools": [function_tool("x_search")]})); + let first = + GrokResponsesRequest::encode(&first, "grok-4.7", &client_key()).expect("initial alias"); + let generated = first.body()["tools"][0]["name"].as_str().expect("alias"); + assert_ne!(generated, "x_search"); + for generated_first in [false, true] { + let mut tools = vec![function_tool("x_search"), function_tool(generated)]; + if generated_first { + tools.reverse(); + } + let request = raw_request( + json!({"tools": tools, "prompt_cache_key": "collision_session", "input": "search"}), + ); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("alias collision"); + let first = encoded.body()["tools"][0]["name"] + .as_str() + .expect("first alias"); + let second = encoded.body()["tools"][1]["name"] + .as_str() + .expect("second alias"); + assert_ne!(first, second); + assert_ne!(first, "x_search"); + assert_ne!(second, "x_search"); + let original_names = if generated_first { + [generated, "x_search"] + } else { + ["x_search", generated] + }; + assert_function_response_restored(&encoded, first, original_names[0]); + assert_function_response_restored(&encoded, second, original_names[1]); + } +} diff --git a/backend/crates/providers/xai/tests/transport/request/mod.rs b/backend/crates/providers/xai/tests/transport/request/mod.rs index 223658c84..3f7813f27 100644 --- a/backend/crates/providers/xai/tests/transport/request/mod.rs +++ b/backend/crates/providers/xai/tests/transport/request/mod.rs @@ -6,6 +6,11 @@ use serde_json::{Map, Value, json}; use provider_xai::{GrokRequestEncodeError, GrokResponsesRequest}; +mod hosted_names; +mod model_identity; +mod parameters; +mod web_search; + fn raw_request(body: Value) -> GenerateRequest { let Value::Object(body) = body else { panic!("request fixture must be an object"); @@ -191,7 +196,7 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { "reasoning.encrypted_content" ])) ); - assert_eq!(body.pointer("/reasoning/effort"), Some(&json!("high"))); + assert_eq!(body.pointer("/reasoning/effort"), Some(&json!("xhigh"))); assert_eq!(body.pointer("/reasoning/summary"), Some(&json!("auto"))); let request = raw_request(json!({ @@ -200,9 +205,9 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { "reasoning": {"effort": "xhigh"} })); let encoded = GrokResponsesRequest::encode(&request, "xai/grok-4.6-latest", &client_key()) - .expect("aliased request"); + .expect("routed model request"); let body = Value::Object(encoded.body().clone()); - assert_eq!(body.pointer("/model"), Some(&json!("grok-4.6"))); + assert_eq!(body.pointer("/model"), Some(&json!("xai/grok-4.6-latest"))); assert_eq!(body.pointer("/reasoning/effort"), Some(&json!("xhigh"))); assert_eq!(body.pointer("/reasoning_effort"), None); assert_eq!(body.pointer("/reasoningEffort"), None); @@ -215,7 +220,10 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { let encoded = GrokResponsesRequest::encode(&request, "grok-composer-2.5-fast", &client_key()) .expect("Composer request"); let body = Value::Object(encoded.body().clone()); - assert_eq!(body.pointer("/reasoning"), None); + assert_eq!( + body.pointer("/reasoning"), + Some(&json!({"effort":"high", "summary":"auto"})) + ); let request = raw_request(json!({ "model": "client-model", @@ -232,7 +240,7 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { body.pointer("/include"), Some(&json!(["reasoning.encrypted_content"])) ); - assert_eq!(body.pointer("/reasoning"), None); + assert_eq!(body.pointer("/reasoning"), Some(&json!({"effort":"max"}))); let request = raw_request(json!({ "model": "client-model", @@ -243,7 +251,7 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { .expect("Composer reasoning without effort"); assert_eq!( Value::Object(encoded.body().clone()).pointer("/reasoning"), - None + Some(&json!({})) ); let request = raw_request(json!({ @@ -255,12 +263,12 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { .expect("unknown provider prefix"); assert_eq!( Value::Object(encoded.body().clone()).pointer("/reasoning/effort"), - None + Some(&json!("max")) ); } #[test] -fn encoder_should_strip_grok_unsupported_fields() { +fn encoder_should_strip_codex_only_fields_without_model_guesses() { let request = raw_request(json!({ "model": "client-model", "input": "hello", @@ -285,16 +293,11 @@ fn encoder_should_strip_grok_unsupported_fields() { })); let encoded = GrokResponsesRequest::encode(&request, "grok-latest", &client_key()) - .expect("sanitized 4.5 request"); + .expect("sanitized request"); let body = Value::Object(encoded.body().clone()); for pointer in [ "/prompt_cache_retention", "/safety_identifier", - "/presence_penalty", - "/presencePenalty", - "/frequency_penalty", - "/frequencyPenalty", - "/stop", "/external_web_access", "/metadata/external_web_access", "/tools/0/external_web_access", @@ -311,12 +314,9 @@ fn encoder_should_strip_grok_unsupported_fields() { let encoded = GrokResponsesRequest::encode(&request, "grok-4.20-reasoning", &client_key()) .expect("sanitized 4.20 request"); let body = Value::Object(encoded.body().clone()); - assert_eq!( - body.pointer("/model"), - Some(&json!("grok-4.20-0309-reasoning")) - ); - assert_eq!(body.pointer("/logprobs"), None); - assert_eq!(body.pointer("/top_logprobs"), None); + assert_eq!(body.pointer("/model"), Some(&json!("grok-4.20-reasoning"))); + assert_eq!(body.pointer("/logprobs"), Some(&json!(true))); + assert_eq!(body.pointer("/top_logprobs"), Some(&json!(5))); } #[test] @@ -1004,7 +1004,7 @@ fn hosted_tool_choice_should_preserve_normalized_web_search_choice() { "model": "client", "input": "search", "tools": [ - {"type": "web_search_preview", "allowed_domains": ["example.com"]}, + {"type": "web_search_preview", "filters": {"allowed_domains": ["example.com"]}}, {"type": "x_search"} ], "tool_choice": {"type": "web_search_preview"} diff --git a/backend/crates/providers/xai/tests/transport/request/model_identity.rs b/backend/crates/providers/xai/tests/transport/request/model_identity.rs new file mode 100644 index 000000000..d6ac9e0d2 --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/request/model_identity.rs @@ -0,0 +1,54 @@ +//! 路由已选定真实模型，xAI 编码器不再解释旧别名或 Provider 前缀 + +use super::*; + +#[test] +fn routed_model_identity_is_preserved_without_provider_local_aliases() { + let request = raw_request(json!({"model": "client-model", "input": "hello"})); + for model in [ + "grok", + "grok-latest", + "grok-build-latest", + "grok-4.5-latest", + "grok-4.6-latest", + "grok-4.3-latest", + "grok-build", + "grok-composer", + "composer-2.5", + "grok-4.20-reasoning", + "grok-4.20-non-reasoning", + "grok-4.20-multi-agent", + "grok-4.20-multi-agent-latest", + "grok-4.7", + "Grok-4.6", + "xai/grok-4.7", + "x-ai/grok-4.7", + "grok/grok-4.7", + "tenant/custom-model", + ] { + let encoded = GrokResponsesRequest::encode(&request, model, &client_key()) + .expect("routed upstream model"); + assert_eq!(encoded.body()["model"], model); + } +} + +#[test] +fn routed_model_identity_does_not_infer_unsupported_fields_from_old_slugs() { + let fields = json!({"presence_penalty": 0.1, "frequency_penalty": 0.2, + "stop": ["done"], "logprobs": true, "top_logprobs": 5}); + let request = raw_request(fields.clone()); + for model in [ + "grok-4.5", + "grok-4.6", + "grok-4.3", + "grok-4.20-reasoning", + "tenant/grok-4.6", + "grok-4.7", + ] { + let encoded = + GrokResponsesRequest::encode(&request, model, &client_key()).expect("request"); + for (key, value) in fields.as_object().expect("fields") { + assert_eq!(encoded.body().get(key), Some(value), "{model} {key}"); + } + } +} diff --git a/backend/crates/providers/xai/tests/transport/request/parameters.rs b/backend/crates/providers/xai/tests/transport/request/parameters.rs new file mode 100644 index 000000000..4f969c9ea --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/request/parameters.rs @@ -0,0 +1,75 @@ +//! 验证推理档位遵守 Responses wire 而不由模型 slug 猜测 + +use super::*; + +#[test] +fn official_reasoning_efforts_are_preserved_for_current_and_future_models() { + for model in [ + "grok-4.7", + "grok-future", + "grok-4.5", + "grok-4.6", + "grok-3-mini", + "grok-composer-2.5-fast", + ] { + for effort in ["none", "minimal", "low", "medium", "high", "xhigh", "max"] { + let request = raw_request(json!({ + "input": "hello", + "reasoning": {"effort": effort, "summary": "auto"} + })); + let encoded = GrokResponsesRequest::encode(&request, model, &client_key()) + .expect("official effort wire"); + assert_eq!( + encoded.body()["reasoning"], + json!({"effort": effort, "summary": "auto"}), + "model {model}" + ); + } + } +} + +#[test] +fn invalid_reasoning_effort_is_rejected_instead_of_silently_dropped() { + for effort in [json!("future-value"), json!(42), json!({}), json!([])] { + let request = raw_request(json!({"reasoning": {"effort": effort}})); + assert!( + matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { + field: "reasoning.effort" + }) + ), + "effort {effort}" + ); + } +} + +#[test] +fn reasoning_effort_keeps_existing_case_and_whitespace_normalization() { + let request = raw_request(json!({"reasoning": {"effort": " XHIGH "}})); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("normalized official effort"); + assert_eq!(encoded.body()["reasoning"]["effort"], "xhigh"); +} + +#[test] +fn reasoning_shape_is_validated_while_absent_effort_uses_server_default() { + for reasoning in [ + Value::Null, + json!({}), + json!({"effort": null}), + json!({"summary":"auto"}), + ] { + let request = raw_request(json!({"reasoning": reasoning})); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("no explicit effort"); + assert_eq!(encoded.body().get("reasoning"), Some(&reasoning)); + } + for reasoning in [json!("high"), json!(42), json!([])] { + let request = raw_request(json!({"reasoning": reasoning})); + assert!(matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { field: "reasoning" }) + )); + } +} diff --git a/backend/crates/providers/xai/tests/transport/request/web_search.rs b/backend/crates/providers/xai/tests/transport/request/web_search.rs new file mode 100644 index 000000000..a050a4117 --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/request/web_search.rs @@ -0,0 +1,136 @@ +//! Web Search 域名过滤的协议转换与字段归属测试 + +use super::*; + +#[test] +fn web_search_filters_preserve_excluded_domains_for_supported_aliases() { + for kind in [ + "web_search", + "web_search_preview", + "web_search_preview_2025_03_11", + "web_search_2025_08_26", + ] { + let request = raw_request(json!({ + "input": "search", + "tools": [{"type": kind, "filters": {"excluded_domains": ["example.com"]}}] + })); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("excluded domains"); + + assert_eq!( + encoded.body()["tools"], + json!([{"type": "web_search", "filters": {"excluded_domains": ["example.com"]}}]), + "{kind}" + ); + } +} + +#[test] +fn web_search_filters_reject_combined_nonempty_allow_and_block_lists() { + for tool in [ + json!({"type": "web_search", "filters": { + "allowed_domains": ["allowed.example"], "excluded_domains": ["excluded.example"] + }}), + json!({"type": "web_search_preview", "allowed_domains": ["allowed.example"], + "filters": {"excluded_domains": ["excluded.example"]} + }), + ] { + let request = raw_request(json!({"input": "search", "tools": [tool]})); + assert!(matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { field: "tools" }) + )); + } +} + +#[test] +fn web_search_filters_keep_empty_lists_unbounded() { + for (filters, expected) in [ + ( + json!({"allowed_domains": [], "excluded_domains": ["example.com"]}), + json!({"excluded_domains": ["example.com"]}), + ), + ( + json!({"allowed_domains": ["example.com"], "excluded_domains": []}), + json!({"allowed_domains": ["example.com"]}), + ), + (json!({"excluded_domains": []}), Value::Null), + ( + json!({"allowed_domains": [], "excluded_domains": []}), + Value::Null, + ), + (json!({"allowed_domains": []}), Value::Null), + ] { + let request = raw_request(json!({ + "input": "search", "tools": [{"type": "web_search", "filters": filters}] + })); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("empty lists are unbounded"); + assert_eq!(encoded.body()["tools"][0]["filters"], expected); + } +} + +#[test] +fn web_search_filters_reject_retired_top_level_allowed_domains() { + for tool in [ + json!({"type": "web_search_preview", "allowed_domains": ["example.com"]}), + json!({"type": "web_search", "allowed_domains": ["example.com"], + "filters": {"allowed_domains": ["example.com"]}}), + json!({"type": "web_search", "allowed_domains": ["one.example"], + "filters": {"allowed_domains": ["other.example"]}}), + json!({"type": "web_search", "allowed_domains": []}), + ] { + let request = raw_request(json!({"input": "search", "tools": [tool]})); + assert!(matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { field: "tools" }) + )); + } +} + +#[test] +fn web_search_filters_reject_invalid_excluded_domains() { + for excluded in [ + Value::Null, + json!("example.com"), + json!([" "]), + json!([3]), + json!(["example.com", false]), + ] { + let request = raw_request(json!({"input": "search", "tools": [{ + "type": "web_search", "filters": {"excluded_domains": excluded} + }]})); + assert!(matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { field: "tools" }) + )); + } +} + +#[test] +fn web_search_filters_only_interpret_owned_tool_fields() { + let schema = json!({"type": "object", "properties": { + "filters": {"type": "string"}, "excluded_domains": {"type": "integer"} + }}); + let opaque = json!({"excluded_domains": "application data"}); + let request = raw_request(json!({ + "input": "search", + "excluded_domains": opaque, + "tools": [ + {"type": "function", "name": "lookup", "parameters": schema}, + {"type": "x_search", "excluded_domains": opaque}, + {"type": "web_search", "excluded_domains": opaque, + "filters": {"future_filter": opaque, "excluded_domains": ["example.com"]}} + ] + })); + let encoded = + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()).expect("field ownership"); + + assert_eq!(encoded.body()["excluded_domains"], opaque); + assert_eq!(encoded.body()["tools"][0]["parameters"], schema); + assert_eq!(encoded.body()["tools"][1]["excluded_domains"], opaque); + assert_eq!( + encoded.body()["tools"][2], + json!({"type": "web_search", "filters": {"excluded_domains": ["example.com"]}}) + ); +} diff --git a/docs/api.md b/docs/api.md index 49afe31d2..2b243ec22 100644 --- a/docs/api.md +++ b/docs/api.md @@ -322,6 +322,10 @@ OAuth 账号默认 `prefer_websocket`，客户端使用 HTTP/SSE 时仍可能选 客户端配置的 `supports_websockets` 只控制第一段连接，不是服务端传输策略开关。 上游在响应终态前发送 Close 1000 仍属于失败，不能按“正常关闭”计为成功 +内建 xAI 适配器不支持 `generate=false` 的非生成预热请求，在发送推理请求前返回 +`400 unsupported_prewarm`，不会将其当作普通生成调用。`generate=true` 或省略该字段正常执行； +客户端可跳过预热继续正式请求，外部插件 adapter 按各自的协议能力处理 + Codex OAuth backend 的候选上游为 WS 时，无 `previous_response_id` 的普通新链若规范化 `response.create` 达到 15 MiB，发送前选择 HTTP/SSE；HTTP 和 WS 入站均适用，下游交付协议不变。 该阈值为已观察到的上游消息大小边界预留余量，不是网关输入长度上限或 OpenAI 公布的统一限制； @@ -359,7 +363,9 @@ API Key 上游返回完整 Codex `models` 目录时沿用该合同；仅返回 未提供的推理能力保持未知，不补充推理档位。 模型别名仅替换 `slug`，不替换上游展示名、提示词、能力或 `priority`；保持原生模型顺序，新增别名附在后面。 目录按当前路由快照的模型存在性及账号模型政策过滤，避免公布已知无法路由的模型；新模型需待后台目录对账后进入列表。 -xAI 没有 Codex 原生目录，继续使用明确的通用画像适配 +xAI 没有 Codex 原生目录，使用官方 CLI proxy 当前目录生成通用画像。推理菜单读取 +`reasoning_efforts` 对象列表，默认值只取菜单中的显式标记；缺失时保持未知，不以首项代替， +也不回退旧字段或其他端点的目录形状 目录账号只能来自本次 Client Key 冻结的账号范围。OpenAI 的 OAuth 目录按账号 ID 排序，使用首个成功读取的 合格账号，最多尝试三个账号；API Key 目录按账号模型权限过滤后聚合。同名模型优先采用 OAuth 原生对象， @@ -380,7 +386,8 @@ xAI 没有 Codex 原生目录，继续使用明确的通用画像适配 Codex 专用目录中的 `context_window` 与 `max_context_window` 分别表示默认上下文窗口和客户端本地 覆盖的上限。OpenAI Provider 原样保留上游对应字段，缺失与 `null` 不互相转换；网关不通过部署配置 覆盖这些值。Codex 客户端配置 `model_context_window` 后，按该值与非空 `max_context_window` 的较小值 -使用窗口；上限为空时保留客户端本地值。xAI 目录只声明一个窗口，其 Provider 继续以该值作为客户端覆盖上限 +使用窗口；上限为空时保留客户端本地值。xAI 默认窗口取 `context_window`，缺失时取 `context_windows` +中的首个有效值；覆盖上限取默认窗口和可选窗口的最大值，只有标量时以该值作为上限 ### 透传与错误恢复 @@ -398,6 +405,16 @@ Responses wire 之间的协议转换层，转换只在 xAI Provider 内完成。 上游结构化错误的 message/code/type 按上述边界交付客户端，其中内嵌的账号指纹 UUID 已脱敏。模型映射是 全局精确映射，未命中时模型名原样交给候选 Provider；分组只限定账号集合，不参与模型改名 +xAI 编码器同样保留路由选定的完整模型名，不内置 `grok`、`grok-latest` 等旧型号别名，也不剥离模型前缀； +需要别名时使用显式模型映射。目录读取失败按错误处理，不用硬编码旧型号补全能力 + +xAI 的 `reasoning.effort` 接受 `none / minimal / low / medium / high / xhigh / max`，去除两端 +空白并转为小写后保留档位，不按模型名称降档或删除；模型是否支持该选择由上游判断。 +未知值或错误类型返回字段错误，其他 reasoning 字段保留。`web_search.filters` 支持 +`allowed_domains` 与 `excluded_domains`，两者不能同时为非空；客户端函数与托管搜索工具 +同名时使用内部别名，并在返回工具调用时还原客户端名称。搜索工具的旧顶层 `allowed_domains` +不再映射到 filters，直接返回请求错误 + OpenAI 明确返回 `server_is_overloaded`、`slow_down` 或模型容量不足错误时，代理在允许安全重放且 尚未交付输出的前提下，先做最多 3 次同账号重试，再通过现有调度换号。没有有效服务器建议时， 按 500ms 起步的指数退避等待，单次最多 8 秒；有效 `Retry-After` 优先，接受秒数、HTTP 日期和零延迟， @@ -982,6 +999,9 @@ OAuth start 使用： `quota_exhausted`。额度观测不会清除凭据过期、无效或封禁事实；这些事实统一投影为 `error`，并由 `errorReason` 区分。额度接口的 401/403 也不足以判定 refresh token 永久失效，credential 终态只由 OAuth refresh 的明确永久错误写入 +- xAI 额度百分比只读取 `creditUsagePercent`，缺失时保持未知，不由旧月额度与用量计算。 + 预付余额保留上游账本符号，非零余额及剩余按量额度可独立证明可用；窗口已过期或缺少可用证据时， + 不凭旧快照恢复账号 - 正常 Responses 请求会解析上游响应的 rate-limit headers，合并进同一 quota 快照并同步状态。Free、 K12 等套餐共用该状态机；套餐只参与账号展示和按套餐隔离的模型目录 cache，不存在 K12 专属额度路径 - 账号详情的 Token 统计、模型排行和列表 Token 汇总优先使用账号级周额度窗口，无可统计的周窗口时 @@ -1801,6 +1821,10 @@ OpenAI 与 xAI 的本地费用估算按实际发送给上游的请求模型（`u 用量计算；客户端请求 A、路由后发送 B 时按 B 计价，响应返回 C 不改变计价模型。实际发送模型缺少 定价时不估算，也不借用响应模型的价格。Provider 明确上报的已计费金额仍优先于本地估算 +xAI 的 `cost_in_usd_ticks=0` 表示上游没有报告费用，不证明请求免费；正数才作为上游账单。 +未报告时，只有已知价格且能覆盖本次费用的请求才使用本地估算，未知价格或无法覆盖的托管工具费用保持未知。 +动态 latest 与旧 beta/internal 别名不借用具体旧型号的内置价格；显式配置的模型价格仍可用于估算 + OpenAI Responses 用量记录的 `serviceTier` 与本地费用估算统一采用 Provider 最终发给上游的请求 `service_tier`，不使用响应档位覆盖或回退。例如发送 `priority`、响应回显 `default` 时，仍显示 `Fast` 并按 Priority 价格估算。未发送档位时，`serviceTier` 保持缺失，展示与估算按标准档处理。 diff --git a/docs/architecture.md b/docs/architecture.md index b88bce5df..f0f8d3fd1 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -250,16 +250,22 @@ HTTP 字段见 [插件 API](api.md#12-插件管理)，更新与数据恢复见 [ | `account` | Provider 账号/credential/quota 值对象、持久化端口与请求级账号选择；`scope` 持有分组、账号目录和冻结账号范围 | | `policy` | Client API Key 准入、原始 Key 设置与客户端版本策略；只使用账号范围、Provider 身份和基础校验 | | `metering` | 标准化 Usage、金额、费用估算与费用明细；不表示账号或开票系统 | -| `upstream` | 跨 Engine、Event、Error 与 Provider 共用的 transport 名称、发送状态和不透明上游值 | +| `upstream` | 跨 Engine、Event、Error 与 Provider 共用的 transport 名称、尝试传输档位、发送状态和不透明上游值 | | `lifecycle` | 取消信号、可选执行截止与租约生命周期、连接注册与 drain 合同 | -| `engine` | attempt、发送/提交屏障、执行编排和持久化调用时序 | -| `routing` | 冻结路由事实、请求计划、Provider 只读目录合同以及运行时快照的表示与编译 | +| `middleware` | 与领域无关的类型化组合器及单次续体 | +| `engine` | attempt、发送/提交屏障、执行编排、具体中间件与上游适配计划、持久化调用时序 | +| `routing` | 冻结路由事实、请求计划、Provider 只读目录合同、扩展集合身份与保活合同、运行时快照及请求设置派生 | | `runtime` | 当前快照的发布、读取、revision 订阅与周期对账任务 | -| `settings` | 请求设置事实及其编译、宿主与 Key 默认值解析、显式覆盖来源；由 `routing` 快照组合并冻结 | +| `settings` | 设置值及纯编译规则；由 `routing` 快照组合并冻结 | `event` 通过 `validation` / `upstream` 使用基础值，不依赖承载原始事件的执行错误；账号值对象和 选择策略通过 `identity` / `account::scope` 使用身份与范围，不依赖路由计划。`routing` 和 `error` 的 -相关公开类型通过 re-export 引用上述定义，类型所有权和 Core 内部依赖归属定义模块。架构测试约束这些叶子依赖 +相关公开类型通过 re-export 引用上述定义，类型所有权和 Core 内部依赖归属定义模块。架构测试约束这些叶子依赖， +并检查一级 owner 的显式 `crate` / `self` / `super` 路径与重导出形成的依赖图；它不替代 Rust 名称解析和行为测试 + +`routing::request_settings` 统一解析宿主与 Key 默认值、显式覆盖来源及请求快照重算，不向全局发布请求改写。 +`routing::extensions` 只表达扩展集合的中立身份、静态目录和租约；`engine::extensions` 按冻结身份弱引用索引 +具体执行计划，插件 Runtime（`gateway-plugin/runtime`）的 `PreparedSet` 负责强持有计划。请求取得的冻结计划同时保活该集合，旧代次随最后一个使用者退出 Provider 模型能力、目录代次与 `ProviderCatalogPort` 由 `routing::catalog` 定义。快照编译和对账只消费 该只读合同，不反向依赖执行注册表；`ProviderRegistry` 实现目录端口，维护唯一的 Provider