diff --git a/.env.example b/.env.example index bc909f3..e3b7b7b 100644 --- a/.env.example +++ b/.env.example @@ -1,4 +1,5 @@ # Synthetic placeholders only. Never commit a real database URL or AWS resource name. +AUTH_MODE=demo DATABASE_URL=postgresql://USER:PASSWORD@HOST:26257/defaultdb?sslmode=verify-full EMBEDDING_PROVIDER=deterministic BEDROCK_MODEL_ID=amazon.titan-embed-text-v2:0 diff --git a/README.md b/README.md index cac96ea..3388d93 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,20 @@ To run only the SAM validation/build gate: .\scripts\verify.ps1 -WithSam ``` +Before a separately approved live synthetic E2E, the repository provides two fail-fast read-only +guards and one bounded runner: + +```powershell +.\scripts\bedrock-quota-preflight.ps1 -Region us-west-2 -Profile +.\scripts\aws-cost-preflight.ps1 -BudgetName -Profile +.\scripts\live-bedrock-e2e.ps1 -ApiUrl -BudgetName -Profile +``` + +The live runner accepts only a short-lived synthetic Cognito ID token from the process environment, +performs exactly two application actions that require embeddings, prints no token or memory content, +and stops before calls when quota or cost checks fail. Preparing the runner does not authorize or +perform a live invocation. + ## CockroachDB integration `schema.sql` defines: @@ -103,12 +117,18 @@ the application automatically. Current CockroachDB documentation confirms `VECTOR`, `<=>` cosine distance, and prefix columns for vector indexes: [Vector Indexes](https://www.cockroachlabs.com/docs/stable/vector-indexes). +The project also used the CockroachDB Cloud Managed MCP Server as a separate read-only development +and operations inspection path. See +[docs/COCKROACHDB_TOOLS.md](docs/COCKROACHDB_TOOLS.md) for the exact operations, evidence, and +runtime boundary. + ## AWS SAM -`template.yaml` includes explicit root/proxy routes, throttling, a retained encrypted/versioned S3 -bucket, exact Secrets Manager access, and a foundation-model-scoped Bedrock permission. No CORS is -enabled; the bundled UI is same-origin. Production authentication is deliberately left to the -deploying environment. +`template.yaml` includes a Cognito user pool/client/domain, OAuth 2.0 Authorization Code + PKCE for +the browser demo, an API Gateway JWT authorizer, explicit public demo/health/auth-config routes, +protected `/v1/*` routes, throttling, a retained encrypted/versioned S3 bucket, exact Secrets +Manager access, and a foundation-model-scoped Bedrock permission. No API CORS is enabled; protected +requests from the bundled UI are same-origin. See [docs/AWS_DEPLOYMENT.md](docs/AWS_DEPLOYMENT.md) for validation, parameters, deployment steps, cost/security gates, and rollback notes. **No AWS deployment is performed by this repository.** @@ -118,6 +138,7 @@ cost/security gates, and rollback notes. **No AWS deployment is performed by thi | Method | Path | Purpose | |---|---|---| | `GET` | `/healthz` | Runtime mode; never returns secrets | +| `GET` | `/auth/config` | Public OAuth client/domain/callback settings; never returns secrets | | `POST` | `/v1/memories/propose` | Evaluate and optionally persist a proposal | | `POST` | `/v1/memories/{id}/approve` | Activate a pending memory | | `POST` | `/v1/memories/search` | Scoped vector retrieval with `max_sensitivity` | @@ -126,9 +147,11 @@ cost/security gates, and rollback notes. **No AWS deployment is performed by thi ## Security boundary -This MVP demonstrates memory governance; it does not provide end-user authentication or -authorization. The demo accepts tenant, actor, and clearance from the request. A production gateway -must derive these fields from verified identity claims and restrict approval/export operations. +AWS mode derives tenant, actor, groups, and clearance from API Gateway-verified Cognito JWT claims. +The required claims are `sub`, `custom:tenant_id`, and optional `custom:clearance` (defaults to +`public`). Approval requires `memoryguard-reviewer` or `memoryguard-admin`; evidence export requires +`memoryguard-auditor` or `memoryguard-admin`. Client-supplied identity fields remain only for local +`AUTH_MODE=demo` compatibility and are ignored in Cognito mode. Never commit a real `.env`, connection string, credential, or customer dataset. See [SECURITY.md](SECURITY.md) for the threat model and production gaps. @@ -136,9 +159,10 @@ Never commit a real `.env`, connection string, credential, or customer dataset. ## Hackathon material - [Devpost draft](docs/DEVPOST.md) -- [90-second demo script](docs/DEMO_SCRIPT.md) +- [75-second demo script](docs/DEMO_SCRIPT.md) - [Architecture and trust boundaries](docs/ARCHITECTURE.md) - [AWS/SAM runbook](docs/AWS_DEPLOYMENT.md) - [Local verification record](docs/VERIFICATION.md) +- [CockroachDB tool usage and MCP proof](docs/COCKROACHDB_TOOLS.md) License: MIT. diff --git a/SECURITY.md b/SECURITY.md index 2bed666..b6ccee9 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -2,8 +2,8 @@ ## Scope -OME MemoryGuard is an MVP policy and persistence layer. It reduces persistent-memory risk but is not -a complete authorization system, DLP product, secrets scanner, or compliance control. +OME MemoryGuard is an MVP policy and persistence layer. It reduces persistent-memory risk and adds +Cognito JWT/group gates, but is not a complete IAM, DLP, secrets-scanning, or compliance product. ## Threats addressed @@ -23,11 +23,13 @@ a complete authorization system, DLP product, secrets scanner, or compliance con - Secrets Manager errors exclude connection-string contents; - S3 evidence keys hash tenant identifiers and use encryption/checksum parameters; - SAM IAM scopes secret, bucket prefix, and foundation model. +- protected AWS routes require API Gateway-verified Cognito JWTs; +- tenant/actor/clearance are server-derived in Cognito mode, and privileged actions require groups. ## Production gaps -- add JWT/OIDC authentication and server-derived tenant, actor, role, and clearance; -- restrict approval and evidence export to privileged identities; +- add a production login UI, account lifecycle/recovery controls, and fine-grained scopes; +- verify group administration and token-revocation behavior in live integration tests; - add policy versioning, revocation, retention, reconciliation, and tamper-evident audit chaining; - add rate limits/WAF controls and structured logs with privacy review; - run live integration, backup/restore, failure-injection, and authorization tests; diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index f16fdfe..5c83744 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -1,29 +1,35 @@ # Architecture and trust boundaries +![OME MemoryGuard architecture](assets/architecture.svg) + +The rendered diagram is optimized for a 3:2 Devpost gallery image. Its editable Mermaid source is +available in [`assets/architecture.mmd`](assets/architecture.mmd). + ## Runtime topology -```text -Client / agent - | - v -API Gateway HTTP API -> AWS Lambda (Mangum + FastAPI) - | - +-------------+-------------+ - | | - v v - Amazon Bedrock Amazon S3 - Titan Embeddings V2 audit metadata bundles - | (AES-256 + checksum) - v - CockroachDB Cloud - - memory_items VECTOR(1024) - - memory_audit_events - - agent_checkpoints - - prefix-filtered cosine ANN index +```mermaid +flowchart LR + A["AI agent or browser"] -->|"Cognito JWT"| B["API Gateway HTTP API"] + B -->|"JWT authorizer"| C["AWS Lambda
FastAPI + Mangum"] + C --> D["Deterministic policy gate"] + D -->|"allowed or review"| E["Amazon Bedrock
Titan Embeddings V2"] + E --> F["CockroachDB Cloud
memory + audit transaction
VECTOR(1024)"] + D -->|"denied metadata only"| F + C -->|"audit metadata export"| G["Amazon S3
AES-256 + checksum"] + H["AWS Secrets Manager"] -->|"database URL at cold start"| C + + subgraph Local["Local deterministic demo"] + I["InMemoryRepository"] + J["Deterministic embeddings"] + end + D -.->|"local mode"| J + J --> I ``` Local mode replaces CockroachDB with `InMemoryRepository` and Bedrock with reproducible deterministic vectors. The health endpoint reports which mode is active without returning configuration values. +The Lambda and Bedrock inference Regions are configured independently: the runtime calls the exact +approved `BEDROCK_REGION`, while CockroachDB and the rest of the dev stack remain in Frankfurt. ## Write path @@ -58,6 +64,13 @@ transaction. If the database event fails after S3 succeeds, reconciliation is re ## Trust boundaries +- API Gateway validates Cognito issuer/audience/signature before protected `/v1/*` routes reach the + application. The app fails closed if the verified claims are absent in `AUTH_MODE=cognito`. +- The browser demo uses Cognito Authorization Code + PKCE. It receives only public OAuth settings + from `/auth/config`, keeps its short-lived ID token in `sessionStorage`, and sends it only to + same-origin `/v1/*` routes. +- Tenant and actor come from `custom:tenant_id` and `sub`; clearance is clamped to the verified + claim. Approval/export require dedicated Cognito groups. - The policy engine is before persistent memory writes, not after retrieval. - The application never logs a database URL and fails closed on malformed Secrets Manager values. - IAM in `template.yaml` is limited to one secret, one S3 object prefix, and one Bedrock foundation @@ -67,13 +80,15 @@ transaction. If the database event fails after S3 succeeds, reconciliation is re ## Known MVP limitations -- There is no built-in authentication/authorization. Tenant, actor, approver, and clearance must be - derived from authenticated claims in production; request fields are only a demo contract. +- Cognito JWT authentication, PKCE browser sign-in, and coarse group authorization are present, but + lifecycle automation, account recovery policy, access-token claim customization, and fine-grained + scopes are not. - Policy pattern matching is a guardrail, not full DLP, malware scanning, or prompt-injection proof. - There is no revoke endpoint, retention worker, checkpoint API, rate-limiter in local mode, or audit signature chain. -- `schema.sql` and the AWS stack are not applied automatically. Live database/AWS behavior requires - a separately approved deployment and integration test. +- `schema.sql` and the AWS stack are not applied automatically by the source tree. They were applied + only to the separately approved, synthetic `ome-memoryguard-dev` environment; production remains + undeployed. - Vector-index feature availability and Bedrock model availability must be verified in the chosen CockroachDB version and AWS Region. diff --git a/docs/AWS_DEPLOYMENT.md b/docs/AWS_DEPLOYMENT.md index e844c68..eb2fa01 100644 --- a/docs/AWS_DEPLOYMENT.md +++ b/docs/AWS_DEPLOYMENT.md @@ -7,6 +7,8 @@ may incur cost. The repository does not deploy automatically. - Python 3.12 and AWS SAM CLI. - An approved AWS account/Region and least-privilege deployment identity. +- Never deploy with the AWS account root identity. Use root only to bootstrap a narrowly scoped + deployer role, then run all stack commands through that role. - A pre-created Secrets Manager secret whose `SecretString` is either the CockroachDB URL or: ```json @@ -14,7 +16,11 @@ may incur cost. The repository does not deploy automatically. ``` - A CockroachDB database where `schema.sql` has been reviewed and applied by an authorized operator. -- Bedrock access to `amazon.titan-embed-text-v2:0` in the selected Region. +- The Lambda package keeps `sslmode=verify-full`. If the URL does not provide `sslrootcert`, the + application uses the packaged `certifi` CA bundle instead of depending on Lambda/libpq trust-store + discovery. An explicit CA path in the URL is preserved. +- Bedrock access to `amazon.titan-embed-text-v2:0` in the approved `BedrockRegion`. The inference + Region may differ from the Lambda Region only after an explicit data-residency review. Do not place the real URL in `samconfig.toml`, terminal history, screenshots, logs, or CloudFormation parameters. Pass only the secret ARN to SAM. @@ -40,23 +46,67 @@ Optional read-only account checks: ```powershell aws sts get-caller-identity aws secretsmanager describe-secret --secret-id --region -aws bedrock get-foundation-model --model-identifier amazon.titan-embed-text-v2:0 --region +aws bedrock get-foundation-model --model-identifier amazon.titan-embed-text-v2:0 --region ``` Do not run `get-secret-value` during routine validation. +Before any paid Bedrock invocation, run the read-only quota preflight for the approved inference +Region. It checks the Titan V2 on-demand requests-per-minute and tokens-per-minute quotas and exits +nonzero when both are not positive: + +```powershell +.\scripts\bedrock-quota-preflight.ps1 -Region us-west-2 +``` + +Several candidate Regions can be checked without invoking a model: + +```powershell +.\scripts\bedrock-quota-preflight.ps1 -Region us-west-2,us-east-1,eu-west-3 +``` + +The script requires an authenticated AWS CLI profile, prints no identity ARN or secret value, and +emits `BEDROCK_QUOTA_READY=true` only for a Region with nonzero RPM and TPM allocations. A passing +quota preflight is necessary but does not itself prove model access or successful inference. + +Confirm that the approved AWS Budget exists, uses USD, is no higher than USD 5, and has not reached +the approved maximum. This call is read-only and prints only the limit and current actual amount: + +```powershell +.\scripts\aws-cost-preflight.ps1 -BudgetName -Profile +``` + +The AWS Budget is an alerting and preflight control, not a hard account spending cap. + +On Windows paths containing non-ASCII characters, package from the isolated ASCII-only staging +copy instead of calling `sam package` directly from the repository: + +```powershell +.\scripts\sam-package.ps1 -S3Bucket -Profile -Region +``` + +The helper excludes `.env`, Git metadata, virtual environments, generated files, bytecode, and +logs. It writes only the generated `.aws-sam/deploy/packaged.yaml` back to the project. + ## Parameters and generated resources | Item | Purpose | |---|---| | `DatabaseUrlSecretArn` | Exact secret readable by the Lambda role | | `BedrockModelId` | Foundation model ID; defaults to Titan Embeddings V2 | +| `BedrockRegion` | Direct regional Bedrock endpoint; defaults to Paris for the isolated dev stack | +| `EnableDevAdminAuth` | Defaults to `false`; set to `true` only for isolated admin-driven smoke tests | +| `DemoCallbackUrl` | Exact HTTPS callback/logout URL registered for the browser demo | +| `MemoryGuardUserPool` | Cognito users with tenant/clearance custom claims | +| `MemoryGuardUserPoolClient` | Public app client; no client secret | +| `MemoryGuardUserPoolDomain` | Cognito prefix domain for OAuth Authorization Code + PKCE | +| `MemoryGuardHttpApi` | Public demo/health routes and JWT-protected `/v1/*` routes | | `MemoryGuardFunction` | Python 3.12 Lambda, 30 s, 512 MB | -| `ServerlessHttpApi` | Root and proxy routes with per-route throttling | +| `MemoryGuardFunctionLogGroup` | Explicit Lambda log group with seven-day dev retention | | `EvidenceBucket` | Private encrypted/versioned S3 bucket retained on deletion | -IAM permits only `GetSecretValue` on the supplied ARN, `InvokeModel` on the selected foundation-model -ARN, and `PutObject` under the generated bucket's `evidence/*` prefix. +IAM permits only `GetSecretValue` on the supplied ARN, `InvokeModel` on the selected model in the +exact `BedrockRegion`, and `PutObject` under the generated bucket's `evidence/*` prefix. ## Deployment (requires explicit approval) @@ -70,23 +120,65 @@ Recommended guided answers: - use a dedicated stack name and approved Region; - supply only `DatabaseUrlSecretArn`, never its value; +- set `BedrockRegion` only to a Region with verified on-demand quota and approved data processing; +- set `DemoCallbackUrl` to the exact deployed demo URL, including its trailing slash; - save configuration only if the file will not capture sensitive values; - require explicit confirmation before CloudFormation changes; - do not disable rollback for a normal deployment. -The default API is public and has no application authorizer. For any real/customer data, add a JWT -authorizer and derive tenant/actor/clearance from verified claims before deployment. +The root demo page, `/healthz`, and `/auth/config` are public. `/v1/*` uses an API Gateway JWT +authorizer backed by the generated Cognito user pool/client. The browser starts Cognito Authorization +Code + PKCE and keeps the returned ID token only in `sessionStorage`; it never handles a Cognito +password or client secret. In `AUTH_MODE=cognito`, the application requires verified claims and +ignores client-supplied tenant/actor/clearance values. Do not create test users with real personal or +customer data. + +For an isolated automated smoke test, deploy with `EnableDevAdminAuth=true`, create a synthetic +user with `MessageAction=SUPPRESS`, and use `ADMIN_USER_PASSWORD_AUTH`. Keep the generated password +and JWT only in process memory, never logs or shell history. Leave the parameter at its default +`false` outside that short-lived dev test. + +## Bounded Bedrock application E2E (requires separate live-call approval) + +After quota and cost preflights pass, place a short-lived synthetic Cognito ID token only in the +current process environment and run: + +```powershell +$env:MEMORYGUARD_E2E_ID_TOKEN = '' +.\scripts\live-bedrock-e2e.ps1 ` + -ApiUrl ` + -BudgetName ` + -BedrockRegion us-west-2 ` + -Profile +Remove-Item Env:MEMORYGUARD_E2E_ID_TOKEN +``` + +The runner uses fixed synthetic content and a unique synthetic namespace. It executes one proposal +and one search (two application actions that require embeddings), then verifies the associated audit +events. It does not create AWS resources, export evidence, print the JWT, print content, or read a +database secret. Provider retries inside the deployed Lambda remain bounded by its runtime +configuration; the USD 5 Budget remains an alert, not a transactional hard cap. ## Post-deployment verification Use synthetic data only: 1. Call `/healthz`; expect `storage=cockroachdb` and `embeddings=bedrock`. -2. Submit one non-sensitive synthetic proposal and verify active + audit rows. -3. Submit one synthetic denied pattern and verify no raw denied row exists. -4. Retrieve within the same tenant/namespace. -5. Export audit metadata and verify S3 encryption/checksum/versioning. -6. Inspect CloudWatch errors and Lambda duration without printing configuration values. +2. Call a protected `/v1/*` route without a token; expect `401`. +3. Open the demo, sign in through Cognito, and verify the callback returns without exposing tokens in + the URL. +4. With an approved synthetic Cognito user, submit one non-sensitive proposal and verify active + + audit rows inside the claim-derived tenant. +5. Attempt a spoofed tenant/client clearance and verify the JWT claims win. +6. Verify reviewer/auditor group gates, then inspect S3 encryption/checksum/versioning. + +The Bedrock client uses bounded adaptive retries for transient throttling. If Bedrock remains +throttled or unavailable after the retry budget, the API returns a sanitized HTTP 503 with +`embedding_provider_unavailable`. The proposal/search fails closed and no memory or audit row is +written. `BedrockRegion` selects a direct regional endpoint, not a cross-Region inference profile. +When it differs from the Lambda Region, request text is processed in that selected Region; use only +approved data and record the residency decision. +7. Inspect CloudWatch errors and Lambda duration without printing token or configuration values. ## Rollback and cleanup diff --git a/docs/COCKROACHDB_TOOLS.md b/docs/COCKROACHDB_TOOLS.md new file mode 100644 index 0000000..719da98 --- /dev/null +++ b/docs/COCKROACHDB_TOOLS.md @@ -0,0 +1,66 @@ +# CockroachDB tool usage + +OME MemoryGuard uses two CockroachDB tools required by the CockroachDB x AWS Hackathon: + +1. CockroachDB Distributed Vector Indexing in the application data path. +2. CockroachDB Cloud Managed MCP Server in the read-only development and operations validation + path. + +This document separates those paths so the submission does not imply that the Managed MCP Server is +a hidden runtime dependency. + +## 1. Distributed Vector Indexing + +CockroachDB is the system of record for governed agent memory. `schema.sql` defines: + +- `memory_items.embedding` as `VECTOR(1024)`; +- `memory_semantic_idx` as a cosine vector index; +- `(tenant_id, namespace)` as index prefix columns; +- `memory_audit_events` and `agent_checkpoints` beside the memory records. + +The repository search query constrains tenant, namespace, lifecycle state, expiry, and sensitivity +before applying cosine-distance (`<=>`) ranking. This keeps security filters in the same database +operation as semantic retrieval. Memory mutations and their audit events share one transaction. + +Public evidence: + +- [`schema.sql`](../schema.sql) +- [`ome_memoryguard/repository.py`](../ome_memoryguard/repository.py) +- [`tests/test_architecture.py`](../tests/test_architecture.py) +- [`tests/test_repository.py`](../tests/test_repository.py) + +## 2. CockroachDB Cloud Managed MCP Server + +During isolated cluster setup and validation, the development/operations agent connected to the +`ome-memoryguard` CockroachDB Cloud cluster through the Managed MCP Server in read-only mode. The +agent used the server's structured operations to: + +- inspect the target cluster with `get_cluster`; +- enumerate databases with `list_databases`; +- enumerate tables with `list_tables`; +- confirm the separate `ome_memoryguard` database and its initial table state. + +The MCP workflow did not expose the application connection string, read local secret files, or +modify database data. Application writes continue to use the narrowly scoped SQL application user; +the MCP connection is an independent operator and judge inspection path. + +### Reproducible read-only judge workflow + +1. In CockroachDB Cloud, open the isolated cluster and choose **Connect** -> **Model Context + Protocol (MCP)**. +2. Copy the generated client configuration and authenticate through CockroachDB OAuth. +3. Select read-only permission when prompted. +4. Ask the connected agent to inspect the cluster, list databases, and list tables in the + `ome_memoryguard` database. +5. Compare the reported structure with [`schema.sql`](../schema.sql). Do not request mutations or + secret values. + +The public repository intentionally contains no OAuth token, cluster credential, database URL, or +raw MCP response containing account identifiers. + +## Verification boundary + +The Managed MCP Server proves agent-operated, read-only control-plane inspection. Distributed Vector +Indexing proves the persistent runtime memory path. Neither proof claims a complete live Bedrock +retrieval while the AWS account's applied Titan embedding quota remains zero; that limitation is +tracked separately in [`VERIFICATION.md`](VERIFICATION.md). diff --git a/docs/DEMO_SCRIPT.md b/docs/DEMO_SCRIPT.md index 256fbf2..1955bea 100644 --- a/docs/DEMO_SCRIPT.md +++ b/docs/DEMO_SCRIPT.md @@ -1,58 +1,75 @@ -# 90-second demo script +# 75-second demo script ## Setup -- Run the app locally in default in-memory mode. -- Open `/` and confirm the runtime badge says `in_memory · deterministic`. -- Keep browser developer tools closed unless a judge asks; the page exposes all proof signals. +- Run the app locally in default in-memory mode and open `/` at 1080p. +- Confirm the runtime badge says `in_memory · deterministic` and the notice says + **Local deterministic demo**. +- Use only the three supplied synthetic scenarios. Do not enter credentials or customer data. +- Keep sanitized JSON collapsed; the decision facts and audit trail carry the story. -## 0–15 seconds — problem +## 0–10 seconds — the problem -“Agents remember generated text for weeks or months. A secret, stale fact, or prompt injection can -become durable context. MemoryGuard turns remembering into a governed transaction.” +Show the opening frame and say: + +> AI agents can remember generated text for months. One secret or prompt injection can become +> durable context. MemoryGuard turns remembering into a governed transaction. Point to **Fail closed**, **Scoped retrieval**, and **Atomic evidence**. -## 15–35 seconds — safe memory +## 10–28 seconds — safe memory + +1. Select **Safe preference** and click **Run policy gate**. +2. Show `active`, risk `15`, `policy_pass`, and **memory + audit**. +3. Click **Retrieve memory**. + +Say: + +> Allowed memory and its audit event share one CockroachDB transaction. Retrieval applies tenant, +> namespace, lifecycle, expiry, and clearance filters before vector ranking. -1. Select **Safe preference**. -2. Click **Run policy gate**. -3. Show `active`, risk score `15`, and `policy_pass`. -4. Click **Retrieve memory**. +## 28–45 seconds — denied injection -Say: “The same tenant and namespace filters used by the CockroachDB vector index are enforced before -ranking. The write and audit event are one transaction.” +1. Select **Secret injection** and click **Run policy gate**. +2. Show `denied`, `possible_secret`, `persistent_prompt_injection`, and **raw content blocked**. +3. Point to the new denied event in the audit trail. -## 35–55 seconds — denied injection +Say: -1. Select **Secret injection**. -2. Click **Run policy gate**. -3. Show `denied`, `possible_secret`, and `persistent_prompt_injection`. -4. Point to the audit timeline. +> The raw input is never persisted. The audit trail receives only reason codes, content length, and +> a SHA-256 fingerprint. -Say: “The denied raw content is never persisted. Only reason codes, length, and a SHA-256 fingerprint -enter the audit trail.” +## 45–66 seconds — human review -## 55–80 seconds — human review and clearance +1. Select **Needs review** and click **Run policy gate**. +2. Show `pending_review`, risk `55`, and click **Approve as human reviewer**. +3. Keep clearance at `confidential`, click **Retrieve memory**, and show the event chain. -1. Select **Needs review**. -2. Click **Run policy gate**; show `pending_review` and risk `55`. -3. Click **Human approve**. -4. Keep retrieval clearance at `confidential` and click **Retrieve memory**. -5. Show `retrieved: 1` and the proposed → approved → retrieved event chain. +Say: -Say: “Approval changes only pending rows. Retrieval still needs an explicit sensitivity ceiling.” +> Sensitive context stays invisible until a human approves it. Approval does not bypass retrieval +> clearance, and every transition remains auditable. -## 80–90 seconds — platform close +## 66–75 seconds — platform close -“CockroachDB supplies transactional state and prefix-filtered vector search. AWS Lambda hosts the API, -Bedrock creates embeddings, Secrets Manager supplies the connection safely, and S3 stores encrypted -audit metadata bundles. The repository separates locally verified proof from deployment claims.” +Show the architecture image and say: -## Expected proof checklist +> AWS provides the protected serverless runtime and embeddings. CockroachDB keeps governed memory, +> vector search, and audit evidence together. This recording uses the deterministic local demo; the +> remaining live Bedrock quota blocker is disclosed in the submission. -- Runtime badge: `in_memory · deterministic` for local demo. -- Safe: `active`, risk `15`, `policy_pass`. -- Denied: HTTP 422, no ID/raw content, reason codes visible. -- Review: `pending_review` → `active` → one retrieval at confidential clearance. +## Expected proof + +- Safe: `active` → `retrieved`, risk `15`, `policy_pass`. +- Denied: HTTP 422, no memory ID or raw content, two reason codes. +- Review: `pending_review` → `active` → `retrieved` at confidential clearance. - Audit: proposed, denied, approved, and retrieved events in reverse chronological order. +- Runtime badge and notice clearly identify local deterministic mode. + +## Recording checklist + +- Record in 1920×1080, 30 fps; target 70–80 seconds. +- Use a clean browser window at 100% zoom and keep the cursor movement deliberate. +- Cut pauses, but do not hide status transitions or the local-mode disclosure. +- End with `docs/assets/architecture.svg`, repository URL, and public demo URL. +- Do not claim complete live Bedrock retrieval until the bounded synthetic E2E passes. diff --git a/docs/DEVPOST.md b/docs/DEVPOST.md index e4e689a..14b9938 100644 --- a/docs/DEVPOST.md +++ b/docs/DEVPOST.md @@ -1,4 +1,20 @@ -# Devpost submission draft +# Devpost submission-ready draft + +## Copy/paste submission fields + +- **Project name:** OME MemoryGuard +- **Tagline:** Governed persistent memory for AI agents on CockroachDB and AWS. +- **One-line value proposition:** Stop poisoned, secret-bearing, or unapproved context before it + becomes durable agent memory. +- **Built with:** CockroachDB Cloud Managed MCP Server, CockroachDB Distributed Vector Indexing, + VECTOR(1024), AWS Lambda, API Gateway, Cognito, Bedrock Titan Text Embeddings V2, Secrets Manager, + S3, AWS SAM, Python, FastAPI, SQLAlchemy. +- **Repository:** https://github.com/Naim-arg/OME-MemoryGuard +- **Public demo:** https://acnvju8rb8.execute-api.eu-central-1.amazonaws.com +- **Team:** Naiman Yelubayev — solo founder +- **Demo disclosure:** The recorded interactive flow uses deterministic local embeddings and + synthetic data. The isolated AWS/CockroachDB stack is verified separately; complete live Bedrock + retrieval is not claimed until AWS activates a non-zero embedding quota. ## Project name @@ -39,35 +55,56 @@ events. `VECTOR(1024)` stores Bedrock Titan embeddings. A cosine vector index pr `tenant_id, namespace` aligns physical search with the isolation predicates in every retrieval query. Transactions keep a memory mutation and its audit evidence together. -For judging, the database can also be inspected independently through CockroachDB tooling without -exposing the application connection string. +The development/operations agent also used the CockroachDB Cloud Managed MCP Server in read-only +mode to inspect the isolated cluster, list databases and tables, and confirm the separate +`ome_memoryguard` database without exposing the application connection string or modifying data. +This is an independent operator and judge inspection path, not a hidden runtime dependency. Full +evidence and the repeatable workflow are documented in +[`docs/COCKROACHDB_TOOLS.md`](COCKROACHDB_TOOLS.md). ## How we use AWS - **AWS Lambda + API Gateway HTTP API:** serverless FastAPI runtime. +- **Amazon Cognito:** JWT authorization and browser Authorization Code + PKCE sign-in. - **Amazon Bedrock:** Titan Text Embeddings V2, normalized 1,024-dimensional vectors. - **AWS Secrets Manager:** database URL loaded at cold start with fail-closed parsing. - **Amazon S3:** encrypted, checksummed audit metadata bundles under hashed tenant partitions. - **AWS SAM:** reproducible infrastructure template with scoped IAM and explicit routes/throttling. -## Demo story (90 seconds) +## Demo story (75 seconds) 1. Store a safe preference; show `active`, risk score, retrieval, and audit event. -2. Submit a fake secret/prompt injection; show `denied` and prove the raw input is absent from the - decision and audit timeline. +2. Submit a fake secret/prompt injection; show `denied`, **raw content blocked**, and prove the raw + input is absent from the decision and audit timeline. 3. Submit confidential context; show it hidden, approve it, retrieve with confidential clearance, and inspect the full event chain. 4. Point to CockroachDB transaction/vector invariants and the AWS runtime badge. See `docs/DEMO_SCRIPT.md` for narration and expected outputs. +## Submission gallery + +1. [Architecture and trust boundaries](assets/architecture.png) — 3:2 gallery render; editable + source: [`architecture.mmd`](assets/architecture.mmd) +2. [Product overview](assets/screenshots/01-overview.png) +3. [Allowed memory and scoped retrieval](assets/screenshots/02-safe-retrieval.png) +4. [Denied prompt injection](assets/screenshots/03-denied-injection.png) +5. [Human review and audit chain](assets/screenshots/04-review-audit.png) + +Suggested captions are maintained with the screenshot assets in +[`assets/screenshots/README.md`](assets/screenshots/README.md). + ## What we are proud of - The security promise is executable: unit/API/fake-engine tests cover denial non-persistence, conditional approval, sensitivity filtering, transaction boundaries, Secrets Manager parsing, and S3 encryption parameters. - The local demo is deterministic and requires no cloud credentials. -- Documentation distinguishes locally verified behavior from deployment-dependent claims. +- The judge-facing decision panel separates outcome, persistence, review, evidence, reason codes, + and risk instead of relying on raw JSON; the sanitized payload remains available on demand. +- The isolated AWS/CockroachDB dev stack is live and JWT-protected. A synthetic proposal persisted, + its audit trail was queried, and its evidence bundle was exported to S3. Bedrock retrieval remains + explicitly incomplete while the applied embedding quota is zero. ## Challenges @@ -83,21 +120,39 @@ vector index instead of being reconstructed across separate stores. ## What's next -- JWT/OIDC authentication with server-derived tenant, actor, role, and clearance. +- Wait for the paid-account quota activation or request an AWS quota review, then rerun the bounded + synthetic Bedrock write/retrieval/audit E2E. At the latest check, the applied Titan V2 RPM quota + in the final `us-west-2` inference Region was still zero. +- Complete the Cognito PKCE callback/sign-out test with an approved synthetic user, then add + automated user lifecycle and account recovery. - Signed/hashed audit chains and export reconciliation. - Revocation, retention workers, checkpoint APIs, and policy versioning. -- Live CockroachDB/AWS deployment validation under a separately approved environment. +- Production hardening and a separate non-development deployment after the synthetic E2E passes. - Retrieval-quality evaluation on representative multilingual memory datasets. ## Verification status -Local unit, API, fake CockroachDB engine, fake S3, compile, UI, and static architecture checks are -part of the repository workflow. Live CockroachDB writes, schema changes, AWS deployment, and paid -Bedrock invocation are intentionally outside local verification and must not be claimed until run. +Local and SAM release gates pass. The isolated Frankfurt AWS stack, Cognito JWT boundary, +CockroachDB schema/connection, public demo/health endpoints, CloudWatch retention, and S3 hardening +were verified live with synthetic data. A proposal, audit query, and evidence export completed. The +final stack targets Bedrock in `us-west-2`, where the latest applied Titan V2 RPM quota remained zero; +therefore a complete Bedrock-backed write/retrieval/audit E2E is not yet claimed. -## Submission placeholders +## Submission links - Repository: https://github.com/Naim-arg/OME-MemoryGuard -- Demo URL: `TODO_AFTER_APPROVED_DEPLOYMENT` -- Demo video: `TODO_AFTER_RECORDING` -- Team members: `TODO_CONFIRM_FOR_SUBMISSION` +- Demo URL: https://acnvju8rb8.execute-api.eu-central-1.amazonaws.com +- Demo video: published on YouTube and embedded in the Devpost submission +- Team members: Naiman Yelubayev — solo founder + +## Submission asset checklist + +| Asset | Status | Required proof | +|---|---|---| +| Repository | Ready | Public source, current MIT license, setup and verification docs | +| Public demo | Ready with disclosed blocker | Public landing/health; protected actions require Cognito | +| 75-second video | Ready and published | Local deterministic flow, three scenarios, audit trail, explicit disclosure | +| Architecture | Ready | Mermaid source plus 3:2 SVG gallery asset | +| Screenshots | Ready | Opening frame, allowed retrieval, denied decision, review/approval audit chain | +| CockroachDB tools | Ready | Distributed Vector Indexing plus read-only Managed MCP validation documented | +| Live Bedrock E2E | Waiting on AWS Support | Quota + USD 5 budget preflight, two bounded synthetic embedding actions | diff --git a/docs/VERIFICATION.md b/docs/VERIFICATION.md index fae3708..6b303e1 100644 --- a/docs/VERIFICATION.md +++ b/docs/VERIFICATION.md @@ -1,14 +1,17 @@ -# Local verification record +# Verification record -Date: 2026-08-04 (Asia/Qyzylorda) +Date: 2026-08-05 (Asia/Qyzylorda) ## Safety boundary - Project root: `ome-memoryguard` only. - A separate Git repository was initialized inside that root. - `.env` was not present or read; `.env.example` was replaced with synthetic placeholders. -- No live CockroachDB connection, schema write, AWS API mutation, deployment, or paid model call was - performed. +- Live work is limited to the isolated `ome-memoryguard-dev` stack in `eu-central-1`, a separate + CockroachDB development database, and synthetic data. +- Secret values, JWTs, and generated passwords were kept out of logs and files. Temporary Cognito + users were deleted after each smoke attempt. +- A monthly USD 5 AWS budget is configured as an alerting control, not a hard service cap. ## Executed checks @@ -17,10 +20,10 @@ Date: 2026-08-04 (Asia/Qyzylorda) | Initial `python -m unittest discover -s tests -v` | PASS, 7 tests | | Embedding provider fail-closed RED/GREEN test | PASS; unknown provider rejected without cloud calls | | API request-boundary RED/GREEN test | PASS; six invalid search/audit cases rejected with 422 | -| Final `python -m unittest discover -s tests -v` | PASS, 34 tests | +| Final `python -m unittest discover -s tests -v` | PASS, 51 tests | | `python -m compileall -q ome_memoryguard tests` | PASS | | `.venv\Scripts\python.exe -m pip check` | PASS, no broken requirements | -| `.venv\Scripts\python.exe -m pytest -q` | PASS, 34 tests; one upstream Starlette deprecation warning | +| `.venv\Scripts\python.exe -m pytest -q` | PASS, 51 tests; one upstream Starlette deprecation warning | | `.venv\Scripts\ruff.exe check .` | PASS, Ruff 0.16.1 | | `.venv\Scripts\cfn-lint.exe template.yaml` | PASS, cfn-lint 1.53.3 | | `scripts\verify.ps1 -WithGitleaks -WithSam` | PASS for all local release gates | @@ -43,13 +46,77 @@ Date: 2026-08-04 (Asia/Qyzylorda) The temporary local server was stopped and port `8765` was verified free after the smoke test. +## Live development verification + +| Check | Result | +|---|---| +| CockroachDB development database | PASS; separate database/user created and `schema.sql` applied | +| Lambda database driver/TLS path | PASS live; psycopg v3 used the packaged certifi CA with `verify-full`, and the protected audit query reached the separate CockroachDB dev database | +| AWS SAM stack | PASS; `ome-memoryguard-dev` is `UPDATE_COMPLETE` in Frankfurt | +| Public health/demo | PASS; both return HTTP 200 | +| Protected API without JWT | PASS; returns HTTP 401 | +| Protected API with synthetic JWT | PASS; issuer, audience, token type, tenant, and clearance claims matched; audit query returned HTTP 200 with an `items` response | +| Cognito test auth | PASS for isolated smoke setup; the synthetic user was deleted and temporary admin auth was disabled afterward | +| Browser OAuth + PKCE deployment | PASS; public config is secret-free, OIDC discovery/JWKS return HTTP 200, and Cognito managed login is reachable | +| CloudWatch logs | PASS; explicit log group with seven-day retention | +| Evidence S3 bucket | PASS; all public-access blocks enabled, AES256 encryption, versioning enabled | +| Bedrock-backed proposal | PARTIAL PASS live; a synthetic proposal was embedded and persisted during regional testing, but the final configured inference Region is `us-west-2` | +| Bedrock-backed search | BLOCKED live; direct and bounded retry calls returned `ThrottlingException` while the applied Titan V2 RPM quota remained zero | +| Safe failure behavior | PASS live and locally; exhausted provider retries return sanitized HTTP 503; local tests prove failed proposals/searches do not persist | +| Regional quota preflight | BLOCKED pending quota activation; the AWS account is now on the paid plan, but the latest `us-west-2` Titan V2 applied RPM quota remained 0. Earlier regional success did not provide a stable allocation for the final stack | + +The public development demo is available at +https://acnvju8rb8.execute-api.eu-central-1.amazonaws.com. Its protected actions require a Cognito +JWT; the page and health endpoint remain public for review. + ## Not executed | Check | Reason | |---|---| -| Live CockroachDB integration | live writes/schema changes forbidden for this task | -| AWS Lambda/Bedrock/S3 integration | deployment/cloud writes and costs forbidden for this task | +| Successful Bedrock-backed write/retrieval/audit E2E | Proposal persistence, audit query, and S3 evidence export passed with synthetic data, but Bedrock-backed retrieval remains blocked until the paid account receives a nonzero applied embedding quota | +| Production deployment and real customer data | Outside the isolated synthetic dev boundary | +| End-to-end browser OAuth callback/sign-out | Interactive browser callback/sign-out was not exercised; the non-interactive synthetic JWT path was verified instead | `tests/test_architecture.py` provides local static assertions for the SAM route/IAM invariants and CockroachDB vector-index/query alignment. SAM validation/build now covers local packaging, but neither local check replaces live authorization, regional service availability, or integration testing. + +## Bedrock quota recheck — 2026-08-06 + +The AWS Console was inspected without invoking a model or changing cloud resources: + +- `ome-memoryguard-dev` remained `UPDATE_COMPLETE` in `eu-central-1`. +- The USD 5 monthly budget was healthy, with USD 0.00 shown as used. +- Titan Text Embeddings V2 on-demand RPM and TPM applied quotas were both 0 in `us-west-2`, + `us-east-1`, and `eu-west-3` (default values shown by AWS were 6,000 RPM and 300,000 TPM). +- No Bedrock request was attempted because the quota preflight could not pass. + +`scripts/bedrock-quota-preflight.ps1` now provides the same fail-fast read-only check through an +authenticated AWS CLI session. Live Bedrock E2E remains blocked until at least one approved Region +reports positive applied RPM and TPM quotas. + +## MG-5 local preparation — 2026-08-06 + +- Added a read-only AWS Budget preflight that rejects a non-USD budget, a configured limit above + USD 5, or actual spend at/above USD 5. It does not create or update budgets. +- Added a future live E2E runner with fixed synthetic content. It runs quota and cost preflights + before one proposal and one search, then verifies proposal/retrieval audit events. +- The runner receives a short-lived synthetic Cognito ID token only from the process environment and + does not print it, print memory content, export evidence, or read database secrets. +- Improved the local decision UI to show persistence, review, evidence, reason codes, and risk as + separate judge-facing signals; sanitized JSON is collapsed by default. +- No live AWS call, CockroachDB write, deployment, or paid action was executed for this MG-5 update. + +## CockroachDB Managed MCP evidence - 2026-08-08 + +- The development/operations agent connected to the isolated `ome-memoryguard` CockroachDB Cloud + cluster through the Managed MCP Server in read-only mode. +- Structured `get_cluster`, `list_databases`, and `list_tables` operations confirmed the target + cluster, the separate `ome_memoryguard` database, and its initial table state. +- No mutation operation was requested. No application connection string, OAuth token, database + credential, or raw account identifier is stored in this repository. +- The MCP path is documented as an independent operator and judge inspection path. The application + runtime continues to use the scoped SQL user and Distributed Vector Indexing described in + `schema.sql`. +- Reproduction steps and public code evidence are in + [`COCKROACHDB_TOOLS.md`](COCKROACHDB_TOOLS.md). diff --git a/docs/assets/architecture.mmd b/docs/assets/architecture.mmd new file mode 100644 index 0000000..e397519 --- /dev/null +++ b/docs/assets/architecture.mmd @@ -0,0 +1,14 @@ +flowchart LR + actor["AI agent or judge"] -->|"proposal + provenance"| api["API Gateway
Cognito JWT"] + api --> guard["OME MemoryGuard
policy + risk + approval"] + guard -->|"allowed / pending"| bedrock["Amazon Bedrock
Titan Embeddings V2"] + bedrock --> crdb["CockroachDB Cloud
memory + audit transaction
VECTOR(1024)"] + guard -->|"denied metadata only"| crdb + guard -->|"audit metadata"| s3["Amazon S3
encrypted evidence bundle"] + secrets["AWS Secrets Manager"] -->|"database URL at cold start"| guard + crdb -->|"tenant + namespace + clearance"| actor + + subgraph local["Local deterministic demo — no cloud credentials"] + local_embed["Deterministic embeddings"] --> local_repo["InMemoryRepository"] + end + guard -.->|"local mode"| local_embed diff --git a/docs/assets/architecture.png b/docs/assets/architecture.png new file mode 100644 index 0000000..c9d3baf Binary files /dev/null and b/docs/assets/architecture.png differ diff --git a/docs/assets/architecture.svg b/docs/assets/architecture.svg new file mode 100644 index 0000000..c889f2a --- /dev/null +++ b/docs/assets/architecture.svg @@ -0,0 +1,81 @@ + + OME MemoryGuard architecture + Policy-gated persistent memory for AI agents using AWS and CockroachDB, with a deterministic local demo path. + + + + + + + + + + + + + + + + + + + + GOVERNED AGENTIC MEMORY + OME MemoryGuard + Every memory proposal is evaluated before it can influence a future action. + + + + AI agent + or judge + PROPOSE / RETRIEVE + + + POLICY BOUNDARY + Risk · Policy · Approval + active · pending_review · denied + + FAIL CLOSED + EXPLAINABLE + + + Amazon Bedrock + Titan Text Embeddings V2 · 1024 dimensions + + + CockroachDB Cloud + memory + audit event in one transaction + TENANT · NAMESPACE · CLEARANCE + + + Amazon S3 + encrypted evidence + + + Secrets Manager + cold-start secret + + + + Cognito JWT + + allowed + + + + DENIED METADATA + + + + + + LOCAL DETERMINISTIC DEMO + Deterministic embeddings → InMemoryRepository + Same policy states and retrieval filters · no cloud credentials · synthetic data only + + + + AWS hosts and protects the runtime. CockroachDB stores governed memory and its audit evidence. + COCKROACHDB × AWS + diff --git a/docs/assets/screenshots/01-overview.png b/docs/assets/screenshots/01-overview.png new file mode 100644 index 0000000..f332760 Binary files /dev/null and b/docs/assets/screenshots/01-overview.png differ diff --git a/docs/assets/screenshots/02-safe-retrieval.png b/docs/assets/screenshots/02-safe-retrieval.png new file mode 100644 index 0000000..228ba53 Binary files /dev/null and b/docs/assets/screenshots/02-safe-retrieval.png differ diff --git a/docs/assets/screenshots/03-denied-injection.png b/docs/assets/screenshots/03-denied-injection.png new file mode 100644 index 0000000..a0d72f9 Binary files /dev/null and b/docs/assets/screenshots/03-denied-injection.png differ diff --git a/docs/assets/screenshots/04-review-audit.png b/docs/assets/screenshots/04-review-audit.png new file mode 100644 index 0000000..42717fa Binary files /dev/null and b/docs/assets/screenshots/04-review-audit.png differ diff --git a/docs/assets/screenshots/README.md b/docs/assets/screenshots/README.md new file mode 100644 index 0000000..75730c5 --- /dev/null +++ b/docs/assets/screenshots/README.md @@ -0,0 +1,12 @@ +# Submission screenshots + +Use these images in this order on Devpost: + +1. `01-overview.png` — opening frame; product promise, runtime mode, and three safety invariants. +2. `02-safe-retrieval.png` — an allowed memory retrieved inside the tenant/namespace boundary. +3. `03-denied-injection.png` — raw content blocked; only reason codes and a fingerprint are audited. +4. `04-review-audit.png` — confidential proposal after human approval, scoped retrieval, and audit chain. + +All screenshots are produced from the deterministic local demo with synthetic data. They are product +evidence, not proof of a live Bedrock invocation. The live AWS/CockroachDB verification scope and the +remaining Bedrock quota blocker are stated in `docs/DEVPOST.md`. diff --git a/ome_memoryguard/api.py b/ome_memoryguard/api.py index 055f70e..2b2979a 100644 --- a/ome_memoryguard/api.py +++ b/ome_memoryguard/api.py @@ -10,9 +10,15 @@ from fastapi.responses import HTMLResponse from pydantic import BaseModel, Field +from .auth import ( + RequestIdentity, + effective_sensitivity, + get_request_identity, + require_any_group, +) from .demo import DEMO_HTML from .domain import MemoryProposal, MemoryStatus, Sensitivity -from .embeddings import embedding_provider_from_env +from .embeddings import EmbeddingUnavailableError, embedding_provider_from_env from .evidence import EvidenceStore, S3EvidenceStore from .policy import MemoryPolicy from .repository import CockroachRepository, InMemoryRepository @@ -63,6 +69,7 @@ def get_evidence_store() -> EvidenceStore | None: ServiceDependency = Annotated[MemoryService, Depends(get_service)] EvidenceDependency = Annotated[EvidenceStore | None, Depends(get_evidence_store)] +IdentityDependency = Annotated[RequestIdentity | None, Depends(get_request_identity)] app = FastAPI(title="OME MemoryGuard", version="0.2.0") @@ -83,10 +90,34 @@ def health() -> dict[str, str]: } +@app.get("/auth/config") +def auth_config() -> dict[str, str | bool]: + """Return only the public OAuth settings required by the browser demo.""" + domain = os.getenv("COGNITO_DOMAIN", "").strip().rstrip("/") + client_id = os.getenv("COGNITO_CLIENT_ID", "").strip() + redirect_uri = os.getenv("COGNITO_CALLBACK_URL", "").strip() + if not (domain and client_id and redirect_uri): + return {"enabled": False} + return { + "enabled": True, + "domain": domain, + "client_id": client_id, + "redirect_uri": redirect_uri, + } + + @app.post("/v1/memories/propose") -def propose(payload: ProposeRequest, service: ServiceDependency) -> dict: +def propose( + payload: ProposeRequest, + service: ServiceDependency, + identity: IdentityDependency, +) -> dict: + proposal = payload.model_dump() + if identity is not None: + proposal["tenant_id"] = identity.tenant_id + proposal["agent_id"] = identity.actor_id try: - record = service.propose(MemoryProposal(**payload.model_dump())) + record = service.propose(MemoryProposal(**proposal)) except MemoryDeniedError as error: raise HTTPException( status_code=422, @@ -96,6 +127,14 @@ def propose(payload: ProposeRequest, service: ServiceDependency) -> dict: "reason_codes": error.decision.reason_codes, }, ) from error + except EmbeddingUnavailableError as error: + raise HTTPException( + status_code=503, + detail={ + "status": "unavailable", + "reason_code": "embedding_provider_unavailable", + }, + ) from error return { "id": str(record.id), "status": record.decision.status.value, @@ -105,24 +144,47 @@ def propose(payload: ProposeRequest, service: ServiceDependency) -> dict: @app.post("/v1/memories/{memory_id}/approve") -def approve(memory_id: UUID, payload: ApproveRequest, service: ServiceDependency) -> dict: +def approve( + memory_id: UUID, + payload: ApproveRequest, + service: ServiceDependency, + identity: IdentityDependency, +) -> dict: + require_any_group(identity, {"memoryguard-reviewer", "memoryguard-admin"}) + tenant_id = identity.tenant_id if identity is not None else payload.tenant_id + approver_id = identity.actor_id if identity is not None else payload.approver_id try: - record = service.approve(payload.tenant_id, memory_id, payload.approver_id) + record = service.approve(tenant_id, memory_id, approver_id) except (KeyError, LookupError) as error: raise HTTPException(status_code=404, detail="memory not found or not pending") from error return {"id": str(record.id), "status": record.decision.status.value} @app.post("/v1/memories/search") -def search(payload: SearchRequest, service: ServiceDependency) -> dict: - rows = service.search( - payload.tenant_id, - payload.namespace, - payload.query, - payload.actor_id, - payload.limit, - payload.max_sensitivity, - ) +def search( + payload: SearchRequest, + service: ServiceDependency, + identity: IdentityDependency, +) -> dict: + tenant_id = identity.tenant_id if identity is not None else payload.tenant_id + actor_id = identity.actor_id if identity is not None else payload.actor_id + try: + rows = service.search( + tenant_id, + payload.namespace, + payload.query, + actor_id, + payload.limit, + effective_sensitivity(payload.max_sensitivity, identity), + ) + except EmbeddingUnavailableError as error: + raise HTTPException( + status_code=503, + detail={ + "status": "unavailable", + "reason_code": "embedding_provider_unavailable", + }, + ) from error return { "items": [ { @@ -140,8 +202,11 @@ def search(payload: SearchRequest, service: ServiceDependency) -> dict: def audit_events( tenant_id: Annotated[str, Query(min_length=1, max_length=100)], service: ServiceDependency, + identity: IdentityDependency, limit: Annotated[int, Query(ge=1, le=100)] = 50, ) -> dict: + if identity is not None: + tenant_id = identity.tenant_id rows = service.audit_events(tenant_id, limit) return { "items": [ @@ -165,14 +230,18 @@ def export_evidence( payload: EvidenceExportRequest, service: ServiceDependency, store: EvidenceDependency, + identity: IdentityDependency, ) -> dict: + require_any_group(identity, {"memoryguard-auditor", "memoryguard-admin"}) if store is None: raise HTTPException(status_code=503, detail="evidence store is not configured") - rows = service.audit_events(payload.tenant_id, payload.limit) + tenant_id = identity.tenant_id if identity is not None else payload.tenant_id + actor_id = identity.actor_id if identity is not None else payload.actor_id + rows = service.audit_events(tenant_id, payload.limit) artifact = { "schema_version": "ome-memoryguard-evidence/v1", "generated_at": datetime.now(UTC).isoformat(), - "tenant_id": payload.tenant_id, + "tenant_id": tenant_id, "events": [ { "id": str(row.id), @@ -187,6 +256,6 @@ def export_evidence( for row in rows ], } - location = store.put(payload.tenant_id, artifact) - service.record_evidence_export(payload.tenant_id, payload.actor_id, len(rows), location) + location = store.put(tenant_id, artifact) + service.record_evidence_export(tenant_id, actor_id, len(rows), location) return {"location": location, "event_count": len(rows)} diff --git a/ome_memoryguard/auth.py b/ome_memoryguard/auth.py new file mode 100644 index 0000000..7bf9efb --- /dev/null +++ b/ome_memoryguard/auth.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any + +from fastapi import HTTPException, Request + +from .domain import Sensitivity + +_SENSITIVITY_RANK = { + Sensitivity.PUBLIC: 0, + Sensitivity.INTERNAL: 1, + Sensitivity.CONFIDENTIAL: 2, + Sensitivity.RESTRICTED: 3, +} + + +@dataclass(frozen=True) +class RequestIdentity: + tenant_id: str + actor_id: str + clearance: Sensitivity + groups: frozenset[str] + + +def _jwt_claims(request: Request) -> dict[str, Any] | None: + event = request.scope.get("aws.event") + if not isinstance(event, dict): + return None + request_context = event.get("requestContext") + if not isinstance(request_context, dict): + return None + authorizer = request_context.get("authorizer") + if not isinstance(authorizer, dict): + return None + jwt = authorizer.get("jwt") + if not isinstance(jwt, dict): + return None + claims = jwt.get("claims") + return claims if isinstance(claims, dict) else None + + +def _claim_text(claims: dict[str, Any], name: str) -> str | None: + value = claims.get(name) + if not isinstance(value, str): + return None + normalized = value.strip() + return normalized or None + + +def _groups_from_claim(value: Any) -> frozenset[str]: + if isinstance(value, list): + return frozenset(item.strip() for item in value if isinstance(item, str) and item.strip()) + if not isinstance(value, str): + return frozenset() + normalized = value.strip().removeprefix("[").removesuffix("]") + return frozenset(item.strip() for item in normalized.split(",") if item.strip()) + + +def get_request_identity(request: Request) -> RequestIdentity | None: + configured_mode = os.getenv("AUTH_MODE") + default_mode = "cognito" if os.getenv("AWS_LAMBDA_FUNCTION_NAME") else "demo" + auth_mode = (configured_mode or default_mode).strip().lower() + if auth_mode == "demo": + return None + if auth_mode != "cognito": + raise HTTPException(status_code=503, detail="authentication mode is not supported") + + claims = _jwt_claims(request) + if claims is None: + raise HTTPException(status_code=401, detail="verified JWT claims are required") + + tenant_id = _claim_text(claims, "custom:tenant_id") + actor_id = _claim_text(claims, "sub") + if tenant_id is None or actor_id is None: + raise HTTPException(status_code=403, detail="required identity claims are missing") + + raw_clearance = _claim_text(claims, "custom:clearance") or Sensitivity.PUBLIC.value + try: + clearance = Sensitivity(raw_clearance.lower()) + except ValueError as error: + raise HTTPException(status_code=403, detail="clearance claim is invalid") from error + + return RequestIdentity( + tenant_id=tenant_id, + actor_id=actor_id, + clearance=clearance, + groups=_groups_from_claim(claims.get("cognito:groups")), + ) + + +def effective_sensitivity( + requested: Sensitivity, identity: RequestIdentity | None +) -> Sensitivity: + if identity is None or _SENSITIVITY_RANK[requested] <= _SENSITIVITY_RANK[identity.clearance]: + return requested + return identity.clearance + + +def require_any_group(identity: RequestIdentity | None, allowed: set[str]) -> None: + if identity is not None and identity.groups.isdisjoint(allowed): + raise HTTPException(status_code=403, detail="insufficient role") diff --git a/ome_memoryguard/demo.py b/ome_memoryguard/demo.py index 2f2d222..07b5433 100644 --- a/ome_memoryguard/demo.py +++ b/ome_memoryguard/demo.py @@ -11,26 +11,28 @@ button,input,select,textarea{font:inherit}button:focus-visible,input:focus-visible,select:focus-visible,textarea:focus-visible{outline:2px solid var(--cyan);outline-offset:2px} main{max-width:1200px;margin:auto;padding:30px 22px 64px}.top{display:flex;justify-content:space-between;align-items:flex-start;gap:24px;margin-bottom:24px}.eyebrow{font:700 12px/1.2 ui-monospace,Consolas;color:var(--mint);letter-spacing:.17em;text-transform:uppercase} h1{font-size:clamp(38px,6vw,68px);line-height:.93;margin:12px 0 16px;letter-spacing:-.055em}.lead{max-width:720px;color:var(--muted);font-size:17px}.badges{display:flex;gap:8px;flex-wrap:wrap;justify-content:flex-end}.badge{border:1px solid #2d7768;background:#0d2b29;color:var(--mint);padding:8px 12px;border-radius:999px;white-space:nowrap;font:700 12px ui-monospace,Consolas}.badge.secondary{border-color:#315468;background:#102633;color:var(--cyan)} - .metrics{display:grid;grid-template-columns:repeat(3,1fr);gap:10px;margin:0 0 18px}.metric{border:1px solid var(--line);background:#0a1720aa;border-radius:14px;padding:13px 15px}.metric b{display:block;font-size:13px}.metric span{color:var(--muted);font-size:12px} + .metrics{display:grid;grid-template-columns:repeat(3,1fr);gap:10px;margin:0 0 12px}.metric{border:1px solid var(--line);background:#0a1720aa;border-radius:14px;padding:13px 15px}.metric b{display:block;font-size:13px}.metric span{color:var(--muted);font-size:12px}.auth-control{padding:7px 11px;border-radius:999px;font:700 12px ui-monospace,Consolas} + .mode-notice{display:flex;gap:10px;align-items:baseline;border:1px solid var(--line);border-radius:12px;padding:11px 14px;margin:0 0 18px;background:#0a1720aa;color:var(--muted);font-size:12px}.mode-notice strong{color:var(--ink);white-space:nowrap}.mode-local{border-color:#315468}.mode-live{border-color:#2d7768}.mode-warning{border-color:#7b5e27;color:#ffd98b} .grid{display:grid;grid-template-columns:1.04fr .96fr;gap:18px}.card{background:linear-gradient(180deg,var(--panel2),var(--panel));border:1px solid var(--line);border-radius:18px;padding:20px;box-shadow:0 18px 55px #0006}h2{font-size:18px;margin:0 0 4px}.hint{color:var(--muted);font-size:12px;margin-bottom:15px} .scenarios{display:flex;gap:8px;flex-wrap:wrap;margin-bottom:15px}button{border:1px solid #2c4b5b;background:#142a37;color:var(--ink);border-radius:10px;padding:10px 13px;font-weight:750;cursor:pointer;transition:.18s transform,.18s border-color,.18s background}button:hover{border-color:var(--mint);transform:translateY(-1px)}button.primary{background:var(--mint);border-color:var(--mint);color:#04221c}button:disabled{opacity:.55;cursor:wait;transform:none} label{display:block;color:var(--muted);font-size:12px;margin:12px 0 6px}.row{display:grid;grid-template-columns:1fr 1fr;gap:10px}input,select,textarea{width:100%;border:1px solid #294653;background:#06131b;color:var(--ink);border-radius:10px;padding:11px 12px}textarea{min-height:124px;resize:vertical}.actions{display:flex;gap:10px;margin-top:14px;flex-wrap:wrap} - .result{min-height:210px;border:1px dashed #385865;border-radius:14px;padding:17px;background:#07141c}.status{font:800 13px ui-monospace,Consolas;text-transform:uppercase;letter-spacing:.13em}.active{color:var(--mint)}.pending_review{color:var(--amber)}.denied,.error{color:var(--red)}pre{white-space:pre-wrap;word-break:break-word;color:#c0d4d8;margin:12px 0 0;font:13px/1.5 ui-monospace,Consolas}.empty{color:var(--muted)} + .result{min-height:260px;border:1px dashed #385865;border-radius:14px;padding:17px;background:#07141c}.decision-top{display:flex;align-items:center;justify-content:space-between;gap:12px}.status{font:800 13px ui-monospace,Consolas;text-transform:uppercase;letter-spacing:.13em}.active{color:var(--mint)}.retrieved{color:var(--cyan)}.pending_review,.unavailable,.no_match{color:var(--amber)}.denied,.error{color:var(--red)}.decision-grid{display:grid;grid-template-columns:repeat(3,1fr);gap:8px;margin:14px 0}.decision-fact{border:1px solid var(--line);border-radius:10px;padding:10px;background:#0b1a23}.decision-fact span{display:block;color:var(--muted);font-size:10px;text-transform:uppercase;letter-spacing:.08em}.decision-fact b{display:block;margin-top:4px;font-size:12px}.reason-list{display:flex;gap:6px;flex-wrap:wrap;margin:12px 0}.reason{padding:5px 8px;border:1px solid #315468;border-radius:999px;background:#0d2531;color:var(--cyan);font:700 10px ui-monospace,Consolas}.payload{margin-top:12px;border-top:1px solid var(--line);padding-top:10px}.payload summary{cursor:pointer;color:var(--muted);font-size:11px}.payload[open] summary{color:var(--cyan)}pre{white-space:pre-wrap;word-break:break-word;color:#c0d4d8;margin:10px 0 0;font:12px/1.5 ui-monospace,Consolas}.empty{color:var(--muted)} .risk{margin-top:14px}.risk-head{display:flex;justify-content:space-between;color:var(--muted);font-size:12px}.risk-track{height:7px;background:#1b303a;border-radius:999px;overflow:hidden;margin-top:6px}.risk-fill{height:100%;background:linear-gradient(90deg,var(--mint),var(--amber),var(--red));border-radius:999px} - .flow{display:grid;grid-template-columns:repeat(4,1fr);gap:8px;margin-top:18px}.flow div{padding:11px;border:1px solid var(--line);border-radius:10px;color:var(--muted);font-size:12px}.flow b{display:block;color:var(--ink);margin-bottom:3px}.timeline{display:grid;gap:9px;max-height:350px;overflow:auto}.event{border-left:3px solid #3d6577;background:#08161e;padding:10px 12px;border-radius:0 9px 9px 0}.event b{font-size:13px}.event small{display:block;color:var(--muted);margin-top:3px}.footer-note{margin-top:14px;color:var(--muted);font-size:12px} - @media(max-width:840px){.grid{grid-template-columns:1fr}.top{display:block}.badges{justify-content:flex-start;margin-top:14px}.metrics{grid-template-columns:1fr}.flow{grid-template-columns:1fr 1fr}.wide{grid-column:auto!important}} + .flow{display:grid;grid-template-columns:repeat(4,1fr);gap:8px;margin-top:18px}.flow div{padding:11px;border:1px solid var(--line);border-radius:10px;color:var(--muted);font-size:12px}.flow b{display:block;color:var(--ink);margin-bottom:3px}.timeline{display:grid;gap:9px;max-height:350px;overflow:auto}.event{border-left:3px solid #3d6577;background:#08161e;padding:10px 12px;border-radius:0 9px 9px 0}.event.event-active,.event.event-allowed{border-left-color:var(--mint)}.event.event-pending_review{border-left-color:var(--amber)}.event.event-denied{border-left-color:var(--red)}.event-head{display:flex;justify-content:space-between;gap:12px;align-items:baseline}.event b{font-size:13px}.event time{font:10px ui-monospace,Consolas;color:var(--muted);white-space:nowrap}.event small{display:block;color:var(--muted);margin-top:3px}.footer-note{margin-top:14px;color:var(--muted);font-size:12px} + @media(max-width:840px){.grid{grid-template-columns:1fr}.top{display:block}.badges{justify-content:flex-start;margin-top:14px}.metrics{grid-template-columns:1fr}.mode-notice{align-items:flex-start;flex-direction:column;gap:3px}.decision-grid{grid-template-columns:1fr}.flow{grid-template-columns:1fr 1fr}.wide{grid-column:auto!important}}
Governed agentic memory

OME
MemoryGuard

A policy gate between an AI agent and persistent memory. Every proposal is risk-scored, explainable, auditable, and tenant-scoped before it can shape a future action.
-
CockroachDB × AWSchecking runtime…
+
CockroachDB × AWSchecking runtime…checking auth…
Fail closedDenied content is fingerprinted, never persisted
Scoped retrievalTenant + namespace + status + expiry + clearance
Atomic evidenceMemory and audit event share one transaction
+
Checking runtime modeActions remain disabled only while a request is running.

1. Propose a memory

Run the same API flow an agent would use.
@@ -40,13 +42,13 @@
-

2. Decision

The API returns only policy metadata for a denied proposal.
Choose a scenario and run the policy gate.
+

2. Decision

Outcome, persistence boundary, review state, and safe evidence are shown separately.
Choose a scenario and run the policy gate.
01 · ProposePurpose and provenance
02 · EvaluateDeterministic policy
03 · PersistCockroachDB vector row
04 · AuditAtomic decision trail
-

3. Audit trail

Denied inputs appear only as reason codes and a SHA-256 fingerprint.
No events yet.
Demo identity and clearance are client-supplied. Production must derive both from authenticated claims.
+

3. Audit trail

Denied inputs appear only as reason codes and a SHA-256 fingerprint.
No events yet.
Protected deployments derive tenant, actor, and clearance from verified Cognito claims; client-supplied identity values are ignored.
''' diff --git a/ome_memoryguard/embeddings.py b/ome_memoryguard/embeddings.py index e21fb96..0ff9d16 100644 --- a/ome_memoryguard/embeddings.py +++ b/ome_memoryguard/embeddings.py @@ -2,10 +2,17 @@ import hashlib import json +import logging import math import os from typing import Protocol +logger = logging.getLogger(__name__) + + +class EmbeddingUnavailableError(RuntimeError): + """Raised when the configured remote embedding provider cannot respond safely.""" + class EmbeddingProvider(Protocol): def embed(self, text: str) -> list[float]: ... @@ -33,20 +40,45 @@ def embed(self, text: str) -> list[float]: class BedrockTitanEmbeddingProvider: def __init__(self, model_id: str | None = None, region: str | None = None) -> None: import boto3 + from botocore.config import Config self.model_id = model_id or os.getenv("BEDROCK_MODEL_ID", "amazon.titan-embed-text-v2:0") - self.client = boto3.client("bedrock-runtime", region_name=region or os.getenv("AWS_REGION")) + inference_region = region or os.getenv("BEDROCK_REGION") or os.getenv("AWS_REGION") + self.region = inference_region + self.client = boto3.client( + "bedrock-runtime", + region_name=inference_region, + config=Config(retries={"mode": "adaptive", "total_max_attempts": 6}), + ) def embed(self, text: str) -> list[float]: + from botocore.exceptions import BotoCoreError, ClientError + payload = {"inputText": text, "dimensions": 1024, "normalize": True} - response = self.client.invoke_model( - modelId=self.model_id, - body=json.dumps(payload), - contentType="application/json", - accept="application/json", - ) - body = json.loads(response["body"].read()) - return [float(value) for value in body["embedding"]] + try: + response = self.client.invoke_model( + modelId=self.model_id, + body=json.dumps(payload), + contentType="application/json", + accept="application/json", + ) + body = json.loads(response["body"].read()) + embedding = [float(value) for value in body["embedding"]] + except (BotoCoreError, ClientError, KeyError, TypeError, ValueError) as error: + error_code = type(error).__name__ + if isinstance(error, ClientError): + error_code = str(error.response.get("Error", {}).get("Code") or error_code) + logger.warning( + "bedrock_embedding_failed error_type=%s error_code=%s model_id=%s region=%s", + type(error).__name__, + error_code, + self.model_id, + self.region, + ) + raise EmbeddingUnavailableError("remote embedding provider unavailable") from error + if not embedding: + raise EmbeddingUnavailableError("remote embedding provider unavailable") + return embedding def embedding_provider_from_env() -> EmbeddingProvider: diff --git a/ome_memoryguard/repository.py b/ome_memoryguard/repository.py index 8b21d65..8cd97a1 100644 --- a/ome_memoryguard/repository.py +++ b/ome_memoryguard/repository.py @@ -5,8 +5,11 @@ from dataclasses import replace from datetime import UTC, datetime from typing import Any, Protocol +from urllib.parse import quote from uuid import UUID +import certifi + from .domain import AuditEvent, MemoryRecord, MemoryStatus, Sensitivity _ALLOWED_SENSITIVITIES = { @@ -22,15 +25,31 @@ def normalize_database_url(database_url: str) -> str: - """Convert a CockroachDB PostgreSQL URL without logging or exposing it.""" + """Select the installed psycopg v3 driver without logging or exposing the URL.""" value = database_url.strip() - if value.startswith("cockroachdb://"): - return value - if value.startswith("postgresql://"): - return value.replace("postgresql://", "cockroachdb://", 1) - if value.startswith("postgres://"): - return value.replace("postgres://", "cockroachdb://", 1) - raise ValueError("DATABASE_URL must use cockroachdb, postgresql, or postgres scheme") + if value.startswith("cockroachdb+psycopg://"): + normalized = value + elif value.startswith("cockroachdb://"): + normalized = value.replace("cockroachdb://", "cockroachdb+psycopg://", 1) + elif value.startswith("postgresql://"): + normalized = value.replace("postgresql://", "cockroachdb+psycopg://", 1) + elif value.startswith("postgres://"): + normalized = value.replace("postgres://", "cockroachdb+psycopg://", 1) + else: + raise ValueError("DATABASE_URL must use cockroachdb, postgresql, or postgres scheme") + + base, separator, query = normalized.partition("?") + if not separator: + return normalized + parameters = { + item.partition("=")[0].strip().lower(): item.partition("=")[2].strip().lower() + for item in query.split("&") + if item + } + if parameters.get("sslmode") == "verify-full" and "sslrootcert" not in parameters: + ca_bundle = quote(certifi.where(), safe="/") + return f"{base}?{query}&sslrootcert={ca_bundle}" + return normalized def _uuid(value: UUID | str | None) -> UUID | None: diff --git a/pyproject.toml b/pyproject.toml index a39f1eb..5f48aab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ license = { text = "MIT" } authors = [{ name = "Naiman Yelubayev" }] dependencies = [ "boto3>=1.35,<2", + "certifi>=2024.8,<2027", "fastapi>=0.115,<1", "mangum>=0.19,<1", "psycopg[binary]>=3.2,<4", diff --git a/requirements.txt b/requirements.txt index 772ff44..c8c3a10 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ boto3>=1.35,<2 +certifi>=2024.8,<2027 fastapi>=0.115,<1 mangum>=0.19,<1 psycopg[binary]>=3.2,<4 diff --git a/scripts/aws-cost-preflight.ps1 b/scripts/aws-cost-preflight.ps1 new file mode 100644 index 0000000..2fc41a7 --- /dev/null +++ b/scripts/aws-cost-preflight.ps1 @@ -0,0 +1,105 @@ +[CmdletBinding()] +param( + [Parameter(Mandatory = $true)] + [ValidatePattern('^[A-Za-z0-9_.-]{1,100}$')] + [string]$BudgetName, + + [AllowEmptyString()] + [string]$Profile = '', + + [ValidateRange(0.01, 5.00)] + [decimal]$MaximumBudgetUsd = 5.00 +) + +$ErrorActionPreference = 'Stop' + +function Resolve-AwsCli { + $Command = Get-Command aws -ErrorAction SilentlyContinue + if ($null -ne $Command) { + return $Command.Source + } + + $PerUserAws = Join-Path $env:LOCALAPPDATA 'Programs\Amazon\AWSCLIV2\aws.exe' + if (Test-Path -LiteralPath $PerUserAws) { + return $PerUserAws + } + + throw 'AWS CLI is not installed or could not be located' +} + +function Invoke-AwsJson { + param( + [Parameter(Mandatory = $true)] + [string]$Aws, + [Parameter(Mandatory = $true)] + [string[]]$Arguments + ) + + $EffectiveArguments = @($Arguments) + if (-not [string]::IsNullOrWhiteSpace($Profile)) { + $EffectiveArguments += @('--profile', $Profile) + } + $Output = & $Aws @EffectiveArguments 2>$null + if ($LASTEXITCODE -ne 0) { + throw 'Unable to read the configured AWS budget' + } + return ($Output -join [Environment]::NewLine) | ConvertFrom-Json +} + +function Convert-ToDecimal { + param( + [Parameter(Mandatory = $true)] + [object]$Value, + [Parameter(Mandatory = $true)] + [string]$Name + ) + + $Parsed = 0.0m + $Success = [decimal]::TryParse( + [string]$Value, + [System.Globalization.NumberStyles]::Float, + [System.Globalization.CultureInfo]::InvariantCulture, + [ref]$Parsed + ) + if (-not $Success) { + throw "AWS returned an invalid $Name amount" + } + return $Parsed +} + +$Aws = Resolve-AwsCli +$IdentityArguments = @('sts', 'get-caller-identity', '--query', 'Account', '--output', 'text') +if (-not [string]::IsNullOrWhiteSpace($Profile)) { + $IdentityArguments += @('--profile', $Profile) +} +$Account = & $Aws @IdentityArguments 2>$null +if ($LASTEXITCODE -ne 0 -or $Account -notmatch '^\d{12}$') { + throw 'AWS CLI is not authenticated with a valid account' +} +Write-Output 'CHECK_AWS_IDENTITY=PASS' + +$Response = Invoke-AwsJson -Aws $Aws -Arguments @( + 'budgets', 'describe-budget', + '--account-id', $Account, + '--budget-name', $BudgetName, + '--output', 'json' +) + +if ($Response.Budget.BudgetLimit.Unit -ne 'USD') { + throw 'The configured AWS budget must use USD' +} + +$Limit = Convert-ToDecimal -Value $Response.Budget.BudgetLimit.Amount -Name 'budget limit' +$Actual = Convert-ToDecimal -Value $Response.Budget.CalculatedSpend.ActualSpend.Amount -Name 'actual spend' +if ($Limit -gt $MaximumBudgetUsd) { + throw 'The configured AWS budget exceeds the approved maximum' +} +if ($Actual -ge $MaximumBudgetUsd) { + throw 'Actual AWS spend has reached the approved maximum' +} + +$LimitText = $Limit.ToString('0.00', [System.Globalization.CultureInfo]::InvariantCulture) +$ActualText = $Actual.ToString('0.00', [System.Globalization.CultureInfo]::InvariantCulture) +Write-Output "AWS_BUDGET_LIMIT_USD=$LimitText" +Write-Output "AWS_BUDGET_ACTUAL_USD=$ActualText" +Write-Output 'CHECK_AWS_BUDGET=PASS' diff --git a/scripts/bedrock-quota-preflight.ps1 b/scripts/bedrock-quota-preflight.ps1 new file mode 100644 index 0000000..6ac0345 --- /dev/null +++ b/scripts/bedrock-quota-preflight.ps1 @@ -0,0 +1,103 @@ +[CmdletBinding()] +param( + [string[]]$Region = @('us-west-2'), + + [AllowEmptyString()] + [string]$Profile = '' +) + +$ErrorActionPreference = 'Stop' +$RpmQuotaCode = 'L-26C560CE' +$TpmQuotaCode = 'L-DE641971' + +function Resolve-AwsCli { + $Command = Get-Command aws -ErrorAction SilentlyContinue + if ($null -ne $Command) { + return $Command.Source + } + + $PerUserAws = Join-Path $env:LOCALAPPDATA 'Programs\Amazon\AWSCLIV2\aws.exe' + if (Test-Path -LiteralPath $PerUserAws) { + return $PerUserAws + } + + throw 'AWS CLI is not installed or could not be located' +} + +function Get-QuotaValue { + param( + [Parameter(Mandatory = $true)] + [string]$Aws, + [Parameter(Mandatory = $true)] + [string]$AwsRegion, + [Parameter(Mandatory = $true)] + [string]$QuotaCode + ) + + $Arguments = @( + 'service-quotas', 'get-service-quota', + '--service-code', 'bedrock', + '--quota-code', $QuotaCode, + '--region', $AwsRegion, + '--query', 'Quota.Value', + '--output', 'text' + ) + if (-not [string]::IsNullOrWhiteSpace($Profile)) { + $Arguments += @('--profile', $Profile) + } + $Raw = & $Aws @Arguments 2>$null + if ($LASTEXITCODE -ne 0) { + throw "Unable to read Bedrock quota $QuotaCode in $AwsRegion" + } + + $Value = 0.0 + $Parsed = [double]::TryParse( + $Raw.Trim(), + [System.Globalization.NumberStyles]::Float, + [System.Globalization.CultureInfo]::InvariantCulture, + [ref]$Value + ) + if (-not $Parsed) { + throw "AWS returned an invalid quota value for $QuotaCode in $AwsRegion" + } + return $Value +} + +$Aws = Resolve-AwsCli +$IdentityArguments = @('sts', 'get-caller-identity', '--query', 'Account', '--output', 'text') +if (-not [string]::IsNullOrWhiteSpace($Profile)) { + $IdentityArguments += @('--profile', $Profile) +} +$Account = & $Aws @IdentityArguments 2>$null +if ($LASTEXITCODE -ne 0 -or [string]::IsNullOrWhiteSpace($Account)) { + throw 'AWS CLI is not authenticated; configure a short-lived deployment profile first' +} +Write-Output 'CHECK_AWS_IDENTITY=PASS' + +$ReadyRegions = @() +foreach ($AwsRegion in $Region) { + if ($AwsRegion -notmatch '^[a-z]{2}(-gov)?-[a-z0-9-]+-[0-9]+$') { + throw "Invalid AWS Region: $AwsRegion" + } + + $Rpm = Get-QuotaValue -Aws $Aws -AwsRegion $AwsRegion -QuotaCode $RpmQuotaCode + $Tpm = Get-QuotaValue -Aws $Aws -AwsRegion $AwsRegion -QuotaCode $TpmQuotaCode + $Ready = $Rpm -gt 0 -and $Tpm -gt 0 + if ($Ready) { + $ReadyRegions += $AwsRegion + } + + $RpmText = $Rpm.ToString([System.Globalization.CultureInfo]::InvariantCulture) + $TpmText = $Tpm.ToString([System.Globalization.CultureInfo]::InvariantCulture) + Write-Output "BEDROCK_QUOTA_REGION=$AwsRegion" + Write-Output "BEDROCK_QUOTA_RPM=$RpmText" + Write-Output "BEDROCK_QUOTA_TPM=$TpmText" + Write-Output "BEDROCK_QUOTA_READY=$($Ready.ToString().ToLowerInvariant())" +} + +if ($ReadyRegions.Count -eq 0) { + throw 'Bedrock Titan V2 on-demand quota is zero in every checked Region' +} + +Write-Output "BEDROCK_READY_REGIONS=$($ReadyRegions -join ',')" +Write-Output 'CHECK_BEDROCK_QUOTA=PASS' diff --git a/scripts/live-bedrock-e2e.ps1 b/scripts/live-bedrock-e2e.ps1 new file mode 100644 index 0000000..aa5ea1b --- /dev/null +++ b/scripts/live-bedrock-e2e.ps1 @@ -0,0 +1,103 @@ +[CmdletBinding()] +param( + [Parameter(Mandatory = $true)] + [ValidatePattern('^https://[A-Za-z0-9.-]+/?$')] + [string]$ApiUrl, + + [Parameter(Mandatory = $true)] + [ValidatePattern('^[A-Za-z0-9_.-]{1,100}$')] + [string]$BudgetName, + + [ValidatePattern('^[a-z]{2}(-gov)?-[a-z0-9-]+-[0-9]+$')] + [string]$BedrockRegion = 'us-west-2', + + [AllowEmptyString()] + [string]$Profile = '' +) + +$ErrorActionPreference = 'Stop' +$ProjectRoot = (Resolve-Path -LiteralPath (Join-Path $PSScriptRoot '..')).Path +$TokenVariable = 'MEMORYGUARD_E2E_ID_TOKEN' +$Token = [Environment]::GetEnvironmentVariable($TokenVariable, 'Process') + +if ([string]::IsNullOrWhiteSpace($Token)) { + throw "$TokenVariable must contain a short-lived synthetic Cognito ID token" +} +if ($Token.Length -gt 16KB -or $Token -notmatch '^[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+$') { + throw "$TokenVariable is not a valid compact JWT" +} + +$QuotaArguments = @('-Region', $BedrockRegion) +if (-not [string]::IsNullOrWhiteSpace($Profile)) { + $QuotaArguments += @('-Profile', $Profile) +} +& (Join-Path $PSScriptRoot 'bedrock-quota-preflight.ps1') @QuotaArguments +if ($LASTEXITCODE -ne 0) { + throw 'Bedrock quota preflight failed' +} + +$CostArguments = @('-BudgetName', $BudgetName, '-MaximumBudgetUsd', 5) +if (-not [string]::IsNullOrWhiteSpace($Profile)) { + $CostArguments += @('-Profile', $Profile) +} +& (Join-Path $PSScriptRoot 'aws-cost-preflight.ps1') @CostArguments +if ($LASTEXITCODE -ne 0) { + throw 'AWS cost preflight failed' +} + +$BaseUrl = $ApiUrl.TrimEnd('/') +$Headers = @{ Authorization = "Bearer $Token" } +$RunSuffix = [DateTimeOffset]::UtcNow.ToUnixTimeSeconds().ToString( + [System.Globalization.CultureInfo]::InvariantCulture +) +$Namespace = "synthetic-e2e-$RunSuffix" +$SyntheticContent = 'Synthetic preference: provide concise weekly status updates.' + +$Proposal = @{ + tenant_id = 'ignored-jwt-tenant' + agent_id = 'ignored-jwt-actor' + namespace = $Namespace + content = $SyntheticContent + purpose = 'bounded synthetic Bedrock E2E' + sensitivity = 'internal' + source_uri = 'synthetic://memoryguard/e2e' +} | ConvertTo-Json -Compress + +$ProposalResult = Invoke-RestMethod -Method Post -Uri "$BaseUrl/v1/memories/propose" ` + -Headers $Headers -ContentType 'application/json' -Body $Proposal +if ($ProposalResult.status -ne 'active' -or [string]::IsNullOrWhiteSpace($ProposalResult.id)) { + throw 'Synthetic proposal did not become active' +} + +$Search = @{ + tenant_id = 'ignored-jwt-tenant' + actor_id = 'ignored-jwt-actor' + namespace = $Namespace + query = 'weekly status preference' + limit = 1 + max_sensitivity = 'internal' +} | ConvertTo-Json -Compress + +$SearchResult = Invoke-RestMethod -Method Post -Uri "$BaseUrl/v1/memories/search" ` + -Headers $Headers -ContentType 'application/json' -Body $Search +if (@($SearchResult.items).Count -ne 1 -or $SearchResult.items[0].id -ne $ProposalResult.id) { + throw 'Synthetic memory was not retrieved by the scoped vector search' +} + +$AuditResult = Invoke-RestMethod -Method Get ` + -Uri "$BaseUrl/v1/audit/events?tenant_id=ignored-jwt-tenant&limit=20" -Headers $Headers +$RelevantEvents = @($AuditResult.items | Where-Object { $_.memory_id -eq $ProposalResult.id }) +$EventTypes = @($RelevantEvents.event_type) +if ('memory_proposed' -notin $EventTypes -or 'memory_retrieved' -notin $EventTypes) { + throw 'The synthetic proposal and retrieval audit events were not both observed' +} + +Write-Output 'E2E_SYNTHETIC_DATA_ONLY=true' +Write-Output 'E2E_MAX_APPLICATION_EMBEDDING_CALLS=2' +Write-Output 'E2E_PROPOSAL_STATUS=active' +Write-Output 'E2E_RETRIEVAL_COUNT=1' +Write-Output "E2E_AUDIT_EVENT_COUNT=$($RelevantEvents.Count)" +Write-Output 'CHECK_LIVE_BEDROCK_E2E=PASS' + +$Token = $null +$Headers.Clear() diff --git a/scripts/sam-package.ps1 b/scripts/sam-package.ps1 new file mode 100644 index 0000000..d29bfab --- /dev/null +++ b/scripts/sam-package.ps1 @@ -0,0 +1,95 @@ +[CmdletBinding()] +param( + [Parameter(Mandatory = $true)] + [ValidatePattern('^[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]$')] + [string]$S3Bucket, + + [AllowEmptyString()] + [string]$Profile = 'ome-memoryguard-deployer', + [string]$Region = 'eu-central-1' +) + +$ErrorActionPreference = 'Stop' +$ProjectRoot = (Resolve-Path -LiteralPath (Join-Path $PSScriptRoot '..')).Path +$TempRoot = [System.IO.Path]::GetFullPath([System.IO.Path]::GetTempPath()) +$StageRoot = Join-Path $TempRoot ("ome-memoryguard-package-" + [guid]::NewGuid().ToString('N')) +$OutputDirectory = Join-Path $ProjectRoot '.aws-sam\deploy' +$OutputTemplate = Join-Path $OutputDirectory 'packaged.yaml' + +function Resolve-SamCli { + $Command = Get-Command sam -ErrorAction SilentlyContinue + if ($null -ne $Command) { + return $Command.Source + } + + $PerUserSam = Join-Path $env:LOCALAPPDATA 'Programs\Amazon\AWSSAMCLI\bin\sam.cmd' + if (Test-Path -LiteralPath $PerUserSam) { + return $PerUserSam + } + + throw 'AWS SAM CLI is not installed or could not be located' +} + +$Sam = Resolve-SamCli +$OldTelemetry = $env:SAM_CLI_TELEMETRY + +try { + New-Item -ItemType Directory -Path $StageRoot | Out-Null + & robocopy $ProjectRoot $StageRoot /E /XD .git .venv .aws-sam __pycache__ /XF .env '*.pyc' '*.pyo' '*.log' /NFL /NDL /NJH /NJS /NC /NS /NP | Out-Null + $RobocopyExit = $LASTEXITCODE + if ($RobocopyExit -gt 7) { + throw "SAM staging copy failed with robocopy exit code $RobocopyExit" + } + + $env:SAM_CLI_TELEMETRY = '0' + Push-Location -LiteralPath $StageRoot + try { + & $Sam build --no-cached --template-file template.yaml + if ($LASTEXITCODE -ne 0) { + throw 'SAM build failed' + } + + $PackageArguments = @( + 'package', + '--template-file', '.aws-sam\build\template.yaml', + '--s3-bucket', $S3Bucket, + '--output-template-file', 'packaged.yaml', + '--region', $Region + ) + if (-not [string]::IsNullOrWhiteSpace($Profile)) { + $PackageArguments += @('--profile', $Profile) + } + + $PackageOutput = & $Sam @PackageArguments 2>&1 + $PackageExitCode = $LASTEXITCODE + if ($PackageExitCode -ne 0) { + $SafeOutput = ($PackageOutput -join [Environment]::NewLine) ` + -replace '\b\d{12}\b', '' ` + -replace 'arn:aws:[^\s]+', '' + Write-Output $SafeOutput + throw 'SAM package failed' + } + + New-Item -ItemType Directory -Path $OutputDirectory -Force | Out-Null + Copy-Item -LiteralPath (Join-Path $StageRoot 'packaged.yaml') -Destination $OutputTemplate -Force + Write-Output 'CHECK_SAM_PACKAGE=PASS' + } + finally { + Pop-Location + } +} +finally { + $env:SAM_CLI_TELEMETRY = $OldTelemetry + + if (Test-Path -LiteralPath $StageRoot) { + $ResolvedStage = [System.IO.Path]::GetFullPath((Resolve-Path -LiteralPath $StageRoot).Path) + $ExpectedPrefix = $TempRoot.TrimEnd([System.IO.Path]::DirectorySeparatorChar) + [System.IO.Path]::DirectorySeparatorChar + if (-not $ResolvedStage.StartsWith($ExpectedPrefix, [System.StringComparison]::OrdinalIgnoreCase)) { + throw "Refusing to remove a staging directory outside the temporary root: $ResolvedStage" + } + if (-not ([System.IO.Path]::GetFileName($ResolvedStage)).StartsWith('ome-memoryguard-package-')) { + throw "Refusing to remove an unexpected staging directory: $ResolvedStage" + } + Remove-Item -LiteralPath $ResolvedStage -Recurse -Force + } +} diff --git a/template.yaml b/template.yaml index 87e8540..0542e0c 100644 --- a/template.yaml +++ b/template.yaml @@ -9,13 +9,150 @@ Parameters: BedrockModelId: Type: String Default: amazon.titan-embed-text-v2:0 + BedrockRegion: + Type: String + Default: eu-west-3 + AllowedPattern: '^[a-z]{2}(-gov)?-[a-z0-9-]+-[0-9]+$' + Description: Approved regional Bedrock endpoint; may differ from the Lambda Region + EnableDevAdminAuth: + Type: String + Default: "false" + AllowedValues: + - "false" + - "true" + Description: Enable Cognito admin password auth only for isolated dev smoke tests + DemoCallbackUrl: + Type: String + Default: http://localhost:8000/ + AllowedPattern: '^(https://.+|http://(localhost|127\.0\.0\.1)(:[0-9]+)?(/.*)?)$' + Description: Exact OAuth callback and logout URL for the browser demo + +Conditions: + EnableDevAdminAuthCondition: !Equals [!Ref EnableDevAdminAuth, "true"] Resources: + MemoryGuardUserPool: + Type: AWS::Cognito::UserPool + DeletionPolicy: Delete + UpdateReplacePolicy: Delete + Properties: + UserPoolName: !Sub "${AWS::StackName}-users" + UsernameConfiguration: + CaseSensitive: false + AutoVerifiedAttributes: + - email + Policies: + PasswordPolicy: + MinimumLength: 12 + RequireLowercase: true + RequireNumbers: true + RequireSymbols: true + RequireUppercase: true + TemporaryPasswordValidityDays: 1 + Schema: + - Name: tenant_id + AttributeDataType: String + Mutable: false + Required: false + StringAttributeConstraints: + MinLength: "1" + MaxLength: "100" + - Name: clearance + AttributeDataType: String + Mutable: true + Required: false + StringAttributeConstraints: + MinLength: "6" + MaxLength: "12" + UserPoolTags: + Project: ome-memoryguard + Environment: dev + + MemoryGuardUserPoolClient: + Type: AWS::Cognito::UserPoolClient + Properties: + ClientName: !Sub "${AWS::StackName}-client" + UserPoolId: !Ref MemoryGuardUserPool + GenerateSecret: false + PreventUserExistenceErrors: ENABLED + EnableTokenRevocation: true + AccessTokenValidity: 15 + IdTokenValidity: 15 + RefreshTokenValidity: 1 + TokenValidityUnits: + AccessToken: minutes + IdToken: minutes + RefreshToken: days + ExplicitAuthFlows: + - ALLOW_USER_SRP_AUTH + - ALLOW_REFRESH_TOKEN_AUTH + - !If + - EnableDevAdminAuthCondition + - ALLOW_ADMIN_USER_PASSWORD_AUTH + - !Ref AWS::NoValue + AllowedOAuthFlowsUserPoolClient: true + AllowedOAuthFlows: + - code + AllowedOAuthScopes: + - openid + - email + - profile + CallbackURLs: + - !Ref DemoCallbackUrl + LogoutURLs: + - !Ref DemoCallbackUrl + SupportedIdentityProviders: + - COGNITO + + MemoryGuardUserPoolDomain: + Type: AWS::Cognito::UserPoolDomain + Properties: + Domain: !Sub "${AWS::StackName}-${AWS::AccountId}" + ManagedLoginVersion: 1 + UserPoolId: !Ref MemoryGuardUserPool + + MemoryGuardReviewerGroup: + Type: AWS::Cognito::UserPoolGroup + Properties: + GroupName: memoryguard-reviewer + Description: May approve memories pending human review + UserPoolId: !Ref MemoryGuardUserPool + + MemoryGuardAuditorGroup: + Type: AWS::Cognito::UserPoolGroup + Properties: + GroupName: memoryguard-auditor + Description: May export tenant-scoped audit evidence + UserPoolId: !Ref MemoryGuardUserPool + + MemoryGuardAdminGroup: + Type: AWS::Cognito::UserPoolGroup + Properties: + GroupName: memoryguard-admin + Description: May approve memories and export audit evidence + UserPoolId: !Ref MemoryGuardUserPool + + MemoryGuardHttpApi: + Type: AWS::Serverless::HttpApi + Properties: + Name: !Sub "${AWS::StackName}-api" + StageName: $default + Auth: + DefaultAuthorizer: MemoryGuardJwtAuthorizer + Authorizers: + MemoryGuardJwtAuthorizer: + IdentitySource: $request.header.Authorization + JwtConfiguration: + issuer: !Sub "https://cognito-idp.${AWS::Region}.${AWS::URLSuffix}/${MemoryGuardUserPool}" + audience: + - !Ref MemoryGuardUserPoolClient + EvidenceBucket: Type: AWS::S3::Bucket DeletionPolicy: Retain UpdateReplacePolicy: Retain Properties: + BucketName: !Sub "${AWS::StackName}-${AWS::AccountId}-${AWS::Region}-evidence" BucketEncryption: ServerSideEncryptionConfiguration: - ServerSideEncryptionByDefault: @@ -31,46 +168,119 @@ Resources: VersioningConfiguration: Status: Enabled + MemoryGuardFunctionLogGroup: + Type: AWS::Logs::LogGroup + DeletionPolicy: Delete + UpdateReplacePolicy: Delete + Properties: + LogGroupName: !Sub "/aws/lambda/${AWS::StackName}-api" + RetentionInDays: 7 + Tags: + - Key: Project + Value: ome-memoryguard + - Key: Environment + Value: dev + + MemoryGuardFunctionRole: + Type: AWS::IAM::Role + Properties: + RoleName: !Sub "${AWS::StackName}-lambda" + AssumeRolePolicyDocument: + Version: "2012-10-17" + Statement: + - Effect: Allow + Principal: + Service: lambda.amazonaws.com + Action: sts:AssumeRole + Policies: + - PolicyName: memoryguard-runtime + PolicyDocument: + Version: "2012-10-17" + Statement: + - Sid: CreateFunctionLogGroup + Effect: Allow + Action: logs:CreateLogGroup + Resource: !Sub "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${AWS::StackName}-api" + - Sid: WriteFunctionLogs + Effect: Allow + Action: + - logs:CreateLogStream + - logs:PutLogEvents + Resource: !Sub "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${AWS::StackName}-api:*" + - Sid: WriteEvidenceBundles + Effect: Allow + Action: s3:PutObject + Resource: !Sub "${EvidenceBucket.Arn}/evidence/*" + - Sid: ReadDatabaseSecret + Effect: Allow + Action: secretsmanager:GetSecretValue + Resource: !Ref DatabaseUrlSecretArn + - Sid: InvokeTitanEmbeddingModel + Effect: Allow + Action: bedrock:InvokeModel + Resource: !Sub "arn:${AWS::Partition}:bedrock:${BedrockRegion}::foundation-model/${BedrockModelId}" + MemoryGuardFunction: Type: AWS::Serverless::Function + DependsOn: + - MemoryGuardFunctionLogGroup Properties: + FunctionName: !Sub "${AWS::StackName}-api" CodeUri: . Handler: ome_memoryguard.lambda_handler.handler Runtime: python3.12 Timeout: 30 MemorySize: 512 + Role: !GetAtt MemoryGuardFunctionRole.Arn Environment: Variables: + AUTH_MODE: cognito + COGNITO_CLIENT_ID: !Ref MemoryGuardUserPoolClient + COGNITO_DOMAIN: !Sub "https://${MemoryGuardUserPoolDomain}.auth.${AWS::Region}.amazoncognito.com" + COGNITO_CALLBACK_URL: !Ref DemoCallbackUrl DATABASE_URL_SECRET_ARN: !Ref DatabaseUrlSecretArn EMBEDDING_PROVIDER: bedrock BEDROCK_MODEL_ID: !Ref BedrockModelId + BEDROCK_REGION: !Ref BedrockRegion EVIDENCE_BUCKET: !Ref EvidenceBucket - Policies: - - Statement: - - Sid: WriteEvidenceBundles - Effect: Allow - Action: s3:PutObject - Resource: !Sub "${EvidenceBucket.Arn}/evidence/*" - - Sid: ReadDatabaseSecret - Effect: Allow - Action: secretsmanager:GetSecretValue - Resource: !Ref DatabaseUrlSecretArn - - Sid: InvokeTitanEmbeddingModel - Effect: Allow - Action: bedrock:InvokeModel - Resource: !Sub "arn:${AWS::Partition}:bedrock:${AWS::Region}::foundation-model/${BedrockModelId}" Events: - RootApi: + DemoRoot: Type: HttpApi Properties: + ApiId: !Ref MemoryGuardHttpApi Path: / - Method: ANY + Method: GET + Auth: + Authorizer: NONE + RouteSettings: + ThrottlingBurstLimit: 20 + ThrottlingRateLimit: 10 + HealthApi: + Type: HttpApi + Properties: + ApiId: !Ref MemoryGuardHttpApi + Path: /healthz + Method: GET + Auth: + Authorizer: NONE + RouteSettings: + ThrottlingBurstLimit: 20 + ThrottlingRateLimit: 10 + AuthConfigApi: + Type: HttpApi + Properties: + ApiId: !Ref MemoryGuardHttpApi + Path: /auth/config + Method: GET + Auth: + Authorizer: NONE RouteSettings: ThrottlingBurstLimit: 20 ThrottlingRateLimit: 10 - ProxyApi: + ProtectedProxyApi: Type: HttpApi Properties: + ApiId: !Ref MemoryGuardHttpApi Path: /{proxy+} Method: ANY RouteSettings: @@ -80,6 +290,18 @@ Resources: Outputs: ApiUrl: Description: MemoryGuard API endpoint - Value: !Sub "https://${ServerlessHttpApi}.execute-api.${AWS::Region}.${AWS::URLSuffix}" + Value: !Sub "https://${MemoryGuardHttpApi}.execute-api.${AWS::Region}.${AWS::URLSuffix}" EvidenceBucketName: Value: !Ref EvidenceBucket + CognitoUserPoolId: + Description: Cognito user pool for MemoryGuard JWTs + Value: !Ref MemoryGuardUserPool + CognitoUserPoolClientId: + Description: Public Cognito app client for MemoryGuard JWTs + Value: !Ref MemoryGuardUserPoolClient + CognitoIssuer: + Description: Expected issuer for MemoryGuard JWTs + Value: !Sub "https://cognito-idp.${AWS::Region}.${AWS::URLSuffix}/${MemoryGuardUserPool}" + CognitoDomain: + Description: Cognito domain used by the browser demo OAuth flow + Value: !Sub "https://${MemoryGuardUserPoolDomain}.auth.${AWS::Region}.amazoncognito.com" diff --git a/tests/test_api.py b/tests/test_api.py index 9233fd1..9cfbe0f 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -4,8 +4,20 @@ from fastapi.testclient import TestClient -from ome_memoryguard.api import app, get_evidence_store, get_service, health -from ome_memoryguard.embeddings import DeterministicEmbeddingProvider +from ome_memoryguard.api import ( + app, + auth_config, + get_evidence_store, + get_request_identity, + get_service, + health, +) +from ome_memoryguard.auth import RequestIdentity +from ome_memoryguard.domain import Sensitivity +from ome_memoryguard.embeddings import ( + DeterministicEmbeddingProvider, + EmbeddingUnavailableError, +) from ome_memoryguard.repository import InMemoryRepository from ome_memoryguard.service import MemoryService @@ -70,6 +82,26 @@ def test_health_reports_actual_loaded_storage_mode(self) -> None: ): self.assertEqual(health()["storage"], "cockroachdb") + def test_public_auth_config_is_disabled_or_complete(self) -> None: + with patch.dict(os.environ, {}, clear=True): + self.assertEqual(auth_config(), {"enabled": False}) + + public_settings = { + "COGNITO_DOMAIN": "https://synthetic.auth.eu-central-1.amazoncognito.com/", + "COGNITO_CLIENT_ID": "synthetic-client-id", + "COGNITO_CALLBACK_URL": "https://example.invalid/", + } + with patch.dict(os.environ, public_settings, clear=True): + self.assertEqual( + auth_config(), + { + "enabled": True, + "domain": "https://synthetic.auth.eu-central-1.amazoncognito.com", + "client_id": "synthetic-client-id", + "redirect_uri": "https://example.invalid/", + }, + ) + def test_search_and_audit_reject_out_of_bounds_inputs(self) -> None: search_payload = { "tenant_id": "acme", @@ -147,6 +179,91 @@ def test_evidence_export_contains_audit_metadata_not_memory_content(self) -> Non self.assertNotIn("content", artifact["events"][0]) self.assertEqual(self.repository.list_events("acme")[0].event_type, "evidence_exported") + def test_verified_identity_overrides_client_tenant_actor_and_clearance(self) -> None: + identity = RequestIdentity( + tenant_id="tenant-secure", + actor_id="subject-123", + clearance=Sensitivity.INTERNAL, + groups=frozenset({"memoryguard-reviewer", "memoryguard-auditor"}), + ) + app.dependency_overrides[get_request_identity] = lambda: identity + + response = self.client.post( + "/v1/memories/propose", + json=self.proposal(tenant_id="spoofed", agent_id="spoofed"), + ) + self.assertEqual(response.status_code, 200) + stored = next(iter(self.repository.memories.values())) + self.assertEqual(stored.proposal.tenant_id, "tenant-secure") + self.assertEqual(stored.proposal.agent_id, "subject-123") + + search = self.client.post( + "/v1/memories/search", + json={ + "tenant_id": "spoofed", + "actor_id": "spoofed", + "namespace": "support", + "query": "Customer prefers concise weekly updates.", + "max_sensitivity": "restricted", + }, + ) + self.assertEqual(search.status_code, 200) + self.assertEqual(len(search.json()["items"]), 1) + + audit = self.client.get("/v1/audit/events", params={"tenant_id": "spoofed"}) + self.assertEqual(audit.status_code, 200) + self.assertGreaterEqual(len(audit.json()["items"]), 1) + + def test_verified_identity_without_reviewer_group_cannot_approve(self) -> None: + identity = RequestIdentity( + tenant_id="tenant-secure", + actor_id="subject-123", + clearance=Sensitivity.CONFIDENTIAL, + groups=frozenset(), + ) + app.dependency_overrides[get_request_identity] = lambda: identity + proposed = self.client.post( + "/v1/memories/propose", + json=self.proposal(sensitivity="confidential"), + ).json() + response = self.client.post( + f"/v1/memories/{proposed['id']}/approve", + json={"tenant_id": "spoofed", "approver_id": "spoofed"}, + ) + self.assertEqual(response.status_code, 403) + + def test_cognito_mode_fails_closed_without_gateway_claims(self) -> None: + with patch.dict(os.environ, {"AUTH_MODE": "cognito"}, clear=True): + response = self.client.post("/v1/memories/propose", json=self.proposal()) + self.assertEqual(response.status_code, 401) + + def test_embedding_outage_returns_503_without_persistence(self) -> None: + unavailable_repository = InMemoryRepository() + unavailable_service = MemoryService( + unavailable_repository, UnavailableEmbeddingProvider() + ) + app.dependency_overrides[get_service] = lambda: unavailable_service + + proposed = self.client.post("/v1/memories/propose", json=self.proposal()) + searched = self.client.post( + "/v1/memories/search", + json={ + "tenant_id": "acme", + "actor_id": "agent-2", + "namespace": "support", + "query": "Customer preference", + }, + ) + + self.assertEqual(proposed.status_code, 503) + self.assertEqual(searched.status_code, 503) + self.assertEqual( + proposed.json()["detail"]["reason_code"], + "embedding_provider_unavailable", + ) + self.assertEqual(unavailable_repository.memories, {}) + self.assertEqual(unavailable_repository.events, []) + class FakeEvidenceStore: def __init__(self) -> None: @@ -157,5 +274,10 @@ def put(self, tenant_id, artifact): return "s3://demo-evidence/evidence/synthetic/bundle.json" +class UnavailableEmbeddingProvider: + def embed(self, text): + raise EmbeddingUnavailableError("synthetic outage") + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_architecture.py b/tests/test_architecture.py index eaa6204..ea0aecd 100644 --- a/tests/test_architecture.py +++ b/tests/test_architecture.py @@ -25,17 +25,58 @@ def test_sam_template_has_scoped_iam_and_both_routes(self) -> None: template = (ROOT / "template.yaml").read_text(encoding="utf-8") self.assertIn("Path: /\n", template) self.assertIn("Path: /{proxy+}", template) - self.assertIn("foundation-model/${BedrockModelId}", template) + self.assertIn("bedrock:${BedrockRegion}::foundation-model/${BedrockModelId}", template) + self.assertIn("BEDROCK_REGION: !Ref BedrockRegion", template) self.assertIn("Resource: !Ref DatabaseUrlSecretArn", template) self.assertIn("Resource: !Sub \"${EvidenceBucket.Arn}/evidence/*\"", template) self.assertIn("DeletionPolicy: Retain", template) self.assertNotIn('Resource: "*"', template) self.assertNotIn("S3CrudPolicy", template) + def test_sam_template_protects_api_with_cognito_jwt(self) -> None: + template = (ROOT / "template.yaml").read_text(encoding="utf-8") + self.assertIn("Type: AWS::Cognito::UserPool", template) + self.assertIn("Type: AWS::Cognito::UserPoolClient", template) + self.assertIn("DefaultAuthorizer: MemoryGuardJwtAuthorizer", template) + self.assertIn("IdentitySource: $request.header.Authorization", template) + self.assertIn("AUTH_MODE: cognito", template) + self.assertIn("Path: /healthz", template) + self.assertIn("Path: /auth/config", template) + self.assertEqual(template.count("Authorizer: NONE"), 3) + self.assertIn("Type: AWS::Cognito::UserPoolDomain", template) + self.assertIn("AllowedOAuthFlowsUserPoolClient: true", template) + self.assertIn("- code", template) + self.assertIn("COGNITO_CALLBACK_URL: !Ref DemoCallbackUrl", template) + self.assertIn("GroupName: memoryguard-reviewer", template) + self.assertIn("GroupName: memoryguard-auditor", template) + self.assertIn("GroupName: memoryguard-admin", template) + self.assertIn('EnableDevAdminAuthCondition: !Equals [!Ref EnableDevAdminAuth, "true"]', template) + self.assertIn("ALLOW_ADMIN_USER_PASSWORD_AUTH", template) + self.assertIn("- !Ref AWS::NoValue", template) + self.assertIn("Type: AWS::Logs::LogGroup", template) + self.assertIn('LogGroupName: !Sub "/aws/lambda/${AWS::StackName}-api"', template) + self.assertIn("RetentionInDays: 7", template) + self.assertIn("MemoryGuardFunctionLogGroup", template) + self.assertIn('RoleName: !Sub "${AWS::StackName}-lambda"', template) + self.assertIn('FunctionName: !Sub "${AWS::StackName}-api"', template) + self.assertIn( + 'BucketName: !Sub "${AWS::StackName}-${AWS::AccountId}-${AWS::Region}-evidence"', + template, + ) + def test_demo_uses_safe_dom_rendering_and_correct_brand_text(self) -> None: demo = (ROOT / "ome_memoryguard" / "demo.py").read_text(encoding="utf-8") self.assertIn("CockroachDB × AWS", demo) self.assertIn("replaceChildren", demo) + self.assertIn("Local deterministic demo", demo) + self.assertIn("Live AWS development mode", demo) + self.assertIn("The action failed closed; no memory write was persisted", demo) + self.assertIn("Approve as human reviewer", demo) + self.assertIn("View sanitized API payload", demo) + self.assertIn("raw content blocked", demo) + self.assertIn("hash + reason codes", demo) + self.assertIn("Scoped retrieval complete", demo) + self.assertIn("retrieval audit", demo) self.assertNotIn("innerHTML", demo) self.assertNotIn("Г—", demo) self.assertNotIn("В·", demo) @@ -49,6 +90,16 @@ def test_local_artifacts_are_ignored(self) -> None: def test_verification_script_covers_local_quality_tools(self) -> None: script = (ROOT / "scripts" / "verify.ps1").read_text(encoding="utf-8") sam_script = (ROOT / "scripts" / "sam-verify.ps1").read_text(encoding="utf-8") + package_script = (ROOT / "scripts" / "sam-package.ps1").read_text(encoding="utf-8") + quota_script = (ROOT / "scripts" / "bedrock-quota-preflight.ps1").read_text( + encoding="utf-8" + ) + cost_script = (ROOT / "scripts" / "aws-cost-preflight.ps1").read_text( + encoding="utf-8" + ) + e2e_script = (ROOT / "scripts" / "live-bedrock-e2e.ps1").read_text( + encoding="utf-8" + ) secret_script = (ROOT / "scripts" / "secret-scan.ps1").read_text(encoding="utf-8") self.assertIn("CHECK_PYTEST=PASS", script) self.assertIn("CHECK_RUFF=PASS", script) @@ -57,6 +108,26 @@ def test_verification_script_covers_local_quality_tools(self) -> None: self.assertIn("secret-scan.ps1", script) self.assertIn("CHECK_SAM_VALIDATE=PASS", sam_script) self.assertIn("CHECK_SAM_BUILD=PASS", sam_script) + self.assertIn("CHECK_SAM_PACKAGE=PASS", package_script) + self.assertIn("/XF .env", package_script) + self.assertIn("ome-memoryguard-package-", package_script) + self.assertIn("L-26C560CE", quota_script) + self.assertIn("L-DE641971", quota_script) + self.assertIn("CHECK_BEDROCK_QUOTA=PASS", quota_script) + self.assertNotIn("get-secret-value", quota_script.lower()) + self.assertIn("budgets', 'describe-budget'", cost_script) + self.assertIn("MaximumBudgetUsd = 5.00", cost_script) + self.assertIn("CHECK_AWS_BUDGET=PASS", cost_script) + self.assertNotIn("create-budget", cost_script.lower()) + self.assertNotIn("update-budget", cost_script.lower()) + self.assertIn("MEMORYGUARD_E2E_ID_TOKEN", e2e_script) + self.assertIn("E2E_SYNTHETIC_DATA_ONLY=true", e2e_script) + self.assertIn("E2E_MAX_APPLICATION_EMBEDDING_CALLS=2", e2e_script) + self.assertIn("CHECK_LIVE_BEDROCK_E2E=PASS", e2e_script) + self.assertIn("bedrock-quota-preflight.ps1", e2e_script) + self.assertIn("aws-cost-preflight.ps1", e2e_script) + self.assertNotIn("get-secret-value", e2e_script.lower()) + self.assertNotIn("write-output $token", e2e_script.lower()) self.assertIn("CHECK_GITLEAKS_GIT=PASS", secret_script) self.assertIn("CHECK_GITLEAKS_DIR=PASS", secret_script) diff --git a/tests/test_auth.py b/tests/test_auth.py new file mode 100644 index 0000000..1dd5833 --- /dev/null +++ b/tests/test_auth.py @@ -0,0 +1,106 @@ +import os +import unittest +from unittest.mock import patch + +from fastapi import HTTPException +from starlette.requests import Request + +from ome_memoryguard.auth import ( + RequestIdentity, + effective_sensitivity, + get_request_identity, + require_any_group, +) +from ome_memoryguard.domain import Sensitivity + + +def request_with_claims(claims: dict | None) -> Request: + scope = {"type": "http", "headers": []} + if claims is not None: + scope["aws.event"] = { + "requestContext": {"authorizer": {"jwt": {"claims": claims}}} + } + return Request(scope) + + +class RequestIdentityTests(unittest.TestCase): + def test_demo_mode_does_not_require_gateway_claims(self) -> None: + with patch.dict(os.environ, {"AUTH_MODE": "demo"}, clear=True): + self.assertIsNone(get_request_identity(request_with_claims(None))) + + def test_cognito_mode_fails_closed_without_verified_claims(self) -> None: + with ( + patch.dict(os.environ, {"AUTH_MODE": "cognito"}, clear=True), + self.assertRaises(HTTPException) as raised, + ): + get_request_identity(request_with_claims(None)) + self.assertEqual(raised.exception.status_code, 401) + + def test_lambda_environment_defaults_to_fail_closed_cognito_mode(self) -> None: + with ( + patch.dict(os.environ, {"AWS_LAMBDA_FUNCTION_NAME": "memoryguard"}, clear=True), + self.assertRaises(HTTPException) as raised, + ): + get_request_identity(request_with_claims(None)) + self.assertEqual(raised.exception.status_code, 401) + + def test_cognito_claims_define_tenant_actor_clearance_and_groups(self) -> None: + claims = { + "sub": "subject-123", + "custom:tenant_id": "tenant-secure", + "custom:clearance": "confidential", + "cognito:groups": "[memoryguard-reviewer,memoryguard-auditor]", + } + with patch.dict(os.environ, {"AUTH_MODE": "cognito"}, clear=True): + identity = get_request_identity(request_with_claims(claims)) + self.assertEqual(identity.tenant_id, "tenant-secure") + self.assertEqual(identity.actor_id, "subject-123") + self.assertEqual(identity.clearance, Sensitivity.CONFIDENTIAL) + self.assertEqual( + identity.groups, + frozenset({"memoryguard-reviewer", "memoryguard-auditor"}), + ) + + def test_invalid_clearance_is_rejected(self) -> None: + claims = { + "sub": "subject-123", + "custom:tenant_id": "tenant-secure", + "custom:clearance": "superuser", + } + with ( + patch.dict(os.environ, {"AUTH_MODE": "cognito"}, clear=True), + self.assertRaises(HTTPException) as raised, + ): + get_request_identity(request_with_claims(claims)) + self.assertEqual(raised.exception.status_code, 403) + + def test_requested_clearance_is_clamped_to_verified_claim(self) -> None: + identity = RequestIdentity( + tenant_id="tenant-secure", + actor_id="subject-123", + clearance=Sensitivity.INTERNAL, + groups=frozenset(), + ) + self.assertEqual( + effective_sensitivity(Sensitivity.RESTRICTED, identity), + Sensitivity.INTERNAL, + ) + self.assertEqual( + effective_sensitivity(Sensitivity.PUBLIC, identity), + Sensitivity.PUBLIC, + ) + + def test_privileged_action_requires_verified_group(self) -> None: + identity = RequestIdentity( + tenant_id="tenant-secure", + actor_id="subject-123", + clearance=Sensitivity.INTERNAL, + groups=frozenset(), + ) + with self.assertRaises(HTTPException) as raised: + require_any_group(identity, {"memoryguard-reviewer"}) + self.assertEqual(raised.exception.status_code, 403) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_config.py b/tests/test_config.py index 6164f88..ff888c7 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -19,7 +19,7 @@ def get_secret_value(self, SecretId: str): class DatabaseSecretTests(unittest.TestCase): def test_plain_url_is_normalized(self) -> None: value = "postgresql://user:placeholder@example.invalid/defaultdb" - self.assertTrue(database_url_from_secret(value).startswith("cockroachdb://")) + self.assertTrue(database_url_from_secret(value).startswith("cockroachdb+psycopg://")) def test_json_secret_requires_database_url_key(self) -> None: synthetic = {"username": "demo", "password": "not-a-real-password"} @@ -35,7 +35,7 @@ def test_loader_uses_exact_secret_arn_without_printing_value(self) -> None: with patch.dict(os.environ, {"DATABASE_URL_SECRET_ARN": arn}, clear=True): self.assertTrue(load_database_url_from_secret(client)) self.assertEqual(client.secret_id, arn) - self.assertTrue(os.environ["DATABASE_URL"].startswith("cockroachdb://")) + self.assertTrue(os.environ["DATABASE_URL"].startswith("cockroachdb+psycopg://")) if __name__ == "__main__": diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py index f471681..ba9469a 100644 --- a/tests/test_embeddings.py +++ b/tests/test_embeddings.py @@ -2,8 +2,12 @@ import unittest from unittest.mock import patch +from botocore.exceptions import ClientError + from ome_memoryguard.embeddings import ( + BedrockTitanEmbeddingProvider, DeterministicEmbeddingProvider, + EmbeddingUnavailableError, embedding_provider_from_env, ) @@ -21,6 +25,48 @@ def test_unknown_provider_fails_closed(self) -> None: ), self.assertRaisesRegex(ValueError, "must be deterministic or bedrock"): embedding_provider_from_env() + def test_bedrock_failure_is_sanitized_and_fails_closed(self) -> None: + provider = BedrockTitanEmbeddingProvider.__new__(BedrockTitanEmbeddingProvider) + provider.model_id = "amazon.titan-embed-text-v2:0" + provider.region = "eu-west-3" + provider.client = FailingBedrockClient() + + with patch( + "ome_memoryguard.embeddings.logger.warning" + ) as warning, self.assertRaisesRegex( + EmbeddingUnavailableError, "remote embedding provider unavailable" + ): + provider.embed("synthetic test content") + + warning.assert_called_once() + logged = " ".join(str(value) for value in warning.call_args.args) + self.assertIn("ThrottlingException", logged) + self.assertNotIn("synthetic", logged) + + def test_bedrock_region_overrides_lambda_region(self) -> None: + with patch.dict( + os.environ, + {"AWS_REGION": "eu-central-1", "BEDROCK_REGION": "eu-west-3"}, + clear=True, + ), patch("boto3.client") as client: + BedrockTitanEmbeddingProvider() + + client.assert_called_once() + self.assertEqual(client.call_args.args, ("bedrock-runtime",)) + self.assertEqual(client.call_args.kwargs["region_name"], "eu-west-3") + self.assertEqual( + client.call_args.kwargs["config"].retries, + {"mode": "adaptive", "total_max_attempts": 6}, + ) + + +class FailingBedrockClient: + def invoke_model(self, **kwargs): + raise ClientError( + {"Error": {"Code": "ThrottlingException", "Message": "synthetic"}}, + "InvokeModel", + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_repository.py b/tests/test_repository.py index 028f106..cd31794 100644 --- a/tests/test_repository.py +++ b/tests/test_repository.py @@ -1,5 +1,6 @@ import unittest from datetime import UTC, datetime +from unittest.mock import patch from ome_memoryguard.domain import ( AuditEvent, @@ -83,15 +84,48 @@ def test_database_url_schemes_are_normalized(self) -> None: suffix = "user:placeholder@example.invalid/defaultdb" self.assertEqual( normalize_database_url("postgresql://" + suffix), - "cockroachdb://" + suffix, + "cockroachdb+psycopg://" + suffix, ) self.assertEqual( normalize_database_url("postgres://" + suffix), - "cockroachdb://" + suffix, + "cockroachdb+psycopg://" + suffix, + ) + self.assertEqual( + normalize_database_url("cockroachdb://" + suffix), + "cockroachdb+psycopg://" + suffix, + ) + self.assertEqual( + normalize_database_url("cockroachdb+psycopg://" + suffix), + "cockroachdb+psycopg://" + suffix, ) with self.assertRaises(ValueError): normalize_database_url("https://example.invalid") + def test_normalized_url_selects_installed_psycopg_driver(self) -> None: + from sqlalchemy import create_engine + + engine = create_engine( + normalize_database_url("postgresql://user:placeholder@example.invalid/defaultdb") + ) + try: + self.assertEqual(engine.dialect.driver, "psycopg") + finally: + engine.dispose() + + def test_verify_full_uses_packaged_ca_without_overriding_explicit_ca(self) -> None: + base = "postgresql://user:placeholder@example.invalid/defaultdb" + with patch( + "ome_memoryguard.repository.certifi.where", return_value="/opt/certifi/cacert.pem" + ): + self.assertTrue( + normalize_database_url(base + "?sslmode=verify-full").endswith( + "sslmode=verify-full&sslrootcert=/opt/certifi/cacert.pem" + ) + ) + explicit = base + "?sslmode=verify-full&sslrootcert=%2Fopt%2Fca.crt" + self.assertTrue(normalize_database_url(explicit).endswith("sslrootcert=%2Fopt%2Fca.crt")) + self.assertNotIn("sslrootcert", normalize_database_url(base + "?sslmode=require")) + def test_save_and_audit_event_share_one_transaction(self) -> None: connection = FakeConnection() engine = FakeEngine(connection)