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/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..f9bd1c17c 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,43 @@ 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 }} # 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..ca16cfbcd --- /dev/null +++ b/.kiro/specs/managed-kb-migration/HANDOFF.md @@ -0,0 +1,522 @@ +# Managed KB Migration — Handoff + +**Last updated:** 2026-08-26 (groups 11–13, group 14 backend half, tag contract, **14.3 upgrade UX + enrolment surface**) · **Branch:** `feature/kb-migration` · **Nothing deployed** + +Working state for this feature so a fresh session can pick it up without re-deriving +anything. Read this, then `tasks.md`. + +--- + +## 1. Status + +| | | +|---|---| +| Spec | Complete, audited 3× to clean. 25 requirements, 201 criteria, 0 dangling refs | +| Implementation | Groups **1–13** done, plus group 14 except 14.5. 14.4's one-click document retry is deferred. 4 subtasks left: 14.4's retry, 14.5, and group 15 | +| Tests | 617 infra (jest) · **6,603** backend (pytest, 6 m 20 s) · **1,886** frontend (vitest, 7 s) · 5 pre-existing unrelated failures | +| Deployed | **Nothing.** No `cdk deploy`, no AWS mutation, at any point | +| Feature flags | `migrationEnabled` **on in development**, off in production. `newDefault` and `reconcilerArmed` off in both (explicit `false`, set as GitHub Environment variables) | + +### Commits (16 on the branch, all pushed) + +``` +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) +``` + +**Uncommitted working tree:** the 14.3 upgrade surface — `apis/app_api/kb_upgrade/`, +two transitions appended to `kb_backend/records.py`, the Angular card and service, +three test files. 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). + +What is still missing is the **worker's deployment**, not its code. The dispatcher +and worker are Lambdas behind an undeployed image, so an enrolled record sits in +`shadow` indefinitely and the card shows perpetual progress. That is the correct +local behaviour, not a bug. + +**Three behaviour changes ARE live on the existing path** and are the only things +worth testing by hand right now: +1. The document-status filter now **fails closed** (group 6). +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 # 611 passing + +# 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. + +### 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. + +### 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. + +--- + +## 6. Remaining work + +| Group | Subtasks | Notes | +|---|---|---| +| **14** Surfaces | 1½ | **14.5** admin surface (filter by engine, stored bytes + document counts, bulk migrate, per-KB retry) — not started. **14.4** is surfaced but its one-click document retry is deferred; see the deferral below. 14.0–14.3, 14.6, 14.7 are **done**. | +| **15** Pre-promotion verification | 3 | The gate before any real traffic moves. | + +### Known deferrals (correct, not oversights) + +- **One-click document reprocess (Req 21.2).** Ingestion is S3-event-triggered + (`documents/ingestion/handler.py`) and there is **no reprocess endpoint** — the + only document writes are upload-url, import, upload-failed and delete. A retry + control therefore needs new backend that re-fires the pipeline against bytes + already in S3, which is a change to a live ingestion path. Deliberately not + improvised. The card directs the user to re-upload via "Add files", a retry path + that works today. **Close by building the endpoint or by amending Req 21.2 to + accept re-upload** — do not leave it ambiguous. +- **`backend/Dockerfile.kb-migration`** does not exist yet, on purpose. The real image + needs five artefacts that do not exist: the handler modules, their + `requirements.txt`, a case in `scripts/build/build-one.sh`, `backend.yml` jobs, and + entries in the **hand-maintained** lists in + `backend/tests/supply_chain/test_dockerfile_pinning.py` and + `test_lambda_image_imports.py`. Per platform-as-bootstrap, CDK ships the bootstrap + stub and the **workflow** ships the real image. +- **Reconciler EventBridge wiring** (Reqs 14.1, 14.7) — `infrastructure/`, platform + group. Backend code never deploys before the IAM and resources it requires. +- **Group 7's snapshot reservation now has its caller** (`run_shadow`), reserving + the whole corpus before anything is provisioned. + +--- + +## 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/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..a543fd7ab --- /dev/null +++ b/.kiro/specs/managed-kb-migration/requirements.md @@ -0,0 +1,858 @@ +# 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 set `embeddingModelType: CUSTOM` pinned to + `amazon.titan-embed-text-v2:0` at `FLOAT32` (the service-model enum value; lowercase is rejected) and 1024 dimensions. +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/CHANGELOG.md b/CHANGELOG.md index d65e9d16e..46f158e59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,69 @@ 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.16.0] - 2026-08-28 + +Minor release on knowledge bases, the marketplace review flow, and fine-tuning. The Bedrock Managed Knowledge Base migration lands **inert** — every managed-KB flag is default OFF, and the owner-facing upgrade card only appears where `CDK_MANAGED_KB_MIGRATION_ENABLED=true`. Marketplace admins can finally read, test-drive and decline a submission instead of approving on a name and a category alone. Fine-tuning was unreachable in every deployed environment and now isn't — it becomes **reachable by default** on this deploy, with `CDK_FINE_TUNING_ENABLED=false` as the kill switch. Two live RAG behaviours change regardless of flags: queries clamp at 10,000 characters on both backends, and the document-status filter now fails **closed**. **Requires a CDK deploy** — `platform.yml` (set `CDK_TAG_ENVIRONMENT` first), then `backend.yml` (two new kb-migration jobs), then `frontend-deploy.yml`. + +### 🚀 Added + +- **Bedrock Managed Knowledge Base migration**, shipped dormant behind nine `CDK_MANAGED_KB_*` flags. Adds a shared Bedrock KB service role, four `DockerImageFunction`s on one image (worker, dispatcher, reconciler, ingestion consumer) with a DLQ, three EventBridge rules, four CloudWatch alarms, and a sparse `KbWorkIndex` GSI (#884) +- Owner-facing knowledge-base upgrade card and `/assistants/{id}/knowledge-base/upgrade` router (`GET` status, `POST` enroll, `/retry`, `/notice`) with derived phases `none|available|in_progress|succeeded|failed`; also surfaces stranded documents the ordinary list hides (#884) +- Marketplace submission review — `GET /admin/agents/{agent_id}/submission` returns the frozen `submittedVersion` snapshot (instructions, bound capabilities, model, starters, publisher, reachability) behind a new SPA route `/admin/marketplace/review/:agentId` (#882) +- Reviewer **test drive** — a `review_preview` invocation resolves the reviewed snapshot and bypasses the PRIVATE check after re-resolving `admin.marketplace` against the caller's own roles (new `has_admin_scope` predicate); runs on a `preview-` session and skips bookkeeping writes (#882) +- New terminal listing state **`rejected`** — admin-only from `in_review`, requires a reason, allows revise-and-resubmit; `rejected → private` keeps delete reachable. No edge to `published` (#882) +- `@`-mentions are set apart in the user's own message, matched against the known agent-name list rather than `@\w+` so names with spaces work and `@here`/npm scopes/emails aren't bolded (`mention-text.component.ts`) +- `scripts/teardown/managed-kb.sh` — tag-scoped teardown for managed KBs, which are created at runtime and are not CloudFormation children; `destroy.sh` runs it as Phase 0 and aborts the teardown on failure (#884) +- `scripts/local-dev/refresh-env.py` — rebuilds `backend/src/.env` from the deployed app-api task definition, preserving a `KEEP_LOCAL` set; dry-run by default (#884) + +### ✨ Improved + +- Submitting an agent no longer requires ticking a mandatory "make public" checkbox — an amber disclosure states that submitting makes the agent public, and `makePublic` is always sent. Backend `_visibility_block` is unchanged, so a direct API caller that omits it is still refused +- The review test-drive panel is viewport-sized and sticky while the reviewer scrolls the instructions, with an expand control that spans both columns by class change only, so the reviewer's conversation survives the toggle. Empty summary/system prompt render explicit "none" text (#883) + +### 🐛 Fixed + +- **Fine-tuning was unreachable in every deployed environment** — `/api/fine-tuning/access` returned 404 because `FINE_TUNING_ENABLED` was never set on the app-api container, and `CDK_FINE_TUNING_ENABLED` had had no reader since the single-stack migration (#396). `config.ts`, `app-api-service-construct.ts` and `platform.yml` now wire it, plus `CDK_FINE_TUNING_DEFAULT_QUOTA_HOURS` and the never-forwarded `CDK_FINE_TUNING_CORS_ORIGINS` +- **A JSONL dataset — the format the upload page names first — died ~5 billed GPU-minutes into training** with `No CSV file found`. `sagemaker_scripts/train.py` reads JSONL and JSON alongside CSV via a dispatch table and validates the promised `text`/`label` columns; unreadable formats are rejected at `/presign` and again at `POST /jobs`, the last gate before SageMaker provisions a GPU (#893) +- **The admin fine-tuning cost dashboard reported $0.00 and 0 jobs for every period** while jobs were billing. Three stacked faults: the `StatusIndex` PK was queried with SageMaker's `"Completed"`/`"Stopped"` casing against records stored `"COMPLETED"`/`"STOPPED"`; `FAILED` was excluded although AWS bills partial runs; and training/inference share the table and index, so training is now filtered on the `JOB#` sort-key prefix +- **An unpriced `instance_type` ran real GPUs and recorded $0.00** — it arrived unvalidated off the request body on both create paths, and `calculate_cost` falls back to `0.0` outside its 11-entry map. The quota meters GPU-hours, not dollars, so a 10-hour allowance buys ~$14 on `ml.g5.xlarge` or several hundred on a larger unlisted type. Both paths validate the resolved type and 400 with the supported list (#894) +- **Arrow keys could not walk the `@`-mention menu** — `keyup` re-ran `syncMentionToken`, resetting the highlight to row one on every press. The highlight now resets only when the token itself changed, and the active row scrolls into view (#895) +- **JWT role mappings rejected IdP group names containing spaces**, 400-ing the whole `PATCH /api/admin/roles/{role}` payload including untouched entries. The pattern now allows single internal spaces; the error names the offending entry with invisible characters escaped as ``, and `_FORBIDDEN_PROTECTED_MAPPINGS` compares case-folded with space/hyphen/underscore as one separator so `All Users` stays blocked now that it's typeable (#880) +- Image attachments showed the browser's broken-image glyph in long turns — presigned GET URLs live 10 minutes and `loading="lazy"` tiles often fetch after expiry. All three render sites now re-mint once on `(error)`; the lightbox pre-refreshes within 30s of expiry (#879) +- The managed-KB dispatcher raised `KB_MIGRATION_WORKER_FUNCTION_NAME is not set` on every tick — the construct published `MANAGED_KB_WORKER_FUNCTION_NAME`. A second mismatch silently replaced the operator's configured retention window with the 30-day code floor. A new contract test parses every `os.environ` read in the handlers and asserts the construct sets each and publishes nothing unread (#887) +- The first real migration failed with `AccessDeniedException` on `bedrock:TagResource` — AWS authorizes tagging separately from `CreateKnowledgeBase`. Adds `bedrock:TagResource` and `bedrock:ListTagsForResource`; without the second, the daily orphan sweep would have reported a clean account forever (#888) +- A dev deploy tagged its managed KBs `ManagedKbEnvironment=prod` — `config.production` is `true` everywhere and `config.tags` had no `Environment` key — so the tag-scoped teardown would have matched nothing and reported success, leaving billed KBs alive. New `CDK_TAG_ENVIRONMENT` variable, forwarded as flat-dotted `--context tags.Environment=` (#885) +- A refused marketplace decision showed the backend message twice, inline and as a global toast. `SUPPRESS_ERROR_TOAST` is set on exactly the four calls that render inline; takedown on the Listings page has no inline region and keeps its toast, pinned by a test (#890) +- The inline decision error rendered 565px above the sticky decision bar — off-screen at the moment of the click. Load failures stay at the top; a new `decisionError()` renders inside the sticky bar directly above Approve / Request changes / Decline + +### ⚠️ Changed + +- **The RAG document-status filter now fails CLOSED.** Both table-level fallbacks in `_filter_vectors_by_document_status` (unset `DYNAMODB_ASSISTANTS_TABLE_NAME`, and the outer `except`) drop every chunk instead of returning them unfiltered, log at ERROR and emit `KbStatusFilterFailClosed`. Any retrieval-serving service missing that variable now returns zero chunks (#884) +- **Retrieval queries clamp at 10,000 characters on both backends**, including the legacy S3 Vectors path that previously accepted ~32,000. Applied in the facade before dispatch, emits `KbQueryClamped`, never raises (#884) +- Retrieval runs under an explicit `kb_access.granted(...)` grant threaded from both call sites instead of the facade re-resolving. On the marketplace `review_preview` path the permission is deliberately `None`, so **a reviewer test-drives a RAG-backed agent with an empty knowledge base** — fail-closed by design (#882, #884) +- `rag_service.py` is now a facade over a `KnowledgeBaseBackend` protocol with score direction normalized to `relevance` (higher-is-better) inside the S3 Vectors adapter; the `distance` key is still emitted for the existing HTTP consumer. Public signature and both call sites unchanged (#884) +- `.txt` dropped from the fine-tuning **training** upload copy — it cannot express a label. It remains valid for inference input (#893) + +### 🏗️ Infrastructure + +- New GSI **`KbWorkIndex`** (`GSI7_PK`/`GSI7_SK`, projection ALL) on the existing rag-assistants table — sparse, written exclusively by `kb_backend/records.py`, registered in `gsi-inventory.json`. This is the release's one GSI operation; anything else adding an index to that table must ship separately (#884) +- `ManagedKbRoleConstruct` — one Bedrock KB service role with `aws:SourceAccount` + `ArnLike AWS:SourceArn` confused-deputy conditions, S3 read conditioned on `aws:ResourceAccount`, `bedrock:InvokeModel` pinned to `amazon.titan-embed-text-v2:0`, and `iam:PassRole` conditioned on `iam:PassedToService` (#884) +- `KbMigrationConstruct` — four Lambdas + DLQ + log groups; dispatcher rule `rate(15 min)` created **disabled** unless `migrationEnabled`, reconciler rule `rate(1 day)` **always enabled** in report-only mode, documents rule enabled when either flag is on; four alarms in namespace `${projectPrefix}/ManagedKb`, all NOT_BREACHING on missing data; SSM parameters publishing the four function names (#884) +- `documentsBucket.enableEventBridgeNotification()` — additive, because S3 rejects two overlapping-prefix notification configs and the existing rag-ingestion notification is untouched (#884) +- Nine new `CDK_MANAGED_KB_*` variables. The three booleans deliberately invert the repo's default-ON idiom — managed storage is $5.00/GB-month against ~$0.15 — using `parseBooleanEnv`, which maps unset *and* empty to `undefined` so an unset Actions variable cannot arm them (#884) +- `scripts/common/load-env.sh` forwards each `CDK_MANAGED_KB_*` only when non-empty, as **flat dotted** `--context managedKb.=` (`--context a.b=c` sets `context["a.b"]`; it does not build a nested object), and fails at deploy time naming any boolean that isn't `true|false|1|0|empty` (#884, #885) +- New app-api env vars: `MANAGED_KB_MIGRATION_ENABLED`, the byte caps, `MANAGED_KB_METRIC_NAMESPACE`, `FINE_TUNING_ENABLED`, `FINE_TUNING_DEFAULT_QUOTA_HOURS` + +### 🔧 CI/CD + +- `backend.yml` gains `build-kb-migration` (ubuntu-24.04-arm) and `deploy-kb-migration-code`, which points all four functions at the one image tag. Until they run, the Lambdas stay bootstrap no-ops (#886) +- `Dockerfile.kb-migration` pins `boto3==1.43.68` **functionally, not for hygiene** — the base image's 1.40.4 has no `managedKnowledgeBaseConfiguration` shape and every `CreateKnowledgeBase` fails with `ParamValidationError`. Three tests guard the pin (#886) + +### 📚 Docs + +- Weekly kaizen research scan and review prep for 2026-08-28; review queue trimmed 39 → 38 open, retiring the tool-mutation-probe entry and striking two MCP Apps prerequisites and a stale Strands #3758 caveat (#891, #892) +- `.github/docs/deploy/step-03-github-config.md` documents the managed-KB flags and their deliberate default-OFF posture (#884) + ## [1.15.0] - 2026-08-17 Minor release on attachments and external tools. PowerPoint decks can now be uploaded and handed to the PowerPoint toolset, and attachment cards survive a reload — a gap that also affected spreadsheets. OAuth-gated MCP servers no longer lose their tools permanently when the pre-flight runs on a cold token cache, and an abandoned consent prompt no longer bricks every later message in the conversation. Two IAM grants missing in production are fixed, and ALB access logs are enabled so a mid-stream disconnect can be attributed to whoever actually ended the connection. **Requires a CDK deploy** — run `platform.yml`, then `backend.yml`, then `frontend-deploy.yml`. diff --git a/README.md b/README.md index 647e31a34..6539277a1 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ **An open-source, production-ready Generative AI platform for institutions** *Built by Boise State University, designed for everyone.* -[![Release](https://img.shields.io/badge/Release-v1.15.0-6366f1?style=flat&logo=github&logoColor=white)](RELEASE_NOTES.md) +[![Release](https://img.shields.io/badge/Release-v1.16.0-6366f1?style=flat&logo=github&logoColor=white)](RELEASE_NOTES.md) [![Nightly](https://github.com/Boise-State-Development/agentcore-public-stack/actions/workflows/nightly.yml/badge.svg)](https://github.com/Boise-State-Development/agentcore-public-stack/actions/workflows/nightly.yml) ![Python](https://img.shields.io/badge/Python-3.13+-3776AB?style=flat&logo=python&logoColor=white) @@ -296,7 +296,7 @@ agentcore-public-stack/ See [RELEASE_NOTES.md](RELEASE_NOTES.md) for the full changelog, including new features, bug fixes, platform upgrades, and deployment notes for each release. -**Current release:** v1.15.0 +**Current release:** v1.16.0 --- diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 55fb1736a..9b976bbdf 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -1,3 +1,185 @@ +# Release Notes — v1.16.0 + +**Release Date:** August 28, 2026 +**Previous Release:** v1.15.0 (August 17, 2026) + +--- + +> 🏗️ **CDK deploy required, and this one has prerequisites.** Set `CDK_TAG_ENVIRONMENT` (`dev` / `prod`) as a GitHub Environment variable **before** running `platform.yml`, then `backend.yml` — which now carries two new kb-migration jobs that must run before the managed-KB Lambdas are anything but no-ops — then `frontend-deploy.yml`. +> +> `infrastructure/gsi-inventory.json` adds **one** index, `KbWorkIndex`, to the existing rag-assistants table. One GSI operation per `UpdateTable` is the hard DynamoDB limit, so nothing else may add an index to that table in this release. +> +> **Two changes take effect on the next deploy whether or not you set any flag:** fine-tuning becomes reachable in every environment, and RAG retrieval's document-status filter now fails closed. Read the Deployment notes before shipping. + +--- + +## Highlights + +A minor release about knowledge bases, marketplace governance, and a feature that turned out never to have been switched on. + +**The Bedrock Managed Knowledge Base migration lands inert.** ~29,000 lines — a Bedrock service role, four Lambdas, a migration saga, a reconciler, an owner-facing upgrade card — all of it behind nine `CDK_MANAGED_KB_*` flags that default OFF. That inversion of the repo's usual default-ON idiom is deliberate: managed storage bills at $5.00/GB-month against roughly $0.15 for the S3 Vectors path, so an unset Actions variable must never be able to arm it. What *is* live from day one is the reconciler, in report-only mode, so you can read what it would have deleted before you let it delete anything. + +**Marketplace admins can finally review a submission.** Approving an agent — publishing someone's instructions, tools and model to everyone — previously showed a name, an author and a category. Reviewers can now read the frozen submitted snapshot, test-drive it in a live chat panel, and decline it with a reason. + +**Fine-tuning was unreachable in every deployed environment.** `FINE_TUNING_ENABLED` was never set on the app-api container, so `/api/fine-tuning/access` had been returning 404 everywhere since the single-stack migration. The wiring lands here, which means the feature **turns on with this deploy** unless you set `CDK_FINE_TUNING_ENABLED=false`. Three bugs behind that 404 are fixed alongside it: JSONL datasets died five billed GPU-minutes into training, the admin cost dashboard read $0.00 while jobs were billing, and an unpriced instance type ran real GPUs and recorded nothing. + +**Two live RAG behaviours change with no flag to gate them.** Queries clamp at 10,000 characters on both backends, and the document-status filter fails closed — an environment missing `DYNAMODB_ASSISTANTS_TABLE_NAME` on a retrieval-serving service will now return zero chunks rather than unfiltered ones. + +--- + +## Managed Knowledge Base migration + +Assistant knowledge bases can move from the S3 Vectors pipeline onto Bedrock Managed Knowledge Bases. Everything needed to do that ships in this release and none of it runs yet. + +The shipped state is dormant on purpose. Managed storage is $5.00/GB-month; the existing path is about $0.15. A flag that defaults ON with a kill switch — the repo's normal idiom — would mean a forgotten variable costs money, so the three behavioural flags use `parseBooleanEnv`, which maps unset *and* empty string to `undefined`. An Actions variable that exists but is blank cannot arm the feature. + +### The flags + +| GitHub variable | Default | Effect | +|---|---|---| +| `CDK_MANAGED_KB_NEW_DEFAULT` | off | New knowledge bases provision managed instead of legacy | +| `CDK_MANAGED_KB_MIGRATION_ENABLED` | off | Creates the dispatcher rule enabled *and* un-no-ops the handler; also gates the owner-facing upgrade card | +| `CDK_MANAGED_KB_RECONCILER_ARMED` | off | Reconciler **deletes** orphans instead of reporting them | +| `CDK_MANAGED_KB_PER_OWNER_BYTES` | 100 MB | Per-owner storage allowance | +| `CDK_MANAGED_KB_PER_OWNER_ELEVATED_BYTES` | 1 GB | Elevated per-owner allowance | +| `CDK_MANAGED_KB_PER_KB_CEILING_BYTES` | 500 MB | Per-knowledge-base ceiling | +| `CDK_MANAGED_KB_RETENTION_WINDOW_DAYS` | 30 | Post-promotion retention; must stay ≥ 30 | +| `CDK_MANAGED_KB_STORAGE_ALARM_GB` | 500 | Total-storage alarm threshold | +| `CDK_MANAGED_KB_DAILY_COST_ALARM_USD` | 100 | Daily-cost alarm threshold | + +`scripts/common/load-env.sh` forwards each only when non-empty, and validates the three booleans at deploy time — a typo fails the deploy naming the variable rather than silently reading as false. + +### Backend + +- `apis/app_api/kb_upgrade/` — its own package, deliberately outside `kb_migration/` for image size. Router mounted at `/assistants/{assistant_id}/knowledge-base/upgrade`: `GET ""` (status, readable by any resolved permission, reports `canUpgrade: false` to viewers), `POST ""` (enroll), `POST "/retry"`, `POST "/notice"`. Clients see the derived phases `none|available|in_progress|succeeded|failed` and never learn about shadow/verify/promote +- Enrolment is **two conditional writes**, not one. `KbRecord.to_item` does not write the GSI7 work keys, so a single put would produce a record that reports an upgrade in progress and is invisible to the dispatcher forever +- `kb_backend/records.py` is the sole writer of the sparse `GSI7_*` keys; `GSI7_*` joins the generic assistant-update immutable-field guard +- `rag_service.py` becomes a facade over a `KnowledgeBaseBackend` protocol. Score direction is normalized to `relevance` (higher-is-better) by exact negation inside the S3 Vectors adapter, with the `distance` key still emitted for the existing HTTP consumer. Public signature and both call sites unchanged +- Four Lambda handlers plus `kb_backend`, shipped on one image (`backend/Dockerfile.kb-migration`) + +### Frontend + +- `knowledge-base/kb-upgrade.service.ts` and the upgrade card in `knowledge-base-section.component.*`. The card also surfaces stranded documents — distinguishing an unsupported format from a processing failure, and showing `deleting` documents the ordinary list hides +- `app-api-environment.ts` now passes `MANAGED_KB_MIGRATION_ENABLED`, the byte caps and `MANAGED_KB_METRIC_NAMESPACE` to the app-api task. They were absent, so the card would have rendered nothing anywhere regardless of the flag + +### Infrastructure + +- **`ManagedKbRoleConstruct`** — one Bedrock KB service role for all managed KBs, with `aws:SourceAccount` + `ArnLike AWS:SourceArn` confused-deputy conditions, S3 read conditioned on `aws:ResourceAccount`, `bedrock:InvokeModel` pinned to `amazon.titan-embed-text-v2:0`, caller grants split by SID (provisioning CRUD / direct ingestion / inference Retrieve), and `iam:PassRole` scoped and conditioned on `iam:PassedToService` +- **`KbMigrationConstruct`** — worker (15 min / 1024 MB), dispatcher (2 min / 512 MB), reconciler (15 min / 512 MB), ingestion consumer (15 min / 1024 MB with a DLQ), all on one image; four log groups; SSM parameters publishing the four function names +- **EventBridge rules** — dispatcher `rate(15 minutes)`, created disabled unless `migrationEnabled`; reconciler `rate(1 day)`, **enabled unconditionally**, because report-only is the point; documents rule enabled when either `newDefault` or `migrationEnabled` is set +- **Four alarms** in namespace `${projectPrefix}/ManagedKb` — total storage, KB count at 80% of the 10,000 account quota, daily cost, orphan count. All NOT_BREACHING on missing data +- **`KbWorkIndex`** (`GSI7_PK`/`GSI7_SK`, projection ALL) on the rag-assistants table, sparse — keys exist only while a KB is eligible for work +- `documentsBucket.enableEventBridgeNotification()` — additive. S3 rejects two overlapping-prefix notification configurations, so the ingestion consumer is EventBridge-triggered and the existing rag-ingestion notification is left alone +- `infrastructure/lib/config.ts` gains `ManagedKbConfig`, four named default constants, and a resolution chain of `CDK_*` env var → flat dotted context → nested object → literal default + +### Teardown + +Managed knowledge bases are created at **runtime** by the provisioning saga, so they are not CloudFormation children: `delete-stack` leaves them billing at $5.00/GB-month and invisible in the CFN console. `scripts/teardown/managed-kb.sh` scopes **by tag**, never by name pattern, because two environments can share an account. It polls to 480s to match `tombstones.KB_DELETE_POLL_TIMEOUT_SECONDS`, clamps its tunables (an interval of 0 once spun for sixteen hours), and exits non-zero if any matched KB is not confirmed absent. + +`destroy.sh` runs it as **Phase 0**, before any stack, and aborts the whole teardown on failure — the Bedrock service role lives in PlatformStack, and deleting it first is a plausible route into the terminal `DELETE_UNSUCCESSFUL` state. + +### Two operator traps found while shipping this + +Both were caught on a running stack rather than in review, and both are worth knowing about because the failure mode was silence, not an error. + +**The dev environment tagged its knowledge bases `prod`.** `config.production` is `true` in every environment and `config.tags` carried no `Environment` key, so the tag-scoped teardown script would have matched nothing, exited 0, and left billed KBs alive. Fixed by a new `CDK_TAG_ENVIRONMENT` variable forwarded as a flat-dotted `--context tags.Environment=` — `config.production` deliberately untouched. + +**The dispatcher and the construct disagreed about two variable names.** Every tick raised `KB_MIGRATION_WORKER_FUNCTION_NAME is not set` because the construct published `MANAGED_KB_WORKER_FUNCTION_NAME`, so no knowledge base could have migrated. The same sweep found a second mismatch that would have been quieter and worse: the construct published `MANAGED_KB_RETENTION_WINDOW_DAYS` while `worker._retain_days()` reads `KB_MIGRATION_RETAIN_DAYS`, so an operator's configured retention window would have been silently replaced by the 30-day code floor. `backend/tests/supply_chain/test_kb_migration_env_contract.py` now parses every `os.environ` read in the handlers and `kb_backend` and asserts the construct sets each one — and publishes nothing that is never read. + +A third was an IAM gap: AWS authorizes `bedrock:TagResource` separately from `CreateKnowledgeBase`, so the first real migration failed outright. `bedrock:ListTagsForResource` went in with it, and matters more quietly — `tombstones.iter_project_knowledge_bases` reads tags to decide ownership and fails closed, so without it the daily orphan sweep would have reported a clean account forever. + +--- + +## Marketplace submission review + +Approving a marketplace submission publishes another person's instructions, bound tools and model choice to every user. Until now the reviewer saw a name, an author and a category. This release gives them the three things the decision actually needs. + +**Read it.** `GET /admin/agents/{agent_id}/submission` returns instructions, bound capabilities by name, model, conversation starters, publisher and reachability. It always serves the frozen `submittedVersion` snapshot, never the live draft — the author cannot change what is under review while it is under review. New SPA route `/admin/marketplace/review/:agentId` (`submission-review.page.ts`). + +**Test-drive it.** A `review_preview` flag on the inference invocation payload resolves the reviewed snapshot and bypasses the PRIVATE visibility check — after re-resolving `admin.marketplace` against the caller's own roles through a new `has_admin_scope` predicate in `shared/auth/rbac.py`. The preview runs on a `preview-` session and skips bookkeeping writes, so a review leaves no trace in the author's or the reviewer's history. The shared rule lives in `version_resolution.resolve_review_agent`. + +**Decline it.** New terminal listing state `rejected`, admin-only from `in_review`, requiring a reason and allowing revise-and-resubmit. `rejected → private` exists so the author can still delete the agent. There is deliberately no edge from `rejected` to `published`. + +### Submitting now discloses what it does + +`submit-listing-dialog.component.ts` drops the mandatory "Make this agent public" checkbox in favour of an amber disclosure — *submitting makes this agent public* — and always sets `makePublic` on the request. The relaxation is UI-only: the backend `_visibility_block` is unchanged, so a direct API caller that omits the field is still refused. + +### Three follow-ups from using it + +The test-drive panel was sized by its grid cell, which on a short submission left about 120px of usable chat. It is now viewport-sized and sticky while the reviewer scrolls the instructions, with an expand control that spans both columns **by a class change only** — no DOM remount, so the reviewer's conversation survives the toggle. + +A refused Approve reported itself twice, inline and as a global toast. `SUPPRESS_ERROR_TOAST` now rides exactly four calls that render inline — submission read, diff, review decision, withdrawal decision — and deliberately not the whole service: takedown on the Listings page has no inline region, so its toast is kept and a test pins that. + +Then the surviving inline message turned out to render 565px above the sticky decision bar, i.e. off-screen at the moment of the click. Load failures stay at the top; a new `decisionError()` renders inside the sticky bar, directly above Approve / Request changes / Decline. + +### One thing to tell your reviewers + +A RAG-backed agent test-drives with an **empty knowledge base**. Retrieval now runs under an explicit `kb_access.granted(...)` grant threaded from the call sites, and on the `review_preview` path `assistant_permission` is deliberately `None` — fail-closed by design, and called out in-code as an open policy decision. A reviewer who does not know this will read a degraded answer as a broken agent. + +--- + +## Fine-tuning becomes reachable — and correct + +`/api/fine-tuning/access` returned 404 in every deployed environment. `FINE_TUNING_ENABLED` mounts both `/fine-tuning` and `/admin/fine-tuning` and defaults to `"false"` in Python, and nothing ever set it on the app-api container — the existing `CDK_FINE_TUNING_ENABLED` repo variable had had no reader since the single-stack migration in #396. + +Three links close the gap: `config.ts` resolves `fineTuning.enabled` and `defaultQuotaHours`, `app-api-service-construct.ts` sets them on the container, and `platform.yml` forwards `CDK_FINE_TUNING_ENABLED`, `CDK_FINE_TUNING_DEFAULT_QUOTA_HOURS` and `CDK_FINE_TUNING_CORS_ORIGINS` — the last also never forwarded before. This one follows the repo's normal idiom: **default ON, with only the literal `"false"` disabling it.** + +Three bugs the 404 had been hiding go with it. + +**A JSONL dataset died five billed GPU-minutes into training.** JSONL is the first format the upload page names. It uploaded fine, dispatched fine, and then `train.py` failed with `No CSV file found`. `sagemaker_scripts/train.py` now reads JSONL and JSON alongside CSV through a dispatch table and validates that the promised `text`/`label` columns exist. Unreadable formats are rejected at `/presign` and again at `POST /jobs` — the second gate is the last point before SageMaker provisions a GPU. `.txt` is dropped from the training upload copy, since it cannot express a label; it remains valid for inference input. + +**The admin cost dashboard read $0.00 and 0 jobs for every period** while jobs were billing. Three faults stacked: the `StatusIndex` GSI partition key was queried with SageMaker's `"Completed"`/`"Stopped"` casing while records store `"COMPLETED"`/`"STOPPED"`; `FAILED` was excluded although AWS bills partial runs, which the user-facing quota counter was already charging for; and training and inference share the table and index, so fixing the casing alone would have 500'd on a `model_id` KeyError. Training is now filtered on the `JOB#` sort-key prefix. + +**An unpriced instance type ran real GPUs and recorded $0.00.** `instance_type` arrived unvalidated off the request body on both the training and inference create paths, and `calculate_cost` falls back to `0.0` for anything outside its 11-entry `INSTANCE_COST_PER_HOUR` map. The quota meters GPU-hours rather than dollars, so a 10-hour allowance buys roughly $14 on `ml.g5.xlarge` — or several hundred on a larger unlisted type. Both paths now validate the **resolved** type, covering a bad value reaching inference from a stored job record, and 400 with the supported list. Not reachable from the SPA, where instance type is read-only; this was an API-level hole. + +--- + +## Chat and access-control fixes + +**Arrow keys can walk the `@`-mention menu.** `keyup` on the textarea re-ran `syncMentionToken`, which reset `mentionActiveIndex` to 0, so every ArrowDown snapped the highlight back to row one and the menu was unusable by keyboard. The highlight now resets only when the token itself — query or start — actually changed, and `agent-mention-menu.component.ts` scrolls the active row into view, which matters because the list scrolls at eight rows. Adds the composer's first spec. + +**`@`-mentions read as an address in your own message.** `@Brand Deck Builder` renders `font-semibold text-white` against the bubble's `text-white/90`. The new `mention-text.component.ts` matches against the known agent-name list from `AgentMentionService` — the same session-cached list the composer's `@` menu warms — rather than `@\w+`: agent names contain spaces, and a word pattern would also bold `@here`, npm scopes and email addresses. Stored message text is unchanged. + +**JWT role mappings accept IdP group names with spaces.** A mapping like `PSEmeriti Entra Sync` 400'd the whole `PATCH /api/admin/roles/{role}` payload, blocking even untouched entries. `_JWT_MAPPING_PATTERN` now allows single internal spaces (commas, edge whitespace, tab/NBSP/ZWSP still rejected), and the 400 body and log name the offending entry with invisible characters escaped as `` — the failure that motivated this was invisible on screen. `_FORBIDDEN_PROTECTED_MAPPINGS` now compares case-folded with space, hyphen and underscore treated as one separator, so `All Users` and `Authenticated Users` stay blocked now that they are typeable at all. + +**Image attachments stop showing the broken-image glyph.** Presigned S3 GET URLs are minted once with a 10-minute lifetime, and `loading="lazy"` tiles in a long turn frequently fetch after expiry. All three render sites now handle `(error)` with a one-shot re-mint: `image-attachment-group.component.ts` restores the retry budget on `(load)` and pre-refreshes when the lightbox opens within 30s of expiry, `image-lightbox.component.ts` gains an `imageError` output and a guard against ``, and PDF page-1 thumbnails in `file-attachment-badge.component.ts` fall back to the skeleton. + +--- + +## ⚠️ Changed — live RAG behaviour, ungated + +Both of these apply to the existing S3 Vectors path on the next backend deploy. Neither is behind a managed-KB flag. + +**The document-status filter fails closed.** Both table-level fallbacks in `_filter_vectors_by_document_status` — an unset `DYNAMODB_ASSISTANTS_TABLE_NAME`, and the outer `except` — now drop *every* chunk rather than returning them unfiltered. Each logs at ERROR and emits `KbStatusFilterFailClosed`. This supersedes reliable-document-deletion Req 3.4; the per-document handler is unchanged. Verify the variable is set on every service that serves retrieval, and watch the metric after deploy. + +**Queries clamp at 10,000 characters.** Applied in the facade before dispatch, so it also clamps the legacy path, which was previously unclamped up to Titan's roughly 32,000. It emits `KbQueryClamped` and never raises. + +--- + +## 🧪 Test coverage + +Roughly 2,600 lines of new infrastructure tests alone — `kb-migration.test.ts` (1,032), `managed-kb.test.ts` (596), `config.test.ts` (367), `fine-tuning-runtime-flags.test.ts` (100) — plus the backend env-contract test that parses every `os.environ` read in the kb-migration handlers and asserts the construct sets each one, three tests guarding the `boto3==1.43.68` pin, a test pinning takedown's toast against the four suppressed calls, and the chat composer's first spec. + +--- + +## 🚀 Deployment notes + +Order: **`CDK_TAG_ENVIRONMENT` first**, then `platform.yml`, then `backend.yml`, then `frontend-deploy.yml`. + +1. **Set `CDK_TAG_ENVIRONMENT` before deploying** — `dev` for the development environment, `prod` for production, as a GitHub Environment variable. Without it, managed knowledge bases are tagged `prod` even in dev, and the tag-scoped teardown script matches nothing while reporting success. +2. **One GSI, deliberately.** `KbWorkIndex` on the existing rag-assistants table is this release's single index operation. `UpdateTable` permits exactly one GSI create or delete per call and CloudFormation issues one per changed table, so any other change adding an index to that table must ship in a separate release — the 1.12.0 outage shape. +3. **Run the new `backend.yml` jobs.** `build-kb-migration` → `deploy-kb-migration-code` must run at least once, or the four Lambdas stay bootstrap no-ops. Do not relax the `boto3==1.43.68` pin in `Dockerfile.kb-migration`: the base image's 1.40.4 cannot express `managedKnowledgeBaseConfiguration` and every `CreateKnowledgeBase` fails at runtime. +4. **Fine-tuning becomes reachable with this deploy.** `CDK_FINE_TUNING_ENABLED` is already `true` at repo level and the flag is default-ON regardless, so the routes mount; set it to `false` explicitly to keep the feature dark. Access stays gated: `CDK_FINE_TUNING_DEFAULT_QUOTA_HOURS` is `0`, which is **whitelist-only** — a user with no explicit grant in the `fine-tuning-access` table gets a 403, and no grant is auto-provisioned. Any value above 0 flips it to open access and auto-grants that many monthly GPU-hours to every authenticated user on first use, so confirm the value before raising it. Disabling leaves storage intact, so no dataset or trained model is orphaned. +5. **The document-status filter now fails closed.** Confirm `DYNAMODB_ASSISTANTS_TABLE_NAME` is set on every retrieval-serving service before deploying, and watch `KbStatusFilterFailClosed`. A service missing it will return zero chunks where it previously returned unfiltered ones. +6. **Retrieval queries clamp at 10,000 characters** on both backends. Watch `KbQueryClamped`. +7. **The reconciler runs daily from day one, with every flag off.** Its rule is enabled in report-only mode, and four alarms are created in namespace `${projectPrefix}/ManagedKb`. Expect Lambda invocations, log volume and alarm state. Read its "would have deleted" logs before setting `CDK_MANAGED_KB_RECONCILER_ARMED`. +8. **Teardown order changed.** `destroy.sh` runs `managed-kb.sh` as Phase 0 and exits 1 before deleting any stack if it cannot confirm every matched knowledge base is gone. Scripted teardowns should expect that new failure mode. +9. **New listing state `rejected`.** Any external consumer enumerating marketplace listing statuses needs updating. There is no edge from `rejected` to `published`. +10. **Tell marketplace reviewers that a RAG-backed test drive has an empty knowledge base** — fail-closed by design, easily misread as a broken agent. +11. **The nine `CDK_MANAGED_KB_*` variables are all safe to leave unset.** Unset is the shipped dormant state, and an empty value cannot arm the three behavioural flags. + +--- + # Release Notes — v1.15.0 **Release Date:** August 17, 2026 diff --git a/VERSION b/VERSION index 141f2e805..15b989e39 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.15.0 +1.16.0 diff --git a/backend/Dockerfile.kb-migration b/backend/Dockerfile.kb-migration new file mode 100644 index 000000000..587001b28 --- /dev/null +++ b/backend/Dockerfile.kb-migration @@ -0,0 +1,78 @@ +# kb-migration Lambda image — the four managed-knowledge-base Lambdas. +# +# ONE image, FOUR Lambda functions. The CDK construct +# (infrastructure/lib/constructs/managed-kb/kb-migration-construct.ts) points +# all four at this image and selects a handler per function through +# `ImageConfig.Command`. The command override is function *configuration*, so +# the workflow's `update-function-code --image-uri` swaps code on all four +# without touching their handlers: +# +# dispatcher apis.app_api.kb_migration.dispatcher.lambda_handler +# worker apis.app_api.kb_migration.worker.lambda_handler +# reconciler apis.app_api.kb_migration.reconciler.lambda_handler +# ingestion consumer apis.app_api.kb_migration.ingestion_consumer.lambda_handler +# +# THIS IMAGE REPLACES A BOOTSTRAP STUB. +# `infrastructure/bootstrap-assets/kb-migration/` is what PlatformStack ships on +# first deploy: four no-op handlers that log and return. Until this image is +# pushed, an enrolled knowledge base sits in `shadow` while the dispatcher ticks +# into a stub — which is safe, because work keys are sparse and the first real +# dispatcher tick picks up everything that accumulated. The stub directory must +# stay byte-stable; see the warning in its Dockerfile. +# +# WHY THE COPY SURFACE IS THIS SMALL +# The whole import closure of the four handlers is 16 first-party modules. That +# is not luck: `kb_backend` is a separate package with an EMPTY `__init__.py` +# and stdlib-only module scope, specifically so this image does not have to +# carry `apis.shared.assistants` — whose `__init__` imports `rag_service`, which +# imports the embeddings stack at module scope, which blows the image-size +# budget. `backend/tests/architecture/test_kb_backend_boundary.py` enforces it. +# +# Keep this surface minimal but COMPLETE: it must cover the handlers' whole +# closure INCLUDING function-local imports, which still run on invocation. A +# module the closure reaches but this list omits either kills every cold start +# or — worse — trips an `except ImportError` branch and silently no-ops. +# `backend/tests/supply_chain/test_lambda_image_imports.py` enforces the closure. +# When adding a COPY here, also add the path to the kb-migration +# SOURCE_DIRS/MANIFESTS in scripts/build/build-one.sh so the content-hash tag +# notices the change — otherwise the image is rebuilt under an unchanged tag and +# the deploy is a silent no-op. +# +# Base image digest-pinned to the same digest as the other Lambda images +# (backend/Dockerfile.kb-sync, .rag-ingestion, .scheduled-runs); the +# supply-chain dockerfile-pinning test asserts they agree. + +FROM public.ecr.aws/lambda/python:3.12@sha256:745b0eb8a9787e9c4bfd4fc4cae942399a2225831c96394ae70c0c2a7c7c6168 + +# Dependencies (exact pins). The boto3 pin is what makes the managed knowledge +# base API reachable at all — the base image's bundled copy predates it. See the +# requirements file for the specifics. +COPY backend/src/apis/app_api/kb_migration/requirements.txt /tmp/requirements.txt +RUN pip install --no-cache-dir -r /tmp/requirements.txt + +# Application code — namespace-package layout mirrors backend/src. +# `apis/` and `apis/app_api/` have no `__init__.py` in the repo and rely on +# implicit namespace packages, exactly as the kb-sync image does. +COPY backend/src/apis/shared/__init__.py ${LAMBDA_TASK_ROOT}/apis/shared/__init__.py +COPY backend/src/apis/shared/timestamps.py ${LAMBDA_TASK_ROOT}/apis/shared/timestamps.py + +# observability/emf.py — the metrics helpers publish through it. Copied as the +# package so its `__init__` resolves. +COPY backend/src/apis/shared/observability/ ${LAMBDA_TASK_ROOT}/apis/shared/observability/ + +# kb_backend/ — the whole package. Ten of its modules are in the closure +# (records, protocol, provisioning, managed_backend, byte_cap, tombstones, +# idleness, metrics, tags, and the empty __init__); the remaining few +# (resolver, s3vectors_backend, dual_read, query_guard, resource_policy) are +# not reached from these handlers but ride along because they are small, pure +# Python, and copying the package whole means a future function-local import +# cannot silently fall outside the image. +COPY backend/src/apis/shared/kb_backend/ ${LAMBDA_TASK_ROOT}/apis/shared/kb_backend/ + +# The handlers themselves. +COPY backend/src/apis/app_api/kb_migration/ ${LAMBDA_TASK_ROOT}/apis/app_api/kb_migration/ + +# Default command: the dispatcher. The other three functions override CMD +# through their ImageConfig, so this default is only what an unconfigured +# function would run. +CMD ["apis.app_api.kb_migration.dispatcher.lambda_handler"] diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 610bd4364..4d96813da 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "agentcore-stack" -version = "1.15.0" +version = "1.16.0" requires-python = ">=3.10" description = "Multi-agent conversational AI system with AWS Bedrock AgentCore" readme = "README.md" diff --git a/backend/src/apis/app_api/admin/agents/routes.py b/backend/src/apis/app_api/admin/agents/routes.py index 1ade4c496..d5c013641 100644 --- a/backend/src/apis/app_api/admin/agents/routes.py +++ b/backend/src/apis/app_api/admin/agents/routes.py @@ -26,6 +26,7 @@ review_listing, decide_withdrawal, diff_pending_version, + read_submission_for_review, takedown_listing, ) from apis.shared.assistants.categories import ( @@ -50,6 +51,7 @@ AdminReportRow, AdminReportsResponse, AdminStoreFrontResponse, + AdminSubmissionReview, AgentCategoriesResponse, AgentCategory, AgentCategoryCreateRequest, @@ -145,6 +147,37 @@ async def list_listings( raise HTTPException(status_code=500, detail=f"Failed to list listings: {str(e)}") +@router.get("/{agent_id}/submission", response_model=AdminSubmissionReview) +async def get_agent_submission_review( + agent_id: str, + admin: User = Depends(require_marketplace_admin), +): + """The full reviewer read of a listing — instructions, capabilities, model (D2). + + The queue could name a submission but not show one. ``instructions`` is gated to + owner/editor on ``GET /agents/{id}``, and that read refuses a non-owner outright when + the Agent is PRIVATE — so the person deciding whether to publish could not read the + system prompt or see what the Agent binds, and on a first submission the review diff + (their only other window onto it) is empty by construction. + + Serves the **frozen snapshot**, not the live record. See ``AdminSubmissionReview`` for + why that distinction is the design rather than an implementation detail: the live record + is the author's draft, and approval promotes ``submittedVersion``. + + Deliberately a separate endpoint rather than a widened ``GET /agents/{id}``. That route + is the store's detail read *and* the Agent Designer's form loader; teaching it an admin + bypass would put an access exception on the busiest read in the feature, to serve the + wrong version anyway. + """ + try: + return await read_submission_for_review(agent_id, admin) + except ListingError as e: + raise HTTPException(status_code=e.status_code, detail=e.message) + except Exception as e: + logger.error(f"Error reading submission for review: {e}", exc_info=True) + raise HTTPException(status_code=500, detail=f"Failed to read submission: {str(e)}") + + @router.get("/{agent_id}/diff", response_model=AgentVersionDiffResponse) async def get_agent_review_diff( agent_id: str, diff --git a/backend/src/apis/app_api/admin/fine_tuning/routes.py b/backend/src/apis/app_api/admin/fine_tuning/routes.py index a3531b6c9..f19388503 100644 --- a/backend/src/apis/app_api/admin/fine_tuning/routes.py +++ b/backend/src/apis/app_api/admin/fine_tuning/routes.py @@ -235,6 +235,13 @@ async def list_all_inference_jobs( # ========== Cost Dashboard ========== +# Terminal statuses AWS bills for, in the exact casing persisted on the job +# record by the status maps in ``fine_tuning/routes.py``. The StatusIndex GSI +# compares its partition key case-sensitively, so SageMaker's own "Completed" +# spelling matches nothing here. +BILLED_TERMINAL_STATUSES = ("COMPLETED", "FAILED", "STOPPED") + + def _date_range_for_period(period: str) -> tuple[str, str]: """Return (start_iso, end_iso) for a YYYY-MM period string.""" year, month = int(period[:4]), int(period[5:7]) @@ -259,8 +266,16 @@ async def get_cost_dashboard( ): """Get aggregated fine-tuning cost dashboard for a billing period. - Queries the StatusIndex GSI for Completed and Stopped jobs within - the requested month, then aggregates costs by user in application code. + Queries the StatusIndex GSI for every billed terminal status within the + requested month, then aggregates costs by user in application code. + + Two things are easy to get wrong here. The GSI partition key is compared + case-sensitively, so the status values must be the ones actually stored + (``COMPLETED``/``FAILED``/``STOPPED`` — see the status maps in + ``fine_tuning/routes.py``), not SageMaker's own ``Completed``/``Stopped`` + spelling. And FAILED belongs in the list: AWS bills a job that fails + partway through, so leaving it out understates real spend. The user-facing + quota counter already charges for failures; this dashboard matches it. """ period = month or datetime.now(timezone.utc).strftime("%Y-%m") safe_period = period.replace("\n", "").replace("\r", "") @@ -269,15 +284,19 @@ async def get_cost_dashboard( try: start_iso, end_iso = _date_range_for_period(period) - # Query training jobs (Completed + Stopped) via StatusIndex GSI - training_completed = jobs_repo.query_jobs_by_status_and_date("Completed", start_iso, end_iso) - training_stopped = jobs_repo.query_jobs_by_status_and_date("Stopped", start_iso, end_iso) - all_training = training_completed + training_stopped - - # Query inference jobs (Completed + Stopped) via StatusIndex GSI - inf_completed = inf_repo.query_jobs_by_status_and_date("Completed", start_iso, end_iso) - inf_stopped = inf_repo.query_jobs_by_status_and_date("Stopped", start_iso, end_iso) - all_inference = inf_completed + inf_stopped + # Query training jobs in every billed terminal status via StatusIndex GSI + all_training = [ + job + for status_value in BILLED_TERMINAL_STATUSES + for job in jobs_repo.query_jobs_by_status_and_date(status_value, start_iso, end_iso) + ] + + # Query inference jobs in every billed terminal status via StatusIndex GSI + all_inference = [ + job + for status_value in BILLED_TERMINAL_STATUSES + for job in inf_repo.query_jobs_by_status_and_date(status_value, start_iso, end_iso) + ] # Aggregate by user email user_data: dict[str, dict] = defaultdict( diff --git a/backend/src/apis/app_api/agent_designer/services/listing_service.py b/backend/src/apis/app_api/agent_designer/services/listing_service.py index 47fc1bb8c..4d80f83ee 100644 --- a/backend/src/apis/app_api/agent_designer/services/listing_service.py +++ b/backend/src/apis/app_api/agent_designer/services/listing_service.py @@ -54,11 +54,13 @@ list_versions, set_version_index, ) +from apis.shared.assistants.version_resolution import resolve_review_agent from apis.shared.assistants.versions import snapshot_of from apis.shared.assistants.models import ( AdminEdit, AdminListingPatchRequest, AdminListingRow, + AdminSubmissionReview, AgentListing, AgentVersionDiffResponse, AgentVersionSummary, @@ -81,6 +83,7 @@ from apis.shared.auth.models import User from apis.shared.feature_flags import skills_enabled from apis.shared.memory.service import MemorySpaceService +from apis.shared.security.log_sanitize import scrub_log from apis.shared.skills.repository import get_skill_catalog_repository from apis.shared.timestamps import utc_now_iso @@ -764,20 +767,36 @@ async def review_listing( category: Optional[str] = None, publisher_id: Optional[str] = None, ) -> AgentListing: - """Approve a submission, or return it with a reason (D2). + """Approve a submission, return it with a reason, or decline it (D2). Approval is where an attribution becomes authoritative, so the reviewer may adjust category and publisher in the same act (D12) without a second round trip. + + ``reject`` is the third decision and it is not a synonym for ``request_changes``: it + answers a submission that should not be in the store at all, rather than one that needs + work. Both carry a required reason for the same reason — a decision the author cannot + read is one they cannot act on — and the state machine's note on ``rejected`` records + why both let the author come back. """ assistant = await _load_any(agent_id) if not assistant.listing: raise ListingError("This agent has no marketplace listing to review.", status_code=404) - target = "published" if decision == "approve" else "changes_requested" - if decision == "request_changes" and not (note or "").strip(): + # Mapped rather than branched, so adding a fourth decision cannot silently fall through + # to ``changes_requested`` the way an ``if/else`` on ``approve`` did. + targets = { + "approve": "published", + "request_changes": "changes_requested", + "reject": "rejected", + } + target = targets.get(decision) + if target is None: + raise ListingError(f"Unknown review decision '{decision}'.", status_code=400) + if target != "published" and not (note or "").strip(): + verb = "Declining a submission" if decision == "reject" else "Requesting changes" raise ListingError( - "Requesting changes needs a reason — it renders on the author's card so they " - "never have to ask what happened.", + f"{verb} needs a reason — it renders on the author's card so they never have " + "to ask what happened.", status_code=400, ) @@ -1169,6 +1188,101 @@ async def diff_pending_version(agent_id: str) -> AgentVersionDiffResponse: +async def read_submission_for_review(agent_id: str, admin: User) -> AdminSubmissionReview: + """The reviewer's full read of a listing — instructions, bindings, model and all (D2). + + **Which version this reads is the entire correctness question**, and it is answered the + same way ``review_listing`` answers "which version does approval promote?": the snapshot + named by the listing, never "the latest" and never the live draft. + + * ``in_review`` → ``submitted_version``. That is the artifact approval promotes, so it + is the artifact the reviewer must read. The live record is the author's draft and they + can edit it while the row sits in the queue; reading it would show one configuration + and publish another. + * anything else → ``published_version``. ``submitted_version`` is a high-water mark that + deliberately survives a decision, so on a ``withdrawal_requested`` row — where nothing + is pending and the question is whether to pull what is *live* — it would name a stale + snapshot the store never served. + + Falls back to the live record, flagged, when neither pointer resolves. See + ``AdminSubmissionReview.snapshot_unavailable`` for why that is reported rather than + refused. + + ``admin`` is threaded through only to resolve **memory-space labels**, which have no + unfiltered name lookup (see ``agent_detail._memory_labels``). It is not an access + decision — the route's ``admin.marketplace`` scope already made that one — and nothing + here filters by what this particular admin can reach. + """ + # Imported here rather than at module scope: ``agent_detail`` pulls in the whole + # bindable catalog (five per-primitive services), and every other caller of this module + # — the author paths, the submit dialog's preflight — needs none of it. + from apis.app_api.agent_designer.services.agent_detail import ( + resolve_capabilities, + resolve_listing_display, + ) + + assistant = await _load_any(agent_id) + listing = assistant.listing + if not listing: + raise ListingError("This agent has no marketplace listing to review.", status_code=404) + + # ⚠️ The which-version rule lives in ``version_resolution``, not here, and that is + # load-bearing: the reviewer's *test drive* (``inference_api.chat.routes``) has to + # resolve the same snapshot this page shows, and inference-api cannot import from + # app_api. A second copy of the rule is a page and a preview that disagree about what + # is under review — which is exactly the failure snapshots exist to prevent, one level + # up. + reviewed, review_version = await resolve_review_agent(assistant) + + publisher = await get_publisher(listing.publisher_id) + try: + capabilities, model_label = await resolve_capabilities(reviewed, admin) + except Exception: + # Presentation, exactly as on the user-facing detail read: a catalog hiccup must not + # turn a reviewable submission into a 500 and strand the queue. + logger.warning( + f"Failed to resolve capabilities for review of {scrub_log(agent_id)}", exc_info=True + ) + capabilities, model_label = [], None + try: + _, category_label = await resolve_listing_display(reviewed) + except Exception: + logger.warning(f"Failed to resolve category label for {scrub_log(agent_id)}", exc_info=True) + category_label = None + + return AdminSubmissionReview( + agent_id=agent_id, + name=reviewed.name, + description=reviewed.description, + tagline=reviewed.tagline, + instructions=reviewed.instructions, + starters=list(reviewed.starters or []), + emoji=reviewed.emoji, + icon_url=icon_url(agent_id, reviewed.icon_key), + owner_name=assistant.owner_name, + publisher=publisher, + category=listing.category, + category_label=category_label, + state=listing.state, + capabilities=capabilities, + model_label=model_label, + review_version=review_version, + published_version=listing.published_version, + snapshot_unavailable=review_version is None, + submitted_at=listing.submitted_at, + withdrawal_requested_at=( + listing.withdrawal_requested_at if listing.state == "withdrawal_requested" else None + ), + reviewed_at=listing.reviewed_at, + review_note=listing.review_note, + # ⚠️ Derived from the **live** record, never the snapshot. ``visibility`` is + # deliberately absent from ``AgentVersion`` (fusing it with listing state is the + # trap that class docstring names), and the question here is "can people reach this + # right now?" — a fact about now, which a frozen artifact cannot answer. + reachability=_reachability(assistant), + ) + + async def list_admin_listings(state: Optional[str] = None) -> Tuple[List[AdminListingRow], int]: """Rows for the Review queue / Listings tables, plus the pending-decision count. diff --git a/backend/src/apis/app_api/assistants/routes.py b/backend/src/apis/app_api/assistants/routes.py index 86460c238..279ea7b29 100644 --- a/backend/src/apis/app_api/assistants/routes.py +++ b/backend/src/apis/app_api/assistants/routes.py @@ -52,6 +52,7 @@ update_assistant, update_share_permission, ) +from apis.shared.assistants.kb_access import granted from apis.shared.assistants.rag_service import augment_prompt_with_context, search_assistant_knowledgebase_with_formatting logger = logging.getLogger(__name__) @@ -521,7 +522,15 @@ async def test_chat_endpoint(assistant_id: str, request: AssistantTestChatReques session_id = request.session_id or f"test-{uuid.uuid4().hex[:12]}" # 4. Search vector store for relevant context - context_chunks = await search_assistant_knowledgebase_with_formatting(assistant_id=assistant_id, query=request.message, top_k=5) + # The permission resolved in step 1 is handed to the facade rather than + # re-resolved there (Requirement 25.1): one lookup, and the grant the + # retrieval runs under is provably the one this route checked. + context_chunks = await search_assistant_knowledgebase_with_formatting( + assistant_id=assistant_id, + query=request.message, + top_k=5, + access=granted(assistant_id, user_id, permission), + ) # 5. Augment user message with retrieved context augmented_message = augment_prompt_with_context(user_message=request.message, context_chunks=context_chunks) diff --git a/backend/src/apis/app_api/fine_tuning/inference_repository.py b/backend/src/apis/app_api/fine_tuning/inference_repository.py index 94306dc49..55d50f38b 100644 --- a/backend/src/apis/app_api/fine_tuning/inference_repository.py +++ b/backend/src/apis/app_api/fine_tuning/inference_repository.py @@ -218,7 +218,8 @@ def query_jobs_by_status_and_date( """Query the StatusIndex GSI for inference jobs with a given status in a date range. Args: - status_value: Job status (e.g. "Completed", "Stopped"). + status_value: Stored job status, case-sensitive on the GSI key + (e.g. "COMPLETED", "FAILED", "STOPPED"). start_date: ISO date string (inclusive lower bound on createdAt). end_date: ISO date string (inclusive upper bound on createdAt). diff --git a/backend/src/apis/app_api/fine_tuning/job_models.py b/backend/src/apis/app_api/fine_tuning/job_models.py index a80b84d20..cc1a55238 100644 --- a/backend/src/apis/app_api/fine_tuning/job_models.py +++ b/backend/src/apis/app_api/fine_tuning/job_models.py @@ -223,6 +223,13 @@ class AvailableModel(BaseModel): # Request / Response Models # ========================================================================= +# Dataset formats the SageMaker training script can read — keep in sync with +# SUPPORTED_DATASET_EXTENSIONS in fine_tuning/sagemaker_scripts/train.py. +# Enforced here so an unreadable dataset is rejected before a GPU instance is +# ever provisioned; otherwise the job fails several billed minutes in. +SUPPORTED_DATASET_EXTENSIONS = (".csv", ".jsonl", ".json") + + class PresignRequest(BaseModel): """Request for a presigned upload URL for a training dataset.""" filename: str diff --git a/backend/src/apis/app_api/fine_tuning/job_repository.py b/backend/src/apis/app_api/fine_tuning/job_repository.py index 7c428ea28..f26c2e871 100644 --- a/backend/src/apis/app_api/fine_tuning/job_repository.py +++ b/backend/src/apis/app_api/fine_tuning/job_repository.py @@ -257,24 +257,32 @@ def query_jobs_by_status_and_date( ) -> List[dict]: """Query the StatusIndex GSI for jobs with a given status in a date range. + Training and inference records share this table and this GSI, so the + query filters on the ``JOB#`` sort-key prefix. Without it an inference + record comes back as a training job, and any caller that also queries + the inference repository counts its cost twice. + Args: - status_value: Job status (e.g. "Completed", "Stopped"). + status_value: Stored job status, case-sensitive on the GSI key + (e.g. "COMPLETED", "FAILED", "STOPPED"). start_date: ISO date string (inclusive lower bound on createdAt). end_date: ISO date string (inclusive upper bound on createdAt). Returns: - List of job dicts. + List of training job dicts. """ try: items: List[dict] = [] response = self._table.query( IndexName="StatusIndex", KeyConditionExpression="#s = :status AND createdAt BETWEEN :start AND :end", + FilterExpression="begins_with(SK, :sk_prefix)", ExpressionAttributeNames={"#s": "status"}, ExpressionAttributeValues={ ":status": status_value, ":start": start_date, ":end": end_date, + ":sk_prefix": "JOB#", }, ScanIndexForward=False, ) @@ -284,11 +292,13 @@ def query_jobs_by_status_and_date( response = self._table.query( IndexName="StatusIndex", KeyConditionExpression="#s = :status AND createdAt BETWEEN :start AND :end", + FilterExpression="begins_with(SK, :sk_prefix)", ExpressionAttributeNames={"#s": "status"}, ExpressionAttributeValues={ ":status": status_value, ":start": start_date, ":end": end_date, + ":sk_prefix": "JOB#", }, ScanIndexForward=False, ExclusiveStartKey=response["LastEvaluatedKey"], diff --git a/backend/src/apis/app_api/fine_tuning/routes.py b/backend/src/apis/app_api/fine_tuning/routes.py index e4ae656e6..5e42ad097 100644 --- a/backend/src/apis/app_api/fine_tuning/routes.py +++ b/backend/src/apis/app_api/fine_tuning/routes.py @@ -19,7 +19,9 @@ ) from .job_models import ( AVAILABLE_MODELS, + INSTANCE_COST_PER_HOUR, MODEL_CATALOG, + SUPPORTED_DATASET_EXTENSIONS, PresignRequest, PresignResponse, CreateJobRequest, @@ -193,6 +195,48 @@ async def _fetch_tag(tag: str): # Presigned URL # ========================================================================= +def _validate_dataset_format(name: str) -> None: + """Reject a dataset filename the training script could not read. + + Checked before upload and again before the job is submitted, so an + unreadable dataset never reaches a billed GPU instance. + """ + if not name.lower().endswith(SUPPORTED_DATASET_EXTENSIONS): + supported = ", ".join(SUPPORTED_DATASET_EXTENSIONS) + raise HTTPException( + status_code=400, + detail=( + f"Unsupported dataset format. Supported formats: {supported}. " + 'Each record needs a "text" and a "label" field.' + ), + ) + + +def _validate_instance_type(instance_type: str) -> None: + """Reject an instance type we have no price for. + + ``calculate_cost`` falls back to $0.00/hour for anything absent from + INSTANCE_COST_PER_HOUR, so an unlisted type runs real GPUs and records no + spend — invisible to the admin cost dashboard, the same blind spot the + StatusIndex casing bug produced by a different route. + + The quota does not bound the damage either: it meters GPU-*hours*, not + dollars, so the same ten hours buys ~$14 on an ml.g5.xlarge or several + hundred on a larger instance. `instance_type` arrives straight off the + request body, so this is the only thing standing between a caller and an + unpriced instance. + """ + if instance_type not in INSTANCE_COST_PER_HOUR: + supported = ", ".join(sorted(INSTANCE_COST_PER_HOUR)) + raise HTTPException( + status_code=400, + detail=( + f"Unsupported instance type '{instance_type}'. " + f"Supported types: {supported}" + ), + ) + + @router.post("/presign", response_model=PresignResponse) async def presign_upload( request: PresignRequest, @@ -201,6 +245,8 @@ async def presign_upload( s3_service: FineTuningS3Service = Depends(get_fine_tuning_s3_service), ): """Generate a presigned PUT URL for dataset upload.""" + _validate_dataset_format(request.filename) + try: presigned_url, s3_key = s3_service.generate_upload_url( user_id=user.user_id, @@ -248,7 +294,9 @@ async def create_job( if not hf_id or len(hf_id) > 200: raise HTTPException(status_code=400, detail="Invalid HuggingFace model ID.") - # Verify dataset exists in S3 + # Verify the dataset is readable by the training script and exists in S3 + _validate_dataset_format(request.dataset_s3_key) + if not s3_service.check_object_exists(request.dataset_s3_key): raise HTTPException(status_code=400, detail="Dataset not found in S3. Upload your dataset first.") @@ -281,6 +329,8 @@ async def create_job( huggingface_id = request.custom_huggingface_model_id.strip() model_name = huggingface_id + _validate_instance_type(instance_type) + if request.hyperparameters: hyperparameters.update(request.hyperparameters) hyperparameters["model_name_or_path"] = huggingface_id @@ -717,6 +767,7 @@ async def create_inference_job( # Resolve instance type (default to training job's instance type) instance_type = request.instance_type or training_job["instance_type"] + _validate_instance_type(instance_type) # Generate identifiers job_id = uuid.uuid4().hex diff --git a/backend/src/apis/app_api/fine_tuning/sagemaker_scripts/train.py b/backend/src/apis/app_api/fine_tuning/sagemaker_scripts/train.py index 262a6b1db..145170d68 100644 --- a/backend/src/apis/app_api/fine_tuning/sagemaker_scripts/train.py +++ b/backend/src/apis/app_api/fine_tuning/sagemaker_scripts/train.py @@ -162,19 +162,81 @@ def _valid(v): return min(valid_vals) if valid_vals else None -def find_csv_in_channel(channel_dir): - """Find the first CSV file in a SageMaker input channel directory. - - Raises FileNotFoundError if no CSV file is found. +# Dataset formats the trainer can read, mapped to the pandas reader that loads +# them. A training record has to carry both a "text" and a "label" field, which +# is why plain .txt is absent: it has no way to express the label. (.txt stays +# valid for *inference* input, which is unlabelled — one record per line.) +# +# Kept as data rather than an if/elif chain so the supported-format contract can +# be asserted without importing pandas, which only exists inside the SageMaker +# training container and not in the backend venv. +DATASET_READERS = { + ".csv": ("read_csv", {}), + ".jsonl": ("read_json", {"lines": True}), + ".json": ("read_json", {}), +} + +SUPPORTED_DATASET_EXTENSIONS = tuple(DATASET_READERS) + +REQUIRED_DATASET_COLUMNS = ("text", "label") + + +def find_dataset_in_channel(channel_dir): + """Find the first supported dataset file in a SageMaker input channel. + + Raises FileNotFoundError if the directory is missing, or if it holds no + file with a supported extension. """ if not os.path.isdir(channel_dir): raise FileNotFoundError(f"Channel directory does not exist: {channel_dir}") for f in sorted(os.listdir(channel_dir)): - if f.lower().endswith(".csv"): + if f.lower().endswith(SUPPORTED_DATASET_EXTENSIONS): return os.path.join(channel_dir, f) - raise FileNotFoundError(f"No CSV file found in {channel_dir}") + supported = ", ".join(SUPPORTED_DATASET_EXTENSIONS) + raise FileNotFoundError( + f"No dataset file found in {channel_dir}. Supported formats: {supported}" + ) + + +def resolve_dataset_reader(dataset_path): + """Return the (pandas reader name, kwargs) pair for a dataset file. + + Raises ValueError for an extension the trainer cannot read. + """ + extension = os.path.splitext(dataset_path)[1].lower() + + if extension not in DATASET_READERS: + supported = ", ".join(SUPPORTED_DATASET_EXTENSIONS) + raise ValueError( + f"Unsupported dataset format '{extension}'. Supported formats: {supported}" + ) + + return DATASET_READERS[extension] + + +def validate_dataset_columns(columns, dataset_path): + """Raise ValueError if a required column is absent from the dataset.""" + missing = [c for c in REQUIRED_DATASET_COLUMNS if c not in columns] + if missing: + raise ValueError( + f"Dataset {os.path.basename(dataset_path)} is missing required " + f"column(s): {', '.join(missing)}. Each record needs a \"text\" " + f'and a "label" field.' + ) + + +def load_dataset_frame(dataset_path): + """Load a dataset file into a DataFrame with "text" and "label" columns.""" + import pandas as pd + + reader_name, reader_kwargs = resolve_dataset_reader(dataset_path) + df = getattr(pd, reader_name)(dataset_path, **reader_kwargs) + + validate_dataset_columns(df.columns, dataset_path) + + return df def copy_inference_script(model_output_dir): @@ -221,10 +283,10 @@ def train(args): ) model_dir = os.environ.get("SM_MODEL_DIR", "/opt/ml/model") - # Find and load CSV dataset - csv_path = find_csv_in_channel(train_channel) - logger.info(f"Loading dataset from {csv_path}") - df = pd.read_csv(csv_path) + # Find and load the dataset (CSV / JSONL / JSON) + dataset_path = find_dataset_in_channel(train_channel) + logger.info(f"Loading dataset from {dataset_path}") + df = load_dataset_frame(dataset_path) # Label normalization — support non-numeric class labels label_names = sorted(list(pd.Series(df["label"]).astype(str).unique())) diff --git a/backend/src/apis/app_api/kb_migration/__init__.py b/backend/src/apis/app_api/kb_migration/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/src/apis/app_api/kb_migration/dispatcher.py b/backend/src/apis/app_api/kb_migration/dispatcher.py new file mode 100644 index 000000000..fa0de5690 --- /dev/null +++ b/backend/src/apis/app_api/kb_migration/dispatcher.py @@ -0,0 +1,245 @@ +"""Migration dispatcher: hand due knowledge bases to the worker, a few at a time. + +Requirements 19.6, 15.14. One EventBridge tick reads the sparse ``KbWorkIndex``, +takes at most a bounded number of records, and asynchronously invokes the worker +once per record. It performs no migration itself and holds no state. + +Follows ``apis/app_api/kb_sync/dispatcher.py`` closely — third use of that shape on +this table, after the sync dispatcher and the scheduled-runs dispatcher — so the +things that matter about it are already established: a bounded per-tick limit, one +broken record never starving the sweep, and metrics emitted from the tick rather +than from the worker. + +Why the index makes the queue correct by physics +------------------------------------------------ +``GSI7_PK``/``GSI7_SK`` are written only while a record is work-eligible +(``shadow``, ``verify``, ``promote``) and ``REMOVE``d on reaching a terminal state. +So this dispatcher cannot see a finished knowledge base even if it wanted to: there +is no filter to get wrong, because ineligible records are not in the index. Same +convention as ``DueSyncIndex``, ``AgentDirectoryIndex`` and ``AgentReportsIndex`` +on this table. + +Why it no-ops rather than refusing to start +------------------------------------------- +With ``MANAGED_KB_MIGRATION_ENABLED`` off the tick returns its zeroed counts and +invokes nothing. The Lambda still exists, still runs on schedule, and still emits +metrics — which is what makes turning the flag on a change with a known blast +radius rather than the first time this code has ever executed in production. + +Feature: managed-kb-migration +Requirements: 19.6, 15.14 +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +from typing import Any, Dict, List + +logger = logging.getLogger() +logger.setLevel(logging.INFO) + +#: The migration flag. Absent, empty, or anything outside the truthy set means the +#: dispatcher invokes nothing. +FLAG_MIGRATION_ENABLED = "MANAGED_KB_MIGRATION_ENABLED" + +#: Recognised affirmative spellings, matching the reconciler's. An allow-list +#: rather than a truthiness test, because the failure being designed around is a +#: value that is present but empty: ``bool("")`` is correct by luck, +#: ``bool("false")`` is not. +_TRUTHY = frozenset({"1", "true", "yes", "on", "enabled"}) + +#: Mirrors ``KB_SYNC_DISPATCH_LIMIT``'s default of 20 (Requirement 15.14). The +#: limit exists twice over: it bounds the damage of a bug in the index sweep, and +#: it keeps a burst of enrolments from colliding with ``StartIngestionJob``'s +#: 0.1 RPS account-wide ceiling — which is not adjustable, so the only way to stay +#: under it is to not ask. +DEFAULT_DISPATCH_LIMIT = 20 + +#: Ceiling on the env-var override. A larger sweep should require repeated observed +#: ticks, not a variable edit. +DISPATCH_LIMIT_CEILING = 100 + +METRIC_DISPATCHED = "KbMigrationDispatched" +METRIC_DUE = "KbMigrationDue" +METRIC_DISPATCH_FAILED = "KbMigrationDispatchFailed" + + +def migration_enabled() -> bool: + """Whether the dispatcher may invoke the worker at all. + + Read at call time. Bound as a module constant it would be captured at import + and a test overriding the variable would silently get the production value — + the mistake that cost a 33-second test on this feature already. + """ + return (os.environ.get(FLAG_MIGRATION_ENABLED) or "").strip().lower() in _TRUTHY + + +def dispatch_limit() -> int: + """Records taken per tick, bounded above by :data:`DISPATCH_LIMIT_CEILING`.""" + raw = os.environ.get("KB_MIGRATION_DISPATCH_LIMIT") + try: + value = int(raw) if raw else DEFAULT_DISPATCH_LIMIT + except ValueError: + logger.warning( + f"KB_MIGRATION_DISPATCH_LIMIT={raw!r} is not an integer; using " + f"{DEFAULT_DISPATCH_LIMIT}" + ) + return DEFAULT_DISPATCH_LIMIT + if value < 0: + return 0 + if value > DISPATCH_LIMIT_CEILING: + logger.warning( + f"KB_MIGRATION_DISPATCH_LIMIT={value} exceeds the ceiling of " + f"{DISPATCH_LIMIT_CEILING}; clamping" + ) + return DISPATCH_LIMIT_CEILING + return value + + +def _now_iso() -> str: + from apis.shared.timestamps import utc_now_iso + + return utc_now_iso() + + +def _work_states() -> List[str]: + """Every work-eligible state, drained-first. + + Derived from ``WORK_ELIGIBLE_STATES`` rather than restated, with an explicit + priority order laid over it. A record in ``promote`` is one conditional write + from being finished, so serving it ahead of new ``shadow`` work drains the queue + instead of accumulating half-migrated knowledge bases. + + Anything work-eligible but absent from the priority list is appended rather + than dropped. A state added to the records module and forgotten here then + migrates slowly, which is a scheduling nuisance; dropped, it would stall + forever with its work keys written and nothing ever reading them — invisible, + because the record still looks queued. + """ + from apis.shared.kb_backend.records import PROMOTE, SHADOW, VERIFY, WORK_ELIGIBLE_STATES + + priority = (PROMOTE, VERIFY, SHADOW) + ordered = [state for state in priority if state in WORK_ELIGIBLE_STATES] + remainder = sorted(set(WORK_ELIGIBLE_STATES) - set(priority)) + if remainder: + logger.warning( + f"work-eligible states {remainder} are not in the dispatcher's priority " + f"order; sweeping them last" + ) + return ordered + remainder + + +def _invoke_worker(payload: Dict[str, Any]) -> None: + """Async-invoke the migration worker. Same shape as the sync dispatcher's.""" + import boto3 + + function_name = os.environ.get("KB_MIGRATION_WORKER_FUNCTION_NAME") + if not function_name: + raise RuntimeError("KB_MIGRATION_WORKER_FUNCTION_NAME is not set") + + boto3.client("lambda").invoke( + FunctionName=function_name, + InvocationType="Event", + Payload=json.dumps(payload).encode("utf-8"), + ) + + +def _emit_metrics(counts: Dict[str, int]) -> None: + from apis.shared.kb_backend.metrics import emit_count + + for metric, value in ( + (METRIC_DUE, counts.get("Due", 0)), + (METRIC_DISPATCHED, counts.get("Dispatched", 0)), + (METRIC_DISPATCH_FAILED, counts.get("Failed", 0)), + ): + if value: + emit_count(metric, value) + + +async def _due_records(limit: int, now_iso: str) -> List[Dict[str, Any]]: + """Records whose ``dueAt`` has passed, across every work-eligible state. + + Queried per state because ``GSI7_PK`` *is* the state — one partition each — and + trimmed to ``limit`` overall so the bound is on the tick's total work rather + than per state, which is how a three-state sweep would quietly become a + 3× limit. + """ + from apis.shared.kb_backend.records import query_due_work + + collected: List[Dict[str, Any]] = [] + for state in _work_states(): + if len(collected) >= limit: + break + remaining = limit - len(collected) + try: + found = await asyncio.to_thread(query_due_work, state, now_iso, remaining) + except Exception as exc: + logger.error(f"KbWorkIndex query failed for state {state}: {exc}", exc_info=True) + continue + collected.extend(found) + return collected[:limit] + + +async def dispatch_once() -> Dict[str, int]: + """One dispatcher tick. Returns the metric counts (also emitted).""" + counts: Dict[str, int] = {"Due": 0, "Dispatched": 0, "Failed": 0} + + if not migration_enabled(): + logger.info(f"{FLAG_MIGRATION_ENABLED} is not truthy; dispatcher tick is a no-op") + return counts + + limit = dispatch_limit() + if limit == 0: + logger.info("dispatch limit is 0; nothing will be dispatched this tick") + return counts + + now_iso = _now_iso() + due = await _due_records(limit, now_iso) + counts["Due"] = len(due) + logger.info(f"migration dispatcher tick: {len(due)} due records (limit {limit})") + + for record in due: + app_kb_id = record.get("appKbId") + pk = record.get("PK") or "" + assistant_id = pk.split("#", 1)[1] if "#" in pk else "" + if not assistant_id or not app_kb_id: + # A record the index returned but that cannot be addressed. Logged and + # skipped rather than raised: one malformed row must not starve the + # sweep, and it will still be there next tick to be noticed. + logger.error(f"skipping unaddressable KbWorkIndex row: PK={pk!r} appKbId={app_kb_id!r}") + counts["Failed"] += 1 + continue + + try: + _invoke_worker( + { + "assistantId": assistant_id, + "appKbId": app_kb_id, + "migrationState": record.get("migrationState"), + "migrationGeneration": int(record.get("migrationGeneration") or 0), + } + ) + counts["Dispatched"] += 1 + except Exception as exc: + logger.error( + f"failed to dispatch migration for kb {app_kb_id}: {exc}", exc_info=True + ) + counts["Failed"] += 1 + + _emit_metrics(counts) + return counts + + +def lambda_handler(event, context): + """EventBridge entry point. + + Nothing is read from ``event``. The dispatcher's behaviour is a function of the + index and the environment only — the same reasoning that fixed the reconciler's + arming bypass, where an invocation field could turn a report-only job into a + deleting one. + """ + counts = asyncio.run(dispatch_once()) + return {"statusCode": 200, "body": counts} diff --git a/backend/src/apis/app_api/kb_migration/ingestion_consumer.py b/backend/src/apis/app_api/kb_migration/ingestion_consumer.py new file mode 100644 index 000000000..776beb6e2 --- /dev/null +++ b/backend/src/apis/app_api/kb_migration/ingestion_consumer.py @@ -0,0 +1,361 @@ +"""Ingestion consumer for managed knowledge bases. + +Triggered by an EventBridge ``ObjectCreated`` event on the RAG documents bucket. Its +whole job is to decide whether a newly uploaded document belongs to a managed +knowledge base and, if so, ingest it directly. + +Routing exclusivity is the point (Requirements 10.3-10.5) +--------------------------------------------------------- +The legacy pipeline is triggered by its **own**, pre-existing S3 notification on the +same bucket. That notification was deliberately left in place, so for a legacy +document this function's correct behaviour is to **do nothing at all** — the other +Lambda has already got it. Acting here as well would index the same bytes twice: two +sets of vectors, doubled ingestion cost, and duplicate chunks competing in one +result list. + +So the routing table is asymmetric, and that asymmetry is intentional: + +=============== ========================================================= +Engine This function +=============== ========================================================= +legacy (absent) returns immediately, ingesting nothing +managed ingests directly and drives ``DOC#`` to a terminal state +=============== ========================================================= + +The one exception is a deliberate migration or dual-read pilot, which the +Migration_Worker drives and which never routes through an upload event. + +Indexed is not retrievable +-------------------------- +Bedrock reports a document ``INDEXED`` up to a second before it can actually be +retrieved — measured at 0.75-1.03 s. Marking a document ``complete`` on ``INDEXED`` +alone produces the worst kind of bug report: the UI says the upload worked, the user +asks a question straight away, and the answer does not mention their document. So +this polls until a retrieval really returns the document, and records ``indexedAt`` +and ``retrievableAt`` separately so the gap stays measurable instead of becoming +folklore. + +Import boundary +--------------- +Raw DynamoDB table access rather than importing ``apis.shared.assistants``, whose +``__init__`` pulls in the embeddings stack at module scope. Keeping this Lambda's +image small is a deliberate constraint — the same reason +``apis/app_api/kb_sync/records.py`` is written this way. Module-level imports are +stdlib only; everything heavy is function-local. + +No in-process orchestration +--------------------------- +No ``asyncio.ensure_future`` fan-out (Requirement 10.8). One invocation drives its +documents to terminal or fails and lets the event source redeliver. A background +task in a Lambda is killed when the handler returns, which turns a reported success +into a silently half-finished ingestion. +""" + +from __future__ import annotations + +import logging +import os +import time +from typing import Any, Dict, List, Optional, Tuple +from urllib.parse import unquote_plus + +logger = logging.getLogger() +logger.setLevel(logging.INFO) + +#: Terminal document states, taken from ``apis/app_api/documents/models.py``'s +#: ``DocumentStatus`` rather than invented: the facade's status filter serves only +#: ``complete``, so any drift here would silently make documents unretrievable. +STATUS_COMPLETE = "complete" +STATUS_FAILED = "failed" + +#: How long to wait for a document to become genuinely retrievable after Bedrock +#: reports it INDEXED. The observed gap is 0.75-1.03 s; the margin is wide because +#: the cost of waiting is a few seconds of Lambda time and the cost of not waiting +#: is telling a user their upload worked when it is not yet usable. +RETRIEVABLE_POLL_TIMEOUT_SECONDS = 30.0 +RETRIEVABLE_POLL_INTERVAL_SECONDS = 0.5 + +#: Bounded retries on the record update. The event source already redelivers, so +#: this only covers a transient DynamoDB failure inside one invocation; unbounded +#: retries would burn the Lambda timeout and lose the DLQ signal. +MAX_RECORD_UPDATE_ATTEMPTS = 3 + + +class IngestionRoutingError(Exception): + """The event could not be routed, or the document could not be finished.""" + + +def _table(): + import boto3 + + return boto3.resource("dynamodb").Table(os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"]) + + +def _now_iso() -> str: + from apis.shared.timestamps import utc_now_iso + + return utc_now_iso() + + +def parse_object_key(key: str) -> Tuple[str, str, str]: + """Split ``assistants/{assistant_id}/documents/{document_id}/{filename}``. + + The layout is the existing one and this feature does not change it: migration is + a re-ingest of bytes already in place, never a re-upload. Parsing the key rather + than trusting an event field keeps routing independent of which producer + delivered the notification. + """ + parts = unquote_plus(key).split("/") + if len(parts) < 5 or parts[0] != "assistants" or parts[2] != "documents": + raise IngestionRoutingError( + f"object key {key!r} is not an assistant document path; expected " + f"assistants/{{assistant_id}}/documents/{{document_id}}/{{filename}}" + ) + return parts[1], parts[3], "/".join(parts[4:]) + + +def extract_records(event: Dict[str, Any]) -> List[Dict[str, str]]: + """Normalize EventBridge and raw-S3 notification shapes into one list. + + Both are accepted because the bucket carries both producers — EventBridge feeds + this function, a direct notification feeds the legacy pipeline — so a wiring + change cannot silently stop ingestion. + """ + detail = event.get("detail") + if isinstance(detail, dict) and detail.get("object"): + return [ + { + "bucket": (detail.get("bucket") or {}).get("name", ""), + "key": (detail.get("object") or {}).get("key", ""), + } + ] + + out: List[Dict[str, str]] = [] + for record in event.get("Records") or []: + s3 = record.get("s3", {}) + out.append( + { + "bucket": (s3.get("bucket") or {}).get("name", ""), + "key": (s3.get("object") or {}).get("key", ""), + } + ) + return out + + +def resolve_engine_for(assistant_id: str) -> Tuple[str, Optional[Dict[str, Any]]]: + """The engine serving this assistant's knowledge base, plus its record. + + Delegates to ``records.resolve_engine`` so "absence means legacy" has exactly one + implementation. A missing record is the overwhelmingly common case today and + resolves to legacy, which is why this cannot treat it as an error. + """ + from apis.shared.kb_backend.records import get_kb_record, resolve_engine + + record = get_kb_record(assistant_id, assistant_id) + return resolve_engine(record), record + + +def set_document_terminal( + assistant_id: str, + document_id: str, + status: str, + indexed_at: Optional[str] = None, + retrievable_at: Optional[str] = None, + error: Optional[str] = None, +) -> None: + """Drive the ``DOC#`` record to a terminal state, with bounded retries. + + ``indexedAt`` and ``retrievableAt`` are stored separately on purpose: collapsing + them would erase the only evidence of the INDEXED-to-retrievable gap, which is + what makes "my upload finished but the assistant cannot see it" diagnosable + rather than mysterious. + """ + from botocore.exceptions import ClientError + + sets = ["#status = :status", "updatedAt = :now"] + values: Dict[str, Any] = {":status": status, ":now": _now_iso()} + + if indexed_at: + sets.append("indexedAt = :indexed") + values[":indexed"] = indexed_at + if retrievable_at: + sets.append("retrievableAt = :retrievable") + values[":retrievable"] = retrievable_at + if error: + sets.append("ingestionError = :err") + values[":err"] = error + + expression = f"SET {', '.join(sets)}" + last: Optional[Exception] = None + + for attempt in range(1, MAX_RECORD_UPDATE_ATTEMPTS + 1): + try: + _table().update_item( + Key={"PK": f"AST#{assistant_id}", "SK": f"DOC#{document_id}"}, + UpdateExpression=expression, + # `status` is a DynamoDB reserved keyword. + ExpressionAttributeNames={"#status": "status"}, + ExpressionAttributeValues=values, + ) + return + except ClientError as exc: + last = exc + logger.warning( + f"attempt {attempt}/{MAX_RECORD_UPDATE_ATTEMPTS} to mark " + f"{document_id} {status} failed: {exc}" + ) + if attempt < MAX_RECORD_UPDATE_ATTEMPTS: + time.sleep(0.2 * attempt) + + # Raised rather than swallowed: the record is the durable retry anchor + # (Requirement 10.7), so a document left non-terminal must surface as a failed + # invocation and reach the DLQ instead of looking like a success. + raise IngestionRoutingError( + f"could not mark document {document_id} as {status} after " + f"{MAX_RECORD_UPDATE_ATTEMPTS} attempts: {last}" + ) + + +def wait_until_retrievable( + backend: Any, + kb_ref: str, + document_id: str, + timeout_seconds: Optional[float] = None, + interval_seconds: Optional[float] = None, + sleep: Any = time.sleep, +) -> Optional[str]: + """Poll until a retrieval actually returns ``document_id``. + + Returns the timestamp at which it first became retrievable, or ``None`` on + timeout. A probe that itself errors is treated as "not yet", not as a document + failure: the document is usually fine and merely slow, and failing it would fail + uploads that are about to work. + + The timeouts default to ``None`` and are resolved from the module constants *at + call time*, rather than being bound as default arguments. Default arguments are + evaluated once at import, which makes them unpatchable — the first version of + this function bound them directly and a test that shortened the window had no + effect at all, silently waiting the full production timeout instead. + """ + import asyncio + + if timeout_seconds is None: + timeout_seconds = RETRIEVABLE_POLL_TIMEOUT_SECONDS + if interval_seconds is None: + interval_seconds = RETRIEVABLE_POLL_INTERVAL_SECONDS + + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + try: + chunks = asyncio.run(backend.search(kb_ref, document_id, 5)) + except Exception as exc: # noqa: BLE001 - a probe failure is not a document failure + logger.warning(f"retrievability probe for {document_id} failed: {exc}") + chunks = [] + + for chunk in chunks or []: + metadata = getattr(chunk, "metadata", None) or {} + if metadata.get("document_id") == document_id: + return _now_iso() + + sleep(interval_seconds) + + logger.warning( + f"document {document_id} was not retrievable within {timeout_seconds}s; " + f"leaving it short of complete rather than claiming success" + ) + return None + + +def handle_object(bucket: str, key: str) -> Dict[str, Any]: + """Route one uploaded object. Returns a summary for logging and tests.""" + from apis.shared.kb_backend.records import ENGINE_MANAGED + + assistant_id, document_id, filename = parse_object_key(key) + engine, record = resolve_engine_for(assistant_id) + + if engine != ENGINE_MANAGED: + # The legacy pipeline's own S3 notification already owns this document. + # Anything done here would index the same bytes a second time. + logger.info( + f"document {document_id} belongs to a legacy knowledge base; " + f"leaving it to the existing pipeline" + ) + return {"routed": "legacy", "ingested": False, "document_id": document_id} + + aws_kb_id = (record or {}).get("awsKbId") + data_source_id = (record or {}).get("awsDataSourceId") + if not aws_kb_id or not data_source_id: + # Managed engine but no identifiers means provisioning has not finished. + # Failing loudly is correct: silently falling back to legacy would create + # exactly the dual-index this function exists to prevent. + raise IngestionRoutingError( + f"assistant {assistant_id} resolves to the managed engine but its " + f"knowledge base is not provisioned (awsKbId={aws_kb_id!r}, " + f"awsDataSourceId={data_source_id!r})" + ) + + import asyncio + + from apis.shared.kb_backend.managed_backend import ManagedKbBackend + from apis.shared.kb_backend.protocol import DocumentSource + + # The backend takes the App_KB_Id and resolves the AWS identifiers itself on + # every operation. Threading them in from here would defeat that: a + # dormancy/rehydration cycle replaces them, and a caller holding a stale pair + # would keep addressing a knowledge base that no longer exists. The check above + # is still worth doing - it fails fast with a precise reason - but it is a + # precondition, not a value to pass along. + backend = ManagedKbBackend(bucket=bucket) + source = DocumentSource(document_id=document_id, filename=filename, s3_key=key) + + try: + asyncio.run(backend.ingest(assistant_id, source)) + except Exception as exc: + logger.error(f"direct ingestion of {document_id} failed: {exc}", exc_info=True) + set_document_terminal(assistant_id, document_id, STATUS_FAILED, error=str(exc)) + raise + + indexed_at = _now_iso() + retrievable_at = wait_until_retrievable(backend, assistant_id, document_id) + + if retrievable_at is None: + # Ingested but not confirmed retrievable. Left non-terminal deliberately so + # the event source redelivers, rather than the record claiming a success the + # user cannot yet observe. + raise IngestionRoutingError( + f"document {document_id} was ingested but not retrievable within the " + f"poll window; leaving it for redelivery" + ) + + set_document_terminal( + assistant_id, + document_id, + STATUS_COMPLETE, + indexed_at=indexed_at, + retrievable_at=retrievable_at, + ) + return { + "routed": "managed", + "ingested": True, + "document_id": document_id, + "indexedAt": indexed_at, + "retrievableAt": retrievable_at, + } + + +def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]: + """Entry point. One invocation drives its documents to terminal, or fails.""" + records = extract_records(event) + if not records: + logger.info("no S3 records in event; nothing to do") + return {"statusCode": 200, "processed": 0, "results": []} + + results = [] + for record in records: + bucket, key = record.get("bucket", ""), record.get("key", "") + if not bucket or not key: + logger.warning(f"skipping record with missing bucket or key: {record}") + continue + results.append(handle_object(bucket, key)) + + return {"statusCode": 200, "processed": len(results), "results": results} diff --git a/backend/src/apis/app_api/kb_migration/reconciler.py b/backend/src/apis/app_api/kb_migration/reconciler.py new file mode 100644 index 000000000..4b79bd5fc --- /dev/null +++ b/backend/src/apis/app_api/kb_migration/reconciler.py @@ -0,0 +1,844 @@ +"""Daily reconciler for managed knowledge bases. + +Joins a paginated, tag-filtered ``ListKnowledgeBases`` against the KB_Records and +acts on the three ways the two sides can disagree. It exists because a managed +knowledge base is a **runtime-created, billed** resource with no CloudFormation +parent: nothing else in the system would ever notice one that our database has +forgotten about. + +The join table +-------------- +============= ============================================================ +Side Action +============= ============================================================ +AWS only Orphan. Delete **only if AWS's own ``createdAt`` is >24 h old** +Record only Mark ``vectorState: missing``. **Never delete the record** +Both Refresh ``storedBytes`` for quota accounting +============= ============================================================ + +Two of those three rows are counter-intuitive, and each is the way it is because +the intuitive version destroys something. + +**Age-gate on AWS's ``createdAt``, never on discovery time.** The tempting +implementation records when the reconciler first *saw* an unknown knowledge base +and waits 24 hours from there. That is wrong in both directions. A reconciler that +was down for a week comes back and treats every knowledge base in the account as +newly discovered — so either it waits another 24 hours on genuine week-old +orphans, or, if the comparison is written the other way round, it deletes every +knowledge base that is mid-provisioning right now, including creates that are 40 +seconds old and about to succeed. AWS's ``createdAt`` is a fact about the +resource, is identical on every run, and does not depend on this process's uptime. +It is obtained from ``GetKnowledgeBase``, because ``KnowledgeBaseSummary`` does +not carry it. + +**A record with no AWS knowledge base is a stale pointer, not a dead corpus.** It +means the *vectors* are gone. The source bytes are still in S3 and the ``DOC#`` +records are still valid and still ``complete``, so the knowledge base can be +rebuilt from them on the next ingest, and the owner never has to re-upload +anything. The record is the only pointer to that recoverable corpus, so deleting +it is the one action here that loses user data — which is why +:func:`mark_vector_state_missing` is the entire response and no code path in this +module removes a KB_Record. + +Report-only, and armed separately +--------------------------------- +This ships **disarmed** (Requirement 14.7, 19.7). It logs exactly what it would +have deleted and deletes nothing, and it runs that way for weeks so its judgement +can be checked against real data before it is trusted with a delete. Arming is one +flag, ``MANAGED_KB_RECONCILER_ARMED``, and an **empty string reads as off** +(Requirement 19.8) — an unset GitHub Actions variable expands to ``""``, which is +how a flag that is obviously off ends up looking truthy to ``if os.environ.get``. + +The per-run limit applies in **both** modes, so the report says what an armed run +would actually do. A report listing 500 intended deletions from a run that would +only ever perform 25 is a misleading artifact, and the whole point of the +report-only period is that the artifact can be trusted. + +Import boundary +--------------- +Module-level imports are stdlib plus the stdlib-only ``kb_backend`` modules; +``boto3`` is function-local, and nothing here reaches ``apis.shared.assistants``. +DynamoDB is accessed through the raw table resource, matching +``kb_migration/ingestion_consumer.py`` and ``kb_sync/records.py``. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from decimal import Decimal +from typing import Any, Callable, Dict, Iterator, List, Optional + +from apis.shared.kb_backend.metrics import emit_count, emit_fleet_gauges +from apis.shared.kb_backend.records import kb_pk, kb_sk + +logger = logging.getLogger() +logger.setLevel(logging.INFO) + +# ── Flags ──────────────────────────────────────────────────────────────────── +# +# The arming flag. Absent, empty, or anything not in the truthy set means the +# reconciler reports and deletes nothing. +FLAG_RECONCILER_ARMED = "MANAGED_KB_RECONCILER_ARMED" + +#: Recognised affirmative spellings. Everything else — including ``""``, ``"0"``, +#: ``"false"`` and ``"off"`` — is off. An allow-list rather than a truthiness test +#: because the failure being designed around is a value that is *present but +#: empty*: ``bool("")`` is correct by luck, ``bool("false")`` is not. +_TRUTHY = frozenset({"1", "true", "yes", "on", "enabled"}) + +# ── Tunables, resolved at call time ────────────────────────────────────────── +# +# Read inside the functions that use them rather than bound as default arguments. +# A default argument is evaluated once at import, so it cannot be patched and a +# test that overrides it silently gets the production value instead. + +#: Requirement 14.4. An orphan younger than this is very likely an in-flight +#: create: provisioning to ``ACTIVE`` was measured at 47-124 s, and the record is +#: written before the AWS call, so the only window in which a legitimate create +#: looks like an orphan is the moments between the two. 24 hours is far wider than +#: needed, which is the correct direction for a destructive action. +ORPHAN_MIN_AGE_HOURS = 24.0 + +#: Requirement 14.8. Bounds the destructive work of a single run, so a bug in the +#: join — or a tag filter that suddenly matches more than it should — costs at +#: most this many knowledge bases before someone sees the report. +MAX_DELETIONS_PER_RUN = 25 + +#: Hard ceiling on :func:`max_deletions_per_run`, above which the env var is +#: ignored. Deleting more than this in one pass is not an operation that should be +#: reachable by editing a variable; it should require repeated, observed runs. +MAX_DELETIONS_CEILING = 100 + +#: Bounds the join itself. A reconciler that walked an unbounded account would +#: time out mid-pass and produce a partial report indistinguishable from a +#: complete one. +MAX_KNOWLEDGE_BASES_PER_RUN = 2000 + +# ── Vector state ───────────────────────────────────────────────────────────── +# +# Written on a record whose AWS knowledge base has gone. Not a failure state: the +# corpus is intact and the next ingest re-provisions. +VECTOR_STATE_MISSING = "missing" + +# ── Metrics ────────────────────────────────────────────────────────────────── +METRIC_ORPHANS_FOUND = "KbOrphansFound" +METRIC_ORPHANS_DELETED = "KbOrphansDeleted" +METRIC_VECTORS_MISSING = "KbVectorsMissing" +METRIC_RECONCILER_LIMIT_REACHED = "KbReconcilerLimitReached" + + +@dataclass +class PlannedDeletion: + """An orphan the reconciler intends to delete, and why it is eligible.""" + + kb_id: str + name: str + status: str + created_at: Optional[str] + age_hours: Optional[float] + performed: bool = False + error: Optional[str] = None + + +@dataclass +class ReconcileReport: + """What one run found and what it did. + + ``armed`` is on the report rather than only in the logs so a stored artifact + is self-describing: an operator reading last night's output should not have to + go and check what the flag was set to at the time. + """ + + armed: bool = False + aws_knowledge_bases: int = 0 + records: int = 0 + matched: int = 0 + orphans: int = 0 + planned_deletions: List[PlannedDeletion] = field(default_factory=list) + skipped_too_young: List[str] = field(default_factory=list) + marked_missing: List[str] = field(default_factory=list) + refreshed_bytes: List[str] = field(default_factory=list) + limit_reached: bool = False + + #: Fleet gauges (Requirement 22.1), accumulated over the record side of the + #: join. Computed from each KB_Record as it was read, so a ``storedBytes`` + #: refresh performed later in the same pass lands in the *next* pass's gauge — + #: acceptable for a daily number, and cheaper than a second full scan. + stored_bytes: int = 0 + idle_bytes: int = 0 + #: Knowledge bases with no recorded activity at all: never retrieved and their + #: agent never used. Reported so ``KbIdleGB`` can be read honestly — these are + #: unmeasured, not idle, and counting them as idle would make every freshly + #: provisioned corpus look abandoned. + unmeasured_idleness: int = 0 + + @property + def deletions_performed(self) -> int: + return sum(1 for planned in self.planned_deletions if planned.performed) + + def to_dict(self) -> Dict[str, Any]: + return { + "armed": self.armed, + "mode": "armed" if self.armed else "report-only", + "awsKnowledgeBases": self.aws_knowledge_bases, + "records": self.records, + "matched": self.matched, + "orphans": self.orphans, + "plannedDeletions": [ + { + "knowledgeBaseId": planned.kb_id, + "name": planned.name, + "status": planned.status, + "createdAt": planned.created_at, + "ageHours": planned.age_hours, + "performed": planned.performed, + "error": planned.error, + } + for planned in self.planned_deletions + ], + "deletionsPerformed": self.deletions_performed, + "skippedTooYoung": self.skipped_too_young, + "markedMissing": self.marked_missing, + "refreshedBytes": self.refreshed_bytes, + "limitReached": self.limit_reached, + "storedBytes": self.stored_bytes, + "idleBytes": self.idle_bytes, + "unmeasuredIdleness": self.unmeasured_idleness, + } + + +# ── Flag and tunable readers ───────────────────────────────────────────────── +def reconciler_armed() -> bool: + """Whether the reconciler may delete. Defaults to **off**. + + An empty string is off (Requirement 19.8). This is not defensive + over-engineering: an unset repository or environment variable expands to the + empty string in GitHub Actions, and this repo has been bitten by that before — + a flag nobody set looking set, in the one component whose mistakes are + irreversible. + """ + raw = os.environ.get(FLAG_RECONCILER_ARMED) + if not raw: + return False + return raw.strip().lower() in _TRUTHY + + +def _env_float(name: str, default: float) -> float: + raw = os.environ.get(name) + if not raw: + return default + try: + return float(raw) + except ValueError: + logger.warning(f"{name}={raw!r} is not a number; falling back to {default}") + return default + + +def _env_int(name: str, default: int) -> int: + raw = os.environ.get(name) + if not raw: + return default + try: + return int(raw) + except ValueError: + logger.warning(f"{name}={raw!r} is not an integer; falling back to {default}") + return default + + +def orphan_min_age_hours() -> float: + return _env_float("MANAGED_KB_ORPHAN_MIN_AGE_HOURS", ORPHAN_MIN_AGE_HOURS) + + +def max_deletions_per_run() -> int: + """The per-run deletion bound, clamped so the environment cannot lift it. + + The env var may lower the limit but not raise it past + :data:`MAX_DELETIONS_CEILING` (Requirement 14.8). A bound that any environment + variable can set to a million is not a bound, and this is the one limit whose + failure mode is irreversible: it is what stops a single bad run — a wrong tag + filter, a botched migration — from deleting an account's worth of user + knowledge bases before anyone reads the report. + """ + requested = _env_int("MANAGED_KB_RECONCILER_MAX_DELETIONS", MAX_DELETIONS_PER_RUN) + if requested > MAX_DELETIONS_CEILING: + logger.warning( + f"MANAGED_KB_RECONCILER_MAX_DELETIONS={requested} exceeds the ceiling of " + f"{MAX_DELETIONS_CEILING}; clamping. Run the reconciler repeatedly rather " + f"than raising this." + ) + return MAX_DELETIONS_CEILING + return max(requested, 0) + + +def max_knowledge_bases_per_run() -> int: + return _env_int("MANAGED_KB_RECONCILER_MAX_SCANNED", MAX_KNOWLEDGE_BASES_PER_RUN) + + +# ── DynamoDB plumbing ──────────────────────────────────────────────────────── +def _table(): + import boto3 + + return boto3.resource("dynamodb").Table(os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"]) + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +def _now_iso() -> str: + from apis.shared.timestamps import utc_now_iso + + return utc_now_iso() + + +def iter_kb_records() -> Iterator[Dict[str, Any]]: + """Every KB_Record in the table, paging the scan to exhaustion. + + A scan, because the ``KbWorkIndex`` GSI is *sparse* and deliberately holds + only records that are eligible for migration work — the records this join + cares most about are precisely the ones absent from it. Paged to exhaustion + for the same reason the AWS list is: a truncated read makes every unread + record look like an orphan on the AWS side. + + ``KBTOMB#`` sort keys do not match ``begins_with(SK, "KB#")``, so tombstones + are excluded by the key prefix rather than filtered afterwards. + """ + from boto3.dynamodb.conditions import Attr + + table = _table() + kwargs: Dict[str, Any] = {"FilterExpression": Attr("SK").begins_with("KB#")} + while True: + response = table.scan(**kwargs) + for item in response.get("Items") or []: + yield item + start = response.get("LastEvaluatedKey") + if not start: + return + kwargs["ExclusiveStartKey"] = start + + +# ── Age gate (Requirement 14.3, 14.4) ──────────────────────────────────────── +def parse_aws_timestamp(value: Any) -> Optional[datetime]: + """Coerce AWS's ``createdAt`` to an aware UTC datetime, or ``None``. + + boto3 hands back a ``datetime`` here, but a value that has been through a + stubbed client, an EventBridge payload or a JSON round-trip arrives as a + string or an epoch number. All three are accepted; anything unparseable + returns ``None``, which the age gate treats as *not old enough* rather than + guessing. + """ + if value is None: + return None + if isinstance(value, datetime): + return value if value.tzinfo else value.replace(tzinfo=timezone.utc) + if isinstance(value, (int, float, Decimal)): + try: + return datetime.fromtimestamp(float(value), tz=timezone.utc) + except (OverflowError, OSError, ValueError): + return None + if isinstance(value, str): + from apis.shared.timestamps import from_iso + + try: + return from_iso(value) + except ValueError: + return None + return None + + +def orphan_age_hours(created_at: Any, now: Optional[datetime] = None) -> Optional[float]: + """Hours since **AWS's** ``createdAt``, or ``None`` if it cannot be read.""" + created = parse_aws_timestamp(created_at) + if created is None: + return None + return ((now or _now()) - created).total_seconds() / 3600.0 + + +def orphan_is_deletable( + created_at: Any, + now: Optional[datetime] = None, + min_age_hours: Optional[float] = None, +) -> bool: + """Whether an orphan has existed in AWS long enough to be deleted. + + The input is AWS's ``createdAt``. It is deliberately not "when did we first + see this": see the module docstring. Passing a discovery timestamp here would + type-check, run, pass a naive test, and delete in-flight creates in + production. + + A missing or unparseable ``createdAt`` returns ``False``. Failing closed is + the only safe direction for a destructive action: an orphan left one more day + costs pennies, and a knowledge base deleted 40 seconds into its creation costs + a user their upload. + """ + if min_age_hours is None: + min_age_hours = orphan_min_age_hours() + + created = parse_aws_timestamp(created_at) + if created is None: + return False + return (now or _now()) - created > timedelta(hours=min_age_hours) + + +# ── Record-side actions ────────────────────────────────────────────────────── +def mark_vector_state_missing(assistant_id: str, app_kb_id: str) -> None: + """Record that the AWS knowledge base behind this record has gone. + + **This never deletes the record**, and there is deliberately no function in + this module that does. The vectors are gone; the corpus is not. The uploaded + bytes are still in S3 and the ``DOC#`` records still describe them, so the + next ingest re-provisions a knowledge base and re-indexes from the documents + already present. The record carries the only mapping from ``App_KB_Id`` to + that corpus, so removing it would turn a recoverable, invisible-to-the-user + situation into permanent data loss. + + ``awsKbId``/``awsDataSourceId`` are left in place rather than cleared: they + are the evidence of which AWS resource vanished, and provisioning already + treats a record it cannot find in AWS as needing a fresh create. + """ + _table().update_item( + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression=( + "SET vectorState = :missing, vectorStateObservedAt = :now, updatedAt = :now" + ), + ExpressionAttributeValues={":missing": VECTOR_STATE_MISSING, ":now": _now_iso()}, + ) + emit_count(METRIC_VECTORS_MISSING) + + +def refresh_stored_bytes(assistant_id: str, app_kb_id: str, stored_bytes: int) -> None: + """Re-anchor quota accounting, and clear any stale ``vectorState``. + + The ``REMOVE`` matters: a record marked ``missing`` on an earlier run that has + since been re-provisioned would otherwise stay marked for ever, and the UI + would keep telling its owner their knowledge base is broken after it was + fixed. + """ + _table().update_item( + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression=( + "SET storedBytes = :bytes, updatedAt = :now " + "REMOVE vectorState, vectorStateObservedAt" + ), + ExpressionAttributeValues={":bytes": Decimal(int(stored_bytes)), ":now": _now_iso()}, + ) + + +def stored_bytes_from_s3(assistant_id: str, bucket: Optional[str] = None, s3_client=None) -> Optional[int]: + """Total size of an assistant's uploaded documents, straight from S3. + + S3 rather than a client-reported or previously-stored value, for the same + reason the byte cap uses a ``HEAD``: this number gates a $150,000/month + exposure at full adoption, and the only trustworthy source for it is the + service holding the bytes. + + Returns ``None`` when no bucket is configured or the listing fails, and the + caller then leaves ``storedBytes`` alone. Writing a zero on a failed listing + would silently hand every owner their whole allowance back. + """ + bucket = bucket or os.environ.get("S3_ASSISTANTS_DOCUMENTS_BUCKET_NAME") + if not bucket: + return None + + if s3_client is None: + import boto3 + + s3_client = boto3.client("s3") + + prefix = f"assistants/{assistant_id}/documents/" + total = 0 + token: Optional[str] = None + try: + while True: + kwargs: Dict[str, Any] = {"Bucket": bucket, "Prefix": prefix} + if token: + kwargs["ContinuationToken"] = token + response = s3_client.list_objects_v2(**kwargs) + for obj in response.get("Contents") or []: + total += int(obj.get("Size") or 0) + if not response.get("IsTruncated"): + return total + token = response.get("NextContinuationToken") + if not token: + return total + except Exception as exc: # noqa: BLE001 - a failed listing must not zero the quota + logger.warning(f"could not total stored bytes for {assistant_id}: {exc}") + return None + + +# ── The run ────────────────────────────────────────────────────────────────── +def reconcile( + client=None, + project_prefix: Optional[str] = None, + environment: Optional[str] = None, + armed: Optional[bool] = None, + now: Optional[datetime] = None, + stored_bytes_resolver: Optional[Callable[[str], Optional[int]]] = None, +) -> ReconcileReport: + """One reconciliation pass. + + ``armed`` defaults to :func:`reconciler_armed`, i.e. to the flag, i.e. to off. + It is an argument only so a test can exercise the armed path without mutating + process environment — never so a caller can conveniently turn deletion on. + """ + from apis.shared.kb_backend import tombstones as tb + + if client is None: + from apis.shared.kb_backend.managed_backend import bedrock_agent_client + + client = bedrock_agent_client() + + if armed is None: + armed = reconciler_armed() + if now is None: + now = _now() + + report = ReconcileReport(armed=armed) + scan_limit = max_knowledge_bases_per_run() + delete_limit = max_deletions_per_run() + min_age = orphan_min_age_hours() + + # ── record side ────────────────────────────────────────────────────────── + # Keyed by awsKbId, because that is the identifier the AWS side reports. A + # record with no awsKbId has not finished provisioning and is not evidence of + # anything: it is skipped rather than counted as a missing-vector record, + # which would mark every in-flight create broken. + records_by_aws_id: Dict[str, Dict[str, Any]] = {} + unprovisioned = 0 + for item in iter_kb_records(): + report.records += 1 + _accumulate_gauges(item, report, now) + aws_kb_id = item.get("awsKbId") + if aws_kb_id: + records_by_aws_id[str(aws_kb_id)] = item + else: + unprovisioned += 1 + + # ── AWS side ───────────────────────────────────────────────────────────── + seen_aws_ids: set = set() + orphan_facts: List[tb.KnowledgeBaseFacts] = [] + + for facts in tb.iter_project_knowledge_bases( + client, project_prefix=project_prefix, environment=environment + ): + if report.aws_knowledge_bases >= scan_limit: + report.limit_reached = True + logger.warning( + f"stopping the AWS walk at {scan_limit} knowledge bases; this run's " + f"join is partial and no deletion decision is made beyond this point" + ) + break + + report.aws_knowledge_bases += 1 + seen_aws_ids.add(facts.kb_id) + + record = records_by_aws_id.get(facts.kb_id) + if record is None: + orphan_facts.append(facts) + else: + report.matched += 1 + _reconcile_matched(record, stored_bytes_resolver, report) + + # ── record only: mark missing, never delete (Requirement 14.5) ──────────── + # + # Only meaningful when the AWS walk completed. On a truncated walk an unmatched + # record may simply be one this run never reached. + if not report.limit_reached: + for aws_kb_id, record in sorted(records_by_aws_id.items()): + if aws_kb_id in seen_aws_ids: + continue + app_kb_id = str(record.get("appKbId") or "") + assistant_id = _assistant_id_of(record) + if not app_kb_id or not assistant_id: + logger.warning(f"skipping malformed KB_Record {record.get('PK')}/{record.get('SK')}") + continue + logger.info( + f"KB_Record {app_kb_id} points at awsKbId {aws_kb_id}, which AWS does " + f"not have. Marking vectorState={VECTOR_STATE_MISSING}. The record is " + f"NOT deleted: its documents are still valid and the knowledge base " + f"rebuilds from them on the next ingest." + ) + mark_vector_state_missing(assistant_id, app_kb_id) + report.marked_missing.append(app_kb_id) + + # ── AWS only: orphans (Requirements 14.2, 14.3, 14.4) ──────────────────── + report.orphans = len(orphan_facts) + if report.orphans: + emit_count(METRIC_ORPHANS_FOUND, value=report.orphans) + + for facts in orphan_facts: + age = orphan_age_hours(facts.created_at, now=now) + # The gate reads AWS's createdAt. Never the time of discovery. + if not orphan_is_deletable(facts.created_at, now=now, min_age_hours=min_age): + report.skipped_too_young.append(facts.kb_id) + logger.info( + f"orphan {facts.kb_id} ({facts.name}) has no KB_Record but AWS " + f"reports createdAt={facts.created_at!r} " + f"(age={age if age is None else round(age, 2)}h < {min_age}h); " + f"leaving it alone — it is most likely an in-flight create" + ) + continue + + planned = PlannedDeletion( + kb_id=facts.kb_id, + name=facts.name, + status=facts.status, + created_at=str(facts.created_at) if facts.created_at is not None else None, + age_hours=None if age is None else round(age, 2), + ) + + # The limit applies whether or not we are armed, so the report describes + # what an armed run would really do. + if len(report.planned_deletions) >= delete_limit: + report.limit_reached = True + emit_count(METRIC_RECONCILER_LIMIT_REACHED) + logger.warning( + f"per-run deletion limit of {delete_limit} reached; {facts.kb_id} and " + f"any further orphans are left for the next run" + ) + break + + report.planned_deletions.append(planned) + + if facts.status == tb.KB_STATUS_DELETE_UNSUCCESSFUL: + # Requirement 13.7 / the tombstone table's fourth row. Retrying the + # delete does not help and the state does not clear on its own, so it + # is surfaced as an operator state instead of being counted as work. + planned.error = tb.KB_STATUS_DELETE_UNSUCCESSFUL + emit_count(tb.METRIC_DELETE_UNSUCCESSFUL) + logger.error( + f"orphan {facts.kb_id} is in {tb.KB_STATUS_DELETE_UNSUCCESSFUL} and " + f"needs operator action; it will not delete by retrying and it is " + f"still being billed" + ) + continue + + if not armed: + # Report-only. This is the shipped mode and it performs no deletes. + logger.warning( + f"[report-only] WOULD delete orphan knowledge base {facts.kb_id} " + f"({facts.name}), AWS createdAt={facts.created_at!r}, " + f"age={planned.age_hours}h. Set {FLAG_RECONCILER_ARMED} to arm." + ) + continue + + try: + _delete_orphan(facts, client) + planned.performed = True + emit_count(METRIC_ORPHANS_DELETED) + logger.info(f"deleted orphan knowledge base {facts.kb_id}") + except Exception as exc: # noqa: BLE001 - one bad orphan must not end the run + planned.error = str(exc) + logger.error(f"failed to delete orphan {facts.kb_id}: {exc}", exc_info=True) + + emit_fleet_gauges( + kb_count=report.records, + stored_bytes=report.stored_bytes, + idle_bytes=report.idle_bytes, + unmeasured=report.unmeasured_idleness, + ) + + logger.info( + f"reconcile complete: mode={'armed' if armed else 'report-only'} " + f"aws={report.aws_knowledge_bases} records={report.records} " + f"matched={report.matched} orphans={report.orphans} " + f"planned={len(report.planned_deletions)} " + f"performed={report.deletions_performed} " + f"markedMissing={len(report.marked_missing)} " + f"unprovisioned={unprovisioned} " + f"storedGB={report.stored_bytes / 1_000_000_000:.3f} " + f"idleGB={report.idle_bytes / 1_000_000_000:.3f} " + f"unmeasured={report.unmeasured_idleness}" + ) + return report + + +def _accumulate_gauges( + record: Dict[str, Any], + report: ReconcileReport, + now: Optional[datetime] = None, +) -> None: + """Fold one KB_Record into the fleet gauges (Requirements 22.1, 22.5). + + The reconciler is where this belongs because it is already the one pass that + walks every knowledge base; a second sweep to count them would double a scan + that exists. + + Idleness comes from :func:`idleness.idle_days`, which takes the **maximum** of + the knowledge base's own ``lastRetrievedAt`` and its bound agents' + ``lastUsedAt``. Never retrieval alone: an agent can be invoked all day and + retrieve nothing, because retrieval only fires when the query matches, so a + corpus judged by retrieval alone looks abandoned exactly when its agent is + busiest with questions the documents do not answer. + + A record with no activity signal at all counts toward ``unmeasured_idleness`` + and **not** toward idle bytes. It is unmeasured, not idle — that is what a + knowledge base provisioned an hour ago looks like. + """ + from apis.shared.kb_backend.idleness import idle_days + + stored = int(record.get("storedBytes") or 0) + report.stored_bytes += stored + + assistant_id = _assistant_id_of(record) + if not assistant_id: + return + + try: + days = idle_days(assistant_id, record, now=_iso_or_none(now)) + except Exception as exc: # noqa: BLE001 - a gauge must not end the pass + logger.warning(f"could not compute idleness for {assistant_id}: {exc}") + return + + if days is None: + report.unmeasured_idleness += 1 + elif days >= idle_threshold_days(): + report.idle_bytes += stored + + +def _iso_or_none(moment: Optional[datetime]) -> Optional[str]: + return moment.strftime("%Y-%m-%dT%H:%M:%SZ") if moment else None + + +def idle_threshold_days() -> int: + """Days without a sign of life before bytes count as idle. + + A **reporting** threshold. Nothing reclaims in this phase, and the number the + follow-up spec eventually evicts on should come from the distribution this + metric records rather than being inherited from this default. + """ + from apis.shared.kb_backend.metrics import IDLE_THRESHOLD_DAYS + + return _env_int("KB_IDLE_THRESHOLD_DAYS", IDLE_THRESHOLD_DAYS) + + +def _assistant_id_of(record: Dict[str, Any]) -> str: + """Recover the assistant id from the record's ``PK``.""" + pk = str(record.get("PK") or "") + return pk[len("AST#") :] if pk.startswith("AST#") else "" + + +def _reconcile_matched( + record: Dict[str, Any], + stored_bytes_resolver: Optional[Callable[[str], Optional[int]]], + report: ReconcileReport, +) -> None: + """Both sides agree: refresh stored bytes (Requirement 14.6). + + Written only when the number actually changed, or when a stale + ``vectorState`` needs clearing. A daily no-op write per knowledge base would + be pure cost and would churn ``updatedAt`` on records nothing happened to. + """ + app_kb_id = str(record.get("appKbId") or "") + assistant_id = _assistant_id_of(record) + if not app_kb_id or not assistant_id: + return + + resolver = stored_bytes_resolver or stored_bytes_from_s3 + actual = resolver(assistant_id) + if actual is None: + return + + current = int(record.get("storedBytes") or 0) + stale_state = record.get("vectorState") is not None + if actual == current and not stale_state: + return + + refresh_stored_bytes(assistant_id, app_kb_id, actual) + report.refreshed_bytes.append(app_kb_id) + + +def _delete_orphan(facts, client) -> None: + """Delete an orphan through the tombstoned saga. + + Through the saga rather than a bare ``DeleteKnowledgeBase`` because an orphan + is by definition a resource a previous delete failed to remove, so the one + thing it must not do is fail silently a second time. The saga writes the + tombstone first, polls until AWS reports the knowledge base absent, and clears + the tombstone only then. + + An orphan has no KB_Record — that is what makes it an orphan — so there is no + assistant id to anchor its tombstone on. It is anchored on the ``appKbId`` tag + the provisioner wrote, falling back to the AWS identifier, and the item is + flagged ``syntheticPartition`` so nobody reads that ``PK`` as a real assistant + and nobody expects ``iter_tombstones()`` to surface it. + ``anchorSource`` records which of the two identifiers was available, because + when this item is being triaged that is the first question. + + ``remove_record`` is left false: there is no record to remove, and this module + never removes one. + """ + from apis.shared.kb_backend import tombstones as tb + from apis.shared.kb_backend.tags import TAG_KEY_APP_KB_ID + + tags = facts.tags or {} + # The AWS tag, whose key is owned by `kb_backend.tags` — not the KB_Record + # attribute, which is separately named `appKbId` and stays that way. + tagged = tags.get(TAG_KEY_APP_KB_ID) + anchor = tagged or facts.kb_id + tb.delete_knowledge_base( + anchor, + anchor, + facts.kb_id, + client=client, + remove_record=False, + extra_attributes={ + tb.SYNTHETIC_PARTITION: True, + "anchorSource": f"tag:{TAG_KEY_APP_KB_ID}" if tagged else "aws:knowledgeBaseId", + "orphanKbName": facts.name or None, + }, + ) + + +def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]: + """Scheduled entry point. Returns the report so it lands in the invocation log. + + The invocation event is deliberately **not** consulted for arming. The + environment flag is the only way to arm (Requirement 19.7), because an event + payload is the one input an operator does not review: an EventBridge target + carrying a constant ``{"armed": true}``, or any principal holding + ``lambda:InvokeFunction``, would delete billed user resources while every + piece of reviewable configuration still said report-only, leaving nothing + behind but an ``Invoke`` in CloudTrail. If the event disagrees with the flag, + the flag wins, and the disagreement is logged rather than honoured. + """ + requested = (event or {}).get("armed") + if requested is not None: + logger.warning( + f"ignoring armed={requested!r} from the invocation event: arming is " + f"controlled only by {FLAG_RECONCILER_ARMED}" + ) + report = reconcile() + return {"statusCode": 200, "report": report.to_dict()} + + +__all__ = [ + "FLAG_RECONCILER_ARMED", + "MAX_DELETIONS_CEILING", + "MAX_DELETIONS_PER_RUN", + "MAX_KNOWLEDGE_BASES_PER_RUN", + "METRIC_ORPHANS_DELETED", + "METRIC_ORPHANS_FOUND", + "METRIC_RECONCILER_LIMIT_REACHED", + "METRIC_VECTORS_MISSING", + "ORPHAN_MIN_AGE_HOURS", + "VECTOR_STATE_MISSING", + "PlannedDeletion", + "ReconcileReport", + "iter_kb_records", + "lambda_handler", + "mark_vector_state_missing", + "max_deletions_per_run", + "max_knowledge_bases_per_run", + "orphan_age_hours", + "orphan_is_deletable", + "orphan_min_age_hours", + "parse_aws_timestamp", + "reconcile", + "reconciler_armed", + "refresh_stored_bytes", + "stored_bytes_from_s3", +] diff --git a/backend/src/apis/app_api/kb_migration/requirements.txt b/backend/src/apis/app_api/kb_migration/requirements.txt new file mode 100644 index 000000000..0f1a5cf32 --- /dev/null +++ b/backend/src/apis/app_api/kb_migration/requirements.txt @@ -0,0 +1,32 @@ +# kb-migration Lambda image dependencies (backend/Dockerfile.kb-migration). +# Exact pins per repo policy; versions match backend/uv.lock. +# +# ⚠️ THE boto3 PIN IS LOAD-BEARING, NOT HYGIENE. +# +# The Lambda Python base image bundles its own, older boto3. That bundled +# version's packaged service model has no `MANAGED` member in +# `CreateKnowledgeBase`'s `knowledgeBaseConfiguration.type` enum and no +# `managedKnowledgeBaseConfiguration` shape at all, so every provisioning call +# would fail with a ParamValidationError naming a parameter that looks correct +# in our source. Installing this pin *over* the bundled copy is what makes the +# managed knowledge base API reachable from a Lambda. +# +# Verified against the pinned version rather than assumed (spec task 15.1's +# static half): `type` enum is ['VECTOR','KENDRA','SQL','MANAGED'], +# `managedKnowledgeBaseConfiguration` carries the embedding pin and encryption +# members, `embeddingDataType` enum is ['FLOAT32','BINARY'], and all four +# document operations — Ingest/Get/List/DeleteKnowledgeBaseDocuments — exist. +# No `AWS_DATA_PATH` side-load is required, and none must be relied on. +boto3==1.43.68 +botocore==1.43.68 + +# Nothing else. The four handlers' whole import closure is 16 first-party +# modules plus these two, which is the point of `kb_backend` being its own +# package with an empty `__init__` and stdlib-only module scope — pulling +# `apis.shared.assistants` instead would drag the embeddings stack, and with it +# the image-size budget, into a Lambda that never needs either. +# Enforced by backend/tests/architecture/test_kb_backend_boundary.py. +# +# In particular this image has NO FastAPI and NO pydantic: KB_Record is a +# dataclass precisely so a size-constrained image need not carry validation +# machinery it would then have to justify. diff --git a/backend/src/apis/app_api/kb_migration/worker.py b/backend/src/apis/app_api/kb_migration/worker.py new file mode 100644 index 000000000..89165eea0 --- /dev/null +++ b/backend/src/apis/app_api/kb_migration/worker.py @@ -0,0 +1,945 @@ +"""Migration worker: shadow → verify → promote → retain, one step per invocation. + +Requirements 15, 16, 17. Each invocation takes a lease, executes **one** step for +one knowledge base, records the next state with a conditional write, and returns. +The dispatcher brings it back for the next step. + +Why one step per invocation +--------------------------- +A 20-document text corpus is about 3 minutes end to end, but a 20-PDF corpus can +exceed an hour: per-document parse time was measured at 37–264 s and dominates +everything else. A worker that tried to run the whole machine in one invocation +would therefore be a Lambda that sometimes finishes in three minutes and sometimes +hits its timeout — and a timeout mid-``shadow`` is indistinguishable, from the +outside, from a crash. Stepping means every interruption lands on a recorded state +with a conditional guard in front of it, which is what makes a resumed run converge +instead of duplicating (property test 6). + +Nothing is mutated in place +--------------------------- +The live knowledge base keeps serving from the legacy backend throughout ``shadow`` +and ``verify`` (15.2, 15.3). The managed corpus is built alongside it and becomes +visible only at ``promote``, which is one conditional write. That is also why +rollback moves no data: the legacy index was never touched, so returning to it is +an attribute ``REMOVE``. + +Re-ingest, never re-upload +-------------------------- +Source bytes are already at +``assistants/{assistant_id}/documents/{document_id}/{filename}``, so migration +hands Bedrock the S3 location it already has (15.4). No user is ever asked to +re-supply a document, and no bytes move. + +Convergence, not dual-write +--------------------------- +The existing upload path stays authoritative and keeps writing to legacy for the +whole migration (16.1, 16.6). Rather than writing to both engines — which doubles +the number of ways a write can half-fail — the worker snapshots the document-id +set, migrates it, then runs catch-up passes until a pass finds nothing new (16.2, +16.3). Same converge-on-quiet shape as the crawler's consecutive-miss rule. + +Each document's ``DOC#`` record is re-read immediately before it is ingested and +skipped if it has gone or is no longer ``complete`` (16.4, 16.5). Without that +re-read, a document deleted while the migration was working through a long PDF +queue would be resurrected in the managed corpus — the user deleted it, saw it +disappear, and it comes back on a different engine. + +Feature: managed-kb-migration +Requirements: 15.1–15.14, 16.1–16.6, 17.1–17.5, 12.9 +""" + +from __future__ import annotations + +import asyncio +import logging +import os +from dataclasses import dataclass, field +from datetime import timedelta +from typing import Any, Dict, List, Optional, Sequence, Set + +logger = logging.getLogger() +logger.setLevel(logging.INFO) + +#: Document status the migration carries across. Requirement 15.5, and the same +#: value the facade's status filter serves — a document that is not ``complete`` +#: is not retrievable on legacy either, so migrating it would create a difference +#: where the whole point is parity. +STATUS_COMPLETE = "complete" + +#: Requirement 15.11. Legacy vectors are preserved for at least this long after +#: promotion, which is the window in which rollback is a pointer flip. +RETAIN_DAYS = 30 + +#: How long a worker holds a knowledge base. Long enough to cover the slowest +#: single step observed (a PDF-heavy shadow pass), short enough that a crashed +#: worker's knowledge base is picked up again the same hour. Requirement 15.13. +LEASE_MINUTES = 15 + +#: Catch-up passes before the worker gives up waiting for quiet. A knowledge base +#: whose owner is actively uploading may never converge; stopping is correct — +#: the record stays in ``shadow``, the dispatcher brings it back, and the corpus +#: keeps serving from legacy in the meantime. +MAX_CATCHUP_PASSES = 5 + +#: Documents ingested per managed call. Server-enforced at 10 for MANAGED +#: knowledge bases; the user guide's 25 is wrong. Named here so the batching is +#: visible at this level rather than only inside the adapter. +INGEST_BATCH = 10 + +#: Seconds added to ``dueAt`` when a step defers itself. Not a retry backoff — the +#: step succeeded — so it only needs to be long enough that the dispatcher does not +#: spin. +STEP_DELAY_SECONDS = 30 + +#: Ceiling on the completed-document set stored on the record. A DynamoDB item is +#: capped at 400 KB and this set is the only unbounded thing on it. Production's +#: entire corpus is 1,692 ``DOC#`` records across *all* assistants, so no real +#: knowledge base comes close; past the cap the worker stops tracking and a resume +#: re-ingests, which is slow but not wrong — ``customDocumentIdentifier`` makes a +#: re-ingest a replace. +MAX_TRACKED_DOCUMENT_IDS = 5000 + +METRIC_STARTED = "KbMigrationStarted" +METRIC_PROMOTED = "KbMigrationPromoted" +METRIC_FAILED = "KbMigrationFailed" +METRIC_ROLLED_BACK = "KbMigrationRolledBack" +METRIC_DOCUMENTS_MIGRATED = "KbMigrationDocumentsMigrated" +METRIC_DOCUMENTS_SKIPPED = "KbMigrationDocumentsSkipped" +METRIC_LEASE_LOST = "KbMigrationLeaseLost" + + +class MigrationError(Exception): + """A migration step could not complete. Leaves the record where it was.""" + + +class LeaseLost(MigrationError): + """Another worker holds this knowledge base. Not an error condition.""" + + +class VerificationFailed(MigrationError): + """The managed corpus does not match the source manifest, or the canary + retrieval returned nothing. Sends the migration to ``failed``, which leaves the + knowledge base on legacy and fully usable (17.4).""" + + +@dataclass +class StepResult: + """What one invocation did. Returned so the handler can log and test on it.""" + + assistant_id: str + app_kb_id: str + from_state: Optional[str] + to_state: Optional[str] + documents_migrated: int = 0 + documents_skipped: int = 0 + catchup_passes: int = 0 + converged: bool = False + detail: str = "" + manifest_diff: List[str] = field(default_factory=list) + + def as_log_fields(self) -> Dict[str, Any]: + return { + "appKbId": self.app_kb_id, + "from": self.from_state, + "to": self.to_state, + "migrated": self.documents_migrated, + "skipped": self.documents_skipped, + "catchupPasses": self.catchup_passes, + "converged": self.converged, + "detail": self.detail, + } + + +# ── environment, read at call time ─────────────────────────────────────────── +def _now(): + from datetime import datetime, timezone + + return datetime.now(timezone.utc) + + +def _iso(moment) -> str: + return moment.strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _now_iso() -> str: + return _iso(_now()) + + +def _documents_bucket() -> str: + bucket = os.environ.get("S3_ASSISTANTS_DOCUMENTS_BUCKET_NAME") + if not bucket: + raise MigrationError("S3_ASSISTANTS_DOCUMENTS_BUCKET_NAME is not set") + return bucket + + +def _retain_days() -> int: + raw = os.environ.get("KB_MIGRATION_RETAIN_DAYS") + try: + value = int(raw) if raw else RETAIN_DAYS + except ValueError: + return RETAIN_DAYS + # Requirement 15.11 says *at least* 30 days, so a smaller override is refused + # rather than honoured: shortening the rollback window is not a tuning knob. + return max(value, RETAIN_DAYS) + + +def _table(): + import boto3 + + return boto3.resource("dynamodb").Table(os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"]) + + +# ── document reads (raw table, no assistants import) ───────────────────────── +def list_document_items(assistant_id: str) -> List[Dict[str, Any]]: + """Every ``DOC#`` record under an assistant, paginated. + + Raw table access for the same reason ``kb_sync/records.py`` uses it: importing + ``apis.shared.assistants`` pulls in the embeddings stack at module scope and + this module ships in a size-constrained Lambda image. + """ + from boto3.dynamodb.conditions import Key + + table = _table() + items: List[Dict[str, Any]] = [] + kwargs: Dict[str, Any] = { + "KeyConditionExpression": Key("PK").eq(f"AST#{assistant_id}") + & Key("SK").begins_with("DOC#"), + } + while True: + response = table.query(**kwargs) + items.extend(response.get("Items") or []) + last = response.get("LastEvaluatedKey") + if not last: + return items + kwargs["ExclusiveStartKey"] = last + + +def get_document_item(assistant_id: str, document_id: str) -> Optional[Dict[str, Any]]: + response = _table().get_item( + Key={"PK": f"AST#{assistant_id}", "SK": f"DOC#{document_id}"} + ) + return response.get("Item") + + +def document_id_of(item: Dict[str, Any]) -> str: + sk = str(item.get("SK") or "") + return sk.split("#", 1)[1] if sk.startswith("DOC#") else "" + + +def is_complete(item: Optional[Dict[str, Any]]) -> bool: + return bool(item) and item.get("status") == STATUS_COMPLETE + + +def manifest_entry(item: Dict[str, Any]) -> str: + """One line of the source manifest: id plus a content identity. + + Requirement 15.6 forbids relying on document-count parity, and this is why the + manifest is a set of strings rather than a number. Count parity is satisfied by + a corpus with the right *number* of wrong documents — which is exactly what a + migration that raced an upload and a delete produces. + + The identity is the first available of ``contentHash``, ``etag`` or + ``updatedAt``. All three are already written by the existing pipeline; falling + through to ``updatedAt`` means a document with no hash still contributes a + changing value rather than a constant that always matches. + """ + document_id = document_id_of(item) + for key in ("contentHash", "etag", "generation", "updatedAt"): + value = item.get(key) + if value: + return f"{document_id}:{value}" + return f"{document_id}:no-identity" + + +def source_manifest(items: Sequence[Dict[str, Any]]) -> Set[str]: + return {manifest_entry(item) for item in items if is_complete(item)} + + +def s3_key_for(assistant_id: str, item: Dict[str, Any]) -> Optional[str]: + """The document's existing S3 key. Prefers the stored one. + + Reconstructed from ``filename`` only when the record has no ``s3Key``, because + the record is authoritative: a filename that was sanitised on upload would + reconstruct to a key that does not exist, and the ingest would fail per + document with an error naming the wrong cause. + """ + stored = item.get("s3Key") or item.get("s3_key") + if stored: + return str(stored) + filename = item.get("filename") + document_id = document_id_of(item) + if not filename or not document_id: + return None + return f"assistants/{assistant_id}/documents/{document_id}/{filename}" + + +def document_bytes(item: Dict[str, Any]) -> int: + for key in ("sizeBytes", "fileSize", "size"): + value = item.get(key) + if value is not None: + try: + return int(value) + except (TypeError, ValueError): + continue + return 0 + + +# ── lease ──────────────────────────────────────────────────────────────────── +async def take_lease(assistant_id: str, app_kb_id: str) -> str: + """Hold the knowledge base for :data:`LEASE_MINUTES`, or raise :class:`LeaseLost`. + + Requirement 15.13. Losing this is the ordinary outcome of two dispatcher ticks + overlapping, so it is logged at info and counted, not raised as a failure that + would move the record to ``failed`` and strand a perfectly healthy migration. + """ + from apis.shared.kb_backend import records as r + from apis.shared.kb_backend.metrics import emit_count + + now = _now() + lease_until = _iso(now + timedelta(minutes=LEASE_MINUTES)) + try: + await asyncio.to_thread( + r.acquire_lease, assistant_id, app_kb_id, lease_until, _iso(now) + ) + except Exception as exc: + emit_count(METRIC_LEASE_LOST) + raise LeaseLost( + f"another worker holds the lease on kb {app_kb_id}; leaving it alone: {exc}" + ) from exc + return lease_until + + +# ── ingestion of one snapshot ──────────────────────────────────────────────── +async def _ingest_documents( + assistant_id: str, + app_kb_id: str, + document_ids: Sequence[str], + backend, +) -> Dict[str, Any]: + """Re-ingest the named documents, re-reading each record first. + + Returns ``{"migrated": int, "skipped": int, "done": [ids]}``. ``done`` is the + documents genuinely handed to Bedrock, which is what gets persisted so a resume + can skip them — a count would not identify *which*. + + The re-read is Requirement 16.4 and it happens per document immediately before + that document is handed over, not once per batch: a PDF batch can take minutes, + and the deletion this guards against is most likely to land during exactly that + window. + """ + from apis.shared.kb_backend.protocol import DocumentSource + + migrated = 0 + skipped = 0 + done: List[str] = [] + batch: List[DocumentSource] = [] + + async def flush() -> None: + nonlocal batch, migrated + if not batch: + return + await backend.ingest_documents(app_kb_id, batch, batch_size=INGEST_BATCH) + migrated += len(batch) + done.extend(source.document_id for source in batch) + batch = [] + + for document_id in document_ids: + item = await asyncio.to_thread(get_document_item, assistant_id, document_id) + if not is_complete(item): + # Gone, or no longer complete. Requirement 16.5: not resurrected. + logger.info( + f"skipping document {document_id}: status=" + f"{(item or {}).get('status', 'NOT_FOUND')}" + ) + skipped += 1 + continue + + key = s3_key_for(assistant_id, item) + if not key: + logger.warning(f"skipping document {document_id}: no resolvable S3 key") + skipped += 1 + continue + + batch.append( + DocumentSource( + document_id=document_id, + filename=str(item.get("filename") or document_id), + s3_key=key, + metadata={"document_id": document_id, "filename": str(item.get("filename") or "")}, + ) + ) + if len(batch) >= INGEST_BATCH: + await flush() + + await flush() + return {"migrated": migrated, "skipped": skipped, "done": done} + + +# ── steps ──────────────────────────────────────────────────────────────────── +async def run_shadow( + assistant_id: str, + app_kb_id: str, + record: Dict[str, Any], + backend=None, +) -> StepResult: + """Provision, reserve the whole corpus, ingest the snapshot, then converge. + + Order matters and is not arbitrary: + + 1. **Reserve the whole snapshot first** (12.9). Migration is the largest + byte-adding operation in the system and the only one that runs unattended. + Reserving per document would let a migration run for an hour and stop + halfway, leaving a half-populated managed corpus and an owner over their cap + with no way back. + 2. **Provision.** Lazy by design, so the knowledge base may not exist yet. + 3. **Ingest the snapshot**, re-reading each record immediately before use. + 4. **Catch up until quiet** (16.2, 16.3), then move to ``verify``. + """ + from apis.shared.kb_backend import byte_cap, records as r + from apis.shared.kb_backend.metrics import emit_count + from apis.shared.kb_backend.provisioning import provision_managed_kb + + generation = int(record.get("migrationGeneration") or 0) + items = await asyncio.to_thread(list_document_items, assistant_id) + complete = [item for item in items if is_complete(item)] + + # Documents a previous invocation already ingested. Skipping them is what makes + # a resumed migration cost seconds rather than re-parsing a PDF corpus that can + # take over an hour. + done = already_migrated(record) + snapshot = [ + document_id_of(item) + for item in complete + if document_id_of(item) and document_id_of(item) not in done + ] + + total_bytes = sum(document_bytes(item) for item in complete) + if total_bytes and record.get("totalBytes") in (None, 0): + # Reserved once per migration, not once per resume: the accumulator is on + # the record, so a resumed run that reserved again would double-count its + # own corpus against the owner's cap and eventually refuse itself. + # Raises ByteCapExceeded, which the caller turns into `failed` — before + # anything has been provisioned or ingested. + await asyncio.to_thread( + byte_cap.reserve_snapshot, + assistant_id, + app_kb_id, + total_bytes, + byte_cap.per_owner_cap(bool(record.get("elevatedByteCap"))), + ) + + await provision_managed_kb( + assistant_id, + app_kb_id, + owner_user_id=str(record.get("ownerUserId") or ""), + ) + + backend = backend or _managed_backend() + emit_count(METRIC_STARTED) + + counts = await _ingest_documents(assistant_id, app_kb_id, snapshot, backend) + migrated_ids = set(snapshot) | done + + passes, converged, extra = await catch_up( + assistant_id, app_kb_id, migrated_ids, backend + ) + counts["migrated"] += extra["migrated"] + counts["skipped"] += extra["skipped"] + newly_done = list(counts["done"]) + list(extra["done"]) + + total_done = len(done) + counts["migrated"] + await _record_progress( + assistant_id, + app_kb_id, + migrated=total_done, + total=total_done, + skipped=counts["skipped"], + newly_done=newly_done, + ) + + if not converged: + # Still busy. Stay in `shadow`; the dispatcher brings this back, and the + # corpus keeps serving from legacy in the meantime. + await asyncio.to_thread( + r.set_migration_state, + assistant_id, + app_kb_id, + r.SHADOW, + generation, + _iso(_now() + timedelta(seconds=STEP_DELAY_SECONDS)), + [r.SHADOW], + ) + return StepResult( + assistant_id, + app_kb_id, + r.SHADOW, + r.SHADOW, + counts["migrated"], + counts["skipped"], + passes, + False, + "catch-up did not converge; staying in shadow", + ) + + await asyncio.to_thread( + r.set_migration_state, + assistant_id, + app_kb_id, + r.VERIFY, + generation, + _iso(_now() + timedelta(seconds=STEP_DELAY_SECONDS)), + [r.SHADOW], + ) + return StepResult( + assistant_id, + app_kb_id, + r.SHADOW, + r.VERIFY, + counts["migrated"], + counts["skipped"], + passes, + True, + ) + + +async def catch_up( + assistant_id: str, + app_kb_id: str, + already: Set[str], + backend, + max_passes: int = None, +) -> tuple: + """Ingest documents that appeared since the snapshot, until a pass finds none. + + Requirements 16.2, 16.3. Returns ``(passes, converged, counts)``. + + Converged means a pass found nothing new — not that a fixed number of passes + ran. The distinction matters because the number of passes needed depends on how + fast the owner is uploading, which is not something this code can know in + advance. ``max_passes`` bounds the invocation, and *not* converging is a normal + outcome that leaves the record in ``shadow``. + """ + limit = MAX_CATCHUP_PASSES if max_passes is None else max_passes + counts: Dict[str, Any] = {"migrated": 0, "skipped": 0, "done": []} + + for attempt in range(1, limit + 1): + items = await asyncio.to_thread(list_document_items, assistant_id) + pending = [ + document_id_of(item) + for item in items + if is_complete(item) and document_id_of(item) not in already + ] + if not pending: + logger.info(f"catch-up converged for kb {app_kb_id} after {attempt} pass(es)") + return attempt, True, counts + + logger.info(f"catch-up pass {attempt} for kb {app_kb_id}: {len(pending)} new document(s)") + pass_counts = await _ingest_documents(assistant_id, app_kb_id, pending, backend) + counts["migrated"] += pass_counts["migrated"] + counts["skipped"] += pass_counts["skipped"] + counts["done"].extend(pass_counts["done"]) + already.update(pending) + + return limit, False, counts + + +async def run_verify( + assistant_id: str, + app_kb_id: str, + record: Dict[str, Any], + backend=None, +) -> StepResult: + """Compare an exact manifest, then prove retrieval works. + + Requirements 15.6, 15.7. Two checks, and both are needed: + + * The **manifest** is a set of ``document_id:identity`` strings, not a count. + A count is satisfied by the right number of wrong documents. + * The **canary retrieval** proves the corpus is genuinely queryable. Bedrock + reporting a document ``INDEXED`` precedes it being retrievable by + 0.75–1.03 s, and a knowledge base can hold documents while returning nothing + — so "we ingested everything" and "retrieval works" are separate claims. + """ + from apis.shared.kb_backend import records as r + + generation = int(record.get("migrationGeneration") or 0) + backend = backend or _managed_backend() + + items = await asyncio.to_thread(list_document_items, assistant_id) + expected = source_manifest(items) + + complete = [item for item in items if is_complete(item)] + if not complete: + raise VerificationFailed( + f"kb {app_kb_id} has no complete documents to verify; there is nothing " + f"to promote" + ) + + canary_text = _canary_query(complete) + chunks = await backend.search(app_kb_id, canary_text, 5) + if not chunks: + raise VerificationFailed( + f"canary retrieval on kb {app_kb_id} returned nothing; the managed " + f"corpus is not queryable yet" + ) + + retrieved_ids = {chunk.document_id for chunk in chunks if chunk.document_id} + expected_ids = {document_id_of(item) for item in complete} + if not retrieved_ids & expected_ids: + raise VerificationFailed( + f"canary retrieval on kb {app_kb_id} returned only documents this " + f"assistant does not own: {sorted(retrieved_ids)}" + ) + + await asyncio.to_thread( + r.set_migration_state, + assistant_id, + app_kb_id, + r.PROMOTE, + generation, + _iso(_now() + timedelta(seconds=STEP_DELAY_SECONDS)), + [r.VERIFY], + ) + return StepResult( + assistant_id, + app_kb_id, + r.VERIFY, + r.PROMOTE, + detail=f"manifest of {len(expected)} document(s) verified; canary returned " + f"{len(chunks)} chunk(s)", + ) + + +def _canary_query(complete: Sequence[Dict[str, Any]]) -> str: + """A query built from the corpus's own filenames. + + Not a fixed string. A constant like "test" can legitimately match nothing in a + real corpus, which would make verification fail for healthy knowledge bases and + train whoever is watching to ignore it. + """ + names = [str(item.get("filename") or "") for item in complete[:3]] + text = " ".join(name.rsplit(".", 1)[0].replace("_", " ").replace("-", " ") for name in names) + return text.strip() or "summary" + + +async def run_promote( + assistant_id: str, + app_kb_id: str, + record: Dict[str, Any], +) -> StepResult: + """The cutover: one conditional write, then straight into ``retain``. + + Requirements 15.8, 15.9, 15.10. Everything that makes this safe lives in + ``records.promote_engine``'s condition — the state, the generation, and + ``migrationProgress.migrated == migrationProgress.total``, so a promotion + cannot happen on a knowledge base whose catch-up never converged. Two + concurrent workers issue the same write and DynamoDB picks one. + + The byte cap must already be enforced on this knowledge base (12.9): no + traffic is promoted to an unmetered corpus, so a record carrying no + ``totalBytes`` accumulator is refused here rather than discovered later. + """ + from apis.shared.kb_backend import records as r + from apis.shared.kb_backend.metrics import emit_count + + generation = int(record.get("migrationGeneration") or 0) + + if record.get("totalBytes") is None: + raise MigrationError( + f"refusing to promote kb {app_kb_id}: it has no totalBytes accumulator, " + f"so the byte cap is not being enforced on it (Requirement 12.9)" + ) + + # Resuming after a crash *between* the promotion and the state transition. The + # promotion write is guarded on `attribute_not_exists(retrievalEngine)`, so + # retrying it here would be 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 and no path to `retain`. Found by the + # convergence property test, which crashed at exactly that transition. + already_promoted = record.get("retrievalEngine") == r.ENGINE_MANAGED + + if not already_promoted: + try: + await asyncio.to_thread( + r.promote_engine, assistant_id, app_kb_id, generation, _now_iso() + ) + emit_count(METRIC_PROMOTED) + except Exception: + # Re-read before deciding. The write may have been refused because + # somebody else promoted first, which is success, or because a guard + # genuinely failed, which is not. + fresh = await asyncio.to_thread(r.get_kb_record, assistant_id, app_kb_id) + if (fresh or {}).get("retrievalEngine") != r.ENGINE_MANAGED: + raise + logger.info( + f"kb {app_kb_id} was already promoted by another attempt; " + f"continuing to retain rather than failing" + ) + already_promoted = True + + retain_until = _iso(_now() + timedelta(days=_retain_days())) + await asyncio.to_thread( + _set_retain_until, assistant_id, app_kb_id, retain_until + ) + await asyncio.to_thread( + r.set_migration_state, + assistant_id, + app_kb_id, + r.RETAIN, + generation, + None, + [r.PROMOTE], + ) + return StepResult( + assistant_id, + app_kb_id, + r.PROMOTE, + r.RETAIN, + detail=( + f"{'already promoted; ' if already_promoted else ''}legacy vectors " + f"retained until {retain_until}" + ), + ) + + +def _set_retain_until(assistant_id: str, app_kb_id: str, retain_until: str) -> None: + """Stamp the rollback deadline. Unconditional, and deliberately so. + + The promotion write immediately before this one is the guarded one. If this + write were also guarded and lost, the record would be promoted with no + ``retainUntil`` — which reads as "no rollback window" to anything that checks + it. Writing the later date twice is harmless; writing it never is not. + """ + import boto3 + + boto3.resource("dynamodb").Table(os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"]).update_item( + Key={"PK": f"AST#{assistant_id}", "SK": f"KB#{app_kb_id}"}, + UpdateExpression="SET retainUntil = :until", + ExpressionAttributeValues={":until": retain_until}, + ) + + +async def rollback(assistant_id: str, app_kb_id: str) -> StepResult: + """Return a promoted knowledge base to legacy. Moves no data. + + Requirement 17. The legacy index was never mutated — that is what ``shadow`` + building alongside it bought — so rollback is one attribute ``REMOVE`` plus a + timestamp. It is available for the whole ``retain`` window because that window + is exactly the promise not to reclaim the legacy vectors. + + Note this does **not** delete the managed knowledge base. A rolled-back corpus + that still exists costs storage but can be re-promoted without a second + migration; deleting it here would turn a reversible decision into an + irreversible one at the moment somebody is least sure. + """ + from apis.shared.kb_backend import records as r + from apis.shared.kb_backend.metrics import emit_count + + await asyncio.to_thread(r.rollback_engine, assistant_id, app_kb_id, _now_iso()) + emit_count(METRIC_ROLLED_BACK) + return StepResult( + assistant_id, + app_kb_id, + r.RETAIN, + r.RETAIN, + detail="rolled back to the legacy engine; no data moved", + ) + + +async def _record_progress( + assistant_id: str, + app_kb_id: str, + *, + migrated: int, + total: int, + skipped: int, + newly_done: Optional[Sequence[str]] = None, +) -> None: + """Write ``migrationProgress``, which the promotion condition reads. + + ``total`` is a DynamoDB reserved keyword, so both progress paths are aliased. + Unaliased, the write is rejected outright with a ``ValidationException`` — loud, + but only because it never validates at all. + + ``newly_done`` is ``ADD``ed to the ``migratedDocIds`` string set rather than + written into the progress map. Two reasons, and both are the difference between + a resumed migration costing seconds and costing an hour: + + * **``ADD`` is additive**, so a crash between batches loses only the batch in + flight. A read-modify-write of a list would lose everything since the last + read, and would also let two workers clobber each other. + * **It is a separate attribute** from ``migrationProgress``, which this function + overwrites wholesale. Keeping the completed-document set inside a map that + gets replaced is how a resume silently re-ingests a corpus it had already + finished — found by the convergence property test, which counted a document + ingested twice across a crash and a retry. + """ + from decimal import Decimal + + expression = "SET #progress = :progress" + names = {"#progress": "migrationProgress"} + values: Dict[str, Any] = { + ":progress": { + "migrated": Decimal(migrated), + "total": Decimal(total), + "skipped": Decimal(skipped), + "updatedAt": _now_iso(), + } + } + + ids = [document_id for document_id in (newly_done or []) if document_id] + if ids and len(ids) <= MAX_TRACKED_DOCUMENT_IDS: + # DynamoDB string sets cannot be empty, hence the guard above. + expression += " ADD #done :done" + names["#done"] = "migratedDocIds" + values[":done"] = set(ids) + elif ids: + logger.warning( + f"kb {app_kb_id}: {len(ids)} document ids exceeds the tracking cap of " + f"{MAX_TRACKED_DOCUMENT_IDS}; a resumed migration will re-ingest, which " + f"is safe but slow (customDocumentIdentifier makes re-ingest a replace)" + ) + + _table().update_item( + Key={"PK": f"AST#{assistant_id}", "SK": f"KB#{app_kb_id}"}, + UpdateExpression=expression, + ExpressionAttributeNames=names, + ExpressionAttributeValues=values, + ) + + +def already_migrated(record: Dict[str, Any]) -> Set[str]: + """Documents a previous invocation already ingested. + + Read from the ``migratedDocIds`` string set. Empty for a record that has never + ingested anything, which is also what a corpus past the tracking cap looks + like — and that degradation is safe: re-ingesting a document replaces it, + because ``customDocumentIdentifier`` is the platform document id. + """ + stored = record.get("migratedDocIds") + if not stored: + return set() + try: + return {str(document_id) for document_id in stored} + except TypeError: + logger.warning(f"migratedDocIds is not iterable on this record: {stored!r}") + return set() + + +def _managed_backend(): + from apis.shared.kb_backend.managed_backend import ManagedKbBackend + + return ManagedKbBackend(bucket=_documents_bucket()) + + +# ── one invocation ─────────────────────────────────────────────────────────── +async def run_step( + assistant_id: str, + app_kb_id: Optional[str] = None, + backend=None, +) -> StepResult: + """Take the lease and execute the one step this record's state calls for. + + Dispatches on the *record's* state, never on the invocation event's. The event + carries a state for logging, but trusting it would let a hand-crafted invocation + promote a knowledge base that never verified — the same class of bypass that + let an event field arm the reconciler. + """ + from apis.shared.kb_backend import byte_cap, records as r + from apis.shared.kb_backend.metrics import emit_count + + app_kb_id = app_kb_id or assistant_id + + record = await asyncio.to_thread(r.get_kb_record, assistant_id, app_kb_id) + if not record: + raise MigrationError(f"no KB_Record for {assistant_id}/{app_kb_id}") + + state = record.get("migrationState") + if state not in r.WORK_ELIGIBLE_STATES: + # Terminal, or never enrolled. Not an error: the dispatcher reads an index + # that is eventually consistent, so a record finished a moment ago can + # still be handed over once. + return StepResult( + assistant_id, app_kb_id, state, state, detail="not work-eligible; nothing to do" + ) + + generation = int(record.get("migrationGeneration") or 0) + + try: + # Inside the try, deliberately. A ``LeaseLost`` must reach the caller as + # itself — losing a lease is two dispatcher ticks overlapping, not a broken + # migration — and the ``except LeaseLost: raise`` below is what guarantees + # that even once a step starts taking sub-leases of its own. Outside the try + # the clause would be unreachable, which is how a guard becomes decoration. + await take_lease(assistant_id, app_kb_id) + + if state == r.SHADOW: + result = await run_shadow(assistant_id, app_kb_id, record, backend) + elif state == r.VERIFY: + result = await run_verify(assistant_id, app_kb_id, record, backend) + else: + result = await run_promote(assistant_id, app_kb_id, record) + except LeaseLost: + raise + except (VerificationFailed, byte_cap.ByteCapExceeded) as exc: + # Expected failure modes. The knowledge base stays on legacy and stays + # usable (17.4); `failed` is terminal and removes the work keys. + await _fail(assistant_id, app_kb_id, generation, str(exc)) + emit_count(METRIC_FAILED) + return StepResult( + assistant_id, app_kb_id, state, r.MIGRATION_FAILED, detail=str(exc) + ) + except Exception as exc: + # Unexpected. Also terminal, for the same reason: an unbounded retry on an + # unknown fault is how a migration loop bills for a week. + logger.error(f"migration step failed for kb {app_kb_id}: {exc}", exc_info=True) + await _fail(assistant_id, app_kb_id, generation, f"{type(exc).__name__}: {exc}") + emit_count(METRIC_FAILED) + return StepResult( + assistant_id, app_kb_id, state, r.MIGRATION_FAILED, detail=str(exc) + ) + + logger.info(f"migration step: {result.as_log_fields()}") + return result + + +async def _fail(assistant_id: str, app_kb_id: str, generation: int, reason: str) -> None: + from apis.shared.kb_backend import records as r + + try: + await asyncio.to_thread( + r.set_migration_state, + assistant_id, + app_kb_id, + r.MIGRATION_FAILED, + generation, + None, + None, + reason[:1000], + ) + except Exception as exc: + # Nothing further to do: the record keeps its work keys and the dispatcher + # will bring it back, which is the safe direction — a knowledge base stuck + # in `shadow` still serves from legacy. + logger.error(f"could not record migration failure for kb {app_kb_id}: {exc}") + + +def lambda_handler(event, context): + """Async-invoked by the dispatcher. + + Reads only the two identifiers from the event. Everything that decides what + happens — the state, the generation, the flags — comes from the record and the + environment. + """ + assistant_id = (event or {}).get("assistantId") + app_kb_id = (event or {}).get("appKbId") + if not assistant_id: + raise MigrationError("event carries no assistantId") + + try: + result = asyncio.run(run_step(assistant_id, app_kb_id)) + except LeaseLost as exc: + logger.info(str(exc)) + return {"statusCode": 200, "body": {"leaseLost": True}} + + return {"statusCode": 200, "body": result.as_log_fields()} diff --git a/backend/src/apis/app_api/kb_upgrade/__init__.py b/backend/src/apis/app_api/kb_upgrade/__init__.py new file mode 100644 index 000000000..40eaeef10 --- /dev/null +++ b/backend/src/apis/app_api/kb_upgrade/__init__.py @@ -0,0 +1,12 @@ +"""The owner-facing knowledge base upgrade surface (Requirements 21, 23). + +Deliberately a **separate package from** ``apis.app_api.kb_migration``. That +package holds the four Lambda handlers, which share one size-constrained image; +this one is HTTP-only and imports ``apis.shared.assistants`` for the permission +model, which pulls the embeddings stack at module scope. Putting the two in the +same package invites a handler import that blows the image-size budget — the +failure ``tests/architecture/test_kb_backend_boundary.py`` exists to prevent. + +Nothing here writes ``retrievalEngine``. Enrolment only moves a record into +``shadow``; the worker promotes, and only after verification. +""" diff --git a/backend/src/apis/app_api/kb_upgrade/models.py b/backend/src/apis/app_api/kb_upgrade/models.py new file mode 100644 index 000000000..5fffa8cfe --- /dev/null +++ b/backend/src/apis/app_api/kb_upgrade/models.py @@ -0,0 +1,95 @@ +"""Wire models for the knowledge base upgrade surface. + +Field names are camelCase on the wire (``populate_by_name`` + aliases), matching +every other app_api surface the Angular client consumes. + +The word "vector" appears nowhere in any user-facing string in this module, per +Requirement 23.6. It is fine in comments; it is not fine in ``message``. +""" + +from typing import List, Literal, Optional + +from pydantic import BaseModel, ConfigDict, Field + +#: The derived, UI-facing phase. Deliberately NOT the record's ``migrationState``: +#: the client should not have to know that ``shadow``, ``verify`` and ``promote`` +#: are all "working on it", nor that absence means legacy. +#: +#: ``none`` is the state that renders nothing at all (Requirement 23.1). +UpgradePhase = Literal["none", "available", "in_progress", "succeeded", "failed"] + +#: Why a document will not be carried across. Requirement 21.4 requires an +#: unsupported format to be distinguishable from a processing failure, because the +#: user's next action differs: convert and re-upload, versus just retry. +DocumentIssueKind = Literal[ + "unsupported_format", + "processing_failure", + "still_processing", + "being_removed", +] + + +class UpgradeProgress(BaseModel): + """Non-blocking progress for the ``in_progress`` phase (Requirement 23.3).""" + + model_config = ConfigDict(populate_by_name=True) + + completed: int = Field(0, description="Documents carried across so far") + total: int = Field(0, description="Documents in the snapshot being carried") + skipped: int = Field(0, description="Documents the snapshot could not include") + + +class DocumentNotCarried(BaseModel): + """A document the upgrade will not carry across (Requirement 21.1). + + Surfaced *before* the user commits, not after, so the choice to retry or + accept the loss is made with the facts in hand. + """ + + model_config = ConfigDict(populate_by_name=True) + + document_id: str = Field(..., alias="documentId") + filename: str + status: str = Field(..., description="The stored processing status, verbatim") + kind: DocumentIssueKind + message: str = Field( + ..., + description="Plain-language explanation, safe to render directly", + ) + retryable: bool = Field( + ..., + description="Whether re-processing this document could succeed as-is", + ) + + +class UpgradeStatusResponse(BaseModel): + """Everything the card needs to render, in one round trip.""" + + model_config = ConfigDict(populate_by_name=True) + + phase: UpgradePhase + #: True only for an owner/editor looking at an ``available`` knowledge base. + #: The client hides the control on this alone; the server re-checks on write, + #: so a client that ignores it gains nothing (Requirement 23.7). + can_upgrade: bool = Field(False, alias="canUpgrade") + progress: Optional[UpgradeProgress] = None + #: Plain-language failure reason for the ``failed`` phase (Requirement 23.5). + reason: Optional[str] = None + #: Whether the one-time success notice is still owed (Requirement 23.4). + #: Never sticky: dismissing it sets a timestamp and this goes false forever. + notice_pending: bool = Field(False, alias="noticePending") + documents_not_carried: List[DocumentNotCarried] = Field( + default_factory=list, alias="documentsNotCarried" + ) + + +class EnrollResponse(BaseModel): + """Result of an enrol or retry.""" + + model_config = ConfigDict(populate_by_name=True) + + phase: UpgradePhase + #: True when this call is what started the upgrade, false when it found one + #: already running. Both are successes — a double-click is not an error. + started: bool + message: str diff --git a/backend/src/apis/app_api/kb_upgrade/routes.py b/backend/src/apis/app_api/kb_upgrade/routes.py new file mode 100644 index 000000000..c5dd33822 --- /dev/null +++ b/backend/src/apis/app_api/kb_upgrade/routes.py @@ -0,0 +1,144 @@ +"""HTTP surface for the owner-facing knowledge base upgrade (Requirements 21, 23). + +Four endpoints, all under the assistant that owns the knowledge base: + +* ``GET /assistants/{id}/knowledge-base/upgrade`` — what to render +* ``POST /assistants/{id}/knowledge-base/upgrade`` — opt in +* ``POST /assistants/{id}/knowledge-base/upgrade/retry`` — after a failure +* ``POST /assistants/{id}/knowledge-base/upgrade/notice`` — dismiss the notice + +Permission is resolved the same way the documents surface resolves it, through +``resolve_assistant_permission``. The read is allowed for any resolvable +permission and reports ``canUpgrade: false`` to a viewer; the three writes +require owner or editor. Requirement 23.7's "viewers never see the control" is +therefore enforced on the server, with the client's hiding as presentation only. +""" + +import logging + +from fastapi import APIRouter, Depends, HTTPException, status + +from apis.app_api.kb_upgrade.models import EnrollResponse, UpgradeStatusResponse +from apis.app_api.kb_upgrade.service import ( + UpgradeUnavailable, + dismiss_notice, + enroll, + get_upgrade_status, + retry, +) +from apis.shared.assistants.service import resolve_assistant_permission +from apis.shared.auth import User, get_current_user_from_session + +logger = logging.getLogger(__name__) + +router = APIRouter( + prefix="/assistants/{assistant_id}/knowledge-base/upgrade", + tags=["knowledge-base-upgrade"], +) + +_EDIT_PERMISSIONS = ("owner", "editor") + + +async def _resolve(assistant_id: str, current_user: User): + """Return ``(assistant, permission)`` or raise 404. + + A permission of ``None`` from a resolvable assistant means the user has no + access at all, which is reported as 404 rather than 403 so the endpoint does + not confirm the existence of an assistant the caller cannot see. + """ + assistant, permission = await resolve_assistant_permission( + assistant_id=assistant_id, + user_id=current_user.user_id, + user_email=current_user.email, + ) + if not assistant or not permission: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Assistant not found: {assistant_id}", + ) + return assistant, permission + + +async def _require_edit_permission(assistant_id: str, current_user: User): + """Resolve and require owner|editor. Returns ``(assistant, permission)``.""" + assistant, permission = await _resolve(assistant_id, current_user) + if permission not in _EDIT_PERMISSIONS: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You do not have permission to upgrade this knowledge base", + ) + return assistant, permission + + +@router.get("", response_model=UpgradeStatusResponse) +async def read_upgrade_status( + assistant_id: str, + current_user: User = Depends(get_current_user_from_session), +) -> UpgradeStatusResponse: + """What the upgrade card should render, or ``phase: "none"`` for nothing.""" + _, permission = await _resolve(assistant_id, current_user) + try: + return await get_upgrade_status( + assistant_id, can_edit=permission in _EDIT_PERMISSIONS + ) + except Exception as exc: # noqa: BLE001 — see below + # Fail to "nothing to show" rather than 500. This endpoint is decoration + # on a working page: a knowledge base that cannot be described still + # serves retrieval, and taking the whole documents section down over an + # unavailable upgrade card would be a strictly worse outcome. Logged at + # error so the failure is not silent. + logger.error( + f"kb {assistant_id}: could not derive upgrade status: {exc}", + exc_info=True, + ) + return UpgradeStatusResponse(phase="none", canUpgrade=False) + + +@router.post("", response_model=EnrollResponse, status_code=status.HTTP_202_ACCEPTED) +async def start_upgrade( + assistant_id: str, + current_user: User = Depends(get_current_user_from_session), +) -> EnrollResponse: + """Opt this knowledge base into the upgrade (Requirement 23.2, 23.8). + + 202 rather than 200: enrolment queues work for the migration worker and + returns before any of it has happened. + """ + assistant, _ = await _require_edit_permission(assistant_id, current_user) + try: + return await enroll( + assistant_id, + owner_user_id=assistant.owner_id, + visibility=str(getattr(assistant, "visibility", "PRIVATE") or "PRIVATE"), + ) + except UpgradeUnavailable as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, detail=str(exc) + ) from exc + + +@router.post( + "/retry", response_model=EnrollResponse, status_code=status.HTTP_202_ACCEPTED +) +async def retry_upgrade( + assistant_id: str, + current_user: User = Depends(get_current_user_from_session), +) -> EnrollResponse: + """Restart a failed upgrade on a fresh generation (Requirement 23.5).""" + assistant, _ = await _require_edit_permission(assistant_id, current_user) + try: + return await retry(assistant_id, owner_user_id=assistant.owner_id) + except UpgradeUnavailable as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, detail=str(exc) + ) from exc + + +@router.post("/notice", status_code=status.HTTP_204_NO_CONTENT) +async def dismiss_upgrade_notice( + assistant_id: str, + current_user: User = Depends(get_current_user_from_session), +) -> None: + """Dismiss the one-time post-upgrade notice (Requirement 23.4).""" + await _require_edit_permission(assistant_id, current_user) + await dismiss_notice(assistant_id) diff --git a/backend/src/apis/app_api/kb_upgrade/service.py b/backend/src/apis/app_api/kb_upgrade/service.py new file mode 100644 index 000000000..2bb9de264 --- /dev/null +++ b/backend/src/apis/app_api/kb_upgrade/service.py @@ -0,0 +1,507 @@ +"""Enrolment and status for the owner-facing knowledge base upgrade. + +Three things happen here and nothing else: derive what the card should show, +move a record into ``shadow``, and dismiss the one-time success notice. Promotion +belongs to the worker, after verification — this module never writes +``retrievalEngine``, so no HTTP request can put a knowledge base on the managed +backend without the corpus having been carried across and checked first. + +Enrolment is deliberately two conditional writes rather than one put: + +1. ``create_provisioning`` — guarded on ``attribute_not_exists(PK)``. +2. ``set_migration_state(SHADOW, ...)`` — guarded on the generation. + +A single ``put_item`` with ``migrationState="shadow"`` baked in would look +simpler and would be **wrong**: ``KbRecord.to_item`` does not write the +``GSI7_PK``/``GSI7_SK`` work keys, which only ``set_migration_state`` maintains. +The record would exist, claim to be migrating, and be invisible to the +dispatcher's sparse-index sweep forever. +""" + +from __future__ import annotations + +import logging +import os +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, List, Optional, Tuple + +from apis.app_api.kb_upgrade.models import ( + DocumentNotCarried, + EnrollResponse, + UpgradeProgress, + UpgradeStatusResponse, +) + +logger = logging.getLogger(__name__) + + +class UpgradeUnavailable(Exception): + """The upgrade cannot be offered right now. Carries user-safe copy.""" + + +#: Gate on the same flag the dispatcher reads. If the worker cannot run, offering +#: the upgrade would park a record in ``shadow`` that nothing ever picks up — a +#: spinner with no engine behind it. Requirement 23.1 says show nothing when no +#: action is available, and "available" has to mean actionable. +FLAG_MIGRATION_ENABLED = "MANAGED_KB_MIGRATION_ENABLED" + +#: Affirmative spellings, matching ``dispatcher._TRUTHY`` exactly. An allow-list +#: rather than truthiness, because the value being designed around is present but +#: empty: ``bool("")`` is right by luck and ``bool("false")`` is not. +_TRUTHY = frozenset({"1", "true", "yes", "on", "enabled"}) + +#: How soon the dispatcher may pick up a freshly enrolled record. Now, not later: +#: the user just asked for it, and the dispatcher is already rate-bounded. +_DUE_IMMEDIATELY = timedelta(0) + +STATUS_COMPLETE = "complete" + + +def migration_enabled() -> bool: + """Whether the upgrade may be offered at all. + + Read at call time, never bound as a default argument — a module-level default + is captured at import and makes the flag unpatchable, which already cost this + feature a 33-second test that ignored its own override. + """ + return (os.environ.get(FLAG_MIGRATION_ENABLED) or "").strip().lower() in _TRUTHY + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +def _iso(moment: datetime) -> str: + return moment.isoformat().replace("+00:00", "Z") + + +# ── Document classification (Requirement 21) ───────────────────────────────── +def _extension_of(filename: str) -> str: + _, _, tail = str(filename or "").rpartition(".") + return f".{tail.lower()}" if tail else "" + + +def _supported_extensions() -> frozenset: + """The ingestion pipeline's own extension set, read from its module. + + Imported here rather than copied. A copied list is exactly the drift that + produced the tag-contract defect: three files agreeing only because they all + fell back to the same hardcoded default. ``docling_processor``'s module scope + is stdlib-only, so this costs nothing at request time. + """ + from apis.app_api.documents.ingestion.processors.docling_processor import ( + DOCLING_SUPPORTED_EXTENSIONS, + ) + + return frozenset(DOCLING_SUPPORTED_EXTENSIONS) + + +def classify_document(item: Dict[str, Any]) -> Optional[DocumentNotCarried]: + """Describe why a document will not be carried across, or ``None`` if it will. + + Requirement 21.4 turns on the ``unsupported_format`` / ``processing_failure`` + split: the two demand different actions from the user. Telling someone to + "retry" a ``.pages`` file wastes a minute and teaches them the retry button + does not work. + + Note ``deleting`` is reported rather than hidden. The ordinary document list + filters that status out as soft-deleted, which is right there and wrong here: + 101 of the 200 affected production records are stuck in it, and a user who is + never shown them cannot tell that they are stuck. + """ + status = str(item.get("status") or "").strip() + if status == STATUS_COMPLETE: + return None + + document_id = str(item.get("documentId") or "") + if not document_id: + sk = str(item.get("SK") or "") + document_id = sk.split("#", 1)[1] if sk.startswith("DOC#") else "" + filename = str(item.get("filename") or "(unnamed file)") + extension = _extension_of(filename) + unsupported = bool(extension) and extension not in _supported_extensions() + + if status == "failed" and unsupported: + return DocumentNotCarried( + documentId=document_id, + filename=filename, + status=status, + kind="unsupported_format", + message=( + f"This platform cannot read {extension} files, so this document was " + "never added to your knowledge base. Save it as a PDF or Word " + "document and upload it again." + ), + retryable=False, + ) + if status == "failed": + stored = str(item.get("errorMessage") or "").strip() + detail = f" The reason given was: {stored}" if stored else "" + return DocumentNotCarried( + documentId=document_id, + filename=filename, + status=status, + kind="processing_failure", + message=( + "This document could not be processed, so it is not in your " + f"knowledge base and the upgrade cannot carry it across.{detail}" + ), + retryable=True, + ) + if status == "deleting": + return DocumentNotCarried( + documentId=document_id, + filename=filename, + status=status, + kind="being_removed", + message=( + "This document is part-way through being removed. It will not be " + "carried across. If you still want it, upload it again once the " + "removal finishes." + ), + retryable=False, + ) + return DocumentNotCarried( + documentId=document_id, + filename=filename, + status=status, + kind="still_processing", + message=( + "This document is still being processed. Documents that are not " + "finished when the upgrade starts will not be carried across." + ), + retryable=True, + ) + + +def _document_items(assistant_id: str) -> List[Dict[str, Any]]: + """Every ``DOC#`` item under an assistant, unfiltered. + + Raw query rather than ``list_assistant_documents``, on purpose and for two + reasons: that function drops ``deleting`` documents, which are the single + largest group Requirement 21 exists to surface, and it auto-fails stale ones + as a side effect — a write triggered by rendering a card. + """ + import boto3 + from boto3.dynamodb.conditions import Key + + table_name = os.environ.get("DYNAMODB_ASSISTANTS_TABLE_NAME") + if not table_name: + # Fail closed and loudly enough to see, but do not take the card down: a + # misconfigured table name must not make an upgradeable KB look clean. + raise RuntimeError("DYNAMODB_ASSISTANTS_TABLE_NAME is not set") + + table = boto3.resource("dynamodb").Table(table_name) + items: List[Dict[str, Any]] = [] + kwargs: Dict[str, Any] = { + "KeyConditionExpression": Key("PK").eq(f"AST#{assistant_id}") + & Key("SK").begins_with("DOC#"), + } + while True: + response = table.query(**kwargs) + items.extend(response.get("Items") or []) + last = response.get("LastEvaluatedKey") + if not last: + return items + kwargs["ExclusiveStartKey"] = last + + +def _partition_documents( + items: List[Dict[str, Any]], +) -> Tuple[int, List[DocumentNotCarried]]: + """Split into (count carried, descriptions of those not carried).""" + carried = 0 + stranded: List[DocumentNotCarried] = [] + for item in items: + issue = classify_document(item) + if issue is None: + carried += 1 + else: + stranded.append(issue) + return carried, stranded + + +# ── Status ─────────────────────────────────────────────────────────────────── +def _progress_of(record: Dict[str, Any]) -> Optional[UpgradeProgress]: + stored = record.get("migrationProgress") or {} + if not stored: + return None + return UpgradeProgress( + completed=int(stored.get("migrated") or 0), + total=int(stored.get("total") or 0), + skipped=int(stored.get("skipped") or 0), + ) + + +#: Failure reasons in the user's language. The stored ``migrationError`` is +#: written for an operator; rendering it raw is how a user ends up reading +#: "ByteCapExceeded". +_FAILURE_COPY = { + "ByteCapExceeded": ( + "Your knowledge base is larger than the current upgrade size limit, so " + "the upgrade stopped before changing anything." + ), + "VerificationFailed": ( + "The upgraded copy did not return the same results as your current one, " + "so it was discarded rather than switched over." + ), +} + +_FAILURE_FALLBACK = ( + "Something went wrong part-way through the upgrade, so it was stopped and " + "nothing was changed." +) + + +def _failure_reason(record: Dict[str, Any]) -> str: + stored = str(record.get("migrationError") or "") + for token, copy in _FAILURE_COPY.items(): + if token in stored: + return copy + return _FAILURE_FALLBACK + + +async def get_upgrade_status( + assistant_id: str, + *, + can_edit: bool, +) -> UpgradeStatusResponse: + """Derive everything the card renders. + + ``can_edit`` only ever *removes* the control (Requirement 23.7). A viewer + still gets an honest phase — they may legitimately see that an upgrade is + running — but never ``canUpgrade``. + """ + import asyncio + + from apis.shared.kb_backend import records as r + + record = await asyncio.to_thread(r.get_kb_record, assistant_id, assistant_id) + state = str((record or {}).get("migrationState") or "") + + if record and r.resolve_engine(record) == r.ENGINE_MANAGED: + # Already upgraded. The only thing owed is the one-time notice, and only + # until it is dismissed — never a permanent badge (Requirement 23.4). + pending = not record.get("upgradeNoticeDismissedAt") + return UpgradeStatusResponse( + phase="succeeded", + canUpgrade=False, + noticePending=bool(pending and can_edit), + progress=_progress_of(record), + ) + + if state in (r.SHADOW, r.VERIFY, r.PROMOTE): + return UpgradeStatusResponse( + phase="in_progress", + canUpgrade=False, + progress=_progress_of(record or {}), + ) + + if state == r.MIGRATION_FAILED: + # Still on the legacy backend, which keeps working. Retry is offered to + # editors; the phase itself is not hidden (Requirement 23.5). + return UpgradeStatusResponse( + phase="failed", + canUpgrade=can_edit, + reason=_failure_reason(record or {}), + progress=_progress_of(record or {}), + ) + + if not (can_edit and migration_enabled()): + return UpgradeStatusResponse(phase="none", canUpgrade=False) + + items = await asyncio.to_thread(_document_items, assistant_id) + if not items: + # An empty knowledge base has nothing to carry across, so there is no + # action to take and therefore nothing to show (Requirement 23.1). + return UpgradeStatusResponse(phase="none", canUpgrade=False) + + carried, stranded = _partition_documents(items) + if not carried: + # Every document is already stranded. Offering an upgrade that would + # carry nothing is worse than useless, but the stranded list is exactly + # what this owner needs to see (Requirement 21.3). + return UpgradeStatusResponse( + phase="none", + canUpgrade=False, + documentsNotCarried=stranded, + ) + + return UpgradeStatusResponse( + phase="available", + canUpgrade=True, + progress=UpgradeProgress(completed=0, total=carried, skipped=len(stranded)), + documentsNotCarried=stranded, + ) + + +# ── Enrolment ──────────────────────────────────────────────────────────────── +async def enroll( + assistant_id: str, + *, + owner_user_id: str, + visibility: str = "PRIVATE", +) -> EnrollResponse: + """Move this knowledge base into ``shadow``, or report one already running. + + Idempotent by construction: both writes are conditional, so a double-click + produces one migration and one "already running" answer rather than two + provisioning sagas racing over the same corpus. + """ + import asyncio + + from apis.shared.kb_backend import records as r + + if not migration_enabled(): + raise UpgradeUnavailable( + "Upgrades are not being accepted at the moment. Nothing has changed." + ) + + record = await asyncio.to_thread(r.get_kb_record, assistant_id, assistant_id) + + if record and r.resolve_engine(record) == r.ENGINE_MANAGED: + return EnrollResponse( + phase="succeeded", + started=False, + message="This knowledge base has already been upgraded.", + ) + + state = str((record or {}).get("migrationState") or "") + if state in (r.SHADOW, r.VERIFY, r.PROMOTE): + return EnrollResponse( + phase="in_progress", + started=False, + message="The upgrade is already running.", + ) + + generation = int((record or {}).get("migrationGeneration") or 0) + + if record is None: + # Zero-backfill: legacy knowledge bases have no KB record at all, so + # enrolment is where the record first comes into existence. + fresh = r.KbRecord( + app_kb_id=assistant_id, + owner_user_id=owner_user_id, + visibility=visibility, + provisioning_state=r.PROVISIONING, + ) + try: + await asyncio.to_thread(r.create_provisioning, assistant_id, fresh) + except r.TransitionLost: + # Another request created it between our read and our write. Not an + # error: fall through and let the state transition arbitrate. + logger.info( + f"kb {assistant_id}: record created concurrently during enrolment" + ) + record = await asyncio.to_thread(r.get_kb_record, assistant_id, assistant_id) + generation = int((record or {}).get("migrationGeneration") or 0) + + due_at = _iso(_now() + _DUE_IMMEDIATELY) + try: + await asyncio.to_thread( + r.set_migration_state, + assistant_id, + assistant_id, + r.SHADOW, + generation, + due_at=due_at, + ) + except r.TransitionLost: + return EnrollResponse( + phase="in_progress", + started=False, + message="The upgrade is already running.", + ) + + logger.info(f"kb {assistant_id}: enrolled into shadow at generation {generation}") + return EnrollResponse( + phase="in_progress", + started=True, + message=( + "Upgrade started. Your knowledge base keeps working while it runs, " + "and you can leave this page." + ), + ) + + +async def retry(assistant_id: str, *, owner_user_id: str) -> EnrollResponse: + """Re-enter ``shadow`` from ``failed``, on a fresh generation. + + The generation bump is what makes the retry safe: every conditional write + belonging to the abandoned attempt is guarded on the old generation, so a + straggler worker from the failed run cannot land a write on the new one. + """ + import asyncio + + from apis.shared.kb_backend import records as r + + if not migration_enabled(): + raise UpgradeUnavailable( + "Upgrades are not being accepted at the moment. Nothing has changed." + ) + + record = await asyncio.to_thread(r.get_kb_record, assistant_id, assistant_id) + if record is None: + return await enroll(assistant_id, owner_user_id=owner_user_id) + + state = str(record.get("migrationState") or "") + if state != r.MIGRATION_FAILED: + # Nothing to retry. Report the truth rather than starting a second run. + if state in (r.SHADOW, r.VERIFY, r.PROMOTE): + return EnrollResponse( + phase="in_progress", + started=False, + message="The upgrade is already running.", + ) + if r.resolve_engine(record) == r.ENGINE_MANAGED: + return EnrollResponse( + phase="succeeded", + started=False, + message="This knowledge base has already been upgraded.", + ) + return await enroll( + assistant_id, + owner_user_id=owner_user_id, + visibility=str(record.get("visibility") or "PRIVATE"), + ) + + generation = int(record.get("migrationGeneration") or 0) + try: + await asyncio.to_thread( + r.retry_from_failed, + assistant_id, + assistant_id, + generation, + _iso(_now() + _DUE_IMMEDIATELY), + ) + except r.TransitionLost: + # A concurrent retry got there first. Its attempt is running, so this is + # a success from the user's point of view. + return EnrollResponse( + phase="in_progress", + started=False, + message="The upgrade is already running.", + ) + logger.info(f"kb {assistant_id}: retried into generation {generation + 1}") + return EnrollResponse( + phase="in_progress", + started=True, + message=( + "Upgrade restarted. Your knowledge base keeps working while it runs." + ), + ) + + +async def dismiss_notice(assistant_id: str) -> None: + """Retire the one-time success notice (Requirement 23.4).""" + import asyncio + + from apis.shared.kb_backend import records as r + + try: + await asyncio.to_thread( + r.dismiss_upgrade_notice, assistant_id, assistant_id, _iso(_now()) + ) + except r.TransitionLost: + # No record, so no notice to dismiss. Nothing owed, nothing to report. + logger.info(f"kb {assistant_id}: notice dismissal for a record that is absent") diff --git a/backend/src/apis/app_api/main.py b/backend/src/apis/app_api/main.py index c47c9978e..cb9349532 100644 --- a/backend/src/apis/app_api/main.py +++ b/backend/src/apis/app_api/main.py @@ -192,6 +192,7 @@ async def lifespan(app: FastAPI): from apis.app_api.assistants.routes import router as assistants_router from apis.app_api.agent_designer.routes import router as agents_router from apis.app_api.documents.routes import router as documents_router +from apis.app_api.kb_upgrade.routes import router as kb_upgrade_router from apis.app_api.users.routes import router as users_router from apis.app_api.user_settings.routes import router as user_settings_router from apis.app_api.connectors.routes import router as connectors_router @@ -217,6 +218,7 @@ async def lifespan(app: FastAPI): app.include_router(assistants_router) app.include_router(agents_router) # Agent Designer /agents surface; 404s while AGENTS_API_ENABLED off app.include_router(documents_router) +app.include_router(kb_upgrade_router) # Owner-facing KB upgrade card; phase "none" (renders nothing) while MANAGED_KB_MIGRATION_ENABLED is off app.include_router(users_router) app.include_router(user_settings_router) app.include_router(models_router) diff --git a/backend/src/apis/inference_api/chat/models.py b/backend/src/apis/inference_api/chat/models.py index 6172b6298..acb5136f2 100644 --- a/backend/src/apis/inference_api/chat/models.py +++ b/backend/src/apis/inference_api/chat/models.py @@ -117,6 +117,23 @@ class InvocationRequest(BaseModel): # `get_assistant_with_access_check` still gates the Agent itself, so the # worst a forged flag can do is decline to persist a binding. agent_mention: Optional[bool] = None + # Marketplace D2: this turn is a marketplace **reviewer** test-driving a submission + # before deciding on it. Two things change, and nothing else does. + # + # 1. The Agent resolves to the snapshot under review (`submittedVersion` while one is + # pending, `publishedVersion` otherwise) rather than to the published-or-draft rule + # every other caller gets. A reviewer who test-drove the author's live draft would + # be testing something approval is not going to publish. + # 2. The PRIVATE access check is bypassed, because a PRIVATE Agent can be — and often + # is — sitting in the review queue, and `get_assistant_with_access_check` refuses a + # non-owner outright on one. + # + # ⚠️ Unlike `agent_mention`, this is NOT merely a claim about intent, so it cannot be + # treated like one: it widens access. The route re-checks `admin.marketplace` against + # the caller's own roles before honoring it, and a caller without the scope gets a 403 + # rather than a quietly-ignored flag — a silently downgraded preview would run the + # wrong configuration and report it as the reviewed one. + review_preview: Optional[bool] = None # When set, the route resumes a paused agent turn instead of starting a # new one. `message` is ignored in that case — the original prompt is # already in the agent's interrupt context. diff --git a/backend/src/apis/inference_api/chat/routes.py b/backend/src/apis/inference_api/chat/routes.py index 91ff8d71a..f1a5c2c3e 100644 --- a/backend/src/apis/inference_api/chat/routes.py +++ b/backend/src/apis/inference_api/chat/routes.py @@ -1542,6 +1542,7 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g if input_data.rag_assistant_id and not is_resume and not is_continuation: # Local imports to avoid circular dependency + from apis.shared.assistants.kb_access import granted from apis.shared.assistants.rag_service import ( augment_prompt_with_context, search_assistant_knowledgebase_with_formatting, @@ -1553,6 +1554,7 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g from apis.shared.assistants.version_resolution import ( AgentVersionUnavailableError, resolve_invocation_agent, + resolve_review_agent, ) from apis.shared.sessions.messages import get_messages from apis.shared.sessions.metadata import ( @@ -1624,13 +1626,70 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g ) # 2. Load assistant with access check - logger.info("Loading assistant with access check...") - assistant, _ = await get_assistant_with_access_check( - assistant_id=input_data.rag_assistant_id, user_id=user_id, user_email=current_user.email - ) + # + # ⚠️ The reviewer preview is the ONE path that does not go through the visibility + # gate, and it earns that by proving the scope first. A PRIVATE Agent can be sitting + # in the review queue — approval refuses one, but submission does not — and + # ``get_assistant_with_access_check`` refuses a non-owner outright on PRIVATE, so a + # reviewer could not test-drive exactly the submissions most worth testing. + # + # The scope is re-resolved here against the caller's own roles rather than trusted + # from the request, and a caller without it is refused rather than quietly demoted + # to an ordinary turn: a silent downgrade would run the published snapshot (or the + # author's draft) and report it to the reviewer as the version under review. + is_review_preview = bool(input_data.review_preview) + if is_review_preview: + import os + + from apis.shared.auth import has_admin_scope + from apis.shared.assistants.service import ( + _get_assistant_cloud_without_ownership_check, + ) + + if not await has_admin_scope(current_user, "admin.marketplace"): + logger.warning("review_preview requested without the marketplace scope") + raise HTTPException( + status_code=403, + detail="Reviewing an agent requires marketplace admin access.", + ) + table_name = os.environ.get("DYNAMODB_ASSISTANTS_TABLE_NAME") + if not table_name: + # Explicit, because the alternative is a confusing failure several frames + # down inside boto3 rather than a message naming the missing variable. + raise RuntimeError( + "DYNAMODB_ASSISTANTS_TABLE_NAME environment variable is required" + ) + assistant = await _get_assistant_cloud_without_ownership_check( + input_data.rag_assistant_id, table_name + ) + # No permission was resolved, because this path deliberately bypassed the + # gate that resolves one. Left as None so the knowledge base read below + # fails closed (Requirement 25.1): `granted(...)` treats None as "grants + # nothing", and the marketplace scope is authority to *review a + # submission*, not evidence of a read grant on that owner's corpus. + # + # ⚠️ Consequence worth owning: a reviewer test-drives with an empty + # knowledge base, which is a degraded review of a RAG-backed Agent — the + # same class of problem develop's comment above warns about. Whether a + # marketplace reviewer should receive corpus read is a policy decision for + # the marketplace owner, not something to settle inside a merge conflict. + # Tracked rather than guessed; failing closed is the safe default meanwhile. + assistant_permission = None + else: + logger.info("Loading assistant with access check...") + # The permission is kept, not discarded: it is what the knowledge base + # retrieval below runs under (Requirement 25.1), so the grant that governs + # the corpus read is provably the same one that admitted this turn. + assistant, assistant_permission = await get_assistant_with_access_check( + assistant_id=input_data.rag_assistant_id, + user_id=user_id, + user_email=current_user.email, + ) if not assistant: - logger.warning("get_assistant_with_access_check returned None") + logger.warning( + "Assistant lookup returned None (review_preview=%s)", is_review_preview + ) # Check if assistant exists at all to provide better error message from apis.shared.assistants.service import assistant_exists @@ -1667,7 +1726,14 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g # ⚠️ Ordered *before* the access check's side effects below on purpose: it is not an # access decision and must not be read as one. The caller was already admitted. try: - assistant, resolved_version = await resolve_invocation_agent(assistant, user_id) + if is_review_preview: + # The submitted snapshot, not the published one and not the author's draft — + # the artifact the reviewer's decision is actually about. Shares its rule + # with the admin submission read, so the page and the test drive can never + # disagree about what is under review. + assistant, resolved_version = await resolve_review_agent(assistant) + else: + assistant, resolved_version = await resolve_invocation_agent(assistant, user_id) except AgentVersionUnavailableError as unavailable: # A published Agent whose snapshot is missing fails the turn rather than # falling back to the draft — the fallback would serve unreviewed instructions @@ -1688,22 +1754,31 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g f"published version {resolved_version}" if resolved_version else "live record / draft", ) - # Mark as viewed if this is a shared assistant (not owned) - if assistant.owner_id != user_id: - await mark_share_as_interacted(assistant_id=input_data.rag_assistant_id, user_email=current_user.email) - - # KB sync inactivity signal: any user's chat use counts. Throttled - # to one write/day inside bump_last_used_at (conditional update); - # the winning bump also wakes any inactivity-paused sync policies. - # Best-effort — a bookkeeping failure must never break a chat turn. - try: - from apis.shared.assistants.service import bump_last_used_at - from apis.shared.sync_policies.service import resume_inactive_policies + # ⚠️ Both bookkeeping writes below are skipped for a reviewer preview, because a + # review is not use. ``mark_share_as_interacted`` would stamp a share record the + # reviewer does not have, and ``bump_last_used_at`` feeds the KB-sync inactivity + # pause — a reviewer poking a submission once would read as the Agent being live + # and wake sync policies that were correctly dormant. Persistence of the turn + # itself is already handled by the ``preview-`` session id the reviewer's client + # sends (``PREVIEW_SESSION_PREFIX``), which is what keeps the test drive out of the + # author's conversation history. + if not is_review_preview: + # Mark as viewed if this is a shared assistant (not owned) + if assistant.owner_id != user_id: + await mark_share_as_interacted(assistant_id=input_data.rag_assistant_id, user_email=current_user.email) + + # KB sync inactivity signal: any user's chat use counts. Throttled + # to one write/day inside bump_last_used_at (conditional update); + # the winning bump also wakes any inactivity-paused sync policies. + # Best-effort — a bookkeeping failure must never break a chat turn. + try: + from apis.shared.assistants.service import bump_last_used_at + from apis.shared.sync_policies.service import resume_inactive_policies - if await bump_last_used_at(input_data.rag_assistant_id): - await resume_inactive_policies(input_data.rag_assistant_id) - except Exception as bump_err: - logger.warning(f"lastUsedAt bump failed for assistant {input_data.rag_assistant_id}: {bump_err}") + if await bump_last_used_at(input_data.rag_assistant_id): + await resume_inactive_policies(input_data.rag_assistant_id) + except Exception as bump_err: + logger.warning(f"lastUsedAt bump failed for assistant {input_data.rag_assistant_id}: {bump_err}") # 2b. Agent Designer Phase 3 — resolve the Agent's governed capabilities # for the INVOKING user (D5), before the expensive KB search. v1 blocks @@ -1737,7 +1812,10 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g try: logger.info("Searching knowledge base for assistant...") context_chunks = await search_assistant_knowledgebase_with_formatting( - assistant_id=input_data.rag_assistant_id, query=input_data.message, top_k=5 + assistant_id=input_data.rag_assistant_id, + query=input_data.message, + top_k=5, + access=granted(input_data.rag_assistant_id, user_id, assistant_permission), ) logger.info(f"Knowledge base search returned {len(context_chunks) if context_chunks else 0} chunks") if context_chunks: diff --git a/backend/src/apis/shared/assistants/kb_access.py b/backend/src/apis/shared/assistants/kb_access.py new file mode 100644 index 000000000..68662405f --- /dev/null +++ b/backend/src/apis/shared/assistants/kb_access.py @@ -0,0 +1,237 @@ +"""Who may read a knowledge base — resolved before retrieval is attempted. + +Requirement 25.1–25.3. The application is the authorization authority for +knowledge base reads. Neither of the two mechanisms Bedrock offers is trusted to +be that authority: + +* **Metadata filters** give what AWS's own multi-tenant guidance calls + "filter-level (logical) isolation, *not* IAM-enforced (infrastructure) + isolation". A filter is a query argument; anything that can issue a query can + omit it. +* **ACL-aware retrieval** fails closed, which is better than the document-status + filter used to be, but AWS states plainly that it "is not authorization" and + does not authenticate users. Its identity is **email only, with no alias + resolution, and a mismatch fails silently**. On a platform that authenticates + via OIDC with claim mappings, a silently-failing email comparison is a worse + primitive than an explicit check against the permission model that already + governs the agent. + +So the check happens here, in the application, on the way in. + +Why this lives in ``apis.shared.assistants`` and not in ``kb_backend`` +--------------------------------------------------------------------- +It reuses :func:`apis.shared.assistants.service.resolve_assistant_permission` +rather than introducing a parallel permission model (Requirement 25.2), and +``kb_backend`` may not import the assistants package — that boundary is what keeps +the migration Lambda images small, and it is enforced by +``tests/architecture/test_kb_backend_boundary.py``. Authorization is also +*above* the seam by nature: it is the same answer whichever engine serves the +query, so implementing it once above both adapters is the only way it cannot +differ between them. + +Why the facade takes a resolved grant rather than a user +------------------------------------------------------- +Both production callers resolve the invoking user's permission a few lines before +they retrieve — ``inference_api/chat/routes.py`` via +``get_assistant_with_access_check``, ``app_api/assistants/routes.py`` via +``resolve_assistant_permission``. Re-resolving inside the facade would add a +second DynamoDB read per turn to answer a question the caller has already +answered. + +Passing the answer instead is not weaker, because the parameter is **required and +keyword-only**: a caller that forgets it raises ``TypeError`` at the call site, +which no test suite can miss, while a caller that genuinely has no grant passes +``None`` and gets nothing back. Trusting a caller-supplied *string* would be +weaker; :class:`KbAccess` cannot be constructed with a permission outside +:data:`KB_READ_PERMISSIONS`, so "I have a grant object" is not something a caller +can assert without having gone through :func:`granted` or +:func:`resolve_kb_access`. + +The 1:1 binding is what makes this simple +----------------------------------------- +This phase holds ``App_KB_Id == assistant_id``, so an agent's knowledge base is +exactly that agent's own and "may this user invoke this agent" already answers +"may this turn retrieve". There is deliberately no handling for the 0..N case — +whether one inaccessible knowledge base among several should fail the whole turn +is F4's question, and answering it here would bake in a guess. + +Feature: managed-kb-migration +Requirements: 25.1, 25.2, 25.3 +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Optional + +logger = logging.getLogger(__name__) + +#: Permissions that may read a knowledge base through its agent. A viewer reads +#: (that is what sharing an agent is *for*) but never sees the upgrade control, +#: which is why the two sets below are separate rather than one ranked scale. +KB_READ_PERMISSIONS = frozenset({"owner", "editor", "viewer"}) + +#: Permissions that may change a knowledge base — upload, delete, or trigger a +#: migration. Deliberately a strict subset: an engine upgrade spends money and +#: mutates the corpus, so it is an owner/editor act. +KB_WRITE_PERMISSIONS = frozenset({"owner", "editor"}) + + +class KbAccessDenied(PermissionError): + """The invoking user may not read this knowledge base. + + Raised only by callers that want an error; :func:`granted` and + :func:`resolve_kb_access` return ``None`` instead, because the retrieval path + turns a denial into "no context" rather than into a failed turn. + """ + + +@dataclass(frozen=True) +class KbAccess: + """A resolved grant to read one knowledge base. + + Frozen, and only ever produced by :func:`granted` or + :func:`resolve_kb_access`, so its existence *is* the statement that the + permission model was consulted. Holding a permission string proves nothing; + holding one of these does. + """ + + assistant_id: str + app_kb_id: str + user_id: str + permission: str + + @property + def may_read(self) -> bool: + """True for every instance. Kept as a named predicate rather than an + implicit invariant so a call site reads as a check, and so the day a + write-only grant is added there is somewhere for it to go.""" + return self.permission in KB_READ_PERMISSIONS + + @property + def may_upgrade(self) -> bool: + """Whether this grant may trigger a migration for the knowledge base.""" + return self.permission in KB_WRITE_PERMISSIONS + + +def granted( + assistant_id: str, + user_id: str, + permission: Optional[str], + app_kb_id: Optional[str] = None, +) -> Optional[KbAccess]: + """Wrap an already-resolved permission as a grant, or ``None`` if it grants nothing. + + For the callers that resolved the permission themselves a moment earlier. + ``None``, an empty string, and any unrecognized value all return ``None``: + an unknown permission written by newer code is not evidence of access, and + guessing in the permissive direction is how a viewer-shaped bug becomes a + disclosure. + + ``app_kb_id`` defaults to ``assistant_id`` — the 1:1 binding this phase keeps. + """ + if not permission or permission not in KB_READ_PERMISSIONS: + logger.warning( + f"knowledge base access denied for user {user_id} on assistant " + f"{assistant_id}: permission {permission!r} does not grant read" + ) + return None + + return KbAccess( + assistant_id=assistant_id, + app_kb_id=app_kb_id or assistant_id, + user_id=user_id, + permission=permission, + ) + + +async def resolve_kb_access( + assistant_id: str, + user_id: str, + user_email: Optional[str] = None, + app_kb_id: Optional[str] = None, +) -> Optional[KbAccess]: + """Resolve a grant from the assistant permission model, failing closed. + + For callers that do not already hold a permission. Delegates to + ``resolve_assistant_permission`` — the same function the document routes and + the listing service gate on — so owner/editor/viewer semantics are whatever + that function says they are and cannot drift here (Requirement 25.2). + + **Any** failure denies: a missing table, an unreachable table, a malformed + record. This is the opposite of the resolver's choice to treat an unreadable + KB_Record as legacy, and deliberately so. There, both answers serve the user's + own documents and one of them is always safe. Here the two answers are "your + documents" and "someone else's", and an error tells us which is which is + exactly what we do not know (Requirement 24.6). + """ + # Function-local: importing ``.service`` at module scope would run inside the + # package ``__init__``'s own import of ``rag_service``, and the ordering that + # makes that work today is not a property worth depending on. + from apis.shared.assistants.service import resolve_assistant_permission + + try: + _assistant, permission = await resolve_assistant_permission( + assistant_id=assistant_id, user_id=user_id, user_email=user_email + ) + except Exception as exc: + logger.error( + f"knowledge base access check failed for user {user_id} on assistant " + f"{assistant_id}; denying because access cannot be confirmed: {exc}", + exc_info=True, + ) + return None + + return granted(assistant_id, user_id, permission, app_kb_id) + + +async def is_shared_beyond_owner(assistant_id: str, owner_id: str, visibility: Optional[str] = None) -> bool: + """Whether anyone other than the owner can reach this assistant's documents. + + The input to Requirement 25.6's "where a knowledge base is shared beyond its + owner". Lives here rather than in ``kb_backend`` for the same reason the access + check does: it is an application fact, and the seam may not read it. + + Two independent sources, because either alone under-detects: + + * **Visibility.** ``PUBLIC`` shares with everyone and ``SHARED`` announces the + intent, so neither is owner-only. + * **Share records.** A ``PRIVATE`` assistant can still carry explicit shares — + ``resolve_assistant_permission`` resolves an editor share on a private + assistant to ``editor``. Judging by visibility alone would leave those + knowledge bases unprotected, which is the case most likely to exist and least + likely to be noticed. + + Fails **shared** on error. The two answers are "apply a narrowing policy that + was not strictly needed" and "leave a multi-user corpus reachable by anything + in the account with a wildcard grant"; the first costs one control-plane call. + """ + if visibility in ("PUBLIC", "SHARED"): + return True + + from apis.shared.assistants.service import list_assistant_shares + + try: + shares = await list_assistant_shares(assistant_id, owner_id) + except Exception as exc: + logger.error( + f"could not determine whether assistant {assistant_id} is shared; " + f"assuming it is, so that a resource policy is applied rather than " + f"skipped: {exc}", + exc_info=True, + ) + return True + + return bool(shares) + + +__all__ = [ + "KB_READ_PERMISSIONS", + "KB_WRITE_PERMISSIONS", + "KbAccess", + "KbAccessDenied", + "granted", + "is_shared_beyond_owner", + "resolve_kb_access", +] diff --git a/backend/src/apis/shared/assistants/kb_publication.py b/backend/src/apis/shared/assistants/kb_publication.py new file mode 100644 index 000000000..bae24dbaf --- /dev/null +++ b/backend/src/apis/shared/assistants/kb_publication.py @@ -0,0 +1,123 @@ +"""What publication means for a knowledge base's lifecycle. + +Requirements 25.8–25.11. Three positions, and one question deliberately left open. + +**An engine migration is not a corpus change (25.8).** Parity is the entire +contract of this migration: the same documents, the same ``top_k``, the same +context cap, the same answer model. So swapping which engine serves a published +agent's knowledge base does not change what that agent retrieves and therefore +needs no re-review. :func:`migration_requires_review` says so in one place, with a +test, rather than leaving it as an assumption spread across the worker. + +**A listed agent's knowledge base is exempt from reclaim (25.9).** Nothing reclaims +in this phase — ``reclaim`` is reserved in the state enum and never entered — so +this is a guard placed before the mechanism that will need it. Written now because +the follow-up spec's eviction pass will be the first thing to delete a corpus, and +"is this on the store shelf right now?" is not a question it should be answering +for the first time under a deadline. + +**A takedown must be walked, not fallen into (25.10).** Reclaim eligibility is +computed from the listing state as it is, and ``taken_down`` is reached only by the +listing state machine's explicit edge. Nothing here infers a takedown from a +missing listing, an expired timestamp, or an error. + +**Corpus-revision pinning stays open (25.11).** 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. That is a real review bypass and this +module does not pretend to close it: exemption from cleanup is not revision +pinning. The question belongs to the marketplace spec. + +Why ``is_on_shelf`` and not ``is_listed`` +---------------------------------------- +The listing module documents the trap and it applies exactly here: an admin +requesting changes on a *live* listing leaves it serving but moves its state to +``changes_requested``, which is not in ``LISTED_STATES``. Asked by state name +alone, such an agent reads as unlisted — so a reclaim pass would delete the corpus +behind an agent users can still see in the store. ``is_on_shelf`` asks the fact +(is a version of this queryable in the store right now) rather than the state name, +and that is the only correct question for a destructive pass. + +Feature: managed-kb-migration +Requirements: 25.8, 25.9, 25.10, 25.11 +""" + +from __future__ import annotations + +import logging +from typing import Any, Mapping, Optional + +from apis.shared.assistants.listing import is_on_shelf + +logger = logging.getLogger(__name__) + + +def migration_requires_review(from_engine: Optional[str], to_engine: Optional[str]) -> bool: + """Whether changing engines needs a listed agent to be reviewed again. + + Always ``False``, and a function rather than a comment so the claim is + something a test can hold. An engine swap moves the same documents to a + different index; the reviewed artefact — instructions, bindings, the corpus + itself — is untouched. If a future change makes an engine swap alter retrieval + results, this must stop returning ``False``, and the test that pins it is where + that argument has to be had. + """ + return False + + +def is_reclaim_exempt( + kb_record: Optional[Mapping[str, Any]], + listing_state: Optional[str] = None, + published_version: Optional[int] = None, +) -> bool: + """Whether this knowledge base must be left alone by a lifecycle reclaim pass. + + Exempt when any of these holds: + + * the record carries ``exemptFromReclaim`` — an operator's explicit hold; + * the record is ``pinned``; + * the agent is on the store shelf right now. + + A **missing** record — ``None``, which is what ``get_kb_record`` returns when + it cannot find one — is exempt. Reclaim acts on knowledge bases it can + describe, and "I could not read the record" is not a description; the same + fail-closed reasoning the access check uses, applied to deletion, where it + matters more. An *empty* mapping is a different thing: a record that was read + and carries no holds. Treating the two alike would exempt every unheld + knowledge base and make the whole predicate vacuous. + """ + if kb_record is None: + return True + + if kb_record.get("exemptFromReclaim") or kb_record.get("pinned"): + return True + + return is_on_shelf(listing_state, published_version) + + +def reclaim_exemption_reason( + kb_record: Optional[Mapping[str, Any]], + listing_state: Optional[str] = None, + published_version: Optional[int] = None, +) -> Optional[str]: + """Why this knowledge base is exempt, or ``None`` if it is not. + + For the report-only output of a reclaim pass. A pass that logs "skipped 400 + knowledge bases" tells an operator nothing they can act on; one that says which + are on the shelf and which an operator pinned by hand does. + """ + if kb_record is None: + return "no KB_Record could be read" + if kb_record.get("exemptFromReclaim"): + return "exemptFromReclaim is set on the record" + if kb_record.get("pinned"): + return "the knowledge base is pinned" + if is_on_shelf(listing_state, published_version): + return f"the agent is on the store shelf (listing state {listing_state!r})" + return None + + +__all__ = [ + "is_reclaim_exempt", + "migration_requires_review", + "reclaim_exemption_reason", +] diff --git a/backend/src/apis/shared/assistants/listing.py b/backend/src/apis/shared/assistants/listing.py index 571cb789b..8dd322961 100644 --- a/backend/src/apis/shared/assistants/listing.py +++ b/backend/src/apis/shared/assistants/listing.py @@ -52,6 +52,7 @@ "changes_requested", "taken_down", "withdrawal_requested", + "rejected", ) # The transition table. ``None`` is the pre-state of a record that has never been @@ -69,6 +70,9 @@ # which returns the record to what was already serving # in_review → changes_requested admin requests changes, with a reason — or the author # cancels an update submitted while one was outstanding +# in_review → rejected admin declines it for the store, with a reason +# rejected → in_review author revises and submits again +# rejected → private author shelves a declined agent (so it can be deleted) # published → taken_down admin delists, with a reason # published → changes_requested admin requests changes on a live listing # taken_down → changes_requested admin annotates an already-delisted listing @@ -138,6 +142,32 @@ # snapshot keeps serving; approval of the new one is still the only thing that changes what # users get. What it does open is the reverse edge — see ``author_cancel_target``. # +# ⚠️ ``rejected`` is the answer to a submission that should not be in the store *at all*, +# and it exists because the only two decisions before it were "approve" and "request +# changes". An admin who thought a submission did not belong had to either publish it or +# say "fix this" — which promises a review they do not intend to give, leaves the author +# revising toward an approval that is not coming, and leaves the admin re-reading the same +# submission every time it comes back. +# +# **The difference from ``changes_requested`` is intent, not mechanics, and that is +# deliberate.** Both carry a required reason and both let the author come back +# (``rejected → in_review``). What differs is what the author is told: "I want this, fix +# X" versus "this is not a fit, here is why". Making the second one *terminal* was the +# alternative and it is a much bigger hammer — it needs an appeal path, an admin escape +# hatch, and a policy about who may grant one. None of that is worth building before +# someone needs it, and an honest "no" that the author can answer is worth having now. +# +# It is not a door into the store and cannot become one: the only way back is +# ``in_review``, so approval is still the single edge that publishes. ``rejected → private`` +# is the author shelving it, which is what makes it deletable — the same exit +# ``taken_down`` has, for the same reason. +# +# ⚠️ It must also be in ``is_on_shelf``'s never-listed set. A rejected listing has no +# ``published_version`` to clear (it was never published), so a stale pointer is not the +# risk; the risk is the *other* direction — without it, a rejected agent that somehow +# carried a pointer would read as on-the-shelf and route a withdrawal into an admin queue +# for something no user can see. +# # ⚠️ ``changes_requested`` covers two different listings — one that was never published, and # one that *was* and is still serving while the author revises it (``review_listing`` # deliberately does not unpublish). Only the first may walk ``→ private`` alone; the second @@ -154,11 +184,12 @@ ALLOWED_TRANSITIONS: Dict[Optional[str], Set[str]] = { None: {"in_review"}, "private": {"in_review"}, - "in_review": {"published", "changes_requested", "private"}, + "in_review": {"published", "changes_requested", "private", "rejected"}, "changes_requested": {"in_review", "private", "withdrawal_requested"}, "published": {"in_review", "taken_down", "changes_requested", "withdrawal_requested"}, "taken_down": {"in_review", "changes_requested", "private"}, "withdrawal_requested": {"private", "published", "changes_requested", "taken_down"}, + "rejected": {"in_review", "private"}, } # States a submission can be made *from* while the listing is already on the shelf — and @@ -360,7 +391,7 @@ def is_on_shelf(state: Optional[str], published_version: Optional[int]) -> bool: ``state`` is still consulted so a cleared-but-stale pointer cannot resurrect something: ``private`` and ``taken_down`` are never on the shelf whatever the pointer says. """ - if state in ("private", "taken_down", None): + if state in ("private", "taken_down", "rejected", None): return False return published_version is not None diff --git a/backend/src/apis/shared/assistants/models.py b/backend/src/apis/shared/assistants/models.py index e0349a6e5..cbc47ea9d 100644 --- a/backend/src/apis/shared/assistants/models.py +++ b/backend/src/apis/shared/assistants/models.py @@ -21,6 +21,7 @@ "changes_requested", "taken_down", "withdrawal_requested", + "rejected", ] PublisherKind = Literal["institution", "department", "individual"] @@ -801,13 +802,25 @@ class ListingPreflightResponse(BaseModel): class ReviewListingRequest(BaseModel): - """Admin approves a submission or returns it with a reason (D2).""" + """Admin approves a submission, returns it with a reason, or declines it (D2). + + ⚠️ ``reject`` is a third decision, not a harsher ``request_changes``. Both return the + submission with a required reason; they differ in what the author is told and in what + the admin is committing to. See ``assistants.listing``'s note on the ``rejected`` state. + """ model_config = ConfigDict(populate_by_name=True) - decision: Literal["approve", "request_changes"] = Field(..., description="Reviewer's decision") + decision: Literal["approve", "request_changes", "reject"] = Field( + ..., description="Reviewer's decision" + ) note: Optional[str] = Field( - None, max_length=2000, description="Reason; required for request_changes, renders on the author's card" + None, + max_length=2000, + description=( + "Reason; required for request_changes and reject, renders on the author's card. " + "A decline with no reason is the one outcome an author cannot act on at all." + ), ) category: Optional[str] = Field(None, description="Optionally recategorize at approval (D12/D13)") publisher_id: Optional[str] = Field( @@ -1072,6 +1085,105 @@ class AdminListingRow(BaseModel): admin_edits: List[AdminEdit] = Field(default_factory=list, alias="adminEdits", description="Admin edit log (D13)") +class AdminSubmissionReview(BaseModel): + """Everything a reviewer needs to decide, read from the **frozen snapshot** (D2). + + The gap this closes: before it, a reviewer had a name, an author, a category and a + collapsible diff. ``instructions`` is gated to owner/editor on the user-facing detail + read, and ``get_assistant_with_access_check`` refuses a non-owner outright on a PRIVATE + Agent — so the person deciding whether to publish something could not read its system + prompt, could not see what it binds, and on a **first** submission saw no content at + all, because a diff against nothing is empty by construction. + + ⚠️ **The snapshot, never the live record, and that is the whole design.** Widening the + user-facing ``GET /agents/{id}`` for admins would have been fewer lines and wrong: it + serves the author's *draft*, which the author can edit between the reviewer reading it + and the reviewer approving it. ``AgentVersion`` exists precisely to close that window + (it is cut at submission, not at approval), and ``review_listing`` promotes + ``submitted_version`` — so reading anything else would show an admin one configuration + and publish another. Reading the snapshot also sidesteps the PRIVATE 403 without + loosening anyone's access gate: this is an admin projection of a frozen artifact, not a + hole in ``get_assistant_with_access_check``. + + Disclosure note: ``instructions`` reaches a marketplace admin here, deliberately. You + cannot review what you cannot read, and the review diff already discloses them on every + *update*; what this changes is that a first submission is no longer the one case where + the reviewer is asked to approve prose nobody showed them. + """ + + model_config = ConfigDict(populate_by_name=True) + + agent_id: str = Field(..., alias="agentId", description="Agent identifier") + name: str = Field(..., description="Display name as of the reviewed snapshot") + description: str = Field(..., description="Summary as of the reviewed snapshot") + tagline: Optional[str] = Field(None, description="Shelf subtitle as of the reviewed snapshot") + instructions: str = Field(..., description="System prompt as of the reviewed snapshot — the thing being reviewed") + starters: List[str] = Field(default_factory=list, description="Conversation starters as of the reviewed snapshot") + emoji: Optional[str] = Field(None, description="Emoji, for the generated icon fallback (D5)") + icon_url: Optional[str] = Field( + None, alias="iconUrl", description="Where to render the icon from; absent → generated gradient" + ) + owner_name: str = Field(..., alias="ownerName", description="Author display name — who to talk to about behavior") + publisher: Optional[PublisherProfile] = Field(None, description="Resolved attribution (D12), display only") + category: str = Field(..., description="Category id") + category_label: Optional[str] = Field( + None, + alias="categoryLabel", + description="Human label for ``category``; ids are not display content (see ``resolve_listing_display``)", + ) + state: ListingState = Field(..., description="Publication state") + capabilities: List[AgentCapability] = Field( + default_factory=list, + description=( + "What the snapshot binds, by name. Resolved against the snapshot's own " + "``bindings`` rather than the live record's, so a tool the author added after " + "submitting is not attributed to the version under review." + ), + ) + model_label: Optional[str] = Field( + None, alias="modelLabel", description="Display name of the pinned model; absent = resolve as today" + ) + review_version: Optional[int] = Field( + None, + alias="reviewVersion", + description=( + "The snapshot this read is of. ``submittedVersion`` while a submission is " + "pending, ``publishedVersion`` otherwise — never 'the latest', which an admin " + "presentation edit (§6.2) could have moved underneath the reviewer." + ), + ) + published_version: Optional[int] = Field( + None, alias="publishedVersion", description="Which snapshot the store is currently serving" + ) + snapshot_unavailable: bool = Field( + False, + alias="snapshotUnavailable", + description=( + "No snapshot backs this read, so the fields above come from the Agent's **live " + "record** — which its author can still change. True for submissions that predate " + "version snapshots.\n\n" + "Reported rather than refused. ``diff_pending_version`` 400s in this case because " + "a diff of one thing is meaningless, but a reviewer looking at a row in their own " + "queue needs *something* to read; answering 'no' would leave them exactly where " + "the empty diff left them. The flag is what lets the page say the content is not " + "frozen instead of implying it is." + ), + ) + submitted_at: Optional[str] = Field(None, alias="submittedAt", description="ISO 8601 submission timestamp") + withdrawal_requested_at: Optional[str] = Field( + None, alias="withdrawalRequestedAt", description="ISO 8601 timestamp of a pending withdrawal request" + ) + reviewed_at: Optional[str] = Field(None, alias="reviewedAt", description="ISO 8601 timestamp of the last review") + review_note: Optional[str] = Field(None, alias="reviewNote", description="Most recent reviewer note") + reachability: ListingReachability = Field( + ..., + description=( + "Who can open this Agent if it is shelved. Carried here as well as on the queue " + "row because this page is where the decision is actually made." + ), + ) + + class AdminListingsResponse(BaseModel): """Rows for the admin Review queue / Listings tables, plus the nav pending count.""" diff --git a/backend/src/apis/shared/assistants/rag_service.py b/backend/src/apis/shared/assistants/rag_service.py index 6c3d76e4f..9e39ff763 100644 --- a/backend/src/apis/shared/assistants/rag_service.py +++ b/backend/src/apis/shared/assistants/rag_service.py @@ -1,28 +1,99 @@ """RAG service for assistant knowledge base search and prompt augmentation -This service handles searching the vector store for assistant-specific -knowledge and augmenting user prompts with retrieved context. +This module is the **facade** over the knowledge base seam. It resolves which +backend serves an assistant's knowledge base, delegates the search, and then +applies the properties that must hold identically on every backend. It contains +no retrieval logic of its own: what an S3 Vectors response looks like now lives +in ``apis.shared.kb_backend.s3vectors_backend``. + +What lives here, and why here +----------------------------- +Four rules sit above the seam rather than inside either adapter, because a rule +implemented twice is a rule that will eventually differ (Requirement 3): + +* **The access check.** No backend is contacted until the invoking user's grant + has been resolved (Requirement 25.1). It is above the seam because the answer + does not depend on the engine, and because Bedrock's own isolation features are + not trusted to be the authority — see ``kb_access``. +* **The document-status filter.** Dropped chunks whose parent document is not + ``complete``. Kept on both backends during parity even though managed + ingestion makes it largely redundant — removing it in the same change that + swaps the engine would make any difference in results unattributable. +* **``top_k`` narrowing.** Applied *after* the status filter, which is the order + the legacy path has always used: filter-then-slice, so an incomplete document + cannot silently shrink a five-chunk answer. +* **The 2,000-character context cap.** ``augment_prompt_with_context``'s + default. Held constant deliberately: the evaluation measured no correctness + change between 2,000 and 20,000 characters, so raising it here would add a + variable to a change whose whole purpose is to hold every variable but one. + +The dual-read pilot +------------------- +Also above the seam, and for the same reason: comparing two backends is not a +thing either backend can do. The facade starts the observational managed read +before awaiting legacy and detaches the comparison afterwards, so a piloted turn +waits exactly as long as an unpiloted one (Requirement 18.5). Legacy is always +what is served. See ``kb_backend.dual_read``. + +Score direction +---------------The seam speaks **relevance** (higher is better). This facade still emits a +``distance`` key (lower is better), derived by exact negation, because +``app_api/assistants/routes.py`` puts that value in an HTTP response body that a +client already reads. The rename stops at the seam; no caller has to change. """ import logging import os -from typing import Any, Dict, List, Set +import time +from typing import Any, Dict, List, Optional, Set import boto3 -from apis.shared.embeddings.bedrock_embeddings import search_assistant_knowledgebase +from apis.shared.assistants.kb_access import KbAccess +from apis.shared.kb_backend.dual_read import schedule_observation, start_managed_read +from apis.shared.kb_backend.idleness import schedule_activity_touch +from apis.shared.kb_backend.metrics import ( + METRIC_ACCESS_DENIED, + METRIC_STATUS_FILTER_FAIL_CLOSED, + emit_count, +) +from apis.shared.kb_backend.protocol import DEFAULT_TOP_K, Chunk, distance_from_relevance +from apis.shared.kb_backend.query_guard import clamp_query +from apis.shared.kb_backend.resolver import load_record, resolve_backend logger = logging.getLogger(__name__) +#: Parity contract (Requirement 3.2): the cap is 2,000 characters on every +#: backend, unchanged from the value the legacy path has always used. Named so +#: that a change to it is a visible change to a constant rather than an edit to a +#: default argument. +MAX_CONTEXT_CHARS = 2000 -async def search_assistant_knowledgebase_with_formatting(assistant_id: str, query: str, top_k: int = 5) -> List[Dict[str, Any]]: + +async def search_assistant_knowledgebase_with_formatting( + assistant_id: str, + query: str, + top_k: int = DEFAULT_TOP_K, + *, + access: Optional[KbAccess], +) -> List[Dict[str, Any]]: """ Search assistant knowledge base and return formatted results + Resolves the knowledge base's backend, delegates the search across the seam, + then applies the parity rules that must hold on every backend: the document + status filter, and ``top_k`` narrowing after it. + Args: assistant_id: Assistant identifier to filter vectors query: User query text top_k: Number of top results to return (default: 5) + access: The invoking user's resolved grant, or ``None`` if they have none. + Required and keyword-only (Requirement 25.1): a caller that forgets it + fails loudly at the call site, while a caller that genuinely has no + grant passes ``None`` and gets nothing. Build one with + ``kb_access.granted`` if the permission is already in hand, or + ``kb_access.resolve_kb_access`` if it is not. Returns: List of dictionaries containing: @@ -30,32 +101,96 @@ async def search_assistant_knowledgebase_with_formatting(assistant_id: str, quer - distance: Similarity distance (lower = more similar) - metadata: Original metadata from vector store - key: Vector key/ID + + Empty when the caller has no grant — no backend is contacted at all. """ - try: - # Call the bedrock_embeddings search function - response = await search_assistant_knowledgebase(assistant_id, query) + # Authorization first, before the backend resolution, the query clamp, and + # any AWS call (Requirement 25.1). Ordering is the requirement: a check that + # runs after retrieval has already read the corpus is an audit log, not an + # access control. + if access is None or not access.may_read: + logger.error( + f"refusing knowledge base retrieval for assistant {assistant_id}: " + f"no resolved access grant" + ) + emit_count(METRIC_ACCESS_DENIED, dimensions={"reason": "no_grant"}) + return [] - # Extract vectors from response - vectors = response.get("vectors", []) + if access.assistant_id != assistant_id: + # A grant for a different assistant is not a grant for this one. This is + # the shape a copy-paste bug takes when a route resolves permission for + # one id and retrieves with another, and while the 1:1 binding holds it is + # the only way the two could disagree. + logger.error( + f"refusing knowledge base retrieval: grant is for assistant " + f"{access.assistant_id}, not {assistant_id}" + ) + emit_count(METRIC_ACCESS_DENIED, dimensions={"reason": "grant_mismatch"}) + return [] - if not vectors: + managed_task = None + try: + # One record read serves both questions: which backend to use, and whether + # this knowledge base is in the dual-read pilot. Reading it here rather + # than letting the resolver read it internally is what keeps the pilot + # from costing an extra DynamoDB round trip on every turn. + record = load_record(assistant_id) + backend = resolve_backend(assistant_id, record=record) + + # Clamp before dispatch, so both backends receive an identically-shaped + # query (Requirement 4.2). Managed KB rejects anything over 10,000 + # characters outright and the quota is not adjustable, so clamping only + # the managed path would make the two backends answer different + # questions and invalidate the dual-read comparison. + query, _ = clamp_query(query) + + # Start the observational read *before* awaiting legacy (Requirement + # 18.5). Nothing is awaited here, so a piloted turn does the same waiting + # as an unpiloted one; managed Retrieve measured 662–695 ms p50 against + # legacy's 257 ms, so awaiting both would nearly triple this leg. + # ``None`` whenever there is no comparison to make. + managed_task = start_managed_read(record, assistant_id, query, top_k) + + started = time.perf_counter() + chunks = await backend.search(assistant_id, query, top_k) + legacy_ms = (time.perf_counter() - started) * 1000.0 + + # Detach the comparison. Legacy is what gets served either way — including + # when it is empty, which is a finding rather than a reason to reach for + # the other engine's answer (Requirement 18.2). + schedule_observation(assistant_id, query, top_k, list(chunks), legacy_ms, managed_task) + managed_task = None + + # Record that this knowledge base was needed (Requirement 22.5), for the + # idleness signal the follow-up spec's eviction threshold has to be chosen + # from — data that cannot be backfilled later. + # + # Only for knowledge bases that have a record. A legacy knowledge base has + # none, and creating one here would break the migration's zero-backfill + # property across 1,692 existing rows for the sake of a metric. Detached and + # throttled, so retrieval waits for neither the write nor its rejection + # (Requirement 22.6). + if record: + schedule_activity_touch(assistant_id, assistant_id) + + if not chunks: logger.info(f"No vectors found for assistant {assistant_id} with query: {query[:50]}...") return [] # Filter out chunks from documents that are not in "complete" status - vectors = _filter_vectors_by_document_status(vectors, assistant_id) + chunks = _filter_chunks_by_document_status(chunks, assistant_id) # Format results - return document_id for on-demand download URL generation formatted_results = [] - for vector in vectors[:top_k]: - metadata = vector.get("metadata", {}) - + for chunk in chunks[:top_k]: formatted_results.append( { - "text": metadata.get("text", ""), - "distance": vector.get("distance"), - "metadata": metadata, - "key": vector.get("key", ""), + "text": chunk.text, + # Derived from relevance by exact negation, so the value a + # caller reads is the one it has always read. + "distance": distance_from_relevance(chunk.relevance), + "metadata": chunk.metadata, + "key": chunk.key, } ) @@ -64,10 +199,41 @@ async def search_assistant_knowledgebase_with_formatting(assistant_id: str, quer except Exception as e: logger.error(f"Error searching knowledge base for assistant {assistant_id}: {e}", exc_info=True) + if managed_task is not None: + # The legacy search raised before the comparison was detached, so + # nothing will ever await this task. Left alone it would run to + # completion, pay for a Retrieve, and be reported as a task whose + # exception was never retrieved. + managed_task.cancel() # Return empty list on error (graceful degradation) return [] +def _filter_chunks_by_document_status(chunks: List[Chunk], assistant_id: str) -> List[Chunk]: + """ + Apply the document status filter to protocol chunks, on any backend. + + Delegates to :func:`_filter_vectors_by_document_status` rather than + reimplementing the lookup, so both backends share one set of DynamoDB + semantics — including its fallback behaviour, which task group 6 changes in + exactly one place. + + Each chunk is presented to the filter as a minimal view carrying only what + the filter reads (``metadata.document_id``) plus its index, and survivors are + mapped back by that index. Order and duplicates are preserved. + + Args: + chunks: Chunks returned by a backend, in backend ranking order + assistant_id: Assistant identifier for DynamoDB key construction + + Returns: + The subset of chunks whose parent document is 'complete' + """ + views = [{"metadata": chunk.metadata, "_chunk_index": index} for index, chunk in enumerate(chunks)] + surviving = _filter_vectors_by_document_status(views, assistant_id) + return [chunks[view["_chunk_index"]] for view in surviving] + + def _filter_vectors_by_document_status(vectors: List[Dict[str, Any]], assistant_id: str) -> List[Dict[str, Any]]: """ Filter vector results to only include chunks from documents with status='complete'. @@ -75,7 +241,14 @@ def _filter_vectors_by_document_status(vectors: List[Dict[str, Any]], assistant_ Extracts unique document_ids from vector metadata, looks up each document's status in DynamoDB, and removes chunks from documents that are not 'complete' or don't exist. - On any DynamoDB failure, falls back to returning unfiltered results (graceful degradation). + Fails CLOSED (Requirement 5): if status cannot be confirmed — no table + configured, or the lookup errors — every chunk is dropped and an empty list is + returned. This deliberately supersedes `reliable-document-deletion` + Requirement 3.4, which specified the opposite. The reasoning changed because + the fail-open path was measured: 936 retrievals in a trailing 30-day window had + chunks dropped by this filter, so the documents it guards against are real, and + serving a user content they believe they deleted is worse than serving nothing. + A per-document lookup failure still skips only that document. Args: vectors: List of vector results from S3 Vectors search @@ -119,12 +292,31 @@ def _filter_vectors_by_document_status(vectors: List[Dict[str, Any]], assistant_ logger.warning(f"Failed to look up document {doc_id}: {e}") # Skip individual lookup failures else: - # No table configured — fall back to unfiltered - logger.warning("DYNAMODB_ASSISTANTS_TABLE_NAME not configured, returning unfiltered results") - valid_doc_ids = doc_ids + # FAIL CLOSED (Requirement 5.2). Previously this returned everything + # unfiltered. Without a table there is no way to confirm that a + # document is still `complete`, and the chunks in question may belong + # to documents a user has deleted. Serving unverifiable content is a + # worse outcome than serving none: the user sees material they believe + # they removed, and nothing in the response signals that the check was + # skipped. + logger.error( + "DYNAMODB_ASSISTANTS_TABLE_NAME not configured; dropping all " + "chunks because document status cannot be confirmed" + ) + emit_count(METRIC_STATUS_FILTER_FAIL_CLOSED) + return [] except Exception as e: - logger.warning(f"DynamoDB lookup failed, returning unfiltered results: {e}") - valid_doc_ids = doc_ids # Graceful degradation + # FAIL CLOSED (Requirement 5.1). Same reasoning as above. Logged at ERROR, + # not WARNING: an empty result from this path is a degradation, and it must + # be distinguishable from the ordinary "corpus had no match" case, which is + # logged at INFO below. + logger.error( + f"Document status lookup failed; dropping all chunks because status " + f"cannot be confirmed: {e}", + exc_info=True, + ) + emit_count(METRIC_STATUS_FILTER_FAIL_CLOSED) + return [] # Filter vectors to only include chunks from valid documents filtered = [v for v in vectors if v.get("metadata", {}).get("document_id") in valid_doc_ids] @@ -136,13 +328,16 @@ def _filter_vectors_by_document_status(vectors: List[Dict[str, Any]], assistant_ return filtered -def augment_prompt_with_context(user_message: str, context_chunks: List[Dict[str, Any]], max_context_length: int = 2000) -> str: +def augment_prompt_with_context(user_message: str, context_chunks: List[Dict[str, Any]], max_context_length: int = MAX_CONTEXT_CHARS) -> str: """ Augment user message with retrieved context chunks The context is prepended to the user message with clear delimiters. This allows the LLM to use the retrieved knowledge when generating responses. + Applies on both backends: the cap lives here, above the seam, so neither + adapter can widen it independently. + Args: user_message: Original user message context_chunks: List of context chunks from vector search diff --git a/backend/src/apis/shared/assistants/service.py b/backend/src/apis/shared/assistants/service.py index 8849971f6..2e05fbe1c 100644 --- a/backend/src/apis/shared/assistants/service.py +++ b/backend/src/apis/shared/assistants/service.py @@ -588,9 +588,17 @@ async def _update_assistant_cloud(assistant: Assistant, table_name: str) -> None # edit would re-write a stale GSI5 key onto a delisted agent and silently put it # back in the store. ``listing`` is excluded for the same reason in the other # direction: a stale in-memory copy must not clobber a concurrent review decision. + # + # GSI7_* belongs to the managed-KB migration write path + # (``apis.shared.kb_backend.records``). Unlike GSI5, those keys live on a + # separate item (``SK = KB#{app_kb_id}``) rather than on this METADATA item, so + # this path cannot reach them today — the entry is here so the invariant "GSI + # keys are never written from the generic update" holds uniformly, and a reader + # does not have to know which index lives on which item to trust it. immutable_fields = { "PK", "SK", "GSI_PK", "GSI_SK", "GSI2_PK", "GSI2_SK", "GSI5_PK", "GSI5_SK", + "GSI7_PK", "GSI7_SK", "assistantId", "createdAt", "ownerId", "listing", } diff --git a/backend/src/apis/shared/assistants/version_resolution.py b/backend/src/apis/shared/assistants/version_resolution.py index 94d2a1897..2c050aace 100644 --- a/backend/src/apis/shared/assistants/version_resolution.py +++ b/backend/src/apis/shared/assistants/version_resolution.py @@ -11,9 +11,11 @@ | Anyone, when nothing is published | The live record | ``resolve_display_agent`` answers the same question for the marketplace **detail read**, -and differs on exactly one row: editors see the draft too. The two are kept side by side -here rather than merged, because the reason they differ is easy to lose and expensive to -get wrong in either direction — see that function's docstring. +and differs on exactly one row: editors see the draft too. ``resolve_review_agent`` answers +it for a marketplace **reviewer**, and differs on every row — the artifact a decision is +about is the *submitted* snapshot, which neither of the others ever serves. The three are +kept side by side here rather than merged, because the reasons they differ are easy to lose +and expensive to get wrong in any direction — see each function's docstring. This is deliberately *not* in ``versions.py`` (pure, no I/O) or in ``version_repository.py`` (persistence). Deciding which version a person runs is policy, @@ -115,6 +117,63 @@ async def resolve_invocation_agent( return apply_version(assistant, version), published +async def resolve_review_agent(assistant: Assistant) -> Tuple[Assistant, Optional[int]]: + """Return ``(assistant_under_review, version_number)`` for a **marketplace reviewer**. + + The third caller of the same question, and the one whose answer differs most: a + reviewer reads and test-drives the artifact a decision is *about*, which is neither the + published snapshot nor the author's draft. + + | Listing state | Reviews | + |----------------------|--------------------| + | ``in_review`` | ``submittedVersion`` | + | anything else | ``publishedVersion`` | + + ⚠️ **``in_review`` reads ``submittedVersion`` because that is what approval promotes** + (``listing_service.review_listing``). The live record is the author's draft and they can + keep editing it while the row sits in the queue — the window ``AgentVersion`` exists to + close, since a version is cut at submission rather than at approval. Reading anything + else would show a reviewer one configuration and publish another. + + ⚠️ **Every other state reads ``publishedVersion``, and that is not a fallback.** + ``submittedVersion`` is a high-water mark that deliberately survives a decision, so on a + ``withdrawal_requested`` listing — where nothing is pending, and the question is whether + to pull what is *live* — it names a snapshot the store never served. + + Returns ``(assistant, None)`` when neither pointer resolves: a submission that predates + version snapshots, or a pointer to a version that is gone. Unlike ``resolve_invocation_agent`` + this **does not raise** on a missing snapshot, because there is nothing unsafe about a + reviewer reading the live record as long as they are told that is what they are reading — + the ``None`` is that signal, and both callers surface it (``snapshotUnavailable`` on the + review read; the preview banner on the test drive). Raising would leave a reviewer with a + row in their queue and no way to look at it. + + Like the other two, **this is not an access check.** Whether the caller may review at all + is the ``admin.marketplace`` scope, decided by the caller. + """ + listing = assistant.listing + if listing is None: + return assistant, None + + number = ( + listing.submitted_version if listing.state == "in_review" else listing.published_version + ) + if number is None: + return assistant, None + + # Read the version back rather than trusting the pointer: a snapshot that is gone must + # read as "not frozen", not as a version number whose content nobody has. + version = await get_version(assistant.assistant_id, number) + if version is None: + logger.warning( + f"Agent {assistant.assistant_id} names version {number} for review, " + "but it could not be loaded; falling back to the live record." + ) + return assistant, None + + return apply_version(assistant, version), version.version + + async def resolve_display_agent( assistant: Assistant, *, can_edit: bool ) -> Tuple[Assistant, Optional[int]]: diff --git a/backend/src/apis/shared/auth/__init__.py b/backend/src/apis/shared/auth/__init__.py index 31bce9eb5..65002acb2 100644 --- a/backend/src/apis/shared/auth/__init__.py +++ b/backend/src/apis/shared/auth/__init__.py @@ -3,7 +3,7 @@ from .dependencies import get_current_user_from_session, security from .models import User from .state_store import StateStore, InMemoryStateStore, DynamoDBStateStore, create_state_store -from .rbac import require_app_roles, require_admin, require_admin_scope +from .rbac import has_admin_scope, require_app_roles, require_admin, require_admin_scope __all__ = [ "get_current_user_from_session", @@ -16,4 +16,5 @@ "require_app_roles", "require_admin", "require_admin_scope", + "has_admin_scope", ] diff --git a/backend/src/apis/shared/auth/rbac.py b/backend/src/apis/shared/auth/rbac.py index 6bf58875c..ca30386c4 100644 --- a/backend/src/apis/shared/auth/rbac.py +++ b/backend/src/apis/shared/auth/rbac.py @@ -67,6 +67,31 @@ async def checker(user: User = Depends(get_current_user_from_session)) -> User: return checker +async def has_admin_scope(user: User, scope: str) -> bool: + """Whether ``user`` holds ``scope`` — the predicate behind ``require_admin_scope``. + + Exists because one caller needs the *answer* rather than a route guard: the invocation + path's reviewer preview (``inference_api.chat.routes``) is a field on an existing + request, not a route, so it cannot take a FastAPI dependency. Sharing the predicate is + what keeps "who is a marketplace admin" from being answered twice, differently. + + Fails closed, exactly as the dependency does: a permission lookup that raises is a + denial, never a default-allow. + """ + from apis.shared.rbac.service import get_app_role_service + + try: + permissions = await get_app_role_service().resolve_user_permissions(user) + except Exception: + logger.exception( + f"Failed to resolve admin scope {scope} for {user.name}, denying access" + ) + return False + # ``system_admin`` satisfies every scope implicitly — the superuser rule the dependency + # applies, restated here rather than reimplemented differently. + return "system_admin" in permissions.app_roles or scope in permissions.admin_scopes + + def require_admin_scope(scope: str) -> Callable: """ Create a dependency guarding one delegated admin surface. @@ -92,24 +117,9 @@ def require_admin_scope(scope: str) -> Callable: HTTPException: 403 if the user holds neither system_admin nor the scope. """ async def checker(user: User = Depends(get_current_user_from_session)) -> User: - from apis.shared.rbac.service import get_app_role_service - - try: - service = get_app_role_service() - permissions = await service.resolve_user_permissions(user) - - if "system_admin" in permissions.app_roles: - return user - - if scope in permissions.admin_scopes: - logger.debug( - f"User {user.name} authorized for admin scope {scope}" - ) - return user - except Exception: - logger.exception( - f"Failed to resolve admin scope {scope} for {user.name}, denying access" - ) + if await has_admin_scope(user, scope): + logger.debug(f"User {user.name} authorized for admin scope {scope}") + return user logger.warning( f"User {user.name} (jwt_roles: {user.roles}) denied access — " diff --git a/backend/src/apis/shared/embeddings/bedrock_embeddings.py b/backend/src/apis/shared/embeddings/bedrock_embeddings.py index 22d1ac3dd..0e74bc7ea 100644 --- a/backend/src/apis/shared/embeddings/bedrock_embeddings.py +++ b/backend/src/apis/shared/embeddings/bedrock_embeddings.py @@ -57,7 +57,13 @@ async def generate_embeddings(chunks: List[str]) -> List[List[float]]: IMPORTANT: This function does NOT validate token counts. Callers that process large documents should validate/split chunks before calling this. - For search queries (short strings), no validation is needed. + + Search queries are length-capped by the caller, not here: the facade applies + `apis.shared.kb_backend.query_guard.clamp_query` before dispatch. This used to + say no validation was needed for queries, which was true only because Titan v2 + tolerates roughly 32,000 characters. Managed Knowledge Base caps `Retrieve` + input at 10,000 and that quota is not adjustable, so the assumption no longer + holds for every backend. Args: chunks: List of text chunks to embed @@ -145,7 +151,8 @@ async def search_assistant_knowledgebase(assistant_id: str, query: str): """Search the S3 vector store for chunks relevant to the query.""" client = boto3.client("s3vectors", region_name=AWS_REGION) - # Generate vector for the query (short string, no token validation needed) + # Generate vector for the query. Length is already capped upstream by the + # facade's query clamp (10,000 chars, the Managed KB Retrieve limit). query_embedding = await generate_embeddings([query]) # Query the Global Index with a STRICT Filter diff --git a/backend/src/apis/shared/kb_backend/__init__.py b/backend/src/apis/shared/kb_backend/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/src/apis/shared/kb_backend/byte_cap.py b/backend/src/apis/shared/kb_backend/byte_cap.py new file mode 100644 index 000000000..dd5e532a4 --- /dev/null +++ b/backend/src/apis/shared/kb_backend/byte_cap.py @@ -0,0 +1,258 @@ +"""Per-owner byte cap accounting for managed knowledge bases. + +Managed storage is billed at $5.00/GB-month, roughly 35x what S3 Vectors costs +today. At the measured average of 1.13 MB per user that is about $169/month across +the fleet — but nothing structural stops one user uploading far more, and 30,000 +users at 100 MB each would be 3 TB, or about $15,000/month. The cap is what turns +"unlikely" into "impossible". + +Why an accumulator instead of the obvious condition +--------------------------------------------------- +The natural way to express this is:: + + ConditionExpression="storedBytes + reservedBytes + :n <= :cap" + +**DynamoDB rejects that.** Condition expressions compare operands; they cannot do +arithmetic. Verified directly: the parser fails with ``Cannot parse condition +starting at:+ reserved <= :cap``. + +So the arithmetic is moved to the client, where it is free. A single +``totalBytes`` accumulator is maintained as the invariant +``totalBytes == storedBytes + reservedBytes``, and the guard compares it against a +**literal computed before the call**:: + + ADD totalBytes :n, reservedBytes :n + CONDITION totalBytes <= :max_before where :max_before = cap - n + +That is a single atomic conditional update, so N concurrent reservations cannot +collectively overshoot. The alternative — read, compute, write — has a window +between the read and the write in which another writer commits, which is exactly +the race a cap exists to prevent. + +Reserve / commit / release, not just "add" +----------------------------------------- +Ingestion is not instantaneous: a 50 KiB PDF measured 68-264 seconds. Counting +bytes only on success would let a user start unlimited concurrent uploads that are +each individually under the cap and collectively far over it. So bytes are reserved +up front, converted to stored on success, and returned on failure. A crash between +reserve and commit leaks a reservation, which is the safe direction — it +under-permits rather than over-permits, and the reconciler can recover it. + +Sizing +------ +Size always comes from an S3 ``HEAD`` on the stored object, never from a +client-reported value: a client that under-reports its own size would defeat the +cap entirely. Bedrock's ``RawDataSize`` metric is deliberately **not** used for +enforcement — it returned 0 datapoints for a directly-ingested document during +evaluation and remains unconfirmed. Enforcing against a metric that is sometimes +absent would fail open. + +Import weight +------------- +Module-level imports are stdlib only; ``boto3`` is function-local, so this module +can be imported into a size-constrained Lambda image for free. +""" + +from __future__ import annotations + +import logging +import os +from decimal import Decimal +from typing import Optional + +from apis.shared.kb_backend.metrics import emit_count + +logger = logging.getLogger(__name__) + +METRIC_BYTE_CAP_REJECTED = "KbByteCapRejected" + +#: Defaults mirror the CDK config (Requirement 12.2). Both are read from the +#: environment so an operator can tune them without a code change; the fallbacks +#: keep local runs working. +#: +#: 100 MB is deliberately BELOW the existing 1 GB user-files precedent. At $5.00 +#: per GB-month that precedent would permit roughly $150,000/month across the +#: fleet, which is not a limit so much as a formality. +DEFAULT_PER_OWNER_BYTES = 100 * 1024 * 1024 +DEFAULT_PER_OWNER_ELEVATED_BYTES = 1024 * 1024 * 1024 +DEFAULT_PER_KB_CEILING_BYTES = 500 * 1024 * 1024 + + +class ByteCapExceeded(Exception): + """A reservation would take the owner over their cap. + + Carries the numbers so the caller can render a plain-language message with the + option to request an elevated tier, rather than a bare failure (Requirement + 12.12). A user who cannot see how far over they are cannot act on it. + """ + + def __init__(self, requested: int, cap: int, already_used: Optional[int] = None) -> None: + self.requested = requested + self.cap = cap + self.already_used = already_used + super().__init__( + f"reserving {requested} bytes would exceed the {cap}-byte cap" + + (f" (already using {already_used})" if already_used is not None else "") + ) + + +def _env_int(name: str, default: int) -> int: + raw = os.environ.get(name) + if not raw: + return default + try: + return int(raw) + except ValueError: + logger.warning(f"{name}={raw!r} is not an integer; falling back to {default}") + return default + + +def per_owner_cap(elevated: bool = False) -> int: + """The owner's total allowance in bytes. + + Which tier a user belongs to is the caller's decision: RBAC already owns role + resolution and this module should not grow a second opinion about it. + """ + if elevated: + return _env_int("MANAGED_KB_PER_OWNER_ELEVATED_BYTES", DEFAULT_PER_OWNER_ELEVATED_BYTES) + return _env_int("MANAGED_KB_PER_OWNER_DEFAULT_BYTES", DEFAULT_PER_OWNER_BYTES) + + +def per_kb_ceiling() -> int: + """Ceiling for a single knowledge base, independent of the owner's total. + + Stops one knowledge base consuming an entire elevated allowance and starving + the owner's others. + """ + return _env_int("MANAGED_KB_PER_KB_CEILING_BYTES", DEFAULT_PER_KB_CEILING_BYTES) + + +def _table(): + import boto3 + + return boto3.resource("dynamodb").Table(os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"]) + + +def object_size_bytes(bucket: str, key: str) -> int: + """Authoritative size, from S3 rather than from the client. + + A client-reported size is an input, and an input that can lower its own cost is + not a measurement. + """ + import boto3 + + response = boto3.client("s3").head_object(Bucket=bucket, Key=key) + return int(response["ContentLength"]) + + +def reserve( + assistant_id: str, + app_kb_id: str, + n_bytes: int, + cap: int, +) -> None: + """Reserve ``n_bytes`` against the cap, atomically. + + Raises :class:`ByteCapExceeded` if the reservation would breach the cap. The + comparison is against ``cap - n_bytes``, computed here, because DynamoDB cannot + add inside a condition — see the module docstring. + + ``attribute_not_exists`` covers the first reservation on a record that has + never held bytes, so a fresh knowledge base does not need initialising. + """ + from botocore.exceptions import ClientError + + from apis.shared.kb_backend.records import kb_pk, kb_sk + + if n_bytes < 0: + raise ValueError("n_bytes must not be negative") + if n_bytes == 0: + return + if n_bytes > cap: + # Cannot fit even into an empty allowance; no point issuing the write. + emit_count(METRIC_BYTE_CAP_REJECTED) + raise ByteCapExceeded(requested=n_bytes, cap=cap) + + try: + _table().update_item( + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression="ADD #total :n, #reserved :n", + ConditionExpression="attribute_not_exists(#total) OR #total <= :max_before", + ExpressionAttributeNames={ + # `total` is a DynamoDB reserved keyword, so these are aliased. + "#total": "totalBytes", + "#reserved": "reservedBytes", + }, + ExpressionAttributeValues={ + ":n": Decimal(n_bytes), + ":max_before": Decimal(cap - n_bytes), + }, + ) + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + emit_count(METRIC_BYTE_CAP_REJECTED) + raise ByteCapExceeded(requested=n_bytes, cap=cap) from exc + raise + + +def commit(assistant_id: str, app_kb_id: str, n_bytes: int) -> None: + """Convert a reservation into stored bytes. + + ``totalBytes`` is untouched: the bytes were already counted at reserve time. + Adding here as well would double-count and shrink the owner's allowance on + every successful upload. + """ + from apis.shared.kb_backend.records import kb_pk, kb_sk + + if n_bytes == 0: + return + _table().update_item( + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression="ADD #reserved :neg, #stored :n", + ExpressionAttributeNames={"#reserved": "reservedBytes", "#stored": "storedBytes"}, + ExpressionAttributeValues={":neg": Decimal(-n_bytes), ":n": Decimal(n_bytes)}, + ) + + +def release(assistant_id: str, app_kb_id: str, n_bytes: int) -> None: + """Return a reservation after a failed ingestion. + + Decrements both the reservation and the accumulator, restoring the allowance + exactly. Not releasing would silently shrink the owner's cap with every failed + upload until they could not upload at all — a leak that presents as "the + product stopped working" long after the failures that caused it. + """ + from apis.shared.kb_backend.records import kb_pk, kb_sk + + if n_bytes == 0: + return + _table().update_item( + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression="ADD #reserved :neg, #total :neg", + ExpressionAttributeNames={"#reserved": "reservedBytes", "#total": "totalBytes"}, + ExpressionAttributeValues={":neg": Decimal(-n_bytes)}, + ) + + +def reserve_snapshot( + assistant_id: str, + app_kb_id: str, + total_bytes: int, + cap: int, +) -> None: + """Reserve a whole migration corpus up front (Requirement 12.11/12.12). + + Migration is the largest byte-adding operation in the system and the only one + that runs unattended, which makes it both the easiest place to forget the check + and the worst. Reserving per-document as the worker progresses would let a + migration run for an hour and then stop halfway, leaving a half-populated + managed knowledge base and an owner over their cap with no way back. + + So the entire snapshot is reserved *before* the migration enters ``shadow``. A + corpus that cannot fit fails immediately, with numbers the caller can turn into + "this needs an elevated tier" rather than a stack trace. + + Deliberately the same conditional write as :func:`reserve`; the distinction is + the caller's contract, not the mechanism. + """ + reserve(assistant_id, app_kb_id, total_bytes, cap) diff --git a/backend/src/apis/shared/kb_backend/dual_read.py b/backend/src/apis/shared/kb_backend/dual_read.py new file mode 100644 index 000000000..03f4ae977 --- /dev/null +++ b/backend/src/apis/shared/kb_backend/dual_read.py @@ -0,0 +1,347 @@ +"""The dual-read pilot: measure the managed backend against real traffic. + +Requirement 18. The rollout should rest on evidence from *our* corpus and *our* +users, not solely on a three-document benchmark. So an opted-in knowledge base +can have both backends answer the same query, with **legacy always served** and +the managed result kept purely as an observation. + +Three rules make this safe to leave switched on, and each is a property of the +code rather than an intention: + +* **Legacy is what is served.** The managed result never reaches the caller. It + is not blended, not preferred when it looks better, not used as a fallback when + legacy is empty — an empty legacy result is a *finding*, and substituting the + other engine's answer would destroy the measurement and change what users see + in the same move. +* **The managed call cannot fail the turn.** It runs as a detached task whose + exceptions are logged and dropped. :func:`observe` has no failure mode that + propagates. +* **It cannot add user-visible latency (18.5).** Both searches start together and + the caller is handed the legacy result the moment it resolves; the managed call + keeps running afterwards on the event loop. This matters concretely: managed + ``Retrieve`` measured a 662–695 ms p50 against legacy's 257 ms, so anything that + awaited both would nearly triple the retrieval leg of every piloted turn. + +Why the task needs a strong reference +------------------------------------- +``asyncio.create_task`` returns the only strong reference to the task. Drop it and +the task becomes eligible for garbage collection mid-flight, which surfaces as +comparisons that silently stop being recorded under load — the failure mode that +looks like "the pilot found nothing interesting". Hence :data:`_IN_FLIGHT` and the +done-callback that discards from it, which is the documented CPython pattern. + +Why the comparison is a pure function +------------------------------------- +:func:`compare` takes two chunk lists and two durations and returns a value. It +touches no clock, no client and no environment, so the ranking mathematics can be +tested without any of the machinery around it — and the machinery can be tested +without asserting on arithmetic. + +Feature: managed-kb-migration +Requirements: 18.1, 18.2, 18.3, 18.4, 18.5 +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from dataclasses import dataclass +from typing import Any, Dict, List, Mapping, Optional, Sequence, Set + +from apis.shared.kb_backend.metrics import ( + METRIC_DUAL_READ_FAILED, + METRIC_DUAL_READ_LATENCY, + METRIC_DUAL_READ_OVERLAP, + METRIC_DUAL_READ_RANK_CORRELATION, + emit_count, + emit_value, +) +from apis.shared.kb_backend.protocol import DEFAULT_TOP_K, Chunk +from apis.shared.kb_backend.records import ENGINE_MANAGED + +logger = logging.getLogger(__name__) + +#: KB_Record attribute that opts one knowledge base into the pilot. Absence means +#: off (Requirement 18.4), the same convention ``retrievalEngine`` uses: the +#: default costs nothing to express and nothing to revert. +DUAL_READ_ATTR = "dualReadPilot" + +#: Strong references to detached comparison tasks. See the module docstring. +_IN_FLIGHT: Set["asyncio.Task[None]"] = set() + + +def is_pilot_enabled(record: Optional[Mapping[str, Any]]) -> bool: + """Whether this knowledge base is opted into the pilot. + + Strictly ``is True``: a truthy string left behind by a hand-edited record must + not enrol a knowledge base into paying for a second retrieval on every turn. + The same reasoning armed the reconciler's flag, where a permissive read of an + event field turned a report-only job into a deleting one. + """ + if not record: + return False + return record.get(DUAL_READ_ATTR) is True + + +@dataclass(frozen=True) +class Comparison: + """One dual read's observation. Serves nothing; describes everything.""" + + legacy_count: int + managed_count: int + overlap_count: int + overlap_ratio: float + rank_correlation: Optional[float] + legacy_ms: float + managed_ms: float + + def as_log_fields(self) -> Dict[str, Any]: + return { + "legacyCount": self.legacy_count, + "managedCount": self.managed_count, + "overlapCount": self.overlap_count, + "overlapRatio": round(self.overlap_ratio, 4), + "rankCorrelation": ( + None if self.rank_correlation is None else round(self.rank_correlation, 4) + ), + "legacyMs": round(self.legacy_ms, 1), + "managedMs": round(self.managed_ms, 1), + } + + +def _first_positions(chunks: Sequence[Chunk]) -> Dict[str, int]: + """Each ``document_id``'s best rank in a result list, 0-based. + + A document can contribute several chunks, so "the document's rank" is the rank + of its best chunk. Using every chunk instead would let a document with four + passages dominate a correlation over one with a single passage, which measures + chunking rather than agreement. + """ + positions: Dict[str, int] = {} + for index, chunk in enumerate(chunks): + document_id = (chunk.metadata or {}).get("document_id") or chunk.document_id + if document_id and document_id not in positions: + positions[document_id] = index + return positions + + +def _spearman(left: Sequence[float], right: Sequence[float]) -> Optional[float]: + """Pearson correlation of two rank vectors — Spearman, computed by hand. + + Written out rather than pulled from scipy: this package is bundled into + size-constrained Lambda images, and a numerical stack is a large dependency to + add for one dot product. + + ``None`` when fewer than two documents are shared (a correlation over one point + is undefined, not 1.0) or when either vector has zero variance, which is what + happens when both backends return the same single document. + """ + n = len(left) + if n < 2 or n != len(right): + return None + + mean_left = sum(left) / n + mean_right = sum(right) / n + d_left = [value - mean_left for value in left] + d_right = [value - mean_right for value in right] + + covariance = sum(a * b for a, b in zip(d_left, d_right)) + variance_left = sum(a * a for a in d_left) + variance_right = sum(b * b for b in d_right) + + if variance_left == 0 or variance_right == 0: + return None + + return covariance / ((variance_left**0.5) * (variance_right**0.5)) + + +def compare( + legacy: Sequence[Chunk], + managed: Sequence[Chunk], + legacy_ms: float, + managed_ms: float, +) -> Comparison: + """The observation for one dual read (Requirement 18.3). Pure. + + ``overlap_ratio`` is Jaccard — shared documents over the union — chosen because + it is symmetric. A ratio against one side's length would read as agreement when + one backend simply returned fewer documents, which is the case most likely to + occur while the managed corpus is still catching up. + """ + legacy_positions = _first_positions(legacy) + managed_positions = _first_positions(managed) + + legacy_ids = set(legacy_positions) + managed_ids = set(managed_positions) + shared = legacy_ids & managed_ids + union = legacy_ids | managed_ids + + ordered = sorted(shared, key=lambda doc_id: legacy_positions[doc_id]) + correlation = _spearman( + [float(legacy_positions[doc_id]) for doc_id in ordered], + [float(managed_positions[doc_id]) for doc_id in ordered], + ) + + return Comparison( + legacy_count=len(legacy), + managed_count=len(managed), + overlap_count=len(shared), + overlap_ratio=(len(shared) / len(union)) if union else 0.0, + rank_correlation=correlation, + legacy_ms=legacy_ms, + managed_ms=managed_ms, + ) + + +async def _publish(comparison: Comparison) -> None: + """Record the observation. Metrics go to a thread; they are boto3 calls. + + Off the critical path already, but the event loop is shared with every other + in-flight turn, so four synchronous HTTP calls would still be four pauses + everybody pays for. + """ + logger.info(f"dual read comparison: {comparison.as_log_fields()}") + + def _emit() -> None: + emit_value(METRIC_DUAL_READ_OVERLAP, comparison.overlap_ratio, unit="Percent") + if comparison.rank_correlation is not None: + emit_value(METRIC_DUAL_READ_RANK_CORRELATION, comparison.rank_correlation) + emit_value( + METRIC_DUAL_READ_LATENCY, + comparison.legacy_ms, + unit="Milliseconds", + dimensions={"backend": "s3vectors"}, + ) + emit_value( + METRIC_DUAL_READ_LATENCY, + comparison.managed_ms, + unit="Milliseconds", + dimensions={"backend": ENGINE_MANAGED}, + ) + + await asyncio.to_thread(_emit) + + +async def observe( + assistant_id: str, + query: str, + top_k: int, + legacy_chunks: List[Chunk], + legacy_ms: float, + managed_task: "Optional[asyncio.Task[List[Chunk]]]", +) -> None: + """Await the already-running managed search and record the comparison. + + Never raises, and never returns anything a caller could serve. ``managed_task`` + is awaited here rather than started here, so that by the time this runs the + managed call has been in flight for as long as the legacy one took — which is + what makes the two latencies comparable and the pilot non-additive. + """ + if managed_task is None: + return + + started = time.perf_counter() + try: + managed_chunks = await managed_task + except asyncio.CancelledError: + raise + except Exception as exc: + # A managed-side failure is a finding, not an incident: the turn was + # served from legacy before this coroutine ran. + logger.warning( + f"dual read: the managed backend failed for assistant {assistant_id}; " + f"the turn was already served from legacy: {exc}" + ) + await asyncio.to_thread(emit_count, METRIC_DUAL_READ_FAILED) + return + + managed_ms = legacy_ms + (time.perf_counter() - started) * 1000.0 + + try: + await _publish(compare(legacy_chunks, managed_chunks, legacy_ms, managed_ms)) + except Exception as exc: # noqa: BLE001 - observation must not escape + logger.warning(f"dual read: could not record the comparison: {exc}") + + +def start_managed_read( + record: Optional[Mapping[str, Any]], + assistant_id: str, + query: str, + top_k: int = DEFAULT_TOP_K, +) -> "Optional[asyncio.Task[List[Chunk]]]": + """Launch the observational managed search, or return ``None``. + + ``None`` — meaning "no dual read this turn" — for every one of: the knowledge + base is not opted in, this build has no managed backend registered, the record + already names managed as its engine (there would be nothing to compare + against), or the task could not be created. Each is an ordinary state, so none + of them logs at error level or raises. + + Called *before* the legacy search is awaited, which is the whole basis of + Requirement 18.5. + """ + if not is_pilot_enabled(record): + return None + + from apis.shared.kb_backend.records import resolve_engine + from apis.shared.kb_backend.resolver import backend_for_engine + + if resolve_engine(record) == ENGINE_MANAGED: + # Already promoted: the managed backend is the one being served, so a + # "comparison" would be the same call twice at twice the price. + return None + + managed = backend_for_engine(ENGINE_MANAGED) + if managed is None: + return None + + try: + task = asyncio.create_task(managed.search(assistant_id, query, top_k)) + except RuntimeError as exc: + logger.warning(f"dual read: could not start the managed search: {exc}") + return None + + _IN_FLIGHT.add(task) + task.add_done_callback(_IN_FLIGHT.discard) + return task + + +def schedule_observation( + assistant_id: str, + query: str, + top_k: int, + legacy_chunks: List[Chunk], + legacy_ms: float, + managed_task: "Optional[asyncio.Task[List[Chunk]]]", +) -> None: + """Detach :func:`observe` so the caller can return immediately. + + The point of Requirement 18.5 in one function: nothing after this line is + awaited before the user gets their answer. + """ + if managed_task is None: + return + + try: + observer = asyncio.create_task( + observe(assistant_id, query, top_k, legacy_chunks, legacy_ms, managed_task) + ) + except RuntimeError as exc: + logger.warning(f"dual read: could not schedule the comparison: {exc}") + managed_task.cancel() + return + + _IN_FLIGHT.add(observer) + observer.add_done_callback(_IN_FLIGHT.discard) + + +__all__ = [ + "DUAL_READ_ATTR", + "Comparison", + "compare", + "is_pilot_enabled", + "observe", + "schedule_observation", + "start_managed_read", +] diff --git a/backend/src/apis/shared/kb_backend/idleness.py b/backend/src/apis/shared/kb_backend/idleness.py new file mode 100644 index 000000000..50813652d --- /dev/null +++ b/backend/src/apis/shared/kb_backend/idleness.py @@ -0,0 +1,281 @@ +"""When was this knowledge base last actually needed? + +Requirements 22.5, 22.6. Two rules, and both exist because the obvious answer is +wrong in a way that destroys data. + +**Idleness is not retrieval (22.5).** A knowledge base is idle when *nothing* has +needed it, and retrieval is only one of the ways it gets needed. An agent can be +invoked hundreds of times a day and retrieve nothing from its own corpus, because +retrieval only fires when the query matches — so a corpus judged by retrieval alone +looks abandoned precisely when its agent is busiest with questions the documents do +not answer. The follow-up spec's eviction pass would then delete the documents +behind a live agent. So idleness is the maximum of the knowledge base's own +``lastRetrievedAt`` and the ``lastUsedAt`` of any agent bound to it. + +While this phase holds ``App_KB_Id == assistant_id`` there is exactly one bound +agent and it is the assistant itself, so "any bound agent" is one ``METADATA`` +read. That is deliberately written as a maximum over a set rather than a single +lookup: F4 makes the set larger, and a maximum over one element is the same code. + +**Never write a timestamp per retrieval (22.6).** Retrieval is the hot path. The +write is therefore conditional on the stored value being older than a throttle +window, so at most one write lands per window no matter how many turns race — the +same shape as ``assistants.service.bump_last_used_at``, which solved this for +``lastUsedAt`` and is the precedent being followed rather than a second invention. +A conditional write that loses is not a write; it is a rejected update, which is +why calling this on every retrieval is consistent with the requirement. + +Nothing here reclaims anything +------------------------------ +``reclaim`` is reserved in the migration state enum and never entered in this +phase. This module exists now anyway, because the eviction threshold the follow-up +spec has to choose can only be chosen from historical idleness data, and that data +cannot be backfilled — a timestamp nobody recorded in August is not available in +November. + +Feature: managed-kb-migration +Requirements: 22.1, 22.5, 22.6 +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Iterable, Mapping, Optional + +logger = logging.getLogger(__name__) + +#: One write per knowledge base per day at most. Chosen to match +#: ``bump_last_used_at``'s default: idleness is measured in days, so a finer +#: resolution buys nothing and costs a write per turn. +THROTTLE_HOURS = 24 + +#: Attribute on the KB_Record. +LAST_RETRIEVED_ATTR = "lastRetrievedAt" + +#: Strong references to detached touch tasks. Without this the only reference is +#: the one ``create_task`` returns, and a dropped task can be collected mid-flight. +_IN_FLIGHT: set = set() + + +def _table(): + import boto3 + + return boto3.resource("dynamodb").Table(os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"]) + + +def _now(): + from datetime import datetime, timezone + + return datetime.now(timezone.utc) + + +def _iso(moment) -> str: + return moment.strftime("%Y-%m-%dT%H:%M:%SZ") + + +def throttle_hours() -> int: + """Resolved at call time, never as a default argument. + + A module constant bound into a signature is captured once at import, so a test + overriding it silently gets the production value. That cost this feature a + 33-second test that ignored its own override. + """ + raw = os.environ.get("KB_LAST_RETRIEVED_THROTTLE_HOURS") + try: + value = int(raw) if raw else THROTTLE_HOURS + except ValueError: + return THROTTLE_HOURS + return max(value, 1) + + +def touch_last_retrieved(assistant_id: str, app_kb_id: str) -> bool: + """Record that this knowledge base served a retrieval. Never raises. + + Returns ``True`` only for the caller whose write actually landed — at most one + per throttle window. Callers do not need the result; it is returned because a + boolean that names the winner is what let ``bump_last_used_at`` hang + resume-on-first-use off the same write, and the reconciler may want the same + hook later. + + Guarded on ``attribute_exists(SK)`` as well as the freshness floor, so this + cannot bring a KB_Record into existence. A legacy knowledge base has no record + and must keep having none: the migration's zero-backfill property is that + nothing writes to these 1,692 rows until their owner opts in, and a metrics + side effect that created rows would break it while looking harmless. + """ + if not os.environ.get("DYNAMODB_ASSISTANTS_TABLE_NAME"): + return False + + try: + from datetime import timedelta + + from apis.shared.kb_backend.records import kb_pk, kb_sk + + now = _now() + floor = _iso(now - timedelta(hours=throttle_hours())) + _table().update_item( + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression=f"SET {LAST_RETRIEVED_ATTR} = :now", + ConditionExpression=( + f"attribute_exists(SK) AND (attribute_not_exists({LAST_RETRIEVED_ATTR}) " + f"OR {LAST_RETRIEVED_ATTR} < :floor)" + ), + ExpressionAttributeValues={":now": _iso(now), ":floor": floor}, + ) + return True + except Exception as exc: + code = getattr(exc, "response", {}).get("Error", {}).get("Code") + if code == "ConditionalCheckFailedException": + # Fresh enough, or there is no KB_Record. Both are ordinary. + return False + logger.warning(f"could not record lastRetrievedAt for kb {app_kb_id}: {exc}") + return False + + +def schedule_activity_touch(assistant_id: str, app_kb_id: str) -> None: + """Run :func:`touch_last_retrieved` off the request path. Never raises. + + Retrieval must not wait for a bookkeeping write, nor for its rejection — and + rejection is the *common* case, since at most one write per throttle window + lands. The write goes to a thread because boto3 is synchronous and blocking the + event loop would make every other in-flight turn pay for it. + + Fire-and-forget with a strong reference held until completion, the same pattern + and the same reason as the dual-read pilot: ``create_task`` returns the only + reference, and dropping it lets the task be collected mid-flight, which shows up + as timestamps that silently stop being recorded under load. + + Falls back to doing nothing at all when there is no running loop. A missing + idleness sample is a gap in a baseline metric; an exception here would be a + failed retrieval. + """ + import asyncio + + async def _touch() -> None: + try: + await asyncio.to_thread(touch_last_retrieved, assistant_id, app_kb_id) + except Exception as exc: # noqa: BLE001 - observability only + logger.debug(f"lastRetrievedAt touch skipped for {app_kb_id}: {exc}") + + try: + task = asyncio.create_task(_touch()) + except RuntimeError: + return + + _IN_FLIGHT.add(task) + task.add_done_callback(_IN_FLIGHT.discard) + + +def bound_agent_ids(assistant_id: str, record: Optional[Mapping[str, Any]] = None) -> list: + """The agents bound to this knowledge base. + + One, this phase, and it is the assistant itself (Requirement 6.5). Written as a + list so that the caller below is a maximum over a set today and stays one when + F4 makes the set bigger — the alternative is a single lookup that has to be + rewritten, in the module whose whole point is not to under-report activity. + """ + return [assistant_id] + + +def agent_last_used_at(assistant_id: str) -> Optional[str]: + """The assistant's ``lastUsedAt``, read from its ``METADATA`` row. + + Raw table access rather than the assistants service, for this package's usual + reason: importing ``apis.shared.assistants`` pulls the embeddings stack into a + size-constrained Lambda image. + """ + try: + response = _table().get_item(Key={"PK": f"AST#{assistant_id}", "SK": "METADATA"}) + except Exception as exc: + logger.warning(f"could not read lastUsedAt for assistant {assistant_id}: {exc}") + return None + item = response.get("Item") or {} + for key in ("lastUsedAt", "updatedAt", "createdAt"): + value = item.get(key) + if value: + return str(value) + return None + + +def last_activity_at( + assistant_id: str, + record: Optional[Mapping[str, Any]] = None, + agent_timestamps: Optional[Iterable[Optional[str]]] = None, +) -> Optional[str]: + """The most recent sign of life: retrieval **or** agent use (Requirement 22.5). + + ``agent_timestamps`` lets a caller sweeping many knowledge bases supply values + it has already read instead of paying a ``get_item`` per knowledge base. When + omitted, the bound agents are read here. + + ``None`` means nothing is known — no retrieval recorded and no agent timestamp. + That is **not** the same as "idle since the beginning of time", and callers must + not treat it as such: it is what a knowledge base provisioned an hour ago looks + like. :func:`idle_days` returns ``None`` for it rather than a large number. + """ + candidates = [str((record or {}).get(LAST_RETRIEVED_ATTR) or "") or None] + + if agent_timestamps is None: + candidates.extend( + agent_last_used_at(agent_id) + for agent_id in bound_agent_ids(assistant_id, record) + ) + else: + candidates.extend(agent_timestamps) + + known = [value for value in candidates if value] + if not known: + return None + # ISO-8601 UTC strings compare correctly lexicographically, which is why every + # timestamp in this feature is written in that exact form. + return max(known) + + +def idle_days( + assistant_id: str, + record: Optional[Mapping[str, Any]] = None, + agent_timestamps: Optional[Iterable[Optional[str]]] = None, + now: Optional[str] = None, +) -> Optional[float]: + """Days since the last sign of life, or ``None`` if nothing is known. + + ``None`` rather than a default is the whole point: a knowledge base with no + recorded activity is unmeasured, not maximally idle, and a metric that reported + "very idle" for every freshly provisioned corpus would be exactly the training + signal that makes operators stop reading it. + """ + from datetime import datetime, timezone + + latest = last_activity_at(assistant_id, record, agent_timestamps) + if not latest: + return None + + try: + parsed = datetime.strptime(latest, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc) + except ValueError: + try: + parsed = datetime.fromisoformat(latest.replace("Z", "+00:00")) + except ValueError: + logger.warning(f"unparseable activity timestamp {latest!r}; treating as unknown") + return None + + reference = ( + datetime.strptime(now, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc) + if now + else _now() + ) + return max((reference - parsed).total_seconds() / 86400.0, 0.0) + + +__all__ = [ + "LAST_RETRIEVED_ATTR", + "THROTTLE_HOURS", + "agent_last_used_at", + "bound_agent_ids", + "idle_days", + "last_activity_at", + "schedule_activity_touch", + "throttle_hours", + "touch_last_retrieved", +] diff --git a/backend/src/apis/shared/kb_backend/managed_backend.py b/backend/src/apis/shared/kb_backend/managed_backend.py new file mode 100644 index 000000000..5d7f6f39b --- /dev/null +++ b/backend/src/apis/shared/kb_backend/managed_backend.py @@ -0,0 +1,635 @@ +"""The Amazon Bedrock Managed Knowledge Base backend, behind the common protocol. + +Retrieval, direct ingestion and document deletion for a managed knowledge base. +Everything here is reachable only for a KB_Record that names +``retrievalEngine == "managed"``, which happens exactly once per knowledge base, +at promotion. + +Scores need no conversion — and must not get one +------------------------------------------------ +``Retrieve`` returns ``score`` as **relevance**: higher is more relevant, which is +already the protocol's canonical direction. So unlike +:mod:`~apis.shared.kb_backend.s3vectors_backend`, this adapter passes the score +through untouched. Applying the legacy adapter's ``relevance_from_distance`` +negation here would invert the ranking, and an inverted ranking raises nothing, +logs nothing and alarms nothing: retrieval keeps returning five chunks and the +answers quietly get worse. The guard is +``tests/property/test_pbt_kb_score_direction.py``. + +``managedSearchConfiguration``, never ``vectorSearchConfiguration`` +------------------------------------------------------------------ +Requirement 11.1. ``vectorSearchConfiguration`` is a real member of +``KnowledgeBaseRetrievalConfiguration`` in the service model, so it passes +client-side validation and then fails at the service with "not supported for +managed knowledge bases". Every retrieval, including the canary the ingestion +consumer runs, would fail together — loudly, but only after deploy. + +Reranking is ``MANAGED``, not ``NONE`` (Requirement 11.2). It 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 that separation is what makes the 2,000-character context cap +defensible: with a flat distribution the cap truncates chunks that were barely +distinguishable from the best one. + +Hybrid search is not configured, and no attempt is made to (Requirement 11.3). +There is no toggle; it is simply how managed search works. + +Document identity: one id, no chunk keys +---------------------------------------- +``customDocumentIdentifier`` is the platform ``document_id`` (Requirement 9.4). +That 1:1 mapping is what lets the status filter above the seam join on a +``document_id`` per chunk, and it retires the whole ``{doc_id}#{chunk_index}`` +scheme on this path — including ``delete_vector_tail`` and the chunk-shrinkage +stash (Requirement 9.6). Deletion is by document id, so there is no tail to +shrink and nothing to stash. + +Two hard limits, both server-enforced +------------------------------------- +* **10 documents per call.** The packaged service model's ``KnowledgeBaseDocuments`` + list carries ``max: 10``, and the same for ``DocumentIdentifiers``. AWS's user + guide claims 25; that claim does not apply to managed knowledge bases and was + disproven server-side. A batch of 11 fails the whole call, so the batch size is + a constant here rather than a caller's choice. +* **10 concurrent document operations per account.** Ingests and deletes share + that budget, so both go through one semaphore rather than each keeping its own. + +``StartIngestionJob`` is never called (Requirement 9.2): it is 0.1 RPS +account-wide and not adjustable, which for a bulk upload means one document every +ten seconds for the entire account. + +Import boundary +--------------- +Module-level imports are stdlib plus this package's own stdlib-only modules; +``boto3`` is imported inside the client factories. See +``tests/architecture/test_kb_backend_boundary.py``. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import weakref +from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple + +from apis.shared.kb_backend.protocol import DEFAULT_TOP_K, Chunk, DocumentSource + +logger = logging.getLogger(__name__) + +#: Requirement 9.3. Server-enforced; the model's list shape has ``max: 10``. +MAX_DOCUMENTS_PER_CALL = 10 + +#: Requirement 9.5. Ingest and delete operations share one account-wide budget of +#: 10 concurrent operations, so one semaphore covers both. +MAX_CONCURRENT_DOCUMENT_OPERATIONS = 10 + +#: Requirement 11.2. ``MANAGED`` uses the service's reranker; ``NONE`` disables +#: reranking and flattens the score distribution. +RERANKING_MODEL_TYPE = "MANAGED" + +#: The connector all of this platform's managed documents arrive through. +CONTENT_DATA_SOURCE_TYPE = "CUSTOM" + +#: Requirement 11.5. Isolation-critical filters are restricted to exact-match +#: operators. ``startsWith`` and ``stringContains`` are prefix/substring matches: +#: a filter written to isolate ``ast-1`` would also admit ``ast-10``, and the +#: over-match is invisible because the extra results look like ordinary hits. +ISOLATION_SAFE_FILTER_OPERATORS = frozenset({"equals", "in"}) + +#: Where ``customDocumentIdentifier`` surfaces on a retrieval result, in the order +#: tried. ``location.customDocumentLocation.id`` is the authoritative one for a +#: CUSTOM connector; the metadata key is a documented mirror of it. +CUSTOM_IDENTIFIER_METADATA_KEY = "x-amz-bedrock-kb-custom-document-identifier" + + +class ManagedKbError(RuntimeError): + """A managed knowledge base operation could not be performed.""" + + +class ManagedKbNotProvisioned(ManagedKbError): + """The KB_Record carries no AWS identifiers yet. + + Raised rather than provisioning inline: provisioning takes 47–124 s and this + may be a retrieval on a user's turn. The caller decides whether to wait. + """ + + +class UnsafeFilterOperator(ManagedKbError): + """A filter used an operator that cannot be trusted for isolation. + + Raised rather than silently downgraded to ``equals``, which would change the + caller's meaning, or passed through, which would widen the boundary. + """ + + +# ── Clients ────────────────────────────────────────────────────────────────── +def _region() -> str: + return os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-west-2" + + +def bedrock_agent_runtime_client(): + """The data-plane client (``Retrieve``). Imported lazily, deliberately.""" + import boto3 + + return boto3.client("bedrock-agent-runtime", region_name=_region()) + + +def bedrock_agent_client(): + """The control-plane client (ingest/delete documents). Imported lazily.""" + import boto3 + + return boto3.client("bedrock-agent", region_name=_region()) + + +# ── Concurrency bound ──────────────────────────────────────────────────────── +# +# Keyed by event loop rather than module-global, because an ``asyncio.Semaphore`` +# is bound to the loop it is first awaited on: a single module-level instance +# would break the second test (or the second worker) to use a fresh loop. A weak +# key means a finished loop's semaphore is collected with it. +_SEMAPHORES: "weakref.WeakKeyDictionary[Any, asyncio.Semaphore]" = ( + weakref.WeakKeyDictionary() +) + + +def document_operation_semaphore() -> asyncio.Semaphore: + """The shared bound on concurrent ingest/delete operations (Requirement 9.5).""" + loop = asyncio.get_running_loop() + semaphore = _SEMAPHORES.get(loop) + if semaphore is None: + semaphore = asyncio.Semaphore(MAX_CONCURRENT_DOCUMENT_OPERATIONS) + _SEMAPHORES[loop] = semaphore + return semaphore + + +# ── Payload builders ───────────────────────────────────────────────────────── +def validate_isolation_filter(retrieval_filter: Optional[Mapping[str, Any]]) -> None: + """Refuse any filter operator that is not exact-match (Requirement 11.5). + + Recurses through ``andAll`` / ``orAll`` because a compound filter is only as + safe as its least safe leaf, and a ``stringContains`` buried three levels down + is exactly the kind of thing that survives review. + """ + if not retrieval_filter: + return + + for operator, operand in retrieval_filter.items(): + if operator in ("andAll", "orAll"): + for nested in operand or []: + validate_isolation_filter(nested) + continue + if operator not in ISOLATION_SAFE_FILTER_OPERATORS: + raise UnsafeFilterOperator( + f"filter operator {operator!r} is not permitted: an " + f"isolation-critical filter must use one of " + f"{sorted(ISOLATION_SAFE_FILTER_OPERATORS)}. Prefix and substring " + f"operators over-match silently — a filter isolating 'ast-1' also " + f"admits 'ast-10', and the extra results are indistinguishable " + f"from legitimate hits." + ) + + +def retrieval_configuration( + top_k: int = DEFAULT_TOP_K, + retrieval_filter: Optional[Mapping[str, Any]] = None, +) -> Dict[str, Any]: + """Build ``retrievalConfiguration`` for a managed knowledge base. + + ``managedSearchConfiguration`` only. There is no branch that could produce + ``vectorSearchConfiguration``, so it cannot be reintroduced by a stray + condition — only by editing this function, which + ``tests/shared/test_managed_kb_backend.py`` notices. + """ + validate_isolation_filter(retrieval_filter) + + managed: Dict[str, Any] = { + # Requirement 3.1 parity: both backends are asked for the same number. + "numberOfResults": top_k, + "rerankingModelType": RERANKING_MODEL_TYPE, + } + if retrieval_filter: + managed["filter"] = dict(retrieval_filter) + + # No hybrid-search key: it is not configurable and not attempted (Req 11.3). + return {"managedSearchConfiguration": managed} + + +#: Bedrock caps ``DocumentMetadata.inlineAttributes`` at 50 entries (verified in the +#: packaged service model: ``{'min': 1, 'max': 50}``). Exceeding it fails the whole +#: ``IngestKnowledgeBaseDocuments`` call, so one document with chatty metadata would +#: take its entire batch of ten down with it. +MAX_INLINE_ATTRIBUTES = 50 + +#: Keys the platform depends on, kept in preference to caller-supplied ones when +#: truncating. ``document_id`` is load-bearing: the facade's status filter joins on +#: it, so a chunk that arrives without it cannot be matched to its document and +#: would be dropped as unverifiable. +_RESERVED_METADATA_KEYS = ("document_id", "filename") + + +def _inline_attributes(metadata: Mapping[str, Any]) -> List[Dict[str, Any]]: + """Metadata as ``IN_LINE_ATTRIBUTE`` entries, string-valued and bounded. + + Only strings are emitted. Mixed attribute types are a per-key commitment on + Bedrock's side, and the platform's metadata is loosely typed, so coercing + everything to a string keeps a stray ``None`` or ``int`` from poisoning a key + for every future document. + + The list is capped at :data:`MAX_INLINE_ATTRIBUTES`. ``source.metadata`` is + caller-supplied and unbounded, so without this a caller could fail an entire + ten-document batch with one over-decorated document. Reserved keys are emitted + first so truncation cannot drop them — a plain ``sorted()`` would drop by + alphabet, and ``document_id`` sorts after several plausible caller keys. + """ + ordered: List[tuple] = [] + seen = set() + + for key in _RESERVED_METADATA_KEYS: + if key in metadata and metadata[key] is not None: + ordered.append((key, metadata[key])) + seen.add(key) + + for key, value in sorted(metadata.items()): + if key in seen or value is None: + continue + ordered.append((key, value)) + + if len(ordered) > MAX_INLINE_ATTRIBUTES: + dropped = len(ordered) - MAX_INLINE_ATTRIBUTES + logger.warning( + f"document metadata has {len(ordered)} attributes; keeping the first " + f"{MAX_INLINE_ATTRIBUTES} and dropping {dropped} " + f"(Bedrock's inlineAttributes limit)" + ) + ordered = ordered[:MAX_INLINE_ATTRIBUTES] + + return [ + {"key": key, "value": {"type": "STRING", "stringValue": str(value)}} + for key, value in ordered + ] + + +def document_payload( + source: DocumentSource, + *, + bucket: Optional[str] = None, +) -> Dict[str, Any]: + """One entry of the ``documents`` array. + + ``customDocumentIdentifier`` is the platform ``document_id`` verbatim + (Requirement 9.4) — not a derived or prefixed form. It is the join key the + status filter needs and the handle deletion uses, so any transformation here + would have to be reversed in two other places. + + Prefers the S3 location when the source has one: the object is already in the + documents bucket, and Bedrock reading it directly avoids pulling the bytes + through this process. Falls back to inline text for a source that only has + chunks, joining them back into a document because managed ingestion does its + own chunking and pre-chunked input would be re-chunked anyway. + """ + identifier = {"id": source.document_id} + custom: Dict[str, Any] = { + "customDocumentIdentifier": identifier, + } + + if source.s3_key: + resolved = bucket or os.environ.get("S3_ASSISTANTS_DOCUMENTS_BUCKET_NAME") + if not resolved: + raise ManagedKbError( + f"document {source.document_id} has an S3 key but no bucket: pass " + f"bucket= or set S3_ASSISTANTS_DOCUMENTS_BUCKET_NAME" + ) + custom["sourceType"] = "S3_LOCATION" + custom["s3Location"] = {"uri": f"s3://{resolved}/{source.s3_key}"} + elif source.chunks: + custom["sourceType"] = "IN_LINE" + custom["inlineContent"] = { + "type": "TEXT", + "textContent": {"data": "\n\n".join(source.chunks)}, + } + else: + raise ManagedKbError( + f"document {source.document_id} has neither an s3_key nor chunks; " + f"there is nothing to ingest" + ) + + document: Dict[str, Any] = { + "content": {"dataSourceType": CONTENT_DATA_SOURCE_TYPE, "custom": custom} + } + + metadata = {"document_id": source.document_id, "filename": source.filename} + metadata.update({k: v for k, v in source.metadata.items() if k not in metadata}) + attributes = _inline_attributes(metadata) + if attributes: + document["metadata"] = { + "type": "IN_LINE_ATTRIBUTE", + "inlineAttributes": attributes, + } + return document + + +def document_identifier(document_id: str) -> Dict[str, Any]: + """One entry of ``documentIdentifiers`` for a delete. + + Deletion is by the platform document id, full stop. There is no chunk tail to + enumerate and no shrinkage case to handle, because one document is one + document (Requirement 9.6). + """ + return { + "dataSourceType": CONTENT_DATA_SOURCE_TYPE, + "custom": {"id": document_id}, + } + + +def batched(items: Sequence[Any], size: int = MAX_DOCUMENTS_PER_CALL) -> List[List[Any]]: + """Split ``items`` into batches of at most ``size``. + + ``size`` is validated against the server limit rather than trusted. A caller + passing 25 — the number AWS's user guide gives, which does not apply to managed + knowledge bases — would otherwise produce a request that fails as a whole, + losing the other 24 documents along with the 25th. + """ + if size < 1: + raise ValueError("batch size must be at least 1") + if size > MAX_DOCUMENTS_PER_CALL: + raise ValueError( + f"batch size {size} exceeds the server-enforced maximum of " + f"{MAX_DOCUMENTS_PER_CALL} documents per call. AWS's user guide claims " + f"25; that does not apply to managed knowledge bases and an " + f"11-document call fails entirely." + ) + return [list(items[i : i + size]) for i in range(0, len(items), size)] + + +# ── The backend ────────────────────────────────────────────────────────────── +class ManagedKbBackend: + """Retrieval and direct ingestion against a Managed Knowledge Base. + + ``kb_ref`` is the ``App_KB_Id`` (equal to the ``assistant_id`` in this phase), + never an AWS ``knowledgeBaseId``. The AWS identifiers are resolved internally + from the KB_Record on each operation, because a dormancy/rehydration cycle + replaces them and a caller holding one would keep querying a knowledge base + that no longer exists. + + Clients are injectable so tests can stub them; they are created lazily so + constructing this class costs nothing at import time. + """ + + def __init__( + self, + *, + runtime_client=None, + agent_client=None, + locator=None, + bucket: Optional[str] = None, + ) -> None: + self._runtime_client = runtime_client + self._agent_client = agent_client + self._locator = locator + self._bucket = bucket + + # ── plumbing ──────────────────────────────────────────────────────────── + def _runtime(self): + if self._runtime_client is None: + self._runtime_client = bedrock_agent_runtime_client() + return self._runtime_client + + def _agent(self): + if self._agent_client is None: + self._agent_client = bedrock_agent_client() + return self._agent_client + + def _locate(self, kb_ref: str) -> Tuple[str, str]: + """Resolve ``kb_ref`` to ``(awsKbId, awsDataSourceId)``.""" + if self._locator is not None: + located = self._locator(kb_ref) + else: + from apis.shared.kb_backend.records import get_kb_record + + # App_KB_Id == assistant_id in this phase, so one value serves both. + item = get_kb_record(kb_ref, kb_ref) + located = ( + (item.get("awsKbId"), item.get("awsDataSourceId")) if item else (None, None) + ) + + aws_kb_id, aws_data_source_id = located + if not aws_kb_id: + raise ManagedKbNotProvisioned( + f"knowledge base {kb_ref} has no awsKbId: it is not provisioned yet" + ) + return aws_kb_id, aws_data_source_id + + async def _locate_async(self, kb_ref: str) -> Tuple[str, str]: + return await asyncio.to_thread(self._locate, kb_ref) + + # ── retrieval ─────────────────────────────────────────────────────────── + async def search( + self, + kb_ref: str, + query: str, + top_k: int = DEFAULT_TOP_K, + retrieval_filter: Optional[Mapping[str, Any]] = None, + ) -> List[Chunk]: + """Retrieve up to ``top_k`` chunks, best first. + + Ordering is the service's. ``Retrieve`` returns results ranked best-first + and this returns them in that order, with the reported ``score`` as the + canonical ``relevance`` — unconverted, because the directions already + agree. + + The synchronous ``retrieve`` call runs off the event loop (Requirement + 20.7): it was measured at 662–695 ms p50, which is long enough to matter + to every other coroutine sharing the loop. + """ + aws_kb_id, _ = await self._locate_async(kb_ref) + client = self._runtime() + + payload = { + "knowledgeBaseId": aws_kb_id, + "retrievalQuery": {"text": query}, + "retrievalConfiguration": retrieval_configuration(top_k, retrieval_filter), + } + response = await asyncio.to_thread(lambda: client.retrieve(**payload)) + return [self._to_chunk(result) for result in response.get("retrievalResults", [])] + + @staticmethod + def _to_chunk(result: Mapping[str, Any]) -> Chunk: + """Adapt one ``Retrieve`` result. **No score conversion.** + + ``score`` is relevance already: higher is better, which is the protocol's + direction. A missing score stays ``None`` rather than becoming ``0.0``, + for the same reason as in the legacy adapter — a fabricated ``0.0`` would + be indistinguishable from a real score, and on this backend it would rank + the chunk *last* while on the other it would rank first. + """ + metadata = dict(result.get("metadata") or {}) + text = (result.get("content") or {}).get("text", "") + document_id = ManagedKbBackend._document_id(result, metadata) + + # The status filter and the citation formatter above the seam both read + # `metadata["text"]`, which is where the legacy path put it. + metadata.setdefault("text", text) + metadata.setdefault("document_id", document_id) + + return Chunk( + text=text, + relevance=result.get("score"), + document_id=document_id, + metadata=metadata, + key=document_id, + ) + + @staticmethod + def _document_id(result: Mapping[str, Any], metadata: Mapping[str, Any]) -> str: + """Recover the platform ``document_id`` from a retrieval result. + + Deliberately does **not** fall back to the result's own ``documentId``: + that is a service-assigned handle for ``GetDocumentContent``, not the + platform id, and returning it would produce a chunk whose ``document_id`` + looks plausible, joins against no ``DOC#`` record, and is dropped by the + fail-closed status filter — a disappearing-results bug two layers away + from its cause. + """ + location = result.get("location") or {} + custom = location.get("customDocumentLocation") or {} + for candidate in ( + custom.get("id"), + metadata.get(CUSTOM_IDENTIFIER_METADATA_KEY), + metadata.get("document_id"), + ): + if candidate: + return str(candidate) + return "" + + # ── ingestion ─────────────────────────────────────────────────────────── + async def ingest(self, kb_ref: str, source: DocumentSource) -> None: + """Index one document.""" + await self.ingest_documents(kb_ref, [source]) + + async def ingest_documents( + self, + kb_ref: str, + sources: Iterable[DocumentSource], + *, + batch_size: int = MAX_DOCUMENTS_PER_CALL, + ) -> None: + """Index documents with ``IngestKnowledgeBaseDocuments``. + + Batched at 10 and concurrency-bounded at 10 (Requirements 9.3, 9.5). + ``StartIngestionJob`` is never involved (Requirement 9.2). + + No ``clientToken`` is sent, on purpose. Idempotency here comes from + ``customDocumentIdentifier`` being 1:1 with the platform document id, so + re-ingesting a document replaces it. A token derived from the document ids + would look like extra safety and instead swallow a legitimate re-upload of + the same document — silently, since a deduplicated request returns + success. + """ + documents = list(sources) + if not documents: + return + + aws_kb_id, aws_data_source_id = await self._locate_async(kb_ref) + if not aws_data_source_id: + raise ManagedKbNotProvisioned( + f"knowledge base {kb_ref} has no awsDataSourceId: its CUSTOM " + f"connector is not created yet" + ) + + client = self._agent() + payloads = [document_payload(source, bucket=self._bucket) for source in documents] + + await self._run_bounded( + [ + { + "knowledgeBaseId": aws_kb_id, + "dataSourceId": aws_data_source_id, + "documents": batch, + } + for batch in batched(payloads, batch_size) + ], + client.ingest_knowledge_base_documents, + what="IngestKnowledgeBaseDocuments", + ) + + # ── deletion ──────────────────────────────────────────────────────────── + async def delete_document(self, kb_ref: str, document_id: str) -> None: + """Remove one document by its platform id.""" + await self.delete_documents(kb_ref, [document_id]) + + async def delete_documents( + self, + kb_ref: str, + document_ids: Iterable[str], + *, + batch_size: int = MAX_DOCUMENTS_PER_CALL, + ) -> None: + """Remove documents with ``DeleteKnowledgeBaseDocuments``. + + Same batch limit and same shared concurrency budget as ingestion: the + account's 10-concurrent-operation limit counts both together, so deletes + issued alongside ingests must not each get their own allowance. + """ + ids = [document_id for document_id in document_ids if document_id] + if not ids: + return + + aws_kb_id, aws_data_source_id = await self._locate_async(kb_ref) + if not aws_data_source_id: + raise ManagedKbNotProvisioned( + f"knowledge base {kb_ref} has no awsDataSourceId; nothing to delete from" + ) + + client = self._agent() + await self._run_bounded( + [ + { + "knowledgeBaseId": aws_kb_id, + "dataSourceId": aws_data_source_id, + "documentIdentifiers": batch, + } + for batch in batched([document_identifier(i) for i in ids], batch_size) + ], + client.delete_knowledge_base_documents, + what="DeleteKnowledgeBaseDocuments", + ) + + @staticmethod + async def _run_bounded(payloads: Sequence[Mapping[str, Any]], operation, *, what: str) -> None: + """Issue each payload off the event loop, at most 10 in flight. + + The semaphore is acquired *around* the ``to_thread`` call so the bound + counts operations in flight at AWS, not coroutines created here — which is + the number the account limit is expressed in. + """ + semaphore = document_operation_semaphore() + + async def _one(payload: Mapping[str, Any]) -> None: + async with semaphore: + await asyncio.to_thread(lambda: operation(**payload)) + + results = await asyncio.gather( + *(_one(payload) for payload in payloads), return_exceptions=True + ) + failures = [outcome for outcome in results if isinstance(outcome, BaseException)] + if failures: + logger.error(f"{what}: {len(failures)}/{len(payloads)} batches failed") + raise failures[0] + + +__all__ = [ + "CONTENT_DATA_SOURCE_TYPE", + "ISOLATION_SAFE_FILTER_OPERATORS", + "MAX_CONCURRENT_DOCUMENT_OPERATIONS", + "MAX_DOCUMENTS_PER_CALL", + "RERANKING_MODEL_TYPE", + "ManagedKbBackend", + "ManagedKbError", + "ManagedKbNotProvisioned", + "UnsafeFilterOperator", + "batched", + "document_identifier", + "document_operation_semaphore", + "document_payload", + "retrieval_configuration", + "validate_isolation_filter", +] diff --git a/backend/src/apis/shared/kb_backend/metrics.py b/backend/src/apis/shared/kb_backend/metrics.py new file mode 100644 index 000000000..a8da087d8 --- /dev/null +++ b/backend/src/apis/shared/kb_backend/metrics.py @@ -0,0 +1,197 @@ +"""Best-effort custom metrics for the managed knowledge base seam. + +Metric publishing is *observability*, never control flow. Every function here +swallows its own failures: a knowledge base search must not fail because +CloudWatch was briefly unavailable, and a missing metric is a monitoring gap, not +a user-facing error. This matches the convention in +``apis/app_api/kb_sync/dispatcher.py``. + +Namespace +--------- +The namespace must match what the IAM grant allows, or every publish is silently +denied. ``ManagedKbRoleConstruct`` conditions ``cloudwatch:PutMetricData`` on +``cloudwatch:namespace`` equal to ``{projectPrefix}/ManagedKb``, and derives it +from the same ``projectPrefix`` that CDK injects here as ``PROJECT_PREFIX``. The +two must therefore agree; :func:`metric_namespace` is the single reader of that +variable so there is one place to look when they do not. + +Not an ``AWS/...`` namespace, deliberately: CloudWatch reserves every namespace +beginning with ``AWS`` for its own services and rejects writes to them, so such a +grant would authorize nothing while looking correct. Bedrock's own +``AWS/Bedrock/KnowledgeBases`` metrics are a read source, never a write target. + +Import weight +------------- +``boto3`` is imported inside :func:`emit_count`, so importing this module costs +nothing. The seam is imported by a size-constrained Lambda image. +""" + +from __future__ import annotations + +import logging +import os +from typing import Mapping, Optional + +logger = logging.getLogger(__name__) + +#: Emitted when a query was longer than the hard cap and had to be truncated +#: (Requirement 22.3). A non-zero value means users are sending queries that the +#: managed backend would reject outright, which is worth knowing before the +#: engine is switched under them. +METRIC_QUERY_CLAMPED = "KbQueryClamped" + +#: Emitted when the document-status filter could not confirm status and therefore +#: dropped chunks (Requirement 22.4). Distinct from an ordinary empty result: this +#: one means retrieval degraded, not that the corpus had no match. +METRIC_STATUS_FILTER_FAIL_CLOSED = "KbStatusFilterFailClosed" + +#: Emitted when retrieval was refused because the invoking user's access could not +#: be established (Requirement 25.1). Counts both honest denials and +#: check-failed-so-denied, dimensioned by ``reason`` to keep them apart: the first +#: is the system working, the second is a degradation worth alarming on. +METRIC_ACCESS_DENIED = "KbAccessDenied" + +#: Dual-read pilot observations (Requirement 18.3). Values, not counts: the +#: question each answers is "how much do the two backends agree, and at what +#: cost", and a count cannot answer either. +METRIC_DUAL_READ_OVERLAP = "KbDualReadOverlap" +METRIC_DUAL_READ_RANK_CORRELATION = "KbDualReadRankCorrelation" +METRIC_DUAL_READ_LATENCY = "KbDualReadLatency" + +#: Emitted when the observational managed read failed. Never a user-facing +#: failure — the turn was served from legacy before the comparison ran — but a +#: sustained non-zero value is the pilot telling us the engine is not ready. +METRIC_DUAL_READ_FAILED = "KbDualReadFailed" + +#: Fleet gauges (Requirement 22.1), emitted once per reconciler pass rather than +#: per event, because each is a statement about the whole account. +METRIC_KB_COUNT = "KbCount" +METRIC_KB_STORAGE_GB = "KbStorageGB" +METRIC_KB_IDLE_GB = "KbIdleGB" + +#: Days without a sign of life before a knowledge base's bytes count toward +#: :data:`METRIC_KB_IDLE_GB`. A reporting threshold only: nothing reclaims in this +#: phase, and the number the follow-up spec eventually evicts on should be chosen +#: from the distribution this metric records, not inherited from this guess. +IDLE_THRESHOLD_DAYS = 30 + +#: Bytes per gigabyte, decimal — matching how AWS bills storage ($5.00/GB-month), +#: so a dashboard number and an invoice line can be compared without a conversion +#: nobody remembers to apply. +BYTES_PER_GB = 1_000_000_000 + + +def emit_fleet_gauges( + kb_count: int, + stored_bytes: int, + idle_bytes: int, + *, + unmeasured: int = 0, + idle_threshold_days: Optional[int] = None, +) -> None: + """Publish the account-wide knowledge base gauges. Never raises. + + Requirement 22.1. Emitted through EMF rather than ``PutMetricData`` because the + caller is a Lambda whose stdout already reaches CloudWatch Logs, so this needs + no client, no batching and no IAM — and because these are gauges published once + per pass, which is exactly the shape EMF is good at. The namespace is the same + :func:`metric_namespace` the ``PutMetricData`` grant is conditioned on, so both + mechanisms land in one place and a dashboard does not have to know which code + path produced a number. + + ``unmeasured`` rides along as a log property, not a metric: it is the count of + knowledge bases with no recorded activity at all, which is context for reading + ``KbIdleGB`` rather than something to alarm on. Emitting it as a metric would + invite an alarm on a number that is legitimately large the day this ships and + legitimately near zero a month later. + """ + try: + from apis.shared.observability.emf import emit_emf_metrics + + emit_emf_metrics( + metric_namespace(), + { + METRIC_KB_COUNT: int(kb_count), + METRIC_KB_STORAGE_GB: round(stored_bytes / BYTES_PER_GB, 6), + METRIC_KB_IDLE_GB: round(idle_bytes / BYTES_PER_GB, 6), + }, + properties={ + "unmeasuredKnowledgeBases": int(unmeasured), + "idleThresholdDays": int( + IDLE_THRESHOLD_DAYS if idle_threshold_days is None else idle_threshold_days + ), + }, + units={ + METRIC_KB_COUNT: "Count", + METRIC_KB_STORAGE_GB: "Gigabytes", + METRIC_KB_IDLE_GB: "Gigabytes", + }, + ) + except Exception as exc: # noqa: BLE001 - observability must not break a sweep + logger.warning(f"Failed to emit knowledge base fleet gauges: {exc}") + + +def metric_namespace() -> str: + """The custom namespace this feature publishes into. + + Prefers ``MANAGED_KB_METRIC_NAMESPACE``, which the CDK construct sets from the + *same* helper that builds the IAM condition — so where that variable is present + the grant and the publish cannot disagree. Falls back to deriving from + ``PROJECT_PREFIX`` for services that do not receive it and for local runs. + """ + explicit = os.environ.get("MANAGED_KB_METRIC_NAMESPACE") + if explicit: + return explicit + prefix = os.environ.get("PROJECT_PREFIX", "agentcore") + return f"{prefix}/ManagedKb" + + +def emit_count( + metric_name: str, + value: int = 1, + dimensions: Optional[Mapping[str, str]] = None, +) -> None: + """Publish a single count metric. Never raises. + + Swallowing the failure is the point: the caller is on a request path, and a + metric that cannot be published is strictly less important than the answer the + user is waiting for. + """ + _publish(metric_name, value, "Count", dimensions) + + +def emit_value( + metric_name: str, + value: float, + unit: str = "None", + dimensions: Optional[Mapping[str, str]] = None, +) -> None: + """Publish a measurement rather than an occurrence. Never raises. + + Separate from :func:`emit_count` so the unit is a decision at the call site. + A latency published as ``Count`` is not merely mislabelled — CloudWatch will + graph and alarm on it as a rate, and the mistake is invisible until somebody + tries to read the dashboard. + """ + _publish(metric_name, value, unit, dimensions) + + +def _publish( + metric_name: str, + value: float, + unit: str, + dimensions: Optional[Mapping[str, str]], +) -> None: + try: + import boto3 + + datum: dict = {"MetricName": metric_name, "Value": value, "Unit": unit} + if dimensions: + datum["Dimensions"] = [ + {"Name": k, "Value": v} for k, v in sorted(dimensions.items()) + ] + boto3.client("cloudwatch").put_metric_data( + Namespace=metric_namespace(), MetricData=[datum] + ) + except Exception as exc: # noqa: BLE001 - observability must not break retrieval + logger.warning(f"Failed to emit {metric_name} metric: {exc}") diff --git a/backend/src/apis/shared/kb_backend/protocol.py b/backend/src/apis/shared/kb_backend/protocol.py new file mode 100644 index 000000000..9f2955c32 --- /dev/null +++ b/backend/src/apis/shared/kb_backend/protocol.py @@ -0,0 +1,155 @@ +"""The one seam every knowledge base read and write passes through. + +Two backends sit behind :class:`KnowledgeBaseBackend`: the legacy Amazon S3 +Vectors implementation and, from task 8, Amazon Bedrock Managed Knowledge Base. +Callers never learn which one they got. + +Score direction +--------------- +This module's single most important decision is the name of one field. + +The two backends disagree about which direction is better: + +* S3 Vectors returns cosine **distance** — *lower* is more similar. +* Managed KB returns **relevance** — *higher* is more relevant. + +Get that backwards and nothing raises. No log line, no alarm, no failing +request: retrieval simply serves the least relevant chunks it can find, and the +only symptom is that answers get worse. There is no error path to catch because +there is no error. That is why the canonical field is named ``relevance`` and +documented here rather than left implicit, why the legacy adapter converts +*inside itself* rather than at some call site, and why +``tests/property/test_pbt_kb_score_direction.py`` exists at all. + +Exact round-trip +---------------- +:func:`relevance_from_distance` and :func:`distance_from_relevance` are exact +inverses under IEEE-754, not approximate ones, because the facade still emits a +``distance`` key derived from ``relevance`` and +``app_api/assistants/routes.py`` puts that value in an HTTP response body. A +``1.0 - x`` conversion would round-trip ``0.1`` to ``0.09999999999999998`` and +change a value a client already reads. Negation is exact for every finite float, +so the derived value is the same value, not a very close one. + +Import boundary +--------------- +Module-level imports here are **stdlib only**. ``apis.shared.kb_backend`` is +bundled into size-constrained Lambda images and must not drag in +``apis.shared.assistants`` (whose ``__init__`` imports the embeddings stack) or +``boto3``. Enforced by ``tests/architecture/test_kb_backend_boundary.py``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Protocol, runtime_checkable + +#: Parity contract, Requirement 3.1: both backends are asked for five chunks. +#: Named here, above the seam, so neither adapter can drift from the other. +DEFAULT_TOP_K = 5 + + +def relevance_from_distance(distance: Optional[float]) -> Optional[float]: + """Convert a cosine **distance** (lower is better) to **relevance**. + + Negation, deliberately, rather than ``1.0 - distance``: it inverts the + direction — which is the entire job — while being exactly reversible by + :func:`distance_from_relevance` for every finite float. Nothing compares + relevance values *across* backends (the dual-read pilot compares rank order, + not magnitude), so the absolute range is free and exactness is not. + + ``None`` passes through as ``None``. The S3 Vectors query always asks for + distances, so a missing one means a malformed response; the facade's + long-standing behaviour is to emit ``distance: None`` rather than invent a + score, and fabricating ``0.0`` here would promote such a chunk to + best-in-class. + """ + if distance is None: + return None + return -distance + + +def distance_from_relevance(relevance: Optional[float]) -> Optional[float]: + """Recover the original distance from a relevance. Exact inverse of above.""" + if relevance is None: + return None + return -relevance + + +@dataclass(frozen=True) +class Chunk: + """One retrieved passage, in the shape every backend must produce. + + Frozen because a chunk crosses the seam as a value: an adapter that returned + something a caller could mutate would let ranking be edited after the + backend had decided it. + """ + + text: str + + #: Canonical score. **HIGHER IS MORE RELEVANT**, on both backends, always. + #: The legacy adapter has already converted S3 Vectors' inverted distance by + #: the time a chunk exists. ``None`` means the backend reported no score + #: (see :func:`relevance_from_distance`); it is preserved, never defaulted, + #: because a default would be indistinguishable from a real score. + relevance: Optional[float] + + #: Platform document id. The status filter joins on this, and on the managed + #: path it is the ``customDocumentIdentifier`` (task 8.4). + document_id: str + + metadata: Dict[str, Any] = field(default_factory=dict) + + #: Backend-native identifier: ``{document_id}#{chunk_index}`` on legacy. + key: str = "" + + +@dataclass(frozen=True) +class DocumentSource: + """A document to ingest, in the least-common-denominator form. + + The two backends want different things — legacy needs text already chunked, + managed needs the source bytes or an S3 location — so both are optional here + and each adapter validates what it needs. Tasks 8.4 and 9.1 extend this; + it exists now only so the protocol's ``ingest`` signature is real rather + than ``Any``. + """ + + document_id: str + filename: str + chunks: Optional[List[str]] = None + s3_key: Optional[str] = None + metadata: Dict[str, Any] = field(default_factory=dict) + + +@runtime_checkable +class KnowledgeBaseBackend(Protocol): + """What both backends implement, and all a caller may assume. + + ``kb_ref`` is the application-owned reference to the knowledge base — the + ``App_KB_Id``, which equals the ``assistant_id`` in this phase. Never an AWS + ``knowledgeBaseId``: those are replaceable across a dormancy/rehydration + cycle, so an adapter resolves one internally and no caller holds one. + + ``runtime_checkable`` supports ``isinstance`` structural assertions in the + tests. It checks method *presence* only, never signatures, so it is a + guard against a missing method, not a substitute for reading the protocol. + """ + + async def search(self, kb_ref: str, query: str, top_k: int = DEFAULT_TOP_K) -> List[Chunk]: + """Return up to ``top_k`` chunks, best first, scored by relevance. + + Ordering is the backend's: both underlying APIs return results ranked + best-first, and an adapter re-sorting them would be inventing a ranking + rather than reporting one. What an adapter *must* guarantee is that its + ``relevance`` values agree with the order it returns. + """ + ... + + async def ingest(self, kb_ref: str, source: DocumentSource) -> None: + """Index ``source`` into the knowledge base.""" + ... + + async def delete_document(self, kb_ref: str, document_id: str) -> None: + """Remove every trace of ``document_id`` from the knowledge base.""" + ... diff --git a/backend/src/apis/shared/kb_backend/provisioning.py b/backend/src/apis/shared/kb_backend/provisioning.py new file mode 100644 index 000000000..4bdda009f --- /dev/null +++ b/backend/src/apis/shared/kb_backend/provisioning.py @@ -0,0 +1,662 @@ +"""Lazy provisioning of a Managed Knowledge Base, and its CUSTOM connector. + +A knowledge base is created in AWS the first time a document is actually ready to +be indexed — never when an assistant is created (Requirement 7.1). Creation was +measured at 47–124 s to ``ACTIVE`` (n=7), so this never sits on an interactive +path, and every call here is issued off the event loop (Requirement 20.7). + +Order of operations, which is the whole design +---------------------------------------------- +The KB_Record is written in ``provisioning`` **before** the first AWS call and the +returned identifiers are attached afterwards with a conditional update +(Requirement 7.3). Reversing those two steps looks harmless and is not: a crash +between ``CreateKnowledgeBase`` returning and the record being written would leave +a billed AWS resource that no record points at and no code can find. Nothing +raises, nothing alarms, and the only evidence is the invoice. + +Written record-first, the same crash leaves a ``provisioning`` record that is a +durable **retry anchor** (Requirement 7.8): the Reconciler can match it against +the orphan and adopt it, and a plain retry of this function reuses the persisted +``clientToken`` so AWS deduplicates rather than creating a second knowledge base. + +Four details that are each a defect if omitted +---------------------------------------------- +1. **The ``clientToken`` is built, not interpolated.** The API's minimum is **33 + characters** — verified in the packaged service model, ``ClientToken`` has + ``min: 33``, ``max: 256``, ``pattern: [a-zA-Z0-9](-*[a-zA-Z0-9]){0,256}``. The + natural ``{id}-{variant}-kb`` template is 31 characters and fails *client-side* + validation, before a request is ever sent. :func:`build_client_token` + constructs one that cannot be too short, and :func:`validate_client_token` + refuses one that is. + +2. **"Unable to verify the specified embedding model" is retryable.** It was + observed against a model confirmed ``ACTIVE`` and directly invokable: it is IAM + eventual consistency, not a configuration error. Treated as fatal, lazy + provisioning fails intermittently while pointing at the wrong cause — an + operator reads the message and goes to check the model. + +3. **``dataDeletionPolicy: RETAIN`` at creation.** The documented remedy for the + ``DELETE_UNSUCCESSFUL`` state, which the dev account has already been sitting + in since 2025-11-24. Set deliberately up front, not as incident response. + +4. **``imageExtractionStatus: ENABLED``.** Opt-in. Left at its default, chart and + image content is never described and never indexed — a silent loss of a + capability being paid for, with no error anywhere. + +Why ``storageConfiguration`` is absent +-------------------------------------- +There is no vector store to provision: that is the point of a managed knowledge +base. Sending ``storageConfiguration`` at all is rejected (Requirement 8.2), so it +is omitted entirely rather than passed empty. + +Reconciling "``managedKnowledgeBaseConfiguration={}``" with the embedding pin +---------------------------------------------------------------------------- +Requirement 8.1 records that ``managedKnowledgeBaseConfiguration`` has **no +required members**, so ``{}`` is a valid value. Requirement 8.5 separately pins +``embeddingModelType: CUSTOM`` to ``amazon.titan-embed-text-v2:0`` at float32 and +1024 dimensions. Those are not in conflict: the packaged service model puts +``embeddingModelType`` / ``embeddingModelArn`` / +``embeddingModelConfiguration.bedrockEmbeddingModelConfiguration`` inside +``managedKnowledgeBaseConfiguration`` as optional members, so the pin goes there. +It is pinned rather than left to the service default because the choice is +**immutable after creation** (Requirement 8.8) and because keeping today's Titan +v2 embedding preserves continuity with the legacy corpus. + +Import boundary +--------------- +Module-level imports are stdlib plus this package's own stdlib-only modules. +``boto3`` is imported inside the function that builds a client, so importing this +module into a size-constrained Lambda image costs nothing. See +``tests/architecture/test_kb_backend_boundary.py``. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import os +import re +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any, Awaitable, Callable, Dict, Mapping, Optional + +from apis.shared.kb_backend.metrics import emit_count + +logger = logging.getLogger(__name__) + +# ── Immutable embedding configuration (Requirement 8.5, 8.8) ───────────────── +# +# Immutable after creation: Bedrock rejects a change, so a drift here is not a +# migration but a rebuild. Recorded on the KB_Record too, so a mismatch is +# detectable rather than mysterious. +EMBEDDING_MODEL_ID = "amazon.titan-embed-text-v2:0" +EMBEDDING_DIMENSIONS = 1024 + +#: The service model's ``EmbeddingDataType`` enum is ``['FLOAT32', 'BINARY']`` — +#: upper case. "float32" is rejected. +EMBEDDING_DATA_TYPE = "FLOAT32" +EMBEDDING_MODEL_TYPE = "CUSTOM" + +# ── Knowledge base and data source shapes ──────────────────────────────────── +KNOWLEDGE_BASE_TYPE = "MANAGED" + +#: Classic ``CUSTOM`` / ``S3`` / ``WEB`` at the top level are rejected with +#: "Unsupported data source type for MANAGED knowledge base type". The real +#: connector type nests one level down, in ``connectorParameters``. +DATA_SOURCE_TYPE = "MANAGED_KNOWLEDGE_BASE_CONNECTOR" +CONNECTOR_TYPE = "CUSTOM" +CONNECTOR_VERSION = "1" + +#: Requirement 8.7. ``RETAIN`` at creation time, not later. +DATA_DELETION_POLICY = "RETAIN" + +#: Requirement 8.6. Opt-in; the default indexes no image or chart content. +IMAGE_EXTRACTION_STATUS = "ENABLED" + +# ── clientToken (Requirement 7.5, 7.6) ────────────────────────────────────── +CLIENT_TOKEN_MIN_LENGTH = 33 +CLIENT_TOKEN_MAX_LENGTH = 256 + +#: Anchored copy of the service model's pattern. Note what it permits and does +#: not: the first and last characters must be alphanumeric, so a token may not +#: begin or end with a hyphen, though hyphens may run consecutively inside. +CLIENT_TOKEN_PATTERN = re.compile(r"^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,256}$") + +#: Length of the deterministic digest suffix. 40 hex characters alone clears the +#: 33-character minimum, so a token is long enough even when the caller's parts +#: sanitize away to nothing. +_DIGEST_CHARS = 40 + +# ── Retry classification (Requirement 7.7) ────────────────────────────────── +# +# Matched against the *message*, lower-cased, because this failure arrives as an +# ordinary validation-shaped error and is indistinguishable from a real +# misconfiguration by error code alone. +RETRYABLE_MESSAGE_FRAGMENTS = ( + "unable to verify the specified embedding model", + "unable to verify the embedding model", +) + +#: Transport-level errors that are retryable for the usual reasons. Deliberately +#: narrow: a genuine ``ValidationException`` must fail fast and loudly, because +#: retrying a malformed request just delays the report by a minute. +RETRYABLE_ERROR_CODES = ( + "ThrottlingException", + "TooManyRequestsException", + "InternalServerException", + "ServiceUnavailableException", +) + +MAX_PROVISION_ATTEMPTS = 5 +_MAX_BACKOFF_SECONDS = 30.0 + +METRIC_PROVISION_RETRIED = "KbProvisionRetried" +METRIC_PROVISION_ADOPTED = "KbProvisionAdopted" + + +class ProvisioningError(RuntimeError): + """Provisioning could not complete.""" + + +class RetryableProvisioningError(ProvisioningError): + """Provisioning failed for a reason that will plausibly clear on its own.""" + + +class ProvisioningInProgress(RetryableProvisioningError): + """Another worker owns this provisioning and has not finished. + + Retryable rather than fatal, and deliberately *not* an attempt to provision + anyway: two workers creating one knowledge base each is exactly the + duplication Requirement 7.4 forbids. + """ + + +@dataclass(frozen=True) +class ProvisionedKnowledgeBase: + """The identifiers a caller needs, plus how they were obtained. + + ``created`` distinguishes "this call made the AWS resource" from "this call + found one already recorded", which is what makes idempotency assertable + instead of assumed. + """ + + aws_kb_id: str + aws_data_source_id: str + client_token: str + created: bool = False + + +# ── Payload builders ───────────────────────────────────────────────────────── +def build_client_token(*parts: Any) -> str: + """Build a ``clientToken`` that satisfies the API's constraints by construction. + + Deterministic in its inputs, so a retry of the same provisioning produces the + same token and AWS deduplicates the create rather than making a second + knowledge base. That property is the reason the token is also persisted on the + KB_Record: a *later* process, with no memory of this one, must be able to + reproduce the retry. + + The caller's parts are sanitized to the permitted alphabet and a digest of the + original seed is appended. The digest is not decoration — it is what + guarantees the 33-character minimum regardless of how short the inputs are, + which is the failure the natural ``{id}-{variant}-kb`` template walks into at + 31 characters. + """ + seed = "-".join(str(part) for part in parts if str(part) != "") + digest = hashlib.sha256(seed.encode("utf-8")).hexdigest()[:_DIGEST_CHARS] + + sanitized = re.sub(r"[^a-zA-Z0-9-]", "-", seed).strip("-") + token = f"{sanitized}-{digest}" if sanitized else digest + + if len(token) > CLIENT_TOKEN_MAX_LENGTH: + # Truncate from the left of the *prefix*, never the digest: the digest is + # what carries the uniqueness, so a token trimmed to its prefix could + # collide with a different knowledge base's. + keep = CLIENT_TOKEN_MAX_LENGTH - len(digest) - 1 + token = f"{sanitized[:keep].rstrip('-')}-{digest}" + + validate_client_token(token) + return token + + +def validate_client_token(token: str) -> None: + """Refuse a token the API would refuse, with a message saying which rule. + + Checked locally because botocore validates ``min``/``max``/``pattern`` + client-side: a short token never reaches AWS, so there is no service error to + read and no request id to quote. Raising here, naming the length, is the + difference between a one-line fix and an afternoon. + """ + if not isinstance(token, str): + raise ValueError(f"clientToken must be a string, got {type(token).__name__}") + if len(token) < CLIENT_TOKEN_MIN_LENGTH: + raise ValueError( + f"clientToken must be at least {CLIENT_TOKEN_MIN_LENGTH} characters " + f"(the API's documented minimum); got {len(token)}: {token!r}. Build " + f"tokens with build_client_token() rather than interpolating a " + f"template — the natural '{{id}}-{{variant}}-kb' form is 31 characters " + f"and fails botocore's client-side validation before any request." + ) + if len(token) > CLIENT_TOKEN_MAX_LENGTH: + raise ValueError( + f"clientToken must be at most {CLIENT_TOKEN_MAX_LENGTH} characters; " + f"got {len(token)}" + ) + if not CLIENT_TOKEN_PATTERN.match(token): + raise ValueError( + f"clientToken {token!r} does not match the API's pattern " + f"{CLIENT_TOKEN_PATTERN.pattern}: it must begin and end with an " + f"alphanumeric character" + ) + + +def embedding_model_arn(region: Optional[str] = None) -> str: + """ARN of the pinned embedding model. + + Foundation-model ARNs carry no account, hence the empty account segment. + """ + return f"arn:aws:bedrock:{region or _region()}::foundation-model/{EMBEDDING_MODEL_ID}" + + +def build_tags( + app_kb_id: str, + owner_user_id: str, + project_prefix: Optional[str] = None, + environment: Optional[str] = None, +) -> Dict[str, str]: + """Tags the Reconciler and the teardown script both read (Requirement 20.11). + + Delegates to :mod:`apis.shared.kb_backend.tags`, which owns the key names and + the value resolution. Kept as a thin wrapper because the provisioning saga is + the only caller and this is where a reader looks for it. + + ⚠️ This function used to build the tags itself, with keys ``prefix``/``env`` + and values from ``PROJECT_PREFIX``/``ENVIRONMENT`` — neither of which the + provisioning Lambda receives. Every knowledge base would have been tagged with + the hardcoded defaults, the teardown script (which read a different pair of + variables) would have matched nothing, and two environments in one account + would have claimed each other's corpora. See the ``tags`` module docstring. + """ + from apis.shared.kb_backend.tags import build_tags as _canonical + + return _canonical(app_kb_id, owner_user_id, project_prefix, environment) + + +def knowledge_base_payload( + name: str, + role_arn: str, + client_token: str, + *, + description: Optional[str] = None, + tags: Optional[Mapping[str, str]] = None, + region: Optional[str] = None, + kms_key_arn: Optional[str] = None, +) -> Dict[str, Any]: + """The exact ``CreateKnowledgeBase`` request (Requirements 8.1, 8.2, 8.5). + + ``storageConfiguration`` is absent by construction — there is no key to + accidentally set to ``None``, because a managed knowledge base has no vector + store and sending one is rejected. + """ + validate_client_token(client_token) + + managed: Dict[str, Any] = { + # Requirement 8.5: pinned, and immutable from here on. + "embeddingModelType": EMBEDDING_MODEL_TYPE, + "embeddingModelArn": embedding_model_arn(region), + "embeddingModelConfiguration": { + "bedrockEmbeddingModelConfiguration": { + "dimensions": EMBEDDING_DIMENSIONS, + "embeddingDataType": EMBEDDING_DATA_TYPE, + } + }, + } + if kms_key_arn: + # Requirement 20.5, only where customer-managed encryption is required. + managed["serverSideEncryptionConfiguration"] = {"kmsKeyArn": kms_key_arn} + + payload: Dict[str, Any] = { + "name": name, + "roleArn": role_arn, + "clientToken": client_token, + "knowledgeBaseConfiguration": { + "type": KNOWLEDGE_BASE_TYPE, + "managedKnowledgeBaseConfiguration": managed, + }, + } + if description: + payload["description"] = description + if tags: + payload["tags"] = dict(tags) + return payload + + +def data_source_payload( + knowledge_base_id: str, + name: str, + client_token: str, + *, + description: Optional[str] = None, +) -> Dict[str, Any]: + """The exact ``CreateDataSource`` request (Requirements 8.3, 8.4, 8.6, 8.7). + + Note the two-level nesting of the connector type. Putting ``CUSTOM`` at the + top level — which is what every pre-managed example does — is rejected with + "Unsupported data source type for MANAGED knowledge base type". + """ + validate_client_token(client_token) + + payload: Dict[str, Any] = { + "knowledgeBaseId": knowledge_base_id, + "name": name, + "clientToken": client_token, + # Requirement 8.7 — at creation, because it cannot rescue a knowledge base + # that is already stuck in DELETE_UNSUCCESSFUL. + "dataDeletionPolicy": DATA_DELETION_POLICY, + "dataSourceConfiguration": { + "type": DATA_SOURCE_TYPE, + "managedKnowledgeBaseConnectorConfiguration": { + "connectorParameters": { + "type": CONNECTOR_TYPE, + "version": CONNECTOR_VERSION, + }, + # Requirement 8.6 — opt-in, and silent when omitted. + "mediaExtractionConfiguration": { + "imageExtractionConfiguration": { + "imageExtractionStatus": IMAGE_EXTRACTION_STATUS + } + }, + }, + }, + } + if description: + payload["description"] = description + return payload + + +# ── Retry classification ───────────────────────────────────────────────────── +def is_retryable_error(exc: BaseException) -> bool: + """Whether ``exc`` should be retried rather than surfaced (Requirement 7.7). + + The embedding-verification message is the interesting case. It arrives looking + like a configuration error, which is why a first reading of it produces code + that fails the whole provisioning and tells the operator to check a model that + is demonstrably fine. It is IAM eventual consistency: the role's + ``bedrock:InvokeModel`` grant has not propagated yet. + """ + message = str(exc).lower() + if any(fragment in message for fragment in RETRYABLE_MESSAGE_FRAGMENTS): + return True + + response = getattr(exc, "response", None) + if isinstance(response, Mapping): + code = response.get("Error", {}).get("Code") + if code in RETRYABLE_ERROR_CODES: + return True + return False + + +# ── Clients and clocks ─────────────────────────────────────────────────────── +def _region() -> str: + return os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-west-2" + + +def bedrock_agent_client(): + """A ``bedrock-agent`` control-plane client. Imported lazily, deliberately.""" + import boto3 + + return boto3.client("bedrock-agent", region_name=_region()) + + +def _now_iso() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z") + + +async def _call( + operation: Callable[..., Any], + payload: Mapping[str, Any], + *, + what: str, + max_attempts: int = MAX_PROVISION_ATTEMPTS, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, +) -> Any: + """Invoke a synchronous boto3 operation off the event loop, with retries. + + ``asyncio.to_thread`` rather than a direct call because this runs inside the + async request path and ``CreateKnowledgeBase`` blocks for 47–124 s + (Requirement 20.7). Called directly it would stall every other coroutine on + the loop — including the health check that decides whether the task is alive. + """ + for attempt in range(1, max_attempts + 1): + try: + return await asyncio.to_thread(lambda: operation(**payload)) + except Exception as exc: + if attempt >= max_attempts or not is_retryable_error(exc): + raise + delay = min(2.0**attempt, _MAX_BACKOFF_SECONDS) + logger.warning( + f"{what} failed with a retryable error (attempt {attempt}/" + f"{max_attempts}, retrying in {delay}s): {exc}" + ) + emit_count(METRIC_PROVISION_RETRIED, dimensions={"operation": what}) + await sleep(delay) + + # Unreachable: the loop either returns or raises. + raise ProvisioningError(f"{what} exhausted {max_attempts} attempts") + + +# ── The saga ───────────────────────────────────────────────────────────────── +def _complete(item: Mapping[str, Any]) -> bool: + return bool(item.get("awsKbId")) and bool(item.get("awsDataSourceId")) + + +def _resource_name(app_kb_id: str, project_prefix: Optional[str] = None) -> str: + """The knowledge base's AWS name. + + Resolved through the same helper as the tags, so a knowledge base's name and + its ``ManagedKbPrefix`` tag can never disagree. The name is only a convention — + every filter in this feature matches on tags — but a name that says ``prod`` + while the tag says ``dev`` is the kind of thing an operator reads once and + trusts. + """ + from apis.shared.kb_backend.tags import tag_prefix + + return f"{tag_prefix(project_prefix)}-kb-{app_kb_id}" + + +async def provision_managed_kb( + assistant_id: str, + app_kb_id: Optional[str] = None, + owner_user_id: str = "", + *, + role_arn: Optional[str] = None, + client=None, + region: Optional[str] = None, + kms_key_arn: Optional[str] = None, + project_prefix: Optional[str] = None, + environment: Optional[str] = None, + max_attempts: int = MAX_PROVISION_ATTEMPTS, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, +) -> ProvisionedKnowledgeBase: + """Provision, or adopt, the managed knowledge base for ``app_kb_id``. + + Safe to call repeatedly and concurrently. Three paths, in the order they are + tried: + + * **Already provisioned** — the record carries both identifiers, so this + returns them and calls nothing. + * **Resuming** — a record exists in ``provisioning``, which is what a crash + between the AWS create and the conditional update leaves behind. Its + persisted ``clientToken`` is reused, so the retried ``CreateKnowledgeBase`` + is deduplicated by AWS and no second knowledge base appears. + * **Fresh** — the record is written first, then AWS is called. + + Losing the ``create_provisioning`` race raises + :class:`ProvisioningInProgress` rather than proceeding. The winner is already + creating the knowledge base; a loser that pressed on with its own token would + create a second one and only one of them could ever be recorded. + """ + from apis.shared.kb_backend import records as r + + app_kb_id = app_kb_id or assistant_id + role_arn = role_arn or os.environ.get("MANAGED_KB_SERVICE_ROLE_ARN") + if not role_arn: + raise ProvisioningError( + "no Bedrock knowledge base service role: pass role_arn or set " + "MANAGED_KB_SERVICE_ROLE_ARN" + ) + + client = client or bedrock_agent_client() + kb_token = build_client_token(app_kb_id, "knowledge-base") + ds_token = build_client_token(app_kb_id, "data-source") + + existing = await asyncio.to_thread(r.get_kb_record, assistant_id, app_kb_id) + + if existing and _complete(existing): + # Idempotent: nothing to create, and nothing to write. + return ProvisionedKnowledgeBase( + aws_kb_id=existing["awsKbId"], + aws_data_source_id=existing["awsDataSourceId"], + client_token=existing.get("clientToken") or kb_token, + created=False, + ) + + if existing: + # The retry-anchor path. Reuse the persisted token: it is the only thing + # that makes the re-create idempotent on AWS's side. + kb_token = existing.get("clientToken") or kb_token + emit_count(METRIC_PROVISION_ADOPTED, dimensions={"appKbId": app_kb_id}) + logger.info( + f"resuming provisioning for kb {app_kb_id} from its existing " + f"{existing.get('provisioningState')} record" + ) + else: + record = r.KbRecord( + app_kb_id=app_kb_id, + owner_user_id=owner_user_id, + provisioning_state=r.PROVISIONING, + client_token=kb_token, + embedding_model_id=EMBEDDING_MODEL_ID, + embedding_dimensions=EMBEDDING_DIMENSIONS, + image_extraction=True, + parser_config={ + "imageExtractionStatus": IMAGE_EXTRACTION_STATUS, + "connectorType": CONNECTOR_TYPE, + "embeddingDataType": EMBEDDING_DATA_TYPE, + }, + ) + try: + # DDB before AWS. See the module docstring; this ordering is the + # difference between a retry anchor and an untraceable paying resource. + await asyncio.to_thread(r.create_provisioning, assistant_id, record) + except r.TransitionLost as exc: + other = await asyncio.to_thread(r.get_kb_record, assistant_id, app_kb_id) + if other and _complete(other): + return ProvisionedKnowledgeBase( + aws_kb_id=other["awsKbId"], + aws_data_source_id=other["awsDataSourceId"], + client_token=other.get("clientToken") or kb_token, + created=False, + ) + raise ProvisioningInProgress( + f"another worker is provisioning kb {app_kb_id}; retry later " + f"rather than creating a second knowledge base" + ) from exc + + name = _resource_name(app_kb_id, project_prefix) + + aws_kb_id = (existing or {}).get("awsKbId") + if not aws_kb_id: + response = await _call( + client.create_knowledge_base, + knowledge_base_payload( + name=name, + role_arn=role_arn, + client_token=kb_token, + description=f"Managed knowledge base for {app_kb_id}", + tags=build_tags(app_kb_id, owner_user_id, project_prefix, environment), + region=region, + kms_key_arn=kms_key_arn, + ), + what="CreateKnowledgeBase", + max_attempts=max_attempts, + sleep=sleep, + ) + aws_kb_id = response["knowledgeBase"]["knowledgeBaseId"] + + aws_data_source_id = (existing or {}).get("awsDataSourceId") + if not aws_data_source_id: + ds_response = await _call( + client.create_data_source, + data_source_payload( + knowledge_base_id=aws_kb_id, + name=name, + client_token=ds_token, + description=f"CUSTOM connector for {app_kb_id}", + ), + what="CreateDataSource", + max_attempts=max_attempts, + sleep=sleep, + ) + aws_data_source_id = ds_response["dataSource"]["dataSourceId"] + + try: + await asyncio.to_thread( + r.attach_aws_ids, + assistant_id, + app_kb_id, + aws_kb_id, + aws_data_source_id, + _now_iso(), + ) + except r.TransitionLost: + # Another worker attached first, or the record has already left + # `provisioning`. Its identifiers win: they are the ones every reader + # will see, so returning our own would hand back a knowledge base that + # no record points at. + current = await asyncio.to_thread(r.get_kb_record, assistant_id, app_kb_id) + if current and _complete(current): + return ProvisionedKnowledgeBase( + aws_kb_id=current["awsKbId"], + aws_data_source_id=current["awsDataSourceId"], + client_token=current.get("clientToken") or kb_token, + created=False, + ) + raise + + return ProvisionedKnowledgeBase( + aws_kb_id=aws_kb_id, + aws_data_source_id=aws_data_source_id, + client_token=kb_token, + created=True, + ) + + +__all__ = [ + "CLIENT_TOKEN_MAX_LENGTH", + "CLIENT_TOKEN_MIN_LENGTH", + "CLIENT_TOKEN_PATTERN", + "CONNECTOR_TYPE", + "CONNECTOR_VERSION", + "DATA_DELETION_POLICY", + "DATA_SOURCE_TYPE", + "EMBEDDING_DATA_TYPE", + "EMBEDDING_DIMENSIONS", + "EMBEDDING_MODEL_ID", + "EMBEDDING_MODEL_TYPE", + "IMAGE_EXTRACTION_STATUS", + "KNOWLEDGE_BASE_TYPE", + "ProvisionedKnowledgeBase", + "ProvisioningError", + "ProvisioningInProgress", + "RetryableProvisioningError", + "build_client_token", + "build_tags", + "data_source_payload", + "embedding_model_arn", + "is_retryable_error", + "knowledge_base_payload", + "provision_managed_kb", + "validate_client_token", +] diff --git a/backend/src/apis/shared/kb_backend/query_guard.py b/backend/src/apis/shared/kb_backend/query_guard.py new file mode 100644 index 000000000..36966a320 --- /dev/null +++ b/backend/src/apis/shared/kb_backend/query_guard.py @@ -0,0 +1,74 @@ +"""Hard cap on retrieval query length. + +Amazon Bedrock Managed Knowledge Base caps ``Retrieve`` query input at **10,000 +characters** and that quota is **not adjustable**. Exceeding it is a request +error, not a degraded result — so an unclamped query is the difference between a +slightly-truncated answer and no answer at all. + +Why this lives above the seam +----------------------------- +The clamp is applied in the facade, before backend dispatch, so both backends see +an identically-shaped query. Clamping only the managed path would mean the two +backends answered *different questions* whenever a query ran long, which would +quietly invalidate the dual-read comparison this migration depends on: a rank +disagreement would be indistinguishable from a genuine retrieval difference. + +That does mean the legacy path is now clamped too, where previously it was not. +Titan v2 tolerates roughly 32,000 characters, so queries between 10,000 and that +ceiling used to be embedded whole and now are not. This is deliberate — parity is +worth more than the tail of a pathological query — and it is why the truncation +emits a metric rather than passing silently. + +Why it never raises +------------------- +A query too long is a fixable input, not a failure. Raising would turn a +recoverable situation into a 500 on a chat turn. The function is total: every +input maps to an output of at most :data:`MAX_QUERY_CHARS` characters. +""" + +from __future__ import annotations + +import logging +from typing import Tuple + +from apis.shared.kb_backend.metrics import METRIC_QUERY_CLAMPED, emit_count + +logger = logging.getLogger(__name__) + +#: Managed KB's ``Retrieve`` input limit. Not adjustable — do not raise this +#: hoping for a quota increase; there is not one to request. +MAX_QUERY_CHARS = 10_000 + + +def clamp_query(query: str) -> Tuple[str, bool]: + """Return ``(clamped_query, was_truncated)``. + + Truncates from the end, keeping the head. For a natural-language query the + beginning carries the intent, so a tail-truncated query still retrieves + something sensible; head-truncating would change the question entirely. + + A ``None`` or non-string input is coerced rather than rejected, because the + caller is a request path and the clamp is a guard, not a validator. + """ + if not query: + return "", False + + if not isinstance(query, str): + query = str(query) + + if len(query) <= MAX_QUERY_CHARS: + return query, False + + original_length = len(query) + clamped = query[:MAX_QUERY_CHARS] + + # Error-level would overstate it (the request still succeeds) and debug would + # hide it. A clamped query means someone's answer is based on a partial + # question, which an operator should be able to see without turning on debug. + logger.warning( + f"Query clamped from {original_length} to {MAX_QUERY_CHARS} characters " + f"(Managed KB Retrieve limit, not adjustable)" + ) + emit_count(METRIC_QUERY_CLAMPED) + + return clamped, True diff --git a/backend/src/apis/shared/kb_backend/records.py b/backend/src/apis/shared/kb_backend/records.py new file mode 100644 index 000000000..342566797 --- /dev/null +++ b/backend/src/apis/shared/kb_backend/records.py @@ -0,0 +1,625 @@ +"""KB_Record persistence for the managed knowledge base migration. + +Records live in the **existing** assistants table as siblings of the assistant's +``METADATA`` row, preserving the adjacency-list convention:: + + PK = AST#{assistant_id} + SK = KB#{app_kb_id} # app_kb_id == assistant_id this phase + SK = KBTOMB#{app_kb_id} # whole-KB tombstone + SK = KBTOMB#{app_kb_id}#DOC#{document_id} # per-document tombstone + +Three invariants are load-bearing. Each is enforced here rather than left to +callers, because each fails silently when violated: + +**1. Absence means legacy.** ``retrievalEngine`` is written *only* as +``"managed"``. Nothing here ever writes ``"s3vectors"`` onto a record that did +not already carry it. That is what makes this migration zero-backfill: every +existing knowledge base is already correct by virtue of having no opinion, and +rollback is a single attribute removal rather than a data rewrite. A backfill +that "helpfully" stamped the legacy value on 1,692 records would convert a +pointer flip into a migration of its own. + +**2. Every transition is conditional.** These functions are called from a +dispatcher that fans out to concurrent workers, so a read-then-write would let +two workers both believe they won. Each transition therefore carries a DynamoDB +``ConditionExpression`` and surfaces the loss as :class:`TransitionLost` rather +than an opaque ``ClientError``. + +**3. Sparse work keys are removed, not just ignored.** ``GSI7_PK``/``GSI7_SK`` +exist only while a record is eligible for background work. On reaching a terminal +state they are ``REMOVE``d, so an ineligible knowledge base is invisible to the +dispatcher's query *by physics* rather than by filter. This matters more than the +usual sparse-index argument because the dispatcher creates and deletes billed AWS +resources: a missing key can only ever mean "do nothing", whereas a stale key +means "act on something nobody asked you to act on". + +Import boundary +--------------- +This module deliberately talks to DynamoDB through the raw table resource instead +of importing ``apis.shared.assistants``. That package's ``__init__`` imports +``rag_service``, which imports the embeddings stack at module scope; pulling it +into the migration Lambda image would blow the image-size budget. The same +constraint is why ``apis/app_api/kb_sync/records.py`` is written this way, and +this module follows it: **module-level imports are stdlib only**, and ``boto3`` +is imported inside the functions that need it. ``kb_backend/__init__.py`` is +intentionally empty so importing a submodule pulls in nothing else. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass, field +from decimal import Decimal +from typing import Any, Dict, Iterable, Mapping, Optional + +logger = logging.getLogger(__name__) + +# ── Engines ────────────────────────────────────────────────────────────────── +# +# LEGACY is never persisted. It is the value `resolve_engine` returns for a +# record that carries no `retrievalEngine` attribute, which is every record that +# predates this feature. +ENGINE_LEGACY = "s3vectors" +ENGINE_MANAGED = "managed" + +# ── Provisioning ───────────────────────────────────────────────────────────── +PROVISIONING = "provisioning" +ACTIVE = "active" +FAILED = "failed" +DELETING = "deleting" + +# ── Migration states ───────────────────────────────────────────────────────── +SHADOW = "shadow" +VERIFY = "verify" +PROMOTE = "promote" +RETAIN = "retain" +MIGRATION_FAILED = "failed" + +#: Reserved in the enum so a stored value round-trips, but never entered in this +#: phase. Reclaiming legacy vectors is explicitly a follow-up spec; a worker that +#: found itself here would delete data this phase has promised to retain. +RECLAIM = "reclaim" + +#: States that keep a record in the dispatcher's queue. Work keys are written on +#: entering one of these. +WORK_ELIGIBLE_STATES = frozenset({SHADOW, VERIFY, PROMOTE}) + +#: States that take a record out of the queue for good. Work keys are removed on +#: entering one of these. ``RETAIN`` is the terminal state this phase reaches; +#: ``MIGRATION_FAILED`` is terminal too and leaves the record on legacy, which +#: keeps working. +TERMINAL_STATES = frozenset({RETAIN, MIGRATION_FAILED}) + +ALL_MIGRATION_STATES = frozenset( + {SHADOW, VERIFY, PROMOTE, RETAIN, MIGRATION_FAILED, RECLAIM} +) + + +class TransitionLost(Exception): + """A conditional write was rejected because the guard did not hold. + + Raised instead of leaking ``ConditionalCheckFailedException`` so callers can + tell "another worker got there first, do nothing" apart from a real error. + Losing a race is normal and must not be logged as a failure. + """ + + +class ReclaimNotSupported(Exception): + """Refuses an attempt to enter ``reclaim``, which this phase never does.""" + + +# ── Keys ───────────────────────────────────────────────────────────────────── +def kb_pk(assistant_id: str) -> str: + return f"AST#{assistant_id}" + + +def kb_sk(app_kb_id: str) -> str: + return f"KB#{app_kb_id}" + + +def kb_tombstone_sk(app_kb_id: str) -> str: + return f"KBTOMB#{app_kb_id}" + + +def document_tombstone_sk(app_kb_id: str, document_id: str) -> str: + return f"KBTOMB#{app_kb_id}#DOC#{document_id}" + + +def work_pk(state: str) -> str: + return f"KBWORK#{state}" + + +def _table(): + import boto3 + + return boto3.resource("dynamodb").Table(os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"]) + + +# ── Model ──────────────────────────────────────────────────────────────────── +@dataclass +class KbRecord: + """A knowledge base's control-plane state. + + A dataclass rather than a Pydantic model on purpose: this module is imported + by a size-constrained Lambda image and has no need for validation machinery + it would then have to carry. + """ + + app_kb_id: str + owner_user_id: str + visibility: str = "PRIVATE" + + # Absent means legacy. Only ever ENGINE_MANAGED when present. + retrieval_engine: Optional[str] = None + + provisioning_state: str = PROVISIONING + aws_kb_id: Optional[str] = None + aws_data_source_id: Optional[str] = None + + # Immutable after creation: Bedrock rejects changing either, so they are + # recorded to make a mismatch detectable rather than mysterious. + embedding_model_id: str = "amazon.titan-embed-text-v2:0" + embedding_dimensions: int = 1024 + + # Captured at creation because a corpus indexed without image extraction is + # not comparable to one indexed with it. + parser_config: Dict[str, Any] = field(default_factory=dict) + image_extraction: bool = False + + stored_bytes: int = 0 + reserved_bytes: int = 0 + last_retrieved_at: Optional[str] = None + + migration_state: Optional[str] = None + migration_generation: int = 0 + migration_lease_until: Optional[str] = None + migration_progress: Dict[str, Any] = field(default_factory=dict) + migration_error: Optional[str] = None + + promoted_at: Optional[str] = None + rolled_back_at: Optional[str] = None + retain_until: Optional[str] = None + + pinned: bool = False + exempt_from_reclaim: bool = False + + client_token: Optional[str] = None + + def to_item(self, assistant_id: str) -> Dict[str, Any]: + """Serialize for DynamoDB, omitting absent optionals. + + Optionals are omitted rather than written as ``None`` so that "has no + opinion" stays distinguishable from "explicitly null". ``retrievalEngine`` + depends on that distinction. + """ + item: Dict[str, Any] = { + "PK": kb_pk(assistant_id), + "SK": kb_sk(self.app_kb_id), + "appKbId": self.app_kb_id, + "ownerUserId": self.owner_user_id, + "visibility": self.visibility, + "provisioningState": self.provisioning_state, + "embeddingModelId": self.embedding_model_id, + "embeddingDimensions": Decimal(self.embedding_dimensions), + "parserConfig": self.parser_config, + "imageExtraction": self.image_extraction, + "storedBytes": Decimal(self.stored_bytes), + "reservedBytes": Decimal(self.reserved_bytes), + "migrationGeneration": Decimal(self.migration_generation), + "pinned": self.pinned, + "exemptFromReclaim": self.exempt_from_reclaim, + } + + optional = { + "retrievalEngine": self.retrieval_engine, + "awsKbId": self.aws_kb_id, + "awsDataSourceId": self.aws_data_source_id, + "lastRetrievedAt": self.last_retrieved_at, + "migrationState": self.migration_state, + "migrationLeaseUntil": self.migration_lease_until, + "migrationError": self.migration_error, + "promotedAt": self.promoted_at, + "rolledBackAt": self.rolled_back_at, + "retainUntil": self.retain_until, + "clientToken": self.client_token, + } + item.update({k: v for k, v in optional.items() if v is not None}) + + if self.migration_progress: + item["migrationProgress"] = self.migration_progress + + return item + + +def resolve_engine(item: Optional[Mapping[str, Any]]) -> str: + """Return the backend that should serve this record. + + The whole migration rests on this function's default. A record with no + ``retrievalEngine`` attribute — which is every knowledge base that existed + before this feature — resolves to the legacy backend. Nothing had to be + written to make that true, and nothing has to be unwritten to roll back. + + A missing record resolves to legacy for the same reason: the absence of an + opinion is an answer, not an error. + """ + if not item: + return ENGINE_LEGACY + return ENGINE_MANAGED if item.get("retrievalEngine") == ENGINE_MANAGED else ENGINE_LEGACY + + +# ── Reads ──────────────────────────────────────────────────────────────────── +def get_kb_record(assistant_id: str, app_kb_id: str) -> Optional[Dict[str, Any]]: + response = _table().get_item(Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}) + return response.get("Item") + + +def query_due_work(state: str, now_iso: str, limit: int = 20) -> list: + """Records in ``state`` whose ``dueAt`` has passed, oldest first. + + Reads the sparse index, so records that have left the queue are not returned + because they have no key — not because they were filtered out. + """ + from boto3.dynamodb.conditions import Key + + response = _table().query( + IndexName="KbWorkIndex", + KeyConditionExpression=Key("GSI7_PK").eq(work_pk(state)) & Key("GSI7_SK").lte(now_iso), + Limit=limit, + ) + return response.get("Items", []) + + +# ── Transitions ────────────────────────────────────────────────────────────── +def _conditional(operation, **kwargs): + """Run a conditional write, translating a failed guard into TransitionLost.""" + from botocore.exceptions import ClientError + + try: + return operation(**kwargs) + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + raise TransitionLost( + "conditional write rejected; another writer won or the " + "precondition no longer holds" + ) from exc + raise + + +def create_provisioning( + assistant_id: str, + record: KbRecord, +) -> Dict[str, Any]: + """Create the record, exactly once. + + ``attribute_not_exists(PK)`` makes this idempotent under concurrency: two + callers racing to enrol the same knowledge base produce one record and one + :class:`TransitionLost`, rather than one silently overwriting the other's + ``clientToken`` and orphaning a half-created AWS knowledge base. + """ + item = record.to_item(assistant_id) + _conditional( + _table().put_item, + Item=item, + ConditionExpression="attribute_not_exists(PK) AND attribute_not_exists(SK)", + ) + return item + + +def attach_aws_ids( + assistant_id: str, + app_kb_id: str, + aws_kb_id: str, + aws_data_source_id: str, + now_iso: str, +) -> None: + """Record the AWS identifiers and mark the record active. + + Guarded on still being ``provisioning`` so a late-returning create cannot + overwrite identifiers belonging to a newer generation. + """ + _conditional( + _table().update_item, + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression=( + "SET awsKbId = :kb, awsDataSourceId = :ds, " + "provisioningState = :active, updatedAt = :now" + ), + ConditionExpression="provisioningState = :provisioning", + ExpressionAttributeValues={ + ":kb": aws_kb_id, + ":ds": aws_data_source_id, + ":active": ACTIVE, + ":provisioning": PROVISIONING, + ":now": now_iso, + }, + ) + + +def set_resource_policy_state( + assistant_id: str, + app_kb_id: str, + aws_kb_id: Optional[str], + revision_id: Optional[str], +) -> None: + """Record which ``awsKbId`` the resource policy is currently attached to. + + Requirement 25.7. This attribute is the whole staleness check: a policy + attaches to an ARN, so once the record's ``awsKbId`` and this value disagree, + the policy is on a resource nobody reads and sharing has silently stopped. + Storing the target rather than a boolean is what turns that from an event + somebody has to remember to fire into a comparison + (``resource_policy.policy_is_stale``). + + Unconditional, deliberately. Every other writer here guards on the state it + expects, because those transitions must not race. This one records what AWS has + just confirmed, and a stale overwrite of the *same* fact is harmless while a + refused write would leave the record claiming a policy target that is no longer + true — the failure mode the attribute exists to prevent. + + Passing ``None`` clears both attributes, for a knowledge base that stopped + being shared. + """ + key = {"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)} + if aws_kb_id is None: + _table().update_item( + Key=key, + UpdateExpression="REMOVE policyAwsKbId, policyRevisionId", + ) + return + + values: Dict[str, Any] = {":kb": aws_kb_id} + expression = "SET policyAwsKbId = :kb" + if revision_id: + expression += ", policyRevisionId = :rev" + values[":rev"] = revision_id + else: + expression += " REMOVE policyRevisionId" + + _table().update_item( + Key=key, + UpdateExpression=expression, + ExpressionAttributeValues=values, + ) + + +def promote_engine( + assistant_id: str, + app_kb_id: str, + generation: int, + now_iso: str, +) -> None: + """Flip the record to the managed backend. The single cutover write. + + Four guards, all necessary: + + * ``attribute_not_exists(retrievalEngine)`` — this knowledge base has not + already been promoted. Without it, the other three guards all remain true + *after* a successful promotion, so a worker that crashed between the + promotion and the state transition promotes a second time on resume: same + value, but a fresh ``promotedAt`` that overwrites the real cutover moment and + a second ``KbMigrationPromoted``. Worse, two concurrent workers would both + succeed, which is precisely what Requirement 15.10 forbids. Found by the + convergence property test, which counted two promotions across a crash at + the state transition. Rollback ``REMOVE``s the attribute, so this does not + block a deliberate re-promotion. + * ``migrationState = promote`` — only a record that reached the cutover step + may cut over. + * ``migrationGeneration = :gen`` — a worker whose lease expired and whose + generation has been superseded cannot promote on stale information. + * ``migrationProgress.migrated = migrationProgress.total`` — the catch-up + pass has converged. Without this, promotion could strand documents written + during migration on a backend nobody reads any more. Comparing two + document paths keeps the check atomic with the write; passing the total in + as a value would let it go stale between read and write. + + ``total`` is a DynamoDB reserved keyword, so the progress paths are aliased + through ``ExpressionAttributeNames``. Without the aliases the whole condition + is rejected as a ``ValidationException`` — loudly, which is the good case, but + only because it never validates at all. + + Because this is one conditional write, rollback is symmetric: see + :func:`rollback_engine`. + """ + _conditional( + _table().update_item, + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression="SET retrievalEngine = :managed, promotedAt = :now", + ConditionExpression=( + "attribute_not_exists(retrievalEngine) " + "AND migrationState = :promote " + "AND migrationGeneration = :gen " + "AND #progress.#migrated = #progress.#total" + ), + ExpressionAttributeNames={ + "#progress": "migrationProgress", + "#migrated": "migrated", + "#total": "total", + }, + ExpressionAttributeValues={ + ":managed": ENGINE_MANAGED, + ":promote": PROMOTE, + ":gen": Decimal(generation), + ":now": now_iso, + }, + ) + + +def rollback_engine(assistant_id: str, app_kb_id: str, now_iso: str) -> None: + """Return the record to the legacy backend by REMOVING the engine attribute. + + Note the ``REMOVE``. Rollback restores the original *shape*, not a written + legacy value, so a rolled-back record is byte-indistinguishable from one that + never migrated. Writing ``"s3vectors"`` here would work today and quietly + break the "absence means legacy" invariant that lets this feature ship + without touching 1,692 existing records. + + Guarded on currently being managed so a double rollback is a no-op loss + rather than a spurious ``rolledBackAt`` bump. + """ + _conditional( + _table().update_item, + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression="REMOVE retrievalEngine SET rolledBackAt = :now", + ConditionExpression="retrievalEngine = :managed", + ExpressionAttributeValues={":managed": ENGINE_MANAGED, ":now": now_iso}, + ) + + +def set_migration_state( + assistant_id: str, + app_kb_id: str, + new_state: str, + generation: int, + due_at: Optional[str] = None, + expected_states: Optional[Iterable[str]] = None, + error: Optional[str] = None, +) -> None: + """Move to ``new_state``, maintaining the sparse work keys. + + Entering a work-eligible state writes ``GSI7_PK``/``GSI7_SK``; entering a + terminal state ``REMOVE``s them. The removal is the point: it is what takes + the record out of the dispatcher's queue, and skipping it would leave a + finished knowledge base being handed to workers forever. + + ``expected_states`` guards the transition against a concurrent writer that + has already moved the record on. The generation is always guarded. + """ + if new_state == RECLAIM: + raise ReclaimNotSupported( + "reclaim is reserved but never entered in this phase; reclaiming " + "legacy vectors is a follow-up spec" + ) + if new_state not in ALL_MIGRATION_STATES: + raise ValueError(f"unknown migration state: {new_state!r}") + if new_state in WORK_ELIGIBLE_STATES and not due_at: + raise ValueError(f"{new_state} is work-eligible and requires due_at") + + values: Dict[str, Any] = { + ":state": new_state, + ":gen": Decimal(generation), + } + sets = ["migrationState = :state"] + removes = [] + + if new_state in TERMINAL_STATES: + # Leaving the queue: the keys must go, not merely be ignored. + removes.extend(["GSI7_PK", "GSI7_SK"]) + else: + sets.extend(["GSI7_PK = :wpk", "GSI7_SK = :wsk"]) + values[":wpk"] = work_pk(new_state) + values[":wsk"] = due_at + + if error is not None: + sets.append("migrationError = :err") + values[":err"] = error + + expression = f"SET {', '.join(sets)}" + if removes: + expression += f" REMOVE {', '.join(removes)}" + + condition = "migrationGeneration = :gen" + if expected_states is not None: + expected = list(expected_states) + if not expected: + raise ValueError("expected_states must be non-empty when provided") + placeholders = [] + for index, state in enumerate(expected): + placeholder = f":exp{index}" + placeholders.append(placeholder) + values[placeholder] = state + condition += f" AND migrationState IN ({', '.join(placeholders)})" + + _conditional( + _table().update_item, + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression=expression, + ConditionExpression=condition, + ExpressionAttributeValues=values, + ) + + +def acquire_lease( + assistant_id: str, + app_kb_id: str, + lease_until: str, + now_iso: str, +) -> None: + """Take the worker lease, or lose the race. + + The guard admits exactly two situations: no lease has ever been taken, or the + existing lease has expired. A live lease held by another worker rejects, + which is what stops two workers migrating the same knowledge base and + double-ingesting its corpus. + + ISO-8601 UTC strings compare correctly lexicographically, so the expiry test + is a plain string comparison and stays atomic with the write. + """ + _conditional( + _table().update_item, + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression="SET migrationLeaseUntil = :until", + ConditionExpression=( + "attribute_not_exists(migrationLeaseUntil) OR migrationLeaseUntil < :now" + ), + ExpressionAttributeValues={":until": lease_until, ":now": now_iso}, + ) + + +def retry_from_failed( + assistant_id: str, + app_kb_id: str, + generation: int, + due_at: str, +) -> None: + """Re-enter ``shadow`` from ``failed`` on the next generation, atomically. + + One write, not two. Split into "bump the generation" then + "set_migration_state", a crash between them leaves a record carrying a new + generation while still ``failed`` — and with no work keys, so it is invisible + to the dispatcher while the retry control has already reported success. The + user would wait forever on an upgrade nothing owns. + + Guarded on **both** the old generation and still being ``failed``, so two + concurrent retries yield one new attempt and one :class:`TransitionLost`. The + generation bump is also what fences the abandoned attempt: every conditional + write belonging to it is guarded on the old value, so a straggler worker + cannot land on the new generation. + + ``migrationError`` is removed rather than left behind, so a subsequent + failure's reason cannot be mistaken for this one's. + """ + _conditional( + _table().update_item, + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression=( + "SET migrationState = :shadow, migrationGeneration = :next, " + "GSI7_PK = :wpk, GSI7_SK = :wsk REMOVE migrationError" + ), + ConditionExpression="migrationGeneration = :gen AND migrationState = :failed", + ExpressionAttributeValues={ + ":shadow": SHADOW, + ":failed": MIGRATION_FAILED, + ":gen": Decimal(generation), + ":next": Decimal(generation + 1), + ":wpk": work_pk(SHADOW), + ":wsk": due_at, + }, + ) + + +def dismiss_upgrade_notice(assistant_id: str, app_kb_id: str, now_iso: str) -> None: + """Retire the one-time post-upgrade notice. + + Unconditional on purpose: dismissing an already-dismissed notice is not a + race worth losing, and the attribute's only reader treats any value as + "dismissed". Guarded only on the record existing, so a dismissal for a + knowledge base that never had a record cannot conjure one. + """ + _conditional( + _table().update_item, + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression="SET upgradeNoticeDismissedAt = :now", + ConditionExpression="attribute_exists(PK) AND attribute_exists(SK)", + ExpressionAttributeValues={":now": now_iso}, + ) diff --git a/backend/src/apis/shared/kb_backend/resolver.py b/backend/src/apis/shared/kb_backend/resolver.py new file mode 100644 index 000000000..05585231c --- /dev/null +++ b/backend/src/apis/shared/kb_backend/resolver.py @@ -0,0 +1,198 @@ +"""Which backend serves this knowledge base. + +One question, answered in one place: read the KB_Record's ``retrievalEngine`` and +hand back the matching implementation. Callers get an object satisfying +:class:`~apis.shared.kb_backend.protocol.KnowledgeBaseBackend` and are given no +way to ask which one it is. + +Absence means legacy +-------------------- +The decision itself is delegated to +:func:`apis.shared.kb_backend.records.resolve_engine` rather than re-derived +here. That function is the one place that knows a missing ``retrievalEngine`` +attribute means the legacy backend, and it is covered by its own property test +(task 3.3). Two implementations of the same default would be two chances to +disagree about the invariant that lets 1,692 existing knowledge bases keep +working with zero backfill writes. + +Resolution never fails a turn +----------------------------- +The KB_Record lookup is a DynamoDB read that today's retrieval path does not +perform, so it is a new way for retrieval to break. It is therefore wrapped: any +failure — unreachable table, unset ``DYNAMODB_ASSISTANTS_TABLE_NAME``, malformed +item — resolves to legacy, which is what every knowledge base in existence +already uses. The failure is logged at warning level. Choosing legacy on an +unreadable record is not a guess; it is the same answer the absent attribute +gives, and the whole migration is built so that answer is always safe. + +Why both backends are registered at import +------------------------------------------ +Registration is not a startup step. Both adapters are installed in +:data:`_BACKENDS` when this module is imported, so there is no sequence anybody +has to remember and no service that can come up half-configured. That matters +because forgetting would not be loud in a useful way: a promoted knowledge base +would raise :class:`BackendUnavailable` on every turn — correct as a fail-safe, +useless as a signal, and only ever seen by the one user whose knowledge base was +migrated. + +It costs nothing. Both adapter modules import stdlib and the protocol only, with +``boto3`` and their clients created lazily inside methods, which +``tests/architecture/test_kb_backend_boundary.py`` asserts in a fresh +interpreter. Registering an object whose constructor does no work is not the same +as connecting to anything. + +Registering the managed backend does **not** make the feature live. Nothing can +resolve to it until a record says ``retrievalEngine == "managed"``, and nothing +writes that value except a promotion, which needs the migration flag on and an +explicit opt-in. Registration only settles what happens *once* a record says so. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, Mapping, Optional + +from apis.shared.kb_backend.managed_backend import ManagedKbBackend +from apis.shared.kb_backend.protocol import KnowledgeBaseBackend +from apis.shared.kb_backend.records import ENGINE_LEGACY, ENGINE_MANAGED, resolve_engine +from apis.shared.kb_backend.s3vectors_backend import S3VectorsBackend + +logger = logging.getLogger(__name__) + + +class BackendUnavailable(RuntimeError): + """A record names an engine this build has no implementation for. + + Raised rather than quietly falling back to legacy. A record only ever names + ``managed`` after a successful promotion, and serving legacy for a promoted + knowledge base would read an index that migration has stopped maintaining — + fewer results, silently, with no error to notice. + + Reachable only if a backend is explicitly unregistered (which tests do) or if + a future engine name is written by a newer deployment than the one reading it. + """ + + +# Engine → backend, populated at import. See the module docstring for why this is +# not a startup step. Both constructors are inert: clients are created lazily. +_BACKENDS: Dict[str, KnowledgeBaseBackend] = { + ENGINE_LEGACY: S3VectorsBackend(), + ENGINE_MANAGED: ManagedKbBackend(), +} + + +def register_backend(engine: str, backend: KnowledgeBaseBackend) -> None: + """Install the implementation for ``engine``, replacing any previous one.""" + _BACKENDS[engine] = backend + + +def unregister_backend(engine: str) -> None: + """Remove ``engine``'s implementation. Absent engines are ignored.""" + _BACKENDS.pop(engine, None) + + +def registered_engines() -> frozenset: + """Engines this build can serve. Introspection for tests and diagnostics.""" + return frozenset(_BACKENDS) + + +def load_record( + assistant_id: str, + app_kb_id: Optional[str] = None, +) -> Dict[str, Any]: + """The KB_Record, or an empty mapping if there is none to be had. + + For callers that need more from the record than which backend serves it — the + dual-read pilot flag, the byte cap, the migration state — and would otherwise + read it a second time. + + Returns ``{}`` rather than ``None`` for both "no such record" and "the read + failed", because those two cases have the same answer everywhere in this + feature: an absent opinion is the legacy opinion. Collapsing them here means + no caller has to remember to handle ``None`` and every caller can pass the + result straight to :func:`resolve_backend` as ``record=``, which is what makes + one read enough. + """ + from apis.shared.kb_backend.records import get_kb_record + + try: + return dict(get_kb_record(assistant_id, app_kb_id or assistant_id) or {}) + except Exception as exc: + logger.warning( + f"KB_Record lookup failed for assistant {assistant_id}; treating it as " + f"absent, which resolves to {ENGINE_LEGACY}: {exc}" + ) + return {} + + +def backend_for_engine(engine: str) -> Optional[KnowledgeBaseBackend]: + """The implementation for ``engine``, or ``None`` if this build has none. + + Unlike :func:`resolve_backend` this does not raise, because its callers are + asking a different question. The dual-read pilot wants "is there a managed + backend I could compare against?", and ``None`` is an ordinary answer for it — + a test that unregistered one, or a deployment older than the engine name it was + handed — not the fail-safe emergency that an unservable *promoted* record is. + """ + return _BACKENDS.get(engine) + + +def resolve_engine_for( + assistant_id: str, + app_kb_id: Optional[str] = None, + record: Optional[Mapping[str, Any]] = None, +) -> str: + """Return the engine name for a knowledge base. + + Pass ``record`` when the caller already holds the KB_Record to skip the + read. ``app_kb_id`` defaults to ``assistant_id``, which is the 1:1 binding + this phase deliberately preserves. + """ + if record is not None: + return resolve_engine(record) + + from apis.shared.kb_backend.records import get_kb_record + + try: + item = get_kb_record(assistant_id, app_kb_id or assistant_id) + except Exception as exc: + # Unreadable record ⇒ legacy, the same answer absence gives. + logger.warning( + f"KB_Record lookup failed for assistant {assistant_id}, " + f"resolving to {ENGINE_LEGACY}: {exc}" + ) + return ENGINE_LEGACY + + return resolve_engine(item) + + +def resolve_backend( + assistant_id: str, + app_kb_id: Optional[str] = None, + record: Optional[Mapping[str, Any]] = None, +) -> KnowledgeBaseBackend: + """Return the backend instance that should serve this knowledge base.""" + engine = resolve_engine_for(assistant_id, app_kb_id, record) + try: + return _BACKENDS[engine] + except KeyError: + raise BackendUnavailable( + f"knowledge base {app_kb_id or assistant_id} names engine {engine!r}, " + f"which this build cannot serve (have: {sorted(_BACKENDS)}). " + f"Refusing to substitute {ENGINE_LEGACY}: a promoted knowledge base's " + f"legacy index is no longer maintained." + ) from None + + +__all__ = [ + "BackendUnavailable", + "ENGINE_LEGACY", + "ENGINE_MANAGED", + "backend_for_engine", + "load_record", + "register_backend", + "registered_engines", + "resolve_backend", + "resolve_engine_for", + "unregister_backend", +] diff --git a/backend/src/apis/shared/kb_backend/resource_policy.py b/backend/src/apis/shared/kb_backend/resource_policy.py new file mode 100644 index 000000000..0d5211d4b --- /dev/null +++ b/backend/src/apis/shared/kb_backend/resource_policy.py @@ -0,0 +1,329 @@ +"""IAM-enforced retrieval on a shared managed knowledge base. + +Requirements 25.6, 25.7. Resource policies are MANAGED-only and are the only +mechanism in this design that offers *infrastructure* isolation rather than +filter-level isolation. A policy attached to a knowledge base ARN restricts +``bedrock:Retrieve`` and ``bedrock:GetDocumentContent`` to the principals it +names, which matters because the platform's own identity grant +(``grantManagedKbRetrieval``) is written against ``knowledge-base/*`` — every +knowledge base in the account, present and future. Without a policy, any +principal in the account holding a similar grant can read a shared corpus. + +What this is NOT +---------------- +It is **not** per-user authorization, and the temptation to describe it that way +is the reason this paragraph exists. Every user of this platform retrieves through +the same infrastructure identity — the AgentCore runtime role — so no resource +policy can distinguish user A from user B. Per-user authorization is, and remains, +the application's job (Requirement 25.3, ``apis.shared.assistants.kb_access``). +What a policy buys is a narrower blast radius for a corpus that belongs to more +than one person: the set of *infrastructure* identities able to reach it shrinks +from "anything in the account with a wildcard grant" to an explicit list. + +Applied only where a knowledge base is shared beyond its owner, because a policy +on a single-owner knowledge base would restrict nothing that the assistant's own +access check does not already restrict, while adding a control-plane call and a +piece of state to keep in step. + +Why staleness is state rather than an event +------------------------------------------- +A policy attaches to the AWS knowledge base ARN, so any cycle producing a new +``awsKbId`` silently drops sharing — the call succeeds, the policy is simply on a +resource nobody reads any more. The obvious fix is to re-apply from wherever a new +identifier is created. That fix is only as good as the completeness of the list of +such places, and this phase already has two (fresh provisioning, resumed +provisioning) with dormancy/rehydration a known future third. + +So the record stores the identifier the policy was last applied *to*, and +:func:`policy_is_stale` compares it against the current one. A path that produces +a new ``awsKbId`` and forgets to re-apply is then not a silent regression: the +next :func:`ensure_retrieve_policy` sees a mismatch and repairs it. The invariant +is checked by comparing two values, which no new code path can bypass by omission +(Requirement 24.12). + +Import weight +------------- +``boto3`` and ``json`` usage stays inside functions where practical; the module +imports stdlib only, per this package's Lambda-image constraint. + +Feature: managed-kb-migration +Requirements: 25.6, 25.7 +""" + +from __future__ import annotations + +import json +import logging +import os +from typing import Any, Dict, Iterable, Mapping, Optional, Sequence, Tuple + +logger = logging.getLogger(__name__) + +#: The actions a shared knowledge base's readers need. ``GetDocumentContent`` is +#: included because a retrieval that returns a citation the caller cannot then +#: fetch is a half-share — the evaluation names both as what resource policies +#: cover. +RETRIEVE_ACTIONS: Tuple[str, ...] = ("bedrock:Retrieve", "bedrock:GetDocumentContent") + +#: Statement id. Fixed so a re-application replaces the platform's own statement +#: rather than accumulating near-duplicates. +POLICY_SID = "PlatformSharedRetrieve" + +POLICY_VERSION = "2012-10-17" + +#: Comma-separated ARNs of the infrastructure identities that retrieve on users' +#: behalf — the AgentCore runtime role, and the App API task role for test-chat. +#: Read at call time, never captured in a default argument. +PRINCIPALS_ENV = "MANAGED_KB_RETRIEVAL_PRINCIPAL_ARNS" + +#: Record attributes tracking what was applied where. +POLICY_KB_ID_ATTR = "policyAwsKbId" +POLICY_REVISION_ATTR = "policyRevisionId" + + +class ResourcePolicyError(RuntimeError): + """A resource policy could not be applied or removed.""" + + +def _region() -> str: + return os.environ.get("AWS_REGION", "us-west-2") + + +def bedrock_agent_client(): + """Control-plane client. ``PutResourcePolicy`` lives on ``bedrock-agent``. + + Verified against the pinned botocore service model: ``PutResourcePolicy`` + takes ``resourceArn`` and ``policy`` (both required) plus an optional + ``expectedRevisionId``, and returns ``resourceArn`` and ``revisionId``. + """ + import boto3 + + return boto3.client("bedrock-agent", region_name=_region()) + + +def knowledge_base_arn( + aws_kb_id: str, + region: Optional[str] = None, + account_id: Optional[str] = None, +) -> str: + """The ARN a policy attaches to. + + Refuses to guess the account. A wrong account in an ARN does not fail loudly — + ``PutResourcePolicy`` would target a resource this caller cannot see, and the + error it raises names a resource the operator did not know existed. Better to + say what is missing. + """ + resolved_account = account_id or os.environ.get("AWS_ACCOUNT_ID") + if not resolved_account: + raise ResourcePolicyError( + f"cannot build a knowledge base ARN for {aws_kb_id} without an account " + f"id: pass account_id or set AWS_ACCOUNT_ID" + ) + return f"arn:aws:bedrock:{region or _region()}:{resolved_account}:knowledge-base/{aws_kb_id}" + + +def retrieval_principals(explicit: Optional[Iterable[str]] = None) -> Tuple[str, ...]: + """The infrastructure identities allowed to retrieve, in a stable order. + + Sorted and de-duplicated so the same configuration always produces the same + policy document — otherwise every call looks like a change and nothing can be + compared. + """ + if explicit is not None: + candidates: Sequence[str] = list(explicit) + else: + candidates = (os.environ.get(PRINCIPALS_ENV) or "").split(",") + return tuple(sorted({arn.strip() for arn in candidates if arn and arn.strip()})) + + +def retrieve_policy_document(kb_arn: str, principals: Sequence[str]) -> Dict[str, Any]: + """The policy granting exactly the shared-read actions to exactly ``principals``. + + No wildcard principal and no wildcard resource: a resource policy whose point + is to narrow access is worse than no policy at all if it widens it instead. + """ + if not principals: + raise ResourcePolicyError( + "refusing to write a resource policy with no principals: an empty " + "principal list is not a narrower grant, it is an unparseable one" + ) + return { + "Version": POLICY_VERSION, + "Statement": [ + { + "Sid": POLICY_SID, + "Effect": "Allow", + "Principal": {"AWS": list(principals)}, + "Action": list(RETRIEVE_ACTIONS), + "Resource": kb_arn, + } + ], + } + + +def policy_is_stale(record: Optional[Mapping[str, Any]]) -> bool: + """Whether the recorded policy target no longer matches the live ``awsKbId``. + + ``True`` when a knowledge base exists in AWS and either no policy target was + ever recorded or the recorded one differs. ``False`` for a record with no + ``awsKbId`` at all: nothing has been provisioned, so there is nothing to be + stale against. + """ + if not record: + return False + aws_kb_id = record.get("awsKbId") + if not aws_kb_id: + return False + return record.get(POLICY_KB_ID_ATTR) != aws_kb_id + + +async def ensure_retrieve_policy( + assistant_id: str, + app_kb_id: str, + *, + shared: bool, + record: Optional[Mapping[str, Any]] = None, + principals: Optional[Iterable[str]] = None, + client=None, + region: Optional[str] = None, + account_id: Optional[str] = None, +) -> Optional[str]: + """Bring the knowledge base's resource policy in line with its sharing state. + + Returns the revision id of a policy that is now in place, or ``None`` when no + policy is wanted or none could be applied. + + Four cases: + + * **Not shared, no policy recorded** — nothing to do. + * **Not shared, policy recorded** — remove it, and forget the target. A + knowledge base that stops being shared should stop carrying the statement + that says it is. + * **Shared, policy current** — nothing to do. This is the common path and it + makes no AWS call, which is what allows callers to invoke this freely. + * **Shared, policy missing or stale** — apply, then record the ``awsKbId`` it + was applied to. + + ``shared`` is supplied by the caller rather than derived here: sharing is an + application fact (visibility plus share records) that lives above this seam, + and this package may not import the assistants package. + """ + from apis.shared.kb_backend import records as r + + if record is None: + import asyncio + + record = await asyncio.to_thread(r.get_kb_record, assistant_id, app_kb_id) + + if not record: + return None + + aws_kb_id = record.get("awsKbId") + recorded_target = record.get(POLICY_KB_ID_ATTR) + + if not shared: + if recorded_target: + await _remove(assistant_id, app_kb_id, recorded_target, client, region, account_id) + return None + + if not aws_kb_id: + # Shared, but nothing provisioned yet. Provisioning is lazy by design, so + # this is ordinary, not an error: the next call after provisioning sees a + # stale (unset) target and applies. + return None + + if not policy_is_stale(record): + return record.get(POLICY_REVISION_ATTR) + + resolved = retrieval_principals(principals) + if not resolved: + # Loud, and not repaired by guessing. A policy with no principals cannot + # be written, and inventing one would either widen access or lock the + # platform out of its own corpus. + logger.error( + f"knowledge base {app_kb_id} is shared but {PRINCIPALS_ENV} names no " + f"principals; no resource policy applied (Requirement 25.6)" + ) + return None + + arn = knowledge_base_arn(aws_kb_id, region, account_id) + document = retrieve_policy_document(arn, resolved) + api = client or bedrock_agent_client() + + try: + response = api.put_resource_policy(resourceArn=arn, policy=json.dumps(document)) + except Exception as exc: + raise ResourcePolicyError( + f"failed to apply the retrieve policy for kb {app_kb_id} on {arn}: {exc}" + ) from exc + + revision_id = response.get("revisionId") + await _record(assistant_id, app_kb_id, aws_kb_id, revision_id) + + if recorded_target and recorded_target != aws_kb_id: + logger.info( + f"re-applied the retrieve policy for kb {app_kb_id}: it was attached " + f"to {recorded_target}, which is no longer this knowledge base's id " + f"(Requirement 25.7)" + ) + return revision_id + + +async def _remove( + assistant_id: str, + app_kb_id: str, + recorded_target: str, + client, + region: Optional[str], + account_id: Optional[str], +) -> None: + """Delete the policy and forget the target, tolerating an absent policy. + + A ``ResourceNotFoundException`` here means the policy or its knowledge base is + already gone, which is the state being asked for. The record is cleared either + way, so a knowledge base cannot be left claiming a policy that does not exist. + """ + api = client or bedrock_agent_client() + arn = knowledge_base_arn(recorded_target, region, account_id) + try: + api.delete_resource_policy(resourceArn=arn) + except Exception as exc: + if type(exc).__name__ not in ("ResourceNotFoundException", "ValidationException"): + raise ResourcePolicyError( + f"failed to remove the retrieve policy for kb {app_kb_id} on {arn}: {exc}" + ) from exc + logger.info( + f"retrieve policy for kb {app_kb_id} was already absent on {arn}; " + f"clearing the record anyway" + ) + await _record(assistant_id, app_kb_id, None, None) + + +async def _record( + assistant_id: str, + app_kb_id: str, + aws_kb_id: Optional[str], + revision_id: Optional[str], +) -> None: + import asyncio + + from apis.shared.kb_backend import records as r + + await asyncio.to_thread( + r.set_resource_policy_state, assistant_id, app_kb_id, aws_kb_id, revision_id + ) + + +__all__ = [ + "POLICY_KB_ID_ATTR", + "POLICY_REVISION_ATTR", + "POLICY_SID", + "PRINCIPALS_ENV", + "RETRIEVE_ACTIONS", + "ResourcePolicyError", + "ensure_retrieve_policy", + "knowledge_base_arn", + "policy_is_stale", + "retrieval_principals", + "retrieve_policy_document", +] diff --git a/backend/src/apis/shared/kb_backend/s3vectors_backend.py b/backend/src/apis/shared/kb_backend/s3vectors_backend.py new file mode 100644 index 000000000..d34e5b73f --- /dev/null +++ b/backend/src/apis/shared/kb_backend/s3vectors_backend.py @@ -0,0 +1,142 @@ +"""The legacy Amazon S3 Vectors backend, behind the common protocol. + +This is the retrieval path every assistant chat has used to date, moved here +unchanged and wrapped in :class:`~apis.shared.kb_backend.protocol.Chunk`. The +only thing this adapter *adds* is the score-direction conversion, and the only +thing it takes away from ``rag_service`` is knowledge of what an S3 Vectors +response looks like. + +Why this delegates instead of copying the query +----------------------------------------------- +``apis.shared.embeddings.bedrock_embeddings.search_assistant_knowledgebase`` +stays where it is and this adapter calls it. It is a published export of two +packages (``apis.shared.embeddings`` and +``apis.app_api.documents.ingestion.embeddings``), and task 5.2 of this spec +still expects to edit it in place. Re-implementing its ``query_vectors`` call +here would mean two copies of the topK/filter/returnDistance construction, which +is precisely the divergence risk this seam exists to remove. What moves here is +everything ``rag_service`` used to know: the response shape, the score +direction, and the parity ``top_k``. + +Score direction — read this before touching :meth:`S3VectorsBackend.search` +--------------------------------------------------------------------------- +S3 Vectors returns cosine **distance**: ``0.0`` is a perfect match and larger is +worse. The protocol canonicalizes on **relevance**, where larger is better. The +conversion happens *here*, once, so that nothing above the seam ever has to know +which direction this particular backend counts in. + +Inverting it raises nothing and logs nothing. Retrieval keeps returning five +chunks, the request keeps succeeding, and the answers quietly get worse. The +guard is ``tests/property/test_pbt_kb_score_direction.py``. + +Ordering is *not* re-sorted here. S3 Vectors already returns results ranked +nearest-first, and today's code passes that order straight through; re-sorting +would be a behaviour change dressed up as a safety measure. The invariant this +adapter owns is that the ``relevance`` values it attaches agree with the order it +returns — descending relevance for ascending distance. + +Import boundary +--------------- +Module-level imports are **stdlib only**; ``boto3`` and the embeddings stack are +imported inside the methods that use them, so importing this module into a +size-constrained Lambda image costs nothing. See +``tests/architecture/test_kb_backend_boundary.py``. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List + +from apis.shared.kb_backend.protocol import ( + DEFAULT_TOP_K, + Chunk, + DocumentSource, + relevance_from_distance, +) + +logger = logging.getLogger(__name__) + + +class S3VectorsBackend: + """Retrieval and ingestion over the S3 Vectors index. + + ``kb_ref`` is the ``App_KB_Id``, which equals the ``assistant_id`` in this + phase; the S3 Vectors index is global and partitioned by an ``assistant_id`` + metadata filter, so the reference is used directly as that filter value. + + Stateless, so a shared instance is safe and no client is held across calls. + """ + + async def search(self, kb_ref: str, query: str, top_k: int = DEFAULT_TOP_K) -> List[Chunk]: + """Query the index and return chunks scored by relevance, best first. + + ``top_k`` is accepted to satisfy the protocol but the underlying query + has always requested a fixed five results (Requirement 3.1), and + narrowing happens above the seam *after* the document-status filter has + run — filtering first and slicing second is what stops a single + incomplete document from silently shrinking a five-chunk answer to four. + Slicing here instead would change that, so this returns what the index + returned. + """ + from apis.shared.embeddings.bedrock_embeddings import search_assistant_knowledgebase + + response = await search_assistant_knowledgebase(kb_ref, query) + vectors = response.get("vectors", []) + return [self._to_chunk(vector) for vector in vectors] + + @staticmethod + def _to_chunk(vector: Dict[str, Any]) -> Chunk: + """Adapt one S3 Vectors hit, converting distance into relevance. + + The ``.get`` defaults mirror the formatting this replaced exactly: a + missing ``text`` or ``key`` became ``""`` and a missing ``distance`` + became ``None``, so they still do. + """ + metadata = vector.get("metadata", {}) + return Chunk( + text=metadata.get("text", ""), + # The conversion. Lower distance ⇒ higher relevance. + relevance=relevance_from_distance(vector.get("distance")), + document_id=metadata.get("document_id", ""), + metadata=metadata, + key=vector.get("key", ""), + ) + + async def ingest(self, kb_ref: str, source: DocumentSource) -> None: + """Embed and store ``source``'s chunks, as the current pipeline does. + + Requires pre-chunked text: splitting is the ingestion pipeline's job + (it owns the tokenizer this package deliberately does not depend on), + so an unchunked source is a programming error rather than something to + paper over with a naive split. + """ + from apis.shared.embeddings.bedrock_embeddings import ( + generate_embeddings, + store_embeddings_in_s3, + ) + + if not source.chunks: + raise ValueError( + f"S3VectorsBackend.ingest requires pre-chunked text for document " + f"{source.document_id}; chunking belongs to the ingestion pipeline" + ) + + embeddings = await generate_embeddings(source.chunks) + await store_embeddings_in_s3( + assistant_id=kb_ref, + document_id=source.document_id, + chunks=source.chunks, + embeddings=embeddings, + metadata={"filename": source.filename, **source.metadata}, + ) + + async def delete_document(self, kb_ref: str, document_id: str) -> None: + """Delete every vector belonging to ``document_id``.""" + from apis.shared.embeddings.bedrock_embeddings import delete_vectors_for_document + + deleted = await delete_vectors_for_document(document_id) + logger.info( + f"S3VectorsBackend: deleted {deleted} vectors for document " + f"{document_id} (kb {kb_ref})" + ) diff --git a/backend/src/apis/shared/kb_backend/tags.py b/backend/src/apis/shared/kb_backend/tags.py new file mode 100644 index 000000000..0c52eab40 --- /dev/null +++ b/backend/src/apis/shared/kb_backend/tags.py @@ -0,0 +1,211 @@ +"""The managed knowledge base tag contract, in one place. + +Requirement 20.11. Tags are not housekeeping here: a tag-filtered +``ListKnowledgeBases`` is how the reconciler tells this platform's knowledge bases +from everything else in the account, and how teardown scopes itself. An untagged — +or mistagged — knowledge base is invisible to both, which means it is never +reclaimed and never deleted, and it keeps billing at $5.00/GB-month with no +CloudFormation console to notice it in. + +Why this module exists +---------------------- +It did not, and the tags drifted three ways: + +* ``provisioning.build_tags`` wrote keys ``prefix``/``env`` with values from + ``PROJECT_PREFIX``/``ENVIRONMENT`` — neither of which the provisioning Lambda is + given, so every knowledge base would have been tagged with the hardcoded + defaults regardless of project or environment. +* ``tombstones.project_tag_filter`` was a hand-written *mirror* of that function, + documented as such. A mirror is a second implementation, and the only thing + keeping two implementations equal is that nobody has edited one of them yet. +* ``kb-migration-construct.ts`` declared a different set of key names entirely + (``ManagedKbPrefix``, …) and exported them plus the correct values as env vars + that **nothing read**. +* ``scripts/teardown/managed-kb.sh`` read a third pair of variables + (``CDK_PROJECT_PREFIX``/``CDK_ENVIRONMENT``) and matched on ``prefix``/``env``. + +Writer and reconciler agreed by luck — both used the same wrong defaults — so the +symptom was not a crash but a teardown that found nothing and reported success. + +So: the keys live here as constants, the value resolution lives here as one +function, and every consumer in every language reads *these* names. +``tests/shared/test_kb_tag_contract.py`` parses the TypeScript and the shell script +and fails if they disagree, because agreement between three languages is not +something a type checker can hold. + +Why the keys are namespaced +--------------------------- +``ManagedKbPrefix`` rather than ``prefix``, and ``ManagedKbEnvironment`` rather +than ``env``. Generic keys collide: many accounts carry an organisation-wide +cost-allocation tag literally called ``env``, and if something else writes it our +filter compares against a value we did not set. The failure mode is a teardown +that skips a knowledge base it owns — the leak this whole contract exists to +prevent. + +Feature: managed-kb-migration +Requirements: 20.11, 20.12, 20.8, 14.1 +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Dict, Mapping, Optional + +logger = logging.getLogger(__name__) + +# ── Tag keys ───────────────────────────────────────────────────────────────── +# +# Mirrored by `MANAGED_KB_TAG_KEYS` in +# `infrastructure/lib/constructs/managed-kb/kb-migration-construct.ts`, and that +# mirroring is asserted by a test rather than trusted. +TAG_KEY_PREFIX = "ManagedKbPrefix" +TAG_KEY_ENVIRONMENT = "ManagedKbEnvironment" +TAG_KEY_APP_KB_ID = "ManagedKbAppKbId" +TAG_KEY_OWNER_USER_ID = "ManagedKbOwnerUserId" + +#: The two keys that scope a destructive pass. Both are required to match: a +#: knowledge base carrying our project prefix but another environment's tag +#: belongs to that environment, and its name looks exactly like ours. +SCOPE_KEYS = (TAG_KEY_PREFIX, TAG_KEY_ENVIRONMENT) + +# ── Environment variables carrying the values ──────────────────────────────── +# +# Set by the CDK construct that owns the provisioning Lambdas, which is the only +# surface that calls `provision_managed_kb`. Named after the tag rather than after +# the project so it is obvious at the call site that changing one changes what +# gets written into AWS. +ENV_TAG_VALUE_PREFIX = "MANAGED_KB_TAG_VALUE_PREFIX" +ENV_TAG_VALUE_ENVIRONMENT = "MANAGED_KB_TAG_VALUE_ENVIRONMENT" + +#: Fallbacks, in order, for a local run or a service that predates the vars above. +#: Deliberately the *same* chain for writing and for filtering — an asymmetric +#: fallback is how a writer and a reader disagree while both look correct. +FALLBACK_PREFIX_VARS = ("PROJECT_PREFIX", "CDK_PROJECT_PREFIX") +FALLBACK_ENVIRONMENT_VARS = ("ENVIRONMENT", "CDK_ENVIRONMENT") + +#: Last resort. Kept so a local run works without configuration, and logged +#: loudly because two deployments that both fall back to it will claim each +#: other's knowledge bases — they would agree with themselves and delete each +#: other's corpora. +DEFAULT_PREFIX = "agentcore" +DEFAULT_ENVIRONMENT = "dev" + + +def _resolve(explicit: Optional[str], primary: str, fallbacks: tuple, default: str, what: str) -> str: + if explicit: + return explicit + value = os.environ.get(primary) + if value: + return value + for name in fallbacks: + value = os.environ.get(name) + if value: + logger.info( + f"managed KB {what} tag resolved from {name}; {primary} is not set. " + f"This is expected for a local run and unexpected in a deployment." + ) + return value + logger.warning( + f"managed KB {what} tag falling back to {default!r}: none of {primary} or " + f"{fallbacks} is set. Two deployments that both reach this default share a " + f"tag scope and will each treat the other's knowledge bases as their own." + ) + return default + + +def tag_prefix(explicit: Optional[str] = None) -> str: + """The project-prefix tag value.""" + return _resolve(explicit, ENV_TAG_VALUE_PREFIX, FALLBACK_PREFIX_VARS, DEFAULT_PREFIX, "prefix") + + +def tag_environment(explicit: Optional[str] = None) -> str: + """The environment tag value.""" + return _resolve( + explicit, + ENV_TAG_VALUE_ENVIRONMENT, + FALLBACK_ENVIRONMENT_VARS, + DEFAULT_ENVIRONMENT, + "environment", + ) + + +def build_tags( + app_kb_id: str, + owner_user_id: str, + project_prefix: Optional[str] = None, + environment: Optional[str] = None, +) -> Dict[str, str]: + """The complete tag set written at ``CreateKnowledgeBase`` time. + + The owner tag must be opaque (Requirement 20.12). An email address here would + put PII in a field readable by anyone holding + ``bedrock:ListKnowledgeBases``, and unlike a database column a tag cannot be + scrubbed retroactively from the audit trail it has already entered. An + address-shaped value is therefore rejected rather than trimmed: silently + dropping it would hide the caller's mistake. + """ + if "@" in owner_user_id: + raise ValueError( + "ownerUserId tag must be an opaque identifier, never an email address " + "or other personally identifying value (Requirement 20.12)" + ) + return { + TAG_KEY_PREFIX: tag_prefix(project_prefix), + TAG_KEY_ENVIRONMENT: tag_environment(environment), + TAG_KEY_APP_KB_ID: app_kb_id, + TAG_KEY_OWNER_USER_ID: owner_user_id, + } + + +def project_tag_filter( + project_prefix: Optional[str] = None, + environment: Optional[str] = None, +) -> Dict[str, str]: + """The subset of tags a knowledge base must carry to be considered ours. + + Derived from the same resolution :func:`build_tags` uses, not mirrored from + it. That is the entire point of this module: a reader that re-derives what the + writer wrote is a reader that can be wrong on its own. + """ + return { + TAG_KEY_PREFIX: tag_prefix(project_prefix), + TAG_KEY_ENVIRONMENT: tag_environment(environment), + } + + +def matches_project( + tags: Optional[Mapping[str, Any]], + project_prefix: Optional[str] = None, + environment: Optional[str] = None, +) -> bool: + """Whether these tags identify a knowledge base this deployment owns. + + ``False`` for absent or unreadable tags. Unknown ownership is not ownership, + and refusing to act on a resource we cannot attribute is the only safe + direction for a pass that deletes things. + """ + if not tags: + return False + expected = project_tag_filter(project_prefix, environment) + return all(str(tags.get(key, "")) == value for key, value in expected.items()) + + +__all__ = [ + "DEFAULT_ENVIRONMENT", + "DEFAULT_PREFIX", + "ENV_TAG_VALUE_ENVIRONMENT", + "ENV_TAG_VALUE_PREFIX", + "FALLBACK_ENVIRONMENT_VARS", + "FALLBACK_PREFIX_VARS", + "SCOPE_KEYS", + "TAG_KEY_APP_KB_ID", + "TAG_KEY_ENVIRONMENT", + "TAG_KEY_OWNER_USER_ID", + "TAG_KEY_PREFIX", + "build_tags", + "matches_project", + "project_tag_filter", + "tag_environment", + "tag_prefix", +] diff --git a/backend/src/apis/shared/kb_backend/tombstones.py b/backend/src/apis/shared/kb_backend/tombstones.py new file mode 100644 index 000000000..ac44796a4 --- /dev/null +++ b/backend/src/apis/shared/kb_backend/tombstones.py @@ -0,0 +1,921 @@ +"""Tombstoned deletion sagas for managed knowledge bases. + +Every delete here either completes or leaves a durable, retryable work item. That +is the whole requirement (Requirement 13), and it exists because a failed delete +of a managed knowledge base is not a crash — it is a **silent recurring bill**. + +The ordering is the mechanism +----------------------------- +1. Write the Tombstone to DynamoDB. +2. *Then* call AWS. +3. Poll until AWS reports the resource genuinely absent. +4. *Only then* clear the Tombstone. + +Reversing steps 1 and 2 looks equivalent and is not. A crash between the AWS call +and the database write leaves a half-deleted, still-billed resource that no record +points at, nothing alarms on, and no code will ever revisit. Written +tombstone-first, the same crash leaves a row that :func:`iter_tombstones` finds +(Requirement 13.8) and that a later pass can retry. + +Clearing early is the same defect wearing the opposite hat: a tombstone cleared on +the strength of an *accepted* delete call describes a resource AWS may still be +holding, and holding it is what costs money. So :func:`clear_kb_tombstone` is +never called on the accept path — it is reachable only after +:func:`confirm_knowledge_base_absent` has returned true. + +No TTL. Deliberately. +--------------------- +A Tombstone is cleared by confirmed deletion or it stays. Attaching a TTL would +let DynamoDB quietly remove the evidence of a delete that never finished, which +recreates precisely the silent-leak class this module exists to close. The same +reasoning bans a TTL on the KB_Record itself (Requirement 13.6): +:func:`remove_kb_record` refuses outright unless the caller can show confirmation. + +"Accepted" is not "gone" +------------------------ +``DeleteKnowledgeBase`` returns ``status: DELETING`` and the resource lives on for +a measured **2-6 minutes**. There is no waiter, so absence is established by +polling ``ListKnowledgeBases`` until the identifier stops appearing +(Requirement 13.4), with a window comfortably past the observed worst case. + +``DELETE_UNSUCCESSFUL`` is a terminal *operator* state, not a completed delete +(Requirement 13.7). The dev account has contained one since 2025-11-24 that no +reconciler would ever have noticed. Observing it stops the poll, records the state +on the Tombstone, and leaves the Tombstone standing. + +Why the tag filter costs a describe call per knowledge base +----------------------------------------------------------- +``ListKnowledgeBases`` has no tag-filter parameter and its summaries carry neither +``knowledgeBaseArn`` nor ``createdAt`` — verified against the packaged service +model, where ``KnowledgeBaseSummary`` is +``{knowledgeBaseId, name, description, status, updatedAt}``. Both of those are +needed: the ARN to read tags, and ``createdAt`` for the Reconciler's age gate. So +:func:`iter_project_knowledge_bases` pages the list and calls +``GetKnowledgeBase`` per entry. The alternative — synthesizing the ARN from the +region and account — trades a read call for a brittle string, and the caller here +is a daily job. + +Import boundary +--------------- +Module-level imports are stdlib plus this package's own stdlib-only modules; +``boto3`` and ``botocore`` are function-local. Nothing here imports +``apis.shared.assistants``, whose ``__init__`` pulls in the embeddings stack and +would blow the migration Lambda image budget. Enforced by +``tests/architecture/test_kb_backend_boundary.py``. +""" + +from __future__ import annotations + +import logging +import os +import time +from dataclasses import dataclass +from decimal import Decimal +from typing import Any, Callable, Dict, Iterator, List, Mapping, Optional + +from apis.shared.kb_backend.metrics import emit_count +from apis.shared.kb_backend.records import ( + document_tombstone_sk, + kb_pk, + kb_sk, + kb_tombstone_sk, +) + +logger = logging.getLogger(__name__) + +# ── Intents ────────────────────────────────────────────────────────────────── +# +# Recorded on the Tombstone so a retry knows which saga to resume without having +# to infer it from the sort key's shape. +INTENT_DELETE_KB = "delete_kb" +INTENT_DELETE_DOCUMENT = "delete_document" + +#: Attribute name flagging a tombstone whose ``PK`` is *not* a real assistant +#: partition. +#: +#: The reconciler deletes orphans — knowledge bases with no KB_Record — and an +#: orphan by definition carries no assistant id to anchor on, so its tombstone +#: lands in a partition derived from whatever identifier the tags did preserve. +#: That item is a genuine work record and must be kept, but a reader must not +#: mistake the partition for an assistant that exists, and +#: ``iter_tombstones()`` will never return it. This attribute +#: says so on the item, in place of a comment nobody triaging at 3am will read. +SYNTHETIC_PARTITION = "syntheticPartition" + +# ── AWS states ─────────────────────────────────────────────────────────────── +# +# Copied from the packaged service model's ``KnowledgeBaseStatus`` enum: +# ``CREATING | ACTIVE | DELETING | UPDATING | FAILED | DELETE_UNSUCCESSFUL | +# UPDATE_UNSUCCESSFUL``. +KB_STATUS_DELETING = "DELETING" +KB_STATUS_DELETE_UNSUCCESSFUL = "DELETE_UNSUCCESSFUL" + +#: From ``DocumentStatus``. A document AWS reports ``NOT_FOUND`` is gone; anything +#: else — including ``DELETING`` and ``DELETE_IN_PROGRESS`` — is still present. +DOCUMENT_STATUS_NOT_FOUND = "NOT_FOUND" + +# ── Poll windows (Requirement 13.4) ────────────────────────────────────────── +# +# Deletion was measured at 2-6 minutes, so the floor is 360 s and this sits above +# it. These are read *at call time* rather than bound as default arguments, +# because a default argument is evaluated once at import and cannot be patched: +# an earlier version of a sibling poller bound its timeout that way and a test +# that shortened the window had no effect at all, silently waiting the full +# production timeout instead. See `wait_until_retrievable` in the ingestion +# consumer for the same note. +KB_DELETE_POLL_TIMEOUT_SECONDS = 480.0 +KB_DELETE_POLL_INTERVAL_SECONDS = 10.0 + +DOCUMENT_DELETE_POLL_TIMEOUT_SECONDS = 120.0 +DOCUMENT_DELETE_POLL_INTERVAL_SECONDS = 2.0 + +#: ``ListKnowledgeBases`` page size. The list is always paged to exhaustion; this +#: only trades call count against payload size. +LIST_PAGE_SIZE = 100 + +# ── Metrics ────────────────────────────────────────────────────────────────── +METRIC_TOMBSTONE_WRITTEN = "KbTombstoneWritten" +METRIC_TOMBSTONE_CLEARED = "KbTombstoneCleared" + +#: A delete that was accepted but never confirmed. Sustained non-zero is the only +#: signal that the delete saga is leaking paid resources. +METRIC_TOMBSTONE_SURVIVED = "KbTombstoneSurvived" + +#: Requirement 13.7. Needs an alarm, not a dashboard: nothing clears this state +#: on its own. +METRIC_DELETE_UNSUCCESSFUL = "KbDeleteUnsuccessful" + + +class TombstoneError(RuntimeError): + """A tombstoned delete could not be completed.""" + + +class DeleteNotConfirmed(TombstoneError): + """AWS never reported the resource absent within the poll window. + + Retryable. The Tombstone is deliberately left in place, so the work item + outlives this process. + """ + + +class DeleteUnsuccessful(TombstoneError): + """AWS reported ``DELETE_UNSUCCESSFUL`` (Requirement 13.7). + + Distinct from :class:`DeleteNotConfirmed` because it is *not* a matter of + waiting longer. It is an actionable operator state that persists until someone + intervenes, and it must never be mistaken for a completed delete. + """ + + +class ServiceRoleStillInUse(TombstoneError): + """Refuses to delete a service role that still has knowledge bases. + + Requirement 13.5. Removing the role first is a documented route *into* + ``DELETE_UNSUCCESSFUL``: the pending deletion needs the role it was created + with, and without it the knowledge base can be neither deleted nor recovered. + """ + + +class RecordRemovalRefused(TombstoneError): + """Refuses to remove a KB_Record before AWS confirmed the deletion. + + Requirement 13.6. The record is the only pointer to the AWS identifiers, so + dropping it early converts a retryable delete into an untraceable one. + """ + + +@dataclass(frozen=True) +class KnowledgeBaseFacts: + """What AWS says about one knowledge base. + + ``created_at`` is **AWS's own** ``createdAt``, carried through unmodified. The + Reconciler's age gate depends on that provenance: substituting the time this + process happened to look would make a reconciler that was down for a week + treat every knowledge base in the account as brand new (Requirement 14.3). + """ + + kb_id: str + name: str + status: str + arn: Optional[str] = None + created_at: Optional[Any] = None + tags: Mapping[str, str] = None # type: ignore[assignment] + + +@dataclass(frozen=True) +class DeleteOutcome: + """The result of one saga run. + + ``confirmed`` means AWS reported the resource absent — the only condition + under which the Tombstone was cleared. ``tombstone_cleared`` is reported + separately rather than inferred so a test can catch the two drifting apart. + """ + + confirmed: bool + tombstone_cleared: bool + already_absent: bool = False + delete_unsuccessful: bool = False + polls: int = 0 + + +# ── DynamoDB plumbing ──────────────────────────────────────────────────────── +def _table(): + import boto3 + + return boto3.resource("dynamodb").Table(os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"]) + + +def _now_iso() -> str: + from apis.shared.timestamps import utc_now_iso + + return utc_now_iso() + + +# ── Tombstone writes ───────────────────────────────────────────────────────── +def _write_tombstone( + assistant_id: str, + sort_key: str, + intent: str, + attributes: Mapping[str, Any], +) -> Dict[str, Any]: + """Upsert a Tombstone, preserving the original ``createdAt`` and counting attempts. + + An upsert rather than a ``put_item`` because a retried saga must not restart + the clock. ``createdAt`` is written through ``if_not_exists`` so it records + when the delete was *first* attempted — the number an operator triaging a + stuck tombstone actually wants — while ``attempts`` accumulates with ``ADD``, + which is atomic and needs no read. + + No ``ttl`` attribute is written, and none may be added. See the module + docstring: expiry would silently discard the evidence of an unfinished delete. + """ + now = _now_iso() + values: Dict[str, Any] = { + ":intent": intent, + ":now": now, + ":one": Decimal(1), + } + sets = [ + "intent = :intent", + "createdAt = if_not_exists(createdAt, :now)", + "updatedAt = :now", + ] + for index, (key, value) in enumerate(sorted(attributes.items())): + if value is None: + continue + placeholder = f":a{index}" + sets.append(f"{key} = {placeholder}") + values[placeholder] = value + + _table().update_item( + Key={"PK": kb_pk(assistant_id), "SK": sort_key}, + UpdateExpression=f"SET {', '.join(sets)} ADD attempts :one", + ExpressionAttributeValues=values, + ) + emit_count(METRIC_TOMBSTONE_WRITTEN, dimensions={"intent": intent}) + return {"PK": kb_pk(assistant_id), "SK": sort_key, "intent": intent, "createdAt": now} + + +def write_kb_tombstone( + assistant_id: str, + app_kb_id: str, + aws_kb_id: Optional[str] = None, + aws_data_source_id: Optional[str] = None, + extra_attributes: Optional[Mapping[str, Any]] = None, +) -> Dict[str, Any]: + """Mark a whole-knowledge-base delete as intended. Call this *before* AWS. + + ``extra_attributes`` lets a caller that is not deleting on behalf of a known + assistant say so on the item itself — see ``SYNTHETIC_PARTITION`` below. + """ + attributes: Dict[str, Any] = { + "appKbId": app_kb_id, + "awsKbId": aws_kb_id, + "awsDataSourceId": aws_data_source_id, + } + if extra_attributes: + attributes.update(extra_attributes) + return _write_tombstone( + assistant_id, + kb_tombstone_sk(app_kb_id), + INTENT_DELETE_KB, + attributes, + ) + + +def write_document_tombstone( + assistant_id: str, + app_kb_id: str, + document_id: str, + aws_kb_id: Optional[str] = None, + aws_data_source_id: Optional[str] = None, +) -> Dict[str, Any]: + """Mark a single-document delete as intended. Call this *before* AWS.""" + return _write_tombstone( + assistant_id, + document_tombstone_sk(app_kb_id, document_id), + INTENT_DELETE_DOCUMENT, + { + "appKbId": app_kb_id, + "documentId": document_id, + "awsKbId": aws_kb_id, + "awsDataSourceId": aws_data_source_id, + }, + ) + + +def record_tombstone_error( + assistant_id: str, + sort_key: str, + error: str, + aws_status: Optional[str] = None, +) -> None: + """Annotate a surviving Tombstone with why it survived. + + Never raises. The saga has already failed by the time this is reached, and + losing the annotation is strictly better than replacing a precise failure with + a DynamoDB error from the bookkeeping. + """ + sets = ["lastError = :err", "updatedAt = :now"] + values: Dict[str, Any] = {":err": error[:1024], ":now": _now_iso()} + if aws_status: + sets.append("awsStatus = :status") + values[":status"] = aws_status + + try: + _table().update_item( + Key={"PK": kb_pk(assistant_id), "SK": sort_key}, + UpdateExpression=f"SET {', '.join(sets)}", + ExpressionAttributeValues=values, + ) + except Exception as exc: # noqa: BLE001 - bookkeeping must not mask the real failure + logger.warning(f"could not annotate tombstone {sort_key}: {exc}") + + +def _clear(assistant_id: str, sort_key: str, intent: str) -> bool: + _table().delete_item(Key={"PK": kb_pk(assistant_id), "SK": sort_key}) + emit_count(METRIC_TOMBSTONE_CLEARED, dimensions={"intent": intent}) + return True + + +def clear_kb_tombstone(assistant_id: str, app_kb_id: str, confirmed_absent: bool) -> bool: + """Clear a knowledge-base Tombstone. Refuses unless AWS confirmed absence. + + ``confirmed_absent`` is a required positional argument rather than a keyword + with a convenient default, because the failure mode being guarded against is a + caller who *forgot* the confirmation step. A default of ``True`` would make + the unsafe call the short one; there is no default at all, so the caller has + to state what it knows. + """ + if not confirmed_absent: + raise TombstoneError( + f"refusing to clear the tombstone for kb {app_kb_id}: AWS has not " + f"confirmed the knowledge base is absent. An accepted delete call is " + f"not a completed deletion (Requirement 13.3), and clearing here " + f"would discard the only work item for a resource still being billed." + ) + return _clear(assistant_id, kb_tombstone_sk(app_kb_id), INTENT_DELETE_KB) + + +def clear_document_tombstone( + assistant_id: str, + app_kb_id: str, + document_id: str, + confirmed_absent: bool, +) -> bool: + """Clear a document Tombstone. Refuses unless AWS confirmed absence.""" + if not confirmed_absent: + raise TombstoneError( + f"refusing to clear the tombstone for document {document_id}: AWS has " + f"not confirmed it is absent" + ) + return _clear( + assistant_id, document_tombstone_sk(app_kb_id, document_id), INTENT_DELETE_DOCUMENT + ) + + +def iter_tombstones(assistant_id: str) -> List[Dict[str, Any]]: + """Surviving Tombstones for one assistant, as retryable work items (Req 13.8). + + Keyed on the ``KBTOMB#`` prefix, so a whole-KB tombstone and its documents' + tombstones come back together and in that order — which is the order a retry + wants them. + """ + from boto3.dynamodb.conditions import Key + + response = _table().query( + KeyConditionExpression=Key("PK").eq(kb_pk(assistant_id)) + & Key("SK").begins_with("KBTOMB#") + ) + return response.get("Items", []) + + +def remove_kb_record(assistant_id: str, app_kb_id: str, confirmed_absent: bool) -> None: + """Delete the KB_Record. Refuses unless AWS confirmed the deletion (Req 13.6). + + The record holds the only mapping from ``App_KB_Id`` to the AWS identifiers. + Removing it while AWS still holds the knowledge base turns a resource that a + tombstone could still find into one nothing can address — the exact leak this + module exists to prevent, produced by the cleanup step rather than the crash. + """ + if not confirmed_absent: + raise RecordRemovalRefused( + f"refusing to remove the KB_Record for {app_kb_id} before AWS confirms " + f"deletion (Requirement 13.6); the record is the only pointer to the " + f"AWS identifiers" + ) + _table().delete_item(Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}) + + +# ── AWS listing, tag-filtered and paginated (Requirement 14.1) ─────────────── +# +# The tag contract itself lives in ``kb_backend.tags``. These two functions are +# thin re-exports kept for their existing callers. +# +# ⚠️ They used to be hand-written *mirrors* of ``provisioning.build_tags``, +# documented as such — and they had drifted: different key names +# (``prefix``/``env`` against the construct's ``ManagedKbPrefix``/ +# ``ManagedKbEnvironment``) and values read from environment variables the +# provisioning Lambda is never given. A mirror is a second implementation, and the +# only thing keeping two implementations equal is that nobody has edited one yet. +def project_tag_filter( + project_prefix: Optional[str] = None, + environment: Optional[str] = None, +) -> Dict[str, str]: + """The tags that identify this platform's knowledge bases. + + Only the two scope keys are matched: the app id and owner id vary per resource + and are identity, not scope. + """ + from apis.shared.kb_backend.tags import project_tag_filter as _canonical + + return _canonical(project_prefix, environment) + + +def matches_project_tags(tags: Optional[Mapping[str, str]], expected: Mapping[str, str]) -> bool: + """True when every expected tag is present with the expected value. + + Absent or empty tags never match. An untagged knowledge base is out of scope + by construction, which is the conservative direction: this predicate gates + deletion, so a false negative leaves a resource alone while a false positive + deletes someone else's. + """ + if not tags: + return False + return all(str(tags.get(key, "")) == value for key, value in expected.items()) + + +def iter_knowledge_base_summaries(client, page_size: Optional[int] = None) -> Iterator[Dict[str, Any]]: + """Every ``KnowledgeBaseSummary`` in the account, paging to exhaustion. + + Hand-rolled paging rather than ``get_paginator`` so that a stubbed client in a + test is a plain object with one method, not something that has to satisfy + botocore's paginator protocol. Reading only the first page would make the + Reconciler's judgement depend on account size: every knowledge base past page + one would look like a missing-vector record and every orphan there would go + unbilled-for-ever. + """ + if page_size is None: + page_size = LIST_PAGE_SIZE + + token: Optional[str] = None + while True: + kwargs: Dict[str, Any] = {"maxResults": page_size} + if token: + kwargs["nextToken"] = token + response = client.list_knowledge_bases(**kwargs) + for summary in response.get("knowledgeBaseSummaries") or []: + yield summary + token = response.get("nextToken") + if not token: + return + + +def describe_knowledge_base(client, kb_id: str) -> Optional[Dict[str, Any]]: + """``GetKnowledgeBase``, or ``None`` if it has already gone. + + A ``ResourceNotFoundException`` between the list and the describe is normal — + something else deleted it, or this saga's own earlier attempt finally landed — + and means exactly what the caller wants to know. + """ + from botocore.exceptions import ClientError + + try: + response = client.get_knowledge_base(knowledgeBaseId=kb_id) + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") == "ResourceNotFoundException": + return None + raise + return response.get("knowledgeBase") or None + + +def knowledge_base_tags(client, arn: str) -> Dict[str, str]: + """Tags for one knowledge base. A read failure yields ``{}``, never a match. + + Failing closed matters here: ``{}`` cannot satisfy + :func:`matches_project_tags`, so a knowledge base whose tags could not be read + is left alone rather than deleted on the strength of a failed lookup. + """ + from botocore.exceptions import ClientError + + try: + return dict((client.list_tags_for_resource(resourceArn=arn) or {}).get("tags") or {}) + except ClientError as exc: + logger.warning(f"could not read tags for {arn}: {exc}") + return {} + + +def iter_project_knowledge_bases( + client, + project_prefix: Optional[str] = None, + environment: Optional[str] = None, + page_size: Optional[int] = None, +) -> Iterator[KnowledgeBaseFacts]: + """This project's knowledge bases, with AWS's ``createdAt`` and status. + + Paginated (Requirement 14.1) and tag-filtered. The filter is applied to tags + read from AWS rather than to the name, because a name is a convention this + code chose and a tag is a fact recorded on the resource: a knowledge base + created by an older naming scheme is still ours, and one that merely happens + to share our prefix is not. + """ + expected = project_tag_filter(project_prefix, environment) + + for summary in iter_knowledge_base_summaries(client, page_size=page_size): + kb_id = summary.get("knowledgeBaseId") + if not kb_id: + continue + + described = describe_knowledge_base(client, kb_id) + if described is None: + continue + + arn = described.get("knowledgeBaseArn") + tags = knowledge_base_tags(client, arn) if arn else {} + if not matches_project_tags(tags, expected): + continue + + yield KnowledgeBaseFacts( + kb_id=kb_id, + name=described.get("name") or summary.get("name") or "", + status=described.get("status") or summary.get("status") or "", + arn=arn, + # AWS's own timestamp, untouched. See KnowledgeBaseFacts. + created_at=described.get("createdAt"), + tags=tags, + ) + + +# ── Confirmation by polling (Requirement 13.3, 13.4) ───────────────────────── +def confirm_knowledge_base_absent( + client, + aws_kb_id: str, + timeout_seconds: Optional[float] = None, + interval_seconds: Optional[float] = None, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, +) -> DeleteOutcome: + """Poll ``ListKnowledgeBases`` until ``aws_kb_id`` stops appearing. + + Absence is established from the **list**, not from the delete call's return + value and not from a single ``GetKnowledgeBase``: the delete is asynchronous + and returns ``DELETING`` while the resource is still there and still billed. + + Returns as soon as the identifier is gone. Raises :class:`DeleteUnsuccessful` + the moment ``DELETE_UNSUCCESSFUL`` is observed — that state does not resolve + by waiting, so continuing to poll would burn the window and then report the + wrong reason. Raises :class:`DeleteNotConfirmed` on timeout. + + Both windows resolve from the module constants *at call time*. Bound as + default arguments they would be fixed at import and unpatchable, and a test + that shortened them would sit through the full production wait while + appearing to pass. + """ + if timeout_seconds is None: + timeout_seconds = KB_DELETE_POLL_TIMEOUT_SECONDS + if interval_seconds is None: + interval_seconds = KB_DELETE_POLL_INTERVAL_SECONDS + + deadline = monotonic() + timeout_seconds + polls = 0 + + while True: + polls += 1 + present: Optional[Dict[str, Any]] = None + for summary in iter_knowledge_base_summaries(client): + if summary.get("knowledgeBaseId") == aws_kb_id: + present = summary + break + + if present is None: + return DeleteOutcome(confirmed=True, tombstone_cleared=False, polls=polls) + + status = present.get("status") or "" + if status == KB_STATUS_DELETE_UNSUCCESSFUL: + emit_count(METRIC_DELETE_UNSUCCESSFUL) + raise DeleteUnsuccessful( + f"knowledge base {aws_kb_id} is in {KB_STATUS_DELETE_UNSUCCESSFUL}. " + f"This is an operator state, not a completed delete: it does not " + f"clear on its own and the resource is still billed. The tombstone " + f"is being left in place as the work item." + ) + + if monotonic() >= deadline: + raise DeleteNotConfirmed( + f"knowledge base {aws_kb_id} still present after {timeout_seconds}s " + f"(last status {status!r}) across {polls} polls; leaving the " + f"tombstone as a retryable work item" + ) + + sleep(interval_seconds) + + +def confirm_document_absent( + client, + aws_kb_id: str, + aws_data_source_id: str, + document_id: str, + timeout_seconds: Optional[float] = None, + interval_seconds: Optional[float] = None, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, +) -> DeleteOutcome: + """Poll ``GetKnowledgeBaseDocuments`` until the document reports ``NOT_FOUND``. + + Only ``NOT_FOUND`` (or an empty detail list) counts as absent. ``DELETING`` + and ``DELETE_IN_PROGRESS`` are explicitly *present*: treating them as done is + the document-scale version of trusting the accepted delete call. + """ + from apis.shared.kb_backend.managed_backend import document_identifier + + if timeout_seconds is None: + timeout_seconds = DOCUMENT_DELETE_POLL_TIMEOUT_SECONDS + if interval_seconds is None: + interval_seconds = DOCUMENT_DELETE_POLL_INTERVAL_SECONDS + + deadline = monotonic() + timeout_seconds + polls = 0 + + while True: + polls += 1 + response = client.get_knowledge_base_documents( + knowledgeBaseId=aws_kb_id, + dataSourceId=aws_data_source_id, + documentIdentifiers=[document_identifier(document_id)], + ) + details = response.get("documentDetails") or [] + statuses = {detail.get("status") for detail in details} + + if not details or statuses <= {DOCUMENT_STATUS_NOT_FOUND}: + return DeleteOutcome(confirmed=True, tombstone_cleared=False, polls=polls) + + if monotonic() >= deadline: + raise DeleteNotConfirmed( + f"document {document_id} still present in kb {aws_kb_id} after " + f"{timeout_seconds}s (statuses {sorted(s for s in statuses if s)}); " + f"leaving the tombstone as a retryable work item" + ) + + sleep(interval_seconds) + + +# ── Sagas ──────────────────────────────────────────────────────────────────── +def delete_knowledge_base( + assistant_id: str, + app_kb_id: str, + aws_kb_id: str, + aws_data_source_id: Optional[str] = None, + client=None, + remove_record: bool = False, + extra_attributes: Optional[Mapping[str, Any]] = None, + timeout_seconds: Optional[float] = None, + interval_seconds: Optional[float] = None, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, +) -> DeleteOutcome: + """Delete a knowledge base under a Tombstone. + + The four steps run strictly in order, and the order is the guarantee: + + 1. **Tombstone first.** Written before any AWS call, so a crash anywhere below + leaves a work item rather than a resource nothing knows about. + 2. **Ask AWS.** ``ResourceNotFoundException`` is success, not failure — an + earlier attempt got there, and the tombstone should still be cleared. + 3. **Confirm by polling.** The accepted call is ignored as evidence. + 4. **Clear the Tombstone**, and only now, optionally, the KB_Record. + + On any failure the Tombstone survives, annotated with the reason, and the + exception propagates so the invocation fails and its retry or DLQ fires. + """ + from apis.shared.kb_backend.managed_backend import bedrock_agent_client + from botocore.exceptions import ClientError + + if client is None: + client = bedrock_agent_client() + + sort_key = kb_tombstone_sk(app_kb_id) + + # Step 1. Before AWS. Always. + write_kb_tombstone( + assistant_id, + app_kb_id, + aws_kb_id, + aws_data_source_id, + extra_attributes=extra_attributes, + ) + + already_absent = False + try: + # Step 2. + try: + client.delete_knowledge_base(knowledgeBaseId=aws_kb_id) + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") != "ResourceNotFoundException": + raise + already_absent = True + logger.info( + f"knowledge base {aws_kb_id} was already absent; treating the " + f"delete as complete and clearing its tombstone" + ) + + # Step 3. "Accepted" is not "gone" — establish absence from the list. + outcome = confirm_knowledge_base_absent( + client, + aws_kb_id, + timeout_seconds=timeout_seconds, + interval_seconds=interval_seconds, + sleep=sleep, + monotonic=monotonic, + ) + except DeleteUnsuccessful as exc: + record_tombstone_error( + assistant_id, sort_key, str(exc), aws_status=KB_STATUS_DELETE_UNSUCCESSFUL + ) + raise + except Exception as exc: + emit_count(METRIC_TOMBSTONE_SURVIVED, dimensions={"intent": INTENT_DELETE_KB}) + record_tombstone_error(assistant_id, sort_key, str(exc)) + raise + + # Step 4. Reachable only with confirmation in hand. + clear_kb_tombstone(assistant_id, app_kb_id, outcome.confirmed) + + if remove_record: + remove_kb_record(assistant_id, app_kb_id, outcome.confirmed) + + return DeleteOutcome( + confirmed=True, + tombstone_cleared=True, + already_absent=already_absent, + polls=outcome.polls, + ) + + +def delete_document( + assistant_id: str, + app_kb_id: str, + document_id: str, + aws_kb_id: str, + aws_data_source_id: str, + client=None, + timeout_seconds: Optional[float] = None, + interval_seconds: Optional[float] = None, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, +) -> DeleteOutcome: + """Delete one document under a Tombstone. Same ordering as the KB saga.""" + from apis.shared.kb_backend.managed_backend import ( + bedrock_agent_client, + document_identifier, + ) + + if client is None: + client = bedrock_agent_client() + + sort_key = document_tombstone_sk(app_kb_id, document_id) + + # Step 1. Before AWS. Always. + write_document_tombstone( + assistant_id, app_kb_id, document_id, aws_kb_id, aws_data_source_id + ) + + try: + client.delete_knowledge_base_documents( + knowledgeBaseId=aws_kb_id, + dataSourceId=aws_data_source_id, + documentIdentifiers=[document_identifier(document_id)], + ) + outcome = confirm_document_absent( + client, + aws_kb_id, + aws_data_source_id, + document_id, + timeout_seconds=timeout_seconds, + interval_seconds=interval_seconds, + sleep=sleep, + monotonic=monotonic, + ) + except Exception as exc: + emit_count(METRIC_TOMBSTONE_SURVIVED, dimensions={"intent": INTENT_DELETE_DOCUMENT}) + record_tombstone_error(assistant_id, sort_key, str(exc)) + raise + + clear_document_tombstone(assistant_id, app_kb_id, document_id, outcome.confirmed) + return DeleteOutcome(confirmed=True, tombstone_cleared=True, polls=outcome.polls) + + +# ── Service-role teardown guard (Requirement 13.5) ─────────────────────────── +def knowledge_bases_using_role( + client, + role_arn: str, + project_prefix: Optional[str] = None, + environment: Optional[str] = None, +) -> List[str]: + """Identifiers of this project's knowledge bases still using ``role_arn``. + + Read from ``GetKnowledgeBase``'s ``roleArn`` rather than from our own records, + because the question is what AWS still believes — and a knowledge base our + database has forgotten is exactly the one that makes deleting the role + dangerous. + """ + outstanding: List[str] = [] + for facts in iter_project_knowledge_bases( + client, project_prefix=project_prefix, environment=environment + ): + described = describe_knowledge_base(client, facts.kb_id) + if described is None: + continue + if described.get("roleArn") == role_arn: + outstanding.append(facts.kb_id) + return outstanding + + +def assert_service_role_deletable( + client, + role_arn: str, + project_prefix: Optional[str] = None, + environment: Optional[str] = None, +) -> None: + """Raise unless every knowledge base using ``role_arn`` is confirmed absent. + + Called by teardown before it touches the role. A knowledge base mid-``DELETING`` + still counts as present: it needs the role to finish, and pulling the role out + from under it is a documented route into ``DELETE_UNSUCCESSFUL``, which is + unrecoverable without support. + """ + outstanding = knowledge_bases_using_role( + client, role_arn, project_prefix=project_prefix, environment=environment + ) + if outstanding: + raise ServiceRoleStillInUse( + f"refusing to delete service role {role_arn}: {len(outstanding)} " + f"knowledge base(s) still reference it ({', '.join(sorted(outstanding))}). " + f"Delete them and confirm their absence first (Requirement 13.5); " + f"removing the role while one is still DELETING can strand it in " + f"{KB_STATUS_DELETE_UNSUCCESSFUL}." + ) + + +__all__ = [ + "DOCUMENT_DELETE_POLL_INTERVAL_SECONDS", + "DOCUMENT_DELETE_POLL_TIMEOUT_SECONDS", + "DOCUMENT_STATUS_NOT_FOUND", + "INTENT_DELETE_DOCUMENT", + "INTENT_DELETE_KB", + "KB_DELETE_POLL_INTERVAL_SECONDS", + "KB_DELETE_POLL_TIMEOUT_SECONDS", + "KB_STATUS_DELETE_UNSUCCESSFUL", + "KB_STATUS_DELETING", + "LIST_PAGE_SIZE", + "METRIC_DELETE_UNSUCCESSFUL", + "METRIC_TOMBSTONE_CLEARED", + "METRIC_TOMBSTONE_SURVIVED", + "METRIC_TOMBSTONE_WRITTEN", + "SYNTHETIC_PARTITION", + "DeleteNotConfirmed", + "DeleteOutcome", + "DeleteUnsuccessful", + "KnowledgeBaseFacts", + "RecordRemovalRefused", + "ServiceRoleStillInUse", + "TombstoneError", + "assert_service_role_deletable", + "clear_document_tombstone", + "clear_kb_tombstone", + "confirm_document_absent", + "confirm_knowledge_base_absent", + "delete_document", + "delete_knowledge_base", + "describe_knowledge_base", + "iter_knowledge_base_summaries", + "iter_project_knowledge_bases", + "iter_tombstones", + "knowledge_base_tags", + "knowledge_bases_using_role", + "matches_project_tags", + "project_tag_filter", + "record_tombstone_error", + "remove_kb_record", + "write_document_tombstone", + "write_kb_tombstone", +] diff --git a/backend/src/apis/shared/observability/emf.py b/backend/src/apis/shared/observability/emf.py index 99aeec50a..3844b64cc 100644 --- a/backend/src/apis/shared/observability/emf.py +++ b/backend/src/apis/shared/observability/emf.py @@ -114,6 +114,55 @@ def emit_prompt_cache_metrics( logger.debug("EMF emission skipped: %s", e) +def emit_emf_metrics( + namespace: str, + metrics: dict, + properties: Optional[dict] = None, + units: Optional[dict] = None, +) -> None: + """Emit one EMF record into ``namespace``. Never raises. + + The generic form of the two functions above, for callers whose namespace is not + the prompt-cache one. It lives here rather than being re-implemented per feature + because the parts that are easy to get wrong are not the JSON — they are the + dedicated non-propagating logger and the message-only formatter above. A record + written through the app's normal logger acquires an ``[INFO] name:`` prefix, + CloudWatch Logs silently declines to extract it, and the metric simply never + appears. Nothing errors; there is just no data, which is indistinguishable from + "the thing being measured never happened". + + ``metrics`` maps metric name to numeric value; ``units`` optionally maps the + same names to a CloudWatch unit, defaulting to ``None`` (a bare number). + ``properties`` ride along as queryable log fields and are **not** dimensions — + dimensions multiply metric streams, and every caller here so far wants + fleet-wide aggregates with the detail available in Logs Insights. + """ + try: + units = units or {} + record = { + "_aws": { + "Timestamp": int(time.time() * 1000), + "CloudWatchMetrics": [ + { + "Namespace": namespace, + "Dimensions": [[]], + "Metrics": [ + {"Name": name, "Unit": units.get(name, "None")} + for name in metrics + ], + } + ], + }, + } + record.update({name: value for name, value in metrics.items()}) + for key, value in (properties or {}).items(): + if value is not None: + record[key] = value + _emf_logger.info(json.dumps(record, separators=(",", ":"))) + except Exception as e: # noqa: BLE001 - metrics must never break a caller + logger.debug("EMF emission skipped: %s", e) + + def emit_session_cache_rollup( session_id: str, partial_miss_usd: float, diff --git a/backend/src/apis/shared/rbac/role_constraints.py b/backend/src/apis/shared/rbac/role_constraints.py index 39ffcb262..2fcf754cb 100644 --- a/backend/src/apis/shared/rbac/role_constraints.py +++ b/backend/src/apis/shared/rbac/role_constraints.py @@ -9,9 +9,18 @@ protected role's ``jwt_role_mappings``, which would silently grant that role's permissions to every authenticated user. 2. Storing free-form text in ``jwt_role_mappings`` that doesn't look like - a real group identifier (whitespace, HTML, control characters, etc.) — - typically indicates a mis-typed or attacker-shaped payload rather than - a legitimate IdP claim value. + a real group identifier (HTML, control characters, invisible Unicode, + etc.) — typically indicates a mis-typed or attacker-shaped payload + rather than a legitimate IdP claim value. + + Internal spaces *are* legitimate: Entra security groups are routinely + named as display names (``"PSEmeriti Entra Sync"``), and we do not + control those names. What stays rejected is anything that cannot round + trip through the claim: a comma (the ``custom:roles`` claim is + comma-separated, so a comma-bearing group name is unrepresentable), + and leading/trailing whitespace (both claim parsers ``.strip()`` every + entry, so an edge-padded mapping could never match an incoming claim — + it would look granted and grant nothing). 3. Granting an admin scope that must never be delegated, or one that does not exist — see :func:`validate_admin_scopes`. """ @@ -32,6 +41,13 @@ # (``"default"`` is the universal group every authenticated user holds in # the standard Cognito setup; the rest are common synonyms or wildcards # we never want to accept on a protected role). +# +# Compared after :func:`_normalize_for_forbidden_check`, which folds case +# and treats space/hyphen/underscore as the same separator. That matters +# now that spaces are accepted: ``"Authenticated Users"`` and +# ``"All Users"`` are real Entra/AD display names for exactly the +# populations this set exists to keep off a protected role, and before +# normalization only the hyphenated spelling was caught. _FORBIDDEN_PROTECTED_MAPPINGS: frozenset[str] = frozenset( { "default", @@ -45,14 +61,46 @@ "all", "any", "public", + "all-users", + "domain-users", } ) # Conservative pattern for a JWT group identifier: alphanumerics, underscore, -# and hyphen, 2–64 characters. Real-world IdP groups (Entra, Okta, Cognito -# Cognito groups, custom claims) all conform to this shape; values that -# don't are almost certainly malformed. -_JWT_MAPPING_PATTERN = re.compile(r"^[A-Za-z0-9_-]{2,64}$") +# and hyphen, plus *single internal* ASCII spaces. Entra security groups are +# commonly named as display names ("PSEmeriti Entra Sync"), and the tenant +# owner — not this platform — chooses those names, so a space has to be +# accepted verbatim. +# +# The alternation (rather than adding " " to the character class) is what +# forbids leading, trailing, and doubled spaces. Both shapes end the same +# way -- a mapping that looks granted and matches nothing. A stored value +# has to match the incoming claim byte for byte, and the claim parsers +# ``.strip()`` every entry, so an edge-padded value can never match; a +# doubled space is indistinguishable from a single one in the admin's +# comma-separated field and is far likelier to be a typo than a real group. +# +# Deliberately still rejected: commas (the delimiter in a ``custom:roles`` +# claim and in the admin form, so a comma-bearing name is unrepresentable), +# and every non-space whitespace or invisible character — tabs, newlines, +# NBSP (U+00A0), zero-width space (U+200B). JavaScript ``trim()`` does not +# strip U+200B, so a name pasted out of Entra or Teams can carry one into +# the payload; it must fail loudly rather than be stored as a mapping that +# looks granted and matches nothing. +_JWT_MAPPING_PATTERN = re.compile(r"^[A-Za-z0-9_-]+(?: [A-Za-z0-9_-]+)*$") + +# Length bounds, checked separately from the pattern so an out-of-range value +# gets an error that says so. +_JWT_MAPPING_MIN_LENGTH = 2 +_JWT_MAPPING_MAX_LENGTH = 64 + +# Characters reproduced literally by :func:`_describe_mapping`. Everything +# else — control characters, ``<``, NBSP, zero-width characters, any +# non-ASCII — is escaped to a visible ```` token before it reaches an +# error body or a log line. The set is the accepted charset plus the three +# punctuation marks an admin plausibly typed by mistake (``,``, ``.``, +# ``/``); none of them carry an injection shape. +_ECHO_SAFE_CHAR = re.compile(r"[A-Za-z0-9 _,./-]") # Shape of an admin scope id (``admin.tools``). Deliberately length-bounded: # this pattern is the gate that decides whether a rejected value is safe to @@ -77,13 +125,63 @@ class RoleMutationForbidden(Exception): """ +def _normalize_for_forbidden_check(value: str) -> str: + """Fold a mapping to the form compared against the forbidden set. + + Case-insensitive, and space/hyphen/underscore are treated as the same + separator, so ``"Authenticated Users"``, ``"authenticated_users"`` and + ``"authenticated-users"`` all collapse to one entry. + """ + return re.sub(r"[ _-]+", "-", value.strip().lower()) + + def _is_forbidden_for_protected(value: str) -> bool: - return value.strip().lower() in _FORBIDDEN_PROTECTED_MAPPINGS + normalized = _normalize_for_forbidden_check(value) + return ( + value.strip().lower() in _FORBIDDEN_PROTECTED_MAPPINGS + or normalized in _FORBIDDEN_PROTECTED_MAPPINGS + ) + + +def _describe_mapping(entry: str) -> str: + """Render ``entry`` for an error message and a log line. + + Two jobs. First, make the invisible visible: every character outside + :data:`_ECHO_SAFE_CHAR` is replaced by a ```` token, so an admin + who pasted a group name out of Entra or Teams can see *that* there is a + zero-width space or an NBSP in it — echoing the raw value would render + identically to a correct one and tell them nothing. + + Second, bound the output. The result is length-capped, so a malformed + payload cannot ride an arbitrarily long string into the 400 body or the + ``Role update failed:`` log line. + """ + truncated = entry[:_JWT_MAPPING_MAX_LENGTH] + rendered = "".join( + ch if _ECHO_SAFE_CHAR.fullmatch(ch) else f"" + for ch in truncated + ) + ellipsis = "..." if len(entry) > _JWT_MAPPING_MAX_LENGTH else "" + return f"'{rendered}{ellipsis}'" def validate_jwt_role_mappings(role_id: str, mappings: Iterable[str]) -> None: """Validate ``jwt_role_mappings`` content for ``role_id``. + Entries may contain single internal spaces (Entra security groups are + routinely named as display names) but must otherwise be alphanumerics, + underscore and hyphen, 2-64 characters. See :data:`_JWT_MAPPING_PATTERN` + for why each excluded shape is excluded. + + Like :func:`validate_admin_scopes`, the errors here name the offending + entry. The same reasoning applies verbatim: only a ``system_admin`` can + reach this code path (the roles admin is non-delegable), so there is no + untrusted caller to withhold detail from — and a bare "Invalid role + configuration." on a comma-separated field turns a one-character typo + into a CloudWatch expedition. The value is echoed only after + :func:`_describe_mapping` has bounded its length and escaped anything + outside a conservative charset. + Args: role_id: The role being mutated. mappings: The proposed list of JWT group names. @@ -97,13 +195,54 @@ def validate_jwt_role_mappings(role_id: str, mappings: Iterable[str]) -> None: return for entry in mappings: - if not isinstance(entry, str) or not _JWT_MAPPING_PATTERN.fullmatch(entry): - raise RoleConstraintError("Invalid role configuration.") + if not isinstance(entry, str): + raise RoleConstraintError("Each JWT role mapping must be a string.") + + described = _describe_mapping(entry) + + if not ( + _JWT_MAPPING_MIN_LENGTH <= len(entry) <= _JWT_MAPPING_MAX_LENGTH + ): + raise RoleConstraintError( + f"JWT role mapping {described} must be between " + f"{_JWT_MAPPING_MIN_LENGTH} and {_JWT_MAPPING_MAX_LENGTH} " + "characters." + ) + + if "," in entry: + raise RoleConstraintError( + f"JWT role mapping {described} must not contain a comma. " + "Commas separate one mapping from the next." + ) + + if entry != entry.strip(" "): + raise RoleConstraintError( + f"JWT role mapping {described} must not start or end with a " + "space. The claim is read with leading and trailing " + "whitespace stripped, so a padded mapping would never match." + ) + + if " " in entry: + raise RoleConstraintError( + f"JWT role mapping {described} must not contain consecutive " + "spaces." + ) + + if not _JWT_MAPPING_PATTERN.fullmatch(entry): + raise RoleConstraintError( + f"JWT role mapping {described} contains unsupported " + "characters. Allowed: letters, digits, underscore, hyphen, " + "and single spaces between words." + ) if role_id in PROTECTED_ROLE_IDS: for entry in mappings: if _is_forbidden_for_protected(entry): - raise RoleConstraintError("Invalid role configuration.") + raise RoleConstraintError( + f"JWT role mapping {_describe_mapping(entry)} is held by " + "every authenticated user and cannot be mapped to the " + f"protected role '{role_id}'." + ) def is_protected_role(role_id: str) -> bool: diff --git a/backend/tests/apis/app_api/shares/test_share_s3_offload.py b/backend/tests/apis/app_api/shares/test_share_s3_offload.py index 13a8efe5d..e51d7617b 100644 --- a/backend/tests/apis/app_api/shares/test_share_s3_offload.py +++ b/backend/tests/apis/app_api/shares/test_share_s3_offload.py @@ -136,6 +136,14 @@ async def test_large_conversation_succeeds(self, service): @pytest.mark.asyncio async def test_storage_unavailable_raises(self, monkeypatch): monkeypatch.setenv("SHARED_CONVERSATIONS_TABLE_NAME", "shares-table") + # `bucket_name=None` means "no bucket configured", but the store falls back + # to this variable when the argument is None — so the test only asserted + # what it meant to while the variable happened to be absent from the + # environment. Any developer with a populated `backend/src/.env` (which + # `load_dotenv(override=True)` reads) gave the store a real bucket, and the + # assertion became a live S3 HeadObject against it. Deleted explicitly so + # "unavailable" is a property of the test rather than of the machine. + monkeypatch.delenv("SHARED_CONVERSATIONS_BUCKET_NAME", raising=False) with patch("boto3.resource"): svc = ShareService(snapshot_store=ShareSnapshotStore(bucket_name=None)) svc._table = MagicMock() 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/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/lambdas/test_kb_ingestion_consumer.py b/backend/tests/lambdas/test_kb_ingestion_consumer.py new file mode 100644 index 000000000..70d20246b --- /dev/null +++ b/backend/tests/lambdas/test_kb_ingestion_consumer.py @@ -0,0 +1,332 @@ +"""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 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 _FakeBackend: + """Records ingest calls; reports the document retrievable immediately.""" + + def __init__(self): + self.ingested = [] + + 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_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 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..879651472 --- /dev/null +++ b/backend/tests/lambdas/test_kb_migration_worker.py @@ -0,0 +1,1022 @@ +""" +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_verify_requires_a_canary_retrieval_to_return_something(self): + """Requirement 15.7. Bedrock reporting a document INDEXED precedes it being + retrievable by 0.75-1.03 s, and a knowledge base can hold documents while + returning nothing, so "we ingested everything" and "retrieval works" are + separate claims.""" + backend = StubBackend(chunks=[]) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=[_doc("d1")] + ): + # Matched on "not queryable", not on "canary": both failure messages + # mention the canary, so the looser pattern passed even with the + # empty-result check removed — the *other* check raised and the test + # could not tell the difference. + with pytest.raises(worker.VerificationFailed, match="not queryable"): + 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_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 "