diff --git a/.dockerignore b/.dockerignore index 9a9a9fb1..98184094 100644 --- a/.dockerignore +++ b/.dockerignore @@ -40,3 +40,5 @@ task_plan.md test_*.py test_*.sh tests/helpers/workbench_preview.py + +.release-wrangler-* diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e7a14409..629cacfa 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,6 +17,7 @@ jobs: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: persist-credentials: false + fetch-depth: 0 - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: "3.11" @@ -50,6 +51,7 @@ jobs: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: persist-credentials: false + fetch-depth: 0 - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: "3.11" @@ -84,6 +86,7 @@ jobs: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: persist-credentials: false + fetch-depth: 0 - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: "3.11" @@ -106,6 +109,7 @@ jobs: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: persist-credentials: false + fetch-depth: 0 - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 with: node-version: "24" @@ -119,11 +123,15 @@ jobs: run: git diff --exit-code -- static/css/style.css - name: Frontend build contracts run: node --test tests/test_frontend_build.mjs + - name: Check release configuration + run: node scripts/key_data_preflight.mjs --check + # Bundle-only checks cannot attest remote deployment prerequisites. The + # normal Wrangler build enforces those before a real upload. - name: Validate Knowledge Worker bundle env: WRANGLER_SEND_METRICS: "false" - run: pnpm exec wrangler deploy --config wrangler.knowledge.jsonc --dry-run --outdir /tmp/knowledge-worker + run: node scripts/check_worker_bundle.mjs --config wrangler.knowledge.jsonc --outdir /tmp/knowledge-worker - name: Validate main Worker bundle env: WRANGLER_SEND_METRICS: "false" - run: pnpm exec wrangler deploy --dry-run --containers-rollout=none --outdir /tmp/main-worker + run: node scripts/check_worker_bundle.mjs --config wrangler.jsonc --outdir /tmp/main-worker diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 2961c79b..c94246d9 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -121,6 +121,8 @@ jobs: run: npm install --global "$(node -p "require('./package.json').packageManager")" - name: Install deployment dependencies run: pnpm install --frozen-lockfile + - name: Check release configuration + run: node scripts/key_data_preflight.mjs --check - name: Check migration files run: node scripts/verify_d1_migrations.mjs --check - name: List pending D1 migrations @@ -136,19 +138,37 @@ jobs: run: | pnpm exec wrangler d1 migrations apply "$D1_DATABASE" --remote node scripts/verify_d1_migrations.mjs "$D1_DATABASE" + - name: Check migrated schema and required bindings + env: + CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} + CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }} + run: node scripts/key_data_preflight.mjs --prepare - name: Deploy Knowledge Worker id: knowledge env: CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }} run: pnpm exec wrangler deploy --config wrangler.knowledge.jsonc --tag "$SHA" --message "$DEPLOY_MESSAGE" + - name: Capture Container rollout baseline + if: vars.DEPLOY_MAIN_WORKER == 'true' + env: + CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} + CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }} + run: node scripts/container_rollout.mjs --prepare --state "$RUNNER_TEMP/container-before.json" --output .release-wrangler-actions.jsonc - name: Deploy main Worker and Container id: main if: vars.DEPLOY_MAIN_WORKER == 'true' env: CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }} - run: pnpm exec wrangler deploy --tag "$SHA" --message "$DEPLOY_MESSAGE" + run: pnpm exec wrangler deploy --config .release-wrangler-actions.jsonc --tag "$SHA" --message "$DEPLOY_MESSAGE" + - name: Verify actual Container rollout + if: vars.DEPLOY_MAIN_WORKER == 'true' + env: + CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} + CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }} + DEPLOY_STATUS_URL: ${{ vars.DEPLOY_STATUS_URL }} + run: node scripts/container_rollout.mjs --verify --state "$RUNNER_TEMP/container-before.json" - name: Verify deployment id: verify env: diff --git a/.gitignore b/.gitignore index e7a12656..2dee1ec3 100644 --- a/.gitignore +++ b/.gitignore @@ -151,3 +151,6 @@ user_data.json .container-build-id # Local browser verification artifacts output/playwright/ + +# Temporary deployment configuration (build arguments only). +.release-wrangler-* diff --git a/Dockerfile b/Dockerfile index c3c2dd47..a90c72ad 100644 --- a/Dockerfile +++ b/Dockerfile @@ -25,6 +25,11 @@ RUN addgroup --system multillm && \ chmod -R a+rX /app && \ chmod 0750 /app/instance /tmp/multillm +# A fresh release identifier prevents stale application-list comparisons from +# treating a reused image as an unchanged Container. It is not a source stamp. +ARG MULTILLM_RELEASE_IMAGE_ID=local +LABEL multillm.release_image_id=$MULTILLM_RELEASE_IMAGE_ID + USER multillm EXPOSE 8080 diff --git a/cloudflare-worker.mjs b/cloudflare-worker.mjs index 2effc922..534d32d0 100644 --- a/cloudflare-worker.mjs +++ b/cloudflare-worker.mjs @@ -1,3 +1,7 @@ +import { BUILD_ID } from "./worker/build-id.mjs"; +import { admitContainerDelegation } from "./worker/peer-controls.mjs"; +import { dispatchFailure } from "./worker/provider-dispatch.mjs"; +import { KEY_CONTROLS } from "./worker/key-controls-capability.mjs"; import { nativeGenerationFetch, withForwardedCorrelation, withNativeMetrics, nativeCacheHeader, nativeGenerationSetup, generationErrorResponse, forwardedGenerationHeaders, runScheduledMaintenance, scheduledMaintenanceEnabled, handleRoleplayContextPageRequest, handleRealtimeRequest } from "./worker/gateway-extensions.mjs"; import { batchesEnabled, runScheduledBatches } from "./worker/batch-jobs.mjs"; import { tickNativeRevisionSync } from "./worker/native-config-sync.mjs"; @@ -282,7 +286,7 @@ function jsonResponse(body, init = {}) { function buildFallbackHealthResponse() { return jsonResponse({ - status: "healthy", + status: "healthy", key_controls: KEY_CONTROLS, mode: "worker-fallback", }); } @@ -2079,6 +2083,10 @@ export default { try { const container = getContainer(env.MULTILLM_PROXY_CONTAINER, "primary"); + if ((apiPath || isKnowledgeEdgePath(requestUrl.pathname)) && !readyPath) { + try { await admitContainerDelegation(request, env, container); } + catch (error) { return applyCorsHeaders(request, dispatchFailure(error), env); } + } const bodyAllowed = request.method !== "GET" && request.method !== "HEAD"; const headers = forwardedGenerationHeaders(request, env); // Client copies of the Access identity headers never reach the Container. @@ -2102,6 +2110,12 @@ export default { ...(bodyAllowed && request.body ? { duplex: "half" } : {}), }); let response = await container.fetch(forwardedRequest); + if (readyPath && response.ok) { + const health = await response.clone().json(); + const releaseHeaders = new Headers(response.headers); + releaseHeaders.set("X-MultiLLM-Worker-Build", BUILD_ID); + response = jsonResponse({ ...health, main_key_controls: KEY_CONTROLS }, { headers: releaseHeaders }); + } if (requestUrl.pathname === "/v1/metrics/prometheus") { response = await withNativeMetrics(response, env); } diff --git a/docs/deployment-cloudflare.md b/docs/deployment-cloudflare.md index 4ab4bfae..74236e4a 100644 --- a/docs/deployment-cloudflare.md +++ b/docs/deployment-cloudflare.md @@ -383,14 +383,12 @@ the code, and nothing ships unless ci passed. 3. Add *Containers Edit* to the API token, and run the `deploy` workflow once from the **Actions** tab to deploy the current head of `main`. -**(b) Workers Builds keeps deploying the main Worker.** Leave `DEPLOY_MAIN_WORKER` unset -(the default) or `false`, and keep the Workers Builds deploy command as -`npx wrangler deploy`; do not change it to `npm run deploy`, which would apply the -migrations and deploy the Knowledge Worker a second time. Actions applies migrations and -deploys the Knowledge Worker after ci passes. Workers Builds starts at push time, so after -a push that adds a migration the new main Worker usually runs for a few minutes before -its tables exist, and `/ready` answers 503 `d1_schema_missing` until they do. Code must -tolerate a missing new table, and a push whose ci fails still ships the main Worker. +**(b) Keeping automatic promotion is not supported for the key-controls rollout.** +Workers Builds can publish the main Worker and Container before migrations or the +private Worker, even when GitHub CI fails. Disconnect it or use `wrangler versions upload` +without automatic promotion, and deploy through `scripts/deploy.sh`. The preflight's +D1, secret-list and deployments-read permissions inside Builds have not been verified. +See [the staged rollout and rollback procedure](key-controls-rollout.md). ## One-time setup @@ -440,13 +438,19 @@ which can show no commit until the next deploy. ## Manual deploys -`npm run deploy` (`scripts/deploy.sh`) runs the same order from your machine: the -migration checks, `wrangler d1 migrations apply`, the remote check, the Knowledge -Worker and the main Worker. From a clean checkout it tags both versions with the commit. -Extra arguments go to the final deploy, for example -`npm run deploy -- --containers-rollout immediate` for a one-shot rollout of a changed -Container image. Use it only when the workflow cannot run, and never while a deploy run -is in progress; with option (a) the next push to `main` replaces it. +Set `D1_DATABASE` explicitly to the database in both selected Worker configurations, +then run `npm run deploy` (`scripts/deploy.sh`). For another environment set +`MAIN_WORKER_CONFIG`, `KNOWLEDGE_WORKER_CONFIG` and `DEPLOY_STATUS_URL` as well. The +script checks migrations, applies them, deploys the private Worker, then the main +Worker and Container. From a clean checkout it tags both versions with the commit. +Extra arguments go to the final deploy. Never run it concurrently with another deploy. + +Each release uses a fresh build-only image identifier. The command does not finish +successfully until application detail shows a newer Container version and image, and +`/ready` returns matching Worker/Container artifact fingerprints. "No changes" is not +rollout proof. See [deployment verification](key-controls-rollout.md#deployment-verification). +Status reads also require `--database` or `D1_DATABASE`; alternate configs are selected +with `--config` and `--knowledge-config`. No default database is queried. **Alerts.** Worker logs are sampled at 100% (`head_sampling_rate: 1`): the Worker writes little besides structured failure lines such as `account_storage_failed`, @@ -471,7 +475,11 @@ With D1, more Gunicorn worker processes behave like more instances: each keeps i ## Rollback -Revert the commit on `main` so the deploy workflow ships the previous code, or use `npx wrangler rollback` for an immediate rollback that lasts until the next deploy. D1 migrations are not rolled back; because they are additive, the previous code keeps working with the newer schema. Container SQLite files are ephemeral and need no rollback; restore D1 data itself with D1 Time Travel. +Use the [Container-first rollback procedure](key-controls-rollout.md#deployment-verification) +until A/B rollback has passed real staging. A Worker-only rollback across earlier strict +RPC decoders can fail unrestricted requests. Keep migrations and stored restrictions; +A refuses controlled dispatch while preserving unrestricted and management access. +Do not clear policies or roll back D1 migrations to make a rollback appear healthy. ## Sources diff --git a/docs/key-controls-rollout.md b/docs/key-controls-rollout.md new file mode 100644 index 00000000..1efa694f --- /dev/null +++ b/docs/key-controls-rollout.md @@ -0,0 +1,115 @@ +# Key controls rollout + +The main Worker, its Container and the private Knowledge Worker deploy independently. +Workers Builds deploys the main Worker and Container without running the GitHub deploy +workflow that deploys the private Knowledge Worker. Do not rely on deployment order. + +## Expand + +Apply migrations 0041–0043, then deploy the expansion revision to all three components. +Their readiness responses advertise `key_controls: {protocol: 1, enforcement: false}`. +The main Worker also exposes the Container readiness result separately. Verify each +component before proceeding; a build tag alone does not prove the running capability. + +Existing accounts have NULL policy columns and retain their existing behavior. Any non-NULL policy value, +including `official_only=0` or an empty preference object, activates the expansion refusal fence. The +expansion accepts additive RPC fields and both account row shapes, sends the legacy +protocol, and does not expose setters for restrictions. Legacy account writes preserve +both policy columns in the same SQL statement, including concurrent writes. + +If a later revision enables an account restriction, rolling back to this revision +refuses provider dispatch with 503 before sending content. This dormant fence includes +native routes, private bindings, Knowledge, Realtime and deferred work. A malformed or +unreadable current policy cannot grant dispatch. Legacy background jobs without an +identity run only while a fresh durable read proves every policy is NULL. Headerless +private calls from a verified expansion Container remain compatible beside controlled +accounts: that Container applies its current-account refusal fence before the hop. +An unreadable or malformed store cannot use this compatibility proof. + +## Contract + +Deploy the enforcement revision only after all three components advertise expansion +support. Verify enforcement capability on every component before enabling restrictions. +Controlled delegation requires a fresh enforcement-capability response from the actual +peer before any content is sent. Unrestricted traffic remains compatible with expansion +peers during independent deployments. Keep deployment preflight enabled. Expansion +already checks the required secret names, both account tables' policy columns and the +complete migration ledger before main upload; enforcement retains those checks. + +Rollback from enforcement to expansion leaves policy columns intact and refuses +controlled accounts. Restore enforcement-capable components to resume those accounts; +never clear restrictions as an automatic rollback step. + +## Staging evidence and management compatibility + +Real staging of the earlier re-land reproduced both incident errors with mixed Worker +and Container versions. The matched original enforcement release was healthy on 45 of +46 requests. Forward negotiation fixed the new-Worker/old-Container mix, but the reverse +rollback still produced roughly 50–55 seconds of chat 503s in two samples. Treat both +RPC directions as compatibility requirements; a passing matched-version test is insufficient. + +The local matrix uses real migrated D1 rows, the account RPC and registered Container +routes. Every A/B combination must keep unrestricted chat and Knowledge working, including +rollback with restricted rows still stored. It also signs in both a controlled account +and an unrestricted administrator, then reloads the complete management account list. +Policy capability checks belong at provider dispatch/delegation, never a management read. +A restricted row must not poison another account's login or dashboard. Management access +does not grant generation authority; A's dispatch fence remains active after sign-in. +These are local compatibility proofs; repeat them on staging before relying on plain rollback. + +## Deployment verification + +Expansion release admission runs from the default main Worker build and both deployment +entry points. Before private upload, `--prepare` verifies the exact local migration ledger, +`control_users` and `intelligence_principals` policy columns, and the required main binding +names (`ADMIN_API_KEY`, `FLASK_SECRET_KEY`, `JWT_SECRET`), without reading secret values. +Before main upload it also verifies a single fully active private deployment tagged with +an expansion-capable source revision. Keep that revision's history available locally. +These deployment checks are separate from runtime capability probes and the final +Container version/image/fingerprint verification. All remote migration commands require +an explicit database argument or `D1_DATABASE`; there is no production fallback. + +Before merging, disconnect Workers Builds or change its deploy command to +`npx wrangler versions upload`, which uploads without promoting traffic. Use +`scripts/deploy.sh` to preserve migrations → private Worker → main Worker → Container. +The preflight's token permissions inside Workers Builds remain untested: they include D1, +secret-list and deployment/version reads, plus Container list/detail for rollout +verification. Do not treat a locally passing preflight as proof that the Builds token +has those permissions. + +Set `D1_DATABASE` explicitly, matching `INTELLIGENCE_DB` in both selected configs. +For another environment also set `MAIN_WORKER_CONFIG`, `KNOWLEDGE_WORKER_CONFIG` and +`DEPLOY_STATUS_URL`. Status reads require `--database` or `D1_DATABASE` and validate +that target against both configurations; there is no implicit production database. +Use `--config` and `--knowledge-config` for status against alternate configurations. + +Wrangler can print "Image already exists remotely" and "no changes" while leaving an +old Container running, even with immediate rollout. Its application list can lag the +application detail. The deploy script records the detail version/image before upload, +adds a fresh build-only image identifier, and polls the detail after deployment. It +requires a strictly newer version and changed image, plus matching actual Worker and +Container artifact fingerprints on `/ready`. A no-op, unavailable probe or mismatch +fails the deployment after a bounded wait. Never enable controls after such a failure. +The GitHub deployment workflow uses the same capture and verification. These checks are +read-only and do not change secret bindings, D1 data or rollout policy. + +Until A/B rollback has passed real staging, use the staging-verified Container-first +rollback: from the newer Worker checkout, deploy a reviewed temporary config referencing +the previous Container image digest (without the Dockerfile build command). Keep the +newer Worker active. Verify that application detail advanced to that image and the +responding Container has the intended previous fingerprint. Only then deploy the previous +Worker with `--containers-rollout=none`. If Wrangler reports no changes, stop and resolve +the image rollout; do not advance to the Worker rollback. Do not disable the normal +preflight for routine deploys; an emergency digest rollback is a separately reviewed +operation. Keep migrations and stored restrictions intact. Once A is installed everywhere, +B→A rollback must retain unrestricted service and refuse controlled dispatch on every path. + +Mixed-load "Container suddenly disconnected" failures also occurred on old code +(roughly 9–14% in that staging load). They are outside this rollout change. Zero policy +errors in the compatibility matrix does not claim zero infrastructure disconnects. + +See [Container deployment semantics](https://developers.cloudflare.com/containers/guides/deploy/) +and [rollout controls](https://developers.cloudflare.com/containers/configuration/rollouts/). + + +Account reads negotiate the additive row protocol only after a successful capability probe. Old Containers receive the original account fields on list reads and NULL-policy username reads, including when a controlled neighbour exists, so dashboard cache reloads and unrestricted sign-in remain available. Legacy username and prefix authentication refuse controlled rows, since either can also feed deferred provider work. Provider forwarding and private outbound admission remain separate from management projection; an old Container cannot use a projected row to dispatch for a controlled account. A Containers request the full policy from compatible peers and retain their dormant refusal fence. A failed or malformed capability probe never downgrades to a legacy row; only the old peer's explicit unsupported-operation response permits that fallback. diff --git a/intelligence-migrations/0041_key_data_guards.sql b/intelligence-migrations/0041_key_data_guards.sql new file mode 100644 index 00000000..31571653 --- /dev/null +++ b/intelligence-migrations/0041_key_data_guards.sql @@ -0,0 +1,3 @@ +-- Key data controls. Apply before the matching release. +ALTER TABLE control_users ADD COLUMN official_only INTEGER CHECK (official_only IN (0, 1)); +ALTER TABLE control_users ADD COLUMN forced_provider_preferences TEXT; diff --git a/intelligence-migrations/0042_usage_token_details.sql b/intelligence-migrations/0042_usage_token_details.sql new file mode 100644 index 00000000..c5416980 --- /dev/null +++ b/intelligence-migrations/0042_usage_token_details.sql @@ -0,0 +1,5 @@ +-- Provider-reported usage details; old event rows keep unknown values. +ALTER TABLE usage_events ADD COLUMN cached_input_tokens INTEGER CHECK (cached_input_tokens >= 0); +ALTER TABLE usage_events ADD COLUMN reasoning_tokens INTEGER CHECK (reasoning_tokens >= 0); +ALTER TABLE usage_daily ADD COLUMN cached_input_tokens INTEGER NOT NULL DEFAULT 0 CHECK (cached_input_tokens >= 0); +ALTER TABLE usage_daily ADD COLUMN reasoning_tokens INTEGER NOT NULL DEFAULT 0 CHECK (reasoning_tokens >= 0); diff --git a/intelligence-migrations/0043_integration_data_policy.sql b/intelligence-migrations/0043_integration_data_policy.sql new file mode 100644 index 00000000..76eecd9c --- /dev/null +++ b/intelligence-migrations/0043_integration_data_policy.sql @@ -0,0 +1,3 @@ +-- Integration credentials use the same explicit data controls as account keys. +ALTER TABLE intelligence_principals ADD COLUMN official_only INTEGER CHECK (official_only IS NULL OR official_only IN (0, 1)); +ALTER TABLE intelligence_principals ADD COLUMN forced_provider_preferences TEXT; diff --git a/package.json b/package.json index f1264d01..43a630cb 100644 --- a/package.json +++ b/package.json @@ -9,7 +9,7 @@ "lint": "ruff check . --select E9,F63,F7,F82", "typecheck": "mypy --ignore-missing-imports --follow-imports=skip --disable-error-code import-untyped --disable-error-code var-annotated --disable-error-code arg-type --disable-error-code assignment --disable-error-code operator app.py routes services providers streaming", "test": "python -m pytest -q", - "test:worker": "node --test tests/test_cloudflare_worker.mjs tests/test_child_safety_worker.mjs tests/test_child_safety_dashboard.mjs tests/test_runtime_feature_flags_worker.mjs tests/test_anthropic_worker.mjs tests/test_haiku_intelligence_worker.mjs tests/test_cascades_worker.mjs tests/test_cascades_dashboard.mjs tests/test_image_quality_worker.mjs tests/test_secret_scan_worker.mjs tests/test_secret_firewall_worker.mjs tests/test_secret_firewall_coverage_worker.mjs tests/test_roleplay_compaction_budget_worker.mjs tests/test_roleplay_compaction_model_worker.mjs tests/test_roleplay_stream_validation_worker.mjs tests/test_aihubmix_worker.mjs tests/test_gpt_image_moderation_worker.mjs tests/test_dashboard_scripts.mjs tests/test_operations_dashboard.mjs tests/test_opencode_reasoning_worker.mjs tests/test_opencode_thinking_worker.mjs tests/test_nanogpt_reasoning_worker.mjs tests/test_nanogpt_speed_routing_worker.mjs tests/test_sse_heartbeat_worker.mjs tests/test_roleplay_worker.mjs tests/test_roleplay_config_worker.mjs tests/test_roleplay_glm_failover_worker.mjs tests/test_roleplay_intelligence_worker.mjs tests/test_roleplay_memory_worker.mjs tests/test_roleplay_compatibility_worker.mjs tests/test_roleplay_history_reconciliation_worker.mjs tests/test_roleplay_output_contract_analysis.mjs tests/test_roleplay_output_contract_cleanup_worker.mjs tests/test_roleplay_output_contract_nonstream_worker.mjs tests/test_roleplay_output_contract_worker.mjs tests/test_roleplay_performance_worker.mjs tests/test_roleplay_quality_routing_worker.mjs tests/test_roleplay_reasoning_worker.mjs tests/test_roleplay_session_storage_worker.mjs tests/test_roleplay_transport_failure_worker.mjs tests/test_roleplay_unlimited_worker.mjs tests/test_roleplay_workbench_worker.mjs tests/test_roleplay_candidate_context_worker.mjs tests/test_roleplay_operator_worker.mjs tests/test_workbench_stream.mjs tests/test_intelligence_worker.mjs tests/test_intelligence_d1_worker.mjs tests/test_intelligence_auth_worker.mjs tests/test_intelligence_outbound_worker.mjs tests/test_intelligence_operator.mjs tests/test_control_users_worker.mjs tests/test_control_audit_worker.mjs tests/test_access_sso_worker.mjs tests/test_d1_schema_worker.mjs tests/test_shadow_eval_worker.mjs tests/test_auto_routes_worker.mjs tests/test_ai_outbound_worker.mjs tests/test_knowledge_*.mjs tests/test_control_state_worker.mjs tests/test_route_health_worker.mjs tests/test_gateway_mcp_worker.mjs tests/test_usage_ledger_worker.mjs tests/test_media_outbound_worker.mjs tests/test_media_files_worker.mjs tests/test_media_jobs_worker.mjs tests/test_tool_repair_worker.mjs tests/test_prometheus_worker.mjs tests/test_evaluation_noise_floor_worker.mjs tests/test_prompt_templates_worker.mjs tests/test_config_snapshots_worker.mjs tests/test_managed_idempotency_worker.mjs tests/test_wave3_hub_worker.mjs tests/test_wave3_transport_worker.mjs tests/test_admission_leases_worker.mjs tests/test_edge_request_metrics_worker.mjs tests/test_request_cancellation_worker.mjs tests/test_mcp_contract_drift_worker.mjs tests/test_zero_content_retention_worker.mjs tests/test_prompt_cache_cost_worker.mjs tests/test_config_revision_sync_worker.mjs tests/test_gateway_alerts_worker.mjs tests/test_session_tiers_worker.mjs tests/test_usage_reservations_worker.mjs tests/test_shared_generation_cache_worker.mjs tests/test_deferred_tools_worker.mjs tests/test_generation_deadline_worker.mjs tests/test_wave4_worker.mjs tests/test_responses_state_worker.mjs tests/test_gateway_batches_worker.mjs tests/test_semantic_generation_cache_worker.mjs tests/test_canary_traffic_worker.mjs tests/test_context_pages_worker.mjs tests/test_observability_exports_worker.mjs tests/test_prompt_injection_worker.mjs tests/test_pii_rehydration_worker.mjs tests/test_wave5_worker.mjs tests/test_usage_receipts_worker.mjs tests/test_learned_cooldown_worker.mjs tests/test_stream_cost_breaker_worker.mjs tests/test_realtime_worker.mjs tests/test_latency_slo_worker.mjs tests/test_context_canary_worker.mjs tests/test_enterprise_contracts_worker.mjs tests/test_wave6_worker.mjs tests/test_wave7_worker.mjs tests/test_payment_billing_worker.mjs tests/test_credits_ledger_worker.mjs tests/test_scim_provisioning_worker.mjs tests/test_saml_federation_worker.mjs tests/test_tenant_governance_worker.mjs tests/test_tenant_hierarchy_worker.mjs", + "test:worker": "node --import ./tests/helpers/network_guard.mjs --test tests/test_cloudflare_worker.mjs tests/test_child_safety_worker.mjs tests/test_child_safety_dashboard.mjs tests/test_runtime_feature_flags_worker.mjs tests/test_anthropic_worker.mjs tests/test_haiku_intelligence_worker.mjs tests/test_cascades_worker.mjs tests/test_cascades_dashboard.mjs tests/test_image_quality_worker.mjs tests/test_secret_scan_worker.mjs tests/test_secret_firewall_worker.mjs tests/test_secret_firewall_coverage_worker.mjs tests/test_roleplay_compaction_budget_worker.mjs tests/test_roleplay_compaction_model_worker.mjs tests/test_roleplay_stream_validation_worker.mjs tests/test_aihubmix_worker.mjs tests/test_gpt_image_moderation_worker.mjs tests/test_dashboard_scripts.mjs tests/test_operations_dashboard.mjs tests/test_opencode_reasoning_worker.mjs tests/test_opencode_thinking_worker.mjs tests/test_nanogpt_reasoning_worker.mjs tests/test_nanogpt_speed_routing_worker.mjs tests/test_sse_heartbeat_worker.mjs tests/test_roleplay_worker.mjs tests/test_roleplay_config_worker.mjs tests/test_roleplay_glm_failover_worker.mjs tests/test_roleplay_intelligence_worker.mjs tests/test_roleplay_memory_worker.mjs tests/test_roleplay_compatibility_worker.mjs tests/test_roleplay_history_reconciliation_worker.mjs tests/test_roleplay_output_contract_analysis.mjs tests/test_roleplay_output_contract_cleanup_worker.mjs tests/test_roleplay_output_contract_nonstream_worker.mjs tests/test_roleplay_output_contract_worker.mjs tests/test_roleplay_performance_worker.mjs tests/test_roleplay_quality_routing_worker.mjs tests/test_roleplay_reasoning_worker.mjs tests/test_roleplay_session_storage_worker.mjs tests/test_roleplay_transport_failure_worker.mjs tests/test_roleplay_unlimited_worker.mjs tests/test_roleplay_workbench_worker.mjs tests/test_roleplay_candidate_context_worker.mjs tests/test_roleplay_operator_worker.mjs tests/test_workbench_stream.mjs tests/test_intelligence_worker.mjs tests/test_intelligence_d1_worker.mjs tests/test_intelligence_auth_worker.mjs tests/test_intelligence_outbound_worker.mjs tests/test_intelligence_operator.mjs tests/test_control_users_worker.mjs tests/test_control_audit_worker.mjs tests/test_access_sso_worker.mjs tests/test_d1_schema_worker.mjs tests/test_shadow_eval_worker.mjs tests/test_auto_routes_worker.mjs tests/test_ai_outbound_worker.mjs tests/test_knowledge_*.mjs tests/test_control_state_worker.mjs tests/test_route_health_worker.mjs tests/test_gateway_mcp_worker.mjs tests/test_usage_ledger_worker.mjs tests/test_media_outbound_worker.mjs tests/test_media_files_worker.mjs tests/test_media_jobs_worker.mjs tests/test_tool_repair_worker.mjs tests/test_prometheus_worker.mjs tests/test_evaluation_noise_floor_worker.mjs tests/test_prompt_templates_worker.mjs tests/test_config_snapshots_worker.mjs tests/test_managed_idempotency_worker.mjs tests/test_wave3_hub_worker.mjs tests/test_wave3_transport_worker.mjs tests/test_admission_leases_worker.mjs tests/test_edge_request_metrics_worker.mjs tests/test_request_cancellation_worker.mjs tests/test_mcp_contract_drift_worker.mjs tests/test_zero_content_retention_worker.mjs tests/test_prompt_cache_cost_worker.mjs tests/test_config_revision_sync_worker.mjs tests/test_gateway_alerts_worker.mjs tests/test_session_tiers_worker.mjs tests/test_usage_reservations_worker.mjs tests/test_shared_generation_cache_worker.mjs tests/test_deferred_tools_worker.mjs tests/test_generation_deadline_worker.mjs tests/test_wave4_worker.mjs tests/test_responses_state_worker.mjs tests/test_gateway_batches_worker.mjs tests/test_semantic_generation_cache_worker.mjs tests/test_canary_traffic_worker.mjs tests/test_context_pages_worker.mjs tests/test_observability_exports_worker.mjs tests/test_prompt_injection_worker.mjs tests/test_pii_rehydration_worker.mjs tests/test_wave5_worker.mjs tests/test_usage_receipts_worker.mjs tests/test_learned_cooldown_worker.mjs tests/test_stream_cost_breaker_worker.mjs tests/test_realtime_worker.mjs tests/test_latency_slo_worker.mjs tests/test_context_canary_worker.mjs tests/test_enterprise_contracts_worker.mjs tests/test_wave6_worker.mjs tests/test_wave7_worker.mjs tests/test_payment_billing_worker.mjs tests/test_credits_ledger_worker.mjs tests/test_scim_provisioning_worker.mjs tests/test_saml_federation_worker.mjs tests/test_tenant_governance_worker.mjs tests/test_tenant_hierarchy_worker.mjs tests/test_frontend_build.mjs tests/test_deploy_pipeline.mjs tests/test_free_routes_worker.mjs tests/test_client_headers_worker.mjs tests/test_key_controls_expand_worker.mjs tests/test_key_data_preflight.mjs", "db:migrate": "python scripts/validate_sqlite_schema.py --apply", "db:validate": "python scripts/validate_sqlite_schema.py", "security:secrets": "python scripts/check_static_secrets.py", @@ -18,7 +18,7 @@ "cf:dry-run": "wrangler deploy --dry-run", "cf:tail": "wrangler tail", "test:knowledge": "node --test tests/test_knowledge_*.mjs", - "test:integration": "node --test tests/test_d1_integration.mjs tests/test_config_revision_sync_integration.mjs", + "test:integration": "node --import ./tests/helpers/network_guard.mjs --test tests/test_d1_integration.mjs tests/test_config_revision_sync_integration.mjs tests/test_key_controls_expand_worker.mjs", "cf:knowledge:dry-run": "wrangler deploy --config wrangler.knowledge.jsonc --dry-run", "cf:knowledge:deploy": "wrangler deploy --config wrangler.knowledge.jsonc", "cf:migrations:check": "node scripts/verify_d1_migrations.mjs --check", diff --git a/providers/data-classes.json b/providers/data-classes.json new file mode 100644 index 00000000..c1ea4dd1 --- /dev/null +++ b/providers/data-classes.json @@ -0,0 +1,48 @@ +{ + "a6api": "reseller", + "aihubmix": "reseller", + "aimlapi": "reseller", + "anthropic": "official", + "azure": "official", + "ce-claude-kiro": "subscription", + "ce-claude-kiro-cheap": "subscription", + "ce-claude-max": "subscription", + "ce-gpt-plus": "subscription", + "ce-gpt-pro": "subscription", + "ce-grok-heavy": "subscription", + "ce-image": "subscription", + "cerebras": "official", + "chutes": "reseller", + "cline-pass": "subscription", + "cloudflare": "official", + "codex-easy": "subscription", + "ephone": "reseller", + "free": "free", + "gemini": "free", + "gemma": "free", + "gguu": "reseller", + "gguu-grok": "reseller", + "googleai": "official", + "groq": "official", + "hyperbolic": "official", + "kimi-code": "subscription", + "latix": "reseller", + "linkapi": "reseller", + "llamacpp": "free", + "lmstudio": "free", + "mimo": "subscription", + "nanogpt": "reseller", + "navyai": "reseller", + "nineteen": "reseller", + "ollama": "free", + "openai": "official", + "opencode": "subscription", + "openrouter": "official", + "palm": "free", + "sambanova": "official", + "scaleway": "official", + "sglang": "free", + "together": "official", + "vllm": "free", + "xai": "official" +} diff --git a/routes/core.py b/routes/core.py index cd4d2013..590135c4 100644 --- a/routes/core.py +++ b/routes/core.py @@ -1,3 +1,4 @@ +from services import key_controls_capability import json import logging import os @@ -529,12 +530,17 @@ def health_check(): response = jsonify( { "status": "healthy", + "key_controls": key_controls_capability.CAPABILITY, "config": { "host": os.environ.get("SERVER_HOST", Config.DEFAULT_HOST), "port": int(os.environ.get("SERVER_PORT", Config.DEFAULT_PORT)), }, } ) + from services.release_status import container_release + stamp = container_release()["build_id"] + if stamp: + response.headers["X-MultiLLM-Container-Build"] = stamp response.headers["Cache-Control"] = "no-store" return response, 200 except Exception as error: diff --git a/routes/media.py b/routes/media.py index 87037629..53af564b 100644 --- a/routes/media.py +++ b/routes/media.py @@ -206,7 +206,9 @@ def create(candidate: str): try: return video_generation.create_job(provider, provider_model, video_request, key, base_url, owner) except APIError as error: - if error.status_code in _PRE_GENERATION_STATUSES: + if error.status_code in _PRE_GENERATION_STATUSES and (error.payload or {}).get("error") not in { + "provider_credential_not_allowed", "attested_credential_unavailable", + "key_data_controls_unavailable", "provider_class_not_allowed"}: raise AutoRouteCandidateUnavailable(error.message) from error raise @@ -282,7 +284,7 @@ def video_content(job_id: str): pass key, base_url = video_credentials(job["p"]) url, headers = video_generation.content_source(job, key, base_url) - return video_generation.stream_content(url, headers) + return video_generation.stream_content(url, headers, provider=job["p"]) @app.route("/v1/media/providers", methods=["GET", "OPTIONS"]) @csrf.exempt diff --git a/routes/media_narration.py b/routes/media_narration.py index b4e5ab62..2910a583 100644 --- a/routes/media_narration.py +++ b/routes/media_narration.py @@ -19,6 +19,7 @@ import re import time from collections.abc import Callable +from services.provider_dispatch import capture_dispatch from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass @@ -295,7 +296,7 @@ def run(self) -> dict: break if shots: with ThreadPoolExecutor(max_workers=min(SHOT_PARALLELISM, len(shots))) as pool: - futures = [pool.submit(copy_current_request_context(lambda shot=shot: self.shot(shot, model))) + futures = [pool.submit(capture_dispatch(copy_current_request_context(lambda shot=shot: self.shot(shot, model)))) for shot in shots] for shot, future in zip(shots, futures, strict=True): entries[shot.index] = future.result() diff --git a/scripts/check_worker_bundle.mjs b/scripts/check_worker_bundle.mjs new file mode 100644 index 00000000..4d1cd53c --- /dev/null +++ b/scripts/check_worker_bundle.mjs @@ -0,0 +1,35 @@ +#!/usr/bin/env node +/** Bundle validation without remote release admission. This command cannot deploy. */ +import { execFileSync } from "node:child_process"; +import { readFileSync, writeFileSync, unlinkSync } from "node:fs"; +import { resolve } from "node:path"; +import { randomUUID } from "node:crypto"; +import { parseArgs } from "node:util"; +import { fileURLToPath, pathToFileURL } from "node:url"; + +const ROOT = new URL("../", import.meta.url); +export function bundleConfig(config) { + const result = structuredClone(config); + delete result.build; + return result; +} +export function bundleArguments(config, outdir, main) { + return ["exec", "wrangler", "deploy", "--config", resolve(config), "--dry-run", + ...(main ? ["--containers-rollout=none"] : []), "--outdir", resolve(outdir)]; +} +function main() { + const { values } = parseArgs({ options: { config: { type: "string" }, outdir: { type: "string" } } }); + if (!["wrangler.jsonc", "wrangler.knowledge.jsonc"].includes(values.config) || !values.outdir) { + throw new Error("Pass --config wrangler.jsonc|wrangler.knowledge.jsonc and --outdir"); + } + const source = JSON.parse(readFileSync(new URL(values.config, ROOT))); + // Keep relative paths rooted beside the real configuration. Never change the + // committed release config or provide a general build-admission bypass flag. + const temporary = new URL(`.wrangler-bundle-${randomUUID()}.jsonc`, ROOT); + writeFileSync(temporary, JSON.stringify(bundleConfig(source)), { flag: "wx", mode: 0o600 }); + try { + execFileSync("pnpm", bundleArguments(fileURLToPath(temporary), values.outdir, values.config === "wrangler.jsonc"), + { cwd: ROOT, stdio: "inherit", env: { ...process.env, WRANGLER_SEND_METRICS: "false" } }); + } finally { unlinkSync(temporary); } +} +if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) main(); diff --git a/scripts/container_rollout.mjs b/scripts/container_rollout.mjs new file mode 100644 index 00000000..8aaea936 --- /dev/null +++ b/scripts/container_rollout.mjs @@ -0,0 +1,87 @@ +#!/usr/bin/env node +/** Verify application detail and the responding image, never the stale list version. */ +import { execFileSync } from "node:child_process"; +import { randomUUID } from "node:crypto"; +import { readFileSync, writeFileSync } from "node:fs"; +import { parseArgs } from "node:util"; +import { pathToFileURL } from "node:url"; +import { deploymentTarget } from "./deployment_target.mjs"; +import { localRelease, parseJson } from "./deploy_status.mjs"; + +const STAMP = /^[0-9a-f]{64}$/; +export function snapshot(application) { + if (!application?.id || !Number.isInteger(application.version) || application.version < 1 + || typeof application.configuration?.image !== "string") throw Error("Container detail has no verified version/image"); + return { id: application.id, version: application.version, image: application.configuration.image }; +} +export function releaseConfig(config, releaseId = randomUUID()) { + if (config.containers?.length !== 1 || config.containers[0].image !== "./Dockerfile") { + throw Error("Verified deployment requires one Dockerfile Container"); + } + return { ...config, containers: config.containers.map(container => ({ ...container, + image_vars: { ...container.image_vars, MULTILLM_RELEASE_IMAGE_ID: releaseId } })) }; +} +export async function verifyRollout({ before, readApplication, probe, expected, waitSeconds = 600, + now = Date.now, sleep = ms => new Promise(resolve => setTimeout(resolve, ms)) }) { + if (!STAMP.test(expected)) throw Error("Missing expected source fingerprint"); + const deadline = now() + waitSeconds * 1000; + for (;;) { + try { + const after = snapshot(await readApplication()); + const ready = await probe(); + if (after.version > before.version && after.image !== before.image && ready.ok + && ready.worker === expected && ready.container === expected) return after; + } catch { /* A swap can temporarily have no live Container; the deadline still applies. */ } + if (now() >= deadline) throw Error("Container rollout unverified: version/image did not advance or Worker/Container fingerprints differ; stop deployment"); + await sleep(Math.min(10000, Math.max(1, deadline - now()))); + } +} +export function captureApplication(target, run) { + const container = target.main.containers?.[0]; + const name = container?.name || `${target.main.name}-${container?.class_name?.toLowerCase()}`; + const applications = parseJson(run(["containers", "list", "--json", "--config", target.mainConfig])); + if (!Array.isArray(applications)) throw Error("Container list is not an array"); + const matches = applications.filter(application => application.name === name); + if (matches.length !== 1) throw Error("Select an existing, unique Container application before deployment"); + // Only use the list for identity. Its version/image may lag application detail. + return snapshot(parseJson(run(["containers", "info", matches[0].id, "--json", "--config", target.mainConfig]))); +} +async function main(argv) { + const { values: options } = parseArgs({ args: argv, options: { + prepare: { type: "boolean" }, verify: { type: "boolean" }, state: { type: "string" }, output: { type: "string" }, + config: { type: "string" }, "knowledge-config": { type: "string" }, wait: { type: "string", default: "600" }, + } }); + if (!options.state || options.prepare === options.verify || !Number.isFinite(Number(options.wait)) || Number(options.wait) < 0) { + throw Error("Use --prepare --state --output , or --verify --state "); + } + const target = deploymentTarget({ mainConfig: options.config, knowledgeConfig: options["knowledge-config"] }); + if (!URL.canParse(target.origin ?? "")) throw Error("Configure DEPLOY_STATUS_URL or WORKBENCH_WORKER_URL"); + const run = args => execFileSync("npx", ["--no-install", "wrangler", ...args], { + encoding: "utf8", stdio: ["ignore", "pipe", "pipe"], timeout: 60000, + env: { ...process.env, WRANGLER_SEND_METRICS: "false", CI: "true" }, + }); + if (options.prepare) { + if (!options.output) throw Error("Missing release config output"); + const before = captureApplication(target, run); + const expected = localRelease(process.cwd(), execFileSync).fingerprint; + if (!STAMP.test(expected ?? "")) throw Error("Could not fingerprint the release"); + writeFileSync(options.output, JSON.stringify(releaseConfig(target.main), null, 2) + "\n"); + writeFileSync(options.state, JSON.stringify({ before, expected })); + return; + } + const { before, expected } = JSON.parse(readFileSync(options.state, "utf8")); + await verifyRollout({ before, expected, waitSeconds: Number(options.wait), + readApplication: async () => parseJson(run(["containers", "info", before.id, "--json", "--config", target.mainConfig])), + probe: async () => { + const response = await fetch(new URL("/ready", target.origin), { redirect: "error", signal: AbortSignal.timeout(30000), + headers: { "Cache-Control": "no-cache" } }); + await response.body?.cancel(); + return { ok: response.ok, worker: response.headers.get("X-MultiLLM-Worker-Build"), + container: response.headers.get("X-MultiLLM-Container-Build") }; + } }); + console.log("Container version/image advanced; responding Worker and Container source fingerprints match."); +} +if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) { + try { await main(process.argv.slice(2)); } + catch { console.error("Container deployment verification failed. Check the selected application detail, rollout and release fingerprints; do not enable restrictions."); process.exitCode = 1; } +} diff --git a/scripts/deploy.sh b/scripts/deploy.sh index 3ecb8ad5..02fd227b 100644 --- a/scripts/deploy.sh +++ b/scripts/deploy.sh @@ -7,7 +7,16 @@ set -euo pipefail cd "$(dirname "$0")/.." -DATABASE="${D1_DATABASE:-multillm-intelligence}" +DATABASE="${D1_DATABASE:?Set D1_DATABASE explicitly to the configured target database}" +MAIN_CONFIG="${MAIN_WORKER_CONFIG:-wrangler.jsonc}" +KNOWLEDGE_CONFIG="${KNOWLEDGE_WORKER_CONFIG:-wrangler.knowledge.jsonc}" +export MAIN_WORKER_CONFIG="$MAIN_CONFIG" KNOWLEDGE_WORKER_CONFIG="$KNOWLEDGE_CONFIG" +# Validate target identity before any migration, upload or remote read. +node --input-type=module -e 'import { deploymentTarget } from "./scripts/deployment_target.mjs"; deploymentTarget({ mainConfig: process.env.MAIN_WORKER_CONFIG, knowledgeConfig: process.env.KNOWLEDGE_WORKER_CONFIG });' +state="$(mktemp)" +release_config="$(mktemp .release-wrangler-XXXXXX.jsonc)" +trap 'rm -f -- "$state" "$release_config"' EXIT +node scripts/container_rollout.mjs --prepare --state "$state" --output "$release_config" --config "$MAIN_CONFIG" --knowledge-config "$KNOWLEDGE_CONFIG" annotations=() if commit="$(git rev-parse HEAD 2>/dev/null)"; then # The deploy build rewrites worker/build-id.mjs, which is not a source change. @@ -18,8 +27,11 @@ if commit="$(git rev-parse HEAD 2>/dev/null)"; then fi fi +node scripts/key_data_preflight.mjs --check node scripts/verify_d1_migrations.mjs --check -npx --no-install wrangler d1 migrations apply "$DATABASE" --remote +npx --no-install wrangler d1 migrations apply "$DATABASE" --remote --config "$MAIN_CONFIG" node scripts/verify_d1_migrations.mjs "$DATABASE" -npx --no-install wrangler deploy --config wrangler.knowledge.jsonc ${annotations[@]+"${annotations[@]}"} -npx --no-install wrangler deploy ${annotations[@]+"${annotations[@]}"} "$@" +node scripts/key_data_preflight.mjs --prepare +npx --no-install wrangler deploy --config "$KNOWLEDGE_CONFIG" ${annotations[@]+"${annotations[@]}"} +npx --no-install wrangler deploy --config "$release_config" ${annotations[@]+"${annotations[@]}"} "$@" +node scripts/container_rollout.mjs --verify --state "$state" --config "$MAIN_CONFIG" --knowledge-config "$KNOWLEDGE_CONFIG" diff --git a/scripts/deploy_status.mjs b/scripts/deploy_status.mjs index 5a04b841..7eadfba2 100644 --- a/scripts/deploy_status.mjs +++ b/scripts/deploy_status.mjs @@ -17,11 +17,13 @@ * main Worker is deployed by Workers Builds, so a mismatch there is only a warning. */ import { execFileSync } from "node:child_process"; -import { appendFileSync, readFileSync } from "node:fs"; +import { appendFileSync } from "node:fs"; import { fileURLToPath, pathToFileURL } from "node:url"; import { parseArgs } from "node:util"; import { localMigrations, pendingMigrations, remoteAppliedMigrations } from "./verify_d1_migrations.mjs"; +import { deploymentTarget } from "./deployment_target.mjs"; + const ROOT = fileURLToPath(new URL("../", import.meta.url)); const COMMIT = /^[0-9a-f]{40}$/; const FINGERPRINT = /^[0-9a-f]{64}$/; @@ -29,7 +31,7 @@ export const WORKERS = Object.freeze({ main: { label: "Main Worker", name: "multillm-proxy", config: "wrangler.jsonc", fingerprinted: true }, knowledge: { label: "Knowledge Worker", name: "multillm-knowledge", config: "wrangler.knowledge.jsonc", fingerprinted: false }, }); -const USAGE = `Usage: node scripts/deploy_status.mjs [--json] [--url ] [--database ] +const USAGE = `Usage: node scripts/deploy_status.mjs [--json] [--url ] --database [--config ] [--knowledge-config ] node scripts/deploy_status.mjs --verify --commit --main-worker actions|workers-builds [--wait ] [--markdown ]`; @@ -174,9 +176,9 @@ export function localRelease(root, exec) { fingerprint: FINGERPRINT.test(fingerprint ?? "") ? fingerprint : null }; } -export async function collectStatus({ run, api, fetchImpl, origin, database, local }) { +export async function collectStatus({ run, api, fetchImpl, origin, database, local, workers: definitions = WORKERS }) { const workers = {}; - for (const [key, worker] of Object.entries(WORKERS)) workers[key] = await workerStatus(worker, { run, api }); + for (const [key, worker] of Object.entries(definitions)) workers[key] = await workerStatus(worker, { run, api }); return { checked_at: new Date().toISOString(), local, workers, d1: d1Status(database, run), ready: await probeReady(origin, fetchImpl) }; } @@ -255,7 +257,7 @@ function describeLocal(local) { function rows(status) { return [ ["Local checkout", describeLocal(status.local)], - ...Object.entries(WORKERS).map(([key, worker]) => [worker.label, describeWorker(status.workers[key], status.local)]), + ...Object.entries(status.workers).map(([key, worker]) => [WORKERS[key].label, describeWorker(worker, status.local)]), ["D1 migrations", describeD1(status.d1)], ["/ready", `${status.ready.status ?? "no response"} ${status.ready.detail} (${status.ready.url})`], ]; @@ -289,18 +291,13 @@ export async function verifyUntilSettled({ collect, evaluateOptions, waitSeconds } } -function defaultOrigin() { - const config = JSON.parse(readFileSync(new URL("../wrangler.jsonc", import.meta.url), "utf8")); - return config.vars?.WORKBENCH_WORKER_URL; -} - async function main(argv) { let options; try { ({ values: options } = parseArgs({ args: argv, options: { json: { type: "boolean" }, verify: { type: "boolean" }, commit: { type: "string" }, url: { type: "string" }, "main-worker": { type: "string" }, wait: { type: "string", default: "0" }, markdown: { type: "string" }, - database: { type: "string", default: "multillm-intelligence" }, help: { type: "boolean" }, + database: { type: "string" }, config: { type: "string" }, "knowledge-config": { type: "string" }, help: { type: "boolean" }, } })); } catch (error) { console.error(`${error.message}\n${USAGE}`); @@ -317,18 +314,21 @@ async function main(argv) { console.error(`--verify needs --commit <40-character sha>; --main-worker is actions or workers-builds.\n${USAGE}`); return 2; } - const origin = options.url || process.env.DEPLOY_STATUS_URL || defaultOrigin(); + let target; + try { target = deploymentTarget({ database: options.database, mainConfig: options.config, knowledgeConfig: options["knowledge-config"] }); } + catch (error) { console.error(error.message); return 2; } + const origin = options.url || target.origin; if (!URL.canParse(origin ?? "")) { console.error(`Pass the Worker origin with --url; ${JSON.stringify(origin ?? null)} is not a URL.`); return 2; } - const run = args => execFileSync("npx", ["--no-install", "wrangler", ...args], { + const run = args => execFileSync("npx", ["--no-install", "wrangler", ...args, ...(args[0] === "d1" ? ["--config", target.mainConfig] : [])], { cwd: ROOT, encoding: "utf8", stdio: ["ignore", "pipe", "pipe"], maxBuffer: 16 * 1024 * 1024, env: { ...process.env, WRANGLER_SEND_METRICS: "false", FORCE_COLOR: "0" }, }); const local = localRelease(ROOT, execFileSync); const collect = () => collectStatus({ run, api: cloudflareApi(process.env), fetchImpl: fetch, origin, - database: options.database, local }); + database: target.database, workers: target.workers, local }); const { status, checks } = options.verify ? await verifyUntilSettled({ collect, evaluateOptions: { commit: options.commit, mainWorker }, waitSeconds: wait, sleep: ms => new Promise(resolve => setTimeout(resolve, ms)), now: Date.now }) diff --git a/scripts/deployment_target.mjs b/scripts/deployment_target.mjs new file mode 100644 index 00000000..9451954c --- /dev/null +++ b/scripts/deployment_target.mjs @@ -0,0 +1,22 @@ +/** Resolve explicit deployment targets before any remote command. */ +import { readFileSync } from "node:fs"; +import { resolve } from "node:path"; + +export function deploymentTarget({ database, mainConfig = "wrangler.jsonc", knowledgeConfig = "wrangler.knowledge.jsonc", + env = process.env, read = path => JSON.parse(readFileSync(path, "utf8")) } = {}) { + mainConfig = resolve(mainConfig); knowledgeConfig = resolve(knowledgeConfig); + const main = read(mainConfig), knowledge = read(knowledgeConfig); + database ||= env.D1_DATABASE; + if (!database) throw Error("Select D1 explicitly with --database or D1_DATABASE; no default database is used"); + const binding = config => config.d1_databases?.find(item => item.binding === "INTELLIGENCE_DB"); + const first = binding(main), second = binding(knowledge); + if (!first?.database_id || first.database_name !== database || second?.database_name !== database + || first.database_id !== second.database_id) throw Error("Selected D1 must match INTELLIGENCE_DB in both configurations"); + if (!main.name || !knowledge.name) throw Error("Both Worker names must be configured"); + return { database, main, knowledge, mainConfig, knowledgeConfig, + origin: env.DEPLOY_STATUS_URL || main.vars?.WORKBENCH_WORKER_URL, + workers: { + main: { label: "Main Worker", name: main.name, config: mainConfig, fingerprinted: true }, + knowledge: { label: "Knowledge Worker", name: knowledge.name, config: knowledgeConfig, fingerprinted: false }, + } }; +} diff --git a/scripts/key_data_preflight.mjs b/scripts/key_data_preflight.mjs new file mode 100644 index 00000000..d4efb1af --- /dev/null +++ b/scripts/key_data_preflight.mjs @@ -0,0 +1,112 @@ +#!/usr/bin/env node +/** Release admission, before uploading a Worker or replacing its Container. */ +import { execFileSync } from "node:child_process"; +import { readFileSync } from "node:fs"; +import { pathToFileURL } from "node:url"; +import { deploymentTarget } from "./deployment_target.mjs"; +import { localMigrations, remoteAppliedMigrations } from "./verify_d1_migrations.mjs"; + +const ROOT = new URL("../", import.meta.url); +const MAIN_CONFIG = process.env.MAIN_WORKER_CONFIG || "wrangler.jsonc"; +const KNOWLEDGE_CONFIG = process.env.KNOWLEDGE_WORKER_CONFIG || "wrangler.knowledge.jsonc"; +const REQUIRED_SECRETS = ["ADMIN_API_KEY", "FLASK_SECRET_KEY", "JWT_SECRET"]; +const REQUIRED_COLUMNS = { + control_users: ["official_only", "forced_provider_preferences", "max_stream_cost_microusd"], + intelligence_principals: ["official_only", "forced_provider_preferences"], +}; +class PreflightError extends Error {} +const requireThat = (condition, message) => { if (!condition) throw new PreflightError(message); }; + +export function checkConfig(main, knowledge) { + const database = config => config.d1_databases?.find(item => item.binding === "INTELLIGENCE_DB"); + const first = database(main), second = database(knowledge); + requireThat(main.vars?.AUTH_STORAGE_BACKEND === "d1", "AUTH_STORAGE_BACKEND must explicitly select d1"); + requireThat(first?.database_id && first.database_name && first.migrations_dir === "intelligence-migrations", + "Main Worker needs the migrated INTELLIGENCE_DB binding"); + requireThat(second?.database_id === first.database_id && second?.database_name === first.database_name + && second?.migrations_dir === first.migrations_dir, "Both Workers must bind the same migrated D1"); + requireThat(main.services?.some(item => item.binding === "KNOWLEDGE_SERVICE" && item.service === knowledge.name), + "KNOWLEDGE_SERVICE must target the configured private Worker"); + requireThat(knowledge.workers_dev === false && knowledge.preview_urls === false, + "Knowledge Worker must remain private"); + requireThat(main.containers?.some(item => item.class_name === "MultiLLMProxyContainer" && item.image === "./Dockerfile") + && main.durable_objects?.bindings?.some(item => item.name === "MULTILLM_PROXY_CONTAINER" + && item.class_name === "MultiLLMProxyContainer"), "Main Worker needs its Container binding and image"); + return first.database_name; +} + +export function checkPrerequisites({ main, knowledge, run, migrations = localMigrations() }) { + const database = checkConfig(main, knowledge); + const applied = remoteAppliedMigrations(database, run); + requireThat(migrations.every(name => applied.includes(name)) && applied.every(name => migrations.includes(name)), + "D1 migration ledger must exactly match this release"); + for (const [table, required] of Object.entries(REQUIRED_COLUMNS)) { + const results = JSON.parse(run(["d1", "execute", database, "--remote", "--json", "--command", `PRAGMA table_info(${table})`])); + requireThat(Array.isArray(results) && results.every(item => item.success === true && Array.isArray(item.results)), + "Account schema could not be verified"); + const columns = results.flatMap(item => item.results).map(item => item.name); + requireThat(required.every(name => columns.includes(name)), "Account policy columns are missing from D1"); + } + const secrets = JSON.parse(run(["secret", "list", "--config", MAIN_CONFIG])); + requireThat(REQUIRED_SECRETS.every(name => main.vars?.[name] || secrets.some(item => item.name === name)), + "Main Worker needs ADMIN_API_KEY, FLASK_SECRET_KEY and JWT_SECRET bindings"); +} + +export function sourceCapability(commit) { + if (!/^[0-9a-f]{40}$/.test(commit)) return null; + try { + const source = execFileSync("git", ["show", `${commit}:worker/key-controls-capability.mjs`], + { cwd: ROOT, encoding: "utf8", stdio: ["ignore", "pipe", "pipe"] }); + const match = /protocol: ([12]), enforcement: (true|false)/.exec(source); + return match ? { protocol: Number(match[1]), enforcement: match[2] === "true" } : null; + } catch { return null; } +} + +export function checkRelease(options) { + checkPrerequisites(options); + const { commit, run } = options; + requireThat(/^[0-9a-f]{40}$/.test(commit), "Release needs a full commit identity"); + const deployment = JSON.parse(run(["deployments", "status", "--json", "--config", KNOWLEDGE_CONFIG])); + requireThat(deployment.versions?.length === 1 && deployment.versions[0].percentage === 100, + "Private Worker must have one fully active version"); + const version = JSON.parse(run(["versions", "view", deployment.versions[0].version_id, "--json", "--config", KNOWLEDGE_CONFIG])); + const capability = (options.sourceCapability ?? sourceCapability)(version.annotations?.["workers/tag"]); + requireThat(capability?.protocol >= 1, "Deploy and verify the expansion private Worker before the main Worker"); +} + +function main(args) { + const mainConfig = JSON.parse(readFileSync(new URL(MAIN_CONFIG, ROOT))); + const knowledge = JSON.parse(readFileSync(new URL(KNOWLEDGE_CONFIG, ROOT))); + try { + checkConfig(mainConfig, knowledge); + if (args[0] === "--check" || process.env.WRANGLER_COMMAND === "dev") { + console.log("Release configuration checked (no deployment admission)."); + return 0; + } + requireThat(["--prepare", "--deploy"].includes(args[0]), "Use --check, --prepare or --deploy"); + deploymentTarget({ mainConfig: MAIN_CONFIG, knowledgeConfig: KNOWLEDGE_CONFIG }); + const git = args => execFileSync("git", args, { cwd: ROOT, encoding: "utf8", stdio: ["ignore", "pipe", "pipe"] }).trim(); + const commit = git(["rev-parse", "HEAD"]); + requireThat(!git(["status", "--porcelain", "--untracked-files=no", "--", ".", ":(exclude)worker/build-id.mjs"]), + "Deploy from a clean committed checkout"); + const run = args => execFileSync("npx", ["--no-install", "wrangler", ...args, ...(args[0] === "d1" ? ["--config", MAIN_CONFIG] : [])], { + cwd: ROOT, encoding: "utf8", stdio: ["ignore", "pipe", "pipe"], timeout: 60000, + env: { ...process.env, WRANGLER_SEND_METRICS: "false" }, + }); + const options = { main: mainConfig, knowledge, commit, run }; + if (args[0] === "--prepare") { + checkPrerequisites(options); + console.log("Release prerequisites checked before private Worker deployment."); + return 0; + } + checkRelease(options); + console.log("Release admitted: migrations, account schema, required bindings and private Worker expansion verified."); + return 0; + } catch (error) { + // CLI output may contain environment details; print only our own diagnostics. + const message = error instanceof PreflightError ? error.message : "Could not verify release prerequisites"; + console.error(`Deployment preflight refused: ${message}. No main upload is admitted.`); + return 1; + } +} +if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) process.exitCode = main(process.argv.slice(2)); diff --git a/scripts/verify_d1_migrations.mjs b/scripts/verify_d1_migrations.mjs index c6f12554..7bc73191 100644 --- a/scripts/verify_d1_migrations.mjs +++ b/scripts/verify_d1_migrations.mjs @@ -2,8 +2,9 @@ * Guards for the D1 migrations in intelligence-migrations/. * * node scripts/verify_d1_migrations.mjs --check Static checks, no network (CI). - * node scripts/verify_d1_migrations.mjs --pending [database] Print migrations not yet applied remotely. - * node scripts/verify_d1_migrations.mjs [database] Fail unless every file is applied remotely. + * node scripts/verify_d1_migrations.mjs --pending Print migrations not yet applied remotely. + * node scripts/verify_d1_migrations.mjs Fail unless every file is applied remotely. + * D1_DATABASE may supply an explicit target; there is no default database. * * The static check fails when a file name is malformed, when numbering has a gap or a * duplicate (two branches that each added the same number), or when a migration drops or @@ -17,7 +18,6 @@ import { readdirSync, readFileSync } from "node:fs"; import { pathToFileURL } from "node:url"; const DIRECTORY = new URL("../intelligence-migrations/", import.meta.url); -const DEFAULT_DATABASE = "multillm-intelligence"; export const MIGRATION_NAME = /^(\d{4})_[a-z0-9_]+\.sql$/; const MARKER = /^--[ \t]*destructive-migration:[ \t]*(\S.*)?$/m; const DESTRUCTIVE = [ @@ -141,11 +141,16 @@ export function localMigrations(directory = DIRECTORY) { } export function remoteAppliedMigrations(database, run = defaultRun) { + if (typeof database !== "string" || !database.trim() || database.startsWith("-")) { + throw new Error("Select a database explicitly with an argument or D1_DATABASE before a remote migration read"); + } return appliedFromWrangler(run(["d1", "execute", database, "--remote", "--json", "--command", "SELECT name FROM d1_migrations ORDER BY id"])); } function defaultRun(args) { + const config = process.env.MAIN_WORKER_CONFIG; + if (config) args = [...args, "--config", config]; return execFileSync("npx", ["--no-install", "wrangler", ...args], { encoding: "utf8", stdio: ["ignore", "pipe", "inherit"] }); } @@ -184,10 +189,26 @@ function remote(database, { listOnly }) { console.log(`All ${files.length} D1 migrations are applied to ${database}.`); } +export function remoteOptions(args, env = process.env) { + const flags = args.filter(arg => arg.startsWith("--")); + const targets = args.filter(arg => !arg.startsWith("--")); + if (flags.some(arg => !["--pending", "--remote"].includes(arg)) || targets.length > 1) { + throw new Error("Use [--pending] [--remote] , or --check for local validation"); + } + const database = targets[0] || env.D1_DATABASE; + if (typeof database !== "string" || !database.trim() || database.startsWith("-")) { + throw new Error("Select a database explicitly with an argument or D1_DATABASE before a remote migration read"); + } + return { database, listOnly: flags.includes("--pending") }; +} + function main(args) { if (args[0] === "--check") return check(); - if (args[0] === "--pending") return remote(args[1] || DEFAULT_DATABASE, { listOnly: true }); - return remote(args[0] || DEFAULT_DATABASE, { listOnly: false }); + const { database, listOnly } = remoteOptions(args); + return remote(database, { listOnly }); } -if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) main(process.argv.slice(2)); +if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) { + try { main(process.argv.slice(2)); } + catch (error) { console.error(error.message); process.exitCode = 1; } +} diff --git a/services/accounted_dispatch.py b/services/accounted_dispatch.py index dba61099..1043a305 100644 --- a/services/accounted_dispatch.py +++ b/services/accounted_dispatch.py @@ -29,6 +29,10 @@ def release_outer_accounting() -> None: def accounted_dispatch(payload: dict, dispatch, *, kind: str, skip_rate: bool = False): user = getattr(g, "authenticated_user", None) or {} model = payload["model"] + from services.key_data_guards import model_permitted + if not model_permitted(user, model): + raise APIError("This API key cannot use this provider data class", 403, + {"error": "provider_class_not_allowed"}) if not key_controls.model_allowed(user, model): raise APIError("This API key is not allowed to use the subrequest model", 403, payload={"error": "model_not_allowed"}) diff --git a/services/auth_service.py b/services/auth_service.py index 81ca3b91..61eca671 100644 --- a/services/auth_service.py +++ b/services/auth_service.py @@ -153,6 +153,8 @@ def _ensure_users_schema(cls, connection: sqlite3.Connection) -> None: "expires_at": "TEXT", "secret_scan_mode": "TEXT", # nosec B105 - a column type, not a password "shadow_eval_rate": "DOUBLE PRECISION", + "official_only": "INTEGER", + "forced_provider_preferences": "TEXT", } for column_name, column_definition in required_columns.items(): if column_name not in columns: @@ -193,7 +195,9 @@ def _create_users_table(cls, connection: sqlite3.Connection, table_name: str = " allowed_ips TEXT, expires_at TEXT, secret_scan_mode TEXT, - shadow_eval_rate DOUBLE PRECISION + shadow_eval_rate DOUBLE PRECISION, + official_only INTEGER, + forced_provider_preferences TEXT ) """ ) @@ -279,7 +283,9 @@ def _backfill_user_metadata(cls, connection: sqlite3.Connection) -> None: @classmethod def _row_to_user(cls, row: sqlite3.Row) -> Dict[str, Any]: + from services.dispatch_policy import envelope, verify return { + "data_policy": verify(row["data_policy"]) if dict(row).get("data_policy") is not None else envelope({"official_only": None, "forced_provider_preferences": None, **dict(row)}), "username": row["username"], "api_key_hash": row["api_key_hash"], "api_key_prefix": row["api_key_prefix"], @@ -293,9 +299,7 @@ def _row_to_user(cls, row: sqlite3.Row) -> Dict[str, Any]: "rotated_at": deserialize_datetime(row["rotated_at"]), "revoked_at": deserialize_datetime(row["revoked_at"]), **stream_cost_breaker.storage_control(row), - **key_controls.from_storage( - {name: key_controls.row_value(row, name) for name in key_controls.CONTROL_FIELDS} - ), + **key_controls.from_storage(row), } @classmethod @@ -311,7 +315,7 @@ def _reload_user_cache(cls) -> None: username, api_key_hash, api_key_prefix, scopes, is_admin, created_at, last_login, last_used_at, last_used_ip, created_by, rotated_at, revoked_at, daily_budget_usd, - monthly_budget_usd, allowed_models, allowed_ips, expires_at, secret_scan_mode, shadow_eval_rate + monthly_budget_usd, allowed_models, allowed_ips, expires_at, secret_scan_mode, shadow_eval_rate, official_only, forced_provider_preferences FROM users ORDER BY username """) @@ -336,7 +340,7 @@ def _load_user_by_username(cls, username: str) -> Optional[Dict[str, Any]]: username, api_key_hash, api_key_prefix, scopes, is_admin, created_at, last_login, last_used_at, last_used_ip, created_by, rotated_at, revoked_at, daily_budget_usd, - monthly_budget_usd, allowed_models, allowed_ips, expires_at, secret_scan_mode, shadow_eval_rate + monthly_budget_usd, allowed_models, allowed_ips, expires_at, secret_scan_mode, shadow_eval_rate, official_only, forced_provider_preferences FROM users WHERE username = ? """), @@ -366,7 +370,7 @@ def _load_users_by_api_key_prefix(cls, api_key_prefix: str) -> List[tuple[str, D username, api_key_hash, api_key_prefix, scopes, is_admin, created_at, last_login, last_used_at, last_used_ip, created_by, rotated_at, revoked_at, daily_budget_usd, - monthly_budget_usd, allowed_models, allowed_ips, expires_at, secret_scan_mode, shadow_eval_rate + monthly_budget_usd, allowed_models, allowed_ips, expires_at, secret_scan_mode, shadow_eval_rate, official_only, forced_provider_preferences FROM users WHERE api_key_prefix = ? AND revoked_at IS NULL ORDER BY username @@ -438,9 +442,9 @@ def _persist_user( username, api_key_hash, api_key_prefix, scopes, is_admin, created_at, last_login, last_used_at, last_used_ip, created_by, rotated_at, revoked_at, daily_budget_usd, - monthly_budget_usd, allowed_models, allowed_ips, expires_at, secret_scan_mode, shadow_eval_rate + monthly_budget_usd, allowed_models, allowed_ips, expires_at, secret_scan_mode, shadow_eval_rate, official_only, forced_provider_preferences ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(username) DO UPDATE SET api_key_hash = excluded.api_key_hash, api_key_prefix = excluded.api_key_prefix, @@ -475,6 +479,7 @@ def _persist_user( serialize_datetime(rotated_at), serialize_datetime(revoked_at), *(stored_controls[name] for name in key_controls.CONTROL_FIELDS), + None, None, ), ) stream_cost_breaker.persist_sql_cap(connection, username, stored_controls) diff --git a/services/cloudflare_ai.py b/services/cloudflare_ai.py index 4079a23b..98972be4 100644 --- a/services/cloudflare_ai.py +++ b/services/cloudflare_ai.py @@ -42,7 +42,9 @@ def post(path: str, payload: dict, timeout) -> requests.Response | Response: with requests.Session() as session: session.trust_env = False try: - response = session.post(ORIGIN + path, json=payload, timeout=timeout, allow_redirects=False, stream=True) + from services.provider_dispatch import provider_http + response = provider_http("cloudflare", "POST", ORIGIN + path, transport=session.request, + binding=True, json=payload, timeout=timeout, stream=True, allow_redirects=False) except requests.exceptions.ConnectionError as error: connect = isinstance(error, requests.exceptions.ConnectTimeout) or "NewConnectionError" in repr(error) return _transport_failure("connect" if connect else "interrupted") diff --git a/services/dispatch_network.py b/services/dispatch_network.py new file mode 100644 index 00000000..ef5c4c1d --- /dev/null +++ b/services/dispatch_network.py @@ -0,0 +1,55 @@ +"""A short-lived network capability issued only after dispatch admission.""" +from contextlib import contextmanager +from contextvars import ContextVar + +_active = ContextVar("admitted_provider_network", default=False) + + +def network_admitted(address=None): + scope = _active.get() + if not scope or not scope["active"]: + return False + destination = scope.get("destination") + if destination is None: + return True + if not isinstance(address, tuple) or address[1] != destination[1]: + return False + if address[:2] == destination: + return True + # Resolve only when a real connection is attempted; mocked transports and + # disabled exporters must retain their existing no-network behavior. + import socket + try: + addresses = scope.get("resolved") + if addresses is None: + addresses = {item[4][:2] for item in socket.getaddrinfo(*destination, type=socket.SOCK_STREAM)} + scope["resolved"] = addresses + return address[:2] in addresses + except OSError: + return False + + +@contextmanager +def admitted_network(): + scope = {"active": True} + token = _active.set(scope) + try: + yield + finally: + scope["active"] = False + _active.reset(token) + + +@contextmanager +def control_network(url): + """Explicit non-provider control traffic, confined to its configured endpoint.""" + from urllib.parse import urlsplit + target = urlsplit(url) + port = target.port or (443 if target.scheme == "https" else 80) + scope = {"active": True, "destination": (target.hostname, port)} + token = _active.set(scope) + try: + yield + finally: + scope["active"] = False + _active.reset(token) diff --git a/services/dispatch_policy.py b/services/dispatch_policy.py new file mode 100644 index 00000000..d9a5bed1 --- /dev/null +++ b/services/dispatch_policy.py @@ -0,0 +1,23 @@ +"""Versioned, explicit data policy carried by verified account records.""" + +from error_handlers import APIError +from services.key_data_guards import stored_controls + +VERSION = 1 +FIELDS = frozenset({"version", "official_only", "forced_provider_preferences"}) + + +def unavailable(): + return APIError("Current key data controls could not be verified", 503, + {"error": "key_data_controls_unavailable"}) + + +def envelope(row): + controls = stored_controls(row) + return {"version": VERSION, **controls} + + +def verify(value): + if not isinstance(value, dict) or set(value) != FIELDS or type(value.get("version")) is not int or value["version"] != VERSION: + raise unavailable() + return envelope(value) diff --git a/services/gemini_credentials.py b/services/gemini_credentials.py new file mode 100644 index 00000000..00bd3fad --- /dev/null +++ b/services/gemini_credentials.py @@ -0,0 +1,149 @@ +"""Credential selection for native and translated Gemini requests.""" + +import logging +import os +import re +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +from flask import g, has_request_context, request + +from error_handlers import APIError +from services.auth_primitives import usable_credential + +logger = logging.getLogger(__name__) + + +def restricted(provider): + from services.key_data_guards import has_data_controls + return provider == "gemini" and has_request_context() and has_data_controls( + getattr(g, "authenticated_user", None) or {}) + + +def _refuse_override(): + raise APIError("Caller provider credentials are not allowed for this key", 403, + {"error": "provider_credential_not_allowed"}) + + +def reject_caller_credentials(provider): + if not restricted(provider): + return + if any(name.lower() == "key" for name in request.args) or "x-goog-api-key" in request.headers: + _refuse_override() + gateway = request.headers.get("x-multillm-api-key") + authorization = request.headers.get("authorization") + if gateway and authorization: + bearer = re.fullmatch(r"Bearer\s+(.+)", authorization.strip(), re.IGNORECASE) + if not bearer or bearer[1].strip() != gateway.strip(): + _refuse_override() + + +def prepare_credentials(url, headers, params, provider, auth_token, auth): + headers, params = dict(headers or {}), dict(params or {}) + if restricted(provider): + reject_caller_credentials(provider) + if any(name.lower() == "key" for name in params) or any( + name.lower() == "key" for name, _ in parse_qsl(urlsplit(url).query, keep_blank_values=True)): + _refuse_override() + # Only this operator binding is covered by the paid-tier attestation. Do not + # recover through a legacy alias, cached credential or caller-supplied key. + key = usable_credential(os.environ.get("GEMINI_API_KEY"), "GEMINI_API_KEY") + if not key: + raise APIError("Attested provider credential is unavailable", 503, + {"error": "attested_credential_unavailable"}) + gateway_authorization = request.headers.get("authorization") + for name, value in headers.items(): + if name.lower() == "x-goog-api-key" and value != key: + _refuse_override() + if name.lower() == "authorization" and value not in {f"Bearer {key}", gateway_authorization}: + _refuse_override() + else: + url, params, key = _legacy_credentials(url, params, provider, auth_token, auth) + headers = {name: value for name, value in headers.items() + if name.lower() not in {"authorization", "x-goog-api-key"}} + headers["x-goog-api-key"] = key + return url, headers, params + + +def _legacy_credentials(url, params, api_provider, auth_token, auth): + # Extract API key from URL parameters and rebuild the URL without it + api_key = None + + # Check if key is in params + if params and 'key' in params: + api_key = params.pop('key') + if isinstance(api_key, (list, tuple)): + api_key = api_key[0] if api_key else None + logger.info(f"Found API key in URL parameters for {api_provider}") + + # If no key in params, check if it's in the URL + else: + parsed_url = urlsplit(url) + query_pairs = parse_qsl(parsed_url.query, keep_blank_values=True) + filtered_query_pairs = [] + for key, value in query_pairs: + if key == 'key' and api_key is None: + api_key = value + logger.info(f"Found API key in URL for {api_provider}") + else: + filtered_query_pairs.append((key, value)) + + if api_key is not None: + url = urlunsplit( + ( + parsed_url.scheme, + parsed_url.netloc, + parsed_url.path, + urlencode(filtered_query_pairs), + parsed_url.fragment, + ) + ) + + # If still no API key, use the auth token passed from make_request + # But first check if it's the admin API key - if so, use the Gemini key from env instead + if not api_key and auth_token: + admin_api_key = os.environ.get('ADMIN_API_KEY') + if auth_token == admin_api_key: + # Using admin key, look up the actual provider key instead + logger.info(f"Admin API key detected - using {api_provider} API key from env instead") + api_key = auth.get_api_key(api_provider) + if not api_key: + # Try alternative keys + if api_provider == 'gemma': + api_key = auth.get_api_key('gemini') + elif api_provider == 'gemini': + api_key = auth.get_api_key('gemma') + else: + # Not the admin key, might be a direct Gemini key + api_key = auth_token + logger.info(f"Using API key from Authorization header for {api_provider}") + + # If still no API key, get from auth service + if not api_key: + # First try with the exact api_provider + api_key = auth.get_api_key(api_provider) + + # If using 'gemma' as provider but no gemma key found, try with 'gemini' as they share the same key + if not api_key and api_provider == 'gemma': + api_key = auth.get_api_key('gemini') + logger.info("Using Gemini API key for Gemma model") + + # If using 'gemini' as provider but no gemini key found, try with 'gemma' as a fallback + elif not api_key and api_provider == 'gemini': + api_key = auth.get_api_key('gemma') + logger.info("Using Gemma API key for Gemini model") + + if not api_key: + raise APIError(f"No API key found for {api_provider}. Please set {api_provider.upper()}_API_KEY in your .env file.", status_code=401) + + logger.info(f"Using API key from AuthService for {api_provider}") + + # Validate API key format - Gemini keys should start with "AIza" + if not api_key.startswith("AIza"): + logger.warning(f"API key for {api_provider} doesn't match expected format (should start with 'AIza')") + # Try to find a valid key from the environment + env_key = os.environ.get(f'{api_provider.upper()}_API_KEY') + if env_key and env_key.startswith("AIza"): + logger.info(f"Found valid {api_provider} API key format in environment, using it instead") + api_key = env_key + + return url, params, api_key diff --git a/services/intelligence_transport.py b/services/intelligence_transport.py index f77e0ab2..84083d39 100644 --- a/services/intelligence_transport.py +++ b/services/intelligence_transport.py @@ -72,7 +72,8 @@ def __init__( self.response = None self._response_lock = threading.Lock() self.max_bytes = max_bytes - self.send, self.record = send, record + from services.provider_dispatch import capture_dispatch + self.send, self.record = capture_dispatch(send), record self.record_response = record_response if not _TRANSPORT_SLOTS.acquire(blocking=False): raise GatewayError( @@ -395,7 +396,12 @@ def start( if content_type: headers["Content-Type"] = content_type - dispatch_data = protect_body(upstream.data if data is None else data, headers, provider=provider) + from services.key_data_guards import guard_request + guarded_data = guard_request(provider, upstream.data if data is None else data) + from services.gemini_credentials import prepare_credentials, restricted + if restricted(provider): + url, headers, _ = prepare_credentials(url, headers, {}, provider, token, self.auth) + dispatch_data = protect_body(guarded_data, headers, provider=provider) scope = selection_context(provider, model, config=self.config) from services.generation_deadline import bounded_timeout, current_deadline diff --git a/services/key_controls_capability.py b/services/key_controls_capability.py new file mode 100644 index 00000000..09cd3e71 --- /dev/null +++ b/services/key_controls_capability.py @@ -0,0 +1,13 @@ +"""The code version advertises rollout support; configuration cannot turn it on.""" +CAPABILITY = {"protocol": 1, "enforcement": False} + + +def policy_present(policy): + return policy.get("official_only") is not None or policy.get("forced_provider_preferences") is not None + + +def expansion_fence(policy): + if not CAPABILITY["enforcement"] and policy_present(policy): + from error_handlers import APIError + raise APIError("This runtime cannot enforce the account policy", 503, {"error": "key_controls_require_upgrade"}) + return policy diff --git a/services/key_data_guards.py b/services/key_data_guards.py new file mode 100644 index 00000000..3acad499 --- /dev/null +++ b/services/key_data_guards.py @@ -0,0 +1,152 @@ +"""Provider eligibility and administrator-owned request preferences.""" + +import json +from copy import deepcopy +from pathlib import Path + +from flask import current_app, g, has_app_context, has_request_context + +from error_handlers import APIError + +DATA_CLASSES = json.loads((Path(__file__).parent.parent / "providers/data-classes.json").read_text()) +MAX_PREFERENCES_BYTES = 4096 + + +def preferences(value): + """Decode storage or API values, rejecting malformed policy rather than dropping it.""" + if value is None: + return None + try: + parsed = json.loads(value) if isinstance(value, str) else value + encoded = json.dumps(parsed, allow_nan=False) + if not isinstance(parsed, dict) or len(encoded.encode()) > MAX_PREFERENCES_BYTES: + raise ValueError() + if parsed.get("data_collection", "deny") not in {"allow", "deny"}: + raise ValueError() + return parsed + except (ValueError, TypeError, RecursionError): + raise APIError("forced_provider_preferences must be a bounded JSON object", 400) from None + + +def validate_controls(payload): + official = payload.get("official_only") + if official is not None and type(official) is not bool: + raise APIError("official_only must be a boolean", 400) + raw = payload.get("forced_provider_preferences") + if raw is not None and not isinstance(raw, dict): + raise APIError("forced_provider_preferences must be a JSON object", 400) + forced = preferences(raw) + return {"official_only": int(official) if official is not None else None, + "forced_provider_preferences": json.dumps(forced, separators=(",", ":")) if forced is not None else None} + + +def stored_controls(row): + def value(name): + try: + return row[name] + except (KeyError, IndexError, TypeError): + raise APIError("Stored key data controls are incomplete", 503) from None + official = value("official_only") + if official is not None and (type(official) not in (bool, int) or official not in (0, 1)): + raise APIError("Stored key data controls are invalid", 503) + raw = value("forced_provider_preferences") + try: + forced = preferences(raw) + except APIError: + raise APIError("Stored key data controls are invalid", 503) from None + return {"official_only": official, + "forced_provider_preferences": json.dumps(forced, separators=(",", ":")) if forced is not None else None} + + +def provider_data_class(provider, config=None): + config = config if config is not None else (current_app.config if has_app_context() else {}) + provider = str(provider).strip().lower() + free = str(config.get("FREE_ROUTE_FREE_TIER_PROVIDERS", "")).replace(",", " ").split() + if provider in free or provider == "free": + return "free" + if provider == "gemini": + return "official" if config.get("GEMINI_DATA_TIER", "free") == "paid" else "free" + return DATA_CLASSES.get(provider, "subscription" if provider.startswith("ce-") else "reseller") + + +def account_controls(user): + """Resolve the canonical envelope, or validate an explicit stored row.""" + from services.dispatch_policy import verify, unavailable + if not isinstance(user, dict): + raise unavailable() + if "data_policy" in user: + return verify(user["data_policy"]) + if "version" in user: + return verify(user) + return stored_controls(user) + + +def has_data_controls(user): + """Any stored non-NULL value represents an administrator-owned data policy.""" + from services.key_controls_capability import policy_present + policy = account_controls(user) if "data_policy" in user or "version" in user else user + return policy_present(policy) + + +def allows_shadow_sampling(user): + """Unverifiable authority cannot authorize retaining request content.""" + try: + return not has_data_controls(account_controls(user)) + except APIError: + return False + + +def model_permitted(user, model, config=None): + from services.key_controls_capability import expansion_fence + policy = account_controls(user) if "data_policy" in user else user + expansion_fence(policy) + if policy.get("official_only") not in (1, True): + return True + provider = str(model or "").split(":", 1)[0].lower() + if provider in {"auto", "cascade", "intelligence"}: + return True # Each concrete candidate is checked before dispatch. + return provider_data_class(provider, config) == "official" + + +def eligible_models(models, user): + """Filter route candidates and refuse before dispatch if none remain.""" + eligible = [model for model in models if model_permitted(user, model)] + if not eligible: + raise APIError("No eligible provider data class remains", 403, {"error": "provider_class_not_allowed"}) + return eligible + + +def require_provider(provider, user=None): + user = user if user is not None else (getattr(g, "authenticated_user", None) or {} if has_request_context() else {}) + if not model_permitted(user, provider + ":*"): + raise APIError("This API key cannot use this provider data class", 403, + {"error": "provider_class_not_allowed"}) + from services.gemini_credentials import reject_caller_credentials + reject_caller_credentials(provider) + + +def merge_preferences(caller, forced): + result = deepcopy(caller) if isinstance(caller, dict) else {} + for name, value in forced.items(): + if name == "data_collection" and result.get(name) == "deny": + continue + result[name] = merge_preferences(result.get(name), value) if isinstance(value, dict) else deepcopy(value) + return result + + +def guard_request(provider, data): + """Final transport boundary covers raw, unified, routed and internal requests.""" + require_provider(provider) + user = getattr(g, "authenticated_user", None) or {} if has_request_context() else {} + policy = account_controls(user) if "data_policy" in user else user + forced = preferences(policy.get("forced_provider_preferences")) + if provider != "openrouter" or forced is None or data is None: + return data + try: + payload = json.loads(data) + if not isinstance(payload, dict): + raise ValueError() + except (ValueError, TypeError): + raise APIError("A JSON object is required for enforced provider preferences", 400) from None + payload["provider"] = merge_preferences(payload.get("provider"), forced) + return json.dumps(payload).encode() diff --git a/services/knowledge_client.py b/services/knowledge_client.py index 8b6586a2..1b8b53b3 100644 --- a/services/knowledge_client.py +++ b/services/knowledge_client.py @@ -132,7 +132,8 @@ def _submit(body, stopped, deadline, results): session.mount("http://", HTTPAdapter(max_retries=0)) if stopped.is_set(): return - with session.post(ENDPOINT, data=body, + from services.provider_dispatch import provider_http + with provider_http("knowledge", "POST", ENDPOINT, transport=session.request, binding=True, data=body, headers={"Content-Type": "application/json", "Accept": "application/json", "Accept-Encoding": "identity"}, timeout=(3, DEADLINE_SECONDS), allow_redirects=False, stream=True) as response: @@ -173,7 +174,9 @@ def dispatch(operation, user, payload=None): stopped = threading.Event() results = queue.Queue(maxsize=1) deadline = time.monotonic() + DEADLINE_SECONDS - thread = threading.Thread(target=_submit, args=(body, stopped, deadline, results), + from services.provider_dispatch import capture_dispatch + thread = threading.Thread(target=capture_dispatch(_submit, reference={"username": user.get("username") or user.get("id"), + "api_key_prefix": user.get("api_key_prefix")}), args=(body, stopped, deadline, results), daemon=True, name="knowledge-request") try: thread.start() diff --git a/services/media_storage.py b/services/media_storage.py index abd4ab1a..ee75717a 100644 --- a/services/media_storage.py +++ b/services/media_storage.py @@ -26,6 +26,7 @@ from error_handlers import APIError from services.media_signing import FILE_ID, file_link_params from services.media_urls import public_https_url +from services.provider_dispatch import provider_http logger = logging.getLogger(__name__) @@ -78,7 +79,7 @@ def fetch_public(url: str, *, max_bytes: int = MAX_FETCH_BYTES, field: str = "im url = public_https_url(url, field) with _session() as session: try: - with session.post(f"{ORIGIN}/v1/fetch", json={"url": url, "max_bytes": max_bytes}, timeout=FETCH_TIMEOUT, + with provider_http("media", "POST", f"{ORIGIN}/v1/fetch", transport=session.request, binding=True, json={"url": url, "max_bytes": max_bytes}, timeout=FETCH_TIMEOUT, allow_redirects=False, stream=True) as response: if response.status_code != 200: code = _error_code(response) @@ -168,7 +169,7 @@ def import_url(file_id: str, url: str, *, owner: str, kind: str, model: str | No """The Worker downloads a provider URL straight into R2.""" try: with _session() as session: - response = session.post(_file_url(file_id, "/import"), timeout=STORE_TIMEOUT, allow_redirects=False, + response = provider_http("media", "POST", _file_url(file_id, "/import"), transport=session.request, binding=True, timeout=STORE_TIMEOUT, allow_redirects=False, json={"url": url, "metadata": {"owner": owner, "kind": kind, **({"model": model} if model else {})}}) return _stored(response) @@ -365,7 +366,10 @@ def video_file_id(job_id: str) -> str: def store_video(file_id: str, url: str, headers: dict, *, owner: str, model: str | None) -> dict | None: """Copy a finished video into R2 once; None when its size is unknown or too large.""" try: - upstream = requests.get(url, headers=headers, timeout=VIDEO_TIMEOUT, stream=True, allow_redirects=True) + from services.provider_assets import asset_http + provider = str(model or "unknown").split(":", 1)[0] + upstream = asset_http(provider, url, headers=headers, timeout=VIDEO_TIMEOUT, stream=True, + reference={"username": owner}) except requests.exceptions.RequestException: return None with upstream: diff --git a/services/native_token_count.py b/services/native_token_count.py index 7ae2dcb7..fb0b38be 100644 --- a/services/native_token_count.py +++ b/services/native_token_count.py @@ -130,7 +130,8 @@ def count_native_tokens( session = ProxyService._get_provider_session(provider, raw_passthrough=True) response = None try: - response = session.post(url=url, headers=headers, json=body, stream=True, + from services.provider_dispatch import provider_http + response = provider_http(provider, "POST", url, transport=session.request, headers=headers, json=body, stream=True, allow_redirects=False, timeout=(5, 10)) status = response.status_code if status >= 400: diff --git a/services/provider_assets.py b/services/provider_assets.py new file mode 100644 index 00000000..3f0698b4 --- /dev/null +++ b/services/provider_assets.py @@ -0,0 +1,40 @@ +"""Bounded asset downloads, with fresh policy admission at every redirect.""" +from urllib.parse import urljoin, urlsplit + +from error_handlers import APIError +from services.provider_credentials import api_origin, origin, strip_credentials +from services.provider_dispatch import provider_http + +MAX_REDIRECTS = 5 + + +def asset_http(provider, url, *, headers=None, reference=None, **options): + from services.key_controls_capability import CAPABILITY + if not CAPABILITY["enforcement"]: + return provider_http(provider, "GET", url, headers=dict(headers or {}), reference=reference, **options) + # Only documented authenticated content APIs may start with credentials. + # Status-returned xAI and private-binding assets are always public downloads. + authenticated = provider in {"openai", "gemini"} and api_origin(provider, url) + headers = dict(headers or {}) + for hop in range(MAX_REDIRECTS + 1): + parsed = urlsplit(url) + if origin(url) is None or parsed.scheme != "https" or parsed.port not in (None, 443): + raise APIError("The provider returned an unsupported video location", 502) + if not authenticated: + headers = strip_credentials(headers) + response = provider_http(provider, "GET", url, headers=headers, reference=reference, + public_asset=not authenticated, **options) + if not 300 <= response.status_code < 400: + return response + location = response.headers.get("Location") + response.close() + if not location or hop == MAX_REDIRECTS: + raise APIError("The video download redirect could not be resolved", 502) + destination = urljoin(url, location) + # Once an API redirects to an asset origin, never reintroduce its secrets, + # including if a subsequent redirect returns to the API origin. + if origin(destination) != origin(url): + authenticated = False + headers = strip_credentials(headers) + url = destination + raise AssertionError("Unreachable redirect limit") diff --git a/services/provider_credentials.py b/services/provider_credentials.py new file mode 100644 index 00000000..cd1db5c4 --- /dev/null +++ b/services/provider_credentials.py @@ -0,0 +1,100 @@ +"""Origin-bound operator credentials; public assets never receive API secrets.""" +import json +import os +import shutil +import subprocess +from pathlib import Path +from urllib.parse import urlsplit + +from flask import current_app, has_app_context + +from error_handlers import APIError +from services.auth_primitives import provider_api_key_env_names, usable_credential + +SECRET_HEADERS = {"authorization", "proxy-authorization", "x-api-key", "x-goog-api-key", "api-key", "cookie"} + + +def origin(url): + try: + parsed = urlsplit(url) + if parsed.username or parsed.password or not parsed.hostname: + return None + return parsed.scheme.lower(), parsed.hostname.lower(), parsed.port or (443 if parsed.scheme == "https" else 80) + except ValueError: + return None + + +def api_origin(provider, url): + from config import Config + bases = current_app.config.get("API_BASE_URLS", Config.API_BASE_URLS) if has_app_context() else Config.API_BASE_URLS + allowed = {origin(bases.get(provider, ""))} + # Native Gemini and Vertex protocols use the same operator API binding. + if provider in {"gemini", "gemma", "palm"}: + allowed.add(origin("https://generativelanguage.googleapis.com")) + allowed.discard(None) + target = origin(url) + return target is not None and target[0] == "https" and target in allowed + + +def strip_credentials(headers): + return {name: value for name, value in headers.items() if name.lower() not in SECRET_HEADERS} + + +def unavailable_credential(): + return APIError("Attested provider credential is unavailable", 503, {"error": "attested_credential_unavailable"}) + + +def _google_binding(): + raw = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON", "").strip() + path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "").strip() + if raw: + return "json", raw + if path: + return "file", Path(path).read_text() + return None + + +def google_operator_token(): + """Refresh the current binding, never the legacy process-wide token cache.""" + try: + binding = _google_binding() + if binding: + from google.oauth2 import service_account + from google.auth.transport.requests import Request + info = json.loads(binding[1]) + if origin(info.get("token_uri", "")) != origin("https://oauth2.googleapis.com"): + raise ValueError() + credentials = service_account.Credentials.from_service_account_info(info, + scopes=["https://www.googleapis.com/auth/cloud-platform"]) + credentials.refresh(Request()) + token = credentials.token + if binding != _google_binding(): + raise ValueError() + else: + binary = shutil.which("gcloud") + if not binary: + raise ValueError() + def run(*args): + return subprocess.run([binary, *args, "--quiet"], capture_output=True, text=True, + check=True, timeout=30).stdout.strip() + account = run("config", "get-value", "account") + if not account or account == "(unset)": + raise ValueError() + token = run("auth", "print-access-token", account) + if _google_binding() is not None or account != run("config", "get-value", "account"): + raise ValueError() + if not isinstance(token, str) or not token.strip(): + raise ValueError() + return token.strip() + except Exception: + raise unavailable_credential() from None + + +def operator_credential(provider): + if provider == "googleai": + return google_operator_token() + names = ("GEMINI_API_KEY",) if provider == "gemini" else provider_api_key_env_names(provider) + key = next((usable_credential(os.environ.get(name), name) for name in names if usable_credential(os.environ.get(name))), None) + if not key: + raise unavailable_credential() + return key diff --git a/services/provider_dispatch.py b/services/provider_dispatch.py new file mode 100644 index 00000000..b9bc0281 --- /dev/null +++ b/services/provider_dispatch.py @@ -0,0 +1,155 @@ +"""The sole provider network boundary, including deferred and non-chat work. + +Only an account reference crosses a thread boundary. Its policy is reread at each +submission; a cached principal, environment administrator or old RPC cannot grant +generation authority. Private service bindings receive the same verified reference. +""" + +import json +from contextvars import ContextVar +from contextlib import contextmanager +from urllib.parse import parse_qsl, urlsplit + +import requests +from flask import current_app, g, has_app_context, has_request_context, request + +from error_handlers import APIError +from services.provider_credentials import api_origin, operator_credential, strip_credentials, unavailable_credential +from services.dispatch_network import admitted_network +from services.dispatch_policy import unavailable, verify +from services.key_controls_capability import CAPABILITY, expansion_fence +from services.key_data_guards import has_data_controls, merge_preferences, model_permitted, preferences + +_reference = ContextVar("provider_principal", default=None) + + +def principal_reference(): + if _reference.get() is not None: + return _reference.get() + user = getattr(g, "authenticated_user", None) if has_request_context() else None + if user: + return {"username": user.get("username") or user.get("id"), "api_key_prefix": user.get("api_key_prefix")} + return _reference.get() + + +@contextmanager +def dispatch_as(reference): + token = _reference.set(reference) + try: + yield + finally: + _reference.reset(token) + + +def current_principal(reference=None): + from services.auth_service import AuthService + from services.key_controls import expired + + reference = reference or principal_reference() + try: + if not isinstance(reference, dict) or not reference.get("username"): + raise unavailable() + if str(reference["username"]).startswith("integration:"): + from services.intelligence_d1_store import request_private_intelligence + response = request_private_intelligence({"version": 1, "operation": "lookup", + "keyPrefix": reference.get("api_key_prefix")}, endpoint="auth") + row = response.get("principal") if isinstance(response, dict) else None + if not row or row.get("id") != reference["username"] or row.get("revokedAt"): + raise unavailable() + user = {"username": row["id"], "api_key_prefix": row["keyPrefix"], + "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}} + + else: + user = AuthService._load_user_by_username(reference["username"]) + if not user or CAPABILITY["enforcement"] and (user.get("revoked_at") or expired(user)): + raise unavailable() + if reference.get("api_key_prefix") and reference["api_key_prefix"] != user.get("api_key_prefix"): + raise unavailable() + policy = verify(user.get("data_policy")) + expansion_fence(policy) + return {**user, **policy} + except APIError: + raise + except Exception: + raise unavailable() from None + + +def capture_dispatch(callback, *, reference=None): + """Capture identity and application, never a policy snapshot, for a new thread.""" + reference = reference or principal_reference() + app = current_app._get_current_object() if has_app_context() else None + + def bound(*args, **kwargs): + token = _reference.set(reference) + try: + if app: + with app.app_context(): + return callback(*args, **kwargs) + return callback(*args, **kwargs) + finally: + _reference.reset(token) + return bound + + +def _credential(provider, user, url, options, *, public_asset=False): + headers = dict(options.get("headers") or {}) + if public_asset: + return {**options, "headers": strip_credentials(headers)} + if not has_data_controls(user): + return options + if any(name.lower() == "key" for name, _ in parse_qsl(urlsplit(url).query, keep_blank_values=True)) or any( + name.lower() == "key" for name in (options.get("params") or {})): + raise APIError("Caller provider credentials are not allowed", 403, {"error": "provider_credential_not_allowed"}) + if has_request_context() and (any(name.lower() == "key" for name in request.args) + or "x-goog-api-key" in request.headers): + raise APIError("Caller provider credentials are not allowed", 403, {"error": "provider_credential_not_allowed"}) + if has_request_context() and request.headers.get("X-MultiLLM-Api-Key"): + authorization = request.headers.get("Authorization") + if authorization and (authorization[7:].strip() if authorization.lower().startswith("bearer ") else authorization.strip()) != request.headers["X-MultiLLM-Api-Key"].strip(): + raise APIError("Caller provider credentials are not allowed", 403, {"error": "provider_credential_not_allowed"}) + # A private binding uses the operator's binding rather than a bearer credential. + if provider == "cloudflare" or urlsplit(url).hostname in {"ai.internal", "knowledge.internal", "media.internal"}: + return options + if not api_origin(provider, url): + raise unavailable_credential() + # Resolve only after current-policy admission, including OAuth token exchange. + with admitted_network(): + key = operator_credential(provider) + headers = strip_credentials(headers) + header = {"gemini": "x-goog-api-key", "anthropic": "x-api-key", "azure": "api-key"}.get(provider, "Authorization") + headers[header] = f"Bearer {key}" if header == "Authorization" else key + return {**options, "headers": headers} + + +def prepare_dispatch(provider, url, options, *, reference=None, binding=False, public_asset=False): + user = current_principal(reference) + if not CAPABILITY["enforcement"]: + return dict(options) + if not binding and not model_permitted(user, provider + ":*"): + raise APIError("This API key cannot use this provider data class", 403, {"error": "provider_class_not_allowed"}) + options = _credential(provider, user, url, dict(options), public_asset=public_asset) + forced = preferences(user["forced_provider_preferences"]) + if provider == "openrouter" and forced is not None and (options.get("data") is not None or "json" in options): + try: + body = options["json"] if "json" in options else json.loads(options["data"]) + if not isinstance(body, dict): + raise ValueError() + except (ValueError, TypeError): + raise APIError("Enforced provider preferences require a JSON object", 400) from None + body = {**body, "provider": merge_preferences(body.get("provider"), forced)} + options = {**options, **({"json": body} if "json" in options else {"data": json.dumps(body).encode()})} + headers = dict(options.get("headers") or {}) + headers = {name: value for name, value in headers.items() if name.lower() not in {"content-length", "x-multillm-dispatch-principal"}} + if binding and CAPABILITY["enforcement"]: + headers["X-MultiLLM-Dispatch-Principal"] = json.dumps({"version": 1, "username": user["username"]}) + return {**options, "headers": headers} + + +def provider_http(provider, method, url, *, transport=None, reference=None, binding=False, public_asset=False, **options): + options = prepare_dispatch(provider, url, options, reference=reference, binding=binding, public_asset=public_asset) + # Redirects are separate dispatches and must never forward an attested credential + # to a provider-selected host without another admission. + if CAPABILITY["enforcement"]: + options["allow_redirects"] = False + with admitted_network(): + return (transport or requests.request)(method=method, url=url, **options) diff --git a/services/proxy_service.py b/services/proxy_service.py index e7b3b841..add230b4 100644 --- a/services/proxy_service.py +++ b/services/proxy_service.py @@ -193,8 +193,9 @@ def probe_nanogpt_key( } with requests.Session() as session: session.cookies.set_policy(_RejectAllCookiesPolicy()) - response = session.get( - f"{base_url.rstrip('/')}/v1/models", + from services.provider_dispatch import provider_http + response = provider_http("nanogpt", "GET", + f"{base_url.rstrip('/')}/v1/models", transport=session.request, headers=headers, params={"detailed": "false"}, timeout=(5, max(1, timeout_seconds)), @@ -1134,7 +1135,8 @@ def _make_request_with_timeout( headers: Dict[str, str], params: Dict[str, Any], data: Optional[bytes], - timeout: Tuple[int, int] + timeout: Tuple[int, int], + api_provider: str = "unknown", ) -> requests.Response: """ Make a single request with a given timeout (not widely used in code below). @@ -1156,7 +1158,8 @@ def _make_request_with_timeout( from services.admission_request import bounded_generation_timeout timeout = bounded_timeout(bounded_generation_timeout(timeout)) - response = session.request( + from services.provider_dispatch import provider_http + response = provider_http(api_provider, transport=session.request, method=method, url=url, headers=headers, @@ -1244,7 +1247,8 @@ def next_attempt(count: int) -> requests.Response: from services.admission_request import bounded_generation_timeout timeout = bounded_timeout(bounded_generation_timeout(timeout)) - response = session.request( + from services.provider_dispatch import provider_http + response = provider_http(api_provider, transport=session.request, method=method, url=url, headers=headers, @@ -2531,94 +2535,9 @@ def _handle_gemini_request( ) request_data.pop("webSearchSpec", None) - # Extract API key from URL parameters and rebuild the URL without it - api_key = None - - # Check if key is in params - if params and 'key' in params: - api_key = params.pop('key') - if isinstance(api_key, (list, tuple)): - api_key = api_key[0] if api_key else None - logger.info(f"Found API key in URL parameters for {api_provider}") - - # If no key in params, check if it's in the URL - else: - parsed_url = urlsplit(url) - query_pairs = parse_qsl(parsed_url.query, keep_blank_values=True) - filtered_query_pairs = [] - for key, value in query_pairs: - if key == 'key' and api_key is None: - api_key = value - logger.info(f"Found API key in URL for {api_provider}") - else: - filtered_query_pairs.append((key, value)) - - if api_key is not None: - url = urlunsplit( - ( - parsed_url.scheme, - parsed_url.netloc, - parsed_url.path, - urlencode(filtered_query_pairs), - parsed_url.fragment, - ) - ) - - # If still no API key, use the auth token passed from make_request - # But first check if it's the admin API key - if so, use the Gemini key from env instead - if not api_key and auth_token: - admin_api_key = os.environ.get('ADMIN_API_KEY') - if auth_token == admin_api_key: - # Using admin key, look up the actual provider key instead - logger.info(f"Admin API key detected - using {api_provider} API key from env instead") - api_key = AuthService.get_api_key(api_provider) - if not api_key: - # Try alternative keys - if api_provider == 'gemma': - api_key = AuthService.get_api_key('gemini') - elif api_provider == 'gemini': - api_key = AuthService.get_api_key('gemma') - else: - # Not the admin key, might be a direct Gemini key - api_key = auth_token - logger.info(f"Using API key from Authorization header for {api_provider}") - - # If still no API key, get from auth service - if not api_key: - # First try with the exact api_provider - api_key = AuthService.get_api_key(api_provider) - - # If using 'gemma' as provider but no gemma key found, try with 'gemini' as they share the same key - if not api_key and api_provider == 'gemma': - api_key = AuthService.get_api_key('gemini') - logger.info("Using Gemini API key for Gemma model") - - # If using 'gemini' as provider but no gemini key found, try with 'gemma' as a fallback - elif not api_key and api_provider == 'gemini': - api_key = AuthService.get_api_key('gemma') - logger.info("Using Gemma API key for Gemini model") - - if not api_key: - raise APIError(f"No API key found for {api_provider}. Please set {api_provider.upper()}_API_KEY in your .env file.", status_code=401) - - logger.info(f"Using API key from AuthService for {api_provider}") - - # Validate API key format - Gemini keys should start with "AIza" - if not api_key.startswith("AIza"): - logger.warning(f"API key for {api_provider} doesn't match expected format (should start with 'AIza')") - # Try to find a valid key from the environment - env_key = os.environ.get(f'{api_provider.upper()}_API_KEY') - if env_key and env_key.startswith("AIza"): - logger.info(f"Found valid {api_provider} API key format in environment, using it instead") - api_key = env_key - - headers["x-goog-api-key"] = api_key - - # Remove the Authorization header as Gemini doesn't use it - if 'Authorization' in headers: - headers.pop('Authorization') - logger.info("Removed Authorization header for Gemini API request") - + from services.gemini_credentials import prepare_credentials + url, headers, params = prepare_credentials(url, headers, params, api_provider, auth_token, AuthService) + # Extract model if it's in the URL path for direct model invocation # Format like /v1beta/models/gemini-1.5-pro:generateContent or /models/gemini-1.5-flash:generateContent model = None @@ -3372,6 +3291,11 @@ def make_request( """ Make a request with retries and error handling """ + from services.key_data_guards import guard_request + data = guard_request(api_provider, data) + from services.gemini_credentials import prepare_credentials, restricted + if restricted(api_provider): + url, headers, params = prepare_credentials(url, headers, params, api_provider, None, AuthService) headers = ( with_opencode_request_session(headers, data) if api_provider == "opencode" diff --git a/services/shadow_eval_contract.py b/services/shadow_eval_contract.py index f098f6c5..352d5189 100644 --- a/services/shadow_eval_contract.py +++ b/services/shadow_eval_contract.py @@ -8,6 +8,7 @@ import uuid from services.secret_scan import scan_payload +from services.key_data_guards import allows_shadow_sampling TASKS = ("coding", "extraction", "writing", "reasoning", "chat") REQUEST_BYTES = 65536 @@ -69,7 +70,7 @@ def task_type(payload): def eligible(payload, user, path, *, random_value=None): model = payload.get("model") rate = user.get("shadow_eval_rate") - return (path in {"/v1/chat/completions", "/v1/responses", "/intelligence/v1/chat/completions", + return (allows_shadow_sampling(user) and path in {"/v1/chat/completions", "/v1/responses", "/intelligence/v1/chat/completions", "/v1/messages", "/optimize/v1/chat/completions"} and isinstance(payload.get("messages"), list) and ((isinstance(model, str) and model.startswith(("auto:", "cascade:"))) @@ -78,6 +79,8 @@ def eligible(payload, user, path, *, random_value=None): def make_sample(payload, user, route, answer, model, latency_ms, usage, *, now=None, finish_reason=None, on_skip=None): + if not allows_shadow_sampling(user): + return None request = {name: payload[name] for name in REQUEST_FIELDS if name in payload} answer = {name: answer[name] for name in ("content", "tool_calls") if name in answer} def reject(reason): diff --git a/services/shadow_eval_data_policy.py b/services/shadow_eval_data_policy.py new file mode 100644 index 00000000..8d073c82 --- /dev/null +++ b/services/shadow_eval_data_policy.py @@ -0,0 +1,28 @@ +"""Recheck source-key authority before retained content leaves the evaluation runner.""" + +from error_handlers import APIError +from services.key_data_guards import has_data_controls + + +def source_allows_replay(sample): + from services.auth_service import AuthService + + try: + user = AuthService._load_user_by_username(sample["key_id"]) + from services.dispatch_policy import verify + policy = verify(user.get("data_policy") if user else None) + return bool(user and not has_data_controls(policy)) + except Exception: + # Missing, malformed or unavailable authority cannot authorize content replay. + return False + + +def guarded_dispatch(sample, dispatch): + def guarded(payload): + if not source_allows_replay(sample): + raise APIError("Source key does not permit shadow replay", 503, + {"error": "shadow_source_not_allowed"}) + from services.provider_dispatch import dispatch_as + with dispatch_as({"username": sample["key_id"]}): + return dispatch(payload) + return guarded diff --git a/services/shadow_eval_runner.py b/services/shadow_eval_runner.py index 63b24635..996b3dd4 100644 --- a/services/shadow_eval_runner.py +++ b/services/shadow_eval_runner.py @@ -15,6 +15,7 @@ from services.shadow_eval_contract import TASKS, clean_usage, encoded, finite from services.shadow_eval_judge import judge_payload, outcome, parse_judgment, random_order, tool_validity from services.shadow_eval_store import ShadowEvalStore +from services.shadow_eval_data_policy import guarded_dispatch, source_allows_replay from services import evaluation_noise_floor as noise_floor PRINCIPAL = "internal:shadow-evaluation" @@ -76,6 +77,7 @@ def replay_payload(sample, candidate): def evaluate(sample, candidate, config, dispatch, *, deadline): + dispatch = guarded_dispatch(sample, dispatch) result = {"sample_id": sample["id"], "task_type": sample["task_type"], "candidate_model": candidate, "candidate_route": candidate, "production_model": sample["production_model"], "outcome": "failed", "judge_model": config["judge_model"], "latencies": {"production": sample["latency_ms"], "candidate": 0, "judges": []}, @@ -146,7 +148,7 @@ def run(app, dispatch, *, store=ShadowEvalStore, options=None): if (count + 1) * units > config["max_replays_per_run"] or time.monotonic() >= deadline: return count sample = store.pending(task, model) - if not sample: + if not sample or not source_allows_replay(sample): continue identifier = uuid.uuid4().hex reserve = noise_floor.claim if options is not None else lambda storage, *args: storage.claim(*args) @@ -155,7 +157,7 @@ def run(app, dispatch, *, store=ShadowEvalStore, options=None): count += 1 progressed = True if options is not None: - result = noise_floor.evaluate_three_arm(sample, model, config, dispatch, deadline=deadline, + result = noise_floor.evaluate_three_arm(sample, model, config, guarded_dispatch(sample, dispatch), deadline=deadline, options=options, run_id=run_id, call=call) else: result = evaluate(sample, model, config, dispatch, deadline=deadline) diff --git a/services/shadow_eval_sampling.py b/services/shadow_eval_sampling.py index 9be52827..3477c3ea 100644 --- a/services/shadow_eval_sampling.py +++ b/services/shadow_eval_sampling.py @@ -11,6 +11,8 @@ from services.shadow_eval_contract import ANSWER_BYTES, REQUEST_BYTES, REQUEST_FIELDS, clean_usage, eligible, encoded, make_sample from services.shadow_eval_store import ShadowEvalStore +from services.shadow_eval_data_policy import source_allows_replay +from services.key_data_guards import account_controls from services.retention_policy import RetentionPolicy, request_policy logger = logging.getLogger(__name__) @@ -62,7 +64,8 @@ def _persist(): sample, retention = sample if not retention.allows_content: continue - ShadowEvalStore.put(sample) + if source_allows_replay(sample): + ShadowEvalStore.put(sample) except Exception as error: count("skipped_error") logger.warning("Shadow sample not stored (%s)", type(error).__name__) @@ -198,7 +201,8 @@ def _sampling_decision(payload): if report["high"] or report["truncated"]: count("skipped_secret" if report["high"] else "skipped_error") return None - return json.loads(encoded(payload)), {"username": user.get("username"), "id": user.get("id")} + return json.loads(encoded(payload)), {"username": user.get("username"), "id": user.get("id"), + "data_policy": {"version": 1, **account_controls(user)}} def sample_success(response, payload=None): @@ -219,7 +223,7 @@ def sample_success(response, payload=None): model = response.headers.get("X-MultiLLM-Auto-Selected-Model") or getattr(g, "multillm_model", None) def capture(answer, selected, usage, finish_reason=None): - if not retention.allows_content: + if not retention.allows_content or not source_allows_replay({"key_id": user.get("username") or user.get("id")}): return sample = make_sample(payload, user, route, answer, selected, (time.monotonic() - started) * 1000, usage, finish_reason=finish_reason, on_skip=count) diff --git a/services/stream_cost_breaker.py b/services/stream_cost_breaker.py index 93029ecc..8ff9f31f 100644 --- a/services/stream_cost_breaker.py +++ b/services/stream_cost_breaker.py @@ -78,7 +78,7 @@ def validated_control(payload) -> dict: return {CAP_FIELD: validate_cap(payload["max_stream_cost_usd"])} if enabled() and "max_stream_cost_usd" in payload else {} def sql_select(query: str) -> str: - return query.replace("shadow_eval_rate\n", "shadow_eval_rate, max_stream_cost_microusd\n") if enabled() else query + return query.replace("forced_provider_preferences\n", "forced_provider_preferences, max_stream_cost_microusd\n") if enabled() else query def ensure_sql_cap(connection) -> None: if enabled() and CAP_FIELD not in {row["name"] for row in connection.execute("PRAGMA table_info(users)")}: diff --git a/services/telemetry_export.py b/services/telemetry_export.py index 09433994..54b22aed 100644 --- a/services/telemetry_export.py +++ b/services/telemetry_export.py @@ -21,6 +21,8 @@ import requests +from services.dispatch_network import control_network + from services.observability_adapters import OBSERVABILITY_EXPORTER logger = logging.getLogger(__name__) @@ -138,8 +140,9 @@ def export_once(self) -> int: if not target: continue try: - response = requests.post(target, json=payload, headers=headers(signal), timeout=TIMEOUT, - allow_redirects=False) + with control_network(target): + response = requests.post(target, json=payload, headers=headers(signal), timeout=TIMEOUT, + allow_redirects=False) if response.status_code >= 300: raise ValueError(f"HTTP {response.status_code}") sent = True diff --git a/services/user_store.py b/services/user_store.py index 7e9224a6..e9198289 100644 --- a/services/user_store.py +++ b/services/user_store.py @@ -52,16 +52,33 @@ def _transient(error): return not isinstance(error, PrivateIntelligenceError) or error.status >= 500 +def _policy_protocol(): + """Send policy extensions only after the peer advertises their decoder.""" + try: + capability = request_private_intelligence({"operation": "capabilities"}, endpoint="users").get("key_controls") + except PrivateIntelligenceError as error: + if error.status in (400, 404): + return {} # Old peers reject this operation, without seeing an extended row request. + raise unavailable() from None + except Exception: + raise unavailable() from None + if (not isinstance(capability, dict) or type(capability.get("protocol")) is not int + or capability["protocol"] not in (1, 2)): + raise unavailable() + return {"data_policy_version": 1} + + def _call(operation, **values): """One private call; reads are retried once after a transient failure. Reads are idempotent, so a single retry absorbs a slow first call after the Container or D1 was idle. Writes are sent once. """ + protocol = _policy_protocol() if operation in READ_OPERATIONS else {} attempts = 2 if operation in READ_OPERATIONS else 1 for attempt in range(1, attempts + 1): try: - return request_private_intelligence({"operation": operation, **values}, endpoint="users") + return request_private_intelligence({"operation": operation, **values, **protocol}, endpoint="users") except Exception as error: cause = (f"{error.status} {error.code}" if isinstance(error, PrivateIntelligenceError) else getattr(error, "code", None) or type(error).__name__) @@ -77,8 +94,16 @@ def _call(operation, **values): def _row(value): + policy = None if isinstance(value, dict): # A Worker released before migrations 0011 and 0013 returns rows without these columns. + from services.dispatch_policy import verify + if value.get("data_policy") is not None: + policy = verify(value["data_policy"]) + if "official_only" in value or "forced_provider_preferences" in value: + from services.key_data_guards import stored_controls + policy = {"version": 1, **stored_controls(value)} + value = {name: item for name, item in value.items() if name in USER_FIELDS} value = {"secret_scan_mode": None, "shadow_eval_rate": None, **value} # nosec B105 - column defaults if ( not isinstance(value, dict) @@ -91,7 +116,7 @@ def _row(value): or value["is_admin"] not in (0, 1) ): raise unavailable() - return value + return {**value, **({"data_policy": policy} if policy is not None else {})} def _rows(value, limit=PAGE_SIZE): diff --git a/services/video_generation.py b/services/video_generation.py index b4c0b32f..b0b18b52 100644 --- a/services/video_generation.py +++ b/services/video_generation.py @@ -150,14 +150,15 @@ def _failure(status: int, message: str, transport: str | None = None) -> Respons return response -def _send(method: str, url: str, *, timeout, **kwargs) -> requests.Response | Response: +def _send(method: str, url: str, *, provider: str, timeout, **kwargs) -> requests.Response | Response: """One provider call without retries: a sent job may already be accepted and billed.""" if "json" in kwargs: kwargs["json"] = protect_payload(kwargs["json"], provider="video") if "data" in kwargs: kwargs["data"] = protect_body(kwargs["data"], kwargs.get("headers"), provider="video") try: - return requests.request(method, url, timeout=timeout, allow_redirects=False, **kwargs) + from services.provider_dispatch import provider_http + return provider_http(provider, method, url, timeout=timeout, allow_redirects=False, **kwargs) except requests.exceptions.ConnectTimeout: return _failure(502, "The video provider could not be reached", "connect") except requests.exceptions.ConnectionError as error: @@ -186,14 +187,14 @@ def create_job(provider: str, model: str, request: VideoRequest, key: str | None "seconds": str(_nearest(request.seconds, (4, 8, 12)))} if request.image_url: body["input_reference"] = {"image_url": request.image_url} - response = _send("POST", f"{base_url.rstrip('/')}/v1/videos", timeout=CREATE_TIMEOUT, json=body, + response = _send("POST", f"{base_url.rstrip('/')}/v1/videos", provider=provider, timeout=CREATE_TIMEOUT, json=body, headers={"Authorization": f"Bearer {key}"}) elif provider == "xai": body = {"model": model, "prompt": request.prompt, "duration": min(15, request.seconds), "aspect_ratio": request.aspect_ratio, "resolution": request.resolution} if request.image_url: body["image"] = {"url": request.image_url} - response = _send("POST", f"{base_url.rstrip('/')}/v1/videos/generations", timeout=CREATE_TIMEOUT, json=body, + response = _send("POST", f"{base_url.rstrip('/')}/v1/videos/generations", provider=provider, timeout=CREATE_TIMEOUT, json=body, headers={"Authorization": f"Bearer {key}"}) elif provider == "gemini": instance: dict[str, object] = {"prompt": request.prompt} @@ -204,7 +205,7 @@ def create_job(provider: str, model: str, request: VideoRequest, key: str | None body = {"instances": [instance], "parameters": { "sampleCount": 1, "aspectRatio": aspect_ratio, "resolution": request.resolution, "durationSeconds": _nearest(request.seconds, (4, 6, 8))}} - response = _send("POST", f"{GEMINI_BASE}/models/{model}:predictLongRunning", timeout=CREATE_TIMEOUT, + response = _send("POST", f"{GEMINI_BASE}/models/{model}:predictLongRunning", provider=provider, timeout=CREATE_TIMEOUT, json=body, headers={"x-goog-api-key": key}) elif provider == "cloudflare": if not cloudflare_ai.enabled(): @@ -240,10 +241,10 @@ def job_status(job: dict, job_id: str, key: str | None, base_url: str) -> dict: if provider == "cloudflare": return {**public, "status": "completed", "progress": 100} if provider in ("openai", "xai"): - response = _send("GET", f"{base_url.rstrip('/')}/v1/videos/{upstream}", timeout=POLL_TIMEOUT, + response = _send("GET", f"{base_url.rstrip('/')}/v1/videos/{upstream}", provider=provider, timeout=POLL_TIMEOUT, headers={"Authorization": f"Bearer {key}"}) else: - response = _send("GET", f"{GEMINI_BASE}/{upstream}", timeout=POLL_TIMEOUT, headers={"x-goog-api-key": key}) + response = _send("GET", f"{GEMINI_BASE}/{upstream}", provider=provider, timeout=POLL_TIMEOUT, headers={"x-goog-api-key": key}) if isinstance(response, Response) or response.status_code >= 500: # A failed status read is not a failed job: report it and let the caller poll again. return {**public, "status": "in_progress", "poll_error": "The provider status could not be read; poll again"} @@ -275,13 +276,13 @@ def content_source(job: dict, key: str | None, base_url: str) -> tuple[str, dict if provider == "openai": return f"{base_url.rstrip('/')}/v1/videos/{upstream}/content", {"Authorization": f"Bearer {key}"} if provider == "xai": - response = _send("GET", f"{base_url.rstrip('/')}/v1/videos/{upstream}", timeout=POLL_TIMEOUT, + response = _send("GET", f"{base_url.rstrip('/')}/v1/videos/{upstream}", provider=provider, timeout=POLL_TIMEOUT, headers={"Authorization": f"Bearer {key}"}) url = ((_json(response).get("video") or {}).get("url") if not isinstance(response, Response) else None) if not url: raise APIError("The video is not ready", status_code=409) return url, {} - response = _send("GET", f"{GEMINI_BASE}/{upstream}", timeout=POLL_TIMEOUT, headers={"x-goog-api-key": key}) + response = _send("GET", f"{GEMINI_BASE}/{upstream}", provider=provider, timeout=POLL_TIMEOUT, headers={"x-goog-api-key": key}) samples = (((_json(response).get("response") or {}).get("generateVideoResponse") or {}).get("generatedSamples") if not isinstance(response, Response) else None) or [] uri = ((samples[0] or {}).get("video") or {}).get("uri") if samples else None @@ -290,14 +291,15 @@ def content_source(job: dict, key: str | None, base_url: str) -> tuple[str, dict return uri, {"x-goog-api-key": key} -def stream_content(url: str, headers: dict) -> Response: +def stream_content(url: str, headers: dict, *, provider: str) -> Response: if not url.startswith("https://"): raise APIError("The provider returned an unsupported video location", status_code=502) try: - upstream = requests.get(url, headers=headers, timeout=DOWNLOAD_TIMEOUT, stream=True, allow_redirects=True) + from services.provider_assets import asset_http + upstream = asset_http(provider, url, headers=headers, timeout=DOWNLOAD_TIMEOUT, stream=True) except requests.exceptions.RequestException: raise APIError("The video could not be downloaded; try again", status_code=502) from None - if upstream.status_code >= 400: + if upstream.status_code != 200: upstream.close() raise APIError("The video is not available", status_code=409 if upstream.status_code in (404, 409) else 502) length = int(upstream.headers.get("Content-Length") or 0) diff --git a/tests/conftest.py b/tests/conftest.py index 3afb448c..d6bdb5ec 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,6 +16,20 @@ import pytest +def pytest_sessionstart(session): + """Audit the extension socket too, including connections from alternate clients.""" + import sys + from services.dispatch_network import network_admitted + + def audit(event, args): + if event == "socket.connect" and not network_admitted(args[1]): + raise AssertionError("Outbound connection outside provider chokepoint") + + # Audit hooks cannot be removed. Keep enforcement active for this process, + # including collection, fixture teardown and background completions. + sys.addaudithook(audit) + + @pytest.fixture(autouse=True) def _fresh_public_capability_cache(): # The models.dev index is cached per process; each test sees its own fixtures. @@ -24,3 +38,22 @@ def _fresh_public_capability_cache(): reset_public_capability_cache() yield reset_public_capability_cache() + + +@pytest.fixture(autouse=True) +def _outbound_connections_require_dispatch(monkeypatch): + """All HTTP clients, including unlisted modules and SDKs, meet this boundary.""" + import socket + from services.dispatch_network import network_admitted + + def guard(original): + def checked(*args, **kwargs): + address = args[1] if args and isinstance(args[0], socket.socket) else (args[0] if args else kwargs.get("address")) + if not network_admitted(address): + raise AssertionError("Outbound connection outside provider chokepoint") + return original(*args, **kwargs) + return checked + + monkeypatch.setattr(socket, "create_connection", guard(socket.create_connection)) + monkeypatch.setattr(socket.socket, "connect", guard(socket.socket.connect)) + monkeypatch.setattr(socket.socket, "connect_ex", guard(socket.socket.connect_ex)) diff --git a/tests/dispatch_fixture.py b/tests/dispatch_fixture.py new file mode 100644 index 00000000..d0a5bf23 --- /dev/null +++ b/tests/dispatch_fixture.py @@ -0,0 +1,20 @@ +"""Explicit durable identity for isolated provider transport tests.""" + +import pytest + +from services.dispatch_policy import envelope +from services.provider_dispatch import dispatch_as + + +@pytest.fixture +def verified_transport_principal(monkeypatch): + from services.auth_service import AuthService + + # Exercise real admission with an explicitly unrestricted stored account. + # Policy and authentication regressions use their own stores instead. + monkeypatch.setattr(AuthService, "_load_user_by_username", lambda username: { + "username": username, + "data_policy": envelope({"official_only": None, "forced_provider_preferences": None}), + }) + with dispatch_as({"username": "synthetic-transport"}): + yield diff --git a/tests/fixtures/secret_firewall_coverage.json b/tests/fixtures/secret_firewall_coverage.json index 02dccef4..ba9bb67b 100644 --- a/tests/fixtures/secret_firewall_coverage.json +++ b/tests/fixtures/secret_firewall_coverage.json @@ -307,6 +307,8 @@ "worker/credits-admission.mjs", "worker/credits-d1.mjs", "worker/d1-schema.mjs", + "worker/dispatch-network.mjs", + "worker/dispatch-policy.mjs", "worker/enterprise-contract.mjs", "worker/exact-generation-cache.mjs", "worker/fallback-page.mjs", @@ -321,6 +323,8 @@ "worker/intelligence-d1.mjs", "worker/intelligence-outbound.mjs", "worker/janitor-reasoning-response.mjs", + "worker/key-controls-capability.mjs", + "worker/key-data-guards.mjs", "worker/knowledge-edge.mjs", "worker/knowledge-outbound.mjs", "worker/knowledge/alexandria/catalogue.mjs", @@ -371,6 +375,7 @@ "worker/knowledge/skills.mjs", "worker/latency-slo.mjs", "worker/learned-cooldown-d1.mjs", + "worker/legacy-binding-policy.mjs", "worker/log.mjs", "worker/managed-state-dispatch.mjs", "worker/media-files.mjs", @@ -385,11 +390,13 @@ "worker/opencode/reasoning-request.mjs", "worker/opencode/reasoning-response.mjs", "worker/payment-webhooks.mjs", + "worker/peer-controls.mjs", "worker/pii-redaction.mjs", "worker/pii-rehydration.mjs", "worker/prometheus.mjs", "worker/prompt-injection-detection.mjs", "worker/prompt-templates-d1.mjs", + "worker/provider-dispatch.mjs", "worker/realtime-metering.mjs", "worker/realtime-session.mjs", "worker/realtime.mjs", diff --git a/tests/helpers/dispatch_fixture.mjs b/tests/helpers/dispatch_fixture.mjs new file mode 100644 index 00000000..53313e04 --- /dev/null +++ b/tests/helpers/dispatch_fixture.mjs @@ -0,0 +1,37 @@ +/** Explicit durable controls for transport tests; never a runtime fallback. */ +export function policyStore(controls = { official_only: null, forced_provider_preferences: null }) { + return { prepare(query) { + if (!query.startsWith("SELECT official_only, forced_provider_preferences FROM control_users")) throw Error(`Unexpected policy query: ${query}`); + return { bind() { return this; }, async first() { return structuredClone(controls); }, async all() { return { results: [structuredClone(controls)] }; } }; + } }; +} +export const dispatchEnv = () => ({ INTELLIGENCE_DB: policyStore() }); + +const policyEnvironments = new WeakSet(); + +/** Independent middleware tests supply a durable policy without emulating unrelated SQL. */ +export function policyFixtureEnv(env = {}) { + if (policyEnvironments.has(env)) return env; + const db = env.INTELLIGENCE_DB; + const policy = policyStore(); + const bound = { ...env, INTELLIGENCE_DB: { ...db, prepare(query) { + if (query.startsWith("SELECT official_only, forced_provider_preferences FROM control_users") + || query.startsWith("SELECT official_only, forced_provider_preferences FROM intelligence_principals")) { + return policy.prepare(query.replace("intelligence_principals", "control_users")); + } + if (!db) throw Error(`Unexpected fixture SQL: ${query}`); + return db.prepare(query); + } } }; + policyEnvironments.add(bound); + return bound; +} + + +/** Real binding fixtures use the same NULL controls in a migrated disposable D1. */ +export async function initializePolicyAccount(mf, username) { + const { applyMigrations } = await import("../d1_migrations.mjs"); + const db = await mf.getD1Database("INTELLIGENCE_DB"); + await applyMigrations(db); + await db.prepare("INSERT INTO control_users(username,api_key_hash,api_key_prefix,scopes,is_admin,created_at) VALUES(?,?,?,'knowledge:read,knowledge:manage',0,'2026-10-01T00:00:00Z')") + .bind(username, "synthetic-hash", "synthetic-prefix").run(); +} diff --git a/tests/helpers/load_cloudflare_worker.mjs b/tests/helpers/load_cloudflare_worker.mjs index eafa4e70..4579d652 100644 --- a/tests/helpers/load_cloudflare_worker.mjs +++ b/tests/helpers/load_cloudflare_worker.mjs @@ -1,6 +1,12 @@ import { readFile } from "node:fs/promises"; +export const { policyStore, policyFixtureEnv } = await import(new URL("./dispatch_fixture.mjs", import.meta.url)); function dataModuleUrl(source) { + source = source.replaceAll('from "./worker/build-id.mjs"', `from "${new URL("../../worker/build-id.mjs", import.meta.url)}"`); + source = source.replaceAll('from "./worker/peer-controls.mjs"', `from "${new URL("../../worker/peer-controls.mjs", import.meta.url)}"`); + source = source.replaceAll('from "./worker/provider-dispatch.mjs"', `from "${new URL("../../worker/provider-dispatch.mjs", import.meta.url)}"`); + source = source.replaceAll('from "./worker/key-controls-capability.mjs"', `from "${new URL("../../worker/key-controls-capability.mjs", import.meta.url)}"`); + source = source.replaceAll('from "../provider-dispatch.mjs"', `from "${new URL("../../worker/provider-dispatch.mjs", import.meta.url)}"`); return `data:text/javascript;base64,${Buffer.from(source, "utf8").toString("base64")}`; } @@ -796,3 +802,12 @@ export async function loadWorkerModule({ transformSource = source => source } = return import(dataModuleUrl(patchedSource)); } + + +/** Header and rendering contracts explicitly use a readable unrestricted store. */ +export async function loadWorkerWithUnrestrictedControls(options) { + const module = await loadWorkerModule(options), worker = module.default; + return { ...module, default: { ...worker, fetch(request, env, ctx) { + return worker.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} diff --git a/tests/helpers/local_miniflare.mjs b/tests/helpers/local_miniflare.mjs new file mode 100644 index 00000000..44155c7d --- /dev/null +++ b/tests/helpers/local_miniflare.mjs @@ -0,0 +1,19 @@ +/** Register only each disposable D1 server's exact loopback port with the guard. */ +export * from "miniflare"; +import { Miniflare as NativeMiniflare } from "miniflare"; +import { registerFixtureOrigin } from "./network_guard.mjs"; +export class Miniflare extends NativeMiniflare { + constructor(...args) { + super(...args); + this.fixtureRegistration = this.ready.then(registerFixtureOrigin); + } + async getD1Database(...args) { + await this.fixtureRegistration; + return super.getD1Database(...args); + } + async dispose() { + const unregister = await this.fixtureRegistration.catch(() => null); + try { return await super.dispose(); } + finally { unregister?.(); } + } +} diff --git a/tests/helpers/network_fixture.mjs b/tests/helpers/network_fixture.mjs new file mode 100644 index 00000000..18a4b07e --- /dev/null +++ b/tests/helpers/network_fixture.mjs @@ -0,0 +1,7 @@ +import { controlNetwork } from "../../worker/dispatch-network.mjs"; +export function mockFetch(t, implementation) { + const previous = globalThis.fetch; + globalThis.fetch = implementation; + t.after(() => { globalThis.fetch = previous; }); +} +export const clientNetwork = (url, run) => controlNetwork(url, run); diff --git a/tests/helpers/network_guard.mjs b/tests/helpers/network_guard.mjs new file mode 100644 index 00000000..8dc9630b --- /dev/null +++ b/tests/helpers/network_guard.mjs @@ -0,0 +1,39 @@ +import net from "node:net"; +/** Suite-wide runtime guard; replacing fetch cannot replace the admission check. */ +import { networkAdmitted, connectionAdmitted } from "../../worker/dispatch-network.mjs"; +const guarded = new WeakSet(); +function wrap(transport) { + if (guarded.has(transport)) return transport; + const checked = function (...args) { + if (!networkAdmitted(args[0])) throw new Error("Outbound connection outside provider chokepoint"); + return transport.apply(this, args); + }; + guarded.add(checked); + return checked; +} +let transport = wrap(globalThis.fetch); +Object.defineProperty(globalThis, "fetch", { configurable: false, get: () => transport, + set: value => { transport = wrap(value); } }); +if (typeof globalThis.connect === "function") { + let connect = wrap(globalThis.connect); + Object.defineProperty(globalThis, "connect", { configurable: false, get: () => connect, + set: value => { connect = wrap(value); } }); +} + +// Only explicitly registered local D1 bridge ports are infrastructure traffic. +const fixturePorts = new Set(); +export function registerFixtureOrigin(url) { + const target = new URL(url); + if (!["127.0.0.1", "[::1]", "localhost"].includes(target.hostname)) throw Error("Non-local test fixture"); + fixturePorts.add(String(target.port)); + return () => fixturePorts.delete(String(target.port)); +} +const socketConnect = net.Socket.prototype.connect; +net.Socket.prototype.connect = function (...args) { + const options = Array.isArray(args[0]) ? args[0][0] : args[0]; + const host = typeof options === "object" ? options.host : args[1]; + const port = typeof options === "object" ? options.port : options; + const fixtureRpc = ["127.0.0.1", "::1", "localhost"].includes(host) && fixturePorts.has(String(port)); + if (!connectionAdmitted(host, port) && !fixtureRpc) throw new Error("Outbound connection outside provider chokepoint"); + return socketConnect.apply(this, args); +}; diff --git a/tests/helpers/policy_test.mjs b/tests/helpers/policy_test.mjs new file mode 100644 index 00000000..1363fb17 --- /dev/null +++ b/tests/helpers/policy_test.mjs @@ -0,0 +1,10 @@ +/** Unit transports have an explicit authenticated durable account. */ +import nativeTest from "node:test"; +import { withDispatchPrincipal } from "../../worker/provider-dispatch.mjs"; +import { dispatchEnv } from "./dispatch_fixture.mjs"; +function policyTest(name, ...args) { + const callback = args.pop(); + return nativeTest(name, ...args, (...values) => withDispatchPrincipal(dispatchEnv(), + { authenticated: true, keyId: "fixture-reader" }, () => callback(...values))); +} +export default Object.assign(policyTest, nativeTest); diff --git a/tests/helpers/provider_binding.mjs b/tests/helpers/provider_binding.mjs new file mode 100644 index 00000000..baa96fbb --- /dev/null +++ b/tests/helpers/provider_binding.mjs @@ -0,0 +1,28 @@ +/** Build every provider SDK/binding double here, including nested SDK handles. */ +import assert from "node:assert/strict"; +import { networkAdmitted } from "../../worker/dispatch-network.mjs"; + +export function providerBinding(binding) { + const wrapped = new WeakMap(); + function guard(value) { + if (!value || Array.isArray(value) || typeof value !== "object" && typeof value !== "function") return value; + if (wrapped.has(value)) return wrapped.get(value); + const proxy = new Proxy(value, { + get(target, name, receiver) { + const member = Reflect.get(target, name, receiver); + return name === "constructor" ? member : guard(member); + }, + apply(target, receiver, args) { + assert.ok(networkAdmitted(), "Provider binding outside provider chokepoint"); + const result = Reflect.apply(target, receiver, args); + // Synchronous SDK factories return handles whose methods need their own + // admission check. Async responses and streams remain ordinary results. + return result && Object.getPrototypeOf(result) === Object.prototype ? guard(result) : result; + }, + }); + wrapped.set(value, proxy); + wrapped.set(proxy, proxy); + return proxy; + } + return guard(binding); +} diff --git a/tests/helpers/roleplay_fixture.mjs b/tests/helpers/roleplay_fixture.mjs index a1f3aada..1862c397 100644 --- a/tests/helpers/roleplay_fixture.mjs +++ b/tests/helpers/roleplay_fixture.mjs @@ -1,4 +1,4 @@ -import { loadRoleplayModule } from "./load_cloudflare_worker.mjs"; +import { loadRoleplayModule, policyStore } from "./load_cloudflare_worker.mjs"; const { RoleplaySession, @@ -60,6 +60,7 @@ export function makeRoleplayEnv(overrides = {}) { const storageBySession = new Map(); const waits = []; const env = { + INTELLIGENCE_DB: policyStore(), ADMIN_API_KEY: "admin-roleplay-key", ROLEPLAY_API_KEY: "janitor-roleplay-key", OPENCODE_GO_API_KEY: "opencode-roleplay-key", diff --git a/tests/helpers/runtime_versions.mjs b/tests/helpers/runtime_versions.mjs new file mode 100644 index 00000000..16945384 --- /dev/null +++ b/tests/helpers/runtime_versions.mjs @@ -0,0 +1,44 @@ +import { execFileSync } from "node:child_process"; +import { build } from "esbuild"; + +export const OLD = "544bf49ba3b3d03d185ec5b5a71b2b1877dda786"; + +/** Compile the actual pinned source tree, without replacing protocol handlers. */ +export async function runtimeBundle(contents, revision = "", platform = "node") { + const root = process.cwd(); + const output = await build({ stdin: { resolveDir: root, contents }, bundle: true, + format: "esm", platform, external: ["node:*"], write: false, + plugins: [{ name: "runtime-version", setup(b) { + b.onResolve({ filter: /^cloudflare:workers$/ }, () => ({ path: "runtime", namespace: "stub" })); + b.onLoad({ filter: /.*/, namespace: "stub" }, () => ({ contents: + "export class DurableObject {constructor(ctx,env){this.ctx=ctx;this.env=env}}; export class WorkflowEntrypoint {}" })); + if (platform === "node") b.onResolve({ filter: /dispatch-network\.mjs$/ }, () => ({ + path: new URL("../../worker/dispatch-network.mjs", import.meta.url).href, external: true, + })); + if (revision) b.onLoad({ filter: /\/worker\/.*\.(mjs|json)$/ }, args => ({ + contents: execFileSync("git", ["show", `${revision}:${args.path.slice(root.length + 1)}`], { encoding: "utf8" }), + loader: args.path.endsWith(".json") ? "json" : "js", + })); + } }] }); + return output.outputFiles[0].text; +} + +export async function runtimeModule(contents, revision) { + const source = await runtimeBundle(contents, revision); + return import(`data:text/javascript;base64,${Buffer.from(source).toString("base64")}`); +} + +export class MemoryStorage { + constructor() { this.rows = new Map(); this.pending = Promise.resolve(); } + async get(k) { return structuredClone(this.rows.get(k)); } + async put(k, v) { + if (typeof k === "string") this.rows.set(k, structuredClone(v)); + else for (const [name, value] of Object.entries(k)) this.rows.set(name, structuredClone(value)); + } + async delete(keys) { for (const key of Array.isArray(keys) ? keys : [keys]) this.rows.delete(key); } + async list({ prefix = "", limit = Infinity, startAfter } = {}) { + return new Map([...this.rows].filter(([k]) => k.startsWith(prefix) && (!startAfter || k > startAfter)) + .sort().slice(0, limit).map(([k, v]) => [k, structuredClone(v)])); + } + transaction(fn) { const result = this.pending.then(() => fn(this)); this.pending = result.catch(() => {}); return result; } +} diff --git a/tests/knowledge_fixture.mjs b/tests/knowledge_fixture.mjs index 5a08500e..d7ff2953 100644 --- a/tests/knowledge_fixture.mjs +++ b/tests/knowledge_fixture.mjs @@ -1,3 +1,4 @@ +import { policyStore } from "./helpers/dispatch_fixture.mjs"; import { KnowledgeAuthority } from "../worker/knowledge/authority.mjs"; import { createArtifact } from "../worker/knowledge/evidence.mjs"; import { defaultPolicy } from "../worker/knowledge/policy.mjs"; @@ -63,7 +64,7 @@ export async function fixture() { async match(key) { return cacheEntries.get(key.url)?.clone(); }, async put(key, value) { cacheEntries.set(key.url, value.clone()); }, }; - const env = { KNOWLEDGE_SNAPSHOTS: {}, KNOWLEDGE_INDEX: {}, EXA_API_KEY: "synthetic-test-key" }; + const env = { INTELLIGENCE_DB: policyStore(), KNOWLEDGE_SNAPSHOTS: {}, KNOWLEDGE_INDEX: {}, EXA_API_KEY: "synthetic-test-key" }; const retrieve = async (provider, intent, { invoke }) => invoke(provider, "lookup", async () => { counts.providers++; return { observations: [{ kind: "source_excerpt", url: source.url, text, provider, freshness: "live" }], warnings: [] }; diff --git a/tests/support/account_envelope_driver.py b/tests/support/account_envelope_driver.py new file mode 100644 index 00000000..0d47914f --- /dev/null +++ b/tests/support/account_envelope_driver.py @@ -0,0 +1,127 @@ +"""Registered Container routes, using actual HTTP to the disposable main Worker.""" +import json +import os +import re +import sys +import tempfile +import types +from pathlib import Path + +root = str(Path(__file__).resolve().parents[2]) +sys.path.insert(0, root) +# The child receives only explicit synthetic configuration, never a host dotenv. +sys.modules["env_loader"] = types.SimpleNamespace(load_runtime_env=lambda *a, **kw: None) +os.environ.update(json.loads(sys.argv[2])) +scratch = tempfile.TemporaryDirectory(prefix="multillm-account-envelope-") +os.environ.update({ + "CONTROL_PLANE_DATABASE_URL": "", "PROVIDER_CATALOG_AUTO_REFRESH": "false", "USAGE_LEDGER_ENABLED": "false", + "AUTH_DB_PATH": scratch.name + "/auth.sqlite3", "MODEL_REGISTRY_DB_PATH": scratch.name + "/registry.sqlite3", + "RATE_LIMIT_DB_PATH": scratch.name + "/limits.sqlite3", +}) + +# Load a pinned earlier runtime in memory for actual mixed-version route probes. +revision = sys.argv[3] if len(sys.argv) > 3 else "" +if revision: + import importlib.abc + import importlib.util + import subprocess + + class RuntimeModules(importlib.abc.MetaPathFinder, importlib.abc.Loader): + def find_spec(self, fullname, path=None, target=None): + if fullname.split(".")[0] not in {"app", "config", "error_handlers", "route_helpers", "services", "routes", "providers", "streaming"}: + return None + name = fullname.replace(".", "/") + for relative, package in [(name + ".py", False), (name + "/__init__.py", True)]: + result = subprocess.run(["git", "show", revision + ":" + relative], cwd=root, capture_output=True, text=True) + if result.returncode == 0: + spec = importlib.util.spec_from_loader(fullname, self, origin=root + "/" + relative, is_package=package) + spec.loader_state = (relative, result.stdout) + return spec + return None + + def create_module(self, spec): + return None + + def exec_module(self, module): + relative, source = module.__spec__.loader_state + module.__file__ = root + "/" + relative + if module.__spec__.submodule_search_locations is not None: + module.__path__ = [root + "/" + relative.rsplit("/", 1)[0]] + exec(compile(source, module.__file__, "exec"), module.__dict__) + + sys.meta_path.insert(0, RuntimeModules()) + +from services import intelligence_d1_store # noqa: E402 + +base = sys.argv[1] +for name, url in intelligence_d1_store._ENDPOINTS.items(): + intelligence_d1_store._ENDPOINTS[name] = base + url.removeprefix("http://intelligence.internal") + +import requests # noqa: E402 + +seen = [] +original = requests.sessions.Session.request + + +def transport(self, method, url, **kwargs): + if url.startswith(base + "/"): + return original(self, method, url, **kwargs) + if url.startswith("https://openrouter.ai/"): + from services.dispatch_network import network_admitted + if revision != "544bf49ba3b3d03d185ec5b5a71b2b1877dda786": + assert network_admitted() + if os.environ.get("ACCOUNT_DRIVER_OUTBOUND_ADMISSION") == "true": + admitted = original(self, "GET", base + "/probe-provider") + if admitted.status_code != 200: + return admitted + payload = kwargs.get("json") or json.loads(kwargs.get("data") or "{}") + if "messages" in payload: + seen.append(payload) + response = requests.Response() + response.status_code = 200 + response.headers["Content-Type"] = "application/json" + response._content = json.dumps({"data": [], "choices": [{ + "message": {"content": "synthetic"}, "finish_reason": "stop", + }]}).encode() + return response + raise AssertionError("Unexpected external request") + + +# Only the external provider is mocked. Account/control requests travel over HTTP. +requests.sessions.Session.request = transport +from app import app # noqa: E402 + +if revision: + from jinja2 import FunctionLoader + + def versioned_template(name): + if ".." in name or name.startswith("/"): + return None + result = subprocess.run(["git", "show", revision + ":templates/" + name], + cwd=root, capture_output=True, text=True) + return result.stdout if result.returncode == 0 else None + + app.jinja_loader = FunctionLoader(versioned_template) +client = app.test_client() +headers = {"Authorization": "Bearer mllm_synthetic_incident_key_long_enough"} +requests_to_run = json.loads(sys.argv[4]) if len(sys.argv) > 4 else [("/v1/models", "get"), ("/v1/chat/completions", "post")] +for operation in requests_to_run: + path, method = operation[:2] + if path == "/login": + username = operation[2] + client = app.test_client() + page = client.get("/login") + token = re.search(r'name="csrf_token" value="([^"]+)"', page.get_data(as_text=True)).group(1) + response = client.post("/login", data={"username": username, "csrf_token": token, + "api_key": "mllm_synthetic_incident_key_long_enough" if username == "synthetic" else os.environ["ADMIN_API_KEY"]}) + print(json.dumps({"path": path, "status": response.status_code, "dispatches": len(seen)})) + continue + arguments = {"json": {"model": "openrouter:openai/gpt-6-luna", "messages": [ + {"role": "user", "content": "synthetic"}, + ]}} if method == "post" else {} + request_headers = {"Authorization": "Bearer " + os.environ["ADMIN_API_KEY"]} if len(operation) > 2 and operation[2] == "admin" else headers + response = getattr(client, method)(path, headers=request_headers, **arguments) + print(json.dumps({"path": path, "status": response.status_code, + "error": response.json.get("error") if isinstance(response.json, dict) else None, + "dispatches": len(seen), "preferences": seen[-1].get("provider") if seen else None})) +scratch.cleanup() diff --git a/tests/support/config_revision_fixtures.mjs b/tests/support/config_revision_fixtures.mjs index 06c1af8b..dfe81b6e 100644 --- a/tests/support/config_revision_fixtures.mjs +++ b/tests/support/config_revision_fixtures.mjs @@ -1,5 +1,5 @@ // Shared D1 fixtures for the config revision Worker and integration suites. -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "../helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../../worker/intelligence-outbound.mjs"; import { handleRevisionedAutoRoutes } from "../../worker/config-revision.mjs"; import { handleAutoRoutesRequest } from "../../worker/auto-routes-d1.mjs"; @@ -28,7 +28,7 @@ export const account = (changes = {}) => ({ username: "alice", api_key_hash: "sy scopes: "chat,models", is_admin: 0, created_at: "2026-10-09T00:00:00+00:00", last_login: null, last_used_at: null, last_used_ip: null, created_by: null, rotated_at: null, revoked_at: null, daily_budget_usd: null, monthly_budget_usd: null, allowed_models: null, allowed_ips: null, - expires_at: null, secret_scan_mode: null, shadow_eval_rate: null, ...changes }); + expires_at: null, secret_scan_mode: null, shadow_eval_rate: null, official_only: null, forced_provider_preferences: null, ...changes }); const usersRequest = body => new Request("http://intelligence.internal/v1/users", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, ...body }), diff --git a/tests/test_ai_outbound_worker.mjs b/tests/test_ai_outbound_worker.mjs index c62d4d7a..e8177ee1 100644 --- a/tests/test_ai_outbound_worker.mjs +++ b/tests/test_ai_outbound_worker.mjs @@ -1,20 +1,22 @@ +import { policyStore } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; import test from "node:test"; +import { providerBinding } from "./helpers/provider_binding.mjs"; import { handleAiOutbound, imageInput } from "../worker/ai-outbound.mjs"; import { collectContainerEnv } from "../worker/container-env.mjs"; function ai(result) { const calls = []; - return { calls, async run(model, input, options) { + return providerBinding({ calls, async run(model, input, options) { calls.push({ model, input, options }); if (result instanceof Error) throw result; return typeof result === "function" ? result(model, input) : result; - } }; + } }); } const post = (path, body, env) => handleAiOutbound(new Request(`http://ai.internal${path}`, { - method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(body) }), env); + method: "POST", headers: { "content-type": "application/json", "X-MultiLLM-Dispatch-Principal": '{"version":1,"username":"fixture-reader"}' }, body: JSON.stringify(body) }), { INTELLIGENCE_DB: policyStore(), ...env }); test("third-party GPT Image models run through AI Gateway with OpenAI parameters", async () => { const env = { AI: ai({ result: { image: "https://images.example/out.png" } }), AI_GATEWAY_ID: "media" }; @@ -30,7 +32,7 @@ test("third-party GPT Image models run through AI Gateway with OpenAI parameters test("an image URL is downloaded when the caller asks for base64", async () => { const original = globalThis.fetch; - globalThis.fetch = async url => { assert.equal(url, "https://images.example/out.png"); return new Response(new Uint8Array([1, 2, 3])); }; + globalThis.fetch = async url => { assert.equal(url instanceof Request ? url.url : url, "https://images.example/out.png"); return new Response(new Uint8Array([1, 2, 3])); }; try { const env = { AI: ai({ image: "https://images.example/out.png" }) }; const body = await (await post("/v1/images/generations", { model: "openai/gpt-image-2", prompt: "x" }, env)).json(); diff --git a/tests/test_app_structure.py b/tests/test_app_structure.py index 3bc1aa77..713907bc 100644 --- a/tests/test_app_structure.py +++ b/tests/test_app_structure.py @@ -1,4 +1,6 @@ import json +import re +import shlex from pathlib import Path import unittest @@ -22,11 +24,28 @@ def test_ci_uses_current_worker_test_path(self): worker_test = package["scripts"]["test:worker"] self.assertIn("npm run test:worker", workflow_text) - self.assertIn("node --test tests/test_cloudflare_worker.mjs", worker_test) + self.assertIn("node --import ./tests/helpers/network_guard.mjs --test tests/test_cloudflare_worker.mjs", worker_test) self.assertIn("tests/test_roleplay_worker.mjs", worker_test) self.assertIn("tests/test_roleplay_memory_worker.mjs", worker_test) self.assertNotIn("node --test test_cloudflare_worker.mjs", workflow_text) + # Standalone Node CI steps must also run in the ordinary Worker suite. + node_job = workflow_text.split("\n node:\n", 1)[1].split("\n integration:\n", 1)[0] + worker_patterns = shlex.split(worker_test)[shlex.split(worker_test).index("--test") + 1:] + covered = {str(path.relative_to(repository)) for pattern in worker_patterns for path in repository.glob(pattern)} + for command in re.findall(r"^\s+run: (.+)$", node_job, re.MULTILINE): + if command.startswith("npm run "): + command = package["scripts"][command.removeprefix("npm run ")] + tokens = shlex.split(command) + if tokens[0] != "node" or "--test" not in tokens: + continue + for pattern in tokens[tokens.index("--test") + 1:]: + files = list(repository.glob(pattern)) + self.assertTrue(files, f"CI test pattern matches no files: {pattern}") + for path in files: + self.assertIn(str(path.relative_to(repository)), covered, + f"Node CI test is missing from test:worker: {path.name}") + if __name__ == "__main__": unittest.main() diff --git a/tests/test_auto_routes_worker.mjs b/tests/test_auto_routes_worker.mjs index efdfaf27..bace0cf9 100644 --- a/tests/test_auto_routes_worker.mjs +++ b/tests/test_auto_routes_worker.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { applyMigrations } from "./d1_migrations.mjs"; diff --git a/tests/test_cascades.py b/tests/test_cascades.py index 97a02887..0630ee07 100644 --- a/tests/test_cascades.py +++ b/tests/test_cascades.py @@ -633,7 +633,9 @@ def test_cascade_is_sampled_once_as_the_outer_request_and_never_per_tier(context submitted = [] monkeypatch.setattr(sampling, "submit", submitted.append) monkeypatch.setattr(sampling.random, "random", lambda: 0.0) - g.authenticated_user = {"username": "synthetic", "is_admin": True, "shadow_eval_rate": 0.2} + g.authenticated_user = {"username": "synthetic", "is_admin": True, "shadow_eval_rate": 0.2, + "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}} + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda _: g.authenticated_user) g.shadow_eval_started = 0.0 @sampling.sample_chat_dispatch diff --git a/tests/test_cascades_worker.mjs b/tests/test_cascades_worker.mjs index 7a15b4de..20a2b928 100644 --- a/tests/test_cascades_worker.mjs +++ b/tests/test_cascades_worker.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { validCascade } from "../worker/cascades-d1.mjs"; import { CORS_EXPOSE_HEADERS } from "../worker/cors-policy.mjs"; diff --git a/tests/test_child_safety_worker.mjs b/tests/test_child_safety_worker.mjs index fafbd730..a46669d3 100644 --- a/tests/test_child_safety_worker.mjs +++ b/tests/test_child_safety_worker.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { handleControlUsersRequest } from "../worker/control-users-d1.mjs"; diff --git a/tests/test_client_headers_worker.mjs b/tests/test_client_headers_worker.mjs index 83d3b6a2..22fecf38 100644 --- a/tests/test_client_headers_worker.mjs +++ b/tests/test_client_headers_worker.mjs @@ -5,6 +5,7 @@ import { collectContainerEnv } from "../worker/container-env.mjs"; import { MAX_SESSION_BODY_BYTES, withOpencodeRequestSession } from "../worker/opencode-session.mjs"; import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; import { completionResponse, makeRoleplayEnv, roleplayRequest, withGlobalFetch } from "./helpers/roleplay_fixture.mjs"; +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; const { default: worker } = await loadWorkerModule(); @@ -80,7 +81,7 @@ test("stalled affinity discovery times out without consuming the original upload test("every direct OpenCode protocol supplies a session without harness headers", async () => { for (const path of ["/opencode/v1/chat/completions", "/opencode/v1/responses", "/opencode/v1/messages", "/opencode/v1/models"]) { - const env = { ADMIN_API_KEY: "admin-test", OPENCODE_GO_API_KEY: "opencode-test", OPENCODE_EDGE_FETCH: "true" }; + const env = policyFixtureEnv({ ADMIN_API_KEY: "admin-test", OPENCODE_GO_API_KEY: "opencode-test", OPENCODE_EDGE_FETCH: "true" }); let seen; const response = await withGlobalFetch(async (request) => { seen = request.headers; @@ -97,6 +98,20 @@ test("every direct OpenCode protocol supplies a session without harness headers" } }); +test("direct OpenCode requests refuse absent or unreadable current controls", async () => { + for (const store of [undefined, { prepare() { throw new Error("Synthetic control-store outage"); } }]) { + const env = { ADMIN_API_KEY: "admin-test", OPENCODE_GO_API_KEY: "opencode-test", OPENCODE_EDGE_FETCH: "true", + ...(store ? { INTELLIGENCE_DB: store } : {}) }; + const response = await withGlobalFetch(() => assert.fail("Unverified request reached the provider"), () => worker.fetch( + new Request("https://proxy.example/opencode/v1/chat/completions", { + method: "POST", headers: { Authorization: "Bearer admin-test", "Content-Type": "application/json" }, + body: '{"input":"Synthetic request"}', + }), env)); + assert.equal(response.status, 503); + assert.equal((await response.json()).error.code, "key_data_controls_unavailable"); + } +}); + test("Codex defaults are overridable, validated, and do not mutate caller headers", () => { const source = new Headers({ "uSeR-aGeNt": "fleet/1", Authorization: "Bearer test" }); const env = { UPSTREAM_DEFAULT_USER_AGENT: "central/2", UPSTREAM_DEFAULT_ORIGINATOR: "central" }; @@ -140,11 +155,11 @@ test("all Worker provider namespaces carry defaults and caller overrides", async requests.push(request); return new Response('{"data":[]}', { headers: { "Content-Type": "application/json" } }); }; - const env = { + const env = policyFixtureEnv({ ADMIN_API_KEY: "admin-test", OPENCODE_GO_API_KEY: "opencode-test", OPENCODE_EDGE_FETCH: "true", LINKAPI_KEY: "linkapi-test", CODEX_EASY_API_KEY: "easy-test", KIMI_CODE_API_KEY: "kimi-test", MULTILLM_PROXY_CONTAINER: { getByName: () => ({ fetch: receive }) }, - }; + }); const isGet = path.endsWith("/models"); const response = await withGlobalFetch(receive, () => worker.fetch(new Request(`https://proxy.example${path}`, { method: isGet ? "GET" : "POST", diff --git a/tests/test_cline_pass_provider.py b/tests/test_cline_pass_provider.py index 56a2e7ef..c52c0646 100644 --- a/tests/test_cline_pass_provider.py +++ b/tests/test_cline_pass_provider.py @@ -41,6 +41,12 @@ } + +import pytest +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + class ClinePassProviderTest(unittest.TestCase): def test_chat_goes_to_the_cline_api_with_a_long_timeout(self): adapter = get_adapter("cline-pass", Config.API_BASE_URLS) diff --git a/tests/test_cloudflare_worker.mjs b/tests/test_cloudflare_worker.mjs index 233c2c60..e843ea5b 100644 --- a/tests/test_cloudflare_worker.mjs +++ b/tests/test_cloudflare_worker.mjs @@ -1,3 +1,4 @@ +import { policyStore } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; import { readFile } from "node:fs/promises"; import test from "node:test"; @@ -5,7 +6,10 @@ import test from "node:test"; import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; const workerModule = await loadWorkerModule(); -const worker = workerModule.default; +const worker = { ...workerModule.default, fetch(request, env, ctx) { + return workerModule.default.fetch(request, new URL(request.url).pathname === "/ready" ? env + : { INTELLIGENCE_DB: policyStore(), ...env }, ctx); +} }; const { MultiLLMProxyContainer } = workerModule; function makeEnv(fetchImpl, envOverrides = {}) { @@ -2435,6 +2439,7 @@ test("worker answers /health directly without touching the container", async () assert.equal(stub.getCalls(), 0); assert.deepEqual(await response.json(), { status: "healthy", + key_controls: {protocol: 1, enforcement: false}, mode: "worker-fallback", }); }); @@ -2446,7 +2451,7 @@ test("worker checks application readiness through the container health endpoint" assert.equal(forwardedUrl.pathname, "/healthz"); assert.equal(request.method, "GET"); assert.equal(request.redirect, "manual"); - return new Response(JSON.stringify({ status: "ready", database: "ok" }), { + return new Response(JSON.stringify({ status: "ready", database: "ok", main_key_controls: {protocol: 1, enforcement: false} }), { status: 200, headers: { "Content-Type": "application/json" }, }); @@ -2463,7 +2468,8 @@ test("worker checks application readiness through the container health endpoint" assert.equal(stub.getCalls(), 1); assert.equal(stub.getStartCalls(), 0); assert.equal(response.headers.get("Access-Control-Allow-Origin"), origin); - assert.deepEqual(await response.json(), { status: "ready", database: "ok" }); + assert.deepEqual(await response.json(), { status: "ready", database: "ok", main_key_controls: {protocol: 1, enforcement: false} }); + assert.match(response.headers.get("X-MultiLLM-Worker-Build"), /^[0-9a-f]{64}$/); }); test("worker relies on container.fetch for startup and readiness", async () => { diff --git a/tests/test_codex_easy_transport.py b/tests/test_codex_easy_transport.py index 98ae915a..1aa01e5e 100644 --- a/tests/test_codex_easy_transport.py +++ b/tests/test_codex_easy_transport.py @@ -11,6 +11,12 @@ import requests + +import pytest +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + class CodexEasyNativeRouteTest(unittest.TestCase): def setUp(self): self.original_env = os.environ.copy() diff --git a/tests/test_config_revision_sync.py b/tests/test_config_revision_sync.py index 0971a4fd..6024acf9 100644 --- a/tests/test_config_revision_sync.py +++ b/tests/test_config_revision_sync.py @@ -311,6 +311,8 @@ def __call__(self, payload, *, endpoint): return {"version": 1, "revisions": {name: state["revisions"].get(name, 0) for name in payload["domains"]}} return {"version": 1, "overrides": []} if endpoint == "users": + if payload["operation"] == "capabilities": + return {"version": 1, "key_controls": {"protocol": 1, "enforcement": False}} assert payload["operation"] == "list", "Enabled authentication must use the installed copy" users = [row for row in state["users"] if payload["after"] is None or row["username"] > payload["after"]] return {"version": 1, "users": users[:payload["limit"]]} diff --git a/tests/test_config_revision_sync_worker.mjs b/tests/test_config_revision_sync_worker.mjs index 3810c679..eeb8b549 100644 --- a/tests/test_config_revision_sync_worker.mjs +++ b/tests/test_config_revision_sync_worker.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { handleRevisionedAutoRoutes, commitRevision, revisionSyncSettings, RevisionConsumer } from "../worker/config-revision.mjs"; import { handleAutoRoutesRequest } from "../worker/auto-routes-d1.mjs"; diff --git a/tests/test_config_snapshots_worker.mjs b/tests/test_config_snapshots_worker.mjs index c15c012f..a123d3f5 100644 --- a/tests/test_config_snapshots_worker.mjs +++ b/tests/test_config_snapshots_worker.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { applyMigrations } from "./d1_migrations.mjs"; diff --git a/tests/test_control_audit_worker.mjs b/tests/test_control_audit_worker.mjs index 28305100..3083029f 100644 --- a/tests/test_control_audit_worker.mjs +++ b/tests/test_control_audit_worker.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleControlUsersRequest } from "../worker/control-users-d1.mjs"; import { applyMigrations } from "./d1_migrations.mjs"; @@ -10,7 +10,7 @@ const user = (username, changes = {}) => ({ username, api_key_hash: hash, api_key_prefix: "mllm_abcdefgh", scopes: "chat,models", is_admin: 0, created_at: "2026-09-26T00:00:00+00:00", last_login: null, last_used_at: null, last_used_ip: null, created_by: "admin", rotated_at: null, revoked_at: null, daily_budget_usd: null, monthly_budget_usd: null, - allowed_models: null, allowed_ips: null, expires_at: null, ...changes, + allowed_models: null, allowed_ips: null, expires_at: null, official_only: null, forced_provider_preferences: null, ...changes, }); const event = (action, outcome, actor, target, detail) => ({ operation: "audit_record", action, outcome, actor, target, detail }); const list = (changes = {}) => ({ operation: "audit_list", actor: null, target: null, action: null, @@ -25,7 +25,7 @@ async function database(t) { const env = { INTELLIGENCE_DB: db, ADMIN_USERNAME: "owner" }; const call = async body => { const response = await handleControlUsersRequest(new Request("http://intelligence.internal/v1/users", { - method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, ...body }) }), env); + method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, data_policy_version: 1, ...body }) }), env); return { status: response.status, body: await response.json() }; }; return { db, call }; @@ -132,7 +132,7 @@ test("an unavailable audit table fails closed without revealing storage details" const env = { INTELLIGENCE_DB: { prepare() { throw new Error("private database detail"); } } }; for (const body of [list(), event("sign_out", "succeeded", "owner", null, null)]) { const response = await handleControlUsersRequest(new Request("http://intelligence.internal/v1/users", { - method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, ...body }) }), env); + method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, data_policy_version: 1, ...body }) }), env); assert.equal(response.status, 503); assert.equal((await response.text()).includes("private database detail"), false); } diff --git a/tests/test_control_state_worker.mjs b/tests/test_control_state_worker.mjs index 7a7a9324..029abdad 100644 --- a/tests/test_control_state_worker.mjs +++ b/tests/test_control_state_worker.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { applyMigrations } from "./d1_migrations.mjs"; diff --git a/tests/test_control_users_worker.mjs b/tests/test_control_users_worker.mjs index 23321a44..54a9443b 100644 --- a/tests/test_control_users_worker.mjs +++ b/tests/test_control_users_worker.mjs @@ -1,6 +1,7 @@ +import { dataPolicy } from "../worker/dispatch-policy.mjs"; import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleControlUsersRequest, keyControlsPermit } from "../worker/control-users-d1.mjs"; import { applyMigrations } from "./d1_migrations.mjs"; @@ -12,7 +13,7 @@ const user = (username, changes = {}) => ({ username, api_key_hash: hash, api_key_prefix: "mllm_abcdefgh", scopes: "knowledge:read", is_admin: 0, created_at: "2026-09-24T00:00:00+00:00", last_login: null, last_used_at: null, last_used_ip: null, created_by: "admin", rotated_at: null, revoked_at: null, daily_budget_usd: null, monthly_budget_usd: null, - allowed_models: null, allowed_ips: null, expires_at: null, secret_scan_mode: null, shadow_eval_rate: null, ...changes, + allowed_models: null, allowed_ips: null, expires_at: null, secret_scan_mode: null, shadow_eval_rate: null, official_only: null, forced_provider_preferences: null, ...changes, }); async function database(options) { @@ -21,7 +22,7 @@ async function database(options) { await applyMigrations(db, options); const call = async body => { const response = await handleControlUsersRequest(new Request("http://intelligence.internal/v1/users", { - method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, ...body }) }), { INTELLIGENCE_DB: db }); + method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, data_policy_version: 1, ...body }) }), { INTELLIGENCE_DB: db }); return { status: response.status, body: await response.json() }; }; return { mf, db, call }; @@ -39,7 +40,7 @@ test("accounts persist, page, match active prefixes and record usage in D1", asy assert.deepEqual((await call({ operation: "list", after: "bob", limit: 2 })).body.users.map(item => item.username), ["carol"]); const active = (await call({ operation: "by_prefix", prefix: "mllm_abcdefgh" })).body.users; assert.deepEqual(active.map(item => item.username), ["alice", "carol"], "revoked accounts never authenticate"); - assert.deepEqual(Object.keys(active[0]).sort(), Object.keys(user("x")).sort()); + assert.deepEqual(Object.keys(active[0]).sort(), [...Object.keys(user("x")), "data_policy"].sort()); assert.equal((await call({ operation: "touch", username: "alice", last_used_at: "2026-09-24T02:00:00+00:00", last_used_ip: "203.0.113.9" })).body.updated, true); assert.equal((await call({ operation: "get", username: "alice" })).body.user.last_used_ip, "203.0.113.9"); assert.equal((await call({ operation: "delete", username: "carol" })).body.deleted, true); @@ -76,11 +77,11 @@ test("the users RPC is private, bounded and never reveals storage errors", async assert.equal((await handleControlUsersRequest(request("http://intelligence.internal/v1/users", "x".repeat(9000)), env)).status, 400); assert.equal((await handleControlUsersRequest(request("http://intelligence.internal/v1/users", "{}"), {})).status, 503); const failed = await handleControlUsersRequest(request("http://intelligence.internal/v1/users", - JSON.stringify({ version: 1, operation: "get", username: "alice" })), env); + JSON.stringify({ version: 1, data_policy_version: 1, operation: "get", username: "alice" })), env); assert.equal(failed.status, 503); assert.equal((await failed.text()).includes("private database detail"), false); const routed = await handleIntelligenceOutbound(request("http://intelligence.internal/v1/users", - JSON.stringify({ version: 1, operation: "get", username: "alice" })), env); + JSON.stringify({ version: 1, data_policy_version: 1, operation: "get", username: "alice" })), env); assert.equal(routed.status, 503, "the outbound router reaches the users domain"); }); @@ -98,7 +99,7 @@ test("the Worker refuses durable admin accounts it did not configure and audits const env = { INTELLIGENCE_DB: db, ADMIN_USERNAME: "owner", ADMIN_USERNAMES: "deputy, auditor" }; const send = async body => { const response = await handleControlUsersRequest(new Request("http://intelligence.internal/v1/users", { - method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, ...body }) }), env); + method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, data_policy_version: 1, ...body }) }), env); return { status: response.status, body: await response.json() }; }; try { @@ -141,7 +142,7 @@ test("a transient D1 error on an account read is retried once", async () => { } }; const response = await handleControlUsersRequest(new Request("http://intelligence.internal/v1/users", { method: "POST", headers: { "content-type": "application/json" }, - body: JSON.stringify({ version: 1, operation: "get", username: "alice" }) }), { INTELLIGENCE_DB: flaky }); + body: JSON.stringify({ version: 1, data_policy_version: 1, operation: "get", username: "alice" }) }), { INTELLIGENCE_DB: flaky }); assert.equal(response.status, 200); assert.equal((await response.json()).user.username, "alice"); } finally { await mf.dispose(); } @@ -153,7 +154,7 @@ test("per-key budgets, allowlists, expiry and address ranges persist and are val const limited = user("limited", { daily_budget_usd: 5, monthly_budget_usd: 50.5, allowed_models: "auto:*,free:*,openai:gpt-4.1", allowed_ips: "203.0.113.0/24,2001:db8::/32", expires_at: "2026-12-31T00:00:00+00:00" }); assert.equal((await call({ operation: "upsert", user: limited })).status, 200); - assert.deepEqual((await call({ operation: "get", username: "limited" })).body.user, limited); + assert.deepEqual((await call({ operation: "get", username: "limited" })).body.user, { ...limited, data_policy: dataPolicy(limited) }); assert.equal((await call({ operation: "by_prefix", prefix: "mllm_abcdefgh" })).body.users[0].allowed_models, "auto:*,free:*,openai:gpt-4.1"); for (const changes of [{ daily_budget_usd: -1 }, { daily_budget_usd: "5" }, { monthly_budget_usd: 2e9 }, @@ -166,13 +167,13 @@ test("per-key budgets, allowlists, expiry and address ranges persist and are val } finally { await mf.dispose(); } }); -test("accounts keep working when code is deployed before the key-control migration", async () => { - const { mf, call } = await database({ skip: ["0007_usage_ledger.sql", "0016_usage_buckets.sql"] }); +test("management reports an unverifiable policy before the key-control migration", async () => { + const { mf, call } = await database({ skip: ["0007_usage_ledger.sql", "0016_usage_buckets.sql", "0042_usage_token_details.sql"] }); try { assert.equal((await call({ operation: "upsert", user: user("alice") })).status, 200); - assert.deepEqual((await call({ operation: "get", username: "alice" })).body.user, user("alice")); - assert.deepEqual((await call({ operation: "by_prefix", prefix: "mllm_abcdefgh" })).body.users, [user("alice")]); - assert.deepEqual((await call({ operation: "list", after: null, limit: 5 })).body.users, [user("alice")]); + assert.deepEqual((await call({ operation: "get", username: "alice" })).body.user, { ...user("alice"), data_policy: null }); + assert.deepEqual((await call({ operation: "by_prefix", prefix: "mllm_abcdefgh" })).body.users, [{ ...user("alice"), data_policy: null }]); + assert.deepEqual((await call({ operation: "list", after: null, limit: 5 })).body.users, [{ ...user("alice"), data_policy: null }]); const refused = await call({ operation: "upsert", user: user("alice", { daily_budget_usd: 1 }) }); assert.equal(refused.status, 503, "controls cannot be stored before the migration"); assert.equal((await call({ operation: "get", username: "alice" })).body.user.daily_budget_usd, null); diff --git a/tests/test_core_security.py b/tests/test_core_security.py index de44ae31..821e24af 100644 --- a/tests/test_core_security.py +++ b/tests/test_core_security.py @@ -216,6 +216,11 @@ def test_health_endpoints_are_public_and_not_cached(self): self.assertEqual(response.get_json()["status"], "healthy") self.assertEqual(response.headers["Cache-Control"], "no-store") + def test_health_stamp_comes_from_the_container_artifact(self): + with patch("services.release_status.container_release", return_value={"build_id": "c" * 64}): + response = self.client.get("/healthz") + self.assertEqual(response.headers["X-MultiLLM-Container-Build"], "c" * 64) + def test_private_dashboard_responses_are_not_cached(self): self._set_admin_session() response = self.client.get("/", follow_redirects=False) diff --git a/tests/test_d1_integration.mjs b/tests/test_d1_integration.mjs index 54da1d5f..090c6ef2 100644 --- a/tests/test_d1_integration.mjs +++ b/tests/test_d1_integration.mjs @@ -8,7 +8,7 @@ import { existsSync } from "node:fs"; import test from "node:test"; import { promisify } from "node:util"; import { build } from "esbuild"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { applyMigrations } from "./d1_migrations.mjs"; @@ -72,10 +72,10 @@ test("key controls and the usage ledger persist through the Worker in D1", async test("a D1 without the account migration fails closed for dashboard keys but not for the environment admin", async t => { // 0007, 0011, 0013 and 0030 alter control_users, so they are skipped with the migration that creates the table; - // 0016 alters 0007's usage_events. + // 0016 and 0042 alter usage tables; 0041 alters control_users. const { url } = await privateStore(t, ["0003_control_users.sql", "0004_control_user_audit.sql", "0005_auto_routes.sql", "0007_usage_ledger.sql", "0011_secret_firewall.sql", "0013_shadow_eval.sql", "0016_usage_buckets.sql", "0039_haiku_primary.sql", - "0030_stream_cost_caps.sql"]); + "0030_stream_cost_caps.sql", "0041_key_data_guards.sql", "0042_usage_token_details.sql"]); assert.deepEqual(await drive(url, "unmigrated"), { dashboard_key: 503, admin: "admin", route: ["gguu:gpt-image-2.5-sunburst", "gguu:gpt-image-2.5"], save_route: 503 }); }); diff --git a/tests/test_d1_schema_worker.mjs b/tests/test_d1_schema_worker.mjs index 853cb454..b0ca6a7f 100644 --- a/tests/test_d1_schema_worker.mjs +++ b/tests/test_d1_schema_worker.mjs @@ -1,7 +1,7 @@ import assert from "node:assert/strict"; import { readFile } from "node:fs/promises"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { d1Readiness, REQUIRED_D1_TABLES } from "../worker/d1-schema.mjs"; import { applyMigrations, migrationNames } from "./d1_migrations.mjs"; @@ -30,7 +30,7 @@ test("/ready reports a D1 schema that was never migrated instead of forwarding", let forwarded = 0; const container = { getByName: () => ({ fetch: async () => { forwarded += 1; return Response.json({ status: "healthy" }); } }) }; // 0007 and 0030 alter control_users, so skipping 0003 skips them too; 0016 alters 0007's usage_events. - const partial = await database(t, { skip: ["0003_control_users.sql", "0005_auto_routes.sql", "0007_usage_ledger.sql", "0011_secret_firewall.sql", "0013_shadow_eval.sql", "0039_haiku_primary.sql", "0016_usage_buckets.sql", "0030_stream_cost_caps.sql"] }); + const partial = await database(t, { skip: ["0003_control_users.sql", "0005_auto_routes.sql", "0007_usage_ledger.sql", "0011_secret_firewall.sql", "0013_shadow_eval.sql", "0039_haiku_primary.sql", "0016_usage_buckets.sql", "0030_stream_cost_caps.sql", "0041_key_data_guards.sql", "0042_usage_token_details.sql"] }); assert.deepEqual(await d1Readiness(partial), { ready: false, checked: true, missing: ["control_users", "auto_routes", "usage_events", "usage_daily", "usage_batches", "shadow_eval_samples", "shadow_eval_results", "shadow_eval_config", "shadow_eval_policy_backups"] }); const blocked = await worker.fetch(new Request("https://gateway.example/ready"), { INTELLIGENCE_DB: partial, MULTILLM_PROXY_CONTAINER: container }); assert.equal(blocked.status, 503); diff --git a/tests/test_deferred_tools_worker.mjs b/tests/test_deferred_tools_worker.mjs index add4d63d..5ad74361 100644 --- a/tests/test_deferred_tools_worker.mjs +++ b/tests/test_deferred_tools_worker.mjs @@ -1,3 +1,4 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; import test from "node:test"; import { DatabaseSync } from "node:sqlite"; @@ -6,7 +7,7 @@ import { DeferredTools, deferredEnabled, readGrantSnapshot, handleDeferredMcp } import { catalogueDigest } from "../worker/knowledge/contract-drift.mjs"; import { rankToolDefinitions } from "../worker/knowledge/skills-index.mjs"; import * as deferred from "../worker/knowledge/deferred-tools.mjs"; -import { handleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; +import { handleKnowledgeEdgeRequest as rawhandleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; import { handleControlStateRequest } from "../worker/control-state-d1.mjs"; import { generateIntegrationCredential } from "../scripts/intelligence_operator.mjs"; @@ -314,3 +315,8 @@ test("loading the module makes no random values, which Workers forbid in global else delete globalThis.crypto.randomUUID; } }); + +function handleKnowledgeEdgeRequest(request, env, ...args) { + if (env.INTELLIGENCE_DB) { env.INTELLIGENCE_DB = policyFixtureEnv(env).INTELLIGENCE_DB; } + return rawhandleKnowledgeEdgeRequest(request, env, ...args); +} diff --git a/tests/test_deploy_pipeline.mjs b/tests/test_deploy_pipeline.mjs index 5b092eb3..b0eab27f 100644 --- a/tests/test_deploy_pipeline.mjs +++ b/tests/test_deploy_pipeline.mjs @@ -2,7 +2,7 @@ import assert from "node:assert/strict"; import { readdirSync, readFileSync } from "node:fs"; import test from "node:test"; -import { checkMigrations, destructiveChanges, localMigrations, numberingProblems, sqlStatements } +import { checkMigrations, destructiveChanges, localMigrations, numberingProblems, sqlStatements, remoteOptions, remoteAppliedMigrations } from "../scripts/verify_d1_migrations.mjs"; import { buildIdFromBundle, cloudflareApi, collectStatus, describeVersion, evaluate, readyVerdict, renderMarkdown, renderText, verifyUntilSettled, wranglerFailure } from "../scripts/deploy_status.mjs"; @@ -14,6 +14,21 @@ const FINGERPRINT = "c".repeat(64); const TOKEN = "synthetic-token-not-real"; const ACCOUNT = "0".repeat(32); +test("every remote migration read requires an explicit target before running Wrangler", () => { + for (const args of [[], ["--pending"], ["--remote"], ["--pending", "--remote"]]) { + assert.throws(() => remoteOptions(args, {}), /explicitly/); + } + assert.deepEqual(remoteOptions(["--remote", "synthetic-db"], {}), { database: "synthetic-db", listOnly: false }); + assert.deepEqual(remoteOptions(["--pending"], { D1_DATABASE: "synthetic-stage-db" }), + { database: "synthetic-stage-db", listOnly: true }); + assert.throws(() => remoteOptions(["--unknown"], {})); + let calls = 0; + for (const database of [undefined, null, "", "--remote"]) { + assert.throws(() => remoteAppliedMigrations(database, () => { calls++; }), /explicitly/); + } + assert.equal(calls, 0); +}); + test("the committed migrations pass the static checks", () => { const names = readdirSync(MIGRATIONS); const result = checkMigrations(names, name => readFileSync(new URL(name, MIGRATIONS), "utf8")); @@ -205,3 +220,62 @@ test("verification retries until the checks settle or the wait runs out", async assert.equal(expired.checks.find(check => check.name === "/ready").state, "fail"); assert.ok(calls >= 2 && calls <= 4); }); + + +test("status requires an explicit matching database and reads Worker names from the chosen configs", async () => { + const { deploymentTarget } = await import("../scripts/deployment_target.mjs"); + const binding = { binding: "INTELLIGENCE_DB", database_name: "synthetic-stage-db", database_id: "synthetic-stage-id" }; + const read = path => ({ name: path.includes("private") ? "synthetic-private" : "synthetic-main", d1_databases: [binding] }); + const options = { mainConfig: "stage-main.jsonc", knowledgeConfig: "stage-private.jsonc", read, env: {} }; + assert.throws(() => deploymentTarget(options), /explicitly/); + assert.throws(() => deploymentTarget({ ...options, database: "different-db" }), /match/); + const target = deploymentTarget({ ...options, env: { D1_DATABASE: "synthetic-stage-db" } }); + assert.equal(target.database, "synthetic-stage-db"); + assert.equal(target.workers.main.name, "synthetic-main"); + assert.equal(target.workers.knowledge.name, "synthetic-private"); +}); + +test("a fresh image build identifier changes each release without changing source configuration", async () => { + const { releaseConfig } = await import("../scripts/container_rollout.mjs"); + const config = { containers: [{ image: "./Dockerfile", image_vars: { EXISTING: "value" } }] }; + assert.notEqual(releaseConfig(config).containers[0].image_vars.MULTILLM_RELEASE_IMAGE_ID, + releaseConfig(config).containers[0].image_vars.MULTILLM_RELEASE_IMAGE_ID); + assert.deepEqual(config, { containers: [{ image: "./Dockerfile", image_vars: { EXISTING: "value" } }] }); + assert.equal(releaseConfig(config).containers[0].image_vars.EXISTING, "value"); +}); + +test("Container capture gets version from detail even when the list is stale", async () => { + const { captureApplication } = await import("../scripts/container_rollout.mjs"); + const target = { mainConfig: "stage.jsonc", main: { name: "synthetic", containers: [{ class_name: "Example" }] } }; + const result = captureApplication(target, args => JSON.stringify(args[1] === "list" + ? [{ id: "application", name: "synthetic-example", version: 7, configuration: { image: "stale" } }] + : { id: "application", version: 8, configuration: { image: "current" } })); + assert.deepEqual(result, { id: "application", version: 8, image: "current" }); +}); + +test("Container rollout requires a newer detail version, new image and both actual fingerprints", async () => { + const { verifyRollout } = await import("../scripts/container_rollout.mjs"); + const before = { id: "application", version: 8, image: "old" }; + const expected = "c".repeat(64); + const good = { id: "application", version: 9, configuration: { image: "new" } }; + const ready = { ok: true, worker: expected, container: expected }; + for (const [application, probe] of [ + [{ ...good, version: 8 }, ready], [ { ...good, configuration: { image: "old" } }, ready ], + [good, { ...ready, container: "d".repeat(64) }], [good, { ...ready, worker: "d".repeat(64) }], + [good, { ...ready, ok: false }], [good, {}], + ]) await assert.rejects(verifyRollout({ before, expected, waitSeconds: 0, + readApplication: async () => application, probe: async () => probe }), /unverified/); + assert.equal((await verifyRollout({ before, expected, waitSeconds: 0, + readApplication: async () => good, probe: async () => ready })).version, 9); +}); + +test("Container verification retries temporary swaps but stops at its deadline", async () => { + const { verifyRollout } = await import("../scripts/container_rollout.mjs"); + const before = { id: "application", version: 8, image: "old" }, expected = "c".repeat(64); + let clock = 0, calls = 0; + const result = await verifyRollout({ before, expected, waitSeconds: 60, now: () => clock, + sleep: async ms => { clock += ms; }, + readApplication: async () => { if (!calls++) throw Error("swap"); return { id: before.id, version: 9, configuration: { image: "new" } }; }, + probe: async () => ({ ok: true, worker: expected, container: expected }) }); + assert.equal(result.version, 9); assert.equal(calls, 2); assert.equal(clock, 10000); +}); diff --git a/tests/test_edge_request_metrics_worker.mjs b/tests/test_edge_request_metrics_worker.mjs index 04e713fb..bb4d3034 100644 --- a/tests/test_edge_request_metrics_worker.mjs +++ b/tests/test_edge_request_metrics_worker.mjs @@ -1,3 +1,6 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +const nativeGenerationFetch = (request, env, ctx, authority, ...rest) => rawNativeGenerationFetch( + request, policyFixtureEnv(env), ctx, { authenticated: true, keyId: authority?.principal?.id ?? "admin", ...authority }, ...rest); import assert from "node:assert/strict"; import test from "node:test"; import { readFile } from "node:fs/promises"; @@ -5,7 +8,7 @@ import { spawn } from "node:child_process"; import { recordNativeUsage } from "../worker/usage-ledger-d1.mjs"; import { createGatewayLifecycle } from "../worker/gateway-lifecycle.mjs"; import { createUsageObserver, observeNativeResponse, renderNativeMetrics } from "../worker/request-telemetry.mjs"; -import { nativeGenerationFetch, withForwardedCorrelation, withNativeMetrics } from "../worker/gateway-extensions.mjs"; +import { nativeGenerationFetch as rawNativeGenerationFetch, withForwardedCorrelation, withNativeMetrics } from "../worker/gateway-extensions.mjs"; // Keep the shared loader untouched; register the owned import in this test's copy. const helperUrl = new URL("./helpers/load_cloudflare_worker.mjs", import.meta.url); @@ -13,14 +16,15 @@ const helper = (await readFile(helperUrl, "utf8")) .replaceAll("import.meta.url", JSON.stringify(helperUrl.href)) .replace("const patchedSource = source", `const patchedSource = source.replace('from "./worker/gateway-extensions.mjs";', 'from "${new URL("../worker/gateway-extensions.mjs", import.meta.url)}";')`); const registeredHelperUrl = `data:text/javascript;base64,${Buffer.from(helper).toString("base64")}`; -const worker = (await (await import(registeredHelperUrl)).loadWorkerModule()).default; +const rawWorker = (await (await import(registeredHelperUrl)).loadWorkerModule()).default; +const worker = { ...rawWorker, fetch: (request, env, ctx) => rawWorker.fetch(request, policyFixtureEnv(env), ctx) }; const encoder = new TextEncoder(); const chunks = values => new ReadableStream({ start(c) { for (const value of values) c.enqueue(encoder.encode(value)); c.close(); } }); function fixture() { const rows = [], pending = []; const env = { NATIVE_EDGE_METRICS_ENABLED: "true", ADMIN_API_KEY: "synthetic-admin", ADMIN_USERNAME: "test-owner", CODEX_EASY_API_KEY: "synthetic-upstream", INTELLIGENCE_DB: { - prepare(sql) { return { bind(...values) { return { sql, values }; } }; }, + prepare(sql) { return { bind(...values) { return { sql, values, first: async () => null }; } }; }, async batch(statements) { rows.push(...JSON.parse(statements[1].values[0])); return [{ meta: { changes: 1 } }]; }, } }; return { env, rows, ctx: { waitUntil(p) { pending.push(p); } }, async flush() { await Promise.all(pending); } }; @@ -223,13 +227,13 @@ test("all finalizers run once even if a collaborator fails", async () => { test("native ledger statements preserve nullable usage and batch idempotency in SQLite", async () => { const statements = []; const env = { INTELLIGENCE_DB: { - prepare(sql) { return { bind(...values) { return { sql, values }; } }; }, + prepare(sql) { return { bind(...values) { return { sql, values, first: async () => null }; } }; }, async batch(batch) { statements.push(...batch); return [{ meta: { changes: 1 } }]; }, } }; const event = { provider: "codex-easy", model: "test-model", principal: "edge:trusted", endpoint: "/codex-easy/v1/chat/completions", status: 520, duration_ms: 9, input_tokens: 0, output_tokens: null, cost_usd: null, cost_basis: null, requestId: crypto.randomUUID() }; await recordNativeUsage(env, event); await recordNativeUsage(env, event); - const schema = (await Promise.all(["0003_control_users.sql", "0007_usage_ledger.sql"].map(name => + const schema = (await Promise.all(["0003_control_users.sql", "0007_usage_ledger.sql", "0042_usage_token_details.sql"].map(name => readFile(new URL(`../intelligence-migrations/${name}`, import.meta.url), "utf8")))).join("\n"); const code = `import json,sqlite3,sys payload=json.loads(sys.argv[1]) diff --git a/tests/test_evaluation_noise_floor_worker.mjs b/tests/test_evaluation_noise_floor_worker.mjs index 33604fb2..17ff1e0e 100644 --- a/tests/test_evaluation_noise_floor_worker.mjs +++ b/tests/test_evaluation_noise_floor_worker.mjs @@ -1,7 +1,7 @@ import assert from 'node:assert/strict'; import { readFileSync } from 'node:fs'; import test from 'node:test'; -import { convertV4MiniflareOptions, Miniflare } from 'miniflare'; +import { convertV4MiniflareOptions, Miniflare } from './helpers/local_miniflare.mjs'; import { applyMigrations } from './d1_migrations.mjs'; import { handleShadowEvalRequest, validResult } from '../worker/shadow-eval-d1.mjs'; diff --git a/tests/test_free_routes_worker.mjs b/tests/test_free_routes_worker.mjs index c97622d9..b709747a 100644 --- a/tests/test_free_routes_worker.mjs +++ b/tests/test_free_routes_worker.mjs @@ -4,6 +4,7 @@ import test from "node:test"; import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; import { collectContainerEnv } from "../worker/container-env.mjs"; +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; const worker = (await loadWorkerModule()).default; const origin = "https://example.test"; @@ -15,7 +16,7 @@ test("deployments preserve operator-owned free-tier settings", async () => { }); function environment(fetch) { - return { MULTILLM_PROXY_CONTAINER: { getByName: () => ({ fetch }) } }; + return policyFixtureEnv({ ADMIN_API_KEY: "test-proxy-key", MULTILLM_PROXY_CONTAINER: { getByName: () => ({ fetch }) } }); } test("free pool preflight is answered at the edge without starting a container", async () => { diff --git a/tests/test_gateway_mcp_worker.mjs b/tests/test_gateway_mcp_worker.mjs index 00d04a92..216cee2a 100644 --- a/tests/test_gateway_mcp_worker.mjs +++ b/tests/test_gateway_mcp_worker.mjs @@ -3,7 +3,7 @@ import test from "node:test"; import { isApiRequestPath } from "../worker/api-paths.mjs"; import { isKnowledgeEdgePath } from "../worker/knowledge-edge.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerWithUnrestrictedControls as loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; const worker = (await loadWorkerModule()).default; const ADMIN_KEY = "synthetic-gateway-mcp-admin"; diff --git a/tests/test_generation_deadline.py b/tests/test_generation_deadline.py index ba2fa9c0..45a21642 100644 --- a/tests/test_generation_deadline.py +++ b/tests/test_generation_deadline.py @@ -11,6 +11,11 @@ from services.resilience_service import ResilienceService + +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + class Clock: def __init__(self): self.now = 100.0 diff --git a/tests/test_gpt_image_moderation_worker.mjs b/tests/test_gpt_image_moderation_worker.mjs index f8465bbc..afa2c725 100644 --- a/tests/test_gpt_image_moderation_worker.mjs +++ b/tests/test_gpt_image_moderation_worker.mjs @@ -1,7 +1,14 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +async function loadWorkerModule(...args) { + const module = await rawLoadWorkerModule(...args); + return { ...module, default: { ...module.default, fetch(request, env, ctx) { + return module.default.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} import assert from "node:assert/strict"; import test from "node:test"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerModule as rawLoadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; const worker = (await loadWorkerModule()).default; diff --git a/tests/test_haiku_intelligence_worker.mjs b/tests/test_haiku_intelligence_worker.mjs index 9f1f3a6e..55124139 100644 --- a/tests/test_haiku_intelligence_worker.mjs +++ b/tests/test_haiku_intelligence_worker.mjs @@ -1,7 +1,7 @@ import assert from "node:assert/strict"; import { readFile } from "node:fs/promises"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { applyMigrations } from "./d1_migrations.mjs"; diff --git a/tests/test_handoff_script.py b/tests/test_handoff_script.py index 1e88832a..5a3295b5 100644 --- a/tests/test_handoff_script.py +++ b/tests/test_handoff_script.py @@ -14,6 +14,8 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from unittest.mock import Mock, patch +from services.dispatch_network import control_network + import pytest from scripts import handoff import handoff_contracts as contracts @@ -83,7 +85,8 @@ def do_POST(self): thread = threading.Thread(target=http.serve_forever, daemon=True) thread.start() state["base"] = f"http://127.0.0.1:{http.server_port}" - yield state + with control_network(state["base"]): + yield state http.shutdown() http.server_close() thread.join(2) diff --git a/tests/test_intelligence_auth_worker.mjs b/tests/test_intelligence_auth_worker.mjs index 4cd639dd..7428d837 100644 --- a/tests/test_intelligence_auth_worker.mjs +++ b/tests/test_intelligence_auth_worker.mjs @@ -1,7 +1,7 @@ import {test} from "node:test"; import assert from "node:assert/strict"; import {readFile} from "node:fs/promises"; -import {convertV4MiniflareOptions,Miniflare} from "miniflare"; +import {convertV4MiniflareOptions,Miniflare} from "./helpers/local_miniflare.mjs"; import {handleIntelligenceAuthRequest} from "../worker/intelligence-auth-d1.mjs"; const hash = "scrypt:32768:8:1$salt123456789012$" + "a".repeat(128); @@ -22,6 +22,8 @@ test("D1 credentials survive new callers, rotate atomically and retain revocatio assert.equal((await call({operation:"rotate",principalId:"integration:omni",expectedVersion:1,keyPrefix:"mllm_intelligence_"+"e".repeat(16),keyHash:hash,scopes:["admin"]})).status,400); assert.equal((await call(provision)).status,409); const lookup = async keyPrefix => (await (await call({operation:"lookup",keyPrefix})).json()).principal; + const added = await readFile(new URL("../intelligence-migrations/0043_integration_data_policy.sql", import.meta.url), "utf8"); + for (const statement of added.split(";").filter(x=>x.trim())) await db.prepare(statement).run(); assert.equal((await lookup(prefix)).id,"integration:omni"); const next = "mllm_intelligence_" + "b".repeat(16); const rotate = {operation:"rotate",principalId:"integration:omni",expectedVersion:1,keyPrefix:next,keyHash:hash}; diff --git a/tests/test_intelligence_cancellation.py b/tests/test_intelligence_cancellation.py index 65242cec..4b04179c 100644 --- a/tests/test_intelligence_cancellation.py +++ b/tests/test_intelligence_cancellation.py @@ -9,6 +9,7 @@ from werkzeug.serving import make_server from services.control_plane_backup import capture +from services.dispatch_network import control_network from services.intelligence_cancellation import CallerCancellation from tests.intelligence_fixtures import IntelligenceApiTestCase, upstream @@ -65,7 +66,7 @@ def blocked_provider(**kwargs): rows = [] try: with self.requests(side_effect=blocked_provider) as send: - with socket.create_connection( + with control_network(f"http://127.0.0.1:{server.server_port}"), socket.create_connection( ("127.0.0.1", server.server_port) ) as client: client.sendall(wire) diff --git a/tests/test_intelligence_d1_worker.mjs b/tests/test_intelligence_d1_worker.mjs index 55f71ee0..eebade48 100644 --- a/tests/test_intelligence_d1_worker.mjs +++ b/tests/test_intelligence_d1_worker.mjs @@ -3,7 +3,7 @@ import { mkdtemp, readFile, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceStoreRequest } from "../worker/intelligence-d1.mjs"; const source = await readFile(new URL("../worker/intelligence-d1.mjs", import.meta.url), "utf8"); diff --git a/tests/test_intelligence_operator.mjs b/tests/test_intelligence_operator.mjs index bedeab92..8a3a5927 100644 --- a/tests/test_intelligence_operator.mjs +++ b/tests/test_intelligence_operator.mjs @@ -3,7 +3,7 @@ import { lstat, mkdtemp, readFile, rm, stat, symlink, writeFile } from "node:fs/ import { tmpdir } from "node:os"; import { dirname, join } from "node:path"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceAuthRequest } from "../worker/intelligence-auth-d1.mjs"; import { generateIntegrationCredential, hashIntegrationKey, openRemoteDatabase, runOperator, @@ -37,7 +37,7 @@ async function database(t) { const mf = new Miniflare(convertV4MiniflareOptions({ modules: true, script: "export default {fetch(){return new Response('ok')}}", d1Databases: ["INTELLIGENCE_DB"] })); t.after(() => mf.dispose()); const db = await mf.getD1Database("INTELLIGENCE_DB"); - for (const name of ["0001_intelligence.sql", "0002_integration_principals.sql"]) { + for (const name of ["0001_intelligence.sql", "0002_integration_principals.sql", "0043_integration_data_policy.sql"]) { const migration = await readFile(new URL(`../intelligence-migrations/${name}`, import.meta.url), "utf8"); for (const sql of migration.split(";").map((part) => part.trim()).filter(Boolean)) await db.prepare(sql).run(); } diff --git a/tests/test_intelligence_outbound_worker.mjs b/tests/test_intelligence_outbound_worker.mjs index cf18c734..3069453d 100644 --- a/tests/test_intelligence_outbound_worker.mjs +++ b/tests/test_intelligence_outbound_worker.mjs @@ -1,3 +1,4 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; import test from "node:test"; import { collectContainerEnv } from "../worker/container-env.mjs"; diff --git a/tests/test_key_controls_expand.py b/tests/test_key_controls_expand.py new file mode 100644 index 00000000..2d22a767 --- /dev/null +++ b/tests/test_key_controls_expand.py @@ -0,0 +1,166 @@ +"""Expansion preserves legacy rows and fences every future controlled dispatch.""" +import subprocess +import sys +import types +from unittest.mock import MagicMock, Mock + +import pytest +from flask import Flask, g + +from error_handlers import APIError +from services.dispatch_policy import envelope +from services.provider_dispatch import capture_dispatch, dispatch_as, provider_http +from services.user_store import USER_FIELDS, _row + + +def account(policy): + return {"username": "synthetic", "data_policy": envelope(policy)} + + +@pytest.mark.parametrize("controls", [ + {"official_only": None, "forced_provider_preferences": None}, + {"official_only": 1, "forced_provider_preferences": None}, + {"official_only": 0, "forced_provider_preferences": None}, + {"official_only": False, "forced_provider_preferences": None}, + {"official_only": None, "forced_provider_preferences": '{}'}, + {"official_only": None, "forced_provider_preferences": '{"data_collection":"deny"}'}, +]) +@pytest.mark.parametrize("provider", ["openai", "gemini", "openrouter", "linkapi", "knowledge", "cloudflare"]) +def test_future_policies_refuse_before_every_transport_but_nulls_keep_working(controls, provider, monkeypatch): + current = account(controls) + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda _: current) + send = Mock() + with dispatch_as({"username": "synthetic"}): + if controls["official_only"] is not None or controls["forced_provider_preferences"] is not None: + with pytest.raises(APIError) as error: + provider_http(provider, "POST", "https://provider.example/generate", transport=send) + assert error.value.status_code == 503 + send.assert_not_called() + else: + provider_http(provider, "POST", "https://provider.example/generate", transport=send) + send.assert_called_once() + + +def test_background_callback_rereads_policy_after_rollback(monkeypatch): + current = account({"official_only": None, "forced_provider_preferences": None}) + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda _: current) + send = Mock() + with Flask(__name__).test_request_context(): + g.authenticated_user = current + callback = capture_dispatch(lambda: provider_http("openai", "POST", "https://provider.example/", transport=send)) + current = account({"official_only": 1, "forced_provider_preferences": None}) + with pytest.raises(APIError) as error: + callback() + assert error.value.status_code == 503 + send.assert_not_called() + + +def test_both_row_shapes_remain_readable_without_a_global_account_list_fence(): + row = dict.fromkeys(USER_FIELDS) + row.update(username="synthetic", api_key_hash="synthetic", api_key_prefix="synthetic", scopes="chat", + is_admin=0, created_at="2026-10-01T00:00:00Z") + assert _row(row) == row + controlled = {"official_only": 1, "forced_provider_preferences": None} + assert _row({**row, **controlled, "data_policy": envelope(controlled), "future_metadata": 1})["data_policy"]["official_only"] == 1 + + +@pytest.mark.parametrize("policy", [{}, {"version": 1, "official_only": 1}, + {"version": 1, "official_only": "invalid", "forced_provider_preferences": None}]) +def test_malformed_current_policy_never_dispatches(policy, monkeypatch): + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda _: {"data_policy": policy}) + send = Mock() + with dispatch_as({"username": "synthetic"}), pytest.raises(APIError) as error: + provider_http("openai", "POST", "https://provider.example/", transport=send) + assert error.value.status_code == 503 + send.assert_not_called() + + +@pytest.mark.parametrize("controls", [{"official_only": 1, "forced_provider_preferences": None}, + {"official_only": 0, "forced_provider_preferences": None}, + {"official_only": False, "forced_provider_preferences": None}, + {"official_only": None, "forced_provider_preferences": '{}'}, + {"official_only": None, "forced_provider_preferences": '{"data_collection":"deny"}'}]) +def test_media_and_both_shadow_dispatches_refuse_future_controls(controls, monkeypatch): + from services.media_storage import fetch_public + from services.shadow_eval_data_policy import guarded_dispatch + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda _: account(controls)) + send = Mock() + session = Mock() + session.__enter__ = Mock(return_value=session) + session.__exit__ = Mock(return_value=False) + monkeypatch.setattr("services.media_storage._session", lambda: session) + with dispatch_as({"username": "synthetic"}), pytest.raises(APIError) as error: + fetch_public("https://images.example.com/asset.png") + assert error.value.status_code == 503 + session.request.assert_not_called() + for role in ("candidate", "judge"): + with pytest.raises(APIError) as error: + guarded_dispatch({"key_id": "synthetic"}, send)({"role": role}) + assert error.value.status_code == 503 + send.assert_not_called() + + +OLD = "544bf49ba3b3d03d185ec5b5a71b2b1877dda786" + + +def pinned_module(name, monkeypatch): + module = types.ModuleType("pinned_" + name) + monkeypatch.setitem(sys.modules, module.__name__, module) + source = subprocess.check_output(["git", "show", f"{OLD}:services/{name}.py"], text=True) + exec(compile(source, f"{OLD}:{name}", "exec"), module.__dict__) + return module + + +@pytest.fixture +def null_principal(monkeypatch): + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda _: { + "username": "synthetic", "data_policy": envelope({"official_only": None, "forced_provider_preferences": None})}) + with dispatch_as({"username": "synthetic"}): + yield + + +def test_old_and_expansion_video_api_operations_do_not_follow_redirects(monkeypatch, null_principal): + from services import video_generation + for module in (pinned_module("video_generation", monkeypatch), video_generation): + send = Mock() + monkeypatch.setattr(module.requests, "request", send) + options = {"provider": "xai"} if module is video_generation else {} + module._send("POST", "https://api.x.ai/v1/videos/generations", timeout=10, **options) + assert send.call_args.kwargs["allow_redirects"] is False + + +def test_old_and_expansion_private_ai_posts_do_not_follow_redirects(monkeypatch, null_principal): + from services import cloudflare_ai + for module in (pinned_module("cloudflare_ai", monkeypatch), cloudflare_ai): + session = MagicMock() + session.__enter__.return_value = session + response = MagicMock(status_code=302, headers={}) + response.iter_content.return_value = iter([b"redirect"]) + session.request.return_value = response + session.post.return_value = response + monkeypatch.setattr(module.requests, "Session", lambda: session) + result = module.post("/v1/images/generations", {"model": "synthetic"}, 10) + assert result.status_code == 302 + send = session.request if module is cloudflare_ai else session.post + assert send.call_args.kwargs["allow_redirects"] is False + + +@pytest.mark.parametrize("redirects", [False, True]) +def test_expansion_boundary_keeps_explicit_python_redirect_mode(redirects, null_principal): + send = Mock() + provider_http("xai", "GET", "https://api.x.ai/v1/videos/synthetic", transport=send, allow_redirects=redirects) + assert send.call_args.kwargs["allow_redirects"] is redirects + + +@pytest.mark.parametrize("controls", [ + {"official_only": 0, "forced_provider_preferences": None}, + {"official_only": False, "forced_provider_preferences": None}, + {"official_only": None, "forced_provider_preferences": {}}, + {"official_only": None, "forced_provider_preferences": '{}'}, +]) +def test_non_null_policies_never_enter_shadow_sampling(controls): + from services.shadow_eval_contract import eligible, make_sample + user = {"username": "synthetic", "shadow_eval_rate": 0.2, **controls} + payload = {"model": "auto:synthetic", "messages": [{"role": "user", "content": "synthetic"}]} + assert not eligible(payload, user, "/v1/chat/completions") + assert make_sample(payload, user, "auto:synthetic", {"content": "synthetic"}, "openai:synthetic", 1, {}) is None diff --git a/tests/test_key_controls_expand_worker.mjs b/tests/test_key_controls_expand_worker.mjs new file mode 100644 index 00000000..bcd61fc7 --- /dev/null +++ b/tests/test_key_controls_expand_worker.mjs @@ -0,0 +1,219 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { readFileSync } from "node:fs"; +import { execFile } from "node:child_process"; +import { promisify } from "node:util"; +import { scryptSync } from "node:crypto"; +import { OLD, runtimeBundle, runtimeModule, MemoryStorage } from "./helpers/runtime_versions.mjs"; +import { Miniflare, convertV4MiniflareOptions } from "./helpers/local_miniflare.mjs"; +import { applyMigrations } from "./d1_migrations.mjs"; +import { collectContainerEnv } from "../worker/container-env.mjs"; +import { handleControlUsersRequest } from "../worker/control-users-d1.mjs"; +import { providerFetch, providerOperation } from "../worker/provider-dispatch.mjs"; +import { controlNetwork } from "../worker/dispatch-network.mjs"; +import { policyStore } from "./helpers/dispatch_fixture.mjs"; + +const run = promisify(execFile); +const config = JSON.parse(readFileSync(new URL("../wrangler.jsonc", import.meta.url))); +const states = [ + { official_only: null, forced_provider_preferences: null }, + { official_only: 1, forced_provider_preferences: null }, + { official_only: 0, forced_provider_preferences: null }, + { official_only: null, forced_provider_preferences: '{}' }, + { official_only: null, forced_provider_preferences: '{"data_collection":"deny"}' }, +]; + +test("expansion refuses future controls at every Worker provider boundary", async () => { + for (const policy of states.slice(1)) for (const provider of ["openai", "gemini", "openrouter", "cloudflare", "linkapi", "knowledge"]) { + const env = { INTELLIGENCE_DB: policyStore(policy) }; + const options = { env, authority: { authenticated: true, keyId: "synthetic" }, send() { assert.fail("No upstream request"); } }; + await assert.rejects(providerFetch(provider, "https://provider.example", {}, options), { status: 503 }); + await assert.rejects(providerOperation(provider, () => assert.fail("No SDK dispatch"), options), { status: 503 }); + } +}); + +test("old and expansion account RPC and Container pairs keep NULL-policy routes working", async t => { + const digest = scryptSync("mllm_synthetic_incident_key_long_enough", "synthetic-salt", 64, + { N: 32768, r: 8, p: 1, maxmem: 64 * 1024 * 1024 }).toString("hex"); + for (const revision of [OLD, ""]) { + const admission = revision ? "" : `import {providerFetch} from './worker/provider-dispatch.mjs';`; + const script = await runtimeBundle(`${admission} import {handleIntelligenceOutbound} from './worker/intelligence-outbound.mjs'; + export default {async fetch(request,env){ + ${revision ? "" : `if(new URL(request.url).pathname==='/probe-provider') { + try { return await providerFetch('openrouter','https://openrouter.ai/api/v1/chat/completions',{}, + {env,authority:{legacyBinding:true},send:()=>Response.json({ok:true})}); } + catch { return Response.json({error:'key_controls_require_upgrade'},{status:503}); } + }`} + return handleIntelligenceOutbound(new Request('http://intelligence.internal'+new URL(request.url).pathname,request),env)}}`, revision, "neutral"); + const mf = new Miniflare(convertV4MiniflareOptions({ modules: true, script, compatibilityFlags: ["nodejs_compat"], + d1Databases: ["INTELLIGENCE_DB"], bindings: config.vars })); + try { + const db = await mf.getD1Database("INTELLIGENCE_DB"); + await applyMigrations(db); + await db.prepare(`INSERT INTO control_users(username,api_key_hash,api_key_prefix,scopes,is_admin,created_at) + VALUES('synthetic',?,?,'chat,models',0,'2026-10-01T00:00:00Z')`) + .bind(`scrypt:32768:8:1$synthetic-salt$${digest}`, "mllm_mllm_syn").run(); + const settings = collectContainerEnv({ ...config.vars, ADMIN_API_KEY: "synthetic-bootstrap-long-enough", + FLASK_SECRET_KEY: "synthetic-session-long-enough", JWT_SECRET: "synthetic-jwt-long-enough", OPENROUTER_API_KEY: "synthetic-operator" }); + for (const container of [OLD, ""]) await t.test(`${revision || "expand"} RPC / ${container || "expand"} Container`, async () => { + const { stdout } = await run(process.env.PYTHON || "python3", ["-I", "tests/support/account_envelope_driver.py", + (await mf.ready).href.replace(/\/$/, ""), JSON.stringify(settings), container, + JSON.stringify([["/v1/models", "get"], ["/v1/chat/completions", "post"], ["/login", "post", "synthetic"], ["/login", "post", "admin"], ["/users", "get"]])], + { cwd: process.cwd(), env: { PATH: process.env.PATH }, timeout: 120000 }); + const rows = stdout.split("\n").filter(line => line.startsWith('{"path":')).map(JSON.parse); + assert.deepEqual(rows.map(row => row.status), [200, 200, 302, 302, 200]); + }); + if (!revision) { + // A future controlled row must never poison login or a full management cache reload. + await db.prepare("UPDATE control_users SET official_only=1 WHERE username='synthetic'").run(); + const { stdout: management } = await run(process.env.PYTHON || "python3", ["-I", "tests/support/account_envelope_driver.py", + (await mf.ready).href.replace(/\/$/, ""), JSON.stringify(settings), "", + JSON.stringify([["/login", "post", "synthetic"], ["/login", "post", "admin"], ["/users", "get"], ["/v1/chat/completions", "post", "admin"]])], + { cwd: process.cwd(), env: { PATH: process.env.PATH }, timeout: 120000 }); + const results = management.split("\n").filter(line => line.startsWith('{"path":')).map(JSON.parse); + assert.deepEqual(results.map(row => row.status), [302, 302, 200, 200]); + assert.ok(results.slice(0, 3).every(row => row.dispatches === 0)); + assert.equal(results.at(-1).dispatches, 1); + assert.equal((await db.prepare("SELECT official_only FROM control_users WHERE username='synthetic'").first()).official_only, 1); + // Use the actual strict decoder and auth/cache paths from the pre-expansion Container. + // A legacy management projection must not turn cached restricted identities into dispatch authority. + for (const policy of states) { + await db.prepare("UPDATE control_users SET official_only=?,forced_provider_preferences=? WHERE username='synthetic'") + .bind(policy.official_only, policy.forced_provider_preferences).run(); + const { stdout: oldManagement } = await run(process.env.PYTHON || "python3", ["-I", "tests/support/account_envelope_driver.py", + (await mf.ready).href.replace(/\/$/, ""), JSON.stringify({ ...settings, ACCOUNT_DRIVER_OUTBOUND_ADMISSION: "true" }), OLD, + JSON.stringify([["/login", "post", "synthetic"], ["/login", "post", "admin"], ["/users", "get"], ["/v1/chat/completions", "post"]])], + { cwd: process.cwd(), env: { PATH: process.env.PATH }, timeout: 120000 }); + const oldRows = oldManagement.split("\n").filter(line => line.startsWith('{"path":')).map(JSON.parse); + assert.deepEqual(oldRows.slice(0, 3).map(row => row.status), [policy === states[0] ? 302 : 503, 302, 200]); + assert.equal(oldRows.at(-1).status, policy === states[0] ? 200 : 503); + assert.equal(oldRows.at(-1).dispatches, policy === states[0] ? 1 : 0); + const current = await db.prepare("SELECT official_only,forced_provider_preferences FROM control_users WHERE username='synthetic'").first(); + assert.deepEqual(current, policy); + } + for (const policy of states.slice(1)) { + await db.prepare("UPDATE control_users SET official_only=?,forced_provider_preferences=? WHERE username='synthetic'") + .bind(policy.official_only, policy.forced_provider_preferences).run(); + const { stdout } = await run(process.env.PYTHON || "python3", ["-I", "tests/support/account_envelope_driver.py", + (await mf.ready).href.replace(/\/$/, ""), JSON.stringify(settings)], + { cwd: process.cwd(), env: { PATH: process.env.PATH }, timeout: 120000 }); + const rows = stdout.split("\n").filter(line => line.startsWith('{"path":')).map(JSON.parse); + assert.equal(rows.at(-1).status, 503); + assert.equal(rows.at(-1).dispatches, 0); + } + } + } finally { await mf.dispose(); } + } +}); + +test("legacy upsert cannot reset a policy enabled between the read and atomic write", async t => { + const mf = new Miniflare(convertV4MiniflareOptions({ modules: true, script: "export default {fetch(){return new Response()}}", d1Databases: ["INTELLIGENCE_DB"] })); + t.after(() => mf.dispose()); + const db = await mf.getD1Database("INTELLIGENCE_DB"); + await applyMigrations(db); + await db.prepare(`INSERT INTO control_users(username,api_key_hash,api_key_prefix,scopes,is_admin,created_at) + VALUES('synthetic','synthetic-hash','synthetic-prefix','chat',0,'2026-10-01T00:00:00Z')`).run(); + const call = (body, database = db) => handleControlUsersRequest(new Request("http://intelligence.internal/v1/users", { + method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ version: 1, ...body }), + }), { INTELLIGENCE_DB: database }); + const legacy = (await (await call({ operation: "get", username: "synthetic", unknown_extension: 1 })).json()).user; + assert.equal(Object.keys(legacy).length, 19); + let interleaved = false; + const racing = { prepare: sql => db.prepare(sql), async batch(statements) { + if (!interleaved) { + interleaved = true; + await db.prepare("UPDATE control_users SET official_only=1,forced_provider_preferences=? WHERE username='synthetic'") + .bind('{"data_collection":"deny"}').run(); + } + return db.batch(statements); + } }; + assert.equal((await call({ operation: "upsert", user: legacy }, racing)).status, 200); + assert.deepEqual(await db.prepare("SELECT official_only,forced_provider_preferences FROM control_users").first(), + { official_only: 1, forced_provider_preferences: '{"data_collection":"deny"}' }); +}); + +test("old and expansion Knowledge pairs keep unrestricted requests identical; rollback refuses controlled content", async () => { + const entry = `export {default,KnowledgeCatalogue} from './worker/knowledge/index.mjs'; + export {defaultPolicy} from './worker/knowledge/policy.mjs'; + export {handleKnowledgeEdgeRequest} from './worker/knowledge-edge.mjs';`; + const modules = new Map([[OLD, await runtimeModule(entry, OLD)], ["", await runtimeModule(entry, "")]]); + const original = globalThis.fetch; + let dispatches = 0; + globalThis.fetch = async () => { dispatches++; return Response.json({ results: [] }); }; + try { + for (const main of [OLD, ""]) for (const privateRef of [OLD, ""]) for (const policy of main === OLD && privateRef === OLD ? states.slice(0, 1) : states) { + dispatches = 0; + const edge = modules.get(main), knowledge = modules.get(privateRef); + const storage = new MemoryStorage(), configuration = knowledge.defaultPolicy(); + configuration.enabled = true; + Object.assign(configuration.providers.exa, { enabled: true, limit: 100, background_limit: 50, hard_limit_confirmed: true, retention_allowed: true }); + await storage.put("policy", configuration); + const env = { ADMIN_API_KEY: "synthetic-bootstrap-long-key", ADMIN_USERNAME: "synthetic", AUTH_STORAGE_BACKEND: "d1", + EXA_API_KEY: "synthetic-operator", INTELLIGENCE_DB: policyStore(policy) }; + const catalogue = new knowledge.KnowledgeCatalogue({ storage }, env); + env.KNOWLEDGE_AUTHORITY = { idFromName: () => 1, get: () => ({ fetch: (url, init) => catalogue.fetch(new Request(url, init)) }) }; + env.KNOWLEDGE_SERVICE = { fetch: (url, init) => knowledge.default.fetch(new Request(url, init), env, { waitUntil() {} }) }; + const invoke = () => edge.handleKnowledgeEdgeRequest(new Request("https://gateway.example/v1/knowledge/native/exa_search", { + method: "POST", headers: { authorization: "Bearer synthetic-bootstrap-long-key", "content-type": "application/json" }, + body: JSON.stringify({ query: "synthetic" }), + }), env); + // Only the pinned baseline predates the runtime test guard. + const response = privateRef ? await controlNetwork("https://api.exa.ai/search", invoke) : await invoke(); + assert.equal(response.status, policy === states[0] ? 200 : 503, `${main || "A"}/${privateRef || "A"}`); + assert.equal(dispatches, policy === states[0] ? 1 : 0); + } + } finally { globalThis.fetch = original; } +}); + + +test("legacy background work runs only while every durable policy remains NULL", async () => { + for (const policy of states) { + let dispatches = 0; + const options = { env: { INTELLIGENCE_DB: policyStore(policy) }, authority: { legacyBinding: true }, + send: () => { dispatches++; return new Response("ok"); } }; + const operation = () => providerFetch("knowledge", "https://provider.example", {}, options); + if (policy === states[0]) assert.equal((await operation()).status, 200); + else await assert.rejects(operation(), { status: 503 }); + assert.equal(dispatches, policy === states[0] ? 1 : 0); + } +}); + + +test("registered Container forwarding refuses controlled accounts before sending any body", async () => { + const { loadWorkerModule } = await import("./helpers/load_cloudflare_worker.mjs"); + const worker = (await loadWorkerModule()).default; + for (const policy of states) for (const path of ["/v1/chat/completions", "/v1/images/generations", "/v1/videos/generations", + "/v1/embeddings", "/v1/audio/speech", "/gemini/v1beta/models/example:generateContent", "/v1/models"]) { + let sent = 0; + const env = { ADMIN_API_KEY: "synthetic-bootstrap-long-key", ADMIN_USERNAME: "synthetic", INTELLIGENCE_DB: policyStore(policy), + MULTILLM_PROXY_CONTAINER: { getByName: () => ({ fetch() { sent++; return Response.json({ ok: true }); } }) } }; + const response = await worker.fetch(new Request("https://gateway.example" + path, { + method: path.endsWith("/models") ? "GET" : "POST", headers: { authorization: "Bearer synthetic-bootstrap-long-key", "content-type": "application/json" }, + ...(path.endsWith("/models") ? {} : { body: JSON.stringify({ model: "openrouter:example", input: "synthetic" }) }), + }), env, { waitUntil() {} }); + assert.equal(response.status, policy === states[0] ? 200 : 503, path); + assert.equal(sent, policy === states[0] ? 1 : 0, path); + } +}); + +test("a fenced expansion Container preserves a NULL private call beside other controlled accounts", async () => { + const { legacyBindingControls } = await import("../worker/legacy-binding-policy.mjs"); + for (const policy of states.slice(1)) { + const database = policyStore(policy); + for (const capability of [{ protocol: 1, enforcement: false }, { protocol: 0 }, null]) { + let probes = 0; + const env = { INTELLIGENCE_DB: database, MULTILLM_PROXY_CONTAINER: { getByName: () => ({ fetch(url, init) { + probes++; assert.equal(new URL(url).pathname, "/healthz"); assert.equal(init.method, "GET"); + return Response.json({ key_controls: capability }); + } }) } }; + if (capability?.protocol === 1) assert.deepEqual(await legacyBindingControls(env), states[0]); + else await assert.rejects(legacyBindingControls(env)); + assert.equal(probes, 1); + } + } + let probes = 0; + await assert.rejects(legacyBindingControls({ INTELLIGENCE_DB: policyStore({}), MULTILLM_PROXY_CONTAINER: { + getByName() { probes++; assert.fail("Unreadable policy must not negotiate a fallback"); }, + } })); + assert.equal(probes, 0); +}); diff --git a/tests/test_key_data_preflight.mjs b/tests/test_key_data_preflight.mjs new file mode 100644 index 00000000..74991cba --- /dev/null +++ b/tests/test_key_data_preflight.mjs @@ -0,0 +1,72 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { readFileSync } from "node:fs"; +import { bundleArguments, bundleConfig } from "../scripts/check_worker_bundle.mjs"; +import { checkConfig, checkPrerequisites, checkRelease } from "../scripts/key_data_preflight.mjs"; + +const main = JSON.parse(readFileSync(new URL("../wrangler.jsonc", import.meta.url))); +const knowledge = JSON.parse(readFileSync(new URL("../wrangler.knowledge.jsonc", import.meta.url))); +const commit = "a".repeat(40); +function release(changes = {}) { + const evidence = { + migrations: [{ name: "0043_integration_data_policy.sql" }], + columns: ["official_only", "forced_provider_preferences", "max_stream_cost_microusd"].map(name => ({ name })), + integrationColumns: ["official_only", "forced_provider_preferences"].map(name => ({ name })), + secrets: ["ADMIN_API_KEY", "FLASK_SECRET_KEY", "JWT_SECRET"].map(name => ({ name, type: "secret_text" })), + versions: [{ version_id: "private-version", percentage: 100 }], tag: commit, capability: { protocol: 1, enforcement: false }, ...changes, + }; + return { main, knowledge, commit, sourceCapability: tag => /^[a-f0-9]{40}$/.test(tag || "") ? evidence.capability : null, migrations: ["0043_integration_data_policy.sql"], run: args => { + if (args[0] === "secret") return JSON.stringify(evidence.secrets); + if (args[0] === "deployments") return JSON.stringify({ versions: evidence.versions }); + if (args[0] === "versions") return JSON.stringify({ annotations: { "workers/tag": evidence.tag } }); + if (args[0] === "d1") return JSON.stringify([{ success: true, + results: args.at(-1).includes("PRAGMA") ? (args.at(-1).includes("intelligence_principals") ? evidence.integrationColumns : evidence.columns) : evidence.migrations }]); + throw new Error("Unexpected release probe"); + } }; +} +test("committed release bindings are consistent", () => assert.equal(checkConfig(main, knowledge), "multillm-intelligence")); +test("absent or mismatched configuration refuses deploy admission", () => { + for (const mutation of [config => { delete config.d1_databases; }, config => { delete config.services; }, + config => { config.vars.AUTH_STORAGE_BACKEND = "sql"; }, config => { config.containers = []; }]) { + const config = structuredClone(main); mutation(config); + assert.throws(() => checkConfig(config, knowledge)); + } + assert.throws(() => checkConfig(main, { ...knowledge, d1_databases: [] })); +}); +test("a complete release is admitted without reading secret values", () => assert.doesNotThrow(() => checkRelease(release()))); +test("prerequisites refuse missing main bindings before updating the private Worker", () => { + assert.throws(() => checkPrerequisites(release({ secrets: [] }))); + assert.doesNotThrow(() => checkPrerequisites(release({ tag: "b".repeat(40) }))); +}); +test("missing credentials, schema, migrations or current private version refuse deployment", () => { + for (const changes of [{ secrets: [] }, { columns: [] }, { integrationColumns: [] }, { migrations: [] }, + { migrations: [{ name: "0043_integration_data_policy.sql" }, { name: "0044_unexpected.sql" }] }, { capability: null }, + { versions: [] }, { versions: [{ version_id: "old", percentage: 50 }, { version_id: "new", percentage: 50 }] }]) { + assert.throws(() => checkRelease(release(changes))); + } + assert.throws(() => checkRelease({ ...release(), run: () => { throw new Error("Unavailable evidence"); } })); +}); +test("independent builds cannot upload through the default build without preflight", () => { + assert.match(main.build.command, /^node scripts\/key_data_preflight.mjs --deploy && /); + const ci = readFileSync(new URL("../.github/workflows/ci.yml", import.meta.url), "utf8"); + assert.equal(ci.split("node scripts/check_worker_bundle.mjs --config").length - 1, 2); + for (const mainWorker of [false, true]) assert.ok(bundleArguments("temporary", "output", mainWorker).includes("--dry-run")); + const injected = bundleArguments("temporary", "--dry-run=false", true); + assert.equal(injected.includes("--dry-run=false"), false); + assert.ok(injected.includes("--dry-run")); + const copy = bundleConfig(main); + assert.equal(Object.hasOwn(copy, "build"), false); + assert.deepEqual(copy, Object.fromEntries(Object.entries(main).filter(([name]) => name !== "build"))); + assert.match(main.build.command, /key_data_preflight/); + const manual = readFileSync(new URL("../scripts/deploy.sh", import.meta.url), "utf8"); + const workflow = readFileSync(new URL("../.github/workflows/deploy.yml", import.meta.url), "utf8"); + assert.ok(manual.indexOf("key_data_preflight.mjs --check") < manual.indexOf("migrations apply")); + assert.ok(workflow.indexOf("key_data_preflight.mjs --check") < workflow.indexOf("Apply D1 migrations")); + assert.ok(workflow.indexOf("key_data_preflight.mjs --prepare") < workflow.indexOf("Deploy Knowledge Worker")); + assert.ok(manual.indexOf("key_data_preflight.mjs --prepare") < manual.indexOf('wrangler deploy --config "$KNOWLEDGE_CONFIG"')); +}); + +test("a different expansion commit is compatible with independently deployed enforcement", () => { + assert.doesNotThrow(() => checkRelease(release({ tag: "b".repeat(40) }))); + assert.throws(() => checkRelease(release({ tag: null }))); +}); diff --git a/tests/test_kimi_code_transport.py b/tests/test_kimi_code_transport.py index 523d473e..181d49fc 100644 --- a/tests/test_kimi_code_transport.py +++ b/tests/test_kimi_code_transport.py @@ -9,6 +9,12 @@ import requests + +import pytest +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + class KimiCodeNativeRouteTest(unittest.TestCase): def setUp(self): self.original_env = os.environ.copy() diff --git a/tests/test_knowledge_alexandria.mjs b/tests/test_knowledge_alexandria.mjs index b2f6c9f1..6243cc0e 100644 --- a/tests/test_knowledge_alexandria.mjs +++ b/tests/test_knowledge_alexandria.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { dispatchKnowledge } from "../worker/knowledge/service.mjs"; import { KnowledgeAuthority } from "../worker/knowledge/authority.mjs"; import { defaultPolicy, validatePolicy } from "../worker/knowledge/policy.mjs"; diff --git a/tests/test_knowledge_authority.mjs b/tests/test_knowledge_authority.mjs index f60948cc..d587e4a2 100644 --- a/tests/test_knowledge_authority.mjs +++ b/tests/test_knowledge_authority.mjs @@ -2,9 +2,9 @@ import assert from "node:assert/strict"; import { mkdtemp, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { build } from "esbuild"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { defaultPolicy, validatePolicy } from "../worker/knowledge/policy.mjs"; import { createArtifact, normalizeSourceText, packEvidence, validateChunk, versionEvidence } from "../worker/knowledge/evidence.mjs"; import { parseQuery, publicHost, publicUrl } from "../worker/knowledge/contracts.mjs"; @@ -22,7 +22,7 @@ const bundled = await build({ stdin: { resolveDir: process.cwd(), sourcefile: "k } } export default { fetch(request, env) { return env.CATALOGUE.get(env.CATALOGUE.idFromName('personal')).fetch(request); } }; -` }, bundle: true, format: "esm", platform: "neutral", external: ["cloudflare:workers"], write: false }); +` }, bundle: true, format: "esm", platform: "neutral", external: ["cloudflare:workers", "node:*"], write: false }); function enabledPolicy(limit = 10) { const { revision, ...policy } = defaultPolicy(); @@ -34,7 +34,7 @@ function enabledPolicy(limit = 10) { async function fixture(t) { const directory = await mkdtemp(join(tmpdir(), "knowledge-authority-")); - const create = () => new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityDate: "2026-07-28", + const create = () => new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityFlags: ["nodejs_compat"], compatibilityDate: "2026-07-28", script: bundled.outputFiles[0].text, resourcePersistencePath: directory, durableObjects: { CATALOGUE: { className: "Catalogue", useSQLite: true } } })); let mf = create(); @@ -71,12 +71,13 @@ test("SQLite Durable Object serializes concurrent reservations and survives reop }))); assert.equal(attempts.filter(item => item.status === 200).length, 5); assert.equal(attempts.filter(item => item.status === 429).length, 20); + const acceptedId = `query-${attempts.findIndex(item => item.status === 200)}`; await f.reopen(); const usage = (await f.call("snapshot")).result.usage.find(item => item.provider === "firecrawl"); assert.equal(usage.total, 5); assert.equal(usage.pending, 5); - assert.equal((await f.call("reserve", { operation_id: "query-0", provider: "firecrawl", background: false })).result.replay, true); - assert.equal((await f.call("reserve", { operation_id: "query-0", provider: "exa", background: false })).status, 409); + assert.equal((await f.call("reserve", { operation_id: acceptedId, provider: "firecrawl", background: false })).result.replay, true); + assert.equal((await f.call("reserve", { operation_id: acceptedId, provider: "exa", background: false })).status, 409); }); test("credential failover persists across SQLite restarts and concurrent selection", async t => { diff --git a/tests/test_knowledge_corpus.mjs b/tests/test_knowledge_corpus.mjs index c9c2c5d7..6ccc7926 100644 --- a/tests/test_knowledge_corpus.mjs +++ b/tests/test_knowledge_corpus.mjs @@ -1,8 +1,9 @@ import assert from "node:assert/strict"; import { createHash } from "node:crypto"; -import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import test from "./helpers/policy_test.mjs"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { KnowledgeCorpus } from "../worker/knowledge/corpus.mjs"; +import { providerBinding } from "./helpers/provider_binding.mjs"; const text = "# Flask 3.1\n\nRequest limits support café and 東京.\n"; const digest = value => createHash("sha256").update(value).digest("hex"); @@ -32,7 +33,7 @@ function fixture() { return bytes ? { size: bytes.byteLength, body: new Response(bytes).body } : null; }, }; - const index = { + const index = providerBinding({ async search(request) { searches.push(request); return { chunks: [{ text, score: 0.8, item: { key: artifact().index_key } }] }; @@ -57,7 +58,7 @@ function fixture() { }; }, }, - }; + }); return { corpus: new KnowledgeCorpus({ KNOWLEDGE_SNAPSHOTS: bucket, KNOWLEDGE_INDEX: index }), objects, uploads, searches, lists, bucket, index, @@ -84,7 +85,7 @@ test("snapshots are immutable, hash checked and readable across corpus instances test("the real local R2 binding honors immutable snapshot preconditions", async t => { const runtime = new Miniflare(convertV4MiniflareOptions({ - cf: false, modules: true, script: "export default { fetch() { return new Response('ok'); } };", + cf: false, modules: true, compatibilityFlags: ["nodejs_compat"], script: "export default { fetch() { return new Response('ok'); } };", r2Buckets: { KNOWLEDGE_SNAPSHOTS: "knowledge-snapshot-test" }, })); t.after(() => runtime.dispose()); @@ -138,14 +139,14 @@ test("AI binding search preserves chunk identity and scores without calling inst const f = fixture(); const calls = []; const corpus = new KnowledgeCorpus({ KNOWLEDGE_INDEX: f.index, KNOWLEDGE_SEARCH_INSTANCE: "knowledge-test", - KNOWLEDGE_SEARCH_AI: { autorag(name) { + KNOWLEDGE_SEARCH_AI: providerBinding({ autorag(name) { assert.equal(name, "knowledge-test"); return { async search(query) { calls.push(query); return { data: [{ filename: artifact().index_key, score: 0.8, content: [{ type: "text", text, score: 0.9 }, { type: "text", text: "東京" }] }] }; } }; - } }, + } }), }); assert.deepEqual(await corpus.search({ query: "request limits", product: "flask" }), [ { text, index_key: artifact().index_key, score: 0.9 }, @@ -163,11 +164,11 @@ test("AI binding failures never fall back or retry and malformed responses fail let calls = 0; let response; const corpus = new KnowledgeCorpus({ KNOWLEDGE_INDEX: f.index, KNOWLEDGE_SEARCH_INSTANCE: "knowledge-test", - KNOWLEDGE_SEARCH_AI: { autorag() { return { async search() { + KNOWLEDGE_SEARCH_AI: providerBinding({ autorag() { return { async search() { calls += 1; if (response instanceof Error) throw response; return response; - } }; } }, + } }; } }), }); response = new Error("Invalid ctx.props: missing accountId or accountTag"); await assert.rejects(corpus.search({ query: "limits" }), /missing accountId/); diff --git a/tests/test_knowledge_edge.mjs b/tests/test_knowledge_edge.mjs index 182dbe79..a0e86101 100644 --- a/tests/test_knowledge_edge.mjs +++ b/tests/test_knowledge_edge.mjs @@ -1,9 +1,10 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; import { readFile } from "node:fs/promises"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { generateIntegrationCredential, hashIntegrationKey } from "../scripts/intelligence_operator.mjs"; -import { handleKnowledgeEdgeRequest, isKnowledgeEdgePath } from "../worker/knowledge-edge.mjs"; +import { handleKnowledgeEdgeRequest as rawhandleKnowledgeEdgeRequest, isKnowledgeEdgePath } from "../worker/knowledge-edge.mjs"; import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; const ORIGIN = "https://gateway.example"; @@ -48,7 +49,7 @@ function environment({ result = { status: "ok" }, status = 200, overrides = {} } ADMIN_API_KEY: ADMIN_KEY, ADMIN_USERNAME: " operator ", AUTH_STORAGE_BACKEND: "d1", - INTELLIGENCE_DB: database(records), + INTELLIGENCE_DB: policyFixtureEnv({ INTELLIGENCE_DB: database(records) }).INTELLIGENCE_DB, KNOWLEDGE_SERVICE: { async fetch(url, init) { assert.equal(url, "http://knowledge.internal/v1/dispatch"); @@ -271,7 +272,7 @@ test("dashboard accounts stored in D1 are verified at the edge with their Knowle return { key, row: { username, api_key_hash: await hashIntegrationKey(key, "saltsaltsalt1234"), api_key_prefix: `mllm_${key.slice(0, 8)}`, scopes, is_admin: 0, created_at: "2026-09-24T00:00:00+00:00", last_login: null, last_used_at: null, last_used_ip: null, created_by: "admin", rotated_at: null, revoked_at: null, daily_budget_usd: null, monthly_budget_usd: null, - allowed_models: null, allowed_ips: null, expires_at: null, ...changes } }; + allowed_models: null, allowed_ips: null, expires_at: null, official_only: null, forced_provider_preferences: null, ...changes } }; }; const reader = await account("reader", "knowledge:read"); const chat = await account("chatter", "chat,models"); @@ -526,3 +527,5 @@ for (const path of ["/v1/knowledge/handoffs?limit=1&limit=2", "/v1/knowledge/han assert.equal(response.status, 400); assert.equal(dispatched.length, 0); }); } + +function handleKnowledgeEdgeRequest(request, env, ...args) { return rawhandleKnowledgeEdgeRequest(request, policyFixtureEnv(env), ...args); } diff --git a/tests/test_knowledge_handoffs.mjs b/tests/test_knowledge_handoffs.mjs index 44fe2c4c..3efafd8c 100644 --- a/tests/test_knowledge_handoffs.mjs +++ b/tests/test_knowledge_handoffs.mjs @@ -1,9 +1,9 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { DatabaseSync } from "node:sqlite"; import { readFileSync } from "node:fs"; import { build } from "esbuild"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { mkdtemp, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -169,9 +169,9 @@ test("real SQLite handoff Durable Objects isolate principals and survive restart const bundled = await build({ stdin: { resolveDir: process.cwd(), contents: ` export { KnowledgeHandoffs } from "./worker/knowledge/index.mjs"; export default { fetch(request, env) { return env.HANDOFFS.get(env.HANDOFFS.idFromName(new URL(request.url).searchParams.get("principal"))).fetch(request); } }; - ` }, bundle: true, format: "esm", platform: "neutral", external: ["cloudflare:workers"], write: false }); + ` }, bundle: true, format: "esm", platform: "neutral", external: ["cloudflare:workers", "node:*"], write: false }); const directory = await mkdtemp(join(tmpdir(), "knowledge-handoffs-")); - const create = () => new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityDate: "2026-07-28", + const create = () => new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityFlags: ["nodejs_compat"], compatibilityDate: "2026-07-28", script: bundled.outputFiles[0].text, resourcePersistencePath: directory, durableObjects: { HANDOFFS: { className: "KnowledgeHandoffs", useSQLite: true } } })); let mf = create(); t.after(async () => { await mf.dispose(); await rm(directory, { recursive: true, force: true }); }); diff --git a/tests/test_knowledge_ingestion.mjs b/tests/test_knowledge_ingestion.mjs index 8285bbbe..2734c218 100644 --- a/tests/test_knowledge_ingestion.mjs +++ b/tests/test_knowledge_ingestion.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { KnowledgeAuthority } from "../worker/knowledge/authority.mjs"; import { defaultPolicy } from "../worker/knowledge/policy.mjs"; import { metered } from "../worker/knowledge/operations.mjs"; diff --git a/tests/test_knowledge_keys.mjs b/tests/test_knowledge_keys.mjs index d0d9c482..ad48f955 100644 --- a/tests/test_knowledge_keys.mjs +++ b/tests/test_knowledge_keys.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { configuredKeys, MAX_KEYS } from "../worker/knowledge/providers/keys.mjs"; import { upstreamError } from "../worker/knowledge/providers/errors.mjs"; import { requestJSON } from "../worker/knowledge/providers/transport.mjs"; diff --git a/tests/test_knowledge_ledger.mjs b/tests/test_knowledge_ledger.mjs index a39633bf..391fdb96 100644 --- a/tests/test_knowledge_ledger.mjs +++ b/tests/test_knowledge_ledger.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { KnowledgeAuthority } from "../worker/knowledge/authority.mjs"; import { fixture } from "./knowledge_fixture.mjs"; diff --git a/tests/test_knowledge_memo_latency.mjs b/tests/test_knowledge_memo_latency.mjs index 1e984da1..f944ea3e 100644 --- a/tests/test_knowledge_memo_latency.mjs +++ b/tests/test_knowledge_memo_latency.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { DatabaseSync } from "node:sqlite"; import { setTimeout as delay } from "node:timers/promises"; import { MemoStore } from "../worker/knowledge/memo-store.mjs"; diff --git a/tests/test_knowledge_memos.mjs b/tests/test_knowledge_memos.mjs index f7c83336..afc8c032 100644 --- a/tests/test_knowledge_memos.mjs +++ b/tests/test_knowledge_memos.mjs @@ -1,8 +1,9 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import { providerBinding } from "./helpers/provider_binding.mjs"; +import test from "./helpers/policy_test.mjs"; import { DatabaseSync } from "node:sqlite"; import { build } from "esbuild"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { mkdtemp, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -320,7 +321,7 @@ test("embedding or memo backend errors degrade to ordinary retrieval", async t = test("Workers AI binding receives bge-m3 request and storage uses waitUntil", async t => { const f = await setup(t); let called; - f.env.KNOWLEDGE_SEARCH_AI = { run: async (...args) => { called = args; return { data: [[0.5, 1, -1]] }; } }; + f.env.KNOWLEDGE_SEARCH_AI = providerBinding({ run: async (...args) => { called = args; return { data: [[0.5, 1, -1]] }; } }); const tasks = []; await f.run(f.query, { embed: undefined, waitUntil: promise => tasks.push(promise) }); await Promise.all(tasks); @@ -395,9 +396,9 @@ test("real KnowledgeMemos SQLite Durable Object survives restart with hit counts const bundled = await build({ stdin: { resolveDir: process.cwd(), contents: ` export { KnowledgeMemos } from './worker/knowledge/index.mjs'; export default { fetch(request, env) { return env.MEMOS.get(env.MEMOS.idFromName('personal')).fetch(request); } }; - ` }, bundle: true, format: "esm", platform: "neutral", external: ["cloudflare:workers"], write: false }); + ` }, bundle: true, format: "esm", platform: "neutral", external: ["cloudflare:workers", "node:*"], write: false }); const directory = await mkdtemp(join(tmpdir(), "knowledge-memos-")); - const create = () => new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityDate: "2026-07-28", + const create = () => new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityFlags: ["nodejs_compat"], compatibilityDate: "2026-07-28", script: bundled.outputFiles[0].text, resourcePersistencePath: directory, durableObjects: { MEMOS: { className: "KnowledgeMemos", useSQLite: true } } })); let mf = create(); t.after(async () => { await mf.dispose(); await rm(directory, { recursive: true, force: true }); }); diff --git a/tests/test_knowledge_native.mjs b/tests/test_knowledge_native.mjs index 90f40051..4b244928 100644 --- a/tests/test_knowledge_native.mjs +++ b/tests/test_knowledge_native.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import { readFile } from "node:fs/promises"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { dispatchNative, NATIVE_OPERATIONS, nativeUnits } from "../worker/knowledge/native.mjs"; import { dispatchKnowledge } from "../worker/knowledge/service.mjs"; diff --git a/tests/test_knowledge_product_site_boundaries.mjs b/tests/test_knowledge_product_site_boundaries.mjs index f6838eb5..9258403f 100644 --- a/tests/test_knowledge_product_site_boundaries.mjs +++ b/tests/test_knowledge_product_site_boundaries.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { createArtifact, dropForeignProviderSections, site } from "../worker/knowledge/evidence.mjs"; import { parseProductSites } from "../worker/knowledge/product-sites.mjs"; import { filterProviderSites } from "../worker/knowledge/provider-site-filter.mjs"; diff --git a/tests/test_knowledge_product_site_latency.mjs b/tests/test_knowledge_product_site_latency.mjs index 122b7099..8054433c 100644 --- a/tests/test_knowledge_product_site_latency.mjs +++ b/tests/test_knowledge_product_site_latency.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { retrieveKnowledge } from "../worker/knowledge/retrieval.mjs"; import { fixture, principal, request } from "./knowledge_fixture.mjs"; diff --git a/tests/test_knowledge_product_site_storage.mjs b/tests/test_knowledge_product_site_storage.mjs index c10142b9..f4752714 100644 --- a/tests/test_knowledge_product_site_storage.mjs +++ b/tests/test_knowledge_product_site_storage.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { fixture, seedProductSitePages } from "./knowledge_fixture.mjs"; import { PAGE_LIMIT } from "../worker/knowledge/product-sites.mjs"; import { lruKey, pageKey } from "../worker/knowledge/product-site-identities.mjs"; diff --git a/tests/test_knowledge_product_sites.mjs b/tests/test_knowledge_product_sites.mjs index 7121d406..daf7f1b5 100644 --- a/tests/test_knowledge_product_sites.mjs +++ b/tests/test_knowledge_product_sites.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { createArtifact, dropForeignProviderSections, site } from "../worker/knowledge/evidence.mjs"; import { parseProductSites, SITE_LIMIT, PAGE_LIMIT } from "../worker/knowledge/product-sites.mjs"; import { filterProviderSites } from "../worker/knowledge/provider-site-filter.mjs"; diff --git a/tests/test_knowledge_provider_runtime.mjs b/tests/test_knowledge_provider_runtime.mjs index 709488e8..60e99733 100644 --- a/tests/test_knowledge_provider_runtime.mjs +++ b/tests/test_knowledge_provider_runtime.mjs @@ -1,23 +1,25 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { build } from "esbuild"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; async function runtime(t, outboundService) { const bundle = await build({ stdin: { resolveDir: process.cwd(), contents: ` + import { withDispatchPrincipal } from "./worker/provider-dispatch.mjs"; + import { policyStore } from "./tests/helpers/dispatch_fixture.mjs"; import { requestJSON, jsonPost } from "./worker/knowledge/providers/transport.mjs"; export default { async fetch() { try { - const result = await requestJSON("firecrawl", "scrape", "https://api.firecrawl.dev/v2/scrape", + const result = await withDispatchPrincipal({ INTELLIGENCE_DB: policyStore() }, { authenticated: true, keyId: "fixture-reader" }, () => requestJSON("firecrawl", "scrape", "https://api.firecrawl.dev/v2/scrape", jsonPost({ url: "https://docs.python.org/" }), { env: { FIRECRAWL_API_KEY: "synthetic-key" }, fetchImpl: fetch, invoke: (_provider, _operation, callback) => callback(), - }); + })); return Response.json(result); } catch (error) { return Response.json({ code: error.code }, { status: 502 }); } }}; - ` }, bundle: true, write: false, format: "esm", platform: "neutral" }); - const mf = new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, script: bundle.outputFiles[0].text, + ` }, bundle: true, write: false, format: "esm", platform: "neutral", external: ["node:*"] }); + const mf = new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityFlags: ["nodejs_compat"], script: bundle.outputFiles[0].text, compatibilityDate: "2026-07-28", outboundService })); t.after(() => mf.dispose()); return mf; diff --git a/tests/test_knowledge_providers.mjs b/tests/test_knowledge_providers.mjs index ad25b717..77e868a8 100644 --- a/tests/test_knowledge_providers.mjs +++ b/tests/test_knowledge_providers.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { PROVIDERS, providerStatus, ProviderError, retrieve } from "../worker/knowledge/providers/index.mjs"; import { sourceLinks, sourceOperation, sourceURL } from "../worker/knowledge/providers/source-policy.mjs"; diff --git a/tests/test_knowledge_retrieval.mjs b/tests/test_knowledge_retrieval.mjs index c77f30d8..def88168 100644 --- a/tests/test_knowledge_retrieval.mjs +++ b/tests/test_knowledge_retrieval.mjs @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { setTimeout as delay } from "node:timers/promises"; import { retrieveKnowledge } from "../worker/knowledge/retrieval.mjs"; import { retrieve } from "../worker/knowledge/providers/index.mjs"; diff --git a/tests/test_knowledge_routes.py b/tests/test_knowledge_routes.py index 2f5bcb33..1a3e2b44 100644 --- a/tests/test_knowledge_routes.py +++ b/tests/test_knowledge_routes.py @@ -344,15 +344,18 @@ def __exit__(self, *_args): def test_private_transport_uses_one_fixed_submission(monkeypatch): + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda username: {"username": username, + "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}}) monkeypatch.setenv("KNOWLEDGE_SERVICE_ENABLED", "true") with patch.object(knowledge_client.requests, "Session") as factory: session = factory.return_value.__enter__.return_value - session.post.return_value = FakeResponse({"version": 1, "result": {"status": "ok"}}) + session.request.return_value = FakeResponse({"version": 1, "result": {"status": "ok"}}) result = knowledge_client.dispatch("context", {"username": "reader", "scopes": ["knowledge:read", "chat"]}, {"query": "limits"}) assert result == {"status": "ok"} - assert session.post.call_count == 1 - assert session.post.call_args.args == ("http://knowledge.internal/v1/dispatch",) - arguments = session.post.call_args.kwargs + assert session.request.call_count == 1 + assert session.request.call_args.kwargs["method"] == "POST" + assert session.request.call_args.kwargs["url"] == "http://knowledge.internal/v1/dispatch" + arguments = session.request.call_args.kwargs assert arguments["allow_redirects"] is False assert session.trust_env is False assert "Authorization" not in arguments["headers"] @@ -451,13 +454,15 @@ def test_alexandria_mcp_surfaces_unknown_cost_as_tool_error(app, keys): def test_alexandria_private_operation_is_allowlisted(monkeypatch): + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda username: {"username": username, + "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}}) monkeypatch.setenv("KNOWLEDGE_SERVICE_ENABLED", "true") with patch.object(knowledge_client.requests, "Session") as factory: session = factory.return_value.__enter__.return_value - session.post.return_value = FakeResponse({"version": 1, "result": {"tools": [], "cost": {"credits": 0, "state": "confirmed"}}}) + session.request.return_value = FakeResponse({"version": 1, "result": {"tools": [], "cost": {"credits": 0, "state": "confirmed"}}}) result = knowledge_client.dispatch("alexandria.search", {"username": "reader", "scopes": ["knowledge:read"]}, {"query": "podcasts"}) assert result["cost"]["credits"] == 0 - assert json.loads(session.post.call_args.kwargs["data"])["operation"] == "alexandria.search" + assert json.loads(session.request.call_args.kwargs["data"])["operation"] == "alexandria.search" def test_provider_tools_forward_native_arguments_over_mcp_and_rest(app, keys): diff --git a/tests/test_knowledge_secret_firewall.mjs b/tests/test_knowledge_secret_firewall.mjs index 449ec4ac..f3eb1b58 100644 --- a/tests/test_knowledge_secret_firewall.mjs +++ b/tests/test_knowledge_secret_firewall.mjs @@ -1,7 +1,8 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { readFileSync } from "node:fs"; -import { handleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; +import { handleKnowledgeEdgeRequest as rawhandleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; import { dispatchKnowledge } from "../worker/knowledge/service.mjs"; const TOKEN = "AK" + "IA" + "AB12CD34EF56GH78"; for (const path of ["/v1/knowledge/context", "/mcp"]) for (const mode of ["observe", "redact", "block", "off"]) test(`Knowledge ${path} forces block except ${mode}`, async () => { @@ -32,3 +33,5 @@ test("Knowledge audit binding reuses the gateway database", () => { assert.equal(binding.database_name, "multillm-intelligence"); assert.equal(binding.database_id, /"database_id":\s*"([^"]+)"/.exec(gateway)[1]); }); + +function handleKnowledgeEdgeRequest(request, env, ...args) { return rawhandleKnowledgeEdgeRequest(request, policyFixtureEnv(env), ...args); } diff --git a/tests/test_knowledge_service.mjs b/tests/test_knowledge_service.mjs index 228247b8..9419c70a 100644 --- a/tests/test_knowledge_service.mjs +++ b/tests/test_knowledge_service.mjs @@ -1,11 +1,12 @@ +import { policyFixtureEnv, initializePolicyAccount } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { build } from "esbuild"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { dispatchKnowledge, maintainKnowledge, scheduleSource } from "../worker/knowledge/service.mjs"; import { KnowledgeAuthority, RECONCILE_BATCH } from "../worker/knowledge/authority.mjs"; import { createArtifact } from "../worker/knowledge/evidence.mjs"; -import { handleKnowledgeOutbound } from "../worker/knowledge-outbound.mjs"; +import { handleKnowledgeOutbound as rawhandleKnowledgeOutbound } from "../worker/knowledge-outbound.mjs"; import { collectContainerEnv } from "../worker/container-env.mjs"; import { fixture, manager, principal } from "./knowledge_fixture.mjs"; @@ -300,7 +301,7 @@ test("Container outbound forwarding accepts only its fixed private target and st return Response.json({ version: 1, result: { ready: false } }); } } }; const body = JSON.stringify(envelope("status")); - const options = { method: "POST", headers: { "content-type": "application/json", "authorization": "Bearer synthetic" }, body }; + const options = { method: "POST", headers: { "content-type": "application/json", "authorization": "Bearer synthetic", "X-MultiLLM-Dispatch-Principal": JSON.stringify({ version: 1, username: manager.id }) }, body }; const result = await handleKnowledgeOutbound(new Request("http://knowledge.internal/v1/dispatch", options), env); assert.equal(result.status, 200); assert.equal(requests.length, 1); @@ -322,11 +323,12 @@ test("Container outbound forwarding accepts only its fixed private target and st test("bundled private Worker serves disabled setup through the real Durable Object binding", async t => { const bundle = await build({ entryPoints: ["worker/knowledge/index.mjs"], bundle: true, write: false, - format: "esm", platform: "neutral", external: ["cloudflare:workers"] }); - const mf = new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityDate: "2026-07-28", script: bundle.outputFiles[0].text, - durableObjects: { KNOWLEDGE_AUTHORITY: { className: "KnowledgeCatalogue", useSQLite: true } } })); + format: "esm", platform: "neutral", external: ["cloudflare:workers", "node:*"] }); + const mf = new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityDate: "2026-07-28", compatibilityFlags: ["nodejs_compat"], script: bundle.outputFiles[0].text, + d1Databases: ["INTELLIGENCE_DB"], durableObjects: { KNOWLEDGE_AUTHORITY: { className: "KnowledgeCatalogue", useSQLite: true } } })); t.after(() => mf.dispose()); - const options = { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(envelope("status")) }; + await initializePolicyAccount(mf, manager.id); + const options = { method: "POST", headers: { "content-type": "application/json", "X-MultiLLM-Dispatch-Principal": JSON.stringify({ version: 1, username: manager.id }) }, body: JSON.stringify(envelope("status")) }; const result = await mf.dispatchFetch("http://knowledge.internal/v1/dispatch", options); assert.equal(result.status, 200); assert.equal(result.headers.get("cache-control"), "no-store"); @@ -336,3 +338,5 @@ test("bundled private Worker serves disabled setup through the real Durable Obje assert.equal(payload.result.ready, false); assert.equal((await mf.dispatchFetch("https://public.example/v1/dispatch", options)).status, 404); }); + +function handleKnowledgeOutbound(request, env, ...args) { return rawhandleKnowledgeOutbound(request, policyFixtureEnv(env), ...args); } diff --git a/tests/test_knowledge_skills.mjs b/tests/test_knowledge_skills.mjs index da4046b0..789829ef 100644 --- a/tests/test_knowledge_skills.mjs +++ b/tests/test_knowledge_skills.mjs @@ -1,11 +1,12 @@ +import { policyFixtureEnv, initializePolicyAccount } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { mkdtemp, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { DatabaseSync } from "node:sqlite"; import { build } from "esbuild"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { SkillsStore } from "../worker/knowledge/skills-store.mjs"; import { SkillsIndex, terms } from "../worker/knowledge/skills-index.mjs"; import { parseFind, parseGet, parseSync, validateSkill, SKILLS_LIMIT, SYNC_REQUEST_BYTES } from "../worker/knowledge/skills-validation.mjs"; @@ -13,8 +14,8 @@ import { quantizeEmbedding } from "../worker/knowledge/memo-store.mjs"; import { digest } from "../worker/knowledge/evidence.mjs"; import { errorReply } from "../worker/knowledge/contracts.mjs"; import { dispatchKnowledge } from "../worker/knowledge/service.mjs"; -import { handleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; -import { handleKnowledgeOutbound } from "../worker/knowledge-outbound.mjs"; +import { handleKnowledgeEdgeRequest as rawhandleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; +import { handleKnowledgeOutbound as rawhandleKnowledgeOutbound } from "../worker/knowledge-outbound.mjs"; async function skill(name = "testing", description = "Test regression behavior", root = "agents", extra = []) { const content = `---\nname: ${name}\ndescription: ${description}\n---\n# ${name}\nTest guide`; @@ -149,14 +150,15 @@ test("real Skills Durable Object survives restart with R2 files and helpful coun const bundled = await build({ stdin: { resolveDir: process.cwd(), contents: ` export { KnowledgeSkills } from './worker/knowledge/index.mjs'; export default { fetch(request, env) { return env.KNOWLEDGE_SKILLS.get(env.KNOWLEDGE_SKILLS.idFromName('personal')).fetch(request); } }; - ` }, bundle: true, format: "esm", platform: "neutral", external: ["cloudflare:workers"], write: false }); + ` }, bundle: true, format: "esm", platform: "neutral", external: ["cloudflare:workers", "node:*"], write: false }); const directory = await mkdtemp(join(tmpdir(), "knowledge-skills-")); - const create = () => new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityDate: "2026-07-28", script: bundled.outputFiles[0].text, - resourcePersistencePath: directory, durableObjects: { KNOWLEDGE_SKILLS: { className: "KnowledgeSkills", useSQLite: true } }, r2Buckets: ["KNOWLEDGE_SNAPSHOTS"] })); + const create = () => new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityFlags: ["nodejs_compat"], compatibilityDate: "2026-07-28", script: bundled.outputFiles[0].text, + resourcePersistencePath: directory, d1Databases: ["INTELLIGENCE_DB"], durableObjects: { KNOWLEDGE_SKILLS: { className: "KnowledgeSkills", useSQLite: true } }, r2Buckets: ["KNOWLEDGE_SNAPSHOTS"] })); let mf = create(); + await initializePolicyAccount(mf, "synthetic-principal"); t.after(async () => { await mf.dispose(); await rm(directory, { recursive: true, force: true }); }); const call = async (operation, payload) => { - const response = await mf.dispatchFetch("http://skills.internal/dispatch", { method: "POST", headers: { "content-type": "application/json" }, + const response = await mf.dispatchFetch("http://skills.internal/dispatch", { method: "POST", headers: { "content-type": "application/json", "X-MultiLLM-Dispatch-Principal": '{"version":1,"username":"synthetic-principal"}' }, body: JSON.stringify({ operation, payload, principal: "synthetic-principal" }) }); assert.equal(response.status, 200); return (await response.json()).result; @@ -394,14 +396,18 @@ test("hybrid semantic confidence requires explicit calibration and fast mode nev test("8 MiB sync ingress and outbound boundary accept exact size and reject one extra byte", async () => { assert.equal(SYNC_REQUEST_BYTES, 8 * 1024 * 1024); const env = { ADMIN_API_KEY: "synthetic-skills-boundary", ADMIN_USERNAME: "operator", KNOWLEDGE_SERVICE: { fetch: async () => Response.json({ version: 1, result: { results: [] } }) } }; - const request = (url, body) => new Request(url, { method: "POST", headers: { authorization: "Bearer synthetic-skills-boundary", "content-type": "application/json" }, body }); + const request = (url, body) => new Request(url, { method: "POST", headers: { authorization: "Bearer synthetic-skills-boundary", "content-type": "application/json", "X-MultiLLM-Dispatch-Principal": '{"version":1,"username":"synthetic-principal"}' }, body }); const raw = JSON.stringify({ skills: [] }); for (const [bytes, expected] of [[SYNC_REQUEST_BYTES, 200], [SYNC_REQUEST_BYTES + 1, 413]]) { const response = await handleKnowledgeEdgeRequest(request("https://gateway.example/v1/knowledge/skills", raw.padEnd(bytes)), env); assert.equal(response.status, expected); - const outbound = await handleKnowledgeOutbound(request("http://knowledge.internal/v1/dispatch", JSON.stringify({ operation: "skills.sync" }).padEnd(bytes)), env); + const outbound = await handleKnowledgeOutbound(request("http://knowledge.internal/v1/dispatch", JSON.stringify({ operation: "skills.sync", principal: { id: "synthetic" } }).padEnd(bytes)), env); assert.equal(outbound.status, expected); } const response = await handleKnowledgeEdgeRequest(request("https://gateway.example/v1/knowledge/skills/find", JSON.stringify({ query: "x" }).padEnd(65537)), env); assert.equal(response.status, 413); }); + +function handleKnowledgeEdgeRequest(request, env, ...args) { return rawhandleKnowledgeEdgeRequest(request, policyFixtureEnv(env), ...args); } + +function handleKnowledgeOutbound(request, env, ...args) { return rawhandleKnowledgeOutbound(request, policyFixtureEnv(env), ...args); } diff --git a/tests/test_knowledge_skills_import.mjs b/tests/test_knowledge_skills_import.mjs index b6ec18dc..5da45db3 100644 --- a/tests/test_knowledge_skills_import.mjs +++ b/tests/test_knowledge_skills_import.mjs @@ -1,5 +1,6 @@ +import { policyStore } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { DatabaseSync } from "node:sqlite"; import { digest } from "../worker/knowledge/evidence.mjs"; import { dispatchKnowledge } from "../worker/knowledge/service.mjs"; @@ -23,7 +24,7 @@ function library(t, clock) { catch (error) { db.exec("ROLLBACK"); throw error; } } }; const files = new Map(); - const env = { KNOWLEDGE_SNAPSHOTS: { put: async (key, bytes) => files.set(key, bytes), delete: async () => {}, get: async key => { + const env = { INTELLIGENCE_DB: policyStore(), KNOWLEDGE_SNAPSHOTS: { put: async (key, bytes) => files.set(key, bytes), delete: async () => {}, get: async key => { const bytes = files.get(key); return bytes && { arrayBuffer: async () => bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength) }; } } }; return { env, store: new SkillsStore(storage, env, clock ? { clock } : {}) }; diff --git a/tests/test_knowledge_skills_market.mjs b/tests/test_knowledge_skills_market.mjs index 64821681..71d5eda9 100644 --- a/tests/test_knowledge_skills_market.mjs +++ b/tests/test_knowledge_skills_market.mjs @@ -1,5 +1,7 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +import { policyStore } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { DatabaseSync } from "node:sqlite"; import { errorReply } from "../worker/knowledge/contracts.mjs"; import { digest } from "../worker/knowledge/evidence.mjs"; @@ -7,7 +9,7 @@ import { dispatchKnowledge } from "../worker/knowledge/service.mjs"; import { SkillsStore } from "../worker/knowledge/skills-store.mjs"; import { readFileSync } from "node:fs"; import { DISCOVER_SOURCES, mergeCandidates, parseDiscover, parsePreview } from "../worker/knowledge/skills-market.mjs"; -import { handleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; +import { handleKnowledgeEdgeRequest as rawhandleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; const json = body => Response.json(body); // Shapes follow each marketplace's public search response as probed on 2026-10-07. @@ -60,7 +62,7 @@ function library(t) { catch (error) { db.exec("ROLLBACK"); throw error; } } }; const files = new Map(); - const env = { KNOWLEDGE_SNAPSHOTS: { put: async (key, bytes) => files.set(key, bytes), delete: async () => {}, get: async key => { + const env = { INTELLIGENCE_DB: policyStore(), KNOWLEDGE_SNAPSHOTS: { put: async (key, bytes) => files.set(key, bytes), delete: async () => {}, get: async key => { const bytes = files.get(key); return bytes && { arrayBuffer: async () => bytes.buffer }; } } }; return { env, store: new SkillsStore(storage, env) }; @@ -72,7 +74,7 @@ async function seed(store) { description: "Postgres patterns for this team", files: [{ path: "SKILL.md", content, sha256: await digest(content) }] }] }); } -const submit = (env, operation, payload, options) => dispatchKnowledge(env, { version: 1, operation, +const submit = (env, operation, payload, options) => dispatchKnowledge({ INTELLIGENCE_DB: policyStore(), ...env }, { version: 1, operation, principal: { id: "synthetic-principal", scopes: ["knowledge:read"] }, payload }, options); test("discover merges marketplaces by repository, withholds flagged skills and reports every source", async t => { @@ -241,3 +243,5 @@ test("edge REST and MCP serve discover and preview, and scan the query before an assert.equal(fake.calls.length, calls); assert.equal((await request("/v1/knowledge/skills/discover", { query: "x", sources: ["npm"] })).status, 400); }); + +function handleKnowledgeEdgeRequest(request, env, ...args) { return rawhandleKnowledgeEdgeRequest(request, policyFixtureEnv(env), ...args); } diff --git a/tests/test_knowledge_ui.mjs b/tests/test_knowledge_ui.mjs index 16742416..e097a920 100644 --- a/tests/test_knowledge_ui.mjs +++ b/tests/test_knowledge_ui.mjs @@ -1,5 +1,6 @@ +import { clientNetwork } from "./helpers/network_fixture.mjs"; import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { api, queryPayload, sourcePayload } from "../static/js/knowledge/api.mjs"; import { policyPayload, renderPolicy } from "../static/js/knowledge/policy.mjs"; import { safeUrl, renderEvidence, renderSources, renderJobs, renderReadiness } from "../static/js/knowledge/render.mjs"; @@ -128,7 +129,7 @@ test("admin transport sends CSRF on mutations without client credentials", async dom(); let call; globalThis.fetch = async (url, options) => { call = { url, options }; return Response.json({ id: "s1" }); }; - await api("sources", { method: "POST", body: { product: "flask" } }); + await clientNetwork("/admin/knowledge/sources", () => api("sources", { method: "POST", body: { product: "flask" } })); assert.equal(call.url, "/admin/knowledge/sources"); assert.equal(call.options.headers.get("X-CSRFToken"), "synthetic-csrf"); assert.equal(call.options.credentials, "same-origin"); @@ -140,7 +141,7 @@ test("admin transport reports a conflict and never retries automatically", async dom(); let calls = 0; globalThis.fetch = async () => { calls += 1; return Response.json({ error: { code: "revision_conflict", message: "Reload the current policy." } }, { status: 409 }); }; - await assert.rejects(api("policy", { method: "PUT", body: {} }), /Reload the current policy/); + await assert.rejects(clientNetwork("/admin/knowledge/policy", () => api("policy", { method: "PUT", body: {} })), /Reload the current policy/); assert.equal(calls, 1); }); diff --git a/tests/test_knowledge_workflow_runtime.mjs b/tests/test_knowledge_workflow_runtime.mjs index 3a320d29..47ae1233 100644 --- a/tests/test_knowledge_workflow_runtime.mjs +++ b/tests/test_knowledge_workflow_runtime.mjs @@ -1,17 +1,19 @@ import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { setTimeout } from "node:timers/promises"; import { build } from "esbuild"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; test("native Workflow accepts ingestion step contracts and publishes once", async t => { const bundle = await build({ stdin: { resolveDir: process.cwd(), contents: ` + import { withDispatchPrincipal } from "./worker/provider-dispatch.mjs"; + import { policyStore } from "./tests/helpers/dispatch_fixture.mjs"; import worker, { KnowledgeCatalogue } from "./worker/knowledge/index.mjs"; import { WorkflowEntrypoint } from "cloudflare:workers"; import { KnowledgeCorpus } from "./worker/knowledge/corpus.mjs"; import { runIngestion } from "./worker/knowledge/ingestion.mjs"; export { KnowledgeCatalogue }; - export default worker; + export default { fetch(request, env, ctx) { return worker.fetch(request, { ...env, INTELLIGENCE_DB: policyStore() }, ctx); } }; export class TestIngestion extends WorkflowEntrypoint { async run(event, step) { const snapshots = new KnowledgeCorpus(this.env); @@ -26,12 +28,12 @@ test("native Workflow accepts ingestion step contracts and publishes once", asyn }, async verifySearchable() { return true; }, }; - return runIngestion(this.env, step, event.payload.job_id, { corpus }); + return withDispatchPrincipal({ ...this.env, INTELLIGENCE_DB: policyStore() }, event.payload.dispatch_principal, () => runIngestion(this.env, step, event.payload.job_id, { corpus })); } } - ` }, bundle: true, write: false, format: "esm", platform: "neutral", external: ["cloudflare:workers"] }); + ` }, bundle: true, write: false, format: "esm", platform: "neutral", external: ["cloudflare:workers", "node:*"] }); let acquisitions = 0; - const mf = new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, script: bundle.outputFiles[0].text, + const mf = new Miniflare(convertV4MiniflareOptions({ cf: false, modules: true, compatibilityFlags: ["nodejs_compat"], script: bundle.outputFiles[0].text, compatibilityDate: "2026-07-28", bindings: { FIRECRAWL_API_KEY: "synthetic-key" }, durableObjects: { KNOWLEDGE_AUTHORITY: { className: "KnowledgeCatalogue", useSQLite: true } }, workflows: { KNOWLEDGE_INGESTION: { name: "test-ingestion", className: "TestIngestion" } }, @@ -46,7 +48,7 @@ test("native Workflow accepts ingestion step contracts and publishes once", asyn t.after(() => mf.dispose()); async function dispatch(operation, payload = {}) { const response = await mf.dispatchFetch("http://knowledge.internal/v1/dispatch", { - method: "POST", headers: { "content-type": "application/json" }, + method: "POST", headers: { "content-type": "application/json", "X-MultiLLM-Dispatch-Principal": '{"version":1,"username":"fixture"}' }, body: JSON.stringify({ version: 1, operation, principal: { id: "fixture", scopes: ["knowledge:manage"] }, payload }), }); const body = await response.json(); diff --git a/tests/test_mcp_contract_drift_worker.mjs b/tests/test_mcp_contract_drift_worker.mjs index 372dd5a2..e11f5540 100644 --- a/tests/test_mcp_contract_drift_worker.mjs +++ b/tests/test_mcp_contract_drift_worker.mjs @@ -1,9 +1,16 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +async function loadWorkerModule(...args) { + const module = await rawLoadWorkerModule(...args); + return { ...module, default: { ...module.default, fetch(request, env, ctx) { + return module.default.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} import assert from "node:assert/strict"; import { readFile } from "node:fs/promises"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { canonicalContract, canonicalValue, contractDigest, catalogueDigest, discoveryContract, digestsEnabled, checkContractPin } from "../worker/knowledge/contract-drift.mjs"; import { dispatchNative } from "../worker/knowledge/native.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerModule as rawLoadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; import { generateIntegrationCredential } from "../scripts/intelligence_operator.mjs"; import { fixture, principal } from "./knowledge_fixture.mjs"; diff --git a/tests/test_media_files_worker.mjs b/tests/test_media_files_worker.mjs index 44a45d1e..8fc8bfcf 100644 --- a/tests/test_media_files_worker.mjs +++ b/tests/test_media_files_worker.mjs @@ -1,3 +1,4 @@ +import { policyFixtureEnv, policyStore } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; import { Buffer } from "node:buffer"; import test from "node:test"; @@ -45,7 +46,8 @@ export function memoryBucket() { } const metadata = value => Buffer.from(JSON.stringify(value)).toString("base64url"); -const outbound = (path, init, env) => handleMediaOutbound(new Request(`http://media.internal${path}`, init), env); +const outbound = (path, init, env) => handleMediaOutbound(new Request(`http://media.internal${path}`, { ...init, + headers: { ...init?.headers, "X-MultiLLM-Dispatch-Principal": '{"version":1,"username":"admin"}' } }), policyFixtureEnv(env)); async function signedLink(secret, id, expires) { return `https://gateway.example/v1/media/files/${id}?expires=${expires}&signature=${ @@ -160,7 +162,7 @@ test("the Worker answers signed links itself and forwards unsigned ones to the C const worker = (await loadWorkerModule()).default; const forwarded = []; const container = { getByName: () => ({ fetch: async req => { forwarded.push(new URL(req.url).pathname); return Response.json({ ok: true }); } }) }; - const env = { MEDIA_BUCKET: memoryBucket(), FLASK_SECRET_KEY: "flask-secret", MULTILLM_PROXY_CONTAINER: container }; + const env = { INTELLIGENCE_DB: policyStore(), MEDIA_BUCKET: memoryBucket(), FLASK_SECRET_KEY: "flask-secret", MULTILLM_PROXY_CONTAINER: container }; await env.MEDIA_BUCKET.put(`media/${FILE}`, new Uint8Array([9]), { httpMetadata: { contentType: "image/png" }, customMetadata: { owner: "a", kind: "image" } }); const link = await signedLink("flask-secret", FILE, Math.floor(Date.now() / 1000) + 600); const served = await worker.fetch(new Request(link), env); diff --git a/tests/test_media_generation.py b/tests/test_media_generation.py index ebc0c0af..b627f3e1 100644 --- a/tests/test_media_generation.py +++ b/tests/test_media_generation.py @@ -223,11 +223,10 @@ def provider(method, url, **kwargs): stream.status_code = 200 stream.raw = __import__("io").BytesIO(b"mp4-bytes") stream.headers["Content-Type"] = "video/mp4" - with patch("services.video_generation.requests.request", side_effect=provider), \ - patch("services.video_generation.requests.get", return_value=stream) as download: + with patch("services.video_generation.requests.request", side_effect=lambda method, url, **kw: stream if url == "https://vidgen.x.ai/v.mp4" else provider(method, url, **kw)) as download: content = self.client.get(f"/v1/videos/{job['id']}/content", headers=ADMIN) self.assertEqual((content.status_code, content.data, content.mimetype), (200, b"mp4-bytes", "video/mp4")) - self.assertEqual(download.call_args.args[0], "https://vidgen.x.ai/v.mp4") + self.assertEqual(download.call_args.kwargs["url"], "https://vidgen.x.ai/v.mp4") def test_video_jobs_belong_to_their_creator_and_tampered_ids_are_unknown(self): from services import video_generation diff --git a/tests/test_media_jobs_worker.mjs b/tests/test_media_jobs_worker.mjs index 9c1fec11..640ac978 100644 --- a/tests/test_media_jobs_worker.mjs +++ b/tests/test_media_jobs_worker.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { collectContainerEnv } from "../worker/container-env.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; diff --git a/tests/test_media_outbound_worker.mjs b/tests/test_media_outbound_worker.mjs index cb81dea8..6bf6cce7 100644 --- a/tests/test_media_outbound_worker.mjs +++ b/tests/test_media_outbound_worker.mjs @@ -1,10 +1,14 @@ +import { policyStore } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; -import test from "node:test"; +import test from "./helpers/policy_test.mjs"; import { fetchPublic, handleMediaOutbound, publicHttpsUrl } from "../worker/media-outbound.mjs"; +import { OLD, runtimeModule } from "./helpers/runtime_versions.mjs"; + +const EXPAND = ""; // The contract layer pins its expansion rollback target here. const fetchImage = body => handleMediaOutbound(new Request("http://media.internal/v1/fetch", { - method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(body) }), {}); + method: "POST", headers: { "content-type": "application/json", "X-MultiLLM-Dispatch-Principal": '{"version":1,"username":"fixture-reader"}' }, body: JSON.stringify(body) }), { INTELLIGENCE_DB: policyStore() }); async function withFetch(handler, run) { const original = globalThis.fetch; @@ -58,3 +62,39 @@ test("the private media host refuses other origins and operations", async () => assert.equal((await handleMediaOutbound(new Request("http://media.internal/v1/unknown", { method: "POST" }), {})).status, 404); assert.equal((await fetchImage({ url: "https://10.1.2.3/a.png" })).status, 400); }); + +test("OLD and expansion media transports retain manual admission on every redirect hop", async () => { + for (const revision of [OLD, EXPAND]) { + const media = await runtimeModule("export {fetchPublic} from './worker/media-outbound.mjs';" + + (revision === OLD ? "" : "export {withDispatchPrincipal} from './worker/provider-dispatch.mjs';"), revision); + for (const destination of ["https://cdn.example.com/final", "http://cdn.example.com/final", "https://127.0.0.1/final", "loop"]) { + const seen = []; + const retrieve = () => media.fetchPublic("https://images.example.com/start", { fetcher: async (url, init) => { + seen.push(url); + // Model an actual fetcher's automatic-follow behavior: admission would be bypassed. + if (init.redirect !== "manual") return new Response("automatically followed"); + if (url.endsWith("/final")) return new Response("image"); + return new Response(null, { status: 302, headers: { location: destination === "loop" ? url : destination } }); + } }); + const response = media.withDispatchPrincipal ? await media.withDispatchPrincipal( + { INTELLIGENCE_DB: policyStore() }, { authenticated: true, keyId: "synthetic" }, retrieve) : await retrieve(); + if (destination === "https://cdn.example.com/final") { + assert.equal(await response.text(), "image"); + assert.deepEqual(seen, ["https://images.example.com/start", destination]); + } else { + assert.equal(response, null); + assert.equal(seen.length, destination === "loop" ? 4 : 1); + } + } + } +}); + +test("expansion providerFetch preserves each caller's redirect mode", async () => { + const { providerFetch } = await runtimeModule("export {providerFetch} from './worker/provider-dispatch.mjs';", EXPAND); + for (const redirect of ["manual", "error", "follow"]) { + await providerFetch("media", new Request("https://images.example.com/asset", { redirect }), {}, { + env: { INTELLIGENCE_DB: policyStore() }, authority: { authenticated: true, keyId: "synthetic" }, + send: request => { assert.equal(request.redirect, redirect); return new Response(); }, + }); + } +}); diff --git a/tests/test_native_token_count.py b/tests/test_native_token_count.py index 7574e1a0..354515ed 100644 --- a/tests/test_native_token_count.py +++ b/tests/test_native_token_count.py @@ -18,7 +18,7 @@ def setUp(self): super().setUp() os.environ["NATIVE_TOKEN_COUNT_ENDPOINTS_JSON"] = json.dumps({"opencode": "/v1/messages/count_tokens"}) self.session = Mock() - self.session.post.return_value = self.upstream({"input_tokens": 123}) + self.session.request.return_value = self.upstream({"input_tokens": 123}) self.session_patch = patch("services.proxy_service.ProxyService._get_provider_session", return_value=self.session) self.session_factory = self.session_patch.start() self.addCleanup(self.session_patch.stop) @@ -109,15 +109,15 @@ def test_native_success_preserves_count_fields_and_strips_model_prefix(self): "tools": [{"name": "f", "input_schema": {"type": "object"}}], "tool_choice": {"type": "auto"}, "thinking": {"type": "enabled", "budget_tokens": 100}, "max_tokens": 999, "stream": True, "temperature": 0.5, "foreign_extra": "omit"} - upstream = self.session.post.return_value + upstream = self.session.request.return_value response = self.post("native", body) self.assertEqual(response.status_code, 200) self.assertEqual(response.json, {"input_tokens": 123}) self.assertEqual(response.headers["X-MultiLLM-Token-Count"], "provider") self.assertNotIn("X-MultiLLM-Token-Count-Fallback", response.headers) self.session_factory.assert_called_once_with("opencode", raw_passthrough=True) - self.session.post.assert_called_once() - kwargs = self.session.post.call_args.kwargs + self.session.request.assert_called_once() + kwargs = self.session.request.call_args.kwargs self.assertEqual(kwargs["url"], "https://opencode.ai/v1/messages/count_tokens") self.assertEqual(kwargs["json"], {key: value for key, value in body.items() if key in {"messages", "system", "tools", "tool_choice", "thinking"}} @@ -143,7 +143,7 @@ def test_header_builder_uses_configured_count_path(self): response = self.post("native", {**self.BODY, "model": "linkapi:claude-model"}) self.assertEqual(response.json, {"input_tokens": 123}) self.assertEqual(prepare.call_args.kwargs["upstream_path"], "/v1/messages/count_tokens") - self.assertEqual(self.session.post.call_args.kwargs["headers"]["X-Api-Key"], "linkapi-provider-key") + self.assertEqual(self.session.request.call_args.kwargs["headers"]["X-Api-Key"], "linkapi-provider-key") def test_unsupported_modes_are_explicit(self): for model in ("auto:route", "minimax-m3", "unknown:model", "opencode:"): @@ -196,10 +196,10 @@ def test_exact_origin_path_replaces_base_prefix(self): for base in ("https://trusted.example/v1", "https://trusted.example/api/coding/v1/", "https://trusted.example:8443/v1"): with self.subTest(base=base), patch.dict(os.environ, {"NATIVE_TOKEN_COUNT_ENDPOINTS_JSON": json.dumps({"opencode": "/api/coding/v1/messages/count_tokens"})}): self.app.config["API_BASE_URLS"] = {"opencode": base} - self.session.post.return_value = self.upstream({"input_tokens": 0}) + self.session.request.return_value = self.upstream({"input_tokens": 0}) self.assertEqual(self.post("native").json, {"input_tokens": 0}) origin = base.split("/v1")[0] if "/api" not in base else "https://trusted.example" - self.assertEqual(self.session.post.call_args.kwargs["url"], origin + "/api/coding/v1/messages/count_tokens") + self.assertEqual(self.session.request.call_args.kwargs["url"], origin + "/api/coding/v1/messages/count_tokens") def test_invalid_trusted_origin_never_dispatches(self): for base in ("file:///v1", "https:///v1", "https://user:password@trusted.example/v1", @@ -215,30 +215,30 @@ def test_upstream_errors_never_fall_back(self): for status in (400, 401, 403, 404, 429, 500, 502, 503, 504): with self.subTest(mode=mode, status=status): upstream = self.upstream({"error": {"message": "provider rejected request"}}, status) - self.session.post.return_value = upstream + self.session.request.return_value = upstream self.assert_error(self.post(mode), status) upstream.close.assert_called_once() - self.assertEqual(self.session.post.call_count, 18) + self.assertEqual(self.session.request.call_count, 18) self.generation.assert_not_called() def test_redirects_are_terminal_protocol_errors(self): upstream = self.upstream(status=307, raw=b"redirect") upstream.headers["Location"] = "https://other.example" - self.session.post.return_value = upstream + self.session.request.return_value = upstream self.assert_error(self.post("auto"), 502) - self.session.post.assert_called_once() + self.session.request.assert_called_once() upstream.close.assert_called_once() def test_timeout_and_transport_failures_do_not_retry(self): for mode in ("auto", "native"): for error, status in ((requests.Timeout("private detail"), 504), (requests.ConnectionError("private detail"), 502)): with self.subTest(mode=mode, error=type(error)): - self.session.post.reset_mock() - self.session.post.side_effect = error + self.session.request.reset_mock() + self.session.request.side_effect = error response = self.post(mode) self.assert_error(response, status) self.assertNotIn("private detail", response.get_data(as_text=True)) - self.session.post.assert_called_once() + self.session.request.assert_called_once() def test_invalid_counts_and_json_are_terminal_errors(self): for body in ({}, [], {"input_tokens": True}, {"input_tokens": "123"}, {"input_tokens": -1}, @@ -246,23 +246,23 @@ def test_invalid_counts_and_json_are_terminal_errors(self): {"input_tokens": float("nan")}, {"input_tokens": 4, "extra": float("nan")}): with self.subTest(body=body): upstream = self.upstream(body) - self.session.post.return_value = upstream + self.session.request.return_value = upstream self.assert_error(self.post("auto"), 502) upstream.close.assert_called_once() upstream = self.upstream(raw=b"not json") - self.session.post.return_value = upstream + self.session.request.return_value = upstream self.assert_error(self.post("native"), 502) upstream.close.assert_called_once() def test_response_limit_and_read_timeout_close_once(self): for raw in (b" " * 65537, b'{"input_tokens":1,"padding":"' + b"x" * 65536 + b'"}'): upstream = self.upstream(raw=raw) - self.session.post.return_value = upstream + self.session.request.return_value = upstream self.assert_error(self.post("native"), 502) upstream.close.assert_called_once() upstream = self.upstream({"input_tokens": 1}) upstream.iter_content = Mock(side_effect=requests.Timeout("read timed out")) - self.session.post.return_value = upstream + self.session.request.return_value = upstream self.assert_error(self.post("auto"), 504) upstream.close.assert_called_once() @@ -271,20 +271,20 @@ def test_wrapped_read_timeout_closes_once(self): upstream = self.upstream({"input_tokens": 1}) upstream.iter_content = Mock(side_effect=requests.ConnectionError(ReadTimeoutError(None, None, "read timeout"))) - self.session.post.return_value = upstream + self.session.request.return_value = upstream self.assert_error(self.post("auto"), 504) - self.session.post.assert_called_once() + self.session.request.assert_called_once() upstream.close.assert_called_once() def test_exact_response_bound_and_chunked_json(self): raw = b'{"input_tokens":42}' upstream = self.upstream(raw=raw + b" " * (65536 - len(raw))) - self.session.post.return_value = upstream + self.session.request.return_value = upstream self.assertEqual(self.post("native").json, {"input_tokens": 42}) upstream.close.assert_called_once() upstream = self.upstream() upstream.iter_content = Mock(return_value=iter([b'{"input_', b'tokens":', b"42}"])) - self.session.post.return_value = upstream + self.session.request.return_value = upstream self.assertEqual(self.post("native").json, {"input_tokens": 42}) upstream.close.assert_called_once() @@ -292,7 +292,7 @@ def test_configured_headers_ignore_caller_credentials(self): headers = {**self.HEADERS, "Authorization": "Bearer untrusted-caller-value", "anthropic-beta": "count-test", "X-MultiLLM-Api-Key": "admin-test-key", "x-api-key": "untrusted-caller-value", "Cookie": "stream=true"} self.assertEqual(self.post("native", headers=headers).status_code, 200) - forwarded = self.session.post.call_args.kwargs["headers"] + forwarded = self.session.request.call_args.kwargs["headers"] self.assertEqual(forwarded["Authorization"], "Bearer opencode-provider-key") self.assertEqual(forwarded["Anthropic-Beta"], "count-test") self.assertNotIn("untrusted-caller-value", str(forwarded)) diff --git a/tests/test_opencode_provider.py b/tests/test_opencode_provider.py index e82a3f2c..374f2b22 100644 --- a/tests/test_opencode_provider.py +++ b/tests/test_opencode_provider.py @@ -10,6 +10,12 @@ from flask import Response + +import pytest +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + class OpenCodeProviderRouteTest(unittest.TestCase): def setUp(self): self.original_env = os.environ.copy() @@ -493,6 +499,11 @@ def setUp(self): sys.modules.pop(module_name, None) self.proxy_module = importlib.import_module("services.proxy_service") + from services.dispatch_policy import envelope + store = patch("services.auth_service.AuthService._load_user_by_username", return_value={ + "username": "synthetic-transport", "data_policy": envelope({"official_only": None, "forced_provider_preferences": None})}) + store.start() + self.addCleanup(store.stop) def tearDown(self): os.environ.clear() diff --git a/tests/test_opencode_reasoning_worker.mjs b/tests/test_opencode_reasoning_worker.mjs index 76d07aa8..0d517d2d 100644 --- a/tests/test_opencode_reasoning_worker.mjs +++ b/tests/test_opencode_reasoning_worker.mjs @@ -1,7 +1,7 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerWithUnrestrictedControls as loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; const worker = (await loadWorkerModule()).default; const OPENCODE_CHAT_URL = diff --git a/tests/test_opencode_thinking_worker.mjs b/tests/test_opencode_thinking_worker.mjs index 4e376277..f99b96a8 100644 --- a/tests/test_opencode_thinking_worker.mjs +++ b/tests/test_opencode_thinking_worker.mjs @@ -1,6 +1,13 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +async function loadWorkerModule(...args) { + const module = await rawLoadWorkerModule(...args); + return { ...module, default: { ...module.default, fetch(request, env, ctx) { + return module.default.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} import assert from "node:assert/strict"; import test from "node:test"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerModule as rawLoadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; import { previewProfile } from "../worker/roleplay/profile-receipt.mjs"; import { applyReasoningPolicy } from "../worker/roleplay/reasoning.mjs"; import { withOpencodeGlmReasoning } from "../worker/opencode/reasoning-request.mjs"; diff --git a/tests/test_payment_billing_worker.mjs b/tests/test_payment_billing_worker.mjs index 42ac8a1c..3ea42c65 100644 --- a/tests/test_payment_billing_worker.mjs +++ b/tests/test_payment_billing_worker.mjs @@ -1,10 +1,11 @@ +import { mockFetch } from "./helpers/network_fixture.mjs"; import assert from "node:assert/strict"; import test from "node:test"; import { readFileSync } from "node:fs"; import { DatabaseSync } from "node:sqlite"; import { createHmac } from "node:crypto"; import { fileURLToPath } from "node:url"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; import { handlePaymentStoreRequest, verifyPaymentSignature, paymentSettings } from "../worker/payment-webhooks.mjs"; @@ -170,7 +171,7 @@ test("Worker forwards unauthenticated webhook raw bytes and signature to the Con const worker = (await loadWorkerModule()).default; const raw = new TextEncoder().encode(' {"fixture":"é"}\r\n'); let forwards = 0; - t.mock.method(globalThis, "fetch", () => assert.fail("unexpected external network")); + mockFetch(t, () => assert.fail("unexpected external network")); const env = { MULTILLM_PROXY_CONTAINER: { getByName(name) { assert.equal(name, "primary"); return { async fetch(request) { diff --git a/tests/test_pii_rehydration_worker.mjs b/tests/test_pii_rehydration_worker.mjs index 6b11ec5e..0388613a 100644 --- a/tests/test_pii_rehydration_worker.mjs +++ b/tests/test_pii_rehydration_worker.mjs @@ -115,7 +115,9 @@ test("event fragments, terminal prefixes, bounds and cancellation", async () => }); test("native roleplay transport redacts after firewall and restores", async () => { - const { attemptRoleplayCandidates } = await import("../worker/roleplay/transport.mjs"); + const { attemptRoleplayCandidates: rawAttemptRoleplayCandidates } = await import("../worker/roleplay/transport.mjs"); + const { policyFixtureEnv } = await import("./helpers/dispatch_fixture.mjs"); + const attemptRoleplayCandidates = (...args) => { args[3] = policyFixtureEnv(args[3]); return rawAttemptRoleplayCandidates(...args); }; const previous = globalThis.fetch; const calls = []; globalThis.fetch = async (_url, init) => { @@ -137,7 +139,9 @@ test("native roleplay transport redacts after firewall and restores", async () = }); test("required failure never dispatches or falls back; best effort preserves payload", async () => { - const { attemptRoleplayCandidates } = await import("../worker/roleplay/transport.mjs"); + const { attemptRoleplayCandidates: rawAttemptRoleplayCandidates } = await import("../worker/roleplay/transport.mjs"); + const { policyFixtureEnv } = await import("./helpers/dispatch_fixture.mjs"); + const attemptRoleplayCandidates = (...args) => { args[3] = policyFixtureEnv(args[3]); return rawAttemptRoleplayCandidates(...args); }; const previous = globalThis.fetch; let calls = 0; globalThis.fetch = async (_url, init) => { calls += 1; return Response.json({ content: JSON.parse(init.body).messages.map(message => message.content).join("") }); }; @@ -168,7 +172,9 @@ test("abort during nonstream restoration clears state and cancels pending reader }); test("restoration failure after output cannot authorize another provider", async () => { - const { attemptRoleplayCandidates } = await import("../worker/roleplay/transport.mjs"); + const { attemptRoleplayCandidates: rawAttemptRoleplayCandidates } = await import("../worker/roleplay/transport.mjs"); + const { policyFixtureEnv } = await import("./helpers/dispatch_fixture.mjs"); + const attemptRoleplayCandidates = (...args) => { args[3] = policyFixtureEnv(args[3]); return rawAttemptRoleplayCandidates(...args); }; const previous = globalThis.fetch; let calls = 0; globalThis.fetch = async () => { calls += 1; return Response.json({ content: "__MLPII_" }); }; diff --git a/tests/test_prometheus_worker.mjs b/tests/test_prometheus_worker.mjs index ad8a6492..cc28b63c 100644 --- a/tests/test_prometheus_worker.mjs +++ b/tests/test_prometheus_worker.mjs @@ -3,7 +3,7 @@ import test, { before } from "node:test"; import { collectContainerEnv } from "../worker/container-env.mjs"; import { resetStatusMemo } from "../worker/status-page.mjs"; import { escapeLabel, renderStatusMetrics } from "../worker/prometheus.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerWithUnrestrictedControls as loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; let worker; before(async () => { worker = (await loadWorkerModule()).default; }); diff --git a/tests/test_prompt_cache_affinity.py b/tests/test_prompt_cache_affinity.py index 070bfc10..12861f62 100644 --- a/tests/test_prompt_cache_affinity.py +++ b/tests/test_prompt_cache_affinity.py @@ -319,6 +319,7 @@ def send(**kwargs): monkeypatch.setattr(app_module.ProxyService, "probe_nanogpt_key", lambda *args: pytest.fail("unexpected live probe")) from services.protocol_translation import responses as responses_translation monkeypatch.setattr(responses_translation, "new_id", lambda prefix: f"{prefix}_synthetic") + monkeypatch.setattr(responses_translation, "unix_time", lambda: 1_700_000_000) CredentialPool.record("ce-gpt-plus", "synthetic-first", 429) client = app.test_client() headers = {"Authorization": "Bearer admin-test-key", "Session-Id": "conversation"} diff --git a/tests/test_prompt_cache_cost_worker.mjs b/tests/test_prompt_cache_cost_worker.mjs index 62c84168..ec61228f 100644 --- a/tests/test_prompt_cache_cost_worker.mjs +++ b/tests/test_prompt_cache_cost_worker.mjs @@ -1,7 +1,7 @@ import assert from 'node:assert/strict'; import test from 'node:test'; import { readFile } from 'node:fs/promises'; -import { convertV4MiniflareOptions, Miniflare } from 'miniflare'; +import { convertV4MiniflareOptions, Miniflare } from './helpers/local_miniflare.mjs'; import { BUCKET_FIELDS, serializeUsageBuckets, recordUsageWithBuckets, bucketInsertStatement } from '../worker/usage-buckets-d1.mjs'; import { validRow } from '../worker/usage-ledger-d1.mjs'; import { applyMigrations } from './d1_migrations.mjs'; diff --git a/tests/test_prompt_injection_worker.mjs b/tests/test_prompt_injection_worker.mjs index aeafbd0a..efc29256 100644 --- a/tests/test_prompt_injection_worker.mjs +++ b/tests/test_prompt_injection_worker.mjs @@ -1,3 +1,5 @@ +import { mockFetch } from "./helpers/network_fixture.mjs"; +import { policyStore } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; import test from "node:test"; import { readFile } from "node:fs/promises"; @@ -86,7 +88,7 @@ function fixture(mode, extra = {}) { async put(key, value) { if (typeof key === "object") for (const [k, v] of Object.entries(key)) rows.set(k, v); else rows.set(key, structuredClone(value)); }, async delete(key) { rows.delete(key); }, async setAlarm() {} }; - const env = { ADMIN_API_KEY: "synthetic-admin", OPENCODE_GO_API_KEY: "synthetic-provider", + const env = { INTELLIGENCE_DB: policyStore(), ADMIN_API_KEY: "synthetic-admin", OPENCODE_GO_API_KEY: "synthetic-provider", SECRET_SCAN_DEFAULT: "off", PROMPT_INJECTION_MODE: mode, PROMPT_INJECTION_THRESHOLD: "3", ROLEPLAY_COMPACT_TRIGGER_TOKENS: "0", ...extra }; env.ROLEPLAY_SESSION = { getByName(id) { @@ -101,7 +103,7 @@ for (const [mode, status, calls, action] of [["off", 200, 1, null], ["log", 200, test(`authorized roleplay ${mode} dispatch and content-free headers/events`, async t => { const f = fixture(mode), bodies = [], logs = []; t.mock.method(console, "info", (...args) => logs.push(args)); - t.mock.method(globalThis, "fetch", async (_url, init) => { + mockFetch(t, async (_url, init) => { bodies.push(JSON.parse(init.body)); return Response.json({ choices: [{ message: { role: "assistant", content: "Fixture answer." }, finish_reason: "stop" }] }); }); @@ -126,7 +128,7 @@ for (const [mode, status, calls, action] of [["off", 200, 1, null], ["log", 200, test("block precedes compaction and zero-retention storage", async t => { const f = fixture("block", { CONTENT_RETENTION_ENABLED: "true", CONTENT_RETENTION_POLICY_JSON: '{"default":"zero"}' }); let calls = 0; - t.mock.method(globalThis, "fetch", async () => { calls += 1; throw new Error("Unexpected dispatch"); }); + mockFetch(t, async () => { calls += 1; throw new Error("Unexpected dispatch"); }); const response = await handleRoleplayEdgeRequest(new Request("https://proxy.example/v1/roleplay", { method: "POST", headers: { Authorization: "Bearer synthetic-admin", "Content-Type": "application/json" }, body: JSON.stringify({ model: "roleplay:auto", session_id: "zero-injection", stream: false, @@ -176,7 +178,7 @@ for (const [mode, text, bodyPolicy, action] of [ test(`roleplay ${mode} ignores body policy for ${JSON.stringify(text)} and preserves forwarding`, async t => { const f = fixture(mode), bodies = []; t.mock.method(console, "info", () => {}); - t.mock.method(globalThis, "fetch", async (_url, init) => { + mockFetch(t, async (_url, init) => { bodies.push(JSON.parse(init.body)); return Response.json({ choices: [{ message: { role: "assistant", content: "Fixture answer." }, finish_reason: "stop" }] }); }); @@ -208,7 +210,7 @@ test("limits include multibyte text and cyclic structures", () => { test("log header survives upstream failure without changing replay permission", async t => { const f = fixture("log"); let calls = 0; t.mock.method(console, "info", () => {}); - t.mock.method(globalThis, "fetch", async () => { calls += 1; return Response.json({ error: "fixture" }, { status: 400 }); }); + mockFetch(t, async () => { calls += 1; return Response.json({ error: "fixture" }, { status: 400 }); }); const response = await handleRoleplayEdgeRequest(new Request("https://proxy.example/v1/roleplay", { method: "POST", headers: { Authorization: "Bearer synthetic-admin", "Content-Type": "application/json" }, body: JSON.stringify({ model: "roleplay:auto", session_id: "log-failure", stream: false, @@ -220,7 +222,7 @@ test("log header survives upstream failure without changing replay permission", test("unauthorized roleplay does not scan or dispatch", async t => { const f = fixture("block"), logs = []; let calls = 0; t.mock.method(console, "info", (...args) => logs.push(args)); - t.mock.method(globalThis, "fetch", async () => { calls += 1; throw new Error("Unexpected dispatch"); }); + mockFetch(t, async () => { calls += 1; throw new Error("Unexpected dispatch"); }); const response = await handleRoleplayEdgeRequest(new Request("https://proxy.example/v1/roleplay", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(payload()) }), f.env); assert.equal(response.status, 401); assert.equal(calls, 0); assert.deepEqual(logs, []); @@ -229,7 +231,7 @@ test("unauthorized roleplay does not scan or dispatch", async t => { test("streaming log decision preserves terminal events and one dispatch", async t => { const f = fixture("log"); let calls = 0; t.mock.method(console, "info", () => {}); - t.mock.method(globalThis, "fetch", async () => { + mockFetch(t, async () => { calls += 1; return new Response('data: {"choices":[{"delta":{"content":"Fixture answer."},"finish_reason":null}]}\n\n' + 'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n', diff --git a/tests/test_prompt_templates_worker.mjs b/tests/test_prompt_templates_worker.mjs index 1e1596cc..027732bb 100644 --- a/tests/test_prompt_templates_worker.mjs +++ b/tests/test_prompt_templates_worker.mjs @@ -4,7 +4,7 @@ import { createHash } from "node:crypto"; import { spawnSync } from "node:child_process"; import { existsSync } from "node:fs"; import { fileURLToPath } from "node:url"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { applyMigrations, migrationNames } from "./d1_migrations.mjs"; diff --git a/tests/test_rate_limit_headers.py b/tests/test_rate_limit_headers.py index dca580b7..7e1ec945 100644 --- a/tests/test_rate_limit_headers.py +++ b/tests/test_rate_limit_headers.py @@ -194,7 +194,8 @@ def test_usage_registered_route_default_and_bounded_snapshot(monkeypatch): monkeypatch.setattr(usage, "usage_history", lambda *args, **kwargs: {}) monkeypatch.setattr(usage.BudgetService, "status", lambda _: {}) monkeypatch.setattr(usage.key_controls, "public", lambda _: { - "allowed_models": None, "allowed_ips": None, "expires_at": None}) + "allowed_models": None, "allowed_ips": None, "expires_at": None, + "official_only": None, "forced_provider_preferences": None}) @app.before_request def authenticate(): g.authenticated_user = USER diff --git a/tests/test_raw_passthrough.py b/tests/test_raw_passthrough.py index be5448f9..24a3b36e 100644 --- a/tests/test_raw_passthrough.py +++ b/tests/test_raw_passthrough.py @@ -11,6 +11,12 @@ from werkzeug.datastructures import MultiDict + +import pytest +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + class RawPassthroughTest(unittest.TestCase): def setUp(self): self.original_env = os.environ.copy() diff --git a/tests/test_realtime_worker.mjs b/tests/test_realtime_worker.mjs index 0776c2fa..9a733962 100644 --- a/tests/test_realtime_worker.mjs +++ b/tests/test_realtime_worker.mjs @@ -1,3 +1,4 @@ +import { policyStore } from "./helpers/dispatch_fixture.mjs"; import test from "node:test"; import assert from "node:assert/strict"; import { DatabaseSync } from "node:sqlite"; @@ -46,6 +47,7 @@ async function fixture(t) { return { success: true, results: [], meta: { changes: Number(stmt.run(...values).changes) } }; } }; }, async batch(statements) { sql.exec("BEGIN"); try { const results = statements.map(stmt => stmt.execute()); sql.exec("COMMIT"); return results; } catch (error) { sql.exec("ROLLBACK"); throw error; } } }; + sql.exec("INSERT INTO control_users (username, api_key_hash, api_key_prefix, scopes, is_admin, created_at) VALUES ('admin', 'synthetic-hash', 'synthetic', 'admin,chat,audio', 1, '2026-10-20T00:00:00Z')"); const server = new Socket(), client = new Socket(), upstream = new Socket(), timers = clock(), pending = [], leases = []; const env = { ...enabled, INTELLIGENCE_DB: db, ADMISSION_COORDINATOR: { getByName() { return { async fetch(req) { const value = await req.json(); leases.push(value); @@ -66,7 +68,7 @@ const done = (id = "r1") => JSON.stringify({ type: "response.done", response: { test("disabled, empty and malformed flags preserve registered Container bytes and headers", async () => { for (const flag of [undefined, "", "false", "garbage"]) { - let calls = 0; const env = { REALTIME_ENABLED: flag, MULTILLM_PROXY_CONTAINER: { getByName() { return { + let calls = 0; const env = { INTELLIGENCE_DB: policyStore(), REALTIME_ENABLED: flag, MULTILLM_PROXY_CONTAINER: { getByName() { return { async fetch(req) { calls++; assert.equal(req.headers.get("upgrade"), "websocket"); return new Response("unchanged 雪", { status: 418, headers: { "x-original": "yes" } }); } }; } } }; const response = await worker.fetch(request(), env, {}); @@ -83,8 +85,9 @@ test("invalid mappings disable once without logging values", () => { } } finally { console.warn = old; } assert.equal(messages.length, 1); assert.ok(!messages.join().includes("password")); }); -test("registered enabled errors never wake Container", async () => { - const env = { ...enabled, MULTILLM_PROXY_CONTAINER: { getByName() { throw Error("Container woke"); } } }; +test("registered enabled errors never wake Container", async t => { + const f = await fixture(t); + const env = { ...f.env, MULTILLM_PROXY_CONTAINER: { getByName() { throw Error("Container woke"); } } }; for (const [req, expected] of [[request({ upgrade: "" }), 426], [request({ authorization: "" }), 401], [new Request("https://gateway.example/v1/realtime?model=openai:unmapped", { headers: { upgrade: "websocket", authorization: "Bearer synthetic-gateway" } }), 403]]) { const response = await worker.fetch(req, env, {}); assert.equal(response.status, expected); assert.ok((await response.json()).error.code); @@ -258,3 +261,69 @@ test("missing lease, storage failure and budget exhaustion never dial", async t assert.ok(!JSON.stringify(await response.json()).includes("private storage details")); } }); + +async function controlledAccount(t, selected = "linkapi:realtime-test") { + const f = await fixture(t), key = "synthetic-controlled-key", salt = "syntheticsalt"; + const hash = `scrypt:32768:8:1$${salt}$${scryptSync(key, salt, 64, { N: 32768, r: 8, p: 1, maxmem: 64 * 1024 * 1024 }).toString("hex")}`; + f.env.AUTH_STORAGE_BACKEND = "d1"; + f.env.REALTIME_PROVIDERS_JSON = JSON.stringify({ [selected]: { url: "wss://approved.example/v1/realtime?model=realtime-test" } }); + f.env.REALTIME_PRICING_JSON = JSON.stringify({ [selected]: prices }); + f.env.LINKAPI_KEY = f.env.OPENROUTER_API_KEY = "synthetic-provider"; + f.sql.prepare(`INSERT INTO control_users(username,api_key_hash,api_key_prefix,scopes,is_admin,created_at) + VALUES (?,?,?,?,0,?)`).run("controlled", hash, `mllm_${key.slice(0, 8)}`, "chat,audio", new Date().toISOString()); + const socket = (token = key) => new Request(`https://gateway.example/v1/realtime?model=${selected}`, { + headers: { upgrade: "websocket", authorization: `Bearer ${token}` } }); + const ticket = () => new Request("https://gateway.example/v1/realtime/client_secrets", { + method: "POST", headers: { authorization: `Bearer ${key}`, "content-type": "application/json" }, + body: JSON.stringify({ model: selected }) }); + return { ...f, socket, ticket }; +} + +test("registered Realtime sockets and ticket issuance enforce provider classes", async t => { + const f = await controlledAccount(t); + f.sql.prepare("UPDATE control_users SET official_only=1 WHERE username='controlled'").run(); + const original = globalThis.fetch; let dials = 0; + globalThis.fetch = async () => { dials++; assert.fail("Restricted Realtime must not dial upstream"); }; + try { + for (const request of [f.socket(), f.ticket()]) { + const response = await worker.fetch(request, f.env, f.ctx); + assert.equal(response.status, 503); + assert.equal((await response.json()).error.code, "key_data_controls_unavailable"); + } + assert.equal(dials, 0); + assert.equal(f.sql.prepare("SELECT COUNT(*) AS n FROM realtime_tickets").get().n, 0); + } finally { globalThis.fetch = original; } +}); + +for (const selected of ["openrouter:realtime-test", "openai:realtime-test"]) { + test(`registered Realtime refuses forced preferences for ${selected.split(":")[0]}`, async t => { + const f = await controlledAccount(t, selected); + f.sql.prepare("UPDATE control_users SET forced_provider_preferences=? WHERE username='controlled'") + .run('{"data_collection":"deny"}'); + for (const request of [f.socket(), f.ticket()]) { + const response = await worker.fetch(request, f.env, f.ctx); + assert.equal(response.status, 503); + assert.equal((await response.json()).error.code, "key_data_controls_unavailable"); + } + assert.equal(f.calls(), 0); + }); +} + +test("ticket consumption rechecks changed data controls before consumption or dispatch", async t => { + const f = await controlledAccount(t); + const issued = await worker.fetch(f.ticket(), f.env, f.ctx); + assert.equal(issued.status, 200); + const { value } = await issued.json(); + for (const [official, preferences, code] of [[1, null, "provider_class_not_allowed"], + [0, '{"data_collection":"deny"}', "forced_provider_preferences_not_supported"], + [0, "invalid", "realtime_auth_unavailable"]]) { + f.sql.prepare("UPDATE control_users SET official_only=?,forced_provider_preferences=? WHERE username='controlled'") + .run(official, preferences); + const response = await worker.fetch(f.socket(value), f.env, f.ctx); + assert.equal(response.status, 503); + assert.ok(["key_data_controls_unavailable", "realtime_auth_unavailable"].includes((await response.json()).error.code)); + assert.equal(f.sql.prepare("SELECT consumed_at FROM realtime_tickets").get().consumed_at, null); + assert.equal(f.sql.prepare("SELECT COUNT(*) AS n FROM realtime_sessions").get().n, 0); + assert.equal(f.calls(), 0); + } +}); diff --git a/tests/test_request_cancellation_worker.mjs b/tests/test_request_cancellation_worker.mjs index b110ba35..318c0ae8 100644 --- a/tests/test_request_cancellation_worker.mjs +++ b/tests/test_request_cancellation_worker.mjs @@ -1,8 +1,10 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +const attemptRoleplayCandidates = (...args) => { args[3] = policyFixtureEnv(args[3]); return rawAttemptRoleplayCandidates(...args); }; import assert from "node:assert/strict"; import test from "node:test"; import { UpstreamCancellation } from "../worker/upstream-cancellation.mjs"; -import { attemptRoleplayCandidates, readBoundedBytes } from "../worker/roleplay/transport.mjs"; +import { attemptRoleplayCandidates as rawAttemptRoleplayCandidates, readBoundedBytes } from "../worker/roleplay/transport.mjs"; async function withGlobalFetch(fetchImpl, callback) { const previous = globalThis.fetch; diff --git a/tests/test_retry_advice.py b/tests/test_retry_advice.py index bb233204..bf23b241 100644 --- a/tests/test_retry_advice.py +++ b/tests/test_retry_advice.py @@ -8,6 +8,11 @@ import requests + +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + @pytest.fixture(autouse=True) def settings(monkeypatch): monkeypatch.delenv("UPSTREAM_RETRY_AFTER_ADVICE_ENABLED", raising=False) diff --git a/tests/test_route_health_worker.mjs b/tests/test_route_health_worker.mjs index af1e81ca..c392d76f 100644 --- a/tests/test_route_health_worker.mjs +++ b/tests/test_route_health_worker.mjs @@ -1,7 +1,7 @@ import assert from "node:assert/strict"; import { readFile } from "node:fs/promises"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { fetchIfRunning, runScheduledHealth, scheduleSettings } from "../worker/health-schedule.mjs"; diff --git a/tests/test_scim_provisioning_worker.mjs b/tests/test_scim_provisioning_worker.mjs index 7073e786..28ee8367 100644 --- a/tests/test_scim_provisioning_worker.mjs +++ b/tests/test_scim_provisioning_worker.mjs @@ -25,7 +25,7 @@ function account(n = 1, changes = {}) { function database(t, migrated = true) { const sql = new DatabaseSync(":memory:"); t.after(() => sql.close()); - for (const file of ["0003_control_users.sql", "0007_usage_ledger.sql", "0011_secret_firewall.sql", "0013_shadow_eval.sql", "0017_control_revisions.sql"]) + for (const file of ["0003_control_users.sql", "0007_usage_ledger.sql", "0011_secret_firewall.sql", "0013_shadow_eval.sql", "0017_control_revisions.sql", "0041_key_data_guards.sql"]) sql.exec(schema(file)); sql.exec("CREATE TABLE retained_usage (username TEXT, amount INTEGER); INSERT INTO retained_usage VALUES ('user-1', 7)"); sql.exec("CREATE TABLE fake_teams (org TEXT, id TEXT, name TEXT, active INTEGER, PRIMARY KEY(org,id))"); diff --git a/tests/test_secret_firewall.py b/tests/test_secret_firewall.py index 813eda07..370b1aa3 100644 --- a/tests/test_secret_firewall.py +++ b/tests/test_secret_firewall.py @@ -11,6 +11,11 @@ TOKEN = "AK" + "IA" + "AB12CD34EF56GH78" + +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + @pytest.fixture(autouse=True) def reset_audit_window(monkeypatch): monkeypatch.setattr(firewall, "_audit_events", OrderedDict()) diff --git a/tests/test_secret_firewall_coverage_worker.mjs b/tests/test_secret_firewall_coverage_worker.mjs index 2a0e9144..8ab8da33 100644 --- a/tests/test_secret_firewall_coverage_worker.mjs +++ b/tests/test_secret_firewall_coverage_worker.mjs @@ -14,7 +14,8 @@ test("new edge dispatches require a firewall coverage review", () => { assert.ok(!/await fetch\(upstreamRequest/.test(source)); const roleplay = readFileSync(new URL("../worker/roleplay/transport.mjs", import.meta.url), "utf8"); assert.ok(roleplay.includes("await firewallFetch(new Request")); - assert.ok(roleplay.includes("await containerNamespace.getByName")); + assert.ok(roleplay.includes("transformed => providerFetch(candidate.provider")); + assert.ok(roleplay.includes('send: async permitted => useOpenCodeContainer ? containerNamespace.getByName("primary").fetch(permitted)')); }); test("new edge routes and modules require reviewing the dispatch inventory", () => { diff --git a/tests/test_secret_firewall_worker.mjs b/tests/test_secret_firewall_worker.mjs index fa225ed1..ab1b4e93 100644 --- a/tests/test_secret_firewall_worker.mjs +++ b/tests/test_secret_firewall_worker.mjs @@ -2,13 +2,13 @@ import assert from "node:assert/strict"; import test from "node:test"; import { setImmediate } from "node:timers"; import { firewallFetch, protectPayload } from "../worker/secret-firewall.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerWithUnrestrictedControls as loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; const TOKEN = "AK" + "IA" + "AB12CD34EF56GH78"; const worker = (await loadWorkerModule()).default; for (const mode of ["off", "observe", "redact", "block"]) test(`direct edge provider ${mode}`, async () => { const events = [], bodies = []; const env = { ADMIN_API_KEY: "synthetic-bootstrap", OPENCODE_GO_API_KEY: "synthetic-upstream", OPENCODE_EDGE_FETCH: "true", - SECRET_SCAN_DEFAULT: mode, INTELLIGENCE_DB: { prepare() { return { bind(...args) { events.push(args); return { run: async () => ({}) }; } }; } } }; + SECRET_SCAN_DEFAULT: mode, INTELLIGENCE_DB: { prepare(sql) { if (sql.startsWith("SELECT official_only")) return { bind() { return { first: async () => ({ official_only: null, forced_provider_preferences: null }) }; } }; return { bind(...args) { events.push(args); return { run: async () => ({}) }; } }; } } }; const original = globalThis.fetch; globalThis.fetch = async req => { bodies.push(await req.text()); return Response.json({ choices: [{ message: { content: "ok" } }] }); }; let response; diff --git a/tests/test_shadow_eval.py b/tests/test_shadow_eval.py index c6b93fdd..dfa590d6 100644 --- a/tests/test_shadow_eval.py +++ b/tests/test_shadow_eval.py @@ -23,6 +23,12 @@ from tests.test_intelligence_policy import candidate, policy +@pytest.fixture(autouse=True) +def sampling_authority(monkeypatch): + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", + lambda username: sampling_user() if username == "sampled-user" else None) + + @pytest.fixture def store(tmp_path, monkeypatch): monkeypatch.setenv("INTELLIGENCE_STORAGE_BACKEND", "") @@ -30,12 +36,18 @@ def store(tmp_path, monkeypatch): monkeypatch.setenv("CONTROL_PLANE_DATABASE_URL", "") monkeypatch.setenv("MODEL_REGISTRY_DB_PATH", str(tmp_path / "eval.sqlite3")) monkeypatch.delenv("SHADOW_EVAL_JUDGE_MODEL", raising=False) + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", + lambda username: {"username": username, "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}} if username == "sampled-user" else None) return Store +def sampling_user(**fields): + return {"username": "sampled-user", "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}, **fields} + + def sample(**changes): value = make_sample({"model": "auto:test", "messages": [{"role": "user", "content": "Synthetic greeting"}]}, - {"username": "sampled-user"}, "auto:test", {"content": "Synthetic answer"}, "openai:production", 100, + sampling_user(), "auto:test", {"content": "Synthetic answer"}, "openai:production", 100, {"prompt_tokens": 4, "completion_tokens": 2}) return {**value, **changes} @@ -50,7 +62,7 @@ def configured(store, **changes): @pytest.mark.parametrize("rate,draw,accepted", [(None, 0, False), (0, 0, False), (0.2, 0.199, True), (0.2, 0.2, False), (0.21, 0, False), (True, 0, False)]) def test_sampling_probability(rate, draw, accepted): payload = {"model": "auto:test", "messages": []} - assert eligible(payload, {"shadow_eval_rate": rate}, "/v1/chat/completions", random_value=draw) is accepted + assert eligible(payload, sampling_user(shadow_eval_rate=rate), "/v1/chat/completions", random_value=draw) is accepted @pytest.mark.parametrize("model,path,expected", [("auto:test", "/v1/chat/completions", True), ("cascade:test", "/v1/chat/completions", True), @@ -60,7 +72,7 @@ def test_sampling_probability(rate, draw, accepted): ("openai:concrete", "/v1/messages", False), ("openai:concrete", "/optimize/v1/chat/completions", False), ("auto:test", "/v1/knowledge", False), ("openai:concrete", "/intelligence/v1/chat/completions", True)]) def test_eligibility(model, path, expected): - assert eligible({"model": model, "messages": []}, {"shadow_eval_rate": 0.1}, path, random_value=0) is expected + assert eligible({"model": model, "messages": []}, sampling_user(shadow_eval_rate=0.1), path, random_value=0) is expected @pytest.mark.parametrize("payload,task", [({"tools": [{}]}, "coding"), ({"messages": [{"content": "```python\nx=1"}]}, "coding"), @@ -73,7 +85,7 @@ def test_task_classification(payload, task): def test_secret_skip_size_caps_and_heuristic_preservation(): payload = {"model": "auto:test", "messages": [{"role": "user", "content": "Hello"}]} - args = ({"username": "sampled-user"}, "auto:test", {"content": "Synthetic answer"}, "openai:production", 1, {}) + args = (sampling_user(), "auto:test", {"content": "Synthetic answer"}, "openai:production", 1, {}) secret = "sk-" + "proj-" + "Abc0123456789" * 6 assert make_sample({**payload, "messages": [{"content": secret}]}, *args) is None assert make_sample({**payload, "metadata": {"private": secret}}, *args) is None @@ -88,14 +100,14 @@ def test_secret_skip_size_caps_and_heuristic_preservation(): @pytest.mark.parametrize("rate", [-0.1, 0.21, True, float("nan"), float("inf"), "bad"]) def test_key_rate_validation(rate): with pytest.raises(APIError): - key_controls.validate({"shadow_eval_rate": rate}) + key_controls.validate(sampling_user(shadow_eval_rate=rate)) def test_key_rate_storage_and_default(): assert key_controls.validate({})["shadow_eval_rate"] is None assert key_controls.validate({"shadow_eval_rate": "0.1"})["shadow_eval_rate"] == 0.1 - assert key_controls.public(key_controls.from_storage({"shadow_eval_rate": 0.1}))["shadow_eval_rate"] == 0.1 - assert key_controls.from_storage({"shadow_eval_rate": 1})["shadow_eval_rate"] is None + assert key_controls.public(key_controls.from_storage({"shadow_eval_rate": 0.1, "official_only": None, "forced_provider_preferences": None}))["shadow_eval_rate"] == 0.1 + assert key_controls.from_storage({"shadow_eval_rate": 1, "official_only": None, "forced_provider_preferences": None})["shadow_eval_rate"] is None def test_store_ttl_cap_and_on_demand_text(store): @@ -313,7 +325,7 @@ def test_sampling_hook_never_fails_response_and_has_no_recursive_internal_sample monkeypatch.setattr("services.shadow_eval_sampling.submit", lambda value: captures.append(value)) @app.post("/v1/chat/completions") def complete(): - g.authenticated_user = {"username": "sampled-user", "shadow_eval_rate": 0.2} + g.authenticated_user = sampling_user(shadow_eval_rate=0.2) return response("Synthetic answer", "openai:production") payload = {"model": "auto:test", "messages": [{"content": "Synthetic greeting"}]} client = app.test_client() @@ -365,7 +377,7 @@ def dispatch(app, auth, metrics, proxy, payload): @app.post("/v1/responses") def respond(): - g.authenticated_user = {"username": "sampled-user", "shadow_eval_rate": 0.2} + g.authenticated_user = sampling_user(shadow_eval_rate=0.2) return dispatch(app, None, None, None, {"model": "auto:test", "messages": [{"role": "user", "content": "Synthetic"}]}) assert app.test_client().post("/v1/responses", json={"model": "auto:test", "input": "Synthetic"}).status_code == 200 @@ -523,7 +535,7 @@ def test_judge_message_budget_keeps_newest_and_whole_request_drops_oldest_first( def test_sample_contract_retains_request_limits_and_finish_reason(store): from services.shadow_eval_contract import valid_sample value = make_sample({"model": "auto:test", "messages": [], "max_tokens": 500, "max_completion_tokens": 600}, - {"username": "sampled-user"}, "auto:test", {"content": "Synthetic"}, "openai:production", 1, {}, finish_reason="length") + sampling_user(), "auto:test", {"content": "Synthetic"}, "openai:production", 1, {}, finish_reason="length") assert value["production_finish_reason"] == "length" assert value["request"]["max_tokens"] == 500 and value["request"]["max_completion_tokens"] == 600 store.put(value) @@ -542,7 +554,7 @@ def test_capture_preserves_nonstream_and_stream_finish_reasons(monkeypatch): init_shadow_sampling(app) payload = {"model": "auto:test", "messages": []} with app.test_request_context("/v1/chat/completions", method="POST", json=payload): - g.authenticated_user = {"username": "sampled-user", "shadow_eval_rate": 0.2} + g.authenticated_user = sampling_user(shadow_eval_rate=0.2) g.shadow_eval_started = time.monotonic() response = app.json.response({"choices": [{"message": {"content": "Synthetic"}, "finish_reason": "length"}], "multillm": {"selected_model": "openai:production"}}) @@ -582,7 +594,7 @@ def chat_dispatch(app, auth, metrics, proxy, payload): return response("Synthetic answer", "openai:production") def outer(): - g.authenticated_user = {"username": "sampled-user", "shadow_eval_rate": 0.2} + g.authenticated_user = sampling_user(shadow_eval_rate=0.2) # Even a routed subrequest with an otherwise eligible payload must be ignored. accounted_dispatch({**payload, "model": "auto:internal"}, lambda data: chat_dispatch(app, None, None, None, data), kind="chat") assert not getattr(g, "gateway_subrequest", False) @@ -629,7 +641,7 @@ def test_sampling_counters_cover_rate_secret_size_queue_error_and_are_bounded(mo payload = {"model": "auto:test", "messages": [{"role": "user", "content": "Synthetic"}]} def visit(body, reply=None): with app.test_request_context("/v1/chat/completions", method="POST", json=body): - g.authenticated_user = {"username": "sampled-user", "shadow_eval_rate": 0.2} + g.authenticated_user = sampling_user(shadow_eval_rate=0.2) g.shadow_eval_started = time.monotonic() reply = reply if reply is not None else response("Synthetic", "openai:production") assert coverage.sample_success(reply, body) is reply @@ -653,7 +665,7 @@ def visit(body, reply=None): def test_answer_secret_and_size_skips_are_visible(coverage): reasons = [] - args = ({"model": "auto:test", "messages": []}, {"username": "synthetic"}, "auto:test") + args = ({"model": "auto:test", "messages": []}, sampling_user(), "auto:test") secret = "sk-" + "proj-" + "Abc0123456789" * 6 assert make_sample(*args, {"content": secret}, "openai:production", 1, {}, on_skip=reasons.append) is None assert make_sample(*args, {"content": "x" * 32768}, "openai:production", 1, {}, on_skip=reasons.append) is None diff --git a/tests/test_shadow_eval_worker.mjs b/tests/test_shadow_eval_worker.mjs index c58cf004..0a2d7133 100644 --- a/tests/test_shadow_eval_worker.mjs +++ b/tests/test_shadow_eval_worker.mjs @@ -1,7 +1,7 @@ import assert from 'node:assert/strict'; import { readFileSync } from 'node:fs'; import test from 'node:test'; -import { convertV4MiniflareOptions, Miniflare } from 'miniflare'; +import { convertV4MiniflareOptions, Miniflare } from './helpers/local_miniflare.mjs'; import { applyMigrations } from './d1_migrations.mjs'; import { handleShadowEvalRequest, validConfig, validSample, validResult } from '../worker/shadow-eval-d1.mjs'; import { runScheduledShadowEval } from '../worker/shadow-eval-schedule.mjs'; diff --git a/tests/test_shadow_sampling_policy.py b/tests/test_shadow_sampling_policy.py new file mode 100644 index 00000000..c7ccdc64 --- /dev/null +++ b/tests/test_shadow_sampling_policy.py @@ -0,0 +1,132 @@ +"""Loaded account envelopes govern capture and queued content persistence.""" +import queue +import time +from unittest.mock import Mock + +import pytest +from flask import Flask, Response, g + +from services.dispatch_policy import envelope +from services.key_data_guards import has_data_controls +from services.shadow_eval_contract import eligible, make_sample +from services import shadow_eval_sampling as sampling +from services.user_store import USER_FIELDS, _row + + +POLICIES = [ + {"official_only": None, "forced_provider_preferences": None}, + {"official_only": 0, "forced_provider_preferences": None}, + {"official_only": None, "forced_provider_preferences": '{}'}, + {"official_only": 1, "forced_provider_preferences": None}, + {"official_only": None, "forced_provider_preferences": '{"data_collection":"deny"}'}, +] +PAYLOAD = {"model": "auto:synthetic", "messages": [{"role": "user", "content": "Synthetic request"}]} + + +def loaded_account(policy): + from services.auth_service import AuthService + raw = dict.fromkeys(USER_FIELDS) + raw.update(username="synthetic", api_key_hash="synthetic", api_key_prefix="synthetic", scopes="chat", + is_admin=0, created_at="2026-10-01T00:00:00Z", shadow_eval_rate=0.2, + **policy, data_policy=envelope(policy)) + user = AuthService._row_to_user(_row(raw)) + assert "official_only" not in user and "forced_provider_preferences" not in user + return user + + +@pytest.fixture +def capture(monkeypatch): + monkeypatch.setenv("CONTENT_RETENTION_ENABLED", "false") + pending = queue.Queue(maxsize=32) + monkeypatch.setattr(sampling, "_QUEUE", pending) + monkeypatch.setattr(sampling, "_WORKER", Mock(is_alive=Mock(return_value=True))) + monkeypatch.setattr(sampling.random, "random", lambda: 0) + stored = Mock() + monkeypatch.setattr(sampling.ShadowEvalStore, "put", stored) + return pending, stored + + +def drain(pending, monkeypatch): + get = pending.get + def next_sample(): + try: + return get(block=False) + except queue.Empty: + raise StopIteration from None + monkeypatch.setattr(pending, "get", next_sample) + with pytest.raises(StopIteration): + sampling._persist() + + +def complete(user, stream=False): + app = Flask(__name__) + with app.test_request_context("/v1/chat/completions", method="POST", json=PAYLOAD): + g.authenticated_user = user + g.shadow_eval_started = time.monotonic() + if stream: + response = Response(iter([b'data: {"choices":[{"delta":{"content":"Synthetic answer"},"finish_reason":"stop"}]}\n\n', + b'data: [DONE]\n\n']), mimetype="text/event-stream") + else: + response = Response('{"choices":[{"message":{"content":"Synthetic answer"},"finish_reason":"stop"}]}', mimetype="application/json") + response.headers["X-MultiLLM-Auto-Selected-Model"] = "openai:synthetic" + assert sampling.sample_success(response) is response + return response + + +@pytest.mark.parametrize("policy", POLICIES) +@pytest.mark.parametrize("stream", [False, True]) +def test_loaded_account_sample_success_captures_and_persists_only_null_policy(policy, stream, capture, monkeypatch): + current = loaded_account(policy) + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda _: current) + unrestricted = policy == POLICIES[0] + assert has_data_controls(current) is not unrestricted + assert eligible(PAYLOAD, current, "/v1/chat/completions", random_value=0) is unrestricted + response = complete(current, stream) + if stream: + list(response.response) # Capture finishes outside the Flask request context. + pending, stored = capture + assert pending.qsize() == int(unrestricted) + drain(pending, monkeypatch) + assert stored.call_count == int(unrestricted) + if unrestricted: + assert stored.call_args.args[0]["request"]["messages"] == PAYLOAD["messages"] + + +@pytest.mark.parametrize("stage", ["capture", "persist"]) +@pytest.mark.parametrize("replacement", [*POLICIES[1:], None, {}]) +def test_policy_change_or_unreadable_authority_stops_deferred_retention(stage, replacement, capture, monkeypatch): + current = loaded_account(POLICIES[0]) + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda _: current) + response = complete(current, stream=True) + if stage == "persist": + list(response.response) + current = loaded_account(replacement) if replacement else replacement + if stage == "capture": + list(response.response) + pending, stored = capture + drain(pending, monkeypatch) + stored.assert_not_called() + + +@pytest.mark.parametrize("policy", [None, {}, {"version": 1}, {"version": 0, **POLICIES[0]}, + {"version": 1, "official_only": None, "forced_provider_preferences": 'null'}]) +def test_malformed_canonical_policy_cannot_fall_back_to_flat_nulls(policy): + user = {"username": "synthetic", "shadow_eval_rate": 0.2, **POLICIES[0], "data_policy": policy} + assert not eligible(PAYLOAD, user, "/v1/chat/completions", random_value=0) + assert make_sample(PAYLOAD, user, "auto:synthetic", {"content": "Synthetic"}, "openai:synthetic", 1, {}) is None + + +@pytest.mark.parametrize("policy", [{}, {"version": 1}, {"version": 0, **POLICIES[0]}]) +def test_account_loading_does_not_replace_malformed_envelopes_with_unrestricted_policy(policy): + from error_handlers import APIError + from services.auth_service import AuthService + with pytest.raises(APIError): + AuthService._row_to_user({"data_policy": policy}) + + +@pytest.mark.parametrize("policy", POLICIES) +def test_gemini_override_admission_uses_canonical_non_null_policy(policy): + from services.gemini_credentials import restricted + with Flask(__name__).test_request_context(): + g.authenticated_user = loaded_account(policy) + assert restricted("gemini") is (policy != POLICIES[0]) diff --git a/tests/test_shared_generation_cache_worker.mjs b/tests/test_shared_generation_cache_worker.mjs index 574ee176..79771860 100644 --- a/tests/test_shared_generation_cache_worker.mjs +++ b/tests/test_shared_generation_cache_worker.mjs @@ -1,3 +1,12 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +const nativeGenerationFetch = (request, env, ctx, authority, ...rest) => rawNativeGenerationFetch( + request, policyFixtureEnv(env), ctx, { authenticated: true, keyId: authority?.principal?.id ?? "admin", ...authority }, ...rest); +async function loadWorkerModule(...args) { + const module = await rawLoadWorkerModule(...args); + return { ...module, default: { ...module.default, fetch(request, env, ctx) { + return module.default.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} import assert from "node:assert/strict"; import test from "node:test"; import { readFileSync } from "node:fs"; @@ -5,14 +14,16 @@ import { DatabaseSync } from "node:sqlite"; import { GenerationCacheD1, cleanupGenerationCache, CACHE_PREFIX } from "../worker/generation-cache-d1.mjs"; import { generationCacheSettings, exactCacheIdentity, eligibleRequest, prepareNativeCache } from "../worker/exact-generation-cache.mjs"; import { completeCacheBody } from "../worker/generation-cache-d1.mjs"; -import { nativeGenerationFetch } from "../worker/gateway-extensions.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { nativeGenerationFetch as rawNativeGenerationFetch } from "../worker/gateway-extensions.mjs"; +import { loadWorkerModule as rawLoadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; import { handleControlStateRequest } from "../worker/control-state-d1.mjs"; const MIGRATION = "0019_generation_cache.sql"; function database() { const db = new DatabaseSync(":memory:"); db.exec(readFileSync(new URL(`../intelligence-migrations/${MIGRATION}`, import.meta.url), "utf8")); + for (const file of ["0003_control_users.sql", "0041_key_data_guards.sql"]) + db.exec(readFileSync(new URL(`../intelligence-migrations/${file}`, import.meta.url), "utf8")); const execute = (sql, args) => { const numbered = /\?[1-9]/.test(sql); const statement = db.prepare(sql.replace(/\?(\d+)/g, (_, number) => `$v${number}`)); diff --git a/tests/test_skills_clients.py b/tests/test_skills_clients.py index 3c9e2a9e..9c3acc04 100644 --- a/tests/test_skills_clients.py +++ b/tests/test_skills_clients.py @@ -1,5 +1,6 @@ """Local-only client tests with temporary libraries and a synthetic HTTP server.""" +from contextvars import copy_context import importlib.util import json from pathlib import Path @@ -11,6 +12,8 @@ import pytest +from services.dispatch_network import control_network + ROOT = Path(__file__).resolve().parents[1] @@ -153,7 +156,8 @@ def log_message(self, *_args): server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() - yield f"http://127.0.0.1:{server.server_port}", received + with control_network(f"http://127.0.0.1:{server.server_port}"): + yield f"http://127.0.0.1:{server.server_port}", received server.shutdown() server.server_close() thread.join() @@ -235,7 +239,9 @@ def test_hook_real_rest_find_uses_fast_limit_three(fake_server, monkeypatch): base_url, received = fake_server monkeypatch.setenv("MULTILLM_BASE_URL", base_url) monkeypatch.setenv("MULTILLM_KNOWLEDGE_API_KEY", "synthetic-hook-key") - result = hook.hint({"prompt": "Test regression behavior"}) + context = copy_context() + result = hook.hint({"prompt": "Test regression behavior"}, + fetcher=lambda *args, **kwargs: context.run(hook.fetch, *args, **kwargs)) assert "Relevant skills: Testing" in result["hookSpecificOutput"]["additionalContext"] assert received == [{"query": "Test regression behavior", "mode": "fast", "limit": 3, "min_confidence": "high"}] diff --git a/tests/test_sse_heartbeat_worker.mjs b/tests/test_sse_heartbeat_worker.mjs index 884fe4f8..2dd93d25 100644 --- a/tests/test_sse_heartbeat_worker.mjs +++ b/tests/test_sse_heartbeat_worker.mjs @@ -5,7 +5,7 @@ import { resolveSseHeartbeatMs, withSseHeartbeat, } from "../worker/sse-heartbeat.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerWithUnrestrictedControls as loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; const encoder = new TextEncoder(); diff --git a/tests/test_stream_cost_breaker.py b/tests/test_stream_cost_breaker.py index b5f565f0..e0d412fb 100644 --- a/tests/test_stream_cost_breaker.py +++ b/tests/test_stream_cost_breaker.py @@ -187,7 +187,7 @@ def test_strict_d1_rows_missing_schema_and_off_wire(monkeypatch): row["max_stream_cost_microusd"]=4 assert breaker.get_user("reader")["max_stream_cost_microusd"]==4 row["unknown"]=None - with pytest.raises(APIError): breaker.get_user("reader") + assert breaker.get_user("reader")["max_stream_cost_microusd"] == 4 # Additive protocol metadata is ignored. row.pop("unknown"); row.pop("max_stream_cost_microusd") monkeypatch.setenv("STREAM_COST_BREAKER_ENABLED","false") captured=[] diff --git a/tests/test_tenant_hierarchy_worker.mjs b/tests/test_tenant_hierarchy_worker.mjs index a3456599..c6bf4512 100644 --- a/tests/test_tenant_hierarchy_worker.mjs +++ b/tests/test_tenant_hierarchy_worker.mjs @@ -1,14 +1,17 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +const nativeGenerationFetch = (request, env, ctx, authority, ...rest) => rawNativeGenerationFetch( + request, policyFixtureEnv(env), ctx, { authenticated: true, keyId: authority?.principal?.id ?? "admin", ...authority }, ...rest); import assert from "node:assert/strict"; import test from "node:test"; import { DatabaseSync } from "node:sqlite"; import { readFileSync } from "node:fs"; import { scryptSync } from "node:crypto"; -import { handleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; +import { handleKnowledgeEdgeRequest as rawhandleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; import { handleRealtimeRequest } from "../worker/realtime.mjs"; import { activeUsersByPrefix, USER_FIELDS } from "../worker/control-users-d1.mjs"; import { handleTenantRequest, organisationsEnabled, resolveAccountTenant, tenantNamespace } from "../worker/tenants-d1.mjs"; import { handleManagedStateRequest } from "../worker/managed-state-dispatch.mjs"; -import { nativeGenerationFetch } from "../worker/gateway-extensions.mjs"; +import { nativeGenerationFetch as rawNativeGenerationFetch } from "../worker/gateway-extensions.mjs"; test("disabled tenant domain and resolver touch no storage", async () => { const db = { prepare() { throw Error("No storage"); } }; @@ -189,12 +192,12 @@ test("native requests ignore headers and body scope and resolve before dispatch" }); test("prefix lookup leaves tenant resolution to the verified principal", async () => { - const row = { username: "alice" }; + const row = { username: "alice", official_only: null, forced_provider_preferences: null }; const db = { prepare(sql) { assert.ok(sql.includes("api_key_prefix=?")); return { bind() { return this; }, async all() { return { results: [row] }; } }; } }; - assert.deepEqual(await activeUsersByPrefix(db, "fixture", { ORGANISATIONS_ENABLED: "true" }), [row]); + assert.deepEqual(await activeUsersByPrefix(db, "fixture", { ORGANISATIONS_ENABLED: "true" }), [{ ...row, data_policy: { version: 1, official_only: null, forced_provider_preferences: null } }]); assert.equal(Object.hasOwn(row, "tenant_context"), false); }); @@ -223,10 +226,10 @@ function authFixture() { async first() { return sql.includes("intelligence_credentials") ? integrationRow : null; }, async run() { calls.writes++; return { meta: { changes: 1 } }; } }; } }; - const env = { ...f.env, INTELLIGENCE_DB: db, AUTH_STORAGE_BACKEND: "d1", ADMIN_API_KEY: "synthetic-bootstrap", + const env = policyFixtureEnv({ ...f.env, INTELLIGENCE_DB: db, AUTH_STORAGE_BACKEND: "d1", ADMIN_API_KEY: "synthetic-bootstrap", ADMIN_USERNAME: "operator", JWT_SECRET: "synthetic-signing", REALTIME_ENABLED: "true", REALTIME_PROVIDERS_JSON: '{"openai:realtime-test":{"url":"wss://approved.example/realtime"}}', - KNOWLEDGE_SERVICE: { async fetch() { calls.dispatches++; return Response.json({ version: 1, result: { status: "ok" } }); } } }; + KNOWLEDGE_SERVICE: { async fetch() { calls.dispatches++; return Response.json({ version: 1, result: { status: "ok" } }); } } }); return { ...f, env, calls }; } const knowledgeRequest = (key, mcp = false) => new Request(mcp ? "https://gateway.example/mcp" : "https://gateway.example/v1/knowledge/artifacts/synthetic", @@ -339,3 +342,5 @@ test("D1 malformed fields and unlisted operations never produce state changes", assert.equal(f.sqlite.prepare("SELECT count(*) AS n FROM tenant_memberships").get().n, 0); } finally { f.close(); } }); + +function handleKnowledgeEdgeRequest(request, env, ...args) { return rawhandleKnowledgeEdgeRequest(request, policyFixtureEnv(env), ...args); } diff --git a/tests/test_tool_repair_worker.mjs b/tests/test_tool_repair_worker.mjs index 5e383ee5..15f276c3 100644 --- a/tests/test_tool_repair_worker.mjs +++ b/tests/test_tool_repair_worker.mjs @@ -2,7 +2,7 @@ import assert from "node:assert/strict"; import test from "node:test"; import { collectContainerEnv } from "../worker/container-env.mjs"; import { CORS_DEFAULT_HEADERS, CORS_EXPOSE_HEADERS } from "../worker/cors-policy.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerWithUnrestrictedControls as loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; test("tool repair default reaches the container", () => { assert.equal(collectContainerEnv({ TOOL_CALL_REPAIR_DEFAULT: "full" }).TOOL_CALL_REPAIR_DEFAULT, "full"); diff --git a/tests/test_transport_resilience.py b/tests/test_transport_resilience.py index 9bd8cfc0..b624c176 100644 --- a/tests/test_transport_resilience.py +++ b/tests/test_transport_resilience.py @@ -8,6 +8,12 @@ import requests + +import pytest +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + class FakeSession: def __init__(self, responses): self.responses = list(responses) diff --git a/tests/test_upstream_outcome.py b/tests/test_upstream_outcome.py index 4d6703e2..01efb963 100644 --- a/tests/test_upstream_outcome.py +++ b/tests/test_upstream_outcome.py @@ -31,6 +31,11 @@ SUCCESS_STATUSES = [200, 201, 204, 206, 299] + +from tests.dispatch_fixture import verified_transport_principal # noqa: F401 + +pytestmark = pytest.mark.usefixtures("verified_transport_principal") + @pytest.fixture def check(): return TestCase() diff --git a/tests/test_usage_ledger_worker.mjs b/tests/test_usage_ledger_worker.mjs index 9b79675a..8c46f8df 100644 --- a/tests/test_usage_ledger_worker.mjs +++ b/tests/test_usage_ledger_worker.mjs @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { handleUsageLedgerRequest, latencyBucket, LATENCY_BUCKETS, validRow } from "../worker/usage-ledger-d1.mjs"; diff --git a/tests/test_usage_receipts_worker.mjs b/tests/test_usage_receipts_worker.mjs index 847caa99..c4f5c9b7 100644 --- a/tests/test_usage_receipts_worker.mjs +++ b/tests/test_usage_receipts_worker.mjs @@ -2,7 +2,7 @@ import assert from "node:assert/strict"; import test from "node:test"; import { readFileSync } from "node:fs"; import { DatabaseSync } from "node:sqlite"; -import { convertV4MiniflareOptions, Miniflare } from "miniflare"; +import { convertV4MiniflareOptions, Miniflare } from "./helpers/local_miniflare.mjs"; import { canonicalBytes, appendReceipt, verifyChain, handleUsageReceiptRequest, handleUsageReceiptStoreRequest } from "../worker/usage-receipts.mjs"; import { handleManagedStateRequest } from "../worker/managed-state-dispatch.mjs"; import { handleUsageLedgerRequest } from "../worker/usage-ledger-d1.mjs"; diff --git a/tests/test_usage_reservations_worker.mjs b/tests/test_usage_reservations_worker.mjs index 6e4e117f..7210e4af 100644 --- a/tests/test_usage_reservations_worker.mjs +++ b/tests/test_usage_reservations_worker.mjs @@ -214,7 +214,11 @@ test("registered native dispatch preserves response bytes and completes the inje const { operation, id: ignored, ...values } = reserve(); return values; } }); - const { nativeGenerationFetch } = await import("../worker/gateway-extensions.mjs"); + const { nativeGenerationFetch: rawNativeGenerationFetch } = await import("../worker/gateway-extensions.mjs"); + const { policyFixtureEnv } = await import("./helpers/dispatch_fixture.mjs"); + const nativeGenerationFetch = (request, env, ctx, authority, ...rest) => rawNativeGenerationFetch( + request, policyFixtureEnv(env), ctx, { authenticated: true, keyId: authority?.principal?.id ?? "admin", ...authority }, ...rest); + const body = '{"choices":[],"usage":{"prompt_tokens":10,"completion_tokens":20}}'; const response = await nativeGenerationFetch(new Request("http://localhost/openai/v1/chat/completions", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ model: "test" }), diff --git a/tests/test_user_store_d1.py b/tests/test_user_store_d1.py index 9805ae28..5d2c3631 100644 --- a/tests/test_user_store_d1.py +++ b/tests/test_user_store_d1.py @@ -8,6 +8,7 @@ from error_handlers import APIError from services import user_store +from services.dispatch_policy import envelope from services.auth_service import AuthService from services.intelligence_d1_store import PrivateIntelligenceError @@ -24,6 +25,8 @@ def __call__(self, payload, *, endpoint): self.calls.append(operation) if self.down: raise PrivateIntelligenceError(503, "storage_unavailable") + if operation == "capabilities": + return {"version": 1, "key_controls": {"protocol": 1, "enforcement": False}} if operation == "list": names = sorted(name for name in self.rows if payload["after"] is None or name > payload["after"]) return {"version": 1, "users": [dict(self.rows[name]) for name in names[:payload["limit"]]]} @@ -35,7 +38,7 @@ def __call__(self, payload, *, endpoint): if row["api_key_prefix"] == payload["prefix"] and row["revoked_at"] is None]} if operation == "upsert": assert set(payload["user"]) == set(user_store.USER_FIELDS) - self.rows[payload["user"]["username"]] = dict(payload["user"]) + self.rows[payload["user"]["username"]] = {**payload["user"], "data_policy": envelope({"official_only": None, "forced_provider_preferences": None, **payload["user"]})} return {"version": 1, "stored": True} if operation == "delete": return {"version": 1, "deleted": self.rows.pop(payload["username"], None) is not None} @@ -193,7 +196,7 @@ def refuse(self, payload, *, endpoint): user_store.get_user("agent") with pytest.raises(APIError): user_store.upsert_user(d1.rows["agent"]) - assert d1.calls == ["get", "upsert"] + assert d1.calls == ["capabilities", "capabilities", "get", "upsert"] def test_account_lookups_get_a_longer_transport_budget(monkeypatch): @@ -276,3 +279,42 @@ def fake_submit(url, body, stopped, deadline, results, slots, success_statuses, assert store.request_private_intelligence({"operation": "policy"}) == {"version": 1, "policy": None}, \ "exhausted account slots leave the store's slots free" assert store.transport_stats()["users"]["slot_exhausted"] == before + 1 + + +def test_policy_extensions_are_sent_only_after_a_readable_capability(monkeypatch): + calls = [] + + def peer(payload, *, endpoint): + calls.append(dict(payload)) + if payload["operation"] == "capabilities": + return {"key_controls": {"protocol": 1, "enforcement": False}} + return {"user": None} + + monkeypatch.setattr(user_store, "request_private_intelligence", peer) + assert user_store.get_user("synthetic") is None + assert calls == [{"operation": "capabilities"}, {"operation": "get", "username": "synthetic", "data_policy_version": 1}] + calls.clear() + + def old_peer(payload, *, endpoint): + calls.append(dict(payload)) + if payload["operation"] == "capabilities": + raise PrivateIntelligenceError(400, "invalid_request") + return {"user": None} + + monkeypatch.setattr(user_store, "request_private_intelligence", old_peer) + assert user_store.get_user("synthetic") is None + assert calls == [{"operation": "capabilities"}, {"operation": "get", "username": "synthetic"}] + + +@pytest.mark.parametrize("capability", [None, {}, {"protocol": 0}, {"protocol": True}]) +def test_malformed_capability_never_downgrades_to_an_unverified_row(monkeypatch, capability): + calls = [] + + def peer(payload, *, endpoint): + calls.append(payload) + return {"key_controls": capability} + + monkeypatch.setattr(user_store, "request_private_intelligence", peer) + with pytest.raises(APIError): + user_store.get_user("synthetic") + assert calls == [{"operation": "capabilities"}] diff --git a/tests/test_wave3_hub_worker.mjs b/tests/test_wave3_hub_worker.mjs index e6e27a84..92a572da 100644 --- a/tests/test_wave3_hub_worker.mjs +++ b/tests/test_wave3_hub_worker.mjs @@ -1,7 +1,16 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +const nativeGenerationFetch = (request, env, ctx, authority, ...rest) => rawNativeGenerationFetch( + request, policyFixtureEnv(env), ctx, { authenticated: true, keyId: authority?.principal?.id ?? "admin", ...authority }, ...rest); +async function loadWorkerModule(...args) { + const module = await rawLoadWorkerModule(...args); + return { ...module, default: { ...module.default, fetch(request, env, ctx) { + return module.default.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} import assert from "node:assert/strict"; import test from "node:test"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; -import { nativeGenerationFetch } from "../worker/gateway-extensions.mjs"; +import { loadWorkerModule as rawLoadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { nativeGenerationFetch as rawNativeGenerationFetch } from "../worker/gateway-extensions.mjs"; import { principalHash, admissionModelGroup } from "../worker/admission-do.mjs"; import { nativeRevisionConsumer } from "../worker/native-config-sync.mjs"; @@ -45,10 +54,10 @@ function fixture(overrides = {}) { : { lease: { lease_id: "a".repeat(32), expires_at: Math.min(Date.now() + 29000, body.deadline_ms) } }) }); } }; } }; env.INTELLIGENCE_DB = { prepare(sql) { return { sql, values: [], bind(...values) { this.values = values; return this; }, - async all() { return { results: [] }; } }; }, async batch(statements) { + async all() { return { results: [] }; }, async first() { return null; } }; }, async batch(statements) { rows.push(...JSON.parse(statements[1].values[0])); return [{ meta: { changes: 1 } }]; } }; - return { env, operations, rows, ctx: { waitUntil(promise) { pending.push(promise); } }, + return { env: policyFixtureEnv(env), operations, rows, ctx: { waitUntil(promise) { pending.push(promise); } }, async flush() { await Promise.all(pending); } }; } async function upstream(fetcher, run) { @@ -141,7 +150,8 @@ test("native revision installers are strict and the injectable clock bounds fres let now = 100, offline = false; const f = fixture({ CONFIG_REVISION_SYNC_ENABLED: "true" }); const db = f.env.INTELLIGENCE_DB; - db.prepare = sql => ({ sql, bind() { return this; }, async all() { + db.prepare = sql => ({ sql, bind() { return this; }, async first() { if (offline) throw Error("offline"); return sql.startsWith("SELECT official_only") + ? { official_only: null, forced_provider_preferences: null } : null; }, async all() { if (offline) throw new Error("offline"); return { results: [] }; } }); diff --git a/tests/test_wave3_transport_worker.mjs b/tests/test_wave3_transport_worker.mjs index 4e79051f..877c00be 100644 --- a/tests/test_wave3_transport_worker.mjs +++ b/tests/test_wave3_transport_worker.mjs @@ -1,3 +1,4 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; import assert from "node:assert/strict"; import test from "node:test"; import { DatabaseSync } from "node:sqlite"; @@ -5,7 +6,7 @@ import { readFile } from "node:fs/promises"; import { createHash } from "node:crypto"; import { handleUsageLedgerRequest, recordNativeUsage } from "../worker/usage-ledger-d1.mjs"; import { BUCKET_FIELDS } from "../worker/usage-buckets-d1.mjs"; -import { handleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; +import { handleKnowledgeEdgeRequest as rawhandleKnowledgeEdgeRequest } from "../worker/knowledge-edge.mjs"; import { dispatchKnowledge } from "../worker/knowledge/service.mjs"; import { HandoffStore } from "../worker/knowledge/handoff-store.mjs"; import { fixture, request as knowledgeQuery } from "./knowledge_fixture.mjs"; @@ -319,3 +320,5 @@ test("public key and route zero prevent durable roleplay recovery and memory rea assert.equal(entry.storage.operations.put, 0); } }); + +function handleKnowledgeEdgeRequest(request, env, ...args) { return rawhandleKnowledgeEdgeRequest(request, policyFixtureEnv(env), ...args); } diff --git a/tests/test_wave4_worker.mjs b/tests/test_wave4_worker.mjs index 1bf1a23f..4b23c8b1 100644 --- a/tests/test_wave4_worker.mjs +++ b/tests/test_wave4_worker.mjs @@ -1,14 +1,23 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +const nativeGenerationFetch = (request, env, ctx, authority, ...rest) => rawNativeGenerationFetch( + request, policyFixtureEnv(env), ctx, { authenticated: true, ...authority }, ...rest); +async function loadWorkerModule(...args) { + const module = await rawLoadWorkerModule(...args); + return { ...module, default: { ...module.default, fetch(request, env, ctx) { + return module.default.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} import test from "node:test"; import assert from "node:assert/strict"; import { DatabaseSync } from "node:sqlite"; import { readFileSync } from "node:fs"; -import { nativeGenerationFetch, forwardedGenerationHeaders, nativeGenerationSetup, generationErrorResponse } from "../worker/gateway-extensions.mjs"; +import { nativeGenerationFetch as rawNativeGenerationFetch, forwardedGenerationHeaders, nativeGenerationSetup, generationErrorResponse } from "../worker/gateway-extensions.mjs"; import { runScheduledMaintenance, scheduledMaintenanceEnabled, collectAlertAggregates } from "../worker/scheduled-maintenance.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { handleAlertState, normalizeAlertRules } from "../worker/alert-delivery.mjs"; import { retentionRequestId } from "../worker/retention-policy.mjs"; import { generationDeadlineHook } from "../worker/generation-deadline.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerModule as rawLoadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; const authority = { route: "/openai/v1/chat/completions", provider: "openai", principal: { id: "admin" } }; const completed = JSON.stringify({ choices: [{ message: { content: "fixture" }, finish_reason: "stop" }], usage: { prompt_tokens: 3, completion_tokens: 2 } }); diff --git a/tests/test_wave5_worker.mjs b/tests/test_wave5_worker.mjs index 4df1e6c4..b6a9a651 100644 --- a/tests/test_wave5_worker.mjs +++ b/tests/test_wave5_worker.mjs @@ -1,15 +1,26 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +const nativeGenerationFetch = (request, env, ctx, authority, ...rest) => rawNativeGenerationFetch( + request, policyFixtureEnv(env), ctx, { authenticated: true, ...authority }, ...rest); +async function loadWorkerModule(...args) { + const module = await rawLoadWorkerModule(...args); + return { ...module, default: { ...module.default, fetch(request, env, ctx) { + return module.default.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} import assert from "node:assert/strict"; import test from "node:test"; import { DatabaseSync } from "node:sqlite"; import { readFileSync } from "node:fs"; -import { nativeGenerationFetch, nativeCacheHeader } from "../worker/gateway-extensions.mjs"; +import { nativeGenerationFetch as rawNativeGenerationFetch, nativeCacheHeader } from "../worker/gateway-extensions.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { runScheduledMaintenance, scheduledMaintenanceEnabled } from "../worker/scheduled-maintenance.mjs"; import { ObservabilityExporter, setObservabilityExporter } from "../worker/observability-export.mjs"; import { guardHash } from "../worker/semantic-generation-cache.mjs"; import { prepareCanary } from "../worker/canary-assignment.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; -import { makeRoleplayEnv, roleplayRequest, handleRoleplayEdgeRequest, withGlobalFetch, completionResponse } from "./helpers/roleplay_fixture.mjs"; +import { loadWorkerModule as rawLoadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { makeRoleplayEnv as rawMakeRoleplayEnv, roleplayRequest, handleRoleplayEdgeRequest, withGlobalFetch, completionResponse } from "./helpers/roleplay_fixture.mjs"; + +const makeRoleplayEnv = overrides => rawMakeRoleplayEnv(policyFixtureEnv(overrides)); const route = "/codex-easy/v1/chat/completions"; const authority = { route, provider: "codex-easy", principal: { id: "admin", keyHash: "a".repeat(64) } }; @@ -87,8 +98,10 @@ test("observability alone flushes content-free native observations through waitU const waits = [], exported = [], env = { OBSERVABILITY_EXPORTERS_JSON: JSON.stringify([{ type: "langfuse", endpoint: "https://collector.invalid/ingest", allowed_origins: ["https://collector.invalid"], credential_env: "FIXTURE_EXPORT_CREDENTIAL" }]), FIXTURE_EXPORT_CREDENTIAL: "synthetic" }; const exporter = new ObservabilityExporter(env, { transport: async (_url, options) => { exported.push(options.body); return new Response(null, { status: 200 }); } }); - setObservabilityExporter(env, exporter); - const res = await nativeGenerationFetch(request(), env, { waitUntil: promise => waits.push(promise) }, authority, response); + const guardedEnv = policyFixtureEnv(env); + setObservabilityExporter(guardedEnv, exporter); + const res = await rawNativeGenerationFetch(request(), guardedEnv, { waitUntil: promise => waits.push(promise) }, + { authenticated: true, ...authority }, response); assert.equal(await res.text(), complete); await Promise.all(waits); assert.equal(exported.length, 1); assert.doesNotMatch(exported[0], /Explain caching|fixture|placeholder/); assert.equal(exporter.status().accepted, 1); @@ -121,7 +134,7 @@ test("semantic maintenance uses at most three bounded pages and is off by defaul function durableFixture(t) { const sql = new DatabaseSync(":memory:"); t.after(() => sql.close()); - for (const name of ["0003_control_users", "0005_auto_routes", "0007_usage_ledger", "0019_generation_cache", + for (const name of ["0003_control_users", "0005_auto_routes", "0007_usage_ledger", "0041_key_data_guards", "0042_usage_token_details", "0019_generation_cache", "0020_usage_reservations", "0024_context_pages", "0025_canary_traffic", "0026_semantic_cache"]) { sql.exec(readFileSync(new URL(`../intelligence-migrations/${name}.sql`, import.meta.url), "utf8")); } diff --git a/tests/test_wave6_worker.mjs b/tests/test_wave6_worker.mjs index dbf4db0f..1f2c9b3b 100644 --- a/tests/test_wave6_worker.mjs +++ b/tests/test_wave6_worker.mjs @@ -1,14 +1,23 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +const nativeGenerationFetch = (request, env, ctx, authority, ...rest) => rawNativeGenerationFetch( + request, policyFixtureEnv(env), ctx, { authenticated: true, ...authority }, ...rest); +async function loadWorkerModule(...args) { + const module = await rawLoadWorkerModule(...args); + return { ...module, default: { ...module.default, fetch(request, env, ctx) { + return module.default.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} import assert from "node:assert/strict"; import test from "node:test"; import { readFile } from "node:fs/promises"; -import { nativeGenerationFetch, withForwardedCorrelation } from "../worker/gateway-extensions.mjs"; +import { nativeGenerationFetch as rawNativeGenerationFetch, withForwardedCorrelation } from "../worker/gateway-extensions.mjs"; import { ObservationWindow } from "../worker/latency-slo.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; import { DatabaseSync } from "node:sqlite"; import { handleReservationsRequest } from "../worker/reservations-d1.mjs"; import { activeUsersByPrefix } from "../worker/control-users-d1.mjs"; import { renderNativeMetrics } from "../worker/request-telemetry.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerModule as rawLoadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; const route = "/codex-easy/v1/chat/completions", now = Date.now(); const auth = { authenticated: true, route, provider: "codex-easy", keyId: "admin", principal: { id: "admin" } }; @@ -46,7 +55,7 @@ test("disabled native streaming and nonstream bodies, headers and account reads for (const streaming of [false, true]) { const upstream = streaming ? source([delta("ok"), "data: [DONE]\n\n"]).response : complete("ok"); const expected = await upstream.clone().text(), headers = [...upstream.headers]; - const res = await nativeGenerationFetch(request({ stream: streaming }), { INTELLIGENCE_DB: { prepare() { assert.fail("storage"); } } }, {}, auth, + const res = await nativeGenerationFetch(request({ stream: streaming }), { INTELLIGENCE_DB: { prepare(sql) { if (sql.startsWith("SELECT official_only")) return { bind() { return { first: async () => null }; } }; assert.fail("storage"); } } }, {}, auth, async req => { assert.deepEqual(await req.json(), { ...payload, stream: streaming }); return upstream; }); assert.equal(await res.text(), expected); assert.deepEqual([...res.headers], headers); } @@ -64,7 +73,7 @@ test("latency rejection runs before admission, cache, reservation, canary and di const store = new ObservationWindow(); store.record("codex-easy:m", { ttft_ms: 500, tokens_per_second: 1, now }); const env = { ...slo(), ...cost, ...canary("block"), USAGE_RESERVATIONS_ENABLED: "true", ADMISSION_ENABLED: "true", GENERATION_CACHE_SHARED_ENABLED: "true", GENERATION_CACHE_BACKEND: "d1-r2", SEMANTIC_CACHE_ENABLED: "true", - INTELLIGENCE_DB: { prepare() { assert.fail("storage before rejection"); } }, + INTELLIGENCE_DB: { prepare(sql) { if (sql.startsWith("SELECT official_only")) return { bind() { return { first: async () => null }; } }; assert.fail("storage before rejection"); } }, ADMISSION_COORDINATOR: { getByName() { assert.fail("admission before rejection"); } } }; const response = await nativeGenerationFetch(request(), env, {}, { ...auth, observations: store, now }, () => assert.fail("upstream")); assert.equal(response.status, 503); assert.equal((await response.json()).error.code, "latency_slo_predicted_miss"); @@ -92,6 +101,7 @@ test("capped native streams reject missing pricing before upstream", async () => test("a capped oversized streaming payload cannot bypass preflight inspection", async () => { const req = request({ stream: true, messages: [{ role: "user", content: "x".repeat(65536) }] }); const response = await nativeGenerationFetch(req, { ...cost, INTELLIGENCE_DB: { prepare(sql) { + if (sql.startsWith("SELECT official_only")) return { bind() { return { first: async () => null }; } }; assert.match(sql, /SELECT max_stream_cost_microusd/); return { bind() { return this; }, first() { return { max_stream_cost_microusd: 10 }; } }; } } }, {}, auth, () => assert.fail("upstream before bounded capped preflight")); @@ -110,7 +120,8 @@ test("stream cost stop delivers one error and one cancel and retains ambiguous a test("cost policy leaves nonstream malformed payloads alone and classifies stopped metrics before counting", async () => { const malformed = new Request("https://provider.invalid/v1/chat/completions", { method: "POST", body: "fixture-invalid" }); - const passthrough = await nativeGenerationFetch(malformed, cost, {}, auth, async req => { + const passthrough = await nativeGenerationFetch(malformed, cost, {}, { ...auth, + principal: { ...auth.principal, max_stream_cost_microusd: null } }, async req => { assert.equal(await req.text(), "fixture-invalid"); return complete("ok"); }); assert.equal(passthrough.status, 200); await passthrough.text(); @@ -129,7 +140,7 @@ for (const mode of ["log", "block"]) test("native canary " + mode + " scans befo try { const response = await nativeGenerationFetch(request({ temperature: 0 }, { "X-MultiLLM-Cache": "on" }), { ...canary(mode), GENERATION_CACHE_SHARED_ENABLED: "true", GENERATION_CACHE_BACKEND: "d1-r2", SEMANTIC_CACHE_ENABLED: "true", - INTELLIGENCE_DB: { prepare() { assert.fail("cache storage"); } } }, {}, auth, async req => { + INTELLIGENCE_DB: { prepare(sql) { if (sql.startsWith("SELECT official_only")) return { bind() { return { first: async () => null }; } }; assert.fail("cache storage"); } } }, {}, auth, async req => { const body = await req.json(); marker = body.messages[0].content.match(/marker: ([a-f0-9]{32})/)[1]; return complete("before " + marker + " after"); }, events(seen)); @@ -245,7 +256,7 @@ function sqliteDatabase(sql, observe = () => {}) { test("combined policies meter before settlement and retain a hold when the cost stop wins", async t => { const sql = new DatabaseSync(":memory:"); t.after(() => sql.close()); - for (const file of ["0003_control_users.sql", "0007_usage_ledger.sql", "0020_usage_reservations.sql"]) + for (const file of ["0003_control_users.sql", "0007_usage_ledger.sql", "0041_key_data_guards.sql", "0042_usage_token_details.sql", "0020_usage_reservations.sql"]) sql.exec(await readFile(new URL("../intelligence-migrations/" + file, import.meta.url), "utf8")); const order = [], seen = [], store = new ObservationWindow(); store.record("codex-easy:m", { ttft_ms: 1, tokens_per_second: 1000, now }); diff --git a/tests/test_wave7_worker.mjs b/tests/test_wave7_worker.mjs index 7ee50c45..fe5d9896 100644 --- a/tests/test_wave7_worker.mjs +++ b/tests/test_wave7_worker.mjs @@ -1,3 +1,12 @@ +import { policyFixtureEnv } from "./helpers/dispatch_fixture.mjs"; +const nativeGenerationFetch = (request, env, ctx, authority, ...rest) => rawNativeGenerationFetch( + request, policyFixtureEnv(env), ctx, { authenticated: true, ...authority }, ...rest); +async function loadWorkerModule(...args) { + const module = await rawLoadWorkerModule(...args); + return { ...module, default: { ...module.default, fetch(request, env, ctx) { + return module.default.fetch(request, policyFixtureEnv(env), ctx); + } } }; +} import assert from "node:assert/strict"; import test from "node:test"; import { DatabaseSync } from "node:sqlite"; @@ -7,7 +16,7 @@ import { createCreditsLifecycle, CreditsAdmissionError } from "../worker/credits import { appendCredit, readCredits, contextOwner } from "../worker/credits-d1.mjs"; import { handleTenantGovernanceRequest } from "../worker/tenant-governance-d1.mjs"; import { handleReservationsRequest } from "../worker/reservations-d1.mjs"; -import { nativeGenerationFetch, nativeGenerationSetup } from "../worker/gateway-extensions.mjs"; +import { nativeGenerationFetch as rawNativeGenerationFetch, nativeGenerationSetup } from "../worker/gateway-extensions.mjs"; import { exactCacheIdentity } from "../worker/exact-generation-cache.mjs"; import { semanticPartition } from "../worker/semantic-generation-cache.mjs"; import { handleIntelligenceOutbound } from "../worker/intelligence-outbound.mjs"; @@ -18,7 +27,7 @@ import { recordNativeUsage } from "../worker/usage-ledger-d1.mjs"; import { tenantStorageKey } from "../worker/tenants-d1.mjs"; import { retentionRequestId } from "../worker/retention-policy.mjs"; import { USER_FIELDS } from "../worker/control-users-d1.mjs"; -import { loadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; +import { loadWorkerModule as rawLoadWorkerModule } from "./helpers/load_cloudflare_worker.mjs"; import { handleManagedStateRequest } from "../worker/managed-state-dispatch.mjs"; const request = path => new Request(`http://intelligence.internal${path}`, { @@ -48,7 +57,7 @@ function fixture(t, flags = {}) { for (const migration of ["0003_control_users.sql", "0007_usage_ledger.sql", "0011_secret_firewall.sql", "0013_shadow_eval.sql", "0017_control_revisions.sql", "0019_generation_cache.sql", "0020_usage_reservations.sql", "0023_idempotency.sql", "0024_context_pages.sql", "0026_semantic_cache.sql", "0028_responses_state.sql", "0033_tenant_hierarchy.sql", "0034_tenant_governance.sql", - "0035_saml_federation.sql", "0036_scim_provisioning.sql", "0037_credits_ledger.sql", "0038_payment_billing.sql"]) + "0035_saml_federation.sql", "0036_scim_provisioning.sql", "0037_credits_ledger.sql", "0038_payment_billing.sql", "0041_key_data_guards.sql", "0042_usage_token_details.sql"]) sql.exec(readFileSync(new URL(`../intelligence-migrations/${migration}`, import.meta.url), "utf8")); const operations = []; let failBatch = false; diff --git a/tests/test_zero_content_retention.py b/tests/test_zero_content_retention.py index 741b9e12..1a64bc9c 100644 --- a/tests/test_zero_content_retention.py +++ b/tests/test_zero_content_retention.py @@ -38,6 +38,12 @@ def test_invalid_config_is_off_without_logging_values(raw, caplog): assert raw not in caplog.text +@pytest.fixture(autouse=True) +def readable_controls(monkeypatch): + monkeypatch.setattr("services.auth_service.AuthService._load_user_by_username", lambda username: { + "username": username, "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}}) + + def app_fixture(monkeypatch): app = Flask(__name__) app.config["TESTING"] = True @@ -49,7 +55,8 @@ def app_fixture(monkeypatch): upstream = Mock() @app.before_request def auth(): - g.authenticated_user = {"id": "test-key", "username": "reader", "shadow_eval_rate": 0.2} + g.authenticated_user = {"id": "test-key", "username": "reader", "shadow_eval_rate": 0.2, + "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}} @app.post("/v1/chat/completions") @chat_cache.cached_chat_completion def complete(): @@ -115,7 +122,7 @@ def test_stream_callback_keeps_request_policy_after_context_and_config_change(mo monkeypatch.setattr(sampling, "submit", lambda sample, **kw: submitted.append((sample, kw))) events = [b'data: {"choices":[{"delta":{"content":"fixture"},"finish_reason":"stop"}]}\n\n', b'data: [DONE]\n\n'] with app.test_request_context("/v1/chat/completions", method="POST", json={"model":"auto:test", "messages":[{"content":"fixture"}]}): - g.authenticated_user = {"id":"test-key", "shadow_eval_rate":0.2} + g.authenticated_user = {"id":"test-key", "shadow_eval_rate":0.2, "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}} g.shadow_eval_started = 0 with patch.object(sampling.random, "random", return_value=0): response = sampling.sample_success(Response(iter(events), mimetype="text/event-stream", @@ -125,7 +132,7 @@ def test_stream_callback_keeps_request_policy_after_context_and_config_change(mo assert list(response.response) == events assert len(submitted) == 1 and submitted[0][1]["retention_policy"] == policy with app.test_request_context("/v1/chat/completions", method="POST", headers={"X-MultiLLM-Retention":"zero"}): - g.authenticated_user = {"id":"test-key", "shadow_eval_rate":0.2} + g.authenticated_user = {"id":"test-key", "shadow_eval_rate":0.2, "data_policy": {"version": 1, "official_only": None, "forced_provider_preferences": None}} untouched = iter(events) private = sampling.sample_success(Response(untouched, mimetype="text/event-stream")) assert private.response is untouched @@ -136,7 +143,7 @@ def test_background_persistence_uses_captured_policy(monkeypatch): monkeypatch.setattr(sampling, "_WORKER", Mock(is_alive=lambda: True)) enqueue = Mock() monkeypatch.setattr(sampling._QUEUE, "put_nowait", enqueue) - sampling.submit({"request": "fixture"}, retention_policy=policy) + sampling.submit({"key_id": "test-key", "request": "fixture"}, retention_policy=policy) queued = enqueue.call_args.args[0] monkeypatch.setenv("CONTENT_RETENTION_POLICY_JSON", '{"default":"zero"}') getter = Mock(side_effect=[queued, StopIteration]) @@ -144,7 +151,7 @@ def test_background_persistence_uses_captured_policy(monkeypatch): monkeypatch.setattr(sampling._QUEUE, "task_done", Mock()) with patch.object(sampling.ShadowEvalStore, "put") as persisted, pytest.raises(StopIteration): sampling._persist() - persisted.assert_called_once_with({"request": "fixture"}) + persisted.assert_called_once_with({"key_id": "test-key", "request": "fixture"}) from tests.unified_api_test_case import UnifiedApiTestCase diff --git a/worker/access-sso.mjs b/worker/access-sso.mjs index 0746a0d1..931215f9 100644 --- a/worker/access-sso.mjs +++ b/worker/access-sso.mjs @@ -1,3 +1,4 @@ +import { controlNetwork } from "./dispatch-network.mjs"; /** * Cloudflare Access single sign-on for the dashboard. On the sign-in path only, the Worker * verifies the Access application token (RS256 against the team's published keys, then @@ -120,9 +121,10 @@ async function boundedText(response, maxBytes) { } async function fetchSigningKeys(teamDomain, fetchImpl) { - const response = await fetchImpl(`${teamDomain}/cdn-cgi/access/certs`, { + const url = `${teamDomain}/cdn-cgi/access/certs`; + const response = await controlNetwork(url, () => fetchImpl(url, { headers: { accept: "application/json" }, redirect: "manual", signal: AbortSignal.timeout(KEYS_TIMEOUT_MS), - }); + })); if (response.status !== 200) { void response.body?.cancel().catch(() => {}); throw new Error(`Access keys answered HTTP ${response.status}`); diff --git a/worker/ai-outbound.mjs b/worker/ai-outbound.mjs index 63dadcaa..a6ba5e11 100644 --- a/worker/ai-outbound.mjs +++ b/worker/ai-outbound.mjs @@ -1,3 +1,4 @@ +import { bindingPrincipal, withDispatchPrincipal, providerFetch, providerOperation, dispatchFailure, DispatchError } from "./provider-dispatch.mjs"; /** * Cloudflare AI for the Container, reachable only through its private outbound handler * (`http://ai.internal`). Third-party models (OpenAI GPT Image, Google Veo) run through @@ -94,7 +95,7 @@ async function imageData(image, responseFormat) { if (typeof image !== "string" || !image) return null; if (/^https:\/\//i.test(image)) { if (responseFormat === "url") return { url: image }; - const response = await fetch(image, { signal: AbortSignal.timeout(30000) }); + const response = await providerFetch("cloudflare", image, { signal: AbortSignal.timeout(30000) }); const length = Number(response.headers.get("content-length") ?? 0); if (!response.ok || length > MAX_FETCHED_IMAGE_BYTES) throw new Error(`Image download failed with HTTP ${response.status}`); const bytes = await response.arrayBuffer(); @@ -106,7 +107,7 @@ async function imageData(image, responseFormat) { } async function run(env, model, input) { - const response = await env.AI.run(model, input, { gateway: { id: env.AI_GATEWAY_ID || "default" } }); + const response = await providerOperation("cloudflare", () => env.AI.run(model, input, { gateway: { id: env.AI_GATEWAY_ID || "default" } })); return isRecord(response?.result) ? response.result : response; } @@ -158,7 +159,7 @@ async function speech(env, body) { if (!encoding) return failure("invalid_request", "Unsupported response_format.", 400); const input = { text: body.input, speaker: AURA_SPEAKERS.has(body.voice) ? body.voice : "luna", encoding, ...(container ? { container } : {}) }; - const result = await env.AI.run(body.model, input, { gateway: { id: env.AI_GATEWAY_ID || "default" } }); + const result = await providerOperation("cloudflare", () => env.AI.run(body.model, input, { gateway: { id: env.AI_GATEWAY_ID || "default" } })); let audio = null; if (result instanceof ReadableStream) audio = result; else if (result instanceof Response) audio = result.body; @@ -195,7 +196,7 @@ async function generateVideo(env, body) { return Response.json({ status: "completed", video_url: result.video, model: body.model }, { headers: { "cache-control": "no-store" } }); } -export async function handleAiOutbound(request, env) { +async function runAiOutbound(request, env) { const url = new URL(request.url); if (url.origin !== "http://ai.internal" || url.search || url.hash || url.username || url.password || request.method !== "POST") { return failure("not_found", "Unknown Cloudflare AI operation.", 404); @@ -222,8 +223,15 @@ export async function handleAiOutbound(request, env) { if (url.pathname === "/v1/videos/generations") return await generateVideo(env, body); return failure("not_found", "Unknown Cloudflare AI operation.", 404); } catch (error) { + if (error instanceof DispatchError) return dispatchFailure(error); const [code, status] = refusal(error); logFailure("cloudflare_ai_failed", error, { model: typeof body.model === "string" ? body.model.slice(0, 80) : "unknown", code }); return failure(code, "Cloudflare AI could not complete the generation.", status); } } + +export async function handleAiOutbound(request, env) { + if (new URL(request.url).origin !== "http://ai.internal") return failure("not_found", "Unknown Cloudflare AI operation.", 404); + try { return await withDispatchPrincipal(env, bindingPrincipal(request), () => runAiOutbound(request, env)); } + catch (error) { return dispatchFailure(error); } +} diff --git a/worker/container-env.mjs b/worker/container-env.mjs index 6687e4fa..d18abb2a 100644 --- a/worker/container-env.mjs +++ b/worker/container-env.mjs @@ -113,6 +113,7 @@ const DIRECT_ENV_KEYS = [ "ADMISSION_AUTHORITY_URL", "FREE_ROUTE_EXTRA_PROVIDERS", "FREE_ROUTE_FREE_TIER_PROVIDERS", + "GEMINI_DATA_TIER", "FREE_ROUTE_PROVIDER_ORDER", "FREE_ROUTE_WORKERSAI_ACCOUNT_ID", "GROQ_API_KEY", diff --git a/worker/control-users-d1.mjs b/worker/control-users-d1.mjs index 530d7d62..8f8559c8 100644 --- a/worker/control-users-d1.mjs +++ b/worker/control-users-d1.mjs @@ -2,6 +2,9 @@ * Durable dashboard accounts (username, key hash and metadata) in D1, reachable only * through the Container's private outbound handler. Fixed statements, no client SQL. */ +import { dataPolicy } from "./dispatch-policy.mjs"; +import { KEY_CONTROLS, expansionFence } from "./key-controls-capability.mjs"; +import { validDataControls } from "./key-data-guards.mjs"; import { STREAM_CAP_FIELD, streamCostEnabled, validStreamCap } from "./stream-cost-breaker.mjs"; import { commitRevision, revisionSyncSettings } from "./config-revision.mjs"; import { logFailure } from "./log.mjs"; @@ -9,7 +12,7 @@ import { AUDIT_OPERATIONS, handleAuditOperation } from "./control-audit-d1.mjs"; export const USER_FIELDS = Object.freeze(["username", "api_key_hash", "api_key_prefix", "scopes", "is_admin", "created_at", "last_login", "last_used_at", "last_used_ip", "created_by", "rotated_at", "revoked_at", - "daily_budget_usd", "monthly_budget_usd", "allowed_models", "allowed_ips", "expires_at", "secret_scan_mode", "shadow_eval_rate"]); + "daily_budget_usd", "monthly_budget_usd", "allowed_models", "allowed_ips", "expires_at", "secret_scan_mode", "shadow_eval_rate", "official_only", "forced_provider_preferences"]); // Migration 0007 added the per-key controls. Until it is applied, reads return the // older columns with no controls, so deploying before migrating cannot lock keys out. const CONTROL_FIELDS = USER_FIELDS.slice(12); @@ -31,7 +34,7 @@ const list = (value, pattern) => value === null || (typeof value === "string" && const reply = (value, status = 200) => Response.json(value.error ? { version: 1, error: { code: value.error, message: "Account storage operation failed" } } : value, { status, headers: { "cache-control": "no-store" } }); -const fields = (body, names) => Object.keys(body).length === names.length && names.every(key => Object.hasOwn(body, key)); +const fields = (body, names) => names.every(key => Object.hasOwn(body, key)); export function streamUserFields(env) { return streamCostEnabled(env) ? [...USER_FIELDS, STREAM_CAP_FIELD] : USER_FIELDS; @@ -42,9 +45,12 @@ export function validUser(user) { if (!validStreamCap(user[STREAM_CAP_FIELD])) return false; user = Object.fromEntries(Object.entries(user).filter(([name]) => name !== STREAM_CAP_FIELD)); } + if (user && Object.hasOwn(user, "data_policy")) { + user = Object.fromEntries(Object.entries(user).filter(([name]) => name !== "data_policy")); + } if (user && !Object.hasOwn(user, "shadow_eval_rate")) user = { ...user, shadow_eval_rate: null }; if (user && !Object.hasOwn(user, "secret_scan_mode")) user = { ...user, secret_scan_mode: null }; - return user !== null && typeof user === "object" && !Array.isArray(user) && fields(user, USER_FIELDS) + return user !== null && typeof user === "object" && !Array.isArray(user) && fields(user, USER_FIELDS) && Object.keys(user).every(name => USER_FIELDS.includes(name)) && text(user.username, 128) && text(user.api_key_hash, 512) && text(user.api_key_prefix, 64) && typeof user.scopes === "string" && user.scopes.length <= 512 && /^[A-Za-z0-9:_.,-]*$/.test(user.scopes) && (user.is_admin === 0 || user.is_admin === 1) && text(user.created_at, 64) @@ -53,6 +59,8 @@ export function validUser(user) { && (user.secret_scan_mode === null || ["off", "observe", "redact", "block"].includes(user.secret_scan_mode)) && (user.shadow_eval_rate === null || (typeof user.shadow_eval_rate === "number" && Number.isFinite(user.shadow_eval_rate) && user.shadow_eval_rate >= 0 && user.shadow_eval_rate <= 0.2)) + && validDataControls(user) + && (user.forced_provider_preferences === null || typeof user.forced_provider_preferences === "string") && budget(user.daily_budget_usd) && budget(user.monthly_budget_usd) && list(user.allowed_models, MODEL_PATTERNS) && list(user.allowed_ips, IP_RANGES) && (user.expires_at === null || (text(user.expires_at, 64) && Number.isFinite(Date.parse(user.expires_at)))); @@ -139,11 +147,16 @@ async function read(query) { /** Run an account query with every column, or with the pre-0007 columns and no controls. */ async function readUsers(query, env = {}) { - if (streamCostEnabled(env)) return read(() => query(streamUserFields(env).join(", "))); - for (const end of [USER_FIELDS.length, USER_FIELDS.length - 1, USER_FIELDS.length - 2, 12]) { + const stamp = row => row && { ...row, data_policy: dataPolicy(row) }; + if (streamCostEnabled(env) || revisionSyncSettings(env).enabled) { + const rows = await read(() => query(streamUserFields(env).join(", "))); + return Array.isArray(rows) ? rows.map(stamp) : stamp(rows); + } + for (const end of [USER_FIELDS.length, USER_FIELDS.length - 1, USER_FIELDS.length - 2, USER_FIELDS.length - 3, USER_FIELDS.length - 4, 12]) { try { const rows = await read(() => query(USER_FIELDS.slice(0, end).join(", "))); - const defaults = row => row && { ...Object.fromEntries(USER_FIELDS.slice(end).map(name => [name, null])), ...row }; + const defaults = row => row && { ...Object.fromEntries(USER_FIELDS.slice(end).map(name => [name, null])), ...row, + data_policy: end === USER_FIELDS.length ? dataPolicy(row) : null }; return Array.isArray(rows) ? rows.map(defaults) : defaults(rows); } catch (error) { if (!missingColumn(error) || end === 12) throw error; @@ -152,9 +165,9 @@ async function readUsers(query, env = {}) { } } -function upsertStatement(db, user, fields) { +function upsertStatement(db, user, fields, policyWrite = false) { return db.prepare(`INSERT INTO control_users (${fields.join(", ")}) VALUES (${fields.map(() => "?").join(", ")}) - ON CONFLICT(username) DO UPDATE SET ${fields.slice(1).map(name => `${name}=excluded.${name}`).join(", ")}`) + ON CONFLICT(username) DO UPDATE SET ${fields.slice(1).filter(name => policyWrite || !["official_only", "forced_provider_preferences"].includes(name)).map(name => `${name}=excluded.${name}`).join(", ")}`) .bind(...fields.map(name => user[name])); } @@ -206,6 +219,13 @@ export async function handleControlUsersRequest(request, env) { if (!body || Array.isArray(body) || typeof body !== "object" || body.version !== 1) return reply({ error: "invalid_request" }, 400); } catch { return reply({ error: "invalid_request" }, 400); } const db = env.INTELLIGENCE_DB; + const policyVersion = body.data_policy_version === 1 ? 1 : null; + delete body.data_policy_version; + if (body.operation === "capabilities") return reply({ version: 1, key_controls: KEY_CONTROLS }); + const accountReply = user => { + if (!user || policyVersion === 1) return user; + return Object.fromEntries(Object.entries(user).filter(([name]) => !["official_only", "forced_provider_preferences", "data_policy"].includes(name))); + }; try { if (AUDIT_OPERATIONS.has(body.operation)) { const result = await handleAuditOperation(db, body); @@ -219,23 +239,28 @@ export async function handleControlUsersRequest(request, env) { ? db.prepare(`SELECT ${columns} FROM control_users ORDER BY username LIMIT ?`).bind(body.limit) : db.prepare(`SELECT ${columns} FROM control_users WHERE username > ? ORDER BY username LIMIT ?`).bind(body.after, body.limit) ).all()).results; - const users = revisionSyncSettings(env).enabled ? await read(() => query(streamUserFields(env).join(", "))) : await readUsers(query, env); - return reply({ version: 1, users }); + const users = await readUsers(query, env); + return reply({ version: 1, users: users.map(user => accountReply(user, policyVersion)) }); } case "get": { if (!fields(body, ["version", "operation", "username"]) || !text(body.username, 128)) break; const user = await readUsers(columns => db.prepare(`SELECT ${columns} FROM control_users WHERE username=?`) .bind(body.username).first(), env); - return reply({ version: 1, user: user ?? null }); + if (policyVersion !== 1 && user) expansionFence(user); + return reply({ version: 1, user: accountReply(user ?? null, policyVersion) }); } case "by_prefix": { if (!fields(body, ["version", "operation", "prefix"]) || !text(body.prefix, 64)) break; - return reply({ version: 1, users: await activeUsersByPrefix(db, body.prefix, env) }); + const users = await activeUsersByPrefix(db, body.prefix, env); + if (policyVersion !== 1) users.forEach(expansionFence); + return reply({ version: 1, users: users.map(accountReply) }); } case "upsert": { - if (!fields(body, ["version", "operation", "user"]) || !validUser(body.user)) break; + if (!fields(body, ["version", "operation", "user"])) break; + body.user = { official_only: null, forced_provider_preferences: null, ...body.user }; + if (!validUser(body.user)) break; if (!streamCostEnabled(env) && body.user[STREAM_CAP_FIELD] != null) return reply({ error: "invalid_request" }, 400); - const user = { shadow_eval_rate: null, secret_scan_mode: null, ...body.user }; + const user = { official_only: null, forced_provider_preferences: null, shadow_eval_rate: null, secret_scan_mode: null, ...body.user }; if (grantsAdmin(user) && !adminUsernames(env).has(user.username)) { logFailure("account_admin_refused", new Error("Administration is not configured for this username")); await audit(db, "upsert", "refused", user).run(); @@ -245,17 +270,17 @@ export async function handleControlUsersRequest(request, env) { // Verify schema even when an omitted cap must be preserved. await db.prepare("SELECT max_stream_cost_microusd FROM control_users LIMIT 1").first(); const columns = Object.hasOwn(user, STREAM_CAP_FIELD) ? streamUserFields(env) : USER_FIELDS; - const statements = [upsertStatement(db, user, columns), audit(db, "upsert", "stored", user)]; + const statements = [upsertStatement(db, user, columns, policyVersion === 1), audit(db, "upsert", "stored", user)]; if (revisionSyncSettings(env).enabled) { if (!await commitRevision(db, ["key_controls", "model_grants"], statements)) return reply({ error: "revision_conflict" }, 409); } else await db.batch(statements); return reply({ version: 1, stored: true }); } // The write and its audit row commit together, or neither does. - for (const end of [USER_FIELDS.length, USER_FIELDS.length - 1, USER_FIELDS.length - 2, 12]) { + for (const end of [USER_FIELDS.length, USER_FIELDS.length - 1, USER_FIELDS.length - 2, USER_FIELDS.length - 3, USER_FIELDS.length - 4, 12]) { if (USER_FIELDS.slice(end).some(name => user[name] !== null)) throw new Error("Account controls require migration"); try { - const statements = [upsertStatement(db, user, USER_FIELDS.slice(0, end)), audit(db, "upsert", "stored", user)]; + const statements = [upsertStatement(db, user, USER_FIELDS.slice(0, end), policyVersion === 1), audit(db, "upsert", "stored", user)]; if (revisionSyncSettings(env).enabled) { if (!await commitRevision(db, ["key_controls", "model_grants"], statements)) return reply({ error: "revision_conflict" }, 409); } else await db.batch(statements); diff --git a/worker/dispatch-network.mjs b/worker/dispatch-network.mjs new file mode 100644 index 00000000..2aa5821b --- /dev/null +++ b/worker/dispatch-network.mjs @@ -0,0 +1,32 @@ +/** Network capabilities last only for one admitted transport operation. */ +import { AsyncLocalStorage } from "node:async_hooks"; +const network = new AsyncLocalStorage(); +export const networkAdmitted = input => { + const scope = network.getStore(); + if (!scope?.active) return false; + if (!scope.url) return true; + if (!input) return false; + try { return new URL(typeof input === "string" || input instanceof URL ? input : input.url, "https://gateway.fixture").href === scope.url; } + catch { return false; } +}; +export const connectionAdmitted = (host, port) => { + const scope = network.getStore(); + if (!scope?.active) return false; + if (!scope.url) return true; + const target = new URL(scope.url); + return host === target.hostname && Number(port) === Number(target.port || (target.protocol === "https:" ? 443 : 80)); +}; +function scoped(scope, run) { + return network.run(scope, () => { + try { + const result = run(); + if (result && typeof result.then === "function") return Promise.resolve(result).finally(() => { scope.active = false; }); + scope.active = false; + return result; + } catch (error) { scope.active = false; throw error; } + }); +} +export const admittedNetwork = run => scoped({ active: true }, run); +/** Certificate discovery, configured content-free alerts and browser test clients + * are control traffic. This scope authorizes only their exact destination. */ +export const controlNetwork = (url, run) => scoped({ active: true, url: new URL(url, "https://gateway.fixture").href }, run); diff --git a/worker/dispatch-policy.mjs b/worker/dispatch-policy.mjs new file mode 100644 index 00000000..ca37f86a --- /dev/null +++ b/worker/dispatch-policy.mjs @@ -0,0 +1,14 @@ +/** A supported version and both explicit fields are mandatory, including nulls. */ +import { parsePreferences, validDataControls } from "./key-data-guards.mjs"; + +export function dataPolicy(user) { + if (!validDataControls(user)) throw new TypeError("Unverifiable data policy"); + return { version: 1, official_only: user.official_only, + forced_provider_preferences: parsePreferences(user.forced_provider_preferences) }; +} +export function verifyDataPolicy(policy) { + if (!policy || policy.version !== 1 || Object.keys(policy).sort().join() !== "forced_provider_preferences,official_only,version") { + throw new TypeError("Unverifiable data policy"); + } + return dataPolicy(policy); +} diff --git a/worker/gateway-extensions.mjs b/worker/gateway-extensions.mjs index 93c39ad0..f5963ce6 100644 --- a/worker/gateway-extensions.mjs +++ b/worker/gateway-extensions.mjs @@ -1,3 +1,5 @@ +import { DispatchError, providerFetch, withDispatchPrincipal, dispatchFailure } from "./provider-dispatch.mjs"; +import { guardNativeRequest } from "./key-data-guards.mjs"; import { governanceEnabled, GovernanceError } from "./tenant-governance-d1.mjs"; import { creditsEnforcement } from "./credits-d1.mjs"; import { CreditsAdmissionError } from "./credits-admission.mjs"; @@ -332,7 +334,20 @@ async function runNativeLifecycle(request, env, ctx, authority, fetcher, registr } export async function nativeGenerationFetch(request, env, ctx, authority, fetcher, collaborators = []) { + const send = fetcher; + fetcher = async (outbound, runtime, current) => { + try { return await providerFetch(current.provider, outbound, {}, { env: runtime, authority: current, + send: checked => send(checked, runtime, current) }); } + catch (error) { if (!(error instanceof DispatchError)) throw error; return dispatchFailure(error); } + }; + return withDispatchPrincipal(env, authority, () => runNativeGenerationFetch(request, env, ctx, authority, fetcher, collaborators)); +} + +async function runNativeGenerationFetch(request, env, ctx, authority, fetcher, collaborators = []) { const path = authority.route ?? new URL(request.url).pathname; + const guarded = await guardNativeRequest(request, env, authority); + if (guarded.response) { authority.deadlineHook?.deadline?.stop(); return guarded.response; } + request = guarded.request; if (!generationPath(request, path)) return fetcher(request, env, authority); let enterprise = authority.authenticated === true && (governanceEnabled(env) || creditsEnforcement(env) !== "off"); if ((organisationsEnabled(env) || enterprise) && authority.authenticated === true) { diff --git a/worker/intelligence-auth-d1.mjs b/worker/intelligence-auth-d1.mjs index 79020c11..44c68e14 100644 --- a/worker/intelligence-auth-d1.mjs +++ b/worker/intelligence-auth-d1.mjs @@ -1,3 +1,5 @@ +import { expansionFence } from "./key-controls-capability.mjs"; +import { dataPolicy } from "./dispatch-policy.mjs"; import { logFailure } from "./log.mjs"; export const INTEGRATION_SCOPES = Object.freeze(["chat", "models", "audio", "embeddings", @@ -40,6 +42,11 @@ export async function lookupIntegrationPrincipal(db, keyPrefix) { c.key_prefix, c.key_hash FROM intelligence_credentials c JOIN intelligence_principals p ON p.id=c.principal_id AND p.version=c.version WHERE c.key_prefix=?`).bind(keyPrefix).first(); + if (row && !row.revoked_at) { + const policy = await db.prepare("SELECT official_only, forced_provider_preferences FROM intelligence_principals WHERE id = ? AND revoked_at IS NULL").bind(row.id).first(); + if (!policy) throw Error("key_data_controls_unavailable"); + dataPolicy(policy); expansionFence(policy); + } return row ? { id: row.id, scopes: JSON.parse(row.scopes), credentialVersion: row.version, createdAt: row.created_at, revokedAt: row.revoked_at, @@ -74,6 +81,14 @@ export async function handleIntelligenceAuthRequest(request, env) { return reply({version: 1, principal: await lookupIntegrationPrincipal(db, body.keyPrefix)}); } if (typeof body.principalId !== "string" || !idPattern.test(body.principalId)) return reply({error: "invalid_request"}, 400); + if (body.operation === "dispatch_policy") { + if (!fields(body, ["version", "operation", "principalId"])) return reply({error: "invalid_request"}, 400); + const row = await db.prepare(`SELECT p.official_only, p.forced_provider_preferences, c.key_prefix + FROM intelligence_principals p JOIN intelligence_credentials c ON c.principal_id=p.id AND c.version=p.version + WHERE p.id=? AND p.revoked_at IS NULL`).bind(body.principalId).first(); + return reply({version: 1, principal: row ? { username: body.principalId, api_key_prefix: row.key_prefix, + data_policy: dataPolicy(row) } : null}); + } const now = new Date().toISOString(); if (body.operation === "provision") { if (!fields(body, ["version", "operation", "principalId", "keyPrefix", "keyHash", "scopes"]) diff --git a/worker/key-controls-capability.mjs b/worker/key-controls-capability.mjs new file mode 100644 index 00000000..74ae012b --- /dev/null +++ b/worker/key-controls-capability.mjs @@ -0,0 +1,9 @@ +/** Code capability, never an operator feature flag. */ +export const KEY_CONTROLS = Object.freeze({ protocol: 1, enforcement: false }); +export const controlled = row => row?.official_only != null || row?.forced_provider_preferences != null; +export function expansionFence(row) { + if (!KEY_CONTROLS.enforcement && controlled(row)) { + const error = new Error("key_controls_require_upgrade"); error.code = "key_controls_require_upgrade"; error.status = 503; throw error; + } + return row; +} diff --git a/worker/key-data-guards.mjs b/worker/key-data-guards.mjs new file mode 100644 index 00000000..6fba1cce --- /dev/null +++ b/worker/key-data-guards.mjs @@ -0,0 +1,90 @@ +import { KEY_CONTROLS, controlled, expansionFence } from "./key-controls-capability.mjs"; +/** Shared provider classes and administrator-owned preferences at native dispatch. */ +import classes from "../providers/data-classes.json" with { type: "json" }; + +const object = value => value !== null && typeof value === "object" && !Array.isArray(value); +export function parsePreferences(value) { + if (value == null) return null; + const parsed = typeof value === "string" ? JSON.parse(value) : value; + if (!object(parsed) || new TextEncoder().encode(JSON.stringify(parsed)).length > 4096 + || (Object.hasOwn(parsed, "data_collection") && !["allow", "deny"].includes(parsed.data_collection))) { + throw new TypeError("Invalid provider preferences"); + } + return parsed; +} + +export function validDataControls(user) { + try { + if (!user || !Object.hasOwn(user, "official_only") || !Object.hasOwn(user, "forced_provider_preferences") + || user.official_only === undefined || user.forced_provider_preferences === undefined) return false; + parsePreferences(user.forced_provider_preferences); + return user.official_only === null || [0, 1].includes(user.official_only); + } catch { return false; } +} + +export function providerDataClass(provider, env = {}) { + provider = String(provider).trim().toLowerCase(); + if (provider === "free" || String(env.FREE_ROUTE_FREE_TIER_PROVIDERS ?? "").split(/[\s,]+/).includes(provider)) return "free"; + if (provider === "gemini") return env.GEMINI_DATA_TIER === "paid" ? "official" : "free"; + return Object.hasOwn(classes, provider) ? classes[provider] : provider.startsWith("ce-") ? "subscription" : "reseller"; +} + +export function providerClassPermitted(user, provider, env = {}) { + return !user?.official_only || providerDataClass(provider, env) === "official"; +} + +export function mergePreferences(caller, forced) { + const entries = new Map(Object.entries(object(caller) ? caller : {})); + for (const [name, value] of Object.entries(forced)) { + if (name === "data_collection" && entries.get(name) === "deny") continue; + entries.set(name, object(value) ? mergePreferences(entries.get(name), value) : structuredClone(value)); + } + return Object.fromEntries(entries); +} + +/** Read current administrator controls without accepting policy from client headers. */ +export async function currentDataControls(env, authority) { + const id = authority.keyId ?? authority.principal?.owner ?? authority.principal?.id; + if (!env.INTELLIGENCE_DB || !id) throw new TypeError("Unverifiable authority"); + const user = id.startsWith("integration:") + ? await env.INTELLIGENCE_DB.prepare("SELECT official_only, forced_provider_preferences FROM intelligence_principals WHERE id = ?" + (KEY_CONTROLS.enforcement ? " AND revoked_at IS NULL" : "")).bind(id).first() + : await env.INTELLIGENCE_DB.prepare( + "SELECT official_only, forced_provider_preferences FROM control_users WHERE username = ?" + (KEY_CONTROLS.enforcement + ? " AND revoked_at IS NULL AND (expires_at IS NULL OR expires_at > ?)" : "")) + .bind(...(KEY_CONTROLS.enforcement ? [id, new Date().toISOString()] : [id])).first(); + if (!validDataControls(user)) throw new TypeError("Unverifiable controls"); + return expansionFence(user); +} + +/** Resolve controls from trusted authority or durable accounts, never client headers. */ +export async function guardNativeRequest(request, env, authority) { + if (authority.authenticated !== true) return { response: Response.json({ error: { code: "key_data_controls_unavailable" } }, { status: 503 }) }; + try { + const user = await currentDataControls(env, authority); + if (!providerClassPermitted(user, authority.provider, env)) { + return { response: Response.json({ error: { code: "provider_class_not_allowed" } }, { status: 403 }) }; + } + const forced = parsePreferences(user.forced_provider_preferences); + if (authority.provider !== "openrouter" || !forced || !request.body) return { request }; + const payload = await request.clone().json(); + if (!object(payload)) return { response: Response.json({ error: { code: "invalid_request" } }, { status: 400 }) }; + payload.provider = mergePreferences(payload.provider, forced); + const headers = new Headers(request.headers); headers.delete("content-length"); + return { request: new Request(request, { headers, body: JSON.stringify(payload) }) }; + } catch { + return { response: Response.json({ error: { code: "key_data_controls_unavailable" } }, { status: 503 }) }; + } +} + +/** Legacy persisted work has no identity. It is safe only before any controls exist. */ +export async function legacyDataControls(env) { + if (!env?.INTELLIGENCE_DB) throw new TypeError("Unverifiable legacy controls"); + const rows = await env.INTELLIGENCE_DB.prepare( + "SELECT official_only, forced_provider_preferences FROM control_users UNION ALL SELECT official_only, forced_provider_preferences FROM intelligence_principals").all(); + if (!Array.isArray(rows?.results)) throw new TypeError("Unverifiable legacy controls"); + for (const row of rows.results) { + if (!validDataControls(row)) throw new TypeError("Unverifiable legacy controls"); + if (controlled(row)) { const error = new TypeError("Legacy controlled accounts"); error.code = "legacy_controlled_accounts"; throw error; } + } + return { official_only: null, forced_provider_preferences: null }; +} diff --git a/worker/knowledge-edge.mjs b/worker/knowledge-edge.mjs index d5863ea9..5348c6ce 100644 --- a/worker/knowledge-edge.mjs +++ b/worker/knowledge-edge.mjs @@ -1,3 +1,4 @@ +import { currentDataControls } from "./key-data-guards.mjs"; /** * Knowledge MCP and REST at the edge for credentials the Worker can verify itself: * the bootstrap ADMIN_API_KEY, durable D1 integration keys and, when accounts live in @@ -197,7 +198,7 @@ async function accountPrincipal(env, key, clientAddress) { * Returns {principal}, {denied: true} for a key that is definitively invalid, or null * when the Container must decide (unverifiable key types, storage or data faults). */ -async function resolvePrincipal(request, env) { +export async function resolvePrincipal(request, env) { const key = requestKey(request.headers); if (!key) return null; if (env.ADMIN_API_KEY && sameSecret(key, env.ADMIN_API_KEY)) { @@ -283,6 +284,7 @@ async function dispatch(env, operation, principal, payload, signal) { if (Buffer.byteLength(body) > (operation === "skills.sync" ? SYNC_REQUEST_BYTES : MAX_REQUEST_BYTES)) { throw new KnowledgeEdgeError("request_too_large", "The Knowledge request exceeds 64 KiB.", 413); } + await currentDataControls(env, { keyId: principal.id }); const deadline = AbortSignal.timeout(DEADLINE_MS); let response; let result; diff --git a/worker/knowledge-outbound.mjs b/worker/knowledge-outbound.mjs index 2ebed15a..416f335c 100644 --- a/worker/knowledge-outbound.mjs +++ b/worker/knowledge-outbound.mjs @@ -9,6 +9,7 @@ export async function handleKnowledgeOutbound(request, env) { if (!env.KNOWLEDGE_SERVICE) fail("knowledge_not_configured", "The private Knowledge service binding is not configured.", 503); const payload = await readJson(request, SYNC_REQUEST_BYTES); if (payload.operation !== "skills.sync" && new TextEncoder().encode(JSON.stringify(payload)).length > 65536) fail("request_too_large", "Request is too large.", 413); + await (await import("./key-data-guards.mjs")).currentDataControls(env, { keyId: payload.principal?.id }); return env.KNOWLEDGE_SERVICE.fetch("http://knowledge.internal/v1/dispatch", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(payload), signal: request.signal, }); diff --git a/worker/knowledge/corpus.mjs b/worker/knowledge/corpus.mjs index 59d711e9..8e67d7a6 100644 --- a/worker/knowledge/corpus.mjs +++ b/worker/knowledge/corpus.mjs @@ -1,3 +1,4 @@ +import { providerOperation } from "../provider-dispatch.mjs"; const MAX_SNAPSHOT_BYTES = 256 * 1024; const MAX_INDEX_RESULTS = 20; const INDEX_PAGE_SIZE = 50; @@ -91,7 +92,7 @@ function itemInfo(value, key) { async function readIndex(callback) { for (let attempt = 0; ; attempt += 1) { - try { return await callback(); } + try { return await providerOperation("cloudflare", callback); } catch (error) { // The binding intermittently rejects its own account context before a metadata read. // Uploads, deletions and billable search requests never use this retry path. @@ -184,13 +185,13 @@ export class KnowledgeCorpus { if (typeof request.product === "string" && request.product && encoder.encode(request.product).byteLength <= 64) retrieval.filters = { product: request.product }; const response = this.searchAI ? await this.searchWithAI(request, retrieval.filters) - : await this.requireIndex().search({ + : await providerOperation("cloudflare", () => this.requireIndex().search({ query: request.query, ai_search_options: { retrieval, query_rewrite: { enabled: false }, reranking: { enabled: false }, cache: { enabled: false }, }, - }); + })); if (!response || !Array.isArray(response.chunks) || response.chunks.length > MAX_INDEX_RESULTS) { throw invalid("invalid_index_response"); } @@ -211,23 +212,23 @@ export class KnowledgeCorpus { } // Select one transport before dispatch. Never retry a billable query through // another binding when the first request has an unknown outcome. - const response = await this.searchAI.autorag(this.searchInstance).search({ + const response = await providerOperation("cloudflare", () => this.searchAI.autorag(this.searchInstance).search({ query: request.query, rewrite_query: false, max_num_results: MAX_INDEX_RESULTS, reranking: { enabled: false }, ...(filters ? { filters: { type: "eq", key: "product", value: filters.product } } : {}), - }); + })); return legacySearchChunks(response); } async uploadRevision(artifact, text) { await checkedBytes(artifact, text); - const result = await this.requireIndex().items.upload(artifact.index_key, text, { + const result = await providerOperation("cloudflare", () => this.requireIndex().items.upload(artifact.index_key, text, { metadata: { artifact_id: artifact.id, content_hash: artifact.content_hash, product: artifact.product, version_kind: artifact.version.kind, version: artifact.version.version ?? "unknown", }, - }); + })); return itemReference(result, artifact.index_key); } @@ -286,7 +287,7 @@ export class KnowledgeCorpus { if (artifact.item_id || this.index?.items) { const item = await this.reconcileRevision(artifact); if (item && artifact.item_id && item.id !== artifact.item_id) throw invalid("invalid_artifact_removal"); - if (item) await this.requireIndex().items.delete(item.id); + if (item) await providerOperation("cloudflare", () => this.requireIndex().items.delete(item.id)); } await this.requireBucket().delete(artifact.snapshot_key); } diff --git a/worker/knowledge/index.mjs b/worker/knowledge/index.mjs index 48642ec7..f97bb9c5 100644 --- a/worker/knowledge/index.mjs +++ b/worker/knowledge/index.mjs @@ -1,3 +1,6 @@ +import { currentDataControls } from "../key-data-guards.mjs"; +import { KEY_CONTROLS } from "../key-controls-capability.mjs"; +import { bindingPrincipal, withDispatchPrincipal } from "../provider-dispatch.mjs"; import { DurableObject, WorkflowEntrypoint } from "cloudflare:workers"; import { KnowledgeAuthority } from "./authority.mjs"; import { dispatchKnowledge, maintainKnowledge } from "./service.mjs"; @@ -73,6 +76,7 @@ export class KnowledgeHandoffs extends DurableObject { export class KnowledgeSkills extends DurableObject { constructor(ctx, env) { super(ctx, env); + this.env = env; this.skills = new SkillsStore(ctx.storage, env); } async fetch(request) { @@ -81,23 +85,30 @@ export class KnowledgeSkills extends DurableObject { const body = await readJson(request, SYNC_REQUEST_BYTES); fields(body, ["operation", "payload", "principal"], ["operation", "payload", "principal"]); if (typeof body.principal !== "string" || !body.principal || body.principal.length > 256) fail("invalid_principal", "Invalid skills principal.", 403); - return reply(await this.skills.call(body.operation, body.payload, body.principal)); + const source = bindingPrincipal(request, body.principal); + if (source.keyId !== body.principal) fail("invalid_principal", "The dispatch principal does not match.", 403); + await currentDataControls(this.env, source); + return await withDispatchPrincipal(this.env, bindingPrincipal(request, body.principal), async () => reply(await this.skills.call(body.operation, body.payload, body.principal))); } catch (error) { return failure("knowledge_skills_failed", error); } } } export class KnowledgeIngestion extends WorkflowEntrypoint { - async run(event, step) { return runIngestion(this.env, step, event.payload.job_id); } + async run(event, step) { return withDispatchPrincipal(this.env, event.payload.dispatch_principal, () => runIngestion(this.env, step, event.payload.job_id)); } } export default { async fetch(request, env, ctx) { try { const url = new URL(request.url); + if (request.method === "GET" && url.origin === "http://knowledge.internal" && url.pathname === "/ready") return Response.json({ key_controls: KEY_CONTROLS }); if (request.method !== "POST" || url.origin !== "http://knowledge.internal" || url.pathname !== "/v1/dispatch" || url.search || url.hash || url.username || url.password) fail("not_found", "Unknown Knowledge route.", 404); const body = await readJson(request, SYNC_REQUEST_BYTES); if (body.operation !== "skills.sync" && new TextEncoder().encode(JSON.stringify(body)).length > 65536) fail("request_too_large", "Request is too large.", 413); + const source = bindingPrincipal(request, body.principal?.id); + if (source.keyId !== body.principal?.id) fail("invalid_principal", "The dispatch principal does not match.", 403); + await currentDataControls(env, source); let decision; const result = await dispatchKnowledge(env, body, { signal: request.signal, waitUntil: promise => ctx.waitUntil(promise), onSecretScan: value => { decision = value; } }); const response = reply(result); @@ -108,6 +119,6 @@ export default { async scheduled(_event, env, ctx) { ctx.waitUntil(maintainKnowledge(env).catch(error => { logFailure("knowledge_maintenance_failed", error); throw error; })); // Separate so an upstream check failure cannot stop catalogue maintenance. - ctx.waitUntil(watchImportedSkills(env).catch(error => logFailure("knowledge_skills_watch_failed", error))); + ctx.waitUntil(withDispatchPrincipal(env, { legacyBinding: true }, () => watchImportedSkills(env)).catch(error => logFailure("knowledge_skills_watch_failed", error))); }, }; diff --git a/worker/knowledge/ingestion.mjs b/worker/knowledge/ingestion.mjs index 259d1081..e20d80fb 100644 --- a/worker/knowledge/ingestion.mjs +++ b/worker/knowledge/ingestion.mjs @@ -1,3 +1,4 @@ +import { dispatchContext, withDispatchPrincipal } from "../provider-dispatch.mjs"; import { publicUrl } from "./contracts.mjs"; import { createArtifact, normalizeSourceText } from "./evidence.mjs"; import { KnowledgeCorpus } from "./corpus.mjs"; @@ -24,7 +25,9 @@ const REFUSED_ERRORS = new Set(["provider_disabled", "allowance_exhausted", "led // false "Worker's code had hung" cancellation (cloudflare/workers-sdk#14959). Results cached // as objects by instances started before this change are still accepted. async function durableStep(step, name, config, callback) { - const value = await step.do(name, config, async () => JSON.stringify(await callback())); + const state = dispatchContext(); + const value = await step.do(name, config, () => withDispatchPrincipal(state?.env, state?.authority, + async () => JSON.stringify(await callback()))); return typeof value === "string" ? JSON.parse(value) : value; } diff --git a/worker/knowledge/memos.mjs b/worker/knowledge/memos.mjs index 3aeb6797..9ab765ae 100644 --- a/worker/knowledge/memos.mjs +++ b/worker/knowledge/memos.mjs @@ -1,3 +1,4 @@ +import { providerOperation } from "../provider-dispatch.mjs"; import { fail } from "./contracts.mjs"; import { resolveRetentionPolicy, retentionAllowsContent } from "../retention-policy.mjs"; import { digest } from "./evidence.mjs"; @@ -65,7 +66,7 @@ class MemoSession { usage.push(entry); try { const result = await bounded(() => options.embed ? options.embed(request.query) - : env.KNOWLEDGE_SEARCH_AI.run("@cf/baai/bge-m3", { text: [request.query] }), this.deadline(MEMO_TIMEOUT_MS)); + : providerOperation("cloudflare", () => env.KNOWLEDGE_SEARCH_AI.run("@cf/baai/bge-m3", { text: [request.query] })), this.deadline(MEMO_TIMEOUT_MS)); const vector = Array.isArray(result) ? result : result?.data?.[0]; const quantized = quantizeEmbedding(vector); if (quantized) entry.outcome = "completed"; diff --git a/worker/knowledge/providers/transport.mjs b/worker/knowledge/providers/transport.mjs index 36a223b4..ceab3047 100644 --- a/worker/knowledge/providers/transport.mjs +++ b/worker/knowledge/providers/transport.mjs @@ -1,3 +1,4 @@ +import { DispatchError, providerFetch, transportOptions } from "../../provider-dispatch.mjs"; import { invalidResponse, ProviderError, upstreamError } from "./errors.mjs"; import { withProviderKey } from "./keys.mjs"; @@ -69,7 +70,7 @@ export async function fetchBounded(provider, url, options, { fetchImpl = fetch, try { if (controller.signal.aborted) throw new Error("aborted"); // Workers supports manual redirects; the non-2xx check rejects them before any follow-up. - const response = await fetchImpl(url, { ...options, redirect: "manual", signal: controller.signal }); + const response = await providerFetch(provider, url, { ...options, redirect: "manual", signal: controller.signal }, { send: async checked => fetchImpl(checked.url, await transportOptions(checked, options)) }); if (response.status < 200 || response.status >= 300) { let body = null; if ([402, 429].includes(response.status) || (provider === "alexandria" && response.status === 403)) { @@ -83,7 +84,7 @@ export async function fetchBounded(provider, url, options, { fetchImpl = fetch, const text = await readBounded(response, provider, maxResponseBytes); return parseResponse(text, response.headers.get("content-type") || "", provider, acceptText); } catch (error) { - if (error instanceof ProviderError) throw error; + if (error instanceof ProviderError || error instanceof DispatchError) throw error; if (controller.signal.aborted) { throw new ProviderError(provider, "provider_timeout", "The knowledge provider request was cancelled or timed out.", 504); } diff --git a/worker/knowledge/service.mjs b/worker/knowledge/service.mjs index 37d6388d..5a4e7613 100644 --- a/worker/knowledge/service.mjs +++ b/worker/knowledge/service.mjs @@ -1,3 +1,4 @@ +import { withDispatchPrincipal } from "../provider-dispatch.mjs"; import { authorize, fail, fields, parseQuery, publicUrl, validId } from "./contracts.mjs"; import { getAuthority } from "./authority-client.mjs"; import { KnowledgeCorpus } from "./corpus.mjs"; @@ -101,7 +102,7 @@ async function artifactResult(env, authority, id, corpus) { return { artifact, text }; } -export async function dispatchKnowledge(env, envelope, options = {}) { +async function runKnowledge(env, envelope, options = {}) { fields(envelope, ["version", "operation", "principal", "payload", "secret_scan_mode", "secret_scan_checked", "retention_policy"], ["version", "operation", "principal", "payload"]); if (envelope.version !== 1 || !OPERATIONS.has(envelope.operation)) fail("unknown_operation", "Unknown Knowledge operation.", 404); const { operation, principal, payload } = envelope; @@ -165,7 +166,10 @@ export async function maintainKnowledge(env, { authority = getAuthority(env), co for (const artifact of await authority.call("artifacts.expired")) { try { const claimed = await authority.call("artifact.expiration_claim", { id: artifact.id, expires_at: artifact.expires_at }); - await corpus.removeArtifact(claimed); + // Fixed storage cleanup uses the current durable operator account; ingestion + // and reconciliation Workflows continue to carry the original source account. + await withDispatchPrincipal(env, { legacyBinding: true }, + () => corpus.removeArtifact(claimed)); await authority.call("artifact.expire", { id: artifact.id, expires_at: artifact.expires_at }); } catch (error) { // The claimed revision stays unreadable until a later run removes it. @@ -180,3 +184,7 @@ export async function maintainKnowledge(env, { authority = getAuthority(env), co .catch(error => logFailure("knowledge_reconcile_not_scheduled", error, { job_id: job.id })); } } + +export async function dispatchKnowledge(env, envelope, options = {}) { + return withDispatchPrincipal(env, { authenticated: true, keyId: envelope?.principal?.id }, () => runKnowledge(env, envelope, options)); +} diff --git a/worker/knowledge/skills-import.mjs b/worker/knowledge/skills-import.mjs index 2af33705..5e8c46bd 100644 --- a/worker/knowledge/skills-import.mjs +++ b/worker/knowledge/skills-import.mjs @@ -1,3 +1,4 @@ +import { providerFetch, transportOptions, withDispatchPrincipal } from "../provider-dispatch.mjs"; import { fail, fields, integer } from "./contracts.mjs"; import { digest } from "./evidence.mjs"; import { CLAWHUB, clawhubFile, directory, frontmatter, REF, repository, REVIEW_PATTERNS, reviewFlags } from "./skills-market.mjs"; @@ -62,7 +63,7 @@ async function get(url, headers, options, maxBytes) { const timer = setTimeout(cancel, options.importTimeoutMs ?? TIMEOUT_MS); try { let response; - try { response = await (options.fetch ?? fetch)(url, { headers, redirect: "manual", signal: controller.signal }); } + try { response = await providerFetch("skills", url, { headers, redirect: "manual", signal: controller.signal }, { send: async checked => (options.fetch ?? fetch)(checked.url, await transportOptions(checked, { headers })) }); } catch { fail(controller.signal.aborted ? "upstream_timeout" : "upstream_unavailable", "The skill source could not be reached.", controller.signal.aborted ? 504 : 502); } if (!response.ok) { await response.body?.cancel(); @@ -338,7 +339,7 @@ export async function checkImports(env, store, options = {}, { force = false, li const due = await store.call("imports.due", { limit, force }); for (const item of due) { let result; - try { result = await checkOne(env, item, options); } + try { result = await withDispatchPrincipal(env, item.origin?.dispatch_principal, () => checkOne(env, item, options)); } catch (error) { result = { status: "error", error: error?.code ?? "upstream_unavailable" }; } await store.call("imports.checked", { skill_id: item.skill_id, pin: pinOf(item.origin), result }); } diff --git a/worker/knowledge/skills-store.mjs b/worker/knowledge/skills-store.mjs index 701157c7..a4d19bf2 100644 --- a/worker/knowledge/skills-store.mjs +++ b/worker/knowledge/skills-store.mjs @@ -1,3 +1,4 @@ +import { providerOperation } from "../provider-dispatch.mjs"; import { fail, fields, isRecord } from "./contracts.mjs"; import { digest } from "./evidence.mjs"; import { parseFind, parseGet, parseSync, validateSkill, skillKey, SKILLS_LIMIT } from "./skills-validation.mjs"; @@ -22,7 +23,7 @@ async function bounded(callback, milliseconds) { export class SkillsStore { constructor(storage, env, { embed, clock = Date.now } = {}) { Object.assign(this, { storage, sql: storage.sql, env, clock }); - this.embed = embed ?? (text => env.KNOWLEDGE_SEARCH_AI?.run("@cf/baai/bge-m3", { text: [text] })); + this.embed = embed ?? (text => providerOperation("cloudflare", () => env.KNOWLEDGE_SEARCH_AI?.run("@cf/baai/bge-m3", { text: [text] }))); this.index = null; this.syncTail = Promise.resolve(); this.pendingR2 = new Map(); diff --git a/worker/legacy-binding-policy.mjs b/worker/legacy-binding-policy.mjs new file mode 100644 index 00000000..6b590360 --- /dev/null +++ b/worker/legacy-binding-policy.mjs @@ -0,0 +1,18 @@ +/** Headerless private calls are compatible only with the expansion Container fence. */ +import { legacyDataControls } from "./key-data-guards.mjs"; +export async function legacyBindingControls(env) { + try { return await legacyDataControls(env); } + catch (error) { + // Store outages and malformed rows never acquire an alternate authority. + if (error.code !== "legacy_controlled_accounts") throw error; + const container = env.MULTILLM_PROXY_CONTAINER?.getByName("primary"); + const response = await container.fetch("https://container.internal/healthz", { + method: "GET", redirect: "manual", signal: AbortSignal.timeout(5000), + }); + const capability = (await response.json()).key_controls; + if (!response.ok || capability?.protocol !== 1 || capability.enforcement !== false) throw Error("Unverifiable legacy sender"); + // A's Container rereads and refuses controlled accounts before this private + // hop. It sends no new headers, so a verified A sender proves this call is NULL. + return { official_only: null, forced_provider_preferences: null }; + } +} diff --git a/worker/media-outbound.mjs b/worker/media-outbound.mjs index 758779a1..423e1217 100644 --- a/worker/media-outbound.mjs +++ b/worker/media-outbound.mjs @@ -1,3 +1,4 @@ +import { bindingPrincipal, withDispatchPrincipal, providerFetch, dispatchFailure, DispatchError } from "./provider-dispatch.mjs"; /** * Media I/O for the Container, reachable only through its private outbound handler * (`http://media.internal`). The Worker fetches caller-supplied image URLs, so a URL can @@ -77,8 +78,7 @@ export async function boundedBytes(response, maxBytes) { export async function fetchPublic(value, { accept = "image/*", fetcher = fetch } = {}) { let url = publicHttpsUrl(value); for (let hop = 0; url && hop <= MAX_REDIRECTS; hop += 1) { - const response = await fetcher(url.href, { redirect: "manual", headers: { accept }, - signal: AbortSignal.timeout(FETCH_TIMEOUT_MS) }); + const response = await providerFetch("media", url.href, { headers: { accept }, redirect: "manual", signal: AbortSignal.timeout(FETCH_TIMEOUT_MS) }, { publicAsset: true, send: checked => fetcher(checked.url, { headers: checked.headers, redirect: checked.redirect, signal: checked.signal }) }); if (!REDIRECT_STATUSES.has(response.status)) return response; const location = response.headers.get("location"); await response.body?.cancel(); @@ -227,7 +227,7 @@ async function fileOperation(request, env, id, suffix) { return mediaFailure("not_found", "Unknown media operation.", 404); } -export async function handleMediaOutbound(request, env) { +async function runMediaOutbound(request, env) { const url = new URL(request.url); if (url.origin !== ORIGIN || url.search || url.hash || url.username || url.password) { return mediaFailure("not_found", "Unknown media operation.", 404); @@ -238,7 +238,18 @@ export async function handleMediaOutbound(request, env) { if (file) return await fileOperation(request, env, file[1], file[2]); return mediaFailure("not_found", "Unknown media operation.", 404); } catch (error) { + if (error instanceof DispatchError) return dispatchFailure(error); logFailure("media_outbound_failed", error, { operation: url.pathname.split("/")[2] ?? "unknown" }); return mediaFailure("media_unavailable", "The media operation failed.", 502); } } + +export async function handleMediaOutbound(request, env) { + if (new URL(request.url).search || new URL(request.url).origin !== ORIGIN) return new Response("Not found", { status: 404 }); + if (request.method !== "POST" && new URL(request.url).pathname === "/v1/fetch") return new Response("Not found", { status: 404 }); + // R2 management is independent of provider dispatch. Only external I/O needs an identity. + const path = new URL(request.url).pathname; + if (path !== "/v1/fetch" && !path.endsWith("/import")) return runMediaOutbound(request, env); + try { return await withDispatchPrincipal(env, bindingPrincipal(request), () => runMediaOutbound(request, env)); } + catch (error) { return dispatchFailure(error); } +} diff --git a/worker/native-config-sync.mjs b/worker/native-config-sync.mjs index 5ce8e4dd..9549e1af 100644 --- a/worker/native-config-sync.mjs +++ b/worker/native-config-sync.mjs @@ -28,6 +28,7 @@ async function loadAccounts(db) { after = row.username; // Bootstrap admin keys use environment authority, never these stored key hashes. controls.set(row.username, Object.freeze({ scopes: row.scopes, allowed_models: row.allowed_models, + official_only: row.official_only, forced_provider_preferences: row.forced_provider_preferences, revoked_at: row.revoked_at, expires_at: row.expires_at, allowed_ips: row.allowed_ips })); } if (results.length < 200) return controls; diff --git a/worker/peer-controls.mjs b/worker/peer-controls.mjs new file mode 100644 index 00000000..5fc1894b --- /dev/null +++ b/worker/peer-controls.mjs @@ -0,0 +1,27 @@ +/** Admission before content crosses an independently deployed runtime boundary. */ +import { resolvePrincipal } from "./knowledge-edge.mjs"; +import { currentDataControls, legacyDataControls } from "./key-data-guards.mjs"; +import { DispatchError } from "./provider-dispatch.mjs"; + +export async function delegationPolicy(request, env) { + const headers = new Headers(request.headers); + // Native protocols carry the same proxy credential under their SDK header. + if (!headers.has("x-multillm-api-key") && !headers.has("authorization")) { + const key = headers.get("x-api-key") || headers.get("x-goog-api-key") + || (new URL(request.url).pathname.startsWith("/linkapi/") ? new URL(request.url).searchParams.get("key") : null); + if (key) headers.set("x-multillm-api-key", key); + } + const resolved = await resolvePrincipal(new Request(request.url, { headers }), env); + if (resolved?.denied) return null; // The existing authentication path returns its normal refusal. + return resolved?.principal + ? currentDataControls(env, { keyId: resolved.principal.id }) : legacyDataControls(env); +} + +export async function admitContainerDelegation(request, env, _container) { + const headers = request.headers; + if (!["authorization", "x-multillm-api-key", "x-api-key", "x-goog-api-key"].some(name => headers.has(name)) + && !(new URL(request.url).pathname.startsWith("/linkapi/") && new URL(request.url).searchParams.has("key"))) return; + + try { await delegationPolicy(request, env); } + catch { throw new DispatchError("key_data_controls_unavailable"); } +} diff --git a/worker/provider-dispatch.mjs b/worker/provider-dispatch.mjs new file mode 100644 index 00000000..dfc9eecf --- /dev/null +++ b/worker/provider-dispatch.mjs @@ -0,0 +1,134 @@ +import { legacyBindingControls } from "./legacy-binding-policy.mjs"; +import { KEY_CONTROLS } from "./key-controls-capability.mjs"; +/** Sole provider network boundary. Deferred work carries identity, never cached policy. */ +import { AsyncLocalStorage } from "node:async_hooks"; +import { currentDataControls, legacyDataControls, mergePreferences, parsePreferences, providerClassPermitted } from "./key-data-guards.mjs"; +import { admittedNetwork } from "./dispatch-network.mjs"; +import { dataPolicy } from "./dispatch-policy.mjs"; + +const context = new AsyncLocalStorage(); +export const dispatchContext = () => context.getStore(); +export const withDispatchPrincipal = (env, authority, run) => context.run({ env, authority }, run); +export class DispatchError extends Error { + constructor(code, status = 503) { super(code); this.code = code; this.status = status; } +} +export const dispatchFailure = error => Response.json({ error: { code: error.code ?? "key_data_controls_unavailable" } }, + { status: error.status ?? 503 }); + +export function bindingPrincipal(request, legacyId) { + if (!request.headers.has("X-MultiLLM-Dispatch-Principal")) { + return legacyId ? { authenticated: true, keyId: legacyId } : { legacyBinding: true }; + } + try { + const value = JSON.parse(request.headers.get("X-MultiLLM-Dispatch-Principal")); + if (value.version !== 1 || typeof value.username !== "string" || !value.username + || Object.keys(value).sort().join() !== "username,version") throw Error(); + return { authenticated: true, keyId: value.username }; + } catch { throw new DispatchError("key_data_controls_unavailable"); } +} + +async function admission(provider, options) { + const state = options.env ? options : dispatchContext(); + try { + if (!KEY_CONTROLS.enforcement && state?.env && (!state.authority || state.authority.legacyBinding)) { + return { ...state, policy: await (state.authority?.legacyBinding ? legacyBindingControls(state.env) : legacyDataControls(state.env)) }; + } + if (!state?.authority || state.authority.authenticated !== true) throw Error(); + const user = await currentDataControls(state.env, state.authority); + const policy = dataPolicy(user); + if (!providerClassPermitted(policy, provider, state.env)) throw new DispatchError("provider_class_not_allowed", 403); + return { ...state, policy }; + } catch (error) { + if (error instanceof DispatchError) throw error; + throw new DispatchError("key_data_controls_unavailable"); + } +} + +const apiOrigins = { + openai: "https://api.openai.com", anthropic: "https://api.anthropic.com", xai: "https://api.x.ai", + gemini: "https://generativelanguage.googleapis.com", gemma: "https://generativelanguage.googleapis.com", + openrouter: "https://openrouter.ai", cerebras: "https://api.cerebras.ai", groq: "https://api.groq.com", + together: "https://api.together.xyz", azure: "https://models.inference.ai.azure.com", + hyperbolic: "https://api.hyperbolic.xyz", sambanova: "https://api.sambanova.ai", scaleway: "https://api.scaleway.ai", + opencode: "https://opencode.ai", nanogpt: "https://nano-gpt.com", linkapi: "https://hk.linkapi.ai", +}; +function apiOrigin(provider, url, env, options) { + const prefix = provider.toUpperCase().replaceAll("-", "_"); + const configured = env[`${prefix}_BASE_URL`] ?? (provider === "opencode" ? env.OPENCODE_GO_BASE_URL : null); + try { + const allowed = [apiOrigins[provider], configured].filter(Boolean).map(value => new URL(value).origin); + if (options.realtime) { + for (const [model, entry] of Object.entries(JSON.parse(env.REALTIME_PROVIDERS_JSON ?? "{}"))) { + if (!model.startsWith(provider + ":")) continue; + const endpoint = new URL(entry.url); + if (endpoint.protocol !== "wss:" || endpoint.username || endpoint.password) continue; + endpoint.protocol = "https:"; + allowed.push(endpoint.origin); + } + } + return url.protocol === "https:" && !url.username && !url.password && allowed.includes(url.origin); + } catch { return false; } +} + +function credentials(provider, request, env, policy, options) { + if (!policy.official_only && policy.forced_provider_preferences === null || provider === "cloudflare") return request; + const url = new URL(request.url); + if ([...url.searchParams.keys()].some(name => name.toLowerCase() === "key")) { + throw new DispatchError("provider_credential_not_allowed", 403); + } + if (!apiOrigin(provider, url, env, options)) throw new DispatchError("attested_credential_unavailable"); + const names = { gemini: ["GEMINI_API_KEY"], opencode: ["OPENCODE_GO_API_KEY", "OPENCODE_API_KEY"] }[provider] + ?? [`${provider.toUpperCase().replaceAll("-", "_")}_API_KEY`]; + const key = names.map(name => env[name]).find(value => typeof value === "string" && value.trim() + && !/^(?:your[-_]|changeme|replace[-_]?me|<)/i.test(value)); + if (!key) throw new DispatchError("attested_credential_unavailable"); + const headers = new Headers(request.headers); + for (const name of ["authorization", "x-api-key", "x-goog-api-key", "api-key"]) headers.delete(name); + const name = { gemini: "x-goog-api-key", anthropic: "x-api-key", azure: "api-key" }[provider] ?? "authorization"; + headers.set(name, name === "authorization" ? `Bearer ${key}` : key); + return new Request(request, { headers }); +} + +export async function providerFetch(provider, input, init = {}, options = {}) { + const state = await admission(provider, options); + if (!KEY_CONTROLS.enforcement) return admittedNetwork(() => (options.send ?? fetch)(new Request(input, init))); + const forced = parsePreferences(state.policy.forced_provider_preferences); + if ((options.realtime || options.unsupportedPreferences) && forced !== null) throw new DispatchError("forced_provider_preferences_not_supported", 403); + let request = new Request(input, init); + if (options.publicAsset) { + const headers = new Headers(request.headers); + for (const name of ["authorization", "proxy-authorization", "x-api-key", "x-goog-api-key", "api-key", "cookie"]) headers.delete(name); + request = new Request(request, { headers }); + } else request = credentials(provider, request, state.env, state.policy, options); + if (provider === "openrouter" && forced !== null && request.body) { + let payload; + try { payload = await request.clone().json(); if (!payload || typeof payload !== "object" || Array.isArray(payload)) throw Error(); } + catch { throw new DispatchError("invalid_request", 400); } + payload.provider = mergePreferences(payload.provider, forced); + const headers = new Headers(request.headers); headers.delete("content-length"); + request = new Request(request, { headers, body: JSON.stringify(payload) }); + } + const headers = new Headers(request.headers); headers.delete("x-multillm-dispatch-principal"); + request = new Request(request, { headers, redirect: "manual" }); + return admittedNetwork(() => (options.send ?? fetch)(request)); +} + +/** SDK/binding calls use the same admission, including embeddings and background work. */ +export async function providerOperation(provider, run, options = {}) { + const state = await admission(provider, options); + if (provider === "openrouter" && state.policy.forced_provider_preferences !== null) { + throw new DispatchError("forced_provider_preferences_not_supported", 403); + } + return admittedNetwork(run); +} + +/** Keep adapter transports compatible while passing only the admitted request. */ +export async function transportOptions(request, original = {}) { + const headers = {}; + for (const [name, value] of request.headers) { + const spelling = Object.keys(original.headers ?? {}).find(key => key.toLowerCase() === name) ?? name; + headers[spelling] = value; + } + return { ...original, method: request.method, headers, body: request.body ? await request.text() : undefined, + signal: request.signal, redirect: request.redirect }; +} diff --git a/worker/realtime.mjs b/worker/realtime.mjs index e1fa50ba..48415ef0 100644 --- a/worker/realtime.mjs +++ b/worker/realtime.mjs @@ -1,4 +1,6 @@ +import { providerFetch } from "./provider-dispatch.mjs"; /** Native gateway authentication, ticket admission and approved Realtime dispatch. */ +import { currentDataControls, parsePreferences, providerClassPermitted } from "./key-data-guards.mjs"; import { Buffer } from "node:buffer"; import { createHash, scrypt, timingSafeEqual } from "node:crypto"; import { activeUsersByPrefix, boundedBody, keyControlsPermit, validUser } from "./control-users-d1.mjs"; @@ -67,13 +69,23 @@ function modelGranted(patterns, model) { return patterns.split(",").some(pattern => new RegExp("^" + pattern.toLowerCase().split("*") .map(part => part.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")).join(".*") + "$").test(model.toLowerCase())); } -function permit(principal, request, model, now) { +async function permit(principal, request, model, now, env) { if (!principal || principal.revoked_at || !keyControlsPermit(principal, request.headers.get("cf-connecting-ip"), now)) { throw new RealtimeError("invalid_api_key", 401); } if (!principal.scopes.includes("chat") || !principal.scopes.includes("audio") || !modelGranted(principal.allowed_models, model)) { throw new RealtimeError("model_not_granted", 403); } + let controls; + try { controls = await currentDataControls(env, { principal, keyId: principal.owner }); } + catch { throw new RealtimeError("key_data_controls_unavailable", 503); } + if (!providerClassPermitted(controls, model.split(":", 1)[0], env)) { + throw new RealtimeError("provider_class_not_allowed", 403); + } + // Realtime frames cannot enforce HTTP provider preferences throughout a session. + if (parsePreferences(controls.forced_provider_preferences) !== null) { + throw new RealtimeError("forced_provider_preferences_not_supported", 403); + } return principal; } function account(row) { @@ -124,7 +136,7 @@ async function authenticate(request, env, model, now) { if (!principal) throw new RealtimeError("invalid_api_key", 401); await tenantPrincipal(env, principal); principal.principal_hash = await principalHash(`realtime-key:${digest(key)}`); - return permit(principal, request, model, now); + return permit(principal, request, model, now, env); } async function signingKey(env) { // Derive a domain-specific key from the existing gateway signing binding. @@ -171,7 +183,7 @@ async function consumeTicket(ticket, request, env, model, now) { if (credential) principal = integration(await lookupIntegrationPrincipal(env.INTELLIGENCE_DB, credential.key_prefix)); } if (!principal || principal.owner !== row.owner || !same(principal.credential_fingerprint, row.credential_fingerprint)) throw new RealtimeError("invalid_api_key", 401); - permit(principal, request, model, now); + await permit(principal, request, model, now, env); const consumed = await env.INTELLIGENCE_DB.prepare("UPDATE realtime_tickets SET consumed_at=? WHERE nonce=? AND consumed_at IS NULL AND expires_at>?") .bind(now, match[2], now).run(); if (consumed.meta.changes !== 1) throw new RealtimeError("invalid_api_key", 401); @@ -240,8 +252,10 @@ async function openSession(request, env, ctx, settings, options) { if (request.signal.aborted || controller.signal.aborted) throw new RealtimeError("realtime_canceled", 499); const url = new URL(settings.providers.get(model)); url.protocol = "https:"; handedOff = true; - const response = await (options.dial ?? fetch)(url.href, { method: "GET", headers: upstreamHeaders(request, key), - redirect: "manual", signal: controller.signal }); + const response = await providerFetch(provider, url.href, { method: "GET", headers: upstreamHeaders(request, key), + redirect: "manual", signal: controller.signal }, { env, authority: { authenticated: true, keyId: principal.owner }, realtime: true, + send: checked => (options.dial ?? fetch)(checked.url, { method: checked.method, headers: checked.headers, + redirect: checked.redirect, signal: checked.signal }) }); socket = response.webSocket; if (response.status !== 101 || !socket) { void response.body?.cancel?.().catch(() => {}); throw new RealtimeError("realtime_upstream_error", 502); } if (request.signal.aborted || controller.signal.aborted) throw new RealtimeError("realtime_canceled", 499); diff --git a/worker/roleplay/credential-health.mjs b/worker/roleplay/credential-health.mjs index 960dfbff..a244bd38 100644 --- a/worker/roleplay/credential-health.mjs +++ b/worker/roleplay/credential-health.mjs @@ -1,3 +1,4 @@ +import { providerFetch } from "../provider-dispatch.mjs"; import { buildProviderHeaders } from "./config.mjs"; const DEFINITE_KEY_FAILURES = new Set([401, 402, 403, 429]); @@ -52,12 +53,13 @@ async function probeCredential(candidate, env, settings, signal) { settings.nanogptKeyCheckTimeoutMs, ); try { - return await fetch(candidate.catalogEndpoint, { + return await providerFetch(candidate.provider, candidate.catalogEndpoint, { method: "GET", headers: buildProviderHeaders(candidate, env), redirect: "manual", signal: controller.signal, - }); + }, { env, unsupportedPreferences: true, authority: { authenticated: true, keyId: env.ADMIN_USERNAME || "admin" }, + send: checked => fetch(checked.url, { method: checked.method, headers: checked.headers, redirect: checked.redirect, signal: checked.signal }) }); } finally { clearTimeout(timeout); signal?.removeEventListener("abort", forwardAbort); diff --git a/worker/roleplay/transport.mjs b/worker/roleplay/transport.mjs index 32a53389..93cb08c0 100644 --- a/worker/roleplay/transport.mjs +++ b/worker/roleplay/transport.mjs @@ -1,3 +1,4 @@ +import { providerFetch, dispatchFailure, DispatchError } from "../provider-dispatch.mjs"; import { UpstreamCancellation, readBoundedUpstreamBytes } from "../upstream-cancellation.mjs"; import { firewallFetch, isSecretScanBlock } from "../secret-firewall.mjs"; import { piiFetch, isPIIFailure } from "../pii-rehydration.mjs"; @@ -581,21 +582,19 @@ async function fetchCandidate(candidate, payload, env, settings, signal, key) { }; cancellation.handoff(); const piiOptions = { route: "/v1/roleplay/chat/completions", keyScope: settings.piiKeyScope ?? "", onDecision: settings.onPIIDecision }; - const submit = checked => piiFetch(checked, env, piiOptions, async transformed => - useOpenCodeContainer ? containerNamespace.getByName("primary").fetch(transformed) - : fetch(transformed.url, { ...requestInit, headers: transformed.headers, body: await transformed.text() })); + const authority = { authenticated: true, keyId: env.ADMIN_USERNAME || "admin" }; + const submit = checked => piiFetch(checked, env, piiOptions, transformed => providerFetch(candidate.provider, transformed, {}, { + env, authority, unsupportedPreferences: true, + send: async permitted => useOpenCodeContainer ? containerNamespace.getByName("primary").fetch(permitted) : fetch(permitted.url, { ...requestInit, headers: permitted.headers, body: await permitted.text() }), + }).catch(error => { if (!(error instanceof DispatchError)) throw error; return dispatchFailure(error); })); const response = useOpenCodeContainer && !resolvePIIPolicy(env, piiOptions) - ? await containerNamespace.getByName("primary").fetch( - new Request( - "https://roleplay.internal/opencode/v1/chat/completions", - requestInit, - ), - ) + ? await providerFetch(candidate.provider, new Request("https://roleplay.internal/opencode/v1/chat/completions", requestInit), {}, { + env, authority, unsupportedPreferences: true, send: checked => containerNamespace.getByName("primary").fetch(checked), + }).catch(error => { if (!(error instanceof DispatchError)) throw error; return dispatchFailure(error); }) : await firewallFetch(new Request(useOpenCodeContainer ? "https://roleplay.internal/opencode/v1/chat/completions" : candidate.endpoint, requestInit), env, { - route: "/v1/roleplay/chat/completions", provider: candidate.provider, - onDecision: settings.onSecretScan, - principal: { id: env.ADMIN_USERNAME || "admin" }, - }, submit); + route: "/v1/roleplay/chat/completions", provider: candidate.provider, + onDecision: settings.onSecretScan, principal: { id: authority.keyId }, + }, submit); cancellation.headersReceived(); const ownedResponse = isSecretScanBlock(response) || isPIIFailure(response) ? response : cancellation.wrapResponse(response); diff --git a/worker/scheduled-maintenance.mjs b/worker/scheduled-maintenance.mjs index 175d264f..b4636ada 100644 --- a/worker/scheduled-maintenance.mjs +++ b/worker/scheduled-maintenance.mjs @@ -1,3 +1,4 @@ +import { controlNetwork } from "./dispatch-network.mjs"; /** Bounded shared-cache maintenance and content-free alert delivery. */ import { cleanupGenerationCache } from "./generation-cache-d1.mjs"; import { generationCacheSettings } from "./exact-generation-cache.mjs"; @@ -34,8 +35,8 @@ export async function collectAlertAggregates(env, rules, now = new Date()) { export function alertTransport(destination, options) { // The delivery owner provides its three-second abort signal and bounds the payload. - return fetch(destination, { method: "POST", headers: options.headers, body: options.body, - redirect: "error", signal: options.signal }); + return controlNetwork(destination, () => fetch(destination, { method: "POST", headers: options.headers, body: options.body, + redirect: "error", signal: options.signal })); } /** Scheduled work is registered only when a feature that needs it is enabled. */ diff --git a/worker/secret-firewall.mjs b/worker/secret-firewall.mjs index c46ce09c..37b502be 100644 --- a/worker/secret-firewall.mjs +++ b/worker/secret-firewall.mjs @@ -1,3 +1,4 @@ +import { providerFetch } from "./provider-dispatch.mjs"; /** Provider dispatch policy. Audit failures never change the upstream request. */ import { redactPayload, scanText } from "./secret-scan.mjs"; @@ -124,8 +125,9 @@ function multipart(bytes, contentType) { }]; } -export async function firewallFetch(request, env, options = {}, fetchImpl = fetch) { - if (!request.body || secretScanMode(env, options.principal) === "off") return fetchImpl(request); +export async function firewallFetch(request, env, options = {}, fetchImpl) { + const send = fetchImpl ?? (checked => providerFetch(options.provider, checked, {}, { env, authority: options })); + if (!request.body || secretScanMode(env, options.principal) === "off") return send(request); let protectedRequest = request, decision; try { const bytes = await boundedBody(request); @@ -159,7 +161,7 @@ export async function firewallFetch(request, env, options = {}, fetchImpl = fetc const reason = error?.message; console.warn(reason === "scan_timeout" || reason === "scan_body_limit" ? reason : "secret_scan_body_unavailable"); } - const upstream = await fetchImpl(protectedRequest); + const upstream = await send(protectedRequest); if (!decision?.header || !upstream.ok) return upstream; const response = new Response(upstream.body, upstream); response.headers.set(SECRET_SCAN_HEADER, decision.header); diff --git a/wrangler.jsonc b/wrangler.jsonc index 9807bd6e..deb6925c 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -4,7 +4,7 @@ "keep_vars": true, "main": "./cloudflare-worker.mjs", "build": { - "command": "python scripts/build_release_metadata.py" + "command": "node scripts/key_data_preflight.mjs --deploy && python scripts/build_release_metadata.py" }, "version_metadata": { "binding": "CF_VERSION_METADATA" diff --git a/wrangler.knowledge.jsonc b/wrangler.knowledge.jsonc index c88b2b1a..4183e10c 100644 --- a/wrangler.knowledge.jsonc +++ b/wrangler.knowledge.jsonc @@ -3,7 +3,7 @@ "name": "multillm-knowledge", "main": "worker/knowledge/index.mjs", "compatibility_date": "2026-07-28", - "compatibility_flags": ["enable_request_signal", "request_signal_passthrough"], + "compatibility_flags": ["nodejs_compat", "enable_request_signal", "request_signal_passthrough"], "workers_dev": false, "preview_urls": false, "d1_databases": [{