diff --git a/.github/ACTIONS-REFERENCE.md b/.github/ACTIONS-REFERENCE.md index 6a48bebcb..62fdfd491 100644 --- a/.github/ACTIONS-REFERENCE.md +++ b/.github/ACTIONS-REFERENCE.md @@ -43,6 +43,7 @@ GitHub provides two mechanisms for storing configuration values: | CDK_FILE_UPLOAD_CORS_ORIGINS | Variable | No | None | Platform | Additional CORS origins for the file upload S3 bucket only (appended to global CORS origins) | | CDK_FILE_UPLOAD_MAX_SIZE_MB | Variable | No | `10` | Platform | Maximum file upload size in megabytes | | CDK_FINE_TUNING_CORS_ORIGINS | Variable | No | None | SageMaker Fine-Tuning | Additional CORS origins for the fine-tuning S3 bucket only (appended to global CORS origins) | +| CDK_FINE_TUNING_ENABLED | Variable | No | `true` | App API | Mounts the `/fine-tuning` and `/admin/fine-tuning` routers (sets the container's `FINE_TUNING_ENABLED`). Default ON; set to `false` as a kill switch. Storage and IAM are provisioned either way, so switching off never orphans datasets or trained models. | | CDK_FINE_TUNING_DEFAULT_QUOTA_HOURS | Variable | No | `0` | App API | Default monthly GPU-hour quota for all authenticated users. `0` = whitelist-only (admin must grant each user). Positive value (e.g. `5`) = open access with that default budget. | | CDK_FRONTEND_BUCKET_NAME | Variable | No | None | Frontend | S3 bucket name for frontend assets (defaults to generated name with account ID) | | CDK_FRONTEND_CORS_ORIGINS | Variable | No | None | Frontend | Additional CORS origins for the frontend SSM export only (appended to global CORS origins) | diff --git a/.github/docs/deploy/step-03-github-config.md b/.github/docs/deploy/step-03-github-config.md index 0a568a3c7..eb120a610 100644 --- a/.github/docs/deploy/step-03-github-config.md +++ b/.github/docs/deploy/step-03-github-config.md @@ -109,6 +109,28 @@ The per-origin cert vars below are **optional overrides** — set one only if yo | `CDK_ARTIFACTS_EXTRA_FRAME_ANCESTORS` | — | Comma-separated extra origins (beyond `https://{CDK_DOMAIN_NAME}`) allowed to embed artifact iframes via CSP `frame-ancestors` — applied to both the CloudFront response-headers policy and the render Lambda. Set to `http://localhost:4200` to point a local SPA at this deployment. **Leave unset in production**: every listed origin can frame your users' artifacts (still render-token gated, but a real loosening on a shared environment). | | `CDK_MCP_SANDBOX_EXTRA_FRAME_ANCESTORS` | — | Comma-separated extra origins (beyond `https://{CDK_DOMAIN_NAME}`) allowed to embed the MCP Apps sandbox proxy via CSP `frame-ancestors`. Set to `http://localhost:4200` to point a local SPA at this deployment. **Leave unset in production.** | | `CDK_FINE_TUNING_CORS_ORIGINS` | — | Comma-separated extra CORS origins for the SageMaker fine-tuning data bucket, beyond `https://{CDK_DOMAIN_NAME}`. Optional — fine-tuning itself is always provisioned. | +| `CDK_FINE_TUNING_ENABLED` | `true` | Mounts the fine-tuning routers. Default ON — leave unset unless you want the kill switch. Note this is a *runtime* flag on the app-api container; the identically-named CDK stack gate was removed in the single-stack migration. | +| `CDK_FINE_TUNING_DEFAULT_QUOTA_HOURS` | `0` | Monthly GPU-hour quota auto-granted to any authenticated user. `0` = whitelist-only (an admin grants each user). A positive value (e.g. `10`) = open access with that budget. | + +### Managed Knowledge Bases + +Every variable below is **optional**. Leave them all unset for the shipped state: the managed knowledge-base backend is deployed but **dormant** — no knowledge base is created managed, no migration runs, and the daily reconciler reports what it *would* delete without deleting anything. + +The three flags are independent opt-ins that each default to **off**. An unset GitHub Variable arrives at the deploy as an empty string, which is read as off — so forgetting one never silently arms it. + +| Variable Name | Default | Description | +|---------------|---------|-------------| +| `CDK_MANAGED_KB_NEW_DEFAULT` | `false` | Set to `true` so newly created knowledge bases are provisioned on the managed backend instead of the legacy one. Existing knowledge bases are untouched. | +| `CDK_MANAGED_KB_MIGRATION_ENABLED` | `false` | Set to `true` to let the background migration worker run at all. While unset, the worker performs no work and its schedule stays disabled. | +| `CDK_MANAGED_KB_RECONCILER_ARMED` | `false` | Set to `true` to let the daily reconciler **delete** orphaned knowledge bases. While unset the reconciler still runs and still logs every deletion it intends to make — review those logs before arming it. | +| `CDK_MANAGED_KB_PER_OWNER_BYTES` | `104857600` (100 MB) | Per-owner stored-bytes cap for the standard role tier, **in bytes**. Deliberately below the 1 GB user-files precedent: at 30,000 users a 1 GB cap permits 30 TB. | +| `CDK_MANAGED_KB_PER_OWNER_ELEVATED_BYTES` | `1073741824` (1 GB) | Per-owner cap for the elevated, admin-granted tier, **in bytes**. | +| `CDK_MANAGED_KB_PER_KB_CEILING_BYTES` | `524288000` (500 MB) | Ceiling for any single knowledge base, **in bytes**, bounding one runaway corpus inside an owner's allowance. | +| `CDK_MANAGED_KB_RETENTION_WINDOW_DAYS` | `30` | How long legacy vector data is kept after a knowledge base is promoted to the managed backend, **in days**, so a rollback stays possible. Do not set below `30`. | +| `CDK_MANAGED_KB_STORAGE_ALARM_GB` | `500` | CloudWatch alarm threshold for **fleet-wide** managed knowledge base storage, **in GB**. The per-owner caps above bound one user; this is the only thing that bounds the whole account. | +| `CDK_MANAGED_KB_DAILY_COST_ALARM_USD` | `100` | CloudWatch alarm threshold for the rolled-up daily Knowledge-Base cost, **in USD**. Set alongside the storage alarm — per-owner caps alone permit roughly two orders of magnitude more spend than expected usage. | + +> Accepted values for the three flags are `true`, `false`, `1`, `0`, or empty (empty means off). Anything else fails fast at deploy time with a message naming the variable. --- diff --git a/.github/docs/deploy/step-05-verify.md b/.github/docs/deploy/step-05-verify.md index 364f9b56e..5ed01dd52 100644 --- a/.github/docs/deploy/step-05-verify.md +++ b/.github/docs/deploy/step-05-verify.md @@ -102,6 +102,67 @@ In order of likelihood: --- +### 6. Subscribe to Platform Alarms (required — not automated) + +The deploy creates one SNS topic that **every** CloudWatch alarm in the stack +publishes to. It has no subscribers until you add them, so until you do this step +the alarms change colour in the console and tell nobody. + +This is deliberately not infrastructure-as-code. Several teams usually need to +hear about failures, and their membership changes far more often than the +infrastructure does — requiring a pull request, a review, and a CloudFormation +deploy to add one email address is how a notification list goes stale and stops +being trusted. Subscribing is a one-line command that touches no code. + +Find the topic and subscribe: + +```bash +PREFIX="your-project-prefix" # the same CDK_PROJECT_PREFIX you deployed with + +TOPIC=$(aws ssm get-parameter \ + --name "/${PREFIX}/observability/alarm-topic-arn" \ + --query Parameter.Value --output text) + +aws sns subscribe \ + --topic-arn "$TOPIC" \ + --protocol email \ + --notification-endpoint platform-team@example.edu +``` + +AWS sends a confirmation email; the subscription is inactive until the recipient +clicks the link. Repeat for each address or distribution list. + +Other useful protocols: + +| Protocol | Use for | +|---|---| +| `email` | A team distribution list. Simplest, and enough for most forks. | +| `https` | PagerDuty, Opsgenie, ServiceNow, or any webhook receiver. | +| `sms` | Genuine paging. Costs per message. | +| `lambda` | Custom routing, e.g. severity-based fan-out or Slack formatting. | + +Verify it took: + +```bash +aws sns list-subscriptions-by-topic --topic-arn "$TOPIC" \ + --query 'Subscriptions[].[Protocol,Endpoint,SubscriptionArn]' --output table +``` + +A `SubscriptionArn` of `PendingConfirmation` means the email has not been +confirmed yet. + +Then open the health dashboard — `{PREFIX}-platform-health` in the CloudWatch +console — and confirm the alarm-status row is populated and green. Row 1 tells +you whether traffic is being served, row 2 tells you why, and row 3 lists every +alarm's current state. + +> **Note on latency alarms:** the chat path uses server-sent events, so response +> times of tens of seconds are normal for a healthy agent turn. Latency alarms are +> deliberately set at 120 seconds. A *drop* in latency can actually mean turns are +> failing early. + +--- + ## You're Done! Your AgentCore Public Stack is deployed and running. Here's what you have: diff --git a/.github/workflows/backend.yml b/.github/workflows/backend.yml index e934884be..d8e60114a 100644 --- a/.github/workflows/backend.yml +++ b/.github/workflows/backend.yml @@ -237,6 +237,89 @@ jobs: - name: Deploy kb-sync worker image run: bash scripts/build/deploy-image-lambda-one.sh kb-sync-worker + build-kb-migration: + name: Build kb-migration image + needs: test-backend + # Native ARM64 runner — all four kb-migration Lambdas are arm64 (see the + # managed-kb CDK construct), matching the kb-sync pattern. + runs-on: ubuntu-24.04-arm + environment: ${{ (github.ref == 'refs/heads/main' && 'production') || 'development' }} + + permissions: + id-token: write + contents: read + + env: + CDK_AWS_REGION: ${{ vars.AWS_REGION }} + CDK_AWS_ACCOUNT: ${{ vars.CDK_AWS_ACCOUNT }} + CDK_PROJECT_PREFIX: ${{ vars.CDK_PROJECT_PREFIX }} + AWS_REGION: ${{ vars.AWS_REGION }} + AWS_ACCOUNT_ID: ${{ vars.CDK_AWS_ACCOUNT }} + AWS_ROLE_ARN: ${{ secrets.AWS_ROLE_ARN }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + + outputs: + image_tag: ${{ steps.build.outputs.image_tag }} + + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: ./.github/actions/build-and-push-image + id: build + with: + image-name: kb-migration + aws-region: ${{ vars.AWS_REGION || 'us-west-2' }} + aws-role-arn: ${{ secrets.AWS_ROLE_ARN }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + + deploy-kb-migration-code: + name: Deploy kb-migration Lambda images + # ONE image, FOUR functions: dispatcher, worker, reconciler and ingestion + # consumer share the kb-migration image and differ only in + # ImageConfig.Command (CDK-owned), so a single job points all four at the + # freshly-built tag. + # + # This is the job that replaces the bootstrap stub + # (infrastructure/bootstrap-assets/kb-migration/) with the real handlers. + # Until it has run once, an enrolled knowledge base sits in `shadow` while + # the dispatcher ticks into a no-op — safe, because the work keys are sparse + # and the first real tick picks up everything that accumulated. + needs: [build-kb-migration, test-backend] + runs-on: ubuntu-24.04 + environment: ${{ (github.ref == 'refs/heads/main' && 'production') || 'development' }} + + permissions: + id-token: write + contents: read + + env: + CDK_AWS_REGION: ${{ vars.AWS_REGION }} + CDK_AWS_ACCOUNT: ${{ vars.CDK_AWS_ACCOUNT }} + CDK_PROJECT_PREFIX: ${{ vars.CDK_PROJECT_PREFIX }} + AWS_REGION: ${{ vars.AWS_REGION }} + AWS_ACCOUNT_ID: ${{ vars.CDK_AWS_ACCOUNT }} + AWS_ROLE_ARN: ${{ secrets.AWS_ROLE_ARN }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: ./.github/actions/configure-aws-credentials + with: + aws-region: ${{ vars.AWS_REGION || 'us-west-2' }} + aws-role-arn: ${{ secrets.AWS_ROLE_ARN }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + - name: Deploy kb-migration dispatcher image + run: bash scripts/build/deploy-image-lambda-one.sh kb-migration-dispatcher + - name: Deploy kb-migration worker image + run: bash scripts/build/deploy-image-lambda-one.sh kb-migration-worker + - name: Deploy kb-migration reconciler image + run: bash scripts/build/deploy-image-lambda-one.sh kb-migration-reconciler + - name: Deploy kb-migration ingestion consumer image + run: bash scripts/build/deploy-image-lambda-one.sh kb-migration-ingestion-consumer + build-scheduled-runs: name: Build scheduled-runs image needs: test-backend diff --git a/.github/workflows/platform.yml b/.github/workflows/platform.yml index 839f7917f..37d1f6298 100644 --- a/.github/workflows/platform.yml +++ b/.github/workflows/platform.yml @@ -57,6 +57,11 @@ jobs: CDK_PROJECT_PREFIX: ${{ vars.CDK_PROJECT_PREFIX }} CDK_DOMAIN_NAME: ${{ vars.CDK_DOMAIN_NAME }} CDK_CORS_ORIGINS: ${{ vars.CDK_CORS_ORIGINS }} + # The Environment tag value. Not cosmetic: it is the filter the managed-KB + # reconciler and scripts/teardown/managed-kb.sh match knowledge bases on, so + # an unset value here made a dev deploy tag its knowledge bases 'prod' while + # teardown looked for 'dev' — matching nothing and reporting success. + CDK_TAG_ENVIRONMENT: ${{ vars.CDK_TAG_ENVIRONMENT }} CDK_VPC_CIDR: ${{ vars.CDK_VPC_CIDR }} CDK_ALB_SUBDOMAIN: ${{ vars.CDK_ALB_SUBDOMAIN }} CDK_CERTIFICATE_ARN: ${{ vars.CDK_CERTIFICATE_ARN }} @@ -128,6 +133,25 @@ jobs: # and see the ones their role grants. Which skills a cohort gets is a # role's `grantedSkills`, managed in the admin roles UI. CDK_SKILLS_ENABLED: ${{ vars.CDK_SKILLS_ENABLED }} + # SageMaker fine-tuning. The tables, bucket, SageMaker role and IAM + # grants deploy unconditionally, but these two decide whether the feature + # is reachable and how it is rationed, and neither was forwarded before — + # so every deployed environment served 404s from `/fine-tuning/*` while + # the repo variable `CDK_FINE_TUNING_ENABLED` read "true". + # + # ENABLED is default ON with a kill switch: unset resolves to empty + # string, which config.ts treats as the default (on). Set it to "false" + # to dark-stop the routes; storage is untouched either way, so no dataset + # or trained model is orphaned. + # + # DEFAULT_QUOTA_HOURS picks the access model: `0` (the default) is + # whitelist-only, where an admin grants each user explicitly; a positive + # value auto-grants that monthly GPU-hour budget to any signed-in user. + CDK_FINE_TUNING_ENABLED: ${{ vars.CDK_FINE_TUNING_ENABLED }} + CDK_FINE_TUNING_DEFAULT_QUOTA_HOURS: ${{ vars.CDK_FINE_TUNING_DEFAULT_QUOTA_HOURS }} + # Extra CORS origins for the fine-tuning data bucket, beyond the global + # CDK_CORS_ORIGINS above. Also never forwarded until now. + CDK_FINE_TUNING_CORS_ORIGINS: ${{ vars.CDK_FINE_TUNING_CORS_ORIGINS }} # Agent Designer /agents surface. Default OFF (opt-in) until the Phase-4 # Designer UI ships — a headless API helps no forker. Set the # `CDK_AGENTS_API_ENABLED` variable to "true" in an environment (e.g. dev) @@ -146,6 +170,63 @@ jobs: # cdk.context.json stay inert. CDK_MCP_TOKEN_ENRICHMENT_ENABLED: ${{ vars.CDK_MCP_TOKEN_ENRICHMENT_ENABLED }} CDK_MCP_TOKEN_ENRICHMENT_CLAIMS: ${{ vars.CDK_MCP_TOKEN_ENRICHMENT_CLAIMS }} + # Managed knowledge bases (.kiro/specs/managed-kb-migration). THREE + # INDEPENDENT OPT-IN flags, all defaulting to OFF — the inverse of the + # kill-switch flags above, and the difference matters here. An unset + # GitHub Actions variable renders as an EMPTY STRING, not as absent, so + # a `!== 'false'` reading of an unset variable would resolve to TRUE and + # arm the feature on every fork. config.ts reads these with + # parseBooleanEnv, which maps both unset and empty to undefined and falls + # through to `false` (Requirement 19.8). Leave all three unset to deploy + # the managed backend without starting a fleet migration. + # + # CDK_MANAGED_KB_NEW_DEFAULT new KBs are created managed + # CDK_MANAGED_KB_MIGRATION_ENABLED the background migrator runs at all + # CDK_MANAGED_KB_RECONCILER_ARMED the daily reconciler DELETES orphans + # rather than only reporting them + # + # reconcilerArmed is the inverted one: the Reconciler is deployed and + # running from day one but DISARMED, so its judgement can be reviewed + # against real data before it deletes anything (Requirements 14.7, 19.7). + CDK_MANAGED_KB_NEW_DEFAULT: ${{ vars.CDK_MANAGED_KB_NEW_DEFAULT }} + CDK_MANAGED_KB_MIGRATION_ENABLED: ${{ vars.CDK_MANAGED_KB_MIGRATION_ENABLED }} + CDK_MANAGED_KB_RECONCILER_ARMED: ${{ vars.CDK_MANAGED_KB_RECONCILER_ARMED }} + # Storage cost controls. Byte_Caps are in BYTES (Requirement 12.2), + # defaulting to 100 MB standard / 1 GB elevated / 500 MB per knowledge + # base; the retention window is in DAYS and must stay >= 30 + # (Requirement 15.11). Leave unset to take those defaults — these exist + # so an environment can tune them without a code change. + CDK_MANAGED_KB_PER_OWNER_BYTES: ${{ vars.CDK_MANAGED_KB_PER_OWNER_BYTES }} + CDK_MANAGED_KB_PER_OWNER_ELEVATED_BYTES: ${{ vars.CDK_MANAGED_KB_PER_OWNER_ELEVATED_BYTES }} + CDK_MANAGED_KB_PER_KB_CEILING_BYTES: ${{ vars.CDK_MANAGED_KB_PER_KB_CEILING_BYTES }} + CDK_MANAGED_KB_RETENTION_WINDOW_DAYS: ${{ vars.CDK_MANAGED_KB_RETENTION_WINDOW_DAYS }} + # Fleet-level alarm thresholds (Requirement 12.13). The Byte_Caps above + # bound ONE owner; these two bound the whole account, which is the gap + # between ~$169/month expected and ~$15,000/month that per-owner caps + # alone permit. Storage is in GB (default 500), daily cost in USD + # (default 100). Leave unset to take those defaults. + CDK_MANAGED_KB_STORAGE_ALARM_GB: ${{ vars.CDK_MANAGED_KB_STORAGE_ALARM_GB }} + CDK_MANAGED_KB_DAILY_COST_ALARM_USD: ${{ vars.CDK_MANAGED_KB_DAILY_COST_ALARM_USD }} + # Observability. All optional — unset means the default in config.ts. + # XRAY_SAMPLING_RATE is a rate (0.0-1.0), not a percentage. + CDK_OBSERVABILITY_ALARM_TOPIC_ENABLED: ${{ vars.CDK_OBSERVABILITY_ALARM_TOPIC_ENABLED }} + CDK_OBSERVABILITY_LOG_RETENTION_DAYS: ${{ vars.CDK_OBSERVABILITY_LOG_RETENTION_DAYS }} + CDK_OBSERVABILITY_ALB_TARGET_5XX_THRESHOLD: ${{ vars.CDK_OBSERVABILITY_ALB_TARGET_5XX_THRESHOLD }} + CDK_OBSERVABILITY_ALB_P99_LATENCY_MS: ${{ vars.CDK_OBSERVABILITY_ALB_P99_LATENCY_MS }} + CDK_OBSERVABILITY_AGENTCORE_LATENCY_MS: ${{ vars.CDK_OBSERVABILITY_AGENTCORE_LATENCY_MS }} + CDK_OBSERVABILITY_AGENTCORE_ERROR_THRESHOLD: ${{ vars.CDK_OBSERVABILITY_AGENTCORE_ERROR_THRESHOLD }} + CDK_OBSERVABILITY_LAMBDA_ERROR_THRESHOLD: ${{ vars.CDK_OBSERVABILITY_LAMBDA_ERROR_THRESHOLD }} + CDK_OBSERVABILITY_LAMBDA_DURATION_PERCENT_OF_TIMEOUT: ${{ vars.CDK_OBSERVABILITY_LAMBDA_DURATION_PERCENT_OF_TIMEOUT }} + CDK_OBSERVABILITY_DYNAMO_THROTTLE_THRESHOLD: ${{ vars.CDK_OBSERVABILITY_DYNAMO_THROTTLE_THRESHOLD }} + CDK_OBSERVABILITY_ECS_CPU_PERCENT: ${{ vars.CDK_OBSERVABILITY_ECS_CPU_PERCENT }} + CDK_OBSERVABILITY_ECS_MEMORY_PERCENT: ${{ vars.CDK_OBSERVABILITY_ECS_MEMORY_PERCENT }} + CDK_OBSERVABILITY_XRAY_SAMPLING_RATE: ${{ vars.CDK_OBSERVABILITY_XRAY_SAMPLING_RATE }} + CDK_OBSERVABILITY_XRAY_SAMPLING_RESERVOIR: ${{ vars.CDK_OBSERVABILITY_XRAY_SAMPLING_RESERVOIR }} + CDK_OBSERVABILITY_XRAY_INSIGHTS_NOTIFICATIONS: ${{ vars.CDK_OBSERVABILITY_XRAY_INSIGHTS_NOTIFICATIONS }} + CDK_OBSERVABILITY_AGENTCORE_APPLICATION_LOGS_ENABLED: ${{ vars.CDK_OBSERVABILITY_AGENTCORE_APPLICATION_LOGS_ENABLED }} + CDK_OBSERVABILITY_PROMPT_CACHE_AVOIDABLE_MISS_THRESHOLD: ${{ vars.CDK_OBSERVABILITY_PROMPT_CACHE_AVOIDABLE_MISS_THRESHOLD }} + CDK_OBSERVABILITY_PROMPT_CACHE_WASTED_USD_THRESHOLD: ${{ vars.CDK_OBSERVABILITY_PROMPT_CACHE_WASTED_USD_THRESHOLD }} + CDK_OBSERVABILITY_PROMPT_CACHE_SESSION_WASTED_USD_THRESHOLD: ${{ vars.CDK_OBSERVABILITY_PROMPT_CACHE_SESSION_WASTED_USD_THRESHOLD }} # Secrets AWS_ROLE_ARN: ${{ secrets.AWS_ROLE_ARN }} AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} diff --git a/.kiro/specs/managed-kb-migration/.config.kiro b/.kiro/specs/managed-kb-migration/.config.kiro new file mode 100644 index 000000000..0b0157564 --- /dev/null +++ b/.kiro/specs/managed-kb-migration/.config.kiro @@ -0,0 +1 @@ +{"specId": "612a1431-367c-427e-8fe0-08872d03a9b9", "workflowType": "design-first", "specType": "feature"} diff --git a/.kiro/specs/managed-kb-migration/HANDOFF.md b/.kiro/specs/managed-kb-migration/HANDOFF.md new file mode 100644 index 000000000..3ccf6c8c4 --- /dev/null +++ b/.kiro/specs/managed-kb-migration/HANDOFF.md @@ -0,0 +1,974 @@ +# Managed KB Migration — Handoff + +**Last updated:** 2026-08-31 14:30 · **Shipped to production in 1.16.0, inert behind flags** · +**A migration has completed end to end in dev; adding a document to it needed one more IAM action** + +Working state for this feature so a fresh session can pick it up without +re-deriving anything. Read this, then `tasks.md`. + +--- + +## 0. Read this first + +Four things invalidate earlier versions of this document: + +1. **It is deployed.** The feature shipped to production in release 1.16.0 and the + platform deploy succeeded on 2026-08-28, so `GSI7`, the Bedrock service role and + the four Lambdas exist in **both** dev and prod. Earlier revisions of this file + said "Nothing deployed"; that is no longer true. +2. **Eleven defects were found only by running it**, each reviewed clean and + deployed clean. They are §5 items 25–36 and they are the most useful part of + this document. Items 32–36 all trace to one root cause — the two engines were + never made exclusive — and are fixed in PR #900. +3. **The `document_id` "known unknown" was a false alarm** and is now resolved with + measurements — see §6. An earlier revision listed it as the top open risk. The + probe was reading facade keys that have never existed. Two genuine findings came + out of checking it (§5.32, §5.33), both still open. +4. **Iterate locally, but do not trust it for IAM.** + `scripts/local-dev/run-kb-migration.py` drives the whole state machine + in-process against dev with your SSO credentials. Three of the + five defects would have been minutes of work instead of a merge → image build → + deploy → 15-minute-tick cycle each. Use it. + + But your SSO identity is **broader than every Lambda role**, so the driver is + structurally blind to IAM gaps — that is how §5.31 shipped after a migration had + already "completed end to end". Get the logic right locally; prove the + permissions by deploying and letting the real roles do the work. + +--- + +## 1. Status + +| | | +|---|---| +| Spec | Complete. Requirement **8.5 was amended by measurement on 2026-08-31** — see §5.29 | +| Implementation | Groups 1–14 except 14.5. A migration has completed `shadow → verify → promote → retain` in dev and serves from the managed backend | +| Tests | 640 infra (jest) · ~6,780 backend (pytest) · 1,936 frontend (vitest) · 5 pre-existing unrelated Strands failures | +| Deployed | **dev and prod.** Flags off in prod; `migrationEnabled` on in dev | +| Open PRs | **#900** — engine exclusivity: the legacy pipeline stands down for a promoted KB (§5.32, §5.34) and deletion propagates to the managed engine (§5.36). Four CI checks green. Merging triggers **both** `backend.yml` (rag-ingestion + kb-sync images) and `platform.yml` (two new IAM grants). · **#899** — this document. · #898 merged as `ef2f4c9e` | +| Uncommitted | none — working tree clean as of 2026-08-31 14:30 | + +### Flag state (GitHub Environment variables) + +| Flag | development | production | +|---|---|---| +| `CDK_MANAGED_KB_MIGRATION_ENABLED` | `true` | `false` | +| `CDK_MANAGED_KB_NEW_DEFAULT` | `false` | `false` | +| `CDK_MANAGED_KB_RECONCILER_ARMED` | `false` | `false` | +| `CDK_TAG_ENVIRONMENT` | `dev` | `prod` | + +`newDefault` has **no reader anywhere in `backend/src`** — "new knowledge bases are +created managed" is design §14.7 steps 5–8, a follow-up spec. Setting it does +nothing, which is worth knowing before someone flips it expecting an effect. + +⚠️ **Production carries every defect fixed after 1.16.0 shipped.** It cannot fire, +because nothing enrols while `migrationEnabled` is false. Do not turn that flag on +in prod until #898 and the uncommitted work have landed and shipped. + +### Commits + +**On `fix/kb-legacy-pipeline-engine-gate` — PR #900, open:** + +``` +e3398f30 propagate document deletion to the managed knowledge base (§5.36) +8a35dbc6 the legacy pipeline stands down for a promoted knowledge base (§5.32, §5.34) +``` + +**Merged as `ef2f4c9e` (was PR #898):** + +``` +fdf15d21 grant bedrock:StartIngestionJob, which authorizes direct ingestion (§5.31) +6420f148 drop the embedding pin, defer verify, complete a migration in dev (§5.29, §5.30) +7542d907 record the knowledge base id before anything else can fail (§5.28) +``` + +**Already merged (16, on develop):** + +``` +45239838 one source of truth for the managed KB tag contract +4acaa8f2 handoff reflects group 14 backend half and three more defects +e59f771c register the managed backend, fleet metrics, tagged teardown (group 14 backend) +d5e56f31 handoff reflects group 13 and four more defects +ee091971 migration dispatcher and the shadow/verify/promote/retain worker (group 13) +53476544 handoff reflects groups 11-12 and two new defects +a361fdd4 opt-in dual-read pilot that legacy always wins (group 12) +a43d80bf app-side authorization, IAM-enforced sharing, publication (group 11) +58f0c6b6 handoff document and accurate task-list state +8079f7e2 tombstone deletion sagas and the report-only reconciler (group 10) +e6936b0b ingestion consumer with exclusive engine routing (group 9) +620fa49c managed KB provisioning, retrieval and direct ingestion (group 8) +d433d6f1 per-owner byte cap with atomic reserve/commit/release (group 7) +f2e86afe clamp retrieval queries and fail closed on status (groups 5, 6) +24689de1 backend abstraction seam behind the retrieval entry point (group 4) +ffa7a408 KB_Record data layer with conditional state transitions (group 3) +5f2c98b1 spec, schema and worker platform (groups 1, 2) +``` + +**Working tree: clean.** An earlier revision listed the 14.3 upgrade surface and +then the 8.5 amendment as uncommitted; both have landed. `apis/app_api/kb_upgrade/` +is merged — see §7 for the file map and §2 for how to run it. + +### Is the feature reachable yet? + +**Yes, end to end, once the flag is on — except that nothing performs the work.** +Group 14.3 closed the last gap in the *control* path: a user can now enrol a +knowledge base, which writes a `KB#` record in `shadow` with the GSI7 work keys. +Before it, nothing wrote either, so every group could have been finished with the +feature unreachable (§5 defect 21). + +The worker's image **is deployed** — PR #886 shipped `Dockerfile.kb-migration` and +all four Lambdas run real handlers. An earlier revision of this section said the +image was undeployed and a `shadow` record would sit forever; that is no longer +true. The dispatcher's rule is `ENABLED` and ticking every 15 minutes in dev. + +⚠️ **That tick is a hazard while #898 is unmerged.** The deployed worker predates +it, so an enrolled record can be picked up by pre-fix code and failed at `verify`. +Always drive a local migration with `--break-lease`, which defers `dueAt` 20 minutes +out so the deployed dispatcher skips it. + +### What works today, verified live in dev + +| | | +|---|---| +| Chat against a promoted KB | **Yes.** `ast-1a90784a7f18` is `retain` / `managed` and serves real chunks | +| Add a document to a promoted KB | **Yes, once #898 deploys.** Blocked before that by §5.31. The consumer's EventBridge rule is `ENABLED`, the bucket has EventBridge notification on, and `grantRetrieval` is attached for the retrievability poll | +| Create a knowledge base that is managed from birth | **No — not implemented.** `newDefault` has **zero readers** in `backend/src` (grep for `NEW_DEFAULT`, `newDefault`, `new_default`: no matches). Design §14.7 steps 5–8, a follow-up spec. The only route onto the managed engine is enrol → migrate | +| Image-only PDFs | **Managed yes, legacy no.** A pure-diagram flowchart fails Docling outright (`zero chunks` → `failed`, permanently) and is served fine by managed via `imageExtractionStatus: ENABLED` (§5.35) | +| Deleting a document | **Propagates to managed only after #900 deploys.** Before that the managed copy is orphaned; the fail-closed status filter is what keeps it unserved (§5.36) | +| Whole local chain | **Yes.** SPA :4200 → app_api :8000 → inference_api :8001. `chat-http.service.ts` posts to `{appApiUrl}/chat/stream`; the app_api proxy forwards to `INFERENCE_API_URL`, which defaults to and is set to `http://localhost:8001`. Retrieval runs at `inference_api/chat/routes.py:1814`, so the RAG code answering a local chat is the code on disk | + +**Three behaviour changes are live on the existing legacy path** regardless of any +flag: +1. The document-status filter now **fails closed** (group 6) — except for the one + fail-open line in §5.33. +2. Retrieval queries are **clamped to 10,000 chars** (group 5). +3. Retrieval requires a resolved access grant (group 11). Both production callers + pass one; the parameter is required and keyword-only, so a third caller added + without one fails at the call site rather than silently serving nothing. + +--- + +## 2. Environment + +macOS host, tooling installed locally. **There is no devcontainer.** +`.kiro/steering/dev-environment.md` describes a different machine (WSL2/nspawn, +`/home/colin/...` paths) — ignore it here. + +```bash +# infrastructure +cd infrastructure && npm run build # tsc +cd infrastructure && npx jest # 640 passing, 30 suites + +# backend +cd backend && uv run python -m pytest tests/ -q # 6 m 20 s, 6,603 passing +cd backend && uv run ruff check + +# frontend +cd frontend/ai.client && npx ng test --watch=false # 1,886 passing, ~7 s +cd frontend/ai.client && npx ng test --watch=false --include="**/kb-upgrade*" +cd frontend/ai.client && npx tsc -p tsconfig.app.json --noEmit +``` + +There is **no eslint config** in the repo despite the steering docs mentioning +ESLint; `npx eslint` fails with "couldn't find an eslint.config.*". Type-check +with `tsc --noEmit` and build with `ng build` instead. + +### Driving a migration locally (do this before deploying anything) + +```bash +cd backend +uv run python ../scripts/local-dev/run-kb-migration.py --show +uv run python ../scripts/local-dev/run-kb-migration.py --break-lease +``` + +Runs the worker's steps in-process against dev with your SSO credentials, so the +whole state machine iterates in seconds. `--break-lease` clears the 15-minute lease +between steps and defers `dueAt` 20 minutes out so the deployed dispatcher does not +race you. + +It needs five variables in `backend/src/.env` that the app-api task definition does +not carry — copy them from the deployed worker Lambda: +`MANAGED_KB_SERVICE_ROLE_ARN`, `MANAGED_KB_TAG_VALUE_PREFIX`, +`MANAGED_KB_TAG_VALUE_ENVIRONMENT`, `MANAGED_KB_METRIC_NAMESPACE`, +`KB_MIGRATION_RETAIN_DAYS`. + +**What it does not prove:** the worker Lambda's IAM role (your SSO identity is +broader), the CDK environment wiring, or the image contents. Those are deploy-time +concerns — check them by deploying. Getting the logic right here first is the point, +and two of the five defects below were IAM/wiring and could only surface that way. + +### Two read-only diagnostics worth knowing before you debug anything + +```bash +cd backend +# Per-document ingestion timing, and WHICH engine did the work. +uv run python ../scripts/local-dev/kb-doc-timings.py + +# The same query through BOTH engines, side by side. +uv run python ../scripts/local-dev/kb-compare-engines.py "CS434" +uv run python ../scripts/local-dev/kb-compare-engines.py -f queries.txt +``` + +`kb-doc-timings.py` derives the engine from the record rather than guessing: +`chunkCount`/`vectorStoreId` are only ever written by the legacy pipeline, +`indexedAt`/`retrievableAt` only by the managed consumer, so a document carrying +both was double-indexed (§5.32). Measured on a 132 KB PDF: legacy `complete` at +30 s, managed retrievable at **95 s**, `INDEXED → retrievable` **0.9 s** — which +independently confirms the evaluation's 0.75–1.03 s. + +⚠️ It reports legacy as `(lost)` for a double-indexed document, and that is a real +limit, not a bug: the managed consumer overwrites `updatedAt`, so the legacy +finish time is unrecoverable afterwards. Capture it live if you need it. + +`kb-compare-engines.py` exploits the `retain` window — promotion moves no data, so +for 30 days **both** indexes hold the corpus and the same query can be put to +both. Marks each chunk `[in ]`/`[CUT]` against `MAX_CONTEXT_CHARS`, which is the +detail that makes the comparison meaningful: only ~2,000 characters reach the +model, so in practice **one or two chunks**, and precision@1 is nearly the whole +game. Absolute scores are not comparable across engines (legacy is a negated +cosine distance); order and **spread** are. + +The clearest measured difference, and the one to reach for first — exact-token +search, where legacy is pure vector and managed is hybrid: + +| query `CS434` | top chunk | spread | +|---|---|---| +| legacy | `SECTION 4: TECHNICAL ELECTIVES…` (does not contain CS434) | **0.0561** | +| managed | `- Algorithms of Machine Learning (3) CS434 - Applied Deep…` | **0.4988** | + +Legacy's five chunks sit within 0.056 of each other — flat, so its ranking is +close to arbitrary. Nine times the separation on managed, with the literal match +first. + +⚠️ **PR #900 ends this trick.** Once the legacy pipeline stands down for a +promoted knowledge base, new uploads land in the managed index only, so a +like-for-like A/B needs two assistants or documents that predate promotion. + +### Running the upgrade UI locally + +```bash +# 1. turn the offer on (LOCAL ONLY — backend/src/.env is gitignored) +echo 'MANAGED_KB_MIGRATION_ENABLED=true' >> backend/src/.env + +# 2. app_api on :8000, reading the dev account's DynamoDB via backend/src/.env +cd backend/src/apis/app_api && uv run python main.py + +# 3. SPA on :4200 — environment.ts already points at localhost:8000 +cd frontend/ai.client && npm run start +``` + +Then edit an assistant that has documents. Without the flag the card renders +nothing at all, which is correct rather than broken. + +⚠️ **The local API writes to the real dev tables.** Enrolling writes a genuine +`KB#{id}` item under `AST#{id}`. Use a throwaway assistant; undo by deleting that +item. The upgrade will sit at "Upgrading…" forever because the worker Lambda is +not deployed — expected, not a bug. + +**Baselines that are NOT your fault:** +- 5 backend failures in `tests/agents/main_agent/{session/test_async_persistence.py,streaming/test_cancellation_state.py}` — Strands SDK contract tests looking for `cancel_signal`/`async_mode` that the installed SDK lacks. Pre-existing, unrelated. +- `ruff check src/ tests/` repo-wide reports **369** pre-existing errors in untouched files. Scope ruff to your own files. + +--- + +## 3. Constraints that will bite you + +### DynamoDB one-GSI-per-deploy limit ⚠️ RELEASE-BLOCKING + +`UpdateTable` permits exactly **one** GSI create/delete per call, and CloudFormation +issues one per changed table. Two indexes on an existing table = failed deploy + full +stack rollback. **This took production down on 2026-08-01 in release 1.12.0.** + +This feature's `GSI7` (`KbWorkIndex`) consumes the **entire** `rag-assistants` GSI +budget for whatever release ships it. If another branch adds a GSI to that table, the +two cannot ship together. + +Guards: `infrastructure/test/gsi-update-limit.test.ts` (generation) and +`scripts/release/check-gsi-update-limit.mjs` (CI, vs `origin/main`). +Regenerate: `cd infrastructure && UPDATE_GSI_INVENTORY=1 npx jest gsi-update-limit` + +### DynamoDB cannot do arithmetic in a ConditionExpression + +`storedBytes + reservedBytes + :n <= :cap` is **rejected** (`Cannot parse condition +starting at:+ reserved <= :cap`). The byte cap therefore keeps a single `totalBytes` +accumulator and compares it against a **client-computed literal** (`cap - n`). One +atomic conditional `ADD`, so concurrent reservations cannot collectively overshoot. +See `byte_cap.py`. + +### DynamoDB reserved keywords bit this feature twice + +`total` and `ttl` are both reserved. Alias via `ExpressionAttributeNames`. The failure +is a `ValidationException`, which is loud — but it also masquerades as a caught +mutation (see §4). + +### Import boundary — the reason `kb_backend` exists as its own package + +`apis.shared.assistants.__init__` imports `rag_service`, which imports the embeddings +stack at module scope. Pulling that into a Lambda image blows the size budget. + +- `kb_backend/__init__.py` is **empty**, deliberately. +- Module-level imports are **stdlib only**; `boto3` and anything heavy is + function-local. +- Enforced by `backend/tests/architecture/test_kb_backend_boundary.py`. +- `apis.shared.embeddings` is a *separate* package and is fine to use. + +### AWS authorizes some Bedrock APIs under a *different* action name ⚠️ + +Three times now, a grant listing exactly the API the code calls has deployed clean, +reviewed clean, and failed on first real use: + +| API called | IAM action actually checked | Symptom | +|---|---|---| +| `CreateKnowledgeBase` with tags | `bedrock:TagResource` | fails the moment a real KB is created (§5.26) | +| (reconciler tag read) | `bedrock:ListTagsForResource` | fails closed — every KB looks untagged, orphan sweep reports a clean account forever (§5.26) | +| `IngestKnowledgeBaseDocuments` | **`bedrock:StartIngestionJob`** | every document upload `AccessDenied` (§5.31) | + +The action name and the API name are not the same namespace. **Check the service +authorization reference or the feature's own prerequisites page before assuming a +grant is complete**, and never infer completeness from a successful local run — SSO +identities are broader than every Lambda role here. + +Corollary: an action in a grant that no code calls is not necessarily dead. Read the +docblock before deleting it. + +### Module constants must be read at call time + +Never `def f(timeout=MODULE_CONSTANT)`. Python binds default arguments once at import, +so the constant becomes unpatchable. This cost a 33-second test that silently ignored +its own override. Use `timeout: Optional[float] = None` and resolve inside. + +### The knowledge-base component spec needs every collaborator stubbed + +`KnowledgeBaseSectionComponent` loads documents, crawls, sync policies and +connectors on hydration. Leave any of those services real and their HTTP requests +stay pending, so `fixture.whenStable()` never settles and **every test in the file +times out at 5 s** with "Test timed out in 5000ms" and no hint as to why. 30 of 31 +failed this way before the stubs went in. + +`quietCollaborators()` in `knowledge-base-section.component.spec.ts` provides all +six (`DocumentService`, `FileSourceService`, `WebSourceService`, +`SyncPolicyService`, `UserConnectorsService`, `OAuthConsentService`). Note +`OAuthConsentService`'s `completion` and `inFlightProviders` must be **signals**, +not plain values — the component calls them. + +--- + +## 4. Mutation testing — the discipline that has repeatedly paid + +Every security- or correctness-relevant assertion in this feature has been verified by +breaking the guard and watching a **specific, correct** test fail. This has caught +**seven** tests that passed with their guard removed. Do not skip it. + +**Four ways a mutation lies to you.** All four have happened here: + +| Trap | Symptom | Fix | +|---|---|---| +| Anchor never matched | reported "caught", file unchanged | `diff` the file; assert the anchor matches exactly once | +| Orphaned expression values | removing a `ConditionExpression` leaves its values unused → `ValidationException` → the **happy-path** test goes red | strip the orphaned values too, so the write genuinely succeeds | +| Syntax error | collection error mistaken for a detection | `ast.parse`/`py_compile` the mutant | +| Wrong test failed | something failed, but not the guard's test | always check **which** test failed by name | + +**Also:** never assert a constant against itself. `assert CAP == module.CAP` is a +tautology that follows the constant wherever it moves. Pin the literal, with a comment +saying why that number is a property of AWS rather than a knob. + +--- + +## 5. Defects found and fixed (do not reintroduce) + +### In my own spec + +1. **`PutMetricData` on a reserved namespace.** Req 20.10 originally scoped it to + `AWS/Bedrock/KnowledgeBases`. AWS reserves every namespace beginning with `AWS` and + rejects writes. The grant would have deployed cleanly and published **nothing**, + forever. Root cause: conflating *reading* Bedrock's own metrics (genuinely in that + namespace) with *writing* ours. Now `{projectPrefix}/ManagedKb`; Req 20.13 appended + for the read grant. +2. **Dead grant on the service role.** Same metric grant was also on the Bedrock + service role, which Bedrock assumes and which never publishes our metrics. Removed; + Req 20.10 now says "calling identities only". +3. **`managedKnowledgeBaseConfiguration={}`.** The shape has no *required* members, + but its only members are the embedding pin and encryption — so a literal `{}` makes + Req 8.5's pin unsatisfiable. "No required members" ≠ "must be empty". +4. **`float32`** → the enum value is **`FLOAT32`**; lowercase is rejected. +5. **Missing gate §14.3.** Authorization/publication was absent entirely while the + test matrix demanded tests for it. Added as Requirement 25 → group 11. +6. **Ordering issue:** tasks 4.5/4.7 reference the managed adapter, which task 8.1 + builds. Resolved with a fake backend conforming to the protocol — legitimate, since + the score conversion is adapter-local and the parity rules belong to the facade. + +### In the code + +7. **Reconciler arming bypass (MAJOR).** `lambda_handler` forwarded an `armed` field + from the invocation event, so an EventBridge target with constant + `{"armed": true}` — or anyone with `lambda:InvokeFunction` — would delete user + knowledge bases while all reviewable config said report-only. The pre-existing test + was named `test_an_event_cannot_arm_by_accident` but only covered the *string* + `"true"`; the boolean that actually armed was untested. +8. **Dispatcher over-grant undetectable.** A test asserted only one statement's shape, + so Bedrock permissions added in a *separate* statement went unnoticed. Now a + whole-role whitelist scan. +9. **Inline metadata unbounded** against a 50-attribute limit, and truncation was + alphabetical — which would have dropped `document_id`, the status filter's join + key. Reserved keys now go first. +10. **Latent config bug, twice.** `--context managedKb.x=…` sets a **flat dotted** + key; a nested-only `tryGetContext('managedKb')?.x` read silently ignores it. Hit + the byte caps and then the alarm thresholds. +11. **`{}` treated as an unreadable record (group 11).** `is_reclaim_exempt` used + `if not kb_record`, which conflated "absent, so fail closed" with "read, no + holds set". Every unheld knowledge base would have been exempt and the whole + predicate vacuous. `None` and `{}` are now distinct. Found by writing the test + first and believing it over the implementation. +12. **Requirement 25.6 had no IAM behind it (group 11).** There was no + resource-policy grant anywhere in the construct, so the sharing code would have + deployed as inert. Same category as defect 1: correct-looking, clean-deploying, + authorizes nothing. Now `grantManagedKbResourcePolicyAdmin`, on its own role, + with a test asserting no retrieval identity ever receives it. +13. **A resumed migration re-ingested everything (group 13).** The + completed-document set lived inside `migrationProgress`, which a later write + replaces wholesale, so a crash near the end of a 25-document corpus re-parsed + all 25 — 37–264 s each. Now a separate `migratedDocIds` string set updated with + `ADD` per batch. Found by the convergence property test counting a document + ingested twice. +14. **`promote_engine` permitted a second promotion (group 13).** Every guard it + had stayed true *after* a successful promotion, so two genuinely concurrent + workers would both succeed — exactly what Req 15.10 forbids. Now guarded on + `attribute_not_exists(retrievalEngine)`; rollback `REMOVE`s it, so a deliberate + re-promotion still works. +15. **Fixing 14 then broke resumption (group 13).** A resume after a successful + promotion had its write refused and marked the migration `failed` — a promoted + knowledge base with no retention window. `run_promote` now treats "already + promoted" as success, re-reading before deciding so a genuine guard failure + still raises. +16. **Four mutation-test lies, in one sitting (group 13).** A limit assertion the + final `[:limit]` trim masked; a derivation whose test was vacuous because the + priority list happened to be complete; a `match=` pattern loose enough that the + *other* check satisfied it; and an `except LeaseLost: raise` that was dead code + because the lease was taken outside the `try`. Each was fixed rather than + annotated. + +--- + +17. **Nothing registered the managed backend (group 14).** `register_backend` + was defined in task 4.2 and called by nothing. All 15 groups could have been + finished with the feature unreachable — a promoted record raises + `BackendUnavailable`, a correct fail-safe and a useless signal. Registration is + now at import, so there is no startup sequence to forget. +18. **A three-defect shell script (group 14).** `scripts/teardown/managed-kb.sh`, + all three found by *running* it: an infinite spin at a zero poll interval that + burned sixteen hours of a test run; `list | cut | grep -q` reporting false + absence when SIGPIPE became the pipeline's status under `pipefail`; and a + swallowed `list-knowledge-bases` failure reporting a clean teardown having + deleted nothing. `set -e` is suspended inside a function called in a condition, + which is why the last one was silent. +19. **Requirement 20.13 existed only as a comment (group 14).** The metrics *read* + grant was described in a comment explaining the write grant and never + implemented, so the reconciler could not have read Bedrock's own `Invocations`. + +--- + +20. **The tag contract had drifted three ways (post-group-14).** The Python wrote + keys `prefix`/`env` from variables the provisioning Lambda never receives; the + reconciler's filter was a documented *mirror* of that writer; the construct + declared different key names and exported the correct values as env vars + **nothing read**; and the teardown script read a third pair. Writer and + reconciler agreed only because both fell back to the same hardcoded defaults, + so the sole symptom was a teardown that matched nothing and reported success. + Now `kb_backend/tags.py` owns the keys and one fallback chain, and + `tests/supply_chain/test_kb_tag_contract.py` parses the TypeScript and the + shell script to assert agreement across all three languages. + + ⚠️ **Tag keys are namespaced** (`ManagedKbPrefix`, not `prefix`) because many + accounts carry an org-wide cost-allocation tag literally called `env`. Note the + KB_Record *attribute* `appKbId` is a different thing from the AWS *tag* + `ManagedKbAppKbId`; only the latter belongs to this contract. + +--- + +21. **Nothing enrolled a knowledge base (group 14.3).** The mirror image of + defect 17, and missed by it. `register_backend` made a promoted record + *servable*; this is about a record ever reaching `shadow` in the first place. + The worker only picks up records already in a migration state and the + dispatcher only sweeps GSI7, so with no enrolment surface both were correct + and inert. Task 14.3 was written as frontend-only, which is how it hid: the + missing piece was an **HTTP surface** nobody had scoped. Now + `apis/app_api/kb_upgrade/`. + +22. **A one-put enrolment would have stranded every knowledge base (group 14.3).** + `KbRecord.to_item` does not write `GSI7_PK`/`GSI7_SK` — only + `set_migration_state` maintains them. So the obvious enrolment (one + `put_item` with `migrationState="shadow"`) yields a record that reports an + upgrade in progress to every surface while being invisible to the dispatcher's + sweep **forever**: a spinner with nothing behind it, and no error anywhere. + Enrolment is therefore `create_provisioning` *then* `set_migration_state`, + both conditional. Two tests pin it, including one asserting the created record + does **not** carry `migrationState`. + +23. **An unrecognised failure reason leaked the operator's string (group 14.3).** + Found by mutation, not review. `_failure_reason` maps known tokens to + plain-language copy, and the test proved that for `ByteCapExceeded` — so + mutating the fallback to `return stored or _FAILURE_FALLBACK` **survived**, and + a user would have read `ClientError: An error occurred + (AccessDeniedException)…` in the card. Testing the mapped path proved nothing + about the unmapped one; that is the whole lesson. + `test_an_unrecognised_failure_does_not_leak_the_operator_string` pins it. + +24. **The upgrade flag never reached the service that reads it (pre-merge).** + Third instance of this feature's signature failure, and the most nearly + shipped. `kb-migration-construct.ts` sets all three `MANAGED_KB_*` booleans on + the four migration Lambdas; `app-api-environment.ts` set the byte caps and the + metric namespace but **not** `MANAGED_KB_MIGRATION_ENABLED` — which + `apis/app_api/kb_upgrade/service.py` reads to decide whether to offer the + upgrade at all. + + Setting the environment variable in GitHub would therefore have changed + nothing: the card would render `phase: "none"` for every user in every + environment, forever, with a clean deploy and no log line. Found only by + tracing where the flag is actually consumed before setting it. + + Now wired, shipped as an explicit `'false'` rather than omitted so the state + is readable in the task definition, and guarded by three tests in + `app-api-environment.test.ts`. The mutation — deleting the line, which is + precisely what the defect was — is caught. + + ⚠️ `managedKb.newDefault` has **no reader anywhere in `backend/src`**. It is + set on the Lambdas' environment and consumed by nothing, because + "new knowledge bases are created managed" is a follow-up spec (design §14.7 + steps 5–8), not this phase. Leave it off; turning it on is a no-op that reads + like a behaviour change. + +--- + +### The six that only running it revealed + +These came out in sequence over 2026-08-27 to 08-31, each one step further into the +saga than the last. Every one reviewed clean, deployed clean, and did nothing or +failed on first real use. If you read only one section of this file, read this one. + +25. **The dispatcher could not find the worker.** The construct set + `MANAGED_KB_WORKER_FUNCTION_NAME`; `dispatcher.py` reads + `KB_MIGRATION_WORKER_FUNCTION_NAME` — the convention its siblings use. Every + tick raised `RuntimeError: ... is not set`, on a fifteen-minute schedule, in + silence unless someone read the logs. `kb-sync` does not have this bug for + exactly one reason: `kb-sync.test.ts` asserts the variable exists. + A second mismatch found by the same sweep: the construct published + `MANAGED_KB_RETENTION_WINDOW_DAYS`, which nothing reads, while + `worker._retain_days()` reads `KB_MIGRATION_RETAIN_DAYS` — so Req 15.11's + configured window was silently replaced by the code's 30-day floor. **Guard:** + `tests/supply_chain/test_kb_migration_env_contract.py` now asserts every + `os.environ` name the handlers read is set by the construct, and that the + construct publishes nothing unread. + +26. **`bedrock:TagResource` was not granted.** `CreateKnowledgeBase` is called + *with* tags and AWS authorises the tagging as a separate action, so the grant + reviewed as complete and failed the moment a real knowledge base was created. + `bedrock:ListTagsForResource` was missing for the same reason, with a quieter + failure: the reconciler fails closed on a tag read, so every knowledge base + would look untagged and the orphan sweep would report a clean account forever. + +27. **`CreateDataSource` ran against a `CREATING` knowledge base.** + `CreateKnowledgeBase` returns before the knowledge base is usable — this + module's own header records 47–124 s to `ACTIVE` — and the code called + `CreateDataSource` immediately. `ConflictException` is deliberately not + retryable and `_call`'s backoff tops out near 60 s, so the wait had to be + explicit. **Why no test caught it:** `FakeBedrockAgent.create_knowledge_base` + returned `status: "ACTIVE"`, which the real API never does. + +28. **The knowledge base id was never recorded until both creates succeeded.** + `attach_aws_ids` needs both identifiers, so a failure between them left a + record with no `awsKbId` — and every later attempt re-entered the create path + and was refused, permanently, because the *name* was taken. **The + `clientToken` does not save this**: AWS idempotency tokens expire within + minutes. Earlier revisions of this document and of the module header claimed + otherwise; they were wrong. Fixed by `records.attach_knowledge_base_id` + (persist immediately) plus adopt-by-name for records already stuck. + **Why no test caught it:** two tests *certified* the bug — + `test_the_record_survives_as_a_discoverable_retry_anchor` asserted + `"awsKbId" not in anchor`, and its sibling relied on the fake modelling + `clientToken` dedup as **permanent** while not modelling name uniqueness at + all. The fake now enforces name uniqueness and treats tokens as expired by + default. + +29. **The embedding pin and managed reranking are mutually exclusive.** Req 8.5 + pinned `titan-embed-text-v2:0` via `embeddingModelType: CUSTOM`; Req 11.2 + requires `rerankingModelType: MANAGED`. AWS rejects the combination, and the + §13 evaluation had measured the two **separately, never together**. Req 8.5 is + now amended: the pin protected a failure mode that cannot occur in managed mode + (we never embed the query — Bedrock embeds both sides), and the evaluation + measured the pin as worth nothing while reranking measurably separates scores. + Confirmed in dev: pinned + `NONE` gives flat 1.00/0.982/0.952; unpinned + + `MANAGED` gives 0.413/0.199. + +30. **`verify` failed a good migration for being asked too early.** The canary + retrieval returned nothing because the freshly-ingested document was not yet + queryable, and that was treated as terminal. Measured: **~45 s** from ingest to + retrievable on a fresh knowledge base, against the docstring's "0.75–1.03 s" + (a warm-knowledge-base figure). `verify` now defers via + `records.defer_verify`, bounded at `MAX_VERIFY_ATTEMPTS`. Adoption also learned + to skip knowledge bases in `DELETING`, found the same way: a local recreate + adopted one mid-delete. + +--- + +### The seventh, from adding a document rather than migrating one + +31. **`bedrock:StartIngestionJob` was not granted, so direct ingestion could not + run at all.** Found by uploading a second document to the already-promoted + knowledge base in dev. The `DOC#` record went to `failed` carrying: + + ``` + AccessDeniedException ... IngestKnowledgeBaseDocuments ... not authorized + to perform: bedrock:StartIngestionJob on resource: knowledge-base/M8WQZVQJ8X + ``` + + AWS authorises `IngestKnowledgeBaseDocuments` under the **adjacent action + name** `bedrock:StartIngestionJob`. Both are listed in one statement in AWS's + direct-ingestion prerequisites + (`bedrock/latest/userguide/kb-direct-ingestion-prereq.html`). + `grantManagedKbDirectIngestion` carried only the name matching the API call, + so it reviewed as complete, deployed clean, and failed on first real use — + identical in shape to §5.26's missing `bedrock:TagResource`. Third occurrence + of that pattern; assume a fourth exists. + + **The worker had the same gap.** It receives the same grant and calls the same + API, so the first migration driven by the *deployed* dispatcher would have + failed the same way. It stayed hidden because every migration to date was + driven by `run-kb-migration.py` under an SSO identity broader than either + Lambda role — the exact limitation §2 names. The local driver cannot find this + class of defect, ever. Only a deployed run can. + + ⚠️ **The action looks like a mistake and is not.** Requirement 9.2 forbids + *calling* `StartIngestionJob` (0.1 RPS account-wide — one document per ten + seconds) and nothing calls it. Holding it is authorisation, not invocation. A + docblock on the grant and a separately-named test carry that reason so the + obvious cleanup fails a test that explains itself. + `bedrock:ListKnowledgeBaseDocuments` is in AWS's example policy and omitted + deliberately: no code path calls it. + + Guards: three tests across `managed-kb.test.ts` and `kb-migration.test.ts`, + one asserting **both** the worker and ingestion-consumer roles carry it. + Mutation-tested — removing the action fails exactly four tests, each named for + the reason. + +--- + +### The eighth through eleventh, all from one evening of running it — fixed in PR #900 + +These four are one root cause wearing four costumes: **the two engines were never +made exclusive.** The consumer stands down for a legacy document; nothing made the +legacy pipeline stand down for a managed one, and nothing propagated a deletion to +the managed engine at all. Every symptom below follows from that. + +32. ✅ **Routing exclusivity was enforced on only one side, so a document added to + a managed knowledge base was indexed twice.** design.md §537 states "a + document is indexed on exactly one backend outside a deliberate migration or + dual-read pilot, so no double-indexing", and task 9.2 claims tests for it. The + consumer does return immediately for a legacy document. But the **legacy + handler had no engine gate at all**, and its `s3:ObjectCreated:*` notification + on `assistants/` is still live alongside the (now enabled) EventBridge rule. + Both fired. + + Visible in real data on `DOC#DOC-dc8b65658e29`: `chunkCount: 8` and + `vectorStoreId: assistants-index` written by the legacy pipeline, while the + managed side failed with §5.31. + + **Why no test caught it:** every exclusivity test in + `test_kb_ingestion_consumer.py` is on the consumer's side + (`test_a_legacy_document_is_not_ingested_here` and siblings). Nothing asserted + the legacy handler skips a managed document, which is precisely the half that + did not exist. A test suite can be thorough about the side that works. + + Fixed: `handler.py` resolves the engine before writing anything and returns + early for `managed`. An earlier revision of this entry said the gate was + missing from "both copies" including + `infrastructure/bootstrap-assets/rag-ingestion/handler.py` — that was + misleading. The bootstrap copy is a 33-line no-op placeholder that indexes + nothing and needs no gate. + +33. ⚠️ **STILL OPEN. The document-status filter has one fail-*open* line in an + otherwise fail-closed function.** `_filter_vectors_by_document_status` opens + with `if not doc_ids: return vectors` — so a batch of chunks carrying no + `document_id` at all bypasses the DynamoDB check entirely and is served + unverified. Every other unprovable path in that function returns `[]` and + emits `METRIC_STATUS_FILTER_FAIL_CLOSED`. + + Not firing today: managed chunks do carry the id (§6 below), and legacy chunks + always have. It predates this feature. But `_document_id` returns `""` when + `location.customDocumentLocation.id` and both metadata mirrors are absent, + which is exactly the input that would trip it, and the failure is silent in + the serving direction. + +34. ✅ **Two writers owned one `status` field, so "ready" was decided by a + coin toss.** The sharpest consequence of §5.32, and worth its own entry + because the duplicate vectors are the cheap half of that defect. + + Both outcomes were observed within an hour of each other in dev: + + | document | what happened | + |---|---| + | `DOC-b5d5d8019f44` | legacy wrote `complete` at **+30 s**; the managed KB could not answer until **+95 s**. Sixty-five seconds of "your document is ready" followed by an answer that does not mention it. | + | `DOC-d637491d6cb1` | Docling produced **zero chunks** → `failed`. Bedrock indexed it fine and served it. It only ended up reading `complete` because the consumer finished **second**. | + + That second row is the alarming one, and the ordering was luck: reverse it — + purely a function of parse time — and a good, retrievable document reads + `failed` permanently, with no retry endpoint to recover it (task 14.4 open). + + It also defeated the exact protection `ingestion_consumer.py` documents in its + header: *"the UI says the upload worked, the user asks a question straight + away, and the answer does not mention their document."* The consumer polls + until the document is genuinely retrievable to prevent that. A second, + ungated writer undid it in one line. + + ⚠️ **The lesson generalises past this feature: any field two components can + write needs a stated owner.** Nobody chose this race; it appeared because a + new writer was added beside an old one and the question was never asked. + +35. ✅ **Bedrock's image extraction works, and it makes the legacy pipeline's + hard failure visible.** Not a defect in the new code — a capability + difference nobody had measured, found while testing an image-only PDF. + + A 4-year curriculum flowchart (465 KB, pure diagram, no text layer) on a + **legacy** assistant: `ValueError: Docling produced zero chunks`, status + `failed`, permanently unusable — `docling_processor.py` sets `do_ocr=False` + and `generate_page_images=False`. On the **promoted** assistant the same file + was retrievable in 94.5 s, and the chunks show Bedrock's vision model output + (` Hierarchical and Timeline Diagram`), including + prerequisite chains that exist nowhere in the document as text. + + `imageExtractionStatus: ENABLED` is set by `provisioning.py` and was confirmed + live on the dev data source. Note the old bundled `aws` CLI does **not** echo + `mediaExtractionConfiguration` back from `get-data-source` — it looked + unset until re-read with the pinned boto3. Do not conclude a field was not + applied from CLI output alone. + +36. ✅ **Deletion was never propagated to the managed engine, so a promoted + corpus could only grow.** `cleanup_service` removed the legacy S3 Vectors copy + and the `DOC#` row and never touched the managed knowledge base. The managed + delete path existed — `kb_backend/tombstones.py` — but its only callers are in + the reconciler, which is report-only with `reconcilerArmed` off. + + Three consequences, none of which raised anything: + + * storage paid for indefinitely at **$5.00/GB-month** against S3 Vectors' + ~$0.15 — orphans on the expensive engine; + * **silent retrieval degradation**, because the status filter runs *after* + retrieval: each orphan consumed a slot in `top_k` and was then dropped, so a + query could return five chunks and the model see two; + * the only thing preventing deleted content from being **served** was the + fail-closed status filter — which made §5.33's one fail-open line + load-bearing in a way it was never designed to be. + + Fixed with a third, engine-gated phase in `cleanup_document_resources`, + conjoined into `all_succeeded` so a failure blocks the hard delete. + + ⚠️ **The gate here is deliberately the opposite of the ingestion gate, and + that asymmetry is the interesting part.** On ingest, an unreadable KB record + resolves to *legacy*: being wrong costs a duplicate index while the consumer + still drives the document to a correct terminal state. On delete it must + **fail**: the `DOC#` row is what the status filter joins against, so reporting + success on a failed managed delete would remove the row *and* leave the + content — the one combination that turns a storage leak into a disclosure. + Same question, opposite answers, for a reason worth keeping. + + IAM again had no code to give it away: neither the app-api task role nor the + kb-sync worker could delete from a managed knowledge base, so this would have + deployed clean and failed on first use. New `grantManagedKbDocumentDeletion`, + narrower than `grantDirectIngestion` on purpose — these callers only remove, + so a bug in the delete path cannot add content and a bug in the ingest path + cannot remove it. Fourth instance of the §3 adjacent-action pattern, so it + carries `StartIngestionJob` pre-emptively rather than waiting to be taught. + + kb-sync needed the image change too: its worker calls + `cleanup_document_resources` after soft-deleting a document whose upstream + source has vanished. + +--- + +## 6. Remaining work + +### Do these first + +| | | +|---|---| +| **Merge #900** | Engine exclusivity, both halves. Triggers `backend.yml` (rebuilds rag-ingestion **and** kb-sync — the content hash moves because `kb_backend` was added to both images' `SOURCE_DIRS`) and `platform.yml` (the two new IAM grants). Wait for the platform deploy before testing a deletion on a promoted assistant, or the delete fails on IAM and the `DOC#` row is deliberately kept | +| **Then re-add a document in dev** | `DOC#DOC-dc8b65658e29` is parked at `failed` from §5.31 and is not retried retroactively — there is no reprocess endpoint (task 14.4). Re-upload; that path works | +| **Then drive a migration from the *deployed* dispatcher, not the local driver** | §5.31 is the proof that the local driver cannot see IAM defects: its SSO identity is broader than either Lambda role. Every remaining unknown in this feature is of that class | + +### Open, in rough order + +| Group | Notes | +|---|---| +| **§5.33** the one fail-open line | The only finding from 2026-08-31 still open. `if not doc_ids: return vectors` in `_filter_vectors_by_document_status`. Make it fail closed with `METRIC_STATUS_FILTER_FAIL_CLOSED` like every other unprovable path in that function, and pin it with a test that mutation-fails. Lower stakes now that §5.36 removes deleted content from the managed engine, but still the one silent-serving path left | +| **engine visibility** | Nothing logs *which* engine served a query — the resolver only logs on failure — so "is the new one actually working?" can only be answered from the KB record. One INFO line in the facade, plus a `Managed`/`Classic` badge in the knowledge base section, both unbuilt. Wanted before a wide rollout, because this feature's whole risk profile is silent regressions | +| **14.4** one-click document retry | Req 21.2. Ingestion is S3-event-triggered and there is no reprocess endpoint, so this needs new backend against a live pipeline. The card currently directs the user to re-upload, which works today. Close it by building the endpoint **or** by amending Req 21.2 to accept re-upload | +| **14.5** admin surface | not started. Filter by engine, stored bytes, document counts, bulk migrate, per-KB retry | +| **15.1** packaged-SDK probe | the *static* half is done and passing (`boto3==1.43.68` carries `MANAGED`, the embedding members, `FLOAT32`, all four document ops, no `AWS_DATA_PATH`). The live half has now effectively been done by hand — a real create → ingest → retrieve → promote succeeded in dev | +| **15.2** ingestion-concurrency probe | unanswered. Do not size a wide fleet migration before it | +| **15.3** full matrix | run it once the above land | + +### RESOLVED — the `document_id` / `relevance` "known unknown" was a false alarm + +An earlier revision of this file reported that +`search_assistant_knowledgebase_with_formatting` returned `document_id: None` and +`relevance: None` after promotion, and inferred that the status filter was either +not running on the managed path or losing the id. **Both inferences were wrong.** +Measured against the live dev knowledge base (`M8WQZVQJ8X`) on 2026-08-31 with the +pinned boto3 1.43.68: + +| Layer | Result | +|---|---| +| Raw `Retrieve` | `location.customDocumentLocation.id = "DOC-ae5cc5434f2d"` on both chunks | +| `ManagedKbBackend._to_chunk` | `document_id='DOC-ae5cc5434f2d'`, `relevance=0.4226 / 0.1616` | +| Facade output | `metadata.document_id = 'DOC-ae5cc5434f2d'`, `distance = −0.4226 / −0.1616` | + +The probe read `result["document_id"]` and `result["relevance"]` at the **top level** +of the facade output. Those keys have never existed. The facade has emitted exactly +four keys — `text`, `distance`, `metadata`, `key` — since the function was first +written (`git log -L` on the `formatted_results.append` block confirms it, back +through `e34d928c`). The id lives at `metadata.document_id`; relevance is exposed as +`distance`, its exact negation. Reading a key outside the contract returns `None` on +**both** backends, so the observation said nothing about the managed path. + +The status filter is genuinely running, not bypassed: it collected +`{DOC-ae5cc5434f2d}`, looked it up, found `status = "complete"`, and kept both +chunks. Verified independently — that `DOC#` record does read `complete`. + +**The lesson is about the probe, not the code.** Asserting on a response shape +nobody had checked against the producing function turned four correct layers into a +reported defect, and it was written up as the highest-priority open item. Confirm +the contract before believing a `None`. (While confirming it, §5.33 turned up as a +genuine finding in the same function.) + +### Known deferrals (correct, not oversights) + +- **Reconciler EventBridge wiring** (Reqs 14.1, 14.7) — the rule exists and is + enabled; `reconcilerArmed` stays off so it reports rather than deletes. +- **Group 7's snapshot reservation** has its caller (`run_shadow`). + +## 7. File map + +``` +.kiro/specs/managed-kb-migration/ requirements.md · design.md · tasks.md · HANDOFF.md +docs/specs/bedrock-managed-kb-evaluation.md the measured source of truth + +backend/src/apis/shared/kb_backend/ + __init__.py EMPTY, deliberately + records.py KB_Record + conditional transitions + protocol.py KnowledgeBaseBackend + frozen Chunk (score = relevance) + resolver.py engine → backend registry; absence ⇒ legacy; load_record + s3vectors_backend.py legacy adapter; converts distance → relevance HERE + managed_backend.py ManagedKbBackend: retrieval + direct ingestion + provisioning.py create saga + CUSTOM connector data source + byte_cap.py reserve / commit / release + tombstones.py delete sagas + resource_policy.py IAM-enforced sharing; staleness is state, not an event + dual_read.py pilot: start early, detach, compare, serve legacy + idleness.py activity = max(retrieval, bound agents' use) + tags.py THE tag contract — keys + value resolution, one place + query_guard.py 10,000-char clamp + metrics.py namespace + best-effort emit_count / emit_value + +backend/src/apis/shared/assistants/ + rag_service.py the FACADE — access gate, dual read, status filter, caps + kb_access.py KbAccess grant; reuses resolve_assistant_permission + kb_publication.py engine swap ≠ corpus change; reclaim exemption + +backend/src/apis/app_api/kb_migration/ + ingestion_consumer.py routes by engine; legacy ⇒ do nothing + reconciler.py daily join, report-only + dispatcher.py sparse-index sweep, bounded, no-ops when the flag is off + worker.py ONE step per invocation, leased, resumable + +backend/src/apis/app_api/kb_upgrade/ the OWNER-FACING surface (HTTP only) + models.py camelCase wire models; UpgradePhase + DocumentIssueKind + service.py phase derivation, enrolment, retry, notice, doc triage + routes.py 4 endpoints; read is any permission, writes are edit-only + NOT in kb_migration/: that package's modules share one + size-constrained Lambda image and this one imports the + embeddings-pulling assistants package + +frontend/ai.client/src/app/knowledge-base/ + kb-upgrade.service.ts fails soft; getStatus resolves to phase 'none' + knowledge-base-section.component.* the card: offer / progress / notice / failure + + the stranded-document disclosure + +scripts/local-dev/ + run-kb-migration.py drive the real worker in-process; --break-lease + kb-doc-timings.py per-document ingest timing + which engine did it + kb-compare-engines.py one query, both engines, side by side + +scripts/teardown/ + managed-kb.sh delete tag-matched KBs BEFORE any stack + +docs/specs/ + managed-kb-cost-attribution.md filter on usagetype, never service code alone + +infrastructure/lib/constructs/managed-kb/ + managed-kb-role-construct.ts Bedrock service role + grant methods + kb-migration-construct.ts 4 Lambdas sharing ONE image + alarms +``` + +### Where authorization lives, and why not in `kb_backend` + +`kb_access` and `kb_publication` sit in `apis.shared.assistants` because they reuse +`resolve_assistant_permission` and `listing.is_on_shelf`, and `kb_backend` may not +import that package. Authorization is above the seam by nature anyway: the answer is +the same whichever engine serves the query, so implementing it once above both +adapters is the only way it cannot differ between them. + +The facade's `access` parameter is **required and keyword-only**. Forgetting it is a +`TypeError` at the call site; a genuine denial passes `None` and fails closed. A +`KbAccess` cannot be built with a permission outside the read set, so holding one is +evidence the permission model was consulted — holding a string is not. + +Reclaim exemption keys on `listing.is_on_shelf`, **never** `is_listed`: an admin +requesting changes on a live listing leaves it serving but moves its state out of +`LISTED_STATES`, so by state name alone a reclaim pass would delete the corpus behind +an agent users can still see in the store. + +### Score direction — the highest-silent-risk detail + +S3 Vectors returns cosine **distance** (lower better). Managed returns **relevance** +(higher better). The protocol canonicalizes on `relevance`; `s3vectors_backend` +converts by **exact negation** (order-preserving and losslessly reversible, unlike +`1-d`); the managed adapter applies **no** conversion. The facade still emits a +derived `distance` key so no caller changed. + +Invert it and nothing raises — retrieval keeps returning five chunks and the answers +quietly get worse. `tests/property/test_pbt_kb_score_direction.py` is the only guard. + +--- + +## 8. Data model + +``` +PK = AST#{assistant_id} +SK = METADATA # the assistant row (pre-existing) +SK = KB#{app_kb_id} # app_kb_id == assistant_id THIS PHASE +SK = KBTOMB#{app_kb_id} # whole-KB tombstone, NO TTL +SK = KBTOMB#{app_kb_id}#DOC#{document_id} # document tombstone, NO TTL + +GSI7 "KbWorkIndex" (projection ALL) — sparse + GSI7_PK = KBWORK#{state} GSI7_SK = {dueAt ISO-8601} +``` + +Keys are written **only** while a record is work-eligible and `REMOVE`d on reaching a +terminal state, so ineligible knowledge bases are invisible to the dispatcher **by +physics** rather than by filter. Third use of this convention on this table +(`DueSyncIndex`, `AgentDirectoryIndex`, `AgentReportsIndex`). + +**Absence means legacy.** `retrievalEngine` is only ever written as `"managed"`. +Nothing writes `"s3vectors"` onto a record that lacked it — that is what makes the +migration zero-backfill across 1,692 existing records and makes rollback a single +attribute `REMOVE` rather than a data rewrite. + +Same convention for two more attributes: + +- `dualReadPilot` — read as `is True`, never truthiness. Absence is off. +- `policyAwsKbId` — the `awsKbId` the resource policy was last applied to. + `policy_is_stale` compares it against the live one, so re-application after a + replacement identifier is a comparison nothing can bypass by omission. + +Source bytes already live at +`assistants/{assistant_id}/documents/{document_id}/{filename}`. Migration is a +**re-ingest**, never a re-upload. diff --git a/.kiro/specs/managed-kb-migration/design.md b/.kiro/specs/managed-kb-migration/design.md new file mode 100644 index 000000000..b6e72a35c --- /dev/null +++ b/.kiro/specs/managed-kb-migration/design.md @@ -0,0 +1,963 @@ +# Design Document: Managed Knowledge Base Migration + +## Overview + +This design replaces the custom RAG retrieval backend (Docling → Titan → +Amazon S3 Vectors) with **Amazon Bedrock Managed Knowledge Base**, one knowledge +base at a time, behind a single abstraction seam, with rollback available at every +step. + +The shape of the change is a **strangler fig**. There are exactly two retrieval +call sites today, both routed through +`search_assistant_knowledgebase_with_formatting`. That function becomes a thin +facade over a `KnowledgeBaseBackend` protocol with two implementations. A +per-knowledge-base discriminator selects which one runs. Nothing above the seam +learns which backend it received. + +Three properties are load-bearing and everything else follows from them: + +1. **Absence is the default.** A knowledge base with no `retrievalEngine` + attribute is a legacy knowledge base. No backfill write is ever required, so a + half-finished rollout cannot half-break the fleet. +2. **The expensive resource is created late and deleted through a tombstone.** + Provisioning is lazy and idempotent; deletion writes a durable marker before it + calls AWS. +3. **Promotion is a single conditional write, and legacy data survives it.** + That makes rollback a pointer flip rather than a data restoration. + +### What is deliberately not here + +Phases 5–8 of the evaluation's §14.7 (managed-by-default, stopping legacy writes, +reclaiming legacy vectors, removing the old pipeline), agentic retrieval, any +change to the 2,000-character context cap, and the 1:1 → 0..N binding change (F4). +See `requirements.md` § "Scope boundary" and § "Non-goals". + +--- + +## Guiding measured constraints + +Every number here is measured in the evaluation, not assumed. They are collected +in one place because they are the reason the design has the shape it does. + +| Constraint | Measured value | Design consequence | +|---|---|---| +| `CreateKnowledgeBase` → ACTIVE | 47–124 s (n=7, median ≈73 s) | Never on an interactive path; lazy provisioning with generous timeouts | +| Per-KB cold first ingest | ~68 s, remarkably constant (68.296/68.232/68.334 s) | A fixed cost of the *knowledge base*, not the document; pay it once, in background | +| Warm ingest, small text | ~2.5 s | Comparable to today; bulk migration is feasible | +| Warm ingest, 50 KiB PDF | 68–264 s | Long tail; ingestion timeouts ≥300 s, treated as background work | +| INDEXED → actually retrievable | 0.75–1.03 s | Two distinct timestamps; poll for retrievable, not indexed | +| `Retrieve` p50 / p95 | 662–695 ms / 762–800 ms | +405 ms p50, +538 ms p95 vs today; acceptable but real TTFT cost | +| `StartIngestionJob` | 0.1 RPS, account-wide, **not adjustable** | Direct ingestion only; never per-document sync jobs | +| `IngestKnowledgeBaseDocuments` | **10 documents max**, server-enforced | Batch at 10, not the 25 the user guide claims | +| Concurrent Ingest+Delete document ops | 10 per account | Fleet migration throughput ceiling ~2 docs/s | +| `Retrieve` query input | 10,000 chars, **not adjustable** | Hard clamp at the seam | +| `Retrieve` RPM per KB | 600 + 25 RPS burst | Safe; per-KB isolation is the main quota win | +| `AgenticRetrieveStream` RPM | **60 per account** | Agentic retrieval cannot be a default path — out of scope | +| Managed storage | $5.00/GB-month | 35× today; byte caps are mandatory, not optional | +| Retrieval | $0.001/query | 2.3% of a $0.044 turn | +| Empty/idle KB | $0.00000203 measured for the month | No per-KB floor; count pressure is near zero | +| KB deletion | 2–6 minutes, async | Poll `ListKnowledgeBases`; "accepted" ≠ "gone" | +| Filter operators | fail **closed** (measured 0 results) | A mistyped filter yields nothing rather than leaking — but see the isolation note below | +| Managed reranking | separates scores 0.89/0.38/0.25/0.21/0.19 vs flat 1.00/0.84/0.78/0.77/0.77 | The reranker is what makes a 2,000-char cap defensible | + +--- + +## Architecture + +### The seam + +``` + inference_api/chat/routes.py app_api/assistants/routes.py + │ │ + └──────────────┬───────────────┘ + ▼ + search_assistant_knowledgebase_with_formatting() + (facade — unchanged public signature) + │ + ┌──────────────────┴──────────────────┐ + │ resolve_backend(app_kb_id) │ + │ reads KB_Record.retrievalEngine │ + │ absent ⇒ "s3vectors" │ + └──────────────────┬──────────────────┘ + ▼ + KnowledgeBaseBackend (Protocol) + │ + ┌───────────────────────┴───────────────────────┐ + ▼ ▼ + S3VectorsBackend ManagedKbBackend + (today's code, moved verbatim, (bedrock-agent + agent-runtime, + distance → relevance conversion) managedSearchConfiguration) + │ │ + S3 Vectors index Managed KB +``` + +New Python package: `backend/src/apis/shared/kb_backend/` + +| Module | Responsibility | +|---|---| +| `protocol.py` | `KnowledgeBaseBackend` Protocol, `Chunk` dataclass | +| `resolver.py` | `retrievalEngine` → backend instance; absence defaults to legacy | +| `s3vectors_backend.py` | Legacy adapter; owns distance → relevance conversion | +| `managed_backend.py` | Managed adapter; owns `managedSearchConfiguration` | +| `query_guard.py` | 10,000-character clamp + truncation metric | +| `records.py` | KB_Record read/write, conditional transitions | +| `provisioning.py` | The provisioning saga | +| `byte_cap.py` | reserve / commit / release | +| `tombstones.py` | Durable delete markers | + +> **Why a top-level package under `shared/`, not under `shared/assistants/`.** +> `kb_sync/records.py` documents that importing `apis.shared.assistants` "drags in +> the embeddings stack", which is why the kb-sync Lambdas use raw table access +> instead. The migration and ingestion Lambdas have that same constraint. Nesting +> the seam inside `assistants/` would force them to trip the very import the +> existing code goes out of its way to avoid. A sibling package keeps +> `apis.shared.kb_backend` importable by both the APIs and the Lambdas. +> +> Two rules make that hold rather than merely intend it: +> 1. `kb_backend/__init__.py` stays **empty** — no re-exports. +> 2. Heavy dependencies (`boto3` clients, the embeddings module) are imported +> **inside functions**, matching the existing convention in `kb_sync/records.py`. +> +> An architecture test asserts that `apis.shared.kb_backend` does not transitively +> import `apis.shared.assistants`, alongside the existing boundary tests in +> `backend/tests/architecture/`. + +> `apis/shared/assistants/vector_search.py` currently exists as a zero-byte +> placeholder. It is unused and unimported; leave it alone rather than repurposing +> it, so the new package's boundaries are unambiguous. + +### Component inventory + +| Component | Type | New or changed | Notes | +|---|---|---|---| +| `kb_backend/` package | library | **new** | The seam | +| `search_assistant_knowledgebase_with_formatting` | function | changed | Becomes a facade; signature preserved | +| `_filter_vectors_by_document_status` | function | changed | Fail closed (Req 5) | +| Ingestion consumer | Lambda | **new** | Replaces orchestration role of the Docling Lambda for managed KBs | +| Existing Docling ingestion Lambda | Lambda | unchanged | Still authoritative for legacy KBs | +| Migration dispatcher | Lambda | **new** | Copies `kb-sync` dispatcher shape | +| Migration worker | Lambda | **new** | Shares one image with the dispatcher | +| Reconciler | Lambda | **new** | Daily, report-only initially | +| KB service role | IAM role | **new** | One role serves many KBs | +| KB_Record | DynamoDB items | **new** | In the existing assistants table | + +### Why reuse the `kb-sync` topology + +`infrastructure/lib/constructs/kb-sync/kb-sync-construct.ts` already implements +exactly the shape this feature needs, and `scheduled-runs-construct.ts` documents +itself as following it closely — so this is the third use of an established +in-repo pattern, not a new invention: + +- two Docker Lambdas sharing **one** image (`backend/Dockerfile.kb-sync`); +- the platform-as-bootstrap pattern — CDK ships a byte-stable stub from + `bootstrap-assets/`, the workflow ships the real image via + `update-function-code`; +- SSM parameters publishing the generated function names so the deploy script can + find them; +- an EventBridge `rate()` schedule into the dispatcher; +- a bounded per-tick dispatch limit (`KB_SYNC_DISPATCH_LIMIT`, default 20). + +The migration Lambdas must also follow `kb_sync/records.py`'s **raw table access** +convention. That file exists for a documented reason: importing +`apis.shared.assistants` drags in the whole embeddings stack, and keeping the +Lambda image small is a deliberate constraint. The migration worker has the same +constraint and takes the same approach. + +--- + +## Data model + +All new items live in the **existing** assistants table +(`boisestateai-v2-rag-assistants`), preserving the adjacency-list convention. + +### KB_Record + +For this phase `App_KB_Id == assistant_id`, so the record is a sibling of +`METADATA` under the assistant's partition. This is exactly the compatible +phase-1 option §14.2 proposes, and it is what `compat.py` already anticipates: +its docstring states that when F4 lands `ref` "becomes a real KB id with no shape +change here". + +``` +PK = AST#{assistant_id} +SK = KB#{app_kb_id} # app_kb_id == assistant_id in this phase +``` + +| Attribute | Type | Notes | +|---|---|---| +| `appKbId` | S | Stable identity. What bindings reference | +| `ownerUserId` | S | For byte accounting and cost attribution. Opaque id, never email/PII | +| `visibility` | S | Mirrors the assistant's visibility in this phase | +| `retrievalEngine` | S | `"managed"`. **Never written as `"s3vectors"`** | +| `provisioningState` | S | `provisioning` / `active` / `failed` / `deleting` | +| `awsKbId` | S | AWS `knowledgeBaseId`. Replaceable. Never in a binding | +| `awsDataSourceId` | S | The `CUSTOM` connector id | +| `embeddingModelId` | S | `amazon.titan-embed-text-v2:0`. **Immutable** | +| `embeddingDimensions` | N | 1024. **Immutable** | +| `parserConfig` | M | Managed-parser settings captured at creation, including `imageExtraction`. Recorded because §14.2 requires immutable choices be persisted, and because a corpus indexed without image extraction is not comparable to one indexed with it | +| `imageExtraction` | BOOL | Convenience mirror of `parserConfig.imageExtraction` for queries | +| `storedBytes` | N | Committed bytes, from S3 `HEAD` | +| `reservedBytes` | N | In-flight reservations | +| `lastRetrievedAt` | S | Throttled write, one winner per 24 h | +| `migrationState` | S | `shadow` / `verify` / `promote` / `retain` / `failed`, plus `reclaim` reserved but never entered in this phase | +| `migrationGeneration` | N | Increments per attempt; guards stale workers | +| `migrationLeaseUntil` | S | Worker lease expiry | +| `migrationProgress` | M | `{migrated, total, lastDocumentId}` | +| `migrationError` | S | Plain-language reason for the UI | +| `promotedAt` / `rolledBackAt` | S | Rollback observation window anchors | +| `retainUntil` | S | Earliest eligible reclaim time | +| `pinned` / `exemptFromReclaim` | BOOL | Lifecycle exemptions | +| `clientToken` | S | Persisted so a retry reuses it | + +### Tombstone + +``` +PK = AST#{assistant_id} +SK = KBTOMB#{app_kb_id} # whole-KB delete +SK = KBTOMB#{app_kb_id}#DOC#{document_id} # document delete +``` + +Carries `intent`, `awsKbId`, `awsDataSourceId`, `createdAt`, `attempts`, +`lastError`. **No TTL** — a tombstone is cleared by confirmed deletion or it stays +as a work item. Letting TTL remove it would recreate the exact silent-leak class +this design exists to close. + +### Sparse GSI for work discovery + +Migration work is discovered through a **sparse** GSI: the key attributes are +written *only while the record is eligible*, so ineligible and pinned knowledge +bases are invisible to the scan **by physics** rather than by filter. + +This is an established convention on this exact table, not a new idea. Three +existing indexes already work this way and say so in their own comments: +`DueSyncIndex` (GSI4, written only while a sync policy is `active`), +`AgentDirectoryIndex` (GSI5, written only while a listing is `published`), and +`AgentReportsIndex` (GSI6, written only while a report is `open`). + +The table currently has **six** GSIs, named `GSI_PK`/`GSI_SK` for the first and +`GSI2_PK`/`GSI2_SK` through `GSI6_PK`/`GSI6_SK` thereafter. The new index is +therefore **GSI7**: + +``` +GSI: KbWorkIndex (partition GSI7_PK, sort GSI7_SK, projection ALL) + GSI7_PK = KBWORK#{state} # e.g. KBWORK#shadow + GSI7_SK = {dueAt ISO-8601} +``` + +When a knowledge base reaches a terminal state, the worker **removes** `GSI7_PK` +and `GSI7_SK`. A bug that fails to remove them causes repeated no-op work bounded +by the per-tick dispatch limit, not a runaway. + +Following the `AgentDirectoryIndex` precedent, the generic assistant-update path +must list `GSI7_*` as immutable, so a routine edit can never resurrect a work key +on a knowledge base that has left the queue. + +--- + +## Backend protocol + +```python +# apis/shared/kb_backend/protocol.py +from dataclasses import dataclass +from typing import Any, Protocol + +@dataclass(frozen=True) +class Chunk: + text: str + relevance: float # canonical: HIGHER IS MORE RELEVANT + document_id: str + metadata: dict[str, Any] + key: str + +class KnowledgeBaseBackend(Protocol): + async def search(self, kb_ref: str, query: str, top_k: int) -> list[Chunk]: ... + async def ingest(self, kb_ref: str, document_id: str, source: "DocumentSource") -> None: ... + async def delete_document(self, kb_ref: str, document_id: str) -> None: ... +``` + +### Score direction — the silent-failure risk + +This is the single most dangerous detail in the migration, because getting it +wrong produces **no error, just worse answers**. + +- S3 Vectors returns cosine **distance**: lower is better. The current formatted + result dict literally has a `"distance"` key, and its docstring says + *"lower = more similar"*. +- Managed KB returns **relevance**: higher is better. The probe measured + `score: 1.0` on an exact hit. + +The protocol canonicalizes on **relevance**. `S3VectorsBackend` performs the +conversion in its adapter, and `ManagedKbBackend` passes through. The facade keeps +emitting a `distance` key for any existing consumer during the transition, derived +from relevance, so no caller breaks on the field rename. + +A test asserts that for the same ordered input both backends rank the known-best +chunk first (Req 2.4). Without it, an inversion is undetectable by any other test +in the suite. + +### Query guard + +```python +MAX_QUERY_CHARS = 10_000 # Managed KB Retrieve cap; NOT adjustable +``` + +Applied in the facade, before backend dispatch, so both backends are protected +identically. Truncation emits a metric and never raises. This replaces the +existing inline comment in `bedrock_embeddings.py` asserting that the query is a +"short string, no token validation needed" — which is true only because Titan v2 +tolerates ~32,000 characters. + +### Retrieval configuration + +`ManagedKbBackend` sends `managedSearchConfiguration`, never +`vectorSearchConfiguration` — the latter is rejected outright for managed +knowledge bases: + +```python +retrievalConfiguration = { + "managedSearchConfiguration": { + "numberOfResults": top_k, # 5, parity + "rerankingModelType": "MANAGED", # NOT "NONE" + # "filter": {...} equals/in only for isolation-critical filters + } +} +``` + +Hybrid search is not configurable for managed knowledge bases and is simply how +managed retrieval works; there is no toggle to set and none is attempted. + +### The document-status filter runs on both backends + +Requirement 3.3 keeps the `status == "complete"` post-filter on the managed path +too, even though managed ingestion makes it largely redundant — because removing it +in the same change that swaps the engine would confound the comparison. Parity +means parity, including the parts that look unnecessary. + +This works on the managed path only because `customDocumentIdentifier` is set to +the platform's `document_id` (Requirement 9.4). The filter needs a `document_id` +per returned chunk; the 1:1 identifier mapping is what supplies it. Without that +mapping there would be nothing to join on, which is a second reason the `CUSTOM` +connector beats pointing a native S3 connector at the prefix. + +The filter is applied in the facade, above the seam, so there is exactly one +implementation and it **fails closed** on both backends (Requirement 5). Its +removal from the managed path is a follow-up-spec decision, made only once managed +is the sole engine. + +--- + +## Authorization, isolation, and publication + +This section closes evaluation gate §14.3. It is the gate most easily mistaken for +already-solved, because Managed KB ships two features whose names suggest they do +more than they do. + +### Three isolation levels, correctly ranked + +| Level | Mechanism | What it actually guarantees | +|---|---|---| +| **Weakest** | Metadata filter (`equals`/`in`) | *Logical* separation only. AWS's own multi-tenant guidance calls this "filter-level (logical) isolation, **not** IAM-enforced (infrastructure) isolation" | +| **Middle** | ACL-aware retrieval | Fails closed, which is better than today's document-status filter — but AWS states plainly that it "is not authorization" and does not authenticate users. Identity is **email only, with no alias resolution, and mismatches fail silently** | +| **Strongest** | One knowledge base per boundary, plus a resource policy | Genuine IAM-enforced `bedrock:Retrieve` / `bedrock:GetDocumentContent` | + +**Design consequence: the app remains the authorization authority.** Neither +metadata filters nor ACL-aware retrieval may be the sole thing standing between one +user's documents and another's. Because this phase keeps `App_KB_Id == +assistant_id`, the per-assistant boundary *is* a per-knowledge-base boundary, which +is the strongest of the three by construction. Filters are used for sub-scoping +within a knowledge base, never as the tenant boundary. + +The email-only identity limitation is why ACL-aware retrieval is **not** adopted in +this phase: this platform authenticates via OIDC with claim mappings, and a +silently-failing email match is a worse primitive than an explicit app-side check. + +### Invocation-time access resolution + +The runtime resolves the invoking user's access to a knowledge base **before** +retrieval, reusing the existing assistant permission model rather than inventing a +parallel one: + +- **owner / editor** — may read, may upload, may trigger an upgrade. +- **viewer** — may read through the agent; never sees the upgrade control. +- **no access** — retrieval is not attempted. + +Because this phase is 1:1, an agent's knowledge base is exactly the agent's own, so +"can this user invoke this agent" already answers "may this user's turn retrieve +from this knowledge base". A turn is never failed because of a knowledge base the +user cannot reach — there is no such case while the relationship stays 1:1. That +changes with F4, which is precisely why F4 is a separate spec: the "one +inaccessible knowledge base among N blocks the whole turn?" question only becomes +real then, and it is recorded here as inherited-open rather than answered +prematurely. + +### Published agents and corpus drift + +A marketplace listing freezes a knowledge base **reference**, not its contents, so +a published agent's answers can change after review without any re-review. This +phase does not solve that, and must not pretend to. It takes the one position that +is safe and reversible: + +- Migration **does not change** what a published agent retrieves — parity is the + whole contract, so an engine swap is not a corpus change and needs no re-review. +- A published agent is **exempt from lifecycle reclaim while listed**, and + `taken_down` requires an explicit transition rather than falling through to + reclaim. +- Whether published agents should pin a corpus revision, require re-review after + content changes, or bind only publisher-managed knowledge bases is an **open + question owned by the marketplace spec**, recorded in "Open questions carried + forward". Exemption from cleanup alone does not close that review bypass, and this + design does not claim it does. + +### Resource policies + +Resource policies are MANAGED-only and are the only mechanism here offering real +infrastructure isolation. This phase creates them only where a knowledge base is +shared beyond its owner. Because they attach to the **AWS knowledge base ARN**, any +cycle producing a new `awsKbId` silently drops sharing — so re-application after +rehydration is a tested invariant, not a runbook note. + +--- + +## Dual-read pilot + +The pilot exists so the rollout rests on evidence from *our* corpus and *our* +users, not solely on a 3-document benchmark. + +```mermaid +sequenceDiagram + participant F as Facade + participant L as S3VectorsBackend + participant M as ManagedKbBackend + participant U as User + + F->>L: search(query) + F->>M: search(query) %% concurrent + L-->>F: chunks (authoritative) + M-->>F: chunks (observation only) + F->>F: log overlap, rank correlation, per-backend latency + F-->>U: LEGACY results +``` + +Rules that make it safe to leave on: + +- **Legacy is always what is served.** The managed result is observation only. +- **The managed call is fire-and-forget with respect to correctness.** A managed + failure or timeout is logged and discarded; it can never fail the turn. +- **It must not add user-visible latency.** The two calls are concurrent and the + response is returned as soon as legacy resolves, so the managed call's 662–695 ms + p50 is not additive. +- **Opt-in per knowledge base, default off**, so pilot cost is bounded and + deliberate. + +Recorded per read: overlap in returned `document_id` values, rank correlation, and +per-backend latency. That is the same measure-first pattern used for the +prompt-cache and document-offload work. + +--- + +## Provisioning saga + +Ordering exists to guarantee that a crash leaves a **retry anchor**, never an +invisible paying resource. + +```mermaid +sequenceDiagram + participant IC as Ingestion Consumer + participant DDB as Assistants Table + participant BA as bedrock-agent + + IC->>DDB: conditional PutItem KB_Record
provisioningState=provisioning
attribute_not_exists(SK) + alt another worker already won + DDB-->>IC: ConditionalCheckFailed + IC->>IC: poll existing record until active + else this worker owns provisioning + DDB-->>IC: ok (clientToken persisted) + IC->>BA: CreateKnowledgeBase(type=MANAGED,
managedKnowledgeBaseConfiguration=embedding pin,
clientToken) + Note over IC,BA: 47-124 s to ACTIVE.
"Unable to verify embedding model"
is IAM eventual consistency -> RETRY + BA-->>IC: knowledgeBaseId + IC->>BA: CreateDataSource(MANAGED_KNOWLEDGE_BASE_CONNECTOR
connectorParameters={type:CUSTOM}
dataDeletionPolicy=RETAIN
imageExtractionStatus=ENABLED) + BA-->>IC: dataSourceId + IC->>DDB: conditional update -> active
attach awsKbId, awsDataSourceId + end +``` + +Five details that are each a defect if omitted: + +1. **DDB before AWS.** A crash after `CreateKnowledgeBase` leaves a + `provisioning` record the Reconciler can match against the orphan, so the + resource is adoptable rather than stranded. +2. **`clientToken` is built, not interpolated.** Minimum length is **33 + characters**; the natural `{id}-{variant}-kb` token is 31 and fails client-side + validation. It is persisted on the record so a retry reuses the same token and + AWS deduplicates. +3. **`dataDeletionPolicy: RETAIN` at creation.** This is the documented remedy for + the `DELETE_UNSUCCESSFUL` state, and the dev account already contains a + knowledge base stuck in it since 2025-11-24. Set it deliberately up front, not + as incident response. +4. **`imageExtractionStatus: ENABLED`.** Opt-in. Left default, chart and image + content is never described and never indexed — a silent loss of the capability + being paid for. +5. **The embedding-model verification failure is retryable.** It was observed + against a model confirmed `ACTIVE` and directly invokable. Treated as fatal, lazy + provisioning fails intermittently while pointing at the wrong cause. + +--- + +## Ingestion control plane + +The browser creates an `uploading` `DOC#` row and receives a presigned S3 PUT. +**There is no upload-complete API call**, so the bucket's `ObjectCreated` +notification remains the only trigger. A durable consumer is therefore required — +not an in-process `asyncio.ensure_future` task. + +```mermaid +sequenceDiagram + participant S3 as Documents Bucket + participant IC as Ingestion Consumer + participant DDB as Assistants Table + participant Old as Docling Pipeline + participant BA as bedrock-agent + + S3->>IC: ObjectCreated + IC->>DDB: read DOC# + KB_Record + alt retrievalEngine absent (legacy) + IC->>Old: existing pipeline (unchanged) + else retrievalEngine == managed + IC->>DDB: reserve bytes (S3 HEAD size) + IC->>IC: provisioning saga if needed + IC->>BA: IngestKnowledgeBaseDocuments
(<=10 docs, customDocumentIdentifier=document_id) + loop until retrievable + IC->>BA: GetKnowledgeBaseDocuments + end + IC->>BA: canary Retrieve (indexed != retrievable) + IC->>DDB: DOC# -> complete, commit bytes + end +``` + +- **Routing is exclusive.** A document is indexed on exactly one backend outside a + deliberate migration or dual-read pilot, so no double-indexing. +- **Two timestamps, not one.** `indexedAt` and `retrievableAt` are recorded + separately; the gap measured 0.75–1.03 s and is a real, distinct event. +- **Timeouts ≥300 s.** A 50 KiB PDF has been observed at 264 s. +- **No chunk-key bookkeeping.** `customDocumentIdentifier = document_id` gives a + 1:1 mapping, which retires the whole `{doc_id}#{chunk_index}` scheme including + `delete_vector_tail` and the chunk-shrinkage stash on the managed path. + +--- + +## Migration state machine + +```mermaid +stateDiagram-v2 + [*] --> legacy: no retrievalEngine + legacy --> shadow: owner opts in + shadow --> verify: all complete docs ingested + verify --> shadow: catch-up found new docs + verify --> promote: manifest match + canary pass + converged + promote --> retain: conditional write succeeded + retain --> reclaim: OUT OF SCOPE (follow-up spec) + shadow --> failed: unrecoverable + verify --> failed: manifest mismatch + failed --> legacy: stays usable, retry offered + retain --> legacy: rollback (pointer flip) +``` + +`retain` is the terminal state this spec reaches. `reclaim` is present in the enum +so the follow-up spec adds a transition rather than a schema change, but nothing +here enters it. + +| Phase | Work | Serving | User sees | +|---|---|---|---| +| `shadow` | Provision KB, re-ingest every `complete` doc from existing S3 keys | **legacy** | "Upgrading — 12 of 40 documents", fully usable | +| `verify` | Exact source manifest compare + canary retrieve | **legacy** | same | +| `promote` | Single conditional write `retrievalEngine="managed"` | managed | one-time success note | +| `retain` | Legacy vectors preserved ≥30 days | managed | nothing | +| `reclaim` | **Out of scope — follow-up spec.** The state exists in the enum and the machine reaches `retain` and stops | managed | nothing | + +### Timing, recomputed from the revised measurements + +The evaluation's §10.3 quoted "a 20-doc assistant ≈ 4 min; 100 docs ≈ 9.5 min", +but those totals were computed from the **superseded** §5 figures (85 s create, +~65 s first ingest, ~5 s each thereafter). Recomputed from §5.1's revised numbers: + +| Corpus | Arithmetic | Total | +|---|---|---| +| 20 small text documents | 73 s + 68 s + 19 × 2.5 s | **~3 min** | +| 100 small text documents | 73 s + 68 s + 99 × 2.5 s | **~6.5 min** | +| 20 native layout PDFs (50 KiB class) | 73 s + 68 s + 19 × (68–264 s) | **~24–86 min** | +| 20 scanned PDFs (260 KiB class) | 73 s + 68 s + 19 × (37–58 s) | **~14–21 min** | + +The two PDF rows are kept separate because the measurements come from two different +document classes and averaging them would invent a number: the 50 KiB *native* +PDF measured 68–264 s, while the 260 KiB *scanned* PDF measured 37–58 s. The larger +file was consistently faster, so size is not the predictor — content structure is. + +⚠️ **The PDF rows are the ones to plan around, and they are absent from the +evaluation's own estimate.** Per-document parse time dominates everything else for a +PDF-heavy corpus — the same 50 KiB PDF took 68 s, 89 s, 99 s and 264 s across four +runs. Progress reporting must therefore be per-document rather than +time-estimated, because a credible ETA cannot be computed up front. + +Fleet ceiling is ~2 documents/second given the 10-concurrent-document-operation +account limit — roughly 85 minutes for 10,000 documents, and that is a floor, not a +forecast, for the same reason. + +### Verification is a manifest, not a count + +Document-count parity would pass while content silently diverged. `verify` +compares an exact manifest of `document_id` + content hash or generation, then +performs at least one canary retrieval proving expected content comes back from +the managed side. Count parity alone is explicitly insufficient. + +### Writes and deletes during migration + +Coexistence is **converge-on-quiet**, not dual-write, so exactly one write path +stays authoritative until promotion: + +1. Uploads keep flowing to legacy as today. +2. The worker snapshots the doc-id set and migrates it. +3. A catch-up pass picks up anything created since the snapshot. +4. Repeat until a pass finds nothing new — the same shape as the crawler's + consecutive-miss rule. Warm ingest is ~2.5 s, so convergence is fast. +5. Every document's `DOC#` record is re-read **immediately before** ingesting it + and skipped if it is gone or no longer `complete`. Without this re-read, a + document deleted mid-migration resurrects in the new knowledge base. +6. Promotion is conditional on a converged pass, so two workers cannot both + promote. + +--- + +## Reconciler + +Runs daily. Joins a paginated, tag-filtered `ListKnowledgeBases` against +KB_Records. + +| Case | Action | +|---|---| +| AWS only | Orphan. Delete **only if the AWS-reported `createdAt` is >24 h old** | +| Record only | Stale pointer. Mark `vectorState: missing`, re-create on next ingest. **Never delete the record** — the documents are still valid | +| Both | Refresh `storedBytes` for quota accounting | +| Tombstone present | Retry the delete; escalate `DELETE_UNSUCCESSFUL` as an operator state | + +**Age-gate on `createdAt`, not on discovery time.** A reconciler that was down for +a week would otherwise wake up and delete every in-flight create. + +**Ships in report-only mode.** It logs what it would have deleted and deletes +nothing. It runs that way for weeks before being armed — the inverted flag +convention the evaluation calls for. The arming flag treats an **empty string as +off**, because the repo has been bitten by empty workflow variables before. + +--- + +## Byte cap accounting + +Storage is 35× more expensive per gigabyte than today. The existing 1 GB-per-user +file precedent, applied here at 30,000 users, is a **$150,000/month** exposure. +This is the only part of the design that can cause real financial damage. + +``` +reserve(owner, bytes) → conditional update, fails if committed + reserved + bytes > cap +commit(owner, bytes) → reserved -= bytes; stored += bytes +release(owner, bytes) → reserved -= bytes (on ingestion failure) +``` + +- Size comes from an **S3 `HEAD`** on the stored object, never from a + client-reported value. +- Reserve is a **conditional** update, so two uploads racing the same remaining + allowance cannot both win. +- The default per-owner cap is **lower** than 1 GB and resolves by role tier. +- `RawDataSize` is **not** used for enforcement: it returned 0 datapoints for a + directly-ingested document over a 60-minute lookback, and the cause is + unconfirmed. It may be used for reporting only. +- Cost-allocation tags are delayed reporting, not enforcement. + +### Concrete defaults + +The evaluation requires "a lower role-tier default" without naming one. Proposed, +and flagged as **requiring product sign-off before implementation**: + +| Tier | Per-owner cap | Worst-case at 30,000 users | +|---|---|---| +| Standard user | **100 MB** | 3 TB → ~$15,000/mo | +| Elevated (opt-in, admin-granted) | **1 GB** | — | +| Per-knowledge-base ceiling | **500 MB** | bounds a single runaway corpus | +| The 1 GB precedent, for contrast | 1 GB for everyone | 30 TB → **~$150,000/mo** | + +100 MB is ~88× the measured average of 1.13 MB per active user, so it is generous +in practice while cutting worst-case exposure 10×. Expected spend at full adoption +on measured behaviour remains ~$169/month; the gap between $169 expected and +$15,000 permitted is exactly why the alarms below are not optional. + +### Enforcement points + +The cap is checked at **every** path that can add bytes to a managed knowledge +base, not just interactive upload: + +1. **Upload** — reserve before ingest, commit on success, release on failure. +2. **Migration re-ingest** — the migration worker reserves for the whole snapshot + before entering `shadow`, and **fails the migration up front** rather than + part-migrating a corpus that will not fit. A knowledge base that exceeds its + owner's cap is surfaced as a plain-language failure with the option to request an + elevated tier. +3. **Rehydration** (follow-up spec) — same reserve path. + +Migration is the easy one to miss and the worst one to miss: it is the single +largest byte-adding operation in the system, and it is the one that runs +unattended. + +### Account-level alarms + +Per-owner caps bound one user. They do not bound the fleet, so gate §14.6 also +requires account-wide guards: + +| Alarm | Threshold | Why | +|---|---|---| +| Total managed KB storage | configurable GB | The only thing standing between expected and permitted spend | +| Managed KB count | 80% of the 10,000 default quota | The quota is adjustable, but capacity requests take lead time | +| `AmazonBedrockAgentCore` Knowledge-Base usagetype daily cost | configurable USD | Catches a cost shape no per-owner cap anticipated | +| `KbOrphansFound` sustained non-zero | any | The delete saga is leaking | + +Alarms use `TreatMissingData.NOT_BREACHING`, matching the posture of the existing +kb-sync, scheduled-runs, and prompt-cache observability constructs. + +### Who consumes retrieval quota + +Requirement 12.10 asks whether the knowledge base **owner** or the **invoking +user** consumes retrieval quota. Half the answer is a fact about AWS rather than a +choice, and it inverts the current model: + +| | Today (S3 Vectors) | Managed KB | +|---|---|---| +| `Retrieve` throughput | 20 rps **account-wide** | 600/min + 25 rps burst, **per knowledge base** | + +So the quota is consumed **per knowledge base**, which means the *owner's* knowledge +base absorbs the throughput of everyone who invokes their agent. The invoking user +does not carry a retrieval allowance of their own. + +**Decision: the owner is the payer, and this is an improvement, not a compromise.** +Today a single hot assistant can exhaust a 20 rps account-wide ceiling and degrade +retrieval for every other user on the platform. Per-KB quotas make that blast radius +one agent instead of the fleet — noisy-neighbour containment we do not currently +have. + +Consequences worth stating, because they follow from the decision rather than from +the implementation: + +* A **published** agent is the case to watch. Its knowledge base is one partition + serving an unbounded audience, so it is the only realistic way to approach + 600/min. Measured headroom is comfortable — ~26 requests/min average on a hot + shared knowledge base against an allowance of 600, about 4% — but the ceiling is + now per-agent and therefore reachable by a single popular agent in a way the + account-wide limit never made obvious. +* Retrieval is billed at **$0.001/query** and that cost attaches to the account, not + to a tenant. Attributing it per invoking user is a cost-reporting question, not a + quota question, and is out of scope here. +* Byte caps are per **owner**, consistent with this: the owner controls the corpus, + so the owner carries both its storage cost and its throughput ceiling. + +No code enforces a per-user retrieval allowance, deliberately. Adding one would +invent a limit AWS does not impose and that the existing per-agent permission model +already bounds. + +--- + +## IAM and encryption + +One Bedrock service role serves many knowledge bases — verified: a second +knowledge base created against the first one's role reached ACTIVE normally. +10,000 knowledge bases do not require 10,000 roles. + +| Control | Shape | +|---|---| +| Confused-deputy guard | `aws:SourceAccount` + `ArnLike` on `AWS:SourceArn` scoped to `knowledge-base/*` | +| PassRole | Caller's `iam:PassRole` conditioned on `iam:PassedToService` | +| S3 | Conditioned on `aws:ResourceAccount` | +| KMS | `serverSideEncryptionConfiguration.kmsKeyArn` where customer-managed keys are required | +| Separation | Provisioner/migrator CRUD, direct-ingestion, and inference `bedrock:Retrieve` scoped independently | +| Metrics (write) | `cloudwatch:PutMetricData` scoped to the non-reserved `{projectPrefix}/ManagedKb` namespace on the **calling identities only** — not the service role, which Bedrock assumes and which never publishes our metrics | +| Metrics (read) | `cloudwatch:GetMetricData` / `GetMetricStatistics` for Bedrock's own `AWS/Bedrock/KnowledgeBases` metrics | +| Async safety | Synchronous boto3 calls from async request paths run off the event loop | + +Three notes worth encoding rather than rediscovering: + +- **Metric publishing is best-effort and permission-gated.** Omit the + `PutMetricData` grant and metrics silently vanish while requests keep + succeeding. CDK assertions cover it. +- **The publish namespace must not begin with `AWS`.** CloudWatch reserves those + for its own services — "You cannot specify a namespace that begins with AWS" — + so `PutMetricData` scoped to `AWS/Bedrock/KnowledgeBases` authorizes nothing + that can ever succeed: a grant that reads as correct and silently does nothing. + Our own metrics (the table under Observability below) go to + `{projectPrefix}/ManagedKb`; the prefix keeps two environments in one account + from blending. Bedrock's `AWS/Bedrock/KnowledgeBases` metrics remain a **read** + source via `GetMetricData` / `GetMetricStatistics` — reading a reserved + namespace is fine, only writing is not. Do not "simplify" the two back into one + namespace. +- **Managed embedding and managed reranking need no Bedrock model access at all.** + Only `CUSTOM` does — and this design pins `CUSTOM` Titan v2 embeddings for + continuity across an immutable choice, so the grant is required. + +### Resource policies and rehydration + +Resource policies are MANAGED-only and give genuine IAM-enforced sharing for +`bedrock:Retrieve` and `bedrock:GetDocumentContent`. They attach to the **AWS +knowledge base ARN**, so any cycle producing a new `awsKbId` silently drops +sharing. Re-application after rehydration is a tested invariant (Req 24.12), not a +runbook step. + +### Teardown + +Managed knowledge bases are runtime-created and are **not** CloudFormation +children. `scripts/teardown/destroy.sh` must list and delete only resources tagged +for the project and environment, **before** deleting their service role and the +platform stack. Ordering is not cosmetic: deleting the role while a knowledge base +is still `DELETING` is a plausible route into `DELETE_UNSUCCESSFUL`, and a role +cannot be deleted until its inline policies are removed. + +--- + +## Observability + +EMF metrics alongside the existing PromptCache metrics. All of the following are +**our own** metrics and publish to `{projectPrefix}/ManagedKb` — never to +`AWS/Bedrock/KnowledgeBases`, which is reserved and rejects writes: + +| Metric | Why | +|---|---| +| `KbCount`, `KbStorageGB` | Leading indicators for the adjustable 10,000 cap and the storage curve; feed the alarms above | +| `KbIdleGB` | **Emitted for baseline only in this phase.** Nothing reclaims yet, but the follow-up spec needs historical idleness data to choose its eviction threshold, and that data cannot be backfilled | +| `KbOrphansFound` | **Sustained non-zero is the only signal the delete saga is leaking** | +| `KbQueryClamped` | Req 4 truncation rate | +| `KbStatusFilterFailClosed` | Req 5 — distinguishes a confirmed-empty result from an unconfirmable one | +| `KbMigration{Started,Promoted,Failed,RolledBack}` | Rollout health | +| `KbDualRead{Overlap,RankCorrelation,Latency}` | The pilot's whole output: how much the two engines agree, and what the managed one costs. Values rather than counts, and `KbDualReadLatency` is dimensioned per backend so the 662–695 ms vs 257 ms gap is measured on our own traffic rather than assumed from the benchmark | +| `KbDualReadFailed` | Managed-side failures during the pilot. Never user-facing — the turn was served from legacy before the comparison ran — but sustained non-zero says the engine is not ready | +| `KbByteCapRejected` | Whether the proposed 100 MB default is actually workable, before it hardens into policy | + +`KbReclaimedGBPerDay` is deliberately **not** emitted: nothing reclaims in this +phase, and a metric that is structurally always zero trains operators to ignore it. +It arrives with the reclaim tier. + +**Idleness** is `max(own lastRetrievedAt, max(lastUsedAt) over bound agents)` — +never retrieval alone, or an actively used agent's knowledge base is evicted +because its queries did not match. `lastRetrievedAt` uses the throttled +conditional write pattern (one winner per 24 h), never a write per retrieval. +Per-knowledge-base `Invocations` from `AWS/Bedrock/KnowledgeBases` is a cheaper +idleness signal and is preferred where it is sufficient. That is a **read** of +Bedrock's own namespace via `cloudwatch:GetMetricData` (Req 20.13), not a publish. + +**Cost attribution filters on `usagetype`.** Managed KB bills under +`AmazonBedrockAgentCore`, so anything keyed on `AmazonBedrock` misses it entirely, +and anything keyed on service code alone blends it into the AgentCore Runtime +memory line that is already 73% of that bill. + +--- + +## Flags and deployment choreography + +**Three** independent flags, all defaulting to off, all treating an empty string as +off: + +| Flag | Controls | This spec | +|---|---|---| +| `MANAGED_KB_NEW_DEFAULT` | new knowledge bases are created managed | ships **off** (phase 5) | +| `MANAGED_KB_MIGRATION_ENABLED` | the background migrator runs at all | ships **off**, enabled per-pilot | +| `MANAGED_KB_RECONCILER_ARMED` | the reconciler deletes, rather than only reporting | ships **off** — report-only for weeks first | + +The third is the inverted-convention flag: the reconciler is *deployed* from day +one but *disarmed*, so its judgement can be reviewed against real data before it is +allowed to delete anything. + +Deployment order is fixed by a hard rule: **backend code must never deploy before +the IAM and resources it requires.** + +1. **Platform** — additive schema, sparse GSI, service role, IAM, Lambda shells, + SSM parameters, teardown support. No behaviour change. +2. **Backend** — seam, both adapters, fail-closed filter, query clamp. All three + flags off, so managed code is dark. +3. **Pilot** — opt-in dual read on selected knowledge bases, still serving legacy. +4. **Opt-in migration** — owner-initiated, with the rollback observation window. + +Steps 5–8 of §14.7 are a follow-up spec. Because all three flags default off, +reaching +them is a configuration change rather than a code change. + +--- + +## UX surfaces + +| State | Surface | +|---|---| +| legacy, no action needed | **nothing** — no badge, no nag. A knowledge base that works needs no UI | +| upgrade available | Inline opt-in card: only benefits the §13 benchmark proved, expected duration, and "your knowledge base keeps working during the upgrade" | +| `shadow` / `verify` | Non-blocking progress ("Upgrading — 12 of 40 documents"); safe to navigate away | +| `promote` succeeded | One-time dismissible note. No permanent badge | +| failed | Plain-language reason + Retry. Stays on legacy, which keeps working. Never a dead end | + +Gated on the existing `_require_edit_permission`; viewers never see the control. +The word "vector" never appears in user-facing copy. No silent auto-migration in +this phase. + +**Admin surface:** knowledge bases filterable by engine, with stored bytes and +document counts, bulk migrate, and per-knowledge-base retry. + +### Surfacing failed and stuck documents + +Migration carries only `complete` documents. Measured against production, 200 of +1,692 `DOC#` records (11.8%) are not `complete`: 101 stuck `deleting`, 95 +`failed`, 4 `uploading`. Silently dropping the 95 failures is correct for the +index and wrong for the user — those people believe their uploads worked. The +upgrade flow surfaces them and offers retry. + +Two related messaging defects are in scope only to the extent of Req 21.4 +(distinguishing an unsupported format from a processing failure). The underlying +`.txt` ingestion bug — the deployed Docling build has no plain-text input format +despite the repo and frontend both advertising support, so a user waits 56 s for a +generic failure — is a **separate pre-existing bug**, not fixed here. + +--- + +## Testing strategy + +Mirrors the house pattern in `reliable-document-deletion`: unit tests plus +`hypothesis` property tests for invariants, with AWS stubbed. + +| Area | Approach | +|---|---| +| Adapter parity + score direction | Same input through both backends; assert identical ranking of a known-best chunk | +| Query clamp | Property: for any query length, output ≤10,000 chars and never raises | +| Fail-closed filter | Simulate table-level failure and missing table name; assert zero chunks | +| Byte cap races | Property: concurrent reserves never let committed total exceed the cap | +| Provisioning idempotency | Two concurrent first-ingests create exactly one KB | +| Crash after AWS create | Record left as a retry anchor; Reconciler adopts rather than duplicating | +| Reconciliation | Record-only and AWS-only cases; age-gate honours AWS `createdAt` | +| Migration interference | Upload and delete during migration; deleted doc never resurrects | +| Mixed deployment | Old and new code serving simultaneously; absent discriminator still resolves legacy | +| Resource policy rehydration | New `awsKbId` re-applies the policy | +| CDK assertions | IAM conditions from Req 20, including `PutMetricData` | +| Teardown | Only tagged resources deleted, and before the role | + +Managed AWS APIs are **stubbed**, never called live, so the suite stays +hermetic and free. + +--- + +## Open questions carried forward + +These remain genuinely open and are recorded so they are not mistaken for +settled: + +1. **Does a knowledge base go cold again after idleness?** In progress in the + evaluation. If a cold penalty exists, owners must be warned before eviction, + because rehydration pays the ~68 s cold-ingest cost. Affects the follow-up + spec's reclaim tier more than this one. +2. **Is there an account-level ingestion-concurrency limit?** The quota page lists + none. Probe with a many-knowledge-base backfill during the pilot before sizing a + wide migration. +3. **Does `RawDataSize` ever publish for directly-ingested documents?** Unconfirmed. + Until it does, byte accounting uses S3 `HEAD` (already the design). +4. **Native Google Drive connector vs the current AgentCore-Identity adapter.** + Never investigated. May sidestep the vault principal-binding dead-end at the + cost of moving token custody into Secrets Manager. +5. **`bedrock:GetDocumentContent` shape, size limits, and cost.** Relevant to + whole-document tasks that chunk retrieval structurally cannot serve. Unverified. diff --git a/.kiro/specs/managed-kb-migration/requirements.md b/.kiro/specs/managed-kb-migration/requirements.md new file mode 100644 index 000000000..7d4b9b4ee --- /dev/null +++ b/.kiro/specs/managed-kb-migration/requirements.md @@ -0,0 +1,892 @@ +# Requirements Document + +## Introduction + +This document specifies the requirements for replacing the platform's custom RAG +pipeline (Docling parse → Titan embed → Amazon S3 Vectors) with **Amazon Bedrock +Managed Knowledge Base** as the retrieval backend for assistant knowledge bases. + +The decision to proceed is grounded in `docs/specs/bedrock-managed-kb-evaluation.md`, +whose §13.4 decision gate was **cleared on 2026-08-14**: on a 9-question benchmark +with every variable held constant, the current pipeline answered 4/9 and managed +answered 9/9. Two document classes moved from unusable to working — native +layout-heavy PDFs (1/3 → 3/3) and scanned/OCR PDFs (0/3 → 3/3). + +The governing principle is **parity first, improvements later**. The user must +perceive nothing from the plumbing swap except the parser quality gain. Every +deliberate quality change that the evaluation identified as available — agentic +retrieval, raising the context cap, 0..N agent-to-KB bindings — is explicitly out +of scope here so that its effect remains attributable to itself. + +Migration is **additive and reversible at every step**. Legacy resources remain in +place for dual reads, rollback, and retention; no legacy resource is removed by +this spec. + +### Scope boundary + +This spec covers phases 1–4 of the evaluation's §14.7 choreography: + +1. additive schema, service role, IAM, worker resources, cleanup support; +2. dual backends dark, with mixed-version compatibility; +3. opted-in dual-read pilot, serving legacy; +4. opt-in migration with a rollback observation window. + +Phases 5–8 (managed-by-default for new KBs, stopping legacy writes, reclaiming +legacy vectors, and final target-state cleanup) are **deliberately deferred to a +follow-up spec**. The flags in Requirement 19 exist so that those phases are +config changes rather than code changes. + +### Non-goals + +The following are explicitly **not** in scope, each for a stated reason: + +- **Agentic retrieval.** Gated on the `AgenticRetrieveStream` account quota of + 60 requests/minute being raised (evaluation §6.4, §13.5 requirement 2). The + user-triggered escalation design in §6.5 is a separate future feature. +- **Raising the 2,000-character context cap.** The §13.6 experiment measured no + correctness change from 2,000 to 20,000 characters on either backend. Holding it + constant is required to keep the swap attributable (§9, §13.5 requirement 3). +- **0..N agent-to-KB bindings (F4).** §10.6 requires that the engine swap and the + binding-cardinality change not be coupled, because a joint failure is + unattributable. This spec lands the `KnowledgeBase` entity record while + preserving 1:1 binding semantics. +- **Routing conversation attachments through Managed KB.** §6.3 rejects this on + four grounds, including that a chat attachment ingested into a shared agent KB + becomes retrievable by every other user of that agent. Attachments remain + session-scoped inline blocks. +- **Native Google Drive connector evaluation.** §11 question 4, never + investigated; remains open. +- **Cleanup of the 101 stuck `deleting` and 95 `failed` legacy documents as a + standalone production migration.** Requirement 21 folds this into the migration + path instead. + +## Glossary + +- **Managed_KB**: An Amazon Bedrock Knowledge Base created with + `type: "MANAGED"`, which provisions no customer-visible vector store. Distinct + SKU from the classic `VECTOR` knowledge base, GA 2026-06-17. +- **Legacy_Backend**: The existing retrieval implementation over Amazon S3 + Vectors, as it exists today in `apis/shared/assistants/rag_service.py` and + `apis/shared/embeddings/bedrock_embeddings.py`. +- **Managed_Backend**: The new retrieval implementation over a Managed_KB. +- **KB_Backend_Protocol**: The Python `Protocol` defining `search`, `ingest`, and + `delete_document`, which both backends satisfy and behind which all callers sit. +- **Retrieval_Engine**: The per-knowledge-base discriminator selecting a backend. + Values are `"s3vectors"` and `"managed"`; **absence means `"s3vectors"`**. +- **KB_Record**: The new DynamoDB entity representing a knowledge base as a + first-class object, keyed by App_KB_Id. +- **App_KB_Id**: The stable, application-owned knowledge base identifier that + agent bindings reference. Never the AWS `knowledgeBaseId`. +- **AWS_KB_Id**: The AWS-assigned `knowledgeBaseId`, which is replaceable across a + dormancy/rehydration cycle and therefore never referenced by a binding. +- **Custom_Connector**: A Managed_KB data source of connector type `CUSTOM`, + nested inside the `MANAGED_KNOWLEDGE_BASE_CONNECTOR` envelope, which accepts + direct document ingestion. +- **Direct_Ingestion**: `IngestKnowledgeBaseDocuments`, which writes documents + into a Custom_Connector without a sync job, bypassing the + `StartIngestionJob` quota. +- **Ingestion_Consumer**: The durable S3 `ObjectCreated` event consumer that + replaces the current Docling ingestion Lambda's orchestration role. +- **Migration_Worker**: The background worker that moves one knowledge base from + Legacy_Backend to Managed_Backend through the Migration_State machine. +- **Migration_State**: The per-knowledge-base lifecycle + `shadow → verify → promote → retain`, plus a terminal `failed` state that returns + the knowledge base to Legacy_Backend, and a `reclaim` state reserved in the enum + but never entered in this phase. +- **Reconciler**: The daily job that joins `ListKnowledgeBases` against KB_Records + to detect orphaned AWS resources and stale pointers. +- **Doc_Status_Filter**: The query-time filter in + `rag_service._filter_vectors_by_document_status` that drops chunks whose parent + document is not `status == "complete"`. +- **Byte_Cap**: The enforced per-owner and per-knowledge-base limit on stored + source bytes. +- **Tombstone**: A durable DynamoDB marker written before an AWS delete call and + cleared only after AWS confirms deletion, so that a crashed delete is a + retryable work item rather than a silent leak. +- **Parity_Contract**: The set of retrieval properties held identical across both + backends so that the swap is perceptually invisible (Requirement 3). +- **Assistants_Table**: The existing DynamoDB table storing assistant and + document records (`PK=AST#{assistant_id}`, `SK=DOC#{document_id}`). + +## Requirements + +### Requirement 1: Backend Abstraction Seam + +**User Story:** As a developer, I want exactly one seam through which all +knowledge base retrieval and ingestion flows, so that the backend can be swapped +per knowledge base without any caller knowing which implementation it received. + +#### Acceptance Criteria + +1. THE system SHALL define a KB_Backend_Protocol in + `backend/src/apis/shared/kb_backend/` exposing `search`, `ingest`, and + `delete_document` operations. + +> Placement note: a top-level package under `shared/`, **not** under +> `shared/assistants/`. `apis/shared/assistants/__init__.py` imports +> `rag_service`, which imports `apis.shared.embeddings.bedrock_embeddings` at +> module level — so importing the assistants package drags in the embeddings stack. +> `kb_sync/records.py` uses raw table access specifically to avoid that, and the +> new Lambdas have the same constraint. Requirement 24.15 enforces the boundary by +> test. +2. THE system SHALL provide two implementations of KB_Backend_Protocol: + Legacy_Backend and Managed_Backend. +3. THE Legacy_Backend SHALL preserve the existing S3 Vectors behaviour, moved + without functional change. +4. WHEN a caller resolves a backend, THE system SHALL select it solely from the + knowledge base's Retrieval_Engine value. +5. THE two existing retrieval call sites (`inference_api/chat/routes.py` and + `app_api/assistants/routes.py`, both via + `search_assistant_knowledgebase_with_formatting`) SHALL be the only callers, + and SHALL NOT branch on backend identity. +6. WHEN a KB_Record has no Retrieval_Engine attribute, THE system SHALL resolve + the backend to Legacy_Backend. +7. THE system SHALL NOT write the value `"s3vectors"` to any record that does not + already carry it, so that backwards compatibility is achieved by absence and + requires zero backfill writes. + +### Requirement 2: Score Direction Canonicalization + +**User Story:** As a user, I want retrieved chunks ranked correctly regardless of +backend, so that answer quality does not silently invert when my knowledge base is +migrated. + +#### Acceptance Criteria + +1. THE KB_Backend_Protocol SHALL define chunk scores as **relevance**, where a + higher value is more relevant. +2. WHEN the Legacy_Backend returns S3 Vectors cosine **distance** values, THE + Legacy_Backend SHALL convert them to relevance before returning them across + the seam. +3. THE Managed_Backend SHALL pass Managed_KB relevance scores through unchanged. +4. THE system SHALL include a test asserting that, for the same ordered input, both + backends rank a known-best chunk first. + +### Requirement 3: Parity Contract + +**User Story:** As a user, I want a migrated knowledge base to behave exactly as it +did before except for parser quality, so that I cannot attribute any regression to +the upgrade. + +#### Acceptance Criteria + +1. THE system SHALL request `top_k = 5` on both backends. +2. THE system SHALL apply a context cap of **2,000 characters** on both backends, + unchanged from today's `max_context_length` default. +3. THE system SHALL retain the Doc_Status_Filter on **both** backends during + parity, even though Managed_Backend makes it redundant. +4. THE system SHALL build citations from the same `context_chunks` structure on + both backends, with the excerpt clip held at 500 characters. +5. THE system SHALL NOT enable agentic retrieval on any path. +6. THE system SHALL NOT alter the answer model, system prompt, or `top_k` as part + of this change. + +### Requirement 4: Query Length Clamp + +**User Story:** As a user, I want a long pasted message to still search my +knowledge base, so that I do not receive a hard failure for asking a long question. + +#### Acceptance Criteria + +1. WHEN a retrieval query is issued, THE system SHALL clamp the query string to at + most **10,000 characters** before it reaches the backend. +2. THE clamp SHALL be applied at the KB_Backend_Protocol seam so that it protects + both backends identically. +3. WHEN a query is clamped, THE system SHALL emit a metric or log record + identifying that truncation occurred. +4. THE clamp SHALL NOT raise an error or fail the turn. +5. THE system SHALL remove the inline assertion in + `apis/shared/embeddings/bedrock_embeddings.py` that no token validation is + needed for the query string. + +> Rationale: Managed_KB caps `Retrieve` query input at 10,000 characters and the +> quota is **not adjustable** (evaluation §6.4). Titan v2's ~32,000-character +> tolerance is the only reason nothing fails today. This is the single finding in +> the evaluation that produces a hard API failure rather than a cost or quality +> effect (§13.5 requirement 4). + +### Requirement 5: Fail-Closed Document Status Filter + +**User Story:** As a user who deleted a document, I want that document's content to +never be retrievable, so that a database problem cannot expose content I removed. + +#### Acceptance Criteria + +1. WHEN the Doc_Status_Filter cannot confirm a document's status because of a + table-level lookup failure, THE system SHALL drop that document's chunks. +2. WHEN the Doc_Status_Filter cannot confirm a document's status because the + documents table name is not configured, THE system SHALL drop all chunks. +3. WHEN the Doc_Status_Filter drops chunks because status could not be confirmed, + THE system SHALL emit a distinct error-level signal separating this case from an + ordinary empty-result case. +4. THE per-document lookup path SHALL continue to fail closed, as it does today. +5. **This requirement supersedes Requirement 3.4 of the + `reliable-document-deletion` spec**, which specified that a DynamoDB error + SHALL fall back to returning unfiltered results. +6. THE change SHALL ship as part of this feature's deployment, not as a standalone + production change. + +> Rationale: evaluation §7.4 documents this as a live fail-open path. §14.4 +> requires the filter fail closed before migration. The prior behaviour was a +> deliberate availability-over-privacy choice; retiring it is therefore a +> supersession and must be recorded as one. + +### Requirement 6: Knowledge Base as a First-Class Entity + +**User Story:** As a developer, I want a knowledge base to be its own record with a +stable identifier, so that the AWS resource behind it can be replaced without +breaking any agent binding. + +#### Acceptance Criteria + +1. THE system SHALL introduce a KB_Record persisted in DynamoDB. +2. THE KB_Record SHALL carry at minimum: App_KB_Id; owner identity; visibility or + ACL state; Retrieval_Engine; provisioning/lifecycle state; AWS_KB_Id; + data-source id; embedding and parser configuration including immutable choices; + stored-byte accounting; `lastRetrievedAt`; Migration_State with generation, + progress, lease, error and rollback timestamps; and pin/retention/exemption + flags. +3. Agent bindings SHALL reference App_KB_Id only. +4. THE system SHALL NOT persist AWS_KB_Id in any binding. +5. FOR this phase, THE system SHALL set `App_KB_Id == assistant_id`, preserving + the existing 1:1 relationship. +6. WHEN no KB_Record exists for an assistant, THE system SHALL treat it as a + virtual legacy S3 Vectors knowledge base and SHALL NOT create a record as a + side effect of a read. +7. THE system SHALL NOT change the cardinality of the agent-to-knowledge-base + relationship, and the existing rejections in `bindable_catalog.py` and + `binding_validation.py` SHALL remain in force. +8. THE test suite SHALL assert that an explicit `knowledge_base` binding is still + rejected and that `bindable_catalog` still returns an empty list for it, so the + 1:1 freeze is enforced by test rather than by intention. + +### Requirement 7: Lazy Provisioning Saga + +**User Story:** As a system operator, I want a knowledge base created in AWS only +when it is first needed and never duplicated, so that we do not pay for empty +resources or strand orphans. + +#### Acceptance Criteria + +1. THE system SHALL NOT call `CreateKnowledgeBase` when an assistant or knowledge + base is created. +2. WHEN the first document for a knowledge base is successfully ready to ingest, + THE system SHALL provision the Managed_KB. +3. THE system SHALL write the KB_Record in a `provisioning` state **before** + calling AWS, and SHALL attach returned identifiers with a conditional write. +4. WHEN two ingestions race to provision the same knowledge base, THE system SHALL + create at most one Managed_KB. +5. THE system SHALL pass a `clientToken` that satisfies the API's **33-character + minimum**, 256-character maximum, and + `[a-zA-Z0-9](-*[a-zA-Z0-9]){0,256}` pattern. +6. THE system SHALL construct the `clientToken` programmatically rather than by + interpolating a template that may fall below the minimum length. +7. WHEN `CreateKnowledgeBase` fails with a message indicating the embedding model + could not be verified, THE system SHALL treat the failure as retryable. +8. WHEN provisioning is interrupted after the AWS call but before the conditional + write, THE KB_Record SHALL remain a durable retry anchor discoverable by the + Reconciler. + +> Rationale: §5.1 measured `CreateKnowledgeBase` → ACTIVE at 47–124 s (n=7), so +> this must never sit on an interactive path. The "embedding model could not be +> verified" failure was observed to be pure IAM eventual consistency against a +> model confirmed ACTIVE and invokable. + +### Requirement 8: Managed Knowledge Base Configuration + +**User Story:** As a system operator, I want each Managed_KB created with the exact +configuration the evaluation validated, so that we do not silently lose a +capability we are paying for. + +#### Acceptance Criteria + +1. THE system SHALL call `CreateKnowledgeBase` with `type: "MANAGED"`, a + `roleArn`, and `managedKnowledgeBaseConfiguration`. + +> Shape note, verified against the packaged botocore service model: +> `managedKnowledgeBaseConfiguration` has **no required members**, but its only +> members are `embeddingModelType`, `embeddingModelArn`, +> `embeddingModelConfiguration` and `serverSideEncryptionConfiguration`. So the +> embedding pin required by criterion 5 below has nowhere else to live, and sending +> a literal `{}` would make that criterion unsatisfiable. "No required members" is +> not the same as "must be empty" — earlier drafts of this spec said `{}`, which is +> why this note exists. +2. THE system SHALL omit `storageConfiguration` entirely. +3. THE system SHALL create its data source with + `dataSourceConfiguration.type = "MANAGED_KNOWLEDGE_BASE_CONNECTOR"` and the + real connector type in + `managedKnowledgeBaseConnectorConfiguration.connectorParameters`. +4. THE system SHALL use connector type `CUSTOM`. +5. THE system SHALL NOT send an embedding pin — no `embeddingModelType`, no + `embeddingModelArn`, no `embeddingModelConfiguration` — and SHALL let Bedrock + choose and manage the embedding model for a managed knowledge base. + + > **Amended 2026-08-31, by measurement.** This criterion previously required + > `embeddingModelType: CUSTOM` pinned to `amazon.titan-embed-text-v2:0` at + > `FLOAT32` and 1024 dimensions. That was carried over from the Legacy_Backend + > without re-deriving it, and it is wrong here for two independent reasons. + > + > **It protected a failure mode that cannot occur.** On S3 Vectors *we* embed + > the user's question (`s3vectors_backend` → `apis.shared.embeddings`), so the + > query model must match the model that indexed the documents or the similarity + > search compares vectors from different spaces. Managed retrieval sends + > `retrievalQuery={"text": …}` and managed ingestion sends `inlineContent` + > text — we never produce a vector. Bedrock embeds both sides itself, so + > consistency is the service's invariant and not ours to get wrong by omission. + > + > **AWS refuses the pin together with managed reranking.** Measured against + > dev, all four combinations: + > + > | Embedding | `rerankingModelType` | Result | + > |---|---|---| + > | `CUSTOM` (pinned) | `MANAGED` | `ValidationException` | + > | `CUSTOM` (pinned) | `NONE` | ok — scores 1.00 / 0.982 / 0.952 (flat) | + > | default (managed) | `MANAGED` | ok — scores 0.413 / 0.199 (separated) | + > | default (managed) | `NONE` | ok | + > + > The pin and Requirement 11.2's reranking are therefore mutually exclusive. + > §13 measured the pin as worth nothing ("identical cold-ingest time and + > identical answer quality to the built-in embedding — 9/9 either way") and + > reranking as worth a great deal ("the reranker is what makes a small context + > cap defensible"). Keeping reranking is the side with evidence behind it. + > + > Requirement 8.8's immutability still holds and now matters more: the choice + > cannot be revisited per knowledge base after creation. + +6. THE system SHALL enable + `mediaExtractionConfiguration.imageExtractionConfiguration.imageExtractionStatus + = ENABLED` on the data source. +7. THE system SHALL set the data source's `dataDeletionPolicy` to `RETAIN` at + creation time. +8. THE system SHALL treat embedding configuration as immutable after creation and + SHALL NOT attempt to change it. +9. A single Bedrock service role SHALL be reusable across many Managed_KBs. + +> Rationale: §11.1 — image extraction is opt-in and silently indexes nothing if +> left default; custom Titan v2 embeddings measured identical cold-ingest time and +> identical 9/9 quality, and preserve continuity with today's embedding across an +> immutable choice; `dataDeletionPolicy: RETAIN` is the documented remedy for the +> `DELETE_UNSUCCESSFUL` state already observed in the dev account. + +### Requirement 9: Direct Document Ingestion + +**User Story:** As a user uploading documents, I want ingestion to keep up with +bulk uploads, so that a large batch is not serialized behind an API quota. + +#### Acceptance Criteria + +1. THE system SHALL ingest documents using Direct_Ingestion into the + Custom_Connector. +2. THE system SHALL NOT use `StartIngestionJob` for per-document ingestion. +3. THE system SHALL send at most **10 documents** per + `IngestKnowledgeBaseDocuments` call. +4. THE system SHALL set `customDocumentIdentifier` to the platform's + `document_id`. +5. THE system SHALL treat concurrent `Ingest` and `Delete` document operations as + limited to **10 per account** and SHALL bound its own concurrency accordingly. +6. THE system SHALL NOT carry forward the `{doc_id}#{chunk_index}` vector-key + bookkeeping, including `delete_vector_tail` and the chunk-shrinkage stash, on + the Managed_Backend path. + +> Rationale: `StartIngestionJob` is 0.1 RPS account-wide and not adjustable +> (§9). The API reference caps the document array at 10; AWS's user guide claim of +> 25 was disproven server-side for managed knowledge bases (§11.1). + +### Requirement 10: Durable Ingestion Control Plane + +**User Story:** As a user, I want an upload to reliably become searchable even if a +worker crashes, so that documents do not silently fail to index. + +#### Acceptance Criteria + +1. THE Ingestion_Consumer SHALL be a durable, retryable compute resource triggered + by the documents bucket's `ObjectCreated` notification. +2. THE Ingestion_Consumer SHALL resolve each document's knowledge base and + Retrieval_Engine before doing any work. +3. WHEN a document belongs to a legacy knowledge base, THE Ingestion_Consumer + SHALL route it to the existing pipeline. +4. WHEN a document belongs to a managed knowledge base, THE Ingestion_Consumer + SHALL route it to Direct_Ingestion. +5. THE system SHALL NOT index the same document on both backends outside of a + deliberate migration or dual-read pilot. +6. THE Ingestion_Consumer SHALL poll until the document is not merely reported + indexed but **actually retrievable**, and SHALL record those as two distinct + timestamps. +7. THE Ingestion_Consumer SHALL update the `DOC#` record to a terminal + complete or failed state with bounded retries and a durable retry anchor. +8. THE system SHALL NOT perform ingestion orchestration in an in-process + `asyncio.ensure_future` task. +9. THE Ingestion_Consumer SHALL tolerate ingestion latency of at least 300 + seconds for a single document. + +> Rationale: §14.1 — the browser creates an `uploading` row and receives a +> presigned PUT; there is no upload-complete API call, so the S3 event remains the +> only trigger. §5.1 measured a fixed per-knowledge-base warm-up of ~68 s and a +> long tail to 264 s on a 50 KiB PDF, so timeouts must be generous. + +### Requirement 11: Managed Retrieval Configuration + +**User Story:** As a user, I want retrieval against a managed knowledge base to use +the correct API shape and managed reranking, so that results are well ordered. + +#### Acceptance Criteria + +1. THE Managed_Backend SHALL use `managedSearchConfiguration` and SHALL NOT send + `vectorSearchConfiguration`. +2. THE Managed_Backend SHALL request managed reranking rather than + `rerankingModelType: NONE`. +3. THE Managed_Backend SHALL NOT attempt to configure or toggle hybrid search. +4. WHEN a metadata filter is applied, THE system SHALL rely on filters failing + **closed**, as measured. +5. THE Managed_Backend SHALL constrain any isolation-critical filter to `equals` + or `in`. + +> Rationale: §5.1 — `vectorSearchConfiguration` is rejected outright for managed +> knowledge bases. §11 question 3 measured `equals`, `startsWith` and +> `stringContains` on an impossible key all returning 0 results, disproving the +> silent-ignore/fail-open claim. §11.1 — managed reranking measurably separates +> scores (0.89/0.38/0.25/0.21/0.19 versus a nearly flat 1.00/0.84/0.78/0.77/0.77 +> without it), and **the reranker is what makes a 2,000-character cap defensible**. + +### Requirement 12: Enforceable Storage Cost Controls + +**User Story:** As a platform owner, I want stored bytes capped per owner before any +managed knowledge base holds production data, so that storage cost cannot grow into +a six-figure monthly exposure. + +#### Acceptance Criteria + +1. THE system SHALL enforce a per-owner Byte_Cap and a per-knowledge-base + Byte_Cap. +2. THE per-owner default SHALL be **100 MB**, an elevated admin-granted tier SHALL + be **1 GB**, and the per-knowledge-base ceiling SHALL be **500 MB**. All three + SHALL be configurable and resolvable by role tier. These values require product + sign-off before implementation. +3. THE system SHALL determine a document's contribution to the Byte_Cap from an S3 + `HEAD` on the stored object, NOT from a client-reported size. +4. THE system SHALL apply byte accounting as an atomic reserve → commit → release + flow. +5. WHEN two uploads race against the same remaining allowance, THE system SHALL + NOT allow the combined committed total to exceed the Byte_Cap. +6. WHEN an ingestion fails, THE system SHALL release the reservation. +7. THE system SHALL NOT depend on the `RawDataSize` CloudWatch metric for + enforcement. +8. THE system SHALL NOT depend on cost-allocation tags for enforcement. +9. THE Byte_Cap SHALL be enforced before any production traffic is promoted to + Managed_Backend. +10. THE system SHALL define and document whether the knowledge base owner or the + invoking user consumes retrieval quota. +11. THE Byte_Cap SHALL be enforced on **every** path that adds bytes to a managed + knowledge base, including the migration re-ingest path, not only interactive + upload. +12. WHEN a knowledge base's corpus would exceed its owner's remaining allowance, + THE Migration_Worker SHALL reserve for the whole snapshot and fail the + migration **before** entering `shadow`, rather than part-migrating a corpus + that cannot fit. +13. THE system SHALL raise account-level alarms on total managed storage, on + managed knowledge base count against the 10,000 quota, on daily + Knowledge-Base `usagetype` cost, and on a sustained non-zero orphan count. +14. THE system SHALL emit a metric when a Byte_Cap reservation is rejected, so the + chosen default can be validated against real behaviour before it hardens into + policy. + +> Rationale: §13.5 requirement 1 — managed storage is $5.00/GB-month against +> ~$0.15/GB-month today, a 35× increase. The existing 1 GB-per-user allowance +> would permit 30,000 GB at full adoption, i.e. **$150,000/month**. This is the +> only finding in the evaluation that can cause real financial damage. +> `RawDataSize` returned 0 datapoints for a directly-ingested document (§11 +> question 2), so it is unproven for this purpose. + +### Requirement 13: Deletion Sagas and Tombstones + +**User Story:** As a system operator, I want every delete to either complete or +leave a retryable work item, so that a failed delete is never a silent paying leak. + +#### Acceptance Criteria + +1. WHEN deleting a knowledge base, data source, or document, THE system SHALL write + a Tombstone **before** calling AWS. +2. THE system SHALL clear the Tombstone only after AWS confirms the resource is + gone. +3. THE system SHALL NOT treat an accepted delete call as a completed deletion. +4. THE system SHALL verify knowledge base deletion by polling until the resource is + absent, tolerating at least 6 minutes. +5. THE system SHALL NOT delete a knowledge base's service role until all of its + knowledge bases are confirmed absent. +6. THE system SHALL NOT remove the last KB_Record, nor allow TTL to remove it, + until AWS confirms deletion. +7. WHEN a knowledge base reports `DELETE_UNSUCCESSFUL`, THE system SHALL surface it + as an actionable operator state rather than a completed delete. +8. A surviving Tombstone SHALL be discoverable as a retryable work item. + +> Rationale: §12 measured deletion taking 2–6 minutes and verified only by polling +> `ListKnowledgeBases`. §12.2 documents a knowledge base stuck in +> `DELETE_UNSUCCESSFUL` since 2025-11-24 that no reconciler would ever notice. + +### Requirement 14: Daily Reconciler + +**User Story:** As a system operator, I want a daily job that finds AWS resources +our database does not know about, so that crash orphans are detected rather than +paid for indefinitely. + +#### Acceptance Criteria + +1. THE Reconciler SHALL run on a schedule and join a paginated, tag-filtered + `ListKnowledgeBases` against KB_Records. +2. WHEN a Managed_KB exists in AWS with no KB_Record, THE Reconciler SHALL treat it + as an orphan. +3. THE Reconciler SHALL age-gate orphan deletion on the **AWS-reported + `createdAt`**, NOT on the time of discovery. +4. THE Reconciler SHALL delete an orphan only when it is older than 24 hours. +5. WHEN a KB_Record references an AWS_KB_Id that does not exist, THE Reconciler + SHALL mark the record's vector state as missing and SHALL NOT delete the + record. +6. WHEN both sides agree, THE Reconciler SHALL refresh stored-byte accounting. +7. THE Reconciler SHALL run in a report-only mode that logs intended deletions + without performing them, and report-only SHALL be the initial deployed mode. +8. THE Reconciler SHALL apply a bounded per-run action limit. + +> Rationale: §7.4 — age-gating on discovery time means a reconciler that was down +> for a week deletes in-flight creates. §7.3 requires shipping in report-only mode +> and arming later, and warns specifically about the empty-string workflow-variable +> case. + +### Requirement 15: Migration State Machine + +**User Story:** As a knowledge base owner, I want my knowledge base upgraded without +downtime and without re-uploading anything, so that the upgrade is invisible until +it succeeds. + +#### Acceptance Criteria + +1. THE system SHALL migrate a knowledge base through Migration_State + `shadow → verify → promote → retain`, with `failed` as a terminal state that + returns the knowledge base to Legacy_Backend. `reclaim` is reserved in the enum + and SHALL NOT be entered in this phase. +2. THE system SHALL NOT mutate a live knowledge base in place. +3. DURING `shadow` and `verify`, THE knowledge base SHALL remain fully usable and + SHALL continue serving from Legacy_Backend. +4. THE system SHALL re-ingest source bytes from their existing S3 location and + SHALL NOT ask the user to re-supply any document. +5. THE system SHALL migrate only documents whose status is `complete`. +6. THE `verify` step SHALL compare an exact source manifest of `document_id` plus + content hash or generation, NOT document-count parity alone. +7. THE `verify` step SHALL perform at least one canary retrieval that confirms + expected content is returned from the Managed_Backend. +8. `promote` SHALL be a single conditional write flipping Retrieval_Engine to + `"managed"`. +9. THE system SHALL NOT promote unless a catch-up pass has converged. +10. WHEN two workers attempt promotion concurrently, THE conditional write SHALL + allow at most one to succeed. +11. DURING `retain`, THE system SHALL preserve legacy vector data for a rollback + window of at least 30 days. +12. THE system SHALL NOT enter `reclaim` for a knowledge base until the retention + window has expired AND that knowledge base has served managed traffic without + a rollback. +13. THE Migration_Worker SHALL take a lease so that one knowledge base is not + migrated concurrently by two workers. +14. THE Migration_Worker SHALL apply a bounded per-tick dispatch limit. + +> Rationale: §10.3. Timing recomputed from §5.1's revised figures (~73 s median +> create + ~68 s first ingest + ~2.5 s per warm small document): a 20-document +> knowledge base is **~3 minutes** and 100 documents **~6.5 minutes**. §10.3's own +> "4 min / 9.5 min" figures were computed from the **superseded** §5 numbers and are +> not used here. For a PDF-heavy corpus, per-document parse time of 37–264 s +> dominates and a 20-PDF knowledge base can exceed an hour — so this is background +> work only, and progress must be reported per-document rather than as an ETA. + +### Requirement 16: Writes and Deletes During Migration + +**User Story:** As a user, I want to keep uploading and deleting documents while my +knowledge base is upgrading, so that the upgrade does not freeze my work or corrupt +the result. + +#### Acceptance Criteria + +1. DURING migration, THE existing upload path SHALL remain authoritative and SHALL + continue writing to Legacy_Backend. +2. THE Migration_Worker SHALL snapshot the document-id set, migrate it, then run a + catch-up pass for documents created since the snapshot. +3. THE Migration_Worker SHALL repeat catch-up passes until a pass finds nothing + new. +4. THE system SHALL re-read each document's `DOC#` record immediately before + ingesting it, and SHALL skip the document if it no longer exists or is no longer + `complete`. +5. THE system SHALL NOT resurrect a document that was deleted mid-migration. +6. THE system SHALL NOT implement dual-write as the coexistence mechanism. + +### Requirement 17: Rollback + +**User Story:** As a knowledge base owner, I want an upgrade to be undoable, so that +a bad outcome is recoverable immediately rather than requiring data restoration. + +#### Acceptance Criteria + +1. THE system SHALL support rollback by writing Retrieval_Engine back to its prior + value. +2. Rollback SHALL NOT move or restore any data. +3. Rollback SHALL be available for the entire `retain` window. +4. WHEN a migration fails at any stage before `promote`, THE knowledge base SHALL + remain on Legacy_Backend and SHALL remain fully usable. +5. THE system SHALL record a rollback timestamp on the KB_Record. + +### Requirement 18: Dual-Read Pilot + +**User Story:** As a platform owner, I want real comparative evidence before +migrating anyone, so that the rollout rests on measurement rather than on the +benchmark alone. + +#### Acceptance Criteria + +1. THE system SHALL support running both backends for the same query on an opted-in + knowledge base. +2. DURING a dual read, THE system SHALL serve results from Legacy_Backend. +3. THE system SHALL record, per dual read, the overlap in returned `document_id` + values, a rank correlation, and per-backend latency. +4. THE dual-read path SHALL be opt-in per knowledge base and SHALL default to off. +5. THE dual-read path SHALL NOT increase user-visible latency beyond the legacy + path's own latency. + +### Requirement 19: Independent Feature Flags + +**User Story:** As a platform operator, I want to ship the managed backend without +starting a fleet migration, so that the two risks are separable. + +#### Acceptance Criteria + +1. THE system SHALL provide a flag controlling whether new knowledge bases are + created managed. +2. THE system SHALL provide a separate flag controlling whether the + Migration_Worker runs at all. +3. THE system SHALL provide a third, separate flag controlling whether the + Reconciler deletes rather than only reporting. +4. THE three flags SHALL be independently settable. +5. ALL three flags SHALL default to off. +6. WHEN the migration flag is off, THE Migration_Worker SHALL perform no work. +7. WHILE the Reconciler arming flag is off, THE Reconciler SHALL log intended + deletions and delete nothing. +8. THE system SHALL treat an empty-string flag value as off. + +### Requirement 20: IAM, Encryption, and Teardown + +**User Story:** As a security engineer, I want least-privilege, confused-deputy-safe +roles and a teardown that removes runtime-created resources, so that the feature +neither over-grants nor leaks resources. + +#### Acceptance Criteria + +1. THE system SHALL define a dedicated Bedrock knowledge base service role. +2. THE service role's trust policy SHALL constrain `aws:SourceAccount` and SHALL + apply an `ArnLike` condition on `AWS:SourceArn` scoped to `knowledge-base/*`. +3. THE caller's `iam:PassRole` grant SHALL be conditioned on + `iam:PassedToService`. +4. S3 access SHALL be conditioned on `aws:ResourceAccount`. +5. WHERE customer-managed encryption is required, THE system SHALL supply + `serverSideEncryptionConfiguration.kmsKeyArn`. +6. THE system SHALL scope provisioner/migrator CRUD, direct-ingestion, and + inference `bedrock:Retrieve` permissions separately. +7. WHEN synchronous AWS SDK calls are made from an async request path, THE system + SHALL execute them off the event loop. +8. THE teardown script SHALL list and delete only resources tagged for the project + and environment, and SHALL do so **before** deleting their service role and the + platform stack. +9. THE system SHALL include CDK assertions covering the IAM conditions in this + requirement. +10. THE system SHALL grant `cloudwatch:PutMetricData` scoped to the + `{projectPrefix}/ManagedKb` custom namespace on the **calling identities only**. + THE namespace SHALL NOT begin with `AWS`. THE Bedrock service role SHALL NOT + receive this grant. +11. WHEN a Managed_KB is created, THE system SHALL tag it with the project prefix, + the environment, the App_KB_Id, and the owner identity. +12. THE owner tag value SHALL be an opaque identifier and SHALL NOT be an email + address or any other personally identifying value. +13. THE identities that read Bedrock's own per-knowledge-base metrics SHALL be + granted `cloudwatch:GetMetricData` and `cloudwatch:GetMetricStatistics`. Those + metrics live in the `AWS/Bedrock/KnowledgeBases` namespace, which is a **read + source only** and is never a `PutMetricData` target under 20.10. + +> Note: tagging is a hard prerequisite, not housekeeping. Requirement 14.1's +> tag-filtered `ListKnowledgeBases` and Requirement 20.8's teardown both read these +> tags; without them the Reconciler cannot distinguish our resources from anything +> else in the account, and teardown cannot scope itself. + +> **Why 20.10's namespace is not an `AWS/...` one, and must not be "fixed" back to +> one.** CloudWatch reserves every namespace beginning with `AWS` for its own +> services: "You cannot specify a namespace that begins with AWS. Namespaces that +> begin with AWS are reserved for use by Amazon Web Services products." A +> `PutMetricData` grant scoped to `AWS/Bedrock/KnowledgeBases` therefore authorizes +> no publish that can ever succeed — it reads as correct in a policy review and +> silently does nothing. 20.10 and 20.13 cover two different directions of traffic +> that were previously conflated: +> +> - **Writing** this platform's OWN metrics (`KbByteCapRejected`, `KbOrphansFound`, +> `KbIdleGB`, `KbCount`, `KbStorageGB`, `KbQueryClamped`, +> `KbStatusFilterFailClosed`, `KbMigration{Started,Promoted,Failed,RolledBack}`) +> needs `PutMetricData` into the non-reserved `{projectPrefix}/ManagedKb` +> namespace (20.10). The project prefix keeps two environments in one account from +> blending their metrics. +> - **Reading** Bedrock's own per-KB metrics (`Invocations`, `ClientErrors`, +> `ServerErrors`, `Throttles`, `TotalIterationCount`, `RawDataSize`) needs +> `GetMetricData` / `GetMetricStatistics` against `AWS/Bedrock/KnowledgeBases` +> (20.13). Reading a reserved namespace is permitted; only writing is not. + +> Rationale: §14.5 and §14.0. Metric publishing is best-effort and +> permission-gated: omit the grant and metrics silently vanish while requests keep +> succeeding. Managed embedding and managed reranking need no Bedrock model access; +> only `CUSTOM` embedding or reranking does — and Requirement 8.5 chooses `CUSTOM` +> embedding, so that grant is required. + +### Requirement 21: Failed and Stuck Legacy Documents + +**User Story:** As a user whose upload failed months ago without telling me, I want +to find out and retry, so that migration does not quietly drop my document. + +#### Acceptance Criteria + +1. WHEN a knowledge base is migrated, THE system SHALL surface to its owner any + document not in `complete` status that will therefore not be carried across. +2. THE system SHALL offer a retry path for such documents. +3. THE system SHALL NOT silently omit non-`complete` documents without surfacing + them. +4. THE system SHALL distinguish, in user-facing messaging, an unsupported file + format from a processing failure. + +> Rationale: §7.4 measured 1,692 `DOC#` records of which 200 (11.8%) are not +> `complete` — 101 stuck `deleting`, 95 `failed`, 4 `uploading`. §10.3 ingests only +> `complete` documents, so migration would silently drop all 95 failures. §11.2 +> documents that the deployed pipeline cannot ingest `.txt` at all despite the repo +> and frontend both advertising support, producing a 56-second wait and a generic +> failure message. + +### Requirement 22: Observability + +**User Story:** As a system operator, I want to see knowledge base count, stored +bytes, orphans and migration progress, so that cost and correctness problems are +visible before they become incidents. + +#### Acceptance Criteria + +1. THE system SHALL emit metrics for at least: knowledge base count, stored + gigabytes, idle gigabytes, orphans found, and Byte_Cap rejections. THE system + SHALL NOT emit a reclaimed-gigabytes metric, because nothing reclaims in this + phase and a structurally-always-zero metric trains operators to ignore it. +2. THE system SHALL emit migration progress and failure counts. +3. THE system SHALL emit a metric when a query is clamped per Requirement 4. +4. THE system SHALL emit a metric when the Doc_Status_Filter drops chunks because + status could not be confirmed per Requirement 5. +5. THE system SHALL derive idleness from the maximum of the knowledge base's own + last-retrieved time and the last-used time of any bound agent, NOT from + retrieval alone. +6. THE system SHALL NOT write a last-retrieved timestamp on every retrieval. +7. THE system SHALL attribute cost by filtering on `usagetype`, NOT on service code + alone. +8. THE system SHALL treat a sustained non-zero orphan count as the signal that the + delete saga is leaking. + +> Rationale: §7.2 — idleness computed from retrieval alone evicts an actively used +> agent's knowledge base because its queries did not match. §7.3 requires a +> throttled conditional write rather than per-retrieval writes; §14.0 notes +> per-knowledge-base `Invocations` is a cheaper idleness signal. §8 — Managed KB +> bills under `AmazonBedrockAgentCore`, so anything keyed on `AmazonBedrock` misses +> it entirely and anything keyed on service code alone blends it into the Runtime +> memory line. + +### Requirement 23: User Experience + +**User Story:** As a knowledge base owner, I want the upgrade explained honestly and +never forced on me, so that I keep working normally and understand what changed. + +#### Acceptance Criteria + +1. WHEN a knowledge base is on Legacy_Backend and no action is required, THE system + SHALL show no badge, banner, or prompt. +2. WHEN an upgrade is available, THE system SHALL present it as an inline, opt-in + control describing only benefits proven by the §13 benchmark and stating that + the knowledge base keeps working during the upgrade. +3. DURING `shadow` and `verify`, THE system SHALL show non-blocking progress and + SHALL allow the user to navigate away. +4. WHEN promotion succeeds, THE system SHALL show a one-time dismissible notice and + SHALL NOT show a permanent badge. +5. WHEN migration fails, THE system SHALL show a plain-language reason and a retry + control, and the knowledge base SHALL remain usable on Legacy_Backend. +6. THE system SHALL NOT use the word "vector" in user-facing copy. +7. THE upgrade control SHALL be gated on existing edit permission, and viewers + SHALL NOT see it. +8. THE system SHALL NOT auto-migrate knowledge bases silently in this phase. +9. THE admin surface SHALL list knowledge bases filterable by engine with stored + bytes and document counts, and SHALL support bulk migrate and per-knowledge-base + retry. + +### Requirement 24: Minimum Test Coverage + +**User Story:** As a reviewer, I want the risky paths covered by tests before +promotion, so that correctness does not rest on manual verification. + +#### Acceptance Criteria + +1. THE test suite SHALL cover adapter parity across both backends, including score + direction. +2. THE test suite SHALL cover create, ingest, and delete idempotency. +3. THE test suite SHALL cover a crash after the AWS create call but before the + database update. +4. THE test suite SHALL cover record-only and AWS-only reconciliation outcomes. +5. THE test suite SHALL cover uploads and deletes occurring during migration. +6. THE test suite SHALL cover fail-closed document status and fail-closed access + checks. +7. THE test suite SHALL cover byte-cap reservation races. +8. THE test suite SHALL cover a mixed old/new deployment serving simultaneously. +9. THE test suite SHALL cover teardown of tagged dynamic resources. +10. THE test suite SHALL include CDK assertions for the IAM conditions in + Requirement 20. +11. THE test suite SHALL stub managed AWS APIs rather than calling them. +12. THE test suite SHALL assert that resource policies are re-applied after a + rehydration that produces a new AWS_KB_Id. +13. THE test suite SHALL assert the presence of the CloudWatch metric permissions + in Requirement 20.10. +14. THE test suite SHALL cover published-agent corpus behaviour, asserting that an + engine swap does not alter what a published agent retrieves and that a listed + agent is exempt from lifecycle reclaim. +15. THE test suite SHALL assert that `apis.shared.kb_backend` does not transitively + import `apis.shared.assistants`, so the Lambda image constraint is enforced by + test rather than by convention. + +### Requirement 25: Authorization, Isolation, and Publication Semantics + +**User Story:** As a user, I want my knowledge base readable only by people who are +allowed to read it, so that sharing an agent does not silently expose my documents. + +#### Acceptance Criteria + +1. THE system SHALL resolve the invoking user's access to a knowledge base **before** + retrieval is attempted. +2. THE system SHALL reuse the existing assistant permission model rather than + introducing a parallel one, so that owner, editor, and viewer semantics are + unchanged. +3. THE system SHALL treat the application as the authoritative authorization layer. +4. THE system SHALL NOT rely on a metadata filter as the tenant boundary. +5. THE system SHALL NOT adopt ACL-aware retrieval as an authorization mechanism in + this phase. +6. WHERE a knowledge base is shared beyond its owner, THE system SHALL apply a + resource policy for IAM-enforced `bedrock:Retrieve`. +7. WHEN a rehydration or replacement produces a new AWS_KB_Id, THE system SHALL + re-apply any resource policy that was attached to the previous identifier. +8. WHEN a knowledge base's engine is migrated, THE system SHALL NOT change what a + published agent retrieves. +9. WHILE an agent is listed in the marketplace, THE system SHALL exempt its + knowledge base from lifecycle reclaim. +10. WHEN a listed agent transitions to `taken_down`, THE system SHALL require an + explicit transition rather than allowing it to fall through to reclaim. +11. THE system SHALL NOT claim to resolve whether published agents pin a corpus + revision; that question is owned by the marketplace spec and remains open. + +> Rationale: closes evaluation gate §14.3. Managed KB ships two features whose names +> overstate what they provide. AWS's multi-tenant guidance calls metadata filtering +> *"filter-level (logical) isolation, not IAM-enforced (infrastructure) isolation"*, +> and states that ACL-aware retrieval *"is not authorization"* and does not +> authenticate users — its identity is **email only, with no alias resolution, and +> mismatches fail silently**. This platform authenticates via OIDC with claim +> mappings, so a silently-failing email match would be a worse primitive than an +> explicit app-side check. Because this phase holds `App_KB_Id == assistant_id`, the +> per-assistant boundary *is* a per-knowledge-base boundary, which is the strongest +> available isolation by construction. Resource policies are MANAGED-only and attach +> to the AWS knowledge base ARN, so a new identifier silently drops sharing (§11.1). diff --git a/.kiro/specs/managed-kb-migration/tasks.md b/.kiro/specs/managed-kb-migration/tasks.md new file mode 100644 index 000000000..ab1587a18 --- /dev/null +++ b/.kiro/specs/managed-kb-migration/tasks.md @@ -0,0 +1,749 @@ +# Implementation Plan: Managed Knowledge Base Migration + +## Overview + +Introduce Amazon Bedrock Managed Knowledge Base as a second retrieval backend +behind a single abstraction seam, then migrate knowledge bases to it one at a time, +opt-in, with rollback available throughout. + +Task order enforces the deployment rule that **backend code never deploys before +the IAM and resources it requires**. Groups 1–2 are platform-only and change no +behaviour. Groups 3–11 land backend code that stays dark behind flags. Groups +12–13 enable the pilot and opt-in migration. Groups 14–15 add the user-facing +surfaces and the pre-promotion verification gate. + +**Scope:** §14.7 phases 1–4 only. Managed-by-default, stopping legacy writes, +reclaiming legacy vectors, and removing the old pipeline are a follow-up spec. +All three flags — managed-default, migration, and reconciler arming — ship **off**. + +## Tasks + +- [x] 1. Platform: additive schema and IAM (no behaviour change) + - [x] 1.1 Add the sparse work-discovery GSI to the assistants table + - In `infrastructure/lib/constructs/rag/rag-data-construct.ts`, add GSI + `KbWorkIndex` with partition key `GSI7_PK` and sort key `GSI7_SK`, both + STRING, `projectionType: ALL` + - **GSI7, not GSI1** — the table already has six indexes using `GSI_PK`/`GSI_SK` + for the first and `GSI2_PK` through `GSI6_PK` thereafter + - Follow the sparse pattern and comment style of the adjacent `DueSyncIndex` + (GSI4), `AgentDirectoryIndex` (GSI5) and `AgentReportsIndex` (GSI6): keys are + written only while the record is eligible, so ineligible and pinned knowledge + bases are invisible to the dispatcher's query by physics rather than by filter + - Add `GSI7_PK` / `GSI7_SK` to the generic assistant-update path's immutable + attribute list, mirroring `GSI5_*`, so a routine edit cannot resurrect a work + key on a knowledge base that has left the queue + - ⚠️ **This consumes the entire `rag-assistants` GSI budget for whichever + release ships it.** DynamoDB's `UpdateTable` permits exactly ONE GSI creation + or deletion per call, and CloudFormation issues one `UpdateTable` per changed + table, so a release that adds a second index to this table fails the deploy + and rolls the whole stack back. This is not theoretical: it took production + down on 2026-08-01 in release 1.12.0, when `AgentDirectoryIndex` and + `AgentReportsIndex` arrived in separate `develop` merges and collapsed into a + single prod update. If any other in-flight spec adds a GSI to + `rag-assistants`, the two must ship in different releases. + - Regenerate the committed inventory after adding the index: + `cd infrastructure && UPDATE_GSI_INVENTORY=1 npx jest gsi-update-limit`, and + confirm the diff is exactly one line. `infrastructure/test/gsi-update-limit.test.ts` + fails until this is done, and `scripts/release/check-gsi-update-limit.mjs` + re-checks it against `origin/main` on PRs into `main`. + - _Requirements: 15.14, 15.13_ + + - [x] 1.2 Create the Bedrock knowledge base service role + - New construct `infrastructure/lib/constructs/managed-kb/managed-kb-role-construct.ts` + - Trust policy: `bedrock.amazonaws.com` with `aws:SourceAccount` equal to the + account and `ArnLike` on `AWS:SourceArn` scoped to `knowledge-base/*` + - Grant S3 read on the documents bucket conditioned on `aws:ResourceAccount` + - Grant `bedrock:InvokeModel` on `amazon.titan-embed-text-v2:0` only (required + because Requirement 8.5 pins `embeddingModelType: CUSTOM`) + - Grant `cloudwatch:PutMetricData` scoped to the non-reserved + `${prefix}/ManagedKb` namespace. NOT `AWS/Bedrock/KnowledgeBases`: CloudWatch + reserves every namespace beginning with `AWS` and rejects writes to them, so + an `AWS/...`-scoped grant authorizes nothing while looking correct. Bedrock's + own `AWS/Bedrock/KnowledgeBases` metrics are a read source (Req 20.13), not a + publish target + - Publish the role ARN to SSM at `/${prefix}/managed-kb/service-role-arn` + - _Requirements: 20.1, 20.2, 20.4, 20.5, 20.10, 8.5, 8.9_ + + - [x] 1.3 Grant caller permissions for provisioning, ingestion, and retrieval + - Separate policy statements with distinct SIDs for: provisioner/migrator CRUD + (`bedrock:CreateKnowledgeBase`, `CreateDataSource`, `DeleteKnowledgeBase`, + `DeleteDataSource`, `ListKnowledgeBases`, `GetKnowledgeBase`), direct + ingestion (`IngestKnowledgeBaseDocuments`, `DeleteKnowledgeBaseDocuments`, + `GetKnowledgeBaseDocuments`), and inference (`bedrock:Retrieve`) + - Add `iam:PassRole` on the service role conditioned on `iam:PassedToService` + equal to `bedrock.amazonaws.com` + - Attach retrieval to the AgentCore Runtime role and the App API task role; + attach CRUD only to the migration Lambdas' roles + - _Requirements: 20.3, 20.6_ + + - [x] 1.4 Write CDK assertions for the IAM conditions + - New `infrastructure/test/managed-kb.test.ts`, following + `infrastructure/test/kb-sync.test.ts` + - Assert the `aws:SourceAccount` and `ArnLike` `AWS:SourceArn` conditions, the + `iam:PassedToService` condition, the `aws:ResourceAccount` S3 condition, and + the presence of the `PutMetricData` grant on the calling identities, and its + **absence** on the service role + - Assert the S3 statement's **Resource** as well as its Condition: + `aws:ResourceAccount` scopes the account, not the bucket, so without a + Resource assertion the grant can widen to every bucket in the account (file + uploads, fine-tuning, artifacts, SPA) with all tests still green + - Assert the `PutMetricData` namespace does not begin with `AWS`, so nobody + reverts it to the reserved `AWS/Bedrock/KnowledgeBases` namespace that + authorizes no publish + - The `PutMetricData` assertion matters because metric publishing is + best-effort: omit the grant and metrics silently vanish while requests keep + succeeding + - _Requirements: 20.9, 24.10, 24.13_ + +- [x] 2. Platform: worker resources and config + - [x] 2.1 Add the migration construct with dispatcher, worker, and reconciler + - New `infrastructure/lib/constructs/managed-kb/kb-migration-construct.ts`, + following `infrastructure/lib/constructs/kb-sync/kb-sync-construct.ts` + - Three DockerImage Lambdas sharing ONE image + (`backend/Dockerfile.kb-migration`) + - Byte-stable bootstrap stub at + `infrastructure/bootstrap-assets/kb-migration/`, per the + platform-as-bootstrap pattern + - Publish generated function names to SSM under `/${prefix}/kb-migration/` + - EventBridge `rate()` schedule into the dispatcher and into the reconciler + - Wire the construct in `infrastructure/lib/platform-stack.ts` + - _Requirements: 14.1, 15.13, 15.14_ + + - [x] 2.2 Add the ingestion consumer Lambda + - Same construct; triggered by the documents bucket `ObjectCreated` + notification, wired in `platform-stack.ts` alongside the existing + notification to avoid a circular dependency + - Timeout ≥300 s (a 50 KiB PDF was measured at 264 s) and a dead-letter queue + - _Requirements: 10.1, 10.9_ + + - [x] 2.3 Add configuration properties and flags + - In `infrastructure/lib/config.ts`, add a `managedKb` section carrying + `newDefault`, `migrationEnabled`, `reconcilerArmed`, per-owner byte cap + defaults by role tier, and the retention window in days + - Follow the 7-step config pattern: `config.ts` interface → `loadConfig` → + construct → `scripts/common/load-env.sh` → `synth.sh` and `deploy.sh` + (identical context flags) → workflow job-level `env:` → GitHub variable + - All three booleans default to **false**, and an empty string resolves to + false + - _Requirements: 19.1, 19.2, 19.3, 19.4, 19.5, 19.8, 12.2, 14.7, 15.11_ + + - [x] 2.4 Add tagging for reconciliation and teardown + - Tag every runtime-created knowledge base with `prefix`, `env`, `appKbId`, and + an opaque `ownerUserId` + - The owner tag must be an opaque identifier, never an email address or other + PII + - This is a hard prerequisite, not housekeeping: the Reconciler's tag-filtered + `ListKnowledgeBases` and the teardown script both read these tags + - _Requirements: 20.11, 20.12_ + + - [x] 2.5 Add account-level alarms + - New alarms in the managed-kb construct on total managed storage, managed + knowledge base count against 80% of the 10,000 quota, daily + Knowledge-Base `usagetype` cost, and sustained non-zero `KbOrphansFound` + - Use `TreatMissingData.NOT_BREACHING`, matching the posture of the existing + kb-sync, scheduled-runs and prompt-cache observability constructs + - Per-owner caps bound one user; these bound the fleet, and the gap between + ~$169/month expected and ~$15,000/month permitted is why they are required + - _Requirements: 12.13_ + +- [x] 3. KB_Record data layer + - [x] 3.1 Define the KB_Record model + - New `backend/src/apis/shared/kb_backend/records.py` + - Keys `PK=AST#{assistant_id}`, `SK=KB#{app_kb_id}`, with + `app_kb_id == assistant_id` in this phase + - Fields per the design's data-model table, including `retrievalEngine`, + `provisioningState`, `awsKbId`, `awsDataSourceId`, immutable embedding + config, `storedBytes`, `reservedBytes`, `lastRetrievedAt`, migration state + with generation and lease, and lifecycle exemption flags + - _Requirements: 6.1, 6.2, 6.5_ + + - [x] 3.2 Implement conditional state transitions + - `create_provisioning`, `attach_aws_ids`, `promote_engine`, + `rollback_engine`, `set_migration_state`, `acquire_lease` + - Every transition uses a DynamoDB condition expression; `promote_engine` is + conditional on converged catch-up so two workers cannot both promote + - Sparse GSI attributes are written on entering an eligible state and + **removed** on reaching a terminal state + - _Requirements: 15.8, 15.10, 15.13, 17.1, 17.5_ + + - [x] 3.3 Write property test for engine resolution by absence + - **Property 1: absence means legacy** + - Using `hypothesis`, for any KB_Record shape with no `retrievalEngine` + attribute, verify resolution returns the legacy backend, and verify no code + path writes the literal `"s3vectors"` to a record that did not already carry + it + - **Validates: Requirements 1.6, 1.7, 6.6** + - File: `backend/tests/property/test_pbt_kb_engine_resolution.py` + + - [x] 3.4 Write unit tests for conditional transitions + - Concurrent `create_provisioning` yields exactly one winner + - Concurrent `promote_engine` yields exactly one winner + - Terminal transitions remove the GSI attributes + - File: `backend/tests/shared/test_kb_records.py` + - _Requirements: 7.4, 15.10, 15.13_ + +- [x] 4. Backend abstraction seam + - [x] 4.1 Define the protocol and canonical chunk shape + - New `backend/src/apis/shared/kb_backend/protocol.py` + - `KnowledgeBaseBackend` Protocol with `search`, `ingest`, `delete_document` + - Frozen `Chunk` dataclass whose score field is named `relevance` and is + documented as higher-is-more-relevant + - _Requirements: 1.1, 2.1_ + + - [x] 4.2 Implement the backend resolver + - New `backend/src/apis/shared/kb_backend/resolver.py` + - Reads `retrievalEngine` from the KB_Record; absence resolves to + `S3VectorsBackend` + - _Requirements: 1.4, 1.6_ + + - [x] 4.3 Extract the legacy backend verbatim + - New `backend/src/apis/shared/kb_backend/s3vectors_backend.py` + - Move the existing S3 Vectors search path from + `apis/shared/assistants/rag_service.py` and + `apis/shared/embeddings/bedrock_embeddings.py` without functional change + - Convert S3 Vectors cosine **distance** to **relevance** inside this adapter + - _Requirements: 1.2, 1.3, 2.2_ + + - [x] 4.4 Convert the entry point into a facade + - In `apis/shared/assistants/rag_service.py`, reduce + `search_assistant_knowledgebase_with_formatting(assistant_id, query, top_k=5)` + to resolve-then-delegate, preserving its public signature + - Keep emitting a `distance` key in the formatted result, derived from + `relevance`, so no existing consumer breaks on the field rename + - Neither of the two call sites + (`inference_api/chat/routes.py`, `app_api/assistants/routes.py`) changes + - _Requirements: 1.5, 3.4_ + + - [x] 4.5 Write property test for score direction equivalence + - **Property 2: ranking is backend-independent** + - Using `hypothesis`, for any list of chunks with distinct scores, verify both + backends return the known-best chunk first after adapter conversion + - This is the only test that can catch a silent ranking inversion; without it + the failure mode produces no error, just worse answers + - **Validates: Requirements 2.1, 2.2, 2.3, 2.4, 24.1** + - File: `backend/tests/property/test_pbt_kb_score_direction.py` + + - [x] 4.6 Apply the document-status filter above the seam, on both backends + - Move the `status == "complete"` post-filter into the facade so there is one + implementation covering both backends + - It works on the managed path only because `customDocumentIdentifier` is the + platform `document_id` (task 8.4); the filter needs a `document_id` per chunk + to join on + - Keep it on the managed path even though managed ingestion makes it largely + redundant — removing it in the same change that swaps the engine would + confound the comparison + - Apply the 2,000-character context cap in the same place, for the same reason + - _Requirements: 3.2, 3.3_ + + - [x] 4.7 Write test for parity properties on the managed path + - Assert `top_k=5`, the 2,000-character cap, the status filter, and the + 500-character citation clip all hold on the managed backend, not just legacy + - _Requirements: 3.1, 3.2, 3.3, 3.4_ + + - [x] 4.8 Write architecture test for the Lambda import constraint + - Assert `apis.shared.kb_backend` does not transitively import + `apis.shared.assistants`, whose `__init__` drags in the embeddings stack + - Add alongside the existing boundary tests in `backend/tests/architecture/` + - Keep `kb_backend/__init__.py` empty and heavy imports function-local, matching + the convention in `kb_sync/records.py` + - _Requirements: 24.15_ + +- [x] 5. Query clamp + - [x] 5.1 Implement the query guard + - New `backend/src/apis/shared/kb_backend/query_guard.py` with + `MAX_QUERY_CHARS = 10_000` + - Applied in the facade before backend dispatch so both backends are protected + identically; never raises + - Emit a `KbQueryClamped` metric on truncation + - _Requirements: 4.1, 4.2, 4.3, 4.4, 22.3_ + + - [x] 5.2 Remove the stale no-validation assertion + - In `apis/shared/embeddings/bedrock_embeddings.py`, delete the inline comment + stating the query is a "short string, no token validation needed" + - It is true only because Titan v2 tolerates ~32,000 characters; Managed KB + caps `Retrieve` input at 10,000 and the limit is not adjustable + - _Requirements: 4.5_ + + - [x] 5.3 Write property test for the clamp + - **Property 3: clamp is total and non-throwing** + - Using `hypothesis`, for any input string of any length, verify the output is + at most 10,000 characters, the function never raises, and a truncation signal + is emitted exactly when the input exceeded the cap + - **Validates: Requirements 4.1, 4.3, 4.4** + - File: `backend/tests/property/test_pbt_kb_query_clamp.py` + +- [x] 6. Fail-closed document status filter + - [x] 6.1 Make the status filter fail closed + - In `apis/shared/assistants/rag_service.py`, change + `_filter_vectors_by_document_status` so both fallback paths drop chunks + instead of returning them unfiltered: + the missing-table-name branch (currently `valid_doc_ids = doc_ids`) and the + outer exception handler (currently `valid_doc_ids = doc_ids # Graceful + degradation`) + - Leave the per-document handler as-is; it already fails closed + - Emit `KbStatusFilterFailClosed` at error level, distinct from an ordinary + empty-result log line + - _Requirements: 5.1, 5.2, 5.3, 5.4, 22.4_ + + - [x] 6.2 Record the supersession in the prior spec + - In `.kiro/specs/reliable-document-deletion/requirements.md`, annotate + Requirement 3.4 as superseded by Requirement 5 of this spec + - That requirement specified the fail-open deliberately, so retiring it is a + supersession and must be recorded rather than silently contradicted + - _Requirements: 5.5_ + + - [x] 6.3 Write property test for fail-closed behaviour + - **Property 4: unconfirmable status never leaks** + - Using `hypothesis`, for any set of vectors and any injected table-level + failure or missing table-name condition, verify zero chunks are returned + - **Validates: Requirements 5.1, 5.2, 24.6** + - File: `backend/tests/property/test_pbt_kb_status_fail_closed.py` + + - [x] 6.4 Update existing tests that assert the fail-open contract + - Search `backend/tests/` for tests asserting unfiltered fallback and invert + their expectations, citing this spec's Requirement 5 + - _Requirements: 5.5_ + +- [x] 7. Byte cap accounting + - [x] 7.1 Implement reserve / commit / release + - New `backend/src/apis/shared/kb_backend/byte_cap.py` + - `reserve` is a conditional update failing when + `storedBytes + reservedBytes + n > cap`; `commit` moves reserved to stored; + `release` returns the reservation on failure + - Resolve the per-owner cap by role tier, defaulting **below** the existing + 1 GB user-files precedent + - Determine size from an S3 `HEAD` on the stored object, never from a + client-reported value + - Do not read `RawDataSize` for enforcement; it returned 0 datapoints for a + directly-ingested document and remains unconfirmed + - _Requirements: 12.1, 12.2, 12.3, 12.4, 12.6, 12.7, 12.8_ + + - [x] 7.2 Document the retrieval-quota payer decision + - Record in the design whether the knowledge base owner or the invoking user + consumes retrieval quota, and implement accordingly + - _Requirements: 12.10_ + + - [x] 7.3 Write property test for reservation races + - **Property 5: the cap is never exceeded under concurrency** + - Using `hypothesis`, for any interleaving of N concurrent reserve/commit + operations against a cap, verify the committed total never exceeds the cap + and released reservations are fully returned + - **Validates: Requirements 12.4, 12.5, 12.6, 24.7** + - File: `backend/tests/property/test_pbt_kb_byte_cap.py` + + - [x] 7.4 Enforce the cap on the migration re-ingest path + - The Migration_Worker reserves for the **whole snapshot** before entering + `shadow`, and fails the migration up front rather than part-migrating a corpus + that cannot fit + - Surface the failure as a plain-language reason with the option to request an + elevated tier + - Migration is the largest byte-adding operation in the system and the only one + that runs unattended, so it is both the easiest and the worst place to omit + the check + - Emit `KbByteCapRejected` on rejection + - _Requirements: 12.11, 12.12, 12.14_ + + - [x] 7.5 Write test for migration byte-cap rejection + - A corpus exceeding the owner's remaining allowance fails before `shadow`, and + leaves no partially-ingested managed knowledge base behind + - _Requirements: 12.11, 12.12_ + +- [x] 8. Managed backend: provisioning and retrieval + - [x] 8.1 Implement the provisioning saga + - New `backend/src/apis/shared/kb_backend/provisioning.py` + - Write the KB_Record in `provisioning` **before** calling AWS; attach returned + ids with a conditional update + - `CreateKnowledgeBase` with `type="MANAGED"`, `roleArn`, and + `managedKnowledgeBaseConfiguration` carrying the embedding pin (it has no + required members, but the pin has nowhere else to live — NOT literally `{}`); + omit `storageConfiguration` entirely + - Build the `clientToken` programmatically to satisfy the **33-character + minimum** and persist it so a retry reuses it — a natural + `{id}-{variant}-kb` token is 31 characters and fails client-side validation + - Treat "unable to verify the specified embedding model" as **retryable**; it + was observed as pure IAM eventual consistency against a model confirmed + ACTIVE and invokable + - _Requirements: 7.1, 7.2, 7.3, 7.4, 7.5, 7.6, 7.7, 7.8, 8.1, 8.2_ + + - [x] 8.2 Create the CUSTOM connector data source + - `dataSourceConfiguration.type = "MANAGED_KNOWLEDGE_BASE_CONNECTOR"` with the + real type in + `managedKnowledgeBaseConnectorConfiguration.connectorParameters` + - `embeddingModelType: CUSTOM` pinned to `amazon.titan-embed-text-v2:0`, + `FLOAT32` (upper-case: that is the service-model enum value), 1024 dimensions + - `mediaExtractionConfiguration.imageExtractionConfiguration.imageExtractionStatus + = ENABLED` — opt-in, and silently indexes no chart or image content if left + default + - `dataDeletionPolicy = RETAIN` at creation, the documented remedy for the + `DELETE_UNSUCCESSFUL` state already present in the dev account + - _Requirements: 8.3, 8.4, 8.5, 8.6, 8.7, 8.8_ + + - [x] 8.3 Implement managed retrieval + - New `backend/src/apis/shared/kb_backend/managed_backend.py` + - Use `managedSearchConfiguration` with `numberOfResults=5` and + `rerankingModelType="MANAGED"`; never send `vectorSearchConfiguration`, which + is rejected outright for managed knowledge bases + - Do not attempt to configure hybrid search; it is not toggleable + - Constrain any isolation-critical filter to `equals` or `in` + - Run synchronous boto3 calls off the event loop + - _Requirements: 11.1, 11.2, 11.3, 11.4, 11.5, 3.1, 20.7_ + + - [x] 8.4 Implement direct ingestion and document delete + - `IngestKnowledgeBaseDocuments` batched at **10 documents maximum**, + server-enforced; the user guide's claim of 25 does not apply to managed + knowledge bases + - `customDocumentIdentifier = document_id`, which retires the + `{doc_id}#{chunk_index}` scheme including `delete_vector_tail` and the + chunk-shrinkage stash on this path + - Never call `StartIngestionJob` — 0.1 RPS account-wide and not adjustable + - Bound concurrency against the 10-per-account concurrent document-operation + limit + - _Requirements: 9.1, 9.2, 9.3, 9.4, 9.5, 9.6_ + + - [x] 8.5 Write unit tests with stubbed AWS APIs + - Stub `bedrock-agent` and the agent runtime client; never call live + - Cover create/ingest/delete idempotency, the 10-document batch boundary, + retryable embedding-verification failure, and `clientToken` length ≥33 + - File: `backend/tests/shared/test_managed_kb_backend.py` + - _Requirements: 24.2, 24.11_ + + - [x] 8.6 Write test for crash between AWS create and record update + - Simulate a crash after `CreateKnowledgeBase` returns but before the + conditional update; verify the record remains a discoverable retry anchor and + that a retry does not create a second knowledge base + - _Requirements: 7.8, 24.3_ + +- [x] 9. Ingestion control plane + - [x] 9.1 Implement the ingestion consumer + - New `backend/src/apis/app_api/kb_migration/ingestion_consumer.py` + - Follow `kb_sync/records.py`'s raw-table-access convention: importing + `apis.shared.assistants` drags in the whole embeddings stack, and keeping the + Lambda image small is a deliberate constraint + - Resolve each document's knowledge base and engine, then route legacy + documents to the existing pipeline and managed documents to direct ingestion + - Never index the same document on both backends outside a deliberate migration + or pilot + - Poll until **actually retrievable**, recording `indexedAt` and + `retrievableAt` as two distinct timestamps + - Update `DOC#` to a terminal state with bounded retries and a durable retry + anchor + - No in-process `asyncio.ensure_future` orchestration + - _Requirements: 10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8_ + + - [x] 9.2 Write unit tests for routing exclusivity + - Legacy document routes to the old pipeline only; managed document routes to + direct ingestion only; neither is double-indexed + - File: `backend/tests/lambdas/test_kb_ingestion_consumer.py` + - _Requirements: 10.3, 10.4, 10.5_ + +- [x] 10. Deletion sagas and reconciler + - [x] 10.1 Implement tombstones + - New `backend/src/apis/shared/kb_backend/tombstones.py` + - Write `KBTOMB#{app_kb_id}` (and the `#DOC#{document_id}` variant) **before** + calling AWS; clear only after AWS confirms absence + - No TTL on tombstones — TTL removal would recreate the silent-leak class this + design exists to close + - Verify knowledge base deletion by polling `ListKnowledgeBases` until the name + disappears, tolerating ≥6 minutes; deletion took 2–6 minutes when measured + - Never delete the service role until all of its knowledge bases are confirmed + absent, and never delete the last KB_Record before AWS confirms + - Surface `DELETE_UNSUCCESSFUL` as an actionable operator state + - _Requirements: 13.1, 13.2, 13.3, 13.4, 13.5, 13.6, 13.7, 13.8_ + + - [x] 10.2 Implement the daily reconciler + - New `backend/src/apis/app_api/kb_migration/reconciler.py` + - Join paginated, tag-filtered `ListKnowledgeBases` against KB_Records + - AWS-only ⇒ orphan, deleted only if the **AWS-reported `createdAt`** is >24 h + old; age-gating on discovery time would make a reconciler that was down for a + week delete every in-flight create + - Record-only ⇒ mark `vectorState: missing` and **never** delete the record; + the documents are still valid + - Both ⇒ refresh `storedBytes` + - Ship in **report-only** mode, which logs intended deletions and deletes + nothing; arming is a separate flag that treats an empty string as off + - Apply a bounded per-run action limit + - _Requirements: 14.1, 14.2, 14.3, 14.4, 14.5, 14.6, 14.7, 14.8, 19.7_ + + - [x] 10.3 Write reconciliation tests + - Record-only and AWS-only outcomes; age-gate honours AWS `createdAt` rather + than discovery time; report-only performs no deletes + - File: `backend/tests/lambdas/test_kb_reconciler.py` + - _Requirements: 24.4_ + +- [x] 11. Authorization, isolation, and publication + - [x] 11.1 Resolve access before retrieval + - In the facade, resolve the invoking user's access to the knowledge base + **before** attempting retrieval, reusing the existing assistant permission + model rather than introducing a parallel one + - Because this phase holds `App_KB_Id == assistant_id`, "can this user invoke + this agent" already answers "may this turn retrieve"; do not build for the + 0..N case, which is F4's problem + - _Requirements: 25.1, 25.2, 25.3_ + + - [x] 11.2 Keep filters out of the tenant boundary + - Do not use a metadata filter as the isolation mechanism; the per-knowledge-base + boundary is the tenant boundary in this phase + - Do not adopt ACL-aware retrieval: its identity is email-only with no alias + resolution and mismatches fail silently, which is a worse primitive than an + explicit app-side check on an OIDC claim-mapped platform + - _Requirements: 25.4, 25.5, 11.5_ + + - [x] 11.3 Apply resource policies for shared knowledge bases + - Where a knowledge base is shared beyond its owner, attach a resource policy + granting IAM-enforced `bedrock:Retrieve` + - Re-apply the policy whenever a new `awsKbId` is produced; policies attach to + the AWS ARN, so a replacement silently drops sharing + - _Requirements: 25.6, 25.7_ + + - [x] 11.4 Preserve published-agent semantics + - An engine migration must not change what a published agent retrieves; parity + is the contract, so a swap is not a corpus change and needs no re-review + - Exempt listed agents' knowledge bases from lifecycle reclaim; `taken_down` + requires an explicit transition rather than falling through + - Do not attempt to resolve corpus-revision pinning; it belongs to the + marketplace spec + - _Requirements: 25.8, 25.9, 25.10, 25.11_ + + - [x] 11.5 Write authorization tests + - Viewer can read through the agent but never sees the upgrade control; a user + with no access never reaches retrieval; access checks fail closed + - Published-agent corpus behaviour and reclaim exemption + - Resource policy is re-applied after a new `awsKbId` + - _Requirements: 24.6, 24.12, 24.14_ + + - [x] 11.6 Write test asserting the 1:1 binding freeze + - Assert an explicit `knowledge_base` binding is still rejected by + `binding_validation.py` and that `bindable_catalog.py` still returns an empty + list for it, so the freeze is enforced by test rather than by intention + - _Requirements: 6.7, 6.8_ + +- [x] 12. Dual-read pilot + - [x] 12.1 Implement opt-in dual read + - In the facade, when a knowledge base is flagged for the pilot, run both + backends for the same query and **serve legacy** + - Record per read: overlap in returned `document_id` values, a rank + correlation, and per-backend latency + - Default off; must not increase user-visible latency beyond the legacy path's + own latency + - _Requirements: 18.1, 18.2, 18.3, 18.4, 18.5_ + + - [x] 12.2 Write dual-read tests + - Legacy results are always the ones served; comparison metrics are emitted; + a managed-side failure does not fail the turn + - File: `backend/tests/shared/test_kb_dual_read.py` + - _Requirements: 18.2, 18.5_ + +- [x] 13. Migration dispatcher and worker + - [x] 13.1 Implement the dispatcher + - New `backend/src/apis/app_api/kb_migration/dispatcher.py`, following + `kb_sync/dispatcher.py` + - Query the sparse `KbWorkIndex`, apply a bounded per-tick dispatch limit + (mirroring `KB_SYNC_DISPATCH_LIMIT`, default 20), and no-op entirely when the + migration flag is off + - _Requirements: 19.6, 15.14_ + + - [x] 13.2 Implement the migration worker state machine + - New `backend/src/apis/app_api/kb_migration/worker.py` + - `shadow`: provision, then re-ingest every `complete` document from its + existing S3 key at `assistants/{assistant_id}/documents/{document_id}/{filename}` + — never ask the user to re-supply anything + - `verify`: compare an exact source manifest of `document_id` + content hash or + generation, **not** document-count parity, then run a canary retrieval + - `promote`: single conditional write of `retrievalEngine="managed"`, only after + a converged catch-up pass **and** only once the Byte_Cap is enforced on this + knowledge base — no traffic is promoted to an unmetered corpus + - `retain`: set `retainUntil` at least 30 days out + - Take a lease so one knowledge base is never migrated by two workers + - _Requirements: 15.1, 15.2, 15.3, 15.4, 15.5, 15.6, 15.7, 15.8, 15.9, 15.11, 15.13, 12.9_ + + - [x] 13.3 Implement catch-up convergence + - Snapshot the doc-id set, migrate, then run catch-up passes until a pass finds + nothing new — the same converge-on-quiet shape as the crawler's + consecutive-miss rule + - Re-read each document's `DOC#` record immediately before ingesting and skip it + if gone or no longer `complete`, so a document deleted mid-migration cannot + resurrect + - Do not implement dual-write; one write path stays authoritative until + promotion + - _Requirements: 16.1, 16.2, 16.3, 16.4, 16.5, 16.6_ + + - [x] 13.4 Implement rollback + - Write `retrievalEngine` back to its prior value and stamp `rolledBackAt`; + move no data + - Available for the entire `retain` window; a pre-promotion failure leaves the + knowledge base on legacy and fully usable + - _Requirements: 17.1, 17.2, 17.3, 17.4, 17.5_ + + - [x] 13.5 Write property test for migration idempotency + - **Property 6: interrupted migration converges without duplication** + - Using `hypothesis`, for any interruption point in the state machine, verify a + resumed run reaches the same terminal state, creates exactly one knowledge + base, and ingests each document at most once + - **Validates: Requirements 15.9, 15.10, 15.13, 7.4** + - File: `backend/tests/property/test_pbt_kb_migration_convergence.py` + + - [x] 13.6 Write tests for interference during migration + - Upload during migration is picked up by catch-up; delete during migration + never resurrects; concurrent promotion attempts yield one winner + - File: `backend/tests/lambdas/test_kb_migration_worker.py` + - _Requirements: 16.2, 16.4, 16.5, 24.5_ + + - [x] 13.7 Write test for resource-policy re-application after rehydration + - A rehydration producing a new `awsKbId` re-applies the resource policy; + policies attach to the AWS ARN, so a new id otherwise silently drops sharing + - _Requirements: 24.12_ + + - [x] 13.8 Write test for mixed old/new deployment + - Old and new code serving simultaneously; a record with no `retrievalEngine` + resolves to legacy under both + - _Requirements: 1.6, 24.8_ + +- [ ] 14. Surfaces, observability, and teardown + - [x] 14.0 Register the managed backend in the resolver + - **Spec gap, found during implementation.** `register_backend` was defined in + task 4.2 and called by nothing; task 8.3's note that it would register the + managed backend was never carried out, and no other task picked it up. Every + group could therefore have been completed with the feature unreachable: a + promoted record raises `BackendUnavailable`, which is a correct fail-safe and + a useless signal — visible only to the single migrated user. + - Registered at **import** rather than by a startup call, so there is no + sequence to remember and no service that can come up half-configured. Free + because both adapters' module bodies are stdlib-only and their clients are + lazy, which `test_kb_backend_boundary.py` now asserts for `managed_backend` + too. + - Does **not** make the feature live: nothing resolves to managed until a + record says so, and only a promotion writes that. + - _Requirements: 1.4, 2.5_ + + - [x] 14.1 Emit EMF metrics + - Alongside the existing PromptCache metrics: `KbCount`, `KbStorageGB`, + `KbIdleGB`, `KbOrphansFound`, `KbQueryClamped`, + `KbStatusFilterFailClosed`, and + `KbMigration{Started,Promoted,Failed,RolledBack}` + - Compute idleness as `max(own lastRetrievedAt, max(lastUsedAt) over bound + agents)`, never retrieval alone, or an actively used agent's knowledge base is + evicted because its queries did not match + - Write `lastRetrievedAt` through a throttled conditional write (one winner per + 24 h), never per retrieval; prefer per-knowledge-base `Invocations` from + `AWS/Bedrock/KnowledgeBases` where sufficient — that is a *read* of Bedrock's + own namespace and needs `cloudwatch:GetMetricData` / + `GetMetricStatistics` (Req 20.13). Our own metrics above publish to + `${prefix}/ManagedKb` (Req 20.10), never into an `AWS/...` namespace + - _Requirements: 22.1, 22.2, 22.5, 22.6, 22.8, 20.13_ + + - [x] 14.2 Document cost attribution + - Record that Managed KB bills under `AmazonBedrockAgentCore` and that queries + must filter on `usagetype` — keying on `AmazonBedrock` misses it entirely, and + keying on service code alone blends it into the AgentCore Runtime memory line + - _Requirements: 22.7_ + + - [x] 14.3 Build the upgrade UX + - In `frontend/ai.client/src/app/knowledge-base/knowledge-base-section.component.ts`, + add the opt-in upgrade card, non-blocking progress, one-time success notice, + and a failure state with a retry control + - Show nothing at all for a legacy knowledge base needing no action + - Never use the word "vector" in user-facing copy + - Gate on the existing `_require_edit_permission`; viewers never see the control + - Angular signals, `OnPush`, Tailwind utilities, both light and dark modes, + WCAG AA + - **Spec gap, found during implementation.** This task was written as + frontend-only, but **nothing enrolled a knowledge base**. The worker picks + up records already in `shadow`; the dispatcher sweeps GSI7; no code path + wrote either. Group 14 could have been called complete with the feature + still unreachable. Required a new HTTP surface: + `backend/src/apis/app_api/kb_upgrade/` (`models.py`, `service.py`, + `routes.py`) — `GET`/`POST` `…/knowledge-base/upgrade`, `POST …/retry`, + `POST …/notice`. + - Kept in its **own package**, not `app_api/kb_migration/`: that package's + modules share one size-constrained Lambda image, and this one imports + `apis.shared.assistants` for the permission model, which pulls the + embeddings stack at module scope. + - Enrolment is **two conditional writes**, not one put. `KbRecord.to_item` + does not write `GSI7_PK`/`GSI7_SK` — only `set_migration_state` maintains + them — so the obvious one-put enrolment produces a record that claims to be + migrating and is invisible to the dispatcher *forever*, behind a spinner + that never moves. Asserted by + `test_enrolment_writes_the_dispatcher_work_keys` and + `test_the_created_record_does_not_claim_to_be_migrating`. + - The **offer is gated on `MANAGED_KB_MIGRATION_ENABLED`**, the dispatcher's + own flag, read at call time with the same allow-list. Offering an upgrade + the worker cannot perform is a spinner with no engine behind it, so + "available" is made to mean actionable. Off ⇒ phase `none` ⇒ renders + nothing, which is also 23.1's required behaviour. + - Two public transitions added to `kb_backend/records.py`: + `retry_from_failed` (one atomic write — generation bump, re-enter `shadow`, + work keys, `REMOVE migrationError`; guarded on the old generation **and** + still being `failed`) and `dismiss_upgrade_notice`. + - Client (`kb-upgrade.service.ts`) **fails soft**: `getStatus` resolves to + phase `none` rather than rejecting, so a broken upgrade endpoint cannot take + down the documents section it decorates. + - _Requirements: 23.1, 23.2, 23.3, 23.4, 23.5, 23.6, 23.7, 23.8_ + + - [ ] 14.4 Surface failed and stuck documents — **surfacing done, one-click + retry deferred** (still open: see the deferral note below) + - During the upgrade flow, list any non-`complete` document that will not be + carried across and offer retry; 200 of 1,692 production `DOC#` records + (11.8%) are affected, including 95 `failed` whose owners believe the uploads + worked + - Distinguish an unsupported file format from a processing failure in messaging + - **Done:** `classify_document` splits `unsupported_format` / + `processing_failure` / `being_removed` / `still_processing` (Req 21.4), and + the card discloses them collapsed above the offer — *before* the user + commits, so the choice to fix or accept the loss is theirs (Reqs 21.1, 21.3). + - The unsupported-format set is **imported** from + `docling_processor.DOCLING_SUPPORTED_EXTENSIONS`, never copied. A copied + list is the tag-contract defect's exact shape. Its module scope is + stdlib-only, so the import is free. + - `deleting` documents are **deliberately surfaced**, though + `list_assistant_documents` filters them out as soft-deleted: they are 101 of + the 200 affected records, and a user never shown them cannot tell they are + stuck. That filter — plus its stale-document auto-fail *write* — is why this + surface runs its own raw `DOC#` query. + - **Deferred, Req 21.2 (one-click retry).** Ingestion is S3-event-triggered + (`documents/ingestion/handler.py`) and no reprocess endpoint exists, so a + retry control needs new backend that re-fires that pipeline for bytes + already in S3. Not built: it is a change to a live ingestion path and was + explicitly deferred rather than improvised. The card currently directs the + user to re-upload via "Add files", which is a retry path that works today + and needs nothing new. **Close this subtask by either building the + reprocess endpoint or amending Req 21.2 to accept re-upload.** + - _Requirements: 21.1, 21.3, 21.4 (21.2 partial — see above)_ + + - [ ] 14.5 Build the admin surface + - Knowledge bases filterable by engine, with stored bytes and document counts, + bulk migrate, and per-knowledge-base retry + - _Requirements: 23.9_ + + - [x] 14.6 Extend teardown for runtime-created resources + - In `scripts/teardown/destroy.sh`, list and delete only knowledge bases tagged + for the project and environment, **before** deleting their service role and + the platform stack + - Poll until each resource is confirmed absent; "delete call accepted" is not + "resource gone" + - _Requirements: 20.8, 13.4, 13.5_ + + - [x] 14.7 Write teardown test + - Only tagged resources are deleted, and the service role is deleted only after + all its knowledge bases are confirmed absent + - _Requirements: 24.9_ + +- [ ] 15. Pre-promotion verification + - [ ] 15.1 Run the packaged-SDK contract probe + - Using the checked-in environment with **no** `AWS_DATA_PATH` override, run a + create → ingest → retrieve smoke probe against dev-ai + - This is the contract test that the pinned `boto3==1.43.68` and its packaged + service model are genuinely sufficient, rather than the side-loaded model the + evaluation used + - _Requirements: 8.1, 9.1, 11.1_ + + - [ ] 15.2 Probe for an account-level ingestion-concurrency limit + - During the pilot, run a many-knowledge-base backfill to determine whether an + account-level ingestion-concurrency limit exists; the quota page lists none + - Do not size a wide fleet migration before this is answered + - _Requirements: 9.5_ + + - [ ] 15.3 Confirm the full test matrix passes + - Run the backend suite, the infrastructure suite, and `mypy`/`ruff` inside the + dev container + - Verify every Requirement 24 item has a corresponding passing test + - _Requirements: 24.1, 24.2, 24.3, 24.4, 24.5, 24.6, 24.7, 24.8, 24.9, 24.10, 24.11, 24.12, 24.13, 24.14, 24.15_ diff --git a/.kiro/specs/reliable-document-deletion/requirements.md b/.kiro/specs/reliable-document-deletion/requirements.md index a8b31dff4..4207026a1 100644 --- a/.kiro/specs/reliable-document-deletion/requirements.md +++ b/.kiro/specs/reliable-document-deletion/requirements.md @@ -53,7 +53,17 @@ This document specifies the requirements for reliable document deletion in the R 1. WHEN the RAG_Search_Service receives vector search results, THE RAG_Search_Service SHALL extract unique document_id values from the result metadata and look up their status in the Assistants_Table. 2. THE RAG_Search_Service SHALL return only chunks from documents where the status equals "complete" in the Assistants_Table. 3. WHEN a document record does not exist in the Assistants_Table for a given document_id, THE RAG_Search_Service SHALL exclude chunks from that document. -4. IF the Assistants_Table lookup fails due to a DynamoDB error, THEN THE RAG_Search_Service SHALL fall back to returning unfiltered vector results. +4. ~~IF the Assistants_Table lookup fails due to a DynamoDB error, THEN THE RAG_Search_Service SHALL fall back to returning unfiltered vector results.~~ + **SUPERSEDED** by Requirement 5 of `.kiro/specs/managed-kb-migration`, which + inverts this to fail **closed**: an unconfirmable status now drops the chunks. + + This was a deliberate choice here, not an oversight, so retiring it is recorded + rather than silently contradicted. What changed is evidence: the fail-open path + was measured in production, and 936 retrievals in a trailing 30-day window had + chunks removed by this filter — so the documents it guards are real, not + hypothetical, and a lookup failure would have served users content they believe + they deleted. The per-document lookup failure in criterion 3 already failed + closed and is unchanged; only the table-level fallback moved. ### Requirement 4: Inline Cleanup with Retries diff --git a/.kiro/steering/observability.md b/.kiro/steering/observability.md new file mode 100644 index 000000000..6205c7819 --- /dev/null +++ b/.kiro/steering/observability.md @@ -0,0 +1,317 @@ +--- +inclusion: fileMatch +fileMatchPattern: ["infrastructure/lib/constructs/observability/*", "infrastructure/test/observability-*"] +--- + +# Observability + +Every CloudWatch alarm in `PlatformStack` publishes to one SNS topic. This doc +covers the rules that keep that true, the gotchas that silently break alarms, and +what to do when one fires. + +## 0. The failure this system was built to prevent + +Before this work the stack had 13 alarms and **not one of them notified anybody**. +Three separate constructs carried a comment saying so. Two of those alarms were +worse than silent: they watched metric names that exist in no CloudWatch +namespace at all, so they sat in `INSUFFICIENT_DATA` from the day they were +created — which an operator reads as *healthy*. + +Both failures share a shape: **nothing errored.** The alarm deployed, evaluated, +and turned green. Every rule below exists because this domain fails quietly, and +quiet failure has to be caught by structure or by a test, never by remembering. + +## 1. Never call `new cloudwatch.Alarm()` — use `AlarmFactory` + +```typescript +const alarms = new AlarmFactory(this, config, props.alarmTopic); + +alarms.alarm('MyAlarmLogicalId', { + name: 'my-service-errors', // NOT alarmName; prefix is applied for you + alarmDescription: 'What broke, and what the first response should be', + metric: someMetric, + threshold: config.observability.someThreshold, + evaluationPeriods: 2, + comparisonOperator: cloudwatch.ComparisonOperator.GREATER_THAN_THRESHOLD, + treatMissingData: cloudwatch.TreatMissingData.NOT_BREACHING, +}); +``` + +The factory attaches both `AlarmActions` and `OKActions` as a consequence of +being used at all. `new cloudwatch.Alarm()` produces a console-only alarm that +looks completely finished, which is why the rule is enforced by a source-level +test rather than convention: + +- `observability-alarm-routing.test.ts` fails if any file under `lib/` calls the + constructor directly (except the factory itself), and fails if any alarm in the + synthesized template lacks actions. + +Use `expressionAlarm()` for metric math so the routing guarantee survives. + +## 2. No `config.production` in observability code + +This repo is forked by many institutions. A fork with one environment should not +have to reason about a `production` boolean, and a fork with three should not be +limited to two. Every tunable is a **single scalar** in `ObservabilityConfig`, +and per-environment differences live in the forker's own deployment config +(GitHub Variables scoped to a GitHub Environment) — not in a ternary here. + +Enforced: `observability-alarm-routing.test.ts` fails on any `config.production` +under `lib/constructs/observability/`. + +Adding a tunable means touching five places. Miss the third and the flag is +accepted then silently ignored: + +1. `OBSERVABILITY_DEFAULT_*` constant in `config.ts`, with the reasoning inline +2. field on `ObservabilityConfig` +3. loader entry using the full precedence chain (see §3) +4. `scripts/common/load-env.sh` → `build_cdk_context_params()` +5. job-level `env:` in `.github/workflows/platform.yml` + +Defaults are **cost-conscious**: they are what a fork inherits when it configures +nothing, so they are the cheapest setting that still leaves alerting useful. +Diagnostic depth is opt-in. The X-Ray sampling default is the clearest case — it +was `1.0` for any fork that never set `production`, meaning a recorded trace for +every single agent invocation at $5/million. + +## 3. The flat dotted context key (this has bitten the repo three times) + +`--context observability.logRetentionDays=90` sets the **flat** key +`context['observability.logRetentionDays']`. It does **not** build a nested +object. Reading only the nested form accepts the operator's flag and ignores it. + +```typescript +logRetentionDays: + parseIntEnv(process.env.CDK_OBSERVABILITY_LOG_RETENTION_DAYS) + ?? parseIntEnv(scope.node.tryGetContext('observability.logRetentionDays')) // ← flat + ?? scope.node.tryGetContext('observability')?.logRetentionDays // ← nested + ?? OBSERVABILITY_DEFAULT_LOG_RETENTION_DAYS, +``` + +Use `parseFloatEnv` for fractional values. `parseIntEnv('0.05')` is `0`, which +would switch X-Ray sampling off entirely rather than setting it to 5%. + +## 4. Gotchas that produce silently-broken alarms + +### The SNS topic must use a customer-managed KMS key + +CloudWatch **cannot** publish to a topic encrypted with the AWS-managed +`alias/aws/sns` key: the publish is made by the `cloudwatch.amazonaws.com` +service principal, and an AWS-managed key's policy cannot be edited to grant it +`kms:GenerateDataKey*`. The alarm goes to ALARM, the console shows it firing, and +the notification is dropped. + +`kms:Decrypt` alone is **not enough** — SNS envelope encryption has the +*publisher* generate the data key, so `GenerateDataKey*` is required too. + +### Dimensions must come from real resources + +A `CPUUtilization` alarm with no dimensions is a valid CloudWatch alarm that +silently averages every ECS service in the account. Same for ALB metrics. Always +derive dimensions from the CDK resource (`service.metricCpuUtilization()`, +`targetGroup.metrics.*`, `table.metric()`), never from a name string. + +### Metric-math alarms cap at 10 metrics + +CloudWatch rejects an alarm whose math expression contains more than 10 +individual metrics. CDK's `table.metricSystemErrorsForOperations()` defaults to +**all 14** DynamoDB operations and throws `TooManyMetricsInMathExpression` at +synth. Pass an explicit operations list. + +Also: CDK deprecates `metricThrottledRequests()` and `metricSystemErrors()` as +returning invalid metrics. Use `table.metric('ReadThrottleEvents')` etc. + +### Don't pass `label` to a metric used in an alarm + +It forces CDK to render the alarm as a `Metrics[]` array instead of flat +`Namespace`/`MetricName`/`ExtendedStatistic` properties. CloudWatch labels +percentile series adequately on its own. + +### Units differ between services + +| Metric | Unit | Threshold handling | +|---|---|---| +| AgentCore `Latency` | **Milliseconds** | use `agentCoreLatencyMs` directly | +| ALB `TargetResponseTime` | **Seconds** | divide by 1000 | + +Both verified with `get-metric-statistics`. Getting this wrong is a 1000x error +in either direction, and neither direction fails loudly. + +## 5. Streaming makes latency a weak signal + +The chat path is SSE. The ALB does not consider a request complete until the +stream closes, so `TargetResponseTime` and AgentCore `Latency` are legitimately +tens of seconds for a healthy turn — and a sudden *drop* can mean turns are +failing early. + +Measured over 14 days in dev: average turn **3.0–4.5s**, daily maxima **16.7s, +16.9s, 24.4s**. The original alarm threshold was 30s, i.e. *below* the observed +maximum, so a healthy long turn could trip it. Latency floors default to 120s. + +Reliable signals on this path are the discrete ones: 5xx counts, unhealthy hosts, +rejected connections, throttles. + +## 6. `treatMissingData` is a per-metric decision + +Never defaulted by the factory, because both answers are correct somewhere: + +- **`NOT_BREACHING`** for error and throttle counts. A service that is not + failing publishes nothing, so absent data is the healthy state. +- **`BREACHING`** for `UnHealthyHostCount` and `RunningTaskCount`. These stop + being published when the service is at zero or deleted. `NOT_BREACHING` would + leave the alarm silent during a total outage — the exact case it exists for. + +## 7. Verify metrics exist before alarming on them + +Documentation is not sufficient evidence that a metric is published, and an alarm +on a non-existent metric is indistinguishable from a healthy one. + +```bash +aws cloudwatch list-metrics --namespace "AWS/Bedrock-AgentCore" \ + | jq -r '.Metrics | group_by(.MetricName) | .[] | + "\(.[0].MetricName) dims=\(map(.Dimensions|map(.Name)|sort)|unique|tostring)"' +``` + +Known facts from that sweep, all pinned by tests: + +- Runtime metrics live in `AWS/Bedrock-AgentCore` (hyphenated), dimensioned + `Resource` + `Operation=InvokeAgentRuntime` + `Name={runtimeName}::DEFAULT`. + The lowercase `bedrock-agentcore` namespace is real but holds only the + OpenTelemetry/Strands *application* metrics. +- **Memory and Gateway publish `Resource` as a full ARN; Code Interpreter + publishes a bare ID.** Passing an ARN for Code Interpreter yields an alarm that + matches nothing. +- `AWS/Cognito` on the **ESSENTIALS** feature plan publishes *only* success + metrics. There is no sign-in failure or throttle metric to alarm on — failure + and threat metrics require the **Plus** plan. The auth-path failure signal is + the token-enrichment Lambda's `Errors` metric instead. +- AgentCore Browser has zero metric streams (provisioned, unused). + +A metric absent from `list-metrics` may simply never have fired. Alarming on it +is still correct — `NOT_BREACHING` keeps it quiet until the first occurrence. + +## 8. Resource budget + +CloudFormation caps a stack at **500 resources**, and this is a deliberate +single-stack architecture with nowhere to spill. Alarms are the largest +discretionary consumer. + +Measured: **308** resources before this work, **~370** after. A guard in +`observability-dynamodb-alarms.test.ts` fails the build above 460, so the ceiling +surfaces while there is still room to react rather than as a failed deploy. + +When budget matters, decide with data rather than by covering every documented +metric. The DynamoDB allocation was cut from 78 alarms to 27 because a live sweep +showed `ReadThrottleEvents`, `WriteThrottleEvents` and `SystemErrors` had **zero +streams** — all tables are on-demand, so throttling had never occurred — while +account-level `UserErrors` had real data nobody was watching. + +Dashboards: the first **3 are free**, then $3/month each. The stack is at exactly +3, which is why the platform dashboard links to the other two instead of +restating their widgets. + +## 9. Log retention + +One value, `observability.logRetentionDays`, applied through +`logRetentionFor(config)`. A source guard fails the build if any construct +hardcodes `retention: logs.RetentionDays.*`. + +The AgentCore Runtime's log group is created by the **AgentCore service**, not +CloudFormation, so a CDK `LogGroup` cannot set its retention — declaring one +would collide on create or manage a second, empty group. An `AwsCustomResource` +calls `logs:PutRetentionPolicy` instead. That API is idempotent *and* creates the +group if absent, which matters on a first deploy when the runtime exists but has +never been invoked. There is deliberately **no `onDelete`**: removing the +retention policy on teardown would revert the group to "keep forever", which is +the cost problem it fixes. + +## 10. Subscriptions are not infrastructure-as-code + +The topic is created by CDK; **subscribers are not**. Several teams need to hear +about failures and their membership changes far more often than the +infrastructure does. Requiring a PR, a review, and a CloudFormation deploy to add +one address is how a notification list goes stale and stops being trusted. + +```bash +TOPIC=$(aws ssm get-parameter --name "/${PREFIX}/observability/alarm-topic-arn" \ + --query Parameter.Value --output text) +aws sns subscribe --topic-arn "$TOPIC" \ + --protocol email --notification-endpoint team@example.edu +``` + +A test asserts zero `AWS::SNS::Subscription` resources exist, so this decision +cannot be quietly reversed. + +## 11. Runbook — first response by alarm + +| Alarm | What it means | First action | +|---|---|---| +| `alb-unhealthy-hosts` | Targets failing health checks, **or none reporting** | `aws ecs describe-services` — are tasks running at all? Then app-api logs. | +| `alb-elb-5xx` | The load balancer itself could not serve | Almost always no healthy target. Check the alarm above first. | +| `alb-target-5xx` | App is reachable and erroring | app-api logs. This is application code. | +| `alb-rejected-connections` | ALB connection limit hit | Users were turned away *before* reaching the app, so nothing is in app logs. Check request volume. | +| `app-api-running-tasks-low` | Fewer tasks than desired | Task failing to start: check stopped-task reason, image pull, subnet IPs. | +| `app-api-memory-high` | Sustained memory pressure | Fargate **kills** a task that exhausts memory. Raise memory or find the leak. | +| `agentcore-system-errors` | AgentCore's fault | Escalate to AWS. Not application code. | +| `agentcore-high-error-rate` | `UserErrors` — our requests are malformed | Recent inference-api deploy? Check payload shape and IAM. | +| `agentcore-throttles` | At the TPS or session quota | Request a quota increase. Will not self-resolve. | +| `agentcore-high-latency` | p99 above 120s | Genuinely hung, not merely slow — 24s is a normal maximum here. | +| `bedrock-tpm-quota-usage` | **Leading** indicator | Request a quota increase *now*, before throttling starts. | +| `bedrock-invocation-throttles` | At a model's TPM/RPM quota | Users see chats that never respond. Quota increase. | +| `agentcore-memory-*` | Memory hot path failing | Users experience an agent that has forgotten the conversation. | +| `agentcore-gateway-*` | MCP calls failing at the gateway | Agents lose tool access. Check gateway targets. | +| `ddb-*-throttle` | Named table throttling | On-demand, so this is a hot partition or an account limit. Compare `ReadThrottleEvents` vs `WriteThrottleEvents` on that table. | +| `ddb-user-errors` | DynamoDB rejecting our requests (4xx) | Application code misusing the API. Account-wide, so use CloudTrail or app logs to find the caller. | +| `lambda-token-enrichment-errors` | **Silent** degradation | Handler is fail-open: logins still work, but MCP tools are losing user-identity claims. | +| `dlq-kb-ingestion-not-empty` | Work accepted then failed every retry | Will **not** self-clear. Inspect, fix, then replay or drain. | +| `prompt-cache-session-partial-miss` | One conversation re-writing its prefix every turn | Use the "Sessions by partial-miss waste" widget to find which session. | + +Start at the **`{prefix}-platform-health`** dashboard: row 1 says whether traffic +is being served, row 2 says why, row 3 shows every alarm's current state. + +## 12. Boise State's own profile + +The committed defaults are what a **fork** should inherit. Boise State authors +this platform and does far more diagnostic work than any deployer of it, so our +values differ — and they live in GitHub Variables scoped to a GitHub Environment, +never in committed code. That separation is what lets both be right at once. + +`scripts/observability/set-bsu-overrides.sh --env ` +applies them. It is **not run by CI**, makes no AWS changes, and requires +confirmation, because it mutates shared repository configuration. Add `--dry-run` +to print the plan (no `gh` auth needed). + +| Field | OSS default | BSU dev | BSU prod | +|---|---|---|---| +| `xraySamplingRate` | `0.01` | `0.5` | `0.1` | +| `xraySamplingReservoir` | `1` | `5` | `2` | +| `xrayInsightsNotifications` | `false` | `true` | `true` | +| `agentCoreApplicationLogsEnabled` | `false` | `true` | **`false`** | +| `logRetentionDays` | `30` | `14` | `90` | +| `agentCoreErrorThreshold` | `10` | `5` | `5` | +| `lambdaErrorThreshold` | `5` | `1` | `3` | +| `albTarget5xxThreshold` | `10` | `5` | `5` | +| `dynamoThrottleThreshold` | `10` | `1` | default | +| `ecsCpuPercent` / `ecsMemoryPercent` | `80` / `85` | default | `75` / `80` | +| `promptCacheAvoidableMissThreshold` | `10` | `5` | default | +| `promptCacheWastedUsdThreshold` | `1` | `0.5` | default | + +Two choices worth understanding rather than copying: + +- **Dev traces at 50%, prod at 10%.** Not a mistake. X-Ray bills per trace + recorded, and prod traffic is orders of magnitude larger, so 10% of prod is far + more traces than 50% of dev. Dev is where we debug; prod is where we pay. +- **`agentCoreApplicationLogsEnabled` stays OFF in production.** Those records + carry every user's prompt and the model's response verbatim — the + highest-volume log source available and a genuine PII surface. Enable it + temporarily for a specific investigation, then turn it back off. + +Retention inverts between the two for the same reason latency floors are high: +dev keeps 14 days because dev noise is not worth a month, prod keeps 90 because a +real incident review reaches back weeks. + +**Verifying a variable took effect.** The synth log prints a line beginning +`Observability:` with the *resolved* values. If it disagrees with the variable you +set, the value is not reaching `--context` — check the deploy job's **job-level** +`env:` block, since `vars.*` in a workflow-level `env:` silently resolves to an +empty string. diff --git a/CHANGELOG.md b/CHANGELOG.md index 6bcde0d6d..cab8e0de4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,206 @@ All notable changes to this project are documented in this file. Format follows For narrative release notes written for operators and product owners, see [RELEASE_NOTES.md](RELEASE_NOTES.md). +## [1.17.0] - 2026-09-02 + +Reliability, security and observability. Every CloudWatch alarm in the stack now notifies somebody — before this the stack had 13 alarms and **none of them were routed**, two of which watched metric names that exist in no namespace and had read as healthy since the day they were created. A production outage post-mortem (session `5f34d2b0`) drives four chat-path changes: Bedrock's transient faults are retried, a retry and a long silence are both visible to the user, and attachments a failed turn never delivered are re-sent. Four security findings are closed, including a High-severity OIDC login CSRF in the BFF auth flow and a privilege-escalating stored XSS in skill resources. The Bedrock Managed Knowledge Base migration — still off by default — gets eleven defects fixed from its first real runs in dev. **Requires a CDK deploy**, and one manual step after it: subscribe your team to the new alarm topic (see [step-05-verify](.github/docs/deploy/step-05-verify.md#6-subscribe-to-platform-alarms-required--not-automated)). + +### 🚀 Added + +- **Single SNS alarm topic** (`{prefix}-alarms`) that every alarm routes to, encrypted with a customer-managed KMS key. The CMK is required, not a preference: CloudWatch cannot publish to a topic encrypted with the AWS-managed `alias/aws/sns` key, and that failure is silent — the alarm fires, the console shows it, the notification is dropped. Topic ARN published to SSM at `/{prefix}/observability/alarm-topic-arn` and as a CfnOutput (#910) +- **`AlarmFactory`** — the only sanctioned way to create an alarm. Attaches `AlarmActions` *and* `OKActions` as a consequence of being used, so an unrouted alarm now requires deliberately bypassing it. A source-level test fails the build if any file under `lib/` calls `new cloudwatch.Alarm()` directly (#910) +- **ALB alarms** (6): ELB 5xx, target 5xx, unhealthy hosts, target connection errors, rejected connections, and a streaming-aware p99 latency floor (#910) +- **ECS service alarms** (3): CPU, memory, and running-task-count below desired (#910) +- **DynamoDB alarms** (27): a combined read+write throttle alarm per table naming that table, plus one account-level `UserErrors` alarm (#910) +- **Lambda alarms** (21): errors and throttles across every runtime function, including `artifact-render`, `rag-ingestion` and the four kb-migration functions which previously had none, plus dead-letter-queue depth on the kb-ingestion DLQ (#910) +- **AI-path alarms** (9): Bedrock invocation throttles, server errors and `EstimatedTPMQuotaUsage` (the only *leading* indicator in the set — visible before throttling starts), AgentCore Memory hot-path errors and throttles, Gateway MCP errors and throttles, Code Interpreter session errors and concurrent-session count (#910) +- **`{prefix}-platform-health` dashboard** — one pane answering "is the platform healthy right now": traffic and errors, then saturation, then every alarm's current state. Links to the two existing dashboards rather than restating them, which keeps the stack at exactly 3 (CloudWatch's free ceiling) (#910) +- **`observability` configuration section** — 18 single-scalar tunables with `CDK_OBSERVABILITY_*` overrides, cost-conscious defaults, and validation that rejects retention values CloudWatch does not accept and X-Ray sampling rates given as percentages (#910) +- **`model_retry` SSE event** — the backend now tells the SPA when it is retrying a failed model call instead of surfacing it. The loading indicator swaps its cycling phrases for a fixed amber notice, cleared on `message_start`/`done` (#905) +- **Stall indicator** — the loading indicator says "Still working…" after 30s of silence and "Still working — this is taking longer than usual." after 90s, driven by a client-side liveness stamp on every stream event. A known retry outranks the stall notice (#907) +- **Attachment recovery** — a turn that dies before the model reads its inline documents now re-sends them on the next turn, via a write-ahead marker on the session row that survives every way a turn can die, including a dropped stream. Bounded to the single following turn, discarded after an hour, and the user's own attachments always win (#905) +- **`BedrockTransientRetryStrategy`** — `ServiceUnavailableException`, `InternalServerException`, `ModelNotReadyException` and `ModelTimeoutException` are now retried under the existing four-attempt backoff. Only when they fire *before* the response stream opens, so a retry can never replay visible output. `RETRY_TRANSIENT_SERVICE_ERRORS=false` restores stock behavior without a deploy (#905) +- `.kiro/steering/observability.md` — the gotchas that silently break alarms, plus a first-response runbook for every alarm (#910) +- `scripts/local-dev/run-kb-migration.py`, `kb-doc-timings.py` and `kb-compare-engines.py` — in-process drivers and read-only diagnostics for the managed-KB state machine (#898, #899) + +### ✨ Improved + +- A Bedrock 503 no longer reads as "I ran into a problem with the AI model". Two classifiers both keyed on `"throttl"` and neither recognized service-unavailable; a shared `is_service_unavailable_error` predicate now says the fault is on the provider's side and that we already retried (#905) + +### ⚠️ Changed + +- **Log retention is one configured value** (`observability.logRetentionDays`, default 30) applied to all 15 log groups through `logRetentionFor(config)`. Previously every construct hardcoded `ONE_WEEK`, except AgentCore Memory which used `ONE_MONTH` — differing silently rather than deliberately. A source guard fails the build on any hardcoded `RetentionDays` (#910) +- **No `config.production` branching in observability code.** This repo is forked by many institutions: a fork with one environment should not have to reason about a `production` boolean, and a fork with three should not be limited to two. Per-environment differences now live in the forker's deployment config as single values. Enforced by test (#910) +- **Skill resource files are served as `attachment`, never `inline`**, with their media type re-derived from the filename at serve time rather than reflected from the stored row — which is what neutralizes rows written before the upload allowlist existed, so no data migration is needed. Both SPA callers fetch these over XHR for an in-app viewer; nothing navigates to the URL (#904) +- **Managed KB no longer pins the embedding model.** `embeddingModelType: CUSTOM` and `rerankingModelType: MANAGED` are mutually exclusive — AWS rejects the combination — and measurement settled it: the pin scored 1.00/0.982/0.952 (flat, unusable for a small context cap) against managed reranking's 0.413/0.199. On managed retrieval Bedrock embeds both sides, so query/index consistency is its invariant, not ours. `embeddingModelId`/`embeddingDimensions` are no longer recorded, since nothing here knows what Bedrock chose (#898) + +### 🐛 Fixed + +- **A completed answer could tell the model it was cut short.** The client's Stop writes `lastTurnInterrupted` immediately, but the server only observes the armed cancel on a lease heartbeat that sleeps 10s *before* its first check — so a turn finishing inside that window completed normally and left the marker behind. The next turn then prepended a note saying the user stopped the previous response and not to resume it, every clause of which was false, and it demonstrably steered the answer. A turn that reaches the end of the success path now clears the marker, whatever the client signalled (#909) +- **A session id could be forked across two users.** `_get_session_by_gsi` returned `None` both for "no such session" and for "exists, but owned by someone else", so opening someone else's `/s/{sessionId}` link created a *second* metadata row under the requester — the `attribute_not_exists(PK)` guard cannot catch it, because the new row has a different PK. Not a confidentiality bug (conversation content lives in AgentCore Memory keyed by actor id, so the second user saw an empty thread), but it duplicated the row, mis-attached spend, and made the original owner's session resolve non-deterministically afterwards. Both GSI lookups now scan for the caller's own row rather than reading `items[0]`, since forked rows already exist (#906) +- **Bedrock's transient service faults were never retried.** Strands' stock retry strategy covers `ModelThrottledException` only, and `BedrockModel` maps exactly one error code to it — every other fault re-raised as a raw botocore `ClientError`, so the configured four-attempt backoff never ran. A prod `ConverseStream` carrying two PDFs failed with `ServiceUnavailableException` after 95.6s, billed 56,440 uncached input tokens, and returned nothing (#905) +- **The two AgentCore Runtime alarms were watching metrics that do not exist.** They used namespace `bedrock-agentcore` with `InvocationCount` / `InvocationErrors` / `InvocationLatency`. Verified against the live account: that namespace is real but holds only the OpenTelemetry/Strands *application* metrics, and those three names exist in no namespace at all. Corrected to `AWS/Bedrock-AgentCore` with the verified `Resource` + `Operation` + `Name` dimension set, and split into four alarms — `SystemErrors` (AWS's fault) separated from `UserErrors` (ours), plus a new throttle alarm (#910) +- **The AgentCore latency alarm would have fired on healthy traffic.** Its 30-second threshold sat *below* the observed maximum: measured over 14 days, average turns run 3.0–4.5s with daily maxima reaching 24.4s, because the chat path is SSE and the runtime does not finish a request until the stream closes. Floors now default to 120s (#910) +- **The `agentcore-observability` dashboard's token-usage widget was always empty** — `InputTokens`/`OutputTokens` do not exist, and the token metrics that do exist in that namespace are Memory-strategy counters dimensioned by `StrategyId`, not model tokens. Removed; the header now points at the prompt-cache dashboard for real token accounting (#910) +- **X-Ray recorded a trace for every single agent invocation** in any deployment that never set `production` — `fixedRate` was `1.0` with a 50/sec reservoir on that branch, at $5 per million traces recorded. Now a single configured value defaulting to 1% with a 1/sec reservoir (#910) +- **The AgentCore Runtime's log group had no retention policy and grew forever.** It is created by the AgentCore service rather than CloudFormation, so a CDK `LogGroup` cannot set it. An `AwsCustomResource` calls `logs:PutRetentionPolicy` instead — idempotent, and it creates the group if the runtime has not yet been invoked (#910) +- **Managed KB: the legacy pipeline never stood down for a promoted knowledge base.** Routing exclusivity was enforced on one side only, so every document added to a promoted KB was indexed twice — and worse, two writers owned one `status` field and the last one won by luck. Observed both ways in dev: a PDF marked `complete` 65s before the managed KB could answer for it, and an image-only PDF marked `failed` by Docling while Bedrock's image extraction had indexed and served it correctly. `handler.py` now resolves the engine before writing any status and returns early for `managed` (#900) +- **Managed KB: deleting a document never removed it from the managed knowledge base.** `cleanup_service` removed the legacy S3 Vectors copy and the `DOC#` row and left the managed copy indexed forever — paid for at $5.00/GB-month against S3 Vectors' ~$0.15, silently consuming `top_k` slots that the status filter then dropped, with the one fail-open branch of that filter left load-bearing. A third engine-gated phase is conjoined into `all_succeeded`, and unlike the ingestion gate an unreadable record **fails** rather than assuming legacy, because reporting success on a failed managed delete would remove the row *and* leave the content (#900) +- **Managed KB: provisioning could strand a knowledge base unrecoverably.** `CreateKnowledgeBase` returns while still `CREATING` (measured 47–124s to `ACTIVE`), and `CreateDataSource` was called immediately; `awsKbId` was only written after *both* creates succeeded, so a failure between them left a record with no identifier and every later attempt was refused permanently on the unique name. Now: an explicit bounded wait for `ACTIVE`, `awsKbId` persisted the moment the create returns (guarded on `attribute_not_exists`), and adopt-by-name recovery on a name-collision `ConflictException` — skipping terminal statuses, so adoption cannot take a knowledge base that is being deleted (#889, #898) +- **Managed KB: ingestion gave up before indexing finished, and re-ingested while it ran.** The consumer polled a *retrieval* for 30s — smaller than the documented lower bound for PDF ingestion — and, because `IngestKnowledgeBaseDocuments` is fire-and-forget, every EventBridge redelivery re-submitted the document and restarted the work it was waiting for. A 1.5MB PDF sat at `uploading` indefinitely with a fully retrievable copy in the KB. `handle_object` now probes `GetKnowledgeBaseDocuments` first and branches on the real `DocumentStatus` enum, with a 600s budget that a cross-language test asserts fits inside the Lambda timeout (#901) +- **Managed KB: `indexedAt` was fabricated** — set from the local clock the moment the ingest call returned, presented as when indexing finished, minutes apart for a large document. Now Bedrock's own `updatedAt` (#901) +- **Managed KB: `verify` failed good migrations for being asked too early.** The canary queried a freshly-ingested document before it was retrievable (measured ~45s on a fresh KB, against a warm 0.75–1.03s figure) and treated the empty result as terminal. Now defers via `records.defer_verify`, bounded by `MAX_VERIFY_ATTEMPTS` (#898) +- **Admin "Discover from server" returned 403 for every IAM-authenticated MCP server.** `POST /admin/tools/discover` signs with the app-api task role, which held `AddPermission`/`RemovePermission`/`GetFunctionUrlConfig` but never `lambda:InvokeFunctionUrl` — so admins had to type each tool name by hand in every environment (#911) + +### 🔒 Security + +- **OIDC login CSRF / session fixation in the BFF auth flow (High).** `GET /auth/login` minted a `state`, stored it server-side, and issued no browser-side material at all — no state cookie, no PKCE, no nonce — and `GET /auth/callback` treated "this state exists in the store" as proof the request continued a login *this* browser started. An attacker could mint a state anonymously, authenticate at the IdP themselves, and lure a victim to `/auth/callback?code=&state=`, silently issuing the victim a live session for the attacker's account (reported against a `system_admin` identity). Login now returns a 32-byte secret in a `__Host-bff_oauth_state` cookie and commits only its SHA-256 digest, checked with `secrets.compare_digest` **before** the state-store lookup so a probe cannot burn an in-flight state. PKCE (S256) and OIDC nonce verification are added end-to-end as defense in depth. A state row carrying no digest fails closed (#903) +- **Privilege-escalating stored XSS in skill resources.** A zero-privilege user could upload bytes labelled `text/html` — the routes persisted the client-supplied multipart Content-Type verbatim and permitted an `.html` filename — and the read routes reflected that type with `Content-Disposition: inline` while the CloudFront `/api/*` behavior carried no response-headers policy. Because app-api shares an origin with the SPA, the file parsed as a top-level HTML document and its inline `" +) + +# Every extension that turns a response body into a scriptable document. +ACTIVE_DOCUMENT_FILENAMES = [ + "vx.html", + "vx.htm", + "vx.xhtml", + "vx.xml", + "vx.svg", + "vx.shtml", + "vx.mhtml", + "vx.xht", +] + + +# --------------------------------------------------------------------------- +# Policy module — the allowlist itself +# --------------------------------------------------------------------------- + + +class TestResourceTypePolicy: + @pytest.mark.parametrize("filename", ACTIVE_DOCUMENT_FILENAMES) + def test_active_document_extensions_are_refused(self, filename): + with pytest.raises(SkillResourceTypeError): + resolve_upload_content_type(filename) + + def test_client_supplied_content_type_is_not_a_parameter(self): + """The signature itself is the control: there is nothing to spoof. + + The vulnerability was that the multipart header reached storage. If a + future refactor reintroduces a caller-supplied type argument, this fails. + """ + import inspect + + params = list(inspect.signature(resolve_upload_content_type).parameters) + assert params == ["filename"] + + def test_allowlist_contains_no_scriptable_media_type(self): + """No allowlist value may be a type a browser parses as a document.""" + forbidden = { + "text/html", + "application/xhtml+xml", + "image/svg+xml", + "text/xml", + "application/xml", + "text/javascript", + "application/javascript", + "application/x-shockwave-flash", + } + assert forbidden.isdisjoint(set(SAFE_EXTENSION_CONTENT_TYPES.values())) + + def test_code_extensions_store_as_inert_plain_text(self): + """D5: script resources are stored, listed, and never executable.""" + for ext in ("js", "mjs", "ts", "py", "sh", "css"): + assert resolve_upload_content_type(f"example.{ext}") == "text/plain" + + def test_ordinary_bundle_types_still_work(self): + assert resolve_upload_content_type("forms.md") == "text/markdown" + assert resolve_upload_content_type("data.json") == "application/json" + assert resolve_upload_content_type("chart.png") == "image/png" + assert resolve_upload_content_type("manual.pdf") == "application/pdf" + + def test_extension_match_is_case_insensitive(self): + """``.HTML`` must not slip past a lowercase-only comparison.""" + with pytest.raises(SkillResourceTypeError): + resolve_upload_content_type("vx.HTML") + assert resolve_upload_content_type("FORMS.MD") == "text/markdown" + + def test_double_extension_resolves_on_the_last_one(self): + """``notes.md.html`` is an HTML file, not a markdown file.""" + with pytest.raises(SkillResourceTypeError): + resolve_upload_content_type("notes.md.html") + + def test_extensionless_filename_is_refused(self): + with pytest.raises(SkillResourceTypeError): + resolve_upload_content_type("payload") + + def test_refusal_message_names_the_extension(self): + """An author who is blocked needs to know what to rename or convert.""" + with pytest.raises(SkillResourceTypeError, match=r"\.html"): + resolve_upload_content_type("vx.html") + + def test_download_type_never_reflects_an_active_type(self): + for filename in ACTIVE_DOCUMENT_FILENAMES: + assert safe_download_content_type(filename) == "application/octet-stream" + + def test_download_headers_are_attachment_nosniff_and_inert_csp(self): + headers = resource_download_headers("vx.html") + assert headers["Content-Disposition"].startswith("attachment;") + assert "inline" not in headers["Content-Disposition"] + assert headers["X-Content-Type-Options"] == "nosniff" + assert "default-src 'none'" in headers["Content-Security-Policy"] + assert "sandbox" in headers["Content-Security-Policy"] + + def test_download_headers_cannot_be_broken_out_of(self): + """A quote or CRLF in a legacy filename must not reach a header value.""" + headers = resource_download_headers('a".md\r\nX-Evil: 1') + value = headers["Content-Disposition"] + assert '"' not in value.removeprefix('attachment; filename="').removesuffix('"') + assert "\r" not in value and "\n" not in value + + +# --------------------------------------------------------------------------- +# Write side — the unprivileged upload route +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def author_client(user_skill_service, author_user, monkeypatch): + """The attacker in the report: authenticated, zero admin privilege.""" + monkeypatch.setattr( + my_skill_routes, "get_user_skill_service", lambda: user_skill_service + ) + app = FastAPI() + app.include_router(my_skill_routes.router) + app.dependency_overrides[get_current_user_from_session] = lambda: author_user + with TestClient(app) as c: + yield c + + +@pytest.fixture() +def admin_client(skill_service, admin_user, monkeypatch): + """The victim in the report: an ``admin.skills`` holder reading a resource.""" + monkeypatch.setattr( + admin_skill_routes, "get_skill_catalog_service", lambda: skill_service + ) + app = FastAPI() + app.include_router(admin_skill_routes.router) + override_admin_auth(app, lambda: admin_user) + return TestClient(app) + + +def _create_my_skill(client, name="Verify XSS Skill QQ") -> str: + resp = client.post( + "/skills/mine", + json={"displayName": name, "description": "security verification"}, + ) + assert resp.status_code == 200, resp.text + return resp.json()["skillId"] + + +def _upload_mine(client, skill_id, filename, body, content_type, kind="reference"): + return client.post( + f"/skills/mine/{skill_id}/resources", + files={"file": (filename, body, content_type)}, + data={"kind": kind}, + ) + + +class TestUnprivilegedUploadIsAllowlisted: + def test_html_upload_is_rejected(self, author_client): + """Step 4 of the report: the upload that planted the payload.""" + skill_id = _create_my_skill(author_client) + resp = _upload_mine( + author_client, skill_id, "vx.html", XSS_PAYLOAD, "text/html" + ) + assert resp.status_code == 400, resp.text + assert "not allowed" in resp.json()["detail"].lower() + + # Nothing was persisted — not in the manifest, not in S3. + manifest = author_client.get(f"/skills/mine/{skill_id}/resources").json() + assert manifest["resources"] == [] + s3 = boto3.client("s3", region_name=AWS_REGION) + keys = [ + o["Key"] + for o in s3.list_objects_v2(Bucket=SKILL_RESOURCES_BUCKET).get( + "Contents", [] + ) + ] + assert not any(k.endswith("vx.html") for k in keys) + + @pytest.mark.parametrize("filename", ACTIVE_DOCUMENT_FILENAMES) + def test_every_active_document_extension_is_rejected( + self, author_client, filename + ): + skill_id = _create_my_skill(author_client) + resp = _upload_mine( + author_client, skill_id, filename, XSS_PAYLOAD, "text/html" + ) + assert resp.status_code == 400, f"{filename} was accepted" + + def test_html_payload_under_a_markdown_name_is_stored_as_markdown( + self, author_client + ): + """Bytes are not sniffed — the *served type* is what disarms them. + + ``.md`` is allowed and HTML bytes inside a markdown file are legitimate + (a documentation snippet), so this upload succeeds. What makes it inert + is that it is stored and served as ``text/markdown`` with ``attachment`` + + ``nosniff``, never as a document. + """ + skill_id = _create_my_skill(author_client) + resp = _upload_mine( + author_client, skill_id, "vx.md", XSS_PAYLOAD, "text/html" + ) + assert resp.status_code == 200, resp.text + assert resp.json()["resources"][0]["contentType"] == "text/markdown" + + def test_declared_content_type_is_ignored_entirely(self, author_client): + """The multipart header is attacker-controlled and must not be stored.""" + skill_id = _create_my_skill(author_client) + resp = _upload_mine( + author_client, skill_id, "notes.md", b"# notes", "text/html" + ) + assert resp.status_code == 200, resp.text + assert resp.json()["resources"][0]["contentType"] == "text/markdown" + + def test_script_kind_does_not_bypass_the_allowlist(self, author_client): + """``kind=script`` is about bundle layout, not about type freedom.""" + skill_id = _create_my_skill(author_client) + resp = _upload_mine( + author_client, + skill_id, + "vx.html", + XSS_PAYLOAD, + "text/html", + kind="script", + ) + assert resp.status_code == 400 + + +class TestAdminUploadIsAllowlisted: + """The admin catalog tier shares the code path and the same policy.""" + + def _create_catalog_skill(self, admin_client, skill_id="pdf_workflows"): + resp = admin_client.post( + "/skills/", + json={ + "skillId": skill_id, + "displayName": "PDF Workflows", + "description": "Fill, merge and split PDFs.", + "instructions": "# PDF Workflows", + }, + ) + assert resp.status_code == 200, resp.text + return skill_id + + def test_html_upload_is_rejected(self, admin_client): + skill_id = self._create_catalog_skill(admin_client) + resp = admin_client.post( + f"/skills/{skill_id}/resources", + files={"file": ("vx.html", XSS_PAYLOAD, "text/html")}, + ) + assert resp.status_code == 400, resp.text + + +# --------------------------------------------------------------------------- +# Read side — a legacy row must be neutralized without a data migration +# --------------------------------------------------------------------------- + + +def _plant_legacy_html_resource(service, skill_id, filename="vx.html"): + """Write a manifest entry the way the *vulnerable* upload path would have. + + Goes straight at the repository and the store, deliberately bypassing + ``add_resource``, because the point is to reproduce a row that already + exists in a deployed environment: ``content_type == "text/html"``. + """ + from apis.shared.skills.models import SkillResourceRef + from apis.shared.skills.resource_store import compute_content_hash + + key = service.resource_store.put( + skill_id=skill_id, + filename=filename, + content=XSS_PAYLOAD, + content_type="text/html", + kind="reference", + ) + ref = SkillResourceRef( + filename=filename, + content_hash=compute_content_hash(XSS_PAYLOAD), + size=len(XSS_PAYLOAD), + content_type="text/html", # the poisoned value + s3_key=key, + kind="reference", + ) + return ref + + +def _assert_response_cannot_execute(resp): + """The response body may be the payload; it may not be a document.""" + assert resp.status_code == 200, resp.text + assert resp.content == XSS_PAYLOAD # bytes are still readable by the SPA + assert "text/html" not in resp.headers["content-type"] + assert resp.headers["content-type"].startswith("application/octet-stream") + assert resp.headers["content-disposition"].startswith("attachment;") + assert "inline" not in resp.headers["content-disposition"] + assert resp.headers["x-content-type-options"] == "nosniff" + assert "default-src 'none'" in resp.headers["content-security-policy"] + + +class TestLegacyRowIsNeutralizedOnRead: + @pytest.mark.asyncio + async def test_admin_read_of_a_legacy_html_row( + self, admin_client, skill_service, admin_user + ): + """Step 5 of the report: the request that executed the script.""" + admin_client.post( + "/skills/", + json={ + "skillId": "legacy_skill", + "displayName": "Legacy", + "description": "d", + "instructions": "i", + }, + ) + ref = _plant_legacy_html_resource(skill_service, "legacy_skill") + await skill_service.repository.update_skill( + "legacy_skill", {"resources": [ref]}, admin_user_id=admin_user.user_id + ) + + resp = admin_client.get("/skills/legacy_skill/resources/vx.html") + _assert_response_cannot_execute(resp) + + @pytest.mark.asyncio + async def test_owner_read_of_a_legacy_html_row( + self, author_client, user_skill_service, author_user + ): + """The owner-tier read route is hardened identically.""" + skill_id = _create_my_skill(author_client, "Legacy Mine") + ref = _plant_legacy_html_resource( + user_skill_service.catalog_service, skill_id + ) + await user_skill_service.repository.update_skill( + skill_id, {"resources": [ref]}, admin_user_id=author_user.user_id + ) + + resp = author_client.get(f"/skills/mine/{skill_id}/resources/vx.html") + _assert_response_cannot_execute(resp) + + +class TestHardenedHeadersOnOrdinaryReads: + def test_markdown_read_keeps_its_type_but_gains_the_headers( + self, author_client + ): + """The SPA reads these over XHR as text, so ``attachment`` is safe.""" + skill_id = _create_my_skill(author_client, "Notes Skill") + _upload_mine(author_client, skill_id, "forms.md", b"# Forms", "text/markdown") + + resp = author_client.get(f"/skills/mine/{skill_id}/resources/forms.md") + assert resp.status_code == 200 + assert resp.content == b"# Forms" + assert resp.headers["content-type"].startswith("text/markdown") + assert resp.headers["content-disposition"].startswith("attachment;") + assert resp.headers["x-content-type-options"] == "nosniff" + + def test_no_read_route_serves_inline(self): + """Pins the header both tiers regressed on, at the source level.""" + import inspect + + for module in (my_skill_routes, admin_skill_routes): + source = inspect.getsource(module) + assert "inline; filename=" not in source, module.__name__ diff --git a/backend/tests/apis/app_api/skills/test_user_skill_service.py b/backend/tests/apis/app_api/skills/test_user_skill_service.py index 509648bb8..b42b21045 100644 --- a/backend/tests/apis/app_api/skills/test_user_skill_service.py +++ b/backend/tests/apis/app_api/skills/test_user_skill_service.py @@ -365,12 +365,17 @@ async def test_admin_catalog_list_excludes_user_skills( async def test_user_skills_cannot_be_granted_to_app_roles( user_skill_service, skill_service, admin_user, author_user ): - """Granting a private user skill to a role would leak it to that role.""" + """Granting a private user skill to a role would leak it to that role. + + The refusal reports "not found" rather than "user-authored": the admin + catalog surface must not confirm the existence of a skill it cannot + govern, so a cross-tier id is indistinguishable from an unknown one. + """ mine = await user_skill_service.create_my_skill( author_user, display_name="Mine", description="d" ) - with pytest.raises(ValueError, match="user-authored"): + with pytest.raises(ValueError, match="not found"): await skill_service.set_roles_for_skill(mine.skill_id, ["some_role"], admin_user) - with pytest.raises(ValueError, match="user-authored"): + with pytest.raises(ValueError, match="not found"): await skill_service.add_roles_to_skill(mine.skill_id, ["some_role"], admin_user) diff --git a/backend/tests/apis/app_api/test_presentation_upload_size_cap.py b/backend/tests/apis/app_api/test_presentation_upload_size_cap.py new file mode 100644 index 000000000..5944fa31b --- /dev/null +++ b/backend/tests/apis/app_api/test_presentation_upload_size_cap.py @@ -0,0 +1,117 @@ +"""Presentations upload under a larger cap than everything else. + +The general 4MB limit is sized for Bedrock's *inline* document budget — it is +the point past which a document block starts risking a ValidationException +mid-stream. A .pptx never enters that path (Bedrock's document-format enum has +no `pptx`; it routes to the PowerPoint tools instead), so the ceiling that +justifies 4MB simply does not apply to it. Corporate templates with imagery +clear 4MB routinely, which made `create_powerpoint_presentation`'s +``template_name`` argument unusable for exactly the files it exists to accept. + +The cap that *does* bind a deck is the Code Interpreter hop: +``_ci_write_bytes`` base64-encodes the whole file into a single ``writeFiles`` +``text`` field (~4/3 inflation). That field is a MaxLenString (100MB), so 25MB +of deck → ~33MB of base64 sits well inside it. + +Both size gates must agree on which cap applies. Historically they were one +constant read from two places; now that the cap depends on the file, a gate +that reads ``max_file_size`` directly rejects a deck the other gate allowed — +so ``max_size_for`` is the single decision point and this pins it. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from apis.app_api.files.service import FileTooLargeError, FileUploadService +from apis.shared.files.models import PresignRequest + +PPTX_MIME = "application/vnd.openxmlformats-officedocument.presentationml.presentation" +GENERAL_CAP = 4 * 1024 * 1024 +PRESENTATION_CAP = 25 * 1024 * 1024 + + +@pytest.fixture +def service(): + repository = MagicMock() + repository.get_user_quota = AsyncMock( + return_value=SimpleNamespace(total_bytes=0) + ) + repository.create_file = AsyncMock() + + s3_client = MagicMock() + s3_client.generate_presigned_url.return_value = "https://example.invalid/put" + + return FileUploadService( + repository=repository, + s3_client=s3_client, + bucket_name="test-bucket", + max_file_size=GENERAL_CAP, + presentation_max_file_size=PRESENTATION_CAP, + ) + + +class TestMaxSizeFor: + def test_ordinary_documents_get_the_general_cap(self, service): + assert service.max_size_for("report.pdf", "application/pdf") == GENERAL_CAP + + def test_presentations_get_the_presentation_cap(self, service): + assert service.max_size_for("deck.pptx", PPTX_MIME) == PRESENTATION_CAP + + def test_extension_alone_is_enough(self, service): + # Some clients send octet-stream for pptx; the cap must not depend on + # the browser getting the MIME right. + assert ( + service.max_size_for("deck.pptx", "application/octet-stream") + == PRESENTATION_CAP + ) + + def test_defaults_match_the_documented_values(self): + # Constructed with no overrides — these are the values the frontend + # constants mirror, and the frontend must never be the larger side. + default = FileUploadService( + repository=MagicMock(), s3_client=MagicMock(), bucket_name="b" + ) + assert default.max_file_size == GENERAL_CAP + assert default.presentation_max_file_size == PRESENTATION_CAP + + +class TestPresignSizeEnforcement: + @pytest.mark.asyncio + async def test_rejects_ordinary_document_above_general_cap(self, service): + request = PresignRequest( + sessionId="s1", + filename="big.pdf", + mimeType="application/pdf", + sizeBytes=GENERAL_CAP + 1, + ) + with pytest.raises(FileTooLargeError) as exc: + await service.request_presigned_url("u1", request) + assert exc.value.max_size == GENERAL_CAP + + @pytest.mark.asyncio + async def test_accepts_deck_above_general_cap(self, service): + # The regression this whole change exists to prevent. + request = PresignRequest( + sessionId="s1", + filename="template.pptx", + mimeType=PPTX_MIME, + sizeBytes=GENERAL_CAP + 1, + ) + response = await service.request_presigned_url("u1", request) + assert response.upload_id + + @pytest.mark.asyncio + async def test_rejects_deck_above_presentation_cap(self, service): + request = PresignRequest( + sessionId="s1", + filename="huge.pptx", + mimeType=PPTX_MIME, + sizeBytes=PRESENTATION_CAP + 1, + ) + with pytest.raises(FileTooLargeError) as exc: + await service.request_presigned_url("u1", request) + # The 400 detail is built from max_size, so the user is told 25MB — + # not the 4MB that does not apply to their file. + assert exc.value.max_size == PRESENTATION_CAP diff --git a/backend/tests/apis/inference_api/test_attachment_recovery.py b/backend/tests/apis/inference_api/test_attachment_recovery.py new file mode 100644 index 000000000..24027c76e --- /dev/null +++ b/backend/tests/apis/inference_api/test_attachment_recovery.py @@ -0,0 +1,61 @@ +"""Recovery of attachments a failed turn consumed without ever being read. + +Inline document bytes are one-shot: `TurnBasedSessionManager._strip_document_bytes` +removes them from restored history (Bedrock rejects duplicate document names), +so a turn that dies before the model reads the documents loses them for good. + +Prod session `5f34d2b0` (2026-08-31) is the incident this covers — a +ConverseStream carrying two PDFs failed with ServiceUnavailableException, the +PDFs were gone, and the model's only recourse was to ask the user to upload +them again. The invocations route now pops a write-ahead marker at turn start +and re-sends those uploads on the very next turn. +""" + +from apis.inference_api.chat.routes import ( + _build_attachment_recovery_note, + _select_recovered_attachments, +) + + +class TestSelectRecoveredAttachments: + def test_re_attaches_when_the_turn_has_none_of_its_own(self): + assert _select_recovered_attachments(["up-1", "up-2"]) == ["up-1", "up-2"] + + def test_nothing_recovered_is_a_noop(self): + assert _select_recovered_attachments([]) == [] + + def test_user_upload_ids_win(self): + """Merging could push a deliberately-attached file past the 5-file + resolver cap, and re-sending the same document twice is a Bedrock + ValidationException.""" + assert _select_recovered_attachments(["up-1"], request_upload_ids=["fresh"]) == [] + + def test_direct_file_content_wins(self): + assert _select_recovered_attachments(["up-1"], request_files=[object()]) == [] + + def test_empty_request_collections_do_not_count_as_attachments(self): + assert _select_recovered_attachments( + ["up-1"], request_upload_ids=[], request_files=[] + ) == ["up-1"] + + def test_returns_a_copy(self): + original = ["up-1"] + result = _select_recovered_attachments(original) + result.append("up-2") + assert original == ["up-1"] + + +class TestBuildAttachmentRecoveryNote: + def test_names_the_files_and_forbids_asking_for_a_re_upload(self): + """The whole failure mode was the model telling the user to re-upload + files the server already had.""" + note = _build_attachment_recovery_note(["a.pdf", "b.pdf"]) + assert note.startswith("") + assert note.endswith("") + assert "a.pdf" in note and "b.pdf" in note + assert "do not ask the user to upload them" in note + + def test_says_the_user_did_not_re_attach_them(self): + note = _build_attachment_recovery_note(["report.pdf"]) + assert "did not attach them again" in note + assert "previous turn" in note diff --git a/backend/tests/apis/inference_api/test_presentation_attachment_carveout.py b/backend/tests/apis/inference_api/test_presentation_attachment_carveout.py new file mode 100644 index 000000000..87655b360 --- /dev/null +++ b/backend/tests/apis/inference_api/test_presentation_attachment_carveout.py @@ -0,0 +1,194 @@ +"""A .pptx attachment must never reach Bedrock as an inline document block. + +This is not the same kind of rule as the tabular carve-out it sits next to. +Spreadsheets are diverted as an *optimization* — an xlsx would technically be +accepted inline, it just inflates past the 4.5MB internal limit and analyzes +worse than pandas would. A pptx is diverted because Bedrock's Converse +``DocumentFormat`` enum has no ``pptx`` member at all: + + pdf | csv | doc | docx | xls | xlsx | html | txt | md + +So an inline deck is an unconditional ValidationException that kills the turn, +at any size, forever — not a threshold we tune. That is why +``is_presentation_file`` is checked BEFORE the size gate: routing a small deck +to the "oversized" bucket would produce a note that misdescribes why it was +skipped, and routing it to `inline` at all is simply broken. + +The upload path and this carve-out are one feature. `.pptx` is in the backend +and frontend upload allowlists only because these tools can receive it; if a +future change re-narrows either allowlist, the create-deck tool's own error +text ("Upload a .pptx template first") becomes a lie again. +""" + +import pytest + +from apis.inference_api.chat.routes import ( + _attachment_marker_names, + _build_attachment_guidance, + _partition_attachments, +) +from apis.shared.files.models import ALLOWED_MIME_TYPES, is_presentation_file + +PPTX_MIME = "application/vnd.openxmlformats-officedocument.presentationml.presentation" +XLSX_MIME = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + + +class _Attachment: + """Minimal stand-in for FileContent — the partition only reads these three.""" + + def __init__(self, filename: str, content_type: str, bytes_: str = ""): + self.filename = filename + self.content_type = content_type + self.bytes = bytes_ + + +class TestIsPresentationFile: + def test_detects_by_mime_type(self): + assert is_presentation_file("anything", PPTX_MIME) is True + + def test_detects_by_extension_when_mime_is_missing(self): + # Browsers and some clients send "" or application/octet-stream for + # pptx; the extension is the fallback, same as the tabular helper. + assert is_presentation_file("deck.pptx", "") is True + assert is_presentation_file("deck.PPTX", "application/octet-stream") is True + + def test_does_not_claim_other_documents(self): + assert is_presentation_file("report.pdf", "application/pdf") is False + assert is_presentation_file("data.xlsx", XLSX_MIME) is False + + def test_legacy_ppt_is_not_claimed(self): + # .ppt (binary, pre-2007) is not in the upload allowlist and + # python-pptx cannot open it — don't silently divert it. + assert is_presentation_file("old.ppt", "application/vnd.ms-powerpoint") is False + + def test_pptx_is_uploadable(self): + # The carve-out is unreachable if the upload gate rejects the file. + assert ALLOWED_MIME_TYPES.get(PPTX_MIME) == "pptx" + + +class TestPartitionAttachments: + def test_pptx_is_diverted_not_inlined(self): + deck = _Attachment("deck.pptx", PPTX_MIME) + inline, tabular, presentations, oversized = _partition_attachments([deck]) + + assert presentations == [deck] + assert inline == [] + assert tabular == [] + assert oversized == [] + + def test_tiny_pptx_still_diverted_never_oversized(self): + # The size gate must not get a vote: a 12-byte deck is still a deck, + # and "too big" would be the wrong explanation for skipping it. + deck = _Attachment("small.pptx", PPTX_MIME, bytes_="AAAA") + inline, _, presentations, oversized = _partition_attachments([deck]) + + assert presentations == [deck] + assert oversized == [] + assert inline == [] + + def test_ordinary_documents_still_inline(self): + pdf = _Attachment("report.pdf", "application/pdf", bytes_="AAAA") + _, _, presentations, _ = _partition_attachments([pdf]) + assert presentations == [] + + def test_mixed_batch_lands_in_the_right_buckets(self): + pdf = _Attachment("report.pdf", "application/pdf", bytes_="AAAA") + sheet = _Attachment("data.xlsx", XLSX_MIME) + deck = _Attachment("deck.pptx", PPTX_MIME) + + inline, tabular, presentations, oversized = _partition_attachments( + [pdf, sheet, deck] + ) + + assert inline == [pdf] + assert tabular == [sheet] + assert presentations == [deck] + assert oversized == [] + + +class TestAttachmentMarkerNames: + """Diverting a file must not erase it from the message it was attached to. + + The `[Attached files: …]` marker is the only link the SPA can replay on + reload — see `_attachment_marker_names`. Deriving it from the inline set + is what made a diverted deck's card vanish from restored history. + """ + + def test_includes_a_diverted_deck(self): + pdf = _Attachment("report.pdf", "application/pdf") + deck = _Attachment("deck.pptx", PPTX_MIME) + assert _attachment_marker_names([pdf, deck], []) == [ + "report.pdf", + "deck.pptx", + ] + + def test_includes_a_diverted_spreadsheet(self): + sheet = _Attachment("data.xlsx", XLSX_MIME) + assert _attachment_marker_names([sheet], []) == ["data.xlsx"] + + def test_a_lone_deck_still_yields_a_name(self): + # Nothing inline at all — the case that previously left no trace. + deck = _Attachment("deck.pptx", PPTX_MIME) + assert _attachment_marker_names([deck], []) == ["deck.pptx"] + + def test_excludes_oversized_files(self): + # Dropped from the turn entirely; the guidance explains their absence, + # so a card promising otherwise would be misleading. + pdf = _Attachment("report.pdf", "application/pdf") + huge = _Attachment("huge.pdf", "application/pdf") + assert _attachment_marker_names([pdf, huge], [huge]) == ["report.pdf"] + + def test_preserves_attachment_order(self): + # Order is deterministic because this text reaches the cacheable + # prefix on later turns. + files = [ + _Attachment("b.pptx", PPTX_MIME), + _Attachment("a.pdf", "application/pdf"), + _Attachment("c.xlsx", XLSX_MIME), + ] + assert _attachment_marker_names(files, []) == ["b.pptx", "a.pdf", "c.xlsx"] + + def test_no_attachments_yields_no_names(self): + assert _attachment_marker_names([], []) == [] + + +class TestAttachmentGuidance: + def test_names_the_deck_and_the_read_tool_when_enabled(self): + deck = _Attachment("deck.pptx", PPTX_MIME) + guidance = _build_attachment_guidance( + [], [deck], [], ["create_powerpoint_presentation"] + ) + + assert "`deck.pptx`" in guidance + assert "read_powerpoint_presentation" in guidance + + def test_tells_the_user_which_toggle_to_flip_when_disabled(self): + # A diverted deck with no tool to read it is a dead end unless the + # note names the toggle — the file is neither inline nor reachable. + deck = _Attachment("deck.pptx", PPTX_MIME) + guidance = _build_attachment_guidance([], [deck], [], ["some_other_tool"]) + + assert "PowerPoint Presentations" in guidance + assert "read_powerpoint_presentation" not in guidance + + @pytest.mark.parametrize("enabled_tools", [None, []]) + def test_no_enabled_tools_is_treated_as_disabled(self, enabled_tools): + deck = _Attachment("deck.pptx", PPTX_MIME) + guidance = _build_attachment_guidance([], [deck], [], enabled_tools) + assert "PowerPoint Presentations" in guidance + + def test_silent_when_nothing_was_diverted(self): + assert _build_attachment_guidance([], [], [], ["create_powerpoint_presentation"]) == "" + + def test_spreadsheet_and_deck_notes_coexist(self): + # Both carve-outs can fire on one turn; neither may swallow the other. + sheet = _Attachment("data.xlsx", XLSX_MIME) + deck = _Attachment("deck.pptx", PPTX_MIME) + guidance = _build_attachment_guidance( + [sheet], [deck], [], ["analyze_spreadsheet", "create_powerpoint_presentation"] + ) + + assert "`data.xlsx`" in guidance + assert "`deck.pptx`" in guidance + assert "analyze_spreadsheet" in guidance + assert "read_powerpoint_presentation" in guidance diff --git a/backend/tests/apis/inference_api/test_turn_lease_release.py b/backend/tests/apis/inference_api/test_turn_lease_release.py new file mode 100644 index 000000000..f996b4d8e --- /dev/null +++ b/backend/tests/apis/inference_api/test_turn_lease_release.py @@ -0,0 +1,163 @@ +"""Turn-lease teardown survives the cancellation that triggers it. + +`_release_turn_lease` runs from the SSE stream generator's `finally`, and the +case it exists for is the one where cancellation is what put it there: the +browser's connection drops and Starlette tears the response down. A bare +`await` in that `finally` never completes, so the release is abandoned — the +lease then survives its full 90s window and the user's resend is rejected as a +duplicate turn (409 → the Runtime's 424 rewrite → "Chat Request Failed"), +seconds after the dropped stream already showed them a network error. + +The mechanism is specifically **anyio cancel scopes**, which is what Starlette +cancels a disconnected `StreamingResponse` with. Unlike a one-shot +`task.cancel()` — which lets a `finally` run its awaits to completion — an +anyio scope is level-triggered: while it is cancelled, *every* checkpoint +inside it raises `CancelledError`, including the ones in cleanup code. Tests +that cancel a bare asyncio task therefore pass with or without the fix and +prove nothing; these use a real cancel scope. + +Observed in prod-ai 2026-08-13: session 938a1e68 acquired a lease at 14:28:44, +was interrupted (`reason=connection_lost`) at 14:30:15 with no release ever +logged, and its resend at 14:31:16 was rejected. Twice in twelve minutes. +""" + +from __future__ import annotations + +import asyncio + +import anyio +import pytest + +import apis.shared.sessions.session_lease as session_lease_module +from apis.inference_api.chat.routes import _release_turn_lease +from apis.shared.sessions.session_lease import SessionLease + + +def _lease() -> SessionLease: + return SessionLease(session_id="s1", user_id="u1", owner="owner-token") + + +@pytest.fixture +def released(monkeypatch: pytest.MonkeyPatch) -> list: + """Record released leases, with a real suspension point in the release. + + The suspension is the whole point: DynamoDB's round-trip is where a + cancelled `finally` abandons the work. + """ + seen: list = [] + + async def _release(lease) -> None: + await asyncio.sleep(0.01) + seen.append(lease) + + monkeypatch.setattr(session_lease_module, "release_session_lease", _release) + return seen + + +async def _stream_dropped_mid_turn(lease, heartbeat_task=None) -> None: + """Reproduce a client disconnect the way Starlette delivers one: cancel the + surrounding anyio scope while the stream generator is suspended, so the + generator's `finally` runs inside an already-cancelled scope.""" + with anyio.CancelScope() as scope: + + async def stream(): + try: + while True: + yield "chunk" + scope.cancel() # the browser goes away + await anyio.sleep(0.01) # raises CancelledError here + finally: + await _release_turn_lease(heartbeat_task, lease) + + async for _ in stream(): + pass + + +class TestReleaseUnderClientDisconnect: + @pytest.mark.asyncio + async def test_lease_is_released_when_the_connection_drops(self, released): + lease = _lease() + await _stream_dropped_mid_turn(lease) + + # The shielded release runs on its own task, outside the cancelled + # scope — give the loop a beat to finish it. + await asyncio.sleep(0.05) + assert released == [lease], ( + "lease was not released on a dropped stream — the user's resend " + "will be rejected as a duplicate turn" + ) + + @pytest.mark.asyncio + async def test_heartbeat_is_stopped_before_release(self, released): + """A surviving heartbeat would keep renewing the lease we just freed.""" + started = asyncio.Event() + + async def _heartbeat() -> None: + started.set() + while True: + await asyncio.sleep(0.01) + + heartbeat_task = asyncio.create_task(_heartbeat()) + await started.wait() + + lease = _lease() + await _stream_dropped_mid_turn(lease, heartbeat_task) + await asyncio.sleep(0.05) + + assert heartbeat_task.cancelled() + assert released == [lease] + + +class TestReleaseUnderTaskCancellation: + @pytest.mark.asyncio + async def test_cancellation_still_propagates(self, released): + """Swallowing CancelledError in the teardown must not swallow the + cancellation that unwound the stream — the coordinator's interruption + arm and Starlette's teardown both depend on it propagating.""" + lease = _lease() + + async def consume(): + async def stream(): + try: + while True: + yield "chunk" + await asyncio.sleep(0.01) + finally: + await _release_turn_lease(None, lease) + + async for _ in stream(): + pass + + task = asyncio.create_task(consume()) + await asyncio.sleep(0.03) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + await asyncio.sleep(0.05) + assert released == [lease] + + +class TestReleaseOnNormalCompletion: + @pytest.mark.asyncio + async def test_lease_is_released_when_the_stream_ends_normally(self, released): + lease = _lease() + + async def stream(): + try: + yield "chunk" + finally: + await _release_turn_lease(None, lease) + + async for _ in stream(): + pass + + assert released == [lease] + + @pytest.mark.asyncio + async def test_no_lease_is_a_noop(self, released): + # Preview sessions and the local no-DynamoDB path never take a lease; + # release_session_lease is itself a no-op on None. + await _release_turn_lease(None, None) + assert released == [None] diff --git a/backend/tests/architecture/test_admin_scope_coverage.py b/backend/tests/architecture/test_admin_scope_coverage.py index ff5566bbd..dec034e20 100644 --- a/backend/tests/architecture/test_admin_scope_coverage.py +++ b/backend/tests/architecture/test_admin_scope_coverage.py @@ -203,8 +203,17 @@ def test_every_mounted_admin_route_has_a_scope_dependency() -> None: blob = "\n".join(sources) # `checker` is the closure returned by require_app_roles / # require_admin_scope; require_marketplace_admin wraps one of them. + # + # Two markers, because the two checkers now resolve permissions differently: + # `require_app_roles` calls `resolve_user_permissions` inline, while + # `require_admin_scope` delegates to the shared `has_admin_scope` predicate (which + # the invocation path's reviewer preview also needs, and which cannot be a FastAPI + # dependency there). This is a source grep one level deep, so a checker that moves + # its resolution behind another name must add that name here — the guarantee is + # unchanged, the string that evidences it is not. governed = ( "resolve_user_permissions" in blob + or "has_admin_scope" in blob or "require_marketplace_admin" in blob or "agent_marketplace_enabled" in blob ) diff --git a/backend/tests/architecture/test_kb_backend_boundary.py b/backend/tests/architecture/test_kb_backend_boundary.py new file mode 100644 index 000000000..ef347b9a9 --- /dev/null +++ b/backend/tests/architecture/test_kb_backend_boundary.py @@ -0,0 +1,194 @@ +"""Import-boundary enforcement for ``apis.shared.kb_backend``. + +``apis/shared/assistants/__init__.py`` imports ``rag_service``, which imports the +embeddings stack at module scope. So importing anything from the assistants +package pulls in that whole tree — and ``kb_backend`` is bundled into +size-constrained Lambda images (the migration dispatcher, worker, reconciler and +ingestion consumer) that deliberately do not carry it. The same constraint is why +``apis/app_api/kb_sync/records.py`` reaches DynamoDB through the raw table +resource instead of the assistants package. + +The dependency is also the wrong way round architecturally: the facade in +``rag_service`` sits *above* the seam and depends on ``kb_backend``. An import in +the other direction would make the two mutually dependent and the seam +meaningless. + +This is checked two ways, because either alone is insufficient: + +* **Statically**, so a *lazy* import inside a function body is caught. A deferred + import does not fail at module load; it fails at call time, in production, in a + Lambda that has been running fine for a week. +* **At runtime in a fresh interpreter**, so a transitive import through some + innocuous-looking third module is caught too. Static analysis cannot see + through an import chain; a subprocess with an empty ``sys.modules`` can. + +Feature: managed-kb-migration +Requirements: 24.15 +""" + +import ast +import subprocess +import sys +from pathlib import Path +from typing import List, Tuple + +import pytest + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent.parent +_BACKEND_SRC = _BACKEND_ROOT / "src" +_KB_BACKEND = _BACKEND_SRC / "apis" / "shared" / "kb_backend" + +#: Modules whose absence from a fresh import is asserted. ``boto3`` is here +#: because it is the single largest dependency these Lambdas would otherwise +#: pay for, and keeping it function-local is the convention this package follows. +_FORBIDDEN_AT_IMPORT_TIME = ("apis.shared.assistants", "boto3") + + +def _extract_imports(filepath: Path) -> List[Tuple[str, int]]: + """Every imported module path in a file, including imports inside functions.""" + try: + tree = ast.parse(filepath.read_text(encoding="utf-8"), filename=str(filepath)) + except (SyntaxError, UnicodeDecodeError): + return [] + + imports: List[Tuple[str, int]] = [] + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + imports.append((alias.name, node.lineno)) + elif isinstance(node, ast.ImportFrom) and node.module: + imports.append((node.module, node.lineno)) + return imports + + +def _kb_backend_files() -> List[Path]: + return sorted(_KB_BACKEND.rglob("*.py")) + + +class TestKbBackendDoesNotImportAssistants: + """No file in kb_backend may import apis.shared.assistants, at any depth.""" + + def test_no_assistants_imports_anywhere(self): + if not _KB_BACKEND.exists(): + pytest.skip("kb_backend package not found") + + violations = [] + for pyfile in _kb_backend_files(): + rel = pyfile.relative_to(_BACKEND_SRC) + for module, lineno in _extract_imports(pyfile): + if module == "apis.shared.assistants" or module.startswith("apis.shared.assistants."): + violations.append(f" {rel}:{lineno} imports '{module}'") + + assert violations == [], ( + "apis.shared.kb_backend must not import apis.shared.assistants " + "(its __init__ pulls in rag_service and the whole embeddings stack, " + "which the migration Lambda images do not carry):\n" + + "\n".join(violations) + + "\n\nNote that a lazy, function-local import does not fix this — it " + "moves the failure from image build to production call time." + ) + + def test_package_init_stays_empty(self): + """An empty ``__init__`` is what makes importing one submodule cheap. + + Re-exporting anything here would mean importing ``kb_backend.records`` + also imports every sibling — including, eventually, the managed backend + and its boto3 client. + """ + init = _KB_BACKEND / "__init__.py" + assert init.exists(), "kb_backend/__init__.py must exist" + assert init.read_text(encoding="utf-8").strip() == "", ( + "kb_backend/__init__.py must stay empty: it is imported by every " + "submodule import, so anything placed here is paid for by all of them" + ) + + +class TestKbBackendFreshImportIsLean: + """Importing a kb_backend submodule must not pull the heavy tree in. + + Each case runs in a fresh interpreter, because by the time this test file + executes, the rest of the suite has already imported both forbidden modules + into ``sys.modules`` — an in-process check would pass no matter what. + """ + + @staticmethod + def _import_and_report(module: str) -> List[str]: + """Import *module* in a subprocess; return which forbidden modules loaded.""" + program = ( + "import sys\n" + f"import {module}\n" + "loaded = [name for name in sys.modules\n" + f" if any(name == f or name.startswith(f + '.') for f in {_FORBIDDEN_AT_IMPORT_TIME!r})]\n" + "print(','.join(sorted(set(loaded))))\n" + ) + result = subprocess.run( + [sys.executable, "-c", program], + capture_output=True, + text=True, + cwd=str(_BACKEND_ROOT), + env={"PYTHONPATH": str(_BACKEND_SRC), "PATH": "/usr/bin:/bin"}, + ) + assert result.returncode == 0, ( + f"importing {module} in a clean interpreter failed:\n{result.stderr}" + ) + return [name for name in result.stdout.strip().split(",") if name] + + def test_records_import_is_stdlib_only(self): + """The constraint as written in task 4.8: records pulls in neither.""" + loaded = self._import_and_report("apis.shared.kb_backend.records") + assert loaded == [], ( + "importing apis.shared.kb_backend.records loaded " + f"{loaded}. Module-level imports in this package must be stdlib " + "only; move boto3 and anything from apis.shared.assistants into the " + "functions that need them." + ) + + @pytest.mark.parametrize( + "module", + [ + "apis.shared.kb_backend.protocol", + "apis.shared.kb_backend.resolver", + "apis.shared.kb_backend.s3vectors_backend", + "apis.shared.kb_backend.managed_backend", + "apis.shared.kb_backend.dual_read", + ], + ) + def test_seam_modules_import_lean(self, module): + """The resolver and both adapters obey the same rule as records. + + The resolver is the one that matters most: the facade imports it on every + retrieval, and it in turn imports every registered backend. If it were + not lean, no submodule of this package could be. + + ``managed_backend`` is on this list because the resolver **registers** it at + import (see the resolver's docstring). That registration is only free while + the adapter's module body stays stdlib-only and its clients stay lazy; the + day someone hoists a ``boto3.client(...)`` to module scope, every Lambda + image carrying any part of this package pays for it. + """ + loaded = self._import_and_report(module) + assert loaded == [], ( + f"importing {module} loaded {loaded}; keep these imports " + "function-local" + ) + + +class TestFacadeDependencyDirectionIsOneWay: + """rag_service depends on kb_backend, never the reverse.""" + + def test_facade_imports_the_seam(self): + """A guard against the facade quietly regrowing its own retrieval path. + + If ``rag_service`` stopped importing the resolver, it would mean the + delegation had been inlined again and the managed backend would be + unreachable — with every legacy test still green. + """ + rag_service = _BACKEND_SRC / "apis" / "shared" / "assistants" / "rag_service.py" + modules = {module for module, _ in _extract_imports(rag_service)} + assert "apis.shared.kb_backend.resolver" in modules, ( + "rag_service must resolve its backend through " + "apis.shared.kb_backend.resolver" + ) + assert "apis.shared.kb_backend.protocol" in modules, ( + "rag_service must use the protocol's canonical chunk shape" + ) diff --git a/backend/tests/documents/test_managed_delete_propagation.py b/backend/tests/documents/test_managed_delete_propagation.py new file mode 100644 index 000000000..8adaacb6a --- /dev/null +++ b/backend/tests/documents/test_managed_delete_propagation.py @@ -0,0 +1,222 @@ +"""Deleting a document must remove it from the engine that actually serves it. + +Before this, `cleanup_service` deleted the legacy S3 Vectors copy and the `DOC#` +row and never touched the managed knowledge base. On a promoted knowledge base +the content therefore stayed indexed forever: + +* it kept being paid for, at $5.00/GB-month against S3 Vectors' ~$0.15; +* every orphan kept consuming a slot in `top_k`, because the status filter runs + *after* retrieval — so a query could return five chunks and the model see two, + with nothing logged and nothing raised; +* and the only thing preventing deleted content from being served was the + fail-closed status filter, which has exactly one fail-open branch. + +The most load-bearing assertion in this file is +`test_a_failed_managed_delete_keeps_the_document_record`: the `DOC#` row is what +the status filter joins against, so reporting success on a failed managed delete +would remove the row *and* leave the content — turning a storage leak into a +disclosure. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import pytest + +from apis.app_api.documents.services import cleanup_service + +ASSISTANT_ID = "ast-del-1" +DOCUMENT_ID = "DOC-del-1" + + +class RecordingManagedBackend: + """Stands in for ManagedKbBackend, recording deletes and optionally failing.""" + + def __init__(self, raises: Optional[BaseException] = None) -> None: + self.deleted: List[str] = [] + self.attempts = 0 + self._raises = raises + + async def delete_document(self, kb_ref: str, document_id: str) -> None: + self.attempts += 1 + if self._raises is not None: + raise self._raises + self.deleted.append(document_id) + + +@pytest.fixture +def managed(monkeypatch: pytest.MonkeyPatch) -> RecordingManagedBackend: + backend = RecordingManagedBackend() + import apis.shared.kb_backend.managed_backend as mb + + monkeypatch.setattr(mb, "ManagedKbBackend", lambda *a, **k: backend) + return backend + + +def _set_engine(monkeypatch: pytest.MonkeyPatch, engine: Optional[str]) -> None: + from apis.shared.kb_backend import records as r + + record: Optional[Dict[str, Any]] + if engine is None: + record = None + else: + record = {"retrievalEngine": engine} if engine != "absent" else {} + + monkeypatch.setattr(r, "get_kb_record", lambda *_: record) + + +def _raise_on_lookup(monkeypatch: pytest.MonkeyPatch, exc: BaseException) -> None: + from apis.shared.kb_backend import records as r + + def _boom(*_: Any): + raise exc + + monkeypatch.setattr(r, "get_kb_record", _boom) + + +async def _delete() -> bool: + return await cleanup_service._delete_managed_documents_with_retries( + ASSISTANT_ID, DOCUMENT_ID, max_retries=3, base_delay=0.0 + ) + + +class TestAPromotedKnowledgeBaseHasTheDocumentRemoved: + @pytest.mark.asyncio + async def test_the_document_is_deleted_from_the_managed_kb( + self, managed: RecordingManagedBackend, monkeypatch + ): + _set_engine(monkeypatch, "managed") + + assert await _delete() is True + assert managed.deleted == [DOCUMENT_ID], ( + "a document deleted by its owner was left in the managed knowledge " + "base; the corpus can only grow and every orphan costs a top_k slot" + ) + + +class TestALegacyKnowledgeBaseIsUntouched: + """No managed copy to remove is success, not failure. + + Returning False here would block `hard_delete_document` for every legacy + document in the system — the deletion would never complete. + """ + + @pytest.mark.asyncio + async def test_an_absent_record_needs_no_managed_delete( + self, managed: RecordingManagedBackend, monkeypatch + ): + _set_engine(monkeypatch, None) + + assert await _delete() is True + assert managed.deleted == [] + + @pytest.mark.asyncio + async def test_a_record_with_no_engine_needs_no_managed_delete( + self, managed: RecordingManagedBackend, monkeypatch + ): + """A migration in flight: `shadow`/`verify` have not promoted anything.""" + _set_engine(monkeypatch, "absent") + + assert await _delete() is True + assert managed.deleted == [] + + +class TestFailuresKeepTheDocumentRecordAlive: + """The `DOC#` row is the safety net, so a failure must not report success.""" + + @pytest.mark.asyncio + async def test_a_failed_managed_delete_keeps_the_document_record( + self, monkeypatch + ): + backend = RecordingManagedBackend(raises=RuntimeError("AccessDenied")) + import apis.shared.kb_backend.managed_backend as mb + + monkeypatch.setattr(mb, "ManagedKbBackend", lambda *a, **k: backend) + _set_engine(monkeypatch, "managed") + + assert await _delete() is False, ( + "a failed managed delete reported success; the caller would then " + "hard-delete the DOC# row that the fail-closed status filter joins " + "against, leaving indexed content with nothing left to hide it" + ) + assert backend.attempts == 3, "every attempt should be retried" + + @pytest.mark.asyncio + async def test_an_unreadable_record_fails_rather_than_assuming_legacy( + self, managed: RecordingManagedBackend, monkeypatch + ): + """The opposite choice from the ingestion gate, on purpose. + + On ingest an unreadable record resolves to legacy, because the cost of + being wrong is a duplicate index while the consumer still finishes the + document. On delete the cost of being wrong is hard-deleting the row that + keeps a still-indexed chunk unserved, so it fails and retries instead. + """ + _raise_on_lookup(monkeypatch, RuntimeError("DynamoDB unavailable")) + + assert await _delete() is False + assert managed.deleted == [] + + @pytest.mark.asyncio + async def test_a_promoted_kb_with_no_aws_ids_is_not_retried_forever( + self, monkeypatch + ): + """Nothing was ever indexed, so there is nothing to remove.""" + from apis.shared.kb_backend.managed_backend import ManagedKbNotProvisioned + + backend = RecordingManagedBackend(raises=ManagedKbNotProvisioned("no awsKbId")) + import apis.shared.kb_backend.managed_backend as mb + + monkeypatch.setattr(mb, "ManagedKbBackend", lambda *a, **k: backend) + _set_engine(monkeypatch, "managed") + + assert await _delete() is True + assert backend.attempts == 1, "not provisioned is terminal, not transient" + + +class TestTheCleanupContractIncludesManagedDeletion: + @pytest.mark.asyncio + async def test_a_failed_managed_delete_blocks_the_hard_delete( + self, monkeypatch + ): + """End to end through `cleanup_document_resources`, not just the helper. + + The helper returning False is only useful if the caller conjoins it. This + is the assertion that would fail if someone computed `all_succeeded` + without the managed phase. + """ + hard_deleted: List[str] = [] + + async def _no_hard_delete(assistant_id: str, document_id: str) -> None: + hard_deleted.append(document_id) + + async def _ok(*_: Any, **__: Any) -> bool: + return True + + monkeypatch.setattr(cleanup_service, "_delete_vectors_with_retries", _ok) + monkeypatch.setattr(cleanup_service, "_delete_s3_with_retries", _ok) + + async def _managed_fails(*_: Any, **__: Any) -> bool: + return False + + monkeypatch.setattr( + cleanup_service, "_delete_managed_documents_with_retries", _managed_fails + ) + + import apis.app_api.documents.services.document_service as ds + + monkeypatch.setattr(ds, "hard_delete_document", _no_hard_delete) + + result = await cleanup_service.cleanup_document_resources( + document_id=DOCUMENT_ID, + assistant_id=ASSISTANT_ID, + s3_key=f"assistants/{ASSISTANT_ID}/documents/{DOCUMENT_ID}/f.pdf", + chunk_count=3, + ) + + assert result is False + assert hard_deleted == [], ( + "the DOC# row was hard-deleted even though the managed knowledge " + "base still holds the content" + ) diff --git a/backend/tests/fine_tuning/test_admin_routes.py b/backend/tests/fine_tuning/test_admin_routes.py index e56e03c2e..eb6bee181 100644 --- a/backend/tests/fine_tuning/test_admin_routes.py +++ b/backend/tests/fine_tuning/test_admin_routes.py @@ -383,3 +383,106 @@ def _raise_403(): client = TestClient(app) resp = client.get("/admin/fine-tuning/inference-jobs") assert resp.status_code == 403 + + +class TestCostDashboard: + """The dashboard reported $0.00 while real jobs were being billed. + + Two causes, both regression-guarded here: the StatusIndex GSI partition key + is compared case-sensitively and was queried with SageMaker's "Completed" + spelling instead of the stored "COMPLETED"; and FAILED jobs were excluded + even though AWS bills a job that dies partway through. + """ + + @staticmethod + def _job(email, status_value, billable, cost): + return { + "email": email, + "status": status_value, + "billable_seconds": billable, + "estimated_cost_usd": cost, + } + + def _client(self, make_user, training_by_status, inference_by_status=None): + app = _create_app() + _override_auth(app, make_user(email="admin@example.com", roles=["Admin"])) + + jobs_repo = MagicMock() + jobs_repo.query_jobs_by_status_and_date.side_effect = ( + lambda status_value, *_: list(training_by_status.get(status_value, [])) + ) + _override_jobs_repo(app, jobs_repo) + + inf_repo = MagicMock() + inf_repo.query_jobs_by_status_and_date.side_effect = ( + lambda status_value, *_: list((inference_by_status or {}).get(status_value, [])) + ) + _override_inf_repo(app, inf_repo) + + return TestClient(app), jobs_repo, inf_repo + + def test_queries_stored_uppercase_statuses(self, make_user): + """SageMaker's "Completed" casing matches nothing on the GSI.""" + client, jobs_repo, inf_repo = self._client(make_user, {}) + + resp = client.get("/admin/fine-tuning/costs?month=2026-08") + + assert resp.status_code == 200 + for repo in (jobs_repo, inf_repo): + queried = { + call.args[0] for call in repo.query_jobs_by_status_and_date.call_args_list + } + assert queried == {"COMPLETED", "FAILED", "STOPPED"} + + def test_aggregates_completed_jobs(self, make_user): + client, _, _ = self._client( + make_user, + {"COMPLETED": [self._job("user@example.com", "COMPLETED", 300, 0.1175)]}, + ) + + body = client.get("/admin/fine-tuning/costs?month=2026-08").json() + + assert body["total_cost_usd"] == pytest.approx(0.1175) + # Rounded to 2dp by the route: 300s = 0.0833h -> 0.08 + assert body["total_gpu_hours"] == pytest.approx(0.08) + assert body["training_job_count"] == 1 + assert body["active_user_count"] == 1 + assert body["users"][0]["email"] == "user@example.com" + + def test_counts_failed_jobs_because_aws_bills_them(self, make_user): + client, _, _ = self._client( + make_user, + {"FAILED": [self._job("user@example.com", "FAILED", 296, 0.1159)]}, + ) + + body = client.get("/admin/fine-tuning/costs?month=2026-08").json() + + assert body["total_cost_usd"] == pytest.approx(0.1159) + assert body["training_job_count"] == 1 + + def test_sums_training_and_inference_per_user(self, make_user): + client, _, _ = self._client( + make_user, + { + "COMPLETED": [self._job("user@example.com", "COMPLETED", 300, 0.1175)], + "FAILED": [self._job("user@example.com", "FAILED", 296, 0.1159)], + }, + {"COMPLETED": [self._job("user@example.com", "COMPLETED", 230, 0.09)]}, + ) + + body = client.get("/admin/fine-tuning/costs?month=2026-08").json() + + assert body["total_cost_usd"] == pytest.approx(0.1175 + 0.1159 + 0.09) + assert body["training_job_count"] == 2 + assert body["inference_job_count"] == 1 + assert body["active_user_count"] == 1 + + def test_requires_admin_role(self): + app = _create_app() + + def _raise_403(): + raise HTTPException(status_code=403, detail="Forbidden") + override_admin_auth(app, _raise_403) + + resp = TestClient(app).get("/admin/fine-tuning/costs") + assert resp.status_code == 403 diff --git a/backend/tests/fine_tuning/test_inference_routes.py b/backend/tests/fine_tuning/test_inference_routes.py index e1ef33ab5..789e33c6f 100644 --- a/backend/tests/fine_tuning/test_inference_routes.py +++ b/backend/tests/fine_tuning/test_inference_routes.py @@ -212,6 +212,36 @@ def test_returns_201_on_success(self, make_user): assert body["job_type"] == "inference" assert body["training_job_id"] == "train-abc123" + def test_rejects_instance_type_with_no_known_price(self, make_user): + """Same blind spot as the training path: unpriced means $0.00 recorded.""" + app = _create_app() + user = make_user(email="user@example.com") + + mock_jobs = MagicMock() + mock_jobs.get_job.return_value = SAMPLE_COMPLETED_TRAINING_JOB + + mock_s3 = MagicMock() + mock_s3.check_object_exists.return_value = True + mock_s3.bucket_name = "test-bucket" + + mock_sm = MagicMock() + + _setup_deps(app, user, SAMPLE_GRANT, mock_jobs, MagicMock(), mock_s3, mock_sm, MagicMock()) + + client = TestClient(app) + resp = client.post( + "/fine-tuning/inference", + json={ + "training_job_id": "train-abc123", + "input_s3_key": "inference-input/user-001/xyz/input.txt", + "instance_type": "ml.p4d.24xlarge", + }, + ) + + assert resp.status_code == 400 + assert "Unsupported instance type" in resp.json()["detail"] + mock_sm.create_transform_job.assert_not_called() + @patch.dict("os.environ", {"PROJECT_PREFIX": "test-prefix"}) def test_transform_job_name_includes_project_prefix(self, make_user): app = _create_app() diff --git a/backend/tests/fine_tuning/test_job_repository.py b/backend/tests/fine_tuning/test_job_repository.py index dfa6fbd7d..f00dfd0ec 100644 --- a/backend/tests/fine_tuning/test_job_repository.py +++ b/backend/tests/fine_tuning/test_job_repository.py @@ -283,3 +283,108 @@ def test_deletes_item(self, jobs_repository): def test_returns_false_for_nonexistent(self, jobs_repository): assert jobs_repository.delete_job("user-001", "nonexistent") is False + + +class TestQueryJobsByStatusAndDate: + """The cost dashboard's GSI query. + + Training and inference records share one table and one StatusIndex, so this + query has to return training rows only — otherwise a caller that also + queries the inference repository counts an inference job's cost twice. + """ + + def _completed_training_job(self, jobs_repository, job_id): + jobs_repository.create_job( + user_id="user-001", + email="alice@example.com", + job_id=job_id, + model_id="electra-tiny", + model_name="ELECTRA Tiny", + dataset_s3_key="datasets/user-001/abc/train.csv", + instance_type="ml.g5.xlarge", + hyperparameters={"epochs": "3"}, + sagemaker_job_name=f"ft-{job_id[:8]}", + output_s3_prefix=f"output/user-001/{job_id}", + ) + jobs_repository.update_job_status( + user_id="user-001", + job_id=job_id, + status="COMPLETED", + billable_seconds=300, + estimated_cost_usd=0.1175, + ) + + def _completed_inference_job(self, inference_repository, job_id): + inference_repository.create_inference_job( + user_id="user-001", + email="alice@example.com", + job_id=job_id, + training_job_id="train-001", + model_name="ELECTRA Tiny", + model_s3_path="s3://bucket/model.tar.gz", + input_s3_key="inference-input/user-001/in.txt", + instance_type="ml.g5.xlarge", + transform_job_name=f"inf-{job_id[:8]}", + output_s3_prefix=f"inference-output/user-001/{job_id}", + ) + inference_repository.update_inference_status( + user_id="user-001", + job_id=job_id, + status="COMPLETED", + billable_seconds=230, + estimated_cost_usd=0.0901, + ) + + @staticmethod + def _range(): + return "2000-01-01T00:00:00+00:00", "2100-01-01T00:00:00+00:00" + + def test_returns_matching_training_job(self, jobs_repository): + job_id = uuid.uuid4().hex + self._completed_training_job(jobs_repository, job_id) + start, end = self._range() + + results = jobs_repository.query_jobs_by_status_and_date("COMPLETED", start, end) + + assert [j["job_id"] for j in results] == [job_id] + + def test_stored_status_casing_is_what_matches(self, jobs_repository): + """SageMaker spells it "Completed"; the record stores "COMPLETED".""" + self._completed_training_job(jobs_repository, uuid.uuid4().hex) + start, end = self._range() + + assert jobs_repository.query_jobs_by_status_and_date("Completed", start, end) == [] + assert len(jobs_repository.query_jobs_by_status_and_date("COMPLETED", start, end)) == 1 + + def test_excludes_inference_jobs(self, jobs_repository, inference_repository): + training_id = uuid.uuid4().hex + self._completed_training_job(jobs_repository, training_id) + self._completed_inference_job(inference_repository, uuid.uuid4().hex) + start, end = self._range() + + results = jobs_repository.query_jobs_by_status_and_date("COMPLETED", start, end) + + assert [j["job_id"] for j in results] == [training_id] + + def test_inference_query_excludes_training_jobs( + self, jobs_repository, inference_repository + ): + self._completed_training_job(jobs_repository, uuid.uuid4().hex) + inference_id = uuid.uuid4().hex + self._completed_inference_job(inference_repository, inference_id) + start, end = self._range() + + results = inference_repository.query_jobs_by_status_and_date( + "COMPLETED", start, end + ) + + assert [j["job_id"] for j in results] == [inference_id] + + def test_excludes_jobs_outside_the_period(self, jobs_repository): + self._completed_training_job(jobs_repository, uuid.uuid4().hex) + + results = jobs_repository.query_jobs_by_status_and_date( + "COMPLETED", "2000-01-01T00:00:00+00:00", "2000-02-01T00:00:00+00:00" + ) + + assert results == [] diff --git a/backend/tests/fine_tuning/test_job_routes.py b/backend/tests/fine_tuning/test_job_routes.py index a3d91db81..4c1f8a892 100644 --- a/backend/tests/fine_tuning/test_job_routes.py +++ b/backend/tests/fine_tuning/test_job_routes.py @@ -119,6 +119,25 @@ def test_returns_200_with_presigned_url(self, make_user): assert "s3_key" in body assert "expires_at" in body + @pytest.mark.parametrize("filename", ["notes.txt", "data.parquet"]) + def test_rejects_format_the_trainer_cannot_read(self, make_user, filename): + """Reject before upload, not several billed GPU-minutes into training.""" + app = _create_app() + user = make_user(email="user@example.com") + + mock_s3 = MagicMock() + _setup_deps(app, user, SAMPLE_GRANT, s3_service=mock_s3) + + client = TestClient(app) + resp = client.post( + "/fine-tuning/presign", + json={"filename": filename, "content_type": "text/plain"}, + ) + + assert resp.status_code == 400 + assert "Unsupported dataset format" in resp.json()["detail"] + mock_s3.generate_upload_url.assert_not_called() + class TestCreateJob: @@ -159,6 +178,99 @@ def test_returns_201_on_success(self, make_user): body = resp.json() assert body["model_id"] == "distilgpt2" + def test_rejects_dataset_the_trainer_cannot_read(self, make_user): + """Last gate before SageMaker: no GPU is provisioned for a doomed job.""" + app = _create_app() + user = make_user(email="user@example.com") + + mock_s3 = MagicMock() + mock_s3.check_object_exists.return_value = True + + mock_sm = MagicMock() + + _setup_deps(app, user, SAMPLE_GRANT, s3_service=mock_s3, sagemaker=mock_sm) + + client = TestClient(app) + resp = client.post( + "/fine-tuning/jobs", + json={ + "model_id": "distilgpt2", + "dataset_s3_key": "datasets/user-001/abc/notes.txt", + }, + ) + + assert resp.status_code == 400 + assert "Unsupported dataset format" in resp.json()["detail"] + mock_sm.create_training_job.assert_not_called() + + def test_rejects_instance_type_with_no_known_price(self, make_user): + """An unpriced instance runs real GPUs and records $0.00 spend. + + calculate_cost falls back to 0.0/hour for anything absent from + INSTANCE_COST_PER_HOUR, and quota meters GPU-hours rather than + dollars, so nothing downstream bounds the cost. + """ + app = _create_app() + user = make_user(email="user@example.com") + + mock_s3 = MagicMock() + mock_s3.check_object_exists.return_value = True + + mock_sm = MagicMock() + + _setup_deps(app, user, SAMPLE_GRANT, s3_service=mock_s3, sagemaker=mock_sm) + + client = TestClient(app) + resp = client.post( + "/fine-tuning/jobs", + json={ + "model_id": "distilgpt2", + "dataset_s3_key": "datasets/user-001/abc/train.csv", + "instance_type": "ml.p4d.24xlarge", + }, + ) + + assert resp.status_code == 400 + assert "Unsupported instance type" in resp.json()["detail"] + mock_sm.create_training_job.assert_not_called() + + def test_accepts_a_priced_instance_type(self, make_user): + app = _create_app() + user = make_user(email="user@example.com") + + mock_s3 = MagicMock() + mock_s3.check_object_exists.return_value = True + mock_s3.get_output_s3_prefix.return_value = "output/user-001/job-abc" + mock_s3.get_output_s3_uri.return_value = "s3://bucket/output/user-001/job-abc" + mock_s3.bucket_name = "test-bucket" + + mock_jobs = MagicMock() + mock_jobs.create_job.return_value = SAMPLE_JOB + mock_jobs.update_job_status.return_value = {**SAMPLE_JOB, "status": "TRAINING"} + + mock_sm = MagicMock() + mock_sm.create_training_job.return_value = {} + + mock_script = MagicMock() + mock_script.ensure_scripts_uploaded.return_value = "s3://test-bucket/scripts/sourcedir.tar.gz" + + _setup_deps( + app, user, SAMPLE_GRANT, mock_jobs, mock_s3, mock_sm, + MagicMock(), mock_script, + ) + + client = TestClient(app) + resp = client.post( + "/fine-tuning/jobs", + json={ + "model_id": "distilgpt2", + "dataset_s3_key": "datasets/user-001/abc/train.csv", + "instance_type": "ml.g5.2xlarge", + }, + ) + + assert resp.status_code == 201 + @patch.dict("os.environ", {"PROJECT_PREFIX": "test-prefix"}) def test_sagemaker_job_name_includes_project_prefix(self, make_user): app = _create_app() diff --git a/backend/tests/fine_tuning/test_train_script.py b/backend/tests/fine_tuning/test_train_script.py index e62896337..6731b34c3 100644 --- a/backend/tests/fine_tuning/test_train_script.py +++ b/backend/tests/fine_tuning/test_train_script.py @@ -6,7 +6,11 @@ from apis.app_api.fine_tuning.sagemaker_scripts.train import ( resolve_max_context_length, - find_csv_in_channel, + find_dataset_in_channel, + load_dataset_frame, + resolve_dataset_reader, + validate_dataset_columns, + SUPPORTED_DATASET_EXTENSIONS, copy_inference_script, DynamoDBProgressCallback, SageMakerLoggingCallback, @@ -59,32 +63,117 @@ def test_uses_model_max_length_as_fallback(self): assert result == 768 -class TestFindCsvInChannel: +class TestFindDatasetInChannel: def test_finds_csv_file(self, tmp_path): csv_file = tmp_path / "dataset.csv" csv_file.write_text("text,label\nhello,1\n") - result = find_csv_in_channel(str(tmp_path)) + result = find_dataset_in_channel(str(tmp_path)) assert result == str(csv_file) - def test_raises_when_no_csv(self, tmp_path): + @pytest.mark.parametrize("filename", ["dataset.jsonl", "dataset.json"]) + def test_finds_json_formats(self, tmp_path, filename): + """The UI offers JSONL/JSON, so the trainer has to find them too.""" + dataset = tmp_path / filename + dataset.write_text('{"text": "hello", "label": "a"}\n') + + result = find_dataset_in_channel(str(tmp_path)) + assert result == str(dataset) + + def test_raises_when_no_supported_dataset(self, tmp_path): txt_file = tmp_path / "readme.txt" - txt_file.write_text("not a csv") + txt_file.write_text("not a dataset") - with pytest.raises(FileNotFoundError, match="No CSV file found"): - find_csv_in_channel(str(tmp_path)) + with pytest.raises(FileNotFoundError, match="No dataset file found"): + find_dataset_in_channel(str(tmp_path)) def test_case_insensitive_extension(self, tmp_path): csv_file = tmp_path / "DATA.CSV" csv_file.write_text("text,label\nhello,1\n") - result = find_csv_in_channel(str(tmp_path)) + result = find_dataset_in_channel(str(tmp_path)) assert result == str(csv_file) def test_raises_when_dir_missing(self): with pytest.raises(FileNotFoundError, match="does not exist"): - find_csv_in_channel("/nonexistent/path") + find_dataset_in_channel("/nonexistent/path") + + +class TestResolveDatasetReader: + """Every format the upload UI accepts must actually be readable. + + A JSONL dataset previously uploaded and dispatched fine, then died on the + GPU several billed minutes in because the trainer only read CSV. These + assert the dispatch table directly so they run without pandas, which + exists only inside the SageMaker training container. + """ + + def test_supports_the_formats_the_ui_offers(self): + assert set(SUPPORTED_DATASET_EXTENSIONS) == {".csv", ".jsonl", ".json"} + + def test_csv_uses_read_csv(self): + assert resolve_dataset_reader("/data/dataset.csv") == ("read_csv", {}) + + def test_jsonl_reads_line_delimited(self): + assert resolve_dataset_reader("/data/dataset.jsonl") == ( + "read_json", + {"lines": True}, + ) + + def test_json_reads_whole_document(self): + assert resolve_dataset_reader("/data/dataset.json") == ("read_json", {}) + + def test_extension_match_is_case_insensitive(self): + assert resolve_dataset_reader("/data/DATA.CSV") == ("read_csv", {}) + + def test_raises_on_unsupported_extension(self): + with pytest.raises(ValueError, match="Unsupported dataset format"): + resolve_dataset_reader("/data/dataset.parquet") + + +class TestValidateDatasetColumns: + + def test_accepts_required_columns(self): + validate_dataset_columns(["text", "label"], "/data/dataset.csv") + + def test_raises_when_label_missing(self): + with pytest.raises(ValueError, match="missing required column"): + validate_dataset_columns(["text"], "/data/dataset.csv") + + def test_raises_when_text_missing(self): + with pytest.raises(ValueError, match="missing required column"): + validate_dataset_columns(["label"], "/data/dataset.csv") + + +class TestLoadDatasetFrame: + """End-to-end load, where pandas is available (the training container).""" + + @pytest.mark.parametrize( + "filename,content", + [ + ("dataset.csv", "text,label\nhello,positive\nbye,negative\n"), + ( + "dataset.jsonl", + '{"text": "hello", "label": "positive"}\n' + '{"text": "bye", "label": "negative"}\n', + ), + ( + "dataset.json", + '[{"text": "hello", "label": "positive"},' + ' {"text": "bye", "label": "negative"}]', + ), + ], + ) + def test_loads_each_supported_format(self, tmp_path, filename, content): + pytest.importorskip("pandas") + path = tmp_path / filename + path.write_text(content) + + df = load_dataset_frame(str(path)) + + assert df["text"].tolist() == ["hello", "bye"] + assert df["label"].tolist() == ["positive", "negative"] class TestCopyInferenceScript: diff --git a/backend/tests/ingestion/test_ingestion_engine_gate.py b/backend/tests/ingestion/test_ingestion_engine_gate.py new file mode 100644 index 000000000..65715f2a5 --- /dev/null +++ b/backend/tests/ingestion/test_ingestion_engine_gate.py @@ -0,0 +1,252 @@ +"""The legacy ingestion pipeline must not touch a promoted knowledge base. + +Routing exclusivity (design §537, Requirement 10.5). Both pipelines are triggered +by the same S3 upload — the legacy `s3:ObjectCreated` notification and the managed +consumer's EventBridge rule — so exactly one of them has to stand down per +document. The consumer already returns immediately for a legacy document; these +tests cover the half that was missing, which is this pipeline standing down for a +*managed* one. + +WHY THE STATUS HALF MATTERS MORE THAN THE DUPLICATE VECTORS +Two writers owning one `status` field means the last writer wins by luck. Both +outcomes were observed in dev before this gate existed: + +* A PDF was marked `complete` by this pipeline 65 s before the managed knowledge + base could answer for it — "your document is ready", then an answer that does + not mention it. +* An image-only PDF that Docling could not parse at all (`Docling produced zero + chunks`) was marked `failed` while the managed knowledge base was serving it + correctly. That one only read `complete` in the end because the managed + consumer happened to finish second and overwrite it. Reverse the finishing + order — entirely a matter of document size and parse time — and a good, + retrievable document reads `failed` forever, with no way to retry it. + +So the assertions below are mostly about what is *not* written. +""" + +from __future__ import annotations + +import json +import sys +import types +from typing import Any, Dict, List, Optional + +import pytest + +from apis.app_api.documents.ingestion import handler as handler_module + +ASSISTANT_ID = "ast-gate-1" +DOCUMENT_ID = "DOC-gate-1" + + +class RecordingStatusManager: + """Captures every status transition the handler attempts.""" + + def __init__(self) -> None: + self.calls: List[str] = [] + + async def mark_chunking(self, **_: Any) -> None: + self.calls.append("chunking") + + async def mark_embedding(self, **_: Any) -> None: + self.calls.append("embedding") + + async def mark_complete(self, **_: Any) -> None: + self.calls.append("complete") + + async def mark_failed(self, **_: Any) -> None: + self.calls.append("failed") + + +@pytest.fixture +def status_manager(monkeypatch: pytest.MonkeyPatch) -> RecordingStatusManager: + """Stand in for the `status` module, which only resolves inside the image. + + `handler.py` does `from status import create_status_manager` as a bare + top-level import because the Dockerfile flattens `documents/ingestion/` onto + LAMBDA_TASK_ROOT. Injecting the module is how a test calls the handler + without reproducing that layout. + """ + recorder = RecordingStatusManager() + fake = types.ModuleType("status") + fake.create_status_manager = lambda: recorder # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "status", fake) + return recorder + + +@pytest.fixture +def no_real_pipeline(monkeypatch: pytest.MonkeyPatch) -> List[str]: + """Replace the Docling/embedding pipeline; record whether it was reached.""" + reached: List[str] = [] + + async def _fake_pipeline(**kwargs: Any) -> None: + reached.append(kwargs.get("document_id", "?")) + + monkeypatch.setattr(handler_module, "_process_document_pipeline", _fake_pipeline) + return reached + + +def _kb_record(engine: Optional[str]) -> Dict[str, Any]: + record: Dict[str, Any] = {"PK": f"AST#{ASSISTANT_ID}", "SK": f"KB#{ASSISTANT_ID}"} + if engine is not None: + record["retrievalEngine"] = engine + return record + + +def _set_record(monkeypatch: pytest.MonkeyPatch, record) -> None: + """Point `records.get_kb_record` at a fixed answer, or make it raise.""" + from apis.shared.kb_backend import records as r + + def _get(_assistant_id: str, _app_kb_id: str): + if isinstance(record, Exception): + raise record + return record + + monkeypatch.setattr(r, "get_kb_record", _get) + + +def _event() -> Dict[str, Any]: + return { + "Records": [ + { + "s3": { + "bucket": {"name": "docs-bucket"}, + "object": { + "key": ( + f"assistants/{ASSISTANT_ID}/documents/" + f"{DOCUMENT_ID}/flowchart.pdf" + ) + }, + } + } + ] + } + + +async def _invoke() -> Dict[str, Any]: + return await handler_module.async_lambda_handler(_event(), None) + + +class TestAPromotedKnowledgeBaseIsLeftAlone: + @pytest.mark.asyncio + async def test_no_document_status_is_written_at_all( + self, status_manager: RecordingStatusManager, no_real_pipeline, monkeypatch + ): + """The whole point. `complete` and `failed` both belong to the consumer.""" + _set_record(monkeypatch, _kb_record("managed")) + + response = await _invoke() + + assert status_manager.calls == [], ( + "the legacy pipeline wrote document status for a knowledge base served " + "by the managed engine; two writers on one field is how a good " + "document ends up reading 'failed'" + ) + assert response["statusCode"] == 200 + assert "Skipped" in json.loads(response["body"])["message"] + + @pytest.mark.asyncio + async def test_the_document_is_not_parsed_or_embedded( + self, status_manager, no_real_pipeline: List[str], monkeypatch + ): + """No Docling time and no S3 Vectors storage for vectors nothing reads.""" + _set_record(monkeypatch, _kb_record("managed")) + + await _invoke() + + assert no_real_pipeline == [] + + @pytest.mark.asyncio + async def test_it_returns_success_so_the_event_is_not_retried( + self, status_manager, no_real_pipeline, monkeypatch + ): + """A skip is a correct outcome, not a failure to redeliver.""" + _set_record(monkeypatch, _kb_record("managed")) + + response = await _invoke() + + assert response["statusCode"] == 200 + assert json.loads(response["body"])["engine"] == "managed" + + +class TestEveryOtherKnowledgeBaseStillRuns: + """The gate keys on the ENGINE, not on the presence of a record. + + During `shadow` and `verify` the legacy path is still authoritative and must + keep working (Requirements 16.1, 16.6). Only `promote` writes + `retrievalEngine`, so a record that exists but names no engine — which is + exactly a migration in flight — must not be skipped. + """ + + @pytest.mark.asyncio + async def test_a_record_with_no_engine_still_runs( + self, status_manager: RecordingStatusManager, no_real_pipeline: List[str], monkeypatch + ): + _set_record(monkeypatch, _kb_record(None)) + + await _invoke() + + assert status_manager.calls == ["chunking"] + assert no_real_pipeline == [DOCUMENT_ID] + + @pytest.mark.asyncio + async def test_a_migration_in_flight_still_runs( + self, status_manager: RecordingStatusManager, no_real_pipeline: List[str], monkeypatch + ): + """`shadow` is not `promoted`. Skipping here would strand live uploads.""" + record = _kb_record(None) + record["migrationState"] = "shadow" + _set_record(monkeypatch, record) + + await _invoke() + + assert status_manager.calls == ["chunking"] + assert no_real_pipeline == [DOCUMENT_ID] + + @pytest.mark.asyncio + async def test_no_record_at_all_still_runs( + self, status_manager: RecordingStatusManager, no_real_pipeline: List[str], monkeypatch + ): + """Every knowledge base that predates this feature. Absence ⇒ legacy.""" + _set_record(monkeypatch, None) + + await _invoke() + + assert status_manager.calls == ["chunking"] + assert no_real_pipeline == [DOCUMENT_ID] + + +class TestAnUnreadableRecordRunsTheLegacyPipeline: + """Fail towards legacy, deliberately — the two errors are not symmetric. + + Wrong towards legacy on a promoted knowledge base costs a duplicate index and + a status race, and the managed consumer still drives the document to a correct + terminal state. Wrong towards skipping on a legacy knowledge base means + nothing indexes the document at all: it sits un-ingested with no error, and + the only way out is a re-upload. The second is much worse. + """ + + @pytest.mark.asyncio + async def test_a_lookup_failure_does_not_skip_the_document( + self, status_manager: RecordingStatusManager, no_real_pipeline: List[str], monkeypatch + ): + _set_record(monkeypatch, RuntimeError("DynamoDB unavailable")) + + await _invoke() + + assert status_manager.calls == ["chunking"], ( + "an unreadable KB record skipped the document; a transient read " + "failure must not silently leave an upload un-ingested" + ) + assert no_real_pipeline == [DOCUMENT_ID] + + @pytest.mark.asyncio + async def test_a_lookup_failure_is_not_reported_as_a_document_failure( + self, status_manager: RecordingStatusManager, no_real_pipeline, monkeypatch + ): + """The user's document is fine; our read of an unrelated row was not.""" + _set_record(monkeypatch, RuntimeError("DynamoDB unavailable")) + + await _invoke() + + assert "failed" not in status_manager.calls diff --git a/backend/tests/lambdas/test_kb_ingestion_consumer.py b/backend/tests/lambdas/test_kb_ingestion_consumer.py new file mode 100644 index 000000000..cf0d07fed --- /dev/null +++ b/backend/tests/lambdas/test_kb_ingestion_consumer.py @@ -0,0 +1,562 @@ +"""Routing exclusivity for the managed-KB ingestion consumer. + +Feature: managed-kb-migration, task 9.2. + +The failure this file exists to prevent is **double indexing**. The legacy pipeline +is driven by its own pre-existing S3 notification on the same bucket, so for a legacy +document the correct behaviour of this consumer is to do nothing whatsoever. If it +ingested as well, the same bytes would be embedded twice: two sets of vectors, +doubled ingestion cost, and duplicate chunks competing inside one result list. None +of that raises an error, which is exactly why it needs a test. + +The routing is therefore deliberately asymmetric, and both halves are asserted: +legacy must ingest NOTHING here, managed must ingest here and NOT fall back. +""" + +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch + +import boto3 +import pytest +from moto import mock_aws + +from apis.app_api.kb_migration import ingestion_consumer as ic + +REGION = "us-east-1" +TABLE = "test-ingestion-consumer" +ASSISTANT_ID = "ast-ing01" +DOCUMENT_ID = "doc-ing01" +BUCKET = "docs-bucket" +KEY = f"assistants/{ASSISTANT_ID}/documents/{DOCUMENT_ID}/report.pdf" + + +@pytest.fixture() +def table(monkeypatch): + monkeypatch.setenv("AWS_DEFAULT_REGION", REGION) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing") + monkeypatch.setenv("AWS_SESSION_TOKEN", "testing") + monkeypatch.setenv("DYNAMODB_ASSISTANTS_TABLE_NAME", TABLE) + + with mock_aws(): + ddb = boto3.client("dynamodb", region_name=REGION) + ddb.create_table( + TableName=TABLE, + KeySchema=[ + {"AttributeName": "PK", "KeyType": "HASH"}, + {"AttributeName": "SK", "KeyType": "RANGE"}, + ], + AttributeDefinitions=[ + {"AttributeName": "PK", "AttributeType": "S"}, + {"AttributeName": "SK", "AttributeType": "S"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + t = boto3.resource("dynamodb", region_name=REGION).Table(TABLE) + t.put_item( + Item={ + "PK": f"AST#{ASSISTANT_ID}", + "SK": f"DOC#{DOCUMENT_ID}", + "status": "uploading", + } + ) + yield t + + +def _seed_kb(table, **overrides): + """A KB_Record for this assistant. No retrievalEngine unless asked.""" + item = {"PK": f"AST#{ASSISTANT_ID}", "SK": f"KB#{ASSISTANT_ID}", "appKbId": ASSISTANT_ID} + item.update(overrides) + table.put_item(Item=item) + + +def _doc(table): + return table.get_item( + Key={"PK": f"AST#{ASSISTANT_ID}", "SK": f"DOC#{DOCUMENT_ID}"} + )["Item"] + + +def _eventbridge_event(key=KEY): + return {"detail": {"bucket": {"name": BUCKET}, "object": {"key": key}}} + + +class _FakeAgentClient: + """The control-plane surface the consumer reads document status from.""" + + def __init__(self, owner): + self._owner = owner + + def get_knowledge_base_documents(self, **kwargs): + self._owner.status_calls += 1 + status = self._owner.next_status() + if status == "NOT_FOUND": + return {"documentDetails": []} + return { + "documentDetails": [ + { + "knowledgeBaseId": "KB123", + "dataSourceId": "DS456", + "status": status, + "identifier": { + "dataSourceType": "CUSTOM", + "custom": {"id": DOCUMENT_ID}, + }, + "updatedAt": datetime(2026, 9, 1, 14, 53, 19, tzinfo=timezone.utc), + } + ] + } + + +class _FakeBackend: + """Models the parts of ManagedKbBackend the consumer actually leans on. + + Deliberately models Bedrock's document STATUS, not just the ingest call. + Ingestion is fire-and-forget: the API returns once the request is accepted and + says nothing about progress, so a fake that only recorded ingests could not + express the state the consumer now has to reason about — and a fake that + reported instant success is what let the 30 s poll window look adequate for + documents that take minutes. + + ``statuses`` is the sequence returned by successive status probes. The default + models a small document: unknown, then indexed. Pass a longer sequence to model + a slow one; the last value repeats forever. + """ + + def __init__(self, statuses=None): + self.ingested = [] + self.status_calls = 0 + self._statuses = list(statuses or ["NOT_FOUND", "INDEXED"]) + self._agent_client = _FakeAgentClient(self) + + def next_status(self): + if len(self._statuses) > 1: + return self._statuses.pop(0) + return self._statuses[0] + + # -- the private surface `document_status` reuses -------------------------- + def _agent(self): + return self._agent_client + + def _locate(self, kb_ref): + return ("KB123", "DS456") + + # -- the protocol surface -------------------------------------------------- + async def ingest(self, kb_ref, source): + self.ingested.append(source.document_id) + return None + + async def search(self, kb_ref, query, top_k=5): + chunk = MagicMock() + chunk.metadata = {"document_id": DOCUMENT_ID} + return [chunk] + + +# --------------------------------------------------------------------------- +# Legacy must not be touched +# --------------------------------------------------------------------------- +class TestLegacyRouting: + def test_a_legacy_document_is_not_ingested_here(self, table): + """No retrievalEngine means legacy, and legacy is somebody else's job.""" + _seed_kb(table) + fake = _FakeBackend() + + with patch("apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake): + result = ic.handle_object(BUCKET, KEY) + + assert result["routed"] == "legacy" + assert result["ingested"] is False + assert fake.ingested == [], "a legacy document was ingested into the managed backend" + + def test_a_document_with_no_kb_record_at_all_is_legacy(self, table): + """The overwhelmingly common case today: no record has ever been written.""" + result = ic.handle_object(BUCKET, KEY) + assert result["routed"] == "legacy" + assert result["ingested"] is False + + def test_a_legacy_document_status_is_left_alone(self, table): + """The legacy pipeline owns the terminal transition for its documents. + + Writing `complete` here would race the other Lambda and could mark a + document ready before its vectors exist. + """ + _seed_kb(table) + ic.handle_object(BUCKET, KEY) + assert _doc(table)["status"] == "uploading" + + @pytest.mark.parametrize("engine", ["s3vectors", "S3Vectors", "MANAGED", "managed ", "", "wat"]) + def test_only_the_exact_managed_literal_routes_to_managed(self, table, engine): + """Exact-match, so a casing slip fails safe. + + Failing safe matters asymmetrically: routing to legacy when it should be + managed leaves the existing pipeline handling it correctly, while routing to + managed when the record is not really migrated ingests into a knowledge base + that may not exist. + """ + _seed_kb(table, retrievalEngine=engine) + fake = _FakeBackend() + with patch("apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake): + result = ic.handle_object(BUCKET, KEY) + assert result["routed"] == "legacy" + assert fake.ingested == [] + + +# --------------------------------------------------------------------------- +# Managed must be ingested here, exactly once +# --------------------------------------------------------------------------- +class TestManagedRouting: + def _seed_managed(self, table): + _seed_kb( + table, + retrievalEngine="managed", + awsKbId="KB123", + awsDataSourceId="DS456", + ) + + def test_a_managed_document_is_ingested_directly(self, table): + self._seed_managed(table) + fake = _FakeBackend() + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + result = ic.handle_object(BUCKET, KEY) + + assert result["routed"] == "managed" + assert result["ingested"] is True + assert fake.ingested == [DOCUMENT_ID] + + def test_a_managed_document_is_ingested_exactly_once(self, table): + """One invocation, one ingest. Duplicate chunks would compete in retrieval.""" + self._seed_managed(table) + fake = _FakeBackend() + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + ic.lambda_handler(_eventbridge_event(), None) + + assert fake.ingested == [DOCUMENT_ID] + + def test_the_document_reaches_complete(self, table): + self._seed_managed(table) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", + return_value=_FakeBackend(), + ): + ic.handle_object(BUCKET, KEY) + + assert _doc(table)["status"] == "complete" + + def test_indexed_and_retrievable_are_recorded_separately(self, table): + """Two timestamps, not one. + + Bedrock reports INDEXED up to a second before a document can actually be + retrieved (measured 0.75-1.03 s). Collapsing them would erase the only + evidence of that gap, which is what makes "my upload finished but the + assistant cannot see it" diagnosable. + """ + self._seed_managed(table) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", + return_value=_FakeBackend(), + ): + result = ic.handle_object(BUCKET, KEY) + + item = _doc(table) + assert "indexedAt" in item + assert "retrievableAt" in item + assert result["indexedAt"] and result["retrievableAt"] + + def test_indexed_at_is_bedrocks_timestamp_not_our_clock(self, table): + """`indexedAt` must be the value Bedrock reports, not the local time. + + The original code set `indexed_at = _now_iso()` immediately after the + ingest call returned. That call is fire-and-forget — it returns when the + request is accepted — so the field recorded "when we asked", presented as + "when indexing finished". For a 1.5 MB PDF in dev those were 5.5 minutes + apart, and because the field was always populated the error was invisible: + the pre-existing test asserted only that the key existed and was truthy, + which a fabricated value satisfies perfectly. + """ + self._seed_managed(table) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", + return_value=_FakeBackend(), + ): + result = ic.handle_object(BUCKET, KEY) + + # The fake reports 2026-09-01T14:53:19+00:00 from GetKnowledgeBaseDocuments. + assert result["indexedAt"].startswith("2026-09-01T14:53:19"), ( + f"indexedAt is {result['indexedAt']!r}, which is not the timestamp " + f"Bedrock reported — it looks like a local clock reading" + ) + assert _doc(table)["indexedAt"].startswith("2026-09-01T14:53:19") + + def test_a_managed_document_never_falls_back_to_legacy(self, table): + """Managed engine but unprovisioned must FAIL, not silently degrade. + + A quiet fallback would hand the document to the legacy pipeline as well, + producing the dual index this whole file guards against. + """ + _seed_kb(table, retrievalEngine="managed") # no awsKbId / awsDataSourceId + + with pytest.raises(ic.IngestionRoutingError, match="not provisioned"): + ic.handle_object(BUCKET, KEY) + + def test_a_failed_ingestion_marks_the_document_failed_and_raises(self, table): + """The record is the retry anchor, so a failure must be visible in both + places: on the document and to the event source.""" + self._seed_managed(table) + + class _Failing(_FakeBackend): + async def ingest(self, kb_ref, source): + raise RuntimeError("bedrock unavailable") + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=_Failing() + ): + with pytest.raises(RuntimeError): + ic.handle_object(BUCKET, KEY) + + item = _doc(table) + assert item["status"] == "failed" + assert "bedrock unavailable" in item["ingestionError"] + + def test_a_document_that_never_becomes_retrievable_is_not_marked_complete(self, table): + """Indexed is not retrievable. Claiming success here is the bug.""" + self._seed_managed(table) + + class _NeverRetrievable(_FakeBackend): + async def search(self, kb_ref, query, top_k=5): + return [] + + # Shrink the poll window: the real 30s default is correct in production + # (the observed gap is ~1s and waiting is cheap) but would add 30s to every + # run of this suite. + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", + return_value=_NeverRetrievable(), + ), patch.object(ic, "RETRIEVABLE_POLL_TIMEOUT_SECONDS", 0.05), patch.object( + ic, "RETRIEVABLE_POLL_INTERVAL_SECONDS", 0.01 + ): + with pytest.raises(ic.IngestionRoutingError, match="not retrievable"): + ic.handle_object(BUCKET, KEY) + + assert _doc(table)["status"] != "complete" + + +# --------------------------------------------------------------------------- +# Event parsing +# --------------------------------------------------------------------------- +class TestEventParsing: + def test_eventbridge_shape_is_understood(self): + records = ic.extract_records(_eventbridge_event()) + assert records == [{"bucket": BUCKET, "key": KEY}] + + def test_raw_s3_notification_shape_is_understood(self): + """Both shapes are accepted so a wiring change cannot silently stop + ingestion — the bucket carries two producers.""" + event = {"Records": [{"s3": {"bucket": {"name": BUCKET}, "object": {"key": KEY}}}]} + assert ic.extract_records(event) == [{"bucket": BUCKET, "key": KEY}] + + def test_an_empty_event_is_a_no_op(self): + assert ic.lambda_handler({}, None)["processed"] == 0 + + def test_a_url_encoded_key_is_decoded(self): + a, d, f = ic.parse_object_key( + "assistants/ast-1/documents/doc-2/my+report+%282024%29.pdf" + ) + assert (a, d) == ("ast-1", "doc-2") + assert f == "my report (2024).pdf" + + def test_a_filename_containing_slashes_is_preserved(self): + _, _, f = ic.parse_object_key("assistants/a/documents/d/sub/dir/file.pdf") + assert f == "sub/dir/file.pdf" + + @pytest.mark.parametrize( + "key", + [ + "wrong/ast-1/documents/doc-2/f.pdf", + "assistants/ast-1/wrong/doc-2/f.pdf", + "assistants/ast-1/documents/doc-2", + "", + ], + ) + def test_a_malformed_key_is_refused(self, key): + """Guessing at a malformed key could ingest one assistant's document into + another's knowledge base.""" + with pytest.raises(ic.IngestionRoutingError): + ic.parse_object_key(key) + + +# --------------------------------------------------------------------------- +# Structural guarantees +# --------------------------------------------------------------------------- +class TestNoInProcessOrchestration: + def test_the_module_does_not_use_ensure_future(self): + """Requirement 10.8. A background task is killed when the Lambda handler + returns, converting a reported success into a half-finished ingestion.""" + import ast + import inspect + + # Parsed, not grepped. A substring check trips on this module's own + # docstring, which explains at length WHY it does not orchestrate in + # process — the first version of this test failed on the prose describing + # the very thing it was verifying the absence of. + tree = ast.parse(inspect.getsource(ic)) + called = { + node.func.attr + for node in ast.walk(tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + } + assert "ensure_future" not in called + assert "create_task" not in called + + +# --------------------------------------------------------------------------- +# A slow document must converge, not die +# --------------------------------------------------------------------------- +class TestSlowIndexingConverges: + """The defect a 1.5 MB PDF exposed in dev on 2026-09-01. + + Ingestion succeeded, but the consumer polled a *retrieval* for 30 s starting + the instant the ingest call returned — before Bedrock had indexed anything. + That poll window was sized against the INDEXED -> retrievable gap + (0.75-1.03 s), while it actually had to cover ingest -> INDEXED -> retrievable, + which the §5.1 benchmark measured at 37-264 s for PDFs. + + Each of the three redeliveries then RE-INGESTED, discarding progress and + restarting the clock. The document reached INDEXED 54 s after the final attempt + was dead-lettered, leaving a fully retrievable document parked at `uploading` + with nothing left to reconcile it — the legacy pipeline no longer writes status + for a promoted knowledge base, so there was no second writer to mask it. + """ + + def _seed_managed(self, table): + _seed_kb(table, retrievalEngine="managed", awsKbId="KB123", awsDataSourceId="DS456") + + def test_a_document_already_being_indexed_is_not_re_ingested(self, table): + """The core fix. Re-submitting restarts the work we are waiting for.""" + self._seed_managed(table) + fake = _FakeBackend(statuses=["IN_PROGRESS"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + with pytest.raises(ic.IngestionRoutingError, match="IN_PROGRESS"): + ic.handle_object(BUCKET, KEY) + + assert fake.ingested == [], ( + "a document Bedrock was already indexing was submitted again; that " + "discards the progress this invocation is waiting on" + ) + + def test_a_document_still_indexing_is_left_non_terminal(self, table): + """Not complete and not failed — the next delivery decides.""" + self._seed_managed(table) + fake = _FakeBackend(statuses=["IN_PROGRESS"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + with pytest.raises(ic.IngestionRoutingError): + ic.handle_object(BUCKET, KEY) + + item = _doc(table) + assert item["status"] not in ("complete", "failed") + assert "indexedAt" not in item, "no timestamp may be invented before indexing" + + def test_a_redelivery_completes_the_document_without_a_second_ingest(self, table): + """Delivery 1 submits and defers; delivery 2 finds it INDEXED and finishes. + + This is the whole convergence property: the document ends up `complete` + having been handed to Bedrock exactly once. + """ + self._seed_managed(table) + + # Delivery 1: never seen, then still working for the whole in-invocation wait. + first = _FakeBackend(statuses=["NOT_FOUND", "IN_PROGRESS"]) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=first + ), patch.object(ic, "INDEXED_POLL_TIMEOUT_SECONDS", 0.01): + with pytest.raises(ic.IngestionRoutingError): + ic.handle_object(BUCKET, KEY) + assert first.ingested == [DOCUMENT_ID] + assert _doc(table)["status"] not in ("complete", "failed") + + # Delivery 2: Bedrock has finished. + second = _FakeBackend(statuses=["INDEXED"]) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=second + ): + ic.handle_object(BUCKET, KEY) + + assert second.ingested == [], "the second delivery must not re-ingest" + item = _doc(table) + assert item["status"] == "complete" + assert item["indexedAt"].startswith("2026-09-01T14:53:19") + + def test_a_small_document_still_finishes_in_one_invocation(self, table): + """The fast path must not regress into waiting for a retry. + + Deferring every document would add a minute of EventBridge backoff to the + common case, which is why the in-invocation wait exists at all. + """ + self._seed_managed(table) + fake = _FakeBackend(statuses=["NOT_FOUND", "INDEXED"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + result = ic.handle_object(BUCKET, KEY) + + assert result["ingested"] is True + assert _doc(table)["status"] == "complete" + assert fake.ingested == [DOCUMENT_ID] + + def test_a_failed_document_is_marked_failed_and_not_retried(self, table): + """Terminal on Bedrock's side. Redelivering cannot help.""" + self._seed_managed(table) + fake = _FakeBackend(statuses=["FAILED"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + result = ic.handle_object(BUCKET, KEY) # must NOT raise + + assert fake.ingested == [] + item = _doc(table) + assert item["status"] == "failed" + assert "FAILED" in item["ingestionError"] + assert result["status"] == "FAILED" + + def test_a_partially_indexed_document_is_treated_as_usable(self, table): + """It IS retrievable, so failing it would hide content the user can see.""" + self._seed_managed(table) + fake = _FakeBackend(statuses=["PARTIALLY_INDEXED"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + ic.handle_object(BUCKET, KEY) + + assert fake.ingested == [], "already indexed, even if partially" + assert _doc(table)["status"] == "complete" + + def test_a_status_probe_failure_does_not_fail_the_document(self, table): + """An unreadable probe means "no evidence", so ingesting is correct.""" + self._seed_managed(table) + + class _ProbeBroken(_FakeBackend): + def _agent(self): + raise RuntimeError("bedrock control plane unavailable") + + fake = _ProbeBroken(statuses=["NOT_FOUND"]) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ), patch.object(ic, "INDEXED_POLL_TIMEOUT_SECONDS", 0.01): + with pytest.raises(ic.IngestionRoutingError): + ic.handle_object(BUCKET, KEY) + + assert fake.ingested == [DOCUMENT_ID], "a probe failure must not block ingestion" + assert _doc(table)["status"] != "failed" diff --git a/backend/tests/lambdas/test_kb_migration_worker.py b/backend/tests/lambdas/test_kb_migration_worker.py new file mode 100644 index 000000000..aa3553169 --- /dev/null +++ b/backend/tests/lambdas/test_kb_migration_worker.py @@ -0,0 +1,1051 @@ +""" +Migration dispatcher and worker: bounded, leased, and safe to interrupt. + +Requirements 15, 16, 19.6. The tests here concentrate on the things that are +invisible when they break: + +* The dispatcher **no-ops when the flag is off**, and "off" includes present but + empty. This is the reconciler-arming defect's shape, and it is worth re-testing + per component because each one reads its own flag. +* The worker dispatches on the **record's** state, never the event's. An event + field that could select `promote` would let a hand-crafted invocation cut a + knowledge base over without it ever verifying. +* A document deleted mid-migration is **not resurrected** — asserted by deleting it + between the snapshot and the ingest, which is the only window where the bug + exists. +* Catch-up **converges on quiet**, not after a fixed number of passes. +* Concurrent promotion yields **one winner**, which is a property of the + conditional write rather than of any locking here. + +Feature: managed-kb-migration +Requirements: 15.4, 15.5, 15.6, 15.7, 15.8, 15.10, 15.13, 15.14, 16.2, 16.3, +16.4, 16.5, 17.1, 17.4, 19.6, 24.5 +""" + +from decimal import Decimal +from typing import Any, Dict, List +from unittest.mock import MagicMock, patch + +import pytest + +from apis.app_api.kb_migration import dispatcher, worker +from apis.shared.kb_backend import records as r +from apis.shared.kb_backend.protocol import Chunk + +ASSISTANT_ID = "ast-migrate-001" +TABLE = "test-assistants" +BUCKET = "test-documents" + +BASE_ENV = { + "DYNAMODB_ASSISTANTS_TABLE_NAME": TABLE, + "S3_ASSISTANTS_DOCUMENTS_BUCKET_NAME": BUCKET, + "AWS_REGION": "us-west-2", +} + + +def _doc(document_id: str, status: str = "complete", size: int = 1024) -> Dict[str, Any]: + return { + "PK": f"AST#{ASSISTANT_ID}", + "SK": f"DOC#{document_id}", + "status": status, + "filename": f"{document_id}.pdf", + "s3Key": f"assistants/{ASSISTANT_ID}/documents/{document_id}/{document_id}.pdf", + "contentHash": f"hash-{document_id}", + "sizeBytes": Decimal(size), + } + + +def _kb_record(state: str = r.SHADOW, **overrides) -> Dict[str, Any]: + record = { + "PK": f"AST#{ASSISTANT_ID}", + "SK": f"KB#{ASSISTANT_ID}", + "appKbId": ASSISTANT_ID, + "ownerUserId": "user-migrate", + "migrationState": state, + "migrationGeneration": Decimal(1), + "totalBytes": Decimal(4096), + "awsKbId": "KB123", + "awsDataSourceId": "DS123", + } + record.update(overrides) + return record + + +async def _async_noop(*args, **kwargs): + """An awaitable that does nothing. + + Used as ``side_effect`` rather than assigning a coroutine to ``return_value``: + a coroutine object assigned that way is created once, so a mock called twice + raises and a mock called never emits "coroutine was never awaited" — noise that + makes a real leak invisible. + """ + return None + + +class StubBackend: + """Records what it was asked to ingest, delete and search.""" + + def __init__(self, chunks: List[Chunk] = None): + self.ingested: List[str] = [] + self.searched: List[str] = [] + self._chunks = chunks if chunks is not None else [ + Chunk(text="hit", relevance=1.0, document_id="d1", metadata={"document_id": "d1"}) + ] + + async def ingest_documents(self, kb_ref, sources, *, batch_size=10): + self.ingested.extend(source.document_id for source in sources) + + async def search(self, kb_ref, query, top_k=5): + self.searched.append(query) + return list(self._chunks) + + async def delete_documents(self, kb_ref, document_ids, *, batch_size=10): # pragma: no cover + raise NotImplementedError + + +# ── Dispatcher ─────────────────────────────────────────────────────────────── +class TestDispatcherFlag: + @pytest.mark.parametrize("value", [None, "", " ", "false", "0", "off", "no", "disabled"]) + def test_anything_but_a_truthy_spelling_is_off(self, value): + """An allow-list, not a truthiness test. The failure being designed around + is a value that is *present but empty*: ``bool("")`` is correct by luck, + ``bool("false")`` is not.""" + env = {} if value is None else {dispatcher.FLAG_MIGRATION_ENABLED: value} + with patch.dict("os.environ", env, clear=True): + assert dispatcher.migration_enabled() is False + + @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", "enabled"]) + def test_affirmative_spellings_are_on(self, value): + with patch.dict( + "os.environ", {dispatcher.FLAG_MIGRATION_ENABLED: value}, clear=True + ): + assert dispatcher.migration_enabled() is True + + @pytest.mark.asyncio + async def test_a_tick_with_the_flag_off_invokes_nothing(self): + with patch.dict("os.environ", {}, clear=True), patch.object( + dispatcher, "_invoke_worker" + ) as invoke, patch.object(dispatcher, "_due_records") as due: + counts = await dispatcher.dispatch_once() + + invoke.assert_not_called() + due.assert_not_called() + assert counts == {"Due": 0, "Dispatched": 0, "Failed": 0} + + +class TestDispatcherLimit: + def test_the_default_matches_the_sync_dispatcher(self): + with patch.dict("os.environ", {}, clear=True): + assert dispatcher.dispatch_limit() == 20 + + def test_an_override_is_honoured(self): + with patch.dict("os.environ", {"KB_MIGRATION_DISPATCH_LIMIT": "5"}, clear=True): + assert dispatcher.dispatch_limit() == 5 + + def test_an_override_above_the_ceiling_is_clamped(self): + """A larger sweep should require repeated observed ticks, not a variable + edit — and `StartIngestionJob` is 0.1 RPS account-wide and not + adjustable, so the only way to stay under it is to not ask.""" + with patch.dict("os.environ", {"KB_MIGRATION_DISPATCH_LIMIT": "5000"}, clear=True): + assert dispatcher.dispatch_limit() == dispatcher.DISPATCH_LIMIT_CEILING + + def test_a_nonsense_override_falls_back_to_the_default(self): + with patch.dict("os.environ", {"KB_MIGRATION_DISPATCH_LIMIT": "lots"}, clear=True): + assert dispatcher.dispatch_limit() == 20 + + @pytest.mark.asyncio + async def test_the_limit_bounds_the_tick_across_all_states_not_per_state(self): + """Three states each honouring the limit would quietly be a 3x limit. + + Asserted on what each **query asked for**, not on the tick's total: the + total is trimmed at the end, so a per-state sweep that read three times the + budget from DynamoDB would still *return* the right number while paying for + three times the reads. + + The first state deliberately returns fewer rows than the limit. That is the + only shape where the bug is observable — if the first query fills the + budget the loop exits either way, which is why the obvious version of this + test passes with the arithmetic removed. + """ + asked: List[int] = [] + rows = [_kb_record(r.SHADOW, appKbId=f"kb-{i}") for i in range(10)] + + def _query(state, now_iso, limit): + asked.append(limit) + # promote yields 2 of the 4 allowed; the rest could fill the tick. + available = 2 if state == r.PROMOTE else 10 + return rows[: min(limit, available)] + + with patch.dict( + "os.environ", + {**BASE_ENV, dispatcher.FLAG_MIGRATION_ENABLED: "true", "KB_MIGRATION_DISPATCH_LIMIT": "4"}, + clear=True, + ), patch("apis.shared.kb_backend.records.query_due_work", side_effect=_query), patch.object( + dispatcher, "_invoke_worker" + ) as invoke, patch.object(dispatcher, "_emit_metrics"): + counts = await dispatcher.dispatch_once() + + assert counts["Due"] == 4 + assert invoke.call_count == 4 + assert asked[0] == 4 + assert asked[1] == 2, ( + f"the second state was asked for {asked[1]} records when only " + f"{4 - 2} of the budget remained; each state is being given the whole " + f"limit ({asked})" + ) + + +class TestDispatcherSweep: + def test_every_work_eligible_state_is_swept(self): + """Derived from ``WORK_ELIGIBLE_STATES``, so a state added there cannot be + silently left unswept — it would stall forever with its work keys written + and nothing reading them.""" + assert set(dispatcher._work_states()) == set(r.WORK_ELIGIBLE_STATES) + + def test_a_state_added_to_the_records_module_is_still_swept(self): + """The assertion above passes today whether or not the derivation exists, + because the priority list happens to name every state. So add one the + dispatcher has never heard of and require it to be swept anyway — which is + the whole point of deriving rather than restating. + """ + extended = frozenset(set(r.WORK_ELIGIBLE_STATES) | {"reindex"}) + with patch.object(r, "WORK_ELIGIBLE_STATES", extended): + states = dispatcher._work_states() + + assert "reindex" in states, ( + "a new work-eligible state is not swept; its records would keep their " + "GSI7 work keys and never be handed to a worker" + ) + # Appended, not promoted ahead of the known order. + assert states[-1] == "reindex" + + def test_promote_is_swept_first(self): + """A record in ``promote`` is one conditional write from finished, so + draining beats starting new shadow work.""" + assert dispatcher._work_states()[0] == r.PROMOTE + + def test_no_terminal_state_is_swept(self): + assert not set(dispatcher._work_states()) & set(r.TERMINAL_STATES) + + @pytest.mark.asyncio + async def test_an_unaddressable_row_does_not_starve_the_sweep(self): + good = _kb_record(r.SHADOW, appKbId="kb-good") + bad = {"SK": "KB#kb-bad", "migrationState": r.SHADOW} # no PK + + calls = {"n": 0} + + def _query(state, now_iso, limit): + calls["n"] += 1 + return [bad, good] if calls["n"] == 1 else [] + + with patch.dict( + "os.environ", + {**BASE_ENV, dispatcher.FLAG_MIGRATION_ENABLED: "true"}, + clear=True, + ), patch("apis.shared.kb_backend.records.query_due_work", side_effect=_query), patch.object( + dispatcher, "_invoke_worker" + ) as invoke, patch.object(dispatcher, "_emit_metrics"): + counts = await dispatcher.dispatch_once() + + assert counts["Failed"] == 1 + assert counts["Dispatched"] == 1 + assert invoke.call_args.args[0]["appKbId"] == "kb-good" + + @pytest.mark.asyncio + async def test_a_failing_index_query_does_not_fail_the_tick(self): + with patch.dict( + "os.environ", + {**BASE_ENV, dispatcher.FLAG_MIGRATION_ENABLED: "true"}, + clear=True, + ), patch( + "apis.shared.kb_backend.records.query_due_work", + side_effect=RuntimeError("dynamodb down"), + ), patch.object(dispatcher, "_emit_metrics"): + counts = await dispatcher.dispatch_once() + + assert counts == {"Due": 0, "Dispatched": 0, "Failed": 0} + + def test_the_handler_reads_nothing_from_the_event(self): + """The reconciler's arming bypass came from forwarding an event field. + Nothing here may select a state, a limit or a knowledge base.""" + seen = {} + + async def _tick(): + seen["called"] = True + return {"Due": 0, "Dispatched": 0, "Failed": 0} + + with patch.object(dispatcher, "dispatch_once", side_effect=_tick) as tick: + dispatcher.lambda_handler({"migrationState": "promote", "armed": True}, None) + + assert seen.get("called") is True + tick.assert_called_once_with() + + +# ── Worker: state selection ────────────────────────────────────────────────── +class TestTheRecordDecidesTheStep: + @pytest.mark.asyncio + async def test_an_event_cannot_select_promote(self): + """A hand-crafted invocation must not be able to cut over a knowledge base + that never verified.""" + record = _kb_record(r.SHADOW) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.get_kb_record", return_value=record + ), patch.object(worker, "take_lease", return_value="later"), patch.object( + worker, "run_shadow" + ) as shadow, patch.object(worker, "run_promote") as promote: + shadow.return_value = worker.StepResult(ASSISTANT_ID, ASSISTANT_ID, r.SHADOW, r.VERIFY) + await worker.run_step(ASSISTANT_ID, ASSISTANT_ID) + + shadow.assert_called_once() + promote.assert_not_called() + + @pytest.mark.asyncio + async def test_a_terminal_record_is_a_no_op(self): + """The index is eventually consistent, so a record that finished a moment + ago can still be handed over once. That is not an error.""" + record = _kb_record(r.RETAIN) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.get_kb_record", return_value=record + ), patch.object(worker, "take_lease") as lease: + result = await worker.run_step(ASSISTANT_ID) + + lease.assert_not_called() + assert result.to_state == r.RETAIN + assert "not work-eligible" in result.detail + + @pytest.mark.asyncio + async def test_a_lost_lease_propagates_rather_than_failing_the_migration(self): + """Requirement 15.13. Two overlapping ticks is ordinary; marking the + migration `failed` because of it would strand a healthy knowledge base.""" + record = _kb_record(r.SHADOW) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.get_kb_record", return_value=record + ), patch( + "apis.shared.kb_backend.records.acquire_lease", + side_effect=RuntimeError("conditional check failed"), + ), patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object(worker, "_fail") as fail: + with pytest.raises(worker.LeaseLost): + await worker.run_step(ASSISTANT_ID) + + fail.assert_not_called() + + +# ── Worker: shadow and catch-up ────────────────────────────────────────────── +class TestShadowAndCatchUp: + @pytest.mark.asyncio + async def test_documents_are_ingested_from_their_existing_s3_keys(self): + """Requirement 15.4: a re-ingest, never a re-upload.""" + docs = [_doc("d1"), _doc("d2")] + backend = StubBackend() + captured = {} + + async def _capture(kb_ref, sources, *, batch_size=10): + captured["sources"] = list(sources) + backend.ingested.extend(s.document_id for s in sources) + + backend.ingest_documents = _capture + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=docs + ), patch.object( + worker, "get_document_item", side_effect=lambda a, d: _doc(d) + ), patch( + "apis.shared.kb_backend.byte_cap.reserve_snapshot" + ), patch( + "apis.shared.kb_backend.provisioning.provision_managed_kb" + ) as provision, patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_record_progress" + ), patch( + "apis.shared.kb_backend.records.set_migration_state" + ): + provision.side_effect = _async_noop + result = await worker.run_shadow( + ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.SHADOW), backend + ) + + assert result.to_state == r.VERIFY + assert sorted(backend.ingested) == ["d1", "d2"] + keys = {s.s3_key for s in captured["sources"]} + assert keys == { + f"assistants/{ASSISTANT_ID}/documents/d1/d1.pdf", + f"assistants/{ASSISTANT_ID}/documents/d2/d2.pdf", + } + + @pytest.mark.asyncio + async def test_only_complete_documents_are_migrated(self): + """Requirement 15.5. A non-complete document is not retrievable on legacy + either, so migrating it would create a difference where the point is + parity.""" + docs = [_doc("d1"), _doc("d2", status="failed"), _doc("d3", status="uploading")] + backend = StubBackend() + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=docs + ), patch.object( + worker, "get_document_item", side_effect=lambda a, d: next( + (x for x in docs if worker.document_id_of(x) == d), None + ) + ), patch( + "apis.shared.kb_backend.byte_cap.reserve_snapshot" + ), patch( + "apis.shared.kb_backend.provisioning.provision_managed_kb" + ) as provision, patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_record_progress" + ), patch( + "apis.shared.kb_backend.records.set_migration_state" + ): + provision.side_effect = _async_noop + await worker.run_shadow(ASSISTANT_ID, ASSISTANT_ID, _kb_record(), backend) + + assert backend.ingested == ["d1"] + + @pytest.mark.asyncio + async def test_a_document_deleted_mid_migration_is_not_resurrected(self): + """Requirements 16.4, 16.5, and the reason the re-read is per document + rather than per batch: a PDF batch takes minutes, and the deletion this + guards against is most likely to land inside exactly that window. + + ``d2`` is in the snapshot but gone by the time its turn comes. + """ + docs = [_doc("d1"), _doc("d2")] + deleted = {"d2"} + backend = StubBackend() + + def _get(assistant_id, document_id): + return None if document_id in deleted else _doc(document_id) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=docs + ), patch.object(worker, "get_document_item", side_effect=_get), patch( + "apis.shared.kb_backend.byte_cap.reserve_snapshot" + ), patch( + "apis.shared.kb_backend.provisioning.provision_managed_kb" + ) as provision, patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_record_progress" + ), patch( + "apis.shared.kb_backend.records.set_migration_state" + ): + provision.side_effect = _async_noop + result = await worker.run_shadow(ASSISTANT_ID, ASSISTANT_ID, _kb_record(), backend) + + assert backend.ingested == ["d1"], "a deleted document was resurrected" + assert result.documents_skipped >= 1 + + @pytest.mark.asyncio + async def test_a_document_that_stopped_being_complete_is_skipped(self): + docs = [_doc("d1")] + backend = StubBackend() + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=docs + ), patch.object( + worker, "get_document_item", return_value=_doc("d1", status="deleting") + ), patch( + "apis.shared.kb_backend.byte_cap.reserve_snapshot" + ), patch( + "apis.shared.kb_backend.provisioning.provision_managed_kb" + ) as provision, patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_record_progress" + ), patch( + "apis.shared.kb_backend.records.set_migration_state" + ): + provision.side_effect = _async_noop + await worker.run_shadow(ASSISTANT_ID, ASSISTANT_ID, _kb_record(), backend) + + assert backend.ingested == [] + + @pytest.mark.asyncio + async def test_the_whole_snapshot_is_reserved_before_anything_is_provisioned(self): + """Requirement 12.9. Reserving per document would let a migration run for + an hour and stop halfway, leaving a half-populated corpus and an owner over + their cap with no way back.""" + order: List[str] = [] + docs = [_doc("d1", size=2048), _doc("d2", size=4096)] + + def _reserve(assistant_id, app_kb_id, total, cap): + order.append(f"reserve:{total}") + + async def _provision(*args, **kwargs): + order.append("provision") + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=docs + ), patch.object( + worker, "get_document_item", side_effect=lambda a, d: _doc(d) + ), patch( + "apis.shared.kb_backend.byte_cap.reserve_snapshot", side_effect=_reserve + ), patch( + "apis.shared.kb_backend.provisioning.provision_managed_kb", side_effect=_provision + ), patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_record_progress" + ), patch( + "apis.shared.kb_backend.records.set_migration_state" + ): + fresh = _kb_record() + fresh.pop("totalBytes") + await worker.run_shadow(ASSISTANT_ID, ASSISTANT_ID, fresh, StubBackend()) + + assert order == ["reserve:6144", "provision"] + + @pytest.mark.asyncio + async def test_an_over_cap_corpus_fails_before_provisioning(self): + from apis.shared.kb_backend.byte_cap import ByteCapExceeded + + async def _provision(*args, **kwargs): # pragma: no cover - must not run + raise AssertionError("provisioned despite the byte cap") + + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.get_kb_record", return_value=_kb_record() + ), patch.object(worker, "take_lease", return_value="later"), patch.object( + worker, "list_document_items", return_value=[_doc("d1")] + ), patch( + "apis.shared.kb_backend.byte_cap.reserve_snapshot", + side_effect=ByteCapExceeded(requested=1, cap=0), + ), patch( + "apis.shared.kb_backend.provisioning.provision_managed_kb", side_effect=_provision + ), patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_fail" + ) as fail: + result = await worker.run_step(ASSISTANT_ID) + + assert result.to_state == r.MIGRATION_FAILED + fail.assert_called_once() + + @pytest.mark.asyncio + async def test_catch_up_converges_on_quiet_not_on_a_pass_count(self): + """Requirement 16.3. A new document appears during the first pass; the + second finds nothing and that is what ends it.""" + backend = StubBackend() + state = {"pass": 0} + + def _list(assistant_id): + state["pass"] += 1 + if state["pass"] == 1: + return [_doc("d1"), _doc("d2")] + return [_doc("d1"), _doc("d2")] + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", side_effect=_list + ), patch.object(worker, "get_document_item", side_effect=lambda a, d: _doc(d)): + passes, converged, counts = await worker.catch_up( + ASSISTANT_ID, ASSISTANT_ID, {"d1"}, backend + ) + + assert converged is True + assert passes == 2 + assert backend.ingested == ["d2"] + + @pytest.mark.asyncio + async def test_a_corpus_that_never_settles_does_not_converge(self): + """And staying in ``shadow`` is the correct outcome: the corpus keeps + serving from legacy while the owner keeps uploading.""" + backend = StubBackend() + counter = {"n": 0} + + def _list(assistant_id): + counter["n"] += 1 + return [_doc(f"d{i}") for i in range(counter["n"] + 1)] + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", side_effect=_list + ), patch.object(worker, "get_document_item", side_effect=lambda a, d: _doc(d)): + passes, converged, _ = await worker.catch_up( + ASSISTANT_ID, ASSISTANT_ID, set(), backend, max_passes=3 + ) + + assert converged is False + assert passes == 3 + + @pytest.mark.asyncio + async def test_an_unconverged_shadow_stays_in_shadow(self): + docs = [_doc("d1")] + transitions: List[str] = [] + + def _set_state(assistant_id, app_kb_id, new_state, generation, due=None, expected=None, error=None): + transitions.append(new_state) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=docs + ), patch.object( + worker, "get_document_item", side_effect=lambda a, d: _doc(d) + ), patch( + "apis.shared.kb_backend.byte_cap.reserve_snapshot" + ), patch( + "apis.shared.kb_backend.provisioning.provision_managed_kb" + ) as provision, patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_record_progress" + ), patch( + "apis.shared.kb_backend.records.set_migration_state", side_effect=_set_state + ), patch.object( + worker, "catch_up", return_value=(5, False, {"migrated": 0, "skipped": 0, "done": []}) + ): + provision.side_effect = _async_noop + result = await worker.run_shadow(ASSISTANT_ID, ASSISTANT_ID, _kb_record(), StubBackend()) + + assert transitions == [r.SHADOW] + assert result.to_state == r.SHADOW + assert result.converged is False + + +# ── Worker: verify ─────────────────────────────────────────────────────────── +class TestVerify: + def test_the_manifest_is_content_identity_not_a_count(self): + """Requirement 15.6. Count parity is satisfied by a corpus with the right + *number* of wrong documents — exactly what a migration that raced an upload + and a delete produces.""" + before = worker.source_manifest([_doc("d1"), _doc("d2")]) + changed = dict(_doc("d2")) + changed["contentHash"] = "hash-d2-edited" + after = worker.source_manifest([_doc("d1"), changed]) + + assert len(before) == len(after) + assert before != after, "the manifest is count-equivalent and cannot see an edit" + + def test_a_document_with_no_hash_still_contributes_a_changing_value(self): + item = {"SK": "DOC#d9", "status": "complete", "updatedAt": "2026-08-01T00:00:00Z"} + assert worker.manifest_entry(item) == "d9:2026-08-01T00:00:00Z" + + @pytest.mark.asyncio + async def test_an_unqueryable_corpus_defers_instead_of_failing(self): + """Requirement 15.7, corrected by measurement. + + "Not queryable yet" is a verification that has not happened, not one that + failed. The docstring's 0.75-1.03 s was measured on a warm knowledge base; + a first ingest into a fresh one took ~45 s in dev, and treating that as + terminal marked a good migration `failed` and showed its owner a retry + button for a problem that resolves itself. + """ + backend = StubBackend(chunks=[]) + deferred = [] + + def _defer(_assistant, _kb, _generation, due_at): + deferred.append(due_at) + return len(deferred) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=[_doc("d1")] + ), patch.object(r, "defer_verify", _defer): + result = await worker.run_verify( + ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.VERIFY), backend + ) + + assert deferred, "did not defer; an early canary would fail the migration" + assert result.to_state == r.VERIFY, "left verify on a deferral" + assert "not queryable" in result.detail + + @pytest.mark.asyncio + async def test_deferring_forever_eventually_fails(self): + """Bounded, because "not queryable" past some point is not latency. + + Matched on the attempt count rather than "not queryable", so this cannot + be satisfied by the deferral path it is meant to sit past. + """ + backend = StubBackend(chunks=[]) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=[_doc("d1")] + ), patch.object( + r, "defer_verify", lambda *a, **k: worker.MAX_VERIFY_ATTEMPTS + 1 + ): + with pytest.raises(worker.VerificationFailed, match="attempts over"): + await worker.run_verify( + ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.VERIFY), backend + ) + + @pytest.mark.asyncio + async def test_verify_rejects_a_canary_that_returns_foreign_documents(self): + backend = StubBackend( + chunks=[ + Chunk( + text="someone else's", + relevance=1.0, + document_id="not-ours", + metadata={"document_id": "not-ours"}, + ) + ] + ) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=[_doc("d1")] + ): + with pytest.raises(worker.VerificationFailed): + await worker.run_verify(ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.VERIFY), backend) + + @pytest.mark.asyncio + async def test_an_empty_corpus_cannot_be_verified(self): + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=[_doc("d1", status="failed")] + ): + with pytest.raises(worker.VerificationFailed, match="nothing"): + await worker.run_verify( + ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.VERIFY), StubBackend() + ) + + @pytest.mark.asyncio + async def test_a_successful_verify_moves_to_promote(self): + backend = StubBackend( + chunks=[ + Chunk(text="hit", relevance=1.0, document_id="d1", metadata={"document_id": "d1"}) + ] + ) + transitions: List[tuple] = [] + + def _set_state(assistant_id, app_kb_id, new_state, generation, due=None, expected=None, error=None): + transitions.append((new_state, expected)) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=[_doc("d1")] + ), patch("apis.shared.kb_backend.records.set_migration_state", side_effect=_set_state): + result = await worker.run_verify( + ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.VERIFY), backend + ) + + assert result.to_state == r.PROMOTE + assert transitions == [(r.PROMOTE, [r.VERIFY])] + + def test_the_canary_query_is_built_from_the_corpus(self): + """Not a fixed string: a constant like "test" can legitimately match + nothing in a real corpus, which would fail healthy knowledge bases and + train whoever is watching to ignore it.""" + query = worker._canary_query([{"filename": "student_handbook.pdf"}]) + assert "student" in query and "handbook" in query + assert ".pdf" not in query + + +# ── Worker: promote and rollback ───────────────────────────────────────────── +class TestPromote: + @pytest.mark.asyncio + async def test_promotion_is_refused_without_a_byte_cap_accumulator(self): + """Requirement 12.9: no traffic is promoted to an unmetered corpus.""" + record = _kb_record(r.PROMOTE) + record.pop("totalBytes") + + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.promote_engine" + ) as promote: + with pytest.raises(worker.MigrationError, match="totalBytes"): + await worker.run_promote(ASSISTANT_ID, ASSISTANT_ID, record) + + promote.assert_not_called() + + @pytest.mark.asyncio + async def test_promotion_writes_once_and_then_retains(self): + calls: List[str] = [] + + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.promote_engine", + side_effect=lambda *a: calls.append("promote"), + ), patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_set_retain_until", side_effect=lambda *a: calls.append("retain_until") + ), patch( + "apis.shared.kb_backend.records.set_migration_state", + side_effect=lambda *a, **k: calls.append("state"), + ): + result = await worker.run_promote(ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.PROMOTE)) + + assert calls == ["promote", "retain_until", "state"] + assert result.to_state == r.RETAIN + + @pytest.mark.asyncio + async def test_concurrent_promotion_yields_one_winner(self): + """Requirement 15.10. The property belongs to the conditional write, so the + test is that the loser's exception is not swallowed into a second success.""" + from botocore.exceptions import ClientError + + winners = {"n": 0} + + def _promote(assistant_id, app_kb_id, generation, now_iso): + winners["n"] += 1 + if winners["n"] > 1: + raise ClientError( + {"Error": {"Code": "ConditionalCheckFailedException"}}, "UpdateItem" + ) + + # The loser re-reads before deciding, because a refused write means either + # "somebody else promoted" (success) or "a guard genuinely failed" (not). + # Here the record is still unpromoted, so the refusal must propagate. + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.promote_engine", side_effect=_promote + ), patch( + "apis.shared.kb_backend.records.get_kb_record", return_value=_kb_record(r.PROMOTE) + ), patch("apis.shared.kb_backend.metrics.emit_count"), patch.object( + worker, "_set_retain_until" + ), patch("apis.shared.kb_backend.records.set_migration_state"): + first = await worker.run_promote(ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.PROMOTE)) + with pytest.raises(ClientError): + await worker.run_promote(ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.PROMOTE)) + + assert first.to_state == r.RETAIN + assert winners["n"] == 2 + + def test_the_retain_window_cannot_be_shortened_below_thirty_days(self): + """Requirement 15.11 says *at least* 30 days. Shortening the rollback + window is not a tuning knob.""" + with patch.dict("os.environ", {"KB_MIGRATION_RETAIN_DAYS": "3"}, clear=True): + assert worker._retain_days() == 30 + with patch.dict("os.environ", {"KB_MIGRATION_RETAIN_DAYS": "90"}, clear=True): + assert worker._retain_days() == 90 + + +class TestRollback: + @pytest.mark.asyncio + async def test_rollback_moves_no_data(self): + """Requirement 17.2. The legacy index was never mutated — that is what + building the managed corpus alongside it bought.""" + touched: List[str] = [] + + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.rollback_engine", + side_effect=lambda *a: touched.append("engine"), + ), patch("apis.shared.kb_backend.metrics.emit_count"): + result = await worker.rollback(ASSISTANT_ID, ASSISTANT_ID) + + assert touched == ["engine"] + assert "no data moved" in result.detail + + @pytest.mark.asyncio + async def test_rollback_does_not_delete_the_managed_knowledge_base(self): + """Deleting it here would turn a reversible decision into an irreversible + one at the moment somebody is least sure.""" + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.rollback_engine" + ), patch("apis.shared.kb_backend.metrics.emit_count"), patch( + "apis.shared.kb_backend.tombstones.delete_knowledge_base", create=True + ) as delete_kb: + await worker.rollback(ASSISTANT_ID, ASSISTANT_ID) + + delete_kb.assert_not_called() + + @pytest.mark.asyncio + async def test_a_pre_promotion_failure_leaves_the_record_on_legacy(self): + """Requirement 17.4. `failed` is terminal and removes the work keys; the + engine attribute was never written, so the knowledge base is still legacy + and still usable.""" + recorded: List[tuple] = [] + + def _set_state(assistant_id, app_kb_id, new_state, generation, due=None, expected=None, error=None): + recorded.append((new_state, error)) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.get_kb_record", return_value=_kb_record(r.VERIFY) + ), patch.object(worker, "take_lease", return_value="later"), patch.object( + worker, "run_verify", side_effect=worker.VerificationFailed("canary empty") + ), patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch( + "apis.shared.kb_backend.records.set_migration_state", side_effect=_set_state + ), patch( + "apis.shared.kb_backend.records.promote_engine" + ) as promote: + result = await worker.run_step(ASSISTANT_ID) + + assert result.to_state == r.MIGRATION_FAILED + assert recorded and recorded[0][0] == r.MIGRATION_FAILED + promote.assert_not_called() + + +class TestResumingWithoutRedoingWork: + """The two behaviours the convergence property test forced into existence.""" + + def test_the_completed_set_is_read_off_the_record(self): + assert worker.already_migrated({}) == set() + assert worker.already_migrated({"migratedDocIds": {"d1", "d2"}}) == {"d1", "d2"} + + def test_a_non_iterable_completed_set_degrades_to_empty(self): + """Re-ingesting is slow, not wrong — ``customDocumentIdentifier`` makes it a + replace — so a malformed attribute must not stop the migration.""" + assert worker.already_migrated({"migratedDocIds": 7}) == set() + + @pytest.mark.asyncio + async def test_a_resumed_shadow_skips_documents_it_already_ingested(self): + """Before this, a crash near the end of a PDF corpus re-parsed all of it — + 37-264 s per document, so an hour of work redone for nothing.""" + docs = [_doc("d1"), _doc("d2"), _doc("d3")] + backend = StubBackend() + record = _kb_record(r.SHADOW, migratedDocIds={"d1", "d2"}) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=docs + ), patch.object( + worker, "get_document_item", side_effect=lambda a, d: _doc(d) + ), patch( + "apis.shared.kb_backend.byte_cap.reserve_snapshot" + ) as reserve, patch( + "apis.shared.kb_backend.provisioning.provision_managed_kb", side_effect=_async_noop + ), patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_record_progress" + ), patch( + "apis.shared.kb_backend.records.set_migration_state" + ): + await worker.run_shadow(ASSISTANT_ID, ASSISTANT_ID, record, backend) + + assert backend.ingested == ["d3"] + # And the corpus is not reserved a second time: the accumulator is on the + # record, so re-reserving would double-count the owner's own corpus against + # their cap until the migration refused itself. + reserve.assert_not_called() + + @pytest.mark.asyncio + async def test_progress_persists_the_ids_not_just_a_count(self): + docs = [_doc("d1"), _doc("d2")] + captured = {} + + async def _progress(assistant_id, app_kb_id, *, migrated, total, skipped, newly_done=None): + captured["newly_done"] = list(newly_done or []) + captured["migrated"] = migrated + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=docs + ), patch.object( + worker, "get_document_item", side_effect=lambda a, d: _doc(d) + ), patch( + "apis.shared.kb_backend.byte_cap.reserve_snapshot" + ), patch( + "apis.shared.kb_backend.provisioning.provision_managed_kb", side_effect=_async_noop + ), patch( + "apis.shared.kb_backend.metrics.emit_count" + ), patch.object( + worker, "_record_progress", side_effect=_progress + ), patch( + "apis.shared.kb_backend.records.set_migration_state" + ): + await worker.run_shadow(ASSISTANT_ID, ASSISTANT_ID, _kb_record(), StubBackend()) + + assert sorted(captured["newly_done"]) == ["d1", "d2"], ( + "a count alone cannot tell a resume *which* documents to skip" + ) + + @pytest.mark.asyncio + async def test_a_record_already_promoted_finishes_instead_of_failing(self): + """The crash window between the promotion write and the state transition. + + The promotion write is guarded on ``attribute_not_exists(retrievalEngine)``, + so retrying it is refused — and treating that refusal as a failure would + mark a migration that actually succeeded as ``failed``, leaving a promoted + knowledge base with no retention window. + """ + record = _kb_record(r.PROMOTE, retrievalEngine="managed") + calls: List[str] = [] + + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.promote_engine", + side_effect=lambda *a: calls.append("promote"), + ), patch("apis.shared.kb_backend.metrics.emit_count"), patch.object( + worker, "_set_retain_until", side_effect=lambda *a: calls.append("retain_until") + ), patch( + "apis.shared.kb_backend.records.set_migration_state", + side_effect=lambda *a, **k: calls.append("state"), + ): + result = await worker.run_promote(ASSISTANT_ID, ASSISTANT_ID, record) + + assert "promote" not in calls, "promoted a second time" + assert calls == ["retain_until", "state"] + assert result.to_state == r.RETAIN + assert "already promoted" in result.detail + + @pytest.mark.asyncio + async def test_a_refused_promotion_on_an_unpromoted_record_still_raises(self): + """So "already promoted" cannot become a blanket swallow of the guard.""" + with patch.dict("os.environ", BASE_ENV, clear=True), patch( + "apis.shared.kb_backend.records.promote_engine", + side_effect=RuntimeError("conditional check failed"), + ), patch( + "apis.shared.kb_backend.records.get_kb_record", return_value=_kb_record(r.PROMOTE) + ), patch("apis.shared.kb_backend.metrics.emit_count"): + with pytest.raises(RuntimeError): + await worker.run_promote(ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.PROMOTE)) + + +class TestRehydrationReappliesTheResourcePolicy: + """Task 13.7 / Requirement 24.12, asserted at the level a rehydration works at. + + A resource policy attaches to the AWS knowledge base ARN. Provisioning that + produces a *new* ``awsKbId`` — a rehydration, or a replacement after a failed + delete — therefore leaves the old policy on a resource nobody reads, and sharing + silently stops. The repair is a state comparison rather than an event, so it + cannot be bypassed by a code path that forgets to fire anything. + """ + + @pytest.mark.asyncio + async def test_a_new_aws_kb_id_makes_the_recorded_policy_stale(self): + from apis.shared.kb_backend.resource_policy import POLICY_KB_ID_ATTR, policy_is_stale + + rehydrated = _kb_record(r.RETAIN, awsKbId="KB-NEW", **{POLICY_KB_ID_ATTR: "KB123"}) + assert policy_is_stale(rehydrated) is True + + @pytest.mark.asyncio + async def test_the_policy_is_reapplied_to_the_new_arn(self): + from apis.shared.kb_backend.resource_policy import ( + POLICY_KB_ID_ATTR, + ensure_retrieve_policy, + ) + + client = MagicMock() + client.put_resource_policy.return_value = {"revisionId": "rev-after-rehydration"} + rehydrated = _kb_record(r.RETAIN, awsKbId="KB-NEW", **{POLICY_KB_ID_ATTR: "KB123"}) + + with patch.dict( + "os.environ", + { + **BASE_ENV, + "AWS_ACCOUNT_ID": "123456789012", + "MANAGED_KB_RETRIEVAL_PRINCIPAL_ARNS": "arn:aws:iam::123456789012:role/runtime", + }, + clear=True, + ), patch("apis.shared.kb_backend.records.set_resource_policy_state") as setter: + revision = await ensure_retrieve_policy( + ASSISTANT_ID, ASSISTANT_ID, shared=True, record=rehydrated, client=client + ) + + assert revision == "rev-after-rehydration" + assert client.put_resource_policy.call_args.kwargs["resourceArn"].endswith( + "knowledge-base/KB-NEW" + ) + setter.assert_called_once_with( + ASSISTANT_ID, ASSISTANT_ID, "KB-NEW", "rev-after-rehydration" + ) + + +# ── Mixed old/new deployment ───────────────────────────────────────────────── +class TestMixedDeployment: + def test_a_record_without_an_engine_resolves_to_legacy(self): + """Requirements 1.6, 24.8. Old and new code serving simultaneously agree, + because "absence means legacy" is a property of the data rather than of the + code version reading it.""" + for item in ({}, None, _kb_record(), {"appKbId": "x", "migrationState": r.SHADOW}): + assert r.resolve_engine(item) == r.ENGINE_LEGACY + + def test_only_an_explicit_managed_value_resolves_to_managed(self): + assert r.resolve_engine({"retrievalEngine": "managed"}) == r.ENGINE_MANAGED + for wrong in ("MANAGED", "Managed", "s3vectors", "", None, True): + assert r.resolve_engine({"retrievalEngine": wrong}) == r.ENGINE_LEGACY + + def test_a_mid_migration_record_still_serves_legacy(self): + """Requirements 15.3, 16.1. `shadow` and `verify` never touch + `retrievalEngine`, so a knowledge base being migrated is indistinguishable + from one that is not, to anything doing retrieval.""" + for state in (r.SHADOW, r.VERIFY, r.PROMOTE): + assert r.resolve_engine(_kb_record(state)) == r.ENGINE_LEGACY diff --git a/backend/tests/lambdas/test_kb_reconciler.py b/backend/tests/lambdas/test_kb_reconciler.py new file mode 100644 index 000000000..52c6dd8e4 --- /dev/null +++ b/backend/tests/lambdas/test_kb_reconciler.py @@ -0,0 +1,1011 @@ +"""Daily reconciler — the join, the age gate, and the disarmed default. + +Feature: managed-kb-migration, task 10.3. +Requirements: 24.4, 14.1-14.8, 19.7, 19.8. + +Three assertions here are the reason the file exists, and each guards a mistake +that a passing test suite would otherwise hide: + +**The age gate reads AWS's ``createdAt``, never discovery time.** Asserted from +both ends. An orphan that AWS says is eight days old is deletable on the *very +first* run that ever sees it — an implementation that started a 24-hour clock at +discovery would leave it, and would then leave it again after any reconciler +outage. And a knowledge base AWS says is 30 seconds old is left alone even though +it is equally newly discovered, because that one is an in-flight create. + +**Record-only marks and never deletes.** A KB_Record whose AWS knowledge base has +gone means the *vectors* are gone. The uploaded bytes are still in S3 and the +``DOC#`` rows still describe them, so the corpus rebuilds on the next ingest and +the owner re-uploads nothing. The record is the only pointer to that corpus, so +deleting it is the single action in this module that would lose user data. + +**Report-only really is a no-op.** The shipped mode logs intended deletions and +issues none, and the arming flag treats an empty string as off — an unset GitHub +Actions variable expands to ``""``. + +No test contacts AWS. DynamoDB is moto; ``bedrock-agent`` is a stub +(Requirement 24.11). +""" + +from datetime import datetime, timedelta, timezone + +import boto3 +import pytest +from moto import mock_aws + +from apis.app_api.kb_migration import reconciler as rec +from apis.shared.kb_backend import tombstones as tomb +from apis.shared.kb_backend import tags as kb_tags +from tests.shared.test_kb_tombstones import FakeBedrockAgent + +REGION = "us-east-1" +TABLE = "test-kb-reconciler" +PREFIX = "testprefix" +ENV = "testenv" +NOW = datetime(2026, 6, 1, 12, 0, 0, tzinfo=timezone.utc) + + +@pytest.fixture() +def table(monkeypatch): + monkeypatch.setenv("AWS_DEFAULT_REGION", REGION) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing") + monkeypatch.setenv("AWS_SESSION_TOKEN", "testing") + monkeypatch.setenv("DYNAMODB_ASSISTANTS_TABLE_NAME", TABLE) + monkeypatch.setenv(kb_tags.ENV_TAG_VALUE_PREFIX, PREFIX) + monkeypatch.setenv(kb_tags.ENV_TAG_VALUE_ENVIRONMENT, ENV) + # Never inherited from the developer's shell: the whole point of the flag is + # that the reconciler is disarmed unless something says otherwise. + monkeypatch.delenv(rec.FLAG_RECONCILER_ARMED, raising=False) + + with mock_aws(): + boto3.client("dynamodb", region_name=REGION).create_table( + TableName=TABLE, + KeySchema=[ + {"AttributeName": "PK", "KeyType": "HASH"}, + {"AttributeName": "SK", "KeyType": "RANGE"}, + ], + AttributeDefinitions=[ + {"AttributeName": "PK", "AttributeType": "S"}, + {"AttributeName": "SK", "AttributeType": "S"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + yield boto3.resource("dynamodb", region_name=REGION).Table(TABLE) + + +@pytest.fixture(autouse=True) +def no_metrics(monkeypatch): + monkeypatch.setattr(rec, "emit_count", lambda *a, **k: None) + monkeypatch.setattr(tomb, "emit_count", lambda *a, **k: None) + + +def _iso(moment): + """The exact timestamp shape this feature writes everywhere. + + Spelled out rather than ``isoformat()`` because every comparison in the + idleness path is lexicographic on this format; an offset-style string would + sort differently and the test would be measuring the wrong thing. + """ + return moment.strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _arn(kb_id): + return f"arn:aws:bedrock:{REGION}:123456789012:knowledge-base/{kb_id}" + + +def _aws_kb(kb_id, created_at, status="ACTIVE", app_kb_id=None): + """One knowledge base as AWS reports it, with AWS's own ``createdAt``.""" + return { + "knowledgeBaseId": kb_id, + "name": f"{PREFIX}-kb-{app_kb_id or kb_id}", + "status": status, + "knowledgeBaseArn": _arn(kb_id), + "roleArn": "arn:aws:iam::123456789012:role/kb", + "createdAt": created_at, + } + + +def _ours(kb_id, app_kb_id): + return { + # Built through the canonical helper, not spelled out: a fixture that + # hardcodes tag keys is a fixture that keeps passing after the keys change + # under it, which is how the three-way drift stayed invisible. + _arn(kb_id): kb_tags.build_tags(app_kb_id, "u-1", PREFIX, ENV) + } + + +def _seed_record(table, assistant_id, aws_kb_id=None, **extra): + item = { + "PK": f"AST#{assistant_id}", + "SK": f"KB#{assistant_id}", + "appKbId": assistant_id, + "retrievalEngine": "managed", + } + if aws_kb_id: + item["awsKbId"] = aws_kb_id + item["awsDataSourceId"] = f"DS{aws_kb_id}" + item.update(extra) + table.put_item(Item=item) + return item + + +def _record(table, assistant_id): + return table.get_item( + Key={"PK": f"AST#{assistant_id}", "SK": f"KB#{assistant_id}"} + ).get("Item") + + +def _run(client, table, **kwargs): + kwargs.setdefault("now", NOW) + kwargs.setdefault("stored_bytes_resolver", lambda _assistant_id: None) + return rec.reconcile(client=client, **kwargs) + + +# ── Requirement 19.7, 19.8: the arming flag ────────────────────────────────── +class TestArmingFlag: + @pytest.mark.parametrize( + "value", + ["", " ", "0", "false", "False", "off", "no", "disabled", "maybe"], + ) + def test_falsy_and_empty_values_are_off(self, monkeypatch, value): + """An **empty string must read as off** (Requirement 19.8). + + An unset GitHub Actions variable expands to ``""``, so a truthiness test + on the raw value is the exact bug this guards. ``"false"`` matters too: + ``bool("false")`` is ``True``. + """ + monkeypatch.setenv(rec.FLAG_RECONCILER_ARMED, value) + assert rec.reconciler_armed() is False + + def test_unset_is_off(self, monkeypatch): + monkeypatch.delenv(rec.FLAG_RECONCILER_ARMED, raising=False) + assert rec.reconciler_armed() is False + + @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", "enabled", " true "]) + def test_affirmative_values_arm(self, monkeypatch, value): + monkeypatch.setenv(rec.FLAG_RECONCILER_ARMED, value) + assert rec.reconciler_armed() is True + + def test_reconcile_defaults_to_the_flag(self, table, monkeypatch): + monkeypatch.setenv(rec.FLAG_RECONCILER_ARMED, "") + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBORPH1", NOW - timedelta(days=8))], + tags=_ours("KBORPH1", "ast-orph1"), + ) + + report = rec.reconcile( + client=client, now=NOW, stored_bytes_resolver=lambda _a: None + ) + + assert report.armed is False + assert report.to_dict()["mode"] == "report-only" + + +# ── Requirement 14.7: report-only deletes nothing ──────────────────────────── +class TestReportOnlyDeletesNothing: + def test_an_eligible_orphan_is_reported_and_not_deleted(self, table): + """The shipped mode. It must plan the deletion and perform none of it.""" + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBORPH1", NOW - timedelta(days=8))], + tags=_ours("KBORPH1", "ast-orph1"), + ) + + report = _run(client, table, armed=False) + + assert report.orphans == 1 + assert [p.kb_id for p in report.planned_deletions] == ["KBORPH1"] + assert report.deletions_performed == 0 + assert client.delete_calls == [], ( + "report-only mode issued a DeleteKnowledgeBase call" + ) + + def test_report_only_makes_no_mutating_call_at_all(self, table): + """Nothing happens: no AWS delete, and no DynamoDB side effect. + + Asserted on the AWS call log rather than on the end state of the table, + because the saga cleans up after itself — a run that wrote a tombstone, + deleted the knowledge base and then cleared the tombstone leaves the table + looking exactly as untouched as a run that did nothing. + """ + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBORPH1", NOW - timedelta(days=8))], + tags=_ours("KBORPH1", "ast-orph1"), + ) + + _run(client, table, armed=False) + + performed = [op for op, _probe in client.observations if op.startswith("delete_")] + assert performed == [], f"report-only mode issued mutating calls: {performed}" + assert tomb.iter_tombstones("ast-orph1") == [] + assert table.scan()["Items"] == [] + + def test_armed_actually_deletes_through_the_saga(self, table): + """The contrast case, so the report-only assertion means something.""" + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBORPH1", NOW - timedelta(days=8), app_kb_id="ast-orph1")], + tags=_ours("KBORPH1", "ast-orph1"), + ) + + report = _run(client, table, armed=True) + + assert client.delete_calls == ["KBORPH1"] + assert report.deletions_performed == 1 + assert report.planned_deletions[0].performed is True + # The saga cleared its own tombstone once AWS confirmed absence. + assert tomb.iter_tombstones("ast-orph1") == [] + + def test_armed_delete_writes_the_tombstone_before_calling_aws(self, table): + """The orphan path must go through the saga, not a bare delete call.""" + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBORPH1", NOW - timedelta(days=8), app_kb_id="ast-orph1")], + tags=_ours("KBORPH1", "ast-orph1"), + probe=lambda: table.get_item( + Key={"PK": "AST#ast-orph1", "SK": "KBTOMB#ast-orph1"} + ).get("Item") + is not None, + ) + + _run(client, table, armed=True) + + assert client.probes_for("delete_knowledge_base") == [True], ( + "the orphan was deleted without a tombstone in place first" + ) + + def test_an_orphan_tombstone_declares_its_partition_synthetic(self, table): + """An orphan has no assistant id, so its ``PK`` is not a real partition. + + The tombstone still has to exist — a delete that fails mid-flight must + leave a work item either way — but it lands under the ``appKbId`` tag + rather than an assistant, so ``iter_tombstones()`` will never + surface it. Unmarked, that item reads as a tombstone for an assistant that + does not exist, which sends whoever is triaging it looking for a record + that was never there. Asserted while the tombstone is still in place, + i.e. from inside the delete call, because a successful saga clears it. + """ + seen = {} + + def probe(): + item = table.get_item( + Key={"PK": "AST#ast-orph1", "SK": "KBTOMB#ast-orph1"} + ).get("Item") + if item: + seen.update(item) + return item is not None + + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBORPH1", NOW - timedelta(days=8), app_kb_id="ast-orph1")], + tags=_ours("KBORPH1", "ast-orph1"), + probe=probe, + ) + + _run(client, table, armed=True) + + assert seen, "no tombstone was ever written for the orphan" + assert seen.get(tomb.SYNTHETIC_PARTITION) is True, ( + f"the orphan tombstone did not declare its partition synthetic: {dict(seen)}" + ) + # And it says which identifier the partition was derived from, which is the + # first thing an operator needs in order to go find the resource. + assert seen.get("anchorSource") == f"tag:{kb_tags.TAG_KEY_APP_KB_ID}" + assert seen.get("awsKbId") == "KBORPH1" + + def test_a_tombstone_for_a_real_record_is_not_marked_synthetic(self, table): + """The contrast case: the marker must distinguish, not decorate everything. + + A record-backed delete anchors on a genuine assistant partition, so the + flag must be absent there — otherwise it carries no information. + """ + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + probe = {} + + def spy(): + probe.update( + table.get_item(Key={"PK": "AST#ast-real", "SK": "KBTOMB#ast-real"}).get("Item") + or {} + ) + return True + + client.probe = spy + tomb.write_kb_tombstone("ast-real", "ast-real", "KBREAL") + spy() + + assert probe, "the control tombstone was not written" + assert tomb.SYNTHETIC_PARTITION not in probe, ( + f"a record-backed tombstone was flagged synthetic: {dict(probe)}" + ) + + +# ── Requirement 14.3, 14.4: the age gate ───────────────────────────────────── +class TestAgeGateUsesAwsCreatedAt: + def test_an_orphan_aws_calls_old_is_deletable_on_its_first_discovery(self, table): + """TRAP: age-gating on discovery time would skip this. + + The reconciler has never seen this knowledge base before — this is its + first ever run. AWS says the resource is eight days old, so it is + immediately eligible. An implementation that stamped a ``firstSeenAt`` and + waited 24 hours from there would report zero planned deletions here, and + would do so again after every reconciler outage. + """ + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBOLD", NOW - timedelta(days=8))], + tags=_ours("KBOLD", "ast-old"), + ) + + report = _run(client, table, armed=False) + + assert [p.kb_id for p in report.planned_deletions] == ["KBOLD"], ( + "an 8-day-old orphan was not eligible on first discovery, which is " + "what age-gating on discovery time looks like" + ) + assert report.skipped_too_young == [] + + def test_a_freshly_created_knowledge_base_is_left_alone(self, table): + """The other half of the trap: newly discovered is not newly created. + + 30 seconds old by AWS's clock — an in-flight create whose record has not + been attached yet. Deleting this is the failure mode that loses a user's + upload mid-provisioning. + """ + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBNEW", NOW - timedelta(seconds=30))], + tags=_ours("KBNEW", "ast-new"), + ) + + report = _run(client, table, armed=True) + + assert report.planned_deletions == [] + assert report.skipped_too_young == ["KBNEW"] + assert client.delete_calls == [], "an in-flight create was deleted" + + def test_the_boundary_is_twenty_four_hours(self, table): + """23 h 59 m survives; 24 h 01 m does not.""" + client = FakeBedrockAgent( + knowledge_bases=[ + _aws_kb("KBJUSTUNDER", NOW - timedelta(hours=23, minutes=59)), + _aws_kb("KBJUSTOVER", NOW - timedelta(hours=24, minutes=1)), + ], + tags={**_ours("KBJUSTUNDER", "a1"), **_ours("KBJUSTOVER", "a2")}, + ) + + report = _run(client, table, armed=False) + + assert [p.kb_id for p in report.planned_deletions] == ["KBJUSTOVER"] + assert report.skipped_too_young == ["KBJUSTUNDER"] + + def test_the_gate_is_a_pure_function_of_the_aws_timestamp(self): + eight_days = NOW - timedelta(days=8) + thirty_seconds = NOW - timedelta(seconds=30) + + assert rec.orphan_is_deletable(eight_days, now=NOW) is True + assert rec.orphan_is_deletable(thirty_seconds, now=NOW) is False + # Identical answer regardless of when it is asked, which is the property a + # discovery-time clock does not have. + assert rec.orphan_is_deletable(eight_days, now=NOW + timedelta(days=30)) is True + + def test_a_missing_created_at_fails_closed(self): + """No timestamp from AWS means no deletion. Never a guess.""" + assert rec.orphan_is_deletable(None, now=NOW) is False + assert rec.orphan_is_deletable("not-a-date", now=NOW) is False + + def test_an_orphan_without_a_created_at_is_not_deleted(self, table): + kb = _aws_kb("KBNODATE", None) + kb.pop("createdAt") + client = FakeBedrockAgent(knowledge_bases=[kb], tags=_ours("KBNODATE", "a3")) + + report = _run(client, table, armed=True) + + assert report.planned_deletions == [] + assert report.skipped_too_young == ["KBNODATE"] + assert client.delete_calls == [] + + @pytest.mark.parametrize( + "created", + [ + datetime(2026, 5, 1, tzinfo=timezone.utc), + "2026-05-01T00:00:00Z", + datetime(2026, 5, 1).timestamp(), + ], + ) + def test_aws_timestamp_shapes_all_parse(self, created): + """boto3 gives a datetime; a stub or a JSON round-trip gives the others.""" + assert rec.parse_aws_timestamp(created) is not None + + def test_min_age_is_read_at_call_time(self, monkeypatch): + """The threshold must be patchable, not frozen into a default argument.""" + created = NOW - timedelta(hours=2) + assert rec.orphan_is_deletable(created, now=NOW) is False + + monkeypatch.setattr(rec, "ORPHAN_MIN_AGE_HOURS", 1.0) + assert rec.orphan_is_deletable(created, now=NOW) is True + + +# ── Requirement 14.5: record-only never deletes the record ─────────────────── +class TestRecordOnlyMarksMissing: + def test_a_stale_pointer_is_marked_not_removed(self, table): + """TRAP: the record is the only pointer to a recoverable corpus. + + The vectors are gone; the documents are not. Deleting the record would + destroy the mapping the rebuild depends on, and the owner would have to + re-upload. + """ + _seed_record(table, "ast-stale", aws_kb_id="KBGONE") + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table, armed=True) + + assert report.marked_missing == ["ast-stale"] + record = _record(table, "ast-stale") + assert record is not None, ( + "the KB_Record was deleted; its documents are still valid and the " + "knowledge base rebuilds from them on the next ingest" + ) + assert record["vectorState"] == rec.VECTOR_STATE_MISSING + assert record["vectorStateObservedAt"] + + def test_the_documents_and_identifiers_are_left_intact(self, table): + """Nothing else about the record is touched, including its ``DOC#`` rows.""" + _seed_record(table, "ast-stale", aws_kb_id="KBGONE") + table.put_item( + Item={"PK": "AST#ast-stale", "SK": "DOC#doc-1", "status": "complete"} + ) + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + _run(client, table, armed=True) + + record = _record(table, "ast-stale") + assert record["awsKbId"] == "KBGONE" + assert record["retrievalEngine"] == "managed" + doc = table.get_item(Key={"PK": "AST#ast-stale", "SK": "DOC#doc-1"})["Item"] + assert doc["status"] == "complete" + + def test_marking_missing_is_not_a_deletion_even_when_armed(self, table): + """Being armed licenses deleting *orphans*, never records.""" + _seed_record(table, "ast-stale", aws_kb_id="KBGONE") + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table, armed=True) + + assert report.deletions_performed == 0 + assert client.delete_calls == [] + assert _record(table, "ast-stale") is not None + + def test_an_unprovisioned_record_is_not_marked_missing(self, table): + """No ``awsKbId`` means provisioning has not finished, not that AWS lost it.""" + _seed_record(table, "ast-provisioning", aws_kb_id=None) + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table, armed=True) + + assert report.marked_missing == [] + assert _record(table, "ast-provisioning").get("vectorState") is None + + def test_a_tombstone_row_is_not_mistaken_for_a_record(self, table): + """``KBTOMB#`` must not be swept up by the ``KB#`` prefix scan.""" + tomb.write_kb_tombstone("ast-t", "ast-t", "KBX", "DSX") + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table, armed=False) + + assert report.records == 0 + assert report.marked_missing == [] + + +# ── Requirement 14.6: both sides agree ─────────────────────────────────────── +class TestBothSidesRefreshStoredBytes: + def test_stored_bytes_is_re_anchored_from_the_resolver(self, table): + _seed_record(table, "ast-both", aws_kb_id="KBBOTH", storedBytes=10) + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBBOTH", NOW - timedelta(days=8))], + tags=_ours("KBBOTH", "ast-both"), + ) + + report = _run(client, table, armed=False, stored_bytes_resolver=lambda _a: 4096) + + assert report.matched == 1 + assert report.orphans == 0 + assert report.refreshed_bytes == ["ast-both"] + assert int(_record(table, "ast-both")["storedBytes"]) == 4096 + + def test_an_unchanged_total_writes_nothing(self, table): + """A daily no-op write per knowledge base would be pure cost.""" + _seed_record(table, "ast-both", aws_kb_id="KBBOTH", storedBytes=4096) + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBBOTH", NOW - timedelta(days=8))], + tags=_ours("KBBOTH", "ast-both"), + ) + + report = _run(client, table, armed=False, stored_bytes_resolver=lambda _a: 4096) + + assert report.refreshed_bytes == [] + + def test_a_failed_size_lookup_leaves_stored_bytes_alone(self, table): + """Writing a zero on a failed listing hands the owner their quota back.""" + _seed_record(table, "ast-both", aws_kb_id="KBBOTH", storedBytes=4096) + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBBOTH", NOW - timedelta(days=8))], + tags=_ours("KBBOTH", "ast-both"), + ) + + _run(client, table, armed=False, stored_bytes_resolver=lambda _a: None) + + assert int(_record(table, "ast-both")["storedBytes"]) == 4096 + + def test_a_recovered_knowledge_base_clears_a_stale_missing_marker(self, table): + """Otherwise the UI keeps reporting a knowledge base broken after the fix.""" + _seed_record( + table, + "ast-both", + aws_kb_id="KBBOTH", + storedBytes=4096, + vectorState=rec.VECTOR_STATE_MISSING, + ) + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBBOTH", NOW - timedelta(days=8))], + tags=_ours("KBBOTH", "ast-both"), + ) + + _run(client, table, armed=False, stored_bytes_resolver=lambda _a: 4096) + + assert _record(table, "ast-both").get("vectorState") is None + + def test_stored_bytes_from_s3_totals_the_prefix(self, table): + class FakeS3: + def list_objects_v2(self, **kwargs): + assert kwargs["Prefix"] == "assistants/ast-s3/documents/" + return {"Contents": [{"Size": 100}, {"Size": 23}], "IsTruncated": False} + + assert rec.stored_bytes_from_s3("ast-s3", bucket="b", s3_client=FakeS3()) == 123 + + def test_stored_bytes_from_s3_returns_none_on_failure(self, table): + class Boom: + def list_objects_v2(self, **kwargs): + raise RuntimeError("access denied") + + assert rec.stored_bytes_from_s3("ast-s3", bucket="b", s3_client=Boom()) is None + + +# ── Requirement 14.8: bounded per-run action limit ─────────────────────────── +class TestPerRunActionLimit: + def _five_orphans(self): + kbs = [_aws_kb(f"KBORPH{i}", NOW - timedelta(days=8)) for i in range(5)] + tags = {} + for i in range(5): + tags.update(_ours(f"KBORPH{i}", f"ast-orph{i}")) + return FakeBedrockAgent(knowledge_bases=kbs, tags=tags) + + def test_the_limit_caps_planned_deletions_in_report_only_mode(self, table, monkeypatch): + """The report must describe what an armed run would really do. + + A report listing five intended deletions from a run that would only ever + perform two is a misleading artifact, and the report-only period exists + precisely so the artifact can be trusted. + """ + monkeypatch.setattr(rec, "MAX_DELETIONS_PER_RUN", 2) + client = self._five_orphans() + + report = _run(client, table, armed=False) + + assert report.orphans == 5 + assert len(report.planned_deletions) == 2 + assert report.limit_reached is True + + def test_the_limit_caps_actual_deletions_when_armed(self, table, monkeypatch): + monkeypatch.setattr(rec, "MAX_DELETIONS_PER_RUN", 2) + client = self._five_orphans() + + report = _run(client, table, armed=True) + + assert len(client.delete_calls) == 2, ( + f"the per-run limit did not bound the deletions: {client.delete_calls}" + ) + assert report.deletions_performed == 2 + assert report.limit_reached is True + + def test_without_the_limit_being_hit_nothing_is_flagged(self, table, monkeypatch): + monkeypatch.setattr(rec, "MAX_DELETIONS_PER_RUN", 25) + client = self._five_orphans() + + report = _run(client, table, armed=False) + + assert len(report.planned_deletions) == 5 + assert report.limit_reached is False + + def test_the_limit_is_read_at_call_time(self, monkeypatch): + assert rec.max_deletions_per_run() == rec.MAX_DELETIONS_PER_RUN + monkeypatch.setattr(rec, "MAX_DELETIONS_PER_RUN", 3) + assert rec.max_deletions_per_run() == 3 + monkeypatch.setenv("MANAGED_KB_RECONCILER_MAX_DELETIONS", "7") + assert rec.max_deletions_per_run() == 7 + + def test_the_environment_can_lower_the_limit_but_not_lift_it(self, monkeypatch): + """A bound an env var can raise without limit is not a bound. + + This is the only limit whose failure mode is irreversible, so the ceiling + has to hold against the variable rather than merely default below it. + """ + monkeypatch.setenv("MANAGED_KB_RECONCILER_MAX_DELETIONS", "3") + assert rec.max_deletions_per_run() == 3, "the env var could not lower the limit" + + monkeypatch.setenv("MANAGED_KB_RECONCILER_MAX_DELETIONS", "1000000") + assert rec.max_deletions_per_run() == rec.MAX_DELETIONS_CEILING, ( + "the environment lifted the per-run deletion bound past its ceiling" + ) + + def test_a_negative_limit_does_not_become_unbounded(self, monkeypatch): + """A negative slice bound would silently mean 'all of them' downstream.""" + monkeypatch.setenv("MANAGED_KB_RECONCILER_MAX_DELETIONS", "-5") + assert rec.max_deletions_per_run() == 0 + + +# ── Requirement 14.1: paginated and tag-filtered ───────────────────────────── +class TestJoinIsPaginatedAndTagFiltered: + def test_orphans_on_later_pages_are_still_found(self, table): + """Reading only page one would make account size decide correctness.""" + kbs = [_aws_kb(f"KBP{i}", NOW - timedelta(days=8)) for i in range(5)] + tags = {} + for i in range(5): + tags.update(_ours(f"KBP{i}", f"ast-p{i}")) + client = FakeBedrockAgent(knowledge_bases=kbs, tags=tags, page_size=2) + + report = _run(client, table, armed=False) + + assert report.aws_knowledge_bases == 5 + assert len(report.planned_deletions) == 5 + + def test_another_projects_knowledge_base_is_invisible(self, table): + client = FakeBedrockAgent( + knowledge_bases=[ + _aws_kb("KBMINE", NOW - timedelta(days=8)), + _aws_kb("KBTHEIRS", NOW - timedelta(days=8)), + ], + tags={ + **_ours("KBMINE", "ast-mine"), + _arn("KBTHEIRS"): kb_tags.build_tags("ast-theirs", "u-2", "other-project", "prod"), + }, + ) + + report = _run(client, table, armed=True) + + assert report.aws_knowledge_bases == 1 + assert client.delete_calls == ["KBMINE"], ( + "the reconciler acted outside its tag scope" + ) + + def test_an_untagged_knowledge_base_is_never_deleted(self, table): + client = FakeBedrockAgent( + knowledge_bases=[_aws_kb("KBBARE", NOW - timedelta(days=8))], tags={} + ) + + report = _run(client, table, armed=True) + + assert report.aws_knowledge_bases == 0 + assert client.delete_calls == [] + + def test_a_truncated_aws_walk_suppresses_missing_vector_marks(self, table, monkeypatch): + """An unmatched record on a partial walk may be one we never reached.""" + monkeypatch.setattr(rec, "MAX_KNOWLEDGE_BASES_PER_RUN", 1) + _seed_record(table, "ast-a", aws_kb_id="KBA") + _seed_record(table, "ast-b", aws_kb_id="KBB") + client = FakeBedrockAgent( + knowledge_bases=[ + _aws_kb("KBA", NOW - timedelta(days=8)), + _aws_kb("KBB", NOW - timedelta(days=8)), + ], + tags={**_ours("KBA", "ast-a"), **_ours("KBB", "ast-b")}, + ) + + report = _run(client, table, armed=True) + + assert report.limit_reached is True + assert report.marked_missing == [] + assert _record(table, "ast-b").get("vectorState") is None + + +# ── Requirement 13.7 seen from the reconciler ──────────────────────────────── +class TestDeleteUnsuccessfulOrphan: + def test_it_is_surfaced_and_not_retried(self, table): + """Retrying does not help and the resource keeps billing.""" + client = FakeBedrockAgent( + knowledge_bases=[ + _aws_kb("KBSTUCK", NOW - timedelta(days=200), status="DELETE_UNSUCCESSFUL") + ], + tags=_ours("KBSTUCK", "ast-stuck"), + ) + + report = _run(client, table, armed=True) + + assert len(report.planned_deletions) == 1 + planned = report.planned_deletions[0] + assert planned.error == tomb.KB_STATUS_DELETE_UNSUCCESSFUL + assert planned.performed is False + assert client.delete_calls == [] + + def test_it_appears_in_the_serialized_report(self, table): + client = FakeBedrockAgent( + knowledge_bases=[ + _aws_kb("KBSTUCK", NOW - timedelta(days=200), status="DELETE_UNSUCCESSFUL") + ], + tags=_ours("KBSTUCK", "ast-stuck"), + ) + + payload = _run(client, table, armed=False).to_dict() + + assert payload["plannedDeletions"][0]["status"] == "DELETE_UNSUCCESSFUL" + assert payload["deletionsPerformed"] == 0 + + +# ── Mixed and degenerate cases ─────────────────────────────────────────────── +class TestMixedRun: + def test_all_three_outcomes_in_one_pass(self, table): + _seed_record(table, "ast-both", aws_kb_id="KBBOTH", storedBytes=1) + _seed_record(table, "ast-stale", aws_kb_id="KBVANISHED") + client = FakeBedrockAgent( + knowledge_bases=[ + _aws_kb("KBBOTH", NOW - timedelta(days=8)), + _aws_kb("KBORPH", NOW - timedelta(days=8)), + ], + tags={**_ours("KBBOTH", "ast-both"), **_ours("KBORPH", "ast-orph")}, + ) + + report = _run(client, table, armed=False, stored_bytes_resolver=lambda _a: 99) + + assert report.records == 2 + assert report.matched == 1 + assert report.orphans == 1 + assert report.marked_missing == ["ast-stale"] + assert report.refreshed_bytes == ["ast-both"] + assert [p.kb_id for p in report.planned_deletions] == ["KBORPH"] + assert _record(table, "ast-stale") is not None + + def test_an_empty_account_and_empty_table_is_a_clean_no_op(self, table): + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table, armed=True) + + assert report.to_dict() == { + "armed": True, + "mode": "armed", + "awsKnowledgeBases": 0, + "records": 0, + "matched": 0, + "orphans": 0, + "plannedDeletions": [], + "deletionsPerformed": 0, + "skippedTooYoung": [], + "markedMissing": [], + "refreshedBytes": [], + "limitReached": False, + # Fleet gauges. Zero here, and asserted as an exact dict on purpose: the + # report is a stored artifact an operator reads, so a field appearing or + # vanishing should be a deliberate change to this list. + "storedBytes": 0, + "idleBytes": 0, + "unmeasuredIdleness": 0, + } + + def test_a_failing_delete_does_not_end_the_run(self, table, monkeypatch): + """One stuck orphan must not stop the reconciler reaching the others.""" + client = FakeBedrockAgent( + knowledge_bases=[ + _aws_kb("KBA", NOW - timedelta(days=8)), + _aws_kb("KBB", NOW - timedelta(days=8)), + ], + tags={**_ours("KBA", "ast-a"), **_ours("KBB", "ast-b")}, + polls_before_gone=10_000, + ) + monkeypatch.setattr(tomb, "KB_DELETE_POLL_TIMEOUT_SECONDS", 0.0) + monkeypatch.setattr(tomb, "KB_DELETE_POLL_INTERVAL_SECONDS", 0.0) + + report = _run(client, table, armed=True) + + assert len(report.planned_deletions) == 2 + assert report.deletions_performed == 0 + assert all(p.error for p in report.planned_deletions) + # And the tombstones survive as work items for the next run. + assert tomb.iter_tombstones("ast-a") + assert tomb.iter_tombstones("ast-b") + + +class TestLambdaHandler: + """The scheduled entry point, and the one input nobody reviews. + + ``lambda_handler`` takes an *event*. An event is not reviewable configuration: + an EventBridge target can carry a constant payload, and any principal with + ``lambda:InvokeFunction`` can supply one. So the flag has to be the only way + to arm (Requirement 19.7) — otherwise deletion of billed user resources is + reachable while every reviewable setting still reads report-only, and the only + trace left is an ``Invoke`` in CloudTrail. + """ + + @pytest.fixture() + def stub_client(self, monkeypatch): + """Make the un-injected client path safe: no AWS, and a delete log to read. + + ``lambda_handler`` deliberately passes no client, so this patches the + factory ``reconcile`` reaches for. Without it the test would try to build + a real ``bedrock-agent`` client (Requirement 24.11). + """ + from apis.shared.kb_backend import managed_backend + + # 2020: comfortably older than the 24h gate against real wall-clock time, + # since lambda_handler passes no ``now``. + client = FakeBedrockAgent( + knowledge_bases=[ + _aws_kb("KBORPH1", datetime(2020, 1, 1, tzinfo=timezone.utc), app_kb_id="ast-orph1") + ], + tags=_ours("KBORPH1", "ast-orph1"), + ) + monkeypatch.setattr(managed_backend, "bedrock_agent_client", lambda: client) + return client + + def test_it_returns_the_serialized_report(self, table, monkeypatch): + monkeypatch.setattr(rec, "reconcile", lambda **kwargs: rec.ReconcileReport(armed=False)) + + result = rec.lambda_handler({}, None) + + assert result["statusCode"] == 200 + assert result["report"]["mode"] == "report-only" + + @pytest.mark.parametrize("payload", [True, "true", 1, "1", "yes"]) + def test_the_event_cannot_arm_the_reconciler(self, table, stub_client, payload): + """A flag-off invocation carrying ``armed`` deletes nothing. + + Parametrised over a real boolean and the string/int spellings alike, + because the boolean is the one that would previously have worked: an + ``isinstance(x, bool)`` override honours ``True`` exactly, so a test that + only passed ``"true"`` proved nothing about the path that actually armed. + """ + result = rec.lambda_handler({"armed": payload}, None) + + assert result["report"]["mode"] == "report-only", ( + f"the event payload armed={payload!r} put the reconciler in armed mode" + ) + assert result["report"]["deletionsPerformed"] == 0 + assert stub_client.delete_calls == [], ( + f"the event payload armed={payload!r} caused a real DeleteKnowledgeBase" + ) + # And the orphan it declined to delete is still reported, so suppressing + # the delete has not also suppressed the finding. + assert result["report"]["orphans"] == 1 + + def test_the_flag_is_what_arms_it(self, table, stub_client, monkeypatch): + """The contrast case: same event, same orphan, flag on — now it deletes. + + Without this, the assertions above would also pass on a reconciler that + could never delete at all. + """ + monkeypatch.setenv(rec.FLAG_RECONCILER_ARMED, "true") + + result = rec.lambda_handler({"armed": False}, None) + + assert result["report"]["mode"] == "armed" + assert stub_client.delete_calls == ["KBORPH1"] + assert result["report"]["deletionsPerformed"] == 1 + + def test_an_ignored_arming_request_is_logged(self, table, stub_client, caplog): + """Silently dropping the field would hide a misconfigured schedule.""" + import logging + + with caplog.at_level(logging.WARNING): + rec.lambda_handler({"armed": True}, None) + + assert any( + "ignoring armed" in r.message and rec.FLAG_RECONCILER_ARMED in r.message + for r in caplog.records + ), f"no warning named the ignored override: {[r.message for r in caplog.records]}" + + +# ── Requirements 22.1, 22.5: the fleet gauges ──────────────────────────────── +class TestFleetGaugeAccumulation: + """The reconciler is where the gauges are computed, because it is already the + one pass that walks every knowledge base.""" + + def test_stored_bytes_sum_across_records(self, table): + _seed_record(table, "ast-a", aws_kb_id="KBA", storedBytes=3_000_000_000) + _seed_record(table, "ast-b", aws_kb_id="KBB", storedBytes=1_000_000_000) + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table) + + assert report.records == 2 + assert report.stored_bytes == 4_000_000_000 + + def test_an_agent_used_today_is_not_idle_however_stale_its_retrievals(self, table): + """Requirement 22.5, end to end through the reconciler. + + The knowledge base was last retrieved from 200 days ago but its agent was + used today — an agent answering questions its documents do not cover. Judged + by retrieval alone its bytes would count as idle and the follow-up spec's + eviction pass would delete a live corpus. + """ + _seed_record( + table, + "ast-busy", + aws_kb_id="KBA", + storedBytes=5_000_000_000, + lastRetrievedAt=_iso(NOW - timedelta(days=200)), + ) + table.put_item( + Item={ + "PK": "AST#ast-busy", + "SK": "METADATA", + "lastUsedAt": _iso(NOW - timedelta(hours=2)), + } + ) + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table) + + assert report.stored_bytes == 5_000_000_000 + assert report.idle_bytes == 0, ( + "a busy agent's corpus was counted as idle; idleness was derived from " + "retrieval alone" + ) + + def test_a_genuinely_dormant_knowledge_base_counts_as_idle(self, table): + """So the test above cannot pass by never counting anything.""" + _seed_record( + table, + "ast-cold", + aws_kb_id="KBA", + storedBytes=2_000_000_000, + lastRetrievedAt=_iso(NOW - timedelta(days=200)), + ) + table.put_item( + Item={ + "PK": "AST#ast-cold", + "SK": "METADATA", + "lastUsedAt": _iso(NOW - timedelta(days=180)), + } + ) + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table) + + assert report.idle_bytes == 2_000_000_000 + + def test_a_knowledge_base_with_no_activity_signal_is_unmeasured_not_idle(self, table): + """What a corpus provisioned an hour ago looks like. Counting it as idle + would report every new knowledge base as abandoned.""" + _seed_record(table, "ast-new", aws_kb_id="KBA", storedBytes=9_000_000_000) + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table) + + assert report.unmeasured_idleness == 1 + assert report.idle_bytes == 0 + + def test_the_gauges_are_emitted_once_per_pass(self, table, monkeypatch): + emitted = [] + monkeypatch.setattr(rec, "emit_fleet_gauges", lambda **kw: emitted.append(kw)) + _seed_record(table, "ast-a", aws_kb_id="KBA", storedBytes=1_000_000_000) + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + _run(client, table) + + assert len(emitted) == 1 + assert emitted[0]["kb_count"] == 1 + assert emitted[0]["stored_bytes"] == 1_000_000_000 + + def test_an_idleness_failure_does_not_end_the_pass(self, table, monkeypatch): + """A gauge is never worth a reconciliation.""" + _seed_record(table, "ast-a", aws_kb_id="KBA", storedBytes=1_000_000_000) + monkeypatch.setattr( + "apis.shared.kb_backend.idleness.idle_days", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")), + ) + client = FakeBedrockAgent(knowledge_bases=[], tags={}) + + report = _run(client, table) + + assert report.records == 1 + assert report.stored_bytes == 1_000_000_000 + + def test_the_idle_threshold_is_resolved_at_call_time(self, monkeypatch): + from apis.shared.kb_backend.metrics import IDLE_THRESHOLD_DAYS + + monkeypatch.delenv("KB_IDLE_THRESHOLD_DAYS", raising=False) + assert rec.idle_threshold_days() == IDLE_THRESHOLD_DAYS + monkeypatch.setenv("KB_IDLE_THRESHOLD_DAYS", "7") + assert rec.idle_threshold_days() == 7 diff --git a/backend/tests/property/test_pbt_cleanup_service.py b/backend/tests/property/test_pbt_cleanup_service.py index cca1f099b..92d429d81 100644 --- a/backend/tests/property/test_pbt_cleanup_service.py +++ b/backend/tests/property/test_pbt_cleanup_service.py @@ -104,6 +104,15 @@ def failing_s3_delete(**kwargs): side_effect=failing_fallback, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.cleanup_service.asyncio.sleep", side_effect=mock_sleep, @@ -208,6 +217,15 @@ async def failing_fallback(*args, **kwargs): side_effect=failing_fallback, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.cleanup_service.asyncio.sleep", new_callable=AsyncMock, @@ -286,6 +304,15 @@ async def succeeding_fallback(*args, **kwargs): side_effect=succeeding_fallback, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.document_service.hard_delete_document", mock_hard_delete, diff --git a/backend/tests/property/test_pbt_kb_byte_cap.py b/backend/tests/property/test_pbt_kb_byte_cap.py new file mode 100644 index 000000000..2a9538b8d --- /dev/null +++ b/backend/tests/property/test_pbt_kb_byte_cap.py @@ -0,0 +1,328 @@ +"""Property-based tests for byte cap accounting. + +Feature: managed-kb-migration + +**Property 5: the cap is never exceeded, under any interleaving.** + +Managed storage costs $5.00/GB-month, so the cap is the only thing standing between +the measured ~$169/month fleet cost and the ~$15,000/month that unbounded uploads +would permit. "Usually holds" is not a cap. + +The property is asserted against real DynamoDB semantics (via moto) rather than +against a Python model of them, because the entire correctness argument rests on +one specific database behaviour: that a conditional ``ADD`` is atomic. A test that +simulated the arithmetic in Python would pass just as happily against a +read-then-write implementation, which is precisely the broken version. + +Why the accumulator matters +--------------------------- +DynamoDB cannot do arithmetic inside a condition expression — verified, it fails to +parse. So the guard compares a single ``totalBytes`` accumulator against a literal +computed before the call (``cap - n``). The invariant +``totalBytes == storedBytes + reservedBytes`` is what makes that sound, and several +tests below assert it directly rather than only checking the total. + +Validates: Requirements 12.4, 12.5, 12.6, 24.7. +""" + +import boto3 +import pytest +from hypothesis import HealthCheck, given, settings, strategies as st +from moto import mock_aws + +from apis.shared.kb_backend import byte_cap as bc +from apis.shared.kb_backend.records import kb_pk, kb_sk + +REGION = "us-east-1" +TABLE = "test-byte-cap" +ASSISTANT_ID = "ast-cap01" +APP_KB_ID = ASSISTANT_ID +CAP = 1000 + +# --------------------------------------------------------------------------- +# Strategies +# --------------------------------------------------------------------------- + +#: Reservation sizes, including 0 (a no-op) and sizes larger than the whole cap. +st_size = st.integers(min_value=0, max_value=CAP + 500) + +#: An arbitrary sequence of reservations. Length and sizes both vary so the +#: sequence sometimes fits entirely, sometimes overruns partway, and sometimes +#: overruns on the very first item. +st_sequence = st.lists(st_size, min_size=1, max_size=15) + + +@pytest.fixture() +def table(monkeypatch): + monkeypatch.setenv("AWS_DEFAULT_REGION", REGION) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing") + monkeypatch.setenv("AWS_SESSION_TOKEN", "testing") + monkeypatch.setenv("DYNAMODB_ASSISTANTS_TABLE_NAME", TABLE) + + with mock_aws(): + ddb = boto3.client("dynamodb", region_name=REGION) + ddb.create_table( + TableName=TABLE, + KeySchema=[ + {"AttributeName": "PK", "KeyType": "HASH"}, + {"AttributeName": "SK", "KeyType": "RANGE"}, + ], + AttributeDefinitions=[ + {"AttributeName": "PK", "AttributeType": "S"}, + {"AttributeName": "SK", "AttributeType": "S"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + t = boto3.resource("dynamodb", region_name=REGION).Table(TABLE) + t.put_item(Item={"PK": kb_pk(ASSISTANT_ID), "SK": kb_sk(APP_KB_ID)}) + yield t + + +def _counters(table): + item = table.get_item(Key={"PK": kb_pk(ASSISTANT_ID), "SK": kb_sk(APP_KB_ID)})["Item"] + return ( + int(item.get("totalBytes", 0)), + int(item.get("reservedBytes", 0)), + int(item.get("storedBytes", 0)), + ) + + +def _reset(table): + table.put_item(Item={"PK": kb_pk(ASSISTANT_ID), "SK": kb_sk(APP_KB_ID)}) + + +# --------------------------------------------------------------------------- +# The cap holds +# --------------------------------------------------------------------------- +@given(sizes=st_sequence) +@settings(max_examples=60, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) +def test_the_cap_is_never_exceeded(table, sizes): + """However the sequence interleaves, the accumulator never passes the cap.""" + _reset(table) + accepted = [] + for n in sizes: + try: + bc.reserve(ASSISTANT_ID, APP_KB_ID, n, CAP) + accepted.append(n) + except bc.ByteCapExceeded: + pass + + total, _, _ = _counters(table) + assert total <= CAP, f"cap breached at {total} > {CAP}" + + total, reserved, _ = _counters(table) + assert total == sum(accepted) + assert reserved == sum(accepted) + + +@given(sizes=st_sequence) +@settings(max_examples=60, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) +def test_the_accumulator_invariant_holds(table, sizes): + """totalBytes == storedBytes + reservedBytes, always. + + This is what makes comparing a single attribute a valid cap check. If the two + ever diverge the guard is measuring something that is not the owner's usage. + """ + _reset(table) + for n in sizes: + try: + bc.reserve(ASSISTANT_ID, APP_KB_ID, n, CAP) + # Commit half the time so both counters move. + if n % 2 == 0: + bc.commit(ASSISTANT_ID, APP_KB_ID, n) + except bc.ByteCapExceeded: + pass + + total, reserved, stored = _counters(table) + assert total == reserved + stored, f"{total} != {reserved} + {stored}" + + +@given(sizes=st.lists(st.integers(min_value=1, max_value=200), min_size=1, max_size=10)) +@settings(max_examples=60, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) +def test_released_reservations_are_fully_returned(table, sizes): + """Release restores the allowance exactly. + + A release that returned less than it reserved would shrink the owner's cap on + every failed upload, presenting weeks later as "uploads stopped working" with + no failing request to point at. + """ + _reset(table) + for n in sizes: + bc.reserve(ASSISTANT_ID, APP_KB_ID, n, CAP) + bc.release(ASSISTANT_ID, APP_KB_ID, n) + + total, reserved, stored = _counters(table) + assert (total, reserved, stored) == (0, 0, 0) + + +@given(sizes=st.lists(st.integers(min_value=1, max_value=100), min_size=1, max_size=8)) +@settings(max_examples=60, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) +def test_commit_does_not_double_count(table, sizes): + """Commit moves bytes; it must not add them again. + + Double-counting on commit would halve every owner's effective allowance, and it + would do so only for *successful* uploads — so the symptom would be that the + cap tightens the more correctly the system works. + """ + _reset(table) + for n in sizes: + bc.reserve(ASSISTANT_ID, APP_KB_ID, n, CAP) + bc.commit(ASSISTANT_ID, APP_KB_ID, n) + + total, reserved, stored = _counters(table) + assert total == sum(sizes) + assert stored == sum(sizes) + assert reserved == 0 + + +# --------------------------------------------------------------------------- +# Boundary and rejection behaviour +# --------------------------------------------------------------------------- +def test_a_reservation_exactly_filling_the_cap_is_allowed(table): + """The cap is inclusive: exactly at the limit is within it.""" + _reset(table) + bc.reserve(ASSISTANT_ID, APP_KB_ID, CAP, CAP) + assert _counters(table)[0] == CAP + + +def test_one_byte_over_is_rejected(table): + _reset(table) + with pytest.raises(bc.ByteCapExceeded): + bc.reserve(ASSISTANT_ID, APP_KB_ID, CAP + 1, CAP) + assert _counters(table)[0] == 0, "a rejected reservation must leave no trace" + + +def test_a_rejected_reservation_does_not_consume_allowance(table): + """The failed attempt must not partially apply. + + An ADD that landed before the condition was evaluated would leak allowance on + every rejection, so a user who hit the cap once could never upload again. + """ + _reset(table) + bc.reserve(ASSISTANT_ID, APP_KB_ID, 900, CAP) + with pytest.raises(bc.ByteCapExceeded): + bc.reserve(ASSISTANT_ID, APP_KB_ID, 200, CAP) + + total, reserved, _ = _counters(table) + assert (total, reserved) == (900, 900) + # And the remaining allowance is still usable. + bc.reserve(ASSISTANT_ID, APP_KB_ID, 100, CAP) + assert _counters(table)[0] == CAP + + +def test_zero_is_a_no_op(table): + _reset(table) + bc.reserve(ASSISTANT_ID, APP_KB_ID, 0, CAP) + assert _counters(table) == (0, 0, 0) + + +def test_a_negative_reservation_is_rejected(table): + """Otherwise 'reserving' a negative size would be a way to mint allowance.""" + _reset(table) + with pytest.raises(ValueError): + bc.reserve(ASSISTANT_ID, APP_KB_ID, -100, CAP) + + +def test_the_exception_carries_the_numbers_for_the_user(table): + """Requirement 12.12 wants a plain-language reason and an upgrade path, which + needs the figures, not just a failure.""" + _reset(table) + with pytest.raises(bc.ByteCapExceeded) as excinfo: + bc.reserve(ASSISTANT_ID, APP_KB_ID, CAP + 1, CAP) + assert excinfo.value.requested == CAP + 1 + assert excinfo.value.cap == CAP + + +# --------------------------------------------------------------------------- +# Migration snapshot (Requirement 12.11/12.12) +# --------------------------------------------------------------------------- +def test_a_snapshot_that_cannot_fit_is_rejected_up_front(table): + """The whole corpus is reserved before migration starts. + + Reserving per-document instead would let a migration run for an hour and stop + halfway, leaving a half-populated managed knowledge base behind. + """ + _reset(table) + with pytest.raises(bc.ByteCapExceeded): + bc.reserve_snapshot(ASSISTANT_ID, APP_KB_ID, CAP * 2, CAP) + assert _counters(table)[0] == 0, "a rejected migration must reserve nothing" + + +def test_a_snapshot_that_fits_reserves_the_whole_corpus(table): + _reset(table) + bc.reserve_snapshot(ASSISTANT_ID, APP_KB_ID, 800, CAP) + total, reserved, _ = _counters(table) + assert (total, reserved) == (800, 800) + + +def test_a_snapshot_is_rejected_when_existing_usage_leaves_no_room(table): + """The interesting case: the corpus fits an empty cap but not this owner's.""" + _reset(table) + bc.reserve(ASSISTANT_ID, APP_KB_ID, 700, CAP) + bc.commit(ASSISTANT_ID, APP_KB_ID, 700) + + with pytest.raises(bc.ByteCapExceeded): + bc.reserve_snapshot(ASSISTANT_ID, APP_KB_ID, 400, CAP) + + assert _counters(table)[0] == 700 + + +# --------------------------------------------------------------------------- +# Cap resolution +# --------------------------------------------------------------------------- +def test_the_default_cap_is_below_the_user_files_precedent(monkeypatch): + """100 MB, deliberately under the existing 1 GB user-files limit. + + At $5.00/GB-month that precedent would permit roughly $150,000/month across the + fleet — a number large enough that it is not really a limit. + """ + monkeypatch.delenv("MANAGED_KB_PER_OWNER_DEFAULT_BYTES", raising=False) + assert bc.per_owner_cap() == 100 * 1024 * 1024 + assert bc.per_owner_cap() < 1024 * 1024 * 1024 + + +def test_the_elevated_tier_is_larger_than_the_default(monkeypatch): + monkeypatch.delenv("MANAGED_KB_PER_OWNER_DEFAULT_BYTES", raising=False) + monkeypatch.delenv("MANAGED_KB_PER_OWNER_ELEVATED_BYTES", raising=False) + assert bc.per_owner_cap(elevated=True) > bc.per_owner_cap() + + +def test_caps_are_overridable_from_the_environment(monkeypatch): + monkeypatch.setenv("MANAGED_KB_PER_OWNER_DEFAULT_BYTES", "12345") + assert bc.per_owner_cap() == 12345 + + +def test_a_malformed_override_falls_back_rather_than_crashing(monkeypatch): + """A typo in an operator-set variable must not take retrieval down.""" + monkeypatch.setenv("MANAGED_KB_PER_OWNER_DEFAULT_BYTES", "not-a-number") + assert bc.per_owner_cap() == 100 * 1024 * 1024 + + +def test_the_per_kb_ceiling_is_below_the_elevated_owner_cap(monkeypatch): + """A single knowledge base must not be able to eat an entire elevated + allowance and starve the owner's others.""" + for var in ( + "MANAGED_KB_PER_KB_CEILING_BYTES", + "MANAGED_KB_PER_OWNER_ELEVATED_BYTES", + ): + monkeypatch.delenv(var, raising=False) + assert bc.per_kb_ceiling() < bc.per_owner_cap(elevated=True) + + +# --------------------------------------------------------------------------- +# Sizing authority +# --------------------------------------------------------------------------- +def test_size_comes_from_s3_not_from_the_caller(monkeypatch): + """A client-reported size is an input, and an input that can lower its own + cost is not a measurement.""" + from unittest.mock import MagicMock, patch + + monkeypatch.setenv("AWS_DEFAULT_REGION", REGION) + with patch("boto3.client") as client: + s3 = MagicMock() + s3.head_object.return_value = {"ContentLength": 4242} + client.return_value = s3 + + assert bc.object_size_bytes("bucket", "key") == 4242 + s3.head_object.assert_called_once_with(Bucket="bucket", Key="key") diff --git a/backend/tests/property/test_pbt_kb_engine_resolution.py b/backend/tests/property/test_pbt_kb_engine_resolution.py new file mode 100644 index 000000000..8f1ccf2be --- /dev/null +++ b/backend/tests/property/test_pbt_kb_engine_resolution.py @@ -0,0 +1,205 @@ +"""Property-based tests for engine resolution by absence. + +Feature: managed-kb-migration + +**Property 1: absence means legacy.** + +This is the invariant the whole migration rests on. Every knowledge base that +existed before this feature carries no ``retrievalEngine`` attribute, and must +resolve to the legacy backend on that basis alone. Two consequences follow, and +both are why this file exists: + +* **No backfill.** 1,692 ``DOC#`` records and their knowledge bases are already + correct without being touched. A migration that had to stamp a value on each + one would be a data migration in its own right, with its own failure modes. +* **Rollback is a pointer flip.** Rolling back ``REMOVE``s the attribute, + restoring the original shape exactly. A rolled-back record is + indistinguishable from one that never migrated. + +Both consequences evaporate the moment any code path writes the literal +``"s3vectors"`` onto a record that did not already carry it. That write would +look harmless, pass a naive test, and convert every future rollback into a +rewrite. The second half of this file exists to make that specific mistake fail +loudly. + +Validates: Requirements 1.6, 1.7, 6.6. +""" + +import json +from typing import Any, Dict, List + +import pytest +from hypothesis import given, settings, strategies as st + +from apis.shared.kb_backend import records as r + +# --------------------------------------------------------------------------- +# Shared Hypothesis strategies +# --------------------------------------------------------------------------- + +st_attribute_name = st.text( + alphabet="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_", + min_size=1, + max_size=24, +) + +st_attribute_value = st.one_of( + st.text(max_size=40), + st.integers(min_value=-1000, max_value=10**9), + st.booleans(), + st.none(), + st.lists(st.text(max_size=10), max_size=4), + st.dictionaries(st.text(min_size=1, max_size=8), st.integers(), max_size=3), +) + +#: An arbitrary stored item that carries no opinion about its engine. Extra keys +#: are deliberately unconstrained: real records accumulate attributes over time +#: and resolution must not depend on which ones happen to be present. +st_item_without_engine = st.dictionaries( + st_attribute_name, st_attribute_value, max_size=12 +).map(lambda d: {k: v for k, v in d.items() if k != "retrievalEngine"}) + +#: Anything that is not the one value we accept. Includes the legacy literal +#: itself: even if some historical record somehow carried "s3vectors", it must +#: resolve to legacy, which it does — but it must never be *written*. +st_non_managed_engine = st.one_of( + st.just(r.ENGINE_LEGACY), + st.just(""), + st.just("Managed"), + st.just("MANAGED"), + st.just("managed "), + st.text(max_size=20).filter(lambda s: s != r.ENGINE_MANAGED), +) + + +# --------------------------------------------------------------------------- +# Property 1: absence means legacy +# --------------------------------------------------------------------------- +@given(item=st_item_without_engine) +@settings(max_examples=200) +def test_any_record_without_the_attribute_resolves_to_legacy(item): + """No matter what else the record contains, a missing engine means legacy.""" + assert "retrievalEngine" not in item + assert r.resolve_engine(item) == r.ENGINE_LEGACY + + +@given(item=st_item_without_engine, engine=st_non_managed_engine) +@settings(max_examples=200) +def test_only_the_exact_managed_literal_selects_the_managed_backend(item, engine): + """Resolution is exact-match, so a typo or casing slip fails safe. + + Failing safe matters asymmetrically here: resolving to legacy when it should + be managed serves slightly worse answers, while resolving to managed when the + record is not really migrated queries a knowledge base that may not exist. + """ + item["retrievalEngine"] = engine + assert r.resolve_engine(item) == r.ENGINE_LEGACY + + +@given(item=st_item_without_engine) +@settings(max_examples=100) +def test_the_managed_literal_selects_managed(item): + """The positive case, so the tests above cannot pass by always returning legacy.""" + item["retrievalEngine"] = r.ENGINE_MANAGED + assert r.resolve_engine(item) == r.ENGINE_MANAGED + + +@pytest.mark.parametrize("empty", [None, {}]) +def test_a_missing_record_resolves_to_legacy(empty): + """Absence of the whole record is an answer too, not an error. + + A knowledge base with no KB_Record is every knowledge base today. + """ + assert r.resolve_engine(empty) == r.ENGINE_LEGACY + + +@given(item=st_item_without_engine) +@settings(max_examples=100) +def test_resolution_does_not_mutate_the_item(item): + """Resolution is a read. A resolver that defaulted the attribute *in place* + would silently create the backfill this design exists to avoid.""" + before = json.dumps(item, sort_keys=True, default=str) + r.resolve_engine(item) + assert json.dumps(item, sort_keys=True, default=str) == before + + +# --------------------------------------------------------------------------- +# Property 1, second half: the legacy literal is never written +# --------------------------------------------------------------------------- +class _RecordingTable: + """Captures write payloads instead of performing them. + + Used rather than moto because the assertion here is about what the module + *sends*, not about what DynamoDB does with it — and because it lets a single + test observe every transition without needing each one's preconditions to + hold. + """ + + def __init__(self) -> None: + self.calls: List[Dict[str, Any]] = [] + + def put_item(self, **kwargs): + self.calls.append(kwargs) + return {} + + def update_item(self, **kwargs): + self.calls.append(kwargs) + return {} + + def serialized(self) -> str: + return json.dumps(self.calls, sort_keys=True, default=str) + + +@pytest.fixture() +def recorder(monkeypatch): + table = _RecordingTable() + monkeypatch.setattr(r, "_table", lambda: table) + return table + + +def _drive_every_write(table_unused) -> None: + """Invoke every write path in the module once.""" + r.create_provisioning( + "ast-1", r.KbRecord(app_kb_id="ast-1", owner_user_id="opaque-owner") + ) + r.attach_aws_ids("ast-1", "ast-1", "kb-1", "ds-1", "2026-08-24T12:00:00Z") + r.promote_engine("ast-1", "ast-1", 0, "2026-08-24T12:00:00Z") + r.rollback_engine("ast-1", "ast-1", "2026-08-24T12:00:00Z") + r.acquire_lease("ast-1", "ast-1", "2026-08-24T13:00:00Z", "2026-08-24T12:00:00Z") + for state in (r.SHADOW, r.VERIFY, r.PROMOTE): + r.set_migration_state("ast-1", "ast-1", state, 0, due_at="2026-08-24T12:00:00Z") + for state in (r.RETAIN, r.MIGRATION_FAILED): + r.set_migration_state("ast-1", "ast-1", state, 0, error="a reason") + + +def test_no_write_path_ever_persists_the_legacy_literal(recorder): + """The load-bearing negative. Every write in the module, inspected. + + If this fails, someone has made legacy an explicitly stored value. The + feature would still appear to work, and the next rollback would stop being a + pointer flip. + """ + _drive_every_write(recorder) + assert recorder.calls, "no writes captured; the fixture is not wired" + + payload = recorder.serialized() + assert r.ENGINE_LEGACY not in payload, ( + f"a write path persists the legacy literal {r.ENGINE_LEGACY!r}; " + "absence must remain the only representation of legacy" + ) + + +def test_rollback_removes_the_attribute_rather_than_setting_it(recorder): + """Rollback must restore the original shape, not write a value.""" + r.rollback_engine("ast-1", "ast-1", "2026-08-24T12:00:00Z") + expression = recorder.calls[0]["UpdateExpression"] + assert "REMOVE retrievalEngine" in expression + assert "retrievalEngine = " not in expression + + +def test_the_only_engine_value_ever_written_is_managed(recorder): + """Complements the negative test: promotion writes exactly one engine value.""" + r.promote_engine("ast-1", "ast-1", 0, "2026-08-24T12:00:00Z") + values = recorder.calls[0]["ExpressionAttributeValues"] + engine_values = [v for v in values.values() if v in (r.ENGINE_MANAGED, r.ENGINE_LEGACY)] + assert engine_values == [r.ENGINE_MANAGED] diff --git a/backend/tests/property/test_pbt_kb_migration_convergence.py b/backend/tests/property/test_pbt_kb_migration_convergence.py new file mode 100644 index 000000000..f8b3adbd4 --- /dev/null +++ b/backend/tests/property/test_pbt_kb_migration_convergence.py @@ -0,0 +1,489 @@ +""" +Property-based tests for migration convergence. + +**Property 6: an interrupted migration converges without duplication** + +For any interruption point in the state machine, a resumed run reaches the same +terminal state, creates exactly **one** knowledge base, promotes exactly **once**, +and leaves each document in the corpus exactly once. + +What "without duplication" can and cannot mean +---------------------------------------------- +A worker can die between a successful ``IngestKnowledgeBaseDocuments`` and the +DynamoDB write that records it, and no transaction spans Bedrock and DynamoDB. So +"each document is ingested at most once" is not achievable, and asserting it would +be asserting something false. Two things *are* achievable, and both are asserted: + +* **Each document appears in the corpus exactly once**, because + ``customDocumentIdentifier`` is the platform document id and a re-ingest + therefore replaces. A migration that derived its own identifier would fail here. +* **Redundant re-ingests are bounded by one batch** — the size of the crash window. + That bound is what proves progress is persisted *as the migration proceeds* + rather than only at the end. It is not a theoretical distinction: this test + initially failed because the completed-document set lived inside the + ``migrationProgress`` map, which a later write replaced wholesale, so a crash + near the end of a 25-document corpus re-ingested all 25. + +Why this needs to be a property rather than a set of cases +---------------------------------------------------------- +The interruption points are not a short list. A migration can be cut off between +any two of: reserving bytes, creating the knowledge base, creating the data source, +writing the AWS identifiers back, ingesting each individual batch, recording +progress, promoting, and stamping the retention window. Enumerating them by hand +produces the cases somebody thought of, and the ones that matter are the ones +nobody did — this feature has already been bitten by a crash window between an AWS +create and the database write that records it. + +So the interruption index is a hypothesis input over the sequence of effects, and +the invariants are asserted after replaying from the start, which is what a retry +actually does. + +The model is deliberately in-memory +----------------------------------- +A fake DynamoDB and a fake Bedrock, both of which enforce the properties that make +convergence possible rather than assuming them: + +* ``create_knowledge_base`` is deduplicated by ``clientToken`` — which is how AWS + behaves, and the reason the worker persists the token before calling AWS. +* ``ingest_knowledge_base_documents`` records every ``customDocumentIdentifier`` + it is handed, so "at most once" is measured over the whole replay rather than + per attempt. + +A test against real clients could not interrupt at a chosen point, and a test with +no model at all would assert only that the code does not raise. + +Feature: managed-kb-migration +**Validates: Requirements 15.9, 15.10, 15.13, 7.4** +""" + +from typing import Any, Dict, List, Optional, Set + +import pytest +from hypothesis import HealthCheck, given, settings, strategies as st + +# --------------------------------------------------------------------------- +# The model +# --------------------------------------------------------------------------- + + +class Interrupted(Exception): + """The simulated crash. Raised at the chosen effect index.""" + + +class Clock: + """Counts effects and raises at the interruption point. + + Every externally-visible side effect passes through :meth:`tick`, so the + interruption index addresses effects rather than lines of code — the unit a + crash actually lands between. + """ + + def __init__(self, interrupt_at: Optional[int] = None): + self.count = 0 + self.interrupt_at = interrupt_at + self.log: List[str] = [] + + def tick(self, what: str) -> None: + self.count += 1 + self.log.append(what) + if self.interrupt_at is not None and self.count == self.interrupt_at: + raise Interrupted(f"crashed at effect {self.count}: {what}") + + +class FakeAws: + """Bedrock's idempotency, modelled rather than assumed.""" + + def __init__(self, clock: Clock): + self.clock = clock + self.kbs_by_token: Dict[str, str] = {} + self.data_sources: Dict[str, str] = {} + #: Every ingest ever accepted, across every attempt. The duplication + #: invariant is measured here. + self.ingest_log: List[str] = [] + #: document_id -> times written. Distinct keys are the corpus; the counts + #: are the redundant work a resume did. + self.corpus: Dict[str, int] = {} + self.next_id = 0 + + def create_knowledge_base(self, client_token: str) -> str: + # Deduplicated by token: this is what makes a retried create safe, and the + # reason the record persists the token *before* the AWS call. + if client_token in self.kbs_by_token: + return self.kbs_by_token[client_token] + self.clock.tick("CreateKnowledgeBase") + self.next_id += 1 + kb_id = f"KB{self.next_id:04d}" + self.kbs_by_token[client_token] = kb_id + return kb_id + + def create_data_source(self, kb_id: str, client_token: str) -> str: + if client_token in self.data_sources: + return self.data_sources[client_token] + self.clock.tick("CreateDataSource") + ds_id = f"DS-{kb_id}" + self.data_sources[client_token] = ds_id + return ds_id + + def ingest(self, document_ids: List[str]) -> None: + self.clock.tick(f"Ingest({','.join(document_ids)})") + self.ingest_log.extend(document_ids) + for document_id in document_ids: + # ``customDocumentIdentifier`` is the platform document id, so a + # re-ingest *replaces* rather than appends. Modelled because it is what + # makes the unavoidable crash window survivable: a worker can die + # between a successful Ingest and the bookkeeping write, and no + # transaction spans Bedrock and DynamoDB. + self.corpus[document_id] = self.corpus.get(document_id, 0) + 1 + + @property + def corpus_document_count(self) -> int: + """Distinct documents in the knowledge base.""" + return len(self.corpus) + + @property + def knowledge_base_count(self) -> int: + return len(set(self.kbs_by_token.values())) + + +class FakeRecord: + """The KB_Record, with the conditional writes that matter.""" + + def __init__(self, clock: Clock, document_ids: List[str]): + self.clock = clock + self.item: Dict[str, Any] = {} + self.documents: Dict[str, str] = {d: "complete" for d in document_ids} + self.promotions = 0 + + # -- reads --------------------------------------------------------------- + def get(self) -> Dict[str, Any]: + return dict(self.item) + + def list_complete(self) -> List[str]: + return sorted(d for d, status in self.documents.items() if status == "complete") + + def status_of(self, document_id: str) -> Optional[str]: + return self.documents.get(document_id) + + # -- writes -------------------------------------------------------------- + def create_provisioning(self, client_token: str) -> None: + if self.item: + return # attribute_not_exists guard: the retry anchor already exists + self.clock.tick("CreateProvisioning") + self.item = { + "clientToken": client_token, + "provisioningState": "provisioning", + "migrationState": "shadow", + "migrationGeneration": 1, + "totalBytes": 0, + } + + def attach_ids(self, kb_id: str, ds_id: str) -> None: + if self.item.get("awsKbId"): + return + self.clock.tick("AttachAwsIds") + self.item["awsKbId"] = kb_id + self.item["awsDataSourceId"] = ds_id + self.item["provisioningState"] = "active" + + def reserve(self, total: int) -> None: + self.clock.tick("ReserveSnapshot") + self.item["totalBytes"] = total + + def set_progress(self, migrated: int, total: int, newly_done: List[str] = None) -> None: + self.clock.tick("SetProgress") + self.item["migrationProgress"] = {"migrated": migrated, "total": total} + if newly_done: + # ADD on a string set: additive, and a *separate attribute* from the + # progress map this write replaces. Modelled that way because the + # first version of this test kept the completed set inside the map, + # the map got overwritten, and the resumed run re-ingested a corpus + # it had already finished. The worker had the same bug. + existing = set(self.item.get("migratedDocIds") or ()) + self.item["migratedDocIds"] = existing | set(newly_done) + + def add_done(self, document_ids: List[str]) -> None: + """The per-batch ADD, which is what survives a crash between batches.""" + if not document_ids: + return + self.clock.tick(f"AddDone({','.join(document_ids)})") + existing = set(self.item.get("migratedDocIds") or ()) + self.item["migratedDocIds"] = existing | set(document_ids) + + def set_state(self, new_state: str, expected: Optional[Set[str]] = None) -> bool: + if expected is not None and self.item.get("migrationState") not in expected: + return False + self.clock.tick(f"SetState({new_state})") + self.item["migrationState"] = new_state + return True + + def promote(self) -> bool: + progress = self.item.get("migrationProgress") or {} + if self.item.get("retrievalEngine"): + # attribute_not_exists(retrievalEngine): already promoted. Every other + # guard stays true after a successful promotion, so without this one a + # crash between the promotion and the state transition promotes twice — + # and two concurrent workers both succeed. + return False + if self.item.get("migrationState") != "promote": + return False + if progress.get("migrated") != progress.get("total"): + # Requirement 15.9 in the model: convergence is part of the condition, + # not a separate check somebody could forget to call. + return False + self.clock.tick("Promote") + self.promotions += 1 + self.item["retrievalEngine"] = "managed" + return True + + +BATCH = 10 + + +def run_migration(record: FakeRecord, aws: FakeAws, clock: Clock) -> str: + """Replay the whole machine from the start. Idempotent by construction. + + This mirrors the real worker's ordering exactly, and the ordering is the thing + under test: the record is written *before* AWS is called, the persisted token is + reused on resume, and every document's status is re-read immediately before it + is ingested. + """ + token = "kb-token-fixed-length-padding-000000" + + # shadow + record.create_provisioning(token) + if not record.item.get("totalBytes"): + record.reserve(len(record.list_complete()) * 1024) + + kb_id = record.item.get("awsKbId") or aws.create_knowledge_base( + record.item.get("clientToken") or token + ) + ds_id = record.item.get("awsDataSourceId") or aws.create_data_source(kb_id, token) + record.attach_ids(kb_id, ds_id) + + already = set(record.item.get("migratedDocIds") or ()) + snapshot = record.list_complete() + pending = [d for d in snapshot if d not in already] + + migrated = set(already) + for start in range(0, len(pending), BATCH): + batch = [ + d + for d in pending[start : start + BATCH] + # Requirement 16.4: re-read immediately before ingesting. + if record.status_of(d) == "complete" + ] + if not batch: + continue + aws.ingest(batch) + migrated.update(batch) + # Persisted per batch, so a crash between batches loses only the batch in + # flight rather than the whole run's progress. + record.add_done(batch) + + # catch-up until quiet + passes = 0 + while passes < 5: + passes += 1 + new = [d for d in record.list_complete() if d not in migrated] + if not new: + break + aws.ingest(new) + migrated.update(new) + record.add_done(new) + + record.set_progress(len(migrated), len(record.list_complete())) + record.set_state("verify", {"shadow"}) + + # verify + record.set_state("promote", {"verify"}) + + # promote. An already-promoted record still finishes: the promotion write is + # guarded on the engine attribute being absent, so a resume after a crash + # between the promotion and the state transition must continue to `retain` + # rather than treat the refusal as a failure. + if record.promote() or record.item.get("retrievalEngine") == "managed": + record.set_state("retain", {"promote"}) + + return record.item.get("migrationState", "") + + +# --------------------------------------------------------------------------- +# Strategies +# --------------------------------------------------------------------------- + +st_document_ids = st.lists( + st.text(alphabet="abcdefghijklmnopqrstuvwxyz0123456789", min_size=1, max_size=6), + min_size=1, + max_size=25, + unique=True, +) + +#: Effects, not lines. A migration of 25 documents produces roughly a dozen; the +#: upper bound is generous so an index past the end simply means "not interrupted", +#: which is a case worth generating too. +st_interrupt_at = st.integers(min_value=1, max_value=30) + + +# --------------------------------------------------------------------------- +# The property +# --------------------------------------------------------------------------- + + +@settings(max_examples=200, deadline=None, suppress_health_check=[HealthCheck.too_slow]) +@given(document_ids=st_document_ids, interrupt_at=st_interrupt_at) +def test_an_interrupted_migration_converges_without_duplication(document_ids, interrupt_at): + """The whole property, in one test. + + Run once with a crash injected at ``interrupt_at``; then run again from the + start, as a retry does. Assert the terminal state, exactly one knowledge base, + and each document ingested at most once across **both** runs. + """ + clock = Clock(interrupt_at=interrupt_at) + aws = FakeAws(clock) + record = FakeRecord(clock, document_ids) + + try: + run_migration(record, aws, clock) + except Interrupted: + pass + + # The retry. No interruption this time. + clock.interrupt_at = None + final_state = run_migration(record, aws, clock) + + assert final_state == "retain", ( + f"a resumed migration did not converge: state={final_state!r}, " + f"effects={clock.log}" + ) + + assert aws.knowledge_base_count == 1, ( + f"{aws.knowledge_base_count} knowledge bases were created; the persisted " + f"clientToken is not deduplicating the retried create" + ) + + counts: Dict[str, int] = {} + for document_id in aws.ingest_log: + counts[document_id] = counts.get(document_id, 0) + 1 + redundant = sum(n - 1 for n in counts.values()) + + # Each document appears in the corpus exactly once. This is the invariant that + # actually matters, and it is real rather than tautological: it holds because + # `customDocumentIdentifier` is the platform document id, so a re-ingest + # replaces. A migration that derived its own identifier would fail here. + assert aws.corpus_document_count == len(document_ids) + assert all(document_id in aws.corpus for document_id in document_ids) + + # Redundant re-ingests are bounded by one batch: the crash window between a + # successful Ingest and the write that records it. No transaction spans Bedrock + # and DynamoDB, so that window cannot be closed — but it can be *bounded*, and + # the bound is what proves progress is persisted per batch. Before the + # completed-document set was persisted, a crash near the end of a 25-document + # corpus re-ingested all 25; this assertion is what caught that. + assert redundant <= BATCH, ( + f"{redundant} redundant ingests after one interruption, which is more than " + f"the single batch that can be in flight; progress is not being persisted " + f"as the migration proceeds. effects={clock.log}" + ) + + assert set(aws.ingest_log) == set(document_ids), ( + "the resumed migration did not end up with every document" + ) + + assert record.promotions == 1, ( + f"promotion happened {record.promotions} times; the conditional write is " + f"not the single cutover" + ) + + +@settings(max_examples=100, deadline=None) +@given(document_ids=st_document_ids, delete_index=st.integers(min_value=0, max_value=24)) +def test_a_document_deleted_mid_migration_is_never_ingested(document_ids, delete_index): + """Requirements 16.4, 16.5, as a property over which document is deleted. + + The deletion lands after the snapshot is taken and before the document's turn + comes, which is the only window in which resurrection is possible. + """ + clock = Clock() + aws = FakeAws(clock) + record = FakeRecord(clock, document_ids) + + victim = document_ids[delete_index % len(document_ids)] + + original_status_of = record.status_of + + def _status_with_deletion(document_id: str): + if document_id == victim: + return None + return original_status_of(document_id) + + record.status_of = _status_with_deletion + record.documents.pop(victim) + + run_migration(record, aws, clock) + + assert victim not in aws.ingest_log, ( + f"document {victim!r} was deleted mid-migration and still reached the " + f"managed corpus" + ) + + +@settings(max_examples=100, deadline=None) +@given(document_ids=st_document_ids) +def test_promotion_is_refused_until_catch_up_converges(document_ids): + """Requirement 15.9, asserted through the promotion condition itself. + + Progress is deliberately left short of the total, as an unconverged catch-up + leaves it. Promotion must be refused — and refused by the condition, so no + caller can reach past it. + """ + clock = Clock() + record = FakeRecord(clock, document_ids) + + record.item = { + "migrationState": "promote", + "migrationProgress": {"migrated": max(len(document_ids) - 1, 0), "total": len(document_ids)}, + } + + assert record.promote() is False + assert record.promotions == 0 + assert "retrievalEngine" not in record.item + + +@settings(max_examples=50, deadline=None) +@given(document_ids=st_document_ids) +def test_only_one_of_two_concurrent_promotions_wins(document_ids): + """Requirement 15.10. The second attempt sees a record no longer in ``promote`` + and is refused, which is what the real conditional write does.""" + clock = Clock() + record = FakeRecord(clock, document_ids) + + total = len(document_ids) + record.item = { + "migrationState": "promote", + "migrationProgress": {"migrated": total, "total": total}, + } + + first = record.promote() + record.set_state("retain", {"promote"}) + second = record.promote() + + assert first is True + assert second is False + assert record.promotions == 1 + + +def test_the_model_can_actually_be_interrupted(): + """Guards the guard. + + If ``Clock.tick`` stopped raising, every property above would pass while + testing nothing but the happy path. So assert that some interruption index + genuinely prevents convergence on the first run. + """ + clock = Clock(interrupt_at=1) + aws = FakeAws(clock) + record = FakeRecord(clock, ["d1", "d2"]) + + with pytest.raises(Interrupted): + run_migration(record, aws, clock) + + assert record.item.get("migrationState") != "retain" diff --git a/backend/tests/property/test_pbt_kb_query_clamp.py b/backend/tests/property/test_pbt_kb_query_clamp.py new file mode 100644 index 000000000..a61ca0a4d --- /dev/null +++ b/backend/tests/property/test_pbt_kb_query_clamp.py @@ -0,0 +1,245 @@ +"""Property-based tests for the retrieval query clamp. + +Feature: managed-kb-migration + +**Property 3: the clamp is total and non-throwing.** + +Managed Knowledge Base rejects a ``Retrieve`` query over 10,000 characters +outright, and the quota is not adjustable. So the clamp sits on a request path +where the only acceptable behaviours are "shortened" or "unchanged" — never +"raised". A clamp that threw would convert a fixable input into a failed chat +turn, which is strictly worse than answering a slightly truncated question. + +"Total" is the load-bearing word: *every* input must map to an output, including +the awkward ones. The strategies below deliberately include empty strings, strings +made entirely of astral-plane characters, and lengths sitting exactly on the +boundary, because those are where a length check written against the wrong unit or +with an off-by-one starts returning 10,001 characters to an API that rejects +10,001 characters. + +Validates: Requirements 4.1, 4.3, 4.4. +""" + +from unittest.mock import patch + +import pytest +from hypothesis import given, settings, strategies as st + +from apis.shared.assistants.kb_access import granted +from apis.shared.kb_backend.query_guard import MAX_QUERY_CHARS, clamp_query + +# --------------------------------------------------------------------------- +# Strategies +# --------------------------------------------------------------------------- + +#: Any text at all, including empty and including characters that are one code +#: point but more than one byte — the clamp counts characters, and a byte-based +#: implementation would pass a naive ASCII-only test. +st_any_text = st.text(max_size=200) + +st_long_text = st.text(min_size=1, max_size=50).map( + lambda s: s * (MAX_QUERY_CHARS // max(len(s), 1) + 2) +) + +st_multibyte_text = st.text( + alphabet=st.characters(min_codepoint=0x1F300, max_codepoint=0x1F5FF), + min_size=1, + max_size=40, +).map(lambda s: s * (MAX_QUERY_CHARS // max(len(s), 1) + 2)) + +#: Lengths straddling the cap, where off-by-one errors live. +st_boundary_length = st.integers( + min_value=MAX_QUERY_CHARS - 2, max_value=MAX_QUERY_CHARS + 2 +) + + +# --------------------------------------------------------------------------- +# Totality and the cap +# --------------------------------------------------------------------------- +@given(query=st_any_text) +@settings(max_examples=200) +def test_short_queries_pass_through_unchanged(query): + """Below the cap the clamp must be the identity, not a normalizer. + + Anything else would silently change what users are asking. + """ + with patch("apis.shared.kb_backend.query_guard.emit_count"): + result, truncated = clamp_query(query) + assert result == query + assert truncated is False + + +@given(query=st.one_of(st_long_text, st_multibyte_text)) +@settings(max_examples=100) +def test_output_never_exceeds_the_cap(query): + """The whole point: the value handed to the backend always fits.""" + with patch("apis.shared.kb_backend.query_guard.emit_count"): + result, truncated = clamp_query(query) + assert len(result) <= MAX_QUERY_CHARS + assert truncated is True + + +@given(length=st_boundary_length) +@settings(max_examples=50) +def test_the_boundary_is_inclusive(length): + """Exactly MAX_QUERY_CHARS is allowed; one more is not. + + Managed KB accepts 10,000 and rejects 10,001, so an off-by-one here is a + request error rather than a shorter answer. + """ + with patch("apis.shared.kb_backend.query_guard.emit_count"): + result, truncated = clamp_query("x" * length) + + assert len(result) == min(length, MAX_QUERY_CHARS) + assert truncated == (length > MAX_QUERY_CHARS) + + +@given(query=st.one_of(st_any_text, st_long_text, st_multibyte_text)) +@settings(max_examples=200) +def test_the_clamp_never_raises(query): + """Totality. A raise here would turn a long question into a failed chat turn.""" + with patch("apis.shared.kb_backend.query_guard.emit_count"): + try: + clamp_query(query) + except Exception as exc: # pragma: no cover - the assertion is the point + pytest.fail(f"clamp_query raised {type(exc).__name__}: {exc}") + + +@given(query=st_long_text) +@settings(max_examples=50) +def test_truncation_keeps_the_head(query): + """Keep the beginning: for a natural-language query that is where the intent + is. Head-truncating would change the question rather than shorten it.""" + with patch("apis.shared.kb_backend.query_guard.emit_count"): + result, _ = clamp_query(query) + assert query.startswith(result) + + +@given(query=st_long_text) +@settings(max_examples=50) +def test_the_clamp_is_idempotent(query): + """Clamping twice equals clamping once, and the second pass reports no + truncation — so a retry does not double-count the metric.""" + with patch("apis.shared.kb_backend.query_guard.emit_count"): + once, first = clamp_query(query) + twice, second = clamp_query(once) + assert twice == once + assert first is True + assert second is False + + +# --------------------------------------------------------------------------- +# The truncation signal +# --------------------------------------------------------------------------- +@given(length=st_boundary_length) +@settings(max_examples=50) +def test_the_metric_is_emitted_exactly_when_truncation_happened(length): + """The signal must track reality in both directions. + + A metric that over-reports trains operators to ignore it; one that + under-reports hides the fact that users are already sending queries the + managed backend would reject. + """ + with patch("apis.shared.kb_backend.query_guard.emit_count") as emit: + _, truncated = clamp_query("x" * length) + + assert truncated == (length > MAX_QUERY_CHARS) + assert emit.called == truncated + + +def test_a_metric_failure_does_not_break_the_clamp(): + """Observability is never control flow: if CloudWatch is down the query still + gets clamped and the search still runs.""" + with patch( + "apis.shared.kb_backend.query_guard.emit_count", + side_effect=RuntimeError("cloudwatch unavailable"), + ): + with pytest.raises(RuntimeError): + # Confirms the patch is actually wired, so the next assertion is not + # vacuous. + clamp_query("x" * (MAX_QUERY_CHARS + 1)) + + # emit_count's real implementation swallows its own failures, which is what + # makes the above impossible in production. Assert that contract directly. + from apis.shared.kb_backend.metrics import emit_count + + with patch("boto3.client", side_effect=RuntimeError("no credentials")): + emit_count("KbQueryClamped") # must not raise + + +@pytest.mark.parametrize("falsy", ["", None]) +def test_empty_input_is_handled_without_a_metric(falsy): + """An empty query is not a truncation.""" + with patch("apis.shared.kb_backend.query_guard.emit_count") as emit: + result, truncated = clamp_query(falsy) + assert result == "" + assert truncated is False + emit.assert_not_called() + + +# --------------------------------------------------------------------------- +# Guards the properties above cannot provide +# +# Every test above refers to MAX_QUERY_CHARS symbolically, so all of them follow +# the constant wherever it goes — raise it to 32,000 and they all still pass while +# the managed backend starts rejecting requests. These three assertions were added +# after mutation testing showed exactly that: three separate mutations survived a +# suite that looked thorough. +# --------------------------------------------------------------------------- +def test_the_cap_is_the_literal_managed_kb_limit(): + """Pinned to 10,000 as a LITERAL, not to the constant. + + This is the one assertion in the file that cannot be satisfied by moving the + constant. 10,000 is Managed KB's `Retrieve` input quota and it is not + adjustable, so this number is a property of AWS, not a tuning knob. Raising it + does not buy longer queries; it buys rejected requests. + """ + assert MAX_QUERY_CHARS == 10_000 + + +@pytest.mark.asyncio +async def test_the_facade_actually_clamps_before_dispatch(): + """The clamp must be WIRED, not merely correct. + + Nothing else in this file would notice if the facade stopped calling + clamp_query: the unit-level properties would all still pass while every long + query went to the backend intact. Asserted by inspecting what the backend + actually received. + """ + from apis.shared.assistants import rag_service + + seen = {} + + class _RecordingBackend: + async def search(self, kb_ref, query, top_k=5): + seen["query"] = query + return [] + + with patch.object(rag_service, "resolve_backend", return_value=_RecordingBackend()), patch.object( + rag_service, "emit_count" + ), patch("apis.shared.kb_backend.query_guard.emit_count"): + await rag_service.search_assistant_knowledgebase_with_formatting( + "ast-1", + "x" * (MAX_QUERY_CHARS + 500), + access=granted("ast-1", "user-clamp", "owner"), + ) + + assert seen["query"] is not None + assert len(seen["query"]) == MAX_QUERY_CHARS, ( + "the facade dispatched an unclamped query; the clamp is dead code" + ) + + +def test_the_metric_namespace_is_not_a_reserved_aws_one(): + """CloudWatch rejects PutMetricData into any namespace beginning with "AWS". + + A reserved namespace would make every publish silently denied — the grant looks + correct, the code looks correct, and no metric ever arrives. The CDK grant + conditions on this same namespace, so the two must agree; this is the backend + half of that assertion. + """ + from apis.shared.kb_backend.metrics import metric_namespace + + ns = metric_namespace() + assert not ns.startswith("AWS"), f"{ns!r} is a reserved namespace; writes are rejected" + assert ns.endswith("/ManagedKb") diff --git a/backend/tests/property/test_pbt_kb_score_direction.py b/backend/tests/property/test_pbt_kb_score_direction.py new file mode 100644 index 000000000..62758ab4e --- /dev/null +++ b/backend/tests/property/test_pbt_kb_score_direction.py @@ -0,0 +1,299 @@ +""" +Property-based tests for score direction across knowledge base backends. + +**Property 2: ranking is backend-independent** + +For any list of chunks with distinct scores, both backends return the known-best +chunk first after adapter conversion, and the ``relevance`` values they attach +agree with the order they return. + +This is the only test in the suite that can catch a silent ranking inversion. +S3 Vectors reports cosine *distance* (lower is better); Managed KB reports +*relevance* (higher is better). If the legacy adapter forwards distance as +relevance, nothing raises: every request still succeeds, still returns five +chunks, and still logs "Found 5 relevant chunks". The only symptom is that the +worst passages are ranked best and answers quietly degrade. There is no error +path, so there is nothing else to assert on. + +Feature: managed-kb-migration +**Validates: Requirements 2.1, 2.2, 2.3, 2.4, 24.1** +""" + +from typing import Any, Dict, List +from unittest.mock import patch + +from hypothesis import given, settings, strategies as st + +from apis.shared.kb_backend.protocol import ( + DEFAULT_TOP_K, + Chunk, + KnowledgeBaseBackend, + distance_from_relevance, + relevance_from_distance, +) +from apis.shared.kb_backend.s3vectors_backend import S3VectorsBackend + +# --------------------------------------------------------------------------- +# Strategies +# --------------------------------------------------------------------------- + +# Cosine distance lives in [0, 2]. Distinct values only: the property is about +# strict ranking, and ties would make "the known-best chunk" ambiguous rather +# than wrong. +st_distances = st.lists( + st.floats(min_value=0.0, max_value=2.0, allow_nan=False, allow_infinity=False), + min_size=2, + max_size=8, + unique=True, +) + +# Bounded so that ``score + 1.0`` is genuinely a larger float. Near the top of +# the double range adding 1.0 is a no-op, which would make the pairwise +# comparison below vacuous rather than false. +st_any_score = st.floats( + min_value=-1e6, max_value=1e6, allow_nan=False, allow_infinity=False +) + + +# --------------------------------------------------------------------------- +# A managed backend stand-in +# --------------------------------------------------------------------------- + + +class FakeManagedBackend: + """A protocol-conforming backend that reports relevance natively. + + Stands in for ``managed_backend.ManagedKbBackend``, which task 8.3 builds. + The property under test is about score *direction* — a per-adapter concern + that is fully determined by whether the adapter converts or passes through — + so a stand-in that passes relevance through unchanged, exactly as Managed KB + requires (Requirement 2.3), exercises the property faithfully. Nothing here + depends on Bedrock's wire format. + """ + + def __init__(self, results: List[Dict[str, Any]]): + # results: [{"document_id", "relevance", "text", "key"}], best first, + # which is the order Bedrock's Retrieve returns. + self._results = results + + async def search(self, kb_ref: str, query: str, top_k: int = DEFAULT_TOP_K) -> List[Chunk]: + return [ + Chunk( + text=result["text"], + # Pass-through. Managed already counts in the canonical direction. + relevance=result["relevance"], + document_id=result["document_id"], + metadata={"document_id": result["document_id"], "text": result["text"]}, + key=result["key"], + ) + for result in self._results + ] + + async def ingest(self, kb_ref: str, source) -> None: # pragma: no cover - unused here + raise NotImplementedError + + async def delete_document(self, kb_ref: str, document_id: str) -> None: # pragma: no cover + raise NotImplementedError + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _s3_vectors_response(distances: List[float]) -> Dict[str, Any]: + """Build an S3 Vectors query response, nearest-first as the API returns it.""" + return { + "vectors": [ + { + "key": f"doc-{index}#0", + "distance": distance, + "metadata": {"document_id": f"doc-{index}", "text": f"passage {index}"}, + } + for index, distance in enumerate(sorted(distances)) + ] + } + + +async def _legacy_search(distances: List[float]) -> List[Chunk]: + response = _s3_vectors_response(distances) + with patch( + "apis.shared.embeddings.bedrock_embeddings.search_assistant_knowledgebase", + return_value=response, + ): + return await S3VectorsBackend().search("ast-1", "a query") + + +def _is_non_increasing(values: List[float]) -> bool: + return all(earlier >= later for earlier, later in zip(values, values[1:])) + + +# --------------------------------------------------------------------------- +# Property 2 +# --------------------------------------------------------------------------- + + +@given(distances=st_distances) +@settings(max_examples=200, deadline=None) +def test_legacy_backend_ranks_known_best_chunk_first(distances): + """ + **Validates: Requirements 2.1, 2.2, 2.4** + + The chunk with the *lowest* S3 Vectors distance is the known-best chunk. After + conversion it must be first in the returned list and must carry the *highest* + relevance. + + The relevance-ordering assertion is the one that catches an inversion. The + positional one does not on its own: the adapter preserves the index's order, + so a chunk stays first whatever score is stapled to it. Only the claim that + scores descend can detect that the numbers now disagree with the order. + """ + import asyncio + + chunks = asyncio.run(_legacy_search(distances)) + + best_distance = min(distances) + best_key = f"doc-{sorted(distances).index(best_distance)}#0" + + assert chunks[0].key == best_key, "known-best chunk is not first" + + relevances = [chunk.relevance for chunk in chunks] + assert _is_non_increasing(relevances), ( + f"relevance must descend with rank, got {relevances}. " + f"A rising sequence means distance was forwarded as relevance: the " + f"ranking is inverted and the worst chunks are being served as the best." + ) + + argmax = max(chunks, key=lambda chunk: chunk.relevance) + assert argmax.key == best_key, ( + f"highest relevance is {argmax.key}, expected the nearest chunk {best_key}" + ) + + +@given(distances=st_distances) +@settings(max_examples=200, deadline=None) +def test_managed_backend_ranks_known_best_chunk_first(distances): + """ + **Validates: Requirements 2.1, 2.3, 2.4** + + The managed backend passes relevance through, so the known-best chunk is the + one with the highest relevance and it must come back first. + """ + import asyncio + + # The same logical corpus, expressed in the managed backend's own units. + scored = sorted( + ( + { + "document_id": f"doc-{index}", + "text": f"passage {index}", + "key": f"doc-{index}#0", + "relevance": relevance_from_distance(distance), + } + for index, distance in enumerate(sorted(distances)) + ), + key=lambda result: result["relevance"], + reverse=True, + ) + + backend = FakeManagedBackend(scored) + chunks = asyncio.run(backend.search("ast-1", "a query")) + + best_key = scored[0]["key"] + + assert chunks[0].key == best_key, "known-best chunk is not first" + + relevances = [chunk.relevance for chunk in chunks] + assert _is_non_increasing(relevances), ( + f"relevance must descend with rank, got {relevances}" + ) + + argmax = max(chunks, key=lambda chunk: chunk.relevance) + assert argmax.key == best_key + + +@given(distances=st_distances) +@settings(max_examples=200, deadline=None) +def test_both_backends_agree_on_ranking(distances): + """ + **Validates: Requirement 2.4** + + Given the same corpus and the same relative scores, both backends must return + the same documents in the same order. This is the parity claim a migration + rests on: a knowledge base that moves engines must not reorder its answers. + """ + import asyncio + + legacy_chunks = asyncio.run(_legacy_search(distances)) + + managed_results = [ + { + "document_id": chunk.document_id, + "text": chunk.text, + "key": chunk.key, + "relevance": chunk.relevance, + } + for chunk in sorted(legacy_chunks, key=lambda chunk: chunk.relevance, reverse=True) + ] + managed_chunks = asyncio.run(FakeManagedBackend(managed_results).search("ast-1", "q")) + + assert [chunk.document_id for chunk in legacy_chunks] == [ + chunk.document_id for chunk in managed_chunks + ], "the two backends ranked the same corpus differently" + + assert [chunk.relevance for chunk in legacy_chunks] == [ + chunk.relevance for chunk in managed_chunks + ], "the two backends scored the same corpus differently" + + +# --------------------------------------------------------------------------- +# The derived distance key must be the same value, not a nearby one +# --------------------------------------------------------------------------- + + +@given(distance=st.floats(min_value=0.0, max_value=2.0, allow_nan=False)) +@settings(max_examples=200, deadline=None) +def test_distance_relevance_round_trip_is_exact(distance): + """ + **Validates: Requirement 2.2** + + The facade derives the ``distance`` it emits from ``relevance``, and that + value reaches an HTTP response body. The conversion must therefore be exactly + reversible, not merely close: a ``1.0 - x`` formulation would turn ``0.1`` + into ``0.09999999999999998`` and change a value clients already read. + """ + assert distance_from_relevance(relevance_from_distance(distance)) == distance + + +@given(score=st_any_score) +@settings(max_examples=200, deadline=None) +def test_conversion_inverts_direction_for_every_score(score): + """ + **Validates: Requirements 2.1, 2.2** + + Direction inversion is the whole contract: for any two distinct distances, + the smaller one must produce the larger relevance. Asserted pointwise against + a second score so no clamping, absolute value, or identity mapping can pass. + """ + other = score + 1.0 # strictly greater distance + assert relevance_from_distance(score) > relevance_from_distance(other), ( + "a nearer chunk (smaller distance) must receive a higher relevance" + ) + + +def test_none_score_is_preserved_not_fabricated(): + """ + **Validates: Requirement 2.2** + + A response without a distance yields ``None``, which the facade emits + verbatim as it always has. Defaulting to ``0.0`` would make an unscored + chunk the best-ranked chunk in the list. + """ + assert relevance_from_distance(None) is None + assert distance_from_relevance(None) is None + + +def test_backends_satisfy_the_protocol(): + """Both implementations structurally conform to KnowledgeBaseBackend.""" + assert isinstance(S3VectorsBackend(), KnowledgeBaseBackend) + assert isinstance(FakeManagedBackend([]), KnowledgeBaseBackend) diff --git a/backend/tests/property/test_pbt_kb_status_fail_closed.py b/backend/tests/property/test_pbt_kb_status_fail_closed.py new file mode 100644 index 000000000..eff74cfd7 --- /dev/null +++ b/backend/tests/property/test_pbt_kb_status_fail_closed.py @@ -0,0 +1,196 @@ +"""Property-based tests for fail-closed document status filtering. + +Feature: managed-kb-migration + +**Property 4: unconfirmable status never leaks.** + +The filter's job is to keep chunks belonging to deleted or half-deleted documents +out of retrieval results. Its old fallback returned everything unfiltered whenever +it could not reach DynamoDB, which meant the guard vanished at exactly the moment +it was most likely to matter — and vanished *silently*, since the response looks +identical either way. + +This inverts that (Requirement 5, superseding `reliable-document-deletion` +Requirement 3.4). The property asserted here is deliberately absolute: no matter +how many chunks, how many distinct documents, or what shape of table-level failure +is injected, the result is empty. There is no "mostly" — a single leaked chunk from +a deleted document is the entire failure mode. + +The per-document lookup failure is a different case and is *not* covered by this +property: that one already skipped only its own document, which is correct, and is +left unchanged. + +Validates: Requirements 5.1, 5.2, 24.6. +""" + +from unittest.mock import MagicMock, patch + +import pytest +from hypothesis import HealthCheck, given, settings, strategies as st + +# Imported at module scope, deliberately, and NOT inside the patched context of a +# test. Importing it lazily made the first-ever run differ from every later one: +# the import itself happened while `boto3.resource` was mocked, so module-level +# import work was performed against a mock exactly once and was then cached in +# sys.modules for the rest of the session. That produced a test that failed on a +# cold run and passed on every warm one — the worst failure mode a guard can have, +# because CI is cold and local re-runs are warm. +from apis.shared.assistants.rag_service import _filter_vectors_by_document_status + +ASSISTANT_ID = "ast-failclosed" + +# --------------------------------------------------------------------------- +# Strategies +# --------------------------------------------------------------------------- + +st_document_id = st.text( + alphabet="abcdefghijklmnopqrstuvwxyz0123456789-", min_size=1, max_size=16 +).map(lambda s: f"doc-{s}") + +#: A non-empty set of vectors spread over an arbitrary number of documents. Both +#: axes matter: the filter dedupes document ids before lookup, so "many chunks, +#: one document" and "one chunk each, many documents" exercise different paths. +st_vectors = st.lists(st_document_id, min_size=1, max_size=12).map( + lambda ids: [ + { + "key": f"vec-{i}", + "distance": 0.1, + "metadata": {"document_id": d, "text": f"chunk {i}", "assistant_id": ASSISTANT_ID}, + } + for i, d in enumerate(ids) + ] +) + +#: Table-level failures. Any exception type, raised from the resource or the +#: table handle — the guard must not depend on recognising a specific error. +st_failure = st.sampled_from( + [ + Exception("DynamoDB unavailable"), + RuntimeError("connection reset"), + ValueError("malformed region"), + KeyError("credentials"), + TimeoutError("timed out"), + ] +) + + +# --------------------------------------------------------------------------- +# The property +# --------------------------------------------------------------------------- +@given(vectors=st_vectors, failure=st_failure) +@settings(max_examples=150, suppress_health_check=[HealthCheck.function_scoped_fixture]) +def test_a_table_level_failure_never_leaks_a_chunk(vectors, failure): + """Any table-level failure, any corpus shape → zero chunks.""" + with patch("apis.shared.kb_backend.metrics.emit_count"), patch( + "apis.shared.assistants.rag_service.emit_count" + ), patch("boto3.resource") as resource, patch.dict( + "os.environ", {"DYNAMODB_ASSISTANTS_TABLE_NAME": "t", "AWS_REGION": "us-west-2"} + ): + dynamo = MagicMock() + dynamo.Table.side_effect = failure + resource.return_value = dynamo + + assert _filter_vectors_by_document_status(vectors, ASSISTANT_ID) == [] + + +@given(vectors=st_vectors) +@settings(max_examples=100, suppress_health_check=[HealthCheck.function_scoped_fixture]) +def test_a_missing_table_name_never_leaks_a_chunk(vectors): + """The other former fail-open path: no table configured → zero chunks.""" + with patch("apis.shared.kb_backend.metrics.emit_count"), patch( + "apis.shared.assistants.rag_service.emit_count" + ), patch("boto3.resource") as resource, patch.dict("os.environ", {}, clear=True): + assert _filter_vectors_by_document_status(vectors, ASSISTANT_ID) == [] + # Never contacted, so this is a guard rather than a failed call. + resource.assert_not_called() + + +@given(vectors=st_vectors, failure=st_failure) +@settings(max_examples=100, suppress_health_check=[HealthCheck.function_scoped_fixture]) +def test_the_degradation_is_always_reported(vectors, failure): + """An empty result from this path must be distinguishable from an empty corpus. + + Without the signal, a total retrieval outage looks exactly like "nobody's + documents matched", which is the kind of failure that survives for weeks. + """ + with patch("apis.shared.assistants.rag_service.emit_count") as emit, patch( + "boto3.resource" + ) as resource, patch.dict( + "os.environ", {"DYNAMODB_ASSISTANTS_TABLE_NAME": "t", "AWS_REGION": "us-west-2"} + ): + dynamo = MagicMock() + dynamo.Table.side_effect = failure + resource.return_value = dynamo + + _filter_vectors_by_document_status(vectors, ASSISTANT_ID) + emit.assert_called_once() + + +# --------------------------------------------------------------------------- +# What must NOT change +# --------------------------------------------------------------------------- +@given(vectors=st_vectors) +@settings(max_examples=50, suppress_health_check=[HealthCheck.function_scoped_fixture]) +def test_a_per_document_failure_still_only_drops_that_document(vectors): + """The inner handler was already correct and is deliberately untouched. + + Inverting the table-level fallback must not be over-applied: one unreadable + document should cost that document, not the whole result. Here every lookup + fails individually, so everything drops — but via the per-document path, which + must NOT report a table-level degradation. + """ + with patch("apis.shared.assistants.rag_service.emit_count") as emit, patch( + "boto3.resource" + ) as resource, patch.dict( + "os.environ", {"DYNAMODB_ASSISTANTS_TABLE_NAME": "t", "AWS_REGION": "us-west-2"} + ): + table = MagicMock() + table.get_item.side_effect = Exception("per-item failure") + dynamo = MagicMock() + dynamo.Table.return_value = table + resource.return_value = dynamo + + assert _filter_vectors_by_document_status(vectors, ASSISTANT_ID) == [] + emit.assert_not_called() + + +@given(vectors=st_vectors) +@settings(max_examples=50, suppress_health_check=[HealthCheck.function_scoped_fixture]) +def test_complete_documents_are_still_returned(vectors): + """The happy path, so the properties above cannot pass by always returning [].""" + with patch("apis.shared.assistants.rag_service.emit_count"), patch( + "boto3.resource" + ) as resource, patch.dict( + "os.environ", {"DYNAMODB_ASSISTANTS_TABLE_NAME": "t", "AWS_REGION": "us-west-2"} + ): + table = MagicMock() + table.get_item.return_value = {"Item": {"status": "complete"}} + dynamo = MagicMock() + dynamo.Table.return_value = table + resource.return_value = dynamo + + assert len(_filter_vectors_by_document_status(vectors, ASSISTANT_ID)) == len(vectors) + + +@pytest.mark.parametrize("status", ["deleting", "failed", "uploading", "chunking"]) +def test_a_non_complete_status_is_excluded(status): + """Unchanged behaviour, pinned: only `complete` is served. + + Production carried 200 of 1,692 document records in a non-complete state + (101 deleting, 95 failed, 4 uploading), so this is the common case, not an edge. + """ + with patch("apis.shared.assistants.rag_service.emit_count"), patch( + "boto3.resource" + ) as resource, patch.dict( + "os.environ", {"DYNAMODB_ASSISTANTS_TABLE_NAME": "t", "AWS_REGION": "us-west-2"} + ): + table = MagicMock() + table.get_item.return_value = {"Item": {"status": status}} + dynamo = MagicMock() + dynamo.Table.return_value = table + resource.return_value = dynamo + + vectors = [ + {"key": "v1", "distance": 0.1, "metadata": {"document_id": "doc-a", "text": "t"}} + ] + assert _filter_vectors_by_document_status(vectors, ASSISTANT_ID) == [] diff --git a/backend/tests/rbac/test_role_mutation_constraints.py b/backend/tests/rbac/test_role_mutation_constraints.py index 23dbad37f..34ab35efc 100644 --- a/backend/tests/rbac/test_role_mutation_constraints.py +++ b/backend/tests/rbac/test_role_mutation_constraints.py @@ -6,6 +6,13 @@ strict format. These checks are enforced at the service layer so they apply regardless of whether the call originates from the admin REST API, a CLI script, or future automation. + +On the format axis: single *internal* spaces are accepted, because real Entra +security groups are named as display names ("PSEmeriti Entra Sync") and the +tenant owner picks those names. Everything that cannot round trip through the +``custom:roles`` claim stays rejected -- commas (the claim delimiter), edge +whitespace (both claim parsers ``.strip()`` every entry, so a padded mapping +could never match), and every non-space whitespace or invisible character. """ from __future__ import annotations @@ -39,7 +46,25 @@ def service(mock_app_role_repo, mock_app_role_cache) -> AppRoleAdminService: @pytest.mark.parametrize( "forbidden", - ["default", "DEFAULT", "Default", "*", "user", "users", "everyone", "anyone", "authenticated", "all"], + [ + "default", + "DEFAULT", + "Default", + "*", + "user", + "users", + "everyone", + "anyone", + "authenticated", + "all", + # Now that spaces are accepted, the ubiquitous groups have a spelling + # that could not previously be typed at all. "All Users" and + # "Authenticated Users" are real Entra/AD display names for exactly + # the populations this rule exists to keep off a protected role. + "All Users", + "Authenticated Users", + "domain users", + ], ) @pytest.mark.asyncio async def test_protected_role_rejects_ubiquitous_jwt_mapping(service, mock_app_role_repo, make_app_role, admin, forbidden: str) -> None: @@ -69,11 +94,14 @@ async def test_protected_role_accepts_specific_group_mapping(service, mock_app_r mock_app_role_repo.get_role.return_value = system_admin_role mock_app_role_repo.update_role.return_value = system_admin_role - updates = AppRoleUpdate(jwt_role_mappings=["system_admin", "platform_admin"]) + updates = AppRoleUpdate( + jwt_role_mappings=["system_admin", "platform_admin", "Platform Admins Entra Sync"] + ) result = await service.update_role("system_admin", updates, admin) assert result is not None assert "platform_admin" in result.jwt_role_mappings + assert "Platform Admins Entra Sync" in result.jwt_role_mappings # --------------------------------------------------------------------------- @@ -115,12 +143,33 @@ async def test_non_protected_role_can_have_default_mapping(service, mock_app_rol "", # empty "x", # too short "a" * 65, # too long - "has spaces", + "a" * 62 + " bb", # 65 chars: the length bound still holds with a space "has/slash", "has.dot", + # A comma is the delimiter in a comma-separated ``custom:roles`` claim + # and in the admin form, so a comma-bearing group name is + # unrepresentable and must stay rejected even now that spaces are not. "has,comma", "" + "A" * 200]) + + with pytest.raises(ValueError) as excinfo: + await service.update_role("standard_user", updates, admin) + + message = str(excinfo.value) + assert "