diff --git a/.github/workflows/enterprise-release.yml b/.github/workflows/enterprise-release.yml new file mode 100644 index 0000000..69b620b --- /dev/null +++ b/.github/workflows/enterprise-release.yml @@ -0,0 +1,116 @@ +name: enterprise release + +on: + push: + tags: ["enterprise-v*"] + +permissions: + contents: read + +jobs: + release: + runs-on: ubuntu-latest + permissions: + contents: write + outputs: + version: ${{ steps.version.outputs.value }} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Validate the component version and create the release + id: version + env: + GH_TOKEN: ${{ github.token }} + run: | + version="${GITHUB_REF_NAME#enterprise-v}" + [[ "$version" =~ ^[0-9]+\.[0-9]+\.[0-9]+(-[a-zA-Z0-9.-]+)?$ ]] + test "$version" = "$(python3 -c 'import tomllib; print(tomllib.load(open("ee/pyproject.toml", "rb"))["project"]["version"])')" + echo "value=$version" >> "$GITHUB_OUTPUT" + gh release view "$GITHUB_REF_NAME" > /dev/null 2>&1 || \ + gh release create "$GITHUB_REF_NAME" --verify-tag --draft --prerelease --generate-notes + + image: + needs: release + runs-on: ubuntu-latest + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + arch: [amd64, arm64] + permissions: + contents: write + packages: write + env: + IMAGE: ghcr.io/getshim/shim-enterprise + VERSION: ${{ needs.release.outputs.version }} + ARCH: ${{ matrix.arch }} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: docker/setup-qemu-action@1f40c72289eff860ee54a304f1438e3cff362e0a # v4 + - uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4 + - uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2 + with: + cosign-release: v3.1.3 + - name: Log in with this repository's package token + env: + GH_TOKEN: ${{ github.token }} + run: echo "$GH_TOKEN" | docker login ghcr.io --username "$GITHUB_ACTOR" --password-stdin + - name: Build and publish the component + id: build + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7 + with: + context: . + file: ee/Dockerfile + platforms: linux/${{ matrix.arch }} + push: true + tags: ghcr.io/getshim/shim-enterprise:${{ env.VERSION }}-${{ matrix.arch }} + provenance: mode=max + sbom: true + - name: Sign the image and its offline delivery files + env: + COSIGN_PRIVATE_KEY: ${{ secrets.RELEASE_COSIGN_PRIVATE_KEY }} + COSIGN_PASSWORD: ${{ secrets.RELEASE_COSIGN_PASSWORD }} + DIGEST: ${{ steps.build.outputs.digest }} + GH_TOKEN: ${{ github.token }} + run: | + test -n "$COSIGN_PRIVATE_KEY" + delivery="shim-enterprise-$VERSION-$ARCH" + mkdir "$delivery" + cosign public-key --key env://COSIGN_PRIVATE_KEY > "$delivery/release.pub" + cosign sign --yes --key env://COSIGN_PRIVATE_KEY --use-signing-config=false --tlog-upload=false "$IMAGE@$DIGEST" + docker pull --platform "linux/$ARCH" "$IMAGE@$DIGEST" + docker save --platform "linux/$ARCH" --output "$delivery/image.tar" "$IMAGE@$DIGEST" + docker buildx imagetools inspect "$IMAGE@$DIGEST" --format '{{json .SBOM}}' | \ + jq -e '.SPDX' > "$delivery/sbom.spdx.json" + cp LICENSE "$delivery/community-LICENSE" + cp NOTICE "$delivery/community-NOTICE" + cp ee/LICENSE "$delivery/enterprise-LICENSE" + cp ee/NOTICE "$delivery/enterprise-NOTICE" + cp ee/docs/OFFLINE_RELEASES.md "$delivery/" + docker buildx imagetools inspect "$IMAGE@$DIGEST" --format '{{json .Provenance}}' > "$delivery/provenance.json" + jq -e '.SLSA != null' "$delivery/provenance.json" > /dev/null + jq -n --arg component enterprise --arg version "$VERSION" --arg platform "linux/$ARCH" \ + --arg source "$IMAGE@$DIGEST" --arg commit "$GITHUB_SHA" \ + '{component:$component,version:$version,platform:$platform,source:$source,commit:$commit}' \ + > "$delivery/image.json" + (cd "$delivery" && sha256sum image.tar image.json sbom.spdx.json provenance.json release.pub community-LICENSE community-NOTICE enterprise-LICENSE enterprise-NOTICE OFFLINE_RELEASES.md > SHA256SUMS) + cosign sign-blob --yes --key env://COSIGN_PRIVATE_KEY --use-signing-config=false --tlog-upload=false \ + --bundle "$delivery/SHA256SUMS.sigstore.json" "$delivery/SHA256SUMS" + cosign verify-blob --key "$delivery/release.pub" --insecure-ignore-tlog \ + --bundle "$delivery/SHA256SUMS.sigstore.json" "$delivery/SHA256SUMS" + tar -cf "$delivery.tar" "$delivery" + gh release upload "$GITHUB_REF_NAME" "$delivery.tar" + + publish: + needs: [release, image] + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - name: Publish only after every platform artifact succeeds + env: + GH_TOKEN: ${{ github.token }} + run: gh release edit "$GITHUB_REF_NAME" --repo "$GITHUB_REPOSITORY" --draft=false diff --git a/.github/workflows/on-prem-test.yml b/.github/workflows/on-prem-test.yml new file mode 100644 index 0000000..121c22d --- /dev/null +++ b/.github/workflows/on-prem-test.yml @@ -0,0 +1,85 @@ +name: Isolated on-prem chart smoke + +on: + workflow_dispatch: + inputs: + dashboard_image: + description: Prebuilt OIDC dashboard image with immutable sha256 digest (runner must have pull access) + type: string + required: true + dashboard_version: + description: Dashboard source commit/version recorded for this acceptance run + type: string + required: true + +permissions: + contents: read + +jobs: + isolated-chart: + runs-on: ubuntu-latest + timeout-minutes: 45 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.13" + - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + with: + version: "0.12.5" + enable-cache: false + - name: Install pinned local cluster tools + shell: bash + run: | + set -euo pipefail + cd "$RUNNER_TEMP" + curl --fail --silent --show-error --location -o kind-linux-amd64 https://github.com/kubernetes-sigs/kind/releases/download/v0.33.0/kind-linux-amd64 + curl --fail --silent --show-error --location -o kind-linux-amd64.sha256sum https://github.com/kubernetes-sigs/kind/releases/download/v0.33.0/kind-linux-amd64.sha256sum + sha256sum --check kind-linux-amd64.sha256sum + sudo install -m 0755 kind-linux-amd64 /usr/local/bin/kind + curl --fail --silent --show-error --location -o helm.tar.gz https://get.helm.sh/helm-v4.2.4-linux-amd64.tar.gz + curl --fail --silent --show-error --location -o helm.sha256 https://get.helm.sh/helm-v4.2.4-linux-amd64.tar.gz.sha256sum + mv helm.tar.gz helm-v4.2.4-linux-amd64.tar.gz + sha256sum --check helm.sha256 + tar -xzf helm-v4.2.4-linux-amd64.tar.gz + sudo install -m 0755 linux-amd64/helm /usr/local/bin/helm + kubectl version --client + - name: Authenticate private dashboard artifact pulls when configured + env: + GHCR_TOKEN: ${{ secrets.ON_PREM_DASHBOARD_GHCR_TOKEN }} + GHCR_USER: ${{ vars.ON_PREM_DASHBOARD_GHCR_USER || github.actor }} + run: | + if [ -n "$GHCR_TOKEN" ]; then + printf '%s' "$GHCR_TOKEN" | docker login ghcr.io --username "$GHCR_USER" --password-stdin + fi + - name: Validate dashboard artifact and build enterprise runtime + env: + DASHBOARD_IMAGE: ${{ inputs.dashboard_image }} + DASHBOARD_VERSION: ${{ inputs.dashboard_version }} + shell: bash + run: | + set -euo pipefail + [[ "$DASHBOARD_IMAGE" =~ @sha256:[0-9a-f]{64}$ ]] + docker pull "$DASHBOARD_IMAGE" + printf 'Dashboard source: %s\n' "$DASHBOARD_VERSION" >> "$GITHUB_STEP_SUMMARY" + docker build --network=default -f ee/Dockerfile --target runtime -t shim-enterprise:isolated-ci . + - name: Run real services on an internet-isolated cluster + env: + DASHBOARD_IMAGE: ${{ inputs.dashboard_image }} + run: | + uv run --locked --all-packages ruff check ee/deploy/test + uv run --locked --all-packages ruff format --check ee/deploy/test + uv run --locked --all-packages python -m pytest -q ee/deploy/test/test_reuse.py + uv run --locked --all-packages python ee/deploy/test/run.py --gateway-image shim-enterprise:isolated-ci --dashboard-image "$DASHBOARD_IMAGE" + - name: Save non-secret evidence + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 + with: + name: isolated-chart-evidence + path: | + /tmp/shim-chart-smoke-*/result.json + /tmp/shim-chart-smoke-*/asset-audit.json + /tmp/shim-chart-smoke-*/resources.txt + if-no-files-found: warn diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 51d9dd2..80b996d 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -245,6 +245,9 @@ jobs: with: version: "0.12.5" enable-cache: true + - uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2 + with: + cosign-release: v3.1.3 - run: uv lock --check - run: uv sync --locked --all-packages - run: uv run --locked --package shim-enterprise alembic -c ee/alembic.ini upgrade head diff --git a/AGENTS.md b/AGENTS.md index d240551..93e5fc1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,7 +1,7 @@ # shim repository rules -Read `docs/CURRENT_ARCHITECTURE.md`, `docs/TARGET_ARCHITECTURE.md`, and -`DEVELOPER_GUIDE.md` before changing runtime boundaries. +Read `docs/CURRENT_ARCHITECTURE.md` and `DEVELOPER_GUIDE.md` before changing +runtime boundaries. ## Ownership diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index d439bcc..59bd777 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -89,7 +89,8 @@ uv run --locked --package shim-enterprise alembic -c ee/alembic.ini upgrade head Never downgrade a production database. Roll production back to the previous schema-compatible application image. -Contact-only plan activation is an enterprise operation: +Operator-managed provisioning and plan activation are enterprise operations. +Follow [the provisioning and existing-customer transition runbook](ee/docs/PROVISIONING.md): ```bash uv run --locked --package shim-enterprise python ee/scripts/activate_plan.py --help @@ -125,6 +126,9 @@ assets and require the exact matching `shim-gateway` version. ## Changing an API or dependency +Diagnostic field semantics and null behavior are documented in +[`ee/docs/DIAGNOSTIC_METADATA.md`](ee/docs/DIAGNOSTIC_METADATA.md). + 1. Trace the route through authentication, `GatewayService`, `GatewayKernel`, provider execution, streaming, and usage finalization. 2. Change the owning product; keep the other product unchanged unless a public diff --git a/README.md b/README.md index f813b59..ebee9d3 100644 --- a/README.md +++ b/README.md @@ -120,6 +120,7 @@ Gemini `generateContent` routes, `/v1/models`, `/v1/scan`, `/health` and `/metrics`. The checked-in contract is [`openapi/community.json`](openapi/community.json). To run from source instead, see [the developer guide](DEVELOPER_GUIDE.md). +For customer-operated enterprise installations, see [deployment and recovery](ee/deploy/README.md). ## Limitations @@ -142,7 +143,6 @@ To run from source instead, see [the developer guide](DEVELOPER_GUIDE.md). - [Developer guide](DEVELOPER_GUIDE.md) - [Current architecture](docs/CURRENT_ARCHITECTURE.md) -- [Target architecture](docs/TARGET_ARCHITECTURE.md) - [Contributing](CONTRIBUTING.md) - [Security policy](https://github.com/GetSHIM/shim/security/policy) diff --git a/architecture/module_ownership.toml b/architecture/module_ownership.toml index a64a0a8..3f14666 100644 --- a/architecture/module_ownership.toml +++ b/architecture/module_ownership.toml @@ -18,7 +18,7 @@ forbidden_public_import_roots = [ # Exact community symbols consumed by enterprise runtime code. Any new entry is # a deliberate expansion of the supported cross-license API. [enterprise_public_api] -"shim.api.v1.chat" = ["router"] +"shim.api.v1.chat" = ["model_record", "router"] "shim.api.v1.gemini" = ["router"] "shim.api.v1.messages" = ["router"] "shim.api.v1.responses" = ["router"] @@ -54,7 +54,7 @@ forbidden_public_import_roots = [ "validate_actor_identity", ] "shim.gateway.kernel.gateway_kernel" = ["GatewayKernel"] -"shim.gateway.kernel.result" = ["AdmissionState", "PreparedInference"] +"shim.gateway.kernel.result" = ["AdmissionState", "PreparedInference", "ProviderTarget", "UNSPECIFIED_PROVIDER_MODEL"] "shim.gateway.pipeline.anthropic_execution" = ["AnthropicExecution"] "shim.gateway.pipeline.authenticate" = ["GatewayRequestMetadata"] "shim.gateway.pipeline.google_execution" = ["GoogleExecution"] @@ -62,7 +62,7 @@ forbidden_public_import_roots = [ "shim.gateway.pipeline.privacy" = ["ScanPrivacyOutcome", "ScanPrivacyStage"] "shim.gateway.request_policy" = ["RequestPolicyContext", "ResolvedRequestPolicy"] "shim.gateway.streaming.finalization" = ["StreamFinalization"] -"shim.gateway.usage" = ["UsageFailureReason", "UsageLimitExceeded"] +"shim.gateway.usage" = ["UsageAuditPersistenceError", "UsageFailureReason", "UsageLimitExceeded"] "shim.observability.logging" = [ "configure_error_reporting", "configure_logging", @@ -184,6 +184,7 @@ public = [ "tests/gateway/test_openai_sdk_transport.py", "tests/gateway/test_provider_error_fidelity.py", "tests/gateway/test_request_policy.py", + "tests/gateway/test_token_count.py", "tests/gateway/test_usage.py", "tests/observability/test_sanitization.py", "tests/secrets/test_credentials.py", @@ -193,8 +194,15 @@ public = [ enterprise = [ "ee/alembic/env.py", "ee/alembic/versions/aa8b038bc50c_architecture_baseline.py", + "ee/alembic/versions/c31b7a91d602_oidc_identity.py", + "ee/alembic/versions/c7a108be3201_teams_roles_and_scoped_key_controls.py", "ee/alembic/versions/f10e4ac92d17_harden_function_search_paths.py", + "ee/alembic/versions/fcc02bbe6443_add_tenant_model_deployment_registry.py", + "ee/deploy/test/fixtures.py", + "ee/deploy/test/run.py", + "ee/deploy/test/test_reuse.py", "ee/scripts/activate_plan.py", + "ee/scripts/offline_bundle.py", "ee/scripts/sign_license.py", "ee/src/shim_enterprise/__init__.py", "ee/src/shim_enterprise/ai_act/__init__.py", @@ -215,7 +223,6 @@ enterprise = [ "ee/src/shim_enterprise/api/v1/management.py", "ee/src/shim_enterprise/api/v1/router.py", "ee/src/shim_enterprise/api/v1/scan.py", - "ee/src/shim_enterprise/api/v1/subscriptions.py", "ee/src/shim_enterprise/application.py", "ee/src/shim_enterprise/billing/__init__.py", "ee/src/shim_enterprise/billing/ledger.py", @@ -287,6 +294,7 @@ enterprise = [ "ee/src/shim_enterprise/secrets/gcp_secret_manager.py", "ee/src/shim_enterprise/secrets/migration.py", "ee/src/shim_enterprise/secrets/store.py", + "ee/src/shim_enterprise/secrets/vault.py", "ee/src/shim_enterprise/services/__init__.py", "ee/src/shim_enterprise/services/gateway/__init__.py", "ee/src/shim_enterprise/services/gateway/enterprise.py", @@ -294,14 +302,19 @@ enterprise = [ "ee/src/shim_enterprise/shared_results/api.py", "ee/src/shim_enterprise/shared_results/models.py", "ee/src/shim_enterprise/tenants/__init__.py", + "ee/src/shim_enterprise/tenants/audit.py", + "ee/src/shim_enterprise/tenants/deployments.py", "ee/src/shim_enterprise/tenants/models.py", + "ee/src/shim_enterprise/tenants/oidc.py", + "ee/src/shim_enterprise/tenants/plans.py", "ee/src/shim_enterprise/tenants/policy.py", "ee/src/shim_enterprise/tenants/service.py", - "ee/src/shim_enterprise/tenants/subscriptions.py", + "ee/src/shim_enterprise/tenants/teams.py", "ee/src/shim_enterprise/workers/__init__.py", "ee/src/shim_enterprise/workers/ai_act.py", "ee/src/shim_enterprise/workers/compliance.py", "ee/src/shim_enterprise/workers/outbox.py", + "ee/src/shim_enterprise/workers/readiness.py", "ee/src/shim_enterprise/workers/reconciliation.py", "ee/tests/ai_act/test_audit_chain.py", "ee/tests/ai_act/test_control_plane.py", @@ -325,6 +338,7 @@ enterprise = [ "ee/tests/gateway/api/test_scan_endpoint.py", "ee/tests/gateway/cache/test_loop_detection.py", "ee/tests/gateway/kernel/test_accounting_coordinator.py", + "ee/tests/gateway/pipeline/test_decisions.py", "ee/tests/gateway/pipeline/test_scan.py", "ee/tests/gateway/pipeline/test_scan_recovery.py", "ee/tests/gateway/privacy/test_chain_store.py", @@ -332,12 +346,18 @@ enterprise = [ "ee/tests/outbox/test_dead_letter.py", "ee/tests/outbox/test_publisher.py", "ee/tests/outbox/test_worker.py", + "ee/tests/scripts/test_offline_bundle.py", "ee/tests/secrets/test_ciphertext_compatibility.py", "ee/tests/secrets/test_store.py", + "ee/tests/secrets/test_vault.py", "ee/tests/shared_results/test_api.py", + "ee/tests/tenants/test_deployments.py", + "ee/tests/tenants/test_oidc.py", + "ee/tests/tenants/test_plans.py", "ee/tests/tenants/test_policy.py", "ee/tests/tenants/test_service.py", - "ee/tests/tenants/test_subscriptions.py", + "ee/tests/tenants/test_teams.py", + "ee/tests/workers/test_readiness.py", ] # Transitional Python files that contain or exercise both ownership regions. diff --git a/architecture/route_profiles.toml b/architecture/route_profiles.toml index 11ae800..c85e5aa 100644 --- a/architecture/route_profiles.toml +++ b/architecture/route_profiles.toml @@ -7,12 +7,17 @@ community = [ { method = "GET", path = "/v1/models/{model_id}" }, { method = "POST", path = "/v1/chat/completions" }, { method = "POST", path = "/v1/messages" }, + { method = "POST", path = "/v1/messages/count_tokens" }, { method = "POST", path = "/v1/responses" }, { method = "POST", path = "/v1/scan" }, { method = "POST", path = "/v1beta/models/{model}:generateContent" }, { method = "POST", path = "/v1beta/models/{model}:streamGenerateContent" }, ] enterprise = [ + { method = "GET", path = "/api/v1/management/model-deployments" }, + { method = "POST", path = "/api/v1/management/model-deployments" }, + { method = "PUT", path = "/api/v1/management/model-deployments/{deployment_id}" }, + { method = "POST", path = "/api/v1/management/model-deployments/{deployment_id}/health" }, { method = "DELETE", path = "/api/v1/compliance/connectors/{connector_id}" }, { method = "DELETE", path = "/api/v1/compliance/forward-targets/{target_id}" }, { method = "DELETE", path = "/api/v1/compliance/oversight/policies/{policy_id}" }, @@ -21,6 +26,10 @@ enterprise = [ { method = "DELETE", path = "/api/v1/management/providers/{secret_id}" }, { method = "DELETE", path = "/api/v1/management/team/invites/{invite_id}" }, { method = "DELETE", path = "/api/v1/management/team/members/{member_id}" }, + { method = "DELETE", path = "/api/v1/management/teams/{team_id}/members/{member_id}" }, + { method = "GET", path = "/api/v1/auth/callback" }, + { method = "GET", path = "/api/v1/auth/login" }, + { method = "GET", path = "/api/v1/auth/session" }, { method = "GET", path = "/api/v1/compliance/audit/logs" }, { method = "GET", path = "/api/v1/compliance/connectors" }, { method = "GET", path = "/api/v1/compliance/connectors/{connector_id}" }, @@ -44,6 +53,8 @@ enterprise = [ { method = "GET", path = "/api/v1/management/subscription" }, { method = "GET", path = "/api/v1/management/team/invites" }, { method = "GET", path = "/api/v1/management/team/members" }, + { method = "GET", path = "/api/v1/management/teams" }, + { method = "GET", path = "/api/v1/management/teams/{team_id}/members" }, { method = "GET", path = "/api/v1/management/tier-info" }, { method = "GET", path = "/api/v1/shared-results/{token}" }, { method = "GET", path = "/health" }, @@ -56,6 +67,7 @@ enterprise = [ { method = "PATCH", path = "/api/v1/management/api-keys/{api_key_id}" }, { method = "PATCH", path = "/api/v1/management/cost/budgets/{budget_id}" }, { method = "PATCH", path = "/api/v1/management/team/members/{member_id}" }, + { method = "POST", path = "/api/v1/auth/logout" }, { method = "POST", path = "/api/v1/compliance/audit/anchor" }, { method = "POST", path = "/api/v1/compliance/audit/verify" }, { method = "POST", path = "/api/v1/compliance/connectors" }, @@ -67,15 +79,17 @@ enterprise = [ { method = "POST", path = "/api/v1/compliance/reports/audit" }, { method = "POST", path = "/api/v1/compliance/reports/kvkk" }, { method = "POST", path = "/api/v1/management/api-keys" }, + { method = "POST", path = "/api/v1/management/api-keys/{api_key_id}/rotate" }, { method = "POST", path = "/api/v1/management/cost/budgets" }, { method = "POST", path = "/api/v1/management/cost/budgets/evaluate" }, { method = "POST", path = "/api/v1/management/providers" }, { method = "POST", path = "/api/v1/management/providers/{secret_id}/verify" }, { method = "POST", path = "/api/v1/management/team/invites" }, { method = "POST", path = "/api/v1/management/team/invites/accept" }, - { method = "POST", path = "/api/v1/webhooks/lemonsqueezy" }, + { method = "POST", path = "/api/v1/management/teams" }, { method = "POST", path = "/v1/chat/completions" }, { method = "POST", path = "/v1/messages" }, + { method = "POST", path = "/v1/messages/count_tokens" }, { method = "POST", path = "/v1/responses" }, { method = "POST", path = "/v1/scan" }, { method = "POST", path = "/v1/shared-results" }, @@ -84,4 +98,6 @@ enterprise = [ { method = "PUT", path = "/api/v1/management/auth/me" }, { method = "PUT", path = "/api/v1/management/providers/{secret_id}" }, { method = "PUT", path = "/api/v1/management/settings/pii" }, + { method = "PUT", path = "/api/v1/management/teams/{team_id}" }, + { method = "PUT", path = "/api/v1/management/teams/{team_id}/members/{member_id}" }, ] diff --git a/cloudbuild.yaml b/cloudbuild.yaml index c1b3c72..0a69edd 100644 --- a/cloudbuild.yaml +++ b/cloudbuild.yaml @@ -143,7 +143,7 @@ steps: - --deploy-health-check - --no-traffic - --set-secrets=DATABASE_URL=${_SECRET_PREFIX}-database-url:1,REDIS_URL=${_SECRET_PREFIX}-redis-url:1,SUPABASE_URL=${_SECRET_PREFIX}-supabase-url:1,SUPABASE_KEY=${_SECRET_PREFIX}-supabase-key:1,SECRET_KEY=${_SECRET_PREFIX}-secret-key:1,SHIM_LICENSE_KEY=${_SECRET_PREFIX}-license-key:1,ENCRYPTION_KEY=${_SECRET_PREFIX}-encryption-key:1,SENTRY_DSN=${_SECRET_PREFIX}-sentry-dsn:1,OTEL_EXPORTER_OTLP_ENDPOINT=${_SECRET_PREFIX}-otel-exporter-otlp-endpoint:1,OTEL_EXPORTER_OTLP_HEADERS=${_SECRET_PREFIX}-otel-exporter-otlp-headers:1 - - --set-env-vars=^@^BACKEND_CORS_ORIGINS=${_FRONTEND_ORIGINS}@ENVIRONMENT=production@SECRET_BACKEND=gcp_secret_manager@GOOGLE_CLOUD_PROJECT=$PROJECT_ID@LOG_LEVEL=INFO@OTEL_SERVICE_NAME=${_RESOURCE_PREFIX}-gateway@DATABASE_POOL_SIZE=${_DATABASE_POOL_SIZE}@DATABASE_MAX_OVERFLOW=${_DATABASE_MAX_OVERFLOW}@LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID=${_LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID}@LEMON_SQUEEZY_SOLO_PRO_YEARLY_VARIANT_ID=${_LEMON_SQUEEZY_SOLO_PRO_YEARLY_VARIANT_ID}@LEMON_SQUEEZY_AGENCY_MONTHLY_VARIANT_ID=${_LEMON_SQUEEZY_AGENCY_MONTHLY_VARIANT_ID}@LEMON_SQUEEZY_AGENCY_YEARLY_VARIANT_ID=${_LEMON_SQUEEZY_AGENCY_YEARLY_VARIANT_ID}@LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL=${_LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL}@LEMON_SQUEEZY_SOLO_PRO_YEARLY_CHECKOUT_URL=${_LEMON_SQUEEZY_SOLO_PRO_YEARLY_CHECKOUT_URL}@LEMON_SQUEEZY_AGENCY_MONTHLY_CHECKOUT_URL=${_LEMON_SQUEEZY_AGENCY_MONTHLY_CHECKOUT_URL}@LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL=${_LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL}@COMPLIANCE_EMAIL_FROM=${_COMPLIANCE_EMAIL_FROM} + - --set-env-vars=^@^BACKEND_CORS_ORIGINS=${_FRONTEND_ORIGINS}@ENVIRONMENT=production@SECRET_BACKEND=gcp_secret_manager@GOOGLE_CLOUD_PROJECT=$PROJECT_ID@LOG_LEVEL=INFO@OTEL_SERVICE_NAME=${_RESOURCE_PREFIX}-gateway@DATABASE_POOL_SIZE=${_DATABASE_POOL_SIZE}@DATABASE_MAX_OVERFLOW=${_DATABASE_MAX_OVERFLOW}@COMPLIANCE_EMAIL_FROM=${_COMPLIANCE_EMAIL_FROM} - id: verify-gateway name: gcr.io/google.com/cloudsdktool/cloud-sdk@sha256:570dc7ce2876a2810dbf80f1d2520c1654aa218a7f1347715243a79c98d166ae @@ -473,14 +473,6 @@ substitutions: _RUNTIME_SERVICE_ACCOUNT: "" _MIGRATION_SERVICE_ACCOUNT: "" _FRONTEND_ORIGINS: https://shim-phi.vercel.app,https://getshim.tech,https://www.getshim.tech - _LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID: "" - _LEMON_SQUEEZY_SOLO_PRO_YEARLY_VARIANT_ID: "" - _LEMON_SQUEEZY_AGENCY_MONTHLY_VARIANT_ID: "" - _LEMON_SQUEEZY_AGENCY_YEARLY_VARIANT_ID: "" - _LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL: "" - _LEMON_SQUEEZY_SOLO_PRO_YEARLY_CHECKOUT_URL: "" - _LEMON_SQUEEZY_AGENCY_MONTHLY_CHECKOUT_URL: "" - _LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL: "" _COMPLIANCE_EMAIL_FROM: "" _VPC_NETWORK: "" _VPC_SUBNET: "" diff --git a/context7.json b/context7.json index 9ceb7d9..31c0527 100644 --- a/context7.json +++ b/context7.json @@ -14,7 +14,6 @@ ".github" ], "excludeFiles": [ - "docs/MIGRATION_PROGRESS.md", "AGENTS.md" ], "rules": [ diff --git a/docker-compose.prod.yml b/docker-compose.prod.yml index 4fa9a9f..d6386f4 100644 --- a/docker-compose.prod.yml +++ b/docker-compose.prod.yml @@ -11,15 +11,6 @@ x-runtime: &runtime ENCRYPTION_KEY: "${ENCRYPTION_KEY:?set ENCRYPTION_KEY}" SUPABASE_URL: "${SUPABASE_URL:?set SUPABASE_URL}" SUPABASE_KEY: "${SUPABASE_KEY:-}" - LEMON_SQUEEZY_SIGNING_SECRET: "${LEMON_SQUEEZY_SIGNING_SECRET:-}" - LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID: "${LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID:-}" - LEMON_SQUEEZY_SOLO_PRO_YEARLY_VARIANT_ID: "${LEMON_SQUEEZY_SOLO_PRO_YEARLY_VARIANT_ID:-}" - LEMON_SQUEEZY_AGENCY_MONTHLY_VARIANT_ID: "${LEMON_SQUEEZY_AGENCY_MONTHLY_VARIANT_ID:-}" - LEMON_SQUEEZY_AGENCY_YEARLY_VARIANT_ID: "${LEMON_SQUEEZY_AGENCY_YEARLY_VARIANT_ID:-}" - LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL: "${LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL:-}" - LEMON_SQUEEZY_SOLO_PRO_YEARLY_CHECKOUT_URL: "${LEMON_SQUEEZY_SOLO_PRO_YEARLY_CHECKOUT_URL:-}" - LEMON_SQUEEZY_AGENCY_MONTHLY_CHECKOUT_URL: "${LEMON_SQUEEZY_AGENCY_MONTHLY_CHECKOUT_URL:-}" - LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL: "${LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL:-}" RESEND_API_KEY: "${RESEND_API_KEY:-}" COMPLIANCE_EMAIL_FROM: "${COMPLIANCE_EMAIL_FROM:-}" SENTRY_DSN: "${SENTRY_DSN:-}" diff --git a/docs/CURRENT_ARCHITECTURE.md b/docs/CURRENT_ARCHITECTURE.md index c27e25e..f880cfe 100644 --- a/docs/CURRENT_ARCHITECTURE.md +++ b/docs/CURRENT_ARCHITECTURE.md @@ -105,7 +105,7 @@ create_enterprise_app |-- ManagedProviderCredentialResolver |-- DurableUsageLifecycle and accounting coordinator |-- enterprise scan pipeline and error composition -|-- management, subscription, shared-result, compliance, and AI Act routes +|-- management, shared-result, compliance, and AI Act routes `-- database, Redis, tracing, metrics, and lifecycle hooks ``` @@ -126,7 +126,7 @@ The exact method/path inventories live in `architecture/route_profiles.toml`. | Profile | Surface | Contract | | --- | --- | --- | | Community | OpenAI Chat and Responses; Anthropic Messages; Gemini generate and stream; model discovery; local scan; health | `openapi/community.json` | -| Enterprise | Community provider routes plus durable scan usage, management, subscriptions, shared results, compliance, and AI Act | `ee/openapi/enterprise.json` | +| Enterprise | Community provider routes plus durable scan usage, management, shared results, compliance, and AI Act | `ee/openapi/enterprise.json` | `/metrics` is intentionally excluded from OpenAPI. Enterprise provider routes must preserve the community provider request, response, selector, error, and @@ -134,6 +134,11 @@ stream contracts while adding enterprise authentication and lifecycle policy. ### Authentication and provider credentials +Enterprise OIDC and Vault deployment contracts are documented in +[`ee/docs/ON_PREM_IDENTITY.md`](../ee/docs/ON_PREM_IDENTITY.md). The on-prem +control plane uses configured issuer/subject identities and server-side Redis +sessions; hosted Supabase remains a separate selected authentication mode. + - OpenAI SDKs carry the shim key in `Authorization: Bearer ...`. - Anthropic SDKs carry the shim key in `x-api-key` on Anthropic routes. - `x-shim-key` is the explicit provider-independent gateway-key header. @@ -149,7 +154,7 @@ never forwarded wholesale. | Provider | Current route family | Native stream terminal | | --- | --- | --- | | OpenAI | `/v1/chat/completions`, `/v1/responses`, `/v1/models` | Chat ends with `[DONE]`; Responses uses named `response.*` events | -| Anthropic | `/v1/messages`, `/v1/models` | Native named events ending in `message_stop` | +| Anthropic | `/v1/messages`, `/v1/messages/count_tokens`, `/v1/models` | Messages use native named events ending in `message_stop`; token counting returns JSON | | Gemini | `/v1beta/models/{model}:generateContent` and stream | Data-only Gemini SSE, without `[DONE]` | OpenAI errors retain the safe `{error: {message, type, param, code}}` shape. @@ -159,8 +164,16 @@ details that could contain credentials or PII are discarded. A stream failure after headers is emitted as a sanitized terminal event. `background=true` Responses requests remain unsupported because shim has no -retrieval lifecycle with which to settle them safely. Explicit model IDs must -exist in the checked-in model and price catalog. +retrieval lifecycle with which to settle them safely. Community model IDs must exist in the checked-in model and price catalog. +Enterprise can resolve tenant aliases through its approved deployment registry; +`MODEL_DEPLOYMENT_REQUIRED=true` disables catalog fallback. Registry targets +reuse the native executions with operator-approved destinations and stored +credential references. Unpriced deployments remain explicit in accounting, and +monetary caps reject them. See [`MODEL_DEPLOYMENTS.md`](../ee/docs/MODEL_DEPLOYMENTS.md). + +Anthropic token counting shares authentication, registry authorization and +privacy transformation. It persists nonbillable enterprise audit preflight and +completion without quota/spend reservations or inference lifecycle settlement. ## Physical ownership @@ -259,8 +272,9 @@ those files in wheel and sdist metadata. Production enterprise boots verify an offline `SHIM_LICENSE_KEY` in `shim_enterprise.core.license`; no other runtime licence validator exists. -The public repository uses this mixed-licence package split. It is live in -production; deployment and verification evidence is recorded in -`MIGRATION_PROGRESS.md`. Production releases are automated from `main` through -the staged, serialized, exact-revision Cloud Build flow documented in -`TARGET_ARCHITECTURE.md`. +Production deployment is triggered by `v..` tags, not by a +merge to `main`. Cloud Build serializes migration and promotion with a shared +lock, validates staged gateway and worker revisions, and restores the captured +traffic splits on promotion failure. Release publication workflows do not deploy. +See [repository release rules](../AGENTS.md#release-and-deployment) and the +[customer-operated deployment guide](../ee/deploy/README.md). diff --git a/docs/MIGRATION_PROGRESS.md b/docs/MIGRATION_PROGRESS.md deleted file mode 100644 index b9e571b..0000000 --- a/docs/MIGRATION_PROGRESS.md +++ /dev/null @@ -1,271 +0,0 @@ -# Community and enterprise migration progress - -Status: migration implementation and initial production rollout complete - -Release branch: `main` - -Initial migrated production commit: `e0c7901` - -Architecture contract: [`TARGET_ARCHITECTURE.md`](./TARGET_ARCHITECTURE.md) - -## Goal - -Deliver one mixed-licence monorepo with: - -- a standalone `shim-gateway` distribution importing `shim`; and -- a `shim-enterprise` distribution under `ee/` importing `shim_enterprise` and - composing the exact matching community package. - -Both products must build and run from one commit and one lockfile. CI must reject -dependency, route, schema, and artifact leakage in either direction. - -## Non-negotiable invariants - -- `shim` never imports `shim_enterprise` or enterprise-only dependencies. -- Enterprise runtime imports only the public `shim` symbols declared in the - ownership manifest. -- Provider-native request, response, error, and stream contracts do not change - accidentally. -- One admitted request makes at most one provider attempt. -- Provider credentials remain invocation-scoped and absent from durable state - and telemetry. -- Enterprise stage order and transaction boundaries remain unchanged. -- PostgreSQL remains enterprise lifecycle and accounting truth; Redis remains - an accelerator and continuation store. -- Table names, Alembic revisions, constraints, functions, triggers, and - production data remain unchanged during package movement. -- Community runtime, artifacts, image, settings, behavior tests, direct - dependency metadata, and OpenAPI contain no enterprise code or assets. -- No compatibility package, generic plugin layer, provider-neutral wire model, - retry layer, microservice split, or duplicated gateway implementation. - -## Status - -| Phase | State | Remaining gate | -| --- | --- | --- | -| 0. Baseline and inventory | Complete | Historical DB-baseline limitation recorded | -| 1. Boundary guardrails | Complete | None | -| 2. Public contracts and enterprise adapters | Complete | None | -| 3. Community vertical slice | Complete | None | -| 4. Complete community API | Complete | None | -| 5. Physical package split | Complete | None | -| 6. Packages and deployment | Complete | None | -| 7. Final verification and handoff | Complete | None | -| Mixed licensing | Complete | None | -| Publication policy | Owner decision | Confirm final holder name and contribution/security channels | - -## Detailed checklist - -### Phase 0: baseline and ownership - -- [x] Create an isolated worktree and `migration/community-enterprise` branch - from `00f3ed7`. -- [x] Record the current and target architecture before movement. -- [x] Map routes, kernel stages, provider execution, streaming, accounting, - tenancy, secrets, migrations, workers, OpenAPI, packaging, and deployment. -- [x] Classify governed Python files as community, enterprise, or intentionally - shared tooling. -- [x] Record baseline Ruff and type results. -- [x] Record that a database-backed test baseline could not be captured before - semantic work because the original Docker environment lacked storage. Current - disposable-database results are tracked below; this historical result cannot - be recreated honestly. - -### Phase 1: executable guardrails - -- [x] Add `architecture/module_ownership.toml` as the file-ownership source of - truth. -- [x] Add exact community and enterprise route profiles. -- [x] Reject public imports of enterprise packages and infrastructure clients. -- [x] Detect absolute, relative, lazy, optional, type-checking, and recognized - dynamic imports with standard-library AST checks. -- [x] Fail closed on unclassified governed files and unresolved dynamic import - targets. -- [x] Separate community and enterprise OpenAPI route inventories. - -### Phase 2: public contracts and explicit composition - -- [x] Split community settings from enterprise settings. -- [x] Remove ORM records, sessions, global enterprise settings, and billing - objects from public request contracts. -- [x] Resolve authentication and tenant policy in short sessions before the - provider hot path. -- [x] Separate public provider/privacy errors from enterprise quota, audit, - tenancy, and reconciliation errors. -- [x] Extract only the contracts required by both products: request policy, - credentials, privacy continuation, usage lifecycle, admission, circuits, and - health. -- [x] Add bounded community adapters and preserve current enterprise adapters. -- [x] Preserve credential precedence, consume-once cleanup, stream - finalization, stage order, and transaction boundaries. -- [x] Break the managed-secret import cycle without adding a service locator or - plugin framework. -- [x] Obtain independent architecture review and close blocking findings. - -### Phase 3: community vertical slice - -- [x] Compose OpenAI Chat JSON and streaming without enterprise imports or - infrastructure. -- [x] Implement optional local shim key authentication and refuse keyless - non-loopback binds. -- [x] Implement invocation-first provider credentials. -- [x] Add bounded local admission, circuit, and privacy-continuation state. -- [x] Add redacted terminal JSONL usage events and nullable cost attribution. -- [x] Add community health and metrics. -- [x] Verify clean community startup with `ee/` absent. - -### Phase 4: complete community API - -- [x] Add OpenAI Responses. -- [x] Add Anthropic Messages and beta selection. -- [x] Add Gemini generate and stream routes. -- [x] Add provider-specific model discovery and exact catalog admission. -- [x] Add provider-free local privacy scan without quota, audit, or persistence - fields. -- [x] Preserve privacy restoration and provider-native streaming across all - protocols. -- [x] Generate `openapi/community.json` and reject enterprise route leakage. -- [x] Run SDK client, credential, error, privacy, stream, scan, catalog, and cost - tests without enterprise fixtures. - -### Phase 5: physical package split - -- [x] Move community runtime to `src/shim` and retain public tests under - `tests/`. -- [x] Move enterprise runtime to `ee/src/shim_enterprise` and enterprise tests - to `ee/tests/`. -- [x] Move enterprise Alembic history, operational script, workers, OpenAPI, - and AI Act assets under `ee/`. -- [x] Rename the application and worker entrypoints to their canonical package - modules. -- [x] Remove the legacy package rather than leave aliases or forwarding modules. -- [x] Preserve the two Alembic revision files, IDs, order, and content apart - from package imports. -- [x] Import every enterprise runtime module and run the moved enterprise suite. -- [x] Confirm all 177 runtime modules are free of import-time cycles. -- [x] Enforce the exact enterprise-to-community leaf module-and-symbol allowlist - in both directions: undeclared imports and stale manifest entries fail. -- [x] Obtain independent approval for the cross-package architecture boundary. -- [x] Obtain independent physical/package and persistence approval. No P0-P2 - finding remains; the accepted P3 enterprise test-only white-box imports do - not expand the runtime API. -- [x] Integrate the coordinated source, test, migration, guard, packaging, and - documentation changes in a runnable Git commit. - -### Phase 6: packages, schemas, images, and deployment - -- [x] Configure `shim-gateway` and workspace member `shim-enterprise` with one - authoritative `uv.lock` and exact matching versions. -- [x] Split direct dependencies and keep managed AWS/Azure secret clients as - enterprise extras. -- [x] Build and inspect both wheels and sdists. -- [x] Install the community wheel alone and both wheels together in clean - environments. -- [x] Generate and check community and enterprise OpenAPI from their own - composition roots. -- [x] Split `.env.example` and `ee/.env.example` by product. -- [x] Add a community-only root Dockerfile and an enterprise `ee/Dockerfile`. -- [x] Point Compose and Cloud Build at the canonical enterprise API, migration, - and worker entrypoints. -- [x] Split CI into isolated community and full enterprise jobs with package, - dependency, image, and OpenAPI boundary checks. -- [x] Build, inspect, and start both final images after all deployment edits are - integrated. -- [x] Validate development and production Compose configurations and all - deployment manifest entrypoints. -- [x] Confirm the enterprise OpenAPI has no schema diff; dashboard client - regeneration is therefore not required for this physical move. - -### Phase 7: final verification and handoff - -- [x] Run the locked full Ruff, formatting, type, architecture, public, and - enterprise test matrix on the integrated diff. -- [x] Rehearse fresh upgrade, current-head check, disposable downgrade, and - re-upgrade with `ee/alembic.ini`. -- [x] Verify the migrated production schema identity, grants, role search path, - Alembic head, gateway health, worker readiness, and background queues after - release. -- [x] Rebuild and inspect both wheels, sdists, images, and OpenAPI documents. -- [x] Start community with `ee/` and enterprise variables absent. -- [x] Start the enterprise API with PostgreSQL and Redis connected. -- [x] Import-smoke all four canonical worker module bootstraps with validated - enterprise settings; do not start their long-running loops in CI. -- [x] Recheck one-attempt, credential redaction, continuation fail-closed, - streaming finalization, and accounting recovery invariants. -- [x] Run `git diff --check`, review the full patch, and remove caches and - generated build artifacts. -- [x] Review the final commits and confirm a clean worktree. -- [x] Obtain final independent architecture, code-quality, artifact, and - security review; close all blocking findings. -- [x] Rewrite repository, developer, current-architecture, target-architecture, - migration, and agent guidance for the implemented layout. -- [x] Commit final integration and mark the technical goal complete. -- [x] Obtain owner authorization before public publication. - -### Mixed-licence completion - -- [x] Add canonical Apache-2.0 terms and scope them outside `ee/`. -- [x] Add canonical Elastic-2.0 terms under `ee/` with the licensor named in `ee/NOTICE`. -- [x] Add matching PEP 639 package metadata and legal-file declarations. -- [x] Include and compare legal files in both wheel and sdist profiles. -- [x] Keep the community image free of enterprise source and install both - separately licensed packages in the enterprise image. -- [x] Omit CLA and runtime licence validation until an approved policy requires - either. - -The holder name may be replaced later by the owner without changing the -architecture. Contribution and security-reporting policy remain publication -decisions, not package-boundary dependencies. - -## Verification evidence - -These results were recorded on the integrated worktree. They must be rerun if -review changes runtime, packaging, or deployment files. - -| Check | Latest evidence | -| --- | --- | -| Full suite | 778 tests passed from the locked all-packages environment against disposable PostgreSQL and Redis | -| Isolated community suite | 403 root tests passed from a community-only locked environment with no enterprise package or infrastructure dependencies | -| Enterprise-free public behavior suite | 356 tests passed in a scrubbed environment without database, Redis, tenant fixtures, or enterprise variables | -| Moved enterprise suite | 375 tests passed; all 105 enterprise runtime modules imported | -| Combined non-architecture suite | 733 tests passed after source and test movement | -| Integrated architecture suites | 45 tests passed across root and enterprise architecture tests | -| Critical runtime invariants | 150 credential, one-attempt, provider, privacy-continuation, streaming, and accounting-recovery tests passed | -| Independent architecture-boundary review | Approved with no blocking finding | -| Independent physical/package and persistence review | Approved with no P0-P2 finding; Alembic blobs byte-identical and history preserved; accepted test-only P3 documented | -| Final principal architecture review | Approved with no remaining P0-P3 finding | -| Final code-quality and security review | Approved with no remaining P0-P3 finding | -| Ruff, format, type, OpenAPI, lock, diff checks | Passed at completed implementation milestones | -| Community artifacts | Wheel and sdist contained `shim`, model data, Apache-2.0 metadata, `LICENSE`, and `NOTICE` with no enterprise paths | -| Workspace artifacts | Both wheel/sdist pairs built; enterprise carried Elastic-2.0 metadata and its exact legal files; enterprise YAML assets remained present | -| Alembic | `head -> base -> head`, history, current, and model import passed on disposable PostgreSQL; revision hashes preserved | -| Community image | Import, filesystem, direct-dependency, Apache-2.0 metadata, and live `/health` checks passed | -| Enterprise image | Both packages and licence regions, migrations, control assets, worker modules, dependency boundary, and live database/Redis `/health` checks passed | -| Enterprise worker bootstraps | All four canonical modules imported from the clean enterprise wheel; image module specs present | -| Compose and deployment manifests | Development and production Compose rendered; candidate staging, serialized exact-revision promotion, runtime health, and guarded restoration assertions passed | -| Initial production rollout | Gateway and four worker pools run commit `e0c7901`; PostgreSQL and Redis health passed and all background queues were clear | -| Split CI profiles | Isolated community and full enterprise jobs include exact licence metadata/file assertions and were locally mirrored | -| Exact cross-package import guard | Passed for 37 leaf modules and 62 symbols; undeclared imports, stale entries, broad imports, and unresolved dynamic imports fail | -| Import-time cycle audit | All 177 community and enterprise runtime modules passed | - -The unchanged baseline migration cannot render offline SQL because its existing -JSONB default renderer requires a live dialect. Online migration checks are the -authoritative gate; fixing that unrelated baseline behavior is outside this -migration. - -## Git milestones - -| Commit | Milestone | -| --- | --- | -| `e582fe5` | Runnable community gateway | -| `aedb485` | Community Responses API | -| `ee11460` | Community Anthropic Messages | -| `ea14fac` | Community Gemini API | -| `59e56ad` | Community privacy scan | -| `a1f829c` | Separate community and enterprise OpenAPI | -| `a004c96` | Separate public and enterprise scan contracts | -| `f8d4e25` | Separate public and enterprise metrics | -| `cddb4d3` | Move community package to `src/shim` | -| `8d00cce` | Move enterprise package, tests, migrations, packaging, and deployment under `ee/` | -| `1a618d0` | Finalize the mixed-licence boundary | -| `55bcefe` | Merge the migration into `main` | diff --git a/docs/TARGET_ARCHITECTURE.md b/docs/TARGET_ARCHITECTURE.md deleted file mode 100644 index fccc8a4..0000000 --- a/docs/TARGET_ARCHITECTURE.md +++ /dev/null @@ -1,321 +0,0 @@ -# shim target architecture - -Status: approved and implemented architecture - -Last reviewed: 2026-08-28 - -This is the architectural contract for the shim backend. Implementation and -release evidence are tracked in -[`MIGRATION_PROGRESS.md`](./MIGRATION_PROGRESS.md). - -## Decision - -Use one public, mixed-licence monorepo containing two Python packages: - -| Product | Distribution | Import package | Intended licence region | -| --- | --- | --- | --- | -| Community | `shim-gateway` | `shim` | Apache-2.0 | -| Enterprise | `shim-enterprise` | `shim_enterprise` | Elastic-2.0, source-available | - -The architecture is a package-modular monolith. It takes Polylith's useful -ideas—small capabilities, explicit composition, and testable boundaries—but -does not adopt Polylith tooling, a component taxonomy, or a package per feature. - -```text -one repository / one lock / one release pair - - community package enterprise package - +---------------------------+ +---------------------------+ - | src/shim |<------------| ee/src/shim_enterprise | - | native provider gateway | allowlist | tenant/durable adapters | - | privacy and local policy | | billing/audit/compliance | - +---------------------------+ +---------------------------+ - - allowed: shim_enterprise -> shim - forbidden: shim -> shim_enterprise -``` - -This keeps cross-product changes atomic and makes the licence boundary visible -in the filesystem. It is not a rewrite, microservice split, plugin platform, or -attempt to hide enterprise source. - -## Goals - -- Make community useful, installable, and runnable without enterprise code or - infrastructure. -- Preserve existing enterprise behavior, schema, accounting, privacy, tenant - isolation, streaming, and workers. -- Keep one implementation of provider transport, privacy, and the inference - kernel. -- Make every cross-region dependency explicit and machine-enforced. -- Produce separate packages, schemas, images, settings, and tests from one - commit and one lockfile. -- Give humans and coding agents an obvious owner and verification path for each - change. - -## Repository layout - -```text -shim/ -|-- src/shim/ -| |-- application.py # community composition root -| |-- api/v1/ # provider-native and local scan routes -| |-- gateway/ # contracts, kernel, stages, streaming -| |-- privacy/ # classification, masking, restoration -| |-- billing/ # public catalog and cost attribution -| |-- secrets/ # invocation/local credential contract -| `-- observability/ # public logging, tracing, metrics -|-- tests/ # community and boundary tests -|-- openapi/community.json -|-- Dockerfile # community image -| -|-- ee/ -| |-- src/shim_enterprise/ -| | |-- application.py # enterprise composition root -| | |-- tenants/ # identity and policy -| | |-- billing/ # ledger, quota, budgets, spend -| | |-- gateway/ # durable adapters and enterprise scan -| | |-- secrets/ # managed secret stores -| | |-- outbox/ # durable external effects -| | |-- compliance/ and ai_act/ -| | `-- workers/ -| |-- tests/ -| |-- alembic.ini and alembic/ -| |-- scripts/ -| |-- openapi/enterprise.json -| |-- pyproject.toml -| |-- Dockerfile -| `-- AGENTS.md -| -|-- architecture/ # executable ownership and route rules -|-- scripts/ # cross-profile export and public catalog -|-- pyproject.toml # community package + workspace -|-- uv.lock # one authoritative lock -`-- AGENTS.md -``` - -The tree is the primary ownership signal. Do not add another layer merely to -mirror this diagram. - -## Runtime architecture - -Both products compose the same public request path: - -```text -native request - | - v -provider route + authentication - | - v -GatewayService - | - v -GatewayKernel - | - +-- policy and principal - +-- admission and loop control - +-- privacy transform - +-- provider-start lifecycle marker - +-- exactly one provider attempt - `-- restore, meter, and finalize - | - v -native JSON or provider-specific SSE -``` - -Provider payloads remain native dictionaries or SDK objects at the boundary. -Public contracts contain plain values, not ORM records, database sessions, -enterprise settings, or provider credentials. - -### Community composition - -```text -create_community_app - = public routes and kernel - + local authentication and policy - + invocation/environment credentials - + bounded in-memory admission, circuits, and continuation state - + redacted local usage events and public cost catalog -``` - -Community is intentionally single-process and ephemeral. It requires no -PostgreSQL, Redis, Supabase, enterprise variable, migration, or worker. Add -distributed community state only for a demonstrated use case. - -### Enterprise composition - -```text -create_enterprise_app - = community kernel and provider transports - + database authentication and tenant policy - + durable quota, spend, audit, outbox, and reconciliation - + Redis acceleration and encrypted continuation state - + managed secrets - + management, compliance, AI Act, and workers -``` - -Enterprise adds behavior through explicit construction. Public code must not -contain commercial feature flags, licence checks, optional enterprise imports, -or plugin discovery. The licence check lives in the enterprise application -factory alone; the community distribution is the unlicensed tier. - -## Capability ownership - -| Capability | `shim` | `shim_enterprise` | -| --- | --- | --- | -| Provider-native HTTP and SDK transport | Owns | Reuses | -| Gateway kernel, streaming, safe errors | Owns | Reuses | -| PII classification, masking, restoration | Owns | Configures durable state/policy | -| Local authentication, scan, usage events | Owns | Replaces with tenant/durable adapters | -| Model catalog and cost attribution | Owns | Reuses for settlement inputs | -| Tenant identity, RBAC, organizations | Does not contain | Owns | -| Ledger, quota, budgets, reconciliation | Does not contain | Owns | -| Audit, outbox, compliance, AI Act | Does not contain | Owns | -| SQLAlchemy models and Alembic history | Does not require | Owns | -| Enterprise management and workers | Does not contain | Owns | - -For an ambiguous capability, choose community only when a local developer can -use it without enterprise infrastructure. Otherwise keep it under `ee/`. - -## Dependency rules - -1. `shim` never imports `shim_enterprise` or enterprise-only dependencies, - including lazy, optional, and type-checking imports. -2. Enterprise runtime imports only the exact leaf `shim` modules and symbols - listed in `architecture/module_ownership.toml`; broad eager facades are not - the cross-licence API. -3. A new cross-region symbol is an architectural API change: add it only with - its consumer and boundary test. -4. Provider SDK objects stay in execution adapters. Provider JSON is not - normalized into a shared request model. -5. Infrastructure differences use the existing narrow contracts for policy, - credentials, continuation state, usage lifecycle, admission, and circuits. -6. Do not create `common`, `utils`, re-export facades, or compatibility packages - to bypass ownership. - -Standard-library AST tests compare the enterprise runtime import graph exactly -with the manifest. They reject undeclared imports, stale entries, star or bare -module imports, and unresolved recognized dynamic imports. The ownership -manifest also rejects unclassified Python files. - -## Data and transaction rules - -Community has no required database. Enterprise owns the current SQLAlchemy -models and complete Alembic history under `ee/`. - -Moving package ownership must not change table names, columns, identifiers, -revision IDs, `down_revision` links, `alembic_version`, indexes, constraints, -triggers, functions, grants, or production data. - -PostgreSQL remains the enterprise accounting truth. Redis may accelerate burst -limits, loop detection, circuit state, policy reads, and privacy continuation; -failure must not make it an alternate ledger. No database transaction spans a -provider call. Production rollback uses a previous schema-compatible image, -not an Alembic downgrade. - -## API and artifact profiles - -| Profile | OpenAPI | Artifact boundary | -| --- | --- | --- | -| Community | `openapi/community.json` | `shim` wheel/sdist and community image; no `ee/`, enterprise package, migrations, settings, or direct dependencies | -| Enterprise | `ee/openapi/enterprise.json` | Exact-version `shim` + `shim-enterprise`, enterprise image, migrations, scripts, and required assets | - -Enterprise provider routes keep the same native wire and stream contracts as -community while adding tenant policy and durable lifecycle behavior. Enterprise -may add routes; community must never expose them. - -The distributions share a release version. `shim-enterprise X.Y.Z` requires -`shim-gateway X.Y.Z`. One `uv.lock` is authoritative for repository development -and CI. - -## Build and deployment - -| Process | Canonical entrypoint | -| --- | --- | -| Community API | `shim serve` | -| Enterprise API | `uvicorn shim_enterprise.application:create_enterprise_app --factory` | -| Enterprise migration | `alembic -c ee/alembic.ini upgrade head` | -| Workers | `python -m shim_enterprise.workers.{outbox,reconciliation,compliance,ai_act}` | - -The root Dockerfile is community-only. `ee/Dockerfile` contains both installed -packages and enterprise runtime assets. Root Compose and Cloud Build remain -enterprise deployment definitions and use only canonical enterprise paths. - -The `main` Cloud Build trigger is the production release path. It builds one -commit-tagged enterprise image, then acquires a generation-guarded Cloud Storage -lock before Alembic or any Cloud Run mutation. It stages commit-named gateway -and worker revisions without traffic and verifies the temporarily tagged -gateway. Each exact worker revision must become active, emit its runtime signal, -and remain error-free before the exact gateway revision receives traffic. A -failed promotion restores only the active splits captured by that build; -successful releases verify the default gateway URL and remove the temporary -release tag. - -## Mixed-licence boundary - -The repository is public, so enterprise is source-available rather than secret. -No committed region may contain credentials, customer data, or confidential -customer integrations. - -The checked-in boundary is authoritative: - -- `LICENSE` contains Apache-2.0 and `NOTICE` applies it outside `ee/`; -- `ee/LICENSE` contains Elastic-2.0 and `ee/NOTICE` identifies the - licensor for files under `ee/`; -- each `pyproject.toml` declares its own SPDX expression and legal files; and -- CI compares the legal files and metadata embedded in wheels and sdists with - their source files. - -A repository archive necessarily contains both licence regions. Community -binary/source artifacts exclude `ee/` and carry only the community terms. The -enterprise distribution carries its own terms and depends on the separately -licensed community distribution. The enterprise image contains both installed -packages and both packages' metadata. - -No CLA is part of this architecture. `create_enterprise_app` verifies an -Ed25519-signed `SHIM_LICENSE_KEY` when `ENVIRONMENT` is `production`, offline -and against a public key shipped in the package. A missing or forged key stops -the boot; an expired key boots for thirty more days behind a warning log and -the `license_days_remaining` gauge. No other bootstrap validator exists. - -## Agent-friendly enforcement - -- `AGENTS.md` states repository-wide invariants; `ee/AGENTS.md` adds persistence - and commercial-boundary rules. -- Ownership and supported imports are data in `architecture/`, not prose-only - conventions. -- Tests live with their licence region and each OpenAPI profile has one - deterministic exporter. -- Composition roots reveal concrete adapters; there is no hidden service - locator or runtime plugin graph. -- One focused change should normally touch one owner. A public contract change - updates both compositions and boundary tests atomically. - -## Non-goals - -- Microservices or independently deployed capability packages. -- Canonical Polylith layout or tooling. -- A generic plugin/factory framework. -- A provider-neutral request schema. -- Compatibility modules for superseded package paths. -- Database redesign during package movement. -- Hiding enterprise source in a separate repository. - -## Definition of done - -Technical completion requires: - -- both packages build and install from one locked commit; -- community starts with `ee/` and enterprise infrastructure absent; -- enterprise API, migrations, and four workers use canonical entrypoints; -- reverse imports, unsupported public imports, route leakage, dependency - leakage, and artifact leakage fail CI; -- both OpenAPI profiles and all public/enterprise behavior tests pass; -- schema identifiers and enterprise accounting invariants remain unchanged; and -- independent architecture, persistence, artifact, and security review has no - blocking finding. - -Licence completion additionally requires exact package metadata and legal-file -checks in built artifacts. The licensor name and contributor/security channels -remain owner-controlled publication details. diff --git a/ee/.env.example b/ee/.env.example index a60c25a..87e0a75 100644 --- a/ee/.env.example +++ b/ee/.env.example @@ -43,6 +43,8 @@ LOOP_WINDOW_SECONDS=300 GATEWAY_RECONCILIATION_GRACE_SECONDS=120 GATEWAY_RECONCILIATION_INTERVAL_SECONDS=30 GATEWAY_RECONCILIATION_BATCH_SIZE=100 +# Optional per-process readiness file; leave unset to disable. Parent directory must exist. +# WORKER_HEARTBEAT_PATH=/tmp/shim-worker.json GATEWAY_OUTBOX_INTERVAL_SECONDS=5 GATEWAY_OUTBOX_BATCH_SIZE=100 GATEWAY_OUTBOX_LEASE_SECONDS=60 @@ -66,13 +68,6 @@ OVERSIGHT_ENABLED=false OVERSIGHT_DEFAULT_TTL_SECONDS=3600 MANUAL_TEST_DASHBOARD_ENABLED=false -# Optional billing integration. -# LEMON_SQUEEZY_SIGNING_SECRET=REPLACE_ME -# LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID=REPLACE_ME -# LEMON_SQUEEZY_SOLO_PRO_YEARLY_VARIANT_ID=REPLACE_ME -# LEMON_SQUEEZY_AGENCY_MONTHLY_VARIANT_ID=REPLACE_ME -# LEMON_SQUEEZY_AGENCY_YEARLY_VARIANT_ID=REPLACE_ME - # Optional observability exporters. # SENTRY_DSN=https://REPLACE_ME # OTEL_EXPORTER_OTLP_ENDPOINT=https://REPLACE_ME diff --git a/ee/Dockerfile b/ee/Dockerfile index a66857d..1ea6c26 100644 --- a/ee/Dockerfile +++ b/ee/Dockerfile @@ -4,6 +4,7 @@ COPY --from=ghcr.io/astral-sh/uv:0.12.5@sha256:e85be844203885286c60ffad8a858d48a ENV UV_COMPILE_BYTECODE=1 \ UV_LINK_MODE=copy \ + UV_NO_CACHE=1 \ UV_PYTHON_DOWNLOADS=0 WORKDIR /app diff --git a/ee/alembic/versions/c31b7a91d602_oidc_identity.py b/ee/alembic/versions/c31b7a91d602_oidc_identity.py new file mode 100644 index 0000000..fea6c00 --- /dev/null +++ b/ee/alembic/versions/c31b7a91d602_oidc_identity.py @@ -0,0 +1,23 @@ +"""Bind customer OIDC issuer and subject without email-based linking.""" + +from alembic import op +import sqlalchemy as sa + +revision = "c31b7a91d602" +down_revision = "c7a108be3201" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column("users", sa.Column("oidc_issuer", sa.String(512), nullable=True)) + op.add_column("users", sa.Column("oidc_subject", sa.String(255), nullable=True)) + op.create_unique_constraint( + "uq_users_oidc_identity", "users", ["oidc_issuer", "oidc_subject"] + ) + + +def downgrade() -> None: + op.drop_constraint("uq_users_oidc_identity", "users", type_="unique") + op.drop_column("users", "oidc_subject") + op.drop_column("users", "oidc_issuer") diff --git a/ee/alembic/versions/c7a108be3201_teams_roles_and_scoped_key_controls.py b/ee/alembic/versions/c7a108be3201_teams_roles_and_scoped_key_controls.py new file mode 100644 index 0000000..78ac70d --- /dev/null +++ b/ee/alembic/versions/c7a108be3201_teams_roles_and_scoped_key_controls.py @@ -0,0 +1,188 @@ +"""teams roles and scoped key controls + +Revision: c7a108be3201 +Parent: f10e4ac92d17 +Created: 2026-09-08 02:49:20.778421 +""" + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +revision: str = "c7a108be3201" +down_revision: str | Sequence[str] | None = "f10e4ac92d17" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + for table, constraint in ( + ("users", "ck_users_role"), + ("organization_invites", "ck_invites_role"), + ): + op.drop_constraint(constraint, table, type_="check") + op.create_check_constraint( + constraint, table, "role IN ('owner', 'admin', 'member', 'auditor')" + ) + op.create_table( + "teams", + sa.Column("id", sa.UUID(), nullable=False), + sa.Column("organization_id", sa.UUID(), nullable=False), + sa.Column("name", sa.String(length=128), nullable=False), + sa.Column("daily_request_limit", sa.Integer(), nullable=True), + sa.Column("monthly_request_limit", sa.Integer(), nullable=True), + sa.Column("monthly_token_limit", sa.Integer(), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.CheckConstraint( + "daily_request_limit IS NULL OR daily_request_limit >= 0", + name="ck_teams_daily_requests", + ), + sa.CheckConstraint( + "monthly_request_limit IS NULL OR monthly_request_limit >= 0", + name="ck_teams_monthly_requests", + ), + sa.CheckConstraint( + "monthly_token_limit IS NULL OR monthly_token_limit >= 0", + name="ck_teams_monthly_tokens", + ), + sa.ForeignKeyConstraint( + ["organization_id"], ["organizations.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("organization_id", "id", name="uq_teams_tenant_id"), + sa.UniqueConstraint("organization_id", "name", name="uq_teams_tenant_name"), + ) + op.create_table( + "team_memberships", + sa.Column("organization_id", sa.UUID(), nullable=False), + sa.Column("team_id", sa.UUID(), nullable=False), + sa.Column("user_id", sa.UUID(), nullable=False), + sa.Column( + "role", sa.String(length=16), server_default="member", nullable=False + ), + sa.Column( + "source", sa.String(length=16), server_default="local", nullable=False + ), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.CheckConstraint( + "role IN ('member', 'team_admin')", name="ck_team_memberships_role" + ), + sa.CheckConstraint( + "source IN ('local', 'oidc')", name="ck_team_memberships_source" + ), + sa.ForeignKeyConstraint( + ["organization_id", "team_id"], + ["teams.organization_id", "teams.id"], + name="fk_team_memberships_tenant_team", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["organization_id", "user_id"], + ["users.organization_id", "users.id"], + name="fk_team_memberships_tenant_user", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["organization_id"], ["organizations.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("organization_id", "team_id", "user_id"), + ) + op.add_column("api_keys", sa.Column("team_id", sa.UUID(), nullable=True)) + op.add_column( + "api_keys", + sa.Column( + "allowed_models", postgresql.JSONB(astext_type=sa.Text()), nullable=True + ), + ) + op.create_foreign_key( + "fk_api_keys_tenant_team", + "api_keys", + "teams", + ["organization_id", "team_id"], + ["organization_id", "id"], + ) + # Billing labels retain their exact spelling and never grant membership. + op.execute( + "INSERT INTO teams (id, organization_id, name) SELECT gen_random_uuid(), organization_id, team FROM api_keys WHERE team IS NOT NULL GROUP BY organization_id, team" + ) + op.add_column("quota_period_usage", sa.Column("team_id", sa.UUID(), nullable=True)) + op.alter_column( + "quota_period_usage", "api_key_id", existing_type=sa.UUID(), nullable=True + ) + op.create_index( + "uq_quota_period_usage_team_scope", + "quota_period_usage", + ["organization_id", "team_id", "period_type", "period_start"], + unique=True, + postgresql_where=sa.text("team_id IS NOT NULL"), + ) + op.create_foreign_key( + "fk_quota_period_usage_org_team", + "quota_period_usage", + "teams", + ["organization_id", "team_id"], + ["organization_id", "id"], + ) + op.create_check_constraint( + "ck_quota_period_usage_single_scope", + "quota_period_usage", + "(api_key_id IS NULL) <> (team_id IS NULL)", + ) + + +def downgrade() -> None: + # Downgrades are for disposable data only; team allocations need the new schema. + op.execute("DELETE FROM quota_period_usage WHERE team_id IS NOT NULL") + for table, constraint in ( + ("users", "ck_users_role"), + ("organization_invites", "ck_invites_role"), + ): + op.execute(f"UPDATE {table} SET role = 'member' WHERE role = 'auditor'") + op.drop_constraint(constraint, table, type_="check") + op.create_check_constraint( + constraint, table, "role IN ('owner', 'admin', 'member')" + ) + op.drop_constraint( + "ck_quota_period_usage_single_scope", "quota_period_usage", type_="check" + ) + op.drop_constraint( + "fk_quota_period_usage_org_team", "quota_period_usage", type_="foreignkey" + ) + op.drop_index( + "uq_quota_period_usage_team_scope", + table_name="quota_period_usage", + postgresql_where=sa.text("team_id IS NOT NULL"), + ) + op.alter_column( + "quota_period_usage", "api_key_id", existing_type=sa.UUID(), nullable=False + ) + op.drop_column("quota_period_usage", "team_id") + op.drop_constraint("fk_api_keys_tenant_team", "api_keys", type_="foreignkey") + op.drop_column("api_keys", "allowed_models") + op.drop_column("api_keys", "team_id") + op.drop_table("team_memberships") + op.drop_table("teams") diff --git a/ee/alembic/versions/fcc02bbe6443_add_tenant_model_deployment_registry.py b/ee/alembic/versions/fcc02bbe6443_add_tenant_model_deployment_registry.py new file mode 100644 index 0000000..674e560 --- /dev/null +++ b/ee/alembic/versions/fcc02bbe6443_add_tenant_model_deployment_registry.py @@ -0,0 +1,88 @@ +"""Add tenant model deployment registry + +Revision: fcc02bbe6443 +Parent: c31b7a91d602 +Created: 2026-09-08 02:51:52.112876 +""" + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "fcc02bbe6443" +down_revision: str | Sequence[str] | None = "c31b7a91d602" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.create_unique_constraint( + "uq_provider_secrets_tenant_id", "provider_secrets", ["organization_id", "id"] + ) + op.create_table( + "model_deployments", + sa.Column("id", sa.UUID(), nullable=False), + sa.Column("organization_id", sa.UUID(), nullable=False), + sa.Column("alias", sa.String(length=200), nullable=False), + sa.Column("provider", sa.String(length=32), nullable=False), + sa.Column("upstream_model", sa.String(length=200), nullable=False), + sa.Column("base_url", sa.String(length=2048), nullable=False), + sa.Column("provider_secret_id", sa.UUID(), nullable=False), + sa.Column("timeout_seconds", sa.Integer(), nullable=False), + sa.Column("deployment_kind", sa.String(length=16), nullable=False), + sa.Column("declared_version", sa.String(length=200), nullable=False), + sa.Column("owner", sa.String(length=200), nullable=False), + sa.Column("enabled", sa.Boolean(), server_default="true", nullable=False), + sa.Column( + "health", sa.String(length=16), server_default="unknown", nullable=False + ), + sa.Column("health_checked_at", sa.DateTime(timezone=True), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.CheckConstraint( + "deployment_kind IN ('internal', 'external')", + name="ck_model_deployments_kind", + ), + sa.CheckConstraint( + "health IN ('unknown', 'healthy', 'unhealthy')", + name="ck_model_deployments_health", + ), + sa.CheckConstraint( + "provider IN ('openai', 'anthropic')", name="ck_model_deployments_provider" + ), + sa.CheckConstraint( + "timeout_seconds BETWEEN 1 AND 300", name="ck_model_deployments_timeout" + ), + sa.ForeignKeyConstraint( + ["organization_id", "provider_secret_id"], + ["provider_secrets.organization_id", "provider_secrets.id"], + name="fk_model_deployments_tenant_secret", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["organization_id"], ["organizations.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "organization_id", "alias", name="uq_model_deployments_tenant_alias" + ), + ) + + +def downgrade() -> None: + op.drop_table("model_deployments") + op.drop_constraint( + "uq_provider_secrets_tenant_id", "provider_secrets", type_="unique" + ) diff --git a/ee/deploy/README.md b/ee/deploy/README.md new file mode 100644 index 0000000..99bcd00 --- /dev/null +++ b/ee/deploy/README.md @@ -0,0 +1,250 @@ +# Customer-operated deployment + +This chart runs the enterprise API, dashboard, outbox, reconciliation, compliance, +and audit-maintenance workers. PostgreSQL and Redis remain customer-operated; +optional single-instance services are provided for disposable installations. +The chart has no downloaded chart dependencies. Use the signed release's image +digests and its matching chart, not a mutable `latest` tag. + +Related guides: [identity and Vault](../docs/ON_PREM_IDENTITY.md), +[organization provisioning](../docs/PROVISIONING.md), [teams and keys](../docs/team-access.md), +[model deployments](../docs/MODEL_DEPLOYMENTS.md), [diagnostic metadata](../docs/DIAGNOSTIC_METADATA.md), +[decision evidence](../docs/POLICY_DECISIONS.md), and [offline releases](../docs/OFFLINE_RELEASES.md). + +## Prerequisites + +- Kubernetes 1.28 or later, Helm, a default storage class when using test services, + and an ingress controller when exposing the supplied Ingress. +- A dedicated namespace, customer DNS/TLS, PostgreSQL, Redis, OIDC, and Vault KV v2. + The pinned test image is Redis Stack; the runtime uses standard Redis commands. +- The current signed `SHIM_LICENSE_KEY`. Licence validation remains offline and + runs only at production API startup; this package adds no capacity terms. +- Read [identity setup](../docs/ON_PREM_IDENTITY.md) for OIDC client registration, + group mapping, owner recovery, workload-key revocation, and Vault policy. +- An independently trusted release-verification public key. Verify the delivered + artifacts before importing images or running their scripts. + +Use one HTTPS origin for the dashboard and gateway. The ingress sends `/v1` and +`/v1beta` to the API, and other paths to the dashboard. The dashboard forwards +`/api/v1` internally at runtime using `SHIM_API_URL`; browser sessions and +management requests stay on the same origin. Build the dashboard with +`NEXT_PUBLIC_AUTH_MODE=oidc`; changing that mode requires a different image. +Configure ingress access logs to omit authentication query strings and cookies. + +## Configuration and first installation + +Create a namespace and an environment file outside source control with mode 0600. +Create `shim-runtime` from that file using +`kubectl -n shim create secret generic shim-runtime --from-env-file=/secure/shim.env`. +Supply these settings (values shown here are examples, not working credentials): + +```dotenv +DATABASE_URL=postgresql+asyncpg://shim:REPLACE@postgres.customer.example/shim +REDIS_URL=redis://redis.customer.example:6379/0 +SECRET_KEY=REPLACE_WITH_A_RANDOM_SECRET_OF_AT_LEAST_32_BYTES +SHIM_LICENSE_KEY=REPLACE_WITH_THE_ISSUED_LICENCE +OIDC_ISSUER_URL=https://identity.customer.example/realms/shim +OIDC_CLIENT_ID=shim +OIDC_CLIENT_SECRET=REPLACE +OIDC_REDIRECT_URI=https://shim.customer.example/api/v1/auth/callback +DASHBOARD_ORIGIN=https://shim.customer.example +OIDC_ORGANIZATION_ID=00000000-0000-4000-8000-000000000001 +OIDC_GROUP_ROLE_MAP={"/shim/owners":"owner","/shim/members":"member","/shim/auditors":"auditor"} +VAULT_ADDR=https://vault.customer.example +MODEL_DEPLOYMENT_ALLOWED_ORIGINS=["https://model-a.customer.example","https://model-b.customer.example"] +``` + +The initial organization UUID above is a placeholder until provisioning below. +The chart sets `ENVIRONMENT=production`, `AUTH_MODE=oidc`, `SECRET_BACKEND=vault`, +and `MODEL_DEPLOYMENT_REQUIRED=true`. It disables Sentry and OTLP export by default. +Do not select development mode to bypass production secret or licence validation. + +Create a `shim-vault-token` Secret with a `token` key, refreshed by your Vault Agent +or secret controller. Mounting a projected Secret permits token rotation without +rebuilding the image. The Vault adapter reads the current token for each operation. +For private CAs, create a ConfigMap containing `ca.pem`, a PEM bundle including all +required roots; the chart sets `SSL_CERT_FILE` and `NODE_EXTRA_CA_CERTS`. Never +disable certificate verification. PostgreSQL TLS configuration remains part of +the database URL/customer database policy. + +Create a private `values.yaml`: + +```yaml +existingSecret: shim-runtime +gateway: + image: registry.customer.example/shim-enterprise@sha256:REPLACE +dashboard: + image: registry.customer.example/shim-dashboard@sha256:REPLACE +vault: + tokenSecretName: shim-vault-token +caBundleConfigMap: shim-ca +imagePullSecrets: + - name: customer-registry +ingress: + enabled: true + className: customer-ingress + host: shim.customer.example + tlsSecretName: shim-tls +networkPolicy: + extraEgress: [] # Add customer service/proxy destination rules before installing. +``` + +The default egress policy permits pods in the same namespace and cluster DNS. +Add only required external namespace/IP and port rules for PostgreSQL, Redis, +OIDC, Vault, model endpoints and any explicitly enabled connector. NetworkPolicy +requires an enforcing CNI; verify actual denied traffic, not just resource creation. +For outbound proxies, add `HTTPS_PROXY`, `HTTP_PROXY`, and `NO_PROXY` to the runtime +Secret; include internal service names/addresses in `NO_PROXY` and allow the proxy +in the network policy. The browser also needs access to its customer OIDC origin. + +```sh +helm upgrade --install shim ./chart --namespace shim --create-namespace \ + --values /secure/values.yaml --wait --wait-for-jobs --timeout 10m +kubectl -n shim exec deployment/shim-gateway -- \ + python ee/scripts/activate_plan.py --create-name 'Customer organization' enterprise +``` + +The migration Job upgrades the database before the applications' schema checks +allow startup. Kubernetes retries failed schema init checks. The migration Job's +default deadline is five minutes; set `migration.activeDeadlineSeconds` from the +measured migration duration and give Helm a longer timeout for application startup. +Provisioning prints the organization UUID; update +`OIDC_ORGANIZATION_ID` in the runtime Secret to that value. Restart all backend +deployments after changing environment-backed configuration: + +```sh +kubectl -n shim rollout restart deployment/shim-gateway deployment/shim-outbox \ + deployment/shim-reconciliation deployment/shim-compliance deployment/shim-ai-act +kubectl -n shim rollout status deployment/shim-gateway --timeout=5m +``` + +For an existing organization, activate its existing UUID instead of creating a +duplicate. Sign in with a mapped owner group, create teams, register two internal +model deployments using Vault-backed credentials, then issue a scoped gateway key. +Exercise one JSON and one streaming request and verify usage and audit visibility. +See [team access](../docs/team-access.md) and [decision evidence](../docs/POLICY_DECISIONS.md). + +## Health and operating limits + +API readiness checks `/health`; dashboard readiness checks `/login`. Each worker +is ready only after a successful processing pass. Failed/partial passes do not +refresh its local heartbeat. Adjust `workerReadinessMaxAgeSeconds` when changing +worker intervals or measured pass duration (defaults allow the hourly audit job). +Readiness is not proof that every outbox message was delivered; monitor backlog, +dead letters, reconciliation lag, database capacity and the existing metrics. +Worker readiness failures do not trigger restart loops during a database outage. + +All application containers run without root, elevated capabilities, a writable +root filesystem, or a Kubernetes service-account token. Writable temporary and +dashboard cache volumes are bounded. Set `gateway.resources`, `dashboard.resources` +and replica counts from measured traffic. Workers remain independent processes. +Compliance ingestion only contacts configured active connectors; leave none active +in an isolated installation. Email, exporters and external model endpoints are +explicit operator choices, not prerequisites for login or inference. + +For disposable local databases, enable `postgres.enabled` and `redis.enabled`; +set `POSTGRES_PASSWORD` in `shim-runtime`, point `DATABASE_URL` at `shim-postgres` +with database/user `shim`, and `REDIS_URL` at `shim-redis:6379`. Their image digests +and persistent storage sizes are configurable. These single instances are not an +HA or database backup solution. + +## Upgrade, backup and recovery + +Before upgrading, record the chart version, both image digests, current Alembic +revision, runtime configuration version and a tested PostgreSQL backup. Preserve +Vault data/keys and the application `SECRET_KEY` using customer backup controls. +Back up Redis when active sessions and encrypted continuation state must survive +recovery; losing it invalidates sessions/continuations and does not erase ledger truth. + +Run the same Helm command with the new verified chart/images. It creates a new +migration Job for the release revision. Check the Job and each rollout, perform +login/inference/audit checks, and record results. Never run an Alembic downgrade +against production. + +For application rollback, select previously verified images compatible with the +current schema and upgrade with `--set migration.enabled=false`. This disables +both the migration Job and the exact-head init check; compatibility must be +established beforehand. Do not use blind Helm rollback to an incompatible schema. +For disaster recovery, restore the backup into a separate database, confirm its +revision and integrity, restore the corresponding secrets, then test a compatible +application before switching customer traffic. Rehearse this on disposable data. + +### PostgreSQL restore rehearsal + +Use a separate database and recovery application, leaving the source database and +customer traffic untouched. These commands use PostgreSQL client tools matching +the source server's major version. Configure private libpq service profiles +(`~/.pg_service.conf`, mode 0600): `shim-source` for the source database, +`shim-restore-admin` for a maintenance database with database-creation permission, +and `shim-restored` for the new `shim_restore` database as the application owner +`shim`. Put passwords in a private `~/.pgpass` or your credential helper; do not +put them in command arguments. Use the customer's TLS verification settings in +all three profiles. The application SQLAlchemy URL (`postgresql+asyncpg://...`) +is not a libpq connection string. + +```sh +umask 077 +PGSERVICE=shim-source pg_dump --format=custom --no-owner --no-acl \ + --file=/secure/shim-before-upgrade.dump +PGSERVICE=shim-source psql --no-psqlrc --tuples-only \ + --command='SELECT version_num FROM alembic_version;' +PGSERVICE=shim-restore-admin createdb --owner=shim shim_restore +pg_restore --exit-on-error --no-owner --no-acl \ + --dbname='service=shim-restored' /secure/shim-before-upgrade.dump +PGSERVICE=shim-restored psql --no-psqlrc --tuples-only \ + --command='SELECT version_num FROM alembic_version;' +``` + +The destination must be new and empty; do not add `--clean` against a live +database. `pg_dump` takes a consistent logical snapshot while the source remains +online. Record its time and archive checksum with the image digests and Alembic +revision. This archive omits cluster-level roles, grants and tablespaces: restore +those through customer database administration. Compare critical record counts +and audit-chain verification with the captured source evidence; merely listing +an archive or seeing `pg_restore` exit successfully is insufficient. + +Create recovery configuration pointing `DATABASE_URL` to `shim_restore`, with an +isolated Redis instance or database index. Preserve the backed-up organization +UUID, `SECRET_KEY`, any existing encryption key, and the corresponding Vault +secret versions. Register a separate recovery dashboard/callback origin with +the customer IdP, and provision its DNS, TLS, namespace secrets and permitted +network destinations. Use the exact chart and application images recorded with +the backup before testing an upgrade: + +```sh +helm upgrade --install shim-restore ./chart --namespace shim-recovery \ + --create-namespace --values /secure/restore-values.yaml \ + --set migration.enabled=false \ + --set workers.outbox=false --set workers.reconciliation=false \ + --set workers.compliance=false --set workers.ai_act=false \ + --wait --timeout 10m +kubectl -n shim-recovery exec deployment/shim-restore-gateway -- \ + alembic -c ee/alembic.ini current --check-heads +``` + +The private restore values must name the recovery runtime Secret and recovery +services; they must not reuse source database endpoints. Workers start disabled: +a restored outbox or active connector can repeat external effects from the backup. +Inspect pending events and connector destinations, isolate external delivery +targets, and establish replay/idempotency handling before enabling each worker. +Verify login, existing +request/ledger/audit visibility, audit-chain verification, Vault-backed model +access, and a new disposable workload key. Revoke that key after the test. +Only then rehearse the new chart/images with migrations enabled and repeat the +checks. An application rollback is a separate upgrade to previously verified, +current-schema-compatible images with `migration.enabled=false`; it does not +restore database contents. A successful same-schema restart does not establish +compatibility across a future schema change. + +Keep the source installation available until recovery evidence is reviewed and +traffic cutover is explicitly scheduled. Backups contain sensitive tenant data: +retain or dispose of the private archive and recovery resources under customer +policy. This logical database rehearsal does not test Vault disaster recovery, +IdP recovery, point-in-time recovery, or a different PostgreSQL major version. + +Troubleshooting starts with `kubectl -n shim get pods,jobs`, the migration Job +logs, and the affected process logs. Schema init failures mean the expected +migration has not completed. Login failures usually indicate issuer/client/redirect, +group mapping, CA, time, or Redis problems. Missing usage/audit projections require +checking outbox readiness and pending/dead-letter events. Never paste credentials, +session cookies, provider payloads or login callback queries into incident records. diff --git a/ee/deploy/chart/.helmignore b/ee/deploy/chart/.helmignore new file mode 100644 index 0000000..478d0a4 --- /dev/null +++ b/ee/deploy/chart/.helmignore @@ -0,0 +1,2 @@ +.DS_Store +*.tgz diff --git a/ee/deploy/chart/Chart.yaml b/ee/deploy/chart/Chart.yaml new file mode 100644 index 0000000..9758101 --- /dev/null +++ b/ee/deploy/chart/Chart.yaml @@ -0,0 +1,7 @@ +apiVersion: v2 +name: shim-enterprise +description: Customer-operated shim gateway, dashboard, and workers +type: application +version: 0.1.0 +appVersion: "0.1.3" +kubeVersion: ">=1.28.0-0" diff --git a/ee/deploy/chart/templates/_helpers.tpl b/ee/deploy/chart/templates/_helpers.tpl new file mode 100644 index 0000000..2cd593b --- /dev/null +++ b/ee/deploy/chart/templates/_helpers.tpl @@ -0,0 +1,75 @@ +{{- define "shim.labels" -}} +app.kubernetes.io/name: shim-enterprise +app.kubernetes.io/instance: {{ .Release.Name }} +{{- end }} + +{{- define "shim.security" -}} +runAsNonRoot: true +runAsUser: 10001 +runAsGroup: 10001 +fsGroup: 10001 +seccompProfile: {type: RuntimeDefault} +{{- end }} + +{{- define "shim.containerSecurity" -}} +allowPrivilegeEscalation: false +readOnlyRootFilesystem: true +capabilities: {drop: [ALL]} +{{- end }} + +{{- define "shim.backendEnv" -}} +envFrom: + - secretRef: + name: {{ required "existingSecret is required" .Values.existingSecret | quote }} +env: + - name: WORKER_HEARTBEAT_PATH + value: /tmp/shim-worker.json +{{- range $key, $value := .Values.gateway.env }} + - name: {{ $key }} + value: {{ $value | toString | quote }} +{{- end }} +{{- if .Values.vault.tokenSecretName }} + - name: VAULT_TOKEN_FILE + value: /var/run/shim-vault/token +{{- end }} +{{- if .Values.caBundleConfigMap }} + - name: SSL_CERT_FILE + value: /var/run/shim-ca/ca.pem +{{- end }} +{{- end }} + +{{- define "shim.mounts" -}} +volumeMounts: + - name: tmp + mountPath: /tmp +{{- if .Values.vault.tokenSecretName }} + - name: vault-token + mountPath: /var/run/shim-vault + readOnly: true +{{- end }} +{{- if .Values.caBundleConfigMap }} + - name: ca-bundle + mountPath: /var/run/shim-ca + readOnly: true +{{- end }} +{{- end }} + +{{- define "shim.volumes" -}} +volumes: + - name: tmp + emptyDir: {sizeLimit: 128Mi} +{{- if .Values.vault.tokenSecretName }} + - name: vault-token + secret: + secretName: {{ .Values.vault.tokenSecretName | quote }} + defaultMode: 0440 + items: + - key: {{ .Values.vault.tokenSecretKey | quote }} + path: token +{{- end }} +{{- if .Values.caBundleConfigMap }} + - name: ca-bundle + configMap: + name: {{ .Values.caBundleConfigMap | quote }} +{{- end }} +{{- end }} diff --git a/ee/deploy/chart/templates/backend.yaml b/ee/deploy/chart/templates/backend.yaml new file mode 100644 index 0000000..6c2cd60 --- /dev/null +++ b/ee/deploy/chart/templates/backend.yaml @@ -0,0 +1,96 @@ +{{- $root := . -}} +{{- $processes := dict "gateway" true -}} +{{- range $name, $enabled := .Values.workers -}} +{{- $_ := set $processes $name $enabled -}} +{{- end -}} +{{- range $name, $enabled := $processes }} +{{- if $enabled }} +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ $root.Release.Name }}-{{ $name | replace "_" "-" }} + labels: + {{- include "shim.labels" $root | nindent 4 }} + app.kubernetes.io/component: {{ $name }} +spec: + replicas: {{ if eq $name "gateway" }}{{ $root.Values.gateway.replicas }}{{ else }}1{{ end }} + selector: + matchLabels: + app.kubernetes.io/instance: {{ $root.Release.Name }} + app.kubernetes.io/component: {{ $name }} + template: + metadata: + labels: + {{- include "shim.labels" $root | nindent 8 }} + app.kubernetes.io/component: {{ $name }} + spec: + automountServiceAccountToken: false + securityContext: + {{- include "shim.security" $root | nindent 8 }} + {{- with $root.Values.imagePullSecrets }} + imagePullSecrets: {{ toYaml . | nindent 8 }} + {{- end }} + {{- if $root.Values.migration.enabled }} + initContainers: + - name: schema-ready + image: {{ required "gateway.image is required" $root.Values.gateway.image | quote }} + command: [alembic, -c, ee/alembic.ini, current, --check-heads] + securityContext: + {{- include "shim.containerSecurity" $root | nindent 12 }} + {{- include "shim.backendEnv" $root | nindent 10 }} + {{- include "shim.mounts" $root | nindent 10 }} + {{- end }} + containers: + - name: {{ $name | replace "_" "-" }} + image: {{ required "gateway.image is required" $root.Values.gateway.image | quote }} + imagePullPolicy: IfNotPresent + {{- if ne $name "gateway" }} + command: [python, -m, {{ printf "shim_enterprise.workers.%s" $name | quote }}] + readinessProbe: + exec: + command: + - python + - -m + - shim_enterprise.workers.readiness + - --path + - /tmp/shim-worker.json + - --worker + - {{ $name | quote }} + - --max-age-seconds + - {{ index $root.Values.workerReadinessMaxAgeSeconds $name | toString | quote }} + periodSeconds: 15 + timeoutSeconds: 3 + {{- else }} + ports: + - {name: http, containerPort: 8000} + startupProbe: + httpGet: {path: /health, port: http} + failureThreshold: 30 + periodSeconds: 5 + readinessProbe: + httpGet: {path: /health, port: http} + periodSeconds: 5 + livenessProbe: + httpGet: {path: /health, port: http} + periodSeconds: 15 + {{- end }} + securityContext: + {{- include "shim.containerSecurity" $root | nindent 12 }} + {{- include "shim.backendEnv" $root | nindent 10 }} + {{- include "shim.mounts" $root | nindent 10 }} + resources: {{ toYaml $root.Values.gateway.resources | nindent 12 }} + {{- include "shim.volumes" $root | nindent 6 }} +{{- end }} +{{- end }} +--- +apiVersion: v1 +kind: Service +metadata: + name: {{ .Release.Name }}-gateway +spec: + selector: + app.kubernetes.io/instance: {{ .Release.Name }} + app.kubernetes.io/component: gateway + ports: + - {name: http, port: 8000, targetPort: http} diff --git a/ee/deploy/chart/templates/dashboard.yaml b/ee/deploy/chart/templates/dashboard.yaml new file mode 100644 index 0000000..3193381 --- /dev/null +++ b/ee/deploy/chart/templates/dashboard.yaml @@ -0,0 +1,87 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ .Release.Name }}-dashboard + labels: + {{- include "shim.labels" . | nindent 4 }} +spec: + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/instance: {{ .Release.Name }} + app.kubernetes.io/component: dashboard + template: + metadata: + labels: + {{- include "shim.labels" . | nindent 8 }} + app.kubernetes.io/component: dashboard + spec: + automountServiceAccountToken: false + securityContext: + {{- include "shim.security" . | nindent 8 }} + {{- with .Values.imagePullSecrets }} + imagePullSecrets: {{ toYaml . | nindent 8 }} + {{- end }} + containers: + - name: dashboard + image: {{ required "dashboard.image is required" .Values.dashboard.image | quote }} + imagePullPolicy: IfNotPresent + securityContext: + {{- include "shim.containerSecurity" . | nindent 12 }} + {{- with .Values.dashboard.existingSecret }} + envFrom: + - secretRef: {name: {{ . | quote }}} + {{- end }} + env: + - {name: HOSTNAME, value: "0.0.0.0"} + - {name: PORT, value: "3000"} + - {name: AUTH_MODE, value: oidc} + - {name: NEXT_TELEMETRY_DISABLED, value: "1"} + - name: SHIM_API_URL + value: {{ printf "http://%s-gateway:8000" .Release.Name | quote }} + {{- range $key, $value := .Values.dashboard.env }} + - name: {{ $key }} + value: {{ $value | toString | quote }} + {{- end }} + {{- if .Values.caBundleConfigMap }} + - {name: NODE_EXTRA_CA_CERTS, value: /var/run/shim-ca/ca.pem} + {{- end }} + ports: + - {name: http, containerPort: 3000} + startupProbe: + httpGet: {path: /login, port: http} + failureThreshold: 30 + periodSeconds: 5 + readinessProbe: + httpGet: {path: /login, port: http} + periodSeconds: 5 + livenessProbe: + tcpSocket: {port: http} + periodSeconds: 15 + resources: {{ toYaml .Values.dashboard.resources | nindent 12 }} + volumeMounts: + - {name: tmp, mountPath: /tmp} + - {name: cache, mountPath: /app/.next/cache} + {{- if .Values.caBundleConfigMap }} + - {name: ca-bundle, mountPath: /var/run/shim-ca, readOnly: true} + {{- end }} + volumes: + - name: tmp + emptyDir: {sizeLimit: 64Mi} + - name: cache + emptyDir: {sizeLimit: 256Mi} + {{- if .Values.caBundleConfigMap }} + - name: ca-bundle + configMap: {name: {{ .Values.caBundleConfigMap | quote }}} + {{- end }} +--- +apiVersion: v1 +kind: Service +metadata: + name: {{ .Release.Name }}-dashboard +spec: + selector: + app.kubernetes.io/instance: {{ .Release.Name }} + app.kubernetes.io/component: dashboard + ports: + - {name: http, port: 3000, targetPort: http} diff --git a/ee/deploy/chart/templates/ingress.yaml b/ee/deploy/chart/templates/ingress.yaml new file mode 100644 index 0000000..21ce779 --- /dev/null +++ b/ee/deploy/chart/templates/ingress.yaml @@ -0,0 +1,38 @@ +{{- if .Values.ingress.enabled }} +apiVersion: networking.k8s.io/v1 +kind: Ingress +metadata: + name: {{ .Release.Name }} + {{- with .Values.ingress.annotations }} + annotations: {{ toYaml . | nindent 4 }} + {{- end }} +spec: + {{- with .Values.ingress.className }} + ingressClassName: {{ . | quote }} + {{- end }} + tls: + - hosts: [{{ required "ingress.host is required" .Values.ingress.host | quote }}] + secretName: {{ required "ingress.tlsSecretName is required" .Values.ingress.tlsSecretName | quote }} + rules: + - host: {{ .Values.ingress.host | quote }} + http: + paths: + - path: /v1 + pathType: Prefix + backend: + service: + name: {{ .Release.Name }}-gateway + port: {number: 8000} + - path: /v1beta + pathType: Prefix + backend: + service: + name: {{ .Release.Name }}-gateway + port: {number: 8000} + - path: / + pathType: Prefix + backend: + service: + name: {{ .Release.Name }}-dashboard + port: {number: 3000} +{{- end }} diff --git a/ee/deploy/chart/templates/migration.yaml b/ee/deploy/chart/templates/migration.yaml new file mode 100644 index 0000000..713487e --- /dev/null +++ b/ee/deploy/chart/templates/migration.yaml @@ -0,0 +1,34 @@ +{{- if .Values.migration.enabled }} +apiVersion: batch/v1 +kind: Job +metadata: + name: {{ .Release.Name }}-migrate-{{ .Release.Revision }} + labels: + {{- include "shim.labels" . | nindent 4 }} +spec: + backoffLimit: 4 + activeDeadlineSeconds: {{ .Values.migration.activeDeadlineSeconds }} + template: + metadata: + labels: + {{- include "shim.labels" . | nindent 8 }} + app.kubernetes.io/component: migration + spec: + restartPolicy: OnFailure + automountServiceAccountToken: false + securityContext: + {{- include "shim.security" . | nindent 8 }} + {{- with .Values.imagePullSecrets }} + imagePullSecrets: {{ toYaml . | nindent 8 }} + {{- end }} + containers: + - name: migrate + image: {{ required "gateway.image is required" .Values.gateway.image | quote }} + command: [alembic, -c, ee/alembic.ini, upgrade, head] + securityContext: + {{- include "shim.containerSecurity" . | nindent 12 }} + {{- include "shim.backendEnv" . | nindent 10 }} + {{- include "shim.mounts" . | nindent 10 }} + resources: {{ toYaml .Values.gateway.resources | nindent 12 }} + {{- include "shim.volumes" . | nindent 6 }} +{{- end }} diff --git a/ee/deploy/chart/templates/network-policy.yaml b/ee/deploy/chart/templates/network-policy.yaml new file mode 100644 index 0000000..1754821 --- /dev/null +++ b/ee/deploy/chart/templates/network-policy.yaml @@ -0,0 +1,27 @@ +{{- if .Values.networkPolicy.enabled }} +apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: {{ .Release.Name }}-egress +spec: + podSelector: + matchLabels: + app.kubernetes.io/instance: {{ .Release.Name }} + policyTypes: [Egress] + egress: + - to: + - podSelector: {} + - to: + - namespaceSelector: + matchLabels: + kubernetes.io/metadata.name: kube-system + podSelector: + matchLabels: + k8s-app: kube-dns + ports: + - {protocol: UDP, port: 53} + - {protocol: TCP, port: 53} + {{- with .Values.networkPolicy.extraEgress }} + {{- toYaml . | nindent 4 }} + {{- end }} +{{- end }} diff --git a/ee/deploy/chart/templates/test-services.yaml b/ee/deploy/chart/templates/test-services.yaml new file mode 100644 index 0000000..7187965 --- /dev/null +++ b/ee/deploy/chart/templates/test-services.yaml @@ -0,0 +1,74 @@ +{{- range $name := list "postgres" "redis" }} +{{- $cfg := index $.Values $name }} +{{- if $cfg.enabled }} +--- +apiVersion: v1 +kind: Service +metadata: + name: {{ $.Release.Name }}-{{ $name }} +spec: + clusterIP: None + selector: + app.kubernetes.io/instance: {{ $.Release.Name }} + app.kubernetes.io/component: {{ $name }} + ports: + - port: {{ if eq $name "postgres" }}5432{{ else }}6379{{ end }} +--- +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: {{ $.Release.Name }}-{{ $name }} +spec: + serviceName: {{ $.Release.Name }}-{{ $name }} + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/instance: {{ $.Release.Name }} + app.kubernetes.io/component: {{ $name }} + template: + metadata: + labels: + {{- include "shim.labels" $ | nindent 8 }} + app.kubernetes.io/component: {{ $name }} + spec: + automountServiceAccountToken: false + {{- with $.Values.imagePullSecrets }} + imagePullSecrets: {{ toYaml . | nindent 8 }} + {{- end }} + containers: + - name: {{ $name }} + image: {{ $cfg.image | quote }} + imagePullPolicy: IfNotPresent + {{- if eq $name "postgres" }} + env: + - {name: POSTGRES_USER, value: shim} + - {name: POSTGRES_DB, value: shim} + - {name: PGDATA, value: /var/lib/postgresql/data/pgdata} + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: {{ $.Values.existingSecret | quote }} + key: POSTGRES_PASSWORD + readinessProbe: + exec: {command: [pg_isready, -U, shim, -d, shim]} + volumeMounts: + - {name: data, mountPath: /var/lib/postgresql/data} + {{- else }} + readinessProbe: + exec: {command: [redis-cli, ping]} + volumeMounts: + - {name: data, mountPath: /data} + {{- end }} + volumeClaimTemplates: + - metadata: + name: data + spec: + accessModes: [ReadWriteOnce] + {{- if $cfg.storageClassName }} + storageClassName: {{ $cfg.storageClassName | quote }} + {{- end }} + resources: + requests: + storage: {{ $cfg.storage | quote }} +{{- end }} +{{- end }} diff --git a/ee/deploy/chart/values.yaml b/ee/deploy/chart/values.yaml new file mode 100644 index 0000000..10d1be9 --- /dev/null +++ b/ee/deploy/chart/values.yaml @@ -0,0 +1,76 @@ +# Supply immutable release image references and pre-created secrets. +existingSecret: "" +imagePullSecrets: [] + +gateway: + image: "" + replicas: 1 + env: + ENVIRONMENT: production + AUTH_MODE: oidc + SECRET_BACKEND: vault + MODEL_DEPLOYMENT_REQUIRED: "true" + SENTRY_DSN: "" + OTEL_EXPORTER_OTLP_ENDPOINT: "" + resources: + requests: {cpu: 100m, memory: 256Mi} + limits: {memory: 1Gi} + +dashboard: + image: "" + existingSecret: "" + env: {} + resources: + requests: {cpu: 100m, memory: 128Mi} + limits: {memory: 512Mi} + +workers: + outbox: true + reconciliation: true + compliance: true + ai_act: true + +workerReadinessMaxAgeSeconds: + outbox: 120 + reconciliation: 180 + compliance: 900 + ai_act: 7500 + +migration: + enabled: true + activeDeadlineSeconds: 300 + +# A Vault Agent or your secret controller refreshes this token Secret. +vault: + tokenSecretName: "" + tokenSecretKey: token + +# PEM bundle including all required public/private roots. Never disable TLS checks. +caBundleConfigMap: "" + +ingress: + enabled: false + className: "" + host: "" + tlsSecretName: "" + annotations: {} + +networkPolicy: + enabled: true + # Add only required IdP, Vault, database, model or proxy destinations. + extraEgress: [] + +# Local test services only; production normally uses customer-operated services. +# DATABASE_URL/REDIS_URL are supplied in existingSecret in either configuration. +postgres: + enabled: false + # PostgreSQL 16 Alpine + image: postgres@sha256:cf78e76683b9ca8c5733cbbdce6c9262b45b6767934dd0a95e671f9a0fc20685 + storage: 5Gi + storageClassName: "" +redis: + enabled: false + # Redis Stack 7.4.0-v8 + image: redis/redis-stack-server@sha256:798ab84d9f266936b034ab11c4d04a2b8e4b441884c5aa7d17ac951eefdf742a + storage: 1Gi + storageClassName: "" diff --git a/ee/deploy/test/README.md b/ee/deploy/test/README.md new file mode 100644 index 0000000..e4c3b23 --- /dev/null +++ b/ee/deploy/test/README.md @@ -0,0 +1,62 @@ +# Isolated chart acceptance + +Run from the backend repository after building the final enterprise runtime and +an OIDC dashboard runtime. The dashboard can come from its private repository; +this harness and its manually dispatched CI workflow never check out that repo +or request a source checkout token. For a private GHCR dashboard image, explicitly +configure repository secret `ON_PREM_DASHBOARD_GHCR_TOKEN` with read access only +to that package and repository variable `ON_PREM_DASHBOARD_GHCR_USER` with its +service-account username. The workflow logs in only to GHCR for image pulling; +public images need no credential. The token must not have source write/admin +permissions. Local runs use the operator's existing Docker registry login. + +```sh +uv run --locked --all-packages python ee/deploy/test/run.py \ + --gateway-image shim-enterprise:local \ + --dashboard-image shim-dashboard:local +``` + +Requires Docker, kind 0.33.0, Helm 4.2.4, kubectl compatible with Kubernetes 1.37, +and the locked Python environment. Images are pulled before isolation. Allow +roughly 8 GB of free Docker storage. Use immutable digest references for supplied +images in CI. `--chart` selects a chart checkout; `--cluster` names a new disposable +cluster. Existing clusters are refused, never replaced implicitly. `--keep` +retains the failed/successful cluster for diagnosis; the default destroys only +the named disposable cluster/network. `--reuse-empty-cluster` resumes a failed +image preload and refuses a cluster with application deployments, StatefulSets, +or Jobs. It is only for a cluster created by this harness before isolation. + +The harness installs the real chart with its gateway, dashboard, four workers, +migration job, and disposable PostgreSQL/Redis. Real Keycloak and Vault run beside +two small native OpenAI/Anthropic HTTP fixtures. A generated two-day CA signs the +IdP and fixture certificates. Vault and the dashboard are exposed through the +fixture's TLS proxy. The gateway uses its actual custom-CA configuration, OIDC +code/PKCE callback, Vault secret storage, registry routing, usage, and audit APIs. +It checks all six chart deployments' readiness and creates/revokes a workload key. +The login client follows Keycloak's real HTML form and secure session cookies; +it does not replace browser accessibility/UI tests. + +The disposable kind node has explicit IPv4/IPv6 OUTPUT and FORWARD firewall +chains. Only established traffic, loopback, and the node/pod/service CIDRs pass; +other destinations are rejected. CoreDNS upstream forwarding is removed. +A positive node public-IP connectivity control runs before isolation, then two +public-IP probes must fail from the gateway pod while internal calls succeed. +This verifies a real external network boundary. It does **not** claim kind's +default CNI enforces the chart's NetworkPolicy, nor prove a customer's firewall. +Docker `--internal` is deliberately avoided: it suppresses kind's published +control-plane API port. Static HTML/CSS/JS/font fetches record the declared asset +hosts separately; this evidence is not JavaScript/browser execution. + +The production licence verifier remains enabled. A separately tagged and labelled +**test-only** image overlays only a newly generated public key. Its short-lived +licence is asserted invalid against the original image's packaged public key; +the original image identity is checked unchanged. Neither private signing keys +nor test public keys enter source control. Never publish or release the test-only +image. This is not an alternate customer licence mechanism. + +Temporary material is private to the local OS user and includes disposable test +secrets/kubeconfig; do not upload it wholesale. Only `result.json`, `asset-audit.json`, and +`resources.txt` are CI artifacts. After review, use `trash` on the printed +temporary directory. Remove the printed test-only Docker image with Docker's +image removal command when no longer needed. Real Entra, customer TLS PKI, and +operator restore/rollback remain separate acceptance scenarios. diff --git a/ee/deploy/test/fixtures.py b/ee/deploy/test/fixtures.py new file mode 100644 index 0000000..ba1db0b --- /dev/null +++ b/ee/deploy/test/fixtures.py @@ -0,0 +1,376 @@ +"""Disposable TLS fixture and smoke client; never include in a release image.""" + +from __future__ import annotations + +from html.parser import HTMLParser +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +from pathlib import Path +import re +import socket +import ssl +import sys +import threading +import time +from urllib.parse import urljoin, urlsplit +from uuid import uuid4 + +import httpx + + +class Fixture(BaseHTTPRequestHandler): + def log_message(self, *_args): + pass # Authorization codes, session cookies, and credentials are not logs. + + def do_GET(self): + self.respond() + + def do_POST(self): + self.respond() + + def do_DELETE(self): + self.respond() + + do_PUT = do_POST + do_PATCH = do_POST + + def respond(self): + body = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + port = self.server.server_port + if port in {8443, 8444}: + if port == 8444: + origin = "http://vault:8200" + elif self.path.startswith(("/v1/", "/v1beta/")): + origin = "http://smoke-gateway:8000" + else: + origin = "http://smoke-dashboard:3000" + headers = { + key: value + for key, value in self.headers.items() + if key.lower() not in {"host", "connection", "content-length"} + } + with httpx.Client(trust_env=False, timeout=30) as client: + result = client.request( + self.command, origin + self.path, headers=headers, content=body + ) + self.send_response(result.status_code) + for key, value in result.headers.multi_items(): + if key.lower() not in { + "connection", + "content-length", + "transfer-encoding", + "content-encoding", + }: + self.send_header(key, value) + output = result.content + else: + assert self.command == "POST", "Only native model POSTs are expected" + payload = json.loads(body) + assert payload["messages"], "Messages required" + if self.path == "/openai/chat/completions": + assert ( + self.headers.get("Authorization") + == "Bearer disposable-model-secret" + ) + if payload.get("stream"): + chunk = { + "id": "chatcmpl-smoke-stream", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": payload["model"], + } + frames = [ + chunk + | { + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "content": "internal openai fixture", + }, + "finish_reason": None, + } + ] + }, + chunk + | { + "choices": [ + {"index": 0, "delta": {}, "finish_reason": "stop"} + ] + }, + chunk + | { + "choices": [], + "usage": { + "prompt_tokens": 7, + "completion_tokens": 4, + "total_tokens": 11, + }, + }, + ] + output = ( + "".join( + "data: " + json.dumps(frame) + "\n\n" for frame in frames + ) + + "data: [DONE]\n\n" + ).encode() + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Content-Length", str(len(output))) + self.end_headers() + self.wfile.write(output) + return + result_body = { + "id": "chatcmpl-smoke", + "object": "chat.completion", + "created": int(time.time()), + "model": payload["model"], + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "internal openai fixture", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 7, + "completion_tokens": 4, + "total_tokens": 11, + }, + } + else: + assert self.path == "/anthropic/v1/messages", self.path + assert self.headers.get("x-api-key") == "disposable-model-secret" + result_body = { + "id": "msg-smoke", + "type": "message", + "role": "assistant", + "model": payload["model"], + "content": [{"type": "text", "text": "internal anthropic fixture"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 7, "output_tokens": 4}, + } + output = json.dumps(result_body).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(output))) + self.end_headers() + self.wfile.write(output) + + +def serve(): + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.minimum_version = ssl.TLSVersion.TLSv1_2 + context.load_cert_chain("/tls/tls.crt", "/tls/tls.key") + for port in (8443, 8444, 8445): + server = ThreadingHTTPServer(("0.0.0.0", port), Fixture) + server.socket = context.wrap_socket(server.socket, server_side=True) + threading.Thread(target=server.serve_forever, daemon=True).start() + threading.Event().wait() + + +class LoginForm(HTMLParser): + def __init__(self): + super().__init__() + self.action = None + self.fields = {} + + def handle_starttag(self, tag, attrs): + values = dict(attrs) + if tag == "form" and values.get("id") == "kc-form-login": + self.action = values["action"] + if tag == "input" and values.get("name") and values.get("type") == "hidden": + self.fields[values["name"]] = values.get("value", "") + + +class AssetLinks(HTMLParser): + def __init__(self): + super().__init__() + self.urls = set() + + def handle_starttag(self, tag, attrs): + values = dict(attrs) + if tag in {"script", "img"} and values.get("src"): + self.urls.add(values["src"]) + if ( + tag == "link" + and values.get("href") + and values.get("rel") in {"stylesheet", "preload", "modulepreload", "icon"} + ): + self.urls.add(values["href"]) + + +def assets(client, response): + parser = AssetLinks() + parser.feed(response.text) + pending = {urljoin(str(response.url), url) for url in parser.urls} + visited = {} + while pending: + url = pending.pop() + if url.startswith("data:") or url in visited: + continue + assert urlsplit(url).netloc == "fixture:8443", f"Unexpected asset host: {url}" + asset = client.get(url) + assert asset.is_success, f"Asset HTTP {asset.status_code}: {url}" + visited[url] = asset.status_code + if "text/css" in asset.headers.get("content-type", ""): + pending.update( + urljoin(url, match.strip("\"' ")) + for match in re.findall(r"url\(([^)]+)\)", asset.text) + ) + assert any(".js" in url for url in visited) and any( + ".css" in url for url in visited + ) + Path("/tmp/shim-asset-audit.json").write_text( + json.dumps( + { + "kind": "static HTML/CSS dependency fetch, not browser execution", + "assets": visited, + }, + indent=2, + ) + ) + print( + f"PASS {len(visited)} declared local HTML/CSS/JS/font assets (static fetch, not browser)", + flush=True, + ) + + +def denied(): + for address in ("1.1.1.1", "8.8.8.8"): + try: + with socket.create_connection((address, 443), timeout=3): + raise AssertionError(f"Public egress is reachable: {address}") + except (TimeoutError, OSError): + pass + print("PASS public-IP egress denied from gateway pod", flush=True) + + +def smoke(): + denied() + context = ssl.create_default_context(cafile="/var/run/shim-ca/ca.pem") + origin = "https://fixture:8443" + with httpx.Client( + verify=context, timeout=30, follow_redirects=True, trust_env=False + ) as client: + response = client.get(origin + "/api/v1/auth/login") + response.raise_for_status() + form = LoginForm() + form.feed(response.text) + assert form.action, "Keycloak login form missing" + response = client.post( + urljoin(str(response.url), form.action), + data=form.fields + | {"username": "pilot", "password": "disposable-pilot-password"}, + ) + response.raise_for_status() + assert response.url.path == "/dashboard", response.url.path + assets(client, response) + assert client.cookies.get("shim_session") + session = client.get(origin + "/api/v1/auth/session") + session.raise_for_status() + assert session.json()["user"]["role"] == "owner" + client.headers["Origin"] = origin + + def api(method, path, **kwargs): + response = client.request(method, origin + "/api/v1" + path, **kwargs) + assert response.is_success, ( + f"{method} {path}: HTTP {response.status_code} {response.text[:300]}" + ) + return response.json() if response.content else None + + created = api( + "POST", "/management/api-keys", json={"name": "isolated chart smoke"} + ) + key = created["plaintext"] + request_ids = set() + aliases = {} + run_id = uuid4().hex[:8] + for provider, upstream in ( + ("openai", "gpt-4o-mini"), + ("anthropic", "claude-3-5-haiku-latest"), + ): + secret = api( + "POST", + "/management/providers", + json={"provider": provider, "key": "disposable-model-secret"}, + ) + alias = aliases[provider] = f"smoke-{provider}-{run_id}" + api( + "POST", + "/management/model-deployments", + json={ + "alias": alias, + "provider": provider, + "upstream_model": upstream, + "base_url": "https://fixture:8445/" + provider, + "provider_secret_id": secret["id"], + "deployment_kind": "internal", + "declared_version": "fixture-v1", + "owner": "isolated-test", + }, + ) + payload = { + "model": alias, + "messages": [{"role": "user", "content": "Hello internal fixture"}], + "max_tokens": 16, + } + path = "/v1/chat/completions" if provider == "openai" else "/v1/messages" + response = client.post( + "http://smoke-gateway:8000" + path, + json=payload, + headers={"x-shim-key": key, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, ( + f"{provider}: {response.status_code} {response.text[:300]}" + ) + assert "internal " in response.text + request_ids.add(response.headers["X-Shim-Request-Id"]) + deadline = time.monotonic() + 60 + while True: + usage = api("GET", "/management/requests") + audit = api("GET", "/compliance/audit/logs") + rows = [row for row in usage["items"] if row["request_id"] in request_ids] + audit_rows = [ + row + for row in audit["items"] + if row["request_id"] in request_ids + and row["event_type"] == "ai_request" + ] + if ( + len(rows) == 2 + and {row["request_id"] for row in audit_rows} == request_ids + ): + assert all( + row["prompt_tokens"] == 7 and row["completion_tokens"] == 4 + for row in rows + audit_rows + ) + assert api("POST", "/compliance/audit/verify", json={})["ok"] + break + assert time.monotonic() < deadline, "Usage/audit did not become visible" + time.sleep(1) + api("DELETE", "/management/api-keys/" + created["id"]) + response = client.post( + "http://smoke-gateway:8000/v1/chat/completions", + json={ + "model": aliases["openai"], + "messages": [{"role": "user", "content": "revoked"}], + }, + headers={"x-shim-key": key}, + ) + assert response.status_code == 401 + api("POST", "/auth/logout") + assert client.get(origin + "/api/v1/auth/session").status_code == 401 + print( + "PASS TLS/custom CA, real Keycloak login, Vault secret writes, both internal models, usage/audit, revoke/logout", + flush=True, + ) + + +if __name__ == "__main__": + {"serve": serve, "smoke": smoke, "denied": denied}[sys.argv[1]]() diff --git a/ee/deploy/test/run.py b/ee/deploy/test/run.py new file mode 100644 index 0000000..7fd2762 --- /dev/null +++ b/ee/deploy/test/run.py @@ -0,0 +1,911 @@ +"""Install and smoke-test the chart in a disposable, internet-isolated kind cluster.""" + +from __future__ import annotations + +import argparse +import base64 +from datetime import UTC, datetime, timedelta +import json +import os +import re +from pathlib import Path +import secrets +import subprocess +import sys +import tempfile +from uuid import UUID, uuid4 + +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ed25519, rsa +from cryptography.x509.oid import NameOID +import yaml + +KEYCLOAK = "quay.io/keycloak/keycloak@sha256:ff4257d0d64efbe99ed1ddfaf07765cc3c36dc7518bf8324d41961327f441c54" +VAULT = "hashicorp/vault@sha256:5520cc26271c024e6ffa45cdf95255bd26b70d71ba4b7e0bc18925bef4128adb" +HERE = Path(__file__).resolve().parent + + +def command(*args, data=None, capture=False, env=None): + result = subprocess.run( + args, input=data, text=True, check=True, capture_output=capture, env=env + ) + return result.stdout.strip() if capture else "" + + +def certificates(directory): + now = datetime.now(UTC) + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + name = x509.Name( + [x509.NameAttribute(NameOID.COMMON_NAME, "Disposable shim smoke CA")] + ) + ca = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - timedelta(minutes=5)) + .not_valid_after(now + timedelta(days=2)) + .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True) + .add_extension( + x509.KeyUsage( + digital_signature=True, + content_commitment=False, + key_encipherment=False, + data_encipherment=False, + key_agreement=False, + key_cert_sign=True, + crl_sign=True, + encipher_only=False, + decipher_only=False, + ), + critical=True, + ) + .add_extension( + x509.SubjectKeyIdentifier.from_public_key(key.public_key()), critical=False + ) + .sign(key, hashes.SHA256()) + ) + leaf_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + leaf = ( + x509.CertificateBuilder() + .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "fixture")])) + .issuer_name(name) + .public_key(leaf_key.public_key()) + .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True) + .add_extension( + x509.AuthorityKeyIdentifier.from_issuer_public_key(key.public_key()), + critical=False, + ) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - timedelta(minutes=5)) + .not_valid_after(now + timedelta(days=2)) + .add_extension( + x509.SubjectAlternativeName( + [x509.DNSName("fixture"), x509.DNSName("keycloak")] + ), + critical=False, + ) + .sign(key, hashes.SHA256()) + ) + (directory / "ca.pem").write_bytes(ca.public_bytes(serialization.Encoding.PEM)) + (directory / "tls.crt").write_bytes(leaf.public_bytes(serialization.Encoding.PEM)) + (directory / "tls.key").write_bytes( + leaf_key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + ) + + +def test_license(directory): + key = ed25519.Ed25519PrivateKey.generate() + (directory / "license_public_key.pem").write_bytes( + key.public_key().public_bytes( + serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo + ) + ) + payload = ( + base64.urlsafe_b64encode( + json.dumps( + { + "customer": "DISPOSABLE TEST ONLY - NEVER RELEASE", + "expires": (datetime.now(UTC) + timedelta(days=1)) + .date() + .isoformat(), + } + ).encode() + ) + .decode() + .rstrip("=") + ) + return ( + payload + + "." + + base64.urlsafe_b64encode(key.sign(payload.encode())).decode().rstrip("=") + ) + + +def kind_system_workload(item, cluster): + metadata = item["metadata"] + namespace = metadata.get("namespace") + name = metadata["name"] + known = { + ("Deployment", "kube-system", "coredns"), + ("Deployment", "local-path-storage", "local-path-provisioner"), + ("DaemonSet", "kube-system", "kindnet"), + ("DaemonSet", "kube-system", "kube-proxy"), + } + if (item["kind"], namespace, name) in known: + return True + if item["kind"] == "ReplicaSet": + return any( + kind == "Deployment" + and namespace == ns + and re.fullmatch(re.escape(deployment) + r"-[a-z0-9]+", name) + for kind, ns, deployment in known + ) + if item["kind"] != "Pod": + return False + owners = metadata.get("ownerReferences", []) + if len(owners) != 1: + return False + owner = owners[0] + if owner["kind"] == "DaemonSet": + return ("DaemonSet", namespace, owner["name"]) in known + if owner["kind"] == "ReplicaSet": + return any( + kind == "Deployment" + and namespace == ns + and re.fullmatch(re.escape(deployment) + r"-[a-z0-9]+", owner["name"]) + for kind, ns, deployment in known + ) + return ( + namespace == "kube-system" + and owner["kind"] == "Node" + and owner["name"] == cluster + "-control-plane" + and name + in { + component + "-" + cluster + "-control-plane" + for component in ( + "etcd", + "kube-apiserver", + "kube-controller-manager", + "kube-scheduler", + ) + } + ) + + +def run(args): + chart = args.chart.resolve() + values = yaml.safe_load((chart / "values.yaml").read_text()) + directory = Path(tempfile.mkdtemp(prefix="shim-chart-smoke-")) + directory.chmod(0o700) + print(f"Disposable test material and diagnostics: {directory}", flush=True) + certificates(directory) + license_token = test_license(directory) + images = [ + args.gateway_image, + args.dashboard_image, + args.node_image, + ] + for image in images: + probe = subprocess.run( + ["docker", "image", "inspect", image], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + if probe.returncode: + command("docker", "pull", image) + original_image = command( + "docker", + "image", + "inspect", + args.gateway_image, + "--format", + "{{.Id}}", + capture=True, + ) + # Confirm the disposable signing key cannot authorize the unmodified image. + verify = "import base64,sys; from importlib.resources import files; from cryptography.hazmat.primitives.serialization import load_pem_public_key; from cryptography.exceptions import InvalidSignature; p,s=sys.stdin.read().strip().split('.'); key=load_pem_public_key(files('shim_enterprise.core').joinpath('license_public_key.pem').read_bytes());\ntry: key.verify(base64.urlsafe_b64decode(s+'='*(-len(s)%4)),p.encode())\nexcept InvalidSignature: print('PASS production key rejects disposable licence')\nelse: raise SystemExit('Production image accepted disposable licence!')" + command( + "docker", + "run", + "--rm", + "-i", + "--entrypoint", + "python", + args.gateway_image, + "-c", + verify, + data=license_token, + ) + overlay = f"shim-on-prem-test-only:{uuid4().hex}" + (directory / "Dockerfile").write_text( + f"FROM {args.gateway_image}\nLABEL shim.test-only=true\nCOPY license_public_key.pem /app/.venv/lib/python3.13/site-packages/shim_enterprise/core/license_public_key.pem\n" + ) + (directory / ".dockerignore").write_text( + "*\n!Dockerfile\n!license_public_key.pem\n" + ) + command("docker", "build", "--network=none", "-t", overlay, str(directory)) + assert original_image == command( + "docker", + "image", + "inspect", + args.gateway_image, + "--format", + "{{.Id}}", + capture=True, + ) + existing = command("kind", "get", "clusters", capture=True).splitlines() + if args.cluster in existing and not args.reuse_empty_cluster: + raise SystemExit( + f"Cluster {args.cluster} already exists; explicitly remove that disposable cluster before running" + ) + network = args.cluster + "-isolated" + kubeconfig = directory / "kubeconfig" + if args.cluster not in existing: + command("docker", "network", "create", network) + subnet = command( + "docker", + "network", + "inspect", + network, + "--format", + "{{(index .IPAM.Config 0).Subnet}}", + capture=True, + ) + environment = os.environ | {"KIND_EXPERIMENTAL_DOCKER_NETWORK": network} + if args.cluster in existing: + kubeconfig.write_text( + command("kind", "get", "kubeconfig", "--name", args.cluster, capture=True) + ) + workloads = json.loads( + command( + "kubectl", + "--kubeconfig", + str(kubeconfig), + "get", + "deployments,statefulsets,jobs,daemonsets,cronjobs,replicasets,pods,persistentvolumeclaims", + "--all-namespaces", + "-o", + "json", + capture=True, + ) + ) + unexpected = [ + item + for item in workloads["items"] + if not kind_system_workload(item, args.cluster) + ] + if unexpected: + raise SystemExit( + "Refusing to reuse a cluster containing application workloads" + ) + try: + if args.cluster not in existing: + command( + "kind", + "create", + "cluster", + "--name", + args.cluster, + "--image", + args.node_image, + "--kubeconfig", + str(kubeconfig), + "--retain", + env=environment, + ) + platform = command( + "docker", + "image", + "inspect", + args.gateway_image, + "--format", + "{{.Os}}/{{.Architecture}}", + capture=True, + ) + runtime_images = {} + image_evidence = {} + for image in [overlay, args.dashboard_image]: + alias = "docker.io/library/shim-smoke-runtime:" + uuid4().hex + runtime_images[image] = alias + print(f"Loading {image} for {platform}", flush=True) + with subprocess.Popen( + ["docker", "image", "save", "--platform", platform, image], + stdout=subprocess.PIPE, + ) as exporting: + subprocess.run( + [ + "docker", + "exec", + "--privileged", + "-i", + args.cluster + "-control-plane", + "ctr", + "--namespace=k8s.io", + "images", + "import", + "--platform", + platform, + "--index-name", + alias, + "--digests", + "--snapshotter=overlayfs", + "-", + ], + stdin=exporting.stdout, + stdout=sys.stdout, + check=True, + ) + exporting.stdout.close() + if exporting.wait() != 0: + raise RuntimeError(f"Image export failed: {image}") + for image in [ + KEYCLOAK, + VAULT, + values["postgres"]["image"], + values["redis"]["image"], + ]: + first = image.split("/", 1)[0] + alias = ( + image + if "." in first + else ("docker.io/" if "/" in image else "docker.io/library/") + image + ) + if "@" in alias: + repository, digest = alias.split("@", 1) + alias = ( + repository.rsplit(":", 1)[0] + "@" + digest + if ":" in repository.rsplit("/", 1)[-1] + else alias + ) + runtime_images[image] = alias + print( + f"Pulling pinned fixture {alias} into node for {platform}", flush=True + ) + command( + "docker", + "exec", + args.cluster + "-control-plane", + "ctr", + "--namespace=k8s.io", + "images", + "pull", + "--platform", + platform, + alias, + capture=True, + ) + for source, alias in runtime_images.items(): + info = json.loads( + command( + "docker", + "exec", + args.cluster + "-control-plane", + "crictl", + "inspecti", + alias, + capture=True, + ) + ) + image_evidence[source] = { + "runtime_reference": alias, + "runtime_image_id": info["status"]["id"], + "platform": platform, + } + + def kubectl(*parts, **kwargs): + return command( + "kubectl", + "--kubeconfig", + str(kubeconfig), + "-n", + "default", + *parts, + **kwargs, + ) + + node = args.cluster + "-control-plane" + command( + "docker", + "exec", + node, + "timeout", + "10", + "bash", + "-c", + "exec 3<>/dev/tcp/1.1.1.1/443", + ) + print("PASS node public-IP connectivity before isolation", flush=True) + for binary, destinations in ( + ("iptables", [subnet, "10.244.0.0/16", "10.96.0.0/12"]), + ("ip6tables", []), + ): + command("docker", "exec", node, binary, "-N", "SHIM_SMOKE_EGRESS") + command( + "docker", + "exec", + node, + binary, + "-A", + "SHIM_SMOKE_EGRESS", + "-m", + "conntrack", + "--ctstate", + "ESTABLISHED,RELATED", + "-j", + "RETURN", + ) + command( + "docker", + "exec", + node, + binary, + "-A", + "SHIM_SMOKE_EGRESS", + "-o", + "lo", + "-j", + "RETURN", + ) + for destination in destinations: + command( + "docker", + "exec", + node, + binary, + "-A", + "SHIM_SMOKE_EGRESS", + "-d", + destination, + "-j", + "RETURN", + ) + command( + "docker", + "exec", + node, + binary, + "-A", + "SHIM_SMOKE_EGRESS", + "-j", + "REJECT", + ) + for chain in ("OUTPUT", "FORWARD"): + command( + "docker", + "exec", + node, + binary, + "-I", + chain, + "1", + "-j", + "SHIM_SMOKE_EGRESS", + ) + # Cluster service DNS remains; no upstream DNS forwarding leaves the node. + dns = json.loads( + kubectl( + "-n", + "kube-system", + "get", + "configmap", + "coredns", + "-o", + "json", + capture=True, + ) + ) + dns["data"]["Corefile"] = re.sub( + r"(?m)^\s*forward \. /etc/resolv.conf(?: \{[^}]*\})?\n", + "\n", + dns["data"]["Corefile"], + ) + assert "forward ." not in dns["data"]["Corefile"] + kubectl("-n", "kube-system", "apply", "-f", "-", data=json.dumps(dns)) + kubectl("-n", "kube-system", "rollout", "restart", "deployment/coredns") + kubectl( + "-n", + "kube-system", + "rollout", + "status", + "deployment/coredns", + "--timeout=120s", + ) + + def apply(document): + kubectl("apply", "-f", "-", data=json.dumps(document)) + + def resource(kind, name, **fields): + return { + "apiVersion": "v1", + "kind": kind, + "metadata": {"name": name}, + **fields, + } + + apply( + resource( + "Secret", + "smoke-tls", + stringData={ + name: (directory / name).read_text() + for name in ("tls.crt", "tls.key") + }, + ) + ) + apply( + resource( + "ConfigMap", + "smoke-ca", + data={"ca.pem": (directory / "ca.pem").read_text()}, + ) + ) + apply( + resource( + "ConfigMap", + "smoke-fixture", + data={"fixtures.py": (HERE / "fixtures.py").read_text()}, + ) + ) + apply( + resource( + "Secret", + "smoke-vault-token", + stringData={"token": "disposable-vault-root-token"}, + ) + ) + realm = { + "realm": "shim", + "enabled": True, + "sslRequired": "all", + "groups": [{"name": "owners"}], + "clients": [ + { + "clientId": "shim", + "secret": "disposable-client-secret", + "enabled": True, + "publicClient": False, + "standardFlowEnabled": True, + "redirectUris": ["https://fixture:8443/api/v1/auth/callback"], + "attributes": { + "pkce.code.challenge.method": "S256", + "post.logout.redirect.uris": "https://fixture:8443/login", + }, + "protocolMappers": [ + { + "name": "groups", + "protocol": "openid-connect", + "protocolMapper": "oidc-group-membership-mapper", + "config": { + "claim.name": "groups", + "full.path": "true", + "id.token.claim": "true", + "access.token.claim": "true", + "userinfo.token.claim": "true", + }, + } + ], + } + ], + "users": [ + { + "username": "pilot", + "email": "pilot@example.com", + "emailVerified": True, + "enabled": True, + "firstName": "Pilot", + "lastName": "Test", + "groups": ["/owners"], + "credentials": [ + { + "type": "password", + "value": "disposable-pilot-password", + "temporary": False, + } + ], + } + ], + } + apply( + resource("ConfigMap", "smoke-realm", data={"realm.json": json.dumps(realm)}) + ) + + def fixture( + name, image, ports, command_parts, env=None, mounts=None, volumes=None + ): + container = { + "name": name, + "image": image, + "imagePullPolicy": "IfNotPresent", + "command": command_parts, + "ports": [{"containerPort": port} for port in ports], + "env": [ + {"name": key, "value": value} for key, value in (env or {}).items() + ], + "volumeMounts": mounts or [], + "readinessProbe": {"tcpSocket": {"port": ports[0]}, "periodSeconds": 3}, + } + apply( + { + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": {"name": name}, + "spec": { + "selector": {"matchLabels": {"fixture": name}}, + "template": { + "metadata": {"labels": {"fixture": name}}, + "spec": { + "automountServiceAccountToken": False, + "containers": [container], + "volumes": volumes or [], + }, + }, + }, + } + ) + apply( + resource( + "Service", + name, + spec={ + "selector": {"fixture": name}, + "ports": [ + { + "name": "port-" + str(port), + "port": port, + "targetPort": port, + } + for port in ports + ], + }, + ) + ) + + tls_mount = {"name": "tls", "mountPath": "/tls", "readOnly": True} + tls_volume = {"name": "tls", "secret": {"secretName": "smoke-tls"}} + fixture( + "keycloak", + runtime_images[KEYCLOAK], + [8443], + [ + "/opt/keycloak/bin/kc.sh", + "start-dev", + "--import-realm", + "--hostname=https://keycloak:8443", + "--https-certificate-file=/tls/tls.crt", + "--https-certificate-key-file=/tls/tls.key", + ], + mounts=[ + tls_mount, + {"name": "realm", "mountPath": "/opt/keycloak/data/import"}, + ], + volumes=[ + tls_volume, + {"name": "realm", "configMap": {"name": "smoke-realm"}}, + ], + ) + fixture( + "vault", + runtime_images[VAULT], + [8200], + [ + "vault", + "server", + "-dev", + "-dev-listen-address=0.0.0.0:8200", + "-dev-root-token-id=disposable-vault-root-token", + ], + env={"SKIP_SETCAP": "true"}, + ) + fixture( + "fixture", + runtime_images[overlay], + [8443, 8444, 8445], + ["python", "/fixture/fixtures.py", "serve"], + mounts=[tls_mount, {"name": "script", "mountPath": "/fixture"}], + volumes=[ + tls_volume, + {"name": "script", "configMap": {"name": "smoke-fixture"}}, + ], + ) + organization = str(uuid4()) + backend = { + "POSTGRES_PASSWORD": secrets.token_urlsafe(24), + "DATABASE_URL": "", + "REDIS_URL": "redis://smoke-redis:6379/0", + "SECRET_KEY": secrets.token_urlsafe(48), + "ENCRYPTION_KEY": base64.urlsafe_b64encode( + secrets.token_bytes(32) + ).decode(), + "SHIM_LICENSE_KEY": license_token, + "AUTH_MODE": "oidc", + "OIDC_ISSUER_URL": "https://keycloak:8443/realms/shim", + "OIDC_CLIENT_ID": "shim", + "OIDC_CLIENT_SECRET": "disposable-client-secret", + "OIDC_REDIRECT_URI": "https://fixture:8443/api/v1/auth/callback", + "DASHBOARD_ORIGIN": "https://fixture:8443", + "OIDC_ORGANIZATION_ID": organization, + "OIDC_GROUP_ROLE_MAP": json.dumps({"/owners": "owner"}), + "SECRET_BACKEND": "vault", + "VAULT_ADDR": "https://fixture:8444", + "VAULT_TOKEN_FILE": "/var/run/shim-vault/token", + "MODEL_DEPLOYMENT_ALLOWED_ORIGINS": json.dumps(["https://fixture:8445"]), + "MODEL_DEPLOYMENT_CA_BUNDLE": "/var/run/shim-ca/ca.pem", + } + backend["DATABASE_URL"] = ( + f"postgresql+asyncpg://shim:{backend['POSTGRES_PASSWORD']}@smoke-postgres:5432/shim" + ) + apply(resource("Secret", "smoke-config", stringData=backend)) + overrides = { + "existingSecret": "smoke-config", + "gateway": {"image": runtime_images[overlay]}, + "dashboard": {"image": runtime_images[args.dashboard_image]}, + "vault": {"tokenSecretName": "smoke-vault-token"}, + "caBundleConfigMap": "smoke-ca", + "postgres": { + "enabled": True, + "storage": "1Gi", + "image": runtime_images[values["postgres"]["image"]], + }, + "redis": { + "enabled": True, + "storage": "1Gi", + "image": runtime_images[values["redis"]["image"]], + }, + } + (directory / "values.json").write_text(json.dumps(overrides)) + command( + "helm", + "upgrade", + "--install", + "smoke", + str(chart), + "--kubeconfig", + str(kubeconfig), + "-f", + str(directory / "values.json"), + "--wait", + "--timeout", + "10m", + ) + for name in ( + "gateway", + "dashboard", + "outbox", + "reconciliation", + "compliance", + "ai-act", + ): + kubectl("rollout", "status", "deployment/smoke-" + name, "--timeout=180s") + # Exercise the documented operator bootstrap, then roll out the real tenant UUID. + organization = kubectl( + "exec", + "deployment/smoke-gateway", + "--", + "python", + "ee/scripts/activate_plan.py", + "--create-name", + "Isolated chart smoke", + "enterprise", + capture=True, + ).splitlines()[-1] + backend["OIDC_ORGANIZATION_ID"] = str(UUID(organization)) + apply(resource("Secret", "smoke-config", stringData=backend)) + for name in ("gateway", "outbox", "reconciliation", "compliance", "ai-act"): + kubectl("rollout", "restart", "deployment/smoke-" + name) + for name in ("gateway", "outbox", "reconciliation", "compliance", "ai-act"): + kubectl("rollout", "status", "deployment/smoke-" + name, "--timeout=180s") + kubectl( + "exec", + "-i", + "deployment/smoke-gateway", + "--", + "python", + "-c", + "import sys; from pathlib import Path; Path('/tmp/fixtures.py').write_text(sys.stdin.read())", + data=(HERE / "fixtures.py").read_text(), + ) + kubectl( + "exec", + "deployment/smoke-gateway", + "--", + "python", + "/tmp/fixtures.py", + "smoke", + ) + (directory / "asset-audit.json").write_text( + kubectl( + "exec", + "deployment/smoke-gateway", + "--", + "cat", + "/tmp/shim-asset-audit.json", + capture=True, + ) + ) + (directory / "result.json").write_text( + json.dumps( + { + "passed": True, + "gateway_source_image": original_image, + "dashboard_image": args.dashboard_image, + "test_overlay_image": overlay, + "images": image_evidence, + "internet_isolation": "Node IPv4/IPv6 OUTPUT+FORWARD firewall; no upstream DNS; gateway public-IP runtime probes", + "production_license_unchanged": True, + }, + indent=2, + ) + ) + print( + "PASS chart installation, six deployments ready, complete isolated smoke", + flush=True, + ) + finally: + if kubeconfig.exists(): + result = subprocess.run( + [ + "kubectl", + "--kubeconfig", + str(kubeconfig), + "get", + "pods,jobs", + "-A", + "-o", + "wide", + ], + capture_output=True, + text=True, + ) + (directory / "resources.txt").write_text(result.stdout + result.stderr) + if not args.keep: + command( + "kind", + "delete", + "cluster", + "--name", + args.cluster, + "--kubeconfig", + str(kubeconfig), + ) + command("docker", "network", "rm", network) + print( + f"Test-only image {overlay}; never publish it. Remove local test material with trash {directory} after review." + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--gateway-image", required=True) + parser.add_argument("--dashboard-image", required=True) + parser.add_argument( + "--node-image", + default="kindest/node:v1.37.0@sha256:a1ed56cfb0e7b93589bdf97c8cd566405a265939e3620fc4f5de89adff580ae5", + ) + parser.add_argument("--chart", type=Path, default=HERE.parent / "chart") + parser.add_argument("--cluster", default="shim-on-prem-smoke") + parser.add_argument("--keep", action="store_true") + parser.add_argument( + "--reuse-empty-cluster", + action="store_true", + help="Continue a failed image preload only; refuses application workloads", + ) + args = parser.parse_args() + if any( + any(char.isspace() for char in image) or image.startswith("-") + for image in (args.gateway_image, args.dashboard_image, args.node_image) + ): + parser.error("Image references must be single Docker image names") + if not re.fullmatch(r"[a-z][a-z0-9-]{0,30}", args.cluster): + parser.error("Use a lowercase disposable cluster name, at most 31 characters") + run(args) + + +if __name__ == "__main__": + main() diff --git a/ee/deploy/test/test_reuse.py b/ee/deploy/test/test_reuse.py new file mode 100644 index 0000000..6cc279e --- /dev/null +++ b/ee/deploy/test/test_reuse.py @@ -0,0 +1,67 @@ +"""Rejected reuse must never reach cleanup, including workloads outside default.""" + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +spec = importlib.util.spec_from_file_location( + "smoke_runner", Path(__file__).with_name("run.py") +) +runner = importlib.util.module_from_spec(spec) +spec.loader.exec_module(runner) + + +def test_rejected_cluster_reuse_never_deletes(monkeypatch, tmp_path): + calls = [] + workload = { + "kind": "Deployment", + "metadata": {"namespace": "customer", "name": "database-ui"}, + } + + def command(*args, **kwargs): + calls.append(args) + if args[:3] == ("kind", "get", "clusters"): + return "existing" + if args[:3] == ("kind", "get", "kubeconfig"): + return "disposable config" + if args[0] == "kubectl": + assert "--all-namespaces" in args + return json.dumps({"items": [workload]}) + return "test-image" + + monkeypatch.setattr(runner, "command", command) + monkeypatch.setattr(runner, "certificates", lambda _: None) + monkeypatch.setattr(runner, "test_license", lambda _: "test-token") + monkeypatch.setattr(runner.tempfile, "mkdtemp", lambda **_: str(tmp_path)) + monkeypatch.setattr( + runner.subprocess, "run", lambda *_, **__: SimpleNamespace(returncode=0) + ) + chart = tmp_path / "chart" + chart.mkdir() + (chart / "values.yaml").write_text("{}") + args = SimpleNamespace( + chart=chart, + gateway_image="gateway", + dashboard_image="dashboard", + node_image="node", + cluster="existing", + reuse_empty_cluster=True, + keep=False, + ) + with pytest.raises(SystemExit, match="application workloads"): + runner.run(args) + assert not any("delete" in args or "rm" in args for args in calls) + assert runner.kind_system_workload( + { + "kind": "Deployment", + "metadata": {"namespace": "kube-system", "name": "coredns"}, + }, + "existing", + ) + assert not runner.kind_system_workload( + {"kind": "Pod", "metadata": {"namespace": "kube-system", "name": "unrelated"}}, + "existing", + ) diff --git a/ee/docs/DIAGNOSTIC_METADATA.md b/ee/docs/DIAGNOSTIC_METADATA.md new file mode 100644 index 0000000..537794e --- /dev/null +++ b/ee/docs/DIAGNOSTIC_METADATA.md @@ -0,0 +1,127 @@ +# Diagnostic metadata + +The gateway records the following observations without archiving prompts, +responses, provider error prose, or credentials. Native responses and the one +provider attempt rule are unchanged. + +| Field | Meaning | Unknown/unavailable value | +| --- | --- | --- | +| `provider_finish_reasons` | Map of native completion facts, preserving candidate indices and provider spelling. See the protocol table below. | `null` when no recognized native completion fact was observed; absent map entries remain unknown. | +| `repeat_chain_length` | Number of matching request-content observations in the configured tenant repeat window, including this request (`1` for the first observation). | `null` when the detector has no observation, including Redis unavailability. | +| `ttft_ms` | Floating-point milliseconds from the provider-start callback, after its durable marker commits, to the first nonempty text, refusal, thinking, code, tool arguments, or supported media content observed after restoration. Uses a monotonic clock. | `null` for JSON responses, missing start time, or streams without supported content. | +| `system_prompt_hash` | `hmac-sha256:v1:` followed by a 64-character digest of explicitly supplied system/developer instructions. | `null` when instructions are absent or inherited from provider-held state. | +| `deployment_kind` | `internal`, `external`, or `unknown`, supplied by trusted deployment resolution. | New unclassified requests use `unknown`; historical rows use `null`. | + +The repeat count is a content-match observation, **not evidence of a retry**. +Its existing matching algorithm selects prompt fields plus provider/model, +sorts JSON keys, applies NFKC normalization, and collapses whitespace. The +window and any process/Redis state loss affect comparability. A denied request +that never acquires a durable lifecycle has no diagnostic projection; its +policy-decision event owns denial evidence. + +TTFT is gateway observation time, not the model's internal computation time or +the client's first-byte time. It excludes request admission and input privacy +processing, includes upstream wait and output restoration, and does not count +headers, SSE comments/heartbeats, roles, empty deltas, usage, or terminal events. +Supported media events are OpenAI audio deltas and partial images, and Gemini +inline media. Media contributes to TTFT without being counted as text tokens. +The existing latency and lifecycle status fields remain independent. + +## Native completion facts + +The JSON and SSE paths use the same extraction rules. Only known native enum +values enter telemetry; malformed, unspecified, or unrecognized future values +remain unknown until support is reviewed. Free-form finish messages and stop +sequence text are never retained. + +| Protocol | Map keys | Examples | +| --- | --- | --- | +| OpenAI Chat | `choices..finish_reason` | `stop`, `length`, `tool_calls`, `content_filter` | +| OpenAI Responses | `status`, `incomplete_details.reason` | `completed`; `incomplete` with `max_output_tokens`; `failed`; `cancelled` | +| Anthropic Messages | `stop_reason` | `end_turn`, `max_tokens`, `tool_use`, `refusal` | +| Gemini | `candidates..finishReason`, `promptFeedback.blockReason` | `STOP`, `MAX_TOKENS`, `SAFETY` | + +Native indices are used when present; otherwise the candidate's array position +is used. Completion facts already observed survive later stream failure or +disconnect. `completed` lifecycle status means the transport completed; a +native truncation/refusal reason can still accompany it. `[DONE]` does not +fabricate a finish reason. Missing provider usage does not erase completion +facts or TTFT; the existing `usage_estimated` field describes accounting fallback. + +## System-instruction hashing + +The enterprise quota reservation computes the digest before input privacy +transformation. Hash material is the JSON array +`["shim.system_prompt.v1", tenant_id, protocol, deployment, instructions]`: + +- Chat: only `role` and `content` from ordered system/developer messages. +- Responses: explicit `instructions` and ordered system/developer `input` messages. +- Messages: explicit `system` value. +- Gemini: explicit `systemInstruction` value. + +Canonical JSON sorts object keys, uses compact separators and ASCII escapes, +and preserves array order, content whitespace, and Unicode without normalization. +The HMAC key is `COMPLIANCE_HASH_SALT`, falling back to `SECRET_KEY`. Keep that +key secret and unique per installation. Comparisons are scoped to that key, +tenant, protocol, deployment identity/kind, and algorithm version; changing the key ends comparability +with earlier digests. The deployment scope uses its stable registry UUID and internal/external/unknown +kind; unregistered requests use a null UUID. Alias, endpoint, upstream model and +declared-version edits preserve comparability for the same deployment ID/kind. +Model names and user conversation content are excluded. +An explicitly empty instruction differs from an absent instruction. Provider-held +prompts, previous responses, and cached instructions are not reconstructed. + +## Persistence and reading + +Request fields enter `request_lifecycle.metadata` during quota reservation. +Completion facts and TTFT join them in the terminal accounting transaction, +before audit/analytics outbox intent is constructed. Terminal replay preserves +the first committed observations. Analytics delivery copies them into +`request_logs.details` with the existing tenant/request idempotency constraint. +No table or column migration is needed for these existing JSONB fields. + +`GET /api/v1/management/requests` exposes the five optional fields on each +request; `/requests/export` includes the same fields in CSV (unknown values +are empty cells, and finish-reason maps are JSON). Audit completion `extra` +carries the same fields. Historical rows and +old outbox messages read as null without invented backfills. The community +JSONL event also contains them, with `system_prompt_hash: null` because community +has no configured installation hashing key. + +## Unpriced deployment costs + +An unknown deployment price is not a free request. A terminal spend settlement +marked `event_metadata.pricing.pricing_resolution = "unknown"` is exposed by the +request API as `cost_usd: null` and `cost_complete: false`; CSV exports use an +empty cost cell and `cost_complete: False`. The request summary reports +`unpriced_requests`, `cost_complete`, and a `settled_spend_usd` subtotal containing +only priced settlements. These checks read the tenant-scoped ledger directly, +so missing projection metadata cannot turn an unknown settlement into zero. + +Overview summary and trend costs are null whenever their period includes an +unpriced settlement, with the same completeness flag and request count. Empty +periods retain known zero costs. + +Billing daily and grouped rows expose null costs for groups with an unpriced +settlement. Billing usage totals are also null when incomplete. CSV exports use +empty cost cells plus completeness/count columns; PDF exports label those +groups `Unknown` with the number of unpriced requests. + +Budget alerts retain a known-settlement subtotal in `current_usd` and label it +`cost_basis: known_settled_spend`, with `cost_complete` and `unpriced_requests`. +Their contributor rows have the same completeness markers. Threshold evaluation +uses known spend and settled tokens; missing prices cannot imply full coverage. + +Analytics `details` and audit completion `extra` also carry `pricing_resolution`. +Their existing numeric cost fields reflect the ledger placeholder when it is +unknown; consumers must inspect that marker. Refunds and requests without a +spend settlement have known settled cost zero. Pricing completeness does not +claim provider-invoice accuracy or actual token measurement; `usage_estimated` +continues to describe token fallback independently. + +To verify the contract, run the streaming/community tests and +`ee/tests/gateway/kernel/test_accounting_coordinator.py` plus +`ee/tests/gateway/api/test_management.py` against disposable PostgreSQL/Redis. + +The metadata uses existing JSONB storage and adds no database transaction or +network call. Measure overhead with representative payloads and concurrency. diff --git a/ee/docs/MODEL_DEPLOYMENTS.md b/ee/docs/MODEL_DEPLOYMENTS.md new file mode 100644 index 0000000..c267e59 --- /dev/null +++ b/ee/docs/MODEL_DEPLOYMENTS.md @@ -0,0 +1,77 @@ +# Model deployments + +Workspace owners and admins register models at **Gateway → Model deployments** +or `/api/v1/management/model-deployments`. Create a stored provider credential +first. The registry stores its tenant-scoped identifier, never credential text. +A deployment has a stable UUID, gateway alias, upstream model, protocol, base +URL, timeout, internal/external classification, owner and declared version. +Updates and health checks produce management audit events. + +Set `MODEL_DEPLOYMENT_REQUIRED=true` for registry-only inference. With `false`, +registered aliases override public-catalog routing. Disabled aliases remain +denied. Gateway-key model allowlists use tenant aliases; unknown aliases cannot +be assigned. Registry-routed Responses requests must include the gateway alias, +including continuations with `previous_response_id`; the continuation identifier +does not select a deployment. Restrict each key to the models its workload needs. Registry +checks cover requests passing through shim, not direct access to other servers. + +The platform operator sets `MODEL_DEPLOYMENT_ALLOWED_ORIGINS` to a JSON list of +exact scheme/host/port origins, for example `["https://models.internal:8443"]`. +API users cannot expand this policy. URL credentials, queries, fragments, +metadata addresses and redirects are rejected. Internal HTTP origins require +explicit approval; use HTTPS in production. Add private certificate authorities +with `MODEL_DEPLOYMENT_CA_BUNDLE`; certificate verification stays enabled. +Enforce DNS and outbound network policy at the deployment boundary as well. + +For OpenAI-compatible deployments, use the API base including `/v1`; for +Anthropic use the server root. shim uses its existing native transports, masks +configured sensitive content before forwarding, and makes one provider attempt. +There is no retry or failover. A five-second model-list health probe records +only HTTP success/failure and never reads an unbounded response body. Health is +an observation, not a request-routing or version-verification guarantee. + +The declared version/hash/digest is operator supplied and included in the +registry policy evidence. Pin the actual serving image and model revision at +the model server; shim does not independently attest which weights it serves. + +Public catalog prices apply to recognized upstream model names. Custom models +are explicitly unpriced: ledger arithmetic uses a zero placeholder with +`pricing_resolution=unknown`, and usage reports carry completeness counters. +A monetary provider limit rejects unpriced inference instead of treating it as +free. Token and request quotas still apply. + +## Verified compatibility + +`uv run --locked python -m pytest -q ee/tests/tenants/test_deployments.py` +uses the actual provider SDK over a mock HTTP transport, two distinct endpoint +origins and PostgreSQL accounting. It verifies OpenAI Chat Completions and +Responses JSON/SSE, upstream model selection, tenant-scoped credentials, +privacy, one-attempt errors, isolated endpoint circuits, model restrictions and +unpriced spending limits. Health tests exercise the bounded model-list probe. + +These checks establish shim's wire behavior. They do not certify every vLLM, +NIM, TGI or Ollama version. Run the same request families against each selected +serving version before enabling it; unsupported protocol capabilities return +native errors rather than an emulated result. + +## Anthropic token counting + +`POST /v1/messages/count_tokens` preserves Anthropic's native request/response +and beta forms. It uses the same gateway authentication, model permissions, +RPM/TPM limits and privacy transformation as Messages. Counts describe the +transformed payload that would be sent upstream, not the original sensitive +text. Counting has a separate repeat identity from inference. Streaming is not +supported by this endpoint. + +Token counting makes one nonbillable provider request. It persists audit preflight before forwarding and an audit +completion/outbox intent with `operation_type=token_count` and +`billable_execution=false`; it never reserves billable quota or provider spend, +starts an inference lifecycle, or creates a settlement. Community mode leaves +its billable JSONL usage stream unchanged. Strict enterprise audit mode fails +closed before forwarding if preflight cannot be persisted, and returns an error if +completion persistence fails after the response. + +Run `uv run --locked python -m pytest -q tests/gateway/test_token_count.py + ee/tests/tenants/test_deployments.py` to check native SDK routing, privacy, +errors and durable nonbillable accounting. Registry tests also exercise +Anthropic Messages JSON and SSE. diff --git a/ee/docs/OFFLINE_RELEASES.md b/ee/docs/OFFLINE_RELEASES.md new file mode 100644 index 0000000..d3a79a1 --- /dev/null +++ b/ee/docs/OFFLINE_RELEASES.md @@ -0,0 +1,164 @@ +# Signed releases and offline bundles + +Enterprise tags are `enterprise-v` and dashboard tags are +`dashboard-v`. These publish artifacts only; they do not match the +production `v..` deployment trigger. Each repository builds +its own image with its own GitHub package token. The dashboard repository and +its packages remain private. Give the release operator explicit read access; +no public workflow checks out the dashboard repository. + +The workflows deliver one archive per `linux/amd64` and `linux/arm64` platform: +image archive, immutable source reference, component version, commit, SPDX +SBOM, BuildKit provenance, public key, signed checksums. They also sign the +registry image digest. Backend **0.1.3** and dashboard **0.1.0** are separate +component versions; a bundle ID such as `example-release-1` identifies their +combination. Publishing requires tag versions to match package metadata. +Workflows are not proof that an unpublished image has been built or tested. + +## Trust and release keys + +Use cosign **3.1.3**. Provision `RELEASE_COSIGN_PRIVATE_KEY` (encrypted PEM) and +`RELEASE_COSIGN_PASSWORD` in each repository's release secrets. Restrict tag +creation and changes to release workflows using repository rules. Release +signing authorization is separate from production deployment authorization. + +Distribute the public key and its SHA-256 file fingerprint over an authenticated +channel before delivery. The operator maintains a local approved key file; a +public key delivered beside artifacts is informational, never a new trust root. +A release signer can approve artifacts, so protect that key accordingly. + +This scheme uses explicit keys, without Fulcio identity, online Rekor lookup, +or timestamp-based trust. `--use-signing-config=false --tlog-upload=false` +disables online signing services in this pinned cosign version. Verification +uses `--insecure-ignore-tlog` intentionally: the trusted local key and signature +remain mandatory. No transparency or signing-time guarantee is claimed. + +Before rotation, authenticate the new key fingerprint through the same operator +channel, document its first approved release, and retain the previous key only +for approved rollback releases. After compromise, revoke the old key in operator +policy and reissue affected deliveries; do not accept whichever key accompanies +a bundle. Pin the expected bundle ID to avoid silently accepting an older valid +release. The current offline startup licence key and release key are independent. + +## Verify a component delivery without internet + +After extracting the delivery into a fresh staging directory, use the approved +key, not `release.pub` from the delivery: + +```sh +cosign verify-blob --key /etc/shim/trust/release.pub --insecure-ignore-tlog \ + --bundle SHA256SUMS.sigstore.json SHA256SUMS +sha256sum --check SHA256SUMS +``` + +Both commands must succeed before reading metadata as trusted or importing the +image. Check `image.json` for the approved component version and platform. +The signed checksums bind the archive, SBOM and provenance to that exact source +reference. BuildKit provenance is covered by the delivery signature. Its builder and source +claims are evidence from the authorized release workflow, not a separate +keyless identity attestation. +On macOS the equivalent checksum command is `shasum -a 256 --check SHA256SUMS`. + +## Assemble a complete installation bundle + +Use Python 3.13, Docker 28+ and cosign 3.1.3 on a connected release workstation. +The assembler is `ee/scripts/offline_bundle.py`, provisioned from reviewed source. +Verify component deliveries first. Resolve images to immutable references; use +only reviewed releases. Package the Helm chart with `helm package` before adding +it to the bundle. Its dependencies must be vendored; no dependency update is run +at install time. + +Create a JSON specification beside the local files (replace every placeholder): + +```json +{ + "bundle_id": "example-release-1", + "platform": "linux/amd64", + "images": [ + {"name": "backend", "version": "0.1.3", "source": "ghcr.io/getshim/shim-enterprise@sha256:<64-hex-digest>"}, + {"name": "dashboard", "version": "0.1.0", "source": "ghcr.io/durthvadr/shim-dashboard@sha256:<64-hex-digest>"} + ], + "files": ["shim-enterprise-0.1.0.tgz", "OFFLINE_RELEASES.md", "OPERATIONS.md", "backend.spdx.json", "dashboard.spdx.json", "backend.provenance.json", "dashboard.provenance.json"] +} +``` + +Also list **every** additional runtime image required by the selected topology: +PostgreSQL, Redis, internal IdP, Vault, internal model server, and any test or +migration image distinct from the backend. Services supplied by the customer +must instead be explicitly recorded in the operator inventory with their local +endpoints. Include their SBOMs, provenance, licences, installation instructions, +CA material and chart dependencies in `files`. No secret credentials belong in +the bundle. Backend/dashboard licences and notices must accompany delivery; +keep Apache-2.0 and Elastic-2.0 materials distinguishable. Include approved +Python/cosign/Helm/Docker installation media if the isolated workstation does not +already have them. Provision the verifier and trusted key independently. + +The explicit inventory is a release-operator responsibility: the assembler does +not discover hidden dependencies or certify that the inventory is complete. +It refuses tag-only image references and existing output directories. + +```sh +python ee/scripts/offline_bundle.py create --spec bundle-spec.json \ + --directory example-release-1 --key /secure/release.key +``` + +Transfer the resulting directory. Its signed manifest contains file SHA-256s, +source registry digests, archive-derived configuration digests, versions and +platform. Manifest schema 2 replaces the pre-release schema 1 `image_id` field; +rebuild older bundles rather than trusting a source daemon's image ID. +No image download occurs during verification or import: + +```sh +python /opt/shim/offline_bundle.py verify --directory example-release-1 \ + --key /etc/shim/trust/release.pub --expected-id example-release-1 +python /opt/shim/offline_bundle.py import --directory example-release-1 \ + --key /etc/shim/trust/release.pub --expected-id example-release-1 +``` + +Import rechecks **all** files before the first Docker load. Use a staging directory +writable only by the operator and do not modify it during verification/import. +Import uses the ID actually returned by Docker, re-exports that object to a +temporary archive, and checks its configuration digest and platform against the +original. The configuration digest includes root filesystem layer IDs. Keep +enough temporary disk space for one image archive. Classic Docker and containerd +can use different local identities, including a newly synthesized manifest ID. +Docker save/load does **not** promise to retain registry RepoDigests; neither the +source image ID nor a config digest is a portable lookup reference. To seed an +internal registry, tag the verified ID printed by import, push it, record the new +registry digest, and configure Helm with that internal immutable reference. +Preserve the original signed manifest and import output as mapping evidence. +For example: + +```sh +docker tag sha256: registry.internal/shim/backend:0.1.3 +docker push registry.internal/shim/backend:0.1.3 +# Record the digest returned by push; use registry.internal/shim/backend@sha256:... +``` + +Install the local chart with `imagePullPolicy: IfNotPresent` and internal digest +references. Configure local PostgreSQL/Redis, OIDC, Vault, model origins and TLS +using the chart's operator instructions. No public service may remain mandatory. +Use the backend image for migrations; back up PostgreSQL before upgrading. +Roll back only to a schema-compatible image; never downgrade the production +schema. Retain the previous verified bundle and internal digest inventory for +recovery. A clean recovery also needs database backups and separately protected +operator secrets; the release bundle contains neither. + +## Verification + +`uv run --locked python -m pytest -q ee/tests/scripts/test_offline_bundle.py` +uses real local cosign keys and rejects archive tampering, manifest tampering, +an unexpected signer, wrong bundle ID, extra files and symlinks. HTTP(S) requests +are directed to a refused local proxy; Docker is replaced with a small archive +fixture. The existing `ee/tests/core/test_license.py` covers missing, invalid, +valid and expiry/grace boundaries. Release verification is separate from the existing startup licence check. + +These checks are not the full air-gap rehearsal. Before declaring offline +readiness, an independent operator must use a clean environment with internet +egress blocked, verify and import the full inventory, install, log in through +the internal IdP, call an internal model, inspect usage/audit records, and restart +services. Record platform, bundle ID, internal digests, exact commands and results. +Real customer TLS/IdP, offline tool provisioning, backup/restore and schema-safe +upgrade/rollback remain deployment acceptance checks. + +BuildKit provenance extraction follows the [Docker CLI reference](https://docs.docker.com/reference/cli/docker/buildx/imagetools/inspect/). diff --git a/ee/docs/ON_PREM_IDENTITY.md b/ee/docs/ON_PREM_IDENTITY.md new file mode 100644 index 0000000..b3f580c --- /dev/null +++ b/ee/docs/ON_PREM_IDENTITY.md @@ -0,0 +1,162 @@ +# Customer identity and secrets + +The connected on-prem composition uses customer OIDC, PostgreSQL, Redis, and +Vault KV v2. Set `AUTH_MODE=oidc` explicitly. Hosted Supabase authentication +remains available with `AUTH_MODE=supabase`; the OIDC path needs neither +`SUPABASE_URL` nor `SUPABASE_KEY` and never initializes a Supabase client. + +## Bootstrap and configuration + +1. Apply migrations and provision an organization using + `python ee/scripts/activate_plan.py --create-name 'Example organization' enterprise`. + Record the printed organization UUID. The operator selects capacity separately; + this guide does not issue a licence or prescribe customer limits. +2. Create a confidential OIDC client with authorization code flow and PKCE S256. + Register exactly `https://shim.internal/api/v1/auth/callback` as a redirect and + `https://shim.internal/login` as a permitted post-logout redirect. +3. Configure the following backend variables through Kubernetes Secrets/config: + +| Variable | Example or meaning | +| --- | --- | +| `AUTH_MODE` | `oidc` | +| `OIDC_ISSUER_URL` | Exact discovery issuer, e.g. `https://identity.internal/realms/customer` | +| `OIDC_CLIENT_ID` / `OIDC_CLIENT_SECRET` | Confidential client credentials | +| `OIDC_REDIRECT_URI` | `https://shim.internal/api/v1/auth/callback` | +| `DASHBOARD_ORIGIN` | `https://shim.internal` | +| `OIDC_ORGANIZATION_ID` | Operator-provisioned organization UUID | +| `OIDC_GROUPS_CLAIM` | Top-level group array claim; defaults to `groups` | +| `OIDC_GROUP_ROLE_MAP` | JSON, e.g. `{"/shim/owners":"owner","/shim/users":"member"}` | +| `OIDC_TEAM_GROUP_MAP` | Optional JSON group-to-team mapping, e.g. `{"/shim/platform":{"team_id":"","role":"team_admin"}}`; teams must already belong to the configured organization | +| `OIDC_SESSION_SECONDS` | Absolute session lifetime; default 28,800, maximum 86,400 | +| `OIDC_REVALIDATE_SECONDS` | Refresh/group revalidation interval; default 60, maximum 300 | +| `OIDC_API_AUDIENCE` | Optional, distinct API audience for direct bearer access; absent disables it | +| `OIDC_API_MAX_TOKEN_SECONDS` | Maximum bearer token lifetime; default 300, maximum 900 | +| `SECRET_KEY` | High-entropy shared server secret; protects login state and encrypts Redis session material | + +All application replicas use the same issuer/client/tenant mapping and +`SECRET_KEY`. Changing that secret invalidates current login/session material. +`ENVIRONMENT=production` requires HTTPS for issuer, dashboard/callback, and Vault, +a supported managed secret backend, and the existing offline licence check. +Only local development can use HTTP. Never disable certificate verification. + +Use a customer CA bundle through `SSL_CERT_FILE` for Python HTTP clients and +`NODE_EXTRA_CA_CERTS` for the dashboard. Standard HTTP(S) proxy and `NO_PROXY` +variables are honored by HTTPX; direct Node-to-API networking must be permitted. +Do not assume a proxy environment variable changes Node fetch routing. + +## Identity and authorization contract + +The issuer and subject bind an OIDC identity to its local user. First login +requires a valid email and `email_verified=true`; it never adopts a user based +on a matching email. Email collisions require operator resolution. A token +cannot select an arbitrary organization. Changing the configured tenant does +not migrate existing identities. A locally deactivated user stays deactivated. + +Only mapped groups grant access. The strongest configured role wins; owner, +admin, member, then auditor. User-supplied profile metadata does not grant roles. +Group removal/downgrade is applied on session refresh, at most +`OIDC_REVALIDATE_SECONDS` after the identity provider reflects the change. +Providers must return a newly signed ID token with current groups on refresh. +A provider without refresh tokens requires sign-in again at revalidation. +Explicit bearer tokens must have the configured API audience, valid signature, +issuer, subject, issue/expiry times, and a lifetime within +`OIDC_API_MAX_TOKEN_SECONDS`; existing tokens may remain valid for that bound. + +The browser receives an opaque HttpOnly, SameSite=Lax session cookie, Secure in +production. Token/refresh material is encrypted in Redis, never returned to +browser JavaScript. State, nonce, and PKCE are handled by Authlib. Cookie-based +mutations require the exact dashboard Origin. Redis/identity outages fail +closed. Logout revokes the local session and uses provider discovery for +end-session navigation without exposing an ID token in the browser URL. + +Recovery uses the customer's identity-provider administrator: restore access to +a configured owner group, then sign in again. No local password backdoor, +email invitation, or implicit first-user privilege escalation is added. Human +OIDC session revocation and application API-key revocation are separate +lifecycles: removing a human from IdP groups does not revoke workload keys they +created. Revoke workload keys through the administration API when retiring an +integration; local account deactivation also blocks its keys. + +To retire a person, a remaining organization owner calls +`DELETE /api/v1/management/team/members/{user_id}` (with their own authenticated +session and the dashboard Origin). This deactivates the local account and +revokes every active key owned by that account in one transaction. It refuses +to remove the last owner: provision another mapped owner first. Remove the +person's IdP groups as well; later IdP login cannot reactivate the local account. +To retire only an integration, call +`DELETE /api/v1/management/api-keys/{api_key_id}`. Verify the retired key returns +HTTP 401 on an authenticated gateway request before closing the offboarding +record. A group-only change intentionally leaves workload keys usable. + +Keycloak: map a group-membership claim to both ID/access tokens and enable +verified email for permitted users. Full group paths avoid colliding leaf names. +For Entra ID, use the tenant-specific issuer and application group claims; +explicitly configure a verified email assertion for provisioning. Group-overage +claims requiring Microsoft Graph are not fetched implicitly: require the group +array in the token. Actual Entra tenant acceptance remains required. LDAP/AD +federation belongs in the chosen IdP. Direct LDAP, SCIM, SAML, and local passwords +are not supported. + +## Vault KV v2 + +Set `SECRET_BACKEND=vault`, `VAULT_ADDR=https://vault.internal`, +`VAULT_KV_MOUNT=secret`, and `VAULT_TOKEN_FILE=/run/vault/token`. +`VAULT_NAMESPACE` is optional. Use Vault Agent Kubernetes/AppRole auto-auth and a +read-only token-sink mount. shim rereads the token on every request so Agent +renewal/replacement requires no process restart. Tokens never appear in secret +references or database records. + +A minimal KV v2 policy for the configured mount is: + +```hcl +path "secret/data/shim/*" { + capabilities = ["create", "read", "update"] +} +path "secret/destroy/shim/*" { + capabilities = ["update"] +} +``` + +References pin mount, opaque tenant namespace, object identifier, and numeric +version. Envelopes also verify tenant and purpose. Rotation creates a new object, +so prior references keep their meaning until the existing committed cleanup +path removes them. Deletion destroys the specified version. Vault owns its own +unseal, storage encryption, HA, backups, and token renewal. Fernet remains +forbidden as a production provider-secret backend. + +## Dashboard and network contract + +Build the enterprise dashboard with `NEXT_PUBLIC_AUTH_MODE=oidc` and run with +`SHIM_API_URL` set to the internal gateway origin. `/api/v1/*` forwards at runtime +only to that configured origin, preserving safe session cookies and redirects. +Browser login/management calls use same-origin paths; no baked public API URL is +used for those calls. Public origin build variables still drive documentation +examples and metadata. Enterprise `/` opens the dashboard; hosted registration, +password reset, invitations, and public playground paths are unavailable. + +Permit dashboard-to-gateway, browser-to-IdP, gateway-to-IdP, database, Redis, +Vault, and explicitly configured internal model endpoints. Leave email/hosted +telemetry credentials unset unless intentionally enabling those destinations. +Run enabled workers with the same internal service configuration. Configure +reverse proxy/ingress access logs to omit authentication query strings; shim +redacts `/api/v1/auth/*` query strings from Uvicorn access logs. Production Next +servers should not run with development request logging. + +## Verification + +Backend checks: `uv run --locked python -m pytest -q ee/tests/tenants/test_oidc.py +ee/tests/secrets/test_vault.py`. These exercise RSA/JWKS validation, Authlib's +actual code/PKCE exchange, tenant/subject binding, email collisions, role removal, +encrypted Redis sessions, Origin rejection, expiry, and logout. + +Browser check: run `npm run test:e2e -- e2e/oidc.spec.ts` with +`PLAYWRIGHT_BASE_URL`, `NEXT_PUBLIC_AUTH_MODE=oidc`, `OIDC_TEST_ISSUER`, +`OIDC_TEST_USERNAME`, and `OIDC_TEST_PASSWORD` against disposable services. Never +use production identities or commit credentials. The test creates/revokes a +key, checks account controls and logout, and rejects unexpected browser hosts. +For local Next development with a numeric loopback URL, bind `next dev` to that +same hostname to avoid Next's dev-origin protection blocking hydration. + +Production HTTPS/CA, network-denied Kubernetes, IdP lifecycle, and an actual +Entra tenant require operator acceptance in the deployment record; a local mock +or HTTP development run does not establish those results. diff --git a/ee/docs/POLICY_DECISIONS.md b/ee/docs/POLICY_DECISIONS.md new file mode 100644 index 0000000..ac7f784 --- /dev/null +++ b/ee/docs/POLICY_DECISIONS.md @@ -0,0 +1,69 @@ +# Gateway decision evidence + +Gateway audit events include `policy_verdicts` for the checks actually evaluated. +Each verdict contains `schema_version`, `rule_id`, `rule_version`, +`policy_version`, `stage`, `outcome`, `reason_code`, and `effective_at` (UTC). +`effective_at` is the evaluation time, not a claim about when a configuration +first became effective. Configuration digests or the locked accounting policy +version identify the evaluated snapshot; rule version 1 identifies the shipped +check semantics. No policy document, prompt, response, credential, or raw error +message is included. + +The event envelope binds all verdicts to the request, tenant, actor type and +API key or authenticated user. An API key identifies the key, including when +shared: it does not identify its owner as the person making the request. +Unrecognized credentials and policy-resolution failures without a verified +tenant remain sanitized authentication/error telemetry; they are never assigned +to a guessed tenant or user. HTTP/body validation before the gateway invocation +also remains outside tenant decision evidence. + +| Rule | Evidence | +| --- | --- | +| `tenant.allowed_providers` | Provider passed or failed the tenant allowlist. | +| `tenant.zero_retention_request` | Required request flags/options passed or failed the gateway's check, or the check was not required. This does not attest to the provider's wider retention practices. | +| `gateway.model_catalog` | Model is supported by the catalog snapshot, or was rejected. Unsupported caller-supplied model text is omitted from rejected enterprise events. | +| `rate.requests`, `rate.tokens`, `rate.repeated_requests` | Configured burst/repeat checks passed, were unlimited, or denied admission. Repeated content does not establish an automatic retry. | +| `quota.requests_and_tokens` | Atomic request/token reservation passed or was rejected, using the policy loaded under the accounting lock. The combined limit is not attributed to a particular counter when the atomic check cannot distinguish it. | +| `privacy.input` | Scrubbing masked data, found no enabled entity, was disabled, blocked unsupported content, or failed closed. | +| `spend.provider_monthly` | Provider spending reservation passed, was unlimited, was rejected, or could not be evaluated. Invocation-scoped BYOK remains outside the stored-provider cap. | +| `gateway.admission` | Other admission validation failed or admission infrastructure was unavailable. | + +`allow` means that individual check passed; it does not imply that the entire +request completed. `mask`, `deny`, `error`, and `skip` are distinct outcomes. +Later stages are absent when an earlier stage stopped execution. Terminal +lifecycle status remains separate from policy results and provider behavior. + +## Durability and failures + +Existing short quota, privacy, and spend transactions retain decision snapshots +in lifecycle metadata. Terminal settlement/refund builds the existing audit +completion/outbox event from those snapshots. Failures before a quota lifecycle +exists create an audit completion and committed outbox intent directly, without +usage charges or provider execution. The identity remains +`request::outbox:audit.completion`. Outbox redelivery uses the existing +tenant/request/event deduplication and hash-chain writer. + +For pre-admission denials, audit mode `off` writes no audit intent. `best_effort` +attempts the transaction and emits a content-free error log if it cannot commit, +preserving the original rejection. `strict` returns `AUDIT_INTENT_FAILED` (503) +when required evidence cannot commit; no provider attempt occurs. + +Admitted accounting retains its existing atomic audit behavior: a required +preflight failure prevents provider execution, and an audit completion failure +rolls back settlement for reconciliation. Accounting truth is never discarded to +make evidence delivery appear successful. A strict failure while finalizing a +denial is surfaced instead of swallowed; best-effort failure retains recovery +and a sanitized error log. Database unavailability cannot guarantee new durable +evidence, and best-effort mode must not be described as lossless. + +After worker delivery, the tenant-scoped `GET /v1/compliance/audit/logs` response +exposes verdicts, caller key/user identity, actor type, and terminal lifecycle +status. Unknown historical fields remain null. The existing chain verifier +continues to verify this evidence. +Decision evidence records existing gateway checks. It is not a configurable +policy engine, content archive, signature or independent trust anchor. + +Verification: `uv run --locked python -m pytest -q + ee/tests/gateway/pipeline/test_decisions.py` covers real quota/spend transactions, +pre-admission denials, masking and privacy rejection, audit failure modes, +redelivery, and chain verification. The repository-wide gate is in `AGENTS.md`. diff --git a/ee/docs/PROVISIONING.md b/ee/docs/PROVISIONING.md new file mode 100644 index 0000000..45c07a5 --- /dev/null +++ b/ee/docs/PROVISIONING.md @@ -0,0 +1,83 @@ +# Operator-managed enterprise plans + +## Provisioning and subscription transition + +Enterprise plans are provisioned by an authorized operator after the existing +commercial approval process. PostgreSQL `organizations.tier` and the existing +`tier_definitions` remain authoritative for access and quota policy. The +enterprise startup licence check is separate and remains offline. + +Lemon Squeezy checkout, portal links, webhooks, and configuration are retired. +No migration changes existing organization tiers, API keys, usage, or billing +history. Historical billing columns and `billing_webhook_receipts` stay in the +schema; do not drop or clear them during this transition. + +## Provision a new organization + +Configure the enterprise environment, apply migrations, and run from the +backend repository using an operator database account: + +```bash +uv run --locked --package shim-enterprise python ee/scripts/activate_plan.py \ + --create-name 'Example organization' enterprise +``` + +The command prints the new organization UUID. It creates the organization, +privacy defaults, and plan in one transaction without an identity account, +password, email, or subscription-service call. Each `--create-name` invocation +creates a separate tenant, even when the display name matches; retain the UUID. +Configure the identity provider for that tenant using this UUID (for OIDC, +`OIDC_ORGANIZATION_ID`) and the deployment's approved owner-group mapping. + +Existing organization activation and subsequent plan changes use its UUID: + +```bash +uv run --locked --package shim-enterprise python ee/scripts/activate_plan.py \ + ORGANIZATION_UUID enterprise +``` + +Allowed tiers remain `free`, `managed`, `agency`, and `enterprise`. The command +locks the organization, updates its current plan/source/status, and applies the +tier to active API keys atomically. Newly issued keys inherit it. Revoked keys +stay revoked; usage counters, reservations, and the immutable ledger are not +reset. Unknown tenants/tiers fail without committing a partial change. + +## Transition an existing customer + +1. Before retiring the old integration in a deployed environment, record each + affected customer's tenant UUID, approved tier, last billing status/period, + external references, and the operator responsible for future plan changes. +2. Reconcile any renewal/cancellation/payment obligations through the existing + commercial process. Removing the integration does not cancel or settle a + subscription at the billing service. +3. Run the activation command with the customer's approved **existing tier**. + Change the tier only when that change has been approved. The command retains + external customer/subscription/variant identifiers, historical period and + cancellation fields, portal references, and webhook receipts. +4. Verify an existing key still authenticates, a new key inherits the tier, and + the dashboard shows the expected entitlements and usage. Operator-managed + status is `active` for paid tiers and `free` for the free tier. +5. Remove the obsolete `LEMON_SQUEEZY_*` deployment values and the billing-service + webhook destination after the agreed cutover. Deploy the matching dashboard + contract with the backend. + +Historical period/cancellation fields are no longer a renewal schedule. An +operator must apply future access changes according to the existing agreement; +this release adds no automatic expiration or new commercial terms. + +## API and verification + +`POST /api/v1/webhooks/lemonsqueezy` is removed. +`GET /api/v1/management/subscription` remains a tenant-scoped read endpoint with +`plan`, `status`, `source`, and `entitlements`. Purchase/portal/renewal fields are +removed from its response. The dashboard displays the configured plan and directs +plan changes to an administrator/contact flow. + +```bash +uv run --locked python -m pytest -q ee/tests/tenants/test_plans.py +``` + +These regression checks cover new provisioning, invalid inputs, transitions +across all existing tiers, active/revoked/new keys, tenant separation, historical +billing records, ledger values, and quota counters. They use disposable +PostgreSQL and make no subscription-service calls. diff --git a/ee/docs/team-access.md b/ee/docs/team-access.md new file mode 100644 index 0000000..9b232a2 --- /dev/null +++ b/ee/docs/team-access.md @@ -0,0 +1,76 @@ +# Teams and gateway key controls + +Teams belong to one organization. A membership names an existing organization +user and grants `member` or `team_admin` access to that team. An organization +role remains `owner`, `admin`, `member`, or `auditor`; a team administrator is +an organization member with a delegated membership, not another global role. + +| Role | Gateway keys | Teams and memberships | Organization changes | +| --- | --- | --- | --- | +| Owner | All organization keys and policies | Create teams, set quotas, assign team administrators | Existing owner permissions | +| Admin | All organization keys and policies | Create teams, set quotas, assign team administrators | Existing admin permissions; cannot promote organization roles | +| Team administrator | Own keys and keys assigned to administered teams; edit those team keys' model policies | Add/remove ordinary members in administered teams | None | +| Member | Own keys; cannot loosen an existing key's model policy or reassign its access team | Read assigned teams | None | +| Auditor | Read key metadata; cannot issue, rotate, revoke, or use keys | Read organization teams and memberships | Read only | + +Existing organization overview and member-directory visibility is preserved. +Team delegation limits mutation authority; it does not introduce a separate +tenant or hide existing organization aggregate dashboards. + +Use **Workspace → Teams** to create a team, configure quotas, and assign +members. Use **Gateway → Keys** to assign a key's access team and model list. +Only organization owners/admins can move a key between teams. A member with +team memberships must select a team when creating a key. Removing membership +denies subsequent inference through that user's team keys; organization +owners/admins retain organization-wide authority. + +## Attribution and migration + +The existing `api_keys.team` value remains a billing label. The new `team_id` +is an explicit access and quota reference. They are intentionally independent. +The migration copies distinct existing labels into organization-owned team +names, preserving the original spelling. It does not assign memberships or +bind existing keys. Administrators explicitly assign access teams after +migration; historical request and billing labels are unchanged. + +## Rotation and model policies + +`POST /api/v1/management/api-keys/{id}/rotate` replaces the key's one-way verifier +in the same database row. The old secret stops authenticating when the +transaction commits. The response shows the new plaintext once. The key ID, +owner, team, expiry, tier, model policy and accumulated usage stay unchanged. +Already admitted requests can finish. Revoked or expired keys cannot rotate. + +`allowed_models` contains exact, case-sensitive public model identifiers +(deployment aliases for registered models). `null` adds no key restriction; +`[]` denies every model. Workspace/provider policy still applies. Only an +organization administrator or the key's team administrator can edit this +policy after creation. Quota admission rechecks current key authorization and +model policy before reserving usage or calling a provider. + +## Quotas + +Team limits are optional nonnegative integers for daily requests, monthly +requests, and monthly tokens. `null` adds no team limit; zero denies admission. +Periods are UTC calendar days and months, beginning at 00:00 UTC. + +Every request must fit both its existing per-key tier limit and its team's +shared allowance. Existing tier limits remain per-key; this change does not +add an organization-wide aggregate quota. Requests reserve one request and +estimated input plus maximum output tokens. Settlement replaces reserved token +usage with actual usage; refund releases both key and team reservations. + +The existing PostgreSQL `quota_period_usage` table stores both scopes. One +usage-ledger event references all allocations, and conditional upserts fence +concurrent admissions. Redis is not a team quota ledger. Updating limits does +not reset consumed usage; lowering a limit below current usage denies further +admissions until usage is released or the next period starts. + +Management changes append audit intent in their own committed transaction. +Quota and key policy changes include non-secret change facts. Identity-provider +memberships are synchronized through the same tenant boundary; local grants +remain authoritative and IdP-managed grants must be changed at the provider. + +Run `uv run --locked python -m pytest -q ee/tests/tenants/test_teams.py` against +the disposable enterprise database to verify authorization, rotation, +membership synchronization and concurrent quota reservations. diff --git a/ee/openapi/enterprise.json b/ee/openapi/enterprise.json index 9bcfa9e..71b7e35 100644 --- a/ee/openapi/enterprise.json +++ b/ee/openapi/enterprise.json @@ -329,6 +329,21 @@ }, "ApiKeyInput": { "properties": { + "allowed_models": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "maxItems": 200, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Allowed Models" + }, "cost_center": { "anyOf": [ { @@ -356,6 +371,18 @@ } ], "title": "Team" + }, + "team_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Id" } }, "required": [ @@ -366,6 +393,21 @@ }, "ApiKeyPatch": { "properties": { + "allowed_models": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "maxItems": 200, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Allowed Models" + }, "cost_center": { "anyOf": [ { @@ -387,6 +429,18 @@ } ], "title": "Team" + }, + "team_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Id" } }, "title": "ApiKeyPatch", @@ -394,6 +448,20 @@ }, "ApiKeyView": { "properties": { + "allowed_models": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Allowed Models" + }, "cost_center": { "anyOf": [ { @@ -457,6 +525,18 @@ ], "title": "Team" }, + "team_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Id" + }, "tier": { "title": "Tier", "type": "string" @@ -471,7 +551,9 @@ "is_active", "expires_at", "cost_center", - "team" + "team", + "team_id", + "allowed_models" ], "title": "ApiKeyView", "type": "object" @@ -699,6 +781,45 @@ }, "AuditLogRead": { "properties": { + "actor": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Actor" + }, + "actor_type": { + "anyOf": [ + { + "enum": [ + "api_key", + "user_jwt", + "internal" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Actor Type" + }, + "api_key_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Api Key Id" + }, "completion_tokens": { "minimum": 0.0, "title": "Completion Tokens", @@ -765,6 +886,17 @@ "title": "Latency Ms", "type": "integer" }, + "lifecycle_status": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Lifecycle Status" + }, "model": { "anyOf": [ { @@ -798,6 +930,14 @@ "title": "Pii Entities", "type": "object" }, + "policy_verdicts": { + "items": { + "additionalProperties": true, + "type": "object" + }, + "title": "Policy Verdicts", + "type": "array" + }, "prev_hash": { "title": "Prev Hash", "type": "string" @@ -1192,10 +1332,22 @@ "title": "Completion Tokens", "type": "integer" }, + "cost_complete": { + "default": true, + "title": "Cost Complete", + "type": "boolean" + }, "cost_usd": { - "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$", - "title": "Cost Usd", - "type": "string" + "anyOf": [ + { + "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Cost Usd" }, "key": { "minLength": 1, @@ -1211,6 +1363,11 @@ "minimum": 0.0, "title": "Request Count", "type": "integer" + }, + "unpriced_requests": { + "default": 0, + "title": "Unpriced Requests", + "type": "integer" } }, "required": [ @@ -1284,6 +1441,11 @@ }, "BillingUsageView": { "properties": { + "cost_complete": { + "default": true, + "title": "Cost Complete", + "type": "boolean" + }, "daily_usage": { "items": { "$ref": "#/components/schemas/DailyUsageView" @@ -1295,8 +1457,20 @@ "$ref": "#/components/schemas/BillingPeriodView" }, "total_cost": { - "title": "Total Cost", - "type": "number" + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Total Cost" + }, + "unpriced_requests": { + "default": 0, + "title": "Unpriced Requests", + "type": "integer" } }, "required": [ @@ -2142,8 +2316,29 @@ "title": "Content", "type": "object" }, + "CountTokensRequest": { + "additionalProperties": { + "$ref": "#/components/schemas/JsonValue" + }, + "title": "CountTokensRequest", + "type": "object" + }, "CreatedApiKey": { "properties": { + "allowed_models": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Allowed Models" + }, "cost_center": { "anyOf": [ { @@ -2211,6 +2406,18 @@ ], "title": "Team" }, + "team_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Id" + }, "tier": { "title": "Tier", "type": "string" @@ -2226,6 +2433,8 @@ "expires_at", "cost_center", "team", + "team_id", + "allowed_models", "plaintext" ], "title": "CreatedApiKey", @@ -2281,7 +2490,8 @@ "enum": [ "owner", "admin", - "member" + "member", + "auditor" ], "title": "Role", "type": "string" @@ -2310,9 +2520,21 @@ "title": "Completion Tokens", "type": "integer" }, + "cost_complete": { + "default": true, + "title": "Cost Complete", + "type": "boolean" + }, "cost_usd": { - "title": "Cost Usd", - "type": "number" + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Cost Usd" }, "date": { "format": "date", @@ -2330,6 +2552,11 @@ "request_count": { "title": "Request Count", "type": "integer" + }, + "unpriced_requests": { + "default": 0, + "title": "Unpriced Requests", + "type": "integer" } }, "required": [ @@ -4548,6 +4775,59 @@ "title": "MediaResolution", "type": "string" }, + "MembershipInput": { + "properties": { + "role": { + "default": "member", + "enum": [ + "member", + "team_admin" + ], + "title": "Role", + "type": "string" + } + }, + "title": "MembershipInput", + "type": "object" + }, + "MembershipView": { + "properties": { + "role": { + "default": "member", + "enum": [ + "member", + "team_admin" + ], + "title": "Role", + "type": "string" + }, + "source": { + "enum": [ + "local", + "oidc" + ], + "title": "Source", + "type": "string" + }, + "team_id": { + "format": "uuid", + "title": "Team Id", + "type": "string" + }, + "user_id": { + "format": "uuid", + "title": "User Id", + "type": "string" + } + }, + "required": [ + "user_id", + "team_id", + "source" + ], + "title": "MembershipView", + "type": "object" + }, "MessagesRequest": { "additionalProperties": { "$ref": "#/components/schemas/JsonValue" @@ -4567,38 +4847,241 @@ "title": "Modality", "type": "string" }, - "ModelListView": { + "ModelDeploymentInput": { + "additionalProperties": false, "properties": { - "data": { - "items": { - "$ref": "#/components/schemas/ModelRecordView" - }, - "title": "Data", - "type": "array" - }, - "object": { - "title": "Object", + "alias": { + "maxLength": 200, + "minLength": 1, + "pattern": "^[A-Za-z0-9][A-Za-z0-9_.:-]*$", + "title": "Alias", "type": "string" - } - }, - "required": [ - "object", - "data" - ], - "title": "ModelListView", - "type": "object" - }, - "ModelRecordView": { - "properties": { - "created": { - "title": "Created", - "type": "integer" }, - "id": { - "title": "Id", + "base_url": { + "maxLength": 2048, + "minLength": 1, + "title": "Base Url", "type": "string" }, - "object": { + "declared_version": { + "maxLength": 200, + "minLength": 1, + "title": "Declared Version", + "type": "string" + }, + "deployment_kind": { + "enum": [ + "internal", + "external" + ], + "title": "Deployment Kind", + "type": "string" + }, + "enabled": { + "default": true, + "title": "Enabled", + "type": "boolean" + }, + "owner": { + "maxLength": 200, + "minLength": 1, + "title": "Owner", + "type": "string" + }, + "provider": { + "enum": [ + "openai", + "anthropic" + ], + "title": "Provider", + "type": "string" + }, + "provider_secret_id": { + "format": "uuid", + "title": "Provider Secret Id", + "type": "string" + }, + "timeout_seconds": { + "default": 60, + "maximum": 300.0, + "minimum": 1.0, + "title": "Timeout Seconds", + "type": "integer" + }, + "upstream_model": { + "maxLength": 200, + "minLength": 1, + "title": "Upstream Model", + "type": "string" + } + }, + "required": [ + "alias", + "provider", + "upstream_model", + "base_url", + "provider_secret_id", + "deployment_kind", + "declared_version", + "owner" + ], + "title": "ModelDeploymentInput", + "type": "object" + }, + "ModelDeploymentView": { + "additionalProperties": false, + "properties": { + "alias": { + "maxLength": 200, + "minLength": 1, + "pattern": "^[A-Za-z0-9][A-Za-z0-9_.:-]*$", + "title": "Alias", + "type": "string" + }, + "base_url": { + "maxLength": 2048, + "minLength": 1, + "title": "Base Url", + "type": "string" + }, + "created_at": { + "format": "date-time", + "title": "Created At", + "type": "string" + }, + "declared_version": { + "maxLength": 200, + "minLength": 1, + "title": "Declared Version", + "type": "string" + }, + "deployment_kind": { + "enum": [ + "internal", + "external" + ], + "title": "Deployment Kind", + "type": "string" + }, + "enabled": { + "default": true, + "title": "Enabled", + "type": "boolean" + }, + "health": { + "enum": [ + "unknown", + "healthy", + "unhealthy" + ], + "title": "Health", + "type": "string" + }, + "health_checked_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Health Checked At" + }, + "id": { + "format": "uuid", + "title": "Id", + "type": "string" + }, + "owner": { + "maxLength": 200, + "minLength": 1, + "title": "Owner", + "type": "string" + }, + "provider": { + "enum": [ + "openai", + "anthropic" + ], + "title": "Provider", + "type": "string" + }, + "provider_secret_id": { + "format": "uuid", + "title": "Provider Secret Id", + "type": "string" + }, + "timeout_seconds": { + "default": 60, + "maximum": 300.0, + "minimum": 1.0, + "title": "Timeout Seconds", + "type": "integer" + }, + "updated_at": { + "format": "date-time", + "title": "Updated At", + "type": "string" + }, + "upstream_model": { + "maxLength": 200, + "minLength": 1, + "title": "Upstream Model", + "type": "string" + } + }, + "required": [ + "alias", + "provider", + "upstream_model", + "base_url", + "provider_secret_id", + "deployment_kind", + "declared_version", + "owner", + "id", + "health", + "health_checked_at", + "created_at", + "updated_at" + ], + "title": "ModelDeploymentView", + "type": "object" + }, + "ModelListView": { + "properties": { + "data": { + "items": { + "$ref": "#/components/schemas/ModelRecordView" + }, + "title": "Data", + "type": "array" + }, + "object": { + "title": "Object", + "type": "string" + } + }, + "required": [ + "object", + "data" + ], + "title": "ModelListView", + "type": "object" + }, + "ModelRecordView": { + "properties": { + "created": { + "title": "Created", + "type": "integer" + }, + "id": { + "title": "Id", + "type": "string" + }, + "object": { "title": "Object", "type": "string" }, @@ -5398,6 +5881,10 @@ }, "OverviewSummaryView": { "properties": { + "cost_complete": { + "title": "Cost Complete", + "type": "boolean" + }, "p95_completed_latency_ms": { "anyOf": [ { @@ -5421,9 +5908,16 @@ "type": "integer" }, "settled_spend_usd": { - "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$", - "title": "Settled Spend Usd", - "type": "string" + "anyOf": [ + { + "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Settled Spend Usd" }, "status_counts": { "$ref": "#/components/schemas/OverviewStatusCountsView" @@ -5445,6 +5939,11 @@ } ], "title": "Technical Success Rate" + }, + "unpriced_requests": { + "minimum": 0.0, + "title": "Unpriced Requests", + "type": "integer" } }, "required": [ @@ -5454,6 +5953,8 @@ "technical_success_rate", "p95_completed_latency_ms", "settled_spend_usd", + "cost_complete", + "unpriced_requests", "status_counts" ], "title": "OverviewSummaryView", @@ -5461,26 +5962,44 @@ }, "OverviewTrendPointView": { "properties": { + "cost_complete": { + "title": "Cost Complete", + "type": "boolean" + }, "requests": { "minimum": 0.0, "title": "Requests", "type": "integer" }, "settled_spend_usd": { - "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$", - "title": "Settled Spend Usd", - "type": "string" + "anyOf": [ + { + "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Settled Spend Usd" }, "start": { "format": "date-time", "title": "Start", "type": "string" + }, + "unpriced_requests": { + "minimum": 0.0, + "title": "Unpriced Requests", + "type": "integer" } }, "required": [ "start", "requests", - "settled_spend_usd" + "settled_spend_usd", + "cost_complete", + "unpriced_requests" ], "title": "OverviewTrendPointView", "type": "object" @@ -6538,6 +7057,10 @@ "title": "Completion Tokens", "type": "integer" }, + "cost_complete": { + "title": "Cost Complete", + "type": "boolean" + }, "p95_completed_latency_ms": { "anyOf": [ { @@ -6595,6 +7118,11 @@ } ], "title": "Technical Success Rate" + }, + "unpriced_requests": { + "minimum": 0.0, + "title": "Unpriced Requests", + "type": "integer" } }, "required": [ @@ -6602,6 +7130,8 @@ "technical_success_rate", "p95_completed_latency_ms", "settled_spend_usd", + "cost_complete", + "unpriced_requests", "prompt_tokens", "completion_tokens", "pii_detected_requests", @@ -6630,16 +7160,43 @@ ], "title": "Cost Center" }, + "cost_complete": { + "title": "Cost Complete", + "type": "boolean" + }, "cost_usd": { - "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$", - "title": "Cost Usd", - "type": "string" + "anyOf": [ + { + "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Cost Usd" }, "created_at": { "format": "date-time", "title": "Created At", "type": "string" }, + "deployment_kind": { + "anyOf": [ + { + "enum": [ + "internal", + "external", + "unknown" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Deployment Kind" + }, "endpoint": { "anyOf": [ { @@ -6687,6 +7244,32 @@ ], "title": "Provider" }, + "provider_finish_reasons": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Provider Finish Reasons" + }, + "repeat_chain_length": { + "anyOf": [ + { + "minimum": 1.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Repeat Chain Length" + }, "request_id": { "title": "Request Id", "type": "string" @@ -6706,6 +7289,17 @@ "title": "Status", "type": "string" }, + "system_prompt_hash": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "System Prompt Hash" + }, "tags": { "items": { "type": "string" @@ -6724,13 +7318,25 @@ ], "title": "Team" }, - "usage_estimated": { - "title": "Usage Estimated", - "type": "boolean" - } - }, - "required": [ - "request_id", + "ttft_ms": { + "anyOf": [ + { + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Ttft Ms" + }, + "usage_estimated": { + "title": "Usage Estimated", + "type": "boolean" + } + }, + "required": [ + "request_id", "created_at", "endpoint", "model", @@ -6739,6 +7345,7 @@ "completion_tokens", "usage_estimated", "cost_usd", + "cost_complete", "latency_ms", "pii_detected", "cost_center", @@ -7758,49 +8365,6 @@ }, "SubscriptionView": { "properties": { - "cancel_at_period_end": { - "title": "Cancel At Period End", - "type": "boolean" - }, - "checkout_urls": { - "additionalProperties": { - "additionalProperties": { - "type": "string" - }, - "propertyNames": { - "enum": [ - "monthly", - "yearly" - ] - }, - "type": "object" - }, - "title": "Checkout Urls", - "type": "object" - }, - "current_period_end": { - "anyOf": [ - { - "format": "date-time", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Current Period End" - }, - "customer_portal_url": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Customer Portal Url" - }, "entitlements": { "additionalProperties": { "type": "boolean" @@ -7838,15 +8402,65 @@ "plan", "status", "source", - "current_period_end", - "cancel_at_period_end", - "entitlements", - "checkout_urls", - "customer_portal_url" + "entitlements" ], "title": "SubscriptionView", "type": "object" }, + "TeamInput": { + "properties": { + "daily_request_limit": { + "anyOf": [ + { + "maximum": 2000000000.0, + "minimum": 0.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Daily Request Limit" + }, + "monthly_request_limit": { + "anyOf": [ + { + "maximum": 2000000000.0, + "minimum": 0.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Monthly Request Limit" + }, + "monthly_token_limit": { + "anyOf": [ + { + "maximum": 2000000000.0, + "minimum": 0.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Monthly Token Limit" + }, + "name": { + "maxLength": 128, + "minLength": 1, + "title": "Name", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "TeamInput", + "type": "object" + }, "TeamInviteInput": { "properties": { "email": { @@ -7858,7 +8472,8 @@ "default": "member", "enum": [ "admin", - "member" + "member", + "auditor" ], "title": "Role", "type": "string" @@ -7920,7 +8535,8 @@ "enum": [ "owner", "admin", - "member" + "member", + "auditor" ], "title": "Role", "type": "string" @@ -7974,7 +8590,8 @@ "enum": [ "owner", "admin", - "member" + "member", + "auditor" ], "title": "Role", "type": "string" @@ -7997,7 +8614,8 @@ "enum": [ "owner", "admin", - "member" + "member", + "auditor" ], "title": "Role", "type": "string" @@ -8009,6 +8627,66 @@ "title": "TeamRolePatch", "type": "object" }, + "TeamView": { + "properties": { + "daily_request_limit": { + "anyOf": [ + { + "maximum": 2000000000.0, + "minimum": 0.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Daily Request Limit" + }, + "id": { + "format": "uuid", + "title": "Id", + "type": "string" + }, + "monthly_request_limit": { + "anyOf": [ + { + "maximum": 2000000000.0, + "minimum": 0.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Monthly Request Limit" + }, + "monthly_token_limit": { + "anyOf": [ + { + "maximum": 2000000000.0, + "minimum": 0.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Monthly Token Limit" + }, + "name": { + "maxLength": 128, + "minLength": 1, + "title": "Name", + "type": "string" + } + }, + "required": [ + "name", + "id" + ], + "title": "TeamView", + "type": "object" + }, "TextResponseFormat": { "additionalProperties": false, "description": "Configuration for text-specific output formatting.", @@ -8737,7 +9415,8 @@ "enum": [ "owner", "admin", - "member" + "member", + "auditor" ], "title": "Role", "type": "string" @@ -9296,6 +9975,108 @@ }, "openapi": "3.1.0", "paths": { + "/api/v1/auth/callback": { + "get": { + "operationId": "callback_api_v1_auth_callback_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Callback", + "tags": [ + "identity" + ] + } + }, + "/api/v1/auth/login": { + "get": { + "operationId": "login_api_v1_auth_login_get", + "parameters": [ + { + "in": "query", + "name": "next", + "required": false, + "schema": { + "default": "/dashboard", + "title": "Next", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Login", + "tags": [ + "identity" + ] + } + }, + "/api/v1/auth/logout": { + "post": { + "operationId": "logout_api_v1_auth_logout_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Logout", + "tags": [ + "identity" + ] + } + }, + "/api/v1/auth/session": { + "get": { + "operationId": "get_session_api_v1_auth_session_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "title": "Response Get Session Api V1 Auth Session Get", + "type": "object" + } + } + }, + "description": "Successful Response" + } + }, + "summary": "Get Session", + "tags": [ + "identity" + ] + } + }, "/api/v1/compliance/audit/anchor": { "post": { "operationId": "trigger_anchor_api_v1_compliance_audit_anchor_post", @@ -10942,19 +11723,41 @@ ] } }, - "/api/v1/management/auth/me": { - "get": { - "operationId": "current_profile_api_v1_management_auth_me_get", + "/api/v1/management/api-keys/{api_key_id}/rotate": { + "post": { + "operationId": "rotate_api_key_api_v1_management_api_keys__api_key_id__rotate_post", + "parameters": [ + { + "in": "path", + "name": "api_key_id", + "required": true, + "schema": { + "format": "uuid", + "title": "Api Key Id", + "type": "string" + } + } + ], "responses": { "200": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/UserView" + "$ref": "#/components/schemas/CreatedApiKey" } } }, "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" } }, "security": [ @@ -10962,11 +11765,37 @@ "HTTPBearer": [] } ], - "summary": "Current Profile", + "summary": "Rotate Api Key", "tags": [ "management" ] - }, + } + }, + "/api/v1/management/auth/me": { + "get": { + "operationId": "current_profile_api_v1_management_auth_me_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/UserView" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "Current Profile", + "tags": [ + "management" + ] + }, "put": { "operationId": "update_profile_api_v1_management_auth_me_put", "requestBody": { @@ -11544,6 +12373,186 @@ ] } }, + "/api/v1/management/model-deployments": { + "get": { + "operationId": "list_model_deployments_api_v1_management_model_deployments_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/ModelDeploymentView" + }, + "title": "Response List Model Deployments Api V1 Management Model Deployments Get", + "type": "array" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "List Model Deployments", + "tags": [ + "management" + ] + }, + "post": { + "operationId": "create_model_deployment_api_v1_management_model_deployments_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ModelDeploymentInput" + } + } + }, + "required": true + }, + "responses": { + "201": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ModelDeploymentView" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "Create Model Deployment", + "tags": [ + "management" + ] + } + }, + "/api/v1/management/model-deployments/{deployment_id}": { + "put": { + "operationId": "update_model_deployment_api_v1_management_model_deployments__deployment_id__put", + "parameters": [ + { + "in": "path", + "name": "deployment_id", + "required": true, + "schema": { + "format": "uuid", + "title": "Deployment Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ModelDeploymentInput" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ModelDeploymentView" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "Update Model Deployment", + "tags": [ + "management" + ] + } + }, + "/api/v1/management/model-deployments/{deployment_id}/health": { + "post": { + "operationId": "check_model_deployment_health_api_v1_management_model_deployments__deployment_id__health_post", + "parameters": [ + { + "in": "path", + "name": "deployment_id", + "required": true, + "schema": { + "format": "uuid", + "title": "Deployment Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ModelDeploymentView" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "Check Model Deployment Health", + "tags": [ + "management" + ] + } + }, "/api/v1/management/overview": { "get": { "operationId": "dashboard_overview_api_v1_management_overview_get", @@ -12645,15 +13654,19 @@ ] } }, - "/api/v1/management/tier-info": { + "/api/v1/management/teams": { "get": { - "operationId": "tier_info_api_v1_management_tier_info_get", + "operationId": "list_teams_api_v1_management_teams_get", "responses": { "200": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/TierView" + "items": { + "$ref": "#/components/schemas/TeamView" + }, + "title": "Response List Teams Api V1 Management Teams Get", + "type": "array" } } }, @@ -12665,32 +13678,29 @@ "HTTPBearer": [] } ], - "summary": "Tier Info", + "summary": "List Teams", "tags": [ "management" ] - } - }, - "/api/v1/shared-results/{token}": { - "get": { - "operationId": "view_shared_result_api_v1_shared_results__token__get", - "parameters": [ - { - "in": "path", - "name": "token", - "required": true, - "schema": { - "title": "Token", - "type": "string" + }, + "post": { + "operationId": "create_team_api_v1_management_teams_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TeamInput" + } } - } - ], + }, + "required": true + }, "responses": { - "200": { + "201": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/SharedResultView" + "$ref": "#/components/schemas/TeamView" } } }, @@ -12707,34 +13717,309 @@ "description": "Validation Error" } }, - "summary": "View Shared Result", + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "Create Team", "tags": [ - "shared-results" + "management" ] } }, - "/api/v1/webhooks/lemonsqueezy": { - "post": { - "operationId": "lemonsqueezy_webhook_api_v1_webhooks_lemonsqueezy_post", + "/api/v1/management/teams/{team_id}": { + "put": { + "operationId": "update_team_api_v1_management_teams__team_id__put", + "parameters": [ + { + "in": "path", + "name": "team_id", + "required": true, + "schema": { + "format": "uuid", + "title": "Team Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TeamInput" + } + } + }, + "required": true + }, "responses": { "200": { "content": { "application/json": { "schema": { - "additionalProperties": { - "type": "string" - }, - "title": "Response Lemonsqueezy Webhook Api V1 Webhooks Lemonsqueezy Post", - "type": "object" + "$ref": "#/components/schemas/TeamView" } } }, "description": "Successful Response" - } - }, - "summary": "Lemonsqueezy Webhook", - "tags": [ - "webhooks" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "Update Team", + "tags": [ + "management" + ] + } + }, + "/api/v1/management/teams/{team_id}/members": { + "get": { + "operationId": "list_memberships_api_v1_management_teams__team_id__members_get", + "parameters": [ + { + "in": "path", + "name": "team_id", + "required": true, + "schema": { + "format": "uuid", + "title": "Team Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/MembershipView" + }, + "title": "Response List Memberships Api V1 Management Teams Team Id Members Get", + "type": "array" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "List Memberships", + "tags": [ + "management" + ] + } + }, + "/api/v1/management/teams/{team_id}/members/{member_id}": { + "delete": { + "operationId": "remove_membership_api_v1_management_teams__team_id__members__member_id__delete", + "parameters": [ + { + "in": "path", + "name": "team_id", + "required": true, + "schema": { + "format": "uuid", + "title": "Team Id", + "type": "string" + } + }, + { + "in": "path", + "name": "member_id", + "required": true, + "schema": { + "format": "uuid", + "title": "Member Id", + "type": "string" + } + } + ], + "responses": { + "204": { + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "Remove Membership", + "tags": [ + "management" + ] + }, + "put": { + "operationId": "set_membership_api_v1_management_teams__team_id__members__member_id__put", + "parameters": [ + { + "in": "path", + "name": "team_id", + "required": true, + "schema": { + "format": "uuid", + "title": "Team Id", + "type": "string" + } + }, + { + "in": "path", + "name": "member_id", + "required": true, + "schema": { + "format": "uuid", + "title": "Member Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MembershipInput" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MembershipView" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "Set Membership", + "tags": [ + "management" + ] + } + }, + "/api/v1/management/tier-info": { + "get": { + "operationId": "tier_info_api_v1_management_tier_info_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TierView" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ], + "summary": "Tier Info", + "tags": [ + "management" + ] + } + }, + "/api/v1/shared-results/{token}": { + "get": { + "operationId": "view_shared_result_api_v1_shared_results__token__get", + "parameters": [ + { + "in": "path", + "name": "token", + "required": true, + "schema": { + "title": "Token", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SharedResultView" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "View Shared Result", + "tags": [ + "shared-results" ] } }, @@ -13144,6 +14429,193 @@ "summary": "Messages" } }, + "/v1/messages/count_tokens": { + "post": { + "operationId": "count_tokens_v1_messages_count_tokens_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CountTokensRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "400": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "401": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "402": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "403": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "404": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "408": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "409": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "413": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "429": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "500": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "502": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "503": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "504": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "529": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + } + }, + "security": [ + { + "ShimAPIKey": [] + }, + { + "HTTPBearer": [] + }, + { + "AnthropicAPIKey": [] + } + ], + "summary": "Count Tokens" + } + }, "/v1/models": { "get": { "operationId": "list_models_v1_models_get", diff --git a/ee/pyproject.toml b/ee/pyproject.toml index 9aeae51..e5ff6a5 100644 --- a/ee/pyproject.toml +++ b/ee/pyproject.toml @@ -23,6 +23,10 @@ dependencies = [ "reportlab>=5,<6", "cryptography>=48.0.1,<51", "google-cloud-secret-manager>=2.29,<3", + "authlib>=1.6,<2", + "itsdangerous>=2.2,<3", + "pyjwt[crypto]>=2.12,<3", + "joserfc>=1.6,<2", ] [project.optional-dependencies] diff --git a/ee/scripts/activate_plan.py b/ee/scripts/activate_plan.py index afb44e3..5ab8eb4 100644 --- a/ee/scripts/activate_plan.py +++ b/ee/scripts/activate_plan.py @@ -1,49 +1,49 @@ -"""Override an organization's plan after commercial approval or support.""" +"""Provision or change an organization's plan after commercial approval.""" import argparse import asyncio from uuid import UUID -from sqlalchemy import select - from shim_enterprise.core.database import AsyncSessionLocal -from shim_enterprise.tenants.models import Organization -from shim_enterprise.tenants.subscriptions import set_organization_tier +from shim_enterprise.tenants.plans import ( + activate_organization_plan, + create_organization_plan, +) -async def activate(organization_id: UUID, tier: str) -> None: +async def activate( + organization_id: UUID | None, + tier: str, + create_name: str | None = None, +) -> UUID: async with AsyncSessionLocal() as session: - organization = ( - await session.execute( - select(Organization) - .where(Organization.id == organization_id) - .with_for_update(of=Organization) - ) - ).scalar_one_or_none() - if organization is None: - raise SystemExit(f"Organization not found: {organization_id}") - await set_organization_tier( - session, - organization, - tier, - status="active" if tier != "free" else "free", - source="operator", - ) - organization.external_customer_id = None - organization.external_subscription_id = None - organization.billing_variant_id = None - organization.current_period_end = None - organization.cancel_at_period_end = False - organization.customer_portal_url = None + try: + if create_name is not None: + organization = await create_organization_plan( + session, create_name, tier + ) + elif organization_id is not None: + organization = await activate_organization_plan( + session, organization_id, tier + ) + else: + raise ValueError("An organization ID or --create-name is required") + except ValueError as exc: + raise SystemExit(str(exc)) from exc await session.commit() + return organization.id def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("organization_id", type=UUID) + parser = argparse.ArgumentParser(description=__doc__) + target = parser.add_mutually_exclusive_group(required=True) + target.add_argument("organization_id", type=UUID, nargs="?") + target.add_argument( + "--create-name", help="Create a new organization without an identity account" + ) parser.add_argument("tier", choices=("free", "managed", "agency", "enterprise")) args = parser.parse_args() - asyncio.run(activate(args.organization_id, args.tier)) + print(asyncio.run(activate(args.organization_id, args.tier, args.create_name))) if __name__ == "__main__": diff --git a/ee/scripts/offline_bundle.py b/ee/scripts/offline_bundle.py new file mode 100644 index 0000000..b930345 --- /dev/null +++ b/ee/scripts/offline_bundle.py @@ -0,0 +1,226 @@ +"""Create and verify a directory of signed offline installation artifacts.""" + +import argparse +import hashlib +import json +import re +import shutil +import subprocess +import tarfile +import tempfile +from pathlib import Path + +RESERVED = {"manifest.json", "manifest.sigstore.json"} +NAME = re.compile(r"[a-zA-Z0-9][a-zA-Z0-9._-]*") +DIGEST = re.compile(r"[a-zA-Z0-9][^\s@]*@sha256:[0-9a-f]{64}") + + +def run(*args: str) -> str: + return subprocess.check_output(args, text=True).strip() + + +def checksum(path: Path) -> str: + with path.open("rb") as stream: + return hashlib.file_digest(stream, "sha256").hexdigest() + + +def filename(value: str) -> str: + if not NAME.fullmatch(value) or value in RESERVED: + raise ValueError(f"Invalid artifact name: {value!r}") + return value + + +def archive_config_digest(path: Path, platform: str) -> str: + """Hash saved configuration bytes, including platform and root filesystem IDs.""" + with tarfile.open(path) as archive: + + def read(name: str) -> bytes: + member = archive.getmember(name) + if not member.isfile() or member.size > 16 * 1024 * 1024: + raise ValueError("Image metadata must be a bounded regular file") + stream = archive.extractfile(member) + assert stream is not None + return stream.read() + + manifest = json.loads(read("manifest.json")) + if len(manifest) != 1 or manifest[0].get("RepoTags"): + raise ValueError("Archive must contain exactly one untagged image") + config_bytes = read(manifest[0]["Config"]) + config = json.loads(config_bytes) + if f"{config['os']}/{config['architecture']}" != platform: + raise ValueError("Archive platform does not match bundle platform") + if config["rootfs"]["type"] != "layers" or not isinstance( + config["rootfs"]["diff_ids"], list + ): + raise ValueError("Invalid image root filesystem metadata") + return "sha256:" + hashlib.sha256(config_bytes).hexdigest() + + +def import_image(directory: Path, image: dict, platform: str) -> None: + archive = directory / image["archive"] + output = run("docker", "load", "--platform", platform, "--input", str(archive)) + loaded = re.findall( + r"^Loaded image ID: (sha256:[0-9a-f]{64})$", output, re.MULTILINE + ) + if len(loaded) != 1: + raise ValueError("Docker did not return exactly one loaded image identity") + # containerd may synthesize a new manifest ID when loading a classic archive. + # Re-export the returned object instead of assuming any portable daemon ID. + with tempfile.TemporaryDirectory(prefix="shim-import-") as temporary: + exported = Path(temporary) / "loaded.tar" + run( + "docker", + "save", + "--platform", + platform, + "--output", + str(exported), + loaded[0], + ) + if archive_config_digest(exported, platform) != image["config_digest"]: + raise ValueError( + "Loaded image configuration or layers do not match archive" + ) + print(f"Imported {image['name']} as {loaded[0]}") + + +def create(spec_path: Path, directory: Path, key: str) -> None: + spec = json.loads(spec_path.read_text()) + filename(spec["bundle_id"]) + if spec["platform"] not in {"linux/amd64", "linux/arm64"}: + raise ValueError("Supported platforms: linux/amd64, linux/arm64") + images = spec["images"] + if not images or len({image["name"] for image in images}) != len(images): + raise ValueError("Images must have unique names") + for image in images: + filename(image["name"]) + filename(image["version"]) + if not DIGEST.fullmatch(image["source"]): + raise ValueError("Every source image must use an immutable sha256 digest") + # A failed creation stays visibly incomplete; never overwrite an existing bundle. + directory.mkdir(parents=True, exist_ok=False) + for source in spec["files"]: + path = (spec_path.parent / source).resolve() + target = directory / filename(path.name) + if not path.is_file() or target.exists(): + raise ValueError(f"Missing or duplicate input: {source}") + shutil.copyfile(path, target) + for image in images: + archive = directory / filename(image["name"] + ".tar") + if archive.exists(): + raise ValueError(f"Duplicate artifact: {archive.name}") + run("docker", "pull", "--platform", spec["platform"], image["source"]) + run( + "docker", + "save", + "--platform", + spec["platform"], + "--output", + str(archive), + image["source"], + ) + image["config_digest"] = archive_config_digest(archive, spec["platform"]) + image["archive"] = archive.name + manifest = { + "schema": 2, + "bundle_id": spec["bundle_id"], + "platform": spec["platform"], + "images": images, + "files": {path.name: checksum(path) for path in sorted(directory.iterdir())}, + } + payload = directory / "manifest.json" + payload.write_text(json.dumps(manifest, indent=2) + "\n") + run( + "cosign", + "sign-blob", + "--yes", + "--key", + key, + "--use-signing-config=false", + "--tlog-upload=false", + "--bundle", + str(directory / "manifest.sigstore.json"), + str(payload), + ) + + +def verify(directory: Path, key: str, expected_id: str) -> dict: + # The key is provisioned separately by the operator, never trusted from the bundle. + if not Path(key).is_file(): + raise ValueError( + "Verification requires a local, independently trusted public key" + ) + paths = list(directory.iterdir()) + if any(path.is_symlink() or not path.is_file() for path in paths): + raise ValueError("Bundle must contain only regular files, without symlinks") + payload = directory / "manifest.json" + run( + "cosign", + "verify-blob", + "--key", + key, + "--insecure-ignore-tlog", + "--bundle", + str(directory / "manifest.sigstore.json"), + str(payload), + ) + manifest = json.loads(payload.read_text()) + if manifest["schema"] != 2 or manifest["bundle_id"] != expected_id: + raise ValueError("Unexpected bundle schema or release ID") + if manifest["platform"] not in {"linux/amd64", "linux/arm64"}: + raise ValueError("Unsupported platform") + files = manifest["files"] + if set(files) | RESERVED != {path.name for path in paths}: + raise ValueError("Bundle file inventory does not match manifest") + for name, digest in files.items(): + path = directory / filename(name) + if checksum(path) != digest: + raise ValueError(f"Checksum mismatch: {name}") + for image in manifest["images"]: + if image["archive"] not in files or not DIGEST.fullmatch(image["source"]): + raise ValueError("Invalid image metadata") + config_digest = archive_config_digest( + directory / image["archive"], manifest["platform"] + ) + if image["config_digest"] != config_digest: + raise ValueError("Image configuration digest does not match archive") + return manifest + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + build = commands.add_parser("create") + build.add_argument("--spec", type=Path, required=True) + for command in ( + build, + commands.add_parser("verify"), + commands.add_parser("import"), + ): + command.add_argument("--directory", type=Path, required=True) + command.add_argument("--key", required=True) + if command is not build: + command.add_argument("--expected-id", required=True) + args = parser.parse_args() + try: + if args.command == "create": + create(args.spec, args.directory, args.key) + else: + manifest = verify(args.directory, args.key, args.expected_id) + if args.command == "import": + for image in manifest["images"]: + import_image(args.directory, image, manifest["platform"]) + print(f"Verified {manifest['bundle_id']} ({manifest['platform']})") + except ( + OSError, + ValueError, + KeyError, + TypeError, + tarfile.TarError, + subprocess.CalledProcessError, + ) as error: + parser.exit(1, f"Bundle operation failed: {error}\n") + + +if __name__ == "__main__": + main() diff --git a/ee/src/shim_enterprise/ai_act/schemas.py b/ee/src/shim_enterprise/ai_act/schemas.py index 93ed13c..63c2f3a 100644 --- a/ee/src/shim_enterprise/ai_act/schemas.py +++ b/ee/src/shim_enterprise/ai_act/schemas.py @@ -7,7 +7,7 @@ from typing import Any, Literal from uuid import UUID -from pydantic import BaseModel, ConfigDict, Field +from pydantic import AliasPath, BaseModel, ConfigDict, Field class OrmReadModel(BaseModel): @@ -20,6 +20,15 @@ class AuditLogRead(OrmReadModel): created_at: datetime event_type: str request_id: str | None = None + api_key_id: UUID | None = None + actor: str | None = None + actor_type: Literal["api_key", "user_jwt", "internal"] | None = Field( + default=None, validation_alias=AliasPath("extra", "actor_type") + ) + lifecycle_status: str | None = Field( + default=None, validation_alias=AliasPath("extra", "lifecycle_status") + ) + policy_verdicts: list[dict[str, Any]] = Field(default_factory=list) model: str | None = None provider: str | None = None gateway_version: str | None = None diff --git a/ee/src/shim_enterprise/api/enterprise_deps.py b/ee/src/shim_enterprise/api/enterprise_deps.py index 28f96c2..777ec48 100644 --- a/ee/src/shim_enterprise/api/enterprise_deps.py +++ b/ee/src/shim_enterprise/api/enterprise_deps.py @@ -15,6 +15,8 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from shim_enterprise.core.database import AsyncSessionLocal, get_db +from shim_enterprise.core.config import settings +from shim_enterprise.tenants.oidc import current_oidc_user from shim.gateway.auth import authentication_error, select_gateway_credential from shim.gateway.contracts.ids import ApiKeyId, UserId from shim.gateway.contracts.principal import AuthenticatedPrincipal @@ -201,6 +203,17 @@ async def get_scan_principal( """Authenticate scan callers without accepting caller-supplied tenancy.""" token = select_gateway_credential(request.headers) + if settings.AUTH_MODE == "oidc" and ( + token is None + or not token.startswith(API_KEY_PREFIX) + and "x-shim-key" not in request.headers + ): + user = await current_oidc_user(request, session, token) + return AuthenticatedPrincipal( + actor_type="user_jwt", + user_id=UserId(user.id), + authenticated_at=datetime.now(timezone.utc), + ) if token is None: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, @@ -234,6 +247,7 @@ async def get_scan_principal( async def get_invite_user( + request: Request, bearer: HTTPAuthorizationCredentials | None = Security(bearer_scheme), session: AsyncSession = Depends(get_db), ) -> User: @@ -243,6 +257,14 @@ async def get_invite_user( headers={"WWW-Authenticate": "Bearer"}, ) + if settings.AUTH_MODE == "oidc": + if "/invites" in request.url.path: + raise HTTPException( + 403, "OIDC membership is managed by your identity provider" + ) + return await current_oidc_user( + request, session, bearer.credentials if bearer else None + ) if bearer is None: raise credentials_exception @@ -254,10 +276,22 @@ async def get_invite_user( async def get_current_user( + request: Request, bearer: HTTPAuthorizationCredentials | None = Security(bearer_scheme), session: AsyncSession = Depends(get_db), ) -> User: - user = await get_invite_user(bearer, session) + user = await get_invite_user(request, bearer, session) + if ( + user.role == "auditor" + and request.method not in {"GET", "HEAD", "OPTIONS"} + and (request.method, request.url.path) + not in { + ("POST", "/api/v1/compliance/audit/verify"), + ("POST", "/api/v1/compliance/reports/audit"), + ("POST", "/api/v1/compliance/reports/kvkk"), + } + ): + raise HTTPException(403, "Auditor access is read-only") if not user.is_active: logger.warning("Rejected deactivated user") raise HTTPException( diff --git a/ee/src/shim_enterprise/api/v1/management.py b/ee/src/shim_enterprise/api/v1/management.py index e766428..dd57b48 100644 --- a/ee/src/shim_enterprise/api/v1/management.py +++ b/ee/src/shim_enterprise/api/v1/management.py @@ -3,16 +3,18 @@ from __future__ import annotations from collections.abc import AsyncIterator +import asyncio import csv from datetime import date, datetime, timedelta, timezone from decimal import Decimal import hashlib import io +import json import logging import secrets from typing import Any, Literal, cast from urllib.parse import urlsplit -from uuid import UUID, uuid4 +from uuid import UUID import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response, status @@ -26,8 +28,9 @@ model_validator, ) from sqlalchemy import case, cast as sql_cast, func, or_, select, update -from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.dialects.postgresql import JSONB, insert from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.exc import IntegrityError from shim_enterprise.api.enterprise_deps import ( get_current_user, @@ -61,21 +64,29 @@ from shim_enterprise.observability.analytics_projection import RequestLog from shim_enterprise.observability.overview import OverviewReadModel from shim_enterprise.outbox.models import OutboxEvent -from shim_enterprise.outbox.publisher import OutboxWriter from shim_enterprise.secrets.migration import assign_secret_reference from shim_enterprise.secrets.store import get_secret_store +from shim_enterprise.tenants.audit import record_management_action as _audit from shim_enterprise.tenants.models import ( ApiKey, + ModelDeployment, OrganizationInvite, Organization, ProviderSecret, TierDefinition, + Team, + TeamMembership, User, ) +from shim_enterprise.tenants.deployments import ( + require_model_aliases, + validate_deployment_url, +) from shim_enterprise.tenants.service import create_api_key as create_tenant_api_key +from shim_enterprise.tenants.teams import member_team_ids, require_team +from shim_enterprise.tenants.service import rotate_api_key as rotate_tenant_api_key from shim_enterprise.tenants.service import ensure_privacy_defaults from shim_enterprise.tenants.service import move_user_from_bootstrap -from shim_enterprise.tenants.subscriptions import checkout_urls router = APIRouter() @@ -141,7 +152,7 @@ class UserView(BaseModel): email: EmailStr full_name: str | None organization_name: str - role: Literal["owner", "admin", "member"] + role: Literal["owner", "admin", "member", "auditor"] is_active: bool is_verified: bool created_at: datetime @@ -152,10 +163,53 @@ class UserPatch(BaseModel): organization_name: str | None = Field(default=None, min_length=1, max_length=200) +class TeamInput(BaseModel): + name: str = Field(min_length=1, max_length=128) + daily_request_limit: int | None = Field(default=None, ge=0, le=2_000_000_000) + monthly_request_limit: int | None = Field(default=None, ge=0, le=2_000_000_000) + monthly_token_limit: int | None = Field(default=None, ge=0, le=2_000_000_000) + + @field_validator("name") + @classmethod + def validate_name(cls, value: str) -> str: + if not value.strip(): + raise ValueError("Team name cannot be blank") + return value.strip() + + +class TeamView(TeamInput): + model_config = ConfigDict(from_attributes=True) + id: UUID + + +class MembershipInput(BaseModel): + role: Literal["member", "team_admin"] = "member" + + +class MembershipView(MembershipInput): + model_config = ConfigDict(from_attributes=True) + user_id: UUID + team_id: UUID + source: Literal["local", "oidc"] + + class ApiKeyInput(BaseModel): name: str = Field(min_length=1, max_length=50) cost_center: str | None = None team: str | None = None + team_id: UUID | None = None + allowed_models: list[str] | None = Field(default=None, max_length=200) + + @field_validator("allowed_models") + @classmethod + def validate_models(cls, value: list[str] | None) -> list[str] | None: + if value is not None and any( + not item or item != item.strip() or len(item) > 200 for item in value + ): + raise ValueError( + "Model identifiers must be nonblank and at most 200 characters" + ) + return list(dict.fromkeys(value)) if value is not None else None @field_validator("cost_center", "team") @classmethod @@ -171,6 +225,19 @@ def validate_attribution(cls, value: str | None) -> str | None: class ApiKeyPatch(BaseModel): cost_center: str | None = None team: str | None = None + team_id: UUID | None = None + allowed_models: list[str] | None = Field(default=None, max_length=200) + + @field_validator("allowed_models") + @classmethod + def validate_models(cls, value: list[str] | None) -> list[str] | None: + if value is not None and any( + not item or item != item.strip() or len(item) > 200 for item in value + ): + raise ValueError( + "Model identifiers must be nonblank and at most 200 characters" + ) + return list(dict.fromkeys(value)) if value is not None else None @field_validator("cost_center", "team") @classmethod @@ -195,6 +262,8 @@ class ApiKeyView(BaseModel): expires_at: datetime | None cost_center: str | None team: str | None + team_id: UUID | None + allowed_models: list[str] | None class CreatedApiKey(ApiKeyView): @@ -258,11 +327,7 @@ class SubscriptionView(BaseModel): plan: Literal["free", "managed", "agency", "enterprise"] status: str source: str | None - current_period_end: datetime | None - cancel_at_period_end: bool entitlements: dict[str, bool] - checkout_urls: dict[str, dict[Literal["monthly", "yearly"], str]] - customer_portal_url: str | None class TeamMemberView(BaseModel): @@ -271,14 +336,14 @@ class TeamMemberView(BaseModel): id: UUID email: EmailStr full_name: str | None - role: Literal["owner", "admin", "member"] + role: Literal["owner", "admin", "member", "auditor"] is_active: bool created_at: datetime class TeamInviteInput(BaseModel): email: EmailStr - role: Literal["admin", "member"] = "member" + role: Literal["admin", "member", "auditor"] = "member" class TeamInviteView(BaseModel): @@ -286,7 +351,7 @@ class TeamInviteView(BaseModel): id: UUID email: EmailStr - role: Literal["owner", "admin", "member"] + role: Literal["owner", "admin", "member", "auditor"] expires_at: datetime accepted_at: datetime | None revoked_at: datetime | None @@ -302,7 +367,7 @@ class AcceptTeamInvite(BaseModel): class TeamRolePatch(BaseModel): - role: Literal["owner", "admin", "member"] + role: Literal["owner", "admin", "member", "auditor"] class NotificationTargetInput(BaseModel): @@ -432,13 +497,17 @@ class DailyUsageView(BaseModel): request_count: int prompt_tokens: int completion_tokens: int - cost_usd: float + cost_usd: float | None + unpriced_requests: int = 0 + cost_complete: bool = True class BillingUsageView(BaseModel): period: BillingPeriodView daily_usage: list[DailyUsageView] - total_cost: float + total_cost: float | None + unpriced_requests: int = 0 + cost_complete: bool = True KNOWN_REQUEST_ACTIVITY_STATUSES = ( @@ -474,13 +543,19 @@ class RequestActivityView(BaseModel): prompt_tokens: int = Field(ge=0) completion_tokens: int = Field(ge=0) usage_estimated: bool - cost_usd: Decimal = Field(ge=0) + cost_usd: Decimal | None = Field(ge=0) + cost_complete: bool latency_ms: int = Field(ge=0) pii_detected: bool tags: list[str] = Field(default_factory=list) cost_center: str | None provider: str | None team: str | None + provider_finish_reasons: dict[str, str] | None = None + repeat_chain_length: int | None = Field(default=None, ge=1) + ttft_ms: float | None = Field(default=None, ge=0) + system_prompt_hash: str | None = None + deployment_kind: Literal["internal", "external", "unknown"] | None = None class RequestActivityStatusCountsView(BaseModel): @@ -500,6 +575,8 @@ class RequestActivitySummaryView(BaseModel): technical_success_rate: float | None = Field(ge=0, le=1) p95_completed_latency_ms: int | None = Field(ge=0) settled_spend_usd: Decimal = Field(ge=0) + cost_complete: bool + unpriced_requests: int = Field(ge=0) prompt_tokens: int = Field(ge=0) completion_tokens: int = Field(ge=0) pii_detected_requests: int = Field(ge=0) @@ -542,14 +619,18 @@ class OverviewSummaryView(BaseModel): policy_rejections: int = Field(ge=0) technical_success_rate: float | None = Field(ge=0, le=1) p95_completed_latency_ms: int | None = Field(ge=0) - settled_spend_usd: Decimal = Field(ge=0) + settled_spend_usd: Decimal | None = Field(ge=0) + cost_complete: bool + unpriced_requests: int = Field(ge=0) status_counts: OverviewStatusCountsView class OverviewTrendPointView(BaseModel): start: datetime requests: int = Field(ge=0) - settled_spend_usd: Decimal = Field(ge=0) + settled_spend_usd: Decimal | None = Field(ge=0) + cost_complete: bool + unpriced_requests: int = Field(ge=0) class OverviewExceptionView(BaseModel): @@ -589,11 +670,13 @@ class OverviewDashboardView(BaseModel): class BillingBreakdownRow(BaseModel): + unpriced_requests: int = 0 + cost_complete: bool = True key: str = Field(min_length=1) request_count: int = Field(ge=0) prompt_tokens: int = Field(ge=0) completion_tokens: int = Field(ge=0) - cost_usd: Decimal = Field(ge=0) + cost_usd: Decimal | None = Field(ge=0) class BillingBreakdownView(BaseModel): @@ -644,11 +727,6 @@ async def get_subscription( raise HTTPException( status_code=503, detail="Organization tier is not configured" ) - purchase_urls = ( - checkout_urls(organization.id, user.id) - if user.role == "owner" and organization.tier == "free" - else {} - ) return SubscriptionView( plan=cast( Literal["free", "managed", "agency", "enterprise"], @@ -656,13 +734,7 @@ async def get_subscription( ), status=organization.billing_status, source=organization.billing_source, - current_period_end=organization.current_period_end, - cancel_at_period_end=organization.cancel_at_period_end, entitlements={key: bool(value) for key, value in tier.features.items()}, - checkout_urls=purchase_urls, - customer_portal_url=( - organization.customer_portal_url if user.role == "owner" else None - ), ) @@ -884,8 +956,15 @@ async def update_team_member( member = await _owned_member(session, user, member_id) if member.role == "owner" and patch.role != "owner": await _protect_last_owner(session, _tenant_id(user)) + previous_role = member.role member.role = patch.role - await _audit(session, user, "tenant.team_role_updated", str(member.id)) + await _audit( + session, + user, + "tenant.team_role_updated", + str(member.id), + details={"before": previous_role, "after": patch.role}, + ) await session.commit() await session.refresh(member) return member @@ -916,6 +995,214 @@ async def remove_team_member( await session.commit() +@router.get("/teams", response_model=list[TeamView]) +async def list_teams( + user: User = Depends(get_current_user), + session: AsyncSession = Depends(get_db), +) -> list[Team]: + statement = select(Team).where(Team.organization_id == _tenant_id(user)) + if user.role not in {"owner", "admin", "auditor"}: + statement = statement.where(Team.id.in_(member_team_ids(user))) + return list((await session.scalars(statement.order_by(Team.name, Team.id))).all()) + + +@router.post("/teams", response_model=TeamView, status_code=201) +async def create_team( + payload: TeamInput, + user: User = Depends(get_org_admin), + session: AsyncSession = Depends(get_db), +) -> Team: + statement = ( + insert(Team) + .values(organization_id=_tenant_id(user), **payload.model_dump()) + .on_conflict_do_nothing() + .returning(Team) + ) + team = await session.scalar(statement) + if team is None: + raise HTTPException( + status_code=409, detail="A team with this name already exists" + ) + await _audit( + session, + user, + "tenant.team_created", + str(team.id), + details={"after": payload.model_dump(mode="json")}, + ) + await session.commit() + await session.refresh(team) + return team + + +@router.put("/teams/{team_id}", response_model=TeamView) +async def update_team( + team_id: UUID, + payload: TeamInput, + user: User = Depends(get_org_admin), + session: AsyncSession = Depends(get_db), +) -> Team: + await require_team(session, user, team_id, administer=True) + duplicate = await session.scalar( + select(Team.id).where( + Team.organization_id == user.organization_id, + Team.name == payload.name, + Team.id != team_id, + ) + ) + if duplicate is not None: + raise HTTPException( + status_code=409, detail="A team with this name already exists" + ) + team = await session.scalar( + select(Team) + .where(Team.organization_id == user.organization_id, Team.id == team_id) + .with_for_update() + ) + if team is None: + raise HTTPException(status_code=404, detail="Team not found") + before = TeamView.model_validate(team).model_dump(mode="json") + for field, value in payload.model_dump().items(): + setattr(team, field, value) + try: + await session.flush() + except IntegrityError as exc: + await session.rollback() + if getattr(exc.orig, "sqlstate", None) == "23505": + raise HTTPException( + status_code=409, detail="A team with this name already exists" + ) from exc + raise + await _audit( + session, + user, + "tenant.team_policy_updated", + str(team_id), + details={"before": before, "after": payload.model_dump(mode="json")}, + ) + await session.commit() + await session.refresh(team) + return team + + +@router.get("/teams/{team_id}/members", response_model=list[MembershipView]) +async def list_memberships( + team_id: UUID, + user: User = Depends(get_current_user), + session: AsyncSession = Depends(get_db), +) -> list[TeamMembership]: + await require_team(session, user, team_id) + return list( + ( + await session.scalars( + select(TeamMembership) + .where( + TeamMembership.organization_id == _tenant_id(user), + TeamMembership.team_id == team_id, + ) + .order_by(TeamMembership.user_id) + ) + ).all() + ) + + +@router.put("/teams/{team_id}/members/{member_id}", response_model=MembershipView) +async def set_membership( + team_id: UUID, + member_id: UUID, + payload: MembershipInput, + user: User = Depends(get_current_user), + session: AsyncSession = Depends(get_db), +) -> TeamMembership: + await require_team(session, user, team_id, administer=True) + member = await _owned_member(session, user, member_id) + existing = await session.scalar( + select(TeamMembership) + .where( + TeamMembership.organization_id == user.organization_id, + TeamMembership.team_id == team_id, + TeamMembership.user_id == member_id, + ) + .with_for_update() + ) + if existing is not None and existing.source == "oidc": + raise HTTPException( + status_code=409, detail="Manage this membership in the identity provider" + ) + if payload.role == "team_admin" or ( + existing is not None and existing.role == "team_admin" + ): + _require_role(user, "owner", "admin") + statement = insert(TeamMembership).values( + organization_id=user.organization_id, + team_id=team_id, + user_id=member.id, + role=payload.role, + source="local", + ) + membership = await session.scalar( + statement.on_conflict_do_update( + index_elements=[ + TeamMembership.organization_id, + TeamMembership.team_id, + TeamMembership.user_id, + ], + set_={"role": payload.role}, + where=TeamMembership.source == "local", + ).returning(TeamMembership) + ) + if membership is None: + raise HTTPException( + status_code=409, detail="Membership changed; reload and try again" + ) + await _audit( + session, + user, + "tenant.team_membership_updated", + str(team_id) + ":" + str(member_id), + ) + await session.commit() + await session.refresh(membership) + return membership + + +@router.delete( + "/teams/{team_id}/members/{member_id}", status_code=204, response_model=None +) +async def remove_membership( + team_id: UUID, + member_id: UUID, + user: User = Depends(get_current_user), + session: AsyncSession = Depends(get_db), +) -> None: + await require_team(session, user, team_id, administer=True) + membership = await session.scalar( + select(TeamMembership) + .where( + TeamMembership.organization_id == user.organization_id, + TeamMembership.team_id == team_id, + TeamMembership.user_id == member_id, + ) + .with_for_update() + ) + if membership is None: + raise HTTPException(status_code=404, detail="Team membership not found") + if membership.source == "oidc": + raise HTTPException( + status_code=409, detail="Manage this membership in the identity provider" + ) + if membership.role == "team_admin": + _require_role(user, "owner", "admin") + await session.delete(membership) + await _audit( + session, + user, + "tenant.team_membership_removed", + str(team_id) + ":" + str(member_id), + ) + await session.commit() + + @router.get("/api-keys", response_model=list[ApiKeyView]) async def list_api_keys( user: User = Depends(get_current_user), @@ -926,8 +1213,13 @@ async def list_api_keys( ApiKey.organization_id == tenant_id, ApiKey.is_active.is_(True), ) - if user.role == "member": - statement = statement.where(ApiKey.user_id == user.id) + if user.role not in {"owner", "admin", "auditor"}: + statement = statement.where( + or_( + ApiKey.user_id == user.id, + ApiKey.team_id.in_(member_team_ids(user, administer=True)), + ) + ) now = datetime.now(timezone.utc) return [ item @@ -943,14 +1235,22 @@ async def create_api_key( session: AsyncSession = Depends(get_db), ) -> CreatedApiKey: _tenant_id(user) + _require_role(user, "owner", "admin", "member") if not user.is_verified: raise HTTPException(status_code=403, detail="Verified email required") + if payload.team_id is not None: + await require_team(session, user, payload.team_id) + elif user.role == "member" and await session.scalar(member_team_ids(user).limit(1)): + raise HTTPException(status_code=403, detail="Choose a team for this API key") + await require_model_aliases(session, _tenant_id(user), payload.allowed_models) plaintext, api_key = await create_tenant_api_key( session, user_id=user.id, name=payload.name, cost_center=payload.cost_center, team=payload.team, + team_id=payload.team_id, + allowed_models=payload.allowed_models, ) await _audit(session, user, "tenant.api_key_created", str(api_key.id)) await session.commit() @@ -961,6 +1261,35 @@ async def create_api_key( ) +@router.post("/api-keys/{api_key_id}/rotate", response_model=CreatedApiKey) +async def rotate_api_key( + api_key_id: UUID, + user: User = Depends(get_current_user), + session: AsyncSession = Depends(get_db), +) -> CreatedApiKey: + api_key = await _owned_api_key(session, user, api_key_id) + if not api_key.is_active or ( + api_key.expires_at is not None + and _aware(api_key.expires_at) <= datetime.now(timezone.utc) + ): + raise HTTPException( + status_code=409, detail="Only active API keys can be rotated" + ) + plaintext = rotate_tenant_api_key(api_key) + await _audit( + session, + user, + "tenant.api_key_rotated", + str(api_key.id), + details={"rotation_policy": "immediate"}, + ) + await session.commit() + await session.refresh(api_key) + return CreatedApiKey( + **ApiKeyView.model_validate(api_key).model_dump(), plaintext=plaintext + ) + + @router.patch("/api-keys/{api_key_id}", response_model=ApiKeyView) async def update_api_key( api_key_id: UUID, @@ -969,9 +1298,26 @@ async def update_api_key( session: AsyncSession = Depends(get_db), ) -> ApiKey: api_key = await _owned_api_key(session, user, api_key_id) + if {"team_id", "allowed_models"} & patch.model_fields_set: + if api_key.team_id is None: + _require_role(user, "owner", "admin") + else: + await require_team(session, user, api_key.team_id, administer=True) + if "team_id" in patch.model_fields_set and patch.team_id != api_key.team_id: + _require_role(user, "owner", "admin") + if patch.team_id is not None: + await require_team(session, user, patch.team_id, administer=True) + if "allowed_models" in patch.model_fields_set: + await require_model_aliases(session, _tenant_id(user), patch.allowed_models) for field, value in patch.model_dump(exclude_unset=True).items(): setattr(api_key, field, value) - await _audit(session, user, "tenant.api_key_updated", str(api_key.id)) + await _audit( + session, + user, + "tenant.api_key_updated", + str(api_key.id), + details={"changes": patch.model_dump(mode="json", exclude_unset=True)}, + ) await session.commit() await session.refresh(api_key) return api_key @@ -1553,13 +1899,24 @@ async def list_requests( prompt_tokens=row.prompt_tokens, completion_tokens=row.completion_tokens, usage_estimated=_request_usage_estimated(row.details), - cost_usd=Decimal(str(cost_usd)), + cost_usd=Decimal(str(cost_usd)) if cost_usd is not None else None, + cost_complete=cost_usd is not None, latency_ms=row.latency_ms, pii_detected=row.pii_detected, tags=list(row.tags or []), cost_center=row.cost_center, provider=_request_provider(row), team=row.team, + **{ + field: (row.details or {}).get(field) + for field in ( + "provider_finish_reasons", + "repeat_chain_length", + "ttft_ms", + "system_prompt_hash", + "deployment_kind", + ) + }, ) for row, cost_usd in rows ], @@ -1656,6 +2013,12 @@ async def content() -> AsyncIterator[bytes]: "tags", "cost_center", "team", + "provider_finish_reasons", + "repeat_chain_length", + "ttft_ms", + "system_prompt_hash", + "deployment_kind", + "cost_complete", ) ) yield output.getvalue().encode("utf-8-sig") @@ -1664,6 +2027,7 @@ async def content() -> AsyncIterator[bytes]: ) try: async for row, cost_usd in result: + details = row.details or {} output.seek(0) output.truncate(0) writer.writerow( @@ -1677,12 +2041,24 @@ async def content() -> AsyncIterator[bytes]: _request_activity_status(row.details), row.prompt_tokens, row.completion_tokens, - Decimal(str(cost_usd)), + Decimal(str(cost_usd)) if cost_usd is not None else None, row.latency_ms, row.pii_detected, ",".join(row.tags or []), row.cost_center, row.team, + ( + json.dumps( + details["provider_finish_reasons"], sort_keys=True + ) + if details.get("provider_finish_reasons") is not None + else None + ), + details.get("repeat_chain_length"), + details.get("ttft_ms"), + details.get("system_prompt_hash"), + details.get("deployment_kind"), + cost_usd is not None, ) ) yield output.getvalue().encode("utf-8") @@ -1720,10 +2096,15 @@ async def billing_usage( ), ) rows = [record.as_public_record() for record in records] + unpriced_requests = sum(record.unpriced_requests for record in records) return BillingUsageView( period=BillingPeriodView(start=start, end=end), daily_usage=[DailyUsageView.model_validate(row) for row in rows], - total_cost=sum(float(record.cost_usd) for record in records), + total_cost=None + if unpriced_requests + else sum(float(record.cost_usd) for record in records), + unpriced_requests=unpriced_requests, + cost_complete=not unpriced_requests, ) @@ -1896,12 +2277,30 @@ def _request_settled_spend(tenant_id: UUID): UsageLedger.organization_id == tenant_id, UsageLedger.request_id == RequestLog.request_id, UsageLedger.event_type == "spend_settlement", + UsageLedger.event_metadata["pricing"]["pricing_resolution"] + .as_string() + .is_distinct_from("unknown"), ) .correlate(RequestLog) .scalar_subquery() ) +def _request_unpriced_spend(tenant_id: UUID): + return ( + select(UsageLedger.id) + .where( + UsageLedger.organization_id == tenant_id, + UsageLedger.request_id == RequestLog.request_id, + UsageLedger.event_type == "spend_settlement", + UsageLedger.event_metadata["pricing"]["pricing_resolution"].as_string() + == "unknown", + ) + .correlate(RequestLog) + .exists() + ) + + def _request_summary_statement(tenant_id: UUID, filters: list[Any]): lifecycle_status = _request_lifecycle_status_expression() usage_estimated = RequestLog.details["usage_estimated"].as_boolean().is_(True) @@ -1921,6 +2320,9 @@ def _request_summary_statement(tenant_id: UUID, filters: list[Any]): return ( select( func.count(RequestLog.id).label("requests"), + func.count(RequestLog.id) + .filter(_request_unpriced_spend(tenant_id)) + .label("unpriced_requests"), *( func.count(RequestLog.id) .filter(lifecycle_status == status_name) @@ -1984,6 +2386,8 @@ def _request_activity_summary(row: Any) -> RequestActivitySummaryView: ), p95_completed_latency_ms=round(float(p95)) if p95 is not None else None, settled_spend_usd=Decimal(str(row.settled_spend_usd or 0)), + cost_complete=not row.unpriced_requests, + unpriced_requests=int(row.unpriced_requests or 0), prompt_tokens=int(row.prompt_tokens or 0), completion_tokens=int(row.completion_tokens or 0), pii_detected_requests=int(row.pii_detected_requests or 0), @@ -1999,7 +2403,10 @@ def _request_rows_statement(tenant_id: UUID, filters: list[Any]): return ( select( RequestLog, - func.coalesce(spend, Decimal("0")).label("cost_usd"), + case( + (_request_unpriced_spend(tenant_id), None), + else_=func.coalesce(spend, Decimal("0")), + ).label("cost_usd"), ) .where(*filters) .order_by(RequestLog.timestamp.desc(), RequestLog.id.desc()) @@ -2024,7 +2431,15 @@ def _billing_breakdown_csv(records: list[Any]) -> bytes: output = io.StringIO() writer = csv.writer(output) writer.writerow( - ("key", "request_count", "prompt_tokens", "completion_tokens", "cost_usd") + ( + "key", + "request_count", + "prompt_tokens", + "completion_tokens", + "cost_usd", + "unpriced_requests", + "cost_complete", + ) ) for record in records: writer.writerow( @@ -2034,7 +2449,9 @@ def _billing_breakdown_csv(records: list[Any]) -> bytes: record.request_count, record.prompt_tokens, record.completion_tokens, - record.cost_usd, + None if record.unpriced_requests else record.cost_usd, + record.unpriced_requests, + record.unpriced_requests == 0, ) ) return output.getvalue().encode("utf-8-sig") @@ -2086,7 +2503,9 @@ def _billing_breakdown_pdf( record.key, str(record.request_count), str(record.prompt_tokens + record.completion_tokens), - str(record.cost_usd), + str(record.cost_usd) + if record.unpriced_requests == 0 + else f"Unknown ({record.unpriced_requests} unpriced)", ] for record in records ], @@ -2106,7 +2525,7 @@ async def _user_view(session: AsyncSession, user: User) -> UserView: email=user.email, full_name=user.full_name, organization_name=tenant.name, - role=cast(Literal["owner", "admin", "member"], user.role), + role=cast(Literal["owner", "admin", "member", "auditor"], user.role), is_active=user.is_active, is_verified=user.is_verified, created_at=user.created_at, @@ -2118,12 +2537,30 @@ async def _owned_api_key( user: User, api_key_id: UUID, ) -> ApiKey: - statement = select(ApiKey).where( - ApiKey.id == api_key_id, - ApiKey.organization_id == _tenant_id(user), + _require_role(user, "owner", "admin", "member") + await session.scalar( + select(Organization.id) + .where(Organization.id == _tenant_id(user)) + .with_for_update() + ) + statement = ( + select(ApiKey) + .where( + ApiKey.id == api_key_id, + ApiKey.organization_id == _tenant_id(user), + ) + .with_for_update() ) - if user.role == "member": - statement = statement.where(ApiKey.user_id == user.id) + if user.role not in {"owner", "admin"}: + statement = statement.where( + or_( + (ApiKey.user_id == user.id) + & ( + ApiKey.team_id.is_(None) | ApiKey.team_id.in_(member_team_ids(user)) + ), + ApiKey.team_id.in_(member_team_ids(user, administer=True)), + ) + ) row = (await session.execute(statement)).scalar_one_or_none() if row is None: raise HTTPException(status_code=404, detail="API key not found") @@ -2230,37 +2667,6 @@ async def _owned_budget( return row -async def _audit( - session: AsyncSession, - user: User, - action: str, - subject_id: str, -) -> None: - tenant_id = _tenant_id(user) - event_id = f"management:{uuid4()}" - now = datetime.now(timezone.utc) - await OutboxWriter().append( - session, - organization_id=TenantId(tenant_id), - values={ - "event_type": "audit.chain_append_requested", - "aggregate_type": "management", - "aggregate_id": event_id, - "idempotency_key": f"{event_id}:audit", - "payload": { - "organization_id": str(tenant_id), - "request_id": event_id, - "event_type": "management_action", - "actor": str(user.id), - "endpoint": action, - "extra": {"subject_id": subject_id}, - }, - "status": "pending", - "next_attempt_at": now, - }, - ) - - async def _validate_targets(targets: list[NotificationTargetInput]) -> None: for target in targets: try: @@ -2401,3 +2807,212 @@ def _validate_sync_window(start: datetime, end: datetime) -> None: status_code=422, detail="synchronous operations are limited to 31 days", ) + + +class ModelDeploymentInput(BaseModel): + model_config = ConfigDict(extra="forbid") + alias: str = Field( + min_length=1, max_length=200, pattern=r"^[A-Za-z0-9][A-Za-z0-9_.:-]*$" + ) + provider: Literal["openai", "anthropic"] + upstream_model: str = Field(min_length=1, max_length=200) + base_url: str = Field(min_length=1, max_length=2048) + provider_secret_id: UUID + timeout_seconds: int = Field(default=60, ge=1, le=300) + deployment_kind: Literal["internal", "external"] + declared_version: str = Field(min_length=1, max_length=200) + owner: str = Field(min_length=1, max_length=200) + enabled: bool = True + + @field_validator("upstream_model", "declared_version", "owner") + @classmethod + def nonblank(cls, value: str) -> str: + if not value.strip() or value != value.strip(): + raise ValueError("Value must be nonblank without surrounding whitespace") + return value + + @field_validator("base_url") + @classmethod + def approved_destination(cls, value: str) -> str: + return validate_deployment_url(value) + + +class ModelDeploymentView(ModelDeploymentInput): + model_config = ConfigDict(from_attributes=True) + + @field_validator("base_url") + @classmethod + def approved_destination(cls, value: str) -> str: + # Operators must still be able to inspect a now-disallowed deployment. + return value + + id: UUID + health: Literal["unknown", "healthy", "unhealthy"] + health_checked_at: datetime | None + created_at: datetime + updated_at: datetime + + +@router.get("/model-deployments", response_model=list[ModelDeploymentView]) +async def list_model_deployments( + user: User = Depends(get_org_admin), + session: AsyncSession = Depends(get_db), +): + return ( + ( + await session.execute( + select(ModelDeployment) + .where( + ModelDeployment.organization_id == _tenant_id(user), + ) + .order_by(ModelDeployment.alias) + ) + ) + .scalars() + .all() + ) + + +@router.post("/model-deployments", response_model=ModelDeploymentView, status_code=201) +async def create_model_deployment( + payload: ModelDeploymentInput, + user: User = Depends(get_org_admin), + session: AsyncSession = Depends(get_db), +): + secret = await _owned_provider_secret(session, user, payload.provider_secret_id) + if secret.provider != payload.provider: + raise HTTPException(422, detail="Credential provider does not match deployment") + row = ModelDeployment(organization_id=_tenant_id(user), **payload.model_dump()) + session.add(row) + try: + await session.flush() + await _audit( + session, + user, + "tenant.model_deployment_created", + str(row.id), + details={"configuration": payload.model_dump(mode="json")}, + ) + await session.commit() + except IntegrityError: + await session.rollback() + raise HTTPException( + 409, detail="Model alias already exists or credential is unavailable" + ) from None + await session.refresh(row) + return row + + +@router.put("/model-deployments/{deployment_id}", response_model=ModelDeploymentView) +async def update_model_deployment( + deployment_id: UUID, + payload: ModelDeploymentInput, + user: User = Depends(get_org_admin), + session: AsyncSession = Depends(get_db), +): + row = await _owned_model_deployment(session, user, deployment_id) + secret = await _owned_provider_secret(session, user, payload.provider_secret_id) + if secret.provider != payload.provider: + raise HTTPException(422, detail="Credential provider does not match deployment") + for field, value in payload.model_dump().items(): + setattr(row, field, value) + row.health, row.health_checked_at = "unknown", None + try: + await _audit( + session, + user, + "tenant.model_deployment_updated", + str(row.id), + details={"configuration": payload.model_dump(mode="json")}, + ) + await session.commit() + except IntegrityError: + await session.rollback() + raise HTTPException( + 409, detail="Model alias already exists or credential is unavailable" + ) from None + await session.refresh(row) + return row + + +@router.post( + "/model-deployments/{deployment_id}/health", response_model=ModelDeploymentView +) +async def check_model_deployment_health( + deployment_id: UUID, + request: Request, + user: User = Depends(get_org_admin), + session: AsyncSession = Depends(get_db), +): + row = await _owned_model_deployment(session, user, deployment_id) + secret = await _owned_provider_secret(session, user, row.provider_secret_id) + try: + base_url = validate_deployment_url(row.base_url) + except ValueError: + raise HTTPException( + 503, detail="Deployment origin is no longer approved" + ) from None + checked_version = row.updated_at + secret_ref, provider, tenant_id = secret.secret_ref, row.provider, _tenant_id(user) + # Finish the read transaction before secret-store or provider I/O. + await session.commit() + healthy = False + try: + async with asyncio.timeout(5): + credential = await get_secret_store().get_secret( + TenantId(tenant_id), + SecretRef(secret_ref), + expected_purpose=f"provider:{provider}:api-key", + ) + headers = ( + {"Authorization": f"Bearer {credential}"} + if provider == "openai" + else {"x-api-key": credential, "anthropic-version": "2023-06-01"} + ) + path = "/models" if provider == "openai" else "/v1/models" + # Stream headers only: an unhealthy server cannot force an unbounded body read. + async with request.app.state.http_client.stream( + "GET", + base_url + path, + headers=headers, + timeout=5, + follow_redirects=False, + ) as response: + healthy = response.status_code == 200 + except (httpx.HTTPError, ValueError, TimeoutError): + healthy = False + row = await _owned_model_deployment(session, user, deployment_id) + if row.updated_at != checked_version: + raise HTTPException( + 409, detail="Deployment changed during health check; check again" + ) + row.health = "healthy" if healthy else "unhealthy" + row.health_checked_at = datetime.now(timezone.utc) + await _audit( + session, + user, + "tenant.model_deployment_health_checked", + str(row.id), + details={"health": row.health, "declared_version": row.declared_version}, + ) + await session.commit() + await session.refresh(row) + return row + + +async def _owned_model_deployment( + session: AsyncSession, user: User, deployment_id: UUID +) -> ModelDeployment: + row = ( + await session.execute( + select(ModelDeployment) + .where( + ModelDeployment.id == deployment_id, + ModelDeployment.organization_id == _tenant_id(user), + ) + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() + if row is None: + raise HTTPException(404, detail="Model deployment not found") + return row diff --git a/ee/src/shim_enterprise/api/v1/router.py b/ee/src/shim_enterprise/api/v1/router.py index 38bc3b2..8e0a6f1 100644 --- a/ee/src/shim_enterprise/api/v1/router.py +++ b/ee/src/shim_enterprise/api/v1/router.py @@ -6,10 +6,12 @@ from shim.api.v1.messages import router as messages_router from shim.api.v1.responses import router as responses_router from shim_enterprise.ai_act.api import router as ai_act_router -from shim_enterprise.api.v1 import management, scan, subscriptions +from shim_enterprise.api.v1 import management, scan from shim_enterprise.compliance.api import router as compliance_router from shim_enterprise.shared_results.api import authenticated_router, public_router +from shim_enterprise.tenants.oidc import router as identity_router + gateway_router = APIRouter() gateway_router.include_router(chat_router) gateway_router.include_router(responses_router) @@ -18,7 +20,7 @@ gateway_router.include_router(authenticated_router) management_router = APIRouter() -management_router.include_router(subscriptions.router) +management_router.include_router(identity_router) management_router.include_router( management.router, prefix="/management", diff --git a/ee/src/shim_enterprise/api/v1/subscriptions.py b/ee/src/shim_enterprise/api/v1/subscriptions.py deleted file mode 100644 index bc2b89c..0000000 --- a/ee/src/shim_enterprise/api/v1/subscriptions.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Public payment-provider webhooks.""" - -from fastapi import APIRouter, Depends, HTTPException, Request -from sqlalchemy.ext.asyncio import AsyncSession - -from shim_enterprise.core.config import settings -from shim_enterprise.core.database import get_db -from shim_enterprise.tenants.subscriptions import ( - process_lemonsqueezy_webhook, - verify_lemonsqueezy_signature, -) - -router = APIRouter(tags=["webhooks"]) - - -@router.post("/webhooks/lemonsqueezy") -async def lemonsqueezy_webhook( - request: Request, - session: AsyncSession = Depends(get_db), -) -> dict[str, str]: - secret = settings.LEMON_SQUEEZY_SIGNING_SECRET - if not secret: - raise HTTPException(status_code=503, detail="Billing webhook is not configured") - raw_body = await request.body() - if not verify_lemonsqueezy_signature( - raw_body, - request.headers.get("x-signature", ""), - secret, - ): - raise HTTPException(status_code=403, detail="Invalid webhook signature") - try: - result = await process_lemonsqueezy_webhook(session, raw_body) - except ValueError as exc: - await session.rollback() - raise HTTPException(status_code=400, detail=str(exc)) from exc - except LookupError as exc: - await session.rollback() - raise HTTPException(status_code=404, detail=str(exc)) from exc - except Exception: - await session.rollback() - raise - return {"status": result} diff --git a/ee/src/shim_enterprise/application.py b/ee/src/shim_enterprise/application.py index a6812fc..0e50253 100644 --- a/ee/src/shim_enterprise/application.py +++ b/ee/src/shim_enterprise/application.py @@ -7,6 +7,8 @@ import logging import httpx +import ssl +from hashlib import sha256 from fastapi import FastAPI, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import Response @@ -45,6 +47,8 @@ get_secret_store, ) from shim_enterprise.services.gateway.enterprise import EnterpriseGatewayService +from shim_enterprise.tenants.oidc import install_oidc +from shim_enterprise.tenants.deployments import DeploymentResolver from shim_enterprise.tenants.policy import ( TenantPolicyService, TenantRequestPolicyResolver, @@ -86,7 +90,9 @@ def create_enterprise_app() -> FastAPI: RequestValidationError, gateway_exception_handler, ) + install_oidc(application) application.state.cache = cache + application.state.model_catalog = DeploymentResolver(AsyncSessionLocal).catalog application.state.gateway_authenticator = DatabaseGatewayAuthenticator( AsyncSessionLocal ) @@ -124,7 +130,11 @@ async def _lifespan(application: FastAPI) -> AsyncIterator[None]: write=settings.OPENAI_WRITE_TIMEOUT_SECONDS, pool=settings.OPENAI_POOL_TIMEOUT_SECONDS, ) - http_client = httpx.AsyncClient(timeout=timeout, follow_redirects=False) + http_client = httpx.AsyncClient( + timeout=timeout, + follow_redirects=False, + verify=ssl.create_default_context(cafile=settings.MODEL_DEPLOYMENT_CA_BUNDLE), + ) application.state.http_client = http_client pii_scrubber = PIIScrubberService() application.state.gateway_service = EnterpriseGatewayService( @@ -165,6 +175,9 @@ def _create_gateway_kernel( "openai", secret_store, AsyncSessionLocal ), circuit=RedisCircuitBreaker("openai", cache=cache), + circuit_for_target=lambda url: RedisCircuitBreaker( + "openai-" + sha256(url.encode()).hexdigest()[:40], cache=cache + ), settings=settings, chain_store=chain_store, **dependencies, @@ -174,6 +187,9 @@ def _create_gateway_kernel( "anthropic", secret_store, AsyncSessionLocal ), circuit=RedisCircuitBreaker("anthropic", cache=cache), + circuit_for_target=lambda url: RedisCircuitBreaker( + "anthropic-" + sha256(url.encode()).hexdigest()[:40], cache=cache + ), settings=settings, **dependencies, ), @@ -188,6 +204,7 @@ def _create_gateway_kernel( }, chain_store=chain_store, policy_resolver=policy_resolver, + prepare_inference=DeploymentResolver(AsyncSessionLocal).resolve, rate_limiter=BurstRateLimiter(cache), loop_detector=LoopDetectionService(cache), loop_repeat_limit=settings.LOOP_REPEAT_LIMIT, diff --git a/ee/src/shim_enterprise/billing/ledger.py b/ee/src/shim_enterprise/billing/ledger.py index 0566263..7a0a210 100644 --- a/ee/src/shim_enterprise/billing/ledger.py +++ b/ee/src/shim_enterprise/billing/ledger.py @@ -23,7 +23,10 @@ SpendPeriodUsage, UsageLedger, ) -from shim_enterprise.gateway.pipeline.audit_intent import AuditIntentRepository +from shim_enterprise.gateway.pipeline.audit_intent import ( + AuditIntentPersistenceError, + AuditIntentRepository, +) from shim_enterprise.observability.lifecycle import RequestLifecycleRepository from shim_enterprise.outbox.publisher import OutboxWriter from shim_enterprise.gateway.pipeline.outbox import ( @@ -154,6 +157,8 @@ class QuotaPolicySnapshot: daily_request_limit: int | None monthly_request_limit: int | None monthly_token_limit: int | None + team_id: UUID | None = None + team_policy: QuotaPolicySnapshot | None = None def __post_init__(self) -> None: limits = ( @@ -193,6 +198,11 @@ class QuotaReservationCommand: tags: tuple[str, ...] = () team: str | None = None stream: bool = False + repeat_chain_length: int | None = None + system_prompt_hash: str | None = None + deployment_kind: Literal["internal", "external", "unknown"] = "unknown" + audit_policy_mode: Literal["off", "best_effort", "strict"] = "off" + policy_verdicts: tuple[dict[str, Any], ...] = () def __post_init__(self) -> None: if self.estimated_input_tokens < 0 or self.maximum_output_tokens < 0: @@ -218,6 +228,7 @@ class SpendReservationCommand: policy: SpendPolicySnapshot input_hash: str | None = None pii_entities: dict[str, int] | None = None + policy_verdicts: tuple[dict[str, Any], ...] = () def __post_init__(self) -> None: if self.estimated_cost_usd < 0: @@ -255,6 +266,9 @@ class FinalizationCommand: output_hash: str | None = None completed_at: datetime | None = None reconciliation_urgent: bool = False + provider_finish_reasons: dict[str, str] | None = None + ttft_ms: float | None = None + policy_verdicts: tuple[dict[str, Any], ...] | None = None def __post_init__(self) -> None: if self.quota_action is TerminalAction.NONE: @@ -329,6 +343,11 @@ async def reserve_quota( "cost_center": command.cost_center, "tags": list(command.tags), "team": command.team, + "repeat_chain_length": command.repeat_chain_length, + "system_prompt_hash": command.system_prompt_hash, + "deployment_kind": command.deployment_kind, + "audit_policy_mode": command.audit_policy_mode, + "policy_verdicts": list(command.policy_verdicts), }, }, ) @@ -462,7 +481,7 @@ async def write_spend_denial_preflight( session, command, lifecycle_status="spend_denied", - usage_summary={"denial_reason": "spend_limit_exceeded"}, + usage_summary={"spend_denied": 1}, ) async def finalize( @@ -573,6 +592,17 @@ async def _finalize_locked( ) if command.provider_model is not None: lifecycle.provider_model = command.provider_model + if not (all_replayed and lifecycle.reconciled_at is not None): + lifecycle.lifecycle_metadata = { + **(lifecycle.lifecycle_metadata or {}), + "provider_finish_reasons": command.provider_finish_reasons, + "ttft_ms": command.ttft_ms, + } + if command.policy_verdicts is not None and not all_replayed: + lifecycle.lifecycle_metadata = { + **(lifecycle.lifecycle_metadata or {}), + "policy_verdicts": list(command.policy_verdicts), + } audit_payload = await self._write_audit_completion( session, lifecycle, @@ -604,6 +634,7 @@ async def _finalize_locked( lifecycle_values: dict[str, object] = { "status": command.lifecycle_status, + "lifecycle_metadata": lifecycle.lifecycle_metadata, "reconciled_at": completed_at, "reconciliation_due_at": None, "terminal_error_code": command.terminal_error_code, @@ -847,17 +878,42 @@ async def _reserve_quota_periods( self, session: AsyncSession, command: QuotaReservationCommand, + ) -> list[dict[str, object]]: + allocations = await self._reserve_scoped_quota_periods( + session, command, command.policy + ) + if ( + command.policy.team_id is not None + and command.policy.team_policy is not None + ): + allocations.extend( + await self._reserve_scoped_quota_periods( + session, + command, + command.policy.team_policy, + team_id=command.policy.team_id, + ) + ) + return allocations + + async def _reserve_scoped_quota_periods( + self, + session: AsyncSession, + command: QuotaReservationCommand, + policy: QuotaPolicySnapshot, + *, + team_id: UUID | None = None, ) -> list[dict[str, object]]: token_delta = command.estimated_input_tokens + command.maximum_output_tokens periods: list[tuple[str, date, date, int | None, int | None, int]] = [] request_date = command.started_at.astimezone(timezone.utc).date() - if command.policy.daily_request_limit is not None: + if policy.daily_request_limit is not None: periods.append( ( "daily", request_date, request_date + timedelta(days=1), - command.policy.daily_request_limit, + policy.daily_request_limit, None, 0, ) @@ -871,8 +927,8 @@ async def _reserve_quota_periods( "monthly", month_start, month_end, - command.policy.monthly_request_limit, - command.policy.monthly_token_limit, + policy.monthly_request_limit, + policy.monthly_token_limit, token_delta, ) ) @@ -896,12 +952,14 @@ async def _reserve_quota_periods( token_delta=tokens, request_limit=request_limit, token_limit=token_limit, + team_id=team_id, ) if row is None: raise QuotaLimitExceeded(f"{period_type} quota exceeded") allocations.append( { "counter_type": "quota", + "team_id": str(team_id) if team_id else None, "period_row_id": str(row.id), "period_type": period_type, "period_start": start.isoformat(), @@ -924,10 +982,12 @@ async def _conditional_quota_upsert( token_delta: int, request_limit: int | None, token_limit: int | None, + team_id: UUID | None = None, ) -> QuotaPeriodUsage | None: statement = insert(QuotaPeriodUsage).values( organization_id=command.tenant_id, - api_key_id=command.api_key_id, + api_key_id=command.api_key_id if team_id is None else None, + team_id=team_id, period_type=period_type, period_start=period_start, period_end=period_end, @@ -940,10 +1000,15 @@ async def _conditional_quota_upsert( statement = statement.on_conflict_do_update( index_elements=[ QuotaPeriodUsage.organization_id, - QuotaPeriodUsage.api_key_id, + QuotaPeriodUsage.api_key_id + if team_id is None + else QuotaPeriodUsage.team_id, QuotaPeriodUsage.period_type, QuotaPeriodUsage.period_start, ], + index_where=QuotaPeriodUsage.team_id.is_not(None) + if team_id is not None + else None, set_={ "reserved_requests": ( QuotaPeriodUsage.reserved_requests + excluded.reserved_requests @@ -1097,8 +1162,6 @@ async def _write_preflight( lifecycle_status: str = "provider_pending", usage_summary: Mapping[str, object] | None = None, ) -> None: - if command.audit_policy_mode == "off": - return lifecycle = await RequestLifecycleRepository.get( session, organization_id=command.tenant_id, @@ -1106,6 +1169,12 @@ async def _write_preflight( ) if lifecycle is None: raise AccountingConflictError("audit preflight requires a lifecycle") + lifecycle.lifecycle_metadata = { + **(lifecycle.lifecycle_metadata or {}), + "policy_verdicts": list(command.policy_verdicts), + } + if command.audit_policy_mode == "off": + return await AuditIntentRepository.create( session, organization_id=command.tenant_id, @@ -1137,54 +1206,81 @@ async def _write_audit_completion( output_hash: str | None, completed_at: datetime, ) -> dict[str, Any] | None: - preflight = await AuditIntentRepository.fetch( - session, - organization_id=TenantId(lifecycle.organization_id), - request_id=RequestId(lifecycle.request_id), - event_type="preflight", - ) - if preflight is None: - return None + try: + preflight = await AuditIntentRepository.fetch( + session, + organization_id=TenantId(lifecycle.organization_id), + request_id=RequestId(lifecycle.request_id), + event_type="preflight", + ) + if preflight is None: + metadata = lifecycle.lifecycle_metadata or {} + mode = metadata.get("audit_policy_mode", "off") + if mode == "off": + return None + # Admission may end before the usual provider-spend preflight exists. + preflight = await AuditIntentRepository.create( + session, + organization_id=TenantId(lifecycle.organization_id), + values={ + "request_id": lifecycle.request_id, + "actor_type": lifecycle.actor_type, + "api_key_id": lifecycle.api_key_id, + "user_id": lifecycle.user_id, + "event_type": "preflight", + "audit_policy_mode": mode, + "pii_entities": metadata.get("pii_entities", {}), + "provider": lifecycle.provider, + "model": lifecycle.requested_model, + "lifecycle_status": "accepted", + }, + ) - intent = audit_completion_intent( - lifecycle, - preflight, - quota_event, - spend_event, - lifecycle_status=lifecycle_status, - output_hash=output_hash, - completed_at=completed_at, - ) - outbox = await OutboxWriter().append( - session, - organization_id=TenantId(lifecycle.organization_id), - values=intent.persistence_values(), - ) - await AuditIntentRepository.create( - session, - organization_id=TenantId(lifecycle.organization_id), - values={ - "request_id": lifecycle.request_id, - "actor_type": lifecycle.actor_type, - "api_key_id": lifecycle.api_key_id, - "user_id": lifecycle.user_id, - "event_type": "completion", - "audit_policy_mode": preflight.audit_policy_mode, - "input_hash": preflight.input_hash, - "output_hash": output_hash, - "pii_entities": dict(preflight.pii_entities or {}), - "provider": lifecycle.provider, - "model": lifecycle.provider_model or lifecycle.requested_model, - "usage_summary": { - "prompt_tokens": quota_event.prompt_tokens, - "completion_tokens": quota_event.completion_tokens, - "total_tokens": quota_event.total_tokens, + intent = audit_completion_intent( + lifecycle, + preflight, + quota_event, + spend_event, + lifecycle_status=lifecycle_status, + output_hash=output_hash, + completed_at=completed_at, + ) + outbox = await OutboxWriter().append( + session, + organization_id=TenantId(lifecycle.organization_id), + values=intent.persistence_values(), + ) + await AuditIntentRepository.create( + session, + organization_id=TenantId(lifecycle.organization_id), + values={ + "request_id": lifecycle.request_id, + "actor_type": lifecycle.actor_type, + "api_key_id": lifecycle.api_key_id, + "user_id": lifecycle.user_id, + "event_type": "completion", + "audit_policy_mode": preflight.audit_policy_mode, + "input_hash": preflight.input_hash, + "output_hash": output_hash, + "pii_entities": dict(preflight.pii_entities or {}), + "provider": lifecycle.provider, + "model": lifecycle.provider_model or lifecycle.requested_model, + "usage_summary": { + "prompt_tokens": quota_event.prompt_tokens, + "completion_tokens": quota_event.completion_tokens, + "total_tokens": quota_event.total_tokens, + }, + "lifecycle_status": lifecycle_status, + "outbox_event_id": outbox.id, }, - "lifecycle_status": lifecycle_status, - "outbox_event_id": outbox.id, - }, - ) - return dict(outbox.payload) + ) + return dict(outbox.payload) + except AuditIntentPersistenceError: + raise + except Exception as error: + raise AuditIntentPersistenceError( + "audit completion persistence failed" + ) from error async def _enqueue_analytics_projection( self, @@ -1412,6 +1508,8 @@ async def _apply_period_transitions( reservation.period_allocations, key=lambda item: ( str(item.get("counter_type")), + # Match admission's key-then-team order to avoid lock inversion. + bool(item.get("team_id")), str(item.get("period_type")), str(item.get("period_start")), str(item.get("period_row_id")), diff --git a/ee/src/shim_enterprise/billing/models.py b/ee/src/shim_enterprise/billing/models.py index 195b127..7ef5075 100644 --- a/ee/src/shim_enterprise/billing/models.py +++ b/ee/src/shim_enterprise/billing/models.py @@ -352,7 +352,7 @@ class UsageLedger(Base): class QuotaPeriodUsage(Base): - """Authoritative daily or monthly API-key quota counters.""" + """Authoritative daily or monthly key or team quota counters.""" __tablename__ = "quota_period_usage" __table_args__ = ( @@ -361,6 +361,24 @@ class QuotaPeriodUsage(Base): ("api_keys.organization_id", "api_keys.id"), name="fk_quota_period_usage_org_api_key", ), + ForeignKeyConstraint( + ("organization_id", "team_id"), + ("teams.organization_id", "teams.id"), + name="fk_quota_period_usage_org_team", + ), + CheckConstraint( + "(api_key_id IS NULL) <> (team_id IS NULL)", + name="ck_quota_period_usage_single_scope", + ), + Index( + "uq_quota_period_usage_team_scope", + "organization_id", + "team_id", + "period_type", + "period_start", + unique=True, + postgresql_where=text("team_id IS NOT NULL"), + ), UniqueConstraint( "organization_id", "api_key_id", @@ -398,9 +416,10 @@ class QuotaPeriodUsage(Base): ForeignKey("organizations.id"), nullable=False, ) - api_key_id: Mapped[UUID] = mapped_column( - PostgreSQLUUID(as_uuid=True), ForeignKey("api_keys.id"), nullable=False + api_key_id: Mapped[UUID | None] = mapped_column( + PostgreSQLUUID(as_uuid=True), ForeignKey("api_keys.id"), nullable=True ) + team_id: Mapped[UUID | None] = mapped_column(PostgreSQLUUID(as_uuid=True)) period_type: Mapped[str] = mapped_column(Text, nullable=False) period_start: Mapped[date] = mapped_column(Date, nullable=False) period_end: Mapped[date] = mapped_column(Date, nullable=False) diff --git a/ee/src/shim_enterprise/billing/read_models.py b/ee/src/shim_enterprise/billing/read_models.py index 76d232a..339ab35 100644 --- a/ee/src/shim_enterprise/billing/read_models.py +++ b/ee/src/shim_enterprise/billing/read_models.py @@ -28,6 +28,7 @@ class DailyUsage: prompt_tokens: int completion_tokens: int cost_usd: Decimal + unpriced_requests: int = 0 def as_public_record(self) -> dict[str, object]: return { @@ -36,7 +37,9 @@ def as_public_record(self) -> dict[str, object]: "request_count": self.request_count, "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, - "cost_usd": float(self.cost_usd), + "cost_usd": None if self.unpriced_requests else float(self.cost_usd), + "unpriced_requests": self.unpriced_requests, + "cost_complete": self.unpriced_requests == 0, } @@ -50,6 +53,7 @@ class BillingBreakdown: prompt_tokens: int completion_tokens: int cost_usd: Decimal + unpriced_requests: int = 0 def as_public_record(self) -> dict[str, object]: return { @@ -57,7 +61,9 @@ def as_public_record(self) -> dict[str, object]: "request_count": self.request_count, "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, - "cost_usd": self.cost_usd, + "cost_usd": None if self.unpriced_requests else self.cost_usd, + "unpriced_requests": self.unpriced_requests, + "cost_complete": self.unpriced_requests == 0, } @@ -81,10 +87,26 @@ async def daily_usage( # are additional provenance facts and must not count the request twice. is_usage = UsageLedger.event_type == "quota_settlement" is_spend = UsageLedger.event_type == "spend_settlement" + unpriced_requests = func.sum( + case( + ( + is_spend + & ( + UsageLedger.event_metadata["pricing"][ + "pricing_resolution" + ].as_string() + == "unknown" + ), + 1, + ), + else_=0, + ) + ).label("unpriced_requests") usage_date = func.date(func.timezone("UTC", UsageLedger.created_at)) statement = ( select( usage_date.label("usage_date"), + unpriced_requests, UsageLedger.requested_model.label("model"), func.sum(case((is_usage, UsageLedger.request_count), else_=0)).label( "request_count" @@ -118,6 +140,7 @@ async def daily_usage( prompt_tokens=int(row.prompt_tokens or 0), completion_tokens=int(row.completion_tokens or 0), cost_usd=Decimal(str(row.cost_usd or 0)), + unpriced_requests=int(row.unpriced_requests or 0), ) for row in rows ] @@ -184,6 +207,21 @@ async def breakdown( is_usage = UsageLedger.event_type == "quota_settlement" is_spend = UsageLedger.event_type == "spend_settlement" + unpriced_requests = func.sum( + case( + ( + is_spend + & ( + UsageLedger.event_metadata["pricing"][ + "pricing_resolution" + ].as_string() + == "unknown" + ), + 1, + ), + else_=0, + ) + ).label("unpriced_requests") request_count = func.sum( case((is_usage, UsageLedger.request_count), else_=0) ).label("request_count") @@ -199,6 +237,7 @@ async def breakdown( statement = ( select( group_key.label("key"), + unpriced_requests, request_count, prompt_tokens, completion_tokens, @@ -225,6 +264,7 @@ async def breakdown( prompt_tokens=int(row.prompt_tokens or 0), completion_tokens=int(row.completion_tokens or 0), cost_usd=Decimal(str(row.cost_usd or 0)), + unpriced_requests=int(row.unpriced_requests or 0), ) for row in rows ] diff --git a/ee/src/shim_enterprise/billing/spend.py b/ee/src/shim_enterprise/billing/spend.py index 3e2dba1..cc8f226 100644 --- a/ee/src/shim_enterprise/billing/spend.py +++ b/ee/src/shim_enterprise/billing/spend.py @@ -69,6 +69,7 @@ class BudgetUsage: cost_usd: Decimal tokens: int top_contributors: tuple[dict[str, object], ...] + unpriced_requests: int = 0 def fraction_of(self, budget: CostBudget) -> Decimal: fractions: list[Decimal] = [] @@ -135,8 +136,16 @@ async def _aggregate( period_start: datetime, ) -> BudgetUsage: spend_filters = self._scope_filters(budget, RequestLifecycle) + unpriced = ( + UsageLedger.event_metadata["pricing"]["pricing_resolution"].as_string() + == "unknown" + ) + known_spend = func.coalesce( + func.sum(UsageLedger.cost_usd).filter(unpriced.is_not(True)), Decimal("0") + ) + unpriced_count = func.count(UsageLedger.id).filter(unpriced) spend_statement = ( - select(func.coalesce(func.sum(UsageLedger.cost_usd), Decimal("0"))) + select(known_spend, unpriced_count) .select_from(UsageLedger) .join( RequestLifecycle, @@ -169,7 +178,8 @@ async def _aggregate( contributor_statement = ( select( contributor.label("cost_center"), - func.sum(UsageLedger.cost_usd).label("cost_usd"), + known_spend.label("cost_usd"), + unpriced_count.label("unpriced_requests"), ) .select_from(UsageLedger) .join( @@ -184,19 +194,28 @@ async def _aggregate( *spend_filters, ) .group_by(contributor) - .order_by(func.sum(UsageLedger.cost_usd).desc()) + .order_by(known_spend.desc()) .limit(3) ) - cost = Decimal(str((await session.execute(spend_statement)).scalar_one())) + cost, unpriced_requests = (await session.execute(spend_statement)).one() tokens = int((await session.execute(quota_statement)).scalar_one()) contributors = tuple( { "cost_center": center or UNTAGGED, "cost_usd": float(Decimal(str(value))), + "cost_complete": not count, + "unpriced_requests": count, } - for center, value in (await session.execute(contributor_statement)).all() + for center, value, count in ( + await session.execute(contributor_statement) + ).all() + ) + return BudgetUsage( + cost_usd=Decimal(str(cost)), + tokens=tokens, + top_contributors=contributors, + unpriced_requests=unpriced_requests, ) - return BudgetUsage(cost_usd=cost, tokens=tokens, top_contributors=contributors) @staticmethod def _scope_filters( @@ -285,6 +304,9 @@ async def _enqueue_alert( "threshold": float(threshold), "percent_used": float(fraction * 100), "current_usd": float(usage.cost_usd), + "cost_basis": "known_settled_spend", + "cost_complete": not usage.unpriced_requests, + "unpriced_requests": usage.unpriced_requests, "limit_usd": ( float(budget.limit_usd) if budget.limit_usd is not None diff --git a/ee/src/shim_enterprise/core/config.py b/ee/src/shim_enterprise/core/config.py index 27aa9de..9fad272 100644 --- a/ee/src/shim_enterprise/core/config.py +++ b/ee/src/shim_enterprise/core/config.py @@ -2,7 +2,9 @@ from __future__ import annotations +from pathlib import Path from typing import Literal, Self +from uuid import UUID from urllib.parse import urlsplit from pydantic import EmailStr, Field, RedisDsn, field_validator, model_validator @@ -12,6 +14,10 @@ class Settings(CommunitySettings): + MODEL_DEPLOYMENT_REQUIRED: bool = False + MODEL_DEPLOYMENT_ALLOWED_ORIGINS: list[str] = [] + MODEL_DEPLOYMENT_CA_BUNDLE: str | None = None + API_PREFIX: Literal["/api/v1"] = "/api/v1" DATABASE_URL: str @@ -26,26 +32,40 @@ class Settings(CommunitySettings): "gcp_secret_manager", "aws_secrets_manager", "azure_key_vault", + "vault", ] = "fernet" - SUPABASE_URL: str + VAULT_ADDR: str | None = None + VAULT_TOKEN_FILE: str | None = None + VAULT_KV_MOUNT: str = Field(default="secret", pattern=r"^[A-Za-z0-9_-]+$") + VAULT_NAMESPACE: str | None = None + + AUTH_MODE: Literal["supabase", "oidc"] = "supabase" + SUPABASE_URL: str | None = None + OIDC_ISSUER_URL: str | None = Field(default=None, max_length=512) + OIDC_CLIENT_ID: str | None = None + OIDC_CLIENT_SECRET: str | None = None + OIDC_REDIRECT_URI: str | None = None + OIDC_ORGANIZATION_ID: UUID | None = None + OIDC_GROUPS_CLAIM: str = "groups" + OIDC_GROUP_ROLE_MAP: dict[str, Literal["owner", "admin", "auditor", "member"]] = ( + Field(default_factory=dict) + ) + OIDC_TEAM_GROUP_MAP: dict[str, dict[str, str]] = Field(default_factory=dict) + OIDC_SESSION_SECONDS: int = Field(default=28_800, ge=60, le=86_400) + OIDC_REVALIDATE_SECONDS: int = Field(default=60, ge=10, le=300) + OIDC_API_AUDIENCE: str | None = None + OIDC_API_MAX_TOKEN_SECONDS: int = Field(default=300, ge=60, le=900) + DASHBOARD_ORIGIN: str | None = None SUPABASE_KEY: str | None = None - LEMON_SQUEEZY_SIGNING_SECRET: str | None = None - LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID: str | None = None - LEMON_SQUEEZY_SOLO_PRO_YEARLY_VARIANT_ID: str | None = None - LEMON_SQUEEZY_AGENCY_MONTHLY_VARIANT_ID: str | None = None - LEMON_SQUEEZY_AGENCY_YEARLY_VARIANT_ID: str | None = None - LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL: str | None = None - LEMON_SQUEEZY_SOLO_PRO_YEARLY_CHECKOUT_URL: str | None = None - LEMON_SQUEEZY_AGENCY_MONTHLY_CHECKOUT_URL: str | None = None - LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL: str | None = None - MANUAL_TEST_DASHBOARD_ENABLED: bool = False SHIM_TEST_USER_EMAIL: str | None = None DEFAULT_MONTHLY_TOKEN_LIMIT: int = Field(default=1_000_000, ge=0) + WORKER_HEARTBEAT_PATH: Path | None = None + GATEWAY_RECONCILIATION_GRACE_SECONDS: int = Field(default=120, ge=30, le=3_600) GATEWAY_RECONCILIATION_INTERVAL_SECONDS: int = Field(default=30, ge=5, le=3_600) GATEWAY_RECONCILIATION_BATCH_SIZE: int = Field(default=100, ge=1, le=1_000) @@ -74,22 +94,28 @@ class Settings(CommunitySettings): model_config = SettingsConfigDict(env_file="ee/.env") - @field_validator( - "LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL", - "LEMON_SQUEEZY_SOLO_PRO_YEARLY_CHECKOUT_URL", - "LEMON_SQUEEZY_AGENCY_MONTHLY_CHECKOUT_URL", - "LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL", - mode="before", - ) + @field_validator("OIDC_TEAM_GROUP_MAP") + @classmethod + def validate_oidc_team_mapping( + cls, value: dict[str, dict[str, str]] + ) -> dict[str, dict[str, str]]: + for group, mapping in value.items(): + if ( + not group + or set(mapping) != {"team_id", "role"} + or mapping["role"] not in {"member", "team_admin"} + ): + raise ValueError( + "OIDC team mappings require group, team_id, and member/team_admin role" + ) + UUID(mapping["team_id"]) + return value + + @field_validator("WORKER_HEARTBEAT_PATH") @classmethod - def validate_checkout_url(cls, value: object) -> object: - if value is None or (isinstance(value, str) and not value.strip()): - return None - if not isinstance(value, str): - raise ValueError("checkout URL must be an absolute HTTPS URL") - parsed = urlsplit(value) - if parsed.scheme != "https" or parsed.hostname is None: - raise ValueError("checkout URL must be an absolute HTTPS URL") + def validate_heartbeat_path(cls, value: Path | None) -> Path | None: + if value is not None and not value.is_absolute(): + raise ValueError("worker heartbeat path must be absolute") return value @field_validator("COMPLIANCE_EMAIL_FROM", mode="before") @@ -106,6 +132,68 @@ def validate_production_settings(self) -> Self: raise ValueError( "gateway reconciliation interval cannot exceed its grace period" ) + if self.AUTH_MODE == "supabase" and not self.SUPABASE_URL: + raise ValueError("supabase authentication requires SUPABASE_URL") + if self.AUTH_MODE == "oidc": + for name in ( + "OIDC_ISSUER_URL", + "OIDC_CLIENT_ID", + "OIDC_CLIENT_SECRET", + "OIDC_REDIRECT_URI", + "OIDC_ORGANIZATION_ID", + "DASHBOARD_ORIGIN", + ): + if not getattr(self, name): + raise ValueError(f"oidc authentication requires {name}") + if self.OIDC_API_AUDIENCE and self.OIDC_API_AUDIENCE == self.OIDC_CLIENT_ID: + raise ValueError( + "OIDC_API_AUDIENCE must differ from the login client to reject ID tokens" + ) + if not self.OIDC_GROUP_ROLE_MAP: + raise ValueError( + "OIDC_GROUP_ROLE_MAP must grant at least one group access" + ) + for name in ("OIDC_ISSUER_URL", "OIDC_REDIRECT_URI", "DASHBOARD_ORIGIN"): + parsed = urlsplit(getattr(self, name)) + if ( + parsed.scheme not in {"http", "https"} + or not parsed.hostname + or parsed.username + or parsed.password + or parsed.query + or parsed.fragment + or (self.ENVIRONMENT == "production" and parsed.scheme != "https") + ): + raise ValueError( + f"{name} must be an absolute URL; production requires HTTPS" + ) + dashboard = urlsplit(self.DASHBOARD_ORIGIN) + callback = urlsplit(self.OIDC_REDIRECT_URI) + if dashboard.path not in {"", "/"} or ( + callback.scheme, + callback.netloc, + callback.path, + ) != (dashboard.scheme, dashboard.netloc, "/api/v1/auth/callback"): + raise ValueError( + "OIDC_REDIRECT_URI must use DASHBOARD_ORIGIN/api/v1/auth/callback" + ) + if self.SECRET_BACKEND == "vault": + if not self.VAULT_ADDR or not self.VAULT_TOKEN_FILE: + raise ValueError("vault requires VAULT_ADDR and VAULT_TOKEN_FILE") + parsed = urlsplit(self.VAULT_ADDR) + if ( + parsed.scheme not in {"https", "http"} + or not parsed.hostname + or parsed.username + or parsed.password + or parsed.query + or parsed.fragment + or parsed.path not in {"", "/"} + or (self.ENVIRONMENT == "production" and parsed.scheme != "https") + ): + raise ValueError( + "VAULT_ADDR must be an origin; production requires HTTPS" + ) if self.ENVIRONMENT == "production" and self.SECRET_BACKEND == "fernet": raise ValueError("production requires a managed secret backend") if self.ENVIRONMENT == "production" and self.MANUAL_TEST_DASHBOARD_ENABLED: diff --git a/ee/src/shim_enterprise/gateway/pipeline/audit_intent.py b/ee/src/shim_enterprise/gateway/pipeline/audit_intent.py index 0835ca3..630b0bb 100644 --- a/ee/src/shim_enterprise/gateway/pipeline/audit_intent.py +++ b/ee/src/shim_enterprise/gateway/pipeline/audit_intent.py @@ -14,16 +14,18 @@ from shim_enterprise.billing.models import AuditIntent from shim_enterprise.core.config import settings from shim_enterprise.gateway.contracts.audit import validate_audit_intent +from shim.gateway.kernel.result import PreparedInference from shim.gateway.contracts.context import AuditPolicy from shim.gateway.contracts.ids import ApiKeyId, RequestId, TenantId, UserId from shim.gateway.contracts.inference import ScanVerdict from shim.gateway.contracts.principal import ActorType +from shim.gateway.usage import UsageAuditPersistenceError from shim.observability.metrics import bounded_label from shim.observability.tracing import start_span from shim_enterprise.outbox.publisher import OutboxWriter -class AuditIntentPersistenceError(RuntimeError): +class AuditIntentPersistenceError(UsageAuditPersistenceError): """A required audit intent could not be persisted durably.""" @@ -240,3 +242,88 @@ async def persist_scan_audit_completion( }, ) return completion.id + + +async def persist_token_count_audit( + session: AsyncSession, + prepared: PreparedInference, + input_tokens: int | None, + completed_at: datetime, +) -> None: + """Record auxiliary provider tokenization without an inference ledger entry.""" + mode = prepared.context.audit_policy.mode + if mode == "off": + return + entities = dict(prepared.privacy.pii_entities) if prepared.privacy else {} + phase = "preflight" if input_tokens is None else "completion" + usage = {"billable_executions": 0} + if input_tokens is not None: + usage["input_tokens"] = input_tokens + event = await OutboxWriter().append( + session, + organization_id=prepared.tenant_id, + values={ + "event_type": "audit.chain_append_requested", + "aggregate_type": "request", + "aggregate_id": str(prepared.request_id), + "idempotency_key": f"request:{prepared.request_id}:outbox:token_count:{phase}", + "status": "pending", + "next_attempt_at": completed_at, + "payload": { + "organization_id": str(prepared.tenant_id), + "api_key_id": str(prepared.api_key_id), + "request_id": str(prepared.request_id), + "event_type": "token_count_started" + if input_tokens is None + else "token_count", + "provider": str(prepared.provider), + "model": prepared.model, + "endpoint": "messages.count_tokens", + "pii_entities": entities, + "pii_detected": bool(entities), + "policy_verdicts": [ + verdict.model_dump(mode="json") + for verdict in prepared.policy_verdicts + ], + "extra": { + **( + {"input_tokens": input_tokens} + if input_tokens is not None + else {} + ), + "operation_type": "token_count", + "deployment_kind": prepared.deployment_kind, + **( + { + "deployment_id": prepared.target.deployment_id, + "declared_version": prepared.target.declared_version, + "provider_model": prepared.target.upstream_model, + } + if prepared.target is not None + else {} + ), + "billable_execution": False, + }, + }, + }, + ) + await AuditIntentRepository.create( + session, + organization_id=prepared.tenant_id, + values={ + "request_id": str(prepared.request_id), + "event_type": phase, + "actor_type": prepared.context.actor_type, + "api_key_id": prepared.api_key_id, + "user_id": None, + "audit_policy_mode": mode, + "pii_entities": entities, + "provider": str(prepared.provider), + "model": prepared.model, + "usage_summary": usage, + "lifecycle_status": "provider_pending" + if input_tokens is None + else "completed", + "outbox_event_id": event.id, + }, + ) diff --git a/ee/src/shim_enterprise/gateway/pipeline/outbox.py b/ee/src/shim_enterprise/gateway/pipeline/outbox.py index 02343c9..789241b 100644 --- a/ee/src/shim_enterprise/gateway/pipeline/outbox.py +++ b/ee/src/shim_enterprise/gateway/pipeline/outbox.py @@ -8,6 +8,9 @@ from types import MappingProxyType from typing import Any, Mapping +from shim.billing.pricing import DEFAULT_PRICE_BOOK +from shim.gateway.kernel.result import PreparedInference + @dataclass(frozen=True, slots=True) class GatewayOutboxIntent: @@ -60,7 +63,9 @@ def audit_completion_intent( "completion_tokens": quota_event.completion_tokens, "pii_detected": bool(preflight.pii_entities), "pii_entities": dict(preflight.pii_entities or {}), - "policy_verdicts": [], + "policy_verdicts": list( + (lifecycle.lifecycle_metadata or {}).get("policy_verdicts", []) + ), "is_cache_hit": lifecycle.cache_status == "hit", "latency_ms": max( 0, @@ -71,6 +76,12 @@ def audit_completion_intent( "audit_event_type": "completion", "lifecycle_status": lifecycle_status, "usage_estimated": estimated, + "pricing_resolution": _pricing_resolution(spend_event), + **{ + field: (lifecycle.lifecycle_metadata or {}).get(field) + for field in _DIAGNOSTIC_FIELDS + }, + "actor_type": lifecycle.actor_type, }, } return GatewayOutboxIntent( @@ -82,6 +93,58 @@ def audit_completion_intent( ) +def rejection_intent(prepared: PreparedInference) -> GatewayOutboxIntent: + """A tenant-authenticated request denied before accounting admission.""" + verdicts = [verdict.model_dump(mode="json") for verdict in prepared.policy_verdicts] + completed_at = prepared.policy_verdicts[-1].effective_at + return GatewayOutboxIntent( + event_type="audit.chain_append_requested", + aggregate_id=str(prepared.request_id), + idempotency_key=f"request:{prepared.request_id}:outbox:audit.completion", + available_at=completed_at, + payload=MappingProxyType( + { + "organization_id": str(prepared.tenant_id), + "request_id": str(prepared.request_id), + "api_key_id": str(prepared.api_key_id) + if prepared.api_key_id is not None + else None, + "actor": str(prepared.context.user_id) + if prepared.context.user_id is not None + else None, + "event_type": "ai_request", + "provider": str(prepared.provider), + "model": prepared.model + if DEFAULT_PRICE_BOOK.supports(prepared.model, str(prepared.provider)) + else None, + "endpoint": prepared.source_endpoint, + "policy_verdicts": verdicts, + "prompt_tokens": 0, + "completion_tokens": 0, + "cost_usd": 0.0, + "latency_ms": max( + 0, + int( + (completed_at - prepared.context.started_at).total_seconds() + * 1000 + ), + ), + "extra": { + "audit_event_type": "completion", + "actor_type": prepared.context.actor_type, + "lifecycle_status": "rejected" + if any( + verdict.outcome == "deny" + for verdict in prepared.policy_verdicts + ) + else "failed", + "admitted": False, + }, + } + ), + ) + + def analytics_terminal_intent( lifecycle: Any, quota_event: Any, @@ -115,12 +178,14 @@ def analytics_terminal_intent( "cost_usd": float( spend_event.cost_usd if spend_event is not None else Decimal("0") ), + "pricing_resolution": _pricing_resolution(spend_event), "lifecycle_status": lifecycle_status, "usage_estimated": quota_event.estimated or (spend_event is not None and spend_event.estimated), "cost_center": lifecycle_metadata.get("cost_center", "untagged"), "team": lifecycle_metadata.get("team"), "tags": list(lifecycle_metadata.get("tags") or []), + **{field: lifecycle_metadata.get(field) for field in _DIAGNOSTIC_FIELDS}, } return GatewayOutboxIntent( event_type=event_type, @@ -129,3 +194,20 @@ def analytics_terminal_intent( payload=MappingProxyType(payload), available_at=completed_at, ) + + +_DIAGNOSTIC_FIELDS = ( + "provider_finish_reasons", + "repeat_chain_length", + "ttft_ms", + "system_prompt_hash", + "deployment_kind", +) + + +def _pricing_resolution(spend_event: Any | None) -> str | None: + if spend_event is None or spend_event.event_type != "spend_settlement": + return None + return ((spend_event.event_metadata or {}).get("pricing") or {}).get( + "pricing_resolution" + ) diff --git a/ee/src/shim_enterprise/gateway/pipeline/quota_reservation.py b/ee/src/shim_enterprise/gateway/pipeline/quota_reservation.py index f62544c..9bb6b18 100644 --- a/ee/src/shim_enterprise/gateway/pipeline/quota_reservation.py +++ b/ee/src/shim_enterprise/gateway/pipeline/quota_reservation.py @@ -3,8 +3,9 @@ from __future__ import annotations from collections.abc import Callable, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace import hashlib +import hmac import json import logging from datetime import datetime, timedelta, timezone @@ -12,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Literal, cast from uuid import UUID +from fastapi import HTTPException from sqlalchemy import desc, func, select, text from shim_enterprise.core.config import settings @@ -49,10 +51,20 @@ from shim_enterprise.observability.lifecycle import RequestLifecycleRepository from shim_enterprise.gateway.pipeline.audit_intent import ( AuditIntentPersistenceError, + persist_token_count_audit, AuditIntentRepository, ) from shim_enterprise.gateway.pipeline.scan_policy import ResolvedScanActor -from shim_enterprise.tenants.models import ApiKey, ProviderSecret, TierDefinition +from shim_enterprise.gateway.pipeline.outbox import rejection_intent +from shim_enterprise.outbox.publisher import OutboxWriter +from shim_enterprise.tenants.models import ( + ApiKey, + ProviderSecret, + TierDefinition, + Team, + TeamMembership, + User, +) from shim_enterprise.observability.enterprise_metrics import ( QUOTA_RESERVATION_TOTAL, USAGE_SETTLEMENT_TOTAL, @@ -72,6 +84,55 @@ class AccountingPersistenceError(PersistenceError): logger = logging.getLogger(__name__) +def _system_prompt_hash(prepared: PreparedInference) -> str | None: + """Hash only explicitly supplied system content, before privacy transformation.""" + + payload = prepared.payload + material: dict[str, Any] = {} + field = { + "chat": None, + "responses": "instructions", + "messages": "system", + "count_tokens": "system", + "generate_content": "systemInstruction", + }[prepared.protocol] + if field is not None and payload.get(field) is not None: + material[field] = payload[field] + if prepared.protocol in {"chat", "responses"}: + messages = payload.get("messages" if prepared.protocol == "chat" else "input") + if isinstance(messages, list): + instructions = [ + {"role": item["role"], "content": item["content"]} + for item in messages + if isinstance(item, dict) + and item.get("role") in ("system", "developer") + and item.get("content") is not None + ] + if instructions: + material["messages"] = instructions + if not material: + return None + canonical = json.dumps( + [ + "shim.system_prompt.v1", + str(prepared.tenant_id), + prepared.protocol, + { + "deployment_id": prepared.target.deployment_id + if prepared.target + else None, + "deployment_kind": prepared.deployment_kind, + }, + material, + ], + sort_keys=True, + ensure_ascii=True, + separators=(",", ":"), + ).encode("utf-8") + key = (settings.COMPLIANCE_HASH_SALT or settings.SECRET_KEY).encode("utf-8") + return "hmac-sha256:v1:" + hmac.digest(key, canonical, "sha256").hex() + + class AccountingPolicyLoader: """Load current quota/spend policy while locking its authoritative row.""" @@ -93,6 +154,68 @@ async def quota( if api_key is None: raise AccountingPersistenceError("accounting API key no longer exists") + now = datetime.now(timezone.utc) + expiry = api_key.expires_at + if expiry is not None and expiry.tzinfo is None: + expiry = expiry.replace(tzinfo=timezone.utc) + if not api_key.is_active or (expiry is not None and expiry <= now): + raise HTTPException(status_code=401, detail="Invalid API Key") + if ( + api_key.allowed_models is not None + and prepared.model not in api_key.allowed_models + ): + raise HTTPException( + status_code=403, + detail={ + "code": "MODEL_NOT_ALLOWED", + "message": "Model is not allowed by API-key policy.", + }, + ) + owner = await session.scalar( + select(User).where( + User.id == api_key.user_id, + User.organization_id == api_key.organization_id, + User.is_active.is_(True), + User.role != "auditor", + ) + ) + if owner is None: + raise HTTPException(status_code=401, detail="Invalid API Key") + team_policy = None + if api_key.team_id is not None: + if owner.role not in {"owner", "admin"} and not await session.scalar( + select(TeamMembership.user_id).where( + TeamMembership.organization_id == api_key.organization_id, + TeamMembership.team_id == api_key.team_id, + TeamMembership.user_id == owner.id, + ) + ): + raise HTTPException( + status_code=403, detail="API-key owner is no longer a team member" + ) + team = await session.scalar( + select(Team) + .where( + Team.organization_id == api_key.organization_id, + Team.id == api_key.team_id, + ) + .with_for_update(read=True) + ) + if team is None: + raise AccountingPersistenceError("accounting team no longer exists") + team_policy = QuotaPolicySnapshot( + version=self._version( + f"team:{team.id}:{team.updated_at}", + ( + team.daily_request_limit, + team.monthly_request_limit, + team.monthly_token_limit, + ), + ), + daily_request_limit=team.daily_request_limit, + monthly_request_limit=team.monthly_request_limit, + monthly_token_limit=team.monthly_token_limit, + ) tier_statement = ( select(TierDefinition) .where(TierDefinition.slug == api_key.tier) @@ -108,10 +231,15 @@ async def quota( policy.monthly_token_limit, ) return QuotaPolicySnapshot( - version=self._version("quota-default", values), + version=self._version( + "quota-default", + (*values, team_policy.version if team_policy else None), + ), daily_request_limit=values[0], monthly_request_limit=values[1], monthly_token_limit=values[2], + team_id=api_key.team_id, + team_policy=team_policy, ) values = ( @@ -122,11 +250,13 @@ async def quota( return QuotaPolicySnapshot( version=self._version( f"quota:{tier.slug}:{getattr(tier, 'updated_at', None)}", - values, + (*values, team_policy.version if team_policy else None), ), daily_request_limit=values[0], monthly_request_limit=values[1], monthly_token_limit=values[2], + team_id=api_key.team_id, + team_policy=team_policy, ) async def spend( @@ -146,6 +276,11 @@ async def spend( .where( ProviderSecret.organization_id == prepared.tenant_id, ProviderSecret.provider == str(prepared.provider), + *( + [ProviderSecret.id == UUID(prepared.target.credential_reference)] + if prepared.target + else [] + ), ) .order_by(desc(ProviderSecret.created_at), desc(ProviderSecret.id)) .limit(1) @@ -204,9 +339,17 @@ async def reserve_quota( admission: AdmissionState, session: AsyncSession, ) -> ReservationResult: + policy: QuotaPolicySnapshot | None = None with start_span("gateway.quota_reservation") as span: try: policy = await self.policy_loader.quota(session, prepared) + prepared.record_verdict( + "quota.requests_and_tokens", + stage="admission", + outcome="allow", + reason_code="QUOTA_RESERVED", + policy_version=policy.version, + ) result = await self.repository.reserve_quota( session, QuotaReservationCommand( @@ -226,7 +369,15 @@ async def reserve_quota( tags=admission.tags, team=prepared.policy.team, stream=prepared.stream, + repeat_chain_length=admission.repeat_chain_length, + system_prompt_hash=_system_prompt_hash(prepared), + deployment_kind=prepared.deployment_kind, policy=policy, + audit_policy_mode=prepared.context.audit_policy.mode, + policy_verdicts=tuple( + verdict.model_dump(mode="json") + for verdict in prepared.policy_verdicts + ), ), ) await session.commit() @@ -234,13 +385,30 @@ async def reserve_quota( QUOTA_RESERVATION_TOTAL.labels(status=outcome).inc() span.set_attribute("status", outcome) return result - except QuotaLimitExceeded: + except (QuotaLimitExceeded, HTTPException) as error: await session.rollback() + access_denied = isinstance(error, HTTPException) + prepared.record_verdict( + "api_key.access" if access_denied else "quota.requests_and_tokens", + stage="admission", + outcome="deny", + reason_code=( + "API_KEY_ACCESS_DENIED" if access_denied else "QUOTA_EXCEEDED" + ), + policy_version=policy.version if policy is not None else None, + ) QUOTA_RESERVATION_TOTAL.labels(status="rejected").inc() span.set_attribute("status", "rejected") raise except Exception as exc: await session.rollback() + prepared.record_verdict( + "quota.requests_and_tokens", + stage="admission", + outcome="error", + reason_code="QUOTA_UNAVAILABLE", + policy_version=policy.version if policy is not None else None, + ) QUOTA_RESERVATION_TOTAL.labels(status="failed").inc() span.set_attribute("status", "failed") if isinstance(exc, AccountingPersistenceError): @@ -263,11 +431,21 @@ async def reserve_spend( prepared, ephemeral_byok, ) + prepared.record_verdict( + "spend.provider_monthly", + stage="provider_spend", + outcome="allow", + reason_code="SPEND_UNLIMITED" + if policy.monthly_limit_usd is None + else "SPEND_RESERVED", + policy_version=policy.version, + ) estimated_cost = compute_cost_usd( - prepared.model, + prepared.pricing_model, prepared.admission.estimated_input_tokens, prepared.admission.maximum_output_tokens, str(prepared.provider), + unpriced=prepared.unpriced, ) input_hash = content_ref( settings.COMPLIANCE_HASH_SALT or settings.SECRET_KEY, @@ -280,13 +458,14 @@ async def reserve_spend( request_id=prepared.request_id, requested_model=prepared.model, provider=provider, - provider_model=prepared.model, + provider_model=prepared.pricing_model, estimated_cost_usd=estimated_cost, pricing_metadata=DEFAULT_PRICE_BOOK.resolved_price_metadata( - prepared.model, + prepared.pricing_model, str(prepared.provider), input_tokens=prepared.admission.estimated_input_tokens, output_tokens=prepared.admission.maximum_output_tokens, + unpriced=prepared.unpriced, ), cache_status="bypass", audit_policy_mode=prepared.context.audit_policy.mode, @@ -295,17 +474,39 @@ async def reserve_spend( pii_entities=dict( cast(Mapping[str, int], privacy_facts["pii_entities"]) ), + policy_verdicts=tuple( + verdict.model_dump(mode="json") + for verdict in prepared.policy_verdicts + ), ) + if prepared.unpriced and policy.monthly_limit_usd is not None: + raise SpendLimitExceeded("MODEL_PRICE_UNKNOWN") result = await self.repository.reserve_provider_spend( session, command, ) await session.commit() return result - except SpendLimitExceeded: + except SpendLimitExceeded as exc: await session.rollback() if command is None: raise + prepared.record_verdict( + "spend.provider_monthly", + stage="provider_spend", + outcome="deny", + reason_code="MODEL_PRICE_UNKNOWN" + if str(exc) == "MODEL_PRICE_UNKNOWN" + else "SPEND_LIMIT_EXCEEDED", + policy_version=command.policy.version, + ) + command = replace( + command, + policy_verdicts=tuple( + verdict.model_dump(mode="json") + for verdict in prepared.policy_verdicts + ), + ) try: await self.repository.write_spend_denial_preflight(session, command) await session.commit() @@ -322,6 +523,13 @@ async def reserve_spend( raise except Exception as exc: await session.rollback() + prepared.record_verdict( + "spend.provider_monthly", + stage="provider_spend", + outcome="error", + reason_code="SPEND_UNAVAILABLE", + policy_version=command.policy.version if command is not None else None, + ) if isinstance(exc, AuditIntentPersistenceError): raise if isinstance(exc, AccountingPersistenceError): @@ -344,6 +552,15 @@ async def record_privacy( values={ "privacy_status": privacy_facts["privacy_status"], "pii_detected": privacy_facts["pii_detected"], + "lifecycle_metadata": RequestLifecycle.lifecycle_metadata.op("||")( + { + "policy_verdicts": [ + verdict.model_dump(mode="json") + for verdict in prepared.policy_verdicts + ], + "pii_entities": dict(prepared.privacy.pii_entities), + } + ), }, ) if lifecycle is None: @@ -534,6 +751,10 @@ async def refund( lifecycle_status=lifecycle_status, terminal_error_code=error_code, terminal_error_message=error_message, + policy_verdicts=tuple( + verdict.model_dump(mode="json") + for verdict in prepared.policy_verdicts + ), ), ) @@ -549,6 +770,81 @@ def __init__( self.accounting = accounting self.session_factory = session_factory + async def reject(self, prepared: PreparedInference) -> None: + mode = prepared.context.audit_policy.mode + if mode == "off": + return + try: + async with self.session_factory() as session: + lifecycle = await RequestLifecycleRepository.get( + session, + organization_id=prepared.tenant_id, + request_id=prepared.request_id, + ) + if lifecycle is not None: + # Admission may have committed before its acknowledgement failed. + verdict = prepared.policy_verdicts[-1] + await self.accounting.refund( + session, + prepared, + spend_reserved=False, + error_code=verdict.reason_code, + lifecycle_status="rejected" + if verdict.outcome == "deny" + else "failed", + ) + return + intent = rejection_intent(prepared) + outbox = await OutboxWriter().append( + session, + organization_id=prepared.tenant_id, + values=intent.persistence_values(), + ) + await AuditIntentRepository.create( + session, + organization_id=prepared.tenant_id, + values={ + "request_id": str(prepared.request_id), + "actor_type": prepared.context.actor_type, + "api_key_id": prepared.api_key_id, + "user_id": prepared.context.user_id, + "event_type": "completion", + "audit_policy_mode": mode, + "provider": str(prepared.provider), + "model": intent.payload["model"], + "lifecycle_status": intent.payload["extra"]["lifecycle_status"], + "outbox_event_id": outbox.id, + }, + ) + await session.commit() + except Exception as error: + if mode == "strict": + raise AuditIntentPersistenceError( + "required rejection audit failed" + ) from error + logger.error( + "Rejection audit could not be persisted type=%s", type(error).__name__ + ) + + async def record_token_count( + self, prepared: PreparedInference, input_tokens: int | None + ) -> None: + if prepared.context.audit_policy.mode == "off": + return + async with self.session_factory() as session: + try: + await persist_token_count_audit( + session, prepared, input_tokens, datetime.now(timezone.utc) + ) + await session.commit() + except Exception as exc: + await session.rollback() + if prepared.context.audit_policy.mode == "strict": + raise AuditIntentPersistenceError( + "required token-count audit failed" + ) from exc + logger.warning("Token-count audit failed type=%s", type(exc).__name__) + async def admit( self, prepared: PreparedInference, @@ -611,6 +907,8 @@ async def finalize( terminal_error_code=terminal.error_code, terminal_error_message=terminal.error_message, output_hash=usage.output_hash, + provider_finish_reasons=usage.provider_finish_reasons, + ttft_ms=usage.ttft_ms, completed_at=terminal.completed_at, ), ) @@ -669,13 +967,26 @@ async def fail( error_message = ( "Request ended before a provider value was delivered." ) + denied = next( + ( + verdict + for verdict in reversed(prepared.policy_verdicts) + if verdict.outcome == "deny" + ), + None, + ) await self.accounting.refund( session, prepared, spend_reserved=spend_reserved, - error_code=error_code, + error_code=denied.reason_code if denied is not None else error_code, error_message=error_message, + lifecycle_status="rejected" if denied is not None else "failed", ) + except AuditIntentPersistenceError: + if prepared.context.audit_policy.mode == "strict": + raise + logger.error("Rejected request audit unavailable; stale recovery retained") except Exception as exc: logger.error( "Durable accounting finalization failed; stale recovery retained " diff --git a/ee/src/shim_enterprise/manual_test_dashboard.py b/ee/src/shim_enterprise/manual_test_dashboard.py index 31d155f..3391c22 100644 --- a/ee/src/shim_enterprise/manual_test_dashboard.py +++ b/ee/src/shim_enterprise/manual_test_dashboard.py @@ -40,7 +40,7 @@ def install_manual_test_dashboard(application: FastAPI) -> None: def manual_test_dashboard() -> HTMLResponse: nonce = secrets.token_urlsafe(24) - supabase_origin = _supabase_origin(settings.SUPABASE_URL) + supabase_origin = _supabase_origin(settings.SUPABASE_URL or "") public_key = _public_supabase_key(settings.SUPABASE_KEY) payload = { "email": settings.SHIM_TEST_USER_EMAIL or "", diff --git a/ee/src/shim_enterprise/observability/analytics_projection.py b/ee/src/shim_enterprise/observability/analytics_projection.py index 1f21062..922d752 100644 --- a/ee/src/shim_enterprise/observability/analytics_projection.py +++ b/ee/src/shim_enterprise/observability/analytics_projection.py @@ -129,6 +129,17 @@ def _projection_values(message: OutboxMessage) -> dict: "provider": payload.get("provider"), "lifecycle_status": payload.get("lifecycle_status"), "usage_estimated": bool(payload.get("usage_estimated")), + "pricing_resolution": payload.get("pricing_resolution"), + **{ + field: payload.get(field) + for field in ( + "provider_finish_reasons", + "repeat_chain_length", + "ttft_ms", + "system_prompt_hash", + "deployment_kind", + ) + }, }, "cost_center": payload.get("cost_center"), "team": payload.get("team"), diff --git a/ee/src/shim_enterprise/observability/overview.py b/ee/src/shim_enterprise/observability/overview.py index e397b6a..6f4c5ca 100644 --- a/ee/src/shim_enterprise/observability/overview.py +++ b/ee/src/shim_enterprise/observability/overview.py @@ -37,7 +37,9 @@ class OverviewSummaryRecord: policy_rejections: int technical_success_rate: float | None p95_completed_latency_ms: int | None - settled_spend_usd: Decimal + settled_spend_usd: Decimal | None + cost_complete: bool + unpriced_requests: int status_counts: dict[str, int] @@ -45,7 +47,9 @@ class OverviewSummaryRecord: class OverviewTrendRecord: start: datetime requests: int - settled_spend_usd: Decimal + settled_spend_usd: Decimal | None + cost_complete: bool + unpriced_requests: int @dataclass(frozen=True, slots=True) @@ -158,6 +162,9 @@ def _summary_statement(tenant_id: UUID, start_at: datetime, end_at: datetime): ) columns = [ func.count(RequestLifecycle.id).label("requests"), + func.count(RequestLifecycle.id) + .filter(_unpriced_spend(tenant_id)) + .label("unpriced_requests"), *( func.count(RequestLifecycle.id) .filter(RequestLifecycle.status == lifecycle_status) @@ -199,6 +206,9 @@ def _trend_statement( select( bucket_start.label("start"), func.count(RequestLifecycle.id).label("requests"), + func.count(RequestLifecycle.id) + .filter(_unpriced_spend(tenant_id)) + .label("unpriced_requests"), func.coalesce( func.sum(func.coalesce(spend, Decimal("0"))), Decimal("0") ).label("settled_spend_usd"), @@ -289,6 +299,21 @@ def _settled_spend(tenant_id: UUID): ) +def _unpriced_spend(tenant_id: UUID): + return ( + select(UsageLedger.id) + .where( + UsageLedger.organization_id == tenant_id, + UsageLedger.request_id == RequestLifecycle.request_id, + UsageLedger.event_type == "spend_settlement", + UsageLedger.event_metadata["pricing"]["pricing_resolution"].as_string() + == "unknown", + ) + .correlate(RequestLifecycle) + .exists() + ) + + def _spend_denied(tenant_id: UUID): return ( select(AuditIntent.id) @@ -332,7 +357,11 @@ def _summary_from_row(row) -> OverviewSummaryRecord: status_counts["completed"] / technical_total if technical_total else None ), p95_completed_latency_ms=round(float(p95)) if p95 is not None else None, - settled_spend_usd=Decimal(str(row.settled_spend_usd or 0)), + settled_spend_usd=None + if row.unpriced_requests + else Decimal(str(row.settled_spend_usd or 0)), + cost_complete=not row.unpriced_requests, + unpriced_requests=int(row.unpriced_requests or 0), status_counts=status_counts, ) @@ -356,7 +385,8 @@ def _fill_trend( values = { _utc(row.start): ( int(row.requests or 0), - Decimal(str(row.settled_spend_usd or 0)), + None if row.unpriced_requests else Decimal(str(row.settled_spend_usd or 0)), + int(row.unpriced_requests or 0), ) for row in rows } @@ -364,8 +394,12 @@ def _fill_trend( step = timedelta(hours=1) if bucket == "hour" else timedelta(days=1) trend = [] while cursor < end_at: - requests, spend = values.get(cursor, (0, Decimal("0"))) - trend.append(OverviewTrendRecord(max(cursor, start_at), requests, spend)) + requests, spend, unpriced = values.get(cursor, (0, Decimal("0"), 0)) + trend.append( + OverviewTrendRecord( + max(cursor, start_at), requests, spend, not unpriced, unpriced + ) + ) cursor += step return trend diff --git a/ee/src/shim_enterprise/outbox/handlers.py b/ee/src/shim_enterprise/outbox/handlers.py index 666b922..bbc47b2 100644 --- a/ee/src/shim_enterprise/outbox/handlers.py +++ b/ee/src/shim_enterprise/outbox/handlers.py @@ -240,6 +240,11 @@ def _budget_text(payload: dict) -> str: return ( f"shim budget {scope}: {payload.get('percent_used')}% used in " f"{payload.get('period')}" + + ( + f" (known spend only; {payload.get('unpriced_requests')} unpriced requests)" + if payload.get("cost_complete") is False + else "" + ) ) diff --git a/ee/src/shim_enterprise/secrets/store.py b/ee/src/shim_enterprise/secrets/store.py index 750b2a1..83fcac5 100644 --- a/ee/src/shim_enterprise/secrets/store.py +++ b/ee/src/shim_enterprise/secrets/store.py @@ -6,6 +6,7 @@ import re from dataclasses import dataclass from typing import Any, Protocol +from uuid import UUID from sqlalchemy import desc, select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker @@ -19,6 +20,7 @@ GCP_V1_PREFIX = "gcpsm:v1:" AWS_V1_PREFIX = "awssm:v1:" AZURE_V1_PREFIX = "azurekv:v1:" +VAULT_V1_PREFIX = "vaultkv:v1:" @dataclass(frozen=True, slots=True) @@ -58,6 +60,14 @@ def parse_secret_ref(secret_ref: SecretRef | str) -> ParsedSecretRef: if not re.fullmatch(r"https://[^/]+/secrets/[^/]+/[A-Za-z0-9-]+", locator): raise ValueError("Azure secret reference must pin a version") return ParsedSecretRef("azure", locator.rsplit("/", 1)[-1], locator, "v1") + if value.startswith(VAULT_V1_PREFIX): + locator = value[len(VAULT_V1_PREFIX) :] + match = re.fullmatch( + r"([A-Za-z0-9_-]+/shim/[a-f0-9]{64}/[a-f0-9]{32})@([1-9][0-9]*)", locator + ) + if match is None: + raise ValueError("Vault secret reference must pin a numeric version") + return ParsedSecretRef("vault", match.group(2), match.group(1), "v1") raise ValueError("Unsupported secret reference") @@ -169,6 +179,10 @@ def get_secret_store() -> SecretStore: ) _store_singleton = AWSSecretsManagerStore() + elif settings.SECRET_BACKEND == "vault": + from shim_enterprise.secrets.vault import VaultSecretStore + + _store_singleton = VaultSecretStore() elif settings.SECRET_BACKEND == "azure_key_vault": from shim_enterprise.secrets.azure_key_vault import AzureKeyVaultStore @@ -199,11 +213,13 @@ async def resolve( self, tenant_id: TenantId, credential: EphemeralProviderCredential | None, + *, + reference: str | None = None, ) -> str | None: if credential is not None and credential.provider != self.provider: raise ValueError("credential does not match the selected provider") injected = credential.consume() if credential is not None else None - if injected: + if injected and reference is None: return injected from shim_enterprise.tenants.models import ProviderSecret @@ -216,6 +232,11 @@ async def resolve( .where( ProviderSecret.organization_id == tenant_id, ProviderSecret.provider == self.provider, + *( + [ProviderSecret.id == UUID(reference)] + if reference is not None + else [] + ), ) .order_by( desc(ProviderSecret.created_at), desc(ProviderSecret.id) diff --git a/ee/src/shim_enterprise/secrets/vault.py b/ee/src/shim_enterprise/secrets/vault.py new file mode 100644 index 0000000..e53a304 --- /dev/null +++ b/ee/src/shim_enterprise/secrets/vault.py @@ -0,0 +1,130 @@ +"""Vault KV v2 with tenant-bound envelopes and immutable version references.""" + +from __future__ import annotations + +import asyncio +import hashlib +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import httpx + +from shim.gateway.contracts.ids import SecretRef, TenantId +from shim_enterprise.core.config import settings +from shim_enterprise.secrets.store import ( + VAULT_V1_PREFIX, + decode_envelope, + encode_envelope, + parse_secret_ref, + validate_write, +) + + +class VaultSecretStore: + backend = "vault" + + def __init__(self, *, transport: httpx.AsyncBaseTransport | None = None) -> None: + self._transport = transport + + async def _call(self, method: str, path: str, **kwargs: Any) -> dict[str, Any]: + if not settings.VAULT_ADDR or not settings.VAULT_TOKEN_FILE: + raise ValueError("Vault address and token file are required") + # Vault Agent renews/replaces the token atomically; never cache its contents. + token = ( + await asyncio.to_thread(Path(settings.VAULT_TOKEN_FILE).read_text) + ).strip() + if not token or "\n" in token or "\r" in token: + raise ValueError("Invalid Vault token file") + headers = {"X-Vault-Token": token} + if settings.VAULT_NAMESPACE: + headers["X-Vault-Namespace"] = settings.VAULT_NAMESPACE + async with httpx.AsyncClient( + base_url=settings.VAULT_ADDR.rstrip("/") + "/v1/", + timeout=10, + follow_redirects=False, + transport=self._transport, + ) as client: + response = await client.request(method, path, headers=headers, **kwargs) + response.raise_for_status() + return response.json() if response.content else {} + + @staticmethod + def _path(tenant_id: TenantId) -> str: + tenant_hash = hashlib.sha256(str(tenant_id).encode()).hexdigest() + return f"{settings.VAULT_KV_MOUNT}/shim/{tenant_hash}/" + + async def _read( + self, tenant_id: TenantId, secret_ref: SecretRef, purpose: str | None + ) -> tuple[dict[str, Any], str, str]: + parsed = parse_secret_ref(secret_ref) + if parsed.backend != self.backend or not parsed.locator.startswith( + self._path(tenant_id) + ): + raise ValueError("Vault reference does not belong to this tenant/backend") + mount, path = parsed.locator.split("/", 1) + result = await self._call( + "GET", f"{mount}/data/{path}", params={"version": parsed.version} + ) + envelope = decode_envelope( + result["data"]["data"]["envelope"], tenant_id, purpose + ) + return envelope, f"{mount}/destroy/{path}", parsed.version + + async def put_secret( + self, + tenant_id: TenantId, + purpose: str, + plaintext: str, + metadata: dict[str, Any] | None = None, + ) -> SecretRef: + validate_write(tenant_id, purpose, plaintext) + locator = self._path(tenant_id) + uuid4().hex + mount, path = locator.split("/", 1) + result = await self._call( + "POST", + f"{mount}/data/{path}", + json={ + "options": {"cas": 0}, + "data": { + "envelope": encode_envelope(tenant_id, purpose, plaintext, metadata) + }, + }, + ) + reference = SecretRef(f"{VAULT_V1_PREFIX}{locator}@{result['data']['version']}") + parse_secret_ref(reference) + return reference + + async def get_secret( + self, + tenant_id: TenantId, + secret_ref: SecretRef, + *, + expected_purpose: str | None = None, + ) -> str: + envelope, _, _ = await self._read(tenant_id, secret_ref, expected_purpose) + return envelope["plaintext"] + + async def rotate_secret( + self, + tenant_id: TenantId, + secret_ref: SecretRef, + new_plaintext: str, + *, + expected_purpose: str | None = None, + ) -> SecretRef: + envelope, _, _ = await self._read(tenant_id, secret_ref, expected_purpose) + # A new path preserves the old reference until the DB/outbox commits rotation. + return await self.put_secret( + tenant_id, envelope["purpose"], new_plaintext, envelope["metadata"] + ) + + async def delete_secret( + self, + tenant_id: TenantId, + secret_ref: SecretRef, + *, + expected_purpose: str | None = None, + ) -> None: + _, path, version = await self._read(tenant_id, secret_ref, expected_purpose) + await self._call("PUT", path, json={"versions": [int(version)]}) diff --git a/ee/src/shim_enterprise/services/gateway/enterprise.py b/ee/src/shim_enterprise/services/gateway/enterprise.py index 480ec07..58c5a04 100644 --- a/ee/src/shim_enterprise/services/gateway/enterprise.py +++ b/ee/src/shim_enterprise/services/gateway/enterprise.py @@ -5,6 +5,7 @@ from typing import Any, Literal from sqlalchemy.ext.asyncio import AsyncSession +from fastapi import HTTPException from starlette.responses import Response from shim_enterprise.billing.ledger import QuotaLimitExceeded, SpendLimitExceeded @@ -46,7 +47,9 @@ async def dispatch_inference( *, payload: dict[str, Any], provider: Literal["openai", "anthropic", "google"], - protocol: Literal["chat", "responses", "messages", "generate_content"], + protocol: Literal[ + "chat", "responses", "messages", "count_tokens", "generate_content" + ], model: str, stream: bool, headers: dict[str, str], @@ -72,7 +75,15 @@ async def dispatch_inference( raise_audit_intent_error() except QuotaLimitExceeded: raise_accounting_limit("MONTHLY_QUOTA_EXCEEDED") - except SpendLimitExceeded: + except SpendLimitExceeded as exc: + if str(exc) == "MODEL_PRICE_UNKNOWN": + raise HTTPException( + 403, + detail={ + "code": "MODEL_PRICE_UNKNOWN", + "message": "A priced model is required to enforce this provider spending limit.", + }, + ) from None raise_accounting_limit("SPEND_LIMIT_EXCEEDED") except TenantPolicyConfigurationError: raise_tenant_policy_error() diff --git a/ee/src/shim_enterprise/tenants/audit.py b/ee/src/shim_enterprise/tenants/audit.py new file mode 100644 index 0000000..c942cdd --- /dev/null +++ b/ee/src/shim_enterprise/tenants/audit.py @@ -0,0 +1,42 @@ +"""Transactional audit intent shared by management and identity synchronization.""" + +from datetime import datetime, timezone +from uuid import uuid4 + +from sqlalchemy.ext.asyncio import AsyncSession + +from shim.gateway.contracts.ids import TenantId +from shim_enterprise.outbox.publisher import OutboxWriter +from shim_enterprise.tenants.models import User + + +async def record_management_action( + session: AsyncSession, + user: User, + action: str, + subject_id: str, + *, + details: dict[str, object] | None = None, +) -> None: + """Append non-secret change facts to the caller's transaction; never commit.""" + event_id = f"management:{uuid4()}" + await OutboxWriter().append( + session, + organization_id=TenantId(user.organization_id), + values={ + "event_type": "audit.chain_append_requested", + "aggregate_type": "management", + "aggregate_id": event_id, + "idempotency_key": f"{event_id}:audit", + "payload": { + "organization_id": str(user.organization_id), + "request_id": event_id, + "event_type": "management_action", + "actor": str(user.id), + "endpoint": action, + "extra": {"subject_id": subject_id, **(details or {})}, + }, + "status": "pending", + "next_attempt_at": datetime.now(timezone.utc), + }, + ) diff --git a/ee/src/shim_enterprise/tenants/deployments.py b/ee/src/shim_enterprise/tenants/deployments.py new file mode 100644 index 0000000..d1503e7 --- /dev/null +++ b/ee/src/shim_enterprise/tenants/deployments.py @@ -0,0 +1,340 @@ +"""Tenant deployment resolution and operator-owned outbound destination policy.""" + +from __future__ import annotations + +from dataclasses import replace +from datetime import datetime, timezone +import ipaddress +from typing import cast, Literal +from urllib.parse import urlsplit + +from fastapi import HTTPException +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from shim.billing.pricing import DEFAULT_PRICE_BOOK +from shim.gateway.kernel.result import ( + PreparedInference, + ProviderTarget, + UNSPECIFIED_PROVIDER_MODEL, +) +from shim.gateway.contracts.principal import AuthenticatedPrincipal +from shim.api.v1.chat import model_record +from shim_enterprise.core.config import settings +from shim_enterprise.tenants.models import ModelDeployment, ApiKey, User +from shim_enterprise.tenants.teams import require_team + + +def validate_deployment_url(url: str) -> str: + """Exact origins are approved by the platform operator, never by API callers.""" + parsed = urlsplit(url) + if ( + parsed.scheme not in {"https", "http"} + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.query + or parsed.fragment + or "\\" in url + or any(ord(char) < 33 for char in url) + ): + raise ValueError( + "Deployment URL must be an absolute HTTP(S) URL without credentials, query or fragment" + ) + port = parsed.port or (443 if parsed.scheme == "https" else 80) + origin = (parsed.scheme, parsed.hostname.casefold(), port) + approved = { + ( + item.scheme, + item.hostname, + item.port or (443 if item.scheme == "https" else 80), + ) + for item in map(urlsplit, settings.MODEL_DEPLOYMENT_ALLOWED_ORIGINS) + if item.path in {"", "/"} + and not item.query + and not item.fragment + and item.username is None + } + if origin not in approved: + raise ValueError("Deployment origin is not approved by the platform operator") + host = parsed.hostname.casefold().rstrip(".") + if host in {"metadata.google.internal", "metadata", "instance-data"}: + raise ValueError("Metadata endpoints are forbidden") + try: + address = ipaddress.ip_address(host) + except ValueError: + pass + else: + if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped: + address = address.ipv4_mapped + if address.is_link_local or address.is_multicast or address.is_unspecified: + raise ValueError("Metadata and non-unicast destinations are forbidden") + return url.rstrip("/") + + +async def require_model_aliases( + session: AsyncSession, tenant_id, aliases: list[str] | None +) -> None: + if not aliases: + return + registered = set( + ( + await session.scalars( + select(ModelDeployment.alias).where( + ModelDeployment.organization_id == tenant_id, + ModelDeployment.alias.in_(aliases), + ModelDeployment.enabled.is_(True), + ) + ) + ).all() + ) + if not settings.MODEL_DEPLOYMENT_REQUIRED: + registered.update( + alias + for alias in aliases + if any( + DEFAULT_PRICE_BOOK.supports(alias, provider) + for provider in ("openai", "anthropic", "google") + ) + ) + if set(aliases) - registered: + raise HTTPException( + 422, detail="Model allowlist contains an unregistered model" + ) + + +class DeploymentResolver: + def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: + self.session_factory = session_factory + + async def resolve(self, prepared: PreparedInference) -> PreparedInference: + async with self.session_factory() as session: + try: + key = await self._active_key(session, prepared.api_key_id) + if key.organization_id != prepared.tenant_id: + raise HTTPException(403, detail="Model access is not allowed") + except HTTPException: + prepared.record_verdict( + "access.api_key", + stage="admission", + outcome="deny", + reason_code="API_KEY_NOT_ACTIVE", + ) + raise + allowed = key.allowed_models + prepared.record_verdict( + "access.key_models", + stage="admission", + outcome="deny" + if allowed is not None and prepared.model not in allowed + else "allow", + reason_code="MODEL_NOT_ALLOWED" + if allowed is not None and prepared.model not in allowed + else "KEY_MODEL_ALLOWED", + policy=allowed, + ) + if allowed is not None and prepared.model not in allowed: + raise HTTPException( + 403, + detail={ + "code": "MODEL_NOT_ALLOWED", + "message": "The gateway key does not permit this model.", + }, + ) + deployment = ( + await session.execute( + select(ModelDeployment).where( + ModelDeployment.organization_id == prepared.tenant_id, + ModelDeployment.alias == prepared.model, + ) + ) + ).scalar_one_or_none() + if deployment is None: + if not settings.MODEL_DEPLOYMENT_REQUIRED and ( + DEFAULT_PRICE_BOOK.supports(prepared.model, str(prepared.provider)) + or ( + prepared.provider == "openai" + and prepared.protocol == "responses" + and prepared.model == UNSPECIFIED_PROVIDER_MODEL + ) + ): + prepared.record_verdict( + "deployment.registry", + stage="admission", + outcome="skip", + reason_code="CATALOG_ROUTING_ENABLED", + policy={"required": False}, + ) + return prepared + prepared.record_verdict( + "deployment.registry", + stage="admission", + outcome="deny", + reason_code="MODEL_NOT_REGISTERED", + policy={"required": settings.MODEL_DEPLOYMENT_REQUIRED}, + ) + raise HTTPException( + 403, + detail={ + "code": "MODEL_NOT_REGISTERED", + "message": "The requested model is not registered for this tenant.", + }, + ) + registry_policy = { + "deployment_id": str(deployment.id), + "updated_at": deployment.updated_at, + "enabled": deployment.enabled, + "provider": deployment.provider, + "version": deployment.declared_version, + } + if not deployment.enabled or deployment.provider != str(prepared.provider): + prepared.record_verdict( + "deployment.registry", + stage="admission", + outcome="deny", + reason_code="MODEL_NOT_ALLOWED", + policy=registry_policy, + ) + raise HTTPException( + 403, + detail={ + "code": "MODEL_NOT_ALLOWED", + "message": "The model is disabled or does not support this provider protocol.", + }, + ) + try: + base_url = validate_deployment_url(deployment.base_url) + except ValueError: + prepared.record_verdict( + "deployment.destination", + stage="admission", + outcome="deny", + reason_code="DEPLOYMENT_NOT_APPROVED", + policy=settings.MODEL_DEPLOYMENT_ALLOWED_ORIGINS, + ) + raise HTTPException( + 503, + detail={ + "code": "DEPLOYMENT_NOT_APPROVED", + "message": "The deployment destination is not approved.", + }, + ) from None + prepared.record_verdict( + "deployment.registry", + stage="admission", + outcome="allow", + reason_code="MODEL_REGISTERED", + policy=registry_policy, + ) + prepared.record_verdict( + "deployment.destination", + stage="admission", + outcome="allow", + reason_code="DEPLOYMENT_APPROVED", + policy=settings.MODEL_DEPLOYMENT_ALLOWED_ORIGINS, + ) + return replace( + prepared, + payload={**prepared.payload, "model": deployment.upstream_model}, + deployment_kind=cast( + Literal["internal", "external"], deployment.deployment_kind + ), + target=ProviderTarget( + deployment_id=str(deployment.id), + base_url=base_url, + upstream_model=deployment.upstream_model, + credential_reference=str(deployment.provider_secret_id), + timeout_seconds=deployment.timeout_seconds, + declared_version=deployment.declared_version, + ), + ) + + async def catalog( + self, principal: AuthenticatedPrincipal, provider: str + ) -> list[dict[str, object]]: + async with self.session_factory() as session: + key = await self._active_key(session, principal.api_key_id) + rows = ( + ( + await session.execute( + select(ModelDeployment) + .where( + ModelDeployment.organization_id == key.organization_id, + ) + .order_by(ModelDeployment.alias) + ) + ) + .scalars() + .all() + ) + records = ( + {} + if settings.MODEL_DEPLOYMENT_REQUIRED + else { + alias: model_record(alias, provider) + for alias in DEFAULT_PRICE_BOOK.models(provider) + } + ) + for row in rows: + records.pop(row.alias, None) + if not row.enabled or row.provider != provider: + continue + records[row.alias] = ( + { + "id": row.alias, + "type": "model", + "display_name": row.alias, + "created_at": row.created_at, + } + if provider == "anthropic" + else { + "id": row.alias, + "object": "model", + "created": int(row.created_at.timestamp()), + "owned_by": row.owner, + } + ) + return [ + records[alias] + for alias in sorted(records) + if key.allowed_models is None or alias in key.allowed_models + ] + + async def _active_key(self, session: AsyncSession, key_id) -> ApiKey: + row = ( + await session.execute( + select(ApiKey, User) + .join( + User, + (User.id == ApiKey.user_id) + & (User.organization_id == ApiKey.organization_id), + ) + .where( + ApiKey.id == key_id, + ApiKey.is_active.is_(True), + User.is_active.is_(True), + ) + ) + ).one_or_none() + if row is None: + raise HTTPException( + 401, + detail={ + "code": "INVALID_API_KEY", + "message": "The gateway key is no longer active.", + }, + ) + key, owner = row + if owner.role == "auditor" or ( + key.expires_at is not None and key.expires_at <= datetime.now(timezone.utc) + ): + raise HTTPException( + 401, + detail={ + "code": "INVALID_API_KEY", + "message": "The gateway key is no longer active.", + }, + ) + if key.team_id is not None: + await require_team(session, owner, key.team_id) + return key diff --git a/ee/src/shim_enterprise/tenants/models.py b/ee/src/shim_enterprise/tenants/models.py index 5f7e28e..0f984ad 100644 --- a/ee/src/shim_enterprise/tenants/models.py +++ b/ee/src/shim_enterprise/tenants/models.py @@ -78,12 +78,15 @@ class Organization(Base, TimestampMixin): class User(Base, TimestampMixin): - """Supabase-authenticated user projected into one mandatory tenant.""" + """Authenticated identity projected into one mandatory tenant.""" __tablename__ = "users" __table_args__ = ( UniqueConstraint("organization_id", "id", name="uq_users_tenant_id"), - CheckConstraint("role IN ('owner', 'admin', 'member')", name="ck_users_role"), + CheckConstraint( + "role IN ('owner', 'admin', 'member', 'auditor')", name="ck_users_role" + ), + UniqueConstraint("oidc_issuer", "oidc_subject", name="uq_users_oidc_identity"), Index("ix_users_organization_id", "organization_id"), ) @@ -93,6 +96,8 @@ class User(Base, TimestampMixin): organization_id: Mapped[UUID] = mapped_column( ForeignKey("organizations.id", ondelete="CASCADE"), nullable=False ) + oidc_issuer: Mapped[str | None] = mapped_column(String(512)) + oidc_subject: Mapped[str | None] = mapped_column(String(255)) email: Mapped[str] = mapped_column(String(320), nullable=False, unique=True) full_name: Mapped[str | None] = mapped_column(String(200)) is_active: Mapped[bool] = mapped_column( @@ -118,7 +123,9 @@ class OrganizationInvite(Base, TimestampMixin): __tablename__ = "organization_invites" __table_args__ = ( - CheckConstraint("role IN ('owner', 'admin', 'member')", name="ck_invites_role"), + CheckConstraint( + "role IN ('owner', 'admin', 'member', 'auditor')", name="ck_invites_role" + ), Index("ix_invites_organization_id", "organization_id"), Index("ix_invites_token_hash", "token_hash", unique=True), ) @@ -145,7 +152,7 @@ class OrganizationInvite(Base, TimestampMixin): class BillingWebhookReceipt(Base, TimestampMixin): - """Durable Lemon Squeezy delivery idempotency record.""" + """Historical billing delivery receipt retained after integration retirement.""" __tablename__ = "billing_webhook_receipts" __table_args__ = ( @@ -193,6 +200,77 @@ class TierDefinition(Base, TimestampMixin): ) +class Team(Base, TimestampMixin): + """Explicit tenant-owned delegation and quota boundary, not a billing label.""" + + __tablename__ = "teams" + __table_args__ = ( + UniqueConstraint("organization_id", "id", name="uq_teams_tenant_id"), + UniqueConstraint("organization_id", "name", name="uq_teams_tenant_name"), + CheckConstraint( + "daily_request_limit IS NULL OR daily_request_limit >= 0", + name="ck_teams_daily_requests", + ), + CheckConstraint( + "monthly_request_limit IS NULL OR monthly_request_limit >= 0", + name="ck_teams_monthly_requests", + ), + CheckConstraint( + "monthly_token_limit IS NULL OR monthly_token_limit >= 0", + name="ck_teams_monthly_tokens", + ), + ) + + id: Mapped[UUID] = mapped_column( + SqlUUID(as_uuid=True), primary_key=True, default=uuid4 + ) + organization_id: Mapped[UUID] = mapped_column( + ForeignKey("organizations.id", ondelete="CASCADE"), nullable=False + ) + name: Mapped[str] = mapped_column(String(128), nullable=False) + daily_request_limit: Mapped[int | None] = mapped_column(Integer) + monthly_request_limit: Mapped[int | None] = mapped_column(Integer) + monthly_token_limit: Mapped[int | None] = mapped_column(Integer) + + +class TeamMembership(Base, TimestampMixin): + """Membership grants access only within the named tenant and team.""" + + __tablename__ = "team_memberships" + __table_args__ = ( + ForeignKeyConstraint( + ["organization_id", "team_id"], + ["teams.organization_id", "teams.id"], + name="fk_team_memberships_tenant_team", + ondelete="CASCADE", + ), + ForeignKeyConstraint( + ["organization_id", "user_id"], + ["users.organization_id", "users.id"], + name="fk_team_memberships_tenant_user", + ondelete="CASCADE", + ), + CheckConstraint( + "role IN ('member', 'team_admin')", name="ck_team_memberships_role" + ), + CheckConstraint( + "source IN ('local', 'oidc')", name="ck_team_memberships_source" + ), + ) + + organization_id: Mapped[UUID] = mapped_column( + ForeignKey("organizations.id", ondelete="CASCADE"), primary_key=True + ) + team_id: Mapped[UUID] = mapped_column(SqlUUID(as_uuid=True), primary_key=True) + user_id: Mapped[UUID] = mapped_column(SqlUUID(as_uuid=True), primary_key=True) + role: Mapped[str] = mapped_column( + String(16), nullable=False, default="member", server_default="member" + ) + source: Mapped[str] = mapped_column( + String(16), nullable=False, default="local", server_default="local" + ) + + class ApiKey(Base, TimestampMixin): """One-way API-key verifier with enforced tenant/user ownership.""" @@ -205,6 +283,11 @@ class ApiKey(Base, TimestampMixin): name="fk_api_keys_tenant_user", ondelete="CASCADE", ), + ForeignKeyConstraint( + ["organization_id", "team_id"], + ["teams.organization_id", "teams.id"], + name="fk_api_keys_tenant_team", + ), Index("ix_api_keys_organization_id", "organization_id"), Index("ix_api_keys_key_hash", "key_hash", unique=True), ) @@ -227,6 +310,8 @@ class ApiKey(Base, TimestampMixin): expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) cost_center: Mapped[str | None] = mapped_column(String(128)) team: Mapped[str | None] = mapped_column(String(128)) + team_id: Mapped[UUID | None] = mapped_column(SqlUUID(as_uuid=True)) + allowed_models: Mapped[list[str] | None] = mapped_column(JSONB) tier: Mapped[str] = mapped_column( ForeignKey("tier_definitions.slug"), nullable=False, default="free" ) @@ -273,6 +358,7 @@ class ProviderSecret(Base, TimestampMixin): __tablename__ = "provider_secrets" __table_args__ = ( + UniqueConstraint("organization_id", "id", name="uq_provider_secrets_tenant_id"), CheckConstraint( "monthly_limit_usd IS NULL OR monthly_limit_usd >= 0", name="ck_provider_secret_monthly_limit", @@ -300,3 +386,59 @@ class ProviderSecret(Base, TimestampMixin): verified_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) organization: Mapped[Organization] = relationship(back_populates="provider_secrets") + + +class ModelDeployment(Base, TimestampMixin): + """Tenant-owned alias bound to an administrator-approved serving endpoint.""" + + __tablename__ = "model_deployments" + __table_args__ = ( + UniqueConstraint( + "organization_id", "alias", name="uq_model_deployments_tenant_alias" + ), + ForeignKeyConstraint( + ["organization_id", "provider_secret_id"], + ["provider_secrets.organization_id", "provider_secrets.id"], + ondelete="RESTRICT", + name="fk_model_deployments_tenant_secret", + ), + CheckConstraint( + "provider IN ('openai', 'anthropic')", name="ck_model_deployments_provider" + ), + CheckConstraint( + "deployment_kind IN ('internal', 'external')", + name="ck_model_deployments_kind", + ), + CheckConstraint( + "health IN ('unknown', 'healthy', 'unhealthy')", + name="ck_model_deployments_health", + ), + CheckConstraint( + "timeout_seconds BETWEEN 1 AND 300", name="ck_model_deployments_timeout" + ), + ) + + id: Mapped[UUID] = mapped_column( + SqlUUID(as_uuid=True), primary_key=True, default=uuid4 + ) + organization_id: Mapped[UUID] = mapped_column( + ForeignKey("organizations.id", ondelete="CASCADE"), nullable=False + ) + alias: Mapped[str] = mapped_column(String(200), nullable=False) + provider: Mapped[str] = mapped_column(String(32), nullable=False) + upstream_model: Mapped[str] = mapped_column(String(200), nullable=False) + base_url: Mapped[str] = mapped_column(String(2048), nullable=False) + provider_secret_id: Mapped[UUID] = mapped_column( + SqlUUID(as_uuid=True), nullable=False + ) + timeout_seconds: Mapped[int] = mapped_column(Integer, nullable=False) + deployment_kind: Mapped[str] = mapped_column(String(16), nullable=False) + declared_version: Mapped[str] = mapped_column(String(200), nullable=False) + owner: Mapped[str] = mapped_column(String(200), nullable=False) + enabled: Mapped[bool] = mapped_column( + Boolean, nullable=False, default=True, server_default="true" + ) + health: Mapped[str] = mapped_column( + String(16), nullable=False, default="unknown", server_default="unknown" + ) + health_checked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) diff --git a/ee/src/shim_enterprise/tenants/oidc.py b/ee/src/shim_enterprise/tenants/oidc.py new file mode 100644 index 0000000..033456a --- /dev/null +++ b/ee/src/shim_enterprise/tenants/oidc.py @@ -0,0 +1,484 @@ +"""Customer OIDC login and identity projection; tokens stay on the server.""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import logging +import secrets +import time +from typing import Any +from urllib.parse import urlencode, urlsplit +from uuid import uuid4 + +from authlib.integrations.base_client import OAuthError +from authlib.integrations.starlette_client import OAuth +from cryptography.fernet import Fernet, InvalidToken +from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi.responses import JSONResponse, RedirectResponse +import httpx +import jwt +from joserfc.errors import JoseError +from pydantic import EmailStr, TypeAdapter, ValidationError +from redis.exceptions import RedisError +from sqlalchemy import func, select +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession +from starlette.middleware.sessions import SessionMiddleware + +from shim_enterprise.core.config import settings +from shim_enterprise.core.database import get_db +from shim_enterprise.tenants.models import Organization, User +from shim_enterprise.tenants.teams import synchronize_oidc_teams + +router = APIRouter(prefix="/auth", tags=["identity"]) +SESSION_COOKIE = "shim_session" +ASYMMETRIC_ALGORITHMS = { + "RS256", + "RS384", + "RS512", + "PS256", + "PS384", + "PS512", + "ES256", + "ES384", + "ES512", + "EdDSA", +} + + +def install_oidc(application: Any) -> None: + if settings.AUTH_MODE != "oidc": + return + logging.getLogger("uvicorn.access").addFilter(_redact_login_query) + oauth = OAuth() + application.state.oidc = oauth.register( + "customer", + client_id=settings.OIDC_CLIENT_ID, + client_secret=settings.OIDC_CLIENT_SECRET, + server_metadata_url=f"{str(settings.OIDC_ISSUER_URL).rstrip('/')}/.well-known/openid-configuration", + client_kwargs={ + "scope": "openid profile email", + "code_challenge_method": "S256", + "timeout": 10, + "follow_redirects": False, + }, + ) + application.add_middleware( + SessionMiddleware, + secret_key=settings.SECRET_KEY, + session_cookie="shim_login", + max_age=300, + same_site="lax", + https_only=settings.ENVIRONMENT == "production", + ) + + +def _redact_login_query(record: logging.LogRecord) -> bool: + if isinstance(record.args, tuple) and len(record.args) == 5: + address, method, path, version, status = record.args + if isinstance(path, str) and path.startswith("/api/v1/auth/"): + record.args = (address, method, path.split("?", 1)[0], version, status) + return True + + +def require_origin(request: Request) -> None: + if request.headers.get("origin") != str(settings.DASHBOARD_ORIGIN).rstrip("/"): + raise HTTPException(403, "Same-origin request required") + + +def _redis(request: Request) -> Any: + redis = request.app.state.cache.redis + if redis is None: + raise HTTPException(503, "Identity session store unavailable") + return redis + + +def _cipher() -> Fernet: + return Fernet( + base64.urlsafe_b64encode( + hashlib.sha256( + ("shim-oidc-session:" + settings.SECRET_KEY).encode() + ).digest() + ) + ) + + +def _session_key(session_id: str) -> str: + return "identity:session:" + hashlib.sha256(session_id.encode()).hexdigest() + + +async def _client(request: Request) -> Any: + if settings.AUTH_MODE != "oidc": + raise HTTPException(404, "OIDC is not configured") + client = request.app.state.oidc + metadata = await client.load_server_metadata() + if metadata.get("issuer") != settings.OIDC_ISSUER_URL: + raise HTTPException(503, "OIDC discovery issuer does not match configuration") + algorithms = metadata.get("id_token_signing_alg_values_supported", ["RS256"]) + algorithms = sorted(set(algorithms) & ASYMMETRIC_ALGORITHMS) + if not algorithms: + raise HTTPException(503, "OIDC requires asymmetric token signatures") + metadata["id_token_signing_alg_values_supported"] = algorithms + return client + + +def _claims_options() -> dict[str, Any]: + return { + "iss": {"essential": True, "value": settings.OIDC_ISSUER_URL}, + "sub": {"essential": True}, + "exp": {"essential": True}, + "iat": {"essential": True}, + } + + +def identity_groups(claims: dict[str, Any]) -> list[str]: + groups = claims.get(settings.OIDC_GROUPS_CLAIM, []) + if not isinstance(groups, list) or not all( + isinstance(group, str) for group in groups + ): + raise HTTPException(403, "Identity groups must be an array of strings") + return groups + + +async def synchronize_user(session: AsyncSession, claims: dict[str, Any]) -> User: + issuer, subject = claims.get("iss"), claims.get("sub") + if ( + issuer != settings.OIDC_ISSUER_URL + or not isinstance(subject, str) + or not subject + or len(subject) > 255 + or not subject.isascii() + ): + raise HTTPException(401, "Invalid OIDC identity") + groups = identity_groups(claims) + roles = [ + settings.OIDC_GROUP_ROLE_MAP[group] + for group in groups + if group in settings.OIDC_GROUP_ROLE_MAP + ] + if not roles: + raise HTTPException(403, "No authorized identity group") + role = next( + role for role in ("owner", "admin", "member", "auditor") if role in roles + ) + # Tenant mutations share this lock order with team/key administration. + if not await session.scalar( + select(Organization.id) + .where(Organization.id == settings.OIDC_ORGANIZATION_ID) + .with_for_update() + ): + raise HTTPException(503, "OIDC organization has not been provisioned") + user = await session.scalar( + select(User) + .where(User.oidc_issuer == issuer, User.oidc_subject == subject) + .with_for_update() + ) + if user is None: + if claims.get("email_verified") is not True: + raise HTTPException(403, "A verified email address is required") + try: + email = str( + TypeAdapter(EmailStr).validate_python(claims.get("email")) + ).casefold() + except ValidationError as exc: + raise HTTPException(403, "A verified email address is required") from exc + if await session.scalar(select(User.id).where(func.lower(User.email) == email)): + raise HTTPException( + 403, "Identity email is already assigned; contact your administrator" + ) + name = claims.get("name") + user = User( + id=uuid4(), + organization_id=settings.OIDC_ORGANIZATION_ID, + oidc_issuer=issuer, + oidc_subject=subject, + email=email, + full_name=name[:200] if isinstance(name, str) else None, + role=role, + is_active=True, + is_verified=True, + ) + session.add(user) + if user.organization_id != settings.OIDC_ORGANIZATION_ID or not user.is_active: + raise HTTPException( + 403, "Identity membership is inactive or belongs to another organization" + ) + user.role = role + try: + await session.flush() + await synchronize_oidc_teams( + session, user, groups, settings.OIDC_TEAM_GROUP_MAP + ) + await session.commit() + except IntegrityError as exc: + await session.rollback() + raise HTTPException( + 409, "Identity provisioning conflicted; sign in again" + ) from exc + return user + + +async def _save_session( + request: Request, key: str, data: dict[str, Any], *, create: bool = False +) -> None: + ttl = int(data["expires_at"] - time.time()) + if ttl <= 0: + raise HTTPException(401, "Session expired") + encrypted = _cipher().encrypt(json.dumps(data).encode()).decode() + if not await _redis(request).set(key, encrypted, ex=ttl, nx=create, xx=not create): + raise HTTPException(401, "Identity session was revoked") + + +@router.get("/login") +async def login(request: Request, next: str = "/dashboard") -> Response: + client = await _client(request) + parsed = urlsplit(next) + if ( + not next.startswith("/") + or next.startswith("//") + or "\\" in next + or parsed.netloc + or parsed.scheme + or any(ord(c) < 32 for c in next) + ): + raise HTTPException(400, "Invalid return path") + request.session.clear() + request.session["next"] = next + return await client.authorize_redirect(request, settings.OIDC_REDIRECT_URI) + + +@router.get("/callback") +async def callback( + request: Request, session: AsyncSession = Depends(get_db) +) -> Response: + try: + client = await _client(request) + token = await client.authorize_access_token( + request, claims_options=_claims_options(), leeway=0 + ) + claims = dict(token["userinfo"]) + await synchronize_user(session, claims) + session_id = secrets.token_urlsafe(32) + data = { + "token": dict(token), + "claims": claims, + "checked_at": time.time(), + "expires_at": time.time() + settings.OIDC_SESSION_SECONDS, + } + if not token.get("refresh_token"): + data["expires_at"] = min( + data["expires_at"], + claims["exp"], + token.get("expires_at", claims["exp"]), + ) + await _save_session(request, _session_key(session_id), data, create=True) + except ( + OAuthError, + JoseError, + jwt.PyJWTError, + KeyError, + ValueError, + httpx.HTTPError, + RedisError, + ) as exc: + request.session.clear() + raise HTTPException(401, "OIDC sign-in failed") from exc + next_path = request.session.get("next", "/dashboard") + request.session.clear() + response = RedirectResponse( + str(settings.DASHBOARD_ORIGIN).rstrip("/") + next_path, status_code=303 + ) + response.set_cookie( + SESSION_COOKIE, + session_id, + httponly=True, + secure=settings.ENVIRONMENT == "production", + samesite="lax", + max_age=int(data["expires_at"] - time.time()), + path="/", + ) + response.headers["Cache-Control"] = "no-store" + return response + + +async def session_claims(request: Request) -> dict[str, Any]: + session_id = request.cookies.get(SESSION_COOKIE) + if not session_id or len(session_id) > 128: + raise HTTPException(401, "Session expired") + if request.method not in {"GET", "HEAD", "OPTIONS"}: + require_origin(request) + key = _session_key(session_id) + redis = _redis(request) + try: + stored = await redis.get(key) + if not stored: + raise HTTPException(401, "Session expired") + data = json.loads(_cipher().decrypt(stored.encode())) + if data["expires_at"] <= time.time(): + raise HTTPException(401, "Session expired") + if ( + time.time() - data["checked_at"] >= settings.OIDC_REVALIDATE_SECONDS + or data["claims"]["exp"] <= time.time() + ): + # One refresh per session prevents concurrent use of a rotated refresh token. + lock = redis.lock(key + ":refresh", timeout=30, blocking=False) + if not await lock.acquire(): + raise HTTPException( + 503, "Session refresh in progress", headers={"Retry-After": "1"} + ) + try: + latest = await redis.get(key) + if not latest: + raise HTTPException(401, "Session expired") + data = json.loads(_cipher().decrypt(latest.encode())) + if ( + time.time() - data["checked_at"] < settings.OIDC_REVALIDATE_SECONDS + and data["claims"]["exp"] > time.time() + ): + return data["claims"] + client = await _client(request) + refresh_token = data["token"].get("refresh_token") + if not refresh_token: + raise HTTPException( + 401, "Sign in again to revalidate identity membership" + ) + token = await client.fetch_access_token( + grant_type="refresh_token", refresh_token=refresh_token + ) + claims = dict( + await client.parse_id_token( + token, nonce=None, claims_options=_claims_options(), leeway=0 + ) + ) + if (claims["sub"], claims["iss"]) != ( + data["claims"]["sub"], + data["claims"]["iss"], + ): + raise HTTPException(401, "Session identity changed") + data.update( + token=dict( + token, refresh_token=token.get("refresh_token", refresh_token) + ), + claims=claims, + checked_at=time.time(), + ) + await _save_session(request, key, data) + except (OAuthError, JoseError, KeyError, ValueError, HTTPException): + await redis.delete(key) + raise HTTPException( + 401, "Identity session is no longer authorized" + ) from None + finally: + await lock.release() + return data["claims"] + except (InvalidToken, ValueError, KeyError) as exc: + raise HTTPException(401, "Invalid identity session") from exc + except (RedisError, httpx.HTTPError) as exc: + raise HTTPException(503, "Identity verification unavailable") from exc + + +async def access_token_claims(request: Request, token: str) -> dict[str, Any]: + if not settings.OIDC_API_AUDIENCE: + raise HTTPException(401, "OIDC bearer access is not configured") + try: + client = await _client(request) + header = jwt.get_unverified_header(token) + if header.get("alg") not in ASYMMETRIC_ALGORITHMS or not header.get("kid"): + raise jwt.InvalidTokenError("Invalid signing header") + jwks = await client.fetch_jwk_set() + key = next( + (key for key in jwks["keys"] if key.get("kid") == header["kid"]), None + ) + if key is None: + jwks = await client.fetch_jwk_set(force=True) + key = next( + (key for key in jwks["keys"] if key.get("kid") == header["kid"]), None + ) + if key is None: + raise jwt.InvalidTokenError("Unknown signing key") + claims = jwt.decode( + token, + jwt.PyJWK.from_dict(key).key, + algorithms=list(ASYMMETRIC_ALGORITHMS), + audience=settings.OIDC_API_AUDIENCE, + issuer=settings.OIDC_ISSUER_URL, + options={"require": ["exp", "iat", "sub", "iss", "aud"]}, + ) + if claims["exp"] - claims["iat"] > settings.OIDC_API_MAX_TOKEN_SECONDS: + raise jwt.InvalidTokenError( + "Access token lifetime exceeds configured revocation bound" + ) + return claims + except (jwt.PyJWTError, ValueError, KeyError, TypeError) as exc: + raise HTTPException(401, "Invalid OIDC access token") from exc + except httpx.HTTPError as exc: + raise HTTPException(503, "Identity verification unavailable") from exc + + +async def current_oidc_user( + request: Request, session: AsyncSession, token: str | None = None +) -> User: + claims = ( + await access_token_claims(request, token) + if token + else await session_claims(request) + ) + return await synchronize_user(session, claims) + + +@router.get("/session") +async def get_session( + request: Request, response: Response, session: AsyncSession = Depends(get_db) +) -> dict[str, Any]: + if settings.AUTH_MODE != "oidc": + raise HTTPException(404, "OIDC is not configured") + user = await current_oidc_user(request, session) + response.headers["Cache-Control"] = "no-store" + return { + "user": { + "id": str(user.id), + "email": user.email, + "user_metadata": {"full_name": user.full_name}, + "role": user.role, + } + } + + +@router.post("/logout") +async def logout(request: Request) -> Response: + require_origin(request) + session_id = request.cookies.get(SESSION_COOKIE) + if session_id: + await _redis(request).delete(_session_key(session_id)) + request.session.clear() + try: + client = await _client(request) + endpoint = (await client.load_server_metadata()).get("end_session_endpoint") + except (httpx.HTTPError, HTTPException, ValueError, OAuthError): + endpoint = None + logout_url = ( + endpoint + + "?" + + urlencode( + { + "client_id": settings.OIDC_CLIENT_ID, + "post_logout_redirect_uri": str(settings.DASHBOARD_ORIGIN).rstrip("/") + + "/login", + } + ) + if endpoint + else "/login" + ) + response = JSONResponse( + {"logout_url": logout_url}, headers={"Cache-Control": "no-store"} + ) + response.delete_cookie( + SESSION_COOKIE, + path="/", + secure=settings.ENVIRONMENT == "production", + httponly=True, + samesite="lax", + ) + return response diff --git a/ee/src/shim_enterprise/tenants/plans.py b/ee/src/shim_enterprise/tenants/plans.py new file mode 100644 index 0000000..04d854c --- /dev/null +++ b/ee/src/shim_enterprise/tenants/plans.py @@ -0,0 +1,59 @@ +"""Operator-managed organization plans and inherited API-key tiers.""" + +from datetime import datetime, timezone +from uuid import UUID, uuid4 + +from sqlalchemy import select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from shim_enterprise.tenants.models import ApiKey, Organization, TierDefinition +from shim_enterprise.tenants.service import ensure_privacy_defaults + + +async def create_organization_plan( + session: AsyncSession, + name: str, + tier: str, +) -> Organization: + name = name.strip() + if not 1 <= len(name) <= 200: + raise ValueError("Organization name must contain 1 to 200 characters") + organization_id = uuid4() + organization = Organization( + id=organization_id, name=name, slug=f"org-{organization_id}" + ) + session.add(organization) + await session.flush() + await ensure_privacy_defaults(session, organization.id) + return await activate_organization_plan(session, organization.id, tier) + + +async def activate_organization_plan( + session: AsyncSession, + organization_id: UUID, + tier: str, +) -> Organization: + if await session.get(TierDefinition, tier) is None: + raise ValueError(f"Unknown tier: {tier}") + organization = await session.scalar( + select(Organization) + .where(Organization.id == organization_id) + .with_for_update(of=Organization) + ) + if organization is None: + raise ValueError(f"Organization not found: {organization_id}") + organization.tier = tier + organization.billing_status = "free" if tier == "free" else "active" + organization.billing_source = "operator" + organization.billing_event_at = datetime.now(timezone.utc) + # Retain legacy billing references and period fields as historical records. + await session.flush() + await session.execute( + update(ApiKey) + .where( + ApiKey.organization_id == organization.id, + ApiKey.is_active.is_(True), + ) + .values(tier=tier) + ) + return organization diff --git a/ee/src/shim_enterprise/tenants/service.py b/ee/src/shim_enterprise/tenants/service.py index e88f0d3..06c7909 100644 --- a/ee/src/shim_enterprise/tenants/service.py +++ b/ee/src/shim_enterprise/tenants/service.py @@ -86,6 +86,8 @@ async def create_api_key( name: str, cost_center: str | None = None, team: str | None = None, + team_id: UUID | None = None, + allowed_models: list[str] | None = None, ) -> tuple[str, ApiKey]: tenant_id = await session.scalar( select(User.organization_id).where(User.id == user_id) @@ -119,12 +121,22 @@ async def create_api_key( is_active=True, cost_center=cost_center, team=team, + team_id=team_id, + allowed_models=allowed_models, ) session.add(api_key) await session.flush() return plaintext, api_key +def rotate_api_key(api_key: ApiKey) -> str: + """Replace the verifier in place, retaining ownership, policy and usage counters.""" + plaintext = f"{API_KEY_PREFIX}{secrets.token_hex(32)}" + api_key.key_hash = _digest_api_key(plaintext) + api_key.prefix = plaintext[:16] + return plaintext + + async def move_user_from_bootstrap( session: AsyncSession, *, @@ -206,9 +218,16 @@ async def authenticate_api_key( ) -> ApiKey | None: if not plaintext.startswith(API_KEY_PREFIX): return None - statement = select(ApiKey).where( - ApiKey.key_hash == _digest_api_key(plaintext), - ApiKey.is_active.is_(True), + statement = ( + select(ApiKey) + .join(User, User.id == ApiKey.user_id) + .where( + ApiKey.key_hash == _digest_api_key(plaintext), + ApiKey.is_active.is_(True), + User.organization_id == ApiKey.organization_id, + User.is_active.is_(True), + User.role != "auditor", + ) ) api_key = (await session.execute(statement)).scalar_one_or_none() if api_key is None: @@ -281,7 +300,7 @@ async def verify(self, token: str) -> Any | None: @staticmethod def _build_client() -> Any: - if not settings.SUPABASE_KEY: + if not settings.SUPABASE_URL or not settings.SUPABASE_KEY: raise RuntimeError("SUPABASE_KEY is required for JWT verification") from supabase import create_client diff --git a/ee/src/shim_enterprise/tenants/subscriptions.py b/ee/src/shim_enterprise/tenants/subscriptions.py deleted file mode 100644 index 538d72a..0000000 --- a/ee/src/shim_enterprise/tenants/subscriptions.py +++ /dev/null @@ -1,396 +0,0 @@ -"""Organization billing state and Lemon Squeezy webhook processing.""" - -from __future__ import annotations - -from datetime import datetime, timezone -import hashlib -import hmac -import json -from typing import Any, Literal -from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit -from uuid import UUID, uuid4 - -from sqlalchemy import select, update -from sqlalchemy.dialects.postgresql import insert -from sqlalchemy.ext.asyncio import AsyncSession - -from shim_enterprise.core.config import settings -from shim_enterprise.tenants.models import ( - ApiKey, - BillingWebhookReceipt, - Organization, - TierDefinition, - User, -) - -ACTIVE_BILLING_STATUSES = {"active", "on_trial"} -DEAD_BILLING_STATUSES = {"expired", "unpaid"} -KNOWN_EVENTS = { - "subscription_created", - "subscription_updated", - "subscription_cancelled", - "subscription_expired", - "subscription_paused", - "subscription_resumed", -} - - -def verify_lemonsqueezy_signature( - payload: bytes, - signature: str, - secret: str, -) -> bool: - if not signature or not secret: - return False - expected = hmac.new(secret.encode(), payload, hashlib.sha256).hexdigest() - try: - return hmac.compare_digest(expected, signature) - except TypeError: - return False - - -def checkout_urls( - organization_id: UUID, - user_id: UUID, -) -> dict[str, dict[Literal["monthly", "yearly"], str]]: - values: dict[ - str, - dict[Literal["monthly", "yearly"], str | None], - ] = { - "managed": { - "monthly": settings.LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL, - "yearly": settings.LEMON_SQUEEZY_SOLO_PRO_YEARLY_CHECKOUT_URL, - }, - "agency": { - "monthly": settings.LEMON_SQUEEZY_AGENCY_MONTHLY_CHECKOUT_URL, - "yearly": settings.LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL, - }, - } - return { - tier: { - period: _with_checkout_identity(url, organization_id, user_id) - for period, url in urls.items() - if url - } - for tier, urls in values.items() - if any(urls.values()) - } - - -async def set_organization_tier( - session: AsyncSession, - organization: Organization, - tier: str, - *, - status: str, - source: str, - event_at: datetime | None = None, -) -> None: - if await session.get(TierDefinition, tier) is None: - raise ValueError(f"Unknown tier: {tier}") - organization.tier = tier - organization.billing_status = status - organization.billing_source = source - organization.billing_event_at = event_at or datetime.now(timezone.utc) - await session.flush() - await session.execute( - update(ApiKey) - .where( - ApiKey.organization_id == organization.id, - ApiKey.is_active.is_(True), - ) - .values(tier=tier) - ) - - -async def process_lemonsqueezy_webhook( - session: AsyncSession, - raw_body: bytes, -) -> str: - try: - payload = json.loads(raw_body) - except (UnicodeDecodeError, json.JSONDecodeError) as exc: - raise ValueError("Invalid JSON payload") from exc - if not isinstance(payload, dict): - raise ValueError("Invalid webhook payload") - - meta = payload.get("meta") - data = payload.get("data") - if not isinstance(meta, dict) or not isinstance(data, dict): - raise ValueError("Invalid webhook payload") - event_name = str(meta.get("event_name") or "") - if event_name not in KNOWN_EVENTS: - return "ignored" - attributes = data.get("attributes") - if not isinstance(attributes, dict): - raise ValueError("Missing subscription attributes") - event_at = _parse_datetime(attributes.get("updated_at")) - if event_at is None: - raise ValueError("Missing subscription update timestamp") - - external_subscription_id = str(data.get("id") or "").strip() or None - organization = await _webhook_organization( - session, - meta, - external_subscription_id, - ) - digest = hashlib.sha256(raw_body).hexdigest() - claimed = ( - await session.execute( - insert(BillingWebhookReceipt) - .values( - id=uuid4(), - organization_id=organization.id, - payload_digest=digest, - event_name=event_name, - external_subscription_id=external_subscription_id, - event_at=event_at, - ) - .on_conflict_do_nothing(index_elements=["payload_digest"]) - .returning(BillingWebhookReceipt.id) - ) - ).scalar_one_or_none() - if claimed is None: - return "duplicate" - - status = _event_status(event_name, attributes) - if _is_stale(organization, event_at, status): - await _mark_processed(session, digest) - await session.commit() - return "stale" - - variant_id = str(attributes.get("variant_id") or "").strip() or None - if status in ACTIVE_BILLING_STATUSES: - tier = _variant_tier(variant_id) - if tier is None: - raise ValueError("Unknown Lemon Squeezy variant") - await set_organization_tier( - session, - organization, - tier, - status=status, - source="lemonsqueezy", - event_at=event_at, - ) - organization.cancel_at_period_end = False - elif status in DEAD_BILLING_STATUSES: - await set_organization_tier( - session, - organization, - "free", - status=status, - source="lemonsqueezy", - event_at=event_at, - ) - organization.cancel_at_period_end = False - else: - organization.billing_status = status - organization.billing_source = "lemonsqueezy" - organization.billing_event_at = event_at - organization.cancel_at_period_end = status == "cancelled" - - organization.external_subscription_id = external_subscription_id - organization.external_customer_id = _string_or_none(attributes.get("customer_id")) - organization.billing_variant_id = variant_id - organization.current_period_end = _parse_datetime( - attributes.get("ends_at") or attributes.get("renews_at") - ) - portal_url = _customer_portal_url(attributes) - if portal_url is not None: - organization.customer_portal_url = portal_url - await _mark_processed(session, digest) - await session.commit() - return "processed" - - -async def _webhook_organization( - session: AsyncSession, - meta: dict[str, Any], - external_subscription_id: str | None, -) -> Organization: - custom_organization_id = _custom_organization_id(meta) - if external_subscription_id is not None: - organization = ( - await session.execute( - select(Organization) - .where( - Organization.external_subscription_id == external_subscription_id - ) - .with_for_update(of=Organization) - ) - ).scalar_one_or_none() - if organization is not None: - if ( - custom_organization_id is not None - and custom_organization_id != organization.id - ): - raise ValueError("Subscription identity mismatch") - return organization - - organization_id, user_id = _verified_checkout_identity(meta) - organization = ( - await session.execute( - select(Organization) - .join(User, User.organization_id == Organization.id) - .where( - Organization.id == organization_id, - User.id == user_id, - User.role == "owner", - User.is_active.is_(True), - ) - .with_for_update(of=Organization) - ) - ).scalar_one_or_none() - if organization is None: - raise LookupError("Organization not found") - return organization - - -def _custom_organization_id(meta: dict[str, Any]) -> UUID | None: - custom_data = meta.get("custom_data") - if not isinstance(custom_data, dict) or not custom_data.get("organization_id"): - return None - try: - return UUID(str(custom_data["organization_id"])) - except ValueError as exc: - raise ValueError("Invalid organization ID") from exc - - -def _verified_checkout_identity(meta: dict[str, Any]) -> tuple[UUID, UUID]: - custom_data = meta.get("custom_data") - if not isinstance(custom_data, dict): - raise ValueError("Invalid checkout identity") - try: - organization_id = UUID(str(custom_data["organization_id"])) - user_id = UUID(str(custom_data["user_id"])) - signature = str(custom_data["checkout_signature"]) - except (KeyError, ValueError) as exc: - raise ValueError("Invalid checkout identity") from exc - expected = _checkout_signature(organization_id, user_id) - if not hmac.compare_digest(expected, signature): - raise ValueError("Invalid checkout identity") - return organization_id, user_id - - -async def _mark_processed(session: AsyncSession, digest: str) -> None: - await session.execute( - update(BillingWebhookReceipt) - .where(BillingWebhookReceipt.payload_digest == digest) - .values(processed_at=datetime.now(timezone.utc)) - ) - - -def _event_status(event_name: str, attributes: dict[str, Any]) -> str: - if event_name == "subscription_cancelled": - return "cancelled" - if event_name == "subscription_expired": - return "expired" - if event_name == "subscription_paused": - return "paused" - if event_name == "subscription_resumed": - return "active" - status = str(attributes.get("status") or "").strip() - if status not in ACTIVE_BILLING_STATUSES | DEAD_BILLING_STATUSES | { - "cancelled", - "past_due", - "paused", - }: - raise ValueError("Unknown subscription status") - return status - - -def _is_stale( - organization: Organization, - event_at: datetime, - new_status: str, -) -> bool: - previous = organization.billing_event_at - if previous is None: - return False - previous = _aware(previous) - event_at = _aware(event_at) - if event_at != previous: - return event_at < previous - severity = { - "active": 0, - "on_trial": 0, - "past_due": 1, - "paused": 2, - "cancelled": 3, - "unpaid": 4, - "expired": 5, - } - return severity.get(new_status, 0) < severity.get( - organization.billing_status, - 0, - ) - - -def _variant_tier(variant_id: str | None) -> str | None: - if variant_id is None: - return None - return { - settings.LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID: "managed", - settings.LEMON_SQUEEZY_SOLO_PRO_YEARLY_VARIANT_ID: "managed", - settings.LEMON_SQUEEZY_AGENCY_MONTHLY_VARIANT_ID: "agency", - settings.LEMON_SQUEEZY_AGENCY_YEARLY_VARIANT_ID: "agency", - }.get(variant_id) - - -def _with_checkout_identity( - url: str, - organization_id: UUID, - user_id: UUID, -) -> str: - parts = urlsplit(url) - query = dict(parse_qsl(parts.query, keep_blank_values=True)) - query["checkout[custom][organization_id]"] = str(organization_id) - query["checkout[custom][user_id]"] = str(user_id) - query["checkout[custom][checkout_signature]"] = _checkout_signature( - organization_id, - user_id, - ) - return urlunsplit((*parts[:3], urlencode(query), parts.fragment)) - - -def _checkout_signature(organization_id: UUID, user_id: UUID) -> str: - identity = f"lemonsqueezy-checkout:v1:{organization_id}:{user_id}".encode() - return hmac.new(settings.SECRET_KEY.encode(), identity, hashlib.sha256).hexdigest() - - -def _customer_portal_url(attributes: dict[str, Any]) -> str | None: - urls = attributes.get("urls") - if not isinstance(urls, dict): - return None - value = _string_or_none(urls.get("customer_portal")) - if value is None: - return None - parts = urlsplit(value) - hostname = (parts.hostname or "").casefold() - if parts.scheme != "https" or not ( - hostname == "lemonsqueezy.com" or hostname.endswith(".lemonsqueezy.com") - ): - return None - return value - - -def _parse_datetime(value: Any) -> datetime | None: - if not isinstance(value, str) or not value: - return None - try: - return _aware(datetime.fromisoformat(value.replace("Z", "+00:00"))) - except ValueError: - return None - - -def _aware(value: datetime) -> datetime: - if value.tzinfo is None: - return value.replace(tzinfo=timezone.utc) - return value.astimezone(timezone.utc) - - -def _string_or_none(value: Any) -> str | None: - if value is None: - return None - text = str(value).strip() - return text or None diff --git a/ee/src/shim_enterprise/tenants/teams.py b/ee/src/shim_enterprise/tenants/teams.py new file mode 100644 index 0000000..68b781b --- /dev/null +++ b/ee/src/shim_enterprise/tenants/teams.py @@ -0,0 +1,111 @@ +"""Tenant-scoped team access and identity-provider membership synchronization.""" + +from uuid import UUID + +from fastapi import HTTPException +from sqlalchemy import delete, select +from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.ext.asyncio import AsyncSession + +from shim_enterprise.tenants.models import Organization, Team, TeamMembership, User + + +def member_team_ids(user: User, *, administer: bool = False): + statement = select(TeamMembership.team_id).where( + TeamMembership.organization_id == user.organization_id, + TeamMembership.user_id == user.id, + ) + if administer: + statement = statement.where(TeamMembership.role == "team_admin") + return statement + + +async def require_team( + session: AsyncSession, user: User, team_id: UUID, *, administer: bool = False +) -> Team: + if administer: + if user.role == "auditor": + raise HTTPException( + status_code=403, detail="Auditors have read-only access" + ) + # The same tenant lock fences role and membership changes before authorization. + await session.scalar( + select(Organization.id) + .where(Organization.id == user.organization_id) + .with_for_update() + ) + statement = select(Team).where( + Team.organization_id == user.organization_id, Team.id == team_id + ) + if user.role not in {"owner", "admin", "auditor"}: + statement = statement.where( + Team.id.in_(member_team_ids(user, administer=administer)) + ) + if administer: + statement = statement.with_for_update() + team = await session.scalar(statement) + if team is None: + raise HTTPException(status_code=404, detail="Team not found") + return team + + +async def synchronize_oidc_teams( + session: AsyncSession, + user: User, + groups: list[str], + mapping: dict[str, dict[str, str]], +) -> None: + """Replace IdP grants atomically; explicit local membership stays authoritative.""" + await session.scalar( + select(Organization.id) + .where(Organization.id == user.organization_id) + .with_for_update() + ) + desired: dict[UUID, str] = {} + for group in groups: + if group not in mapping: + continue + config = mapping[group] + team_id, role = UUID(config["team_id"]), config["role"] + if role not in {"member", "team_admin"}: + raise ValueError("Invalid OIDC team role") + if desired.get(team_id) != "team_admin": + desired[team_id] = role + owned = set( + ( + await session.scalars( + select(Team.id).where( + Team.organization_id == user.organization_id, Team.id.in_(desired) + ) + ) + ).all() + ) + if owned != set(desired): + raise ValueError("OIDC team does not belong to the configured organization") + await session.execute( + delete(TeamMembership).where( + TeamMembership.organization_id == user.organization_id, + TeamMembership.user_id == user.id, + TeamMembership.source == "oidc", + TeamMembership.team_id.not_in(desired), + ) + ) + for team_id, role in sorted(desired.items()): + statement = insert(TeamMembership).values( + organization_id=user.organization_id, + team_id=team_id, + user_id=user.id, + role=role, + source="oidc", + ) + await session.execute( + statement.on_conflict_do_update( + index_elements=[ + TeamMembership.organization_id, + TeamMembership.team_id, + TeamMembership.user_id, + ], + set_={"role": role}, + where=TeamMembership.source == "oidc", + ) + ) diff --git a/ee/src/shim_enterprise/workers/ai_act.py b/ee/src/shim_enterprise/workers/ai_act.py index e51917a..6867f3a 100644 --- a/ee/src/shim_enterprise/workers/ai_act.py +++ b/ee/src/shim_enterprise/workers/ai_act.py @@ -19,6 +19,7 @@ from shim_enterprise.core.database import AsyncSessionLocal, engine from shim.observability.logging import configure_error_reporting, configure_logging from shim.observability.tracing import configure_tracing, shutdown_tracing +from shim_enterprise.workers.readiness import write_heartbeat logger = logging.getLogger(__name__) @@ -30,6 +31,7 @@ class MaintenanceSummary: oversight_created: int = 0 oversight_expired: int = 0 archive_eligible: int = 0 + errors: int = 0 class AuditMaintenanceWorker: @@ -46,10 +48,10 @@ def __init__( async def run_once(self, *, anchor_date: date | None = None) -> MaintenanceSummary: async with self.session_factory.begin() as session: - anchored = ( + anchored, errors = ( await self._anchor_tenants(session, anchor_date) if settings.AI_ACT_AUDIT_ANCHOR_ENABLED - else 0 + else (0, 0) ) created = {"created": 0} expired = {"expired": 0} @@ -59,6 +61,7 @@ async def run_once(self, *, anchor_date: date | None = None) -> MaintenanceSumma archive = await archive_expired(session) return MaintenanceSummary( anchored_tenants=anchored, + errors=errors, oversight_created=int(created["created"]), oversight_expired=int(expired["expired"]), archive_eligible=int(archive["eligible"]), @@ -68,7 +71,7 @@ async def _anchor_tenants( self, session: AsyncSession, anchor_date: date | None, - ) -> int: + ) -> tuple[int, int]: target = anchor_date or datetime.now(timezone.utc).date() - timedelta(days=1) start = datetime.combine(target, datetime.min.time(), tzinfo=timezone.utc) tenant_ids = ( @@ -79,21 +82,24 @@ async def _anchor_tenants( ) ) ).scalars() - anchored = 0 + anchored = errors = 0 for tenant_id in tenant_ids: try: async with session.begin_nested(): anchor = await write_anchor(session, tenant_id, target) except Exception as exc: + errors += 1 logger.error("Tenant anchor failed type=%s", type(exc).__name__) else: anchored += int(anchor is not None) - return anchored + return anchored, errors async def run(self, stop_event: asyncio.Event) -> None: while not stop_event.is_set(): try: summary = await self.run_once() + if not summary.errors: + write_heartbeat(settings.WORKER_HEARTBEAT_PATH, "ai_act") logger.info("Audit maintenance completed summary=%s", summary) except asyncio.CancelledError: raise diff --git a/ee/src/shim_enterprise/workers/compliance.py b/ee/src/shim_enterprise/workers/compliance.py index d0c1eeb..67ace75 100644 --- a/ee/src/shim_enterprise/workers/compliance.py +++ b/ee/src/shim_enterprise/workers/compliance.py @@ -18,6 +18,7 @@ from shim_enterprise.core.database import AsyncSessionLocal, engine from shim.observability.logging import configure_error_reporting, configure_logging from shim.observability.tracing import configure_tracing, shutdown_tracing +from shim_enterprise.workers.readiness import write_heartbeat logger = logging.getLogger(__name__) @@ -75,6 +76,8 @@ async def run(self, stop_event: asyncio.Event) -> None: while not stop_event.is_set(): try: summary = await self.run_once() + if not summary.errors: + write_heartbeat(settings.WORKER_HEARTBEAT_PATH, "compliance") logger.info("Compliance sweep completed summary=%s", summary) except asyncio.CancelledError: raise diff --git a/ee/src/shim_enterprise/workers/outbox.py b/ee/src/shim_enterprise/workers/outbox.py index 7e6b36e..c4c5448 100644 --- a/ee/src/shim_enterprise/workers/outbox.py +++ b/ee/src/shim_enterprise/workers/outbox.py @@ -33,6 +33,7 @@ ) from shim_enterprise.outbox.models import OutboxEvent from shim_enterprise.outbox.publisher import OutboxMessage, OutboxPublisher +from shim_enterprise.workers.readiness import write_heartbeat logger = logging.getLogger(__name__) @@ -254,7 +255,9 @@ async def run(self, stop_event: asyncio.Event) -> None: logger.info("Outbox worker started worker_id=%s", self.worker_id) while not stop_event.is_set(): try: - await self.run_once() + summary = await self.run_once() + if not (summary.failed or summary.dead_lettered or summary.lease_lost): + write_heartbeat(settings.WORKER_HEARTBEAT_PATH, "outbox") except asyncio.CancelledError: raise except Exception as exc: diff --git a/ee/src/shim_enterprise/workers/readiness.py b/ee/src/shim_enterprise/workers/readiness.py new file mode 100644 index 0000000..580f9a6 --- /dev/null +++ b/ee/src/shim_enterprise/workers/readiness.py @@ -0,0 +1,72 @@ +"""Local successful-pass heartbeat; independent of application configuration.""" + +import argparse +import json +import logging +import math +import os +from pathlib import Path +import tempfile +import time + +logger = logging.getLogger(__name__) + + +def write_heartbeat(path: Path | None, worker: str) -> None: + if path is None: + return + temporary = None + try: + with tempfile.NamedTemporaryFile( + mode="w", dir=path.parent, delete=False + ) as file: + temporary = Path(file.name) + json.dump({"worker": worker, "monotonic_success": time.monotonic()}, file) + os.replace(temporary, path) + except (OSError, ValueError) as exc: + logger.error("Worker heartbeat write failed type=%s", type(exc).__name__) + finally: + if temporary is not None: + try: + temporary.unlink(missing_ok=True) + except (OSError, ValueError) as exc: + logger.error( + "Worker heartbeat cleanup failed type=%s", type(exc).__name__ + ) + + +def is_ready(path: Path, worker: str, max_age_seconds: float) -> bool: + if ( + not path.is_absolute() + or not math.isfinite(max_age_seconds) + or max_age_seconds <= 0 + ): + return False + try: + payload = json.loads(path.read_text()) + if not isinstance(payload, dict) or payload.get("worker") != worker: + return False + timestamp = payload.get("monotonic_success") + if type(timestamp) not in (int, float): + return False + age = time.monotonic() - timestamp + return math.isfinite(age) and 0 <= age <= max_age_seconds + except (OSError, ValueError, OverflowError): + return False + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--path", type=Path, required=True) + parser.add_argument( + "--worker", + choices=("outbox", "reconciliation", "compliance", "ai_act"), + required=True, + ) + parser.add_argument("--max-age-seconds", type=float, required=True) + args = parser.parse_args() + raise SystemExit(0 if is_ready(args.path, args.worker, args.max_age_seconds) else 1) + + +if __name__ == "__main__": + main() diff --git a/ee/src/shim_enterprise/workers/reconciliation.py b/ee/src/shim_enterprise/workers/reconciliation.py index 0b23000..fd2c232 100644 --- a/ee/src/shim_enterprise/workers/reconciliation.py +++ b/ee/src/shim_enterprise/workers/reconciliation.py @@ -15,6 +15,7 @@ from shim_enterprise.gateway.pipeline.reconciliation import ScanReconciler from shim.observability.logging import configure_error_reporting, configure_logging from shim.observability.tracing import configure_tracing, shutdown_tracing, start_span +from shim_enterprise.workers.readiness import write_heartbeat logger = logging.getLogger(__name__) @@ -77,6 +78,7 @@ async def run(self, stop_event: asyncio.Event) -> None: while not stop_event.is_set(): try: recovered = await self.run_once() + write_heartbeat(settings.WORKER_HEARTBEAT_PATH, "reconciliation") if recovered: logger.warning( "Recovered stale gateway requests count=%s", recovered diff --git a/ee/tests/billing/test_spend.py b/ee/tests/billing/test_spend.py index 6daa015..dd8929b 100644 --- a/ee/tests/billing/test_spend.py +++ b/ee/tests/billing/test_spend.py @@ -20,6 +20,7 @@ validate_budget_notification_config, ) from shim.gateway.contracts.ids import TenantId +from shim_enterprise.outbox.handlers import _budget_text def test_header_tags_define_primary_and_complete_attribution() -> None: @@ -513,6 +514,7 @@ async def execute(statement): prompt_tokens=20, completion_tokens=5, cost_usd=Decimal("0.12345678"), + unpriced_requests=0, ) ] ) @@ -637,3 +639,32 @@ async def test_billing_breakdown_reads_settlements_without_request_log( group_by="tag", limit=100, ) + + +@pytest.mark.asyncio +async def test_budget_alert_labels_incomplete_known_spend(monkeypatch) -> None: + append = AsyncMock() + monkeypatch.setattr("shim_enterprise.billing.spend.OutboxWriter.append", append) + await BudgetEvaluator._enqueue_alert( + AsyncMock(), + SimpleNamespace( + id=uuid4(), + organization_id=uuid4(), + scope_type="org", + scope_value=None, + limit_usd=10, + limit_tokens=None, + notify_targets=[{"kind": "webhook"}], + ), + BudgetUsage(Decimal("8"), 50, (), unpriced_requests=2), + fraction=Decimal("0.8"), + threshold=Decimal("0.8"), + period_key="2026-09", + now=datetime.now(timezone.utc), + ) + payload = append.await_args.kwargs["values"]["payload"] + assert payload["current_usd"] == 8 + assert payload["cost_basis"] == "known_settled_spend" + assert payload["cost_complete"] is False + assert payload["unpriced_requests"] == 2 + assert "known spend only; 2 unpriced requests" in _budget_text(payload) diff --git a/ee/tests/core/test_enterprise_config.py b/ee/tests/core/test_enterprise_config.py index 51e7812..5749f1d 100644 --- a/ee/tests/core/test_enterprise_config.py +++ b/ee/tests/core/test_enterprise_config.py @@ -12,7 +12,6 @@ "DATABASE_URL", "REDIS_URL", "SECRET_KEY", - "SUPABASE_URL", } ENTERPRISE_REQUIRED_VALUES = { "DATABASE_URL": "postgresql+asyncpg://test:test@localhost/test", diff --git a/ee/tests/gateway/api/test_deps.py b/ee/tests/gateway/api/test_deps.py index ed66ddb..0855021 100644 --- a/ee/tests/gateway/api/test_deps.py +++ b/ee/tests/gateway/api/test_deps.py @@ -378,7 +378,7 @@ async def test_verified_identity_sync_failure_is_service_unavailable( ) calls = ( - (enterprise_deps.get_current_user, (bearer, session)), + (enterprise_deps.get_current_user, (_request(), bearer, session)), ( enterprise_deps.get_scan_principal, ( @@ -409,7 +409,7 @@ async def test_identity_provider_failure_is_service_unavailable( ) calls = ( - (enterprise_deps.get_current_user, (bearer, session)), + (enterprise_deps.get_current_user, (_request(), bearer, session)), ( enterprise_deps.get_scan_principal, ( @@ -445,3 +445,28 @@ def unavailable(_token: str): ) with pytest.raises(RuntimeError, match="auth service unavailable"): await unavailable_verifier.verify("valid-looking-token") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path,allowed", + [ + ("/api/v1/compliance/audit/verify", True), + ("/api/v1/compliance/reports/audit", True), + ("/api/v1/compliance/reports/kvkk", True), + ("/api/v1/management/api-keys", False), + ("/api/v1/compliance/reports/kvkk/anything", False), + ], +) +async def test_auditor_only_allows_read_only_posts(monkeypatch, path, allowed): + user = SimpleNamespace(role="auditor", is_active=True) + monkeypatch.setattr( + enterprise_deps, "get_invite_user", AsyncMock(return_value=user) + ) + request = Request({"type": "http", "method": "POST", "path": path, "headers": []}) + if allowed: + assert await enterprise_deps.get_current_user(request, None, None) is user + else: + with pytest.raises(HTTPException) as error: + await enterprise_deps.get_current_user(request, None, None) + assert error.value.status_code == 403 diff --git a/ee/tests/gateway/api/test_management.py b/ee/tests/gateway/api/test_management.py index 0a1bc91..d18ebae 100644 --- a/ee/tests/gateway/api/test_management.py +++ b/ee/tests/gateway/api/test_management.py @@ -1,3 +1,4 @@ +import csv from datetime import datetime, timedelta, timezone from decimal import Decimal from types import SimpleNamespace @@ -276,6 +277,7 @@ async def test_request_activity_is_tenant_scoped_filterable_and_safe() -> None: policy_failed=1, p95_completed_latency_ms=401.2, settled_spend_usd=Decimal("1.25000000"), + unpriced_requests=0, ) row = SimpleNamespace( id=uuid4(), @@ -352,13 +354,27 @@ async def test_request_activity_is_tenant_scoped_filterable_and_safe() -> None: "completion_tokens", "usage_estimated", "cost_usd", + "cost_complete", "latency_ms", "pii_detected", "tags", "cost_center", "provider", "team", + "provider_finish_reasons", + "repeat_chain_length", + "ttft_ms", + "system_prompt_hash", + "deployment_kind", } + assert page.items[0].provider_finish_reasons is None + assert page.items[0].repeat_chain_length is None + assert page.items[0].ttft_ms is None + assert page.items[0].system_prompt_hash is None + assert page.items[0].deployment_kind is None + assert page.items[0].cost_complete is True + assert page.summary.cost_complete is True + assert page.summary.unpriced_requests == 0 assert page.items[0].provider == "openai" assert page.items[0].usage_estimated is False assert page.items[0].team == "platform" @@ -413,6 +429,7 @@ def test_request_activity_summary_has_null_technical_metrics_without_denominator policy_failed=0, p95_completed_latency_ms=None, settled_spend_usd=Decimal("0"), + unpriced_requests=0, ) ) @@ -548,7 +565,14 @@ async def test_request_export_streams_all_filtered_rows_and_neutralizes_formulas timestamp=datetime(2026, 7, 1, tzinfo=timezone.utc), path="/v1/responses", model="gpt-5-nano", - details={"provider": "openai", "lifecycle_status": "completed"}, + details={ + "provider": "openai", + "lifecycle_status": "completed", + "provider_finish_reasons": {"status": "incomplete"}, + "repeat_chain_length": 2, + "ttft_ms": 42.5, + "deployment_kind": "internal", + }, prompt_tokens=10, completion_tokens=2, latency_ms=100, @@ -583,6 +607,13 @@ async def test_request_export_streams_all_filtered_rows_and_neutralizes_formulas assert "'=unsafe" in content assert "'+formula" in content assert "'@ops" in content + exported = list(csv.DictReader(content.splitlines()))[0] + assert exported["provider_finish_reasons"] == '{"status": "incomplete"}' + assert exported["repeat_chain_length"] == "2" + assert exported["ttft_ms"] == "42.5" + assert exported["deployment_kind"] == "internal" + assert exported["system_prompt_hash"] == "" + assert exported["cost_complete"] == "True" rows.close.assert_awaited_once() statement = session.stream.await_args.args[0] compiled = statement.compile(dialect=postgresql.dialect()) @@ -688,6 +719,7 @@ def test_billing_exports_render_safe_csv_and_pdf() -> None: prompt_tokens=20, completion_tokens=5, cost_usd=Decimal("0.12345678"), + unpriced_requests=0, ) start = datetime(2026, 7, 1, tzinfo=timezone.utc) end = datetime(2026, 7, 2, tzinfo=timezone.utc) diff --git a/ee/tests/gateway/api/test_management_overview.py b/ee/tests/gateway/api/test_management_overview.py index 2dd39d8..f4bbc14 100644 --- a/ee/tests/gateway/api/test_management_overview.py +++ b/ee/tests/gateway/api/test_management_overview.py @@ -34,6 +34,8 @@ def _empty_summary() -> OverviewSummaryRecord: technical_success_rate=None, p95_completed_latency_ms=None, settled_spend_usd=Decimal("0"), + cost_complete=True, + unpriced_requests=0, status_counts={ "completed": 0, "provider_error": 0, @@ -68,6 +70,7 @@ def test_overview_metric_definitions_exclude_policy_and_client_outcomes() -> Non policy_failed=1, p95_completed_latency_ms=401.2, settled_spend_usd=Decimal("1.25000000"), + unpriced_requests=0, ) summary = _summary_from_row(row) @@ -78,6 +81,11 @@ def test_overview_metric_definitions_exclude_policy_and_client_outcomes() -> Non assert summary.technical_success_rate == pytest.approx(6 / 9) assert summary.p95_completed_latency_ms == 401 assert summary.settled_spend_usd == Decimal("1.25000000") + row.unpriced_requests = 1 + incomplete = _summary_from_row(row) + assert incomplete.settled_spend_usd is None + assert incomplete.cost_complete is False + assert incomplete.unpriced_requests == 1 assert _exception_category("failed", spend_denied=True) == "policy_rejection" assert _exception_category("client_disconnected", False) == "client_cancelled" assert _exception_category("provider_error", False) == "technical_failure" @@ -89,6 +97,7 @@ def test_overview_trend_zero_fills_utc_buckets() -> None: start=datetime(2026, 8, 2), requests=2, settled_spend_usd=Decimal("0.5"), + unpriced_requests=1, ) ] @@ -105,6 +114,9 @@ def test_overview_trend_zero_fills_utc_buckets() -> None: "2026-08-03T00:00:00+00:00", ] assert [point.requests for point in trend] == [0, 2, 0] + assert [point.settled_spend_usd for point in trend] == [0, None, 0] + assert [point.cost_complete for point in trend] == [True, False, True] + assert [point.unpriced_requests for point in trend] == [0, 1, 0] def test_overview_trend_does_not_label_partial_bucket_before_period() -> None: diff --git a/ee/tests/gateway/kernel/test_accounting_coordinator.py b/ee/tests/gateway/kernel/test_accounting_coordinator.py index 615818c..7affe7b 100644 --- a/ee/tests/gateway/kernel/test_accounting_coordinator.py +++ b/ee/tests/gateway/kernel/test_accounting_coordinator.py @@ -1,10 +1,13 @@ from __future__ import annotations from contextlib import asynccontextmanager +from dataclasses import replace from datetime import datetime, timedelta, timezone from decimal import Decimal import json -from types import SimpleNamespace +import csv +import io +from types import MethodType, SimpleNamespace from unittest.mock import AsyncMock, patch from uuid import UUID, uuid4 @@ -13,13 +16,20 @@ from sqlalchemy.exc import DBAPIError, IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from shim_enterprise.billing.spend import BudgetEvaluator from shim_enterprise.core.config import settings +from shim_enterprise.api.v1 import management import shim.gateway.pipeline.postprocess as postprocess_module -from shim.gateway.kernel.result import UNSPECIFIED_PROVIDER_MODEL +from shim.gateway.kernel.result import ( + PreparedInference, + ProviderTarget, + UNSPECIFIED_PROVIDER_MODEL, +) from shim_enterprise.gateway.pipeline.quota_reservation import ( AccountingPersistenceError, DurableAccountingCoordinator, DurableUsageLifecycle, + _system_prompt_hash, ) from shim_enterprise.gateway.pipeline.audit_intent import AuditIntentPersistenceError from shim.gateway.pipeline.provider_execution import ( @@ -48,19 +58,26 @@ UsageLedger, ) from shim_enterprise.observability.lifecycle import RequestLifecycleRepository +from shim_enterprise.observability import analytics_projection, overview from shim_enterprise.outbox.models import OutboxEvent +from shim_enterprise.outbox.publisher import OutboxMessage from shim.privacy.classification import content_ref from shim_enterprise.tenants.models import ApiKey, Organization, User def _prepared(audit_mode: str = "best_effort") -> SimpleNamespace: - return SimpleNamespace( + prepared = SimpleNamespace( + policy_verdicts=[], tenant_id=uuid4(), api_key_id=uuid4(), request_id=f"req_{uuid4().hex}", provider="openai", protocol="chat", model="gpt-5.6-luna", + pricing_model="gpt-5.6-luna", + target=None, + deployment_kind="unknown", + unpriced=False, stream=False, context=SimpleNamespace( audit_policy=SimpleNamespace(mode=audit_mode), @@ -77,6 +94,9 @@ def _prepared(audit_mode: str = "best_effort") -> SimpleNamespace: ), ) + prepared.record_verdict = MethodType(PreparedInference.record_verdict, prepared) + return prepared + def _failure_state( provider_started: bool, @@ -114,6 +134,165 @@ def _postprocessor(usage) -> ResponsePostprocessor: ) +def test_system_prompt_hash_is_keyed_scoped_and_excludes_conversation( + monkeypatch, +) -> None: + prepared = _prepared() + prepared.payload = { + "messages": [ + {"role": "system", "content": "private instruction"}, + {"role": "user", "content": "user one"}, + ] + } + monkeypatch.setattr(settings, "COMPLIANCE_HASH_SALT", "installation-key-one") + original = _system_prompt_hash(prepared) + assert original is not None and original.startswith("hmac-sha256:v1:") + assert "private instruction" not in original + prepared.payload["messages"][1]["content"] = "user two" + prepared.payload["messages"][0] = { + "content": "private instruction", + "role": "system", + } + assert _system_prompt_hash(prepared) == original + prepared.payload["messages"][0]["content"] += " " + assert _system_prompt_hash(prepared) != original + prepared.payload["messages"][0]["content"] = "private instruction" + monkeypatch.setattr(settings, "COMPLIANCE_HASH_SALT", "installation-key-two") + assert _system_prompt_hash(prepared) != original + monkeypatch.setattr(settings, "COMPLIANCE_HASH_SALT", "installation-key-one") + prepared.tenant_id = uuid4() + assert _system_prompt_hash(prepared) != original + prepared.target = ProviderTarget( + str(uuid4()), "https://one.invalid", "one", "secret", 30, "1" + ) + prepared.deployment_kind = "internal" + deployment_hash = _system_prompt_hash(prepared) + prepared.model = "renamed-alias" + prepared.target = replace( + prepared.target, + base_url="https://two.invalid", + upstream_model="two", + declared_version="2", + ) + assert _system_prompt_hash(prepared) == deployment_hash + prepared.target = replace(prepared.target, deployment_id=str(uuid4())) + assert _system_prompt_hash(prepared) != deployment_hash + deployment_hash = _system_prompt_hash(prepared) + prepared.deployment_kind = "external" + assert _system_prompt_hash(prepared) != deployment_hash + prepared.payload = {"previous_response_id": "resp_inherited"} + assert _system_prompt_hash(prepared) is None + + +@pytest.mark.parametrize( + ("protocol", "payload"), + [ + ("responses", {"instructions": "private instruction"}), + ( + "responses", + {"input": [{"role": "developer", "content": "private instruction"}]}, + ), + ("messages", {"system": [{"type": "text", "text": "private instruction"}]}), + ("count_tokens", {"system": "private instruction"}), + ( + "generate_content", + {"systemInstruction": {"parts": [{"text": "private instruction"}]}}, + ), + ], +) +def test_system_prompt_hash_covers_explicit_native_instructions( + protocol, payload +) -> None: + prepared = _prepared() + prepared.protocol = protocol + prepared.payload = payload + assert _system_prompt_hash(prepared) is not None + + +@pytest.mark.asyncio +async def test_diagnostic_metadata_survives_terminal_and_outbox_replay( + db, test_api_key, monkeypatch +) -> None: + repository = DurableAccountingRepository() + request_id = f"req_diagnostics_{uuid4().hex}" + started_at = datetime.now(timezone.utc) + await repository.reserve_quota( + db, + QuotaReservationCommand( + tenant_id=test_api_key.organization_id, + api_key_id=test_api_key.id, + request_id=request_id, + requested_model="gpt-5.6-luna", + source_endpoint="chat.completions", + started_at=started_at, + reconciliation_due_at=started_at + timedelta(minutes=2), + estimated_input_tokens=20, + maximum_output_tokens=30, + policy=QuotaPolicySnapshot("test", None, None, None), + repeat_chain_length=2, + system_prompt_hash="hmac-sha256:v1:" + "a" * 64, + deployment_kind="internal", + ), + ) + command = FinalizationCommand( + tenant_id=test_api_key.organization_id, + request_id=request_id, + quota_action=TerminalAction.SETTLE, + prompt_tokens=20, + completion_tokens=3, + estimated=True, + lifecycle_status="client_disconnected", + provider_finish_reasons={"choices.1.finish_reason": "length"}, + ttft_ms=125.5, + ) + await repository.finalize(db, command) + assert (await repository.finalize(db, command)).replayed + await db.flush() + lifecycle = ( + await db.execute( + select(RequestLifecycle).where(RequestLifecycle.request_id == request_id) + ) + ).scalar_one() + expected = { + "repeat_chain_length": 2, + "system_prompt_hash": "hmac-sha256:v1:" + "a" * 64, + "deployment_kind": "internal", + "provider_finish_reasons": {"choices.1.finish_reason": "length"}, + "ttft_ms": 125.5, + } + assert all( + lifecycle.lifecycle_metadata[key] == value for key, value in expected.items() + ) + event = ( + await db.execute( + select(OutboxEvent).where( + OutboxEvent.aggregate_id == request_id, + OutboxEvent.event_type == "analytics.request_failed", + ) + ) + ).scalar_one() + assert all(event.payload[key] == value for key, value in expected.items()) + + @asynccontextmanager + async def session_scope(): + yield db + + monkeypatch.setattr(analytics_projection, "AsyncSessionLocal", session_scope) + message = OutboxMessage.from_event(event) + await analytics_projection.project_request(message) + await analytics_projection.project_request(message) + projected = ( + await db.execute( + select(analytics_projection.RequestLog).where( + analytics_projection.RequestLog.request_id == request_id + ) + ) + ).scalar_one() + assert all(projected.details[key] == value for key, value in expected.items()) + assert projected.details["lifecycle_status"] == "client_disconnected" + assert projected.details["usage_estimated"] is True + + async def _create_tenant( session: AsyncSession, label: str, @@ -209,7 +388,7 @@ async def test_unspecified_reservation_is_conservative_and_nonnull() -> None: ) session = SimpleNamespace(commit=AsyncMock(), rollback=AsyncMock()) prepared = _prepared() - prepared.model = UNSPECIFIED_PROVIDER_MODEL + prepared.model = prepared.pricing_model = UNSPECIFIED_PROVIDER_MODEL await DurableAccountingCoordinator( repository=repository, @@ -347,6 +526,8 @@ async def test_disconnected_stream_settles_reserved_usage() -> None: terminal = usage.finalize.await_args.args[1] assert usage.finalize.await_args.args[0] is prepared assert terminal.terminal_status == "client_disconnected" + assert terminal.usage.provider_finish_reasons is None + assert terminal.usage.ttft_ms is None @pytest.mark.asyncio @@ -408,6 +589,7 @@ async def test_urgent_reconciliation_signal_commits_short_transaction() -> None: async def test_privacy_facts_commit_before_openai_execution() -> None: session = SimpleNamespace(commit=AsyncMock(), rollback=AsyncMock()) checked = SimpleNamespace( + policy_verdicts=[], tenant_id=uuid4(), request_id=f"req_{uuid4().hex}", privacy=PrivacyOutcome( @@ -423,7 +605,13 @@ async def test_privacy_facts_commit_before_openai_execution() -> None: ) as update_lifecycle: await DurableAccountingCoordinator().record_privacy(checked, session) - assert update_lifecycle.await_args.kwargs["values"] == { + values = dict(update_lifecycle.await_args.kwargs["values"]) + assert "policy_verdicts" in str( + values.pop("lifecycle_metadata") + .compile(compile_kwargs={"literal_binds": False}) + .params + ) + assert values == { "privacy_status": "scrubbed", "pii_detected": True, } @@ -501,7 +689,7 @@ async def execute(**kwargs): @pytest.mark.asyncio async def test_nonstream_settlement_uses_the_priced_response_model() -> None: prepared = _prepared() - prepared.model = UNSPECIFIED_PROVIDER_MODEL + prepared.model = prepared.pricing_model = UNSPECIFIED_PROVIDER_MODEL usage = SimpleNamespace(finalize=AsyncMock()) await _postprocessor(usage).finalize( @@ -525,7 +713,7 @@ async def test_nonstream_settlement_uses_the_priced_response_model() -> None: @pytest.mark.asyncio async def test_nonstream_settlement_does_not_trust_a_cheaper_response_model() -> None: prepared = _prepared() - prepared.model = "gpt-5.6" + prepared.model = prepared.pricing_model = "gpt-5.6" usage = SimpleNamespace(finalize=AsyncMock()) await _postprocessor(usage).finalize( @@ -549,7 +737,7 @@ async def test_failed_response_without_a_requested_model_uses_conservative_price None ): prepared = _prepared() - prepared.model = UNSPECIFIED_PROVIDER_MODEL + prepared.model = prepared.pricing_model = UNSPECIFIED_PROVIDER_MODEL usage = SimpleNamespace(finalize=AsyncMock()) await _postprocessor(usage).finalize( @@ -565,6 +753,8 @@ async def test_failed_response_without_a_requested_model_uses_conservative_price assert terminal.usage.settlement_cost_usd > 0 assert terminal.usage.pricing_metadata["pricing_resolution"] == "conservative_max" assert terminal.terminal_status == "provider_error" + assert terminal.usage.provider_finish_reasons == {"status": "failed"} + assert terminal.usage.ttft_ms is None @pytest.mark.asyncio @@ -681,6 +871,7 @@ async def session_scope(): spend_reserved=True, error_code="REQUEST_ABORTED", error_message="Request ended before a provider value was delivered.", + lifecycle_status="failed", ) @@ -709,6 +900,7 @@ async def session_scope(): spend_reserved=True, error_code="PROVIDER_UNAVAILABLE", error_message="The provider rejected the request without usage.", + lifecycle_status="failed", ) @@ -742,6 +934,7 @@ async def session_scope(): spend_reserved=spend_reserved, error_code="PROVIDER_USAGE_UNAVAILABLE", error_message="Provider execution began but usage could not be verified.", + lifecycle_status="failed", ) @@ -814,6 +1007,7 @@ async def session_scope(): spend_reserved=False, error_code="ADMISSION_ABORTED", error_message="Request ended before a provider value was delivered.", + lifecycle_status="failed", ) @@ -975,9 +1169,11 @@ async def test_failure_state_uses_durable_provider_marker( @pytest.mark.asyncio +@pytest.mark.parametrize("pricing_resolution", ["catalog", "unknown"]) async def test_spend_pricing_metadata_survives_terminal_fallback( db, test_api_key, + pricing_resolution, ) -> None: repository = DurableAccountingRepository() request_id = f"req_pricing_metadata_{uuid4().hex}" @@ -1004,7 +1200,7 @@ async def test_spend_pricing_metadata_survives_terminal_fallback( ) pricing_metadata = { "catalog_version": "catalog-v1", - "pricing_resolution": "catalog", + "pricing_resolution": pricing_resolution, "input_per_million": "0.2", "output_per_million": "1.2", } @@ -1052,6 +1248,98 @@ async def test_spend_pricing_metadata_survives_terminal_fallback( ) assert len(events) == 2 assert all(event.event_metadata["pricing"] == pricing_metadata for event in events) + outbox_event = ( + await db.execute( + select(OutboxEvent).where( + OutboxEvent.aggregate_id == request_id, + OutboxEvent.event_type == "analytics.request_completed", + ) + ) + ).scalar_one() + assert outbox_event.payload["pricing_resolution"] == pricing_resolution + values = analytics_projection._projection_values( + OutboxMessage.from_event(outbox_event) + ) + assert values["details"]["pricing_resolution"] == pricing_resolution + db.add(analytics_projection.RequestLog(**values)) + await db.flush() + page = await management.list_requests( + start=None, + end=None, + status_filter=None, + model=None, + request_id=request_id, + pii_detected=None, + tag=None, + cost_center=None, + limit=10, + offset=0, + user=SimpleNamespace(organization_id=test_api_key.organization_id), + session=db, + ) + summary_row = ( + await db.execute( + overview._summary_statement( + test_api_key.organization_id, + datetime.now(timezone.utc) - timedelta(days=1), + datetime.now(timezone.utc) + timedelta(days=1), + ).where(RequestLifecycle.request_id == request_id) + ) + ).one() + overview_summary = overview._summary_from_row(summary_row) + assert overview_summary.requests == 1 + assert overview_summary.cost_complete is (pricing_resolution != "unknown") + assert overview_summary.unpriced_requests == (pricing_resolution == "unknown") + assert overview_summary.settled_spend_usd == ( + None if pricing_resolution == "unknown" else Decimal("0.00004") + ) + budget_usage = await BudgetEvaluator()._aggregate( + db, + SimpleNamespace(organization_id=test_api_key.organization_id, scope_type="org"), + datetime.now(timezone.utc) - timedelta(days=1), + ) + assert budget_usage.unpriced_requests == (pricing_resolution == "unknown") + assert budget_usage.cost_usd == ( + 0 if pricing_resolution == "unknown" else Decimal("0.00004") + ) + assert budget_usage.top_contributors[0]["cost_complete"] is ( + pricing_resolution != "unknown" + ) + start = datetime.now(timezone.utc) - timedelta(days=1) + end = datetime.now(timezone.utc) + timedelta(days=1) + user = SimpleNamespace(organization_id=test_api_key.organization_id) + billing = await management.billing_usage(start, end, user, db) + expected_cost = None if pricing_resolution == "unknown" else 0.00004 + assert billing.total_cost == expected_cost + assert billing.daily_usage[0].cost_usd == expected_cost + assert billing.cost_complete is (pricing_resolution != "unknown") + assert billing.unpriced_requests == (pricing_resolution == "unknown") + breakdown = await management.billing_breakdown(start, end, "model", 100, user, db) + assert breakdown.rows[0].cost_usd == ( + None if pricing_resolution == "unknown" else Decimal("0.00004") + ) + exported = await management.export_billing_breakdown( + start, end, "model", "csv", user, db + ) + csv_row = next( + csv.DictReader(io.StringIO(bytes(exported.body).decode("utf-8-sig"))) + ) + assert csv_row["cost_usd"] == ( + "" if pricing_resolution == "unknown" else "0.00004000" + ) + assert csv_row["cost_complete"] == str(pricing_resolution != "unknown") + assert csv_row["unpriced_requests"] == str(int(pricing_resolution == "unknown")) + assert page.total == 1 + assert page.items[0].cost_complete is (pricing_resolution != "unknown") + assert page.summary.cost_complete is (pricing_resolution != "unknown") + if pricing_resolution == "unknown": + assert page.items[0].cost_usd is None + assert page.summary.unpriced_requests == 1 + assert page.summary.settled_spend_usd == 0 + else: + assert page.items[0].cost_usd == Decimal("0.00004") + assert page.summary.unpriced_requests == 0 + assert page.summary.settled_spend_usd == Decimal("0.00004") @pytest.mark.asyncio @@ -1516,10 +1804,19 @@ async def test_quota_policy_uses_shared_tier_lock_and_exclusive_key_lock(): ) session = SimpleNamespace( + scalar=AsyncMock(return_value=SimpleNamespace(role="member")), execute=AsyncMock( side_effect=[ SimpleNamespace( - scalar_one_or_none=lambda: SimpleNamespace(tier="free") + scalar_one_or_none=lambda: SimpleNamespace( + tier="free", + expires_at=None, + is_active=True, + allowed_models=None, + user_id=uuid4(), + organization_id=uuid4(), + team_id=None, + ) ), SimpleNamespace( scalar_one_or_none=lambda: SimpleNamespace( @@ -1530,7 +1827,7 @@ async def test_quota_policy_uses_shared_tier_lock_and_exclusive_key_lock(): ) ), ] - ) + ), ) policy = await AccountingPolicyLoader().quota(session, _prepared()) statements = [ diff --git a/ee/tests/gateway/pipeline/test_decisions.py b/ee/tests/gateway/pipeline/test_decisions.py new file mode 100644 index 0000000..b7ae85c --- /dev/null +++ b/ee/tests/gateway/pipeline/test_decisions.py @@ -0,0 +1,525 @@ +from datetime import datetime, timezone +from decimal import Decimal +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import pytest_asyncio +from uuid import uuid4 +from fastapi import HTTPException +from sqlalchemy import delete, select, text +from sqlalchemy.ext.asyncio import async_sessionmaker + +from shim.gateway.admission import LoopDetectionResult +from shim.gateway.contracts.context import AuditPolicy, TenantPolicy, TierPolicy +from shim.gateway.contracts.principal import AuthenticatedPrincipal +from shim.gateway.kernel.gateway_kernel import GatewayKernel +from shim.gateway.pipeline.authenticate import GatewayInvocation, GatewayRequestMetadata +from shim.gateway.pipeline.provider_execution import ProviderNonStream +from shim.gateway.request_policy import RequestPolicyContext, ResolvedRequestPolicy +from shim_enterprise.ai_act.audit_writer import write_audit_row +from shim_enterprise.ai_act.api import list_audit_logs +from shim_enterprise.ai_act.models import AIActAuditLog +from shim_enterprise.ai_act.verify import verify_chain +from shim_enterprise.billing.ledger import ( + QuotaLimitExceeded, + QuotaPolicySnapshot, + SpendLimitExceeded, + SpendPolicySnapshot, +) +from shim_enterprise.billing.models import ( + AuditIntent, + QuotaPeriodUsage, + RequestLifecycle, + SpendPeriodUsage, + UsageLedger, +) +from shim_enterprise.tenants.models import ApiKey, Organization, User +from shim_enterprise.gateway.pipeline.audit_intent import AuditIntentPersistenceError +from shim_enterprise.gateway.pipeline.quota_reservation import ( + AccountingPolicyLoader, + DurableAccountingCoordinator, + DurableUsageLifecycle, +) +from shim_enterprise.outbox.handlers import append_audit_chain +from shim_enterprise.outbox.models import OutboxEvent +from shim_enterprise.outbox.publisher import OutboxMessage, OutboxWriter + + +@pytest_asyncio.fixture +async def decision_db(async_engine): + factory = async_sessionmaker(async_engine, expire_on_commit=False, autoflush=False) + tenant_id, user_id = uuid4(), uuid4() + async with factory.begin() as setup: + await setup.execute( + text( + "INSERT INTO tier_definitions (slug, name, rate_limit_rpm, rate_limit_tpm, monthly_request_limit, monthly_token_limit, features) VALUES ('free', 'Free', 60, 15000, 1000, 1000000, '{}') ON CONFLICT (slug) DO NOTHING" + ) + ) + setup.add( + Organization( + id=tenant_id, name="Decision test", slug=f"decisions-{tenant_id}" + ) + ) + await setup.flush() + setup.add( + User(id=user_id, organization_id=tenant_id, email=f"{user_id}@example.com") + ) + await setup.flush() + key = ApiKey( + id=uuid4(), + organization_id=tenant_id, + user_id=user_id, + key_hash=uuid4().hex, + prefix="sk-decisions", + tier="free", + is_active=True, + ) + setup.add(key) + try: + async with factory() as session: + yield session, key + finally: + async with factory.begin() as cleanup: + for model in ( + AIActAuditLog, + AuditIntent, + OutboxEvent, + UsageLedger, + RequestLifecycle, + QuotaPeriodUsage, + SpendPeriodUsage, + ApiKey, + User, + ): + await cleanup.execute( + delete(model).where(model.organization_id == tenant_id) + ) + await cleanup.execute( + delete(Organization).where(Organization.id == tenant_id) + ) + + +class Scrubber: + def scrub(self, value, _config, **_kwargs): + if "private@example.com" in value: + return value.replace("private@example.com", ""), { + "": "private@example.com" + } + return value, {} + + +def gateway(db, key, case, *, audit_mode="strict"): + session_scope = async_sessionmaker(db.bind, expire_on_commit=False, autoflush=False) + + policy = ResolvedRequestPolicy( + tenant_id=key.organization_id, + tenant_policy=TenantPolicy( + allowed_providers=("google",) if case == "provider" else (), + require_zero_retention=case == "retention", + ), + tier_policy=TierPolicy(rate_limit_rpm=60), + audit_policy=AuditPolicy(mode=audit_mode), + request_policy=RequestPolicyContext("test-key-hash", "free"), + pii_config={"EMAIL_ADDRESS": True}, + ) + usage = DurableUsageLifecycle( + DurableAccountingCoordinator( + policy_loader=SimpleNamespace( + quota=AsyncMock( + return_value=QuotaPolicySnapshot( + "quota-v1", None, 0 if case == "quota" else 1000, 1000000 + ) + ), + spend=AsyncMock( + return_value=SpendPolicySnapshot( + "spend-v1", Decimal("0") if case == "spend" else None + ) + ), + ) + ), + session_scope, + ) + + async def execute(*, invocation, prepared, provider_start_callback): + await provider_start_callback() + assert "private@example.com" not in json.dumps(prepared.payload) + return ProviderNonStream( + payload={ + "model": prepared.model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 1}, + }, + request_id="provider-id", + ) + + execution = SimpleNamespace( + pii_scrubber=Scrubber(), execute=AsyncMock(side_effect=execute) + ) + kernel = GatewayKernel( + {"openai": execution}, + chain_store=SimpleNamespace(), + policy_resolver=SimpleNamespace(resolve=AsyncMock(return_value=policy)), + rate_limiter=SimpleNamespace(allow=AsyncMock(return_value=case != "rate")), + loop_detector=SimpleNamespace( + check_exact_repeat=AsyncMock(return_value=LoopDetectionResult("SAFE", 1)) + ), + loop_repeat_limit=3, + loop_window_seconds=60, + cost_tag_max_length=64, + usage=usage, + ) + model = "private@example.com" if case == "model" else "gpt-5.6-luna" + content = ( + [{"type": "image_url", "image_url": {"url": "data:private"}}] + if case == "privacy" + else "private@example.com" + if case == "mask" + else "hello" + ) + invocation = GatewayInvocation( + principal=AuthenticatedPrincipal( + actor_type="api_key", + api_key_id=key.id, + authenticated_at=datetime.now(timezone.utc), + ), + payload={ + "model": model, + "messages": [{"role": "user", "content": content}], + "max_completion_tokens": 4, + }, + provider="openai", + protocol="chat", + model=model, + stream=False, + headers={"x-provider-key": "credential-sentinel"}, + provider_credential=None, + metadata=GatewayRequestMetadata(endpoint="/v1/chat/completions"), + ) + return kernel, invocation, execution + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revoked", [False, True]) +async def test_key_access_denial_is_not_recorded_as_exhausted_quota( + decision_db, revoked +): + db, key = decision_db + stored = await db.get(ApiKey, key.id) + stored.allowed_models = ["different-model"] + stored.is_active = not revoked + await db.commit() + kernel, invocation, execution = gateway(db, key, "allow") + kernel.usage.accounting.policy_loader = AccountingPolicyLoader() + with pytest.raises(HTTPException) as caught: + await kernel._execute(invocation) + assert caught.value.status_code == (401 if revoked else 403) + execution.execute.assert_not_awaited() + event = await db.scalar( + select(OutboxEvent).where(OutboxEvent.organization_id == key.organization_id) + ) + denied = [v for v in event.payload["policy_verdicts"] if v["outcome"] == "deny"] + assert [(v["rule_id"], v["reason_code"]) for v in denied] == [ + ("api_key.access", "API_KEY_ACCESS_DENIED") + ] + assert ( + await db.scalar( + select(RequestLifecycle.id).where( + RequestLifecycle.organization_id == key.organization_id + ) + ) + is None + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("case", "rule"), + [ + ("provider", "tenant.allowed_providers"), + ("retention", "tenant.zero_retention_request"), + ("model", "gateway.model_catalog"), + ("rate", "rate.requests"), + ("quota", "quota.requests_and_tokens"), + ], +) +async def test_pre_admission_denial_is_durable_private_and_idempotent( + decision_db, monkeypatch, case, rule +): + db, test_api_key = decision_db + kernel, invocation, execution = gateway(db, test_api_key, case) + with pytest.raises((HTTPException, QuotaLimitExceeded)): + await kernel._execute(invocation) + execution.execute.assert_not_awaited() + event = ( + await db.execute( + select(OutboxEvent).where( + OutboxEvent.organization_id == test_api_key.organization_id + ) + ) + ).scalar_one() + assert event.payload["actor"] is None + assert event.payload["api_key_id"] == str(test_api_key.id) + assert event.payload["extra"]["lifecycle_status"] == "rejected" + assert event.payload["extra"]["admitted"] is False + denied = next( + verdict + for verdict in event.payload["policy_verdicts"] + if verdict["outcome"] == "deny" + ) + assert denied["rule_id"] == rule + assert denied["schema_version"] == denied["rule_version"] == 1 + assert denied["effective_at"] and denied["policy_version"] + assert "private@example.com" not in json.dumps(event.payload) + assert "credential-sentinel" not in json.dumps(event.payload) + assert ( + not ( + await db.execute( + select(UsageLedger).where( + UsageLedger.organization_id == test_api_key.organization_id + ) + ) + ) + .scalars() + .all() + ) + assert ( + not ( + await db.execute( + select(RequestLifecycle).where( + RequestLifecycle.organization_id == test_api_key.organization_id + ) + ) + ) + .scalars() + .all() + ) + assert ( + len( + ( + await db.execute( + select(AuditIntent).where( + AuditIntent.organization_id == test_api_key.organization_id + ) + ) + ) + .scalars() + .all() + ) + == 1 + ) + + async def append(context): + return await write_audit_row(context, db, deduplicate=True) + + monkeypatch.setattr( + "shim_enterprise.ai_act.audit_writer.append_audit_row_deduplicated", append + ) + message = OutboxMessage.from_event(event) + await append_audit_chain(message) + await append_audit_chain(message) + verification = await verify_chain(db, test_api_key.organization_id) + assert verification["ok"] is True + assert verification["rows_checked"] == 1 + page = await list_audit_logs( + request_id=event.aggregate_id, + event_type=None, + start=None, + end=None, + limit=50, + offset=0, + current_user=SimpleNamespace(organization_id=test_api_key.organization_id), + session=db, + ) + assert page.total == 1 + assert page.items[0].policy_verdicts == event.payload["policy_verdicts"] + assert page.items[0].api_key_id == test_api_key.id + assert page.items[0].actor is None + assert page.items[0].actor_type == "api_key" + assert page.items[0].lifecycle_status == "rejected" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("case", "expected_status", "expected_pii"), + [ + ("allow", "completed", "allow"), + ("mask", "completed", "mask"), + ("privacy", "rejected", "deny"), + ("spend", "rejected", "allow"), + ], +) +async def test_admitted_decisions_survive_accounting_finalization( + decision_db, case, expected_status, expected_pii +): + db, test_api_key = decision_db + kernel, invocation, execution = gateway(db, test_api_key, case) + if expected_status == "rejected": + with pytest.raises((HTTPException, SpendLimitExceeded)): + await kernel._execute(invocation) + execution.execute.assert_not_awaited() + else: + response = await kernel._execute(invocation) + assert response.status_code == 200 + execution.execute.assert_awaited_once() + lifecycle = ( + await db.execute( + select(RequestLifecycle).where( + RequestLifecycle.organization_id == test_api_key.organization_id + ) + ) + ).scalar_one() + event = ( + await db.execute( + select(OutboxEvent).where( + OutboxEvent.organization_id == test_api_key.organization_id, + OutboxEvent.event_type == "audit.chain_append_requested", + ) + ) + ).scalar_one() + assert ( + lifecycle.status + == event.payload["extra"]["lifecycle_status"] + == expected_status + ) + verdicts = { + verdict["rule_id"]: verdict for verdict in event.payload["policy_verdicts"] + } + assert verdicts["quota.requests_and_tokens"]["policy_version"] == "quota-v1" + assert verdicts["privacy.input"]["outcome"] == expected_pii + if case != "privacy": + assert verdicts["spend.provider_monthly"]["policy_version"] == "spend-v1" + assert verdicts["spend.provider_monthly"]["outcome"] == ( + "deny" if case == "spend" else "allow" + ) + assert event.payload["actor"] is None + assert "private@example.com" not in json.dumps(event.payload) + assert "credential-sentinel" not in json.dumps(event.payload) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["strict", "best_effort", "off"]) +async def test_denial_audit_failure_obeys_mode_without_provider_execution( + decision_db, monkeypatch, mode, caplog +): + db, test_api_key = decision_db + kernel, invocation, execution = gateway(db, test_api_key, "rate", audit_mode=mode) + append = AsyncMock(side_effect=RuntimeError("private-audit-failure-sentinel")) + monkeypatch.setattr(OutboxWriter, "append", append) + with pytest.raises( + AuditIntentPersistenceError if mode == "strict" else HTTPException + ): + await kernel._execute(invocation) + execution.execute.assert_not_awaited() + assert "private-audit-failure-sentinel" not in caplog.text + if mode == "off": + append.assert_not_awaited() + else: + append.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["strict", "best_effort"]) +async def test_admitted_denial_outbox_failure_preserves_atomicity( + decision_db, monkeypatch, mode, caplog +): + db, key = decision_db + kernel, invocation, execution = gateway(db, key, "privacy", audit_mode=mode) + append = OutboxWriter.append + + async def fail_audit(writer, session, *, organization_id, values): + if values["event_type"] == "audit.chain_append_requested": + raise RuntimeError("private-audit-failure-sentinel") + return await append( + writer, session, organization_id=organization_id, values=values + ) + + monkeypatch.setattr(OutboxWriter, "append", fail_audit) + with pytest.raises( + AuditIntentPersistenceError if mode == "strict" else HTTPException + ): + await kernel._execute(invocation) + execution.execute.assert_not_awaited() + lifecycle = ( + await db.execute( + select(RequestLifecycle).where( + RequestLifecycle.organization_id == key.organization_id + ) + ) + ).scalar_one() + assert lifecycle.status == "accepted" + assert lifecycle.reconciliation_due_at is not None + assert lifecycle.reconciled_at is None + ledger = ( + ( + await db.execute( + select(UsageLedger).where( + UsageLedger.organization_id == key.organization_id + ) + ) + ) + .scalars() + .all() + ) + assert [entry.event_type for entry in ledger] == ["quota_reservation"] + assert "private-audit-failure-sentinel" not in caplog.text + + +@pytest.mark.asyncio +async def test_uncertain_admission_acknowledgement_reuses_durable_lifecycle( + decision_db, monkeypatch +): + db, key = decision_db + kernel, invocation, execution = gateway(db, key, "allow") + admit = kernel.usage.admit + + async def lost_ack(prepared, admission): + await admit(prepared, admission) + raise RuntimeError("lost-ack-private") + + monkeypatch.setattr(kernel.usage, "admit", lost_ack) + with pytest.raises(RuntimeError, match="lost-ack-private"): + await kernel._execute(invocation) + execution.execute.assert_not_awaited() + lifecycle = ( + await db.execute( + select(RequestLifecycle).where( + RequestLifecycle.organization_id == key.organization_id + ) + ) + ).scalar_one() + event = ( + await db.execute( + select(OutboxEvent).where( + OutboxEvent.organization_id == key.organization_id, + OutboxEvent.event_type == "audit.chain_append_requested", + ) + ) + ).scalar_one() + assert lifecycle.status == event.payload["extra"]["lifecycle_status"] == "failed" + assert ( + event.payload["policy_verdicts"][-1]["reason_code"] == "ADMISSION_UNAVAILABLE" + ) + assert "lost-ack-private" not in json.dumps(event.payload) + ledger = ( + ( + await db.execute( + select(UsageLedger).where( + UsageLedger.organization_id == key.organization_id + ) + ) + ) + .scalars() + .all() + ) + assert {entry.event_type for entry in ledger} == { + "quota_reservation", + "quota_refund", + } diff --git a/ee/tests/scripts/test_offline_bundle.py b/ee/tests/scripts/test_offline_bundle.py new file mode 100644 index 0000000..5cf5576 --- /dev/null +++ b/ee/tests/scripts/test_offline_bundle.py @@ -0,0 +1,174 @@ +"""Real local signatures; Docker is replaced only for small archive fixtures.""" + +import hashlib +import io +import json +import os +from pathlib import Path +import runpy +import shutil +import subprocess +import sys +import tarfile + +import pytest + + +def image_archive(path, oci, layer_digest="b" * 64): + layers = ["sha256:" + layer_digest] + config = json.dumps( + { + "os": "linux", + "architecture": "amd64", + "rootfs": {"type": "layers", "diff_ids": layers}, + } + ).encode() + digest = "sha256:" + hashlib.sha256(config).hexdigest() + config_path = ( + ("blobs/sha256/" + digest.removeprefix("sha256:")) + if oci + else digest.removeprefix("sha256:") + ".json" + ) + files = { + config_path: config, + "manifest.json": json.dumps( + [{"Config": config_path, "RepoTags": None}] + ).encode(), + } + # containerd can synthesize an ID absent from a classic Docker archive. + loaded_id = "sha256:" + "d" * 64 + if oci: + manifest = json.dumps({"config": {"digest": digest}}).encode() + loaded_id = "sha256:" + hashlib.sha256(manifest).hexdigest() + files["blobs/sha256/" + loaded_id.removeprefix("sha256:")] = manifest + files["index.json"] = json.dumps( + {"manifests": [{"digest": loaded_id}]} + ).encode() + with tarfile.open(path, "w") as archive: + for name, data in files.items(): + member = tarfile.TarInfo(name) + member.size = len(data) + archive.addfile(member, io.BytesIO(data)) + return digest, loaded_id + + +@pytest.mark.skipif(shutil.which("cosign") is None, reason="Install cosign 3.1.3") +@pytest.mark.parametrize("oci", [False, True]) +def test_offline_bundle_rejects_tampering_before_import(tmp_path, monkeypatch, oci): + monkeypatch.setenv("COSIGN_PASSWORD", "") + monkeypatch.setenv("HTTPS_PROXY", "http://127.0.0.1:9") + monkeypatch.setenv("HTTP_PROXY", "http://127.0.0.1:9") + monkeypatch.setenv("NO_PROXY", "") + for name in ("trusted", "wrong"): + subprocess.run( + [ + "cosign", + "generate-key-pair", + "--output-key-prefix", + str(tmp_path / name), + ], + check=True, + capture_output=True, + ) + script = Path(__file__).parents[2] / "scripts" / "offline_bundle.py" + scope = runpy.run_path(str(script)) + original_run = scope["run"] + docker_calls = [] + fixture = tmp_path / "fixture.tar" + config_digest, loaded_id = image_archive(fixture, oci) + archive_bytes = fixture.read_bytes() + exported_bytes = archive_bytes + source_index_id = "sha256:" + "a" * 64 + assert loaded_id != source_index_id + + def run(*args): + if args[0] != "docker": + return original_run(*args) + docker_calls.append(args) + if args[1] == "save": + Path(args[5]).write_bytes( + exported_bytes if args[-1] == loaded_id else archive_bytes + ) + elif args[1] == "load": + return f"Loaded image ID: {loaded_id}" + else: + assert args[1] == "pull", "Do not inspect non-portable daemon IDs" + return "" + + monkeypatch.setitem(scope["create"].__globals__, "run", run) + (tmp_path / "chart.tgz").write_bytes(b"test-chart") + spec = tmp_path / "spec.json" + spec.write_text( + json.dumps( + { + "bundle_id": "pilot-1", + "platform": "linux/amd64", + "images": [ + { + "name": "backend", + "version": "0.1.3", + "source": "example.test/backend@sha256:" + "a" * 64, + } + ], + "files": ["chart.tgz"], + } + ) + ) + directory = tmp_path / "bundle" + scope["create"](spec, directory, str(tmp_path / "trusted.key")) + verify = scope["verify"] + key = str(tmp_path / "trusted.pub") + assert ( + verify(directory, key, "pilot-1")["images"][0]["config_digest"] == config_digest + ) + with pytest.raises(ValueError, match="platform"): + scope["archive_config_digest"](fixture, "linux/arm64") + with pytest.raises(subprocess.CalledProcessError): + verify(directory, str(tmp_path / "wrong.pub"), "pilot-1") + with pytest.raises(ValueError, match="release ID"): + verify(directory, key, "old-release") + manifest_path = directory / "manifest.json" + manifest = manifest_path.read_bytes() + manifest_path.write_bytes(manifest + b" ") + with pytest.raises(subprocess.CalledProcessError): + verify(directory, key, "pilot-1") + manifest_path.write_bytes(manifest) + archive = directory / "backend.tar" + archive.write_bytes(b"tampered") + docker_calls.clear() + monkeypatch.setattr( + sys, + "argv", + [ + str(script), + "import", + "--directory", + str(directory), + "--key", + key, + "--expected-id", + "pilot-1", + ], + ) + with pytest.raises(SystemExit) as error: + scope["main"]() + assert error.value.code == 1 + assert docker_calls == [] + archive.write_bytes(archive_bytes) + extra = directory / "extra" + extra.write_text("unlisted") + with pytest.raises(ValueError, match="inventory"): + verify(directory, key, "pilot-1") + # Replace the extra file with a symlink, preserving the fixture for cleanup. + extra.rename(tmp_path / "extra") + extra.symlink_to(archive) + with pytest.raises(ValueError, match="symlinks"): + verify(directory, key, "pilot-1") + extra.rename(tmp_path / "extra-link") + scope["main"]() + assert any(args[1] == "load" for args in docker_calls) + image_archive(fixture, oci, layer_digest="c" * 64) + exported_bytes = fixture.read_bytes() + with pytest.raises(SystemExit): + scope["main"]() + assert os.environ["HTTPS_PROXY"] == "http://127.0.0.1:9" diff --git a/ee/tests/secrets/test_store.py b/ee/tests/secrets/test_store.py index 376bdba..24b7492 100644 --- a/ee/tests/secrets/test_store.py +++ b/ee/tests/secrets/test_store.py @@ -404,22 +404,6 @@ def test_settings_accept_csv_or_json_lists(value: str, expected: list[str]) -> N assert configured.BACKEND_CORS_ORIGINS == expected -@pytest.mark.parametrize( - "url", - ["http://store.lemonsqueezy.com/buy/variant", "/buy/variant", "https:///buy"], -) -def test_checkout_urls_require_absolute_https(url: str) -> None: - with pytest.raises(ValidationError, match="absolute HTTPS URL"): - Settings( - DATABASE_URL="postgresql+asyncpg://test:test@localhost/test", - REDIS_URL="redis://localhost:6379/0", - SECRET_KEY="test-secret-key-value", - SUPABASE_URL="https://example.supabase.co", - LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL=url, - _env_file=None, - ) - - class RotationStore: def __init__(self, rotated_reference: SecretRef) -> None: self.rotated_reference = rotated_reference diff --git a/ee/tests/secrets/test_vault.py b/ee/tests/secrets/test_vault.py new file mode 100644 index 0000000..5d92d70 --- /dev/null +++ b/ee/tests/secrets/test_vault.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import json +from uuid import uuid4 + +import httpx +import pytest + +from shim.gateway.contracts.ids import TenantId +from shim_enterprise.core.config import settings +from shim_enterprise.secrets.store import parse_secret_ref +from shim_enterprise.secrets.vault import VaultSecretStore + + +@pytest.mark.asyncio +async def test_vault_pinned_rotation_tenant_purpose_and_agent_token_reload( + tmp_path, monkeypatch +): + token_file = tmp_path / "vault-token" + token_file.write_text("first-token") + monkeypatch.setattr(settings, "VAULT_ADDR", "https://vault.internal") + monkeypatch.setattr(settings, "VAULT_TOKEN_FILE", str(token_file)) + secrets = {} + seen_tokens = [] + + def vault(request): + seen_tokens.append(request.headers["X-Vault-Token"]) + path = request.url.path + if request.method == "POST": + body = json.loads(request.content) + assert body["options"] == {"cas": 0} + secrets[path] = body["data"] + return httpx.Response(200, json={"data": {"version": 1}}) + if request.method == "GET": + assert request.url.params["version"] == "1" + if path not in secrets: + return httpx.Response(404) + return httpx.Response(200, json={"data": {"data": secrets[path]}}) + assert request.method == "PUT" and json.loads(request.content) == { + "versions": [1] + } + del secrets[path.replace("/destroy/", "/data/")] + return httpx.Response(204) + + store = VaultSecretStore(transport=httpx.MockTransport(vault)) + tenant = TenantId(uuid4()) + original = await store.put_secret(tenant, "provider:openai", "sk-one") + assert parse_secret_ref(original).version == "1" + assert "sk-one" not in original + token_file.write_text("renewed-token") + rotated = await store.rotate_secret( + tenant, original, "sk-two", expected_purpose="provider:openai" + ) + assert original != rotated + assert await store.get_secret(tenant, original) == "sk-one" + assert await store.get_secret(tenant, rotated) == "sk-two" + calls = len(seen_tokens) + with pytest.raises(ValueError, match="tenant"): + await store.get_secret(TenantId(uuid4()), original) + assert len(seen_tokens) == calls + with pytest.raises(ValueError, match="purpose"): + await store.get_secret(tenant, original, expected_purpose="wrong") + await store.delete_secret(tenant, original) + with pytest.raises(httpx.HTTPStatusError): + await store.get_secret(tenant, original) + assert seen_tokens[0] == "first-token" + assert set(seen_tokens[1:]) == {"renewed-token"} diff --git a/ee/tests/tenants/test_deployments.py b/ee/tests/tenants/test_deployments.py new file mode 100644 index 0000000..83cb6a0 --- /dev/null +++ b/ee/tests/tenants/test_deployments.py @@ -0,0 +1,809 @@ +"""Registered destinations reuse native transports, credentials and accounting.""" + +from contextlib import asynccontextmanager +import asyncio +from datetime import datetime, timedelta, timezone +from decimal import Decimal +import hashlib +import io +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock +from uuid import uuid4 +from urllib.parse import urlsplit + +import httpx +import pytest +import pytest_asyncio +from fastapi import HTTPException +from pydantic import ValidationError +from sqlalchemy import select, delete +from sqlalchemy.ext.asyncio import async_sessionmaker + +from shim.application import create_community_app +from shim.core.circuit_breaker import InMemoryCircuitBreaker +from shim.core.community_config import CommunitySettings +from shim.gateway.contracts.context import AuditPolicy, TenantPolicy, TierPolicy +from shim.gateway.contracts.principal import AuthenticatedPrincipal +from shim.gateway.request_policy import RequestPolicyContext, ResolvedRequestPolicy +from shim_enterprise.api.enterprise_deps import DatabaseGatewayAuthenticator +from shim_enterprise.api.v1.management import ( + ModelDeploymentInput, + ModelDeploymentView, + create_model_deployment, + check_model_deployment_health, +) +from shim_enterprise.billing.models import ( + UsageLedger, + RequestLifecycle, + QuotaPeriodUsage, + SpendPeriodUsage, + AuditIntent, +) +from shim_enterprise.billing.read_models import BillingReadModels +from shim_enterprise.core.config import settings +from shim_enterprise.gateway.pipeline.quota_reservation import ( + DurableAccountingCoordinator, + DurableUsageLifecycle, +) +from shim_enterprise.outbox.models import OutboxEvent +from shim_enterprise.services.gateway.enterprise import EnterpriseGatewayService +from shim_enterprise.secrets.store import ManagedProviderCredentialResolver +from shim_enterprise.tenants.deployments import ( + DeploymentResolver, + validate_deployment_url, +) +from shim_enterprise.tenants.models import ( + ModelDeployment, + Organization, + ProviderSecret, + User, + ApiKey, +) + + +@pytest_asyncio.fixture +async def db(async_engine): + factory = async_sessionmaker(async_engine, expire_on_commit=False) + async with factory() as session: + session.info["tenant_id"] = uuid4() + yield session + await session.rollback() + for model in ( + ModelDeployment, + AuditIntent, + OutboxEvent, + UsageLedger, + RequestLifecycle, + QuotaPeriodUsage, + SpendPeriodUsage, + ApiKey, + ProviderSecret, + User, + ): + await session.execute( + delete(model).where(model.organization_id == session.info["tenant_id"]) + ) + await session.execute( + delete(Organization).where(Organization.id == session.info["tenant_id"]) + ) + await session.commit() + + +@pytest_asyncio.fixture +async def test_org(db): + row = Organization( + id=db.info["tenant_id"], name="Registry test", slug=f"registry-{uuid4().hex}" + ) + db.add(row) + await db.flush() + return row + + +@pytest.fixture +def origins(monkeypatch): + monkeypatch.setattr( + settings, + "MODEL_DEPLOYMENT_ALLOWED_ORIGINS", + ["https://a.internal", "https://b.internal"], + ) + monkeypatch.setattr(settings, "MODEL_DEPLOYMENT_REQUIRED", True) + + +def test_destinations_only_use_operator_origins(origins, monkeypatch): + assert validate_deployment_url("https://a.internal/v1/") == "https://a.internal/v1" + for url in [ + "https://a.internal.evil/v1", + "http://a.internal/v1", + "https://key@a.internal/v1", + "https://a.internal/v1?key=x", + "https://a.internal/v1#x", + "https://a.internal\\@evil/v1", + ]: + with pytest.raises(ValueError): + validate_deployment_url(url) + monkeypatch.setattr( + settings, "MODEL_DEPLOYMENT_ALLOWED_ORIGINS", ["http://169.254.169.254"] + ) + with pytest.raises(ValueError, match="forbidden"): + validate_deployment_url("http://169.254.169.254/latest/meta-data") + + +async def _deployments(db, key): + rows = [] + for host, model in [("a", "custom-model-v1"), ("b", "gpt-5.6-luna")]: + secret = ProviderSecret( + id=uuid4(), + organization_id=key.organization_id, + provider="openai", + secret_ref=f"reference-{host}", + secret_backend="fernet", + secret_version="v2", + masked_key="masked", + ) + db.add(secret) + await db.flush() + row = ModelDeployment( + organization_id=key.organization_id, + alias=f"internal-{host}", + provider="openai", + upstream_model=model, + base_url=f"https://{host}.internal/v1", + provider_secret_id=secret.id, + timeout_seconds=5, + deployment_kind="internal", + declared_version="sha256:operator-declared", + owner="Platform", + enabled=True, + ) + db.add(row) + rows.append(row) + await db.flush() + return rows + + +@asynccontextmanager +async def _gateway(db, key, handler): + await db.commit() + factory = async_sessionmaker(db.bind, expire_on_commit=False) + resolver = DeploymentResolver(factory) + store = SimpleNamespace( + get_secret=AsyncMock( + side_effect=lambda tenant, ref, **kw: f"key-{str(ref)[-1]}" + ) + ) + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as outbound: + app = create_community_app( + CommunitySettings(_env_file=None, SHIM_API_KEY="sk-shim-architecture-test"), + http_client=outbound, + event_stream=io.StringIO(), + ) + async with app.router.lifespan_context(app): + app.state.gateway_authenticator = DatabaseGatewayAuthenticator(factory) + app.state.model_catalog = resolver.catalog + kernel = app.state.gateway_service.kernel + app.state.gateway_service = EnterpriseGatewayService(kernel, AsyncMock()) + kernel.policy_resolver = SimpleNamespace( + resolve=AsyncMock( + return_value=ResolvedRequestPolicy( + tenant_id=key.organization_id, + tenant_policy=TenantPolicy(), + tier_policy=TierPolicy(), + audit_policy=AuditPolicy(mode="strict"), + request_policy=RequestPolicyContext( + rate_limit_key_hash=key.key_hash, tier="free" + ), + pii_config={"email": True}, + ) + ) + ) + kernel.prepare_inference = resolver.resolve + lifecycle = DurableUsageLifecycle(DurableAccountingCoordinator(), factory) + kernel.usage = kernel.postprocessor.usage = lifecycle + circuits = { + host + ".internal": InMemoryCircuitBreaker(failure_threshold=1) + for host in ("a", "b") + } + for provider in ("openai", "anthropic"): + execution = kernel.executions[provider] + execution.credential_resolver = ManagedProviderCredentialResolver( + provider, store, factory + ) + execution.circuit_for_target = lambda url: circuits[ + urlsplit(url).hostname + ] + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app), + base_url="http://shim.test", + headers={ + "x-shim-key": "sk-shim-architecture-test", + "x-provider-key": "ignored-override", + }, + ) as client: + yield client, app, store + + +def _success(request, stream, protocol): + body = json.loads(request.content) + model = body["model"] + assert "alice@example.com" not in request.content.decode() + assert request.headers["authorization"] == f"Bearer key-{request.url.host[0]}" + if protocol == "responses": + payload = { + "id": f"resp_{uuid4().hex}", + "object": "response", + "created_at": 1, + "model": model, + "status": "completed", + "output": [], + "usage": {"input_tokens": 4, "output_tokens": 2, "total_tokens": 6}, + } + event = { + "type": "response.completed", + "sequence_number": 0, + "response": payload, + } + content = f"event: response.completed\ndata: {json.dumps(event)}\n\n" + else: + payload = { + "id": "chat-1", + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6}, + } + content = f"data: {json.dumps(payload)}\n\ndata: [DONE]\n\n" + return ( + httpx.Response(200, text=content, headers={"content-type": "text/event-stream"}) + if stream + else httpx.Response(200, json=payload) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("protocol", ["chat", "responses"]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_two_endpoints_native_json_sse_credentials_and_unpriced_cost( + db, test_api_key, origins, protocol, stream +): + rows = await _deployments(db, test_api_key) + calls = [] + + def upstream(request): + calls.append(request) + return _success(request, stream, protocol) + + async with _gateway(db, test_api_key, upstream) as (client, _, _store): + catalog = await client.get("/v1/models") + assert [record["id"] for record in catalog.json()["data"]] == [ + "internal-a", + "internal-b", + ] + for row in rows: + payload = {"model": row.alias, "stream": stream} + payload.update( + {"input": "alice@example.com"} + if protocol == "responses" + else {"messages": [{"role": "user", "content": "alice@example.com"}]} + ) + response = await client.post( + "/v1/responses" if protocol == "responses" else "/v1/chat/completions", + json=payload, + ) + assert response.status_code == 200, response.text + if stream: + assert "text/event-stream" in response.headers["content-type"] + else: + assert response.json()["model"] == row.upstream_model + assert [request.url.host for request in calls] == ["a.internal", "b.internal"] + settlements = ( + ( + await db.execute( + select(UsageLedger) + .where( + UsageLedger.organization_id == test_api_key.organization_id, + UsageLedger.event_type == "spend_settlement", + ) + .order_by(UsageLedger.requested_model) + ) + ) + .scalars() + .all() + ) + assert len(settlements) == 2 + assert settlements[0].cost_usd == 0 + assert settlements[0].event_metadata["pricing"]["pricing_resolution"] == "unknown" + assert "input_per_million" not in settlements[0].event_metadata["pricing"] + assert ( + settlements[1].cost_usd > 0 and settlements[1].provider_model == "gpt-5.6-luna" + ) + now = datetime.now(timezone.utc) + daily = await BillingReadModels().daily_usage( + db, + tenant_id=test_api_key.organization_id, + start_at=now - timedelta(days=1), + end_at=now + timedelta(days=1), + ) + assert sum(row.unpriced_requests for row in daily) == 1 + assert any(not row.as_public_record()["cost_complete"] for row in daily) + + +@pytest.mark.asyncio +async def test_registry_denials_are_audited_and_money_cap_cannot_be_bypassed( + db, test_api_key, origins +): + rows = await _deployments(db, test_api_key) + secret = await db.get(ProviderSecret, rows[0].provider_secret_id) + secret.monthly_limit_usd = Decimal("1") + await db.flush() + calls = [] + async with _gateway(db, test_api_key, lambda request: calls.append(request)) as ( + client, + _, + _, + ): + omitted = await client.post("/v1/responses", json={"input": "hello"}) + assert omitted.status_code == 403 + payload = {"model": "missing", "messages": []} + assert ( + await client.post("/v1/chat/completions", json=payload) + ).status_code == 403 + capped = await client.post( + "/v1/chat/completions", json={**payload, "model": "internal-a"} + ) + assert capped.status_code == 403, capped.text + assert capped.json()["error"]["code"] == "MODEL_PRICE_UNKNOWN" + test_api_key.allowed_models = [] + await db.commit() + denied = await client.post( + "/v1/chat/completions", json={**payload, "model": "internal-b"} + ) + assert denied.status_code == 403 + assert (await client.get("/v1/models")).json()["data"] == [] + assert calls == [] + intents = ( + ( + await db.execute( + select(OutboxEvent).where( + OutboxEvent.organization_id == test_api_key.organization_id + ) + ) + ) + .scalars() + .all() + ) + reasons = { + verdict["reason_code"] + for event in intents + for verdict in event.payload.get("policy_verdicts", []) + } + assert { + "MODEL_NOT_REGISTERED", + "MODEL_PRICE_UNKNOWN", + "MODEL_NOT_ALLOWED", + } <= reasons + + +@pytest.mark.asyncio +async def test_unhealthy_endpoint_does_not_retry_or_open_other_endpoint_circuit( + db, test_api_key, origins +): + await _deployments(db, test_api_key) + calls = [] + + def upstream(request): + calls.append(request) + if request.url.host == "b.internal": + return httpx.Response( + 503, json={"error": {"message": "key-b private provider detail"}} + ) + return _success(request, False, "chat") + + async with _gateway(db, test_api_key, upstream) as (client, _, _): + for alias, expected in [ + ("internal-b", 503), + ("internal-b", 503), + ("internal-a", 200), + ]: + response = await client.post( + "/v1/chat/completions", + json={ + "model": alias, + "messages": [{"role": "user", "content": "hello"}], + }, + ) + assert response.status_code == expected, response.text + assert "private provider detail" not in response.text + assert [request.url.host for request in calls] == ["b.internal", "a.internal"] + + +@pytest.mark.asyncio +async def test_registry_tenant_isolation_and_credential_fk(db, test_api_key, origins): + rows = await _deployments(db, test_api_key) + other = Organization(name="Other", slug=f"other-{uuid4().hex}") + db.add(other) + await db.flush() + owner = User( + id=uuid4(), + email=f"{uuid4().hex}@example.com", + organization_id=other.id, + is_active=True, + role="owner", + ) + db.add(owner) + await db.flush() + key = ApiKey( + user_id=owner.id, + organization_id=other.id, + key_hash=hashlib.sha256(uuid4().bytes).hexdigest(), + prefix="other", + tier="free", + is_active=True, + ) + db.add(key) + await db.flush() + resolver = DeploymentResolver( + async_sessionmaker(await db.connection(), expire_on_commit=False) + ) + principal = AuthenticatedPrincipal( + actor_type="api_key", + api_key_id=key.id, + authenticated_at=datetime.now(timezone.utc), + ) + assert await resolver.catalog(principal, "openai") == [] + payload = ModelDeploymentInput( + alias="other-alias", + provider="openai", + upstream_model="local", + base_url="https://a.internal/v1", + provider_secret_id=rows[0].provider_secret_id, + deployment_kind="internal", + declared_version="1", + owner="Platform", + ) + with pytest.raises(HTTPException) as error: + await create_model_deployment(payload, owner, db) + assert error.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_health_headers_are_bounded_and_output_view_survives_policy_change( + db, test_api_key, test_user_with_org, origins, monkeypatch +): + rows = await _deployments(db, test_api_key) + store = SimpleNamespace(get_secret=AsyncMock(return_value="health-key")) + monkeypatch.setattr( + "shim_enterprise.api.v1.management.get_secret_store", lambda: store + ) + + def upstream(request): + assert request.method == "GET" and request.url.path == "/v1/models" + assert not db.in_transaction() + return httpx.Response(503) + + async with httpx.AsyncClient(transport=httpx.MockTransport(upstream)) as client: + request = SimpleNamespace( + app=SimpleNamespace(state=SimpleNamespace(http_client=client)) + ) + result = await check_model_deployment_health( + rows[0].id, request, test_user_with_org, db + ) + assert result.health == "unhealthy" and result.health_checked_at is not None + monkeypatch.setattr(settings, "MODEL_DEPLOYMENT_ALLOWED_ORIGINS", []) + assert ( + ModelDeploymentView.model_validate(result).base_url == "https://a.internal/v1" + ) + with pytest.raises(ValidationError): + ModelDeploymentInput.model_validate( + ModelDeploymentView.model_validate(result).model_dump( + exclude={ + "id", + "health", + "health_checked_at", + "created_at", + "updated_at", + } + ) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["json", "sse", "count_tokens"]) +async def test_registered_anthropic_native_messages_and_nonbillable_token_count( + db, test_api_key, origins, operation +): + rows = await _deployments(db, test_api_key) + row = rows[0] + secret = await db.get(ProviderSecret, row.provider_secret_id) + secret.provider = row.provider = "anthropic" + row.base_url = "https://a.internal" + row.upstream_model = "private-claude" + await db.flush() + calls = [] + + def upstream(request): + calls.append(request) + assert request.headers["x-api-key"] == "key-a" + assert json.loads(request.content)["model"] == "private-claude" + assert "alice@example.com" not in request.content.decode() + if operation == "count_tokens": + assert request.url.path == "/v1/messages/count_tokens" + return httpx.Response( + 200, json={"input_tokens": 17}, headers={"request-id": "count-1"} + ) + assert request.url.path == "/v1/messages" + message = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "private-claude", + "content": [], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 4, "output_tokens": 2}, + } + if operation == "json": + return httpx.Response(200, json=message) + events = [ + { + "type": "message_start", + "message": { + **message, + "stop_reason": None, + "usage": {"input_tokens": 4, "output_tokens": 0}, + }, + }, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, + ] + return httpx.Response( + 200, + text="".join( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" + for event in events + ), + headers={"content-type": "text/event-stream"}, + ) + + async with _gateway(db, test_api_key, upstream) as (client, _, _): + payload = { + "model": row.alias, + "messages": [{"role": "user", "content": "alice@example.com"}], + } + if operation != "count_tokens": + payload.update(max_tokens=10, stream=operation == "sse") + response = await client.post( + "/v1/messages/count_tokens" + if operation == "count_tokens" + else "/v1/messages", + json=payload, + ) + assert response.status_code == 200, response.text + if operation == "sse": + assert "event: message_stop" in response.text + elif operation == "count_tokens": + assert response.json() == {"input_tokens": 17} + assert response.headers["request-id"] == "count-1" + else: + assert response.json()["model"] == "private-claude" + assert len(calls) == 1 + ledger = ( + await db.scalars( + select(UsageLedger).where( + UsageLedger.organization_id == test_api_key.organization_id + ) + ) + ).all() + if operation == "count_tokens": + assert ledger == [] + assert ( + await db.scalars( + select(RequestLifecycle).where( + RequestLifecycle.organization_id == test_api_key.organization_id + ) + ) + ).all() == [] + intent = ( + await db.scalars( + select(AuditIntent).where( + AuditIntent.organization_id == test_api_key.organization_id, + AuditIntent.event_type == "completion", + ) + ) + ).one() + assert intent.usage_summary == {"input_tokens": 17, "billable_executions": 0} + event = await db.get(OutboxEvent, intent.outbox_event_id) + assert event.payload["event_type"] == "token_count" + assert event.payload["extra"]["billable_execution"] is False + assert event.payload["extra"]["deployment_id"] == str(row.id) + assert event.payload["extra"]["deployment_kind"] == "internal" + assert "alice@example.com" not in json.dumps(event.payload) + else: + assert ( + len([entry for entry in ledger if entry.event_type == "spend_settlement"]) + == 1 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["preflight", "completion"]) +async def test_strict_token_count_audit_failure_never_creates_billable_usage( + db, test_api_key, origins, monkeypatch, phase +): + from shim_enterprise.gateway.pipeline.audit_intent import AuditIntentRepository + + rows = await _deployments(db, test_api_key) + secret = await db.get(ProviderSecret, rows[0].provider_secret_id) + secret.provider = rows[0].provider = "anthropic" + rows[0].base_url = "https://a.internal" + await db.flush() + original_create = AuditIntentRepository.create + + async def create(session, *, organization_id, values): + if values["event_type"] == phase: + raise RuntimeError("audit storage unavailable") + return await original_create( + session, organization_id=organization_id, values=values + ) + + monkeypatch.setattr(AuditIntentRepository, "create", create) + calls = [] + + def upstream(request): + calls.append(request) + return httpx.Response(200, json={"input_tokens": 17}) + + async with _gateway(db, test_api_key, upstream) as (client, _, _): + response = await client.post( + "/v1/messages/count_tokens", json={"model": rows[0].alias, "messages": []} + ) + assert response.status_code == 503 + assert len(calls) == (0 if phase == "preflight" else 1) + assert ( + await db.scalars( + select(UsageLedger).where( + UsageLedger.organization_id == test_api_key.organization_id + ) + ) + ).all() == [] + assert ( + await db.scalars( + select(RequestLifecycle).where( + RequestLifecycle.organization_id == test_api_key.organization_id + ) + ) + ).all() == [] + + +@pytest.mark.asyncio +async def test_registry_management_requires_admin_and_audits_configuration( + db, test_api_key, test_user_with_org, origins +): + from fastapi import FastAPI + from shim_enterprise.api.enterprise_deps import get_current_user + from shim_enterprise.api.v1.management import router + from shim_enterprise.core.database import get_db + + rows = await _deployments(db, test_api_key) + app = FastAPI() + app.include_router(router) + app.dependency_overrides[get_current_user] = lambda: test_user_with_org + app.dependency_overrides[get_db] = lambda: db + payload = { + "alias": "registered-model", + "provider": "openai", + "upstream_model": "custom-v1", + "base_url": "https://a.internal/v1", + "provider_secret_id": str(rows[0].provider_secret_id), + "deployment_kind": "internal", + "declared_version": "sha256:v1", + "owner": "Platform", + } + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app), base_url="http://test" + ) as client: + assert (await client.get("/model-deployments")).status_code == 403 + assert ( + await client.post("/model-deployments", json=payload) + ).status_code == 403 + test_user_with_org.role = "owner" + created = await client.post("/model-deployments", json=payload) + assert created.status_code == 201, created.text + deployment_id = created.json()["id"] + updated = await client.put( + f"/model-deployments/{deployment_id}", + json={**payload, "declared_version": "sha256:v2", "enabled": False}, + ) + assert updated.status_code == 200 and updated.json()["id"] == deployment_id + invalid = await client.post( + "/api-keys", + json={"name": "unknown model", "allowed_models": ["missing-alias"]}, + ) + assert invalid.status_code == 422 + disabled = await client.post( + "/api-keys", + json={"name": "disabled model", "allowed_models": ["registered-model"]}, + ) + assert disabled.status_code == 422 + known = await client.post( + "/api-keys", json={"name": "known model", "allowed_models": ["internal-a"]} + ) + assert known.status_code == 200, known.text + assert known.json()["allowed_models"] == ["internal-a"] + events = ( + await db.scalars( + select(OutboxEvent).where( + OutboxEvent.organization_id == test_api_key.organization_id + ) + ) + ).all() + configs = { + event.payload["endpoint"]: event.payload["extra"]["configuration"] + for event in events + if "configuration" in event.payload.get("extra", {}) + } + assert configs["tenant.model_deployment_created"]["declared_version"] == "sha256:v1" + assert configs["tenant.model_deployment_updated"]["declared_version"] == "sha256:v2" + assert configs["tenant.model_deployment_created"]["provider_secret_id"] == str( + rows[0].provider_secret_id + ) + assert all("key" not in config for config in configs.values()) + + +@pytest.mark.asyncio +async def test_disabled_registry_alias_cannot_fall_back_to_public_catalog( + db, test_api_key, origins, monkeypatch +): + rows = await _deployments(db, test_api_key) + rows[1].alias, rows[1].enabled = "gpt-5.6-luna", False + await db.flush() + monkeypatch.setattr(settings, "MODEL_DEPLOYMENT_REQUIRED", False) + calls = [] + async with _gateway(db, test_api_key, lambda request: calls.append(request)) as ( + client, + _, + _, + ): + catalog = await client.get("/v1/models") + assert "gpt-5.6-luna" not in {row["id"] for row in catalog.json()["data"]} + denied = await client.post( + "/v1/chat/completions", json={"model": "gpt-5.6-luna", "messages": []} + ) + assert denied.status_code == 403 + assert calls == [] + + +@pytest.mark.asyncio +async def test_health_probe_has_a_wall_clock_deadline( + db, test_api_key, test_user_with_org, origins, monkeypatch +): + rows = await _deployments(db, test_api_key) + real_timeout = asyncio.timeout + monkeypatch.setattr( + "shim_enterprise.api.v1.management.asyncio.timeout", + lambda seconds: real_timeout(0.01), + ) + + async def delayed_secret(*args, **kwargs): + await asyncio.sleep(1) + pytest.fail("health deadline was not enforced") + + monkeypatch.setattr( + "shim_enterprise.api.v1.management.get_secret_store", + lambda: SimpleNamespace(get_secret=delayed_secret), + ) + result = await check_model_deployment_health( + rows[0].id, SimpleNamespace(), test_user_with_org, db + ) + assert result.health == "unhealthy" diff --git a/ee/tests/tenants/test_oidc.py b/ee/tests/tenants/test_oidc.py new file mode 100644 index 0000000..755e329 --- /dev/null +++ b/ee/tests/tenants/test_oidc.py @@ -0,0 +1,369 @@ +from __future__ import annotations + +import base64 +import hashlib +import json +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock +from urllib.parse import parse_qs, urlsplit +from uuid import uuid4 + +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import FastAPI, HTTPException, Request +import httpx +import jwt +import pytest +from pydantic import ValidationError + +from shim_enterprise.cache.redis_index import CacheService +from shim_enterprise.core.config import Settings, settings +from shim_enterprise.core.database import get_db +from shim_enterprise.tenants import oidc + + +@pytest.fixture +def oidc_config(monkeypatch): + values = dict( + AUTH_MODE="oidc", + OIDC_ISSUER_URL="https://identity.internal/realm", + OIDC_CLIENT_ID="dashboard", + OIDC_CLIENT_SECRET="customer-secret", + OIDC_API_AUDIENCE="shim-api", + OIDC_REDIRECT_URI="https://shim.internal/api/v1/auth/callback", + DASHBOARD_ORIGIN="https://shim.internal", + OIDC_ORGANIZATION_ID=uuid4(), + OIDC_GROUP_ROLE_MAP={"/shim/owners": "owner", "/shim/members": "member"}, + OIDC_REVALIDATE_SECONDS=60, + ) + for name, value in values.items(): + monkeypatch.setattr(settings, name, value) + return values + + +def test_oidc_configuration_needs_no_supabase(oidc_config): + values = dict( + oidc_config, + DATABASE_URL="postgresql+asyncpg://test:test@localhost/test", + REDIS_URL="redis://localhost/0", + SECRET_KEY="test-secret-key-value", + SUPABASE_URL=None, + SUPABASE_KEY=None, + _env_file=None, + ) + assert Settings(**values).SUPABASE_URL is None + for invalid in ( + {"OIDC_GROUP_ROLE_MAP": {}}, + {"OIDC_TEAM_GROUP_MAP": {"group": {"team_id": "invalid", "role": "member"}}}, + {"OIDC_TEAM_GROUP_MAP": {"group": {"team_id": str(uuid4()), "role": "owner"}}}, + {"OIDC_API_AUDIENCE": "dashboard"}, + {"OIDC_REDIRECT_URI": "https://other.internal/api/v1/auth/callback"}, + {"AUTH_MODE": "supabase"}, + ): + with pytest.raises(ValidationError): + Settings(**(values | invalid)) + + +@pytest.mark.asyncio +async def test_identity_binding_role_removal_and_email_collision( + db, test_org, test_user_with_org, test_tier, monkeypatch, oidc_config +): + monkeypatch.setattr(settings, "OIDC_ORGANIZATION_ID", test_org.id) + claims = dict( + iss=settings.OIDC_ISSUER_URL, + sub="subject", + email=f"{uuid4()}@example.com", + email_verified=True, + groups=["/shim/owners"], + ) + user = await oidc.synchronize_user(db, claims) + assert user.organization_id == test_org.id + assert user.role == "owner" + assert (user.oidc_issuer, user.oidc_subject) == (claims["iss"], "subject") + same_user = await oidc.synchronize_user( + db, claims | {"groups": ["/shim/members"], "email": "changed@example.com"} + ) + assert same_user.id == user.id and same_user.role == "member" + for changed in ( + {"groups": []}, + {"groups": "/shim/owners"}, + {"iss": "https://other.internal"}, + {"sub": "collision", "email": test_user_with_org.email}, + {"sub": "new", "email_verified": False}, + ): + with pytest.raises(HTTPException): + await oidc.synchronize_user(db, claims | changed) + from shim_enterprise.api.v1.management import remove_team_member + from shim_enterprise.tenants.service import authenticate_api_key, create_api_key + + plaintext, _ = await create_api_key(db, user_id=user.id, name="workload") + with pytest.raises(HTTPException): + await oidc.synchronize_user(db, claims | {"groups": []}) + assert await authenticate_api_key(db, plaintext) is not None + test_user_with_org.role = "owner" + await remove_team_member(user.id, test_user_with_org, db) + assert await authenticate_api_key(db, plaintext) is None + with pytest.raises(HTTPException, match="inactive"): + await oidc.synchronize_user(db, claims) + + +@pytest.mark.asyncio +async def test_signed_api_tokens_reject_issuer_audience_expiry_signature_and_long_lifetime( + oidc_config, +): + private = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwk = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private.public_key())) | { + "kid": "test-key" + } + client = SimpleNamespace( + load_server_metadata=AsyncMock( + return_value={"issuer": settings.OIDC_ISSUER_URL} + ), + fetch_jwk_set=AsyncMock(return_value={"keys": [jwk]}), + ) + request = Request( + {"type": "http", "app": SimpleNamespace(state=SimpleNamespace(oidc=client))} + ) + now = int(time.time()) + claims = { + "iss": settings.OIDC_ISSUER_URL, + "sub": "user", + "aud": "shim-api", + "iat": now, + "exp": now + 300, + } + + def signed(data): + return jwt.encode(data, private, algorithm="RS256", headers={"kid": "test-key"}) + + assert (await oidc.access_token_claims(request, signed(claims)))["sub"] == "user" + for data in ( + claims | {"iss": "https://attacker.internal"}, + claims | {"aud": "dashboard"}, + claims | {"exp": now - 1}, + claims | {"exp": now + 301}, + {key: value for key, value in claims.items() if key != "exp"}, + ): + with pytest.raises(HTTPException) as error: + await oidc.access_token_claims(request, signed(data)) + assert error.value.status_code == 401 + invalid = jwt.encode( + claims, + rsa.generate_private_key(public_exponent=65537, key_size=2048), + algorithm="RS256", + headers={"kid": "test-key"}, + ) + with pytest.raises(HTTPException): + await oidc.access_token_claims(request, invalid) + + +@pytest.mark.asyncio +async def test_authorization_code_pkce_cookie_refresh_csrf_and_logout( + db, test_org, monkeypatch, oidc_config +): + monkeypatch.setattr(settings, "OIDC_ORGANIZATION_ID", test_org.id) + monkeypatch.setattr(settings, "ENVIRONMENT", "production") + private = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwk = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private.public_key())) | { + "kid": "test-key" + } + login_parameters = {} + groups = ["/shim/owners"] + subject = uuid4().hex + + def provider(request): + if request.url.path.endswith("openid-configuration"): + return httpx.Response( + 200, + json={ + "issuer": settings.OIDC_ISSUER_URL, + "authorization_endpoint": "https://identity.internal/authorize", + "token_endpoint": "https://identity.internal/token", + "jwks_uri": "https://identity.internal/jwks", + "id_token_signing_alg_values_supported": ["RS256", "HS256"], + "end_session_endpoint": "https://identity.internal/logout", + }, + ) + if request.url.path == "/jwks": + return httpx.Response(200, json={"keys": [jwk]}) + assert request.url.path == "/token" + parameters = parse_qs(request.content.decode()) + now = int(time.time()) + claims = dict( + iss=settings.OIDC_ISSUER_URL, + sub=subject, + aud="dashboard", + iat=now, + exp=now + 300, + email=f"{subject}@example.com", + email_verified=True, + groups=groups, + ) + if parameters["grant_type"] == ["authorization_code"]: + verifier = parameters["code_verifier"][0] + challenge = ( + base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()) + .decode() + .rstrip("=") + ) + assert challenge == login_parameters["code_challenge"][0] + claims["nonce"] = login_parameters["nonce"][0] + else: + assert parameters["grant_type"] == ["refresh_token"] + return httpx.Response( + 200, + json={ + "access_token": "provider-access-token", + "refresh_token": "provider-refresh-token", + "token_type": "Bearer", + "expires_in": 300, + "id_token": jwt.encode( + claims, private, algorithm="RS256", headers={"kid": "test-key"} + ), + }, + ) + + cache = CacheService() + await cache.connect() + app = FastAPI() + app.state.cache = cache + oidc.install_oidc(app) + app.state.oidc.client_kwargs["transport"] = httpx.MockTransport(provider) + app.include_router(oidc.router, prefix="/api/v1") + + async def database(): + yield db + + app.dependency_overrides[get_db] = database + key = None + try: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="https://shim.internal" + ) as browser: + for path in ("//evil.example", "/\\evil.example", "https://evil.example"): + assert ( + await browser.get("/api/v1/auth/login", params={"next": path}) + ).status_code == 400 + response = await browser.get( + "/api/v1/auth/login", params={"next": "/dashboard/workspace/settings"} + ) + login_parameters.update( + parse_qs(urlsplit(response.headers["location"]).query) + ) + assert login_parameters["code_challenge_method"] == ["S256"] + assert ( + await browser.get( + "/api/v1/auth/callback", params={"code": "one", "state": "forged"} + ) + ).status_code == 401 + response = await browser.get( + "/api/v1/auth/login", params={"next": "/dashboard/workspace/settings"} + ) + login_parameters.update( + parse_qs(urlsplit(response.headers["location"]).query) + ) + response = await browser.get( + "/api/v1/auth/callback", + params={"code": "one", "state": login_parameters["state"][0]}, + ) + assert response.status_code == 303, response.text + assert response.headers["location"].endswith( + "/dashboard/workspace/settings" + ) + cookie = response.headers["set-cookie"] + assert ( + "HttpOnly" in cookie and "Secure" in cookie and "SameSite=lax" in cookie + ) + assert ( + "provider-access-token" not in cookie + and "provider-refresh-token" not in cookie + ) + key = oidc._session_key(browser.cookies[oidc.SESSION_COOKIE]) + stored = await cache.redis.get(key) + assert "provider-refresh-token" not in stored + assert (await browser.get("/api/v1/auth/session")).json()["user"][ + "role" + ] == "owner" + assert ( + await browser.post( + "/api/v1/auth/logout", headers={"origin": "https://evil.example"} + ) + ).status_code == 403 + groups[:] = ["/shim/members"] + data = json.loads(oidc._cipher().decrypt(stored.encode())) + data["checked_at"] -= 61 + await cache.redis.set( + key, oidc._cipher().encrypt(json.dumps(data).encode()).decode(), ex=300 + ) + assert (await browser.get("/api/v1/auth/session")).json()["user"][ + "role" + ] == "member" + groups.clear() + stored = await cache.redis.get(key) + data = json.loads(oidc._cipher().decrypt(stored.encode())) + data["checked_at"] -= 61 + await cache.redis.set( + key, oidc._cipher().encrypt(json.dumps(data).encode()).decode(), ex=300 + ) + assert (await browser.get("/api/v1/auth/session")).status_code == 403 + response = await browser.post( + "/api/v1/auth/logout", headers={"origin": "https://shim.internal"} + ) + assert response.status_code == 200 + assert response.json()["logout_url"].startswith( + "https://identity.internal/logout?" + ) + assert await cache.redis.get(key) is None + assert (await browser.get("/api/v1/auth/session")).status_code == 401 + browser.cookies.set(oidc.SESSION_COOKIE, "revoked-session") + monkeypatch.setattr( + oidc, + "_client", + AsyncMock(side_effect=httpx.ConnectError("IdP unavailable")), + ) + response = await browser.post( + "/api/v1/auth/logout", headers={"origin": "https://shim.internal"} + ) + assert response.status_code == 200 and response.json() == { + "logout_url": "/login" + } + assert "Max-Age=0" in response.headers["set-cookie"] + finally: + if key: + await cache.redis.delete(key, key + ":refresh") + await cache.close() + + +@pytest.mark.asyncio +async def test_login_projects_and_removes_oidc_team_membership( + db, test_org, monkeypatch, oidc_config +): + from sqlalchemy import select + from shim_enterprise.tenants.models import Team, TeamMembership + + team = Team(id=uuid4(), organization_id=test_org.id, name="OIDC team") + db.add(team) + await db.flush() + monkeypatch.setattr(settings, "OIDC_ORGANIZATION_ID", test_org.id) + monkeypatch.setattr( + settings, + "OIDC_TEAM_GROUP_MAP", + {"/shim/team": {"team_id": str(team.id), "role": "team_admin"}}, + ) + claims = dict( + iss=settings.OIDC_ISSUER_URL, + sub="team-subject", + email=f"{uuid4()}@example.com", + email_verified=True, + groups=["/shim/members", "/shim/team"], + ) + user = await oidc.synchronize_user(db, claims) + membership = await db.scalar( + select(TeamMembership).where(TeamMembership.user_id == user.id) + ) + assert membership.role == "team_admin" and membership.source == "oidc" + await oidc.synchronize_user(db, claims | {"groups": ["/shim/members"]}) + assert ( + await db.scalar(select(TeamMembership).where(TeamMembership.user_id == user.id)) + is None + ) diff --git a/ee/tests/tenants/test_plans.py b/ee/tests/tenants/test_plans.py new file mode 100644 index 0000000..fb6d818 --- /dev/null +++ b/ee/tests/tenants/test_plans.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +from datetime import date, datetime, timedelta, timezone +from decimal import Decimal +from types import SimpleNamespace +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest +from sqlalchemy import select +from sqlalchemy.dialects import postgresql + +from shim_enterprise.api.v1 import management +from shim_enterprise.billing.models import QuotaPeriodUsage, UsageLedger +from shim_enterprise.tenants.models import ( + BillingWebhookReceipt, + Organization, + OrganizationPIIConfig, + User, +) +from shim_enterprise.tenants.plans import ( + activate_organization_plan, + create_organization_plan, +) +from shim_enterprise.tenants.service import ( + authenticate_api_key, + create_api_key, + ensure_privacy_defaults, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target_tier", ["free", "managed", "agency", "enterprise"]) +async def test_operator_transition_preserves_access_accounting_and_history( + db, test_org, test_user_with_org, target_tier: str +) -> None: + plaintext, api_key = await create_api_key( + db, user_id=test_user_with_org.id, name="Existing key" + ) + _, revoked_key = await create_api_key( + db, user_id=test_user_with_org.id, name="Revoked key" + ) + revoked_key.is_active = False + test_org.tier = api_key.tier = "managed" + test_org.billing_status = "active" + test_org.billing_source = "lemonsqueezy" + legacy = { + "external_customer_id": f"customer-{uuid4()}", + "external_subscription_id": f"subscription-{uuid4()}", + "billing_variant_id": "legacy-variant", + "current_period_end": datetime(2026, 9, 1, tzinfo=timezone.utc), + "cancel_at_period_end": True, + "customer_portal_url": "https://example.test/historical-portal", + } + for name, value in legacy.items(): + setattr(test_org, name, value) + receipt = BillingWebhookReceipt( + organization_id=test_org.id, + payload_digest=uuid4().hex, + event_name="subscription_created", + external_subscription_id=test_org.external_subscription_id, + event_at=datetime.now(timezone.utc), + ) + ledger = UsageLedger( + organization_id=test_org.id, + api_key_id=api_key.id, + request_id=f"test-{uuid4()}", + requested_model="gpt-4o-mini", + event_type="adjustment_credit", + idempotency_key=f"test-{uuid4()}", + cost_usd=Decimal("1.25"), + ) + quota = QuotaPeriodUsage( + organization_id=test_org.id, + api_key_id=api_key.id, + period_type="monthly", + period_start=date(2026, 9, 1), + period_end=date(2026, 10, 1), + settled_requests=7, + settled_tokens=900, + reserved_requests=1, + reserved_tokens=100, + ) + other_org = Organization(name="Other tenant", slug=f"other-{uuid4()}") + db.add_all([receipt, ledger, quota, other_org]) + await db.flush() + + await activate_organization_plan(db, test_org.id, target_tier) + await db.commit() + for row in (test_org, api_key, revoked_key, receipt, ledger, quota, other_org): + await db.refresh(row) + + assert test_org.tier == api_key.tier == target_tier + assert test_org.billing_status == ("free" if target_tier == "free" else "active") + assert test_org.billing_source == "operator" + assert await authenticate_api_key(db, plaintext) is api_key + assert revoked_key.is_active is False and revoked_key.tier == "free" + assert other_org.tier == "free" + assert {name: getattr(test_org, name) for name in legacy} == legacy + assert receipt.external_subscription_id == legacy["external_subscription_id"] + assert ledger.cost_usd == Decimal("1.25") + assert (quota.settled_requests, quota.settled_tokens) == (7, 900) + assert (quota.reserved_requests, quota.reserved_tokens) == (1, 100) + _, new_key = await create_api_key(db, user_id=test_user_with_org.id, name="New key") + assert new_key.tier == target_tier + view = await management.get_subscription(test_user_with_org, db) + assert view.plan == target_tier + assert set(view.model_dump()) == {"plan", "status", "source", "entitlements"} + + +@pytest.mark.asyncio +async def test_new_customer_provisioning_and_invalid_plan(db, test_org) -> None: + with pytest.raises(ValueError, match="Unknown tier"): + await activate_organization_plan(db, test_org.id, "missing-tier") + assert test_org.tier == "free" + with pytest.raises(ValueError, match="Organization not found"): + await activate_organization_plan(db, uuid4(), "enterprise") + with pytest.raises(ValueError, match="Organization name"): + await create_organization_plan(db, " ", "enterprise") + + created = await create_organization_plan(db, " Pilot bank ", "enterprise") + await db.commit() + assert created.name == "Pilot bank" and created.tier == "enterprise" + assert created.billing_source == "operator" + assert ( + await db.scalar(select(User.id).where(User.organization_id == created.id)) + is None + ) + assert ( + await db.scalar( + select(OrganizationPIIConfig.id).where( + OrganizationPIIConfig.organization_id == created.id + ) + ) + is not None + ) + other = await create_organization_plan(db, "Pilot bank", "free") + assert other.id != created.id + assert created.tier == "enterprise" + + +@pytest.mark.asyncio +async def test_invite_acceptance_moves_only_an_empty_verified_bootstrap( + db, + test_org, + test_user_with_org, + monkeypatch: pytest.MonkeyPatch, +) -> None: + test_user_with_org.role = "owner" + await activate_organization_plan(db, test_org.id, "agency") + target_id = uuid4() + temporary_org = Organization( + id=uuid4(), + name="Temporary", + slug=f"temporary-{target_id}", + ) + invited = User( + id=target_id, + organization_id=temporary_org.id, + email=f"invited-{target_id}@example.com", + role="owner", + is_active=True, + is_verified=True, + ) + db.add_all([temporary_org, invited]) + await db.flush() + await ensure_privacy_defaults(db, temporary_org.id) + monkeypatch.setattr(management, "_audit", AsyncMock()) + + created = await management.create_team_invite( + management.TeamInviteInput(email=invited.email, role="member"), + test_user_with_org, + db, + ) + receipt = BillingWebhookReceipt( + organization_id=temporary_org.id, + payload_digest=uuid4().hex, + event_name="subscription_created", + event_at=datetime.now(timezone.utc), + ) + db.add(receipt) + await db.flush() + with pytest.raises(management.HTTPException, match="Leave or empty"): + await management.accept_team_invite( + management.AcceptTeamInvite(token=created.token), + invited, + db, + ) + assert await db.get(Organization, temporary_org.id) is not None + await db.delete(receipt) + await db.flush() + + accepted = await management.accept_team_invite( + management.AcceptTeamInvite(token=created.token), + invited, + db, + ) + + assert accepted.organization_id == test_org.id + assert accepted.role == "member" + assert await db.get(Organization, temporary_org.id) is None + with pytest.raises(management.HTTPException, match="invalid or expired"): + await management.accept_team_invite( + management.AcceptTeamInvite(token=created.token), + invited, + db, + ) + + +@pytest.mark.asyncio +async def test_invite_acceptance_locks_destination_before_revalidating_invite( + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_id = uuid4() + destination_id = uuid4() + user = SimpleNamespace( + id=uuid4(), + organization_id=source_id, + email="invited@example.com", + is_verified=True, + ) + invite = SimpleNamespace( + id=uuid4(), + organization_id=destination_id, + email=user.email, + role="member", + accepted_at=None, + revoked_at=None, + expires_at=datetime.now(timezone.utc) + timedelta(days=1), + ) + session = SimpleNamespace( + scalar=AsyncMock(return_value=destination_id), + execute=AsyncMock( + side_effect=[ + SimpleNamespace(), + SimpleNamespace(scalar_one_or_none=lambda: invite), + ] + ), + commit=AsyncMock(), + refresh=AsyncMock(), + ) + move = AsyncMock(return_value=user) + monkeypatch.setattr(management, "_require_entitlement", AsyncMock()) + monkeypatch.setattr(management, "move_user_from_bootstrap", move) + monkeypatch.setattr(management, "_audit", AsyncMock()) + + accepted = await management.accept_team_invite( + management.AcceptTeamInvite(token="x" * 32), + user, + session, + ) + + statements = [call.args[0] for call in session.execute.await_args_list] + assert "FOR UPDATE OF organizations" in str( + statements[0].compile(dialect=postgresql.dialect()) + ) + assert "FOR UPDATE OF organization_invites" in str( + statements[1].compile(dialect=postgresql.dialect()) + ) + assert accepted is user + move.assert_awaited_once_with( + session, + user_id=user.id, + source_organization_id=source_id, + destination_organization_id=destination_id, + role="member", + ) + + +@pytest.mark.asyncio +async def test_last_owner_cannot_be_removed(db, test_org, test_user_with_org) -> None: + test_user_with_org.role = "owner" + with pytest.raises(management.HTTPException, match="needs an owner"): + await management._protect_last_owner(db, test_org.id) + + +@pytest.mark.asyncio +async def test_free_plan_cannot_add_team_members( + db, + test_user_with_org, +) -> None: + test_user_with_org.role = "owner" + with pytest.raises(management.HTTPException, match="Plan upgrade"): + await management.create_team_invite( + management.TeamInviteInput(email="new-member@example.com"), + test_user_with_org, + db, + ) diff --git a/ee/tests/tenants/test_subscriptions.py b/ee/tests/tenants/test_subscriptions.py deleted file mode 100644 index 39c922b..0000000 --- a/ee/tests/tenants/test_subscriptions.py +++ /dev/null @@ -1,514 +0,0 @@ -from __future__ import annotations - -import hashlib -import hmac -import json -from datetime import datetime, timedelta, timezone -from types import SimpleNamespace -from unittest.mock import AsyncMock -from urllib.parse import parse_qs, urlsplit -from uuid import uuid4 - -import pytest -from sqlalchemy.dialects import postgresql - -from shim_enterprise.api.v1 import management -from shim_enterprise.tenants import subscriptions -from shim_enterprise.tenants.models import BillingWebhookReceipt, Organization, User -from shim_enterprise.tenants.service import create_api_key, ensure_privacy_defaults -from shim_enterprise.tenants.subscriptions import ( - checkout_urls, - process_lemonsqueezy_webhook, - set_organization_tier, - verify_lemonsqueezy_signature, -) - - -def _event( - organization_id, - user_id, - *, - event_name: str, - status: str, - updated_at: str, - variant_id: str = "managed-variant", -) -> bytes: - return json.dumps( - { - "meta": { - "event_name": event_name, - "custom_data": { - "organization_id": str(organization_id), - "user_id": str(user_id), - "checkout_signature": subscriptions._checkout_signature( - organization_id, - user_id, - ), - }, - }, - "data": { - "id": "subscription-123", - "attributes": { - "customer_id": "customer-123", - "variant_id": variant_id, - "status": status, - "updated_at": updated_at, - "renews_at": "2026-09-01T00:00:00Z", - "urls": { - "customer_portal": ( - "https://app.lemonsqueezy.com/my-orders/example" - ) - }, - }, - }, - }, - separators=(",", ":"), - ).encode() - - -def test_webhook_signature_is_timing_safe() -> None: - payload = b"signed" - secret = "test-signing-secret" - signature = hmac.new(secret.encode(), payload, hashlib.sha256).hexdigest() - - assert verify_lemonsqueezy_signature(payload, signature, secret) - assert not verify_lemonsqueezy_signature(payload, "invalid", secret) - - -@pytest.mark.asyncio -async def test_plan_changes_propagate_and_webhooks_are_ordered( - db, - test_org, - test_user_with_org, - monkeypatch: pytest.MonkeyPatch, -) -> None: - test_user_with_org.role = "owner" - monkeypatch.setattr( - "shim_enterprise.tenants.subscriptions.settings.LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID", - "managed-variant", - ) - monkeypatch.setattr( - "shim_enterprise.tenants.subscriptions.settings.LEMON_SQUEEZY_AGENCY_YEARLY_VARIANT_ID", - "agency-variant", - ) - plaintext, api_key = await create_api_key( - db, - user_id=test_user_with_org.id, - name="Inherited plan", - ) - assert plaintext.startswith("sk-shim-") - assert api_key.tier == "free" - - created = _event( - test_org.id, - test_user_with_org.id, - event_name="subscription_created", - status="active", - updated_at="2026-08-01T12:00:00Z", - ) - assert await process_lemonsqueezy_webhook(db, created) == "processed" - assert await process_lemonsqueezy_webhook(db, created) == "duplicate" - await db.refresh(test_org) - await db.refresh(api_key) - assert test_org.tier == "managed" - assert api_key.tier == "managed" - - agency = _event( - test_org.id, - test_user_with_org.id, - event_name="subscription_updated", - status="active", - updated_at="2026-08-01T13:00:00Z", - variant_id="agency-variant", - ) - assert await process_lemonsqueezy_webhook(db, agency) == "processed" - await db.refresh(test_org) - await db.refresh(api_key) - assert test_org.tier == "agency" - assert api_key.tier == "agency" - - expired = _event( - test_org.id, - test_user_with_org.id, - event_name="subscription_expired", - status="expired", - updated_at="2026-08-02T12:00:00Z", - ) - assert await process_lemonsqueezy_webhook(db, expired) == "processed" - stale_active = _event( - test_org.id, - test_user_with_org.id, - event_name="subscription_updated", - status="active", - updated_at="2026-08-01T14:00:00Z", - ) - assert await process_lemonsqueezy_webhook(db, stale_active) == "stale" - await db.refresh(test_org) - await db.refresh(api_key) - assert test_org.tier == "free" - assert api_key.tier == "free" - - -@pytest.mark.asyncio -async def test_checkout_identity_cannot_be_redirected_to_another_tenant( - db, - test_org, - test_user_with_org, - monkeypatch: pytest.MonkeyPatch, -) -> None: - test_user_with_org.role = "owner" - victim = Organization( - id=uuid4(), - name="Victim", - slug=f"victim-{uuid4().hex}", - ) - db.add(victim) - await db.flush() - monkeypatch.setattr( - "shim_enterprise.tenants.subscriptions.settings.LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL", - "https://getshim.lemonsqueezy.com/buy/solo", - ) - monkeypatch.setattr( - "shim_enterprise.tenants.subscriptions.settings.LEMON_SQUEEZY_SOLO_PRO_MONTHLY_VARIANT_ID", - "managed-variant", - ) - url = checkout_urls(test_org.id, test_user_with_org.id)["managed"]["monthly"] - custom_data = { - key.removeprefix("checkout[custom][").removesuffix("]"): values[-1] - for key, values in parse_qs(urlsplit(url).query).items() - if key.startswith("checkout[custom][") - } - - forged = json.loads( - _event( - test_org.id, - test_user_with_org.id, - event_name="subscription_created", - status="active", - updated_at="2026-08-01T12:00:00Z", - ) - ) - forged["meta"]["custom_data"] = { - **custom_data, - "organization_id": str(victim.id), - } - forged_body = json.dumps(forged, separators=(",", ":")).encode() - - with pytest.raises(ValueError, match="Invalid checkout identity"): - await process_lemonsqueezy_webhook(db, forged_body) - assert victim.tier == "free" - - valid = json.loads( - _event( - test_org.id, - test_user_with_org.id, - event_name="subscription_created", - status="active", - updated_at="2026-08-01T12:00:00Z", - ) - ) - valid["meta"]["custom_data"] = custom_data - valid_body = json.dumps(valid, separators=(",", ":")).encode() - - assert await process_lemonsqueezy_webhook(db, valid_body) == "processed" - await db.refresh(test_org) - assert test_org.tier == "managed" - - mismatched = json.loads(valid_body) - mismatched["meta"]["event_name"] = "subscription_updated" - mismatched["meta"]["custom_data"]["organization_id"] = str(victim.id) - mismatched["data"]["attributes"]["updated_at"] = "2026-08-01T13:00:00Z" - with pytest.raises(ValueError, match="Subscription identity mismatch"): - await process_lemonsqueezy_webhook( - db, - json.dumps(mismatched, separators=(",", ":")).encode(), - ) - - established = json.loads(valid_body) - established["meta"].pop("custom_data") - established["meta"]["event_name"] = "subscription_updated" - established["data"]["attributes"]["updated_at"] = "2026-08-01T13:00:00Z" - assert ( - await process_lemonsqueezy_webhook( - db, - json.dumps(established, separators=(",", ":")).encode(), - ) - == "processed" - ) - assert victim.tier == "free" - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("role", "is_active"), - [("member", True), ("owner", False)], -) -async def test_new_checkout_requires_an_active_owner( - db, - test_org, - test_user_with_org, - role: str, - is_active: bool, -) -> None: - test_user_with_org.role = role - test_user_with_org.is_active = is_active - await db.flush() - meta = { - "custom_data": { - "organization_id": str(test_org.id), - "user_id": str(test_user_with_org.id), - "checkout_signature": subscriptions._checkout_signature( - test_org.id, - test_user_with_org.id, - ), - } - } - - with pytest.raises(LookupError, match="Organization not found"): - await subscriptions._webhook_organization(db, meta, "new-subscription") - - -@pytest.mark.asyncio -async def test_operator_plan_is_inherited_by_new_keys( - db, - test_org, - test_user_with_org, -) -> None: - await set_organization_tier( - db, - test_org, - "managed", - status="active", - source="operator", - ) - _, api_key = await create_api_key( - db, - user_id=test_user_with_org.id, - name="Managed", - ) - assert api_key.tier == "managed" - - -def test_checkout_urls_include_cadence_and_tenant_identity( - monkeypatch: pytest.MonkeyPatch, -) -> None: - organization_id = uuid4() - user_id = uuid4() - for name in ( - "LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL", - "LEMON_SQUEEZY_SOLO_PRO_YEARLY_CHECKOUT_URL", - "LEMON_SQUEEZY_AGENCY_MONTHLY_CHECKOUT_URL", - "LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL", - ): - monkeypatch.setattr( - f"shim_enterprise.tenants.subscriptions.settings.{name}", None - ) - assert checkout_urls(organization_id, user_id) == {} - monkeypatch.setattr( - "shim_enterprise.tenants.subscriptions.settings.LEMON_SQUEEZY_SOLO_PRO_MONTHLY_CHECKOUT_URL", - "https://getshim.lemonsqueezy.com/buy/solo?discount=launch", - ) - monkeypatch.setattr( - "shim_enterprise.tenants.subscriptions.settings.LEMON_SQUEEZY_AGENCY_YEARLY_CHECKOUT_URL", - "https://getshim.lemonsqueezy.com/buy/agency-yearly", - ) - - urls = checkout_urls(organization_id, user_id) - - assert set(urls) == {"managed", "agency"} - assert set(urls["managed"]) == {"monthly"} - assert "discount=launch" in urls["managed"]["monthly"] - identity = f"lemonsqueezy-checkout:v1:{organization_id}:{user_id}".encode() - expected_signature = hmac.new( - subscriptions.settings.SECRET_KEY.encode(), - identity, - hashlib.sha256, - ).hexdigest() - for plan_urls in urls.values(): - for url in plan_urls.values(): - custom = parse_qs(urlsplit(url).query) - assert custom["checkout[custom][organization_id]"] == [str(organization_id)] - assert custom["checkout[custom][user_id]"] == [str(user_id)] - assert custom["checkout[custom][checkout_signature]"] == [ - expected_signature - ] - - -@pytest.mark.asyncio -async def test_only_free_owners_receive_checkout_links( - db, - test_org, - test_user_with_org, - monkeypatch: pytest.MonkeyPatch, -) -> None: - test_user_with_org.role = "owner" - configured = {"managed": {"monthly": "https://store.example.test/solo"}} - monkeypatch.setattr(management, "checkout_urls", lambda *_: configured) - - free = await management.get_subscription(test_user_with_org, db) - assert free.checkout_urls == configured - - await set_organization_tier( - db, - test_org, - "managed", - status="active", - source="lemonsqueezy", - ) - paid = await management.get_subscription(test_user_with_org, db) - assert paid.checkout_urls == {} - - -@pytest.mark.asyncio -async def test_invite_acceptance_moves_only_an_empty_verified_bootstrap( - db, - test_org, - test_user_with_org, - monkeypatch: pytest.MonkeyPatch, -) -> None: - test_user_with_org.role = "owner" - await set_organization_tier( - db, - test_org, - "agency", - status="active", - source="operator", - ) - target_id = uuid4() - temporary_org = Organization( - id=uuid4(), - name="Temporary", - slug=f"temporary-{target_id}", - ) - invited = User( - id=target_id, - organization_id=temporary_org.id, - email=f"invited-{target_id}@example.com", - role="owner", - is_active=True, - is_verified=True, - ) - db.add_all([temporary_org, invited]) - await db.flush() - await ensure_privacy_defaults(db, temporary_org.id) - monkeypatch.setattr(management, "_audit", AsyncMock()) - - created = await management.create_team_invite( - management.TeamInviteInput(email=invited.email, role="member"), - test_user_with_org, - db, - ) - receipt = BillingWebhookReceipt( - organization_id=temporary_org.id, - payload_digest=uuid4().hex, - event_name="subscription_created", - event_at=datetime.now(timezone.utc), - ) - db.add(receipt) - await db.flush() - with pytest.raises(management.HTTPException, match="Leave or empty"): - await management.accept_team_invite( - management.AcceptTeamInvite(token=created.token), - invited, - db, - ) - assert await db.get(Organization, temporary_org.id) is not None - await db.delete(receipt) - await db.flush() - - accepted = await management.accept_team_invite( - management.AcceptTeamInvite(token=created.token), - invited, - db, - ) - - assert accepted.organization_id == test_org.id - assert accepted.role == "member" - assert await db.get(Organization, temporary_org.id) is None - with pytest.raises(management.HTTPException, match="invalid or expired"): - await management.accept_team_invite( - management.AcceptTeamInvite(token=created.token), - invited, - db, - ) - - -@pytest.mark.asyncio -async def test_invite_acceptance_locks_destination_before_revalidating_invite( - monkeypatch: pytest.MonkeyPatch, -) -> None: - source_id = uuid4() - destination_id = uuid4() - user = SimpleNamespace( - id=uuid4(), - organization_id=source_id, - email="invited@example.com", - is_verified=True, - ) - invite = SimpleNamespace( - id=uuid4(), - organization_id=destination_id, - email=user.email, - role="member", - accepted_at=None, - revoked_at=None, - expires_at=datetime.now(timezone.utc) + timedelta(days=1), - ) - session = SimpleNamespace( - scalar=AsyncMock(return_value=destination_id), - execute=AsyncMock( - side_effect=[ - SimpleNamespace(), - SimpleNamespace(scalar_one_or_none=lambda: invite), - ] - ), - commit=AsyncMock(), - refresh=AsyncMock(), - ) - move = AsyncMock(return_value=user) - monkeypatch.setattr(management, "_require_entitlement", AsyncMock()) - monkeypatch.setattr(management, "move_user_from_bootstrap", move) - monkeypatch.setattr(management, "_audit", AsyncMock()) - - accepted = await management.accept_team_invite( - management.AcceptTeamInvite(token="x" * 32), - user, - session, - ) - - statements = [call.args[0] for call in session.execute.await_args_list] - assert "FOR UPDATE OF organizations" in str( - statements[0].compile(dialect=postgresql.dialect()) - ) - assert "FOR UPDATE OF organization_invites" in str( - statements[1].compile(dialect=postgresql.dialect()) - ) - assert accepted is user - move.assert_awaited_once_with( - session, - user_id=user.id, - source_organization_id=source_id, - destination_organization_id=destination_id, - role="member", - ) - - -@pytest.mark.asyncio -async def test_last_owner_cannot_be_removed(db, test_org, test_user_with_org) -> None: - test_user_with_org.role = "owner" - with pytest.raises(management.HTTPException, match="needs an owner"): - await management._protect_last_owner(db, test_org.id) - - -@pytest.mark.asyncio -async def test_free_plan_cannot_add_team_members( - db, - test_user_with_org, -) -> None: - test_user_with_org.role = "owner" - with pytest.raises(management.HTTPException, match="Plan upgrade"): - await management.create_team_invite( - management.TeamInviteInput(email="new-member@example.com"), - test_user_with_org, - db, - ) diff --git a/ee/tests/tenants/test_teams.py b/ee/tests/tenants/test_teams.py new file mode 100644 index 0000000..c48e311 --- /dev/null +++ b/ee/tests/tenants/test_teams.py @@ -0,0 +1,387 @@ +"""Team boundaries use real PostgreSQL authorization and reservation transactions.""" + +import asyncio +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from fastapi import FastAPI, HTTPException +from httpx import ASGITransport, AsyncClient +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import async_sessionmaker + +from shim_enterprise.api.enterprise_deps import get_current_user +from shim_enterprise.api.v1 import management +from shim_enterprise.billing.ledger import ( + DurableAccountingRepository, + FinalizationCommand, + QuotaLimitExceeded, + QuotaPolicySnapshot, + QuotaReservationCommand, + TerminalAction, +) +from shim_enterprise.billing.models import ( + QuotaPeriodUsage, + RequestLifecycle, + UsageLedger, +) +from shim_enterprise.core.database import get_db +from shim_enterprise.gateway.pipeline.quota_reservation import AccountingPolicyLoader +from shim_enterprise.outbox.models import OutboxEvent +from shim_enterprise.tenants.models import ( + ApiKey, + Organization, + Team, + TeamMembership, + User, +) +from shim_enterprise.tenants.service import authenticate_api_key, create_api_key +from shim_enterprise.tenants.teams import synchronize_oidc_teams + + +@pytest.mark.asyncio +async def test_team_admin_isolation_auditor_and_immediate_rotation( + db, test_user_with_org +): + owner = test_user_with_org + owner.role = "owner" + admin = User( + id=uuid4(), + organization_id=owner.organization_id, + email=f"admin-{uuid4()}@example.com", + role="member", + is_active=True, + is_verified=True, + ) + auditor = User( + id=uuid4(), + organization_id=owner.organization_id, + email=f"audit-{uuid4()}@example.com", + role="auditor", + is_active=True, + is_verified=True, + ) + first = Team(id=uuid4(), organization_id=owner.organization_id, name="First") + second = Team(id=uuid4(), organization_id=owner.organization_id, name="Second") + db.add_all([admin, auditor, first, second]) + await db.flush() + db.add( + TeamMembership( + organization_id=owner.organization_id, + team_id=first.id, + user_id=admin.id, + role="team_admin", + ) + ) + await db.flush() + plaintext, key = await create_api_key( + db, + user_id=owner.id, + name="Team key", + team="historical-label", + team_id=first.id, + allowed_models=["internal-model"], + ) + _, other = await create_api_key( + db, user_id=owner.id, name="Other team", team_id=second.id + ) + + app = FastAPI() + app.include_router(management.router) + current = admin + app.dependency_overrides[get_current_user] = lambda: current + app.dependency_overrides[get_db] = lambda: db + async with AsyncClient( + transport=ASGITransport(app), base_url="http://test" + ) as client: + assert (await client.delete(f"/api-keys/{other.id}")).status_code == 404 + assert ( + await client.put( + f"/teams/{second.id}/members/{admin.id}", json={"role": "member"} + ) + ).status_code == 404 + assert ( + await client.put(f"/teams/{first.id}", json={"name": "Edited"}) + ).status_code == 403 + assert ( + await client.put( + f"/teams/{first.id}/members/{owner.id}", json={"role": "team_admin"} + ) + ).status_code == 403 + rotated = await client.post(f"/api-keys/{key.id}/rotate") + assert rotated.status_code == 200, rotated.text + payload = rotated.json() + assert payload["id"] == str(key.id) + assert payload["allowed_models"] == ["internal-model"] + assert payload["team"] == "historical-label" + assert payload["team_id"] == str(first.id) + assert await authenticate_api_key(db, plaintext) is None + assert (await authenticate_api_key(db, payload["plaintext"])).id == key.id + assert "plaintext" not in (await client.get("/api-keys")).json()[0] + current = auditor + assert ( + await client.post("/api-keys", json={"name": "forbidden"}) + ).status_code == 403 + assert (await client.post(f"/api-keys/{key.id}/rotate")).status_code == 403 + assert ( + await client.put( + f"/teams/{first.id}/members/{owner.id}", json={"role": "member"} + ) + ).status_code == 403 + assert (await client.get("/teams")).status_code == 200 + + event = await db.scalar( + select(OutboxEvent).where(OutboxEvent.organization_id == owner.organization_id) + ) + assert event.payload["endpoint"] == "tenant.api_key_rotated" + assert payload["plaintext"] not in str(event.payload) + + +@pytest.mark.asyncio +async def test_team_mapping_never_grants_cross_tenant_and_revokes_only_oidc( + db, test_user_with_org +): + user = test_user_with_org + teams = [ + Team(id=uuid4(), organization_id=user.organization_id, name=f"Team {i}") + for i in range(2) + ] + db.add_all(teams) + await db.flush() + db.add( + TeamMembership( + organization_id=user.organization_id, + team_id=teams[0].id, + user_id=user.id, + role="member", + source="local", + ) + ) + await db.flush() + mapping = { + "local": {"team_id": str(teams[0].id), "role": "team_admin"}, + "mapped": {"team_id": str(teams[1].id), "role": "team_admin"}, + } + await synchronize_oidc_teams(db, user, ["local", "mapped"], mapping) + memberships = list( + ( + await db.scalars( + select(TeamMembership).where(TeamMembership.user_id == user.id) + ) + ).all() + ) + assert {(row.source, row.role) for row in memberships} == { + ("local", "member"), + ("oidc", "team_admin"), + } + await synchronize_oidc_teams(db, user, [], mapping) + remaining = list( + ( + await db.scalars( + select(TeamMembership).where(TeamMembership.user_id == user.id) + ) + ).all() + ) + assert len(remaining) == 1 and remaining[0].source == "local" + with pytest.raises(ValueError, match="organization"): + await synchronize_oidc_teams( + db, + user, + ["outside"], + {"outside": {"team_id": str(uuid4()), "role": "member"}}, + ) + + +@pytest.mark.asyncio +async def test_model_and_membership_denial_precedes_quota(db, test_user_with_org): + user = test_user_with_org + team = Team(id=uuid4(), organization_id=user.organization_id, name="Restricted") + db.add(team) + await db.flush() + _, key = await create_api_key( + db, + user_id=user.id, + name="Restricted", + team_id=team.id, + allowed_models=["internal"], + ) + prepared = SimpleNamespace( + tenant_id=user.organization_id, api_key_id=key.id, model="external" + ) + with pytest.raises(HTTPException) as model_error: + await AccountingPolicyLoader().quota(db, prepared) + assert model_error.value.status_code == 403 + assert model_error.value.detail["code"] == "MODEL_NOT_ALLOWED" + prepared.model = "internal" + with pytest.raises(HTTPException, match="team member"): + await AccountingPolicyLoader().quota(db, prepared) + assert ( + await db.scalar( + select(QuotaPeriodUsage.id).where( + QuotaPeriodUsage.organization_id == user.organization_id + ) + ) + is None + ) + + +@pytest.mark.asyncio +async def test_concurrent_team_quota_reserves_and_refunds_all_scopes(async_engine): + factory = async_sessionmaker(async_engine, expire_on_commit=False) + organization_id, user_id, team_id = uuid4(), uuid4(), uuid4() + async with factory.begin() as session: + session.add( + Organization( + id=organization_id, + name="Quota race", + slug=f"quota-race-{organization_id}", + ) + ) + await session.flush() + session.add( + User( + id=user_id, + organization_id=organization_id, + email=f"quota-race-{user_id}@example.com", + role="owner", + is_active=True, + is_verified=True, + ) + ) + session.add( + Team( + id=team_id, + organization_id=organization_id, + name="Concurrent", + monthly_request_limit=3, + monthly_token_limit=30, + ) + ) + await session.flush() + keys = [ + ( + await create_api_key( + session, user_id=user_id, name=f"key {i}", team_id=team_id + ) + )[1].id + for i in range(8) + ] + + repository = DurableAccountingRepository() + + async def reserve(key_id): + now = datetime.now(timezone.utc) + request_id = f"req_team_{uuid4().hex}" + async with factory() as session: + policy = await AccountingPolicyLoader().quota( + session, + SimpleNamespace( + tenant_id=organization_id, api_key_id=key_id, model="internal" + ), + ) + # This key's limit is stricter than the team limit, and both count. + policy = QuotaPolicySnapshot( + version=policy.version, + daily_request_limit=None, + monthly_request_limit=1, + monthly_token_limit=10, + team_id=team_id, + team_policy=policy.team_policy, + ) + try: + result = await repository.reserve_quota( + session, + QuotaReservationCommand( + tenant_id=organization_id, + api_key_id=key_id, + request_id=request_id, + requested_model="internal", + source_endpoint="chat.completions", + started_at=now, + reconciliation_due_at=now + timedelta(minutes=2), + estimated_input_tokens=2, + maximum_output_tokens=8, + policy=policy, + ), + ) + await session.commit() + return request_id, result + except QuotaLimitExceeded: + await session.rollback() + return None + + try: + results = await asyncio.gather(*(reserve(key) for key in keys)) + admitted = [result for result in results if result is not None] + assert len(admitted) == 3 + assert all(len(result.period_allocations) == 2 for _, result in admitted) + used_key = keys[next(index for index, result in enumerate(results) if result)] + assert await reserve(used_key) is None + async with factory.begin() as session: + await repository.finalize( + session, + FinalizationCommand( + tenant_id=organization_id, + request_id=admitted[0][0], + quota_action=TerminalAction.REFUND, + lifecycle_status="failed", + terminal_error_code="REQUEST_ABORTED", + completed_at=datetime.now(timezone.utc), + ), + ) + assert await reserve(used_key) is not None + async with factory() as session: + counter = await session.scalar( + select(QuotaPeriodUsage).where(QuotaPeriodUsage.team_id == team_id) + ) + assert counter.reserved_requests == 3 and counter.reserved_tokens == 30 + async with factory.begin() as session: + settled = await repository.finalize( + session, + FinalizationCommand( + tenant_id=organization_id, + request_id=admitted[1][0], + quota_action=TerminalAction.SETTLE, + prompt_tokens=2, + completion_tokens=3, + lifecycle_status="completed", + completed_at=datetime.now(timezone.utc), + ), + ) + replay = await repository.finalize( + session, + FinalizationCommand( + tenant_id=organization_id, + request_id=admitted[1][0], + quota_action=TerminalAction.SETTLE, + prompt_tokens=2, + completion_tokens=3, + lifecycle_status="completed", + ), + ) + assert replay.replayed and replay.quota_event_id == settled.quota_event_id + async with factory() as session: + counter = await session.scalar( + select(QuotaPeriodUsage).where(QuotaPeriodUsage.team_id == team_id) + ) + assert (counter.reserved_requests, counter.settled_requests) == (2, 1) + assert (counter.reserved_tokens, counter.settled_tokens) == (20, 5) + finally: + async with factory.begin() as session: + for model in ( + OutboxEvent, + UsageLedger, + RequestLifecycle, + QuotaPeriodUsage, + ApiKey, + TeamMembership, + Team, + User, + ): + await session.execute( + delete(model).where(model.organization_id == organization_id) + ) + await session.execute( + delete(Organization).where(Organization.id == organization_id) + ) diff --git a/ee/tests/workers/test_readiness.py b/ee/tests/workers/test_readiness.py new file mode 100644 index 0000000..804550a --- /dev/null +++ b/ee/tests/workers/test_readiness.py @@ -0,0 +1,128 @@ +import asyncio +import json +from pathlib import Path +import subprocess +import sys + +import pytest + +from shim_enterprise.workers import ( + ai_act, + compliance, + outbox, + readiness, + reconciliation, +) + + +def test_heartbeat_file_and_cli(tmp_path, monkeypatch, caplog): + path = tmp_path / "heartbeat.json" + monkeypatch.setattr(readiness.time, "monotonic", lambda: 100.0) + assert not readiness.is_ready(path, "outbox", 10) + readiness.write_heartbeat(None, "outbox") + readiness.write_heartbeat(path, "outbox") + assert readiness.is_ready(path, "outbox", 10) + assert not readiness.is_ready(path, "compliance", 10) + for timestamp in [float("nan"), float("inf"), 101, 89, True, None, "100", 10**1000]: + path.write_text( + json.dumps({"worker": "outbox", "monotonic_success": timestamp}) + ) + assert not readiness.is_ready(path, "outbox", 10) + for payload in ["{", "[]", "null", "{}"]: + path.write_text(payload) + assert not readiness.is_ready(path, "outbox", 10) + for age in [0, -1, float("nan"), float("inf")]: + assert not readiness.is_ready(path, "outbox", age) + assert not readiness.is_ready(Path("relative"), "outbox", 10) + readiness.write_heartbeat(tmp_path / "missing" / "secret-path", "outbox") + assert "Worker heartbeat write failed" in caplog.text + assert "secret-path" not in caplog.text + monkeypatch.undo() + readiness.write_heartbeat(path, "outbox") + command = [ + sys.executable, + "-m", + "shim_enterprise.workers.readiness", + "--path", + str(path), + "--worker", + "outbox", + "--max-age-seconds", + "10", + ] + assert subprocess.run(command, capture_output=True).returncode == 0 + path.write_text("invalid") + assert subprocess.run(command, capture_output=True).returncode == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "module,worker_class,summary", + [ + (outbox, outbox.OutboxWorker, outbox.WorkerPass(0, 0, 0, 0, 0)), + (outbox, outbox.OutboxWorker, outbox.WorkerPass(1, 0, 1, 0, 0)), + (outbox, outbox.OutboxWorker, outbox.WorkerPass(1, 0, 0, 1, 0)), + (outbox, outbox.OutboxWorker, outbox.WorkerPass(1, 0, 0, 0, 1)), + (compliance, compliance.ComplianceSweepWorker, compliance.SweepSummary()), + ( + compliance, + compliance.ComplianceSweepWorker, + compliance.SweepSummary(errors=1), + ), + (ai_act, ai_act.AuditMaintenanceWorker, ai_act.MaintenanceSummary()), + (ai_act, ai_act.AuditMaintenanceWorker, ai_act.MaintenanceSummary(errors=1)), + (reconciliation, reconciliation.ReconciliationWorker, 0), + ], +) +async def test_only_successful_pass_refreshes_heartbeat( + tmp_path, monkeypatch, module, worker_class, summary +): + path = tmp_path / "heartbeat.json" + monkeypatch.setattr(module.settings, "WORKER_HEARTBEAT_PATH", path) + worker = object.__new__(worker_class) + worker.interval_seconds = 1 + worker.worker_id = "test" + stop = asyncio.Event() + + async def run_once(): + stop.set() + return summary + + worker.run_once = run_once + await worker.run(stop) + failures = sum( + getattr(summary, key, 0) + for key in ("failed", "dead_lettered", "lease_lost", "errors") + ) + assert path.exists() == (failures == 0) + stop.clear() + + async def fail(): + stop.set() + raise RuntimeError("unavailable") + + worker.run_once = fail + previous = path.read_bytes() if path.exists() else None + await worker.run(stop) + assert (path.read_bytes() if path.exists() else None) == previous + + +@pytest.mark.asyncio +async def test_anchor_failure_is_reported(monkeypatch): + from contextlib import asynccontextmanager + from types import SimpleNamespace + from unittest.mock import AsyncMock + + @asynccontextmanager + async def transaction(): + yield + + session = SimpleNamespace( + execute=AsyncMock(return_value=SimpleNamespace(scalars=lambda: ["tenant"])), + begin_nested=transaction, + ) + monkeypatch.setattr( + ai_act, "write_anchor", AsyncMock(side_effect=OSError("secret")) + ) + worker = object.__new__(ai_act.AuditMaintenanceWorker) + assert await worker._anchor_tenants(session, None) == (0, 1) diff --git a/openapi/community.json b/openapi/community.json index fdd7c98..aa8df2b 100644 --- a/openapi/community.json +++ b/openapi/community.json @@ -886,6 +886,13 @@ "title": "Content", "type": "object" }, + "CountTokensRequest": { + "additionalProperties": { + "$ref": "#/components/schemas/JsonValue" + }, + "title": "CountTokensRequest", + "type": "object" + }, "Delivery": { "description": "Delivery mode for the generated content.", "enum": [ @@ -5628,6 +5635,193 @@ "summary": "Messages" } }, + "/v1/messages/count_tokens": { + "post": { + "operationId": "count_tokens_v1_messages_count_tokens_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CountTokensRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "400": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "401": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "402": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "403": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "404": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "408": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "409": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "413": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "429": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "500": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "502": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "503": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "504": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + }, + "529": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AnthropicErrorResponse" + } + } + }, + "description": "Sanitized Anthropic-compatible error." + } + }, + "security": [ + { + "ShimAPIKey": [] + }, + { + "HTTPBearer": [] + }, + { + "AnthropicAPIKey": [] + } + ], + "summary": "Count Tokens" + } + }, "/v1/models": { "get": { "operationId": "list_models_v1_models_get", diff --git a/src/shim/api/deps.py b/src/shim/api/deps.py index b715398..7c4ba35 100644 --- a/src/shim/api/deps.py +++ b/src/shim/api/deps.py @@ -78,7 +78,9 @@ async def dispatch_gateway_inference( request: Request, payload: dict[str, Any], provider: Literal["openai", "anthropic", "google"], - protocol: Literal["chat", "responses", "messages", "generate_content"], + protocol: Literal[ + "chat", "responses", "messages", "count_tokens", "generate_content" + ], model: str, stream: bool, gateway_service: GatewayService, diff --git a/src/shim/api/v1/chat.py b/src/shim/api/v1/chat.py index 2bd3355..18e65de 100644 --- a/src/shim/api/v1/chat.py +++ b/src/shim/api/v1/chat.py @@ -103,12 +103,14 @@ def model_record(model_id: str, provider: str = "openai") -> dict[str, object]: def _anthropic_model_page( - model_ids: tuple[str, ...], + records: list[dict[str, object]], *, after_id: str | None, before_id: str | None, limit: int, ) -> dict[str, object]: + model_ids = tuple(str(record["id"]) for record in records) + by_id = {str(record["id"]): record for record in records} if after_id is not None and before_id is not None: raise HTTPException(status_code=400, detail="Use only one model cursor.") if after_id is not None or before_id is not None: @@ -132,7 +134,7 @@ def _anthropic_model_page( page = model_ids[:limit] has_more = len(page) < len(model_ids) return { - "data": [model_record(model_id, "anthropic") for model_id in page], + "data": [by_id[model_id] for model_id in page], "has_more": has_more, "first_id": page[0] if page else None, "last_id": page[-1] if page else None, @@ -145,6 +147,7 @@ def _anthropic_model_page( responses=MODEL_ERROR_RESPONSES, ) async def list_models( + request: Request, _principal: AuthenticatedPrincipal = Depends(get_anthropic_authenticated_principal), anthropic_version: str | None = Header(None, alias="anthropic-version"), client_version: str | None = None, @@ -153,9 +156,8 @@ async def list_models( limit: int = Query(20, ge=1, le=1_000), ): if anthropic_version is not None: - model_ids = DEFAULT_PRICE_BOOK.models("anthropic") return _anthropic_model_page( - model_ids, + await _model_records(request, _principal, "anthropic"), after_id=after_id, before_id=before_id, limit=limit, @@ -165,7 +167,7 @@ async def list_models( return {"models": []} return { "object": "list", - "data": [model_record(model_id) for model_id in DEFAULT_PRICE_BOOK.prices], + "data": await _model_records(request, _principal, "openai"), } @@ -176,20 +178,33 @@ async def list_models( ) async def retrieve_model( model_id: str, + request: Request, _principal: AuthenticatedPrincipal = Depends(get_anthropic_authenticated_principal), anthropic_version: str | None = Header(None, alias="anthropic-version"), ): provider = "anthropic" if anthropic_version is not None else "openai" - try: - return model_record(model_id, provider) - except ValueError: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={ - "code": "MODEL_NOT_FOUND", - "message": "The requested model is not in the public catalog.", - }, - ) from None + for record in await _model_records(request, _principal, provider): + if record["id"] == model_id: + return record + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={ + "code": "MODEL_NOT_FOUND", + "message": "The requested model is not available.", + }, + ) + + +async def _model_records( + request: Request, principal: AuthenticatedPrincipal, provider: str +) -> list[dict[str, object]]: + catalog = getattr(request.app.state, "model_catalog", None) + if catalog is not None: + return await catalog(principal, provider) + return [ + model_record(model_id, provider) + for model_id in DEFAULT_PRICE_BOOK.models(provider) + ] @router.post( diff --git a/src/shim/api/v1/messages.py b/src/shim/api/v1/messages.py index 2e52e71..caf5ea0 100644 --- a/src/shim/api/v1/messages.py +++ b/src/shim/api/v1/messages.py @@ -71,3 +71,36 @@ async def messages( gateway_service=gateway_service, principal=principal, ) + + +class CountTokensRequest(ProviderRequest): + @model_validator(mode="after") + def validate_routing_fields(self) -> Self: + self.require_nonempty_string("model") + self.require("messages", list) + if "stream" in self.root: + raise ValueError("count_tokens does not support streaming") + return self + + +@router.post("/messages/count_tokens", responses=ANTHROPIC_ERROR_RESPONSES) +async def count_tokens( + request: Request, + payload: CountTokensRequest, + gateway_service: GatewayService = Depends(get_gateway_service), + principal: AuthenticatedPrincipal = Depends(get_anthropic_authenticated_principal), +): + provider_payload = payload.provider_payload() + user_profile_id = request.headers.get("anthropic-user-profile-id") + if user_profile_id is not None: + provider_payload["user_profile_id"] = user_profile_id + return await dispatch_gateway_inference( + request=request, + payload=provider_payload, + provider="anthropic", + protocol="count_tokens", + model=payload.require_nonempty_string("model"), + stream=False, + gateway_service=gateway_service, + principal=principal, + ) diff --git a/src/shim/billing/pricing.py b/src/shim/billing/pricing.py index 9ab7fdc..1b477df 100644 --- a/src/shim/billing/pricing.py +++ b/src/shim/billing/pricing.py @@ -188,11 +188,18 @@ def resolved_price_metadata( *, input_tokens: int, output_tokens: int, + unpriced: bool = False, ) -> dict[str, str | int]: metadata: dict[str, str | int] = { "catalog_version": self.version, "provider": provider, } + if unpriced: + return { + **metadata, + "provider_model": model or "", + "pricing_resolution": "unknown", + } if model == UNSPECIFIED_PROVIDER_MODEL and provider == "openai": metadata["pricing_resolution"] = "conservative_max" resolved_model, price = max( @@ -363,7 +370,12 @@ def compute_cost_usd( prompt_tokens: int, completion_tokens: int, provider: str = "openai", + *, + unpriced: bool = False, ) -> Decimal: """Return the deterministic provider cost for a token pair.""" + if unpriced: + # Ledger arithmetic needs a numeric placeholder; metadata must mark it unknown. + return Decimal("0") return DEFAULT_PRICE_BOOK.compute(model, prompt_tokens, completion_tokens, provider) diff --git a/src/shim/gateway/api/errors.py b/src/shim/gateway/api/errors.py index 9fced97..eaa3ccb 100644 --- a/src/shim/gateway/api/errors.py +++ b/src/shim/gateway/api/errors.py @@ -170,7 +170,7 @@ def _gateway_provider( headers: Mapping[str, str], ) -> str | None: normalized = path.rstrip("/") - if normalized == "/v1/messages": + if normalized in {"/v1/messages", "/v1/messages/count_tokens"}: return "anthropic" if normalized in {"/v1/chat/completions", "/v1/responses"}: return "openai" diff --git a/src/shim/gateway/kernel/gateway_kernel.py b/src/shim/gateway/kernel/gateway_kernel.py index d8e49ce..d5c554d 100644 --- a/src/shim/gateway/kernel/gateway_kernel.py +++ b/src/shim/gateway/kernel/gateway_kernel.py @@ -2,12 +2,12 @@ from __future__ import annotations -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping import logging from typing import Any from fastapi import HTTPException -from starlette.responses import Response +from starlette.responses import Response, JSONResponse from shim.core.middleware import AsyncRateLimiter from shim.gateway.admission import LoopDetector @@ -18,10 +18,16 @@ from shim.gateway.pipeline.provider_spend import ProviderSpendStage from shim.gateway.pipeline.provider_execution import ( ProviderCallError, + ProviderNonStream, ProviderExecutionStage, ) from shim.gateway.streaming import StreamSession -from shim.gateway.usage import UsageFailureReason, UsageLifecycle, UsageLimitExceeded +from shim.gateway.usage import ( + UsageAuditPersistenceError, + UsageFailureReason, + UsageLifecycle, + UsageLimitExceeded, +) from shim.observability.metrics import REQUESTS_TOTAL, bounded_label from shim.observability.tracing import start_span from shim.privacy.continuation import PrivacyContinuationStore @@ -49,6 +55,8 @@ def __init__( usage: UsageLifecycle, heartbeat_interval_seconds: float = 30, output_hash_salt: str | None = None, + prepare_inference: Callable[[PreparedInference], Awaitable[PreparedInference]] + | None = None, ) -> None: self.executions = dict(executions) if not self.executions or not set(self.executions) <= _PROVIDERS: @@ -66,6 +74,7 @@ def __init__( ) self.chain_store = chain_store self.policy_resolver = policy_resolver + self.prepare_inference = prepare_inference async def execute(self, invocation: GatewayInvocation) -> Response: endpoint = bounded_label("endpoint", invocation.metadata.endpoint) @@ -119,10 +128,30 @@ async def _execute( retryable=True, provider=invocation.provider, ) - prepared = await run_stage( - AuthenticateStage(self.policy_resolver), - invocation, - ) + authenticate_stage = AuthenticateStage(self.policy_resolver) + try: + prepared = await run_stage(authenticate_stage, invocation) + except BaseException: + if authenticate_stage.prepared is not None: + await self.usage.reject(authenticate_stage.prepared) + raise + if self.prepare_inference is not None: + try: + prepared = await self.prepare_inference(prepared) + except BaseException: + if not any( + verdict.outcome in {"deny", "error"} + for verdict in prepared.policy_verdicts + ): + prepared.record_verdict( + "deployment.registry", + stage="admission", + outcome="error", + reason_code="DEPLOYMENT_REGISTRY_UNAVAILABLE", + ) + await self.usage.reject(prepared) + raise + if prepared_observer is not None: prepared_observer(prepared) admission_stage = AdmissionStage( @@ -136,9 +165,25 @@ async def _execute( ) try: prepared = await run_stage(admission_stage, prepared) - except BaseException: + except BaseException as error: if admission_stage.reserved: await self._fail_safely(prepared, reason="admission_aborted") + else: + if not any( + verdict.outcome in {"deny", "error"} + for verdict in prepared.policy_verdicts + ): + prepared.record_verdict( + "gateway.admission", + stage="admission", + outcome="deny" + if isinstance(error, HTTPException) and error.status_code < 500 + else "error", + reason_code="INVALID_REQUEST" + if isinstance(error, HTTPException) and error.status_code < 500 + else "ADMISSION_UNAVAILABLE", + ) + await self.usage.reject(prepared) raise stream_session: StreamSession | None = None @@ -150,6 +195,32 @@ async def _execute( ), prepared, ) + if prepared.protocol == "count_tokens": + await self.usage.record_token_count(prepared, None) + + async def token_count_started() -> None: + pass + + counted = await execution.execute( + invocation=invocation, + prepared=prepared, + provider_start_callback=token_count_started, + ) + assert isinstance(counted, ProviderNonStream) + input_tokens = counted.payload.get("input_tokens") + if ( + not isinstance(input_tokens, int) + or isinstance(input_tokens, bool) + or input_tokens < 0 + ): + raise ProviderCallError( + 502, "PROVIDER_UNAVAILABLE", False, provider="anthropic" + ) + await self.usage.record_token_count(prepared, input_tokens) + headers = {"X-Shim-Request-Id": str(prepared.request_id)} + if counted.request_id: + headers["request-id"] = counted.request_id + return JSONResponse(content=counted.payload, headers=headers) await self.usage.record_privacy(prepared) prepared = await run_stage( ProviderSpendStage(invocation, self.usage), @@ -183,7 +254,20 @@ async def _execute( and error.status_code < 500 else "request_aborted" ) - await self._fail_safely(prepared, reason=reason) + if prepared.protocol != "count_tokens": + await self._fail_safely(prepared, reason=reason) + elif not isinstance(error, UsageAuditPersistenceError): + if not any( + verdict.outcome in {"deny", "error"} + for verdict in prepared.policy_verdicts + ): + prepared.record_verdict( + "gateway.token_count", + stage="privacy", + outcome="error", + reason_code="TOKEN_COUNT_UNAVAILABLE", + ) + await self.usage.reject(prepared) raise async def _fail_safely( @@ -194,5 +278,7 @@ async def _fail_safely( ) -> None: try: await self.usage.fail(prepared, reason=reason) + except UsageAuditPersistenceError: + raise except Exception as exc: logger.error("Usage recovery failed type=%s", type(exc).__name__) diff --git a/src/shim/gateway/kernel/result.py b/src/shim/gateway/kernel/result.py index 795e204..cb07680 100644 --- a/src/shim/gateway/kernel/result.py +++ b/src/shim/gateway/kernel/result.py @@ -2,24 +2,57 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field +from datetime import datetime, timezone +from hashlib import sha256 +import json from typing import Any, Literal +from pydantic import AwareDatetime, Field + from shim.billing.pricing import ( + DEFAULT_PRICE_BOOK, UNSPECIFIED_PROVIDER_MODEL as UNSPECIFIED_PROVIDER_MODEL, ) from shim.gateway.contracts.context import GatewayContext +from shim.gateway.contracts import FrozenContractModel from shim.gateway.contracts.ids import ProviderId from shim.gateway.request_policy import RequestPolicyContext as _RequestPolicyContext from shim.privacy.policies import PrivacyOutcome +class PolicyVerdict(FrozenContractModel): + """Content-free evidence of one evaluated gateway rule, scoped by its envelope.""" + + schema_version: Literal[1] = 1 + rule_id: str = Field(pattern=r"^[a-z][a-z0-9_.]{0,95}$") + rule_version: Literal[1] = 1 + policy_version: str = Field(min_length=1, max_length=128) + stage: Literal["authentication", "admission", "privacy", "provider_spend"] + outcome: Literal["allow", "mask", "deny", "error", "skip"] + reason_code: str = Field(pattern=r"^[A-Z][A-Z0-9_]{0,95}$") + effective_at: AwareDatetime + + +@dataclass(frozen=True) +class ProviderTarget: + """Operator-approved invocation destination; never populated from request JSON.""" + + deployment_id: str + base_url: str + upstream_model: str + credential_reference: str + timeout_seconds: float + declared_version: str + + @dataclass(frozen=True) class AdmissionState: estimated_input_tokens: int maximum_output_tokens: int cost_center: str tags: tuple[str, ...] + repeat_chain_length: int | None = None @dataclass(frozen=True) @@ -29,13 +62,56 @@ class PreparedInference: context: GatewayContext payload: dict[str, Any] provider: ProviderId - protocol: Literal["chat", "responses", "messages", "generate_content"] + protocol: Literal[ + "chat", "responses", "messages", "count_tokens", "generate_content" + ] model: str stream: bool policy: _RequestPolicyContext pii_config: dict[str, bool] | None admission: AdmissionState | None = None privacy: PrivacyOutcome | None = None + deployment_kind: Literal["internal", "external", "unknown"] = "unknown" + target: ProviderTarget | None = None + policy_verdicts: list[PolicyVerdict] = field(default_factory=list) + + def record_verdict( + self, + rule_id: str, + *, + stage: Literal["authentication", "admission", "privacy", "provider_spend"], + outcome: Literal["allow", "mask", "deny", "error", "skip"], + reason_code: str, + policy: object = None, + policy_version: str | None = None, + ) -> None: + # The request-local list survives immutable stage replacements and exceptions. + self.policy_verdicts[:] = [ + verdict for verdict in self.policy_verdicts if verdict.rule_id != rule_id + ] + self.policy_verdicts.append( + PolicyVerdict( + rule_id=rule_id, + policy_version=policy_version + or sha256( + json.dumps(policy, sort_keys=True, default=str).encode() + ).hexdigest(), + stage=stage, + outcome=outcome, + reason_code=reason_code, + effective_at=datetime.now(timezone.utc), + ) + ) + + @property + def pricing_model(self) -> str: + return self.target.upstream_model if self.target is not None else self.model + + @property + def unpriced(self) -> bool: + return self.target is not None and not DEFAULT_PRICE_BOOK.supports( + self.pricing_model, str(self.provider) + ) @property def request_id(self): @@ -55,5 +131,6 @@ def source_endpoint(self) -> str: "chat": "chat.completions", "responses": "responses", "messages": "messages", + "count_tokens": "messages.count_tokens", "generate_content": "generateContent", }[self.protocol] diff --git a/src/shim/gateway/pipeline/admission.py b/src/shim/gateway/pipeline/admission.py index 17b9d7b..1db91b1 100644 --- a/src/shim/gateway/pipeline/admission.py +++ b/src/shim/gateway/pipeline/admission.py @@ -68,9 +68,23 @@ async def run(self, value: PreparedInference) -> PreparedInference: and value.protocol == "responses" and value.model == UNSPECIFIED_PROVIDER_MODEL ) - if not unspecified_openai_response_model and not DEFAULT_PRICE_BOOK.supports( - value.model, str(value.provider) - ): + model_denied = ( + value.target is None + and not unspecified_openai_response_model + and not DEFAULT_PRICE_BOOK.supports(value.model, str(value.provider)) + ) + value.record_verdict( + "gateway.model_catalog", + stage="admission", + outcome="skip" + if value.target is not None + else ("deny" if model_denied else "allow"), + reason_code="DEPLOYMENT_AUTHORIZED" + if value.target is not None + else ("MODEL_NOT_PRICED" if model_denied else "MODEL_SUPPORTED"), + policy_version=DEFAULT_PRICE_BOOK.version, + ) + if model_denied: raise HTTPException( status_code=400, detail={ @@ -86,7 +100,7 @@ async def run(self, value: PreparedInference) -> PreparedInference: else None ) model_output_limit = DEFAULT_PRICE_BOOK.maximum_output_tokens( - value.model, + value.pricing_model, str(value.provider), ) # Reserve the ceiling when omitted without changing provider defaults. @@ -103,11 +117,16 @@ async def run(self, value: PreparedInference) -> PreparedInference: ), ("provider_default", model_output_limit), ) + if value.protocol == "count_tokens": + per_candidate_output_tokens = 0 minimum_output_tokens = ( 0 - if value.provider == "anthropic" - and value.protocol == "messages" - and output_token_field == "max_tokens" + if value.protocol == "count_tokens" + or ( + value.provider == "anthropic" + and value.protocol == "messages" + and output_token_field == "max_tokens" + ) else 1 ) if ( @@ -134,30 +153,38 @@ async def run(self, value: PreparedInference) -> PreparedInference: output_tokens = per_candidate_output_tokens * candidate_count(value) tier = value.context.tier_policy key_hash = value.policy.rate_limit_key_hash - if tier.rate_limit_rpm is not None and not await self.rate_limiter.allow( - key_hash, - limit=tier.rate_limit_rpm, - window_seconds=60, + for dimension, limit, key, amount in ( + ("requests", tier.rate_limit_rpm, key_hash, 1), + ("tokens", tier.rate_limit_tpm, f"tpm:{key_hash}", input_tokens), ): - raise HTTPException( - status_code=429, - detail={"code": "RATE_LIMIT_EXCEEDED", "dimension": "requests"}, + denied = limit is not None and not await self.rate_limiter.allow( + key, + limit=limit, + window_seconds=60, + amount=amount, ) - if tier.rate_limit_tpm is not None and not await self.rate_limiter.allow( - f"tpm:{key_hash}", - limit=tier.rate_limit_tpm, - window_seconds=60, - amount=input_tokens, - ): - raise HTTPException( - status_code=429, - detail={"code": "RATE_LIMIT_EXCEEDED", "dimension": "tokens"}, + value.record_verdict( + f"rate.{dimension}", + stage="admission", + outcome="deny" if denied else "allow" if limit is not None else "skip", + reason_code="RATE_LIMIT_EXCEEDED" + if denied + else "RATE_LIMIT_PASSED" + if limit is not None + else "RATE_LIMIT_UNLIMITED", + policy={"limit": limit, "window_seconds": 60}, ) + if denied: + raise HTTPException( + status_code=429, + detail={"code": "RATE_LIMIT_EXCEEDED", "dimension": dimension}, + ) repeat_material = _repeat_material( { **payload, "model": value.model, "provider": str(value.provider), + "protocol": value.protocol, } ) self.loop_result = await self.loop_detector.check_exact_repeat( @@ -166,7 +193,20 @@ async def run(self, value: PreparedInference) -> PreparedInference: limit=self.loop_repeat_limit, window_seconds=self.loop_window_seconds, ) - if self.loop_result.status == "BLOCKED": + repeated = self.loop_result.status == "BLOCKED" + value.record_verdict( + "rate.repeated_requests", + stage="admission", + outcome="deny" if repeated else "allow", + reason_code="REPEATED_REQUEST_LIMIT_EXCEEDED" + if repeated + else "REPEAT_CHECK_PASSED", + policy={ + "limit": self.loop_repeat_limit, + "window_seconds": self.loop_window_seconds, + }, + ) + if repeated: raise HTTPException( status_code=429, detail={ @@ -184,9 +224,11 @@ async def run(self, value: PreparedInference) -> PreparedInference: maximum_output_tokens=output_tokens, cost_center=attribution.cost_center, tags=attribution.tags, + repeat_chain_length=self.loop_result.chain_length or None, ) - await self.usage.admit(value, admission) - self.reserved = True + if value.protocol != "count_tokens": + await self.usage.admit(value, admission) + self.reserved = True return replace(value, admission=admission) def trace_metadata(self, output: PreparedInference) -> Mapping[str, TraceValue]: diff --git a/src/shim/gateway/pipeline/anthropic_execution.py b/src/shim/gateway/pipeline/anthropic_execution.py index d3f47c5..658a86f 100644 --- a/src/shim/gateway/pipeline/anthropic_execution.py +++ b/src/shim/gateway/pipeline/anthropic_execution.py @@ -47,6 +47,7 @@ def __init__( *, credential_resolver: ProviderCredentialResolver, circuit: CircuitBreaker, + circuit_for_target: Callable[[str], CircuitBreaker] | None = None, settings: CommunitySettings, http_client: httpx.AsyncClient, pii_scrubber: PIIScrubberService | None = None, @@ -55,6 +56,7 @@ def __init__( self.pii_scrubber = pii_scrubber or PIIScrubberService() self.http_client = http_client self.circuit = circuit + self.circuit_for_target = circuit_for_target self.settings = settings self.timeout = httpx.Timeout( connect=settings.ANTHROPIC_CONNECT_TIMEOUT_SECONDS, @@ -70,12 +72,22 @@ async def execute( prepared: PreparedInference, provider_start_callback: Callable[[], Awaitable[None]], ) -> ProviderNonStream | ProviderStream: + circuit = ( + self.circuit_for_target(prepared.target.base_url) + if prepared.target is not None and self.circuit_for_target is not None + else self.circuit + ) if prepared.privacy is None: raise RuntimeError("privacy stage must run before Anthropic execution") try: api_key = await self.credential_resolver.resolve( prepared.tenant_id, invocation.provider_credential, + **( + {"reference": prepared.target.credential_reference} + if prepared.target + else {} + ), ) except Exception: raise ProviderCallError( @@ -91,7 +103,7 @@ async def execute( False, provider="anthropic", ) - if not await self.circuit.acquire_call(): + if not await circuit.acquire_call(): raise ProviderCallError( 503, "PROVIDER_UNAVAILABLE", @@ -101,19 +113,27 @@ async def execute( try: client = AsyncAnthropic( api_key=api_key, - base_url=self.settings.ANTHROPIC_BASE_URL - or "https://api.anthropic.com", - timeout=self.timeout, + base_url=prepared.target.base_url + if prepared.target + else (self.settings.ANTHROPIC_BASE_URL or "https://api.anthropic.com"), + timeout=prepared.target.timeout_seconds + if prepared.target + else self.timeout, max_retries=0, http_client=self.http_client, ) await provider_start_callback() except BaseException: - await self.circuit.release_probe() + await circuit.release_probe() raise try: beta = _beta_enabled(invocation) - create = client.beta.messages.create if beta else client.messages.create + messages = client.beta.messages if beta else client.messages + create = ( + messages.count_tokens + if prepared.protocol == "count_tokens" + else messages.create + ) kwargs = sdk_create_kwargs( create, prepared.payload, @@ -131,12 +151,12 @@ async def execute( for item in headers["anthropic-beta"].split(",") if item.strip() ] - result = await create(**kwargs) + result: Any = await create(**kwargs) except asyncio.CancelledError: - await self.circuit.release_probe() + await circuit.release_probe() raise except Exception as exc: - await self._record_error(exc) + await self._record_error(exc, circuit) raise _public_error(exc) from None if prepared.stream: @@ -153,17 +173,17 @@ async def close_stream() -> None: pass finally: if not state["recorded"]: - await self.circuit.release_probe() + await circuit.release_probe() return ProviderStream( - events=self._stream(result, prepared, state, close_stream), + events=self._stream(result, prepared, state, close_stream, circuit), request_id=_stream_request_id(result), close=close_stream, ) payload = _dump_sdk(result) if _is_anthropic_failure(payload): - await self.circuit.record_failure() + await circuit.record_failure() raise ProviderCallError( 502, "PROVIDER_UNAVAILABLE", @@ -171,7 +191,7 @@ async def close_stream() -> None: provider="anthropic", request_id=getattr(result, "_request_id", None), ) - await self.circuit.record_success() + await circuit.record_success() payload = restore_anthropic_payload( payload, prepared.privacy.verification_map, @@ -188,6 +208,7 @@ async def _stream( prepared: PreparedInference, state: dict[str, bool], close_stream: Callable[[], Awaitable[None]], + circuit: CircuitBreaker, ) -> AsyncIterator[bytes]: assert prepared.privacy is not None restorer = AnthropicStreamRestorer( @@ -199,7 +220,7 @@ async def _stream( for payload in restorer.restore_events(_dump_sdk(event)): if _is_anthropic_failure(payload): if not state["recorded"]: - await self.circuit.record_failure() + await circuit.record_failure() state["recorded"] = True yield _error_event( ProviderCallError( @@ -212,13 +233,13 @@ async def _stream( ) return if payload.get("type") == "message_stop" and not state["recorded"]: - await self.circuit.record_success() + await circuit.record_success() state["recorded"] = True yield encode_responses_event(payload) if payload.get("type") == "message_stop": return if not state["recorded"]: - await self.circuit.record_failure() + await circuit.record_failure() state["recorded"] = True yield _error_event( ProviderCallError( @@ -231,21 +252,21 @@ async def _stream( except (asyncio.CancelledError, GeneratorExit): raise except Exception as exc: - await self._record_error(exc) + await self._record_error(exc, circuit) state["recorded"] = True yield _error_event(_public_error(exc)) finally: await close_stream() - async def _record_error(self, exc: Exception) -> None: + async def _record_error(self, exc: Exception, circuit: CircuitBreaker) -> None: if ( isinstance(exc, APIStatusError) and 400 <= exc.status_code < 500 and exc.status_code not in {408, 409, 429} ): - await self.circuit.record_success() + await circuit.record_success() else: - await self.circuit.record_failure() + await circuit.record_failure() def _dump_sdk(value: Any) -> dict[str, Any]: diff --git a/src/shim/gateway/pipeline/authenticate.py b/src/shim/gateway/pipeline/authenticate.py index 1c9f83d..41df9b9 100644 --- a/src/shim/gateway/pipeline/authenticate.py +++ b/src/shim/gateway/pipeline/authenticate.py @@ -6,7 +6,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Literal -from uuid import UUID, uuid4 +from uuid import uuid4 from fastapi import HTTPException @@ -14,7 +14,7 @@ GatewayContext, PrivacyPolicy, ) -from shim.gateway.contracts.ids import ApiKeyId, ProviderId, RequestId +from shim.gateway.contracts.ids import ProviderId, RequestId from shim.gateway.contracts.principal import AuthenticatedPrincipal from shim.gateway.kernel.result import PreparedInference from shim.gateway.kernel.stage import TraceValue @@ -37,7 +37,9 @@ class GatewayInvocation: principal: AuthenticatedPrincipal payload: dict[str, Any] provider: Literal["openai", "anthropic", "google"] - protocol: Literal["chat", "responses", "messages", "generate_content"] + protocol: Literal[ + "chat", "responses", "messages", "count_tokens", "generate_content" + ] model: str stream: bool headers: dict[str, str] @@ -55,6 +57,7 @@ def __init__( policy_resolver: RequestPolicyResolver, ) -> None: self.policy_resolver = policy_resolver + self.prepared: PreparedInference | None = None async def run(self, value: GatewayInvocation) -> PreparedInference: policy = await self.policy_resolver.resolve(value.principal) @@ -66,8 +69,12 @@ async def run(self, value: GatewayInvocation) -> PreparedInference: request_id=request_id, tenant_id=policy.tenant_id, actor_type=value.principal.actor_type, - api_key_id=ApiKeyId(UUID(str(value.principal.api_key_id))), - user_id=None, + api_key_id=value.principal.api_key_id, + user_id=( + value.principal.user_id + if value.principal.actor_type == "user_jwt" + else None + ), endpoint=metadata.endpoint, started_at=started_at, tier_policy=policy.tier_policy, @@ -78,9 +85,31 @@ async def run(self, value: GatewayInvocation) -> PreparedInference: ), audit_policy=policy.audit_policy, ) - if policy.tenant_policy.allowed_providers and value.provider not in { + self.prepared = prepared = PreparedInference( + context=context, + payload=dict(value.payload), + protocol=value.protocol, + model=value.model, + stream=value.stream, + policy=policy.request_policy, + pii_config=policy.pii_config, + provider=ProviderId(value.provider), + ) + provider_denied = bool( + policy.tenant_policy.allowed_providers + ) and value.provider not in { str(provider) for provider in policy.tenant_policy.allowed_providers - }: + } + prepared.record_verdict( + "tenant.allowed_providers", + stage="authentication", + outcome="deny" if provider_denied else "allow", + reason_code="PROVIDER_NOT_ALLOWED" + if provider_denied + else "PROVIDER_ALLOWED", + policy=policy.tenant_policy.allowed_providers, + ) + if provider_denied: raise HTTPException( status_code=403, detail={ @@ -88,10 +117,30 @@ async def run(self, value: GatewayInvocation) -> PreparedInference: "message": f"{value.provider.title()} is not allowed by tenant policy.", }, ) - if ( + retention_denied = ( policy.tenant_policy.require_zero_retention and not _zero_retention_requested(value) - ): + ) + prepared.record_verdict( + "tenant.zero_retention_request", + stage="authentication", + outcome=( + "deny" + if retention_denied + else "allow" + if policy.tenant_policy.require_zero_retention + else "skip" + ), + reason_code=( + "ZERO_RETENTION_REQUIRED" + if retention_denied + else "ZERO_RETENTION_REQUEST_CHECKED" + if policy.tenant_policy.require_zero_retention + else "ZERO_RETENTION_NOT_REQUIRED" + ), + policy=policy.tenant_policy.require_zero_retention, + ) + if retention_denied: raise HTTPException( status_code=422, detail={ @@ -99,16 +148,7 @@ async def run(self, value: GatewayInvocation) -> PreparedInference: "message": "Tenant policy requires a provider-enforceable zero-retention request.", }, ) - return PreparedInference( - context=context, - payload=dict(value.payload), - protocol=value.protocol, - model=value.model, - stream=value.stream, - policy=policy.request_policy, - pii_config=policy.pii_config, - provider=ProviderId(value.provider), - ) + return prepared def trace_metadata(self, output: PreparedInference) -> Mapping[str, TraceValue]: return { diff --git a/src/shim/gateway/pipeline/openai_execution.py b/src/shim/gateway/pipeline/openai_execution.py index 07b9f99..718822b 100644 --- a/src/shim/gateway/pipeline/openai_execution.py +++ b/src/shim/gateway/pipeline/openai_execution.py @@ -51,6 +51,7 @@ def __init__( *, credential_resolver: ProviderCredentialResolver, circuit: CircuitBreaker, + circuit_for_target: Callable[[str], CircuitBreaker] | None = None, settings: CommunitySettings, http_client: httpx.AsyncClient, chain_store: PrivacyContinuationStore, @@ -61,6 +62,7 @@ def __init__( self.http_client = http_client self.chain_store = chain_store self.circuit = circuit + self.circuit_for_target = circuit_for_target self.settings = settings self.timeout = httpx.Timeout( connect=settings.OPENAI_CONNECT_TIMEOUT_SECONDS, @@ -76,30 +78,44 @@ async def execute( prepared: PreparedInference, provider_start_callback: Callable[[], Awaitable[None]], ) -> ProviderNonStream | ProviderStream: + circuit = ( + self.circuit_for_target(prepared.target.base_url) + if prepared.target is not None and self.circuit_for_target is not None + else self.circuit + ) if prepared.privacy is None: raise RuntimeError("privacy stage must run before OpenAI execution") try: api_key = await self.credential_resolver.resolve( prepared.tenant_id, invocation.provider_credential, + **( + {"reference": prepared.target.credential_reference} + if prepared.target + else {} + ), ) except Exception: raise _error(503, "PROVIDER_UNAVAILABLE", False) from None if not api_key: raise _error(503, "PROVIDER_NOT_CONFIGURED", False) - if not await self.circuit.acquire_call(): + if not await circuit.acquire_call(): raise _error(503, "PROVIDER_UNAVAILABLE", True) try: client = AsyncOpenAI( api_key=api_key, - base_url=self.settings.OPENAI_BASE_URL or "https://api.openai.com/v1", - timeout=self.timeout, + base_url=prepared.target.base_url + if prepared.target + else (self.settings.OPENAI_BASE_URL or "https://api.openai.com/v1"), + timeout=prepared.target.timeout_seconds + if prepared.target + else self.timeout, max_retries=0, http_client=self.http_client, ) await provider_start_callback() except BaseException: - await self.circuit.release_probe() + await circuit.release_probe() raise try: create = ( @@ -116,10 +132,10 @@ async def execute( kwargs["extra_headers"] = headers result = await create(**kwargs) except asyncio.CancelledError: - await self.circuit.release_probe() + await circuit.release_probe() raise except Exception as exc: - await self._record_error(exc) + await self._record_error(exc, circuit) raise _public_error(exc) from None if prepared.stream: @@ -137,12 +153,12 @@ async def close_stream() -> None: pass finally: if not state["recorded"]: - await self.circuit.release_probe() + await circuit.release_probe() events = ( - self._responses_stream(result, prepared, state, close_stream) + self._responses_stream(result, prepared, state, close_stream, circuit) if prepared.protocol == "responses" - else self._chat_stream(result, prepared, state, close_stream) + else self._chat_stream(result, prepared, state, close_stream, circuit) ) return ProviderStream( events=events, @@ -154,7 +170,7 @@ async def close_stream() -> None: if _is_openai_failure(payload) and not ( prepared.protocol == "responses" and payload.get("status") == "failed" ): - await self.circuit.record_failure() + await circuit.record_failure() raise _error( 502, "PROVIDER_UNAVAILABLE", @@ -162,9 +178,9 @@ async def close_stream() -> None: request_id=getattr(result, "_request_id", None), ) if prepared.protocol == "responses" and payload.get("status") == "failed": - await self.circuit.record_failure() + await circuit.record_failure() else: - await self.circuit.record_success() + await circuit.record_success() response_id = payload.get("id") if prepared.protocol == "responses" and isinstance(response_id, str): await self.chain_store.save( @@ -190,6 +206,7 @@ async def _responses_stream( prepared: PreparedInference, state: dict[str, bool], close_stream: Callable[[], Awaitable[None]], + circuit: CircuitBreaker, ) -> AsyncIterator[bytes]: assert prepared.privacy is not None restorer = OpenAIStreamRestorer( @@ -210,7 +227,7 @@ async def _responses_stream( ) if _is_openai_failure(payload) and event_type != "response.failed": if not state["recorded"]: - await self.circuit.record_failure() + await circuit.record_failure() state["recorded"] = True yield _responses_error_event( "PROVIDER_UNAVAILABLE", @@ -237,15 +254,15 @@ async def _responses_stream( event_type in {"response.completed", "response.incomplete"} and not state["recorded"] ): - await self.circuit.record_success() + await circuit.record_success() state["recorded"] = True elif ( event_type in {"error", "response.failed"} and not state["recorded"] ): - await self.circuit.record_failure() + await circuit.record_failure() state["recorded"] = True elif event_type == "response.cancelled" and not state["recorded"]: - await self.circuit.record_success() + await circuit.record_success() state["recorded"] = True yield encode_responses_event(restored) if event_type in { @@ -257,13 +274,13 @@ async def _responses_stream( }: return if not state["recorded"]: - await self.circuit.record_failure() + await circuit.record_failure() state["recorded"] = True except (asyncio.CancelledError, GeneratorExit): raise except PrivacyContinuationUnavailableError: if not state["recorded"]: - await self.circuit.record_success() + await circuit.record_success() state["recorded"] = True yield _responses_error_event( "PRIVACY_STATE_UNAVAILABLE", @@ -271,7 +288,7 @@ async def _responses_stream( next_sequence_number, ) except Exception as exc: - await self._record_error(exc) + await self._record_error(exc, circuit) state["recorded"] = True error = _public_error(exc) yield _responses_error_event( @@ -288,6 +305,7 @@ async def _chat_stream( prepared: PreparedInference, state: dict[str, bool], close_stream: Callable[[], Awaitable[None]], + circuit: CircuitBreaker, ) -> AsyncIterator[bytes]: assert prepared.privacy is not None restorer = OpenAIStreamRestorer( @@ -301,7 +319,7 @@ async def _chat_stream( payload = _dump_sdk(chunk) if _is_openai_failure(payload): if not state["recorded"]: - await self.circuit.record_failure() + await circuit.record_failure() state["recorded"] = True yield _chat_error_event( _error( @@ -314,14 +332,14 @@ async def _chat_stream( return finished_choices.update(_finished_chat_choices(payload)) if len(finished_choices) >= expected_choices and not state["recorded"]: - await self.circuit.record_success() + await circuit.record_success() state["recorded"] = True restored = restorer.restore_chat_chunk(payload) yield encode_data(restored) if len(finished_choices) >= expected_choices: yield b"data: [DONE]\n\n" else: - await self.circuit.record_failure() + await circuit.record_failure() state["recorded"] = True yield _chat_error_event(_error(502, "PROVIDER_UNAVAILABLE", False)) except (asyncio.CancelledError, GeneratorExit): @@ -330,21 +348,21 @@ async def _chat_stream( if state["recorded"] and not isinstance(exc, ValueError): yield b"data: [DONE]\n\n" return - await self._record_error(exc) + await self._record_error(exc, circuit) state["recorded"] = True yield _chat_error_event(_public_error(exc)) finally: await close_stream() - async def _record_error(self, exc: Exception) -> None: + async def _record_error(self, exc: Exception, circuit: CircuitBreaker) -> None: if ( isinstance(exc, APIStatusError) and exc.status_code < 500 and exc.status_code not in {408, 409, 429} ): - await self.circuit.record_success() + await circuit.record_success() else: - await self.circuit.record_failure() + await circuit.record_failure() def _dump_sdk(value: Any) -> dict[str, Any]: diff --git a/src/shim/gateway/pipeline/postprocess.py b/src/shim/gateway/pipeline/postprocess.py index fe0ab2a..b788bce 100644 --- a/src/shim/gateway/pipeline/postprocess.py +++ b/src/shim/gateway/pipeline/postprocess.py @@ -20,7 +20,7 @@ StreamSession, StreamTerminalStatus, ) -from shim.gateway.streaming.meter import StreamUsageSnapshot +from shim.gateway.streaming.meter import StreamUsageSnapshot, native_finish_reasons from shim.gateway.usage import UsageLifecycle from shim.observability.metrics import ( PROVIDER_LATENCY_MS, @@ -75,6 +75,7 @@ async def finalize( ) -> JSONResponse | StreamingResponse: if isinstance(response, ProviderStream): assert stream_session is not None + stream_session.meter.started_at_monotonic = response.started_at_monotonic stream_session.bind( response.events, close=response.close, @@ -115,13 +116,14 @@ async def finalize( and lifecycle_status == "completed" and isinstance(response_model, str) and DEFAULT_PRICE_BOOK.supports(response_model, provider) - else prepared.model + else prepared.pricing_model ) settlement_cost = compute_cost_usd( settlement_model, prompt_tokens, completion_tokens, provider=provider, + unpriced=prepared.unpriced, ) if response.latency_ms is not None: labels = { @@ -154,8 +156,12 @@ async def finalize( provider, input_tokens=prompt_tokens, output_tokens=completion_tokens, + unpriced=prepared.unpriced, ), estimated=not fully_actual, + provider_finish_reasons=native_finish_reasons( + response.payload, provider=provider + ), output_hash=( content_ref( self.output_hash_salt, bytes(gateway_response.body).decode() @@ -216,7 +222,8 @@ def observe_terminal(terminal_status: str) -> None: return StreamSession( meter=StreamMeter( provider=str(prepared.provider), - requested_model=prepared.model, + requested_model=prepared.pricing_model, + unpriced=prepared.unpriced, prompt_tokens_estimated=prepared.admission.estimated_input_tokens, expected_candidates=candidate_count(prepared), output_hash_salt=self.output_hash_salt, diff --git a/src/shim/gateway/pipeline/privacy.py b/src/shim/gateway/pipeline/privacy.py index e9ebedc..c6fc573 100644 --- a/src/shim/gateway/pipeline/privacy.py +++ b/src/shim/gateway/pipeline/privacy.py @@ -149,6 +149,40 @@ def __init__( self.chain_store = chain_store async def run(self, value: PreparedInference) -> PreparedInference: + try: + prepared = await self._run(value) + except BaseException as error: + value.record_verdict( + "privacy.input", + stage="privacy", + outcome="deny" + if isinstance(error, HTTPException) and error.status_code < 500 + else "error", + reason_code="PRIVACY_POLICY_BLOCKED" + if isinstance(error, HTTPException) and error.status_code < 500 + else "PRIVACY_UNAVAILABLE", + policy=value.pii_config, + ) + raise + assert prepared.privacy is not None + prepared.record_verdict( + "privacy.input", + stage="privacy", + outcome="mask" + if prepared.privacy.pii_detected + else "skip" + if value.context.privacy_policy.pii_mode == "disabled" + else "allow", + reason_code="PII_MASKED" + if prepared.privacy.pii_detected + else "PII_DISABLED" + if value.context.privacy_policy.pii_mode == "disabled" + else "PII_NOT_DETECTED", + policy=value.pii_config, + ) + return prepared + + async def _run(self, value: PreparedInference) -> PreparedInference: parent_map: dict[str, str] = {} previous_response_id = ( value.payload.get("previous_response_id") diff --git a/src/shim/gateway/pipeline/provider_execution.py b/src/shim/gateway/pipeline/provider_execution.py index 078299b..3f74128 100644 --- a/src/shim/gateway/pipeline/provider_execution.py +++ b/src/shim/gateway/pipeline/provider_execution.py @@ -45,6 +45,7 @@ class ProviderStream: request_id: str | None close: Callable[[], Awaitable[None]] prefetched_events: tuple[bytes, ...] = () + started_at_monotonic: float | None = None class ProviderExecutionStage: @@ -61,8 +62,12 @@ def __init__( self.usage = usage async def run(self, value: PreparedInference) -> ProviderNonStream | ProviderStream: + provider_started_at: float | None = None + async def mark_started() -> None: + nonlocal provider_started_at await self.usage.mark_provider_started(value) + provider_started_at = perf_counter() started_at = perf_counter() try: @@ -76,7 +81,7 @@ async def mark_started() -> None: raise if isinstance(output, ProviderNonStream): return replace(output, latency_ms=(perf_counter() - started_at) * 1_000) - return output + return replace(output, started_at_monotonic=provider_started_at) def trace_metadata( self, diff --git a/src/shim/gateway/pipeline/provider_spend.py b/src/shim/gateway/pipeline/provider_spend.py index b3e48eb..58f74bd 100644 --- a/src/shim/gateway/pipeline/provider_spend.py +++ b/src/shim/gateway/pipeline/provider_spend.py @@ -23,7 +23,9 @@ async def run(self, value: PreparedInference) -> PreparedInference: credential = self.invocation.provider_credential await self.usage.reserve_provider_spend( value, - ephemeral_byok=credential is not None and credential.available(), + ephemeral_byok=value.target is None + and credential is not None + and credential.available(), ) return value diff --git a/src/shim/gateway/streaming/meter.py b/src/shim/gateway/streaming/meter.py index 95842e4..532f160 100644 --- a/src/shim/gateway/streaming/meter.py +++ b/src/shim/gateway/streaming/meter.py @@ -5,8 +5,10 @@ import json import hashlib import math +from collections.abc import Callable, Mapping from dataclasses import dataclass from decimal import Decimal +from time import perf_counter from typing import Any, Literal from shim.billing.pricing import ( @@ -48,6 +50,8 @@ class StreamUsageSnapshot: pricing_metadata: dict[str, str | int] estimated: bool output_hash: str | None + provider_finish_reasons: dict[str, str] | None = None + ttft_ms: float | None = None class StreamMeter: @@ -66,12 +70,16 @@ def __init__( prompt_tokens_estimated: int, expected_candidates: int = 1, output_hash_salt: str | None = None, + started_at_monotonic: float | None = None, + monotonic_clock: Callable[[], float] = perf_counter, + unpriced: bool = False, ) -> None: if prompt_tokens_estimated < 0: raise ValueError("stream token estimates must be nonnegative") if expected_candidates < 1: raise ValueError("expected stream candidates must be positive") self.provider = provider + self.unpriced = unpriced self.requested_model = requested_model self.prompt_tokens_estimated = prompt_tokens_estimated self.expected_candidates = expected_candidates @@ -80,6 +88,10 @@ def __init__( self.response_model: str | None = None self.emitted_output_characters = 0 self.terminal_hint: StreamTerminalHint | None = None + self.provider_finish_reasons: dict[str, str] = {} + self.started_at_monotonic = started_at_monotonic + self._monotonic_clock = monotonic_clock + self.ttft_ms: float | None = None self._finished_candidates: set[int] = set() self._sse_buffer = b"" self._output_hasher = hashlib.sha256() if output_hash_salt is not None else None @@ -147,6 +159,7 @@ def snapshot(self) -> StreamUsageSnapshot: prompt, completion, provider=self.provider, + unpriced=self.unpriced, ) return StreamUsageSnapshot( prompt_tokens=prompt, @@ -158,6 +171,7 @@ def snapshot(self) -> StreamUsageSnapshot: self.provider, input_tokens=prompt, output_tokens=completion, + unpriced=self.unpriced, ), estimated=estimated, output_hash=( @@ -165,6 +179,8 @@ def snapshot(self) -> StreamUsageSnapshot: if self._output_hasher is not None and self._emitted_wire_bytes > 0 else None ), + provider_finish_reasons=dict(self.provider_finish_reasons) or None, + ttft_ms=self.ttft_ms, ) def _observe_sse_event(self, event_text: str) -> None: @@ -195,10 +211,25 @@ def _observe_sse_event(self, event_text: str) -> None: self._capture_terminal_hint(payload_type, payload) self._capture_response_model(payload) self._capture_usage(payload) - self.emitted_output_characters += self._output_delta_characters( + self.provider_finish_reasons.update( + native_finish_reasons(payload, provider=self.provider) or {} + ) + output_characters = self._output_delta_characters( payload_type, payload, ) + self.emitted_output_characters += output_characters + if ( + self.ttft_ms is None + and self.started_at_monotonic is not None + and ( + output_characters > 0 + or _initial_or_media_content(payload_type, payload) + ) + ): + self.ttft_ms = max( + 0.0, (self._monotonic_clock() - self.started_at_monotonic) * 1_000 + ) def _capture_response_model(self, payload: dict[str, Any]) -> None: candidates = [payload] @@ -441,6 +472,161 @@ def _sum_optional_counts(*values: int | None) -> int | None: return sum(present) if present else None +def _initial_or_media_content(event_type: str, payload: Mapping[str, Any]) -> bool: + """Recognize content readiness without treating opaque media as text tokens.""" + + if event_type == "content_block_start": + block = payload.get("content_block") + if isinstance(block, Mapping) and any( + isinstance(block.get(field), str) and block[field] + for field in ("text", "thinking") + ): + return True + field = { + "response.audio.delta": "delta", + "response.image_generation_call.partial_image": "partial_image_b64", + }.get(event_type) + if field is not None and isinstance(payload.get(field), str) and payload[field]: + return True + choices = payload.get("choices") + if isinstance(choices, list): + for choice in choices: + delta = choice.get("delta") if isinstance(choice, Mapping) else None + audio = delta.get("audio") if isinstance(delta, Mapping) else None + if ( + isinstance(audio, Mapping) + and isinstance(audio.get("data"), str) + and audio["data"] + ): + return True + candidates = payload.get("candidates") + if isinstance(candidates, list): + for candidate in candidates: + content = ( + candidate.get("content") if isinstance(candidate, Mapping) else None + ) + parts = content.get("parts") if isinstance(content, Mapping) else None + if not isinstance(parts, list): + continue + for part in parts: + media = part.get("inlineData") if isinstance(part, Mapping) else None + if ( + isinstance(media, Mapping) + and isinstance(media.get("data"), str) + and media["data"] + ): + return True + return False + + +# Only native enum values enter telemetry; free-form provider fields may contain PII. +_CHAT_REASONS = frozenset( + {"stop", "length", "tool_calls", "content_filter", "function_call"} +) +_MESSAGE_REASONS = frozenset( + { + "end_turn", + "max_tokens", + "stop_sequence", + "tool_use", + "pause_turn", + "refusal", + "model_context_window_exceeded", + } +) +_GOOGLE_REASONS = frozenset( + { + "STOP", + "MAX_TOKENS", + "SAFETY", + "RECITATION", + "LANGUAGE", + "OTHER", + "BLOCKLIST", + "PROHIBITED_CONTENT", + "SPII", + "MALFORMED_FUNCTION_CALL", + "IMAGE_SAFETY", + "UNEXPECTED_TOOL_CALL", + "IMAGE_PROHIBITED_CONTENT", + "NO_IMAGE", + "IMAGE_RECITATION", + "IMAGE_OTHER", + } +) +_GOOGLE_BLOCK_REASONS = frozenset( + { + "SAFETY", + "OTHER", + "BLOCKLIST", + "PROHIBITED_CONTENT", + "IMAGE_SAFETY", + "MODEL_ARMOR", + "JAILBREAK", + } +) + + +def native_finish_reasons( + payload: Mapping[str, Any], *, provider: str +) -> dict[str, str] | None: + """Keep native, candidate-indexed completion facts separate from transport status.""" + + reasons: dict[str, str] = {} + if provider in {"openai", "google"}: + collection, field, allowed = ( + ("candidates", "finishReason", _GOOGLE_REASONS) + if provider == "google" + else ("choices", "finish_reason", _CHAT_REASONS) + ) + candidates = payload.get(collection) + if isinstance(candidates, list): + for position, candidate in enumerate(candidates): + if not isinstance(candidate, Mapping): + continue + index = candidate.get("index", position) + reason = candidate.get(field) + if ( + isinstance(index, int) + and not isinstance(index, bool) + and 0 <= index < 10_000 + and isinstance(reason, str) + and reason in allowed + ): + reasons[f"{collection}.{index}.{field}"] = reason + if provider == "openai": + response = payload.get("response", payload) + if isinstance(response, Mapping): + status = response.get("status") + if isinstance(status, str) and status in { + "completed", + "failed", + "cancelled", + "incomplete", + }: + reasons["status"] = status + incomplete = response.get("incomplete_details") + reason = ( + incomplete.get("reason") if isinstance(incomplete, Mapping) else None + ) + if isinstance(reason, str) and reason in { + "max_output_tokens", + "content_filter", + }: + reasons["incomplete_details.reason"] = reason + elif provider == "anthropic": + for value in (payload, payload.get("delta"), payload.get("message")): + reason = value.get("stop_reason") if isinstance(value, Mapping) else None + if isinstance(reason, str) and reason in _MESSAGE_REASONS: + reasons["stop_reason"] = reason + elif provider == "google": + feedback = payload.get("promptFeedback") + reason = feedback.get("blockReason") if isinstance(feedback, Mapping) else None + if isinstance(reason, str) and reason in _GOOGLE_BLOCK_REASONS: + reasons["promptFeedback.blockReason"] = reason + return reasons or None + + def _google_content_strings(value: Any, *, content: bool = False) -> list[str]: if isinstance(value, str): return [value] if content else [] diff --git a/src/shim/gateway/usage.py b/src/shim/gateway/usage.py index 816a2eb..bd68800 100644 --- a/src/shim/gateway/usage.py +++ b/src/shim/gateway/usage.py @@ -31,7 +31,13 @@ class UsageLimitExceeded(RuntimeError): """An authoritative usage policy denied admission.""" +class UsageAuditPersistenceError(RuntimeError): + """Required audit evidence could not be committed before a response.""" + + class UsageLifecycle(Protocol): + async def reject(self, prepared: PreparedInference) -> None: ... + async def admit( self, prepared: PreparedInference, @@ -40,6 +46,10 @@ async def admit( async def record_privacy(self, prepared: PreparedInference) -> None: ... + async def record_token_count( + self, prepared: PreparedInference, input_tokens: int | None + ) -> None: ... + async def reserve_provider_spend( self, prepared: PreparedInference, @@ -107,9 +117,28 @@ async def admit( ) -> None: pass + async def reject(self, prepared: PreparedInference) -> None: + self._write( + prepared, + outcome="rejected", + completed_at=datetime.now(timezone.utc), + prompt_tokens=0, + completion_tokens=0, + cost_usd=Decimal("0"), + model=prepared.model + if DEFAULT_PRICE_BOOK.supports(prepared.model, str(prepared.provider)) + else "unsupported", + estimated=False, + ) + async def record_privacy(self, prepared: PreparedInference) -> None: pass + async def record_token_count( + self, prepared: PreparedInference, input_tokens: int | None + ) -> None: + pass + async def reserve_provider_spend( self, prepared: PreparedInference, @@ -149,6 +178,8 @@ async def finalize( ), model=usage.provider_model, estimated=usage.estimated, + provider_finish_reasons=usage.provider_finish_reasons, + ttft_ms=usage.ttft_ms, ) async def fail( @@ -194,6 +225,8 @@ def _write( cost_usd: Decimal | None, model: str, estimated: bool, + provider_finish_reasons: dict[str, str] | None = None, + ttft_ms: float | None = None, ) -> None: event = { "version": 1, @@ -211,11 +244,23 @@ def _write( "completion_tokens": completion_tokens, "estimated_cost_usd": str(cost_usd) if cost_usd is not None else None, "estimated": estimated, + "provider_finish_reasons": provider_finish_reasons, + "ttft_ms": ttft_ms, + "repeat_chain_length": ( + prepared.admission.repeat_chain_length + if prepared.admission is not None + else None + ), + "system_prompt_hash": None, + "deployment_kind": prepared.deployment_kind, "privacy_counts": ( dict(prepared.privacy.pii_entities) if prepared.privacy is not None else {} ), + "policy_verdicts": [ + verdict.model_dump(mode="json") for verdict in prepared.policy_verdicts + ], } line = json.dumps(event, ensure_ascii=False, separators=(",", ":")) try: diff --git a/src/shim/observability/metrics.py b/src/shim/observability/metrics.py index 2c85386..80adafe 100644 --- a/src/shim/observability/metrics.py +++ b/src/shim/observability/metrics.py @@ -15,6 +15,7 @@ "/v1/chat/completions", "/v1/responses", "/v1/messages", + "/v1/messages/count_tokens", "/v1beta/models/*:generateContent", "/v1beta/models/*:streamGenerateContent", "/v1/scan", @@ -93,7 +94,14 @@ ), "method": frozenset({"GET", "POST"}), "protocol": frozenset( - {"chat", "responses", "messages", "generate_content", "scan"} + { + "chat", + "responses", + "messages", + "count_tokens", + "generate_content", + "scan", + } ), "actor_type": frozenset({"api_key", "user_jwt", "internal"}), "source_endpoint": frozenset( diff --git a/src/shim/secrets/credentials.py b/src/shim/secrets/credentials.py index aada224..440dbd9 100644 --- a/src/shim/secrets/credentials.py +++ b/src/shim/secrets/credentials.py @@ -65,6 +65,8 @@ async def resolve( self, tenant_id: TenantId, credential: EphemeralProviderCredential | None, + *, + reference: str | None = None, ) -> str | None: ... @@ -87,8 +89,12 @@ async def resolve( self, tenant_id: TenantId, credential: EphemeralProviderCredential | None, + *, + reference: str | None = None, ) -> str | None: del tenant_id + if reference is not None: + raise ValueError("local credentials do not support managed references") if credential is not None and credential.provider != self.provider: raise ValueError("credential does not match the selected provider") injected = credential.consume() if credential is not None else None diff --git a/src/shim/services/gateway/service.py b/src/shim/services/gateway/service.py index 0dd98ef..6202a5b 100644 --- a/src/shim/services/gateway/service.py +++ b/src/shim/services/gateway/service.py @@ -30,7 +30,9 @@ async def dispatch_inference( *, payload: dict[str, Any], provider: Literal["openai", "anthropic", "google"], - protocol: Literal["chat", "responses", "messages", "generate_content"], + protocol: Literal[ + "chat", "responses", "messages", "count_tokens", "generate_content" + ], model: str, stream: bool, headers: dict[str, str], diff --git a/tests/architecture/test_module_ownership.py b/tests/architecture/test_module_ownership.py index e7fdf6e..4dd0f64 100644 --- a/tests/architecture/test_module_ownership.py +++ b/tests/architecture/test_module_ownership.py @@ -12,6 +12,7 @@ OWNERS = ("public", "enterprise", "split") PYTHON_ROOTS = ( "ee/alembic", + "ee/deploy/test", "ee/scripts", "ee/src/shim_enterprise", "ee/tests", @@ -331,6 +332,7 @@ def test_python_ownership_regions_are_canonical() -> None: "public": ("scripts/", "src/shim/", "tests/"), "enterprise": ( "ee/alembic/", + "ee/deploy/test/", "ee/scripts/", "ee/src/shim_enterprise/", "ee/tests/", diff --git a/tests/community/test_application.py b/tests/community/test_application.py index ec534b3..bc0fd40 100644 --- a/tests/community/test_application.py +++ b/tests/community/test_application.py @@ -120,6 +120,11 @@ async def handler(request: httpx.Request) -> httpx.Response: event = json.loads(lines[0]) assert event["outcome"] == "completed" assert event["privacy_counts"] == {"EMAIL_ADDRESS": 1} + assert event["provider_finish_reasons"] == {"choices.0.finish_reason": "stop"} + assert event["ttft_ms"] is None + assert event["repeat_chain_length"] == 1 + assert event["deployment_kind"] == "unknown" + assert event["system_prompt_hash"] is None assert EMAIL not in lines[0] assert GATEWAY_KEY not in lines[0] assert PROVIDER_KEY not in lines[0] @@ -237,7 +242,11 @@ async def handler(request: httpx.Request) -> httpx.Response: await upstream.aclose() lines = events.getvalue().splitlines() assert len(lines) == 1 - assert json.loads(lines[0])["outcome"] == "completed" + event = json.loads(lines[0]) + assert event["outcome"] == "completed" + assert event["provider_finish_reasons"] == {"choices.0.finish_reason": "stop"} + assert event["ttft_ms"] >= 0 + assert event["repeat_chain_length"] == 1 @pytest.mark.asyncio @@ -352,6 +361,7 @@ def test_community_factory_has_no_enterprise_import_or_environment_dependency() "/health", "/v1/chat/completions", "/v1/messages", + "/v1/messages/count_tokens", "/v1/models", "/v1/models/{model_id}", "/v1/responses", diff --git a/tests/community/test_models.py b/tests/community/test_models.py index a76568c..210a091 100644 --- a/tests/community/test_models.py +++ b/tests/community/test_models.py @@ -131,7 +131,7 @@ async def test_missing_models_raise_native_sanitized_sdk_errors() -> None: assert openai_error.value.response.json() == { "error": { - "message": "The requested model is not in the public catalog.", + "message": "The requested model is not available.", "type": "not_found_error", "param": None, "code": "MODEL_NOT_FOUND", @@ -141,7 +141,7 @@ async def test_missing_models_raise_native_sanitized_sdk_errors() -> None: "type": "error", "error": { "type": "not_found_error", - "message": "The requested model is not in the public catalog.", + "message": "The requested model is not available.", }, } assert openai_missing not in openai_error.value.response.text diff --git a/tests/gateway/kernel/test_gateway_kernel.py b/tests/gateway/kernel/test_gateway_kernel.py index 3ed1626..fa2c2ff 100644 --- a/tests/gateway/kernel/test_gateway_kernel.py +++ b/tests/gateway/kernel/test_gateway_kernel.py @@ -36,7 +36,7 @@ async def test_kernel_runs_the_authoritative_stage_order( monkeypatch: pytest.MonkeyPatch, ) -> None: order: list[str] = [] - prepared = SimpleNamespace(stream=False) + prepared = SimpleNamespace(stream=False, protocol="chat") provider_output = object() response = Response("ok") @@ -187,7 +187,7 @@ async def test_kernel_maps_provider_failures_to_usage_reason( status_code: int, expected_reason: str, ) -> None: - prepared = SimpleNamespace(stream=False) + prepared = SimpleNamespace(stream=False, protocol="chat") failure = ProviderCallError( status_code=status_code, error_code="PROVIDER_UNAVAILABLE", @@ -214,7 +214,7 @@ async def run_stage(stage, _value): async def test_kernel_maps_post_reservation_admission_abort( monkeypatch: pytest.MonkeyPatch, ) -> None: - prepared = SimpleNamespace(stream=False) + prepared = SimpleNamespace(stream=False, protocol="chat") failure = RuntimeError("admission interrupted") async def run_stage(stage, _value): @@ -237,7 +237,7 @@ async def run_stage(stage, _value): async def test_recovery_session_failure_does_not_mask_original_error( monkeypatch: pytest.MonkeyPatch, ) -> None: - prepared = SimpleNamespace(stream=False) + prepared = SimpleNamespace(stream=False, protocol="chat") failure = RuntimeError("admission interrupted") async def run_stage(stage, _value): diff --git a/tests/gateway/pipeline/test_admission.py b/tests/gateway/pipeline/test_admission.py index 6e1fde7..1d99d47 100644 --- a/tests/gateway/pipeline/test_admission.py +++ b/tests/gateway/pipeline/test_admission.py @@ -128,6 +128,7 @@ async def test_admission_bounds_provider_payloads_and_output_limits() -> None: "key-hash", limit=60, window_seconds=60, + amount=1, ) usage.admit.assert_awaited_once() diff --git a/tests/gateway/pipeline/test_authenticate.py b/tests/gateway/pipeline/test_authenticate.py index 1a7adaf..1b1cbba 100644 --- a/tests/gateway/pipeline/test_authenticate.py +++ b/tests/gateway/pipeline/test_authenticate.py @@ -82,6 +82,16 @@ async def test_authenticate_uses_session_free_policy_values_only() -> None: assert not hasattr(prepared, "db") assert not hasattr(prepared.policy, "__dict__") policy_resolver.resolve.assert_awaited_once_with(principal) + second = await AuthenticateStage(policy_resolver).run(invocation) + assert second.request_id != prepared.request_id + assert second.policy_verdicts is not prepared.policy_verdicts + second.record_verdict( + "privacy.input", + stage="privacy", + outcome="deny", + reason_code="PRIVACY_POLICY_BLOCKED", + ) + assert all(verdict.outcome != "deny" for verdict in prepared.policy_verdicts) @pytest.mark.asyncio diff --git a/tests/gateway/streaming/test_meter.py b/tests/gateway/streaming/test_meter.py index 1ce1714..00405f0 100644 --- a/tests/gateway/streaming/test_meter.py +++ b/tests/gateway/streaming/test_meter.py @@ -1,11 +1,13 @@ from __future__ import annotations from decimal import Decimal +import json import pytest from shim.billing.pricing import UNSPECIFIED_PROVIDER_MODEL from shim.gateway.streaming import StreamMeter +from shim.gateway.streaming.meter import native_finish_reasons def meter( @@ -272,3 +274,138 @@ def test_timeout_codes_are_independent_of_error_prose(event): stream_meter = meter() stream_meter.observe_sse(event) assert stream_meter.terminal_hint == "timeout" + + +@pytest.mark.parametrize( + ("provider", "payload", "expected"), + [ + ( + "openai", + { + "choices": [ + {"index": 2, "finish_reason": "length"}, + {"index": 0, "finish_reason": "stop"}, + ] + }, + {"choices.2.finish_reason": "length", "choices.0.finish_reason": "stop"}, + ), + ( + "openai", + { + "status": "incomplete", + "incomplete_details": {"reason": "max_output_tokens"}, + }, + {"status": "incomplete", "incomplete_details.reason": "max_output_tokens"}, + ), + ("anthropic", {"stop_reason": "refusal"}, {"stop_reason": "refusal"}), + ( + "google", + {"candidates": [{"index": 1, "finishReason": "SAFETY"}]}, + {"candidates.1.finishReason": "SAFETY"}, + ), + ( + "google", + {"promptFeedback": {"blockReason": "SAFETY"}}, + {"promptFeedback.blockReason": "SAFETY"}, + ), + ("openai", {"choices": [{"finish_reason": "private-provider-text"}]}, None), + ("openai", {"error": {"message": "private-provider-text"}}, None), + ( + "google", + {"candidates": [{"finishReason": "FINISH_REASON_UNSPECIFIED"}]}, + None, + ), + ("openai", {}, None), + ], +) +def test_native_finish_facts_match_json_and_sse(provider, payload, expected) -> None: + stream_meter = meter(provider) + stream_meter.observe_sse(f"data: {json.dumps(payload)}\n\n".encode()) + + assert native_finish_reasons(payload, provider=provider) == expected + assert stream_meter.snapshot().provider_finish_reasons == expected + assert stream_meter.snapshot().ttft_ms is None + + +def test_ttft_ignores_heartbeats_metadata_empty_deltas_and_usage() -> None: + now = 10.0 + stream_meter = StreamMeter( + provider="openai", + requested_model="gpt-5.6-luna", + prompt_tokens_estimated=5, + started_at_monotonic=now, + monotonic_clock=lambda: now, + ) + stream_meter.observe_sse( + b': heartbeat\n\nevent: response.created\ndata: {"response":{"status":"in_progress"}}\n\n' + b'data: {"choices":[{"delta":{"role":"assistant","content":""}}],"usage":{"prompt_tokens":5}}\n\n' + ) + assert stream_meter.snapshot().ttft_ms is None + now = 10.125 + stream_meter.observe_sse(b'data: {"choices":[{"delta":{"content":"hel') + assert stream_meter.snapshot().ttft_ms is None + stream_meter.observe_sse(b'lo"}}]}\n\n') + now = 11.0 + stream_meter.observe_sse( + b'data: {"choices":[{"index":0,"finish_reason":"length"}]}\n\ndata: [DONE]\n\n' + ) + snapshot = stream_meter.snapshot() + assert snapshot.ttft_ms == 125.0 + assert snapshot.provider_finish_reasons == {"choices.0.finish_reason": "length"} + assert snapshot.estimated is True + assert stream_meter.terminal_hint == "completed" + + +@pytest.mark.parametrize( + ("provider", "payload"), + [ + ("openai", {"type": "response.audio.delta", "delta": "opaque-audio"}), + ( + "openai", + { + "type": "response.image_generation_call.partial_image", + "partial_image_b64": "opaque-image", + }, + ), + ("openai", {"choices": [{"delta": {"audio": {"data": "opaque-audio"}}}]}), + ( + "google", + { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "data": "opaque-image", + "mimeType": "image/png", + } + } + ] + } + } + ] + }, + ), + ( + "anthropic", + { + "type": "content_block_start", + "content_block": {"type": "text", "text": "hello"}, + }, + ), + ], +) +def test_initial_and_media_content_sets_ttft_without_changing_token_estimates( + provider, payload +) -> None: + stream_meter = StreamMeter( + provider=provider, + requested_model="gpt-5.6-luna", + prompt_tokens_estimated=5, + started_at_monotonic=10.0, + monotonic_clock=lambda: 10.25, + ) + stream_meter.observe_sse(f"data: {json.dumps(payload)}\n\n".encode()) + assert stream_meter.snapshot().ttft_ms == 250.0 + assert stream_meter.snapshot().completion_tokens == 0 diff --git a/tests/gateway/test_anthropic_sdk_transport.py b/tests/gateway/test_anthropic_sdk_transport.py index f37f2f6..d312521 100644 --- a/tests/gateway/test_anthropic_sdk_transport.py +++ b/tests/gateway/test_anthropic_sdk_transport.py @@ -36,6 +36,8 @@ def _prepared(payload: dict, mapping: dict[str, str] | None = None): return SimpleNamespace( payload=payload, + protocol="messages", + target=None, tenant_id=TenantId(UUID("11111111-1111-1111-1111-111111111111")), stream=bool(payload.get("stream")), privacy=PrivacyOutcome( diff --git a/tests/gateway/test_openai_sdk_transport.py b/tests/gateway/test_openai_sdk_transport.py index b10fbd5..6ca2626 100644 --- a/tests/gateway/test_openai_sdk_transport.py +++ b/tests/gateway/test_openai_sdk_transport.py @@ -68,6 +68,7 @@ def _prepared( ) return SimpleNamespace( payload=payload, + target=None, tenant_id=TenantId(UUID(tenant)), protocol=protocol, stream=bool(payload.get("stream")), @@ -1016,6 +1017,7 @@ async def chunks(): ), state, close_stream, + execution.circuit, ) ] ) @@ -1082,6 +1084,7 @@ async def chunks(): ), {"closed": False, "recorded": False}, AsyncMock(), + execution.circuit, ) ] ) diff --git a/tests/gateway/test_token_count.py b/tests/gateway/test_token_count.py new file mode 100644 index 0000000..b1153e8 --- /dev/null +++ b/tests/gateway/test_token_count.py @@ -0,0 +1,176 @@ +"""Anthropic token counting shares auth/privacy but never the inference ledger.""" + +import io +import json +from unittest.mock import AsyncMock + +import httpx +import pytest +from anthropic import AsyncAnthropic, AuthenticationError + +from shim.application import create_community_app +from shim.core.community_config import CommunitySettings + + +@pytest.mark.asyncio +@pytest.mark.parametrize("beta", [False, True]) +async def test_native_token_count_scrubs_and_does_not_execute_or_settle(beta): + calls = [] + + def upstream(request): + calls.append(request) + assert request.url.path == "/v1/messages/count_tokens" + assert request.headers["x-api-key"] == "provider-secret" + assert "alice@example.com" not in request.content.decode() + assert json.loads(request.content)["messages"][0]["role"] == "user" + return httpx.Response( + 200, json={"input_tokens": 17}, headers={"request-id": "count-1"} + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(upstream)) as outbound: + app = create_community_app( + CommunitySettings(_env_file=None, SHIM_API_KEY="gateway-secret-12345"), + http_client=outbound, + event_stream=io.StringIO(), + ) + async with app.router.lifespan_context(app): + usage = AsyncMock() + app.state.gateway_service.kernel.usage = usage + async with httpx.AsyncClient(transport=httpx.ASGITransport(app)) as inbound: + sdk = AsyncAnthropic( + api_key="gateway-secret-12345", + base_url="http://shim.test", + http_client=inbound, + max_retries=0, + default_headers={"x-provider-key": "provider-secret"}, + ) + messages = sdk.beta.messages if beta else sdk.messages + result = await messages.count_tokens( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Contact alice@example.com"}], + ) + assert result.input_tokens == 17 + assert result._request_id == "count-1" + sdk.api_key = "invalid" + with pytest.raises(AuthenticationError): + await sdk.messages.count_tokens( + model="claude-sonnet-4-6", messages=[] + ) + assert len(calls) == 1 + assert [call.args[1] for call in usage.record_token_count.await_args_list] == [ + None, + 17, + ] + for name in ( + "admit", + "record_privacy", + "reserve_provider_spend", + "mark_provider_started", + "finalize", + "fail", + ): + getattr(usage, name).assert_not_awaited() + + +@pytest.mark.asyncio +async def test_token_count_errors_are_native_sanitized_and_never_retried(): + calls = [] + + def upstream(request): + calls.append(request) + return httpx.Response( + 500, + json={ + "type": "error", + "error": { + "type": "api_error", + "message": "provider-secret alice@example.com", + }, + }, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(upstream)) as outbound: + app = create_community_app( + CommunitySettings(_env_file=None, SHIM_API_KEY="gateway-secret-12345"), + http_client=outbound, + event_stream=io.StringIO(), + ) + async with app.router.lifespan_context(app): + usage = AsyncMock() + app.state.gateway_service.kernel.usage = usage + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app), base_url="http://shim.test" + ) as inbound: + headers = { + "x-api-key": "gateway-secret-12345", + "x-provider-key": "provider-secret", + } + payload = {"model": "claude-sonnet-4-6", "messages": []} + failed = await inbound.post( + "/v1/messages/count_tokens", headers=headers, json=payload + ) + invalid = await inbound.post( + "/v1/messages/count_tokens", + headers=headers, + json={**payload, "stream": True}, + ) + assert len(calls) == 1 + assert failed.status_code == 500 + assert failed.json()["type"] == "error" + assert ( + "provider-secret" not in failed.text and "alice@example.com" not in failed.text + ) + assert invalid.status_code == 422 and invalid.json()["type"] == "error" + usage.admit.assert_not_awaited() + usage.finalize.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_token_count_does_not_consume_message_repeat_allowance(): + def upstream(request): + if request.url.path.endswith("/count_tokens"): + return httpx.Response(200, json={"input_tokens": 2}) + return httpx.Response( + 200, + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 0}, + }, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(upstream)) as outbound: + app = create_community_app( + CommunitySettings( + _env_file=None, SHIM_API_KEY="gateway-secret-12345", LOOP_REPEAT_LIMIT=2 + ), + http_client=outbound, + event_stream=io.StringIO(), + ) + async with app.router.lifespan_context(app): + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app), + base_url="http://shim.test", + headers={ + "x-api-key": "gateway-secret-12345", + "x-provider-key": "provider-secret", + }, + ) as inbound: + payload = { + "model": "claude-sonnet-4-6", + "messages": [{"role": "user", "content": "hello"}], + } + counted = await inbound.post("/v1/messages/count_tokens", json=payload) + generated = await inbound.post( + "/v1/messages", json={**payload, "max_tokens": 10} + ) + repeated = await inbound.post( + "/v1/messages", json={**payload, "max_tokens": 10} + ) + assert counted.status_code == generated.status_code == 200 + assert repeated.status_code == 429 diff --git a/tests/gateway/test_usage.py b/tests/gateway/test_usage.py index 8193bcf..03c886c 100644 --- a/tests/gateway/test_usage.py +++ b/tests/gateway/test_usage.py @@ -17,14 +17,19 @@ def _prepared(*, model: str = "gpt-5.6-luna") -> SimpleNamespace: started_at = datetime.now(timezone.utc) - timedelta(milliseconds=12) return SimpleNamespace( + policy_verdicts=[], request_id="req_local", provider="openai", model=model, + pricing_model=model, + target=None, + unpriced=False, tenant_id="tenant-private", api_key_id="key-private", headers={"authorization": "credential-private"}, context=SimpleNamespace(started_at=started_at), - admission=SimpleNamespace(estimated_input_tokens=11), + admission=SimpleNamespace(estimated_input_tokens=11, repeat_chain_length=1), + deployment_kind="unknown", payload={"messages": [{"content": "secret-body"}]}, privacy=PrivacyOutcome( action=PrivacyAction.SCRUBBED, @@ -87,6 +92,12 @@ async def test_local_usage_writes_one_exact_redacted_terminal_event() -> None: "estimated_cost_usd", "estimated", "privacy_counts", + "provider_finish_reasons", + "ttft_ms", + "repeat_chain_length", + "system_prompt_hash", + "deployment_kind", + "policy_verdicts", } latency_ms = event.pop("latency_ms") assert event == { @@ -100,6 +111,12 @@ async def test_local_usage_writes_one_exact_redacted_terminal_event() -> None: "estimated_cost_usd": "0.0000106", "estimated": False, "privacy_counts": {"EMAIL_ADDRESS": 1}, + "provider_finish_reasons": None, + "ttft_ms": None, + "repeat_chain_length": 1, + "system_prompt_hash": None, + "deployment_kind": "unknown", + "policy_verdicts": [], } assert latency_ms >= 0 @@ -133,6 +150,8 @@ async def test_local_failure_writes_one_terminal_event() -> None: event = json.loads(stream.getvalue()) assert event["outcome"] == "provider_rejected_without_usage" assert event["completion_tokens"] == 0 + assert event["provider_finish_reasons"] is None + assert event["ttft_ms"] is None assert len(stream.getvalue().splitlines()) == 1 diff --git a/uv.lock b/uv.lock index df37e8b..364c0f0 100644 --- a/uv.lock +++ b/uv.lock @@ -87,6 +87,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/66/b25ccb84a246b470eb943b0107c07edcae51804912b824054b3413995a10/asyncpg-0.31.0-cp313-cp313-win_amd64.whl", hash = "sha256:dc5f2fa9916f292e5c5c8b2ac2813763bcd7f58e130055b4ad8a0531314201ab", size = 596569, upload-time = "2025-11-24T23:26:16.189Z" }, ] +[[package]] +name = "authlib" +version = "1.8.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "joserfc" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f1/51/bc1729d3cfdc214b4935f4e886e4dd443c3065fd8e1e66423fe84b490f81/authlib-1.8.0.tar.gz", hash = "sha256:f3ecd5f1da737262fb53bf1a4d95c4ea1ad9dd509316587a255c99ab1838a4f0", size = 177759, upload-time = "2026-08-30T12:12:34.833Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/c6/6f124bcfbbfb20fba22c939b4e43a06dccfc0e1ca20e5634ca573cb1e271/authlib-1.8.0-py2.py3-none-any.whl", hash = "sha256:88aebbd9af6757e14e912d5dc007ae1dc1f3e27e3b2152ce7c552ee2c3b3c121", size = 260804, upload-time = "2026-08-30T12:12:33.162Z" }, +] + [[package]] name = "azure-core" version = "1.41.0" @@ -695,6 +708,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/aa/0aca39a37d3c7eb941ba736ede56d689e7be91cab5d9ca846bde3999eba6/isodate-0.7.2-py3-none-any.whl", hash = "sha256:28009937d8031054830160fce6d409ed342816b543597cece116d966c6d99e15", size = 22320, upload-time = "2024-10-08T23:04:09.501Z" }, ] +[[package]] +name = "itsdangerous" +version = "2.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9c/cb/8ac0172223afbccb63986cc25049b154ecfb5e85932587206f42317be31d/itsdangerous-2.2.0.tar.gz", hash = "sha256:e0050c0b7da1eea53ffaf149c0cfbb5c6e2e2b69c4bef22c81fa6eb73e5f6173", size = 54410, upload-time = "2024-04-16T21:28:15.614Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/96/92447566d16df59b2a776c0fb82dbc4d9e07cd95062562af01e408583fc4/itsdangerous-2.2.0-py3-none-any.whl", hash = "sha256:c6242fc49e35958c8b15141343aa660db5fc54d4f13a1db01a3f5891b98700ef", size = 16234, upload-time = "2024-04-16T21:28:14.499Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -738,6 +760,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, ] +[[package]] +name = "joserfc" +version = "1.7.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/94/80fea1514b7c6d7d37804d3fe9ca81455f633347fc98731bd71ffe1faa17/joserfc-1.7.5.tar.gz", hash = "sha256:d5ff536e658e17664f8c1b1ab60dc4aa62aa973fcef1edd33cc44bda45d6f5ea", size = 234990, upload-time = "2026-08-29T13:05:42.057Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/c5/82addfd375e5ee6520644e0553e4aadde92d668c4fc99cc716d337fe7bb3/joserfc-1.7.5-py3-none-any.whl", hash = "sha256:add2c2c84e8373b084d526a8b53daba5d7a513a118cd2dcd9fc9f979d0922159", size = 71269, upload-time = "2026-08-29T13:05:40.718Z" }, +] + [[package]] name = "mako" version = "1.3.12" @@ -1580,13 +1614,17 @@ source = { editable = "ee" } dependencies = [ { name = "alembic" }, { name = "asyncpg" }, + { name = "authlib" }, { name = "cryptography" }, { name = "email-validator" }, { name = "fastapi" }, { name = "google-cloud-secret-manager" }, { name = "httpx" }, + { name = "itsdangerous" }, + { name = "joserfc" }, { name = "prometheus-client" }, { name = "pydantic" }, + { name = "pyjwt", extra = ["crypto"] }, { name = "pyyaml" }, { name = "redis" }, { name = "reportlab" }, @@ -1609,6 +1647,7 @@ secrets-azure = [ requires-dist = [ { name = "alembic", specifier = ">=1.17.2,<2" }, { name = "asyncpg", specifier = ">=0.31,<1" }, + { name = "authlib", specifier = ">=1.6,<2" }, { name = "azure-identity", marker = "extra == 'secrets-azure'", specifier = ">=1.25,<2" }, { name = "azure-keyvault-secrets", marker = "extra == 'secrets-azure'", specifier = ">=4.10,<5" }, { name = "boto3", marker = "extra == 'secrets-aws'", specifier = ">=1.42,<2" }, @@ -1617,8 +1656,11 @@ requires-dist = [ { name = "fastapi", specifier = ">=0.109,<1" }, { name = "google-cloud-secret-manager", specifier = ">=2.29,<3" }, { name = "httpx", specifier = ">=0.28,<1" }, + { name = "itsdangerous", specifier = ">=2.2,<3" }, + { name = "joserfc", specifier = ">=1.6,<2" }, { name = "prometheus-client", specifier = ">=0.25,<1" }, { name = "pydantic", specifier = ">=2.13,<3" }, + { name = "pyjwt", extras = ["crypto"], specifier = ">=2.12,<3" }, { name = "pyyaml", specifier = ">=6.0.2,<7" }, { name = "redis", specifier = ">=5.0.1,<9" }, { name = "reportlab", specifier = ">=5,<6" },