diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..962c890 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,84 @@ +name: Regression and package checks + +on: + push: + pull_request: + workflow_dispatch: + +permissions: + contents: read + +env: + PIP_DISABLE_PIP_VERSION_CHECK: "1" + LITELLM_LOCAL_MODEL_COST_MAP: "True" + DO_NOT_TRACK: "1" + +jobs: + sqlite-and-package: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.13"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - name: Install test and build dependencies, including the MCP extra + run: python -m pip install -e '.[dev,mcp]' build + - name: Check recorded shared-core contracts + run: python scripts/check-shared-core.py --checkout . + - name: Run full SQLite suite + run: python scripts/check-ci.py --backend sqlite --output reports/sqlite.xml + - name: Run controlled retrieval regression benchmark + run: python -m evals.retrieval --output reports/retrieval.json + - name: Build sdist, then wheel; certify installed package and MCP + run: python scripts/check-package.py --output-dir reports/package + - name: Retain verification reports and package artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: sqlite-package-python-${{ matrix.python-version }} + path: reports/ + retention-days: 14 + + postgres: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.13"] + services: + postgres: + image: pgvector/pgvector:pg16 + env: + POSTGRES_DB: engram_test + POSTGRES_USER: engram_test + POSTGRES_PASSWORD: engram-test-only + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U engram_test -d engram_test" + --health-interval 5s + --health-timeout 5s + --health-retries 12 + env: + ENGRAM_TEST_POSTGRES_DSN: postgresql://engram_test:engram-test-only@127.0.0.1:5432/engram_test + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - run: python -m pip install -e '.[dev,mcp]' + - name: Require ready pgvector and run full suite with PostgreSQL tests enabled + run: python scripts/check-ci.py --backend postgres --output reports/postgres.xml + - name: Retain PostgreSQL verification report + if: always() + uses: actions/upload-artifact@v4 + with: + name: postgres-python-${{ matrix.python-version }} + path: reports/ + retention-days: 14 diff --git a/README.md b/README.md index 9bc7f4a..3221ae5 100644 --- a/README.md +++ b/README.md @@ -91,11 +91,9 @@ There are two storage options: - SQLite (local) — zero setup, tables auto-created on first run - PostgreSQL (local/remote) — requires creating the database first -For PostgreSQL, create the database and enable pgvector before starting the server: +For PostgreSQL, install pgvector on the database server and create the database first: ```sql CREATE DATABASE engram; -\c engram -CREATE EXTENSION IF NOT EXISTS vector; ``` Then set the environment variables: @@ -104,7 +102,14 @@ ENGRAM_STORAGE_BACKEND=postgres ENGRAM_POSTGRES_DSN=postgresql://user:password@host:5432/engram ``` -All tables are auto-created on first server start for both backends. +Engram enables the `vector` extension before opening its connection pool and creates +tables on first startup. The Docker Compose pgvector image already includes the +extension files, so `docker compose up` works with a fresh volume. + +On managed PostgreSQL, an administrator may need to run +`CREATE EXTENSION IF NOT EXISTS vector` in the Engram database before the first +startup if the application role cannot install extensions. Existing extensions and +data are preserved across restarts; do not delete your database volume to upgrade. ### 1. Configure the Server @@ -149,15 +154,11 @@ ENGRAM_CURATOR_SLOW_PATH_MODEL=claude-haiku-4-5 # For complex conflicts # ─── Embedding ────────────────────────────────────────────── ENGRAM_EMBEDDING_MODEL=text-embedding-3-small -# ─── Consolidation ────────────────────────────────────────── -ENGRAM_CONSOLIDATION_TRIGGER=every_10_commits # or "daily", "manual" -ENGRAM_FAST_DECAY_RATE=0.97 # Per-day salience decay for normal bullets -ENGRAM_SLOW_DECAY_RATE=0.995 # Per-day decay for decisions/schemas (slower) +# Consolidation is invoked through the SDK or REST API. +# This server does not schedule consolidation from environment variables. -# ─── Auth (disabled for local dev, enable for production) ─── -ENGRAM_AUTH_ENABLED=false -# COGNITO_USER_POOL_ID=us-west-2_xxx -# COGNITO_APP_CLIENT_ID=xxx +# The OSS server has no built-in account authentication. +# Cloud account/API-key authentication belongs to the hosted edition. ``` ### 2. Start the Server @@ -769,33 +770,59 @@ Capacity management with pressure levels: ### MCP Server (Claude) +For **Engram Cloud**, add `https://api.engram.so/mcp/sse` in Claude's +**Settings → Connectors → Add custom connector**, then sign in to Engram and +authorize access. See [Cloud setup](https://app.engram.so/setup). + +For a **local stdio bridge**, install from a checkout containing this fix using +Python 3.11 or later. The published `engram-contextdb` 0.4.5 wheel lacks the +module entrypoint; use the corrected source checkout until a fixed release is +available. The distribution is `engram-contextdb`; imports remain `engram`. + ```bash -pip install engram[mcp] -python -m engram.integrations.mcp_server +cd /absolute/path/to/engram +python3 -m venv .venv +.venv/bin/python -m pip install ".[mcp]" ``` -Configure in Claude Code: +Start your self-hosted Engram HTTP server separately as described in Quick Start. +The MCP process connects to that server; it does not start the database API. + +Configure Claude Code with the absolute path to the installed interpreter: + ```bash -claude mcp add engram -- python -m engram.integrations.mcp_server +claude mcp add --transport stdio \ + --env ENGRAM_API_URL=http://localhost:5820 \ + engram -- /absolute/path/to/engram/.venv/bin/python -m engram.integrations.mcp_server ``` -Or in `claude_desktop_config.json`: +For Claude Desktop, open **Settings → Developer → Edit Config** and merge the +`engram` entry into the existing `mcpServers` object. Keep one JSON root object: + ```json { "mcpServers": { "engram": { - "command": "python", + "command": "/absolute/path/to/engram/.venv/bin/python", "args": ["-m", "engram.integrations.mcp_server"], "env": { - "ENGRAM_API_URL": "http://localhost:5820", - "ENGRAM_API_KEY": "eng_sk_..." + "ENGRAM_API_URL": "http://localhost:5820" } } } } ``` -Tools: `engram_list_contexts`, `engram_recall`, `engram_commit`, `engram_decide`, `engram_health`, `engram_create_context` +Replace the executable path with your actual absolute path (on Windows, use the +venv's `Scripts\\python.exe` and escape backslashes in JSON). Restart Claude +Desktop after saving. For Cloud access through this bridge, set `ENGRAM_API_URL` +to `https://api.engram.so` and add `ENGRAM_API_KEY` with your Engram API key to +`env`. Self-hosted OSS has no account authentication and does not require a key. +The bridge forwards a configured key as an HTTP Bearer token. + +Tools: `engram_list_contexts`, `engram_recall`, `engram_commit`, `engram_decide`, +`engram_health`, `engram_create_context`; the original unprefixed tools remain +available. Ask Claude to list your contexts to verify an actual API call. ### OpenAI Function Calling @@ -880,14 +907,6 @@ All settings via environment variables or `.env` file: | `ENGRAM_VALIDITY_GATE_MODEL` | `claude-haiku-4-5` | Model used by the validity gate when enabled | | **Embedding** | | | | `ENGRAM_EMBEDDING_MODEL` | `text-embedding-3-small` | For bullet similarity search **and** worked-example retrieval over the activity ledger | -| **Consolidation** | | | -| `ENGRAM_CONSOLIDATION_TRIGGER` | `every_10_commits` | When to run consolidation | -| `ENGRAM_FAST_DECAY_RATE` | `0.97` | Per-day salience decay (normal bullets) | -| `ENGRAM_SLOW_DECAY_RATE` | `0.995` | Per-day decay (decisions/schemas) | -| **Auth** | | | -| `ENGRAM_AUTH_ENABLED` | `false` | Enable multi-tenant auth | -| `COGNITO_USER_POOL_ID` | — | AWS Cognito user pool | -| `COGNITO_APP_CLIENT_ID` | — | AWS Cognito app client | ## Development @@ -969,3 +988,4 @@ please email logos@engram.so (company + use case). For large-scale deployments (>20,000 external users), we offer paid support and an enterprise agreement — contact logos@engram.so. +Upgrade requirements and preserved interfaces are documented in [Compatibility and upgrade notes](docs/backward-compatibility.md). diff --git a/docs/backward-compatibility.md b/docs/backward-compatibility.md new file mode 100644 index 0000000..27ae13e --- /dev/null +++ b/docs/backward-compatibility.md @@ -0,0 +1,37 @@ +# Compatibility and upgrade notes + +The supported upgrade target is the existing REST, MCP and Python SDK API with +populated SQLite or PostgreSQL databases. The compatibility audit compares the +previous Git revision with the corrected source: all 34 existing HTTP operations, +40 schema definitions and 33 public SDK methods retain their existing fields, +enum values, documented responses and method signatures. All 353 existing public Python definitions remain present. Database upgrades add +columns and indexes without replacing account-independent context, memory, +activity or receipt IDs. No developer or customer database was upgraded during +validation. + +Built-in SQLite and PostgreSQL support explicit transaction-scoped storage. +Out-of-tree storage implementations must implement `StorageBackend.transaction` +before adopting these atomic mutation paths; unsupported implementations fail +clearly. This is an extension-interface requirement. Existing renderers and +completion adapters without a per-call model argument remain supported. OSS v0.5 +core memory, worked examples, validity checks and reconsolidation remain intact. + +Incorrect or unsafe behavior is intentionally corrected: cross-context writes +are rejected, active capacity is checked against the actual committed state, +failed mutations cannot leave partial audit/activity records, and unavailable or +invalid similarity comparisons cannot justify destructive matching. Configured +reflector models and deduplication thresholds are now honored. + +Historical audit rows that lack a generated target ID or required inverse state +cannot be reliably undone. Rollback refuses those batches before a partial write +rather than claiming success. An old content update without a saved embedding +clears the mismatched vector on rollback; it cannot recreate that historical +embedding. Old receipt rows cannot recover consumption history never recorded. +See [embedding-backfill.md](embedding-backfill.md) for the optional, explicit, +bounded repair of missing vectors using the correct historical embedding model. + +The distribution is `engram-contextdb`; the Python import and CLI remain `engram`. +MCP 1.x is the supported dependency range. Installed-package checks build from an +sdist, install the resulting wheel outside the checkout, run the CLI, and exercise +the stdio protocol. These checks do not publish a corrected package or validate +native third-party sign-in. All fixes remain local until released. diff --git a/docs/embedding-backfill.md b/docs/embedding-backfill.md new file mode 100644 index 0000000..073848d --- /dev/null +++ b/docs/embedding-backfill.md @@ -0,0 +1,38 @@ +# Repair missing embeddings + +Use the corrected source checkout to repair existing active memories that have no embedding. This does not repeat extraction, change text, or create new memories. The current database schema uses 1536-dimensional embeddings; choose the same embedding model used for the existing vectors. + +Preview one context in an existing SQLite database: + +```bash +.venv/bin/python -m engram.maintenance.backfill_embeddings \ + --sqlite /absolute/path/to/engram.db \ + --context-id CONTEXT_UUID --limit 100 +``` + +For PostgreSQL, put the existing database DSN in an environment variable and pass its **name**: + +```bash +.venv/bin/python -m engram.maintenance.backfill_embeddings \ + --postgres-dsn-env ENGRAM_MAINTENANCE_DSN \ + --context-id CONTEXT_UUID --limit 100 +``` + +Preview opens the existing database read-only. It does not initialize schemas, create the vector extension, change SQLite journal mode, load `.env`, or call an embedding provider. SQLite may create WAL coordination sidecars when opening an existing WAL database; its database content and journal mode remain unchanged. There is no default database target. A missing database or context fails instead of creating one. + +To save embeddings, add the explicit apply and provider options to either command: + +```bash + --apply --embedding-model text-embedding-3-small \ + --embedding-key-env ENGRAM_EMBEDDING_API_KEY +``` + +Set the named provider-key environment variable before running apply. Apply makes provider requests and may incur usage charges. The limit is between 1 and 1000 and bounds attempted memories per invocation; it defaults to 100. The command is an operator tool using database access, not a new unauthenticated HTTP endpoint. + +The JSON summary reports eligible, selected, attempted, updated, skipped, and failed counts. Exit status is 0 for a successful preview/apply, 1 for an operational or per-memory failure, and 2 for invalid command options. Errors identify a memory and a reason without printing provider messages, credentials, or memory text. + +Each vector is computed outside a database transaction. Before saving, the command rechecks context, content, lifecycle, and whether a vector is still missing. Concurrent edits or an already-completed repair are skipped. Only the embedding column changes; IDs, text, timestamps, metadata, counters, and audit history remain intact. Archived and inactive memories are excluded. Both backends validate the float32 representation used by PostgreSQL, rejecting empty, nonfinite, zero, underflow-to-zero, overflowing, and wrong-dimension vectors. Squared norms must also remain finite and nonzero in float32 so repaired vectors work with cosine search; vector scale is preserved. + +Repairs commit one memory at a time. If a later provider request fails, earlier successful repairs remain saved. Re-run the same command to repair remaining items; completed items are skipped. This tool cannot determine the model used for historical vectors, so model compatibility must be selected by the operator. + +Validation is isolated: `pytest tests/test_embedding_backfill.py` uses temporary SQLite databases and fake providers. Setting `ENGRAM_TEST_POSTGRES_DSN` to a disposable test admin database also exercises PostgreSQL in newly created random databases. No repair was run against an existing developer or customer database during implementation. diff --git a/docs/retrieval-evaluation.md b/docs/retrieval-evaluation.md new file mode 100644 index 0000000..ccd18fc --- /dev/null +++ b/docs/retrieval-evaluation.md @@ -0,0 +1,100 @@ +# Retrieval regression evaluation + +Run from a source checkout after installing the project dependencies: + +```sh +python -m evals.retrieval --output retrieval-results.json +pytest tests/test_retrieval_evaluation.py tests/test_materialization.py +``` + +The command prints JSON, optionally writes the same JSON file, and exits **1** +when any scenario fails. It creates and deletes a fresh temporary SQLite database. +It does not load `.env`, application configuration, or existing customer data. +Default embeddings are deterministic text-to-vector fixtures; no embedding or +completion provider is contacted. GPT uses its normal tiktoken estimator when +available (precache its tokenizer assets for a fully disconnected environment), +otherwise its documented character-count fallback. + +This is a **controlled regression benchmark**, not measured real-provider +retrieval quality. The committed `evals/corpus.json` contains curated queries, +texts, graded relevance judgments, synthetic vectors, and per-case thresholds. +Both editions share the corpus, runner and test contract. The engine, scoring, +SQLite vector search, renderer and persisted materialization receipts are real. +The OSS edition retains its MMR ranking and v0.5 features; paid ranking is unchanged. + +## Cases and interpretation + +Each of ten scenarios runs against Claude, GPT and generic renderers: + +| Case | What it checks | +| --- | --- | +| Semantic over salience | A relevant low-salience fact ranks ahead of an unrelated high-salience fact. | +| Irrelevant distractors | Two credential-management facts outrank unrelated facts. | +| Duplicate and diverse | Graded judgments cover a paraphrase and a distinct rollback fact. | +| Lifecycle exclusion | Archived and inactive facts never appear. | +| Tenant isolation | Another context's highly similar fact never appears. | +| Oversized candidate | A long leading candidate does not prevent packing a smaller useful fact. | +| Tight budget receipt | Complete rendered memories match receipt IDs within the budget. | +| Empty context | No IDs are invented; empty judgments do not inflate metric averages. | +| Legacy oversized | Legacy concepts follow the same packing guarantees. | +| Schema summary | Schema IDs match complete rendered summaries. | + +The JSON contains per-scenario/per-renderer `recall_at_k`, `ndcg_at_k`, packed +`ranked_ids`, `rendered_ids`, forbidden-ID leakage, estimated token usage, +receipt/budget consistency flags, and named failures. Recall counts unique +relevant IDs in the first k selections divided by all positively judged IDs. +nDCG uses graded gain `2^grade - 1` and logarithmic rank discount. Repeated IDs +earn credit once. Order is the packer's selection order; renderers may group +the resulting text by concept type. Irrelevant facts may appear after k; this +corpus does not claim that retrieval filters every irrelevant candidate. + +Recall must be 1.0 on every judged case. nDCG must be 1.0 except the duplicate/ +diverse case, whose 0.9 floor permits the existing edition-specific ranking +tradeoff. Every case requires zero forbidden leakage, consistent persisted +receipts, and output within its own renderer's token estimator. Empty judgments +have null ranking metrics and are excluded from averages. A report with no +judged results or missing results cannot pass. Aggregate means are descriptive; +one failed scenario still fails the gate. + +The tests deliberately degrade vectors, introduce a forbidden eligible fact, +and supply only empty judgments, then verify a nonzero CLI exit. They also check +metric arithmetic, duplicate credit, schema/legacy/bullet packing across small +budget boundaries, and OSS core memory, annotations and worked examples. +OSS-only tests skip on paid without changing its public API. + +## Packing behavior + +Selection accounts for the final renderer format, including headers, closing +tags, separators and annotations. Oversized candidates are skipped while +smaller successors are considered. Receipts record only selected complete +memories; duplicate text is tracked through IDs, not substring matching. +XML, Markdown and plain-text formats remain intact. + +Core memory takes priority, then intent, worked examples and ranked memories. +Oversized core memory or intent text is shortened with a truncation marker; +examples, schemas and facts are included whole or omitted. A budget too small +for even an empty format returns an empty string. Tokens are estimates, not a +guarantee about an external model's tokenizer; empty text is zero tokens. + +## Optional provider evaluation + +Provider use requires explicit opt-in, a model and a named credential variable: + +```sh +python -m evals.retrieval --live-embeddings \ + --embedding-model text-embedding-3-small --api-key-env OPENAI_API_KEY \ + --output live-retrieval-results.json +``` + +This makes billable embedding calls for distinct curated corpus texts and queries +once each, caches them for the run, and still uses only a new scratch database. +It never requests completions. Provider failures or invalid vectors fail the +command rather than silently measuring the engine's salience fallback. Results +are labeled `live_provider_labeled_corpus` and identify the requested embedding +model. The same small corpus and thresholds apply; a provider result is not a +general quality estimate or evidence about customer workloads. Live mode is +never part of the default test or CI command. + +The standalone live command prevents implicit SDK `.env` loading and suppresses +raw provider diagnostic output. Failures report only their exception type in +JSON, so credentials included in a provider exception cannot leak through logs. diff --git a/docs/shared-core-release.md b/docs/shared-core-release.md new file mode 100644 index 0000000..8cf051f --- /dev/null +++ b/docs/shared-core-release.md @@ -0,0 +1,126 @@ +# Shared-core port and release checks + +The OSS distribution (`engram-contextdb`, GitHub) and paid distribution (`engram`, +GitLab) share storage and retrieval contracts. Their complete core sources are not +identical: OSS v0.5 has reconsolidation deltas, rollback metadata, and an optional +validity gate; paid v0.4 has account/authentication, API-key, tenant, and hosted MCP +extensions. Port behavior deliberately instead of copying entire core directories. + +`shared-core-manifest.json` records a small set of byte-identical helpers/tests and +named regression contracts. It also enumerates intentionally distinct modules. +Hashes enforce recorded ports within each repository's CI. The two-checkout check +additionally verifies that selected files and manifests actually agree. Contract +names ensure regression cases are retained; passing names/hashes does **not** +replace executing both editions' tests. + +## Run locally + +Use an isolated environment for each edition and Python version (3.11 and 3.13). +Install `.[dev,mcp]` for OSS and `.[dev]` for paid, plus `build`. CI uses those same +extras, so absent MCP dependencies cannot turn the package gate into a skip. + +```bash +python -m pip install -e '.[dev,mcp]' build # OSS +# python -m pip install -e '.[dev]' build # paid +python scripts/check-shared-core.py --checkout . +python scripts/check-ci.py --backend sqlite --output reports/sqlite.xml +python -m evals.retrieval --output reports/retrieval.json +python scripts/check-package.py --output-dir reports/package +``` + +The suite wrapper runs from a temporary directory with a temporary pytest base, +ignoring checkout `.env` and inherited application/provider configuration. The +SQLite gate permits missing-DSN skips only for PostgreSQL cases. Both gates also +permit a narrow list of exact test-name/reason pairs for inapplicable scenarios +(independent `:memory:` databases, SQLite worker-queue tests on PostgreSQL, and +OSS-only extensions in paid). Each is printed separately. The package check +builds an sdist, builds a wheel from that sdist, installs the wheel into a fresh +venv outside the checkout, checks installed-module provenance, starts the actual +`engram` console script against temporary SQLite and calls `/health`, then runs +all six MCP stdio tests. Those tests include a real subprocess initialization, +tool discovery, and tool invocation against a loopback HTTP fixture. All six must +run with zero skips. No embedding/LLM provider is called. Retained wheel artifacts +are verification outputs; these workflows contain no publishing jobs. + +PostgreSQL verification requires a **disposable** pgvector server with a test role +allowed to create/drop databases and roles. Every database integration fixture +creates its own random test database; never supply a customer/production DSN. + +```bash +docker run --detach --rm --name engram-ci-pg \ + -e POSTGRES_DB=engram_test -e POSTGRES_USER=engram_test \ + -e POSTGRES_PASSWORD=engram-test-only -p 127.0.0.1:55432:5432 \ + pgvector/pgvector:pg16 +ENGRAM_TEST_POSTGRES_DSN='postgresql://engram_test:engram-test-only@127.0.0.1:55432/engram_test' \ + python scripts/check-ci.py --backend postgres --output reports/postgres.xml +docker stop engram-ci-pg +``` + +The wrapper fails immediately when the DSN is missing, waits for connectivity, +checks pgvector availability and test privileges, and rejects any unexpected skip, +including a missing-DSN or missing-MCP skip in the PostgreSQL run. It also requires the vector-roundtrip/restart case to +execute. A broken service cannot produce a green pipeline by skipping coverage. + +## Port discipline + +Run the following from either checkout, with OSS first and paid second: + +```bash +python scripts/check-shared-core.py /path/to/oss /path/to/paid +python scripts/check-shared-core.py --self-test /path/to/oss /path/to/paid +``` + +The self-test copies only selected files into disposable directories, introduces +real file drift and removes a required test name, and confirms both fail. It does +not edit either checkout. A single-repo CI run cannot prove the other repository +was updated; the two-checkout command is a required port/release check. + +For each shared behavior change: + +1. Record the source edition, source commit or patch identifier, counterpart + commit/patch, affected contract, and any intentionally different implementation. +2. Port the change and its regression cases. Add new shared helpers/corpus files + to `identical_files` only when both implementations are ready and intentionally + identical. Use a `contracts` entry for shared test names when implementation or + assertions must differ. Update the two manifests together. +3. Execute both editions on SQLite and PostgreSQL, both supported CI Python + versions, retrieval evaluations, and installed-wheel/MCP checks. Record actual + counts, versions, report paths, and failures; do not describe configured hosted + jobs as executed pipelines. +4. Review the diff, then refresh recorded hashes only after the port is complete: + + ```bash + python scripts/check-shared-core.py --refresh /path/to/oss /path/to/paid + python scripts/check-shared-core.py --self-test /path/to/oss /path/to/paid + ``` + +`--refresh` refuses missing contracts, different manifests, or unequal selected +files. It updates both manifests, never source files. Review hash changes as part +of the port; updating hashes is not evidence that tests passed. + +## Release evidence checklist + +Before approving a release, record the following in the release review: + +- Source revisions/patches and port provenance for both editions; two-checkout + shared-core result and deliberate-drift self-test result. +- Both Python versions' SQLite and PostgreSQL reports; pgvector/PostgreSQL version + and the fact that only disposable databases were used; any failed checks. +- Both controlled retrieval reports and corpus/threshold revision. These are + deterministic regression measurements, not real-provider quality claims. +- Both sdist-to-wheel artifacts, installed CLI/import checks, and six-case MCP + smoke reports; actual interpreter/dependency versions used. +- Migration/backward-compatibility review (nullable historical receipt markers, + transaction behavior, missing-embedding repair) and rollback/retry behavior. +- README/SDK/MCP examples and release snippets checked against the installed + package; edition names/versions and intentional differences checked. +- Hosted pipeline URLs and results **when actually run**; record public release + and deployment separately. Local equivalents do not establish + hosted runner/service availability or branch-protection settings. + +Workflow service configuration follows the official +[GitHub PostgreSQL service-container guide](https://docs.github.com/en/actions/tutorials/use-containerized-services/create-postgresql-service-containers) +and [GitLab PostgreSQL service guide](https://docs.gitlab.com/ci/services/postgres/). +GitHub runner jobs use the published loopback port; GitLab container jobs use the +`postgres` service alias. Both run the same readiness/suite wrapper and a Python +3.11/3.13 matrix. Jobs upload reports and package artifacts only. diff --git a/engram/core/consolidation.py b/engram/core/consolidation.py index 229de6f..fd763ab 100644 --- a/engram/core/consolidation.py +++ b/engram/core/consolidation.py @@ -15,10 +15,9 @@ import time import uuid from datetime import datetime, timezone -from typing import Any from engram.core.concurrency import ContextLockManager -from engram.core.delta import DeltaEngine +from engram.core.delta import DeltaEngine, capacity_checked_transaction from engram.core.events import EventBus from engram.core.models import ( Bullet, @@ -30,8 +29,9 @@ DeltaOpType, DeltaSource, LifecycleConfig, - SchemaNode, + LifecycleState, ) +from engram.core.similarity import validated_cosine_similarity from engram.llm.adapter import LLMAdapter from engram.storage.base import StorageBackend @@ -149,6 +149,27 @@ async def consolidate( return report + @staticmethod + async def _locked_active_bullets( + storage: StorageBackend, context_id: str, **filters + ) -> list[Bullet]: + candidates = await storage.list_bullets(context_id, include_archived=False, **filters) + current = [] + # Deterministic row-lock order avoids lock inversions across candidates. + for bullet_id in sorted({bullet.id for bullet in candidates}): + bullet = await storage.get_bullet(bullet_id) + if (bullet is None or bullet.context_id != context_id or not bullet.is_active + or bullet.is_archived or bullet.lifecycle_state != LifecycleState.ACTIVE): + continue + if filters.get("section") is not None and bullet.section != filters["section"]: + continue + if (filters.get("bullet_type") is not None + and bullet.bullet_type != filters["bullet_type"]): + continue + current.append(bullet) + current.sort(key=lambda bullet: (bullet.salience, bullet.created_at), reverse=True) + return current + async def _apply_forgetting_curve( self, context_id: str, config: ConsolidationConfig ) -> int: @@ -158,34 +179,35 @@ async def _apply_forgetting_curve( Concepts that ARE recalled get their decay reset (spaced repetition). DECISION and SCHEMA bullets have slower decay rates. """ - bullets = await self.storage.list_bullets(context_id, include_archived=False) - decayed_count = 0 - now = _utcnow() + async with self.storage.transaction(context_id) as tx: + bullets = await self._locked_active_bullets(tx, context_id) + decayed_count = 0 + now = _utcnow() - for bullet in bullets: - reference_time = bullet.last_recalled_at or bullet.created_at - days_since_active = max(0, (now - reference_time).total_seconds() / 86400) + for bullet in bullets: + reference_time = bullet.last_recalled_at or bullet.created_at + days_since_active = max(0, (now - reference_time).total_seconds() / 86400) - if days_since_active <= 1: - continue + if days_since_active <= 1: + continue - bt = bullet.bullet_type.value if hasattr(bullet.bullet_type, 'value') else str(bullet.bullet_type) - if bt in ("decision", "principle"): - decay_rate = config.slow_decay_rate - elif bt == "exception": - decay_rate = 0.99 - else: - decay_rate = config.fast_decay_rate + bt = bullet.bullet_type.value if hasattr(bullet.bullet_type, 'value') else str(bullet.bullet_type) + if bt in ("decision", "principle"): + decay_rate = config.slow_decay_rate + elif bt == "exception": + decay_rate = 0.99 + else: + decay_rate = config.fast_decay_rate - new_salience = bullet.salience * (decay_rate ** days_since_active) - new_salience = max(config.min_salience, new_salience) + new_salience = bullet.salience * (decay_rate ** days_since_active) + new_salience = max(config.min_salience, new_salience) - if abs(new_salience - bullet.salience) > 0.001: - bullet.salience = new_salience - await self.storage.update_bullet(bullet) - decayed_count += 1 + if abs(new_salience - bullet.salience) > 0.001: + bullet.salience = new_salience + await tx.update_bullet(bullet) + decayed_count += 1 - return decayed_count + return decayed_count async def _semantic_dedup( self, context_id: str, config: ConsolidationConfig @@ -196,30 +218,39 @@ async def _semantic_dedup( If cosine similarity > threshold (default 0.92), merge by keeping the more specific/higher-salience bullet. """ - bullets = await self.storage.list_bullets(context_id, include_archived=False) - merged_ids: set[str] = set() - merge_count = 0 - - for i, bullet_a in enumerate(bullets): - if bullet_a.id in merged_ids or bullet_a.embedding is None: - continue - for bullet_b in bullets[i + 1:]: - if bullet_b.id in merged_ids or bullet_b.embedding is None: + # Candidate reads and donor removals share the same context transaction. + # A failed write cannot leave usage counted in both donor and survivor. + async with self.storage.transaction(context_id) as tx: + bullets = await self._locked_active_bullets(tx, context_id) + merged_ids: set[str] = set() + operations: list[DeltaOperation] = [] + for i, bullet_a in enumerate(bullets): + if bullet_a.id in merged_ids or bullet_a.embedding is None: continue - sim = _cosine_similarity(bullet_a.embedding, bullet_b.embedding) - if sim >= config.dedup_threshold: - # Keep the one with higher salience - keep, remove = (bullet_a, bullet_b) if bullet_a.salience >= bullet_b.salience else (bullet_b, bullet_a) - keep.recall_count += remove.recall_count - keep.hit_count += remove.hit_count - keep.miss_count += remove.miss_count - keep.salience = max(keep.salience, remove.salience) - await self.storage.update_bullet(keep) - await self.storage.remove_bullet(remove.id) + for bullet_b in bullets[i + 1:]: + if bullet_b.id in merged_ids or bullet_b.embedding is None: + continue + sim = validated_cosine_similarity(bullet_a.embedding, bullet_b.embedding) + if sim is None or sim < config.dedup_threshold: + continue + keep, remove = ( + (bullet_a, bullet_b) if bullet_a.salience >= bullet_b.salience + else (bullet_b, bullet_a) + ) + operations.append(DeltaOperation( + op_type=DeltaOpType.MERGE_BULLETS, + target_ids=[keep.id, remove.id], + source=DeltaSource.CONSOLIDATION, + reasoning="Semantic deduplication during consolidation", + )) merged_ids.add(remove.id) - merge_count += 1 - - return merge_count + if remove.id == bullet_a.id: + break + if operations: + await DeltaEngine(tx).apply_batch(DeltaBatch( + context_id=context_id, operations=operations, trigger="consolidation", + )) + return len(operations) async def _induce_schemas( self, context_id: str, config: ConsolidationConfig @@ -255,28 +286,42 @@ async def _induce_schemas( + "\n".join(f"- {c}" for c in bullet_contents) + "\n\nDescribe the abstract pattern they share in 1-2 sentences." ), - system="You are a knowledge organizer. Respond with a concise description only.", + system=( + "You are a knowledge organizer. Respond with a concise description only." + ), temperature=0.0, max_tokens=200, ) except Exception: description = f"Pattern in {section} ({len(section_list)} bullets)" - schema = SchemaNode( - name=section, - description=description.strip(), - instance_count=len(section_list), - confidence=min(1.0, len(section_list) / 10), - bullet_ids=[b.id for b in section_list], - ) - await self.storage.add_schema(context_id, schema) - - # Link bullets to schema - for bullet in section_list: - bullet.schema_id = schema.id - await self.storage.update_bullet(bullet) - - schemas_formed += 1 + expected_contents = {bullet.id: bullet.content for bullet in section_list} + async with self.storage.transaction(context_id) as tx: + if any(schema.name == section for schema in await tx.list_schemas(context_id)): + continue + current = await self._locked_active_bullets(tx, context_id, section=section) + eligible = [bullet for bullet in current if bullet.schema_id is None + and expected_contents.get(bullet.id) == bullet.content] + if len(eligible) < config.schema_min_instances: + continue + schema_id = str(uuid.uuid4())[:8] + await DeltaEngine(tx).apply_batch(DeltaBatch( + context_id=context_id, trigger="consolidation_schema", + operations=[DeltaOperation( + op_type=DeltaOpType.ADD_SCHEMA, target_id=schema_id, + content=section, reasoning=description.strip(), + source=DeltaSource.CONSOLIDATION, + )], + )) + schema = await tx.get_schema(schema_id) + schema.instance_count = len(eligible) + schema.confidence = min(1.0, len(eligible) / 10) + schema.bullet_ids = [bullet.id for bullet in eligible] + await tx.update_schema(schema) + for bullet in eligible: + bullet.schema_id = schema.id + await tx.update_bullet(bullet) + schemas_formed += 1 return schemas_formed @@ -290,30 +335,31 @@ async def _archive_stale( in archive_days_threshold+ days AND not a DECISION or PRINCIPLE type. v0.3: Uses storage.archive_bullet() for proper lifecycle state transitions. """ - bullets = await self.storage.list_bullets(context_id, include_archived=False) - archived_count = 0 - now = _utcnow() - - for bullet in bullets: - bt = bullet.bullet_type.value if hasattr(bullet.bullet_type, 'value') else str(bullet.bullet_type) - if bt in ("decision", "principle"): - continue + async with self.storage.transaction(context_id) as tx: + bullets = await self._locked_active_bullets(tx, context_id) + archived_count = 0 + now = _utcnow() + + for bullet in bullets: + bt = bullet.bullet_type.value if hasattr(bullet.bullet_type, 'value') else str(bullet.bullet_type) + if bt in ("decision", "principle"): + continue - if bullet.salience >= config.archive_salience_threshold: - continue + if bullet.salience >= config.archive_salience_threshold: + continue - reference_time = bullet.last_recalled_at or bullet.created_at - days_inactive = (now - reference_time).total_seconds() / 86400 - if days_inactive < config.archive_days_threshold: - continue + reference_time = bullet.last_recalled_at or bullet.created_at + days_inactive = (now - reference_time).total_seconds() / 86400 + if days_inactive < config.archive_days_threshold: + continue - success = await self.storage.archive_bullet( - context_id, bullet.id, reason="consolidation_stale" - ) - if success: - archived_count += 1 + success = await tx.archive_bullet( + context_id, bullet.id, reason="consolidation_stale" + ) + if success: + archived_count += 1 - return archived_count + return archived_count async def _promote_to_principles( self, context_id: str, config: ConsolidationConfig @@ -364,17 +410,35 @@ async def _promote_to_principles( except Exception: continue - # Create principle bullet - principle_bullet = Bullet( - section=section, - content=principle.strip(), - bullet_type=BulletType.PRINCIPLE, - source_type="consolidation", - salience=0.8, - confidence=min(1.0, len(facts) / 10), - ) - await self.storage.add_bullet(context_id, principle_bullet) - promoted += 1 + expected_contents = {bullet.id: bullet.content for bullet in facts} + async with capacity_checked_transaction(self.storage, context_id) as tx: + current = await self._locked_active_bullets(tx, context_id, section=section) + if any(bullet.bullet_type == BulletType.PRINCIPLE for bullet in current): + continue + eligible = [bullet for bullet in current + if bullet.bullet_type == BulletType.FACT and bullet.hit_count >= 2 + and expected_contents.get(bullet.id) == bullet.content] + if len(eligible) < config.schema_min_instances or not principle.strip(): + continue + context = await tx.get_context(uuid.UUID(context_id)) + if context is not None: + maximum = context.lifecycle_config.max_active_bullets + capacity = await tx.get_capacity_status(context_id, maximum) + if capacity.active_bullet_count >= maximum: + continue + bullet_id = str(uuid.uuid4())[:8] + await DeltaEngine(tx).apply_batch(DeltaBatch( + context_id=context_id, trigger="consolidation_promotion", + operations=[DeltaOperation( + op_type=DeltaOpType.ADD_BULLET, target_id=bullet_id, + section=section, content=principle.strip(), bullet_type="principle", + source=DeltaSource.CONSOLIDATION, confidence=min(1.0, len(eligible) / 10), + )], + )) + principle_bullet = await tx.get_bullet(bullet_id) + principle_bullet.salience = 0.8 + await tx.update_bullet(principle_bullet) + promoted += 1 return promoted diff --git a/engram/core/delta.py b/engram/core/delta.py index a60b82a..6e58dbe 100644 --- a/engram/core/delta.py +++ b/engram/core/delta.py @@ -9,9 +9,11 @@ import logging import uuid +from contextlib import asynccontextmanager from datetime import datetime, timezone from typing import Any +from engram.core.exceptions import CapacityExceededError from engram.core.models import ( ActionType, Activity, @@ -21,6 +23,7 @@ DeltaOperation, DeltaOpType, DeltaSource, + LifecycleState, SchemaNode, SourceType, cap_core_memory, @@ -34,6 +37,32 @@ def _utcnow() -> datetime: return datetime.now(timezone.utc) +@asynccontextmanager +async def capacity_checked_transaction(storage: StorageBackend, context_id: str): + """Atomically validate actual capacity around a public mutation and ledger.""" + async with storage.transaction(context_id) as tx: + try: + context_uuid = uuid.UUID(context_id) + except (ValueError, AttributeError): + context = None + else: + context = await tx.get_context(context_uuid) + before_count = 0 + max_bullets = None + if context is not None: + max_bullets = context.lifecycle_config.max_active_bullets + before = await tx.get_capacity_status(context_id, max_bullets) + before_count = before.active_bullet_count + yield tx + if max_bullets is not None: + after = await tx.get_capacity_status(context_id, max_bullets) + # Existing overfull contexts may still be corrected or reduced. + # Missing/duplicate removes do not create fictitious capacity. + if (after.active_bullet_count > max_bullets + and after.active_bullet_count > before_count): + raise CapacityExceededError(context_id, after.active_bullet_count, max_bullets) + + class DeltaEngine: """Applies delta operations atomically to the concept graph. @@ -45,6 +74,35 @@ def __init__(self, storage: StorageBackend) -> None: self.storage = storage async def apply_batch(self, batch: DeltaBatch) -> DeltaBatch: + # Work on a copy so a failed attempt cannot poison retry snapshots/ids. + working = batch.model_copy(deep=True) + async with self.storage.transaction(batch.context_id) as tx: + result = await DeltaEngine(tx)._apply_batch(working) + # Successful calls historically update the caller's batch and operation + # objects. Publish only after commit so failures remain safe to retry. + for original, applied in zip(batch.operations, result.operations): + for field in type(applied).model_fields: + setattr(original, field, getattr(applied, field)) + for field in type(result).model_fields: + if field != "operations": + setattr(batch, field, getattr(result, field)) + return batch + + async def _validate_operation_context(self, context_id: str, op: DeltaOperation) -> None: + """A batch lock is not authorization to write globally addressed rows.""" + if op.op_type in (DeltaOpType.ADD_SCHEMA, DeltaOpType.UPDATE_SCHEMA): + schema = await self.storage.get_schema(op.target_id) if op.target_id else None + if schema is not None and schema.context_id != context_id: + raise ValueError("Delta target does not belong to the batch context") + return + targets = op.target_ids if op.op_type == DeltaOpType.MERGE_BULLETS else [op.target_id] + for target in dict.fromkeys(targets or []): + if target: + bullet = await self.storage.get_bullet(target) + if bullet is not None and bullet.context_id != context_id: + raise ValueError("Delta target does not belong to the batch context") + + async def _apply_batch(self, batch: DeltaBatch) -> DeltaBatch: """Apply a batch of delta operations atomically. Returns the batch with updated stats. @@ -55,6 +113,7 @@ async def apply_batch(self, batch: DeltaBatch) -> DeltaBatch: merged = 0 for op in batch.operations: + await self._validate_operation_context(batch.context_id, op) match op.op_type: case DeltaOpType.ADD_BULLET: await self._apply_add_bullet(batch.context_id, op) @@ -98,31 +157,89 @@ async def apply_batch(self, batch: DeltaBatch) -> DeltaBatch: return batch async def rollback_batch(self, delta_batch_id: str) -> bool: + batch = await self.storage.get_delta_batch(delta_batch_id) + if batch is None: + return False + async with self.storage.transaction(batch.context_id) as tx: + return await DeltaEngine(tx)._rollback_batch(delta_batch_id) + + async def _rollback_batch(self, delta_batch_id: str) -> bool: """Roll back a delta batch by applying inverse operations.""" batch = await self.storage.get_delta_batch(delta_batch_id) if batch is None: return False + # Preflight every inverse before any write. Historical records may + # lack generated identities or snapshots; partial rollback is unsafe. + for op in batch.operations: + await self._validate_operation_context(batch.context_id, op) + snapshot = op.rollback_state or op.previous_state + if snapshot is not None and snapshot.get("noop") is True: + continue + if op.op_type == DeltaOpType.ADD_BULLET: + if not op.target_id: + return False + elif op.op_type == DeltaOpType.UPDATE_BULLET: + if not op.target_id or not snapshot or not any( + field in snapshot + for field in ("content", "embedding", "salience", "confidence", "section") + ): + return False + elif op.op_type == DeltaOpType.REMOVE_BULLET: + if not op.target_id or snapshot is None or "is_active" not in snapshot: + return False + elif op.op_type == DeltaOpType.MERGE_BULLETS: + if snapshot is None or not isinstance(snapshot.get("bullets"), dict): + return False + elif op.op_type == DeltaOpType.ADD_SCHEMA: + if not op.target_id or not snapshot or snapshot.get("schema_created") is not True: + return False + elif op.op_type == DeltaOpType.UPDATE_SCHEMA: + if not op.target_id or snapshot is None or "description" not in snapshot: + return False + elif op.op_type == DeltaOpType.UPDATE_CORE_MEMORY: + if snapshot is None or "core_memory" not in snapshot: + return False + elif op.op_type == DeltaOpType.RECONSOLIDATE_BULLET: + if not op.target_id or snapshot is None or not all( + field in snapshot for field in ("recall_count", "hit_count", "miss_count", "salience") + ): + return False + # Apply inverse operations in reverse order. Newer ops use rollback_state # for the snapshot; older records (pre-rollback_state) fall back to # previous_state for backward compatibility. for op in reversed(batch.operations): + await self._validate_operation_context(batch.context_id, op) + snapshot = op.rollback_state or op.previous_state + if snapshot is not None and snapshot.get("noop") is True: + continue + # ADD has no previous state. Its persisted target identifies the + # new bullet even for callers that did not provide an id. + if op.op_type == DeltaOpType.ADD_BULLET: + if op.target_id: + await self.storage.remove_bullet(op.target_id) + continue + if op.op_type == DeltaOpType.ADD_SCHEMA: + if op.content and op.target_id: + await self.storage.remove_schema(batch.context_id, op.target_id) + continue snapshot = op.rollback_state or op.previous_state if snapshot is None: continue match op.op_type: - case DeltaOpType.ADD_BULLET: - # Undo add → remove - if op.target_id: - await self.storage.remove_bullet(op.target_id) - case DeltaOpType.UPDATE_BULLET: # Undo update → restore previous state if op.target_id: bullet = await self.storage.get_bullet(op.target_id) if bullet: - bullet.content = snapshot.get("content", bullet.content) + restored_content = snapshot.get("content", bullet.content) + if "embedding" in snapshot or restored_content != bullet.content: + # Legacy snapshots lack vectors; clear the current + # one if it describes the text being undone. + bullet.embedding = snapshot.get("embedding") + bullet.content = restored_content bullet.salience = snapshot.get("salience", bullet.salience) bullet.confidence = snapshot.get("confidence", bullet.confidence) bullet.section = snapshot.get("section", bullet.section) @@ -133,7 +250,7 @@ async def rollback_batch(self, delta_batch_id: str) -> bool: if op.target_id: bullet = await self.storage.get_bullet(op.target_id) if bullet: - bullet.is_active = True + bullet.is_active = snapshot.get("is_active", True) await self.storage.update_bullet(bullet) case DeltaOpType.UPDATE_CORE_MEMORY: @@ -154,6 +271,28 @@ async def rollback_batch(self, delta_batch_id: str) -> bool: ) await self.storage.update_bullet(bullet) + case DeltaOpType.UPDATE_SCHEMA: + if op.target_id and "description" in snapshot: + schema = await self.storage.get_schema(op.target_id) + if schema is not None: + schema.description = snapshot["description"] + await self.storage.update_schema(schema) + + case DeltaOpType.MERGE_BULLETS: + for bullet_id, previous in snapshot.get("bullets", {}).items(): + bullet = await self.storage.get_bullet(bullet_id) + if bullet is None: + continue + if bullet.context_id != batch.context_id: + raise ValueError("Delta snapshot does not belong to the batch context") + for field in ( + "content", "embedding", "recall_count", "hit_count", + "miss_count", "salience", "is_active", + ): + if field in previous: + setattr(bullet, field, previous[field]) + await self.storage.update_bullet(bullet) + return True async def _apply_add_bullet(self, context_id: str, op: DeltaOperation) -> None: @@ -166,10 +305,12 @@ async def _apply_add_bullet(self, context_id: str, op: DeltaOperation) -> None: } source_type = source_map.get(op.source, SourceType.REFLECTION) if op.source else SourceType.REFLECTION + op.target_id = op.target_id or str(uuid.uuid4())[:8] bullet = Bullet( - id=op.target_id or str(uuid.uuid4())[:8], + id=op.target_id, section=op.section or "general", content=op.content or "", + embedding=op.embedding, bullet_type=BulletType(op.bullet_type) if op.bullet_type else BulletType.FACT, source_type=source_type, salience=op.confidence, @@ -181,23 +322,30 @@ async def _apply_add_bullet(self, context_id: str, op: DeltaOperation) -> None: async def _apply_update_bullet(self, op: DeltaOperation) -> None: if not op.target_id: + op.rollback_state = {"noop": True} return bullet = await self.storage.get_bullet(op.target_id) if bullet is None: + op.rollback_state = {"noop": True} return - # Snapshot for rollback. Only captured once — if apply runs twice - # (retry/idempotency), the original pre-apply state is preserved. - if op.rollback_state is None: - op.rollback_state = { - "content": bullet.content, - "salience": bullet.salience, - "confidence": bullet.confidence, - "section": bullet.section, - } + # Each committed application needs its own inverse. Failed attempts + # operate on a copy, so they cannot overwrite the caller's snapshot. + op.rollback_state = { + "content": bullet.content, + "embedding": bullet.embedding, + "salience": bullet.salience, + "confidence": bullet.confidence, + "section": bullet.section, + } if op.content is not None: + if op.content != bullet.content: + # A missing new embedding must not leave an old-content vector. + bullet.embedding = op.embedding bullet.content = op.content + if op.embedding is not None: + bullet.embedding = op.embedding if op.section is not None: bullet.section = op.section if op.confidence is not None: @@ -206,30 +354,50 @@ async def _apply_update_bullet(self, op: DeltaOperation) -> None: async def _apply_remove_bullet(self, op: DeltaOperation) -> None: if not op.target_id: + op.rollback_state = {"noop": True} return bullet = await self.storage.get_bullet(op.target_id) if bullet: - if op.rollback_state is None: - op.rollback_state = {"is_active": bullet.is_active} + op.rollback_state = {"is_active": bullet.is_active} await self.storage.remove_bullet(op.target_id) + else: + op.rollback_state = {"noop": True} async def _apply_merge_bullets(self, context_id: str, op: DeltaOperation) -> None: """Merge multiple bullets into one — keep the most specific/highest salience.""" + # Mark even a no-op merge as known; older records lack inverse evidence. + op.rollback_state = {"bullets": {}} if not op.target_ids or len(op.target_ids) < 2: return bullets_to_merge = [] - for bid in op.target_ids: + for bid in dict.fromkeys(op.target_ids): bullet = await self.storage.get_bullet(bid) - if bullet: + if (bullet is not None and bullet.is_active and not bullet.is_archived + and bullet.lifecycle_state == LifecycleState.ACTIVE): bullets_to_merge.append(bullet) if len(bullets_to_merge) < 2: return + op.rollback_state = { + "bullets": { + bullet.id: { + field: getattr(bullet, field) + for field in ( + "content", "embedding", "recall_count", "hit_count", + "miss_count", "salience", "is_active", + ) + } + for bullet in bullets_to_merge + } + } + # Keep the bullet with highest salience, deactivate others best = max(bullets_to_merge, key=lambda b: b.salience) if op.content: + if op.content != best.content: + best.embedding = op.embedding best.content = op.content best.recall_count = sum(b.recall_count for b in bullets_to_merge) best.hit_count = sum(b.hit_count for b in bullets_to_merge) @@ -243,20 +411,28 @@ async def _apply_merge_bullets(self, context_id: str, op: DeltaOperation) -> Non async def _apply_add_schema(self, context_id: str, op: DeltaOperation) -> None: if not op.content: + op.rollback_state = {"noop": True} return + op.target_id = op.target_id or str(uuid.uuid4())[:8] schema = SchemaNode( + id=op.target_id, name=op.content, description=op.reasoning, ) await self.storage.add_schema(context_id, schema) + op.rollback_state = {"schema_created": True} async def _apply_update_schema(self, op: DeltaOperation) -> None: if not op.target_id: + op.rollback_state = {"noop": True} return schema = await self.storage.get_schema(op.target_id) if schema and op.content: + op.rollback_state = {"description": schema.description} schema.description = op.content await self.storage.update_schema(schema) + else: + op.rollback_state = {"noop": True} async def _apply_update_core_memory( self, context_id: str, op: DeltaOperation, @@ -267,6 +443,7 @@ async def _apply_update_core_memory( rollback_state so a retry doesn't overwrite the original snapshot. """ if op.content is None: + op.rollback_state = {"noop": True} return # Enforce the ≤512-token bound here, at the canonical mutation point, so # it holds for every caller of the delta op — not just the ingestion @@ -276,8 +453,10 @@ async def _apply_update_core_memory( ctx = await self.storage.get_context(uuid.UUID(context_id)) except (ValueError, AttributeError): ctx = None - if ctx is not None and op.rollback_state is None: - op.rollback_state = {"core_memory": ctx.core_memory} + if ctx is None: + op.rollback_state = {"noop": True} + return + op.rollback_state = {"core_memory": ctx.core_memory} await self.storage.update_core_memory(context_id, capped) async def _apply_reconsolidate_bullet(self, op: DeltaOperation) -> None: @@ -290,9 +469,11 @@ async def _apply_reconsolidate_bullet(self, op: DeltaOperation) -> None: a retry/idempotent re-apply reads the same deltas, not the snapshot. """ if not op.target_id or op.previous_state is None: + op.rollback_state = {"noop": True} return bullet = await self.storage.get_bullet(op.target_id) if bullet is None: + op.rollback_state = {"noop": True} return # Don't reinforce a bullet that was archived or deactivated between the # materialization and this feedback commit — get_bullet returns rows @@ -302,17 +483,17 @@ async def _apply_reconsolidate_bullet(self, op: DeltaOperation) -> None: "Skipping reconsolidation of inactive/archived bullet %s", op.target_id, ) + op.rollback_state = {"noop": True} return deltas = op.previous_state - if op.rollback_state is None: - op.rollback_state = { - "recall_count": bullet.recall_count, - "hit_count": bullet.hit_count, - "miss_count": bullet.miss_count, - "salience": bullet.salience, - "last_recalled_at": bullet.last_recalled_at.isoformat() - if bullet.last_recalled_at else None, - } + op.rollback_state = { + "recall_count": bullet.recall_count, + "hit_count": bullet.hit_count, + "miss_count": bullet.miss_count, + "salience": bullet.salience, + "last_recalled_at": bullet.last_recalled_at.isoformat() + if bullet.last_recalled_at else None, + } bullet.recall_count += int(deltas.get("recall_delta", 0)) bullet.hit_count += int(deltas.get("hit_delta", 0)) bullet.miss_count += int(deltas.get("miss_delta", 0)) diff --git a/engram/core/graph.py b/engram/core/graph.py index bb0e4d8..2c28255 100644 --- a/engram/core/graph.py +++ b/engram/core/graph.py @@ -32,7 +32,7 @@ async def invalidate_concepts( for cid in concept_ids: concept = await self.storage.get_concept(cid) - if concept is None or not concept.is_valid: + if concept is None or not concept.is_valid or concept.context_id != context_id: continue concept.is_valid = False concept.invalidated_at = now @@ -70,8 +70,9 @@ async def get_concept_neighborhood( continue visited.add(nid) concept = await self.storage.get_concept(nid) - if concept and concept.is_valid: - result.append(concept) + if concept is None or not concept.is_valid or concept.context_id != context_id: + continue + result.append(concept) edges = await self.storage.get_edges(context_id, node_id=nid) for edge in edges: next_frontier.add(edge.from_node) diff --git a/engram/core/ingestion.py b/engram/core/ingestion.py index e8f5654..f073539 100644 --- a/engram/core/ingestion.py +++ b/engram/core/ingestion.py @@ -13,13 +13,14 @@ import hashlib import json import logging +import math import uuid from datetime import datetime, timezone from typing import Any from engram.core.concurrency import ContextLockManager from engram.core.config import IngestionConfig -from engram.core.delta import DeltaEngine +from engram.core.delta import DeltaEngine, capacity_checked_transaction from engram.core.events import EventBus from engram.core.exceptions import CapacityExceededError, IngestionError from engram.core.models import ( @@ -39,12 +40,22 @@ SourceType, cap_core_memory, ) -from engram.llm.adapter import LLMAdapter +from engram.core.similarity import validated_cosine_similarity +from engram.llm.adapter import LLMAdapter, complete_with_model from engram.storage.base import StorageBackend logger = logging.getLogger(__name__) +async def _embed_content(llm: LLMAdapter, content: str) -> list[float] | None: + """Embed new memory text without making an unavailable provider block writes.""" + try: + return await llm.embed(content) + except Exception as exc: + logger.warning("Bullet embedding failed: %s", exc) + return None + + def _utcnow() -> datetime: return datetime.now(timezone.utc) @@ -160,12 +171,13 @@ async def reflect( f"{context_text}{core_text}{feedback_text}\n\nRaw input:\n{raw_input}" ) - raw_response = await self.llm.complete( + raw_response = await complete_with_model( + self.llm, + model=model_override or self.config.reflector_model, prompt=prompt, system=REFLECTOR_SYSTEM_PROMPT, temperature=0.0, response_format="json", - model=model_override, ) try: @@ -270,6 +282,7 @@ async def curate( target_id=str(uuid.uuid4())[:8], section="prediction_errors", content=error, + embedding=await _embed_content(self.llm, error), bullet_type="exception", reasoning="Prediction error detected by Reflector — high salience signal", source=DeltaSource.REFLECTOR, @@ -285,6 +298,7 @@ async def curate( target_id=str(uuid.uuid4())[:8], section="failure_modes", content=failure, + embedding=await _embed_content(self.llm, failure), bullet_type="warning", reasoning="Failure mode identified by Reflector", source=DeltaSource.REFLECTOR, @@ -300,6 +314,7 @@ async def curate( target_id=str(uuid.uuid4())[:8], section="strategies", content=strategy, + embedding=await _embed_content(self.llm, strategy), bullet_type="strategy", reasoning="Successful strategy recorded by Reflector", source=DeltaSource.REFLECTOR, @@ -430,7 +445,7 @@ async def _process_insight( embedding: list[float] | None = None try: embedding = await self.llm.embed(insight.content) - threshold = 0.85 if is_episodic else 0.92 + threshold = 0.85 if is_episodic else self.ingestion_config.curator_dedup_threshold similar = await self.storage.find_similar_bullets( context_id, embedding, limit=5, threshold=threshold, ) @@ -440,6 +455,9 @@ async def _process_insight( if (b.bullet_type.value if hasattr(b.bullet_type, "value") else str(b.bullet_type)) == "episodic" ] + similar = [(bullet, score) for bullet, score in similar + if math.isfinite(score) + and validated_cosine_similarity(embedding, bullet.embedding) is not None] if similar: existing_bullet, score = similar[0] # Keep the more specific (longer) phrasing of the two. @@ -453,6 +471,8 @@ async def _process_insight( op_type=DeltaOpType.UPDATE_BULLET, target_id=existing_bullet.id, content=merged_content, + embedding=(embedding if merged_content == insight.content + else existing_bullet.embedding), reasoning=f"Merged similar insight into existing bullet (cosine={score:.3f})", source=DeltaSource.CURATOR, confidence=max(existing_bullet.confidence, insight.novelty), @@ -481,6 +501,7 @@ async def _process_insight( target_id=str(uuid.uuid4())[:8], section=insight.suggested_section, content=insight.content, + embedding=embedding, bullet_type=insight.insight_type if insight.insight_type in [bt.value for bt in BulletType] else "fact", reasoning=( f"Contradicts existing bullet {existing_bullet.id}: " @@ -506,6 +527,7 @@ async def _process_insight( target_id=str(uuid.uuid4())[:8], section=insight.suggested_section, content=insight.content, + embedding=embedding, bullet_type=bullet_type, reasoning=f"New insight (novelty={insight.novelty:.2f}): {insight.evidence[:80]}", source=DeltaSource.CURATOR, @@ -545,14 +567,13 @@ def _detect_contradiction( 1. Find semantically similar bullets (same topic) 2. Check for negation signals (opposing conclusions) """ - from engram.core.similarity import cosine_similarity as _cosine_similarity - for bullet in existing_bullets: if bullet.embedding is None: continue - sim = _cosine_similarity(new_embedding, bullet.embedding) + sim = validated_cosine_similarity(new_embedding, bullet.embedding) # High similarity + negation signal = contradiction - if sim >= 0.75 and self._has_negation_signal(new_content, bullet.content): + if (sim is not None and sim >= 0.75 + and self._has_negation_signal(new_content, bullet.content)): return (bullet, sim) return None @@ -573,12 +594,24 @@ async def curate_re_extraction( This is conservative by default — high-value bullets (high salience + hit rate) are kept even if the new model doesn't reproduce them. """ - from engram.core.similarity import cosine_similarity as _cosine_similarity - operations: list[DeltaOperation] = [] matched_old_ids: set[str] = set() + def valid_vector(vector) -> bool: + return validated_cosine_similarity(vector, vector) is not None + + uncertain_old_ids = { + bullet.id for bullet in old_bullets if not valid_vector(bullet.embedding) + } + for insight in new_reflection.new_insights: + # Re-extraction must preserve identity even when an old vector is + # missing/corrupt or the embedding provider is unavailable. + exact = [bullet for bullet in old_bullets + if bullet.content.strip().casefold() == insight.content.strip().casefold()] + if exact: + matched_old_ids.update(bullet.id for bullet in exact) + continue best_match = None best_similarity = 0.0 @@ -587,16 +620,23 @@ async def curate_re_extraction( except Exception: new_embedding = None - if new_embedding is not None: + if valid_vector(new_embedding): for old_bullet in old_bullets: if old_bullet.id in matched_old_ids: continue - if old_bullet.embedding is None: + if (not valid_vector(old_bullet.embedding) + or len(new_embedding) != len(old_bullet.embedding)): + uncertain_old_ids.add(old_bullet.id) + continue + sim = validated_cosine_similarity(new_embedding, old_bullet.embedding) + if sim is None: + uncertain_old_ids.add(old_bullet.id) continue - sim = _cosine_similarity(new_embedding, old_bullet.embedding) if sim > best_similarity: best_similarity = sim best_match = old_bullet + else: + uncertain_old_ids.update(bullet.id for bullet in old_bullets) if best_match and best_similarity > 0.92: # Match found — minor update if content differs meaningfully @@ -606,6 +646,7 @@ async def curate_re_extraction( op_type=DeltaOpType.UPDATE_BULLET, target_id=best_match.id, content=insight.content, + embedding=new_embedding, reasoning=( f"Re-extraction improved this bullet " f"(similarity: {best_similarity:.2f})" @@ -625,6 +666,7 @@ async def curate_re_extraction( target_id=str(uuid.uuid4())[:8], section=insight.suggested_section, content=insight.content, + embedding=new_embedding, bullet_type=bullet_type, reasoning=( f"New insight discovered by re-extraction " @@ -635,7 +677,7 @@ async def curate_re_extraction( # Old bullets with no match in new extraction — consider removal for old_bullet in old_bullets: - if old_bullet.id not in matched_old_ids: + if old_bullet.id not in matched_old_ids and old_bullet.id not in uncertain_old_ids: # Only remove if the old bullet has low salience and low hit rate # High-value bullets are kept even if re-extraction doesn't reproduce them if old_bullet.effective_salience < 0.3 and old_bullet.hit_rate < 0.3: @@ -719,6 +761,13 @@ async def commit( "Duplicate raw input detected (hash=%s), returning existing batch", input_hash[:16], ) + # Ingestion may have committed before feedback failed. Retry + # the idempotent receipt transaction without another extraction. + if materialization_id and feedback: + await self._reconsolidate( + materialization_id, feedback, expected_context_id=ctx_id_str, + agent_id=agent_id, session_id=session_id, + ) return existing_batch # Build context summary + load core memory for the Reflector @@ -759,39 +808,6 @@ async def commit( session_id=str(session_id) if session_id else None, )) - # v0.3: Capacity check before applying deltas - net_adds = sum( - 1 for op in batch.operations if op.op_type == DeltaOpType.ADD_BULLET - ) - sum( - 1 for op in batch.operations if op.op_type == DeltaOpType.REMOVE_BULLET - ) - if net_adds > 0: - context = await self.storage.get_context(context_id) - if context is not None: - max_bullets = context.lifecycle_config.max_active_bullets - capacity = await self.storage.get_capacity_status( - ctx_id_str, max_bullets - ) - if capacity.pressure_level == "full": - raise CapacityExceededError( - ctx_id_str, capacity.active_bullet_count, max_bullets - ) - - # Phase 3: Apply deltas (under lock if available) - if self.lock_manager is not None: - async with self.lock_manager.acquire(ctx_id_str): - self._revalidate_deltas(batch) - batch = await self.delta_engine.apply_batch(batch) - else: - batch = await self.delta_engine.apply_batch(batch) - - # v0.4: Collect bullet IDs produced by ADD_BULLET operations - bullet_ids_produced = [ - op.target_id - for op in batch.operations - if op.op_type == DeltaOpType.ADD_BULLET and op.target_id - ] - # v0.5: embed raw input so future materializations can retrieve it as a # DC-style worked example. Best-effort — recall path tolerates absence. raw_input_embedding: list[float] | None = None @@ -800,42 +816,70 @@ async def commit( except Exception as exc: logger.warning("Raw-input embedding failed (worked-example retrieval disabled for this commit): %s", exc) - # Record activity WITH raw input preservation (v0.4) + embedding (v0.5) - activity = Activity( - agent_id=agent_id, - session_id=session_id, - action_type=ActionType.FACT_LEARNED, - summary=( - f"Ingested {content_type.value}: +{batch.bullets_added} bullets, " - f"~{batch.bullets_updated} updated, -{batch.bullets_removed} removed, " - f"⊕{batch.bullets_merged} merged" - ), - delta_batch_id=batch.id, - # v0.4: Raw input preservation - raw_input=content, - raw_input_hash=input_hash, - content_type=content_type.value, - source_agent_model=source_model, - feedback=feedback.model_dump() if feedback else None, - # v0.4: Extraction metadata - extraction_model=self.ingestion_config.reflector_model, - extraction_prompt_version=self.ingestion_config.reflector_prompt_version, - bullet_ids_produced=bullet_ids_produced, - # v0.5: worked-example retrieval - raw_input_embedding=raw_input_embedding, - ) - await self.storage.add_activity(context_id, activity) + async def persist_commit(batch: DeltaBatch) -> tuple[DeltaBatch, bool]: + async with capacity_checked_transaction(self.storage, ctx_id_str) as tx: + # Another worker may have committed this input while reflection + # ran. Serialize the final hash check with graph + ledger writes. + prior = await tx.get_raw_input_by_hash(ctx_id_str, input_hash) + if prior is not None and prior.delta_batch_id: + existing = await tx.get_delta_batch(prior.delta_batch_id) + if existing is not None: + return existing, False + if self.lock_manager is not None: + self._revalidate_deltas(batch) + batch = await DeltaEngine(tx).apply_batch(batch) + + # v0.4: Collect bullet IDs produced by ADD_BULLET operations + bullet_ids_produced = [ + op.target_id + for op in batch.operations + if op.op_type == DeltaOpType.ADD_BULLET and op.target_id + ] + + # Record activity WITH raw input preservation (v0.4) + embedding (v0.5) + activity = Activity( + agent_id=agent_id, + session_id=session_id, + action_type=ActionType.FACT_LEARNED, + summary=( + f"Ingested {content_type.value}: +{batch.bullets_added} bullets, " + f"~{batch.bullets_updated} updated, -{batch.bullets_removed} removed, " + f"⊕{batch.bullets_merged} merged" + ), + delta_batch_id=batch.id, + # v0.4: Raw input preservation + raw_input=content, + raw_input_hash=input_hash, + content_type=content_type.value, + source_agent_model=source_model, + feedback=feedback.model_dump() if feedback else None, + # v0.4: Extraction metadata + extraction_model=self.ingestion_config.reflector_model, + extraction_prompt_version=self.ingestion_config.reflector_prompt_version, + bullet_ids_produced=bullet_ids_produced, + # v0.5: worked-example retrieval + raw_input_embedding=raw_input_embedding, + ) + await tx.add_activity(context_id, activity) + return batch, True + + if self.lock_manager is not None: + async with self.lock_manager.acquire(ctx_id_str): + batch, persisted = await persist_commit(batch) + else: + batch, persisted = await persist_commit(batch) # Reconsolidation: if this commit references a materialization, update # bullet stats via a delta batch (audit-clean, rollback-able). if materialization_id and feedback: await self._reconsolidate( materialization_id, feedback, + expected_context_id=ctx_id_str, agent_id=agent_id, session_id=session_id, ) # v0.3: Emit event - if self.event_bus is not None: + if persisted and self.event_bus is not None: self.event_bus.emit( ctx_id_str, event_type="commit", @@ -863,14 +907,15 @@ async def add_bullet_directly( session_id: str | None = None, ) -> tuple[Bullet, DeltaBatch]: """Add a single bullet directly (bypass Reflector).""" - # v0.3: Capacity check + # Preserve the cheap full-context rejection before a provider call; + # the transaction below rechecks actual rows to close the race. context = await self.storage.get_context(uuid.UUID(context_id)) if context is not None: max_bullets = context.lifecycle_config.max_active_bullets capacity = await self.storage.get_capacity_status(context_id, max_bullets) if capacity.pressure_level == "full": raise CapacityExceededError( - context_id, capacity.active_bullet_count, max_bullets + context_id, capacity.active_bullet_count, max_bullets, ) bullet_id = str(uuid.uuid4())[:8] @@ -879,6 +924,7 @@ async def add_bullet_directly( target_id=bullet_id, section=section, content=content, + embedding=await _embed_content(self.llm, content), bullet_type=bullet_type.value, reasoning="Direct bullet addition (bypass Reflector)", source=DeltaSource.USER, @@ -892,21 +938,20 @@ async def add_bullet_directly( trigger="direct_add", ) - # v0.3: Apply under lock if available + async def persist_add(): + async with capacity_checked_transaction(self.storage, context_id) as tx: + applied = await DeltaEngine(tx).apply_batch(batch) + bullet = await tx.get_bullet(bullet_id) + if bullet is not None: + bullet.salience = salience + await tx.update_bullet(bullet) + return bullet, applied + if self.lock_manager is not None: async with self.lock_manager.acquire(context_id): - batch = await self.delta_engine.apply_batch(batch) + bullet, batch = await persist_add() else: - batch = await self.delta_engine.apply_batch(batch) - - bullet = await self.storage.get_bullet(bullet_id) - if bullet: - bullet.salience = salience - try: - bullet.embedding = await self.llm.embed(content) - except Exception: - pass - await self.storage.update_bullet(bullet) + bullet, batch = await persist_add() # v0.3: Emit event if self.event_bus is not None: @@ -939,6 +984,7 @@ async def record_decision( target_id=decision_id, section="decisions", content=f"{decision}. Rationale: {rationale}", + embedding=await _embed_content(self.llm, f"{decision}. Rationale: {rationale}"), bullet_type="decision", reasoning=f"Decision recorded: {decision}", source=DeltaSource.USER, @@ -953,6 +999,7 @@ async def record_decision( target_id=str(uuid.uuid4())[:8], section="decisions", content=f"Alternative considered: {alt}", + embedding=await _embed_content(self.llm, f"Alternative considered: {alt}"), bullet_type="fact", reasoning=f"Alternative for decision: {decision}", source=DeltaSource.USER, @@ -966,109 +1013,73 @@ async def record_decision( operations=operations, trigger="decision", ) - batch = await self.delta_engine.apply_batch(batch) - - activity = Activity( - agent_id=agent_id, - session_id=session_id, - action_type=ActionType.DECISION_MADE, - summary=f"Decision: {decision}", - delta_batch_id=batch.id, - ) - await self.storage.add_activity(context_id, activity) + async with capacity_checked_transaction(self.storage, ctx_id_str) as tx: + batch = await DeltaEngine(tx).apply_batch(batch) + activity = Activity( + agent_id=agent_id, + session_id=session_id, + action_type=ActionType.DECISION_MADE, + summary=f"Decision: {decision}", + delta_batch_id=batch.id, + ) + await tx.add_activity(context_id, activity) return decision_id, batch async def _reconsolidate( self, materialization_id: str, feedback: ExecutionFeedback, agent_id: str | None = None, session_id: uuid.UUID | None = None, + expected_context_id: str | None = None, ) -> DeltaBatch | None: - """Post-recall reconsolidation, routed through the DeltaEngine. - - Emits one RECONSOLIDATE_BULLET op per bullet in the recall, applied as - a single batch so the audit trail / rollback story is consistent with - every other mutation. (README claim: "all mutations through deltas.") - """ + """Consume feedback and its delta audit atomically across workers.""" record = await self.storage.get_materialization(materialization_id) if record is None: - logger.warning("Materialization %s not found for reconsolidation", materialization_id) return None - - # Idempotency: a materialization may be reconsolidated exactly once. - # Without this, a replayed materialization_id (retry, at-least-once - # delivery, or two commits citing one recall) double-counts hits and - # compounds salience. - if record.reconsolidated_at is not None: - logger.info( - "Materialization %s already reconsolidated at %s; skipping replay", - materialization_id, record.reconsolidated_at, - ) - return None - - # Decide the per-bullet effect from the outcome once. - if feedback.outcome == FeedbackOutcome.SUCCESS: - recall_delta, hit_delta, miss_delta, mult = 1, 1, 0, 1.05 - elif feedback.outcome == FeedbackOutcome.FAILURE: - recall_delta, hit_delta, miss_delta, mult = 1, 0, 1, 0.95 - else: - recall_delta, hit_delta, miss_delta, mult = 1, 0, 0, 1.00 - - ops: list[DeltaOperation] = [] - for bullet_id in record.bullets_included: - ops.append(DeltaOperation( - op_type=DeltaOpType.RECONSOLIDATE_BULLET, - target_id=bullet_id, - reasoning=( - f"Reconsolidation from materialization {materialization_id} " - f"(outcome={feedback.outcome.value})" - ), - source=DeltaSource.REFLECTOR, - confidence=0.9, - agent_id=agent_id, - session_id=str(session_id) if session_id else None, - # previous_state carries the *input* deltas; DeltaEngine swaps - # it for a rollback snapshot after applying. - previous_state={ - "recall_delta": recall_delta, - "hit_delta": hit_delta, - "miss_delta": miss_delta, - "salience_multiplier": mult, - "outcome": feedback.outcome.value, - }, - )) - if not ops: + if expected_context_id is not None and record.context_id != expected_context_id: return None - - batch = DeltaBatch( - context_id=record.context_id, - operations=ops, - trigger="reconsolidation", - ) - if self.lock_manager is not None: - async with self.lock_manager.acquire(record.context_id): - # Re-check under the lock so two concurrent replays can't both - # pass the pre-lock idempotency guard and double-apply. - fresh = await self.storage.get_materialization(materialization_id) - if fresh is not None and fresh.reconsolidated_at is not None: - logger.info( - "Materialization %s reconsolidated concurrently; skipping", - materialization_id, - ) - return None - batch = await self.delta_engine.apply_batch(batch) - await self.storage.mark_materialization_reconsolidated( - materialization_id, _utcnow(), - ) - else: - batch = await self.delta_engine.apply_batch(batch) - await self.storage.mark_materialization_reconsolidated( - materialization_id, _utcnow(), - ) - - logger.info( - "Reconsolidation: %d bullets updated from materialization %s (outcome=%s)", - len(ops), materialization_id, feedback.outcome.value, - ) + async with self.storage.transaction(record.context_id) as tx: + fresh = await tx.get_materialization(materialization_id) + if ( + fresh is None or fresh.context_id != record.context_id + or fresh.reconsolidated_at is not None + or (expected_context_id is not None and fresh.context_id != expected_context_id) + ): + return None + if feedback.outcome == FeedbackOutcome.SUCCESS: + recall_delta, hit_delta, miss_delta, mult = 1, 1, 0, 1.05 + elif feedback.outcome == FeedbackOutcome.FAILURE: + recall_delta, hit_delta, miss_delta, mult = 1, 0, 1, 0.95 + else: + recall_delta, hit_delta, miss_delta, mult = 1, 0, 0, 1.00 + operations = [] + for bullet_id in sorted(set(fresh.bullets_included)): + bullet = await tx.get_bullet(bullet_id) + if ( + bullet is None or bullet.context_id != fresh.context_id + or not bullet.is_active or bullet.is_archived + or bullet.lifecycle_state.value != "active" + ): + continue + operations.append(DeltaOperation( + op_type=DeltaOpType.RECONSOLIDATE_BULLET, + target_id=bullet_id, + reasoning=f"Reconsolidation from materialization {materialization_id}", + source=DeltaSource.REFLECTOR, + confidence=0.9, + agent_id=agent_id, + session_id=str(session_id) if session_id else None, + previous_state={ + "recall_delta": recall_delta, "hit_delta": hit_delta, + "miss_delta": miss_delta, "salience_multiplier": mult, + "outcome": feedback.outcome.value, + }, + )) + batch = None + if operations: + batch = await DeltaEngine(tx).apply_batch(DeltaBatch( + context_id=fresh.context_id, operations=operations, trigger="reconsolidation", + )) + await tx.mark_materialization_reconsolidated(materialization_id, _utcnow()) return batch @staticmethod diff --git a/engram/core/materialization.py b/engram/core/materialization.py index a963310..800bc77 100644 --- a/engram/core/materialization.py +++ b/engram/core/materialization.py @@ -292,75 +292,50 @@ async def _materialize_bullets( # 10 near-duplicates of the top-scoring bullet. ordered_bullets = _mmr_order(bullets, scores, mmr_lambda) - # Convert top bullets to ConceptNodes for renderer (backward compat with renderers) - selected_concepts: list[ConceptNode] = [] - selected_bullet_ids: list[str] = [] + # Renderers make the budget decision using the final formatted text. + # Keep ID mappings rather than guessing inclusion from content strings: + # duplicate facts and schemas can legitimately have identical text. + candidates: list[ConceptNode] = [] + bullet_ids_by_concept: dict[uuid.UUID, str] = {} + schema_ids_by_concept: dict[uuid.UUID, str] = {} usage_stats: dict[str, str] = {} - total_tokens = 0 - - intent_budget = 0 - if intent or core_memory: - intent_text = renderer.render([], intent, token_budget, core_memory=core_memory) - intent_budget = renderer.estimate_tokens(intent_text) - - # Reserve some headroom for worked examples (rendered at the end). - worked_example_budget = 0 - if worked_examples: - wx_text_estimate = sum( - renderer.estimate_tokens((ex.get("input") or "") + (ex.get("output") or "")) - + 30 - for ex in worked_examples + for schema in schemas[:5]: + concept = ConceptNode( + type=ConceptType.PATTERN, + content=f"[Pattern: {schema.name}] {schema.description}", + salience=schema.confidence, ) - worked_example_budget = wx_text_estimate - - remaining = max(0, token_budget - intent_budget - worked_example_budget) - - # Include schema summaries first (token-efficient) - schema_ids: list[str] = [] - if schemas: - for schema in schemas[:5]: - schema_text = f"[Pattern: {schema.name}] {schema.description}" - est = renderer.estimate_tokens(schema_text) + 10 - if total_tokens + est <= remaining: - selected_concepts.append(ConceptNode( - type=ConceptType.PATTERN, - content=schema_text, - salience=schema.confidence, - )) - total_tokens += est - schema_ids.append(schema.id) - - # Pack bullets into remaining budget + candidates.append(concept) + schema_ids_by_concept[concept.id] = schema.id + for bullet in ordered_bullets: - est = renderer.estimate_tokens(bullet.content) + 10 - if total_tokens + est > remaining: - break - concept_type = self._bullet_type_to_concept_type( - bullet.bullet_type.value if hasattr(bullet.bullet_type, 'value') else str(bullet.bullet_type) - ) - selected_concepts.append(ConceptNode( - type=concept_type, + concept = ConceptNode( + type=self._bullet_type_to_concept_type( + bullet.bullet_type.value if hasattr(bullet.bullet_type, "value") + else str(bullet.bullet_type) + ), content=bullet.content, salience=bullet.salience, confidence=bullet.confidence, - )) - selected_bullet_ids.append(bullet.id) + ) + candidates.append(concept) + bullet_ids_by_concept[concept.id] = bullet.id if include_usage_stats and (bullet.recall_count or bullet.hit_count): - # "(used N×, success Y/Z)" surfaces reinforcement signal to the consumer. - # Denominator floors at max(recall_count, hit_count) so we never render - # "success 3/0" when hit_count was incremented out-of-band from recall. denom = max(bullet.recall_count, bullet.hit_count) - usage_stats[bullet.content] = ( + usage_stats[str(concept.id)] = ( f"(used {bullet.recall_count}×, success {bullet.hit_count}/{denom})" ) - total_tokens += est - rendered = renderer.render( - selected_concepts, intent, token_budget, + rendered, selected_concepts = renderer.render_with_selection( + candidates, intent, token_budget, core_memory=core_memory, worked_examples=worked_examples or None, usage_stats=usage_stats or None, ) + selected_bullet_ids = [bullet_ids_by_concept[c.id] for c in selected_concepts + if c.id in bullet_ids_by_concept] + schema_ids = [schema_ids_by_concept[c.id] for c in selected_concepts + if c.id in schema_ids_by_concept] actual_tokens = renderer.estimate_tokens(rendered) coverage = len(selected_bullet_ids) / len(bullets) if bullets else 0.0 @@ -394,27 +369,10 @@ async def _materialize_concepts( sorted_concepts = sorted(scores.items(), key=lambda x: x[1], reverse=True) id_to_concept = {c.id: c for c in concepts} - selected: list[ConceptNode] = [] - total_tokens = 0 - - intent_budget = 0 - if intent or core_memory: - intent_text = renderer.render([], intent, token_budget, core_memory=core_memory) - intent_budget = renderer.estimate_tokens(intent_text) - remaining = token_budget - intent_budget - - for concept_id, score in sorted_concepts: - concept = id_to_concept.get(concept_id) - if concept is None: - continue - est = renderer.estimate_tokens(concept.content) + 10 - if total_tokens + est > remaining: - break - selected.append(concept) - total_tokens += est - - rendered = renderer.render( - selected, intent, token_budget, + candidates = [id_to_concept[concept_id] for concept_id, _ in sorted_concepts + if concept_id in id_to_concept] + rendered, selected = renderer.render_with_selection( + candidates, intent, token_budget, core_memory=core_memory, worked_examples=worked_examples or None, ) diff --git a/engram/core/models.py b/engram/core/models.py index fde46f9..d761d6b 100644 --- a/engram/core/models.py +++ b/engram/core/models.py @@ -266,6 +266,7 @@ class DeltaOperation(BaseModel): section: str | None = None content: str | None = None bullet_type: str | None = None + embedding: list[float] | None = None reasoning: str = "" source: DeltaSource = DeltaSource.CURATOR confidence: float = 0.5 @@ -466,6 +467,7 @@ class ConceptNode(BaseModel): """ id: uuid.UUID = Field(default_factory=uuid.uuid4) + context_id: uuid.UUID | None = None type: ConceptType content: str = Field(description="Concise natural language, 1-3 sentences max") embedding: list[float] | None = Field(default=None) diff --git a/engram/core/re_extraction.py b/engram/core/re_extraction.py index 48cda0f..9e59ecb 100644 --- a/engram/core/re_extraction.py +++ b/engram/core/re_extraction.py @@ -122,8 +122,7 @@ async def re_extract( return self._build_preview(all_operations, activities) # Apply all deltas - from engram.core.delta import DeltaEngine - delta_engine = DeltaEngine(self.storage) + from engram.core.delta import DeltaEngine, capacity_checked_transaction batch = DeltaBatch( context_id=context_id, @@ -131,28 +130,33 @@ async def re_extract( trigger="re_extraction", ) + async def persist_re_extraction(): + async with capacity_checked_transaction(self.storage, context_id) as tx: + result_batch = await DeltaEngine(tx).apply_batch(batch) + # Record re-extraction event in activity ledger + import uuid as _uuid + await tx.add_activity( + _uuid.UUID(context_id) if len(context_id) >= 32 else _uuid.uuid4(), + Activity( + agent_id="engram-system", + action_type=ActionType.RE_EXTRACTION_RAN, + summary=( + f"Re-extracted with {request.reflector_model}: " + f"{len(activities)} activities processed, " + f"{len(all_operations)} delta operations" + ), + delta_batch_id=result_batch.id, + extraction_model=request.reflector_model, + ), + ) + + return result_batch + if self.lock_manager is not None: async with self.lock_manager.acquire(context_id): - result_batch = await delta_engine.apply_batch(batch) + result_batch = await persist_re_extraction() else: - result_batch = await delta_engine.apply_batch(batch) - - # Record re-extraction event in activity ledger - import uuid as _uuid - await self.storage.add_activity( - _uuid.UUID(context_id) if len(context_id) >= 32 else _uuid.uuid4(), - Activity( - agent_id="engram-system", - action_type=ActionType.RE_EXTRACTION_RAN, - summary=( - f"Re-extracted with {request.reflector_model}: " - f"{len(activities)} activities processed, " - f"{len(all_operations)} delta operations" - ), - delta_batch_id=result_batch.id, - extraction_model=request.reflector_model, - ), - ) + result_batch = await persist_re_extraction() # Emit event if self.event_bus is not None: diff --git a/engram/core/similarity.py b/engram/core/similarity.py index af9ec82..7e2fa86 100644 --- a/engram/core/similarity.py +++ b/engram/core/similarity.py @@ -17,3 +17,23 @@ def cosine_similarity(a: list[float] | None, b: list[float] | None) -> float: if not na or not nb: return 0.0 return dot / (na * nb) + + +def validated_cosine_similarity(a, b) -> float | None: + """Return None when comparison is unknown, never a false finite nonmatch.""" + try: + if not a or not b or len(a) != len(b): + return None + if not all(math.isfinite(value) for value in (*a, *b)): + return None + norm_a, norm_b = math.hypot(*a), math.hypot(*b) + denominator = norm_a * norm_b + if not math.isfinite(denominator) or denominator <= 0: + return None + dot = sum(x * y for x, y in zip(a, b)) + similarity = dot / denominator + if not math.isfinite(similarity): + return None + return max(-1.0, min(1.0, similarity)) + except (TypeError, ValueError, OverflowError): + return None diff --git a/engram/integrations/mcp_server.py b/engram/integrations/mcp_server.py index 6236dba..092a2f6 100644 --- a/engram/integrations/mcp_server.py +++ b/engram/integrations/mcp_server.py @@ -1,29 +1,14 @@ -"""MCP (Model Context Protocol) server for Claude integration. - -v0.2: Added save_bullet, list_bullets, consolidate, get_health tools. -Updated save_context to show bullet-based results. -v0.4: Added re_extract_context and get_ingestion_config tools. -v0.4.1: Added engram_ prefixed tool names per app.engram.so spec. - -Run locally: python -m engram.integrations.mcp_server -Run hosted: Served at mcp.engram.so for cloud customers - -Add to Claude: - claude mcp add engram -- python -m engram.integrations.mcp_server - -Or configure in claude_desktop_config.json: - { - "mcpServers": { - "engram": { - "command": "python", - "args": ["-m", "engram.integrations.mcp_server"], - "env": { - "ENGRAM_API_URL": "http://localhost:5820", - "ENGRAM_API_KEY": "eng_sk_..." - } - } - } - } +"""MCP tools for Claude and other Model Context Protocol clients. + +Run with an absolute virtualenv Python path: + /absolute/path/to/engram/.venv/bin/python -m engram.integrations.mcp_server + +The stdio bridge reads ENGRAM_API_URL (default http://localhost:5820) and +ENGRAM_API_KEY (optional for self-hosted OSS, required for Engram Cloud). +Install from this corrected checkout as documented in the README. The published +engram-contextdb 0.4.5 wheel does not yet include the stdio entrypoint fix. + +Cloud custom connector URL: https://api.engram.so/mcp/sse """ from __future__ import annotations @@ -31,6 +16,7 @@ import asyncio import json import logging +import os import uuid from typing import Any @@ -47,20 +33,30 @@ MCP_AVAILABLE = False -def create_mcp_server(engram_url: str = "http://localhost:5820") -> Any: - """Create an MCP server that exposes Engram operations as tools for Claude. +def create_mcp_server( + engram_url: str = "http://localhost:5820", + client: Any = None, +) -> Any: + """Expose Engram operations as MCP tools. + + ``client`` may be an authenticated SDK client supplied by a hosted transport; + its lifetime remains the caller's responsibility. Otherwise, ``engram_url`` + is used to construct a client. The stdio runner supplies and closes its own + client, including its configured API key. - Requires the `mcp` extra: pip install engram[mcp] + Install MCP support from the checkout: python -m pip install ".[mcp]" """ if not MCP_AVAILABLE: raise ImportError( - "MCP SDK not installed. Install with: pip install engram[mcp]" + "MCP SDK not installed. Install from the source checkout with: " + 'python -m pip install ".[mcp]"' ) from engram.sdk.client import Engram server = Server("engram") - client = Engram(url=engram_url) + if client is None: + client = Engram(url=engram_url) @server.list_tools() async def list_tools() -> list[Tool]: @@ -527,7 +523,7 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]: return [ TextContent( type="text", - text=f"Decision recorded. Bullet ID: {result.get('decision_id', 'N/A')}", + text=f"Decision recorded. Bullet ID: {result.get('concept_id', 'N/A')}", ) ] @@ -689,7 +685,7 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]: return [ TextContent( type="text", - text=f"Decision recorded (id: {result.get('decision_id', 'unknown')})", + text=f"Decision recorded (id: {result.get('concept_id', 'unknown')})", ) ] @@ -736,12 +732,25 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]: async def run_mcp_server(engram_url: str = "http://localhost:5820") -> None: - """Run the MCP server over stdio.""" + """Run the stdio bridge, honoring ENGRAM_API_URL and ENGRAM_API_KEY.""" if not MCP_AVAILABLE: raise ImportError( - "MCP SDK not installed. Install with: pip install engram[mcp]" + "MCP SDK not installed. Install from the source checkout with: " + 'python -m pip install ".[mcp]"' ) - server = create_mcp_server(engram_url) - async with stdio_server() as (read_stream, write_stream): - await server.run(read_stream, write_stream, server.create_initialization_options()) + from engram.sdk.client import Engram + + url = os.environ.get("ENGRAM_API_URL", engram_url) + api_key = os.environ.get("ENGRAM_API_KEY") + client = Engram(url=url, api_key=api_key) + try: + server = create_mcp_server(client=client) + async with stdio_server() as (read_stream, write_stream): + await server.run(read_stream, write_stream, server.create_initialization_options()) + finally: + await client.close() + + +if __name__ == "__main__": + asyncio.run(run_mcp_server()) diff --git a/engram/llm/adapter.py b/engram/llm/adapter.py index 41e9fb3..f742b52 100644 --- a/engram/llm/adapter.py +++ b/engram/llm/adapter.py @@ -3,6 +3,7 @@ from __future__ import annotations import abc +import inspect import json import logging from typing import Any @@ -37,6 +38,21 @@ async def embed(self, text: str) -> list[float]: """Generate an embedding vector for the given text.""" +async def complete_with_model(llm: LLMAdapter, *, model: str, **kwargs: Any) -> str: + """Use per-call routing when supported, retaining pre-model custom adapters. + + Older third-party adapters remain responsible for their own model choice. + Inspect before calling so a provider's TypeError is never mistaken for an + unsupported argument and retried as a second billable completion. + """ + parameters = inspect.signature(llm.complete).parameters + if "model" in parameters or any( + parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values() + ): + kwargs["model"] = model + return await llm.complete(**kwargs) + + class LiteLLMAdapter(LLMAdapter): """Universal LLM adapter using LiteLLM to support any provider.""" diff --git a/engram/llm/explicit.py b/engram/llm/explicit.py new file mode 100644 index 0000000..77ab0fb --- /dev/null +++ b/engram/llm/explicit.py @@ -0,0 +1,61 @@ +"""Isolated provider configuration for standalone maintenance/evaluation commands. + +LiteLLM has process-global configuration, loads .env in development mode, and +can print diagnostics. This context is for standalone commands, never request +handlers or concurrent application tasks. Their only output should be their +own structured, sanitized result. +""" + +from __future__ import annotations + +import logging +import os +from contextlib import contextmanager, redirect_stderr, redirect_stdout +from io import TextIOBase + + +class _Discard(TextIOBase): + @property + def encoding(self) -> str: + return "utf-8" + + def write(self, text: str) -> int: + return len(text) + + def flush(self) -> None: + pass + + +@contextmanager +def explicit_embedding_runtime(): + """Disable implicit dotenv loading and discard raw SDK diagnostics. + + Restore ambient process settings on exit, including on import/provider + failure. Import the SDK only here, after setting production mode. Callers + must enter this context only for an explicitly requested provider operation. + """ + settings = {"LITELLM_MODE": "PRODUCTION", "LITELLM_LOCAL_MODEL_COST_MAP": "True"} + previous = {key: os.environ.get(key) for key in settings} + previous_logging = logging.root.manager.disable + os.environ.update(settings) + logging.disable(logging.CRITICAL) + try: + with redirect_stdout(_Discard()), redirect_stderr(_Discard()): + import litellm + + previous_debug = litellm.suppress_debug_info + previous_verbose = litellm.set_verbose + litellm.suppress_debug_info = True + litellm.set_verbose = False + try: + yield + finally: + litellm.suppress_debug_info = previous_debug + litellm.set_verbose = previous_verbose + finally: + logging.disable(previous_logging) + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value diff --git a/engram/maintenance/__init__.py b/engram/maintenance/__init__.py new file mode 100644 index 0000000..a20d04f --- /dev/null +++ b/engram/maintenance/__init__.py @@ -0,0 +1 @@ +"""Explicit maintenance commands for existing Engram databases.""" diff --git a/engram/maintenance/backfill_embeddings.py b/engram/maintenance/backfill_embeddings.py new file mode 100644 index 0000000..af4897c --- /dev/null +++ b/engram/maintenance/backfill_embeddings.py @@ -0,0 +1,230 @@ +"""Repair missing memory vectors without re-extracting or changing memory text. + +The command defaults to a read-only preview. It never loads .env, initializes a +schema, or chooses a database implicitly. Provider calls require --apply. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import os +import re +import uuid +from collections.abc import Sequence +from contextlib import nullcontext +from typing import Any + +from engram.core.models import Bullet, LifecycleState +from engram.llm.adapter import LiteLLMAdapter, LLMAdapter +from engram.llm.explicit import explicit_embedding_runtime +from engram.storage.base import StorageBackend +from engram.storage.postgres import PostgresBackend +from engram.storage.sqlite import SQLiteBackend + +EMBEDDING_DIMENSIONS = 1536 # Matches the current PostgreSQL vector columns. +MAX_LIMIT = 1000 + + +def _eligible(bullet: Bullet, context_id: str) -> bool: + return ( + bullet.context_id == context_id + and bullet.is_active + and not bullet.is_archived + and bullet.lifecycle_state == LifecycleState.ACTIVE + and bullet.embedding is None + ) + + +def _normalize_vector(vector: Any) -> list[float] | None: + if not isinstance(vector, Sequence) or isinstance(vector, (str, bytes)): + return None + if len(vector) != EMBEDDING_DIMENSIONS: + return None + try: + import struct + + if any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + for value in vector + ): + return None + # PostgreSQL stores float32. Validate the representation that will be + # persisted, including underflow-to-zero and overflow, on both backends. + normalized = list(struct.unpack(f"{EMBEDDING_DIMENSIONS}f", struct.pack( + f"{EMBEDDING_DIMENSIONS}f", *vector, + ))) + # Cosine search also accumulates squared components in float32. A + # representable nonzero component can still square to zero or infinity, + # making the saved vector unusable for search. Preserve provider scale + # but reject that representation before marking a memory repaired. + squared_norm = 0.0 + for value in normalized: + squared = struct.unpack("f", struct.pack("f", value * value))[0] + squared_norm = struct.unpack("f", struct.pack("f", squared_norm + squared))[0] + return normalized if math.isfinite(squared_norm) and squared_norm > 0 else None + except (OverflowError, TypeError, ValueError, struct.error): + return None + + +async def backfill_embeddings( + storage: StorageBackend, + context_id: uuid.UUID, + *, + llm: LLMAdapter | None = None, + apply: bool = False, + limit: int = 100, +) -> dict[str, Any]: + """Preview or repair up to ``limit`` eligible memories in one context. + + Provider calls happen outside transactions. Each conditional vector write + is atomic; successful earlier repairs remain useful if another item fails. + Re-running skips repaired items. Error details intentionally omit provider + messages, which can contain credentials or memory text. + """ + if not 1 <= limit <= MAX_LIMIT: + raise ValueError(f"limit must be between 1 and {MAX_LIMIT}") + if apply and llm is None: + raise ValueError("apply requires an embedding adapter") + if await storage.get_context(context_id) is None: + raise ValueError("context does not exist") + ctx = str(context_id) + eligible = sorted( + (bullet for bullet in await storage.list_bullets(ctx) if _eligible(bullet, ctx)), + key=lambda bullet: bullet.id, + ) + selected = eligible[:limit] + result: dict[str, Any] = { + "context_id": ctx, + "dry_run": not apply, + "eligible": len(eligible), + "selected": len(selected), + "attempted": 0, + "updated": 0, + "skipped": 0, + "failed": 0, + "errors": [], + } + if not apply: + return result + + assert llm is not None + for original in selected: + result["attempted"] += 1 + try: + vector = await llm.embed(original.content) + except Exception: + result["failed"] += 1 + result["errors"].append({"bullet_id": original.id, "reason": "provider_error"}) + continue + vector = _normalize_vector(vector) + if vector is None: + result["failed"] += 1 + result["errors"].append({"bullet_id": original.id, "reason": "invalid_vector"}) + continue + try: + async with storage.transaction(ctx) as tx: + current = await tx.get_bullet(original.id) + if ( + current is None + or not _eligible(current, ctx) + or current.content != original.content + ): + updated = False + else: + updated = await tx.update_bullet_embedding_if_missing( + ctx, current.id, original.content, list(vector), + ) + result["updated" if updated else "skipped"] += 1 + except Exception: + result["failed"] += 1 + result["errors"].append({"bullet_id": original.id, "reason": "storage_error"}) + return result + + +def _positive_limit(value: str) -> int: + try: + parsed = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("limit must be an integer") from exc + if not 1 <= parsed <= MAX_LIMIT: + raise argparse.ArgumentTypeError(f"limit must be between 1 and {MAX_LIMIT}") + return parsed + + +def _env_name(value: str) -> str: + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", value): + raise argparse.ArgumentTypeError("expected an environment-variable name") + return value + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + database = parser.add_mutually_exclusive_group(required=True) + database.add_argument("--sqlite", metavar="EXISTING_PATH", help="existing SQLite database") + database.add_argument( + "--postgres-dsn-env", type=_env_name, metavar="ENV_NAME", + help="environment variable containing the existing PostgreSQL database DSN", + ) + parser.add_argument("--context-id", type=uuid.UUID, required=True) + parser.add_argument("--limit", type=_positive_limit, default=100) + parser.add_argument("--apply", action="store_true", help="call the provider and save vectors") + parser.add_argument("--embedding-model", help="explicit embedding model for --apply") + parser.add_argument( + "--embedding-key-env", type=_env_name, metavar="ENV_NAME", + help="environment variable containing the embedding provider API key", + ) + return parser + + +async def _run(args: argparse.Namespace) -> int: + if args.sqlite: + storage = SQLiteBackend(args.sqlite, read_only=not args.apply, existing_only=True) + else: + storage = PostgresBackend( + os.environ[args.postgres_dsn_env], read_only=not args.apply, existing_only=True, + ) + try: + with explicit_embedding_runtime() if args.apply else nullcontext(): + llm = None + if args.apply: + llm = LiteLLMAdapter( + embedding_model=args.embedding_model, + embedding_api_key=os.environ[args.embedding_key_env], + ) + result = await backfill_embeddings( + storage, args.context_id, llm=llm, apply=args.apply, limit=args.limit, + ) + print(json.dumps(result, sort_keys=True)) + return 1 if result["failed"] else 0 + finally: + await storage.close() + + +def main(argv: Sequence[str] | None = None) -> int: + parser = _parser() + args = parser.parse_args(argv) + if args.postgres_dsn_env and not os.environ.get(args.postgres_dsn_env): + parser.error("the PostgreSQL DSN environment variable must be set") + if args.apply: + if not args.embedding_model or not args.embedding_key_env: + parser.error("--apply requires --embedding-model and --embedding-key-env") + if not os.environ.get(args.embedding_key_env): + parser.error("the embedding API key environment variable must be set") + try: + return asyncio.run(_run(args)) + except ValueError: + print(json.dumps({"error": "invalid_context_or_options"})) + return 1 + except Exception: + # A raw driver error can include the DSN, SQL, or a provider credential. + print(json.dumps({"error": "database_operation_failed"})) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/engram/renderers/base.py b/engram/renderers/base.py index 7e22e33..05ad509 100644 --- a/engram/renderers/base.py +++ b/engram/renderers/base.py @@ -3,6 +3,8 @@ from __future__ import annotations import abc +from collections.abc import Callable +from typing import Any from engram.core.models import ConceptNode, IntentAnchor @@ -26,9 +28,124 @@ def render( - core_memory: Mem-α-style always-in-context summary (rendered first). - worked_examples: DC-style nearest prior-input/output pairs (rendered last, adjacent to the question). - - usage_stats: maps concept.content → "(used N×, success Y/Z)" suffix. + - usage_stats: maps str(concept.id) → "(used N×, success Y/Z)" suffix. + Content-keyed annotations remain supported for existing callers. """ @abc.abstractmethod def estimate_tokens(self, text: str) -> int: """Estimate the number of tokens in a string.""" + + @staticmethod + def _usage_suffix(concept: ConceptNode, usage_stats: dict[str, str] | None) -> str: + """Prefer identity so duplicate text cannot share another memory's stats.""" + if not usage_stats: + return "" + annotation = usage_stats.get(str(concept.id), usage_stats.get(concept.content, "")) + return f" {annotation}" if annotation else "" + + def _render_full( + self, + concepts: list[ConceptNode], + intent: IntentAnchor | None, + token_budget: int, + **options: Any, + ) -> str: + # Compatibility for third-party renderers implementing the original API. + return self.render(concepts, intent, token_budget, **options) + + def render_with_selection( + self, + concepts: list[ConceptNode], + intent: IntentAnchor | None, + token_budget: int, + **options: Any, + ) -> tuple[str, list[ConceptNode]]: + """Pack complete concepts against the final format, returning exact IDs. + + Counts wrappers, separators, tags and optional annotations with this + renderer's estimator. Oversized candidates are skipped, never a reason + to stop considering smaller successors. Priority is core memory, intent, + worked examples, then the supplied candidate order. Oversized core and + intent text is shortened; examples and concepts remain whole. + """ + budget = max(0, token_budget) + if not budget: + return "", [] + supplied = dict(options) + selected_options = dict(options) + if "core_memory" in options: + selected_options["core_memory"] = "" + if "worked_examples" in options: + selected_options["worked_examples"] = None + selected_intent = None + + def full(items: list[ConceptNode], anchor: IntentAnchor | None) -> str: + # Internal full renderers do not trim. The large budget keeps + # original third-party renderer implementations compatible too. + return self._render_full(items, anchor, 2**63 - 1, **selected_options) + + def fits(text: str) -> bool: + return self.estimate_tokens(text) <= budget + + def prefix(text: str, render: Callable[[str], str]) -> str: + if fits(render(text)): + return text + low, high = 0, len(text) + best = "" + while low <= high: + middle = (low + high) // 2 + candidate = text[:middle].rstrip() + " …[truncated]" if middle else "" + if fits(render(candidate)): + best = candidate + low = middle + 1 + else: + high = middle - 1 + return best + + if supplied.get("core_memory"): + + def render_core(value: str) -> str: + selected_options["core_memory"] = value + return full([], None) + + selected_options["core_memory"] = prefix(supplied["core_memory"], render_core) + + if intent: + if fits(full([], intent)): + selected_intent = intent + else: + anchor = intent.model_copy(update={"success_criteria": [], "constraints": []}) + objective = prefix( + anchor.objective, + lambda value: full([], anchor.model_copy(update={"objective": value})), + ) + anchor = anchor.model_copy(update={"objective": objective}) + if objective and fits(full([], anchor)): + selected_intent = anchor + for field in ("success_criteria", "constraints"): + for value in getattr(intent, field): + trial = selected_intent.model_copy( + update={field: [*getattr(selected_intent, field), value]} + ) + if fits(full([], trial)): + selected_intent = trial + + if supplied.get("worked_examples"): + examples = [] + for example in supplied["worked_examples"]: + selected_options["worked_examples"] = [*examples, example] + if fits(full([], selected_intent)): + examples.append(example) + selected_options["worked_examples"] = examples or None + + selected: list[ConceptNode] = [] + rendered = full([], selected_intent) + if not fits(rendered): + return "", [] + for concept in concepts: + candidate = full([*selected, concept], selected_intent) + if fits(candidate): + selected.append(concept) + rendered = candidate + return rendered, selected diff --git a/engram/renderers/claude.py b/engram/renderers/claude.py index d893ef6..e98ad1a 100644 --- a/engram/renderers/claude.py +++ b/engram/renderers/claude.py @@ -17,6 +17,20 @@ def render( core_memory: str = "", worked_examples: list[dict] | None = None, usage_stats: dict[str, str] | None = None, + ) -> str: + return self.render_with_selection( + concepts, intent, token_budget, + core_memory=core_memory, worked_examples=worked_examples, usage_stats=usage_stats, + )[0] + + def _render_full( + self, + concepts: list[ConceptNode], + intent: IntentAnchor | None, + token_budget: int, + core_memory: str = "", + worked_examples: list[dict] | None = None, + usage_stats: dict[str, str] | None = None, ) -> str: sections: list[str] = [""] @@ -63,8 +77,6 @@ def render( } ] - current_tokens = self.estimate_tokens("\n".join(sections)) - for section_name, section_concepts in [ ("key_decisions", decisions), ("relevant_facts", facts), @@ -78,19 +90,7 @@ def render( if not section_concepts: continue block = self._render_section(section_name, section_concepts, usage_stats) - block_tokens = self.estimate_tokens(block) - if current_tokens + block_tokens > token_budget: - # Trim section to fit budget - block = self._render_section_trimmed( - section_name, - section_concepts, - token_budget - current_tokens - 20, - ) - if block: - sections.append(block) - break sections.append(block) - current_tokens += block_tokens if worked_examples: sections.append(" ") @@ -114,7 +114,7 @@ def render( def estimate_tokens(self, text: str) -> int: """Approximate token count: ~4 characters per token for Claude.""" - return len(text) // 4 + 1 + return len(text) // 4 + 1 if text else 0 def _render_section( self, name: str, concepts: list[ConceptNode], @@ -124,24 +124,7 @@ def _render_section( for c in concepts: confidence = f' confidence="{c.confidence:.1f}"' if c.confidence < 1.0 else "" tags = f' tags="{",".join(c.domain_tags)}"' if c.domain_tags else "" - suffix = f" {usage_stats[c.content]}" if usage_stats and c.content in usage_stats else "" + suffix = self._usage_suffix(c, usage_stats) lines.append(f" {c.content}{suffix}") lines.append(f" ") return "\n".join(lines) - - def _render_section_trimmed( - self, name: str, concepts: list[ConceptNode], budget_tokens: int - ) -> str | None: - if budget_tokens <= 0: - return None - lines = [f" <{name}>"] - current = self.estimate_tokens(lines[0]) - for c in concepts: - line = f" {c.content}" - line_tokens = self.estimate_tokens(line) - if current + line_tokens > budget_tokens: - break - lines.append(line) - current += line_tokens - lines.append(f" ") - return "\n".join(lines) if len(lines) > 2 else None diff --git a/engram/renderers/generic.py b/engram/renderers/generic.py index 30e5b96..477079f 100644 --- a/engram/renderers/generic.py +++ b/engram/renderers/generic.py @@ -2,7 +2,7 @@ from __future__ import annotations -from engram.core.models import ConceptNode, ConceptType, IntentAnchor +from engram.core.models import ConceptNode, IntentAnchor from engram.renderers.base import ContextRenderer @@ -17,6 +17,20 @@ def render( core_memory: str = "", worked_examples: list[dict] | None = None, usage_stats: dict[str, str] | None = None, + ) -> str: + return self.render_with_selection( + concepts, intent, token_budget, + core_memory=core_memory, worked_examples=worked_examples, usage_stats=usage_stats, + )[0] + + def _render_full( + self, + concepts: list[ConceptNode], + intent: IntentAnchor | None, + token_budget: int, + core_memory: str = "", + worked_examples: list[dict] | None = None, + usage_stats: dict[str, str] | None = None, ) -> str: sections: list[str] = [] @@ -34,19 +48,10 @@ def render( sections.append("CONSTRAINTS: " + "; ".join(intent.constraints)) sections.append("") - current_tokens = self.estimate_tokens("\n".join(sections)) - for concept in concepts: - usage = ( - f" {usage_stats[concept.content]}" - if usage_stats and concept.content in usage_stats else "" - ) + usage = self._usage_suffix(concept, usage_stats) line = f"[{concept.type.value.upper()}] {concept.content}{usage}" - line_tokens = self.estimate_tokens(line) - if current_tokens + line_tokens > token_budget: - break sections.append(line) - current_tokens += line_tokens if worked_examples: sections.append("") @@ -62,4 +67,4 @@ def render( return "\n".join(sections) def estimate_tokens(self, text: str) -> int: - return len(text) // 4 + 1 + return len(text) // 4 + 1 if text else 0 diff --git a/engram/renderers/gpt.py b/engram/renderers/gpt.py index 50e3b95..bbdaade 100644 --- a/engram/renderers/gpt.py +++ b/engram/renderers/gpt.py @@ -17,6 +17,20 @@ def render( core_memory: str = "", worked_examples: list[dict] | None = None, usage_stats: dict[str, str] | None = None, + ) -> str: + return self.render_with_selection( + concepts, intent, token_budget, + core_memory=core_memory, worked_examples=worked_examples, usage_stats=usage_stats, + )[0] + + def _render_full( + self, + concepts: list[ConceptNode], + intent: IntentAnchor | None, + token_budget: int, + core_memory: str = "", + worked_examples: list[dict] | None = None, + usage_stats: dict[str, str] | None = None, ) -> str: sections: list[str] = [] @@ -50,26 +64,18 @@ def render( ] categorized_types = {ct for _, ct in groups} - current_tokens = self.estimate_tokens("\n".join(sections)) - for heading, concept_type in groups: items = [c for c in concepts if c.type == concept_type] if not items: continue block = self._render_group(heading, items, usage_stats) - block_tokens = self.estimate_tokens(block) - if current_tokens + block_tokens > token_budget: - break sections.append(block) - current_tokens += block_tokens # Remaining types other = [c for c in concepts if c.type not in categorized_types] if other: block = self._render_group("Other Context", other, usage_stats) - block_tokens = self.estimate_tokens(block) - if current_tokens + block_tokens <= token_budget: - sections.append(block) + sections.append(block) if worked_examples: sections.append("\n## Worked Examples\n") @@ -97,7 +103,7 @@ def estimate_tokens(self, text: str) -> int: enc = tiktoken.encoding_for_model("gpt-4o") return len(enc.encode(text)) except Exception: - return len(text) // 4 + 1 + return len(text) // 4 + 1 if text else 0 def _render_group( self, heading: str, concepts: list[ConceptNode], @@ -106,6 +112,6 @@ def _render_group( lines = [f"\n## {heading}\n"] for c in concepts: tag_str = f" `{', '.join(c.domain_tags)}`" if c.domain_tags else "" - usage = f" {usage_stats[c.content]}" if usage_stats and c.content in usage_stats else "" + usage = self._usage_suffix(c, usage_stats) lines.append(f"- {c.content}{tag_str}{usage}") return "\n".join(lines) diff --git a/engram/server/routes/bullets.py b/engram/server/routes/bullets.py index b8468d1..fda4676 100644 --- a/engram/server/routes/bullets.py +++ b/engram/server/routes/bullets.py @@ -82,7 +82,7 @@ async def get_bullet( """Get a single bullet by ID.""" storage = request.app.state.storage bullet = await storage.get_bullet(bullet_id) - if bullet is None: + if bullet is None or bullet.context_id != str(context_id): raise HTTPException(status_code=404, detail=f"Bullet {bullet_id} not found") data = bullet.model_dump(mode="json") diff --git a/engram/server/routes/concepts.py b/engram/server/routes/concepts.py index c3ec2e9..9077333 100644 --- a/engram/server/routes/concepts.py +++ b/engram/server/routes/concepts.py @@ -93,7 +93,7 @@ async def record_decision( raise HTTPException(status_code=404, detail=f"Context {context_id} not found") ingestion = request.app.state.ingestion - concept_id = await ingestion.record_decision( + concept_id, _ = await ingestion.record_decision( context_id=context_id, decision=req.decision, rationale=req.rationale, diff --git a/engram/server/routes/contexts.py b/engram/server/routes/contexts.py index b66ca63..e85a65f 100644 --- a/engram/server/routes/contexts.py +++ b/engram/server/routes/contexts.py @@ -167,7 +167,7 @@ async def invalidate_concepts( if req.bullet_ids: for bullet_id in req.bullet_ids: bullet = await storage.get_bullet(bullet_id) - if bullet: + if bullet and bullet.context_id == str(context_id): bullet.is_active = False await storage.update_bullet(bullet) result["invalidated_bullets"].append(bullet_id) diff --git a/engram/server/routes/deltas.py b/engram/server/routes/deltas.py index e234106..30285f5 100644 --- a/engram/server/routes/deltas.py +++ b/engram/server/routes/deltas.py @@ -35,7 +35,7 @@ async def get_delta( """Get a specific delta batch by ID.""" storage = request.app.state.storage batch = await storage.get_delta_batch(delta_id) - if batch is None: + if batch is None or batch.context_id != str(context_id): raise HTTPException(status_code=404, detail=f"Delta batch {delta_id} not found") return batch.model_dump(mode="json") @@ -50,6 +50,10 @@ async def rollback_delta( if ctx is None: raise HTTPException(status_code=404, detail=f"Context {context_id} not found") + batch = await storage.get_delta_batch(delta_id) + if batch is None or batch.context_id != str(context_id): + raise HTTPException(status_code=404, detail=f"Delta batch {delta_id} not found") + delta_engine = request.app.state.delta_engine success = await delta_engine.rollback_batch(delta_id) if not success: diff --git a/engram/server/routes/lifecycle.py b/engram/server/routes/lifecycle.py index 0fd5408..9a83fd5 100644 --- a/engram/server/routes/lifecycle.py +++ b/engram/server/routes/lifecycle.py @@ -10,10 +10,12 @@ import asyncio import json import uuid -from datetime import datetime +from datetime import datetime, timezone -from fastapi import APIRouter, HTTPException, Request +from fastapi import APIRouter, Body, HTTPException, Request from fastapi.responses import StreamingResponse +from fastapi.exceptions import RequestValidationError +from pydantic import ValidationError from engram.core.config import IngestionConfig from engram.core.exceptions import CapacityExceededError @@ -66,24 +68,26 @@ async def restore_bullet(context_id: str, bullet_id: str, request: Request) -> d """Restore an archived bullet to active state.""" storage = request.app.state.storage - # Check capacity before restoring - ctx = await storage.get_context(uuid.UUID(context_id)) - if ctx is not None: - capacity = await storage.get_capacity_status( - context_id, ctx.lifecycle_config.max_active_bullets - ) - if capacity.pressure_level == "full": + # Serialize capacity and restore with other context mutations across workers. + async with storage.transaction(context_id) as tx: + # Check capacity before restoring + ctx = await tx.get_context(uuid.UUID(context_id)) + if ctx is not None: + capacity = await tx.get_capacity_status( + context_id, ctx.lifecycle_config.max_active_bullets + ) + if capacity.pressure_level == "full": + raise HTTPException( + status_code=409, + detail="Cannot restore: context at capacity", + ) + + bullet = await tx.restore_bullet(context_id, bullet_id) + if bullet is None: raise HTTPException( - status_code=409, - detail="Cannot restore: context at capacity", + status_code=404, + detail=f"Bullet {bullet_id} not found or not in archived state", ) - - bullet = await storage.restore_bullet(context_id, bullet_id) - if bullet is None: - raise HTTPException( - status_code=404, - detail=f"Bullet {bullet_id} not found or not in archived state", - ) return {"restored": True, "bullet_id": bullet_id} @@ -162,6 +166,9 @@ async def sync_context( if since: try: since_dt = datetime.fromisoformat(since) + # Legacy clients send naive ISO timestamps; interpret those as UTC. + if since_dt.tzinfo is None: + since_dt = since_dt.replace(tzinfo=timezone.utc) batches = [b for b in batches if b.timestamp > since_dt] except ValueError: raise HTTPException(status_code=400, detail="Invalid 'since' timestamp") @@ -213,19 +220,21 @@ async def get_ingestion_config(request: Request) -> dict: @config_router.put("/config/ingestion") -async def update_ingestion_config(request: Request) -> dict: +async def update_ingestion_config(request: Request, body: dict = Body(...)) -> dict: """Update server-level ingestion configuration. Note: changes only affect NEW commits. Existing bullets are NOT retroactively re-extracted. Use the re-extract endpoint for that. """ - body = await request.json() current = getattr(request.app.state, "ingestion_config", IngestionConfig()) # Merge: only update fields that are present in the request body updated_data = current.model_dump() updated_data.update(body) - new_config = IngestionConfig(**updated_data) + try: + new_config = IngestionConfig(**updated_data) + except ValidationError as exc: + raise RequestValidationError(exc.errors(), body=body) from exc # Store updated config on app state request.app.state.ingestion_config = new_config @@ -236,6 +245,15 @@ async def update_ingestion_config(request: Request) -> dict: ingestion.ingestion_config = new_config if hasattr(ingestion, "reflector") and ingestion.reflector is not None: ingestion.reflector.config = new_config + if getattr(ingestion, "curator", None) is not None: + ingestion.curator.ingestion_config = new_config + + re_extraction = getattr(request.app.state, "re_extraction", None) + if re_extraction is not None: + if getattr(re_extraction, "reflector", None) is not None: + re_extraction.reflector.config = new_config + if getattr(re_extraction, "curator", None) is not None: + re_extraction.curator.ingestion_config = new_config return { "updated": True, diff --git a/engram/server/routes/schemas.py b/engram/server/routes/schemas.py index 7d99cbc..ca684c3 100644 --- a/engram/server/routes/schemas.py +++ b/engram/server/routes/schemas.py @@ -30,6 +30,6 @@ async def get_schema( """Get a single schema by ID.""" storage = request.app.state.storage schema = await storage.get_schema(schema_id) - if schema is None: + if schema is None or schema.context_id != str(context_id): raise HTTPException(status_code=404, detail=f"Schema {schema_id} not found") return schema.model_dump(mode="json") diff --git a/engram/storage/base.py b/engram/storage/base.py index f0dc7ed..6d9e4c3 100644 --- a/engram/storage/base.py +++ b/engram/storage/base.py @@ -3,6 +3,7 @@ from __future__ import annotations import abc +from contextlib import asynccontextmanager import uuid from datetime import datetime from typing import Any @@ -34,6 +35,18 @@ async def initialize(self) -> None: async def close(self) -> None: """Release connections and resources.""" + @asynccontextmanager + async def transaction(self, context_id: str): + """Yield a task-owned backend for one atomic context mutation.""" + raise NotImplementedError(f"{type(self).__name__} does not support transactions") + yield # pragma: no cover + + async def update_bullet_embedding_if_missing( + self, context_id: str, bullet_id: str, content: str, embedding: list[float], + ) -> bool: + """Repair only a missing vector, retaining content, metadata and timestamps.""" + raise NotImplementedError(f"{type(self).__name__} does not support embedding repair") + # ── Context CRUD ─────────────────────────────────────────────────── @abc.abstractmethod @@ -135,6 +148,10 @@ async def list_schemas(self, context_id: str) -> list[SchemaNode]: async def update_schema(self, schema: SchemaNode) -> SchemaNode: """Update an existing schema.""" + async def remove_schema(self, context_id: str, schema_id: str) -> bool: + """Remove a schema and clear its bullet links within the given context.""" + raise NotImplementedError(f"{type(self).__name__} does not support schema removal") + # ── Delta History (v0.2) ─────────────────────────────────────────── @abc.abstractmethod diff --git a/engram/storage/postgres.py b/engram/storage/postgres.py index d887b80..0cb0b72 100644 --- a/engram/storage/postgres.py +++ b/engram/storage/postgres.py @@ -5,9 +5,13 @@ from __future__ import annotations +import asyncio +import copy +import hashlib import json import logging import uuid +from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone from typing import Any @@ -31,6 +35,9 @@ SlotDefinition, ) from engram.storage.base import StorageBackend +from engram.storage.transactions import ( + PostgresTransactionConnection, TransactionScope, finish, guarded_backend, +) logger = logging.getLogger(__name__) @@ -39,6 +46,17 @@ def _utcnow() -> datetime: return datetime.now(timezone.utc) +def _vector_to_list(value) -> list[float] | None: + """Normalize pgvector 0.5 Vector and earlier numpy-backed codec results.""" + if value is None: + return None + if hasattr(value, "to_list"): + return value.to_list() + if hasattr(value, "tolist"): + return value.tolist() + return list(value) + + def _get_rowcount(status: str) -> int: """Extract affected row count from asyncpg status string (e.g., 'UPDATE 3').""" try: @@ -47,22 +65,107 @@ def _get_rowcount(status: str) -> int: return 0 +@guarded_backend class PostgresBackend(StorageBackend): """PostgreSQL-based storage for production deployment with pgvector similarity search.""" - def __init__(self, dsn: str) -> None: + def __init__( + self, dsn: str, *, read_only: bool = False, existing_only: bool = False, + ) -> None: self.dsn = dsn + self.read_only = read_only + self.existing_only = existing_only or read_only self._pool: asyncpg.Pool | None = None - - async def _get_pool(self) -> asyncpg.Pool: - if self._pool is None: - self._pool = await asyncpg.create_pool( - self.dsn, - min_size=2, - max_size=10, - init=self._init_connection, - ) - return self._pool + self._pool_lock = asyncio.Lock() + self._scope: TransactionScope | None = None + self._connection = None + self._active_transactions: set[asyncio.Task] = set() + + @asynccontextmanager + async def _operation(self, name: str): + if self._scope is not None: + self._scope.check() + if name in {"initialize", "close"}: + raise RuntimeError(f"Cannot {name} a transaction handle") + elif asyncio.current_task() in self._active_transactions: + raise RuntimeError("Use the scoped transaction handle, not its parent") + elif name == "initialize" and self._active_transactions: + raise RuntimeError("Cannot initialize storage during a transaction") + yield + + async def _get_pool(self): + if self._scope is not None: + self._scope.check() + return PostgresTransactionConnection(self._connection, self._scope) + async with self._pool_lock: + if self._pool is None: + if not self.existing_only: + # Codecs require vector to exist before the first pooled connection. + bootstrap = await asyncpg.connect(self.dsn) + try: + await bootstrap.execute("CREATE EXTENSION IF NOT EXISTS vector") + finally: + await finish(bootstrap.close()) + options = {} + if self.read_only: + options["server_settings"] = {"default_transaction_read_only": "on"} + # Keep the awaitable Pool object while it initializes: asyncpg + # may open earlier connections before a later init hook fails. + pool = asyncpg.create_pool( + self.dsn, min_size=2, max_size=10, init=self._init_connection, **options, + ) + try: + await pool + except BaseException: + pool.terminate() + raise + self._pool = pool + return self._pool + + @asynccontextmanager + async def transaction(self, context_id: str): + if self.read_only: + raise RuntimeError("Read-only storage cannot start a write transaction") + if self._scope is not None: + self._scope.check() + if str(context_id) != self._scope.context_id: + raise ValueError("Nested transactions must use the same context") + scoped = copy.copy(self) + scoped._scope = TransactionScope(str(context_id), asyncio.current_task()) + try: + async with self._connection.transaction(): + yield scoped + finally: + scoped._scope.active = False + return + task = asyncio.current_task() + if task in self._active_transactions: + raise RuntimeError("Nested transactions require the scoped transaction handle") + pool = await self._get_pool() + self._active_transactions.add(task) + try: + async with pool.acquire() as connection: + scoped = copy.copy(self) + scoped._connection = connection + scoped._scope = TransactionScope(str(context_id), task) + try: + async with connection.transaction(): + # Stable across processes; collisions merely serialize unrelated contexts. + key = int.from_bytes( + hashlib.sha256(str(context_id).encode()).digest()[:8], "big", signed=True, + ) + await connection.execute("SELECT pg_advisory_xact_lock($1)", key) + # Take the parent lock before child rows. Otherwise a + # cascading context delete can deadlock against the + # parent FK check when this transaction saves its audit. + await connection.fetchrow( + "SELECT id FROM contexts WHERE id=$1 FOR KEY SHARE", str(context_id), + ) + yield scoped + finally: + scoped._scope.active = False + finally: + self._active_transactions.discard(task) @staticmethod async def _init_connection(conn: asyncpg.Connection) -> None: @@ -71,12 +174,12 @@ async def _init_connection(conn: asyncpg.Connection) -> None: # ── Lifecycle ────────────────────────────────────────────────────── async def initialize(self) -> None: + if self.read_only or self.existing_only: + raise RuntimeError("Existing-only storage cannot initialize or migrate") pool = await self._get_pool() async with pool.acquire() as conn: async with conn.transaction(): - await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - await conn.execute(""" CREATE TABLE IF NOT EXISTS contexts ( id TEXT PRIMARY KEY, @@ -288,16 +391,9 @@ async def initialize(self) -> None: await conn.execute( "CREATE INDEX IF NOT EXISTS idx_activities_timestamp ON activities(timestamp)" ) - await conn.execute( - "CREATE INDEX IF NOT EXISTS idx_bullets_lifecycle ON bullets(lifecycle_state)" - ) await conn.execute( "CREATE INDEX IF NOT EXISTS idx_bullets_archived_at ON bullets(archived_at)" ) - await conn.execute( - "CREATE INDEX IF NOT EXISTS idx_activities_raw_input_hash " - "ON activities(context_id, raw_input_hash)" - ) # Schema migrations for existing databases await self._migrate_v03(pool) @@ -318,11 +414,22 @@ async def _column_exists( return row is not None async def _migrate_v03(self, pool: asyncpg.Pool) -> None: + # ALTER and archive backfill must commit together. A failed backfill + # cannot leave a default-active column that makes the retry skip repair. + async with self.transaction("__engram_lifecycle_migration__") as tx: + await self._migrate_v03_columns(await tx._get_pool()) + + async def _migrate_v03_columns(self, pool: asyncpg.Pool) -> None: """Add v0.3 lifecycle columns if they don't exist.""" if not await self._column_exists(pool, "bullets", "lifecycle_state"): await pool.execute( "ALTER TABLE bullets ADD COLUMN lifecycle_state TEXT DEFAULT 'active'" ) + # Pre-v0.3 archives already carry is_archived. Preserve that state + # when adding the lifecycle column so restores/purges still work. + await pool.execute( + "UPDATE bullets SET lifecycle_state='archived' WHERE is_archived=TRUE" + ) logger.info("Added lifecycle_state column to bullets") if not await self._column_exists(pool, "bullets", "archive_reason"): @@ -387,9 +494,12 @@ async def _migrate_v05(self, pool: asyncpg.Pool) -> None: logger.info("Added reconsolidated_at column to materializations") async def close(self) -> None: - if self._pool is not None: - await self._pool.close() - self._pool = None + async with self._pool_lock: + if self._pool is not None: + try: + await finish(self._pool.close()) + finally: + self._pool = None # ── Context CRUD ─────────────────────────────────────────────────── @@ -562,9 +672,25 @@ async def add_bullet(self, context_id: str, bullet: Bullet) -> Bullet: ) return bullet + async def update_bullet_embedding_if_missing( + self, context_id: str, bullet_id: str, content: str, embedding: list[float], + ) -> bool: + pool = await self._get_pool() + status = await pool.execute( + "UPDATE bullets SET embedding=$1 WHERE id=$2 AND context_id=$3 AND content=$4 " + "AND embedding IS NULL AND is_active=TRUE AND is_archived=FALSE " + "AND lifecycle_state='active'", + np.array(embedding, dtype=np.float32), bullet_id, context_id, content, + ) + return _get_rowcount(status) > 0 + async def get_bullet(self, bullet_id: str) -> Bullet | None: pool = await self._get_pool() - row = await pool.fetchrow("SELECT * FROM bullets WHERE id = $1", bullet_id) + query = "SELECT * FROM bullets WHERE id = $1" + if self._scope is not None: + # Guard read-modify-write against direct archive/update operations too. + query += " FOR UPDATE" + row = await pool.fetchrow(query, bullet_id) return self._row_to_bullet(row) if row else None async def list_bullets( @@ -639,6 +765,7 @@ async def find_similar_bullets( AND is_active = TRUE AND is_archived = FALSE AND embedding IS NOT NULL + AND (embedding <=> $1::vector) BETWEEN 0 AND 2 AND 1 - (embedding <=> $1::vector) >= $3 ORDER BY similarity DESC LIMIT $4 @@ -706,6 +833,22 @@ async def update_schema(self, schema: SchemaNode) -> SchemaNode: schema.updated_at = now return schema + async def remove_schema(self, context_id: str, schema_id: str) -> bool: + """Remove one schema and clear its local bullet links atomically.""" + async with self.transaction(context_id) as tx: + schema = await tx.get_schema(schema_id) + if schema is None or schema.context_id != context_id: + return False + pool = await tx._get_pool() + await pool.execute( + "UPDATE bullets SET schema_id=NULL WHERE context_id=$1 AND schema_id=$2", + context_id, schema_id, + ) + status = await pool.execute( + "DELETE FROM schemas WHERE context_id=$1 AND id=$2", context_id, schema_id, + ) + return _get_rowcount(status) > 0 + # ── Delta History (v0.2) ─────────────────────────────────────────── async def save_delta_batch(self, delta_batch: DeltaBatch) -> DeltaBatch: @@ -814,48 +957,52 @@ async def restore_bullet(self, context_id: str, bullet_id: str) -> Bullet | None return await self.get_bullet(bullet_id) async def purge_bullet(self, context_id: str, bullet_id: str) -> bool: - pool = await self._get_pool() - # Delete connected edges first - await pool.execute( - "DELETE FROM edges WHERE context_id=$1 AND (from_node=$2 OR to_node=$3)", - context_id, bullet_id, bullet_id, - ) - status = await pool.execute( - "DELETE FROM bullets WHERE id=$1 AND context_id=$2", - bullet_id, context_id, - ) - return _get_rowcount(status) > 0 + # Keep candidate eligibility and all dependent deletions together. + async with self.transaction(context_id) as tx: + pool = await tx._get_pool() + # Delete connected edges first + await pool.execute( + "DELETE FROM edges WHERE context_id=$1 AND (from_node=$2 OR to_node=$3)", + context_id, bullet_id, bullet_id, + ) + status = await pool.execute( + "DELETE FROM bullets WHERE id=$1 AND context_id=$2", + bullet_id, context_id, + ) + return _get_rowcount(status) > 0 async def purge_expired_archives( self, context_id: str, purge_after_days: int = 180 ) -> int: - pool = await self._get_pool() - cutoff = _utcnow() - timedelta(days=purge_after_days) - - # Get bullet IDs to purge (for edge cleanup) - rows = await pool.fetch( - "SELECT id FROM bullets WHERE context_id=$1 AND lifecycle_state='archived' " - "AND archived_at IS NOT NULL AND archived_at < $2", - context_id, cutoff, - ) - bullet_ids = [row["id"] for row in rows] + # Keep candidate eligibility and all dependent deletions together. + async with self.transaction(context_id) as tx: + pool = await tx._get_pool() + cutoff = _utcnow() - timedelta(days=purge_after_days) + + # Get bullet IDs to purge (for edge cleanup) + rows = await pool.fetch( + "SELECT id FROM bullets WHERE context_id=$1 AND lifecycle_state='archived' " + "AND archived_at IS NOT NULL AND archived_at < $2 ORDER BY id FOR UPDATE", + context_id, cutoff, + ) + bullet_ids = [row["id"] for row in rows] - if not bullet_ids: - return 0 + if not bullet_ids: + return 0 - # Delete connected edges using ANY() - await pool.execute( - "DELETE FROM edges WHERE context_id=$1 AND " - "(from_node = ANY($2) OR to_node = ANY($2))", - context_id, bullet_ids, - ) + # Delete connected edges using ANY() + await pool.execute( + "DELETE FROM edges WHERE context_id=$1 AND " + "(from_node = ANY($2) OR to_node = ANY($2))", + context_id, bullet_ids, + ) - # Delete the bullets - status = await pool.execute( - "DELETE FROM bullets WHERE context_id=$1 AND id = ANY($2)", - context_id, bullet_ids, - ) - return _get_rowcount(status) + # Delete the bullets + status = await pool.execute( + "DELETE FROM bullets WHERE context_id=$1 AND id = ANY($2)", + context_id, bullet_ids, + ) + return _get_rowcount(status) async def get_archived_bullets( self, context_id: str, offset: int = 0, limit: int = 50 @@ -1004,6 +1151,7 @@ async def find_similar_concepts( WHERE context_id = $2 AND is_valid = TRUE AND embedding IS NOT NULL + AND (embedding <=> $1::vector) BETWEEN 0 AND 2 AND 1 - (embedding <=> $1::vector) >= $3 ORDER BY similarity DESC LIMIT $4 @@ -1195,7 +1343,7 @@ async def find_similar_activities( @staticmethod def _row_to_bullet(row: asyncpg.Record) -> Bullet: - keys = row.keys() + keys = set(row.keys()) lifecycle_state_raw = row["lifecycle_state"] if "lifecycle_state" in keys else "active" archive_reason = row["archive_reason"] if "archive_reason" in keys else None emb = row["embedding"] @@ -1203,7 +1351,7 @@ def _row_to_bullet(row: asyncpg.Record) -> Bullet: id=row["id"], context_id=row["context_id"], section=row["section"], content=row["content"], bullet_type=row["bullet_type"], source_type=row["source_type"], - embedding=emb.tolist() if emb is not None else None, + embedding=_vector_to_list(emb), hit_count=row["hit_count"], miss_count=row["miss_count"], recall_count=row["recall_count"], last_recalled_at=row["last_recalled_at"], @@ -1252,7 +1400,8 @@ def _row_to_concept(row: asyncpg.Record) -> ConceptNode: emb = row["embedding"] return ConceptNode( id=uuid.UUID(row["id"]), type=row["type"], content=row["content"], - embedding=emb.tolist() if emb is not None else None, + context_id=uuid.UUID(row["context_id"]), + embedding=_vector_to_list(emb), confidence=row["confidence"], salience=row["salience"], created_at=row["created_at"], updated_at=row["updated_at"], @@ -1279,7 +1428,7 @@ def _row_to_edge(row: asyncpg.Record) -> ConceptEdge: @staticmethod def _row_to_activity(row: asyncpg.Record) -> Activity: - keys = row.keys() + keys = set(row.keys()) # v0.4 fields — backward compatible via key check raw_input = row["raw_input"] if "raw_input" in keys else None raw_input_hash = row["raw_input_hash"] if "raw_input_hash" in keys else None diff --git a/engram/storage/sqlite.py b/engram/storage/sqlite.py index bbf06dd..54eb513 100644 --- a/engram/storage/sqlite.py +++ b/engram/storage/sqlite.py @@ -2,10 +2,14 @@ from __future__ import annotations +import asyncio +import copy import json import logging import uuid +from contextlib import asynccontextmanager from datetime import datetime, timezone +from pathlib import Path from typing import Any import aiosqlite @@ -27,6 +31,9 @@ ) from engram.core.similarity import cosine_similarity as _cosine_similarity from engram.storage.base import StorageBackend +from engram.storage.transactions import ( + SQLiteTransactionConnection, TaskGate, TransactionScope, finish, guarded_backend, +) logger = logging.getLogger(__name__) @@ -35,24 +42,118 @@ def _utcnow() -> datetime: return datetime.now(timezone.utc) +@guarded_backend class SQLiteBackend(StorageBackend): """SQLite-based storage for local development without Docker dependencies.""" - def __init__(self, db_path: str = "./engram.db") -> None: + def __init__( + self, db_path: str = "./engram.db", *, read_only: bool = False, + existing_only: bool = False, + ) -> None: self.db_path = db_path + self.read_only = read_only + self.existing_only = existing_only or read_only self._db: aiosqlite.Connection | None = None - - async def _get_db(self) -> aiosqlite.Connection: - if self._db is None: - self._db = await aiosqlite.connect(self.db_path) - self._db.row_factory = aiosqlite.Row - await self._db.execute("PRAGMA journal_mode=WAL") - await self._db.execute("PRAGMA foreign_keys=ON") - return self._db + self._connect_lock = asyncio.Lock() + self._gate = TaskGate() + self._scope: TransactionScope | None = None + self._transaction_owner: asyncio.Task | None = None + + @asynccontextmanager + async def _operation(self, name: str): + if self._scope is not None: + self._scope.check() + if name in {"initialize", "close"}: + raise RuntimeError(f"Cannot {name} a transaction handle") + yield + return + if self._transaction_owner is asyncio.current_task(): + raise RuntimeError("Use the scoped transaction handle, not its parent") + outermost = self._gate.owner is not asyncio.current_task() + async with self._gate.hold(): + try: + yield + except BaseException: + # Ordinary methods also must not leave an implicit write open + # for a later unrelated caller to commit after cancellation. + if outermost and self._db is not None: + await finish(self._db.rollback()) + raise + + async def _get_db(self): + if self._scope is not None: + self._scope.check() + return SQLiteTransactionConnection(self._db, self._scope) + async with self._connect_lock: + if self._db is None: + path = self.db_path + if self.existing_only: + if path == ":memory:": + raise ValueError("Existing-only storage requires an existing SQLite file") + mode = "ro" if self.read_only else "rw" + path = Path(path).expanduser().resolve().as_uri() + f"?mode={mode}" + db = await aiosqlite.connect(path, uri=self.existing_only, timeout=30) + try: + db.row_factory = aiosqlite.Row + if not self.read_only: + await db.execute("PRAGMA journal_mode=WAL") + await db.execute("PRAGMA foreign_keys=ON") + except BaseException: + await finish(db.close()) + raise + self._db = db + return self._db + + @asynccontextmanager + async def transaction(self, context_id: str): + if self.read_only: + raise RuntimeError("Read-only storage cannot start a write transaction") + if self._scope is not None: + self._scope.check() + if str(context_id) != self._scope.context_id: + raise ValueError("Nested transactions must use the same context") + db = self._db + savepoint = "engram_" + uuid.uuid4().hex + scoped = copy.copy(self) + scoped._scope = TransactionScope(str(context_id), asyncio.current_task()) + try: + await db.execute(f"SAVEPOINT {savepoint}") + yield scoped + except BaseException: + await finish(db.execute(f"ROLLBACK TO SAVEPOINT {savepoint}")) + await finish(db.execute(f"RELEASE SAVEPOINT {savepoint}")) + raise + else: + await finish(db.execute(f"RELEASE SAVEPOINT {savepoint}")) + finally: + scoped._scope.active = False + return + if self._transaction_owner is asyncio.current_task(): + raise RuntimeError("Nested transactions require the scoped transaction handle") + async with self._gate.hold(): + db = await self._get_db() + scoped = copy.copy(self) + scoped._scope = TransactionScope(str(context_id), asyncio.current_task()) + self._transaction_owner = asyncio.current_task() + try: + # Acquire the database writer lock before any reads, including + # across independent connections/processes using the same file. + try: + await db.execute("BEGIN IMMEDIATE") + yield scoped + await finish(db.commit()) + except BaseException: + await finish(db.rollback()) + raise + finally: + scoped._scope.active = False + self._transaction_owner = None # ── Lifecycle ────────────────────────────────────────────────────── async def initialize(self) -> None: + if self.read_only or self.existing_only: + raise RuntimeError("Existing-only storage cannot initialize or migrate") db = await self._get_db() await db.executescript( """ @@ -227,6 +328,12 @@ async def initialize(self) -> None: logger.info("SQLite storage initialized at %s", self.db_path) async def _migrate_v03(self, db: aiosqlite.Connection) -> None: + # ALTER and archive backfill must commit together. A failed backfill + # cannot leave a default-active column that makes the retry skip repair. + async with self.transaction("__engram_lifecycle_migration__") as tx: + await self._migrate_v03_columns(await tx._get_db()) + + async def _migrate_v03_columns(self, db: aiosqlite.Connection) -> None: """Add v0.3 lifecycle columns if they don't exist.""" # Check if lifecycle_state column exists on bullets cursor = await db.execute("PRAGMA table_info(bullets)") @@ -236,6 +343,11 @@ async def _migrate_v03(self, db: aiosqlite.Connection) -> None: await db.execute( "ALTER TABLE bullets ADD COLUMN lifecycle_state TEXT DEFAULT 'active'" ) + # Pre-v0.3 archives already carry is_archived. Preserve that state + # when adding the lifecycle column so restores/purges still work. + await db.execute( + "UPDATE bullets SET lifecycle_state='archived' WHERE is_archived=1" + ) logger.info("Added lifecycle_state column to bullets") if "archive_reason" not in cols: @@ -318,8 +430,10 @@ async def _migrate_v05(self, db: aiosqlite.Connection) -> None: async def close(self) -> None: if self._db is not None: - await self._db.close() - self._db = None + try: + await finish(self._db.close()) + finally: + self._db = None # ── Context CRUD ─────────────────────────────────────────────────── @@ -491,6 +605,18 @@ async def add_bullet(self, context_id: str, bullet: Bullet) -> Bullet: await db.commit() return bullet + async def update_bullet_embedding_if_missing( + self, context_id: str, bullet_id: str, content: str, embedding: list[float], + ) -> bool: + db = await self._get_db() + cursor = await db.execute( + "UPDATE bullets SET embedding=? WHERE id=? AND context_id=? AND content=? " + "AND embedding IS NULL AND is_active=1 AND is_archived=0 AND lifecycle_state='active'", + (json.dumps(embedding), bullet_id, context_id, content), + ) + await db.commit() + return cursor.rowcount > 0 + async def get_bullet(self, bullet_id: str) -> Bullet | None: db = await self._get_db() cursor = await db.execute("SELECT * FROM bullets WHERE id = ?", (bullet_id,)) @@ -626,6 +752,22 @@ async def update_schema(self, schema: SchemaNode) -> SchemaNode: schema.updated_at = datetime.fromisoformat(now) return schema + async def remove_schema(self, context_id: str, schema_id: str) -> bool: + """Remove one schema and clear its local bullet links atomically.""" + async with self.transaction(context_id) as tx: + schema = await tx.get_schema(schema_id) + if schema is None or schema.context_id != context_id: + return False + db = await tx._get_db() + await db.execute( + "UPDATE bullets SET schema_id=NULL WHERE context_id=? AND schema_id=?", + (context_id, schema_id), + ) + cursor = await db.execute( + "DELETE FROM schemas WHERE context_id=? AND id=?", (context_id, schema_id), + ) + return cursor.rowcount > 0 + # ── Delta History (v0.2) ─────────────────────────────────────────── async def save_delta_batch(self, delta_batch: DeltaBatch) -> DeltaBatch: @@ -736,52 +878,56 @@ async def restore_bullet(self, context_id: str, bullet_id: str) -> Bullet | None return await self.get_bullet(bullet_id) async def purge_bullet(self, context_id: str, bullet_id: str) -> bool: - db = await self._get_db() - # Delete connected edges first - await db.execute( - "DELETE FROM edges WHERE context_id=? AND (from_node=? OR to_node=?)", - (context_id, bullet_id, bullet_id), - ) - result = await db.execute( - "DELETE FROM bullets WHERE id=? AND context_id=?", - (bullet_id, context_id), - ) - await db.commit() - return result.rowcount > 0 # type: ignore[union-attr] + # Keep candidate eligibility and all dependent deletions together. + async with self.transaction(context_id) as tx: + db = await tx._get_db() + # Delete connected edges first + await db.execute( + "DELETE FROM edges WHERE context_id=? AND (from_node=? OR to_node=?)", + (context_id, bullet_id, bullet_id), + ) + result = await db.execute( + "DELETE FROM bullets WHERE id=? AND context_id=?", + (bullet_id, context_id), + ) + await db.commit() + return result.rowcount > 0 # type: ignore[union-attr] async def purge_expired_archives( self, context_id: str, purge_after_days: int = 180 ) -> int: - db = await self._get_db() - from datetime import timedelta - cutoff = (_utcnow() - timedelta(days=purge_after_days)).isoformat() - - # Get bullet IDs to purge (for edge cleanup) - cursor = await db.execute( - "SELECT id FROM bullets WHERE context_id=? AND lifecycle_state='archived' " - "AND archived_at IS NOT NULL AND archived_at < ?", - (context_id, cutoff), - ) - bullet_ids = [row[0] for row in await cursor.fetchall()] + # Keep candidate eligibility and all dependent deletions together. + async with self.transaction(context_id) as tx: + db = await tx._get_db() + from datetime import timedelta + cutoff = (_utcnow() - timedelta(days=purge_after_days)).isoformat() + + # Get bullet IDs to purge (for edge cleanup) + cursor = await db.execute( + "SELECT id FROM bullets WHERE context_id=? AND lifecycle_state='archived' " + "AND archived_at IS NOT NULL AND archived_at < ?", + (context_id, cutoff), + ) + bullet_ids = [row[0] for row in await cursor.fetchall()] - if not bullet_ids: - return 0 + if not bullet_ids: + return 0 - # Delete connected edges - placeholders = ",".join("?" * len(bullet_ids)) - await db.execute( - f"DELETE FROM edges WHERE context_id=? AND " - f"(from_node IN ({placeholders}) OR to_node IN ({placeholders}))", - [context_id] + bullet_ids + bullet_ids, - ) + # Delete connected edges + placeholders = ",".join("?" * len(bullet_ids)) + await db.execute( + f"DELETE FROM edges WHERE context_id=? AND " + f"(from_node IN ({placeholders}) OR to_node IN ({placeholders}))", + [context_id] + bullet_ids + bullet_ids, + ) - # Delete the bullets - result = await db.execute( - f"DELETE FROM bullets WHERE context_id=? AND id IN ({placeholders})", - [context_id] + bullet_ids, - ) - await db.commit() - return result.rowcount or 0 # type: ignore[union-attr] + # Delete the bullets + result = await db.execute( + f"DELETE FROM bullets WHERE context_id=? AND id IN ({placeholders})", + [context_id] + bullet_ids, + ) + await db.commit() + return result.rowcount or 0 # type: ignore[union-attr] async def get_archived_bullets( self, context_id: str, offset: int = 0, limit: int = 50 @@ -1173,6 +1319,7 @@ def _row_to_delta_batch(row: aiosqlite.Row) -> DeltaBatch: def _row_to_concept(row: aiosqlite.Row) -> ConceptNode: return ConceptNode( id=uuid.UUID(row["id"]), type=row["type"], content=row["content"], + context_id=uuid.UUID(row["context_id"]), embedding=json.loads(row["embedding"]) if row["embedding"] else None, confidence=row["confidence"], salience=row["salience"], created_at=datetime.fromisoformat(row["created_at"]), diff --git a/engram/storage/transactions.py b/engram/storage/transactions.py new file mode 100644 index 0000000..a3a0cec --- /dev/null +++ b/engram/storage/transactions.py @@ -0,0 +1,135 @@ +"""Explicit, task-owned transaction handles shared by the storage backends.""" + +from __future__ import annotations + +import asyncio +import functools +import inspect +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Any + + +@dataclass +class TransactionScope: + context_id: str + owner: asyncio.Task + active: bool = True + + def check(self) -> None: + if not self.active: + raise RuntimeError("Transaction handle has expired") + if self.owner is not asyncio.current_task(): + raise RuntimeError("Transaction handle belongs to another task") + + +class TaskGate: + """A reentrant task-owned gate; ownership is never inherited by child tasks.""" + + def __init__(self) -> None: + self.lock = asyncio.Lock() + self.owner: asyncio.Task | None = None + + @asynccontextmanager + async def hold(self): + task = asyncio.current_task() + if self.owner is task: + yield + return + async with self.lock: + self.owner = task + try: + yield + finally: + self.owner = None + + +def guarded_backend(cls): + """Guard every public coroutine, including edition-specific auth methods. + + Transactions themselves are async context managers and implement their own + guard. Private SQL helpers run only inside a guarded public operation. + """ + for name, method in tuple(vars(cls).items()): + if name.startswith("_") or not inspect.iscoroutinefunction(method): + continue + + def wrap(fn, method_name): + @functools.wraps(fn) + async def guarded(self, *args, **kwargs): + async with self._operation(method_name): + return await fn(self, *args, **kwargs) + + return guarded + + setattr(cls, name, wrap(method, name)) + return cls + + +async def finish(awaitable): + """Finish connection cleanup before releasing ownership, even if cancelled.""" + task = asyncio.ensure_future(awaitable) + cancellation = None + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError as exc: + cancellation = exc + result = task.result() + if cancellation is not None: + raise cancellation + return result + + +class SQLiteTransactionConnection: + def __init__(self, connection, scope: TransactionScope): + self.connection = connection + self.scope = scope + + async def execute(self, *args, **kwargs): + self.scope.check() + return await self.connection.execute(*args, **kwargs) + + async def executemany(self, *args, **kwargs): + self.scope.check() + return await self.connection.executemany(*args, **kwargs) + + async def commit(self): + self.scope.check() # only the outer transaction may commit + + +class PostgresTransactionConnection: + """The pool-shaped executor expected by existing PostgreSQL methods.""" + + def __init__(self, connection, scope: TransactionScope): + self.connection = connection + self.scope = scope + + async def execute(self, *args, **kwargs): + self.scope.check() + return await self.connection.execute(*args, **kwargs) + + async def executemany(self, *args, **kwargs): + self.scope.check() + return await self.connection.executemany(*args, **kwargs) + + async def fetch(self, *args, **kwargs): + self.scope.check() + return await self.connection.fetch(*args, **kwargs) + + async def fetchrow(self, *args, **kwargs): + self.scope.check() + return await self.connection.fetchrow(*args, **kwargs) + + async def fetchval(self, *args, **kwargs): + self.scope.check() + return await self.connection.fetchval(*args, **kwargs) + + @asynccontextmanager + async def acquire(self): + self.scope.check() + yield self + + def transaction(self, **kwargs: Any): + self.scope.check() + return self.connection.transaction(**kwargs) diff --git a/evals/__init__.py b/evals/__init__.py new file mode 100644 index 0000000..4cc65fb --- /dev/null +++ b/evals/__init__.py @@ -0,0 +1 @@ +"""Source-checkout retrieval regression evaluations.""" diff --git a/evals/corpus.json b/evals/corpus.json new file mode 100644 index 0000000..7226d53 --- /dev/null +++ b/evals/corpus.json @@ -0,0 +1,451 @@ +{ + "version": 1, + "benchmark": "controlled_regression", + "description": "Curated judgments and synthetic vectors; not measured real-provider retrieval quality.", + "renderers": [ + "claude", + "gpt", + "generic" + ], + "cases": [ + { + "id": "semantic_over_salience", + "query": "How do we retry a throttled request?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [ + { + "id": "retry", + "text": "Retry throttled requests with exponential backoff.", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.1 + }, + { + "id": "office", + "text": "The office kitchen stocks coffee beans.", + "vector": [ + 0, + 1, + 0 + ], + "salience": 1 + } + ], + "relevance": { + "retry": 3 + }, + "k": 1, + "token_budget": 500, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 1 + } + }, + { + "id": "irrelevant_distractors", + "query": "How are API credentials stored and rotated?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [ + { + "id": "store", + "text": "Store API credentials in the secret manager.", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.5 + }, + { + "id": "rotate", + "text": "Rotate API credentials every thirty days.", + "vector": [ + 0.95, + 0.312, + 0 + ], + "salience": 0.5 + }, + { + "id": "lunch", + "text": "Lunch is served at noon.", + "vector": [ + 0, + 1, + 0 + ], + "salience": 1 + }, + { + "id": "colors", + "text": "The landing page uses a blue background.", + "vector": [ + 0, + 0, + 1 + ], + "salience": 1 + } + ], + "relevance": { + "store": 3, + "rotate": 2 + }, + "k": 2, + "token_budget": 500, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 1 + } + }, + { + "id": "duplicate_and_diverse", + "query": "How do we make releases recoverable?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [ + { + "id": "canary", + "text": "Deploy a canary before the full rollout.", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.5 + }, + { + "id": "duplicate", + "text": "Run a canary release before broad deployment.", + "vector": [ + 0.999, + 0.0447, + 0 + ], + "salience": 0.5 + }, + { + "id": "rollback", + "text": "Retain the prior image for immediate rollback.", + "vector": [ + 0.8, + 0, + 0.6 + ], + "salience": 0.5 + } + ], + "relevance": { + "canary": 3, + "duplicate": 1, + "rollback": 2 + }, + "k": 3, + "token_budget": 500, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 0.9 + } + }, + { + "id": "lifecycle_exclusion", + "query": "Which support escalation policy is current?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [ + { + "id": "current", + "text": "Current support escalations go to the on-call lead.", + "vector": [ + 0.9, + 0.436, + 0 + ], + "salience": 0.5 + }, + { + "id": "archived", + "text": "Archived escalations go to the former manager.", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.5, + "is_archived": true + }, + { + "id": "inactive", + "text": "Inactive escalations go to the retired queue.", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.5, + "is_active": false + } + ], + "relevance": { + "current": 3 + }, + "k": 1, + "token_budget": 500, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 1 + }, + "forbidden_ids": [ + "archived", + "inactive" + ] + }, + { + "id": "tenant_isolation", + "query": "Where is this project deployed?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [ + { + "id": "local", + "text": "This project deploys in the west region.", + "vector": [ + 0.9, + 0.436, + 0 + ], + "salience": 0.5 + }, + { + "id": "foreign", + "text": "Another tenant deploys in the secret east region.", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.5, + "foreign": true + } + ], + "relevance": { + "local": 3 + }, + "k": 1, + "token_budget": 500, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 1 + }, + "forbidden_ids": [ + "foreign" + ] + }, + { + "id": "oversized_candidate", + "query": "What is the safe fallback during an outage?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [ + { + "id": "oversized", + "text": "Historical outage log: Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. Repeated diagnostic details. ", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.5 + }, + { + "id": "fallback", + "text": "Fallback to the read-only replica.", + "vector": [ + 0.9, + 0.436, + 0 + ], + "salience": 0.5 + } + ], + "relevance": { + "fallback": 3 + }, + "k": 1, + "token_budget": 70, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 1 + } + }, + { + "id": "tight_budget_receipt", + "query": "How should we release the service?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [ + { + "id": "first", + "text": "Use canary releases.", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.5 + }, + { + "id": "second", + "text": "Monitor error rates.", + "vector": [ + 0.9, + 0.436, + 0 + ], + "salience": 0.5 + }, + { + "id": "third", + "text": "Keep rollback ready.", + "vector": [ + 0.8, + 0, + 0.6 + ], + "salience": 0.5 + } + ], + "relevance": { + "first": 3 + }, + "k": 1, + "token_budget": 45, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 1 + } + }, + { + "id": "empty_context", + "query": "What facts have been recorded?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [], + "relevance": {}, + "k": 1, + "token_budget": 500, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 1 + }, + "expect_empty": true + }, + { + "id": "legacy_oversized", + "query": "What is the database recovery action?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [ + { + "id": "legacy_large", + "text": "Old recovery log: Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. Diagnostic detail. ", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.5 + }, + { + "id": "legacy_small", + "text": "Restore the verified backup.", + "vector": [ + 0.9, + 0.436, + 0 + ], + "salience": 0.5 + } + ], + "relevance": { + "legacy_small": 3 + }, + "k": 1, + "token_budget": 70, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 1 + }, + "kind": "legacy" + }, + { + "id": "schema_summary", + "query": "What is our reliability pattern?", + "query_vector": [ + 1, + 0, + 0 + ], + "items": [ + { + "id": "schema_fact", + "text": "Verify the failover before declaring recovery.", + "vector": [ + 1, + 0, + 0 + ], + "salience": 0.5 + } + ], + "relevance": { + "schema_fact": 3 + }, + "k": 1, + "token_budget": 500, + "thresholds": { + "recall_at_k": 1, + "ndcg_at_k": 1 + }, + "schemas": [ + { + "id": "recovery_pattern", + "name": "Recovery", + "description": "Fail over, verify, and report recovery." + } + ] + } + ] +} diff --git a/evals/retrieval.py b/evals/retrieval.py new file mode 100644 index 0000000..1fc7c97 --- /dev/null +++ b/evals/retrieval.py @@ -0,0 +1,302 @@ +"""Controlled retrieval regression benchmark using real SQLite and materialization. + +The default uses offline embedding fixtures. Synthetic vectors control ranking +independently of provider changes; the text judgments also support live embeddings. +""" + +from __future__ import annotations + +import argparse +import asyncio +import inspect +import json +import math +import os +import tempfile +import uuid +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from engram.core.materialization import MaterializationEngine +from engram.core.models import Bullet, ConceptNode, ConceptType, Context, IntentAnchor, SchemaNode +from engram.llm.adapter import LLMAdapter +from engram.storage.sqlite import SQLiteBackend + +CORPUS = Path(__file__).with_name("corpus.json") + + +class FixtureEmbeddings(LLMAdapter): + """Exact text lookup: unknown text fails rather than silently inventing vectors.""" + + def __init__(self, vectors: dict[str, list[float]]) -> None: + self.vectors = vectors + + async def embed(self, text: str) -> list[float]: + return self.vectors[text] + + async def complete(self, *args: Any, **kwargs: Any) -> str: + raise AssertionError("Retrieval evaluations never request completions") + + +def ranking_metrics(ranked: list[str], relevance: dict[str, int], k: int) -> dict: + """Binary recall and graded nDCG; duplicate IDs earn no additional credit. + + Empty judgments return null, and are excluded from aggregate averages. + Empty behavior is checked separately so an all-empty corpus cannot pass. + """ + if k <= 0 or any(grade < 0 for grade in relevance.values()): + raise ValueError("k must be positive and relevance grades nonnegative") + relevant = {key for key, grade in relevance.items() if grade > 0} + if not relevant: + return {"recall_at_k": None, "ndcg_at_k": None} + seen = set() + dcg = 0.0 + for rank, key in enumerate(ranked[:k]): + grade = relevance.get(key, 0) if key not in seen else 0 + dcg += (2**grade - 1) / math.log2(rank + 2) + seen.add(key) + ideal = sum( + (2**grade - 1) / math.log2(rank + 2) + for rank, grade in enumerate(sorted(relevance.values(), reverse=True)[:k]) + ) + return {"recall_at_k": len(seen & relevant) / len(relevant), "ndcg_at_k": dcg / ideal} + + +def assess_result( + case: dict, + result: dict, + receipt: Any, + renderer: Any, + id_map: dict[str, str], + schema_map: dict[str, str], +) -> dict: + """Assess the actual packed order, full rendered texts, and persisted receipt.""" + raw_ids = [str(key) for key in result["bullets_included"] + result["concepts_included"]] + reverse_ids = {value: key for key, value in id_map.items()} + ranked = [reverse_ids.get(key, key) for key in raw_ids] + text = result["rendered_text"] + rendered_ids = [item["id"] for item in case["items"] if item["text"] in text] + forbidden = set(case.get("forbidden_ids", [])) + leakage = sorted(forbidden & (set(ranked) | set(rendered_ids))) + metrics = ranking_metrics(ranked, case["relevance"], case["k"]) + errors = [] + for metric, threshold in case["thresholds"].items(): + if metrics[metric] is not None and metrics[metric] + 1e-12 < threshold: + errors.append(metric) + if leakage: + errors.append("forbidden_id_leakage") + if len(raw_ids) != len(set(raw_ids)) or set(ranked) != set(rendered_ids): + errors.append("rendered_id_consistency") + expected_schemas = {key for key, content in schema_map.items() if content in text} + if set(result["schemas_included"]) != expected_schemas: + errors.append("schema_consistency") + if ( + receipt is None + or receipt.bullets_included != result["bullets_included"] + or receipt.token_count != result["token_count"] + ): + errors.append("persisted_receipt_consistency") + actual_tokens = renderer.estimate_tokens(text) + if result["token_count"] != actual_tokens or actual_tokens > case["token_budget"]: + errors.append("token_budget_consistency") + if case.get("expect_empty") and (ranked or rendered_ids): + errors.append("expected_empty") + return { + "case": case["id"], + "k": case["k"], + **metrics, + "ranked_ids": ranked, + "rendered_ids": rendered_ids, + "forbidden_id_leakage": leakage, + "token_count": actual_tokens, + "token_budget": case["token_budget"], + "token_budget_consistent": "token_budget_consistency" not in errors, + "receipt_consistent": not any( + error in errors + for error in ( + "rendered_id_consistency", + "schema_consistency", + "persisted_receipt_consistency", + ) + ), + "passed": not errors, + "failures": errors, + } + + +async def run_case( + case: dict, target_model: str, storage: SQLiteBackend, provider: LLMAdapter | None = None +) -> dict: + vectors = {case["query"]: case["query_vector"]} + vectors.update({item["text"]: item["vector"] for item in case["items"]}) + if provider is not None: + # Prefetch all provider vectors, including the query. Provider errors + # must fail the benchmark, not trigger the engine's salience fallback. + vectors = {text: await provider.embed(text) for text in vectors} + dimensions = {len(vector) for vector in vectors.values()} + if ( + len(dimensions) != 1 + or not next(iter(dimensions)) + or any( + not all(math.isfinite(value) for value in vector) or not any(vector) + for vector in vectors.values() + ) + ): + raise ValueError("Provider returned invalid or inconsistent embedding vectors") + adapter = FixtureEmbeddings(vectors) + context = Context(name="retrieval-eval", intent=IntentAnchor(objective="Evaluate retrieval")) + foreign = Context(name="foreign-eval", intent=IntentAnchor(objective="Tenant isolation")) + await storage.create_context(context) + await storage.create_context(foreign) + id_map = {} + timestamp = datetime(2025, 1, 1, tzinfo=UTC) + for item in case["items"]: + target = foreign.id if item.get("foreign") else context.id + key = str(uuid.uuid4()) + id_map[item["id"]] = key + if case.get("kind") == "legacy": + await storage.add_concept( + target, + ConceptNode( + id=uuid.UUID(key), + type=ConceptType.FACT, + content=item["text"], + embedding=vectors[item["text"]], + salience=item["salience"], + created_at=timestamp, + ), + ) + else: + await storage.add_bullet( + str(target), + Bullet( + id=key, + content=item["text"], + embedding=vectors[item["text"]], + salience=item["salience"], + created_at=timestamp, + is_active=item.get("is_active", True), + is_archived=item.get("is_archived", False), + ), + ) + schema_map = {} + for item in case.get("schemas", []): + schema = SchemaNode(name=item["name"], description=item["description"]) + await storage.add_schema(str(context.id), schema) + schema_map[schema.id] = f"[Pattern: {schema.name}] {schema.description}" + engine = MaterializationEngine(storage, adapter) + options = {} + if "include_worked_examples" in inspect.signature(engine.materialize).parameters: + options["include_worked_examples"] = False + result = await engine.materialize( + context.id, + query=case["query"], + target_model=target_model, + token_budget=case["token_budget"], + include_intent=False, + recency_weight=0, + **options, + ) + receipt = await storage.get_materialization(result["materialization_id"]) + assessment = assess_result( + case, result, receipt, engine._get_renderer(target_model), id_map, schema_map + ) + assessment["renderer"] = target_model + return assessment + + +def summarize(results: list[dict], expected_count: int) -> dict: + judged = [row for row in results if row["recall_at_k"] is not None] + return { + "passed": bool(judged) + and len(results) == expected_count + and all(row["passed"] for row in results), + "scenario_renderer_count": len(results), + "judged_count": len(judged), + "mean_recall_at_k": sum(row["recall_at_k"] for row in judged) / len(judged) + if judged + else None, + "mean_ndcg_at_k": sum(row["ndcg_at_k"] for row in judged) / len(judged) if judged else None, + "forbidden_id_leakage_count": sum(len(row["forbidden_id_leakage"]) for row in results), + } + + +async def evaluate(corpus: dict, provider: LLMAdapter | None = None) -> dict: + # No configured paths, environment loading, or application startup: this + # runner can only touch a brand-new scratch database, including in live mode. + with tempfile.TemporaryDirectory(prefix="engram-retrieval-") as directory: + storage = SQLiteBackend(str(Path(directory) / "evaluation.db")) + await storage.initialize() + try: + live = provider is not None + if provider is not None: + texts = dict.fromkeys( + text + for case in corpus["cases"] + for text in [case["query"], *(item["text"] for item in case["items"])] + ) + provider = FixtureEmbeddings({text: await provider.embed(text) for text in texts}) + results = [ + await run_case(case, renderer, storage, provider) + for case in corpus["cases"] + for renderer in corpus["renderers"] + ] + finally: + await storage.close() + summary = summarize(results, len(corpus["cases"]) * len(corpus["renderers"])) + return { + "benchmark": "live_provider_labeled_corpus" if live else "controlled_regression", + "corpus_version": corpus["version"], + "quality_claim": "Curated scratch data; offline vectors do not measure provider quality", + "summary": summary, + "results": results, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, help="Write machine-readable JSON (also printed)") + parser.add_argument("--corpus", type=Path, default=CORPUS) + parser.add_argument( + "--live-embeddings", + action="store_true", + help="Explicitly allow billable provider embeddings of curated texts only", + ) + parser.add_argument("--embedding-model") + parser.add_argument("--api-key-env", help="Environment variable holding the provider key") + args = parser.parse_args(argv) + if args.live_embeddings: + if not args.embedding_model or not args.api_key_env or not os.environ.get(args.api_key_env): + parser.error("live mode requires --embedding-model and a populated --api-key-env") + elif args.embedding_model or args.api_key_env: + parser.error("provider options require --live-embeddings") + try: + corpus = json.loads(args.corpus.read_text()) + if args.live_embeddings: + from engram.llm.adapter import LiteLLMAdapter + from engram.llm.explicit import explicit_embedding_runtime + + with explicit_embedding_runtime(): + provider = LiteLLMAdapter( + embedding_model=args.embedding_model, + embedding_api_key=os.environ[args.api_key_env], + ) + report = asyncio.run(evaluate(corpus, provider)) + report["embedding_model"] = args.embedding_model + else: + report = asyncio.run(evaluate(corpus)) + except Exception as exc: + # Provider exception messages can contain request credentials. Never + # print them; the exception type is enough to diagnose a failed gate. + report = {"summary": {"passed": False}, "error_type": type(exc).__name__} + encoded = json.dumps(report, indent=2) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(encoded) + print(encoded, end="") + return 0 if report["summary"]["passed"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 6154dee..516ad2d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["hatchling"] build-backend = "hatchling.build" [project] -name = "engram" +name = "engram-contextdb" version = "0.5.0" description = "A brain-inspired, portable context database for AI agents" readme = "README.md" @@ -48,7 +48,7 @@ dev = [ "mypy>=1.13.0", ] mcp = [ - "mcp>=1.0.0", + "mcp>=1.0.0,<2.0.0", ] [project.scripts] @@ -68,3 +68,6 @@ select = ["E", "F", "I", "N", "W", "UP"] [tool.mypy] python_version = "3.11" strict = true + +[tool.hatch.build.targets.wheel] +packages = ["engram"] diff --git a/scripts/check-ci.py b/scripts/check-ci.py new file mode 100644 index 0000000..439d356 --- /dev/null +++ b/scripts/check-ci.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +"""Run the full regression suite with an explicit SQLite or pgvector CI gate.""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import subprocess +import sys +import tempfile +import tomllib +import xml.etree.ElementTree as ET +from pathlib import Path + + +def inapplicable_cases(project: str) -> set[tuple[str, str, str]]: + """Exact exceptions for tests whose scenario cannot exist on that backend.""" + cases = { + ("test_atomic_mutations", + "test_independent_engines_preserve_distinct_receipt_increments[:memory:]", + "independent in-memory instances intentionally have separate databases"), + ("test_atomic_mutations", "test_two_processes_consume_one_receipt_once[:memory:]", + "process race requires persistent shared storage"), + } + for phase in ("begin", "ordinary_write", "savepoint"): + cases.add(("test_storage_transactions", + "test_sqlite_cancellation_during_queued_sql_does_not_leak_transaction" + f"[postgres-{phase}]", "SQLite worker queue regression")) + if project == "engram": + for renderer in ("claude", "gpt", "generic"): + cases.add(("test_retrieval_evaluation", + f"test_duplicate_content_keeps_selected_bullets_usage_stats[{renderer}]", + "Usage annotations are an OSS v0.5 extension")) + for kind in ("bullets", "legacy", "empty"): + for renderer in ("claude", "gpt", "generic"): + cases.add(("test_retrieval_evaluation", + "test_oss_core_memory_worked_examples_fit_tight_budget" + f"[{kind}-{renderer}]", + "OSS v0.5 extensions are intentionally edition-specific")) + return cases + + +async def wait_for_postgres(dsn: str) -> None: + import asyncpg + + connection = None + for attempt in range(45): + try: + connection = await asyncpg.connect(dsn, timeout=2) + break + except (TimeoutError, OSError, asyncpg.PostgresError): + if attempt == 44: + raise RuntimeError( + "Disposable PostgreSQL did not become ready; DSN redacted" + ) from None + await asyncio.sleep(1) + assert connection is not None + try: + available = await connection.fetchval( + "SELECT EXISTS (SELECT 1 FROM pg_available_extensions WHERE name = 'vector')" + ) + capable = await connection.fetchval( + "SELECT rolsuper OR (rolcreatedb AND rolcreaterole) " + "FROM pg_roles WHERE rolname = current_user" + ) + if not available: + raise RuntimeError("PostgreSQL service must include the pgvector extension") + if not capable: + raise RuntimeError("Disposable test DSN must permit CREATE DATABASE and CREATE ROLE") + finally: + await connection.close() + print("Disposable PostgreSQL is ready; pgvector and test privileges verified", flush=True) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--backend", choices=("sqlite", "postgres"), required=True) + parser.add_argument("--output", type=Path, required=True, help="JUnit XML output path") + args = parser.parse_args() + checkout = Path(__file__).resolve().parents[1] + report = args.output.resolve() + report.parent.mkdir(parents=True, exist_ok=True) + dsn = os.environ.get("ENGRAM_TEST_POSTGRES_DSN", "") + if args.backend == "postgres": + if not dsn: + parser.error("--backend postgres requires ENGRAM_TEST_POSTGRES_DSN; never skip CI") + asyncio.run(wait_for_postgres(dsn)) + # Tests run away from .env; discard application/provider configuration too. + keep = ("PATH", "HOME", "TMPDIR", "TEMP", "TMP", "SYSTEMROOT", "LANG", "LC_ALL") + env = {name: os.environ[name] for name in keep if name in os.environ} + env.update({ + "PYTHONPATH": str(checkout), "PYTHONNOUSERSITE": "1", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", "DO_NOT_TRACK": "1", + "OTEL_SDK_DISABLED": "true", + }) + if args.backend == "postgres": + env["ENGRAM_TEST_POSTGRES_DSN"] = dsn + with tempfile.TemporaryDirectory(prefix="engram-ci-tests-") as temporary: + work = Path(temporary).resolve() + if work.is_relative_to(checkout): + raise RuntimeError("Use a TMPDIR outside the source checkout") + result = subprocess.run([ + sys.executable, "-m", "pytest", "-q", "-c", str(checkout / "pyproject.toml"), + "-o", "cache_dir=" + str(work / "pytest-cache"), + "--basetemp", str(work / "pytest-temp"), "--junitxml", str(report), + str(checkout / "tests"), + ], cwd=work, env=env) + if result.returncode: + raise SystemExit(result.returncode) + cases = ET.parse(report).findall(".//testcase") + if not cases: + raise SystemExit("CI must execute tests") + project = tomllib.loads((checkout / "pyproject.toml").read_text())["project"]["name"] + allowed = inapplicable_cases(project) + unexpected_skips = [] + postgres_unconfigured = 0 + for case in cases: + skipped = case.find("skipped") + if skipped is None: + continue + reason = skipped.get("message", "") + name = case.get("name", "unknown") + key = (case.get("classname", "").rsplit(".", 1)[-1], name, reason) + if key in allowed: + print(f"Explicitly inapplicable: {name}: {reason}", flush=True) + elif (args.backend == "sqlite" and reason in { + "requires a disposable PostgreSQL admin DSN", + "requires disposable PostgreSQL admin DSN", + } and ("[postgres" in name or name == "test_fresh_postgres_vector_roundtrip_and_restart")): + postgres_unconfigured += 1 + else: + unexpected_skips.append(name) + if unexpected_skips: + raise SystemExit("Unexpected skipped tests: " + ", ".join(unexpected_skips)) + if args.backend == "postgres" and not any( + case.get("name") == "test_fresh_postgres_vector_roundtrip_and_restart" + and case.find("skipped") is None for case in cases + ): + raise SystemExit("Required PostgreSQL vector/restart integration test did not execute") + if postgres_unconfigured: + print(f"SQLite-only run: {postgres_unconfigured} PostgreSQL cases require the separate " + "PostgreSQL gate", flush=True) + print(f"{args.backend} gate passed: {len(cases)} cases; no unexpected skips", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/check-package.py b/scripts/check-package.py new file mode 100644 index 0000000..c2a8a86 --- /dev/null +++ b/scripts/check-package.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""Build from an sdist and smoke-test the installed wheel outside the checkout. + +Run with the Python interpreter being certified. Requires `build` in that +interpreter; the fresh smoke venv installs its own package and test dependencies. +No provider credentials are inherited and all HTTP traffic is loopback only. +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import socket +import subprocess +import sys +import tempfile +import time +import tomllib +import urllib.error +import urllib.request +import venv +import xml.etree.ElementTree as ET +from pathlib import Path + + +def clean_env() -> dict[str, str]: + # Preserve transport/trust configuration needed by pip, never API credentials. + names = ( + "PATH", "HOME", "TMPDIR", "TEMP", "TMP", "SYSTEMROOT", "LANG", "LC_ALL", + "SSL_CERT_FILE", "SSL_CERT_DIR", "REQUESTS_CA_BUNDLE", "PIP_INDEX_URL", + "PIP_EXTRA_INDEX_URL", "PIP_TRUSTED_HOST", "PIP_CERT", "PIP_CACHE_DIR", + ) + env = {name: os.environ[name] for name in names if name in os.environ} + env.update({ + "PYTHONNOUSERSITE": "1", "PYTHONUNBUFFERED": "1", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", "DO_NOT_TRACK": "1", + "OTEL_SDK_DISABLED": "true", + }) + return env + + +def run(command: list[str], *, cwd: Path, env: dict[str, str]) -> None: + print("+", " ".join(command), flush=True) + subprocess.run(command, cwd=cwd, env=env, check=True) + + +def check_cli(executable: Path, work: Path, env: dict[str, str]) -> None: + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + cli_env = { + **env, "ENGRAM_HOST": "127.0.0.1", "ENGRAM_PORT": str(port), + "ENGRAM_STORAGE_BACKEND": "sqlite", + "ENGRAM_SQLITE_PATH": str(work / "cli-smoke.sqlite3"), + "ENGRAM_AUTH_ENABLED": "false", + } + with (work / "cli.log").open("w+") as log: + process = subprocess.Popen([str(executable)], cwd=work, env=cli_env, + stdout=log, stderr=subprocess.STDOUT) + try: + deadline = time.monotonic() + 45 + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError("Installed engram CLI exited before serving /health") + try: + with urllib.request.urlopen( + f"http://127.0.0.1:{port}/health", timeout=1 + ) as response: + assert response.status == 200 + assert json.load(response)["status"] == "ok" + print("Installed engram CLI /health passed", flush=True) + break + except (OSError, urllib.error.URLError): + time.sleep(0.2) + else: + raise RuntimeError("Installed engram CLI did not become ready in 45 seconds") + except BaseException: + log.seek(0) + print(log.read(), file=sys.stderr) + raise + finally: + process.terminate() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, + help="Copy sdist, wheel, and smoke JUnit report here on success") + args = parser.parse_args() + checkout = Path(__file__).resolve().parents[1] + project = tomllib.loads((checkout / "pyproject.toml").read_text())["project"] + env = clean_env() + with tempfile.TemporaryDirectory(prefix="engram-wheel-check-") as temporary: + work = Path(temporary).resolve() + if work.is_relative_to(checkout): + raise RuntimeError("Use a TMPDIR outside the source checkout") + dist = work / "dist" + # Explicitly build a source distribution, then a wheel from that archive. + run([sys.executable, "-m", "build", "--sdist", "--outdir", str(dist), + str(checkout)], cwd=work, env=env) + sdists = list(dist.glob("*.tar.gz")) + assert len(sdists) == 1, sdists + run([sys.executable, "-m", "pip", "wheel", "--no-deps", "--wheel-dir", + str(dist), str(sdists[0])], cwd=work, env=env) + wheels = list(dist.glob("*.whl")) + assert len(wheels) == 1, wheels + smoke_env = work / "venv" + # Match `python -m venv` on POSIX. Managed macOS interpreters may depend + # on loader paths next to the original executable and cannot be copied. + venv.EnvBuilder(with_pip=True, symlinks=os.name != "nt").create(smoke_env) + bindir = smoke_env / ("Scripts" if os.name == "nt" else "bin") + python = bindir / ("python.exe" if os.name == "nt" else "python") + # OSS publishes MCP as an extra; explicitly certify that extra. + extra = "[mcp]" if "mcp" in project.get("optional-dependencies", {}) else "" + run([str(python), "-m", "pip", "install", str(wheels[0]) + extra, + "pytest>=8", "pytest-asyncio>=0.24"], cwd=work, env=env) + run([str(python), "-c", ( + "import pathlib,sys,engram,mcp; " + "from engram.server.app import create_app; " + "from engram.storage.sqlite import SQLiteBackend; " + "from engram.storage.postgres import PostgresBackend; " + "from engram.integrations.mcp_server import create_mcp_server; " + "assert pathlib.Path(engram.__file__).resolve().is_relative_to(" + "pathlib.Path(sys.prefix).resolve()), engram.__file__; " + "print('Installed package imports:', engram.__file__)" + )], cwd=work, env=env) + check_cli(bindir / ("engram.exe" if os.name == "nt" else "engram"), work, env) + # Copy only the test, so pytest cannot import the checkout's package. + shutil.copy2(checkout / "tests/test_mcp_stdio.py", work / "test_mcp_stdio.py") + report = work / "mcp-smoke.xml" + run([str(python), "-m", "pytest", "-q", "-o", "asyncio_mode=auto", + "--junitxml", str(report), "test_mcp_stdio.py"], cwd=work, env=env) + cases = ET.parse(report).findall(".//testcase") + if len(cases) != 6 or any(case.find("skipped") is not None for case in cases): + raise RuntimeError("Installed MCP smoke must execute all six cases with zero skips") + if args.output_dir: + output = args.output_dir.resolve() + output.mkdir(parents=True, exist_ok=True) + for artifact in [*sdists, *wheels, report]: + shutil.copy2(artifact, output / artifact.name) + print(f"Package check passed: {project['name']} {project['version']}, " + f"Python {sys.version.split()[0]}, CLI, imports, six MCP cases", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/check-shared-core.py b/scripts/check-shared-core.py new file mode 100644 index 0000000..573edec --- /dev/null +++ b/scripts/check-shared-core.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +"""Check selected shared contracts, not whole-core equality across editions.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +import shutil +import tempfile +import tomllib +from pathlib import Path + +MANIFEST = Path("shared-core-manifest.json") + + +def digest(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def inspect(checkout: Path, manifest: dict, *, check_hashes: bool = True) -> list[str]: + errors = [] + project = tomllib.loads((checkout / "pyproject.toml").read_text())["project"]["name"] + if project not in manifest["editions"]: + errors.append(f"{checkout}: unrecognized distribution {project!r}") + for entry in manifest["identical_files"]: + path = checkout / entry["path"] + if not path.is_file(): + errors.append(f"{checkout}: missing shared file {entry['path']}") + elif check_hashes and digest(path) != entry["sha256"]: + errors.append(f"{checkout}: unported/unrecorded drift in {entry['path']}") + for contract in manifest["contracts"]: + path = checkout / contract["path"] + if not path.is_file(): + errors.append(f"{checkout}: missing {contract['id']} at {contract['path']}") + continue + try: + tree = ast.parse(path.read_text()) + except SyntaxError: + errors.append(f"{checkout}: invalid Python in {contract['path']}") + continue + functions = {node.name for node in ast.walk(tree) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))} + for required in contract["tests"]: + if required not in functions: + errors.append(f"{checkout}: {contract['id']} missing test {required}") + return errors + + +def compare(oss: Path, paid: Path, manifest: dict, *, check_hashes: bool = True) -> list[str]: + errors = inspect(oss, manifest, check_hashes=check_hashes) + errors += inspect(paid, manifest, check_hashes=check_hashes) + for path, expected in [(oss, "engram-contextdb"), (paid, "engram")]: + name = tomllib.loads((path / "pyproject.toml").read_text())["project"]["name"] + if name != expected: + errors.append(f"{path}: expected {expected}, got {name}; arguments are OSS then paid") + peer_manifest = json.loads((paid / MANIFEST).read_text()) + if peer_manifest != manifest: + errors.append("Edition manifests differ; port and review the manifest in both checkouts") + for entry in manifest["identical_files"]: + left, right = oss / entry["path"], paid / entry["path"] + if left.is_file() and right.is_file() and left.read_bytes() != right.read_bytes(): + errors.append(f"Cross-edition mismatch: {entry['path']} ({entry['purpose']})") + return errors + + +def self_test(oss: Path, paid: Path, manifest: dict) -> None: + """Prove that a real file drift and a missing contract both fail, in copies.""" + with tempfile.TemporaryDirectory(prefix="engram-port-check-") as temporary: + paths = {"pyproject.toml", str(MANIFEST)} + paths.update(item["path"] for item in manifest["identical_files"]) + paths.update(item["path"] for item in manifest["contracts"]) + copies = [] + for source, label in [(oss, "oss"), (paid, "paid")]: + target = Path(temporary) / label + for relative in paths: + destination = target / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source / relative, destination) + copies.append(target) + changed = copies[1] / manifest["identical_files"][0]["path"] + changed.write_bytes(changed.read_bytes() + b"\n# deliberate drift\n") + assert any("mismatch" in error for error in compare(*copies, manifest)) + contract = manifest["contracts"][0] + contract_path = copies[0] / contract["path"] + contract_path.write_text(contract_path.read_text().replace( + "def " + contract["tests"][0] + "(", "def removed_contract(", 1 + )) + assert any("missing test" in error for error in compare(*copies, manifest)) + print("Self-test passed: file drift and missing contract rejected in disposable copies") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("oss", type=Path, nargs="?", help="OSS checkout (first of two paths)") + parser.add_argument("paid", type=Path, nargs="?", help="Paid checkout (second of two paths)") + parser.add_argument("--checkout", type=Path, help="Check one checkout against pinned manifest") + parser.add_argument("--refresh", action="store_true", + help="After review, refresh hashes in BOTH checkouts; requires equal files") + parser.add_argument("--self-test", action="store_true", + help="Also prove deliberate drift fails using disposable copies") + args = parser.parse_args() + pair = args.oss is not None and args.paid is not None + if bool(args.checkout) == pair or ((args.oss is None) != (args.paid is None)): + parser.error("provide either --checkout PATH or OSS_PATH PAID_PATH") + if (args.refresh or args.self_test) and not pair: + parser.error("--refresh and --self-test require both checkout paths") + root = (args.checkout or args.oss).resolve() + manifest = json.loads((root / MANIFEST).read_text()) + if manifest.get("version") != 1: + parser.error("unsupported shared-core manifest version") + if pair: + oss, paid = args.oss.resolve(), args.paid.resolve() + if oss == paid: + parser.error("OSS and paid checkout must be different paths") + errors = compare(oss, paid, manifest, check_hashes=not args.refresh) + else: + errors = inspect(root, manifest) + if errors: + raise SystemExit("\n".join(errors) + "\nPort the change, run both suites, and record " + "provenance before refreshing hashes; see docs/shared-core-release.md.") + if args.refresh: + for entry in manifest["identical_files"]: + entry["sha256"] = digest(oss / entry["path"]) + payload = json.dumps(manifest, indent=2) + "\n" + for checkout in [oss, paid]: + (checkout / MANIFEST).write_text(payload) + print("Refreshed shared-file hashes in both manifests") + if args.self_test: + self_test(oss, paid, manifest) + print(f"Shared-core check passed: {len(manifest['identical_files'])} pinned files, " + f"{len(manifest['contracts'])} regression contracts") + + +if __name__ == "__main__": + main() diff --git a/shared-core-manifest.json b/shared-core-manifest.json new file mode 100644 index 0000000..c2ec89a --- /dev/null +++ b/shared-core-manifest.json @@ -0,0 +1,342 @@ +{ + "version": 1, + "editions": { + "engram-contextdb": { + "host": "GitHub", + "line": "0.5", + "role": "OSS" + }, + "engram": { + "host": "GitLab", + "line": "0.4", + "role": "paid" + } + }, + "identical_files": [ + { + "path": "engram/core/concurrency.py", + "sha256": "f8492fb72d49e22498200bd7f78a8247c48c5658c09d11222e6728ee3a8df93d", + "purpose": "Task-local context lock helper" + }, + { + "path": "engram/core/events.py", + "sha256": "3cdcc74bb9645f9870135e4f7749171ab7e4d592b2742eba6a01d70262cb0d1e", + "purpose": "Shared event subscription and dispatch behavior" + }, + { + "path": "tests/test_mcp_stdio.py", + "sha256": "11a689881611d97b82ee60e10df5670fcc669bf2cef2e6cf7796fbb6f6409ea3", + "purpose": "Real stdio protocol, client ownership, decision ID contract" + }, + { + "path": "tests/test_context_isolation.py", + "sha256": "3380585d73dbf5ffbd3affb99d4aa6d1468d7ac1b5295101185f6a1fc9f28dcd", + "purpose": "Cross-context mutation and recall isolation" + }, + { + "path": "tests/test_core_regressions.py", + "sha256": "e02310b56c8f52083cf88c4ad6689f28160fbab950877f752bdb7cbe42ed2386", + "purpose": "Embedding consistency, rollback, and searchable writes" + }, + { + "path": "tests/test_postgres_startup.py", + "sha256": "4261b6f5d7f8bf35c32364c70eee20f0723f56a54c547191a4079359c35fc4af", + "purpose": "pgvector startup, pool retry, roundtrip and restart" + }, + { + "path": "scripts/check-ci.py", + "sha256": "75872c3628c0c080b5736b48a103a9bd9774947f1c9c3a0924236763407d4fe0", + "purpose": "Database readiness and non-skipping full-suite gate" + }, + { + "path": "scripts/check-package.py", + "sha256": "3fda9e99d01a1430457c23fb74afddeb7a6417064c9f40cca337d9e8c80d3029", + "purpose": "sdist-to-wheel, fresh installed CLI/MCP certification" + }, + { + "path": "scripts/check-shared-core.py", + "sha256": "32bf08671006a83553c04cecde4bff0398941c0c3bfd6eb228c6f09de51d9bea", + "purpose": "Shared contract and port drift enforcement" + }, + { + "path": "engram/storage/transactions.py", + "sha256": "a4628aabd329e10bf72fd75675187276c0cc25e118bf1f408661ea37272bd097", + "purpose": "Task-owned scoped handles and cancellation cleanup" + }, + { + "path": "engram/llm/explicit.py", + "sha256": "d9174c1cd014dc695b2dbd96ed4527e2efe51de2d46311659236d4653068dc49", + "purpose": "Explicit standalone provider configuration and structured output" + }, + { + "path": "engram/maintenance/__init__.py", + "sha256": "bbc60a576d962bd576b91b1e0ef4813b38a3382822f57b8de42b8299524c3799", + "purpose": "Maintenance command package" + }, + { + "path": "engram/maintenance/backfill_embeddings.py", + "sha256": "e8c8d73a6919bb5c03bc4f4fbc1b1f0b4e972f64f21d9ded2a80567cdc0eaf1e", + "purpose": "Bounded conditional missing-vector repair" + }, + { + "path": "evals/__init__.py", + "sha256": "be8459d5d5ff8d649403318cc48a34489f4ef441eec47462a426ff926fc3d1d8", + "purpose": "Controlled regression benchmark package" + }, + { + "path": "evals/retrieval.py", + "sha256": "df41070de957291537153aa6fae443fac7ec23886c6e6c5d60869d56a381d4e8", + "purpose": "Retrieval measurements and nonvacuous pass/fail gate" + }, + { + "path": "evals/corpus.json", + "sha256": "e4d1c994b2d4f51d970e3eda13649043e7c9a485f13988eb785ac78f47fddca9", + "purpose": "Curated retrieval relevance judgments and thresholds" + }, + { + "path": "tests/test_embedding_backfill.py", + "sha256": "d126cf02bc4b788835828028515c61209469b6ed298ef32e0e68081b9f140ada", + "purpose": "Read-only preview, repair races, numeric validation and real adapter safety" + }, + { + "path": "tests/test_retrieval_evaluation.py", + "sha256": "831a86d329adc79894343605b02d20ac0a1903170cb2b38e0e2d1901a4a04c25", + "purpose": "Packing fidelity, metric math and explicit-provider output" + }, + { + "path": "tests/test_storage_transactions.py", + "sha256": "18ec729edba17b92b7316faad6597d246fbd81da97e175acabc89d27d44f4843", + "purpose": "Scoped ownership, savepoints and cancellation contracts" + }, + { + "path": "tests/test_atomic_mutations.py", + "sha256": "062c4ac8580db03cbe792a9b961ae8692c10cd9600200b03eb23947019933ab7", + "purpose": "Batch/feedback atomicity and independent worker behavior" + }, + { + "path": "tests/test_pgvector_compatibility.py", + "sha256": "90e57a0fb01e31aa324cf638c2fd41eb9defc95f7f052091d61c5407c629fa16", + "purpose": "Current and older pgvector result conversion" + }, + { + "path": "tests/test_postgres_pool_lifecycle.py", + "sha256": "2f29cd7b0d105c81d6c79b8a123bc013959687f85d763507d73db5fafbd1aeb1", + "purpose": "Actual partial-pool failure/cancellation cleanup and retry" + }, + { + "path": "tests/test_existing_core_correctness.py", + "sha256": "f834019bf41ed4a1e60d3c6a76027e812ed8089d35213f3e95600216a1a90e14", + "purpose": "Context boundaries, reversible merges, populated upgrades and atomic input history" + }, + { + "path": "tests/test_ingestion_config_compatibility.py", + "sha256": "5ea54cec62b24755b9a248b82068b53c491f8954dae57b7a331515fefab3bc60", + "purpose": "Configured model and dedup routing with legacy adapter compatibility" + }, + { + "path": "tests/test_schema_rollback.py", + "sha256": "7ad18dddc365738ee099e210abfe7b02b9c15e75d990479c0d1f70b45b493332", + "purpose": "Reversible schema mutations and honest historical rollback" + }, + { + "path": "tests/test_server_request_regressions.py", + "sha256": "f97b004e33a19077f3b9d0f72b69f0f74fba2f878fb02d63bfd20a4e4d75a59f", + "purpose": "Legacy timestamp and config error contracts" + }, + { + "path": "tests/test_sdk_compatibility.py", + "sha256": "e97666b5e3f079b0b2eef38838beb2eaaf1cf7eb98f365eefb10b33e8252badc", + "purpose": "Existing Python SDK calls against real REST handlers" + }, + { + "path": "tests/test_public_mutation_atomicity.py", + "sha256": "c92e81b0d4b18bc2a74a7afb044fc3293fc5d61c871a8489c4bfcb7049659952", + "purpose": "Public mutation ledger/capacity atomicity, invalid vector ranking and historical rollback" + }, + { + "path": "engram/core/similarity.py", + "sha256": "7bd4c83a27c81508ce351a9fdf9aa0aa07568f1c792cdd2b08b54c4ec14903ec", + "purpose": "Legacy cosine helper plus validated comparison for destructive decisions" + }, + { + "path": "tests/test_consolidation_consistency.py", + "sha256": "6fbdd08b3d7681a2d84ef8df2f12f31e2fb1f573eae60eca63112a31a675d12f", + "purpose": "Consolidation concurrency, uncertain comparisons, historical inverses and migration failure recovery" + }, + { + "path": "tests/test_archive_lifecycle_atomicity.py", + "sha256": "6859d57525eb1c4f6d5c057c05d4705d28b2865b3fa1fd411040453617eb74f2", + "purpose": "Purge rollback, restore-versus-expiry serialization and REST capacity races" + } + ], + "contracts": [ + { + "id": "context-isolation", + "path": "tests/test_context_isolation.py", + "tests": [ + "test_child_get_requires_matching_context", + "test_foreign_rollback_cannot_mutate", + "test_feedback_ignores_foreign_materialization", + "test_feedback_checks_each_bullet_context" + ] + }, + { + "id": "embedding-and-rollback", + "path": "tests/test_core_regressions.py", + "tests": [ + "test_commit_semantic_search_and_paraphrase_dedup", + "test_auxiliary_memories_are_searchable", + "test_update_vector_and_rollback_match_content", + "test_legacy_rollback_does_not_keep_new_content_vector" + ] + }, + { + "id": "mcp-installed-protocol", + "path": "tests/test_mcp_stdio.py", + "tests": [ + "test_stdio_subprocess_initializes_and_forwards_configured_url_and_key", + "test_stdio_closes_owned_sdk_client_on_shutdown", + "test_decision_tools_display_id_returned_by_api" + ] + }, + { + "id": "postgres-startup", + "path": "tests/test_postgres_startup.py", + "tests": [ + "test_extension_exists_before_pool_registers_codecs", + "test_failed_pool_creation_is_not_cached", + "test_fresh_postgres_vector_roundtrip_and_restart" + ] + }, + { + "id": "atomic-mutations", + "path": "tests/test_atomic_mutations.py", + "tests": [ + "test_failed_batch_leaves_no_mutations_or_audit", + "test_failed_retry_uses_immediate_pre_retry_snapshot", + "test_feedback_failure_rolls_back_stats_and_audit", + "test_failed_rollback_is_atomic", + "test_two_processes_consume_one_receipt_once", + "test_public_commit_retry_finishes_feedback_without_reextracting" + ] + }, + { + "id": "scoped-transactions", + "path": "tests/test_storage_transactions.py", + "tests": [ + "test_commit_and_rollback", + "test_scoped_handles_reject_parent_wrong_task_and_expired_use", + "test_nested_savepoint_rollback", + "test_cancellation_rolls_back_and_connection_is_reusable" + ] + }, + { + "id": "embedding-repair", + "path": "tests/test_embedding_backfill.py", + "tests": [ + "test_preview_never_calls_provider_or_transaction", + "test_repairs_only_embedding_and_rerun_is_idempotent", + "test_independent_repair_workers_save_only_one_vector", + "test_real_adapter_invalid_model_does_not_load_dotenv_or_corrupt_json", + "test_every_saved_repair_is_searchable_by_its_own_vector" + ] + }, + { + "id": "retrieval-evaluation", + "path": "tests/test_retrieval_evaluation.py", + "tests": [ + "test_metric_math_and_duplicate_credit", + "test_controlled_corpus_gate", + "test_cli_failure_gate", + "test_receipts_match_complete_rendered_content", + "test_duplicate_content_keeps_selected_bullets_usage_stats" + ] + }, + { + "id": "postgres-pool-lifecycle", + "path": "tests/test_postgres_pool_lifecycle.py", + "tests": [ + "test_partial_pool_initialization_releases_connections_and_retries" + ] + }, + { + "id": "public-mutation-contracts", + "path": "tests/test_public_mutation_atomicity.py", + "tests": [ + "test_activity_failure_rolls_back_every_public_entrypoint", + "test_direct_add_salience_failure_rolls_back_bullet_and_audit", + "test_concurrent_direct_add_checks_capacity_after_embedding", + "test_other_public_entrypoints_reject_capacity_overflow", + "test_snapshotless_historical_merge_refuses_entire_rollback", + "test_zero_vector_does_not_hide_valid_similarity_at_limit_one", + "test_curator_rejects_nonfinite_custom_storage_scores" + ] + }, + { + "id": "consolidation-consistency", + "path": "tests/test_consolidation_consistency.py", + "tests": [ + "test_reextract_preserves_identity_or_retains_uncertain_comparison", + "test_reextract_failed_provider_cannot_authorize_removal", + "test_legacy_missing_identity_or_inverse_declines_whole_rollback", + "test_schema_induction_preserves_changes_during_provider_call", + "test_schema_and_links_rollback_together", + "test_concurrent_schema_induction_creates_one_schema", + "test_consolidation_rechecks_locked_current_rows", + "test_promotion_rechecks_current_qualification_duplicates_and_capacity", + "test_lifecycle_migration_backfill_failure_is_retry_safe" + ] + }, + { + "id": "archive-lifecycle", + "path": "tests/test_archive_lifecycle_atomicity.py", + "tests": [ + "test_purge_failure_preserves_bullet_and_edges", + "test_successful_concurrent_restore_is_not_purged", + "test_concurrent_restores_respect_context_capacity" + ] + } + ], + "intentionally_distinct": [ + { + "paths": [ + "engram/core/config.py", + "engram/core/models.py", + "engram/core/delta.py", + "engram/core/ingestion.py" + ], + "reason": "OSS 0.5 reconsolidation/rollback metadata and validity gate; paid 0.4 feedback behavior remains supported" + }, + { + "paths": [ + "engram/auth/", + "engram/server/", + "engram/integrations/mcp_server.py" + ], + "reason": "Paid authentication, API keys, tenant, and hosted MCP extensions; shared stdio protocol is tested" + }, + { + "paths": [ + "engram/storage/base.py", + "engram/storage/sqlite.py", + "engram/storage/postgres.py" + ], + "reason": "Shared transactional behavior with edition-specific persisted model/schema fields; enforce behavior through both suites" + }, + { + "paths": [ + "pyproject.toml", + ".github/workflows/ci.yml", + ".gitlab-ci.yml" + ], + "reason": "Distinct distributions/dependencies and hosting systems with equivalent verification commands" + }, + { + "paths": [ + "engram/core/materialization.py", + "engram/renderers/" + ], + "reason": "OSS core memory, worked examples, MMR and usage statistics remain supported; shared final-budget and receipt contracts are tested in both editions" + } + ] +} diff --git a/tests/test_api.py b/tests/test_api.py index 0745add..a8b0714 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -5,12 +5,35 @@ import pytest from fastapi.testclient import TestClient -from engram.server.app import create_app +from engram.core.config import Settings +from engram.server import app as app_module +from engram.storage.sqlite import SQLiteBackend + + +class OfflineLLM: + """API tests must never call a provider or use developer credentials.""" + + async def embed(self, text): + import hashlib + + digest = hashlib.sha256(text.encode()).digest() + return [value / 255.0 for value in digest] + + async def complete(self, *args, **kwargs): + raise AssertionError("Unexpected LLM completion in API test") @pytest.fixture -def app(): - return create_app() +def app(tmp_path, monkeypatch): + monkeypatch.setattr( + app_module, "_create_storage", + lambda: SQLiteBackend(str(tmp_path / "api-test.db")), + ) + monkeypatch.setattr(app_module, "_create_llm", OfflineLLM) + settings = Settings(_env_file=None) + monkeypatch.setattr(app_module, "get_settings", lambda: settings) + monkeypatch.setattr("engram.core.config.get_settings", lambda: settings) + return app_module.create_app() @pytest.fixture diff --git a/tests/test_archive_lifecycle_atomicity.py b/tests/test_archive_lifecycle_atomicity.py new file mode 100644 index 0000000..6135c05 --- /dev/null +++ b/tests/test_archive_lifecycle_atomicity.py @@ -0,0 +1,155 @@ +"""Archive expiry, explicit purge and REST restores preserve concurrent state.""" + +import asyncio +import uuid +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace + +import aiosqlite +import asyncpg +import pytest +from fastapi import HTTPException +from starlette.requests import Request +from test_storage_transactions import storage as transaction_storage + +from engram.core.models import Bullet, ConceptEdge, EdgeType +from engram.server.routes.lifecycle import restore_bullet +from engram.storage.postgres import PostgresBackend +from engram.storage.sqlite import SQLiteBackend + +storage = transaction_storage + + +async def archived_with_edge(backend, context_id): + first, second = uuid.uuid4(), uuid.uuid4() + await backend.add_bullet(context_id, Bullet(id=str(first), content="archived")) + await backend.add_bullet(context_id, Bullet(id=str(second), content="other")) + await backend.archive_bullet(context_id, str(first)) + bullet = await backend.get_bullet(str(first)) + bullet.archived_at = datetime.now(UTC) - timedelta(days=200) + await backend.update_bullet(bullet) + edge = await backend.add_edge( + uuid.UUID(context_id), + ConceptEdge(from_node=first, to_node=second, type=EdgeType.RELATED_TO), + ) + return str(first), edge.id + + +def sql_connection_type(backend): + return aiosqlite.Connection if isinstance(backend, SQLiteBackend) else asyncpg.Connection + + +@pytest.mark.parametrize("mode", ["expired", "direct"]) +async def test_purge_failure_preserves_bullet_and_edges(storage, monkeypatch, mode): + backend, context_id = storage + bullet_id, edge_id = await archived_with_edge(backend, context_id) + cls = sql_connection_type(backend) + execute = cls.execute + + async def fail(self, query, *args, **kwargs): + if query.startswith("DELETE FROM bullets"): + raise RuntimeError("injected bullet deletion failure") + return await execute(self, query, *args, **kwargs) + + with monkeypatch.context() as patch: + patch.setattr(cls, "execute", fail) + with pytest.raises(RuntimeError, match="injected bullet deletion"): + if mode == "expired": + await backend.purge_expired_archives(context_id, 180) + else: + await backend.purge_bullet(context_id, bullet_id) + assert await backend.get_bullet(bullet_id) is not None + assert [edge.id for edge in await backend.get_edges(uuid.UUID(context_id))] == [edge_id] + assert await backend.purge_bullet(context_id, bullet_id) + assert await backend.get_bullet(bullet_id) is None + assert await backend.get_edges(uuid.UUID(context_id)) == [] + + +@pytest.mark.parametrize("storage", ["file", "postgres"], indirect=True) +async def test_successful_concurrent_restore_is_not_purged(storage, monkeypatch): + backend, context_id = storage + bullet_id, edge_id = await archived_with_edge(backend, context_id) + peer = ( + SQLiteBackend(backend.db_path, existing_only=True) + if isinstance(backend, SQLiteBackend) + else PostgresBackend(backend.dsn, existing_only=True) + ) + at_delete, release = asyncio.Event(), asyncio.Event() + cls = sql_connection_type(backend) + execute = cls.execute + purge_task = None + + async def pause(self, query, *args, **kwargs): + if asyncio.current_task() is purge_task and query.startswith("DELETE FROM edges"): + at_delete.set() + await release.wait() + return await execute(self, query, *args, **kwargs) + + monkeypatch.setattr(cls, "execute", pause) + purge_task = asyncio.create_task(backend.purge_expired_archives(context_id, 180)) + restore_task = None + try: + await asyncio.wait_for(at_delete.wait(), timeout=5) + restore_task = asyncio.create_task(peer.restore_bullet(context_id, bullet_id)) + await asyncio.wait({restore_task}, timeout=0.25) + release.set() + purged, restored = await asyncio.wait_for( + asyncio.gather(purge_task, restore_task), timeout=10 + ) + if restored is not None: + assert purged == 0 + current = await backend.get_bullet(bullet_id) + assert current is not None and not current.is_archived + assert [e.id for e in await backend.get_edges(uuid.UUID(context_id))] == [edge_id] + else: + assert purged == 1 + assert await backend.get_bullet(bullet_id) is None + finally: + release.set() + await asyncio.gather( + *[t for t in (purge_task, restore_task) if t is not None], return_exceptions=True + ) + await peer.close() + + +@pytest.mark.parametrize("storage", ["file", "postgres"], indirect=True) +async def test_concurrent_restores_respect_context_capacity(storage, monkeypatch): + backend, context_id = storage + context = await backend.get_context(uuid.UUID(context_id)) + context.lifecycle_config.max_active_bullets = 1 + await backend.update_context(context) + for bullet_id in ("one", "two"): + await backend.add_bullet(context_id, Bullet(id=bullet_id, content=bullet_id)) + await backend.archive_bullet(context_id, bullet_id) + peer = ( + SQLiteBackend(backend.db_path, existing_only=True) + if isinstance(backend, SQLiteBackend) + else PostgresBackend(backend.dsn, existing_only=True) + ) + capacity = type(backend).get_capacity_status + both_read = asyncio.Barrier(2) + + async def synchronized_read(self, *args, **kwargs): + value = await capacity(self, *args, **kwargs) + if self._scope is None: + await both_read.wait() + return value + + monkeypatch.setattr(type(backend), "get_capacity_status", synchronized_read) + + async def restore(store, bullet_id): + request = Request({"type": "http", "app": SimpleNamespace( + state=SimpleNamespace(storage=store)), "state": {"user_id": context.owner}}) + try: + return await restore_bullet(context_id, bullet_id, request) + except HTTPException as exc: + return exc.status_code + + try: + results = await asyncio.wait_for(asyncio.gather( + restore(backend, "one"), restore(peer, "two")), timeout=10) + assert sum(not isinstance(result, int) for result in results) == 1 + assert 409 in results + assert len(await backend.list_bullets(context_id)) == 1 + finally: + await peer.close() diff --git a/tests/test_atomic_mutations.py b/tests/test_atomic_mutations.py new file mode 100644 index 0000000..912830a --- /dev/null +++ b/tests/test_atomic_mutations.py @@ -0,0 +1,429 @@ +"""Fault injection checks persisted atomicity, not just operation counters.""" + +import asyncio +import os +import sys +from pathlib import Path + +import pytest +from test_storage_transactions import storage as transaction_storage + +from engram.core.delta import DeltaEngine +from engram.core.ingestion import IngestionEngine +from engram.core.models import ( + Bullet, + DeltaBatch, + DeltaOperation, + DeltaOpType, + ExecutionFeedback, + FeedbackOutcome, + MaterializationRecord, +) +from engram.storage.postgres import PostgresBackend +from engram.storage.sqlite import SQLiteBackend + +storage = transaction_storage + + +@pytest.mark.parametrize("failure", ["second_write", "audit"]) +async def test_failed_batch_leaves_no_mutations_or_audit(storage, monkeypatch, failure): + backend, context_id = storage + batch = DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id="a", content="first"), + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id="b", content="second"), + ], + ) + name = "add_bullet" if failure == "second_write" else "save_delta_batch" + original = getattr(type(backend), name) + + async def fail(self, *args, **kwargs): + if failure == "audit" or args[1].id == "b": + raise RuntimeError("injected failure") + return await original(self, *args, **kwargs) + + monkeypatch.setattr(type(backend), name, fail) + with pytest.raises(RuntimeError, match="injected failure"): + await DeltaEngine(backend).apply_batch(batch) + assert await backend.list_bullets(context_id) == [] + assert await backend.get_delta_batch(batch.id) is None + assert batch.bullets_added == 0 + monkeypatch.setattr(type(backend), name, original) + result = await DeltaEngine(backend).apply_batch(batch) + assert result.bullets_added == 2 + + +async def test_failed_retry_uses_immediate_pre_retry_snapshot(storage, monkeypatch): + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="target", content="initial")) + batch = DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation( + op_type=DeltaOpType.UPDATE_BULLET, target_id="target", content="applied" + ), + ], + ) + original = type(backend).save_delta_batch + + async def fail(self, batch): + raise RuntimeError("audit unavailable") + + monkeypatch.setattr(type(backend), "save_delta_batch", fail) + with pytest.raises(RuntimeError): + await DeltaEngine(backend).apply_batch(batch) + assert batch.operations[0].previous_state is None + if hasattr(batch.operations[0], "rollback_state"): + assert batch.operations[0].rollback_state is None + monkeypatch.setattr(type(backend), "save_delta_batch", original) + independent = await backend.get_bullet("target") + independent.content = "independent update" + await backend.update_bullet(independent) + await DeltaEngine(backend).apply_batch(batch) + await DeltaEngine(backend).rollback_batch(batch.id) + assert (await backend.get_bullet("target")).content == "independent update" + + +@pytest.mark.parametrize("failure", ["marker", "second_bullet"]) +async def test_feedback_failure_rolls_back_stats_and_audit(storage, monkeypatch, failure): + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="a", content="first")) + await backend.add_bullet(context_id, Bullet(id="b", content="second")) + record = MaterializationRecord(context_id=context_id, bullets_included=["a", "b"]) + await backend.save_materialization(record) + name = "mark_materialization_reconsolidated" if failure == "marker" else "update_bullet" + original = getattr(type(backend), name) + + async def fail(self, *args): + if failure == "marker" or args[0].id == "b": + raise RuntimeError("receipt unavailable") + return await original(self, *args) + + monkeypatch.setattr(type(backend), name, fail) + engine = IngestionEngine(backend, None) + feedback = ExecutionFeedback(outcome=FeedbackOutcome.SUCCESS) + with pytest.raises(RuntimeError, match="receipt unavailable"): + await engine._reconsolidate(record.id, feedback, expected_context_id=context_id) + assert [(b.hit_count, b.recall_count) for b in await backend.list_bullets(context_id)] == [ + (0, 0) + ] * 2 + assert (await backend.get_materialization(record.id)).reconsolidated_at is None + assert await backend.list_delta_batches(context_id) == [] + monkeypatch.setattr(type(backend), name, original) + await engine._reconsolidate(record.id, feedback, expected_context_id=context_id) + await engine._reconsolidate(record.id, feedback, expected_context_id=context_id) + assert [(b.hit_count, b.recall_count) for b in await backend.list_bullets(context_id)] == [ + (1, 1) + ] * 2 + + +async def test_failed_rollback_is_atomic(storage, monkeypatch): + backend, context_id = storage + for bid in ("a", "b"): + await backend.add_bullet(context_id, Bullet(id=bid, content="before")) + batch = DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.UPDATE_BULLET, target_id=bid, content="after") + for bid in ("a", "b") + ], + ) + await DeltaEngine(backend).apply_batch(batch) + original = type(backend).update_bullet + + async def fail(self, bullet): + if bullet.id == "a": + raise RuntimeError("rollback interrupted") + return await original(self, bullet) + + monkeypatch.setattr(type(backend), "update_bullet", fail) + with pytest.raises(RuntimeError): + await DeltaEngine(backend).rollback_batch(batch.id) + assert {b.content for b in await backend.list_bullets(context_id)} == {"after"} + + +async def test_independent_engines_preserve_distinct_receipt_increments(storage): + backend, context_id = storage + if isinstance(backend, SQLiteBackend) and backend.db_path == ":memory:": + pytest.skip("independent in-memory instances intentionally have separate databases") + second = ( + SQLiteBackend(backend.db_path, existing_only=True) + if isinstance(backend, SQLiteBackend) + else PostgresBackend(backend.dsn, existing_only=True) + ) + try: + await backend.add_bullet(context_id, Bullet(id="shared", content="fact")) + receipts = [ + MaterializationRecord(context_id=context_id, bullets_included=["shared"]) + for _ in range(2) + ] + for record in receipts: + await backend.save_materialization(record) + feedback = ExecutionFeedback(outcome=FeedbackOutcome.SUCCESS) + await asyncio.gather( + IngestionEngine(backend, None)._reconsolidate( + receipts[0].id, feedback, expected_context_id=context_id + ), + IngestionEngine(second, None)._reconsolidate( + receipts[1].id, feedback, expected_context_id=context_id + ), + ) + assert (await backend.get_bullet("shared")).hit_count == 2 + finally: + await second.close() + + +WORKER = r""" +import asyncio, sys +from engram.core.config import Settings +import engram.core.config as config +from engram.core.ingestion import IngestionEngine +from engram.core.models import ExecutionFeedback, FeedbackOutcome +from engram.storage.sqlite import SQLiteBackend +from engram.storage.postgres import PostgresBackend +config.get_settings = lambda: Settings(_env_file=None, auth_enabled=False) +async def main(): + backend = (SQLiteBackend(sys.argv[2], existing_only=True) if sys.argv[1] == "sqlite" + else PostgresBackend(sys.argv[2], existing_only=True)) + try: + print("ready", flush=True) + await asyncio.to_thread(sys.stdin.readline) + await IngestionEngine(backend, None)._reconsolidate( + sys.argv[3], ExecutionFeedback(outcome=FeedbackOutcome.SUCCESS), + expected_context_id=sys.argv[4]) + finally: + await backend.close() +asyncio.run(main()) +""" + + +async def test_two_processes_consume_one_receipt_once(storage): + backend, context_id = storage + if isinstance(backend, SQLiteBackend) and backend.db_path == ":memory:": + pytest.skip("process race requires persistent shared storage") + await backend.add_bullet(context_id, Bullet(id="shared", content="fact", salience=0.5)) + record = MaterializationRecord(context_id=context_id, bullets_included=["shared"]) + await backend.save_materialization(record) + kind = "sqlite" if isinstance(backend, SQLiteBackend) else "postgres" + target = backend.db_path if kind == "sqlite" else backend.dsn + environment = dict(os.environ, PYTHONPATH=str(Path(__file__).resolve().parents[1])) + children = [] + try: + for _ in range(2): + children.append( + await asyncio.create_subprocess_exec( + sys.executable, + "-c", + WORKER, + kind, + target, + record.id, + context_id, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=environment, + ) + ) + for child in children: + assert await asyncio.wait_for(child.stdout.readline(), 10) == b"ready\n" + for child in children: + child.stdin.write(b"go\n") + await child.stdin.drain() + for child in children: + _, stderr = await asyncio.wait_for(child.communicate(), 30) + assert child.returncode == 0, stderr.decode() + finally: + for child in children: + if child.returncode is None: + child.kill() + await child.wait() + await backend.close() + await backend.initialize() + bullet = await backend.get_bullet("shared") + assert (bullet.recall_count, bullet.hit_count, bullet.miss_count) == (1, 1, 0) + assert bullet.salience == pytest.approx(0.525) + assert (await backend.get_materialization(record.id)).reconsolidated_at is not None + + +async def test_schema_and_bullet_rollback_together_when_audit_fails(storage, monkeypatch): + backend, context_id = storage + batch = DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.ADD_SCHEMA, target_id="schema", content="group"), + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id="fact", content="fact"), + ], + ) + + async def fail(self, batch): + raise RuntimeError("audit unavailable") + + monkeypatch.setattr(type(backend), "save_delta_batch", fail) + with pytest.raises(RuntimeError, match="audit unavailable"): + await DeltaEngine(backend).apply_batch(batch) + assert await backend.list_schemas(context_id) == [] + assert await backend.list_bullets(context_id) == [] + assert await backend.get_delta_batch(batch.id) is None + + +async def test_public_commit_retry_finishes_feedback_without_reextracting(storage, monkeypatch): + import uuid + + from engram.core.config import IngestionConfig + from engram.core.models import Context, IntentAnchor, Reflection, ReflectionInsight + + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="recalled", content="Previously recalled fact")) + receipt = MaterializationRecord(context_id=context_id, bullets_included=["recalled"]) + await backend.save_materialization(receipt) + + class CountingLLM: + completions = 0 + + async def complete(self, *args, **kwargs): + self.completions += 1 + return Reflection( + new_insights=[ReflectionInsight(content="Newly extracted fact")], confidence=0.99 + ).model_dump_json() + + async def embed(self, text): + return [1.0] + [0.0] * 1535 + + llm = CountingLLM() + engine = IngestionEngine( + backend, llm, ingestion_config=IngestionConfig(max_reflection_rounds=1) + ) + original = type(backend).mark_materialization_reconsolidated + + async def fail(self, *args): + raise RuntimeError("receipt unavailable") + + monkeypatch.setattr(type(backend), "mark_materialization_reconsolidated", fail) + arguments = dict( + context_id=uuid.UUID(context_id), + agent_id="retry-agent", + session_id=uuid.uuid4(), + content="Identical request with feedback", + materialization_id=receipt.id, + feedback=ExecutionFeedback(outcome=FeedbackOutcome.SUCCESS), + ) + with pytest.raises(RuntimeError, match="receipt unavailable"): + await engine.commit(**arguments) + ingestion_batches = await backend.list_delta_batches(context_id) + assert len(ingestion_batches) == 1 + assert (await backend.get_bullet("recalled")).hit_count == 0 + assert (await backend.get_materialization(receipt.id)).reconsolidated_at is None + assert {bullet.content for bullet in await backend.list_bullets(context_id)} == { + "Previously recalled fact", + "Newly extracted fact", + } + + monkeypatch.setattr(type(backend), "mark_materialization_reconsolidated", original) + for _ in range(3): + retried = await engine.commit(**arguments) + assert retried.id == ingestion_batches[0].id + assert llm.completions == 1 + recalled = await backend.get_bullet("recalled") + assert (recalled.hit_count, recalled.recall_count) == (1, 1) + assert (await backend.get_materialization(receipt.id)).reconsolidated_at is not None + assert len(await backend.list_bullets(context_id)) == 2 + assert len(await backend.list_activities(uuid.UUID(context_id))) == 1 + batches = await backend.list_delta_batches(context_id) + if hasattr(DeltaOpType, "RECONSOLIDATE_BULLET"): + reconsolidations = [batch for batch in batches if batch.trigger == "reconsolidation"] + assert len(reconsolidations) == 1 + assert len(batches) == 2 + for operation in reconsolidations[0].operations: + assert operation.agent_id == arguments["agent_id"] + assert operation.session_id == str(arguments["session_id"]) + else: + assert len(batches) == 1 + + foreign_context = await backend.create_context( + Context(name="foreign", intent=IntentAnchor(objective="isolated")) + ) + await backend.add_bullet(str(foreign_context.id), Bullet(id="foreign", content="Foreign fact")) + foreign_receipt = MaterializationRecord( + context_id=str(foreign_context.id), bullets_included=["foreign"] + ) + await backend.save_materialization(foreign_receipt) + arguments["materialization_id"] = foreign_receipt.id + await engine.commit(**arguments) + assert (await backend.get_materialization(foreign_receipt.id)).reconsolidated_at is None + assert (await backend.get_bullet("foreign")).hit_count == 0 + assert llm.completions == 1 + + +@pytest.mark.parametrize("operation", ["feedback", "batch"]) +async def test_context_delete_waits_for_atomic_child_mutations(storage, monkeypatch, operation): + import uuid + + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="child", content="fact")) + receipt = MaterializationRecord(context_id=context_id, bullets_included=["child"]) + await backend.save_materialization(receipt) + fetched = asyncio.Event() + release = asyncio.Event() + original = type(backend).get_bullet + + async def pause_after_child_lock(self, bullet_id): + bullet = await original(self, bullet_id) + if self._scope is not None and bullet_id == "child": + fetched.set() + await release.wait() + return bullet + + monkeypatch.setattr(type(backend), "get_bullet", pause_after_child_lock) + if operation == "feedback": + mutation = IngestionEngine(backend, None)._reconsolidate( + receipt.id, + ExecutionFeedback(outcome=FeedbackOutcome.SUCCESS), + expected_context_id=context_id, + ) + else: + mutation = DeltaEngine(backend).apply_batch( + DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation( + op_type=DeltaOpType.UPDATE_BULLET, + target_id="child", + content="updated", + ) + ], + ) + ) + writer = asyncio.create_task(mutation) + deletion = None + try: + await asyncio.wait_for(fetched.wait(), timeout=5) + deletion = asyncio.create_task(backend.delete_context(uuid.UUID(context_id))) + if isinstance(backend, PostgresBackend): + pool = await backend._get_pool() + for _ in range(100): + waiting = await pool.fetchval( + "SELECT EXISTS(SELECT 1 FROM pg_stat_activity WHERE datname=current_database() " + "AND query LIKE 'DELETE FROM contexts%' AND wait_event_type='Lock')" + ) + if waiting: + break + await asyncio.sleep(0.01) + assert waiting, "delete did not reach its database lock" + else: + await asyncio.sleep(0.02) + assert not deletion.done() + release.set() + await asyncio.wait_for(asyncio.gather(writer, deletion), timeout=10) + assert await backend.get_context(uuid.UUID(context_id)) is None + assert await backend.list_bullets(context_id) == [] + assert await backend.list_delta_batches(context_id) == [] + assert await backend.get_materialization(receipt.id) is None + finally: + release.set() + tasks = [task for task in (writer, deletion) if task is not None] + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/tests/test_consolidation_consistency.py b/tests/test_consolidation_consistency.py new file mode 100644 index 0000000..4dfe184 --- /dev/null +++ b/tests/test_consolidation_consistency.py @@ -0,0 +1,411 @@ +"""Consolidation and historical rollback never invent state from stale evidence.""" + +import uuid +from datetime import UTC, datetime, timedelta +from unittest.mock import AsyncMock + +import pytest +from test_storage_transactions import storage as transaction_storage + +from engram.core.consolidation import ConsolidationEngine +from engram.core.delta import DeltaEngine +from engram.core.ingestion import CuratorEngine +from engram.core.models import ( + Bullet, + ConsolidationConfig, + DeltaBatch, + DeltaOperation, + DeltaOpType, + Reflection, + ReflectionInsight, +) +from engram.storage.sqlite import SQLiteBackend + +storage = transaction_storage +VECTOR = [1.0] + [0.0] * 1535 + + +@pytest.mark.parametrize("embedding", [None, [float("nan"), 1.0], [0.0, 0.0]]) +@pytest.mark.parametrize("same_text", [True, False]) +async def test_reextract_preserves_identity_or_retains_uncertain_comparison(embedding, same_text): + old = Bullet( + id="old", content="Historical fact", embedding=embedding, salience=0.1, miss_count=5 + ) + llm = AsyncMock() + llm.embed.return_value = [1.0, 0.0] + result = await CuratorEngine(None, llm).curate_re_extraction( + Reflection( + new_insights=[ + ReflectionInsight(content=" historical FACT " if same_text else "new fact") + ] + ), + [old], + "context", + ) + assert all(op.op_type != DeltaOpType.REMOVE_BULLET for op in result.operations) + if same_text: + assert result.operations == [] + + +async def test_reextract_failed_provider_cannot_authorize_removal(): + old = Bullet( + id="old", content="Historical fact", embedding=[1.0, 0.0], salience=0.1, miss_count=5 + ) + llm = AsyncMock() + llm.embed.side_effect = RuntimeError("provider unavailable") + result = await CuratorEngine(None, llm).curate_re_extraction( + Reflection(new_insights=[ReflectionInsight(content="new fact")]), + [old], + "context", + ) + assert all(op.op_type != DeltaOpType.REMOVE_BULLET for op in result.operations) + + +@pytest.mark.parametrize( + "kind", [DeltaOpType.ADD_BULLET, DeltaOpType.UPDATE_BULLET, DeltaOpType.REMOVE_BULLET] +) +async def test_legacy_missing_identity_or_inverse_declines_whole_rollback(storage, kind): + backend, context_id = storage + for bid in ("unknown", "known"): + await backend.add_bullet(context_id, Bullet(id=bid, content=bid)) + unknown = DeltaOperation( + op_type=kind, + target_id=None if kind == DeltaOpType.ADD_BULLET else "unknown", + content="unknown", + ) + batch = DeltaBatch( + context_id=context_id, + operations=[ + unknown, + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id="known", content="known"), + ], + ) + await backend.save_delta_batch(batch) + assert not await DeltaEngine(backend).rollback_batch(batch.id) + assert (await backend.get_bullet("known")).is_active + assert (await backend.get_bullet("unknown")).is_active + + +async def test_new_explicit_noops_remain_rollbackable(storage): + backend, context_id = storage + operations = [ + DeltaOperation(op_type=kind, target_id="missing") + for kind in ( + DeltaOpType.UPDATE_BULLET, + DeltaOpType.REMOVE_BULLET, + DeltaOpType.UPDATE_SCHEMA, + ) + ] + operations.extend( + [ + DeltaOperation(op_type=DeltaOpType.ADD_SCHEMA), + DeltaOperation( + op_type=DeltaOpType.MERGE_BULLETS, target_ids=["missing", "also-missing"] + ), + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id="new", content="new"), + ] + ) + engine = DeltaEngine(backend) + batch = await engine.apply_batch(DeltaBatch(context_id=context_id, operations=operations)) + assert await engine.rollback_batch(batch.id) + assert not (await backend.get_bullet("new")).is_active + + +async def _facts(backend, context_id, count=3): + for i in range(count): + await backend.add_bullet( + context_id, + Bullet( + id=f"fact{i}", content=f"Fact {i}", section="facts", embedding=VECTOR, hit_count=2 + ), + ) + + +async def test_schema_induction_preserves_changes_during_provider_call(storage): + backend, context_id = storage + await _facts(backend, context_id, 4) + + async def complete(**kwargs): + await backend.archive_bullet(context_id, "fact0") + current = await backend.get_bullet("fact1") + current.content = "newer content" + current.hit_count = 9 + await backend.update_bullet(current) + return "Old pattern" + + llm = AsyncMock() + llm.complete.side_effect = complete + count = await ConsolidationEngine(backend, llm)._induce_schemas( + context_id, ConsolidationConfig() + ) + assert count == 0 + assert (await backend.get_bullet("fact0")).is_archived + assert (await backend.get_bullet("fact1")).content == "newer content" + assert (await backend.get_bullet("fact1")).hit_count == 9 + assert await backend.list_schemas(context_id) == [] + + +async def test_schema_and_links_rollback_together(storage, monkeypatch): + backend, context_id = storage + await _facts(backend, context_id) + llm = AsyncMock() + llm.complete.return_value = "pattern" + original = type(backend).update_bullet + + async def fail(self, bullet): + if bullet.id == "fact0": + raise RuntimeError("link unavailable") + return await original(self, bullet) + + monkeypatch.setattr(type(backend), "update_bullet", fail) + with pytest.raises(RuntimeError, match="link unavailable"): + await ConsolidationEngine(backend, llm)._induce_schemas(context_id, ConsolidationConfig()) + assert await backend.list_schemas(context_id) == [] + assert all(b.schema_id is None for b in await backend.list_bullets(context_id)) + assert await backend.list_delta_batches(context_id) == [] + + +async def test_concurrent_schema_induction_creates_one_schema(storage): + import asyncio + + backend, context_id = storage + await _facts(backend, context_id) + arrived = 0 + ready = asyncio.Event() + + async def complete(**kwargs): + nonlocal arrived + arrived += 1 + if arrived == 2: + ready.set() + await asyncio.wait_for(ready.wait(), 5) + return "pattern" + + llm = AsyncMock() + llm.complete.side_effect = complete + engines = [ConsolidationEngine(backend, llm) for _ in range(2)] + counts = await asyncio.gather( + *(engine._induce_schemas(context_id, ConsolidationConfig()) for engine in engines) + ) + assert sorted(counts) == [0, 1] + assert len(await backend.list_schemas(context_id)) == 1 + + +@pytest.mark.parametrize("phase", ["forgetting", "archive", "dedup"]) +async def test_consolidation_rechecks_locked_current_rows(storage, monkeypatch, phase): + backend, context_id = storage + old = datetime.now(UTC) - timedelta(days=90) + await backend.add_bullet( + context_id, + Bullet( + id="high", + content="old", + embedding=VECTOR, + salience=0.1 if phase == "archive" else 0.9, + created_at=old, + ), + ) + await backend.add_bullet( + context_id, Bullet(id="low", content="other", embedding=VECTOR, salience=0.5) + ) + stale = await backend.list_bullets(context_id) + if phase == "archive": + current = await backend.get_bullet("high") + current.salience = 0.95 + current.hit_count = 10 + current.last_recalled_at = datetime.now(UTC) + await backend.update_bullet(current) + else: + await backend.archive_bullet(context_id, "high") + original = type(backend).list_bullets + first = True + + async def stale_once(self, *args, **kwargs): + nonlocal first + if first: + first = False + return [b.model_copy(deep=True) for b in stale] + return await original(self, *args, **kwargs) + + monkeypatch.setattr(type(backend), "list_bullets", stale_once) + engine = ConsolidationEngine(backend, None) + method = { + "forgetting": engine._apply_forgetting_curve, + "archive": engine._archive_stale, + "dedup": engine._semantic_dedup, + }[phase] + await method(context_id, ConsolidationConfig()) + high = await backend.get_bullet("high") + if phase == "archive": + assert not high.is_archived + assert high.hit_count == 10 + assert high.salience == pytest.approx(0.95) + else: + assert high.is_archived + assert (await backend.get_bullet("low")).is_active + + +@pytest.mark.parametrize("change", ["archive", "capacity", "existing"]) +async def test_promotion_rechecks_current_qualification_duplicates_and_capacity(storage, change): + backend, context_id = storage + await _facts(backend, context_id) + if change == "capacity": + context = await backend.get_context(uuid.UUID(context_id)) + context.lifecycle_config.max_active_bullets = 3 + await backend.update_context(context) + + async def complete(**kwargs): + if change == "archive": + await backend.archive_bullet(context_id, "fact0") + elif change == "existing": + await backend.add_bullet( + context_id, + Bullet( + id="existing-principle", + content="existing", + section="facts", + bullet_type="principle", + ), + ) + return "New principle" + + llm = AsyncMock() + llm.complete.side_effect = complete + assert ( + await ConsolidationEngine(backend, llm)._promote_to_principles( + context_id, ConsolidationConfig() + ) + == 0 + ) + principles = await backend.list_bullets(context_id, bullet_type="principle") + assert len(principles) == (1 if change == "existing" else 0) + + +async def test_lifecycle_migration_backfill_failure_is_retry_safe(storage, monkeypatch): + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="archived", content="preserve")) + await backend.archive_bullet(context_id, "archived") + if isinstance(backend, SQLiteBackend): + import aiosqlite + + db = await backend._get_db() + await db.execute("DROP INDEX idx_bullets_lifecycle") + await db.execute("ALTER TABLE bullets DROP COLUMN lifecycle_state") + await db.execute("ALTER TABLE bullets DROP COLUMN archive_reason") + await db.commit() + connection_type = aiosqlite.Connection + else: + import asyncpg + + pool = await backend._get_pool() + await pool.execute("ALTER TABLE bullets DROP COLUMN lifecycle_state") + await pool.execute("ALTER TABLE bullets DROP COLUMN archive_reason") + connection_type = asyncpg.Connection + original = connection_type.execute + + async def fail(self, query, *args, **kwargs): + if "UPDATE bullets SET lifecycle_state='archived'" in query: + raise RuntimeError("backfill interrupted") + return await original(self, query, *args, **kwargs) + + with monkeypatch.context() as local: + local.setattr(connection_type, "execute", fail) + with pytest.raises(RuntimeError, match="backfill interrupted"): + await backend.initialize() + await backend.initialize() + assert [b.id for b in await backend.get_archived_bullets(context_id)] == ["archived"] + assert (await backend.restore_bullet(context_id, "archived")).content == "preserve" + + +async def test_explicit_zero_confidence_update_and_rollback(storage): + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="zero", content="same", confidence=0.8)) + before = await backend.get_bullet("zero") + engine = DeltaEngine(backend) + batch = await engine.apply_batch( + DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.UPDATE_BULLET, target_id="zero", confidence=0.0), + ], + ) + ) + assert (await backend.get_bullet("zero")).confidence == 0.0 + assert await engine.rollback_batch(batch.id) + assert (await backend.get_bullet("zero")).confidence == before.confidence + + +@pytest.mark.parametrize("invalid", [[0.0, 0.0], [float("nan"), 1.0], [1.79e308, 1.79e308], [1.0]]) +async def test_destructive_comparison_rejects_unknown_vectors_at_zero_threshold( + storage, monkeypatch, invalid +): + backend, context_id = storage + for bid in ("one", "two"): + await backend.add_bullet(context_id, Bullet(id=bid, content=bid, embedding=VECTOR)) + original = type(backend).get_bullet + + async def malformed(self, bullet_id): + bullet = await original(self, bullet_id) + bullet.embedding = invalid if bullet_id == "one" else [1.0, 0.0] + return bullet + + monkeypatch.setattr(type(backend), "get_bullet", malformed) + assert ( + await ConsolidationEngine(backend, None)._semantic_dedup( + context_id, ConsolidationConfig(dedup_threshold=0.0) + ) + == 0 + ) + assert (await backend.get_bullet("one")).is_active + assert (await backend.get_bullet("two")).is_active + + +@pytest.mark.parametrize("storage", ["postgres"], indirect=True) +async def test_postgres_archive_between_candidate_list_and_lock_preserves_active_donor( + storage, monkeypatch +): + from engram.storage.postgres import PostgresBackend + + backend, context_id = storage + await backend.add_bullet( + context_id, Bullet(id="high", content="high", embedding=VECTOR, salience=0.9) + ) + await backend.add_bullet( + context_id, Bullet(id="low", content="low", embedding=VECTOR, salience=0.5) + ) + second = PostgresBackend(backend.dsn, existing_only=True) + original = type(backend).list_bullets + injected = False + + async def archive_between(self, *args, **kwargs): + nonlocal injected + bullets = await original(self, *args, **kwargs) + if self._scope is not None and not injected: + injected = True + assert await second.archive_bullet(context_id, "high") + return bullets + + monkeypatch.setattr(type(backend), "list_bullets", archive_between) + try: + assert ( + await ConsolidationEngine(backend, None)._semantic_dedup( + context_id, ConsolidationConfig() + ) + == 0 + ) + assert (await backend.get_bullet("high")).is_archived + assert (await backend.get_bullet("low")).is_active + finally: + await second.close() + + +async def test_reused_noop_operation_captures_fresh_inverse_when_target_appears(storage): + backend, context_id = storage + engine = DeltaEngine(backend) + op = DeltaOperation(op_type=DeltaOpType.UPDATE_BULLET, target_id="later", content="changed") + await engine.apply_batch(DeltaBatch(context_id=context_id, operations=[op])) + await backend.add_bullet(context_id, Bullet(id="later", content="original")) + changed = await engine.apply_batch(DeltaBatch(context_id=context_id, operations=[op])) + assert (await backend.get_bullet("later")).content == "changed" + assert await engine.rollback_batch(changed.id) + assert (await backend.get_bullet("later")).content == "original" diff --git a/tests/test_context_isolation.py b/tests/test_context_isolation.py new file mode 100644 index 0000000..797f628 --- /dev/null +++ b/tests/test_context_isolation.py @@ -0,0 +1,154 @@ +"""Context boundaries must hold even when an item id belongs to another context.""" + +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException + +from engram.core.delta import DeltaEngine +from engram.core.graph import ConceptGraph +from engram.core.ingestion import IngestionEngine +from engram.core.models import ( + Bullet, + ConceptEdge, + ConceptNode, + ConceptType, + Context, + DeltaBatch, + DeltaOperation, + DeltaOpType, + EdgeType, + ExecutionFeedback, + FeedbackOutcome, + IntentAnchor, + InvalidateRequest, + MaterializationRecord, + SchemaNode, +) +from engram.server.routes.bullets import get_bullet +from engram.server.routes.contexts import invalidate_concepts +from engram.server.routes.deltas import get_delta, rollback_delta +from engram.server.routes.schemas import get_schema +from engram.storage.sqlite import SQLiteBackend + + +@pytest.fixture +async def contexts(tmp_path): + storage = SQLiteBackend(str(tmp_path / "isolation.db")) + await storage.initialize() + a = Context(name="a", owner="owner", intent=IntentAnchor(objective="a")) + b = Context(name="b", owner="owner", intent=IntentAnchor(objective="b")) + await storage.create_context(a) + await storage.create_context(b) + request = SimpleNamespace( + state=SimpleNamespace(user_id="owner"), + app=SimpleNamespace( + state=SimpleNamespace( + storage=storage, + delta_engine=DeltaEngine(storage), + graph=ConceptGraph(storage), + ) + ), + ) + yield storage, a, b, request + await storage.close() + + +@pytest.mark.parametrize("kind", ["bullet", "schema", "delta"]) +async def test_child_get_requires_matching_context(contexts, kind): + storage, a, b, request = contexts + if kind == "bullet": + obj = await storage.add_bullet(str(a.id), Bullet(content="private")) + getter = get_bullet + elif kind == "schema": + obj = await storage.add_schema(str(a.id), SchemaNode(name="private", description="pattern")) + getter = get_schema + else: + obj = await storage.save_delta_batch(DeltaBatch(context_id=str(a.id))) + getter = get_delta + with pytest.raises(HTTPException) as exc: + await getter(b.id, obj.id, request) + assert exc.value.status_code == 404 + assert (await getter(a.id, obj.id, request))["id"] == obj.id + + +async def test_foreign_rollback_cannot_mutate(contexts): + storage, a, b, request = contexts + bullet = await storage.add_bullet(str(a.id), Bullet(content="before")) + batch = DeltaBatch( + context_id=str(a.id), + operations=[ + DeltaOperation(op_type=DeltaOpType.UPDATE_BULLET, target_id=bullet.id, content="after") + ], + ) + await request.app.state.delta_engine.apply_batch(batch) + with pytest.raises(HTTPException) as exc: + await rollback_delta(b.id, batch.id, request) + assert exc.value.status_code == 404 + assert (await storage.get_bullet(bullet.id)).content == "after" + await rollback_delta(a.id, batch.id, request) + assert (await storage.get_bullet(bullet.id)).content == "before" + + +async def test_invalidation_ignores_foreign_concepts_and_bullets(contexts): + storage, a, b, request = contexts + own = await storage.add_concept(a.id, ConceptNode(type=ConceptType.FACT, content="own")) + foreign = await storage.add_concept(b.id, ConceptNode(type=ConceptType.FACT, content="foreign")) + own_bullet = await storage.add_bullet(str(a.id), Bullet(content="own")) + foreign_bullet = await storage.add_bullet(str(b.id), Bullet(content="foreign")) + await invalidate_concepts( + a.id, + InvalidateRequest( + concept_ids=[own.id, foreign.id], + bullet_ids=[own_bullet.id, foreign_bullet.id], + reason="test", + ), + request, + ) + assert not (await storage.get_concept(own.id)).is_valid + assert (await storage.get_concept(foreign.id)).is_valid + assert not (await storage.get_bullet(own_bullet.id)).is_active + assert (await storage.get_bullet(foreign_bullet.id)).is_active + + +async def test_neighborhood_does_not_follow_foreign_nodes(contexts): + storage, a, b, request = contexts + own = await storage.add_concept(a.id, ConceptNode(type=ConceptType.FACT, content="own")) + foreign = await storage.add_concept(b.id, ConceptNode(type=ConceptType.FACT, content="foreign")) + await storage.add_edge( + a.id, ConceptEdge(from_node=own.id, to_node=foreign.id, type=EdgeType.RELATED_TO) + ) + graph = request.app.state.graph + assert await graph.get_concept_neighborhood(a.id, foreign.id) == [] + assert [c.id for c in await graph.get_concept_neighborhood(a.id, own.id)] == [own.id] + + +async def test_feedback_ignores_foreign_materialization(contexts): + storage, a, b, _ = contexts + bullet = await storage.add_bullet(str(b.id), Bullet(content="foreign")) + rec = MaterializationRecord(context_id=str(b.id), bullets_included=[bullet.id], token_count=1) + await storage.save_materialization(rec) + engine = IngestionEngine(storage, None) + await engine._reconsolidate( + rec.id, ExecutionFeedback(outcome=FeedbackOutcome.SUCCESS), expected_context_id=str(a.id) + ) + assert (await storage.get_bullet(bullet.id)).hit_count == 0 + await engine._reconsolidate( + rec.id, ExecutionFeedback(outcome=FeedbackOutcome.SUCCESS), expected_context_id=str(b.id) + ) + assert (await storage.get_bullet(bullet.id)).hit_count == 1 + + +async def test_feedback_checks_each_bullet_context(contexts): + storage, a, b, _ = contexts + own = await storage.add_bullet(str(a.id), Bullet(content="own")) + foreign = await storage.add_bullet(str(b.id), Bullet(content="foreign")) + rec = MaterializationRecord( + context_id=str(a.id), bullets_included=[own.id, foreign.id], token_count=1 + ) + await storage.save_materialization(rec) + await IngestionEngine(storage, None)._reconsolidate( + rec.id, ExecutionFeedback(outcome=FeedbackOutcome.SUCCESS), expected_context_id=str(a.id) + ) + assert (await storage.get_bullet(own.id)).hit_count == 1 + assert (await storage.get_bullet(foreign.id)).hit_count == 0 diff --git a/tests/test_core_regressions.py b/tests/test_core_regressions.py new file mode 100644 index 0000000..5d5ce82 --- /dev/null +++ b/tests/test_core_regressions.py @@ -0,0 +1,253 @@ +"""Persisted-state regressions for semantic ingestion and delta rollback.""" + +import json +from types import SimpleNamespace + +import pytest + +from engram.core.config import IngestionConfig +from engram.core.delta import DeltaEngine +from engram.core.ingestion import CuratorEngine, IngestionEngine +from engram.core.models import ( + Bullet, + Context, + DeltaBatch, + DeltaOperation, + DeltaOpType, + IntentAnchor, + RecordDecisionRequest, + Reflection, + ReflectionInsight, +) +from engram.llm.adapter import LLMAdapter +from engram.server.routes.concepts import record_decision +from engram.storage.sqlite import SQLiteBackend + + +class SemanticLLM(LLMAdapter): + def __init__(self): + self.reflection = Reflection() + self.vectors = {} + self.fail_embeddings = False + + async def complete(self, *args, **kwargs): + return self.reflection.model_dump_json() + + async def embed(self, text): + if self.fail_embeddings: + raise RuntimeError("Embedding provider unavailable") + return self.vectors.get(text, [1.0, 0.0]) + + +@pytest.fixture +async def state(tmp_path): + storage = SQLiteBackend(str(tmp_path / "core.db")) + await storage.initialize() + context = Context(name="regression", intent=IntentAnchor(objective="test")) + await storage.create_context(context) + llm = SemanticLLM() + engine = IngestionEngine( + storage, llm, ingestion_config=IngestionConfig(max_reflection_rounds=1) + ) + yield storage, context, llm, engine + await storage.close() + + +async def test_commit_semantic_search_and_paraphrase_dedup(state): + storage, context, llm, engine = state + short = "Postgres supports our application" + rich = "PostgreSQL supports our application with indexed vector search" + llm.vectors = {short: [1.0, 0.0], rich: [0.98, 0.2]} + llm.reflection = Reflection(new_insights=[ReflectionInsight(content=short)]) + await engine.commit(context.id, "agent", "first raw input") + hits = await storage.find_similar_bullets(str(context.id), [1.0, 0.0], threshold=0.9) + assert len(hits) == 1 + original_id = hits[0][0].id + + llm.reflection = Reflection(new_insights=[ReflectionInsight(content=rich)]) + batch = await engine.commit(context.id, "agent", "second raw input") + bullets = await storage.list_bullets(str(context.id)) + assert len(bullets) == 1 + assert bullets[0].id == original_id + assert bullets[0].content == rich + assert bullets[0].embedding == llm.vectors[rich] + # Stored deltas survive JSON serialization and retain the matching vector. + saved = await storage.get_delta_batch(batch.id) + assert json.loads(saved.model_dump_json())["operations"][0]["embedding"] == llm.vectors[rich] + + +async def test_shorter_paraphrase_keeps_survivor_content_and_vector(state): + storage, context, llm, engine = state + long_text = "Postgres is fast for indexed vector search" + await storage.add_bullet( + str(context.id), Bullet(id="survivor", content=long_text, embedding=[1.0, 0.0]) + ) + llm.vectors["Postgres is fast"] = [0.98, 0.2] + llm.reflection = Reflection(new_insights=[ReflectionInsight(content="Postgres is fast")]) + await engine.commit(context.id, "agent", "shorter paraphrase") + bullet = await storage.get_bullet("survivor") + assert bullet.content == long_text + assert bullet.embedding == [1.0, 0.0] + + +async def test_auxiliary_memories_are_searchable(state): + storage, context, llm, engine = state + llm.reflection = Reflection( + strategies_that_worked=["Benchmark before choosing"], + failure_modes=["Requests timed out without retry"], + prediction_errors=["Throughput was lower than expected"], + ) + await engine.commit(context.id, "agent", "auxiliary reflection") + bullets = await storage.list_bullets(str(context.id)) + assert len(bullets) == 3 + assert all(b.embedding == [1.0, 0.0] for b in bullets) + + +async def test_re_extraction_add_and_update_vectors(state): + storage, context, llm, _ = state + old = Bullet(id="old", content="Old wording", embedding=[1.0, 0.0]) + await storage.add_bullet(str(context.id), old) + llm.vectors = {"Improved wording": [0.96, 0.28], "Independent fact": [0.0, 1.0]} + reflection = Reflection( + new_insights=[ + ReflectionInsight(content="Improved wording"), + ReflectionInsight(content="Independent fact"), + ] + ) + batch = await CuratorEngine(storage, llm).curate_re_extraction( + reflection, [old], str(context.id) + ) + await DeltaEngine(storage).apply_batch(batch) + bullets = await storage.list_bullets(str(context.id)) + assert len(bullets) == 2 + assert {b.content: b.embedding for b in bullets} == llm.vectors + + +async def test_decisions_and_direct_add_persist_vectors_in_delta(state): + storage, context, llm, engine = state + _, batch = await engine.record_decision(context.id, "Use Postgres", "Reliable", ["SQLite"], "a") + direct, direct_batch = await engine.add_bullet_directly(str(context.id), "Manual fact") + assert direct.embedding == [1.0, 0.0] + assert len(await storage.find_similar_bullets(str(context.id), [1.0, 0.0])) == 3 + for item in (batch, direct_batch): + saved = await storage.get_delta_batch(item.id) + assert all(op.embedding == [1.0, 0.0] for op in saved.operations) + + +async def test_decision_route_returns_the_persisted_bullet_id(state): + storage, context, _, engine = state + request = SimpleNamespace( + state=SimpleNamespace(user_id=context.owner), + app=SimpleNamespace(state=SimpleNamespace(storage=storage, ingestion=engine)), + ) + response = await record_decision( + context.id, + RecordDecisionRequest(decision="Use Postgres", rationale="Reliable", agent_id="test"), + request, + ) + data = response.model_dump() if hasattr(response, "model_dump") else response + bullet = await storage.get_bullet(data["concept_id"]) + assert bullet is not None + assert bullet.content == "Use Postgres. Rationale: Reliable" + assert bullet.context_id == str(context.id) + + +async def test_embedding_failure_keeps_commit_available(state): + storage, context, llm, engine = state + llm.fail_embeddings = True + llm.reflection = Reflection(new_insights=[ReflectionInsight(content="A durable fact")]) + await engine.commit(context.id, "agent", "content") + bullets = await storage.list_bullets(str(context.id)) + assert len(bullets) == 1 + assert bullets[0].embedding is None + + +@pytest.mark.parametrize("embedding", [[0.0, 1.0], None]) +async def test_update_vector_and_rollback_match_content(state, embedding): + storage, context, _, _ = state + await storage.add_bullet( + str(context.id), Bullet(id="updatable", content="Original", embedding=[1.0, 0.0]) + ) + engine = DeltaEngine(storage) + batch = DeltaBatch( + context_id=str(context.id), + operations=[ + DeltaOperation( + op_type=DeltaOpType.UPDATE_BULLET, + target_id="updatable", + content="Changed", + embedding=embedding, + ) + ], + ) + await engine.apply_batch(batch) + updated = await storage.get_bullet("updatable") + assert updated.content == "Changed" + assert updated.embedding == embedding # Failed embedding must clear an obsolete vector. + await engine.rollback_batch(batch.id) + restored = await storage.get_bullet("updatable") + assert restored.content == "Original" + assert restored.embedding == [1.0, 0.0] + + +async def test_metadata_update_retains_vector_and_legacy_delta_loads(state): + storage, context, _, _ = state + await storage.add_bullet( + str(context.id), Bullet(id="metadata", content="Original", embedding=[1.0, 0.0]) + ) + legacy_op = DeltaOperation.model_validate( + {"op_type": "update_bullet", "target_id": "metadata", "section": "new-section"} + ) + await DeltaEngine(storage).apply_batch( + DeltaBatch(context_id=str(context.id), operations=[legacy_op]) + ) + updated = await storage.get_bullet("metadata") + assert updated.section == "new-section" + assert updated.embedding == [1.0, 0.0] + + +async def test_legacy_rollback_does_not_keep_new_content_vector(state): + storage, context, _, _ = state + await storage.add_bullet( + str(context.id), Bullet(id="legacy", content="New content", embedding=[0.0, 1.0]) + ) + legacy = DeltaBatch.model_validate( + { + "context_id": str(context.id), + "operations": [ + { + "op_type": "update_bullet", + "target_id": "legacy", + "content": "New content", + "previous_state": {"content": "Original content"}, + } + ], + } + ) + await storage.save_delta_batch(legacy) + await DeltaEngine(storage).rollback_batch(legacy.id) + restored = await storage.get_bullet("legacy") + assert restored.content == "Original content" + assert restored.embedding is None + + +@pytest.mark.parametrize("target_id", ["explicit", None]) +async def test_add_rollback_without_snapshot_in_mixed_batch(state, target_id): + storage, context, _, _ = state + await storage.add_bullet(str(context.id), Bullet(id="existing", content="Original")) + batch = DeltaBatch( + context_id=str(context.id), + operations=[ + DeltaOperation( + op_type=DeltaOpType.UPDATE_BULLET, target_id="existing", content="Changed" + ), + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id=target_id, content="New"), + ], + ) + engine = DeltaEngine(storage) + await engine.apply_batch(batch) + saved = await storage.get_delta_batch(batch.id) + assert saved.operations[1].target_id is not None + assert await engine.rollback_batch(batch.id) + active = await storage.list_bullets(str(context.id)) + assert [(b.id, b.content) for b in active] == [("existing", "Original")] diff --git a/tests/test_embedding_backfill.py b/tests/test_embedding_backfill.py new file mode 100644 index 0000000..bb00126 --- /dev/null +++ b/tests/test_embedding_backfill.py @@ -0,0 +1,376 @@ +"""Exercise repair only against disposable databases and fake providers.""" + +import asyncio +import hashlib +import json +import logging +import os +import sqlite3 +import sys +import uuid +from pathlib import Path +from unittest.mock import AsyncMock +from urllib.parse import urlsplit, urlunsplit + +import asyncpg +import pytest + +from engram.core.models import Bullet, Context, IntentAnchor +from engram.maintenance import backfill_embeddings as repair_module +from engram.maintenance.backfill_embeddings import EMBEDDING_DIMENSIONS, backfill_embeddings +from engram.storage.postgres import PostgresBackend +from engram.storage.sqlite import SQLiteBackend + +VECTOR = [1.0] + [0.0] * (EMBEDDING_DIMENSIONS - 1) + + +@pytest.fixture(params=["sqlite", "postgres"]) +async def repair_storage(request, tmp_path): + if request.param == "sqlite": + backend = SQLiteBackend(str(tmp_path / "repair.db")) + await backend.initialize() + try: + yield backend + finally: + await backend.close() + return + dsn = os.environ.get("ENGRAM_TEST_POSTGRES_DSN") + if not dsn: + pytest.skip("requires a disposable PostgreSQL admin DSN") + admin = await asyncpg.connect(dsn) + name = "engram_backfill_" + uuid.uuid4().hex + target = urlunsplit(urlsplit(dsn)._replace(path="/" + name)) + backend = PostgresBackend(target) + try: + await admin.execute(f'CREATE DATABASE "{name}"') + await backend.initialize() + yield backend + finally: + await backend.close() + await admin.execute(f'DROP DATABASE IF EXISTS "{name}" WITH (FORCE)') + await admin.close() + + +async def seed(storage): + context = await storage.create_context( + Context(name="repair", intent=IntentAnchor(objective="preserve existing memories")) + ) + bullet = await storage.add_bullet( + str(context.id), Bullet(id="a-missing", content="Use dark mode", hit_count=3, salience=0.7) + ) + return context, bullet + + +async def test_preview_never_calls_provider_or_transaction(repair_storage, monkeypatch): + context, bullet = await seed(repair_storage) + llm = AsyncMock() + transaction = AsyncMock(side_effect=AssertionError("preview attempted a transaction")) + monkeypatch.setattr(repair_storage, "transaction", transaction) + result = await backfill_embeddings(repair_storage, context.id, llm=llm) + assert result["dry_run"] and result["eligible"] == result["selected"] == 1 + assert result["attempted"] == result["updated"] == result["failed"] == 0 + llm.embed.assert_not_called() + transaction.assert_not_called() + assert (await repair_storage.get_bullet(bullet.id)).embedding is None + + +async def test_repairs_only_embedding_and_rerun_is_idempotent(repair_storage): + context, bullet = await seed(repair_storage) + before = (await repair_storage.get_bullet(bullet.id)).model_dump() + llm = AsyncMock() + llm.embed.return_value = VECTOR + first = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True) + assert first["updated"] == 1 and first["failed"] == 0 + after = (await repair_storage.get_bullet(bullet.id)).model_dump() + assert after.pop("embedding") == pytest.approx(VECTOR) + before.pop("embedding") + assert after == before # Includes timestamps, lifecycle, counters, sources. + assert await repair_storage.list_delta_batches(str(context.id)) == [] + again = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True) + assert again["eligible"] == again["attempted"] == 0 + llm.embed.assert_awaited_once_with(bullet.content) + + +async def test_limit_context_and_lifecycle_are_respected(repair_storage): + context, _ = await seed(repair_storage) + for bid in ("b-missing", "c-archived", "d-inactive"): + await repair_storage.add_bullet(str(context.id), Bullet(id=bid, content=bid)) + await repair_storage.archive_bullet(str(context.id), "c-archived") + await repair_storage.remove_bullet("d-inactive") + foreign = await repair_storage.create_context( + Context(name="foreign", intent=IntentAnchor(objective="untouched")) + ) + await repair_storage.add_bullet(str(foreign.id), Bullet(id="foreign", content="Private")) + llm = AsyncMock() + llm.embed.return_value = VECTOR + result = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True, limit=1) + assert result["eligible"] == 2 and result["updated"] == result["selected"] == 1 + for bid in ("b-missing", "c-archived", "d-inactive", "foreign"): + assert (await repair_storage.get_bullet(bid)).embedding is None + + +@pytest.mark.parametrize("race", ["text", "archive", "already_repaired"]) +async def test_provider_result_is_skipped_after_concurrent_change(repair_storage, race): + context, bullet = await seed(repair_storage) + + async def embed(_): + if race == "archive": + await repair_storage.archive_bullet(str(context.id), bullet.id) + else: + current = await repair_storage.get_bullet(bullet.id) + if race == "text": + current.content = "A newer instruction" + else: + current.embedding = [0.0, 1.0] + [0.0] * (EMBEDDING_DIMENSIONS - 2) + await repair_storage.update_bullet(current) + return VECTOR + + llm = AsyncMock() + llm.embed.side_effect = embed + result = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True) + assert result["skipped"] == 1 and result["updated"] == result["failed"] == 0 + after = await repair_storage.get_bullet(bullet.id) + if race == "text": + assert after.content == "A newer instruction" and after.embedding is None + elif race == "archive": + assert after.is_archived and after.embedding is None + else: + assert after.embedding[1] == 1.0 + + +@pytest.mark.parametrize("invalid", [ + [], [1.0], [float("nan")] * 1536, [0.0] * 1536, + [1e-300] + [0.0] * 1535, [10**1000] + [0.0] * 1535, + [1e100] + [0.0] * 1535, [1e-30] + [0.0] * 1535, [1e20] + [0.0] * 1535, +]) +async def test_invalid_provider_vectors_are_not_written(repair_storage, invalid): + context, bullet = await seed(repair_storage) + llm = AsyncMock() + llm.embed.return_value = invalid + result = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True) + assert result["failed"] == 1 and result["errors"][0]["reason"] == "invalid_vector" + assert (await repair_storage.get_bullet(bullet.id)).embedding is None + + +async def test_every_saved_repair_is_searchable_by_its_own_vector(repair_storage): + context, _ = await seed(repair_storage) + for bid in ("b-underflow", "c-overflow"): + await repair_storage.add_bullet(str(context.id), Bullet(id=bid, content=bid)) + llm = AsyncMock() + llm.embed.side_effect = [VECTOR, [1e-30] + [0.0] * 1535, [1e20] + [0.0] * 1535] + result = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True) + assert (result["updated"], result["failed"]) == (1, 2) + for bullet in await repair_storage.list_bullets(str(context.id)): + if bullet.embedding is not None: + matches = await repair_storage.find_similar_bullets( + str(context.id), bullet.embedding, threshold=0.99, + ) + assert any(found.id == bullet.id for found, _ in matches) + retry = await backfill_embeddings(repair_storage, context.id) + assert retry["eligible"] == 2 # Rejected vectors remain repairable. + + +async def test_provider_failure_preserves_progress_and_redacts_message(repair_storage): + context, bullet = await seed(repair_storage) + await repair_storage.add_bullet(str(context.id), Bullet(id="b-second", content="Second")) + llm = AsyncMock() + llm.embed.side_effect = [RuntimeError("credential-should-never-be-printed"), VECTOR] + result = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True) + assert (result["updated"], result["failed"]) == (1, 1) + assert "credential-should-never-be-printed" not in json.dumps(result) + assert (await repair_storage.get_bullet(bullet.id)).embedding is None + llm.embed.side_effect = None + llm.embed.return_value = VECTOR + retry = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True) + assert retry["updated"] == 1 and retry["failed"] == 0 + + +async def test_failed_vector_write_rolls_back_and_can_be_retried(repair_storage, monkeypatch): + context, bullet = await seed(repair_storage) + cls = type(repair_storage) + original = cls.update_bullet_embedding_if_missing + + async def fail_after_write(self, *args): + await original(self, *args) + raise RuntimeError("injected after SQL") + + llm = AsyncMock() + llm.embed.return_value = VECTOR + with monkeypatch.context() as patch: + patch.setattr(cls, "update_bullet_embedding_if_missing", fail_after_write) + result = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True) + assert result["failed"] == 1 + assert (await repair_storage.get_bullet(bullet.id)).embedding is None + retry = await backfill_embeddings(repair_storage, context.id, llm=llm, apply=True) + assert retry["updated"] == 1 + + +async def test_read_only_backend_previews_existing_database(repair_storage): + context, _ = await seed(repair_storage) + if isinstance(repair_storage, SQLiteBackend): + reader = SQLiteBackend(repair_storage.db_path, read_only=True, existing_only=True) + else: + reader = PostgresBackend(repair_storage.dsn, read_only=True, existing_only=True) + try: + assert (await backfill_embeddings(reader, context.id))["eligible"] == 1 + with pytest.raises(Exception): + await reader.update_bullet_embedding_if_missing( + str(context.id), "a-missing", "Use dark mode", VECTOR, + ) + finally: + await reader.close() + + +async def test_independent_repair_workers_save_only_one_vector(repair_storage): + context, bullet = await seed(repair_storage) + if isinstance(repair_storage, SQLiteBackend): + other = SQLiteBackend(repair_storage.db_path, existing_only=True) + else: + other = PostgresBackend(repair_storage.dsn, existing_only=True) + arrived = asyncio.Event() + calls = 0 + + async def embed(_): + nonlocal calls + calls += 1 + if calls == 2: + arrived.set() + await arrived.wait() + return VECTOR + + llm = AsyncMock() + llm.embed.side_effect = embed + try: + results = await asyncio.wait_for(asyncio.gather( + backfill_embeddings(repair_storage, context.id, llm=llm, apply=True), + backfill_embeddings(other, context.id, llm=llm, apply=True), + ), timeout=15) + finally: + await other.close() + assert sorted(result["updated"] for result in results) == [0, 1] + assert sum(result["skipped"] for result in results) == 1 + assert all(result["failed"] == 0 for result in results) + assert (await repair_storage.get_bullet(bullet.id)).embedding == pytest.approx(VECTOR) + + +@pytest.mark.parametrize("limit", [0, -1, 1001]) +async def test_invalid_limits_fail_before_opening_storage(limit): + storage = AsyncMock() + with pytest.raises(ValueError, match="limit"): + await backfill_embeddings(storage, uuid.uuid4(), limit=limit) + storage.get_context.assert_not_called() + + +async def run_cli(*args, env=None): + process = await asyncio.create_subprocess_exec( + sys.executable, "-m", "engram.maintenance.backfill_embeddings", *args, + stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=env, + ) + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=20) + return process.returncode, stdout.decode(), stderr.decode() + + +async def test_cli_preview_does_not_change_sqlite_file_or_journal_mode(tmp_path): + path = tmp_path / "existing.db" + store = SQLiteBackend(str(path)) + await store.initialize() + context, _ = await seed(store) + await store.close() + with sqlite3.connect(path) as connection: + assert connection.execute("PRAGMA journal_mode=DELETE").fetchone()[0] == "delete" + digest = hashlib.sha256(path.read_bytes()).hexdigest() + code, stdout, _ = await run_cli("--sqlite", str(path), "--context-id", str(context.id)) + assert code == 0 and json.loads(stdout)["dry_run"] + assert hashlib.sha256(path.read_bytes()).hexdigest() == digest + assert not Path(str(path) + "-wal").exists() + + +async def test_cli_missing_database_is_not_created_and_apply_requires_explicit_provider(tmp_path): + target = tmp_path / "must-not-exist.db" + args = ("--sqlite", str(target), "--context-id", str(uuid.uuid4())) + code, stdout, _ = await run_cli(*args) + assert code == 1 and json.loads(stdout)["error"] == "database_operation_failed" + assert not target.exists() + code, _, stderr = await run_cli(*args, "--apply") + assert code == 2 and "requires --embedding-model" in stderr + assert not target.exists() + + +async def test_cli_requires_a_database_target(): + code, _, stderr = await run_cli("--context-id", str(uuid.uuid4())) + assert code == 2 and "required" in stderr + + +async def test_cli_apply_failure_returns_nonzero_without_raw_provider_log( + tmp_path, monkeypatch, capsys, caplog, +): + path = tmp_path / "repair.db" + storage = SQLiteBackend(str(path)) + await storage.initialize() + context, _ = await seed(storage) + await storage.close() + secret = "a-provider-key-that-must-not-be-logged" + + async def fail(_): + logging.getLogger("engram.llm.adapter").error("provider said %s", secret) + raise RuntimeError(secret) + + llm = AsyncMock() + llm.embed.side_effect = fail + monkeypatch.setattr(repair_module, "LiteLLMAdapter", lambda **kwargs: llm) + monkeypatch.setenv("ENGRAM_REPAIR_TEST_KEY", secret) + args = repair_module._parser().parse_args([ + "--sqlite", str(path), "--context-id", str(context.id), "--apply", + "--embedding-model", "test-only", "--embedding-key-env", "ENGRAM_REPAIR_TEST_KEY", + ]) + assert await repair_module._run(args) == 1 + captured = capsys.readouterr() + assert json.loads(captured.out)["failed"] == 1 + assert secret not in captured.out + captured.err + caplog.text + + +async def test_real_adapter_invalid_model_does_not_load_dotenv_or_corrupt_json(tmp_path): + path = tmp_path / "repair.db" + storage = SQLiteBackend(str(path)) + await storage.initialize() + context, _ = await seed(storage) + await storage.close() + (tmp_path / ".env").write_text("ENGRAM_DOTENV_CANARY=must-not-load\n") + script = ''' +import os, socket, sys +import dotenv +calls = [] +def forbid_dotenv(*args, **kwargs): + calls.append("dotenv") + raise AssertionError("implicit dotenv access") +def forbid_network(*args, **kwargs): + calls.append("network") + raise AssertionError("network access") +dotenv.load_dotenv = forbid_dotenv +socket.socket.connect = forbid_network +from engram.maintenance.backfill_embeddings import main +code = main(["--sqlite", sys.argv[1], "--context-id", sys.argv[2], "--apply", + "--embedding-model", "intentionally-unsupported-test-model", + "--embedding-key-env", "ENGRAM_REPAIR_TEST_KEY"]) +assert calls == [], calls +assert "ENGRAM_DOTENV_CANARY" not in os.environ +assert os.environ["LITELLM_MODE"] == "DEV" +raise SystemExit(code) +''' + env = { + **os.environ, "ENGRAM_REPAIR_TEST_KEY": "fake-provider-key-canary", + "LITELLM_MODE": "DEV", + "PYTHONPATH": str(Path(__file__).resolve().parents[1]), + } + env.pop("ENGRAM_DOTENV_CANARY", None) + process = await asyncio.create_subprocess_exec( + sys.executable, "-c", script, str(path), str(context.id), + stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + cwd=tmp_path, env=env, + ) + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=30) + assert process.returncode == 1, stderr.decode() + result = json.loads(stdout) + assert result["failed"] == 1 and result["errors"][0]["reason"] == "provider_error" + assert b"fake-provider-key-canary" not in stdout + stderr + assert b"Provider List" not in stdout + stderr diff --git a/tests/test_existing_core_correctness.py b/tests/test_existing_core_correctness.py new file mode 100644 index 0000000..a205fb7 --- /dev/null +++ b/tests/test_existing_core_correctness.py @@ -0,0 +1,417 @@ +"""Existing mutation paths preserve context boundaries and historical state.""" + +import pytest +from test_storage_transactions import storage as transaction_storage + +from engram.core.consolidation import ConsolidationEngine +from engram.core.delta import DeltaEngine +from engram.core.models import ( + Bullet, + ConsolidationConfig, + Context, + DeltaBatch, + DeltaOperation, + DeltaOpType, + IntentAnchor, + SchemaNode, +) + +storage = transaction_storage +VECTOR = [1.0] + [0.0] * 1535 + + +@pytest.mark.parametrize( + "kind", ["update_bullet", "remove_bullet", "merge_bullets", "update_schema"] +) +async def test_foreign_delta_target_rejects_whole_batch(storage, kind): + backend, context_id = storage + foreign = await backend.create_context( + Context(name="foreign", intent=IntentAnchor(objective="private")) + ) + await backend.add_bullet(str(foreign.id), Bullet(id="foreign", content="private")) + await backend.add_schema( + str(foreign.id), SchemaNode(id="foreign-schema", name="private", description="private") + ) + await backend.add_bullet(context_id, Bullet(id="local", content="local")) + op = DeltaOperation(op_type=DeltaOpType(kind), target_id="foreign", content="changed") + if kind == "merge_bullets": + op.target_ids = ["local", "foreign"] + elif kind == "update_schema": + op.target_id = "foreign-schema" + batch = DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation( + op_type=DeltaOpType.ADD_BULLET, target_id="before-invalid", content="must roll back" + ), + op, + ], + ) + with pytest.raises(ValueError, match="context"): + await DeltaEngine(backend).apply_batch(batch) + assert await backend.get_bullet("before-invalid") is None + assert (await backend.get_bullet("foreign")).content == "private" + assert (await backend.get_bullet("foreign")).is_active + assert (await backend.get_schema("foreign-schema")).description == "private" + assert await backend.get_delta_batch(batch.id) is None + + +async def test_remove_rollback_preserves_already_inactive_bullet(storage): + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="inactive", content="old", is_active=False)) + engine = DeltaEngine(backend) + batch = await engine.apply_batch( + DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.REMOVE_BULLET, target_id="inactive"), + ], + ) + ) + assert await engine.rollback_batch(batch.id) + assert not (await backend.get_bullet("inactive")).is_active + + +async def test_merge_rollback_restores_content_vectors_and_usage(storage): + backend, context_id = storage + before = [ + Bullet( + id="first", content="first", embedding=VECTOR, salience=0.8, recall_count=4, hit_count=2 + ), + Bullet( + id="second", + content="second", + embedding=VECTOR, + salience=0.4, + recall_count=3, + miss_count=2, + ), + ] + for bullet in before: + await backend.add_bullet(context_id, bullet) + before = [await backend.get_bullet(bullet.id) for bullet in before] + engine = DeltaEngine(backend) + batch = await engine.apply_batch( + DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation( + op_type=DeltaOpType.MERGE_BULLETS, + target_ids=["first", "second"], + content="combined", + ), + ], + ) + ) + assert len(await backend.list_bullets(context_id)) == 1 + assert await engine.rollback_batch(batch.id) + for expected in before: + loaded = await backend.get_bullet(expected.id) + assert loaded.is_active + for field in ( + "content", + "embedding", + "salience", + "recall_count", + "hit_count", + "miss_count", + ): + assert getattr(loaded, field) == getattr(expected, field) + + +async def test_duplicate_merge_ids_do_not_double_usage(storage): + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="one", content="one", recall_count=2)) + await backend.add_bullet(context_id, Bullet(id="two", content="two", recall_count=3)) + await DeltaEngine(backend).apply_batch( + DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.MERGE_BULLETS, target_ids=["one", "one", "two"]), + ], + ) + ) + assert sum(b.recall_count for b in await backend.list_bullets(context_id)) == 5 + + +async def test_legacy_rollback_cannot_mutate_foreign_context(storage): + backend, context_id = storage + foreign = await backend.create_context( + Context(name="foreign", intent=IntentAnchor(objective="private")) + ) + await backend.add_bullet(str(foreign.id), Bullet(id="foreign", content="private")) + batch = DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation( + op_type=DeltaOpType.UPDATE_BULLET, + target_id="foreign", + previous_state={"content": "injected"}, + ), + ], + ) + await backend.save_delta_batch(batch) + with pytest.raises(ValueError, match="context"): + await DeltaEngine(backend).rollback_batch(batch.id) + assert (await backend.get_bullet("foreign")).content == "private" + + +async def test_successful_apply_preserves_caller_batch_and_operation_identity(storage): + backend, context_id = storage + op = DeltaOperation(op_type=DeltaOpType.ADD_BULLET, content="legacy caller") + batch = DeltaBatch(context_id=context_id, operations=[op]) + result = await DeltaEngine(backend).apply_batch(batch) + assert result is batch + assert batch.operations[0] is op + assert op.target_id is not None + assert batch.bullets_added == 1 + assert (await backend.get_bullet(op.target_id)).content == "legacy caller" + + +async def test_dedup_failure_does_not_duplicate_statistics(storage, monkeypatch): + backend, context_id = storage + await backend.add_bullet( + context_id, + Bullet(id="first", content="first", embedding=VECTOR, salience=0.8, recall_count=4), + ) + await backend.add_bullet( + context_id, + Bullet(id="second", content="second", embedding=VECTOR, salience=0.4, recall_count=3), + ) + + async def fail(self, bullet_id): + raise RuntimeError("donor removal unavailable") + + monkeypatch.setattr(type(backend), "remove_bullet", fail) + with pytest.raises(RuntimeError, match="donor removal"): + await ConsolidationEngine(backend, None)._semantic_dedup(context_id, ConsolidationConfig()) + assert [(b.id, b.recall_count) for b in await backend.list_bullets(context_id)] == [ + ("first", 4), + ("second", 3), + ] + assert await backend.list_delta_batches(context_id) == [] + + +async def test_dedup_persists_reversible_audit(storage): + backend, context_id = storage + for i, salience in enumerate((0.8, 0.6, 0.4)): + await backend.add_bullet( + context_id, + Bullet( + id=str(i), content=str(i), embedding=VECTOR, salience=salience, recall_count=i + 1 + ), + ) + count = await ConsolidationEngine(backend, None)._semantic_dedup( + context_id, ConsolidationConfig() + ) + assert count == 2 + assert [(b.id, b.recall_count) for b in await backend.list_bullets(context_id)] == [("0", 6)] + batches = await backend.list_delta_batches(context_id) + assert len(batches) == 1 + assert await DeltaEngine(backend).rollback_batch(batches[0].id) + assert [(b.id, b.recall_count) for b in await backend.list_bullets(context_id)] == [ + ("0", 1), + ("1", 2), + ("2", 3), + ] + + +async def test_existing_archive_lifecycle_survives_old_schema_upgrade(storage): + backend, context_id = storage + from engram.storage.sqlite import SQLiteBackend + + await backend.add_bullet(context_id, Bullet(id="archive", content="historical")) + assert await backend.archive_bullet(context_id, "archive") + # Simulate a populated pre-lifecycle database, retaining its old archive flag. + if isinstance(backend, SQLiteBackend): + db = await backend._get_db() + await db.execute("DROP INDEX idx_bullets_lifecycle") + await db.execute("ALTER TABLE bullets DROP COLUMN lifecycle_state") + await db.execute("ALTER TABLE bullets DROP COLUMN archive_reason") + await db.commit() + else: + pool = await backend._get_pool() + await pool.execute("ALTER TABLE bullets DROP COLUMN lifecycle_state") + await pool.execute("ALTER TABLE bullets DROP COLUMN archive_reason") + await backend.initialize() + archived = await backend.get_archived_bullets(context_id) + assert [b.id for b in archived] == ["archive"] + restored = await backend.restore_bullet(context_id, "archive") + assert restored is not None + assert restored.content == "historical" + assert restored.is_active and not restored.is_archived + # Idempotent initialization must not reclassify an existing lifecycle. + await backend.initialize() + assert (await backend.get_bullet("archive")).lifecycle_state.value == "active" + + +def _offline_ingestion(backend, reflect=None): + from unittest.mock import AsyncMock + + from engram.core.ingestion import IngestionEngine + from engram.core.models import Reflection + + engine = IngestionEngine(backend, AsyncMock()) + engine.llm.embed.return_value = VECTOR + engine.reflector.reflect = AsyncMock(side_effect=reflect, return_value=Reflection()) + + async def curate(context_id, **kwargs): + return DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, content="extracted"), + ], + ) + + engine.curator.curate = AsyncMock(side_effect=curate) + return engine + + +async def test_failed_activity_write_does_not_orphan_ingestion_delta(storage, monkeypatch): + import uuid + + backend, context_id = storage + engine = _offline_ingestion(backend) + original = type(backend).add_activity + + async def fail(self, *args): + raise RuntimeError("activity ledger unavailable") + + monkeypatch.setattr(type(backend), "add_activity", fail) + with pytest.raises(RuntimeError, match="activity ledger"): + await engine.commit(uuid.UUID(context_id), "agent", "raw input") + assert await backend.list_bullets(context_id) == [] + assert await backend.list_delta_batches(context_id) == [] + monkeypatch.setattr(type(backend), "add_activity", original) + await engine.commit(uuid.UUID(context_id), "agent", "raw input") + assert len(await backend.list_bullets(context_id)) == 1 + + +async def test_concurrent_identical_commits_share_one_delta_and_activity(storage): + import asyncio + import uuid + + from engram.core.models import Reflection + + backend, context_id = storage + arrived = 0 + both_reflecting = asyncio.Event() + + async def reflect(**kwargs): + nonlocal arrived + arrived += 1 + if arrived == 2: + both_reflecting.set() + await asyncio.wait_for(both_reflecting.wait(), 5) + return Reflection() + + from unittest.mock import Mock + + engines = [_offline_ingestion(backend, reflect) for _ in range(2)] + events = Mock() + for engine in engines: + engine.event_bus = events + batches = await asyncio.gather( + *[ + engine.commit(uuid.UUID(context_id), "agent", "identical raw input") + for engine in engines + ] + ) + events.emit.assert_called_once() + assert batches[0].id == batches[1].id + assert len(await backend.list_bullets(context_id)) == 1 + assert len(await backend.get_activities_with_raw_input(context_id)) == 1 + + +@pytest.mark.parametrize("pathological", [float("nan"), 1e308]) +async def test_consolidation_does_not_merge_nonfinite_historical_similarity( + storage, monkeypatch, pathological +): + backend, context_id = storage + for bullet_id in ("first", "second"): + await backend.add_bullet( + context_id, Bullet(id=bullet_id, content=bullet_id, embedding=VECTOR) + ) + original = type(backend).list_bullets + + async def historical_vectors(self, *args, **kwargs): + bullets = await original(self, *args, **kwargs) + for bullet in bullets: + # Historical/custom storage can contain values PostgreSQL now + # rejects; an invalid similarity must never authorize deletion. + bullet.embedding = [pathological] + [0.0] * 1535 + return bullets + + monkeypatch.setattr(type(backend), "list_bullets", historical_vectors) + original_get = type(backend).get_bullet + + async def historical_vector(self, bullet_id): + bullet = await original_get(self, bullet_id) + if bullet is not None: + bullet.embedding = [pathological] + [0.0] * 1535 + return bullet + + monkeypatch.setattr(type(backend), "get_bullet", historical_vector) + assert ( + await ConsolidationEngine(backend, None)._semantic_dedup(context_id, ConsolidationConfig()) + == 0 + ) + assert (await backend.get_bullet("first")).is_active + assert (await backend.get_bullet("second")).is_active + assert await backend.list_delta_batches(context_id) == [] + + +@pytest.mark.parametrize("missing_remove", [False, True]) +async def test_commit_rejects_actual_capacity_overflow_atomically(storage, missing_remove): + import uuid + from unittest.mock import AsyncMock + + from engram.core.exceptions import CapacityExceededError + + backend, context_id = storage + context = await backend.get_context(uuid.UUID(context_id)) + context.lifecycle_config.max_active_bullets = 1 + await backend.update_context(context) + engine = _offline_ingestion(backend) + operations = [ + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id=name, content=name) + for name in ("one", "two") + ] + if missing_remove: + operations.append(DeltaOperation(op_type=DeltaOpType.REMOVE_BULLET, target_id="absent")) + engine.curator.curate = AsyncMock( + return_value=DeltaBatch(context_id=context_id, operations=operations) + ) + with pytest.raises(CapacityExceededError): + await engine.commit(uuid.UUID(context_id), "agent", "over capacity") + assert await backend.list_bullets(context_id) == [] + assert await backend.list_delta_batches(context_id) == [] + assert await backend.get_activities_with_raw_input(context_id) == [] + + +@pytest.mark.parametrize("mode", ["replace-full", "reduce-overfull", "edit-overfull"]) +async def test_commit_accepts_capacity_preserving_or_reducing_changes(storage, mode): + import uuid + from unittest.mock import AsyncMock + + backend, context_id = storage + context = await backend.get_context(uuid.UUID(context_id)) + context.lifecycle_config.max_active_bullets = 1 + await backend.update_context(context) + await backend.add_bullet(context_id, Bullet(id="one", content="one")) + if mode != "replace-full": + await backend.add_bullet(context_id, Bullet(id="two", content="two")) + operations = [DeltaOperation(op_type=DeltaOpType.REMOVE_BULLET, target_id="one")] + if mode == "replace-full": + operations.append(DeltaOperation(op_type=DeltaOpType.ADD_BULLET, content="replacement")) + elif mode == "edit-overfull": + operations = [ + DeltaOperation(op_type=DeltaOpType.UPDATE_BULLET, target_id="one", content="corrected") + ] + engine = _offline_ingestion(backend) + engine.curator.curate = AsyncMock( + return_value=DeltaBatch(context_id=context_id, operations=operations) + ) + await engine.commit(uuid.UUID(context_id), "agent", "capacity preserving change") + assert len(await backend.list_bullets(context_id)) == (2 if mode == "edit-overfull" else 1) + assert len(await backend.get_activities_with_raw_input(context_id)) == 1 diff --git a/tests/test_ingestion_config_compatibility.py b/tests/test_ingestion_config_compatibility.py new file mode 100644 index 0000000..d64e027 --- /dev/null +++ b/tests/test_ingestion_config_compatibility.py @@ -0,0 +1,88 @@ +"""Configured ingestion behavior reaches providers without breaking old adapters.""" + +from unittest.mock import AsyncMock + +import pytest +from test_storage_transactions import storage as transaction_storage + +from engram.core.config import IngestionConfig +from engram.core.ingestion import CuratorEngine, ReflectorEngine +from engram.core.models import Bullet, DeltaOpType, Reflection, ReflectionInsight + +storage = transaction_storage + + +@pytest.mark.parametrize("override", [None, "explicit-model"]) +async def test_reflector_routes_configured_and_explicit_models(override): + llm = AsyncMock() + llm.complete.return_value = "{}" + config = IngestionConfig(reflector_model="configured-model") + await ReflectorEngine(llm, config).reflect("raw input", model_override=override) + assert llm.complete.await_args.kwargs["model"] == (override or "configured-model") + + +async def test_legacy_adapter_completion_signature_still_works(): + class LegacyAdapter: + async def complete(self, prompt, system=None, temperature=0, response_format=None): + return '{"confidence": 0.75}' + + result = await ReflectorEngine(LegacyAdapter()).reflect("legacy input") + assert result.confidence == 0.75 + + +@pytest.mark.parametrize( + "threshold,expected", [(0.85, DeltaOpType.UPDATE_BULLET), (0.95, DeltaOpType.ADD_BULLET)] +) +async def test_configured_curator_threshold_changes_actual_dedup(storage, threshold, expected): + backend, context_id = storage + await backend.add_bullet( + context_id, Bullet(id="stored", content="stored knowledge", embedding=[1.0] + [0.0] * 1535) + ) + llm = AsyncMock() + llm.embed.return_value = [0.9, 0.19**0.5] + [0.0] * 1534 + config = IngestionConfig(curator_dedup_threshold=threshold) + if hasattr(config, "validity_gate_enabled"): + config.validity_gate_enabled = False + curator = CuratorEngine(backend, llm, ingestion_config=config) + result = await curator.curate( + context_id, + Reflection( + new_insights=[ + ReflectionInsight(content="a longer paraphrase of stored knowledge", novelty=0.9), + ] + ), + ) + assert len(result.operations) == 1 + assert result.operations[0].op_type == expected + + +async def test_per_call_model_does_not_change_adapter_default(monkeypatch): + import sys + from types import SimpleNamespace + + from engram.llm.adapter import LiteLLMAdapter, complete_with_model + + completion = AsyncMock( + return_value=SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="{}"))] + ) + ) + monkeypatch.setitem(sys.modules, "litellm", SimpleNamespace(acompletion=completion)) + adapter = LiteLLMAdapter(model="adapter-default") + await complete_with_model(adapter, model="one-call-only", prompt="first") + await adapter.complete(prompt="second") + assert [call.kwargs["model"] for call in completion.await_args_list] == [ + "one-call-only", + "adapter-default", + ] + assert adapter.model == "adapter-default" + + +async def test_provider_type_error_is_not_retried_as_legacy_adapter(): + from engram.llm.adapter import complete_with_model + + llm = AsyncMock() + llm.complete.side_effect = TypeError("provider response malformed") + with pytest.raises(TypeError, match="provider response malformed"): + await complete_with_model(llm, model="configured", prompt="input") + llm.complete.assert_awaited_once() diff --git a/tests/test_mcp_stdio.py b/tests/test_mcp_stdio.py new file mode 100644 index 0000000..3bd8212 --- /dev/null +++ b/tests/test_mcp_stdio.py @@ -0,0 +1,181 @@ +"""Exercise the real stdio entrypoint without a running Engram deployment.""" + +import asyncio +import json +import os +import sys +from contextlib import asynccontextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Thread +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +pytest.importorskip("mcp") + +from mcp.types import CallToolRequest, CallToolRequestParams + +from engram.integrations import mcp_server +from engram.sdk import client as sdk_client + + +@pytest.mark.asyncio +async def test_stdio_subprocess_initializes_and_forwards_configured_url_and_key(): + requests = [] + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + requests.append((self.path, self.headers.get("Authorization"))) + body = json.dumps([{"id": "test-context", "name": "Local fixture"}]).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + http_server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = Thread(target=http_server.serve_forever, daemon=True) + thread.start() + env = { + **os.environ, + "ENGRAM_API_URL": f"http://127.0.0.1:{http_server.server_port}/fixture", + "ENGRAM_API_KEY": "eng_test_stdio_only", + } + process = await asyncio.create_subprocess_exec( + sys.executable, "-m", "engram.integrations.mcp_server", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=env, + ) + + async def send(message): + process.stdin.write((json.dumps(message) + "\n").encode()) + await process.stdin.drain() + + async def response(request_id): + while True: + line = await asyncio.wait_for(process.stdout.readline(), timeout=20) + assert line, "MCP server exited without a protocol response" + message = json.loads(line) + if message.get("id") == request_id: + assert "error" not in message, message + return message["result"] + + try: + await send({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": { + "protocolVersion": "2024-11-05", "capabilities": {}, + "clientInfo": {"name": "engram-regression", "version": "1.0"}, + }, + }) + initialized = await response(1) + assert initialized["serverInfo"]["name"] == "engram" + await send({"jsonrpc": "2.0", "method": "notifications/initialized"}) + await send({"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}) + tools = await response(2) + assert {tool["name"] for tool in tools["tools"]} == { + "recall_context", "save_context", "save_bullet", "record_decision", + "list_contexts", "get_health", "archive_bullet", "get_lifecycle", + "re_extract_context", "get_ingestion_config", "engram_list_contexts", + "engram_recall", "engram_commit", "engram_decide", "engram_health", + "engram_create_context", + } + await send({ + "jsonrpc": "2.0", "id": 3, "method": "tools/call", + "params": {"name": "engram_list_contexts", "arguments": {}}, + }) + result = await response(3) + assert "Local fixture" in result["content"][0]["text"] + assert requests == [("/fixture/contexts", "Bearer eng_test_stdio_only")] + process.stdin.close() + assert await asyncio.wait_for(process.wait(), timeout=10) == 0 + finally: + if process.returncode is None: + process.kill() + await process.wait() + await asyncio.to_thread(http_server.shutdown) + http_server.server_close() + thread.join(timeout=2) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail", [False, True]) +async def test_stdio_closes_owned_sdk_client_on_shutdown(monkeypatch, fail): + client = SimpleNamespace(close=AsyncMock()) + created = [] + + def make_client(**kwargs): + created.append(kwargs) + return client + + server = SimpleNamespace( + run=AsyncMock(side_effect=RuntimeError("transport failed") if fail else None), + create_initialization_options=lambda: "options", + ) + injected = [] + + def make_server(*args, **kwargs): + injected.append(kwargs.get("client")) + return server + + @asynccontextmanager + async def transport(): + yield "read", "write" + + monkeypatch.setenv("ENGRAM_API_URL", "http://127.0.0.1:12345") + monkeypatch.setenv("ENGRAM_API_KEY", "eng_test_lifecycle_only") + monkeypatch.setattr(sdk_client, "Engram", make_client) + monkeypatch.setattr(mcp_server, "create_mcp_server", make_server) + monkeypatch.setattr(mcp_server, "stdio_server", transport) + if fail: + with pytest.raises(RuntimeError, match="transport failed"): + await mcp_server.run_mcp_server() + else: + await mcp_server.run_mcp_server() + assert created == [{"url": "http://127.0.0.1:12345", "api_key": "eng_test_lifecycle_only"}] + assert injected == [client] + client.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_injected_sdk_client_is_used_and_remains_caller_owned(monkeypatch): + client = SimpleNamespace( + list_contexts=AsyncMock(return_value=[{"id": "injected", "name": "Injected"}]), + close=AsyncMock(), + ) + + def unexpected_client(**kwargs): + pytest.fail("An injected client must not create another SDK client") + + monkeypatch.setattr(sdk_client, "Engram", unexpected_client) + server = mcp_server.create_mcp_server(client=client) + async with server.lifespan(server): + result = await server.request_handlers[CallToolRequest](CallToolRequest( + method="tools/call", + params=CallToolRequestParams(name="engram_list_contexts", arguments={}), + )) + assert "Injected" in result.root.content[0].text + client.list_contexts.assert_awaited_once_with(owner=None) + client.close.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tool_name", ["record_decision", "engram_decide"]) +async def test_decision_tools_display_id_returned_by_api(tool_name): + client = SimpleNamespace(record_decision=AsyncMock(return_value={ + "concept_id": "returned-decision-id", "status": "recorded", + })) + server = mcp_server.create_mcp_server(client=client) + result = await server.request_handlers[CallToolRequest](CallToolRequest( + method="tools/call", + params=CallToolRequestParams(name=tool_name, arguments={ + "context_id": "context", "decision": "Use SQLite", "rationale": "Local prototype", + }), + )) + assert "returned-decision-id" in result.root.content[0].text diff --git a/tests/test_pgvector_compatibility.py b/tests/test_pgvector_compatibility.py new file mode 100644 index 0000000..b51a73a --- /dev/null +++ b/tests/test_pgvector_compatibility.py @@ -0,0 +1,12 @@ +"""Read compatibility for pgvector codec versions without dependency pinning.""" + +import numpy as np +import pytest +from pgvector import Vector + +from engram.storage.postgres import _vector_to_list + + +@pytest.mark.parametrize("value", [None, [0.5, 1.0], np.array([0.5, 1.0]), Vector([0.5, 1.0])]) +def test_vector_conversion(value): + assert _vector_to_list(value) == (None if value is None else [0.5, 1.0]) diff --git a/tests/test_postgres_pool_lifecycle.py b/tests/test_postgres_pool_lifecycle.py new file mode 100644 index 0000000..9ba1992 --- /dev/null +++ b/tests/test_postgres_pool_lifecycle.py @@ -0,0 +1,75 @@ +"""A partially initialized asyncpg pool must never strand server sessions.""" + +import asyncio +import os +import uuid +from urllib.parse import urlsplit, urlunsplit + +import asyncpg +import pytest + +from engram.storage.postgres import PostgresBackend + + +@pytest.mark.parametrize( + "failure", ["exception", "cancellation"], ids=["postgres-exception", "postgres-cancellation"] +) +async def test_partial_pool_initialization_releases_connections_and_retries(monkeypatch, failure): + admin_dsn = os.environ.get("ENGRAM_TEST_POSTGRES_DSN") + if not admin_dsn: + pytest.skip("requires a disposable PostgreSQL admin DSN") + admin = await asyncpg.connect(admin_dsn) + database = "engram_pool_" + uuid.uuid4().hex + dsn = urlunsplit(urlsplit(admin_dsn)._replace(path="/" + database)) + backend = PostgresBackend(dsn) + original = backend._init_connection + second_started = asyncio.Event() + first_pid = None + calls = 0 + + async def initialize(conn): + nonlocal calls, first_pid + calls += 1 + await original(conn) + if calls == 1: + first_pid = conn.get_server_pid() + else: + second_started.set() + if failure == "exception": + raise RuntimeError("second connection failed") + await asyncio.Event().wait() + + task = None + try: + await admin.execute(f'CREATE DATABASE "{database}"') + monkeypatch.setattr(backend, "_init_connection", initialize) + task = asyncio.create_task(backend._get_pool()) + await asyncio.wait_for(second_started.wait(), timeout=10) + if failure == "cancellation": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + with pytest.raises(RuntimeError, match="second connection failed"): + await task + assert backend._pool is None + assert first_pid is not None + for _ in range(50): + remaining = await admin.fetchval( + "SELECT count(*) FROM pg_stat_activity WHERE datname=$1", database + ) + if remaining == 0: + break + await asyncio.sleep(0.02) + assert remaining == 0, "partially initialized pool leaked database sessions" + monkeypatch.setattr(backend, "_init_connection", original) + assert await backend._get_pool() is backend._pool + await backend.close() + assert backend._pool is None + finally: + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await backend.close() + await admin.execute(f'DROP DATABASE IF EXISTS "{database}" WITH (FORCE)') + await admin.close() diff --git a/tests/test_postgres_startup.py b/tests/test_postgres_startup.py new file mode 100644 index 0000000..09d3bed --- /dev/null +++ b/tests/test_postgres_startup.py @@ -0,0 +1,136 @@ +"""Fresh-database startup regressions for GitHub issue #3. + +Set ENGRAM_TEST_POSTGRES_DSN to a disposable PostgreSQL admin database to +also run the integration test. It creates and removes its own temporary DB. +""" + +import os +import uuid +from unittest.mock import AsyncMock, Mock +from urllib.parse import urlsplit, urlunsplit + +import asyncpg +import pytest + +from engram.core.models import Bullet, Context, IntentAnchor +from engram.storage.postgres import PostgresBackend + + +class AwaitablePool: + """Match asyncpg's synchronous factory returning an awaitable Pool.""" + + def __init__(self): + self.initialize = None + self.close = AsyncMock() + self.terminate = Mock() + + def __await__(self): + async def initialize(): + if self.initialize is not None: + await self.initialize() + return self + + return initialize().__await__() + + +async def test_extension_exists_before_pool_registers_codecs(monkeypatch): + events = [] + connection = AsyncMock() + connection.execute.side_effect = lambda sql: events.append(sql) + connection.close.side_effect = lambda: events.append("closed") + monkeypatch.setattr(asyncpg, "connect", AsyncMock(return_value=connection)) + pool = AwaitablePool() + + async def register(conn): + assert events == ["CREATE EXTENSION IF NOT EXISTS vector", "closed"] + events.append("registered") + + def create_pool(*args, **kwargs): + pool.initialize = lambda: kwargs["init"](AsyncMock()) + return pool + + monkeypatch.setattr("engram.storage.postgres.register_vector", register) + create = Mock(side_effect=create_pool) + monkeypatch.setattr(asyncpg, "create_pool", create) + backend = PostgresBackend("postgresql://test") + assert await backend._get_pool() is pool + assert await backend._get_pool() is pool + create.assert_called_once() + connection.close.assert_awaited_once() + await backend.close() + + +async def test_failed_extension_setup_closes_connection_and_can_retry(monkeypatch): + connection = AsyncMock() + connection.execute.side_effect = asyncpg.InsufficientPrivilegeError("install vector first") + monkeypatch.setattr(asyncpg, "connect", AsyncMock(return_value=connection)) + create = Mock(return_value=AwaitablePool()) + monkeypatch.setattr(asyncpg, "create_pool", create) + backend = PostgresBackend("postgresql://test") + with pytest.raises(asyncpg.InsufficientPrivilegeError, match="install vector first"): + await backend.initialize() + connection.close.assert_awaited_once() + create.assert_not_called() + assert backend._pool is None + + connection.execute.side_effect = None + await backend._get_pool() + create.assert_called_once() + assert connection.close.await_count == 2 + await backend.close() + + +async def test_failed_pool_creation_is_not_cached(monkeypatch): + connection = AsyncMock() + monkeypatch.setattr(asyncpg, "connect", AsyncMock(return_value=connection)) + create = Mock(side_effect=RuntimeError("pool failed")) + monkeypatch.setattr(asyncpg, "create_pool", create) + backend = PostgresBackend("postgresql://test") + with pytest.raises(RuntimeError, match="pool failed"): + await backend._get_pool() + connection.close.assert_awaited_once() + assert backend._pool is None + + +@pytest.mark.skipif( + not os.environ.get("ENGRAM_TEST_POSTGRES_DSN"), + reason="requires a disposable PostgreSQL admin DSN", +) +async def test_fresh_postgres_vector_roundtrip_and_restart(): + admin_dsn = os.environ["ENGRAM_TEST_POSTGRES_DSN"] + admin = await asyncpg.connect(admin_dsn) + db_name = "engram_test_" + uuid.uuid4().hex + role_name = "engram_test_" + uuid.uuid4().hex + parts = urlsplit(admin_dsn) + dsn = urlunsplit(parts._replace(path="/" + db_name)) + backend = PostgresBackend(dsn) + try: + await admin.execute(f'CREATE DATABASE "{db_name}"') + await backend.initialize() + ctx = await backend.create_context( + Context(name="Startup test", intent=IntentAnchor(objective="Vector roundtrip")) + ) + embedding = [1.0] + [0.0] * 1535 + bullet = Bullet(content="Persist across restart", embedding=embedding) + await backend.add_bullet(str(ctx.id), bullet) + await backend.close() + + await backend.initialize() + stored = await backend.get_bullet(bullet.id) + assert stored.embedding == pytest.approx(embedding) + matches = await backend.find_similar_bullets(str(ctx.id), embedding, threshold=0.99) + assert matches[0][0].id == bullet.id + + # Existing extension does not require the app role to be a superuser. + await admin.execute(f'CREATE ROLE "{role_name}"') + async with backend._pool.acquire() as conn: + await conn.execute(f'SET ROLE "{role_name}"') + try: + await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + finally: + await conn.execute("RESET ROLE") + finally: + await backend.close() + await admin.execute(f'DROP DATABASE IF EXISTS "{db_name}" WITH (FORCE)') + await admin.execute(f'DROP ROLE IF EXISTS "{role_name}"') + await admin.close() diff --git a/tests/test_public_mutation_atomicity.py b/tests/test_public_mutation_atomicity.py new file mode 100644 index 0000000..cb313fb --- /dev/null +++ b/tests/test_public_mutation_atomicity.py @@ -0,0 +1,226 @@ +"""All public ingestion writes share atomic metadata and capacity boundaries.""" + +import asyncio +import uuid +from unittest.mock import AsyncMock + +import pytest +from test_storage_transactions import storage as transaction_storage + +from engram.core.delta import DeltaEngine +from engram.core.exceptions import CapacityExceededError +from engram.core.ingestion import CuratorEngine, IngestionEngine +from engram.core.models import ( + ActionType, + Activity, + Bullet, + ConceptNode, + ConceptType, + DeltaBatch, + DeltaOperation, + DeltaOpType, + ReExtractionRequest, + Reflection, + ReflectionInsight, +) +from engram.core.re_extraction import ReExtractionEngine + +storage = transaction_storage +VECTOR = [1.0] + [0.0] * 1535 + + +async def _engine(backend, context_id): + llm = AsyncMock() + + async def embed(text): + # Parent-handle reads are rejected inside the transaction, making this + # assert that providers are called before acquiring the scoped handle. + await backend.list_bullets(context_id) + return VECTOR + + llm.embed.side_effect = embed + return IngestionEngine(backend, llm) + + +async def _reextract(backend, context_id, count=1): + await backend.add_activity( + uuid.UUID(context_id), + Activity( + agent_id="old", + action_type=ActionType.FACT_LEARNED, + summary="original", + raw_input="original raw", + ), + ) + reflector = AsyncMock() + reflector.reflect.return_value = Reflection() + curator = AsyncMock() + curator.curate_re_extraction.return_value = DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, content=f"new {i}") for i in range(count) + ], + ) + return ReExtractionEngine(reflector, curator, backend) + + +@pytest.mark.parametrize("entrypoint", ["decision", "reextract"]) +async def test_activity_failure_rolls_back_every_public_entrypoint( + storage, monkeypatch, entrypoint +): + backend, context_id = storage + engine = ( + await _engine(backend, context_id) + if entrypoint == "decision" + else await _reextract(backend, context_id) + ) + + async def fail(self, *args): + raise RuntimeError("activity unavailable") + + monkeypatch.setattr(type(backend), "add_activity", fail) + with pytest.raises(RuntimeError, match="activity unavailable"): + if entrypoint == "decision": + await engine.record_decision( + uuid.UUID(context_id), "decision", "reason", ["alternative"], "agent" + ) + else: + await engine.re_extract( + context_id, ReExtractionRequest(reflector_model="test", dry_run=False) + ) + assert await backend.list_bullets(context_id) == [] + assert await backend.list_delta_batches(context_id) == [] + activities = await backend.get_activities_with_raw_input(context_id) + assert len(activities) == (1 if entrypoint == "reextract" else 0) + + +async def test_direct_add_salience_failure_rolls_back_bullet_and_audit(storage, monkeypatch): + backend, context_id = storage + engine = await _engine(backend, context_id) + + async def fail(self, bullet): + raise RuntimeError("salience unavailable") + + monkeypatch.setattr(type(backend), "update_bullet", fail) + with pytest.raises(RuntimeError, match="salience unavailable"): + await engine.add_bullet_directly(context_id, "fact", salience=0.9) + assert await backend.list_bullets(context_id) == [] + assert await backend.list_delta_batches(context_id) == [] + + +async def test_concurrent_direct_add_checks_capacity_after_embedding(storage): + backend, context_id = storage + context = await backend.get_context(uuid.UUID(context_id)) + context.lifecycle_config.max_active_bullets = 1 + await backend.update_context(context) + arrived = 0 + both = asyncio.Event() + + async def embed(text): + nonlocal arrived + arrived += 1 + if arrived == 2: + both.set() + await asyncio.wait_for(both.wait(), 5) + return VECTOR + + llm = AsyncMock() + llm.embed.side_effect = embed + engine = IngestionEngine(backend, llm) + results = await asyncio.gather( + *[engine.add_bullet_directly(context_id, text, salience=0.8) for text in ("one", "two")], + return_exceptions=True, + ) + assert sum(isinstance(result, CapacityExceededError) for result in results) == 1 + assert sum(isinstance(result, tuple) for result in results) == 1 + assert len(await backend.list_bullets(context_id)) == 1 + assert len(await backend.list_delta_batches(context_id)) == 1 + + +@pytest.mark.parametrize("entrypoint", ["decision", "reextract"]) +async def test_other_public_entrypoints_reject_capacity_overflow(storage, entrypoint): + backend, context_id = storage + context = await backend.get_context(uuid.UUID(context_id)) + context.lifecycle_config.max_active_bullets = 1 + await backend.update_context(context) + engine = ( + await _engine(backend, context_id) + if entrypoint == "decision" + else await _reextract(backend, context_id, count=2) + ) + with pytest.raises(CapacityExceededError): + if entrypoint == "decision": + await engine.record_decision( + uuid.UUID(context_id), "decision", "reason", ["alternative"], "agent" + ) + else: + await engine.re_extract( + context_id, ReExtractionRequest(reflector_model="test", dry_run=False) + ) + assert await backend.list_bullets(context_id) == [] + assert await backend.list_delta_batches(context_id) == [] + + +async def test_snapshotless_historical_merge_refuses_entire_rollback(storage): + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="keep", content="keep")) + batch = DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.MERGE_BULLETS, target_ids=["lost", "survivor"]), + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id="keep", content="keep"), + ], + ) + await backend.save_delta_batch(batch) + assert not await DeltaEngine(backend).rollback_batch(batch.id) + assert (await backend.get_bullet("keep")).is_active + + +@pytest.mark.parametrize("kind", ["bullet", "concept"]) +async def test_zero_vector_does_not_hide_valid_similarity_at_limit_one(storage, kind): + backend, context_id = storage + if kind == "bullet": + await backend.add_bullet( + context_id, Bullet(id="invalid", content="invalid", embedding=[0.0] * 1536) + ) + await backend.add_bullet(context_id, Bullet(id="valid", content="valid", embedding=VECTOR)) + matches = await backend.find_similar_bullets(context_id, VECTOR, limit=1, threshold=0.9) + else: + await backend.add_concept( + uuid.UUID(context_id), + ConceptNode(type=ConceptType.FACT, content="invalid", embedding=[0.0] * 1536), + ) + await backend.add_concept( + uuid.UUID(context_id), + ConceptNode(type=ConceptType.FACT, content="valid", embedding=VECTOR), + ) + matches = await backend.find_similar_concepts( + uuid.UUID(context_id), VECTOR, limit=1, threshold=0.9 + ) + assert len(matches) == 1 + assert matches[0][0].content == "valid" + assert matches[0][1] == pytest.approx(1.0) + + +@pytest.mark.parametrize("score", [float("nan"), float("inf")]) +async def test_curator_rejects_nonfinite_custom_storage_scores(storage, monkeypatch, score): + backend, context_id = storage + existing = await backend.add_bullet( + context_id, Bullet(id="keep", content="unrelated old", embedding=VECTOR) + ) + + async def invalid(*args, **kwargs): + return [(existing, score)] + + monkeypatch.setattr(type(backend), "find_similar_bullets", invalid) + llm = AsyncMock() + llm.embed.return_value = VECTOR + curator = CuratorEngine(backend, llm) + op = await curator._process_insight( + context_id, + ReflectionInsight(content="different new fact", novelty=0.9), + [existing], + "agent", + None, + ) + assert op.op_type == DeltaOpType.ADD_BULLET diff --git a/tests/test_retrieval_evaluation.py b/tests/test_retrieval_evaluation.py new file mode 100644 index 0000000..623ee30 --- /dev/null +++ b/tests/test_retrieval_evaluation.py @@ -0,0 +1,371 @@ +"""Real-storage retrieval gates and exact renderer packing regressions.""" + +import copy +import inspect +import json +import math +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from engram.core.materialization import MaterializationEngine +from engram.core.models import ( + ActionType, + Activity, + Bullet, + ConceptNode, + ConceptType, + Context, + IntentAnchor, + SchemaNode, +) +from engram.storage.sqlite import SQLiteBackend +from evals.retrieval import ( + CORPUS, + FixtureEmbeddings, + assess_result, + evaluate, + ranking_metrics, + summarize, +) + + +def test_metric_math_and_duplicate_credit(): + relevance = {"best": 3, "next": 1} + assert ranking_metrics(["best", "next"], relevance, 2) == {"recall_at_k": 1, "ndcg_at_k": 1} + actual = ranking_metrics(["next", "best"], relevance, 2) + assert actual["ndcg_at_k"] == pytest.approx((1 + 7 / math.log2(3)) / (7 + 1 / math.log2(3))) + assert ranking_metrics(["best", "best"], relevance, 2)["recall_at_k"] == 0.5 + assert ranking_metrics([], relevance, 2) == {"recall_at_k": 0, "ndcg_at_k": 0} + assert ranking_metrics([], {}, 2) == {"recall_at_k": None, "ndcg_at_k": None} + with pytest.raises(ValueError): + ranking_metrics([], relevance, 0) + + +def test_gate_cannot_pass_without_judgments_or_missing_results(): + assert not summarize([], 0)["passed"] + empty = dict(passed=True, recall_at_k=None, ndcg_at_k=None, forbidden_id_leakage=[]) + assert not summarize([empty], 1)["passed"] + judged = dict(empty, recall_at_k=1, ndcg_at_k=1) + assert not summarize([judged], 2)["passed"] + + +@pytest.mark.asyncio +async def test_controlled_corpus_gate(): + report = await evaluate(json.loads(CORPUS.read_text())) + assert report["benchmark"] == "controlled_regression" + assert report["summary"]["scenario_renderer_count"] >= 24 + assert report["summary"]["judged_count"] > 0 + assert report["summary"]["passed"], [row for row in report["results"] if not row["passed"]] + + +@pytest.mark.asyncio +async def test_opt_in_provider_vectors_are_cached_and_validated(): + corpus = json.loads(CORPUS.read_text()) + corpus["cases"] = corpus["cases"][:1] + case = corpus["cases"][0] + vectors = {case["query"]: case["query_vector"]} + vectors.update({item["text"]: item["vector"] for item in case["items"]}) + + class RecordingProvider(FixtureEmbeddings): + calls = [] + + async def embed(self, text): + self.calls.append(text) + return await super().embed(text) + + provider = RecordingProvider(vectors) + report = await evaluate(corpus, provider) + assert report["benchmark"] == "live_provider_labeled_corpus" + assert report["summary"]["passed"] + assert sorted(provider.calls) == sorted(vectors) + vectors[case["query"]] = [float("nan"), 0, 0] + with pytest.raises(ValueError, match="invalid"): + await evaluate(corpus, provider) + + +@pytest.mark.parametrize("degradation", ["ranking", "leakage", "empty"]) +def test_cli_failure_gate(tmp_path, degradation): + corpus = json.loads(CORPUS.read_text()) + case = copy.deepcopy(corpus["cases"][0]) + if degradation == "ranking": + case["items"][0]["vector"], case["items"][1]["vector"] = ( + case["items"][1]["vector"], + case["items"][0]["vector"], + ) + elif degradation == "leakage": + # Deliberately put a forbidden fact in the eligible context. This + # simulates a storage isolation regression; the gate must catch it. + case["forbidden_ids"] = [case["items"][0]["id"]] + else: + case = next(row for row in corpus["cases"] if row.get("expect_empty")) + corpus["cases"] = [case] + corpus["renderers"] = ["generic"] + path = tmp_path / "degraded.json" + path.write_text(json.dumps(corpus)) + output = tmp_path / "result.json" + run = subprocess.run( + [sys.executable, "-m", "evals.retrieval", "--corpus", str(path), "--output", str(output)], + capture_output=True, + text=True, + cwd=Path(__file__).resolve().parents[1], + timeout=30, + ) + assert run.returncode == 1, run.stdout + run.stderr + report = json.loads(output.read_text()) + assert not report["summary"]["passed"] + if degradation != "empty": + failure = "recall_at_k" if degradation == "ranking" else "forbidden_id_leakage" + assert failure in report["results"][0]["failures"] + + +def test_live_cli_redacts_real_adapter_errors(tmp_path): + canary = "retrieval-fake-provider-key-canary" + # Execute the real LiteLLMAdapter against a fake SDK in a separate process. + # Import output, SDK output and the adapter's exception log must all stay + # out of the machine-readable command output. No external call is possible. + (tmp_path / "litellm.py").write_text( + "import sys\n" + "suppress_debug_info = False\n" + "set_verbose = True\n" + f"print({canary!r})\n" + "async def aembedding(**kwargs):\n" + " print(kwargs['api_key'])\n" + " print(kwargs['api_key'], file=sys.stderr)\n" + " raise RuntimeError(kwargs['api_key'])\n" + ) + output = tmp_path / "live-failure.json" + run = subprocess.run( + [ + sys.executable, + "-m", + "evals.retrieval", + "--live-embeddings", + "--embedding-model", + "fake-model", + "--api-key-env", + "EVAL_TEST_API_KEY", + "--output", + str(output), + ], + env={**os.environ, "PYTHONPATH": str(tmp_path), "EVAL_TEST_API_KEY": canary}, + capture_output=True, + text=True, + cwd=Path(__file__).resolve().parents[1], + timeout=30, + ) + assert run.returncode == 1 + assert canary not in run.stdout + run.stderr + report = json.loads(run.stdout) + assert report == json.loads(output.read_text()) + assert report["summary"]["passed"] is False + assert report["error_type"] == "LLMAdapterError" + + +def test_offline_cli_does_not_import_provider_sdk(tmp_path): + (tmp_path / "litellm.py").write_text("raise AssertionError('Provider SDK imported offline')\n") + run = subprocess.run( + [sys.executable, "-m", "evals.retrieval"], + env={**os.environ, "PYTHONPATH": str(tmp_path)}, + capture_output=True, + text=True, + cwd=Path(__file__).resolve().parents[1], + timeout=30, + ) + assert run.returncode == 0, run.stderr + assert json.loads(run.stdout)["summary"]["passed"] + + +@pytest.mark.parametrize("target", ["claude", "gpt", "generic"]) +@pytest.mark.parametrize("kind", ["bullets", "legacy", "schema"]) +@pytest.mark.asyncio +async def test_receipts_match_complete_rendered_content(tmp_path, target, kind): + storage = SQLiteBackend(str(tmp_path / "receipt.db")) + await storage.initialize() + try: + context = Context(name="packing", intent=IntentAnchor(objective="Pack memories")) + await storage.create_context(context) + engine = MaterializationEngine(storage, FixtureEmbeddings({"query": [1, 0]})) + texts = {} + for index in range(3): + content = f"Unique evidence {index}." + if kind == "legacy": + item = ConceptNode( + type=ConceptType.FACT, + content=content, + embedding=[1, 0], + salience=1 - index * 0.1, + ) + await storage.add_concept(context.id, item) + else: + item = Bullet(content=content, embedding=[1, 0], salience=1 - index * 0.1) + await storage.add_bullet(str(context.id), item) + texts[str(item.id)] = content + schemas = {} + if kind == "schema": + item = SchemaNode(name="Packing", description="Preserve all selected evidence.") + await storage.add_schema(str(context.id), item) + schemas[item.id] = f"[Pattern: {item.name}] {item.description}" + options = ( + {"include_worked_examples": False} + if "include_worked_examples" in (inspect.signature(engine.materialize).parameters) + else {} + ) + renderer = engine._get_renderer(target) + # Sweep boundaries around section wrappers and multiple candidate sizes. + for budget in [1, 10, 25, 30, 35, 40, 45, 55, 70, 100]: + result = await engine.materialize( + context.id, + query="query", + target_model=target, + token_budget=budget, + include_intent=False, + **options, + ) + included = [ + str(key) for key in result["bullets_included"] + result["concepts_included"] + ] + assert set(included) == { + key for key, text in texts.items() if text in result["rendered_text"] + } + assert set(result["schemas_included"]) == { + key for key, text in schemas.items() if text in result["rendered_text"] + } + receipt = await storage.get_materialization(result["materialization_id"]) + assert receipt.bullets_included == result["bullets_included"] + assert receipt.token_count == result["token_count"] <= budget + assert result["token_count"] == renderer.estimate_tokens(result["rendered_text"]) + finally: + await storage.close() + + +@pytest.mark.parametrize("target", ["claude", "gpt", "generic"]) +@pytest.mark.parametrize("kind", ["bullets", "legacy", "empty"]) +@pytest.mark.asyncio +async def test_oss_core_memory_worked_examples_fit_tight_budget(tmp_path, target, kind): + if "core_memory" not in Context.model_fields: + pytest.skip("OSS v0.5 extensions are intentionally edition-specific") + storage = SQLiteBackend(str(tmp_path / "extensions.db")) + await storage.initialize() + try: + context = Context( + name="extensions", + core_memory="Persistent core summary. " * 20, + intent=IntentAnchor(objective="Important objective. " * 20), + ) + await storage.create_context(context) + bullet = Bullet(content="Short actionable fact.", embedding=[1, 0]) + if kind == "bullets": + await storage.add_bullet(str(context.id), bullet) + elif kind == "legacy": + await storage.add_concept( + context.id, + ConceptNode(type=ConceptType.FACT, content=bullet.content, embedding=[1, 0]), + ) + activity = Activity( + agent_id="eval", + action_type=ActionType.MATERIALIZATION_OCCURRED, + summary="Prior useful input", + raw_input="Prior worked input. " * 40, + raw_input_embedding=[1, 0], + bullet_ids_produced=[bullet.id], + ) + await storage.add_activity(context.id, activity) + engine = MaterializationEngine(storage, FixtureEmbeddings({"query": [1, 0]})) + # Real retrieval, including worked-example fetching from stored activity. + for budget in [1, 30, 80, 200, 1000]: + result = await engine.materialize( + context.id, + query="query", + token_budget=budget, + target_model=target, + include_usage_stats=True, + ) + assert result["token_count"] <= budget + if bullet.id in result["bullets_included"]: + assert bullet.content in result["rendered_text"] + assert "Prior worked input." in result["rendered_text"] + assert "Persistent core summary." in result["rendered_text"] + finally: + await storage.close() + + +@pytest.mark.parametrize("target", ["claude", "gpt", "generic"]) +def test_complete_duplicate_candidates_and_oversized_extensions(target): + renderer = MaterializationEngine._get_renderer(target) + large = ConceptNode(type=ConceptType.DECISION, content="Large decision. " * 200) + small = ConceptNode(type=ConceptType.FACT, content="Small complete fact.") + duplicate = small.model_copy(update={"id": large.id}) + single_text = renderer.render([small], None, 1000) + budget = renderer.estimate_tokens(single_text) + options = {} + if "core_memory" in inspect.signature(renderer.render).parameters: + options = {"worked_examples": [{"input": "Oversized worked example. " * 200}]} + text, selected = renderer.render_with_selection( + [large, small, duplicate], None, budget, **options + ) + assert [item.id for item in selected] == [small.id] + assert small.content in text + assert "Oversized worked example." not in text + assert renderer.estimate_tokens(text) <= budget + + +@pytest.mark.parametrize("target", ["claude", "gpt", "generic"]) +@pytest.mark.asyncio +async def test_duplicate_content_keeps_selected_bullets_usage_stats(tmp_path, target): + if "include_usage_stats" not in inspect.signature(MaterializationEngine.materialize).parameters: + pytest.skip("Usage annotations are an OSS v0.5 extension") + storage = SQLiteBackend(str(tmp_path / "duplicate-usage.db")) + await storage.initialize() + try: + context = Context(name="duplicate usage", intent=IntentAnchor(objective="Keep provenance")) + await storage.create_context(context) + first = Bullet(content="Identical fact.", salience=1, recall_count=2, hit_count=2) + second = Bullet(content=first.content, salience=0.1, recall_count=99, hit_count=99) + for bullet in (first, second): + await storage.add_bullet(str(context.id), bullet) + renderer = MaterializationEngine._get_renderer(target) + concept = ConceptNode(type=ConceptType.FACT, content=first.content, confidence=0.5) + # Existing callers can still supply annotations keyed by content. + legacy_text = renderer.render( + [concept], None, 1000, usage_stats={first.content: "(used 99×, success 99/99)"} + ) + assert "(used 99×, success 99/99)" in legacy_text + budget = 19 if target == "generic" else renderer.estimate_tokens(legacy_text) + engine = MaterializationEngine(storage, FixtureEmbeddings({})) + result = await engine.materialize( + context.id, + target_model=target, + token_budget=budget, + include_intent=False, + include_worked_examples=False, + include_usage_stats=True, + recency_weight=0, + ) + assert result["bullets_included"] == [first.id] + assert "(used 2×, success 2/2)" in result["rendered_text"] + assert "99" not in result["rendered_text"] + assert result["token_count"] <= budget + receipt = await storage.get_materialization(result["materialization_id"]) + assert receipt.bullets_included == [first.id] + finally: + await storage.close() + + +def test_inconsistent_receipt_is_detected(): + case = json.loads(CORPUS.read_text())["cases"][0] + result = dict( + bullets_included=["real-id"], + concepts_included=[], + schemas_included=[], + rendered_text="", + token_count=0, + ) + renderer = MaterializationEngine._get_renderer("generic") + assessment = assess_result(case, result, None, renderer, {"retry": "real-id"}, {}) + assert not assessment["passed"] + assert "rendered_id_consistency" in assessment["failures"] + assert "persisted_receipt_consistency" in assessment["failures"] diff --git a/tests/test_schema_rollback.py b/tests/test_schema_rollback.py new file mode 100644 index 0000000..de96bcd --- /dev/null +++ b/tests/test_schema_rollback.py @@ -0,0 +1,151 @@ +"""Schema mutations are reversible without guessing historical state.""" + +import pytest +from test_storage_transactions import storage as transaction_storage + +from engram.core.delta import DeltaEngine +from engram.core.models import ( + Bullet, + Context, + DeltaBatch, + DeltaOperation, + DeltaOpType, + IntentAnchor, + SchemaNode, +) + +storage = transaction_storage + + +@pytest.mark.parametrize("explicit_id", [None, "chosen-schema"]) +async def test_added_schema_persists_target_and_rolls_back_without_dangling_links( + storage, explicit_id +): + backend, context_id = storage + engine = DeltaEngine(backend) + batch = await engine.apply_batch( + DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation( + op_type=DeltaOpType.ADD_SCHEMA, + target_id=explicit_id, + content="new schema", + reasoning="description", + ), + ], + ) + ) + schema_id = batch.operations[0].target_id + assert schema_id is not None + if explicit_id: + assert schema_id == explicit_id + assert (await backend.get_schema(schema_id)).description == "description" + assert (await backend.get_delta_batch(batch.id)).operations[0].target_id == schema_id + await backend.add_bullet( + context_id, Bullet(id="linked", content="retained memory", schema_id=schema_id) + ) + assert await engine.rollback_batch(batch.id) + assert await backend.get_schema(schema_id) is None + assert (await backend.get_bullet("linked")).schema_id is None + assert (await backend.get_bullet("linked")).content == "retained memory" + + +async def test_schema_update_restores_description_without_changing_other_fields(storage): + backend, context_id = storage + before = SchemaNode(id="existing", name="name", description="before", instance_count=3) + await backend.add_schema(context_id, before) + engine = DeltaEngine(backend) + batch = await engine.apply_batch( + DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation( + op_type=DeltaOpType.UPDATE_SCHEMA, target_id=before.id, content="after" + ), + ], + ) + ) + assert (await backend.get_schema(before.id)).description == "after" + assert await engine.rollback_batch(batch.id) + restored = await backend.get_schema(before.id) + assert (restored.description, restored.name, restored.instance_count) == ("before", "name", 3) + + +@pytest.mark.parametrize("kind", [DeltaOpType.ADD_SCHEMA, DeltaOpType.UPDATE_SCHEMA]) +async def test_legacy_schema_without_snapshot_declines_entire_rollback(storage, kind): + backend, context_id = storage + # Old ADD_SCHEMA ignored caller target_id and generated an unrecorded ID. + # This target may therefore identify an unrelated, pre-existing schema. + await backend.add_schema( + context_id, SchemaNode(id="unrelated", name="keep", description="keep") + ) + await backend.add_bullet(context_id, Bullet(id="keep-bullet", content="keep")) + batch = DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=kind, target_id="unrelated", content="legacy unknown state"), + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id="keep-bullet", content="keep"), + ], + ) + await backend.save_delta_batch(batch) + assert not await DeltaEngine(backend).rollback_batch(batch.id) + assert (await backend.get_bullet("keep-bullet")).is_active + assert (await backend.get_schema("unrelated")).description == "keep" + + +async def test_schema_rollback_rejects_foreign_target_atomically(storage): + backend, context_id = storage + foreign = await backend.create_context( + Context(name="foreign", intent=IntentAnchor(objective="private")) + ) + await backend.add_schema( + str(foreign.id), SchemaNode(id="foreign-schema", name="private", description="private") + ) + await backend.add_bullet(context_id, Bullet(id="keep", content="keep")) + op = DeltaOperation( + op_type=DeltaOpType.ADD_SCHEMA, + target_id="foreign-schema", + content="private", + previous_state={"schema_created": True}, + ) + batch = DeltaBatch( + context_id=context_id, + operations=[ + op, + DeltaOperation(op_type=DeltaOpType.ADD_BULLET, target_id="keep", content="keep"), + ], + ) + await backend.save_delta_batch(batch) + with pytest.raises(ValueError, match="context"): + await DeltaEngine(backend).rollback_batch(batch.id) + assert (await backend.get_bullet("keep")).is_active + assert await backend.get_schema("foreign-schema") is not None + + +async def test_schema_rollback_failure_restores_schema_and_bullet_links(storage, monkeypatch): + backend, context_id = storage + engine = DeltaEngine(backend) + await backend.add_schema(context_id, SchemaNode(id="old", name="old", description="before")) + batch = await engine.apply_batch( + DeltaBatch( + context_id=context_id, + operations=[ + DeltaOperation(op_type=DeltaOpType.UPDATE_SCHEMA, target_id="old", content="after"), + DeltaOperation(op_type=DeltaOpType.ADD_SCHEMA, content="new"), + ], + ) + ) + # Locate by name, independent of ordering/legacy missing delta target. + new_id = next(s.id for s in await backend.list_schemas(context_id) if s.name == "new") + await backend.add_bullet(context_id, Bullet(id="linked", content="keep", schema_id=new_id)) + + async def fail(self, schema): + raise RuntimeError("schema update unavailable") + + monkeypatch.setattr(type(backend), "update_schema", fail) + with pytest.raises(RuntimeError, match="schema update"): + await engine.rollback_batch(batch.id) + assert await backend.get_schema(new_id) is not None + assert (await backend.get_bullet("linked")).schema_id == new_id + assert (await backend.get_schema("old")).description == "after" diff --git a/tests/test_sdk_compatibility.py b/tests/test_sdk_compatibility.py new file mode 100644 index 0000000..ea2201f --- /dev/null +++ b/tests/test_sdk_compatibility.py @@ -0,0 +1,77 @@ +"""Exercise legacy public SDK methods against real REST routes and SQLite.""" + +import httpx +import pytest + +from engram.core.config import Settings +from engram.sdk.client import Engram +from engram.server import app as app_module +from engram.storage.sqlite import SQLiteBackend + + +class OfflineProvider: + async def embed(self, text): + return [1.0] + [0.0] * 1535 + + async def complete(self, *args, **kwargs): + return '{"new_insights": [], "confidence": 1.0}' + + +@pytest.mark.parametrize("id_form", ["uuid", "string"]) +async def test_legacy_sdk_roundtrip_preserves_request_and_response_contracts( + tmp_path, monkeypatch, id_form +): + settings = Settings(_env_file=None, auth_enabled=False) + monkeypatch.setattr(app_module, "get_settings", lambda: settings) + monkeypatch.setattr("engram.core.config.get_settings", lambda: settings) + monkeypatch.setattr( + app_module, "_create_storage", lambda: SQLiteBackend(str(tmp_path / "sdk.db")) + ) + monkeypatch.setattr(app_module, "_create_llm", OfflineProvider) + app = app_module.create_app() + async with app.router.lifespan_context(app): + async with Engram(url="http://test") as sdk: + await sdk._client.aclose() + sdk._client = httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://test" + ) + context = await sdk.create_context( + "SDK compatibility", {"objective": "Keep supported clients working"} + ) + context_id = context.id if id_form == "uuid" else str(context.id) + assert (await sdk.get_context(context_id)).id == context.id + assert str(context.id) in {item["id"] for item in await sdk.list_contexts()} + added = await sdk.add_bullet(context_id, "The existing SDK contract remains supported.") + bullet_id = added["bullet_id"] + assert (await sdk.get_bullet(context_id, bullet_id))["content"].startswith( + "The existing SDK" + ) + assert bullet_id in {item["id"] for item in await sdk.list_bullets(context_id)} + decision = await sdk.record_decision( + context_id, "Keep API field names", "Preserve existing integrations" + ) + assert (await sdk.get_bullet(context_id, decision["concept_id"]))[ + "bullet_type" + ] == "decision" + materialized = await sdk.materialize( + context_id, query="existing SDK", token_budget=2000 + ) + assert isinstance(materialized["rendered_text"], str) + assert materialized["materialization_id"] + assert bullet_id in materialized["bullets_included"] + assert isinstance(await sdk.recall(context_id, "existing SDK"), str) + committed = await sdk.commit( + context_id, "legacy-sdk", "No new insights in this fixture." + ) + assert {"delta_batch_id", "activity_id", "bullets_added"} <= committed.keys() + assert (await sdk.archive_bullet(context_id, bullet_id))["archived"] is True + assert bullet_id in {item["id"] for item in await sdk.list_archived_bullets(context_id)} + assert (await sdk.restore_bullet(context_id, bullet_id))["restored"] is True + assert "capacity" in await sdk.get_lifecycle(context_id) + assert "delta_batches" in await sdk.sync(context_id, since="2020-01-01T00:00:00") + assert isinstance(await sdk.get_activity(context_id), list) + assert "reflector_model" in await sdk.get_ingestion_config() + assert (await sdk.update_ingestion_config(curator_dedup_threshold=0.9))[ + "updated" + ] is True + assert (await sdk.purge_context(context_id))["purged"] is True diff --git a/tests/test_server_request_regressions.py b/tests/test_server_request_regressions.py new file mode 100644 index 0000000..14e8f63 --- /dev/null +++ b/tests/test_server_request_regressions.py @@ -0,0 +1,129 @@ +"""Regression tests for supported REST request shapes and server resource cleanup.""" + +import json +from datetime import UTC, datetime +from types import SimpleNamespace + +import httpx +import pytest +from fastapi import FastAPI + +from engram.core.models import Context, DeltaBatch, IntentAnchor +from engram.server.routes.lifecycle import router +from engram.storage.sqlite import SQLiteBackend + + +@pytest.mark.parametrize( + "since", ["2020-01-01T00:00:00", "2020-01-01T00:00:00Z", "2019-12-31T19:00:00-05:00"] +) +async def test_sync_accepts_iso_timestamps_with_and_without_timezone(tmp_path, since): + storage = SQLiteBackend(str(tmp_path / "sync.db")) + app = FastAPI() + app.state.storage = storage + app.include_router(router, prefix="/contexts") + + @app.middleware("http") + async def local_user(request, call_next): + request.state.user_id = "default" + return await call_next(request) + + try: + await storage.initialize() + context = await storage.create_context( + Context(name="Sync", intent=IntentAnchor(objective="Test")) + ) + batch = DeltaBatch(context_id=str(context.id), timestamp=datetime(2021, 1, 1, tzinfo=UTC)) + await storage.save_delta_batch(batch) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app, raise_app_exceptions=False), + base_url="http://test", + ) as client: + response = await client.post(f"/contexts/{context.id}/sync", params={"since": since}) + assert response.status_code == 200, response.text + assert [item["id"] for item in response.json()["delta_batches"]] == [batch.id] + finally: + await storage.close() + + +async def test_sync_invalid_timestamp_returns_existing_400(tmp_path): + storage = SQLiteBackend(str(tmp_path / "sync.db")) + app = FastAPI() + app.state.storage = storage + app.include_router(router, prefix="/contexts") + + @app.middleware("http") + async def local_user(request, call_next): + request.state.user_id = "default" + return await call_next(request) + + try: + await storage.initialize() + context = await storage.create_context( + Context(name="Sync", intent=IntentAnchor(objective="Test")) + ) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app, raise_app_exceptions=False), + base_url="http://test", + ) as client: + response = await client.post( + f"/contexts/{context.id}/sync", params={"since": "invalid"} + ) + assert response.status_code == 400 + finally: + await storage.close() + + +@pytest.mark.parametrize("body", [{"max_reflection_rounds": "not-a-number"}, [], None]) +async def test_invalid_ingestion_config_returns_validation_error_without_mutation( + body, monkeypatch +): + from engram.core.config import IngestionConfig, Settings + from engram.server.routes.lifecycle import config_router + + monkeypatch.setattr("engram.core.config.get_settings", lambda: Settings(_env_file=None)) + app = FastAPI() + app.state.ingestion_config = IngestionConfig() + before = app.state.ingestion_config.model_dump() + app.include_router(config_router) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app, raise_app_exceptions=False), base_url="http://test" + ) as client: + response = await client.put( + "/config/ingestion", + content=json.dumps(body), + headers={"Content-Type": "application/json"}, + ) + assert response.status_code == 422, response.text + assert app.state.ingestion_config.model_dump() == before + + +async def test_ingestion_config_updates_future_commit_and_reextraction_engines(monkeypatch): + from engram.core.config import IngestionConfig, Settings + from engram.server.routes.lifecycle import config_router + + monkeypatch.setattr("engram.core.config.get_settings", lambda: Settings(_env_file=None)) + app = FastAPI() + config = IngestionConfig() + app.state.ingestion_config = config + app.state.ingestion = SimpleNamespace( + ingestion_config=config, + reflector=SimpleNamespace(config=config), + curator=SimpleNamespace(ingestion_config=config), + ) + app.state.re_extraction = SimpleNamespace( + reflector=SimpleNamespace(config=config), + curator=SimpleNamespace(ingestion_config=config), + ) + app.include_router(config_router) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://test" + ) as client: + response = await client.put("/config/ingestion", json={"curator_dedup_threshold": 0.8}) + assert response.status_code == 200 + updated = app.state.ingestion_config + assert updated.curator_dedup_threshold == 0.8 + assert app.state.ingestion.ingestion_config is updated + assert app.state.ingestion.reflector.config is updated + assert app.state.ingestion.curator.ingestion_config is updated + assert app.state.re_extraction.reflector.config is updated + assert app.state.re_extraction.curator.ingestion_config is updated diff --git a/tests/test_storage_transactions.py b/tests/test_storage_transactions.py new file mode 100644 index 0000000..9149076 --- /dev/null +++ b/tests/test_storage_transactions.py @@ -0,0 +1,227 @@ +"""Storage transactions isolate tasks and retain rollback boundaries.""" + +import asyncio +import os +import uuid +from urllib.parse import urlsplit, urlunsplit + +import asyncpg +import pytest + +from engram.core.config import Settings +from engram.core.models import Bullet, Context, IntentAnchor +from engram.storage.postgres import PostgresBackend +from engram.storage.sqlite import SQLiteBackend + + +@pytest.fixture(params=[":memory:", "file", "postgres"]) +async def storage(request, tmp_path, monkeypatch): + monkeypatch.setattr( + "engram.core.config.get_settings", lambda: Settings(_env_file=None, auth_enabled=False) + ) + if request.param == "postgres": + admin_dsn = os.environ.get("ENGRAM_TEST_POSTGRES_DSN") + if not admin_dsn: + pytest.skip("requires disposable PostgreSQL admin DSN") + admin = await asyncpg.connect(admin_dsn) + database = "engram_tx_" + uuid.uuid4().hex + dsn = urlunsplit(urlsplit(admin_dsn)._replace(path="/" + database)) + backend = PostgresBackend(dsn) + try: + await admin.execute(f'CREATE DATABASE "{database}"') + await backend.initialize() + context = await backend.create_context( + Context(name="tx", intent=IntentAnchor(objective="test")) + ) + yield backend, str(context.id) + finally: + await backend.close() + await admin.execute(f'DROP DATABASE IF EXISTS "{database}" WITH (FORCE)') + await admin.close() + return + path = ":memory:" if request.param == ":memory:" else str(tmp_path / "transactions.db") + backend = SQLiteBackend(path) + await backend.initialize() + context = await backend.create_context( + Context(name="tx", intent=IntentAnchor(objective="test")) + ) + yield backend, str(context.id) + await backend.close() + + +async def test_commit_and_rollback(storage): + backend, context_id = storage + async with backend.transaction(context_id) as tx: + await tx.add_bullet(context_id, Bullet(id="committed", content="first")) + with pytest.raises(ValueError): + async with backend.transaction(context_id) as tx: + await tx.add_bullet(context_id, Bullet(id="rolledback", content="second")) + raise ValueError("abort") + assert [b.id for b in await backend.list_bullets(context_id)] == ["committed"] + + +async def test_scoped_handles_reject_parent_wrong_task_and_expired_use(storage): + backend, context_id = storage + async with backend.transaction(context_id) as tx: + with pytest.raises(RuntimeError): + await backend.list_bullets(context_id) + with pytest.raises(RuntimeError): + async with backend.transaction(context_id): + pass + with pytest.raises(RuntimeError): + await asyncio.create_task(tx.list_bullets(context_id)) + with pytest.raises(RuntimeError): + await tx.close() + with pytest.raises(RuntimeError): + await tx.list_bullets(context_id) + + +async def test_nested_savepoint_rollback(storage): + backend, context_id = storage + async with backend.transaction(context_id) as tx: + await tx.add_bullet(context_id, Bullet(id="outer", content="kept")) + with pytest.raises(ValueError): + async with tx.transaction(context_id) as nested: + await nested.add_bullet(context_id, Bullet(id="inner", content="discarded")) + raise ValueError("rollback nested only") + assert [b.id for b in await tx.list_bullets(context_id)] == ["outer"] + with pytest.raises(ValueError): + async with tx.transaction("different"): + pass + assert [b.id for b in await backend.list_bullets(context_id)] == ["outer"] + + +async def test_other_task_cannot_read_or_commit_uncommitted_state(storage): + backend, context_id = storage + entered = asyncio.Event() + release = asyncio.Event() + + async def writer(): + with pytest.raises(ValueError): + async with backend.transaction(context_id) as tx: + await tx.add_bullet(context_id, Bullet(id="uncommitted", content="discarded")) + entered.set() + await release.wait() + raise ValueError("abort") + + task = asyncio.create_task(writer()) + await entered.wait() + outside = asyncio.create_task( + backend.add_bullet(context_id, Bullet(id="outside", content="kept")) + ) + await asyncio.sleep(0) + if isinstance(backend, SQLiteBackend): + assert not outside.done() + release.set() + await asyncio.gather(task, outside) + assert [b.id for b in await backend.list_bullets(context_id)] == ["outside"] + + +async def test_cancellation_rolls_back_and_connection_is_reusable(storage): + backend, context_id = storage + entered = asyncio.Event() + + async def writer(): + async with backend.transaction(context_id) as tx: + await tx.add_bullet(context_id, Bullet(id="cancelled", content="discarded")) + entered.set() + await asyncio.Event().wait() + + task = asyncio.create_task(writer()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert await backend.list_bullets(context_id) == [] + async with backend.transaction(context_id) as tx: + await tx.add_bullet(context_id, Bullet(id="after", content="usable")) + assert (await backend.get_bullet("after")).content == "usable" + + +@pytest.mark.parametrize("during", ["begin", "ordinary_write", "savepoint"]) +async def test_sqlite_cancellation_during_queued_sql_does_not_leak_transaction( + storage, monkeypatch, during +): + backend, context_id = storage + if not isinstance(backend, SQLiteBackend): + pytest.skip("SQLite worker queue regression") + db = await backend._get_db() + original = db.execute + entered = asyncio.Event() + + async def delayed(sql, *args, **kwargs): + result = await original(sql, *args, **kwargs) + matches = ( + (during == "begin" and sql == "BEGIN IMMEDIATE") + or (during == "ordinary_write" and sql.startswith("INSERT INTO bullets")) + or (during == "savepoint" and sql.startswith("SAVEPOINT")) + ) + if matches: + entered.set() + await asyncio.Event().wait() + return result + + monkeypatch.setattr(db, "execute", delayed) + + async def operation(): + if during == "ordinary_write": + await backend.add_bullet(context_id, Bullet(id="cancelled", content="discarded")) + else: + async with backend.transaction(context_id) as tx: + if during == "savepoint": + async with tx.transaction(context_id): + pass + + task = asyncio.create_task(operation()) + await asyncio.wait_for(entered.wait(), 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + monkeypatch.setattr(db, "execute", original) + assert not db.in_transaction + assert await backend.list_bullets(context_id) == [] + async with backend.transaction(context_id) as tx: + await tx.add_bullet(context_id, Bullet(id="next", content="usable")) + + +async def test_read_modify_write_serializes_concurrent_archive(storage): + backend, context_id = storage + await backend.add_bullet(context_id, Bullet(id="active", content="fact")) + if isinstance(backend, PostgresBackend): + other = PostgresBackend(backend.dsn, existing_only=True) + elif backend.db_path != ":memory:": + other = SQLiteBackend(backend.db_path, existing_only=True) + else: + other = backend + fetched = asyncio.Event() + release = asyncio.Event() + + async def reinforce(): + async with backend.transaction(context_id) as tx: + bullet = await tx.get_bullet("active") + fetched.set() + await release.wait() + bullet.hit_count += 1 + await tx.update_bullet(bullet) + + writer = asyncio.create_task(reinforce()) + await asyncio.wait_for(fetched.wait(), timeout=5) + archive = asyncio.create_task(other.archive_bullet(context_id, "active")) + try: + # An archive must wait for the locked read, otherwise the stale full-row + # stats update would silently resurrect the archived bullet. + await asyncio.sleep(0.05) + assert not archive.done() + release.set() + await asyncio.wait_for(asyncio.gather(writer, archive), timeout=5) + after = await backend.get_bullet("active") + assert after.is_archived + assert after.hit_count == 1 + finally: + release.set() + for task in (writer, archive): + if not task.done(): + task.cancel() + await asyncio.gather(writer, archive, return_exceptions=True) + if other is not backend: + await other.close() diff --git a/tests/test_v04_ingestion.py b/tests/test_v04_ingestion.py index 2ddc514..8159853 100644 --- a/tests/test_v04_ingestion.py +++ b/tests/test_v04_ingestion.py @@ -363,6 +363,7 @@ async def test_low_value_unmatched_bullets_removed( # Create a low-value bullet low_value_bullet = Bullet( content="Some unimportant note from old extraction", + embedding=[1.0, 0.0], section="notes", bullet_type=BulletType.FACT, salience=0.1, @@ -372,6 +373,7 @@ async def test_low_value_unmatched_bullets_removed( ) low_value_bullet = await storage.add_bullet(ctx_id, low_value_bullet) + mock_llm.embed = AsyncMock(return_value=[0.0, 1.0]) curator = CuratorEngine(storage, mock_llm) new_reflection = Reflection( new_insights=[ diff --git a/tests/test_validity_gate.py b/tests/test_validity_gate.py index 283eb56..14ebcbb 100644 --- a/tests/test_validity_gate.py +++ b/tests/test_validity_gate.py @@ -206,13 +206,14 @@ async def complete(self, prompt, system=None, temperature=0.0, llm = ModelTrackingLLM([_reflection_response(insights), verdicts]) cfg = IngestionConfig( enable_validity_gate=True, + reflector_model="reflector-model-abc", validity_gate_model="judge-model-xyz", ) engine = IngestionEngine(storage, llm, ingestion_config=cfg) await engine.commit(context_id=context_id, agent_id="t", content="raw") - # First call is the Reflector (no override → None); second is the judge. - assert llm.models_seen == [None, "judge-model-xyz"] + # Each stage must use its own configured model, independent of adapter defaults. + assert llm.models_seen == ["reflector-model-abc", "judge-model-xyz"] # ── 5. Non-ADD ops pass through unchanged ────────────────────────────────────