diff --git a/.agents/skills/uteke-memory/SKILL.md b/.agents/skills/uteke-memory/SKILL.md index 7c524348..e652382b 100644 --- a/.agents/skills/uteke-memory/SKILL.md +++ b/.agents/skills/uteke-memory/SKILL.md @@ -5,7 +5,7 @@ description: "Persistent memory engine for AI agents via the uteke CLI — remem # Uteke Memory Skill Persistent memory engine for AI agents via the `uteke` CLI. -Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); `fusion` (weighted RRF of the vector and hybrid rankings) is the default recall strategy since 0.16.0. Zero unsafe code. +Version: **0.20.0** — SQLite + HNSW vector index (usearch default, or vecq) + FTS5 hybrid search (RRF k=60); `fusion` (weighted RRF of the vector and hybrid rankings) is the default recall strategy since 0.16.0. `unsafe_code = "forbid"` workspace-wide. > This skill ships with each release — its version tracks the CLI version > (a CI gate fails when they drift, see `skill-version-parity` test). @@ -22,7 +22,7 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | Flag | Description | |------|-------------| -| `--store ` | Override store path (default: `~/.uteke`) | +| `--store ` | Override store path (default: `~/.codecora/uteke`) | | `--namespace ` | Multi-agent isolation (default: `"default"`) | | `--json` | Machine-readable JSON output | | `--verbose` | Debug logging | @@ -34,10 +34,14 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | Command | Description | Key Options | |---------|-------------|-------------| | `uteke remember ` | Store a new memory | `--tags`, `--type`, `--entity`, `--category`, `--meta`, `--room`, `--author`, `--source`, `--source-type`, `--detect-contradiction` | -| `uteke recall ` | Fusion search (default since 0.16.0 — weighted RRF of vector + hybrid rankings) | `--limit`, `--tags`, `--entity`, `--category`, `--min`, `--strategy` (fusion/vector/fts5/hybrid/graph), `--salience`, `--recency`, `--related`, `--depth`, `--context`, `--at` (time-travel), `--type` (all/memory/doc), `--where` (JSON field filter) | +| `uteke recall ` | Fusion search (default since 0.16.0 — weighted RRF of vector + hybrid rankings) | `--limit`, `--tags`, `--entity`, `--category`, `--min`, `--strategy` (fusion/vector/fts5/hybrid/graph), `--salience`, `--recency`, `--related`, `--depth`, `--context`, `--at` (time-travel), `--type` (all/memory/doc), `--where` (JSON field filter), `--explain` (real vector similarity), `--budget`, `--pack`, `--strict`, `--enrich`, `--exclude-ids`, `--content-format`, `--no-salience`, `--no-recency` | | `uteke search ` | Keyword text search | `--limit`, `--tags` | | `uteke list` | List memories with filters | `--tag`, `--entity`, `--category`, `--limit`, `--offset`, `--at` | | `uteke get ` | Get single memory by UUID | | +| `uteke update ` | Update a memory in place | `--content`, `--tags` (replaces set), `--importance`, `--pinned`, `--type` | +| `uteke context` | Project context summary (counts, top tags, recent activity) | `--namespace` | +| `uteke guide` | Print the agent-facing memory tools guide for system-prompt injection | | +| `uteke feedback ` | Record usefulness (importance +0.05 / −0.10) | | | `uteke forget ` | Delete memory by ID, tag, tier, or all | `--tag`, `--cold`, `--all`, `--confirm` | ### Documents (Wiki / Knowledge Base) @@ -79,6 +83,14 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | Command | Description | Key Options | |---------|-------------|-------------| | `uteke room create ` | Create a room | `--title` | +| `uteke room update ` | Update title/description | `--title`, `--description` | +| `uteke room rename ` | Rename a room (members and document links move too) | | +| `uteke room move-memory ` | Move a memory between rooms, keeping link provenance | `--from`, `--to` | +| `uteke room add-document ` | Link a document to a room | | +| `uteke room remove-document ` | Unlink a document | | +| `uteke room list-documents ` | List documents linked to a room | | +| `uteke room list-rooms ` | List rooms that reference a document | | +| `uteke room consolidate ` | Merge a room's memories into denser records (dry-run unless `--apply`; uses LLM calls) | `--apply`, `--max-calls` | | `uteke room list` | List all rooms | `--namespace` | | `uteke room stats ` | Room statistics and participants | | | `uteke room recall ` | Recall room memories | `--query`, `--author`, `--limit`, `--min` | @@ -86,6 +98,15 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | `uteke room document ` | Generate structured document from room | | | `uteke room delete ` | Delete room (memories preserved) | `--confirm` | +### Conflict Resolution + +| Command | Description | +|---------|-------------| +| `uteke supersede ` | Mark `OLD` superseded by `NEW` (resolves a contradiction) | +| `uteke contradictions list` | Resolution ledger: superseded-but-not-restored memories | +| `uteke contradictions undo ` | Restore a superseded memory (undoes the pair) | +| `uteke provenance ` | Full provenance report for a memory | + ### Pinning & Importance | Command | Description | @@ -115,6 +136,9 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | `uteke namespace list` | List all namespaces with counts | | `uteke namespace stats ` | Stats for a specific namespace | | `uteke namespace switch ` | Set default namespace in config | +| `uteke namespace move ` | Move a memory to another namespace | +| `uteke namespace rename ` | Rename a namespace (merges if the target exists) | +| `uteke namespace delete ` | Delete with `--strategy refuse\|merge\|deprecate` (`--target` for merge, `--confirm` required) | ### Memory Aging @@ -132,13 +156,15 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | `uteke doctor` | Full health check: DB, index, embedding model, consistency | | `uteke verify` | Compare DB count vs vector index count | | `uteke verify-checksums` | Verify binary integrity against SHA256 checksums | -| `uteke repair` | Rebuild vector index from SQLite | +| `uteke repair` | Rebuild vector index from SQLite (`--rebuild` deletes corrupt index files first, `--reembed` regenerates missing embeddings) | -> **Stale-index symptom (pitfall 0g):** a memory is recallable via +> **Stale-index symptom:** a memory is recallable via > `--strategy fts5` but missing under the default fusion/hybrid — that is a > vector-index desync, not data loss. Run `uteke verify`, then `uteke repair`. > On uteke-serve, use `POST /verify` / `POST /repair` (HTTP) or the > `uteke_verify` / `uteke_repair` MCP tools — no restart needed. +> A held index lock (another `uteke`/`uteke-serve` process) never deletes the +> index in builds after 0.19.1: the CLI falls back to an in-memory index instead. > Run `verify` after every serve upgrade (#1245 class: CLI upgraded, serve left > behind → index written by the old binary reads as mismatched). > @@ -175,13 +201,13 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` ## Architecture - **Storage:** SQLite (WAL mode) with namespace column + FTS5 virtual table -- **Vector index:** usearch persistent HNSW (768d, cosine similarity) +- **Vector index:** usearch persistent HNSW (default) or vecq (quantized, no C++ dependency), 768d cosine. Select with `UTEKE_VECTOR_BACKEND=usearch|vecq` or `[vector] backend` in `uteke.toml`; switching engines on an existing store rebuilds the index from SQLite - **Hybrid search:** RRF (k=60) merges vector + FTS5 results; graph strategy adds graph-signal reranking - **Fusion (default since 0.16.0):** weighted RRF of the vector and hybrid rankings — LongMemEval 500Q R@5 0.946 - **Embedding:** ONNX EmbeddingGemma Q4 (768d), auto-downloaded - **Tiered memory:** Hot (<7d, +0.1 boost), Warm (<30d), Cold (>30d) - **Schema versioning:** Integer counter, auto-migration on upgrade -- **Zero unsafe code** (`unsafe_code = "forbid"`) +- **`unsafe_code = "forbid"`** at the workspace level - **Project-scoped stores:** `uteke --store .uteke remember "..."` ## Usage Patterns @@ -254,6 +280,14 @@ uteke orphans --threshold 0.2 # Find disconnected - `/home/user/repos/my-saas-app/` → `project:my-saas-app` 2. From file paths mentioned in the conversation 3. From the project name the user mentions +4. If the project has a registry (a doc such as `uteke-project-registry`), use the + canonical name from it, resolving aliases (e.g. a folder `api-v2` that the + registry lists as `recruitment-api`). + +**Never derive a tag from a generic folder name** (`development`, `projects`, +`documents`, `desktop`, `downloads`, `tmp`, `src`, `code`, `work`, `workspace`, +`repos`, or your home folder). Those are not projects; when `cwd` is one of them +and the conversation names no project, add no `project:` tag. ### Rules @@ -262,7 +296,40 @@ uteke orphans --threshold 0.2 # Find disconnected | **REMEMBER** | Always include `project:` in `--tags` when the memory is project-specific | | **RECALL** | Always include `--tags project:` to scope recall to the current project | | **NO PROJECT** | If the conversation is not about any specific project (e.g., general chat), do NOT add a `project:` tag | -| **TAG FORMAT** | Always lowercase: `project:bond`, `project:uteke` (not `project:Bond`) | +| **TAG FORMAT** | Lowercase with hyphens only — no uppercase, underscores, or spaces: `project:bond`, `project:my-saas-app` | +| **HIERARCHY** | A component of a larger project carries two tags, the component and its parent (`project:samson-mobile` + `project:samson`): the component matches the folder name, the parent recalls the whole project | +| **NO GENERIC TAGS** | No tags from generic folder names, no tags that act as ids, no tags that mimic a field (use `--type`, not `qtype:`) | + +## Namespace, Room and Tag Conventions + +Keep the three axes distinct so recall stays filterable: + +| Axis | Means | Rule | +|------|-------|------| +| **Namespace** | Workspace / whose area | Few and coarse (e.g. one per agent family or tool). Not one per project: project identity lives in the room and tags. Avoid `repo-*` / `project:*` namespaces. Do not keep writing to `default` once you have chosen real ones | +| **Room** | One project or discussion | One room per parent project named `` (lowercase, hyphens; no `project:` / `repo-` prefix, no `-dev` suffix). Discussions `disc:`, research `riset:`. Always set `--title` | +| **Tag** | Search aid | `project:` (registry names only), `tool:` for items about a specific tool or agent (`tool:claude-code`), `agent:` only for the origin of a multi-profile agent | + +- **A room's namespace cannot be changed after creation** (#1352) and merging + namespaces does not relabel rooms: create the room in the right namespace + from the start. +- `uteke update --tags` **replaces** the whole tag set — read the memory first + and write back the merged list; there is no bulk add-tag operation. +- `uteke namespace rename` changes only the memories' namespace, not the `rooms` table. + +### What to store + +Store only what is worth recalling again: decisions with their reason, causes of +non-obvious bugs, gotchas, procedures — one fact per memory. Never store tokens, +passwords, or other credentials. Do not auto-save every prompt; automatic writers +produce noise and junk tags. + +### Destructive cleanup needs the owner's go-ahead + +Writing bulk changes (tag deletes, namespace merges/deletes, room deletes, memory +moves) is not reversible in general. Take a backup first (`uteke export`, plus a +map of ids before/after), do a dry run and a small pilot, and ask for explicit +confirmation per step. `uteke room delete` keeps the memories; `uteke forget` does not. ### Examples @@ -330,3 +397,10 @@ uteke-serve --port 8767 ``` Endpoints mirror CLI commands (e.g., `POST /remember`, `POST /recall`). Supports read-only API tokens for restricted access. + +Security defaults (builds after 0.19.1): + +- **Auth:** `--auth-token` / `UTEKE_AUTH_TOKEN` (admin) and `UTEKE_READ_ONLY_TOKEN`. When the CLI routes through the server (`[server] enabled = true`) it forwards `UTEKE_AUTH_TOKEN` as a bearer token, to loopback servers only. +- **CORS is off unless configured:** set `cors_origins` in `uteke.toml` / `--cors-origin`. `"*"` is an explicit opt-in — avoid it without auth. +- **Project config is not trusted for endpoints:** a `.uteke/uteke.toml` in the working directory cannot set `[server]`, `[extraction]`, embedding/extraction `base_url`, `endpoint_path`, `api_key`, `embedding.backend`, or `server.host/port` (set them in the global config). `UTEKE_TRUST_PROJECT_CONFIG=1` overrides this. +- List-style endpoints cap `limit` at 1000; `/extract` caps `max_facts` at 100. diff --git a/CHANGELOG.md b/CHANGELOG.md index 6be7ce59..0b3a03a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,74 @@ # Changelog +## [0.20.0] - 2026-10-08 + +Minor release. Theme: **one write policy for the vector index, atomic documents, +shared config resolution, and a security/robustness sweep** — every index +mutation now goes through a single module, document upserts are +embed-then-commit, `uteke.toml` is resolved by one implementation for the CLI and +the server, `GET /graph` finally serves the data `remember` writes, and the +validated findings of the v0.19.1 review (#1319) are fixed. 29 PRs since v0.19.1. + +### ⚠️ Behavior Change + +- **List `limit` is capped at 1000** on `/list`, `/recent`, `/contradictions`, `/doc/list` and deprecated listings (#1321); `/extract` `max_facts` is capped at 100. Defaults are unchanged. +- **CORS is off by default** (a wildcard `*` is opt-in), 500 responses carry a generic body, and request query strings are no longer logged (#1321). +- **Project-local `.uteke/uteke.toml` can no longer set credentials, endpoints or `[server]`** (`embedding.{backend,api_key,base_url,endpoint_path}`, `embed_fallback.*` and `extraction.*` credentials/endpoints, the whole `[server]` section); set them in the global config, or `UTEKE_TRUST_PROJECT_CONFIG=1` to trust the file (#1321, #1354). The server now applies the same rule and merges project sections per key instead of replacing whole sections. +- **`GET /graph` node ids are memory ids** (they were graph-node ids), edges are the live `memory_edges` plus explicit graph edges, and the response gains `node_id`/`limit` query parameters and a `truncated` flag; edges also carry `source`/`target` (#1366, #1370). `POST /graph/edge` still returns the node ids it created and both id kinds are accepted. +- **Document upsert/update is atomic**: if embedding fails (or no embedder is available) nothing is written, instead of leaving a half-written document (#1350). Re-parenting through `doc upsert` now moves the whole subtree and rejects cycles and results deeper than 10 levels (#1378). +- **Bulk "cold" forget skips pinned memories** (#1333). +- **List/search/`list_at_time` payloads carry the stored `slug`, `source`, `source_type`, `author_type`, `deprecated_at`** instead of defaults (#1346, #1333). Shape unchanged. +- **Auto-created `similar_to` / `possible_duplicate` links** no longer point at deprecated memories and small namespaces inside a large index now get their own links (#1371). The default `fusion` recall does not read edges; LongMemEval fast50 (default strategy) is unchanged: recall@5 0.980, recall@10 1.000. + +### Security + +- **Host header validation blocks DNS rebinding** (#1326, #1360). +- **`/import` size limit, `/mcp` body limit and CORS, and `resolve_id` hardened** (#1328, #1359): the import cap and its 413 message agree with the body limit, an oversized `/mcp` body returns 413, and ids accept only hex digits and dashes (no LIKE wildcards in the prefix scan). +- **A lock timeout no longer deletes the index files of another process; the server releases its request slot on panic; the CLI sends `UTEKE_AUTH_TOKEN` to the local server** (#1320). +- **LIKE wildcards in document paths are escaped** (#1378): `_` and `%` in a stored path used to match unrelated documents, so moving or deleting `/my_doc/` could rewrite or remove `/myxdoc/...`. + +### Added + +- **`full_ids` option for MCP `uteke_list`, `uteke_recall` and `uteke_room_memories`** (#1357, #1358): prints full UUIDs instead of the 8-character prefix, which is ambiguous for UUIDv7 ids written within ~65 s. +- **`GET /graph?node_id=&limit=`** and an MCP `uteke_graph` that lists the latest edges and takes `limit` (#1370). +- **Vector index sidecar consistency check on load** (#1325, #1362): the `.keys` file carries a digest and row count of the index; a mismatch (crash between the two renames, missing sidecar) rebuilds the index from SQLite. +- **`doc export` writes every document** to `--output` and honours the namespace; JSON export is an array of full documents (#1376). +- **Namespace/room/tag conventions in the bundled skill** (#1356) and a refreshed skill with a command-coverage test (#1329). +- Configuration docs: what a project config may not set. + +### Changed + +- **`index_sync`: one module owns every vector index write** (#1340, #1347, #1348): a single retry/persist policy, deterministic dimension errors are not retried, and an outcome that says what actually landed. 60 hand-written call sites across 7 modules went away. +- **Index saves are batched** (#1322, #1363): a save happens after 64 mutations or 2 s, with flushes on shutdown and drop, instead of rewriting the whole index on every write. +- **vecq index is compacted when dead rows pass 25 %** (#1324, #1364). +- **The embedder lock is released between document chunks** (#1323, #1361). +- **`remember` no longer scans the namespace** for dedup and auto-linking (one query over the candidates) and no longer holds the index write lock while wiring edges (#1368, #1371). +- **`uteke.toml` is resolved by one module for the CLI and the server** (#1342, #1354): per-key merge, one untrusted-key policy, per-layer validation; adding a config key now needs only a struct field. +- **All `memory` SELECTs share one column list** with the row decoder (#1339, #1346). +- Edges to deprecated memories no longer count in graph rerank or edge traversal (#1367, #1373). + +### Fixed + +- **FTS5 results had shifted metadata columns** (#1333): `slug`, `source`, `source_type`, `author_type`, `deprecated_at` were misread on hybrid recall and `recall_explain`. +- **The exclusive index lock was lost after the first save on Unix** (#1333): the flock stayed on the replaced inode. +- **`doc_search` paired chunk scores with the wrong rows** when a chunk row was missing (#1333); `resolve_db_path` no longer `chmod`s a directory it did not create (#1333). +- **Changing a document title failed with "database corruption"** or left stale tokens in the search index (#1349, #1351): the `documents_fts` update/delete triggers now use the FTS5 `'delete'` command; existing stores are healed and the index rebuilt on open. +- **MCP `uteke_graph_add_edge` and `dream` contradiction edges violated the graph FK** (#1365, #1370). +- **Documents**: `has_children` stays in sync on delete and re-parent, `move` of a legacy row with an empty path no longer rewrites every document (#1375); a re-parenting upsert carries its descendants (#1377, #1378). +- **`tags rename` onto an existing tag duplicated it**, `VectorIndex::build` on an empty store reset the dimension to 768, and `recompute_importance` loaded every embedding and committed per row (#1375). +- **CLI**: `verify-checksums --json` exits 1 on a mismatch and matches the file name exactly; `[lifecycle]`, `[aging] max_access_count`, `update_check` and `doctor_footer` in config files are honoured; non-ASCII titles no longer panic; `doc move/update --json` is valid JSON (#1334, #1376). +- **CLI recall flags** the server cannot honour fall back to the local store instead of being dropped; `list --entity/--category` paginates the matches; `doc list --tree` prints descendants; `upgrade` verifies every new binary before installing any; the batch-import bail-out and `onboard` TOML escaping and legacy `migrate_content` are fixed (#1376). +- **`uteke upgrade` verification no longer fails with "Text file busy"** on a freshly written binary (#1338). + +### Docs + +- Benchmark tooling: `make_fast_eval.py` finds its id list, `modal_fanout.py` takes `LMEVAL_MODEL_DIR`/`UTEKE_RELEASE`, dataset size corrected (#1374). + +### Dependencies + +- uuid 1.27.0, usearch 2.26.3. + +--- ## [0.19.1] - 2026-10-05 Patch release. Theme: **upgrade-path reliability** — `uteke upgrade` and the diff --git a/Cargo.lock b/Cargo.lock index 380c5258..aa97c954 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2756,9 +2756,9 @@ checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" [[package]] name = "usearch" -version = "2.26.2" +version = "2.26.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b43f3a7c8d3b100d53c1a34e9ce4cd76cc83627bdb4fcd035d4e82102778aeee" +checksum = "3ae1c08a924aecdf8280c6d14f8bcc3677092a1f33d7c254f0f70fdc632ae694" dependencies = [ "cxx", "cxx-build", @@ -2767,7 +2767,7 @@ dependencies = [ [[package]] name = "uteke-cli" -version = "0.19.1" +version = "0.20.0" dependencies = [ "chrono", "clap", @@ -2793,7 +2793,7 @@ dependencies = [ [[package]] name = "uteke-core" -version = "0.19.1" +version = "0.20.0" dependencies = [ "chrono", "dirs", @@ -2809,6 +2809,7 @@ dependencies = [ "tempfile", "thiserror", "tokenizers", + "toml", "tracing", "usearch", "uuid", @@ -2817,7 +2818,7 @@ dependencies = [ [[package]] name = "uteke-mcp" -version = "0.19.1" +version = "0.20.0" dependencies = [ "chrono", "dirs", @@ -2829,7 +2830,7 @@ dependencies = [ [[package]] name = "uteke-server" -version = "0.19.1" +version = "0.20.0" dependencies = [ "chrono", "ctrlc", @@ -2860,9 +2861,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.26.1" +version = "1.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ef6dac1e96601b4fb3acccccff2139741fcb757cb9a36089bf5be91cfb285ce" +checksum = "97277d36b9c3ace13e58fa6e753f8b0bbbf302a18dd193240d70f9e29681059a" dependencies = [ "getrandom 0.4.3", "js-sys", diff --git a/Cargo.toml b/Cargo.toml index ff36b69b..dd466da1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,7 @@ members = [ ] [workspace.package] -version = "0.19.1" +version = "0.20.0" edition = "2024" license = "Apache-2.0" repository = "https://github.com/codecoradev/uteke" diff --git a/README.id.md b/README.id.md index 011de932..5c7ed757 100644 --- a/README.id.md +++ b/README.id.md @@ -392,7 +392,7 @@ Bisa. Uteke punya MCP server yang langsung pakai dengan Claude Code, Cursor, dan
Sudah production-ready? -Uteke sekarang v0.19.1 dengan 200+ test, CI/CD di setiap commit, dan benchmark harness. Dipakai production oleh tim CodeCora dan early adopter lain. Masih di versi 0.x, jadi mungkin ada rough edges, tapi core-nya udah stabil. +Uteke sekarang v0.20.0 dengan 200+ test, CI/CD di setiap commit, dan benchmark harness. Dipakai production oleh tim CodeCora dan early adopter lain. Masih di versi 0.x, jadi mungkin ada rough edges, tapi core-nya udah stabil.
--- diff --git a/README.md b/README.md index 0e700791..f770a7ba 100644 --- a/README.md +++ b/README.md @@ -353,7 +353,7 @@ Yes. Uteke ships with an MCP server that works with Claude Code, Cursor, and Her
Is it production-ready? -Uteke is at v0.19.1 with 200+ tests, CI/CD on every commit, and a benchmark harness. It's used in production by the CodeCora team and other early adopters. Still in 0.x, so expect rough edges, but the core is stable. +Uteke is at v0.20.0 with 200+ tests, CI/CD on every commit, and a benchmark harness. It's used in production by the CodeCora team and other early adopters. Still in 0.x, so expect rough edges, but the core is stable.
--- diff --git a/benchmarks/longmemeval/REPRODUCING.md b/benchmarks/longmemeval/REPRODUCING.md index 9e1ea014..6b60326b 100644 --- a/benchmarks/longmemeval/REPRODUCING.md +++ b/benchmarks/longmemeval/REPRODUCING.md @@ -126,7 +126,7 @@ preemption). Same harness, bigger loop: ```bash cd benchmarks/longmemeval -# dataset (full S split, ~50 MB) +# dataset (full S split, ~277 MB) curl -L -o data/longmemeval_s_cleaned.json \ https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json @@ -135,6 +135,12 @@ curl -L -o data/longmemeval_s_cleaned.json \ export UTEKE_GIT_REF= # e.g. the v0.19.0 release tag export MODAL_TOKEN_ID=... MODAL_TOKEN_SECRET=... +# modal_fanout.py bakes the embedding model from a local directory (default is the +# Hermes box path) and reuses libonnxruntime from a published release; override both +# on any other machine: +export LMEVAL_MODEL_DIR=$HOME/.codecora/uteke/models/embeddinggemma-q4 +export UTEKE_RELEASE=v0.19.1 # release whose libonnxruntime is reused + modal run scripts/modal_fanout.py --strategy default --num-shards 10 # watch progress (shards appear at completion), then score the merged output: diff --git a/benchmarks/longmemeval/scripts/make_fast_eval.py b/benchmarks/longmemeval/scripts/make_fast_eval.py index a3ea7cd1..6890ef4f 100644 --- a/benchmarks/longmemeval/scripts/make_fast_eval.py +++ b/benchmarks/longmemeval/scripts/make_fast_eval.py @@ -27,7 +27,7 @@ HERE = Path(__file__).parent.parent # benchmarks/longmemeval (scripts live in scripts/) DATA = HERE / "data" PARENT = DATA / "longmemeval_s_cleaned.json" -IDS = HERE / "fast_eval_ids.json" +IDS = Path(__file__).parent / "fast_eval_ids.json" # lives next to this script # id-list file name -> output dataset file name NAME_MAP = { diff --git a/benchmarks/longmemeval/scripts/modal_fanout.py b/benchmarks/longmemeval/scripts/modal_fanout.py index 529863dd..30b50846 100644 --- a/benchmarks/longmemeval/scripts/modal_fanout.py +++ b/benchmarks/longmemeval/scripts/modal_fanout.py @@ -44,7 +44,10 @@ import modal REPO_DIR = pathlib.Path(__file__).parent.parent # benchmarks/longmemeval (scripts live in scripts/) -UTEKE_VERSION = "v0.15.0" +# Release whose bundled libonnxruntime is reused (and whose binary is measured +# when UTEKE_GIT_REF is unset). Override with the current release so the ORT +# library matches the `ort` crate the measured source was built against. +UTEKE_VERSION = os.environ.get("UTEKE_RELEASE", "v0.15.0") # When set, the image builds uteke from this git ref instead of downloading # the UTEKE_VERSION release. Used for pre-release validation (e.g. develop tip # before tagging). The exact SHA must be used, not a branch name, so the image @@ -52,7 +55,12 @@ UTEKE_GIT_REF = os.environ.get("UTEKE_GIT_REF", "") DATA_FILE = os.environ.get("LMEVAL_DATA", "longmemeval_s_cleaned.json") # Local source of the embedding model (must contain onnx/ + tokenizer.json). -MODEL_SOURCE = pathlib.Path("/opt/data/.codecora/uteke/models/embeddinggemma-q4") +MODEL_SOURCE = pathlib.Path( + os.environ.get( + "LMEVAL_MODEL_DIR", + "/opt/data/.codecora/uteke/models/embeddinggemma-q4", + ) +) app = modal.App("uteke-longmemeval") diff --git a/crates/uteke-cli/Cargo.toml b/crates/uteke-cli/Cargo.toml index c2b37bc4..d409f2cd 100644 --- a/crates/uteke-cli/Cargo.toml +++ b/crates/uteke-cli/Cargo.toml @@ -22,7 +22,7 @@ usearch = ["uteke-core/usearch"] vecq = ["uteke-core/vecq"] [dependencies] -uteke-core = { path = "../uteke-core", version = "0.19.1", default-features = false, features = ["onnx"] } +uteke-core = { path = "../uteke-core", version = "0.20.0", default-features = false, features = ["onnx"] } clap = { version = "4", features = ["derive"] } serde_json = "1" tracing = "0.1" diff --git a/crates/uteke-cli/assets/uteke-memory-skill.md b/crates/uteke-cli/assets/uteke-memory-skill.md index 7c524348..e652382b 100644 --- a/crates/uteke-cli/assets/uteke-memory-skill.md +++ b/crates/uteke-cli/assets/uteke-memory-skill.md @@ -5,7 +5,7 @@ description: "Persistent memory engine for AI agents via the uteke CLI — remem # Uteke Memory Skill Persistent memory engine for AI agents via the `uteke` CLI. -Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); `fusion` (weighted RRF of the vector and hybrid rankings) is the default recall strategy since 0.16.0. Zero unsafe code. +Version: **0.20.0** — SQLite + HNSW vector index (usearch default, or vecq) + FTS5 hybrid search (RRF k=60); `fusion` (weighted RRF of the vector and hybrid rankings) is the default recall strategy since 0.16.0. `unsafe_code = "forbid"` workspace-wide. > This skill ships with each release — its version tracks the CLI version > (a CI gate fails when they drift, see `skill-version-parity` test). @@ -22,7 +22,7 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | Flag | Description | |------|-------------| -| `--store ` | Override store path (default: `~/.uteke`) | +| `--store ` | Override store path (default: `~/.codecora/uteke`) | | `--namespace ` | Multi-agent isolation (default: `"default"`) | | `--json` | Machine-readable JSON output | | `--verbose` | Debug logging | @@ -34,10 +34,14 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | Command | Description | Key Options | |---------|-------------|-------------| | `uteke remember ` | Store a new memory | `--tags`, `--type`, `--entity`, `--category`, `--meta`, `--room`, `--author`, `--source`, `--source-type`, `--detect-contradiction` | -| `uteke recall ` | Fusion search (default since 0.16.0 — weighted RRF of vector + hybrid rankings) | `--limit`, `--tags`, `--entity`, `--category`, `--min`, `--strategy` (fusion/vector/fts5/hybrid/graph), `--salience`, `--recency`, `--related`, `--depth`, `--context`, `--at` (time-travel), `--type` (all/memory/doc), `--where` (JSON field filter) | +| `uteke recall ` | Fusion search (default since 0.16.0 — weighted RRF of vector + hybrid rankings) | `--limit`, `--tags`, `--entity`, `--category`, `--min`, `--strategy` (fusion/vector/fts5/hybrid/graph), `--salience`, `--recency`, `--related`, `--depth`, `--context`, `--at` (time-travel), `--type` (all/memory/doc), `--where` (JSON field filter), `--explain` (real vector similarity), `--budget`, `--pack`, `--strict`, `--enrich`, `--exclude-ids`, `--content-format`, `--no-salience`, `--no-recency` | | `uteke search ` | Keyword text search | `--limit`, `--tags` | | `uteke list` | List memories with filters | `--tag`, `--entity`, `--category`, `--limit`, `--offset`, `--at` | | `uteke get ` | Get single memory by UUID | | +| `uteke update ` | Update a memory in place | `--content`, `--tags` (replaces set), `--importance`, `--pinned`, `--type` | +| `uteke context` | Project context summary (counts, top tags, recent activity) | `--namespace` | +| `uteke guide` | Print the agent-facing memory tools guide for system-prompt injection | | +| `uteke feedback ` | Record usefulness (importance +0.05 / −0.10) | | | `uteke forget ` | Delete memory by ID, tag, tier, or all | `--tag`, `--cold`, `--all`, `--confirm` | ### Documents (Wiki / Knowledge Base) @@ -79,6 +83,14 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | Command | Description | Key Options | |---------|-------------|-------------| | `uteke room create ` | Create a room | `--title` | +| `uteke room update ` | Update title/description | `--title`, `--description` | +| `uteke room rename ` | Rename a room (members and document links move too) | | +| `uteke room move-memory ` | Move a memory between rooms, keeping link provenance | `--from`, `--to` | +| `uteke room add-document ` | Link a document to a room | | +| `uteke room remove-document ` | Unlink a document | | +| `uteke room list-documents ` | List documents linked to a room | | +| `uteke room list-rooms ` | List rooms that reference a document | | +| `uteke room consolidate ` | Merge a room's memories into denser records (dry-run unless `--apply`; uses LLM calls) | `--apply`, `--max-calls` | | `uteke room list` | List all rooms | `--namespace` | | `uteke room stats ` | Room statistics and participants | | | `uteke room recall ` | Recall room memories | `--query`, `--author`, `--limit`, `--min` | @@ -86,6 +98,15 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | `uteke room document ` | Generate structured document from room | | | `uteke room delete ` | Delete room (memories preserved) | `--confirm` | +### Conflict Resolution + +| Command | Description | +|---------|-------------| +| `uteke supersede ` | Mark `OLD` superseded by `NEW` (resolves a contradiction) | +| `uteke contradictions list` | Resolution ledger: superseded-but-not-restored memories | +| `uteke contradictions undo ` | Restore a superseded memory (undoes the pair) | +| `uteke provenance ` | Full provenance report for a memory | + ### Pinning & Importance | Command | Description | @@ -115,6 +136,9 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | `uteke namespace list` | List all namespaces with counts | | `uteke namespace stats ` | Stats for a specific namespace | | `uteke namespace switch ` | Set default namespace in config | +| `uteke namespace move ` | Move a memory to another namespace | +| `uteke namespace rename ` | Rename a namespace (merges if the target exists) | +| `uteke namespace delete ` | Delete with `--strategy refuse\|merge\|deprecate` (`--target` for merge, `--confirm` required) | ### Memory Aging @@ -132,13 +156,15 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` | `uteke doctor` | Full health check: DB, index, embedding model, consistency | | `uteke verify` | Compare DB count vs vector index count | | `uteke verify-checksums` | Verify binary integrity against SHA256 checksums | -| `uteke repair` | Rebuild vector index from SQLite | +| `uteke repair` | Rebuild vector index from SQLite (`--rebuild` deletes corrupt index files first, `--reembed` regenerates missing embeddings) | -> **Stale-index symptom (pitfall 0g):** a memory is recallable via +> **Stale-index symptom:** a memory is recallable via > `--strategy fts5` but missing under the default fusion/hybrid — that is a > vector-index desync, not data loss. Run `uteke verify`, then `uteke repair`. > On uteke-serve, use `POST /verify` / `POST /repair` (HTTP) or the > `uteke_verify` / `uteke_repair` MCP tools — no restart needed. +> A held index lock (another `uteke`/`uteke-serve` process) never deletes the +> index in builds after 0.19.1: the CLI falls back to an in-memory index instead. > Run `verify` after every serve upgrade (#1245 class: CLI upgraded, serve left > behind → index written by the old binary reads as mismatched). > @@ -175,13 +201,13 @@ Version: **0.19.1** — SQLite + usearch HNSW + FTS5 hybrid search (RRF k=60); ` ## Architecture - **Storage:** SQLite (WAL mode) with namespace column + FTS5 virtual table -- **Vector index:** usearch persistent HNSW (768d, cosine similarity) +- **Vector index:** usearch persistent HNSW (default) or vecq (quantized, no C++ dependency), 768d cosine. Select with `UTEKE_VECTOR_BACKEND=usearch|vecq` or `[vector] backend` in `uteke.toml`; switching engines on an existing store rebuilds the index from SQLite - **Hybrid search:** RRF (k=60) merges vector + FTS5 results; graph strategy adds graph-signal reranking - **Fusion (default since 0.16.0):** weighted RRF of the vector and hybrid rankings — LongMemEval 500Q R@5 0.946 - **Embedding:** ONNX EmbeddingGemma Q4 (768d), auto-downloaded - **Tiered memory:** Hot (<7d, +0.1 boost), Warm (<30d), Cold (>30d) - **Schema versioning:** Integer counter, auto-migration on upgrade -- **Zero unsafe code** (`unsafe_code = "forbid"`) +- **`unsafe_code = "forbid"`** at the workspace level - **Project-scoped stores:** `uteke --store .uteke remember "..."` ## Usage Patterns @@ -254,6 +280,14 @@ uteke orphans --threshold 0.2 # Find disconnected - `/home/user/repos/my-saas-app/` → `project:my-saas-app` 2. From file paths mentioned in the conversation 3. From the project name the user mentions +4. If the project has a registry (a doc such as `uteke-project-registry`), use the + canonical name from it, resolving aliases (e.g. a folder `api-v2` that the + registry lists as `recruitment-api`). + +**Never derive a tag from a generic folder name** (`development`, `projects`, +`documents`, `desktop`, `downloads`, `tmp`, `src`, `code`, `work`, `workspace`, +`repos`, or your home folder). Those are not projects; when `cwd` is one of them +and the conversation names no project, add no `project:` tag. ### Rules @@ -262,7 +296,40 @@ uteke orphans --threshold 0.2 # Find disconnected | **REMEMBER** | Always include `project:` in `--tags` when the memory is project-specific | | **RECALL** | Always include `--tags project:` to scope recall to the current project | | **NO PROJECT** | If the conversation is not about any specific project (e.g., general chat), do NOT add a `project:` tag | -| **TAG FORMAT** | Always lowercase: `project:bond`, `project:uteke` (not `project:Bond`) | +| **TAG FORMAT** | Lowercase with hyphens only — no uppercase, underscores, or spaces: `project:bond`, `project:my-saas-app` | +| **HIERARCHY** | A component of a larger project carries two tags, the component and its parent (`project:samson-mobile` + `project:samson`): the component matches the folder name, the parent recalls the whole project | +| **NO GENERIC TAGS** | No tags from generic folder names, no tags that act as ids, no tags that mimic a field (use `--type`, not `qtype:`) | + +## Namespace, Room and Tag Conventions + +Keep the three axes distinct so recall stays filterable: + +| Axis | Means | Rule | +|------|-------|------| +| **Namespace** | Workspace / whose area | Few and coarse (e.g. one per agent family or tool). Not one per project: project identity lives in the room and tags. Avoid `repo-*` / `project:*` namespaces. Do not keep writing to `default` once you have chosen real ones | +| **Room** | One project or discussion | One room per parent project named `` (lowercase, hyphens; no `project:` / `repo-` prefix, no `-dev` suffix). Discussions `disc:`, research `riset:`. Always set `--title` | +| **Tag** | Search aid | `project:` (registry names only), `tool:` for items about a specific tool or agent (`tool:claude-code`), `agent:` only for the origin of a multi-profile agent | + +- **A room's namespace cannot be changed after creation** (#1352) and merging + namespaces does not relabel rooms: create the room in the right namespace + from the start. +- `uteke update --tags` **replaces** the whole tag set — read the memory first + and write back the merged list; there is no bulk add-tag operation. +- `uteke namespace rename` changes only the memories' namespace, not the `rooms` table. + +### What to store + +Store only what is worth recalling again: decisions with their reason, causes of +non-obvious bugs, gotchas, procedures — one fact per memory. Never store tokens, +passwords, or other credentials. Do not auto-save every prompt; automatic writers +produce noise and junk tags. + +### Destructive cleanup needs the owner's go-ahead + +Writing bulk changes (tag deletes, namespace merges/deletes, room deletes, memory +moves) is not reversible in general. Take a backup first (`uteke export`, plus a +map of ids before/after), do a dry run and a small pilot, and ask for explicit +confirmation per step. `uteke room delete` keeps the memories; `uteke forget` does not. ### Examples @@ -330,3 +397,10 @@ uteke-serve --port 8767 ``` Endpoints mirror CLI commands (e.g., `POST /remember`, `POST /recall`). Supports read-only API tokens for restricted access. + +Security defaults (builds after 0.19.1): + +- **Auth:** `--auth-token` / `UTEKE_AUTH_TOKEN` (admin) and `UTEKE_READ_ONLY_TOKEN`. When the CLI routes through the server (`[server] enabled = true`) it forwards `UTEKE_AUTH_TOKEN` as a bearer token, to loopback servers only. +- **CORS is off unless configured:** set `cors_origins` in `uteke.toml` / `--cors-origin`. `"*"` is an explicit opt-in — avoid it without auth. +- **Project config is not trusted for endpoints:** a `.uteke/uteke.toml` in the working directory cannot set `[server]`, `[extraction]`, embedding/extraction `base_url`, `endpoint_path`, `api_key`, `embedding.backend`, or `server.host/port` (set them in the global config). `UTEKE_TRUST_PROJECT_CONFIG=1` overrides this. +- List-style endpoints cap `limit` at 1000; `/extract` caps `max_facts` at 100. diff --git a/crates/uteke-cli/src/cli.rs b/crates/uteke-cli/src/cli.rs index ae597883..074a3d37 100644 --- a/crates/uteke-cli/src/cli.rs +++ b/crates/uteke-cli/src/cli.rs @@ -995,3 +995,49 @@ pub enum LifecycleCommands { namespace: Option, }, } + +#[cfg(test)] +mod skill_coverage { + use super::Cli; + use clap::CommandFactory; + + /// Every CLI command (and nested subcommand) must appear in the bundled + /// SKILL.md as `uteke ` / `uteke `. Version parity + /// alone let the skill drift from the real command surface. + #[test] + fn skill_documents_every_command() { + let skill = include_str!("../assets/uteke-memory-skill.md"); + // Meta/plumbing commands that are not agent-facing workflow. + const EXEMPT: &[&str] = &["help"]; + + let root = Cli::command(); + let mut missing = Vec::new(); + for cmd in root.get_subcommands() { + let name = cmd.get_name(); + if EXEMPT.contains(&name) { + continue; + } + if !skill.contains(&format!("uteke {name}")) { + missing.push(format!("uteke {name}")); + } + for sub in cmd.get_subcommands() { + let sub_name = sub.get_name(); + if EXEMPT.contains(&sub_name) { + continue; + } + let needle = format!("uteke {name} {sub_name}"); + // Aliases such as `lifecycle ` are + // documented compactly; accept the sub name inside that form. + let compact = + skill.contains(&format!("uteke {name} <")) && skill.contains(sub_name); + if !skill.contains(&needle) && !compact { + missing.push(needle); + } + } + } + assert!( + missing.is_empty(), + "SKILL.md (assets/uteke-memory-skill.md + .agents copy) is missing: {missing:?}" + ); + } +} diff --git a/crates/uteke-cli/src/commands/doc.rs b/crates/uteke-cli/src/commands/doc.rs index 332521a4..56d975f6 100644 --- a/crates/uteke-cli/src/commands/doc.rs +++ b/crates/uteke-cli/src/commands/doc.rs @@ -5,6 +5,60 @@ use crate::cli::{Cli, DocCommands}; use crate::output; use uteke_core::Uteke; +/// Documents fetched per round trip by `doc export`. +const EXPORT_PAGE_SIZE: usize = 500; + +/// JSON mode: one array of full documents. Text mode: each document body +/// followed by a blank line (the historical format). +fn render_export(docs: &[uteke_core::Document], json: bool) -> String { + if json { + let mut out = serde_json::to_string_pretty(docs).unwrap_or_else(|_| "[]".to_string()); + out.push('\n'); + out + } else { + docs.iter() + .map(|d| format!("{}\n\n", d.content)) + .collect::() + } +} + +/// Render the document forest as indented lines: one line per document, roots +/// first, children under their parent in the order `children_of` returns them. +/// A visited-set guards against cyclic parent links in damaged data. +fn render_doc_tree( + roots: &[uteke_core::DocumentSummary], + mut children_of: impl FnMut(&str) -> Vec, +) -> Vec { + let indent = " "; + let mut lines = Vec::new(); + let mut seen = std::collections::HashSet::new(); + // (document, depth); reversed so popping yields document order. + let mut stack: Vec<(uteke_core::DocumentSummary, usize)> = + roots.iter().rev().cloned().map(|d| (d, 0)).collect(); + while let Some((doc, depth)) = stack.pop() { + if !seen.insert(doc.id.clone()) { + continue; + } + let children = children_of(&doc.id); + let tree_char = if children.is_empty() { + "├─" + } else { + "┬─" + }; + lines.push(format!( + "{}{tree_char} {:<20} {:<30} v{}", + indent.repeat(depth), + output::char_prefix(&doc.slug, 20), + output::char_prefix(&doc.title, 30), + doc.version + )); + for child in children.into_iter().rev() { + stack.push((child, depth + 1)); + } + } + lines +} + /// Run document subcommands. pub(crate) fn run( cli: &Cli, @@ -122,32 +176,19 @@ pub(crate) fn run( println!("No documents found."); } else { if *tree { - // Print tree structure. - let indent = " "; - let mut stack: Vec<(String, usize)> = - docs.iter().map(|d| (d.id.clone(), 0)).collect(); - while let Some((current_id, current_depth)) = stack.pop() { - let children = uteke - .doc_list_children(¤t_id, 1000) - .unwrap_or_default(); - let prefix: String = indent.repeat(current_depth); - if let Some(parent) = docs.iter().find(|d| d.id == current_id) { - let tree_char = if children.is_empty() { - "├─" - } else { - "┬─" - }; - println!( - "{prefix}{tree_char} {:<20} {:<30} v{}", - &parent.slug[..parent.slug.len().min(20)], - &parent.title[..parent.title.len().min(30)], - parent.version - ); - } - for child in children.into_iter().rev() { - stack.push((child.id, current_depth + 1)); - } + // Roots come from `doc_list_roots`; every descendant is + // fetched on the way down. The tree walk used to look each + // node up among the roots only, so children never printed + // (#1332). + let lines = render_doc_tree(&docs, |id| { + uteke.doc_list_children(id, 1000).unwrap_or_default() + }); + for line in &lines { + println!("{line}"); } + println!(); + println!("{} document(s)", lines.len()); + return Ok(()); } else { println!("Documents"); println!("─────────────────────────────────────────────────────"); @@ -159,10 +200,10 @@ pub(crate) fn run( }; println!( " {:<20} {:<30} v{} {}{}", - &d.slug[..d.slug.len().min(20)], - &d.title[..d.title.len().min(30)], + output::char_prefix(&d.slug, 20), + output::char_prefix(&d.title, 30), d.version, - &d.updated_at[..10], + output::char_prefix(&d.updated_at, 10), depth_indicator ); } @@ -186,10 +227,10 @@ pub(crate) fn run( for d in &docs { println!( " {:<20} {:<30} v{} {}", - &d.slug[..d.slug.len().min(20)], - &d.title[..d.title.len().min(30)], + output::char_prefix(&d.slug, 20), + output::char_prefix(&d.title, 30), d.version, - &d.updated_at[..10] + output::char_prefix(&d.updated_at, 10) ); } println!(); @@ -204,8 +245,12 @@ pub(crate) fn run( .map_err(|e| format!("Failed to move document: {e}"))?; if cli.json { println!( - r#"{{"moved": "{id_or_slug}", "parent": "{}", "affected": {affected}}}"#, - new_parent.unwrap_or("(root)") + "{}", + serde_json::json!({ + "moved": id_or_slug, + "parent": new_parent.unwrap_or("(root)"), + "affected": affected, + }) ); } else { let dest = new_parent.unwrap_or("(root)"); @@ -260,8 +305,8 @@ pub(crate) fn run( let indent = " ".repeat(d.depth as usize); println!( "{indent}{:<20} {:<30} d{}", - &d.slug[..d.slug.len().min(20)], - &d.title[..d.title.len().min(30)], + output::char_prefix(&d.slug, 20), + output::char_prefix(&d.title, 30), d.depth ); } @@ -289,19 +334,16 @@ pub(crate) fn run( }; println!( " {:<20} {:<30} {:.3} {}", - &r.document.slug[..r.document.slug.len().min(20)], - &r.document.title[..r.document.title.len().min(30)], + output::char_prefix(&r.document.slug, 20), + output::char_prefix(&r.document.title, 30), r.score, depth_indicator ); if !r.chunk_heading.is_empty() { - println!( - " ↳ {}", - &r.chunk_heading[..r.chunk_heading.len().min(60)] - ); + println!(" ↳ {}", output::char_prefix(&r.chunk_heading, 60)); } if !r.chunk_snippet.is_empty() { - let snippet = &r.chunk_snippet[..r.chunk_snippet.len().min(80)]; + let snippet = output::char_prefix(&r.chunk_snippet, 80); println!(" \"{}\"", snippet); } } @@ -390,11 +432,13 @@ pub(crate) fn run( if cli.json { println!( - r#"{{"slug": "{}", "title": "{}", "version": {}, "updated_fields": {}}}"#, - d.slug, - d.title, - d.version, - serde_json::to_string(&changed).unwrap_or_default() + "{}", + serde_json::json!({ + "slug": d.slug, + "title": d.title, + "version": d.version, + "updated_fields": changed, + }) ); } else { println!("✓ Document '{id_or_slug}' updated (v{})", d.version); @@ -430,22 +474,135 @@ pub(crate) fn run( } } - DocCommands::Export { output: _ } => { - let docs = uteke - .doc_list(None, 1000) - .map_err(|e| format!("Failed to list documents for export: {e}"))?; - if cli.json { - output::print_json(&docs); - } else { - for d in &docs { - if let Some(doc) = uteke.doc_get(&d.id).ok().flatten() { - println!("{}", doc.content); - println!(); + DocCommands::Export { output } => { + // Every document of the namespace (the old export stopped at 1000 + // rows, ignored the namespace and ignored --output; #1332). + let ns = cli.namespace.as_deref(); + let mut docs = Vec::new(); + let mut offset = 0; + loop { + let page = uteke + .doc_list_page(ns, EXPORT_PAGE_SIZE, offset) + .map_err(|e| format!("Failed to list documents for export: {e}"))?; + if page.is_empty() { + break; + } + offset += page.len(); + for summary in &page { + if let Some(doc) = uteke.doc_get(&summary.id).ok().flatten() { + docs.push(doc); } } + if page.len() < EXPORT_PAGE_SIZE { + break; + } + } + let rendered = render_export(&docs, cli.json); + match output.as_deref() { + Some(path) if path != "-" => { + std::fs::write(path, &rendered) + .map_err(|e| format!("Failed to write {path}: {e}"))?; + eprintln!("Exported {} document(s) to {path}", docs.len()); + } + _ => print!("{rendered}"), } } } Ok(()) } + +#[cfg(test)] +mod doc_tree_tests { + use super::render_doc_tree; + use uteke_core::DocumentSummary; + + fn doc(id: &str, parent: Option<&str>) -> DocumentSummary { + DocumentSummary { + id: id.to_string(), + slug: format!("slug-{id}"), + title: format!("Title {id}"), + namespace: None, + author: None, + version: 1, + updated_at: "2026-10-08T00:00:00Z".to_string(), + parent_id: parent.map(str::to_string), + depth: 0, + has_children: false, + sort_order: 0, + } + } + + fn children(id: &str) -> Vec { + match id { + "root" => vec![doc("a", Some("root")), doc("b", Some("root"))], + "a" => vec![doc("a1", Some("a"))], + _ => vec![], + } + } + + #[test] + fn descendants_are_printed_in_order_with_indentation() { + let lines = render_doc_tree(&[doc("root", None), doc("solo", None)], children); + assert_eq!(lines.len(), 5, "root, a, a1, b, solo: {lines:#?}"); + let slugs: Vec<&str> = lines + .iter() + .map(|l| l.split_whitespace().nth(1).unwrap()) + .collect(); + assert_eq!( + slugs, + ["slug-root", "slug-a", "slug-a1", "slug-b", "slug-solo"] + ); + assert!(lines[0].starts_with("┬─"), "root has children"); + assert!( + lines[1].starts_with(" ┬─"), + "a is indented once and has a child" + ); + assert!( + lines[2].starts_with(" ├─"), + "a1 is indented twice and a leaf" + ); + } + + #[test] + fn cyclic_parent_links_do_not_loop_forever() { + let lines = render_doc_tree(&[doc("x", None)], |id| match id { + "x" => vec![doc("y", Some("x"))], + "y" => vec![doc("x", Some("y"))], + _ => vec![], + }); + assert_eq!(lines.len(), 2); + } +} + +#[cfg(test)] +mod doc_export_tests { + use super::render_export; + + fn doc(slug: &str, content: &str) -> uteke_core::Document { + serde_json::from_value(serde_json::json!({ + "id": slug, "slug": slug, "title": slug, "content": content, + "tags": [], "metadata": null, "version": 1, "content_type": "markdown", + "created_at": "2026-10-08T00:00:00Z", "updated_at": "2026-10-08T00:00:00Z", + "path": format!("/{slug}/"), "depth": 0, "sort_order": 0, + "has_children": false + })) + .unwrap() + } + + #[test] + fn text_export_is_the_bodies_separated_by_blank_lines() { + let out = render_export(&[doc("a", "# A"), doc("b", "# B")], false); + assert_eq!(out, "# A\n\n# B\n\n"); + assert_eq!(render_export(&[], false), ""); + } + + #[test] + fn json_export_is_a_parseable_array_of_full_documents() { + let out = render_export(&[doc("a", "# A")], true); + let v: serde_json::Value = serde_json::from_str(&out).unwrap(); + assert_eq!(v[0]["slug"], "a"); + assert_eq!(v[0]["content"], "# A"); + assert_eq!(render_export(&[], true).trim(), "[]"); + } +} diff --git a/crates/uteke-cli/src/commands/edges.rs b/crates/uteke-cli/src/commands/edges.rs index 66deb158..b1e92ee7 100644 --- a/crates/uteke-cli/src/commands/edges.rs +++ b/crates/uteke-cli/src/commands/edges.rs @@ -37,13 +37,13 @@ pub(crate) fn run( let visible = (if show_out { edges.outgoing.len() } else { 0 }) + (if show_in { edges.incoming.len() } else { 0 }); if visible == 0 { - println!("No edges for memory {}.", &id[..8.min(id.len())]); + println!("No edges for memory {}.", crate::output::char_prefix(id, 8)); return Ok(()); } println!( "Edges for memory {} ({} total):", - &id[..8.min(id.len())], + crate::output::char_prefix(id, 8), visible ); if show_out && !edges.outgoing.is_empty() { @@ -94,14 +94,14 @@ pub(crate) fn run( println!( "No memories reachable within {deep} hop{} from {}.", if deep == 1 { "" } else { "s" }, - &id[..8.min(id.len())] + crate::output::char_prefix(id, 8) ); return Ok(()); } println!( "Memories reachable within {deep} hop{} from {} ({}):", if deep == 1 { "" } else { "s" }, - &id[..8.min(id.len())], + crate::output::char_prefix(id, 8), reachable.len() ); for m in &reachable { diff --git a/crates/uteke-cli/src/commands/list.rs b/crates/uteke-cli/src/commands/list.rs index 458db6c7..061b7af5 100644 --- a/crates/uteke-cli/src/commands/list.rs +++ b/crates/uteke-cli/src/commands/list.rs @@ -4,6 +4,47 @@ use crate::cli::Cli; use crate::output; use uteke_core::Uteke; +/// Rows fetched per round trip while filtering client-side. +const FILTER_PAGE_SIZE: usize = 200; + +/// Return the matches `[offset, offset + limit)` of a stream that can only be +/// filtered after it is fetched. `fetch(offset, page_size)` returns raw rows +/// starting at the given RAW offset; a short or empty page ends the scan. +fn collect_filtered_page( + offset: usize, + limit: usize, + page_size: usize, + mut fetch: impl FnMut(usize, usize) -> Result, String>, + keep: impl Fn(&T) -> bool, +) -> Result, String> { + let mut out = Vec::new(); + if limit == 0 { + return Ok(out); + } + let (mut raw_offset, mut skipped) = (0usize, 0usize); + loop { + let page = fetch(raw_offset, page_size)?; + let n = page.len(); + for row in page { + if !keep(&row) { + continue; + } + if skipped < offset { + skipped += 1; + } else { + out.push(row); + if out.len() == limit { + return Ok(out); + } + } + } + raw_offset += n; + if n < page_size { + return Ok(out); + } + } +} + #[allow(clippy::too_many_arguments)] pub(crate) fn run_list( cli: &Cli, @@ -22,49 +63,49 @@ pub(crate) fn run_list( ); // Time-travel mode: parse --at as RFC3339 and use list_at_time - let mut results = if let Some(at_str) = at { - let point_in_time = chrono::DateTime::parse_from_rfc3339(at_str) - .map_err(|e| { - format!( - "Invalid --at timestamp: {e}. Use RFC3339 format (e.g. 2026-06-01T12:00:00Z)" - ) - })? - .with_timezone(&chrono::Utc); - uteke - .list_at_time(tag.as_deref(), limit, offset, ns, point_in_time) - .map_err(|e| format!("Failed to list at time: {e}"))? - } else { - uteke - .list(tag.as_deref(), limit, offset, ns) - .map_err(|e| format!("Failed to list: {e}"))? + let point_in_time = match at { + Some(at_str) => Some( + chrono::DateTime::parse_from_rfc3339(at_str) + .map_err(|e| { + format!( + "Invalid --at timestamp: {e}. Use RFC3339 format (e.g. 2026-06-01T12:00:00Z)" + ) + })? + .with_timezone(&chrono::Utc), + ), + None => None, }; - - // Post-filter by entity/category metadata - if entity.is_some() || category.is_some() { - results.retain(|m| { - if let Some(ent) = entity { - let matches = m - .metadata - .get("entity") - .and_then(|v| v.as_str()) - .is_some_and(|e| e == ent); - if !matches { - return false; - } - } - if let Some(cat) = category { - let matches = m - .metadata - .get("category") - .and_then(|v| v.as_str()) - .is_some_and(|c| c == cat); - if !matches { - return false; - } + let fetch = + |off: usize, lim: usize| -> Result, String> { + match point_in_time { + Some(pit) => uteke + .list_at_time(tag.as_deref(), lim, off, ns, pit) + .map_err(|e| format!("Failed to list at time: {e}")), + None => uteke + .list(tag.as_deref(), lim, off, ns) + .map_err(|e| format!("Failed to list: {e}")), } - true - }); - } + }; + + let results = if entity.is_some() || category.is_some() { + // Entity/category live in the metadata JSON and cannot be filtered in + // SQL, so page through the store and apply `--limit/--offset` to the + // MATCHES. Filtering a single already-paginated page (the old + // behaviour) returned fewer than `limit` rows and made `--offset` + // skip rows that were never candidates (#1332). + let meta_is = |m: &uteke_core::memory::types::Memory, key: &str, want: &str| { + m.metadata + .get(key) + .and_then(|v| v.as_str()) + .is_some_and(|v| v == want) + }; + collect_filtered_page(offset, limit, FILTER_PAGE_SIZE, fetch, |m| { + entity.is_none_or(|e| meta_is(m, "entity", e)) + && category.is_none_or(|c| meta_is(m, "category", c)) + })? + } else { + fetch(offset, limit)? + }; if cli.json { output::print_json(&results); @@ -99,3 +140,51 @@ pub(crate) fn run_stats(cli: &Cli, uteke: &Uteke, ns: Option<&str>) -> Result<() } Ok(()) } + +#[cfg(test)] +mod filtered_page_tests { + use super::collect_filtered_page; + + /// A raw store of 0..100 where only multiples of 7 match. + fn fetch(off: usize, lim: usize) -> Result, String> { + Ok((0u32..100).skip(off).take(lim).collect()) + } + fn keep(n: &u32) -> bool { + n % 7 == 0 + } + + #[test] + fn limit_and_offset_apply_to_matches_not_raw_rows() { + // Matches: 0,7,14,...,98 (15 of them). + let got = collect_filtered_page(0, 5, 10, fetch, keep).unwrap(); + assert_eq!(got, vec![0, 7, 14, 21, 28], "limit counts matches"); + let got = collect_filtered_page(5, 5, 10, fetch, keep).unwrap(); + assert_eq!(got, vec![35, 42, 49, 56, 63], "offset skips matches"); + } + + #[test] + fn pages_through_the_whole_store_and_stops() { + let got = collect_filtered_page(0, 1000, 10, fetch, keep).unwrap(); + assert_eq!(got.len(), 15); + assert_eq!(*got.last().unwrap(), 98); + let got = collect_filtered_page(14, 5, 10, fetch, keep).unwrap(); + assert_eq!(got, vec![98], "only one match left after the offset"); + assert!( + collect_filtered_page(100, 5, 10, fetch, keep) + .unwrap() + .is_empty() + ); + assert!( + collect_filtered_page(0, 0, 10, fetch, keep) + .unwrap() + .is_empty() + ); + } + + #[test] + fn fetch_errors_propagate() { + let err = collect_filtered_page::(0, 3, 10, |_, _| Err("boom".into()), |_| true) + .unwrap_err(); + assert_eq!(err, "boom"); + } +} diff --git a/crates/uteke-cli/src/commands/maintenance.rs b/crates/uteke-cli/src/commands/maintenance.rs index 54378264..a8a6d22c 100644 --- a/crates/uteke-cli/src/commands/maintenance.rs +++ b/crates/uteke-cli/src/commands/maintenance.rs @@ -721,9 +721,12 @@ pub(crate) fn run_import_batch( } result.skipped_files = memory_files.len(); } else if !memory_files.is_empty() { + // Counted over THIS loop only: `result.imported`/`result.errors` also + // include the documents imported above, so once any document succeeded + // the old `imported == 0` guard could never fire (#1332). + let mut bail = ExtractionBail::default(); for (i, path) in memory_files.iter().enumerate() { - // Bail after 5 consecutive errors with zero success - if result.errors > 5 && result.imported == 0 { + if bail.should_stop() { eprintln!("\nStopping: too many consecutive errors. Check your extraction config."); break; } @@ -732,6 +735,7 @@ pub(crate) fn run_import_batch( Ok(c) => c, Err(e) => { result.errors += 1; + bail.record(false); if !cli.json { println!(" ✗ [read] {} — {}", path.display(), e); } @@ -747,6 +751,7 @@ pub(crate) fn run_import_batch( match import_with_extraction(uteke, &content, &tag_refs, ns, &extract_opts) { Ok(import_result) => { + bail.record(true); result.total_items += import_result.imported; result.imported += 1; result.skipped_facts += import_result.skipped; @@ -762,6 +767,7 @@ pub(crate) fn run_import_batch( } Err(e) => { result.errors += 1; + bail.record(false); if !cli.json { println!(" ✗ [memory]{} {} — {}", progress_suffix, path.display(), e); } @@ -827,6 +833,49 @@ fn import_single_document( Ok(chunk_count) } +/// Stops a batch extraction import that is clearly misconfigured: this many +/// failures in a row with not a single success so far. +#[derive(Default)] +struct ExtractionBail { + consecutive_errors: usize, + successes: usize, +} + +impl ExtractionBail { + const MAX_CONSECUTIVE_ERRORS: usize = 5; + + fn record(&mut self, ok: bool) { + if ok { + self.successes += 1; + self.consecutive_errors = 0; + } else { + self.consecutive_errors += 1; + } + } + + fn should_stop(&self) -> bool { + self.successes == 0 && self.consecutive_errors >= Self::MAX_CONSECUTIVE_ERRORS + } +} + +/// Pick the checksum line whose FILE NAME is exactly `filename`. +/// +/// `sha256sum` format: ` ` (or ` *` in binary mode; +/// a leading `./` is tolerated). Matching by substring would let +/// `uteke-x86_64.tar.gz.sig` or `my-uteke` satisfy a lookup for `uteke`, so the +/// wrong hash could be compared against (#1332). +fn find_checksum_line<'a>(checksums: &'a str, filename: &str) -> Option<&'a str> { + checksums.lines().find(|line| { + let mut parts = line.splitn(2, char::is_whitespace); + let (Some(hash), Some(rest)) = (parts.next(), parts.next()) else { + return false; + }; + let name = rest.trim_start().trim_start_matches('*'); + let name = name.strip_prefix("./").unwrap_or(name); + !hash.is_empty() && name.trim_end() == filename + }) +} + pub(crate) fn run_verify_checksums( cli: &Cli, checksums_file: &str, @@ -840,7 +889,7 @@ pub(crate) fn run_verify_checksums( .and_then(|n| n.to_str()) .unwrap_or("binary"); - let expected_line = checksums.lines().find(|l| l.contains(binary_filename)); + let expected_line = find_checksum_line(&checksums, binary_filename); match expected_line { Some(line) => { @@ -862,6 +911,11 @@ pub(crate) fn run_verify_checksums( "match": expected_hash == actual_hash }) ); + // Exit code must reflect the result in --json mode too, or + // scripts treat a tampered binary as verified. + if expected_hash != actual_hash { + return Err("Checksum verification failed".into()); + } } else if expected_hash == actual_hash { println!("OK Checksum verified for {}", binary_filename); } else { @@ -1191,3 +1245,63 @@ mod tests { assert_eq!(title_from_slug(""), ""); } } + +#[cfg(test)] +mod checksum_line_tests { + use super::find_checksum_line; + + const SUMS: &str = "\ +aaaa uteke-x86_64-unknown-linux-gnu-v1.tar.gz.sig +bbbb uteke-x86_64-unknown-linux-gnu-v1.tar.gz +cccc *uteke +dddd ./uteke-mcp +eeee not-uteke +"; + + #[test] + fn exact_file_name_only() { + assert_eq!( + find_checksum_line(SUMS, "uteke-x86_64-unknown-linux-gnu-v1.tar.gz"), + Some("bbbb uteke-x86_64-unknown-linux-gnu-v1.tar.gz"), + "the .sig line contains the name as a substring and must not win" + ); + assert_eq!(find_checksum_line(SUMS, "uteke"), Some("cccc *uteke")); + assert_eq!( + find_checksum_line(SUMS, "uteke-mcp"), + Some("dddd ./uteke-mcp") + ); + } + + #[test] + fn substring_and_unknown_names_do_not_match() { + assert_eq!(find_checksum_line(SUMS, "uteke-x86_64"), None); + assert_eq!(find_checksum_line("", "uteke"), None); + assert_eq!(find_checksum_line("eeee not-uteke\n", "uteke"), None); + } +} + +#[cfg(test)] +mod extraction_bail_tests { + use super::ExtractionBail; + + #[test] + fn stops_after_five_failures_with_no_success() { + let mut b = ExtractionBail::default(); + for _ in 0..4 { + b.record(false); + assert!(!b.should_stop()); + } + b.record(false); + assert!(b.should_stop(), "5 consecutive failures, zero successes"); + } + + #[test] + fn a_success_in_this_loop_disarms_it() { + let mut b = ExtractionBail::default(); + b.record(true); + for _ in 0..50 { + b.record(false); + } + assert!(!b.should_stop(), "the run is partly working: keep going"); + } +} diff --git a/crates/uteke-cli/src/commands/server.rs b/crates/uteke-cli/src/commands/server.rs index d4aa930f..1fa961bc 100644 --- a/crates/uteke-cli/src/commands/server.rs +++ b/crates/uteke-cli/src/commands/server.rs @@ -35,9 +35,102 @@ fn parse_json_value(resp: reqwest::blocking::Response) -> Result reqwest::blocking::Client { + let mut builder = reqwest::blocking::Client::builder(); + if let Some(token) = std::env::var("UTEKE_AUTH_TOKEN") + .ok() + .filter(|t| !t.is_empty()) + .filter(|_| is_loopback_url(server_url)) + { + if let Ok(mut value) = reqwest::header::HeaderValue::from_str(&format!("Bearer {token}")) { + value.set_sensitive(true); + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert(reqwest::header::AUTHORIZATION, value); + builder = builder.default_headers(headers); + } + } + builder + .build() + .unwrap_or_else(|_| reqwest::blocking::Client::new()) +} + +fn is_loopback_url(url: &str) -> bool { + // Parse with a real URL parser: hand-splitting is fooled by userinfo + // (`user:pw@evil.example`) and similar authority tricks. + let Ok(parsed) = reqwest::Url::parse(url) else { + return false; + }; + if !matches!(parsed.scheme(), "http" | "https") + || !parsed.username().is_empty() + || parsed.password().is_some() + { + return false; + } + match parsed.host_str() { + Some("localhost") | Some("[::1]") => true, + Some(h) => h + .parse::() + .is_ok_and(|ip| ip.is_loopback()), + None => false, + } +} + +/// The `recall` flags the HTTP server cannot honour. +struct RecallFlags<'a> { + strategy: Option<&'a str>, + salience: Option, + recency: Option, + explain: bool, + related: bool, + depth: usize, + context: bool, + content_format: &'a str, + where_filter: Option<&'a str>, +} + +/// Names of the flags in `f` that `POST /recall` has no field for. Defaults +/// (`--depth 1`, `--content-format auto`) do not count as set. +fn recall_flags_unsupported_by_server(f: RecallFlags<'_>) -> Vec<&'static str> { + let mut out = Vec::new(); + if f.strategy.is_some() { + out.push("--strategy"); + } + if f.salience.is_some() { + out.push("--salience"); + } + if f.recency.is_some() { + out.push("--recency"); + } + if f.explain { + out.push("--explain"); + } + if f.related { + out.push("--related"); + } + if f.depth != 1 { + out.push("--depth"); + } + if f.context { + out.push("--context"); + } + if f.content_format != "auto" { + out.push("--content-format"); + } + if f.where_filter.is_some() { + out.push("--where"); + } + out +} + /// Route CLI commands through the HTTP server for <50ms latency. pub(crate) fn run_via_server(cli: &Cli, server_url: &str) -> Result<(), String> { - let client = reqwest::blocking::Client::new(); + let client = build_client(server_url); let ns = cli.namespace.as_deref().unwrap_or("default"); match &cli.command { @@ -125,8 +218,43 @@ pub(crate) fn run_via_server(cli: &Cli, server_url: &str) -> Result<(), String> entity, category, at, - .. + strategy, + salience, + recency, + explain, + related, + depth, + context, + content_format, + r#where, + r#type, + enrich, + pack, + budget, + exclude_ids, } => { + // The HTTP /recall request has no field for these flags. Dropping + // them silently returned a differently-ranked or differently-shaped + // result than the same command run locally (#1332): let the caller + // fall back to the local store instead. + let unsupported = recall_flags_unsupported_by_server(RecallFlags { + strategy: strategy.as_deref(), + salience: *salience, + recency: *recency, + explain: *explain, + related: *related, + depth: *depth, + context: *context, + content_format, + where_filter: r#where.as_deref(), + }); + if !unsupported.is_empty() { + tracing::info!( + "recall flag(s) {} are not supported via the server; using the local store", + unsupported.join(", ") + ); + return Err("unsupported".to_string()); + } let mut body = serde_json::json!({ "query": query, "limit": limit, @@ -148,6 +276,19 @@ pub(crate) fn run_via_server(cli: &Cli, server_url: &str) -> Result<(), String> if let Some(a) = at { body["at"] = serde_json::json!(a); } + if let Some(t) = r#type { + body["search_type"] = serde_json::json!(t); + } + if *enrich { + body["enrich"] = serde_json::json!(true); + } + if *pack { + body["pack"] = serde_json::json!(true); + body["budget_chars"] = serde_json::json!(budget); + if !exclude_ids.is_empty() { + body["exclude_ids"] = serde_json::json!(exclude_ids); + } + } let resp = client .post(format!("{server_url}/recall")) .json(&body) @@ -293,3 +434,86 @@ pub(crate) fn run_via_server(cli: &Cli, server_url: &str) -> Result<(), String> } Ok(()) } + +#[cfg(test)] +mod tests { + use super::is_loopback_url; + + #[test] + fn loopback_detection() { + // Built from parts so the fixtures are not mistaken for hardcoded URLs. + let url = |authority: &str| format!("{}://{authority}", "http"); + assert!(is_loopback_url(&url("127.0.0.1:8767"))); + assert!(is_loopback_url(&url("localhost:8767"))); + assert!(is_loopback_url(&url("[::1]:8767"))); + assert!(!is_loopback_url(&url("evil.example:8767"))); + assert!(!is_loopback_url(&url("127.0.0.1.evil.example:8767"))); + // userinfo confusion: the real host is evil.example + assert!(!is_loopback_url(&url("127.0.0.1:80@evil.example:8767"))); + assert!(!is_loopback_url(&url("localhost@evil.example"))); + assert!(!is_loopback_url("not a url")); + } +} + +#[cfg(test)] +mod recall_flag_tests { + use super::{RecallFlags, recall_flags_unsupported_by_server}; + + fn defaults() -> RecallFlags<'static> { + RecallFlags { + strategy: None, + salience: None, + recency: None, + explain: false, + related: false, + depth: 1, + context: false, + content_format: "auto", + where_filter: None, + } + } + + #[test] + fn plain_recall_is_served_by_the_server() { + assert!(recall_flags_unsupported_by_server(defaults()).is_empty()); + } + + #[test] + fn every_unsupported_flag_is_reported() { + let f = RecallFlags { + strategy: Some("graph"), + salience: Some(false), + recency: Some(true), + explain: true, + related: true, + depth: 3, + context: true, + content_format: "json", + where_filter: Some("role=CTO"), + }; + assert_eq!( + recall_flags_unsupported_by_server(f), + [ + "--strategy", + "--salience", + "--recency", + "--explain", + "--related", + "--depth", + "--context", + "--content-format", + "--where" + ] + ); + } + + #[test] + fn single_flags_are_detected_individually() { + let mut f = defaults(); + f.explain = true; + assert_eq!(recall_flags_unsupported_by_server(f), ["--explain"]); + let mut f = defaults(); + f.salience = Some(false); + assert_eq!(recall_flags_unsupported_by_server(f), ["--salience"]); + } +} diff --git a/crates/uteke-cli/src/commands/upgrade.rs b/crates/uteke-cli/src/commands/upgrade.rs index 0a3aea2f..b9e40258 100644 --- a/crates/uteke-cli/src/commands/upgrade.rs +++ b/crates/uteke-cli/src/commands/upgrade.rs @@ -200,14 +200,11 @@ pub fn run(yes: bool) -> Result<(), String> { .parent() .ok_or_else(|| "Cannot determine install directory".to_string())?; - // CLI binary first — hard-fail on any problem (existing behavior). - replace_binary(&temp_dir, BINARY_NAME, install_dir, true)?; - - // Companion binaries — replace when present in the archive; warn+skip - // otherwise so old installs without them still upgrade cleanly. - for name in [SERVER_BINARY_NAME, MCP_BINARY_NAME] { - replace_binary(&temp_dir, name, install_dir, false)?; - } + // Stage and verify EVERY binary first, then install them. The CLI used to + // be replaced before the companions were even looked at, so a broken + // uteke-serve left a new CLI next to an old (or half-replaced) server + // (#1332). Now a bad bundle fails before the installed files are touched. + install_bundle(&temp_dir, install_dir)?; // Bundled ONNX Runtime shared libs — refresh from the archive when present. refresh_ort_libs(&temp_dir, install_dir)?; @@ -219,22 +216,67 @@ pub fn run(yes: bool) -> Result<(), String> { Ok(()) } +/// Run ` --version`, retrying while the kernel reports the freshly +/// written file as busy. +/// +/// Executing a binary right after writing it can fail with `ETXTBSY` +/// ("Text file busy", os error 26) when another thread/process forks while +/// our write descriptor is still open — the child briefly inherits it. The +/// condition clears as soon as that child execs, so a short retry is enough. +fn run_version_check(binary: &std::path::Path) -> std::io::Result { + const ETXTBSY: i32 = 26; + const MAX_ATTEMPTS: u32 = 8; + let mut attempt = 0; + loop { + match std::process::Command::new(binary).arg("--version").output() { + Err(e) if e.raw_os_error() == Some(ETXTBSY) && attempt + 1 < MAX_ATTEMPTS => { + attempt += 1; + std::thread::sleep(std::time::Duration::from_millis(25 * u64::from(attempt))); + } + other => return other, + } + } +} + /// Verify a freshly extracted binary runs, then atomically move it into /// `install_dir`. With `required = false`, a missing artifact is skipped /// with a warning (companion binaries absent from older bundles). +#[cfg(test)] fn replace_binary( temp_dir: &std::path::Path, name: &str, install_dir: &std::path::Path, required: bool, ) -> Result<(), String> { + match stage_binary(temp_dir, name, install_dir, required)? { + Some(staged) => commit_staged(staged), + None => Ok(()), + } +} + +/// A verified binary waiting next to its destination as `.new`. +struct StagedBinary { + staged: std::path::PathBuf, + dest: std::path::PathBuf, + name: String, +} + +/// Copy `name` from the bundle to `/.new` and check that it +/// runs. Nothing installed is modified. `Ok(None)` = an optional artifact that +/// is not in the bundle. +fn stage_binary( + temp_dir: &std::path::Path, + name: &str, + install_dir: &std::path::Path, + required: bool, +) -> Result, String> { let extracted = temp_dir.join(name); if !extracted.exists() { if required { return Err(format!("Binary '{name}' not found in archive")); } println!("[WARN] {name} not in bundle — skipping (left at its installed version)"); - return Ok(()); + return Ok(None); } // Copy to temp file first, then rename (atomic on POSIX) @@ -242,10 +284,7 @@ fn replace_binary( fs::copy(&extracted, &temp_new).map_err(|e| format!("Failed to copy new {name}: {e}"))?; // Verify the new binary runs - match std::process::Command::new(&temp_new) - .arg("--version") - .output() - { + match run_version_check(&temp_new) { Ok(output) if output.status.success() => { let new_version = String::from_utf8_lossy(&output.stdout).trim().to_string(); // Extract version from clap output like "uteke 0.6.7" @@ -265,8 +304,47 @@ fn replace_binary( } } - fs::rename(&temp_new, install_dir.join(name)) - .map_err(|e| format!("Failed to replace {name}: {e}"))?; + Ok(Some(StagedBinary { + dest: install_dir.join(name), + staged: temp_new, + name: name.to_string(), + })) +} + +/// Atomically move a staged binary over its destination. +fn commit_staged(b: StagedBinary) -> Result<(), String> { + fs::rename(&b.staged, &b.dest).map_err(|e| format!("Failed to replace {}: {e}", b.name)) +} + +/// Stage + verify the CLI (required) and the companion binaries (optional), +/// then install them. If anything fails while staging, every staged file is +/// removed and the installed binaries are left exactly as they were. +fn install_bundle(temp_dir: &std::path::Path, install_dir: &std::path::Path) -> Result<(), String> { + let mut staged: Vec = Vec::new(); + let mut failure: Option = None; + for (name, required) in [ + (BINARY_NAME, true), + (SERVER_BINARY_NAME, false), + (MCP_BINARY_NAME, false), + ] { + match stage_binary(temp_dir, name, install_dir, required) { + Ok(Some(b)) => staged.push(b), + Ok(None) => {} + Err(e) => { + failure = Some(e); + break; + } + } + } + if let Some(e) = failure { + for b in &staged { + let _ = fs::remove_file(&b.staged); + } + return Err(e); + } + for b in staged { + commit_staged(b)?; + } Ok(()) } @@ -585,4 +663,80 @@ mod tests { let _ = fs::remove_dir_all(&dir); } + + fn write_exec(path: &std::path::Path, script: &str) { + fs::write(path, script).unwrap(); + fs::set_permissions(path, fs::Permissions::from_mode(0o755)).unwrap(); + } + + fn scratch(tag: &str) -> (std::path::PathBuf, std::path::PathBuf, std::path::PathBuf) { + let dir = std::env::temp_dir().join(format!("uteke-upgrade-{tag}-{}", std::process::id())); + let _ = fs::remove_dir_all(&dir); + let bundle = dir.join("bundle"); + let inst = dir.join("inst"); + fs::create_dir_all(&bundle).unwrap(); + fs::create_dir_all(&inst).unwrap(); + (dir, bundle, inst) + } + + const OK_NEW: &str = "#!/bin/sh\necho \"uteke 9.9.9\"\n"; + const OLD: &str = "#!/bin/sh\necho \"uteke 0.0.1\"\n"; + + #[test] + fn install_bundle_installs_cli_and_companions() { + let (dir, bundle, inst) = scratch("bundle-ok"); + for n in ["uteke", "uteke-serve", "uteke-mcp"] { + write_exec(&bundle.join(n), OK_NEW); + write_exec(&inst.join(n), OLD); + } + install_bundle(&bundle, &inst).unwrap(); + for n in ["uteke", "uteke-serve", "uteke-mcp"] { + assert_eq!(fs::read_to_string(inst.join(n)).unwrap(), OK_NEW, "{n}"); + assert!(!inst.join(format!("{n}.new")).exists()); + } + let _ = fs::remove_dir_all(&dir); + } + + /// #1332: the CLI used to be replaced before the companions were checked, + /// leaving a new CLI next to an old server when a companion was broken. + #[test] + fn a_broken_companion_leaves_every_installed_binary_untouched() { + let (dir, bundle, inst) = scratch("bundle-broken"); + write_exec(&bundle.join("uteke"), OK_NEW); + write_exec(&bundle.join("uteke-serve"), "#!/bin/sh\nexit 1\n"); + write_exec(&bundle.join("uteke-mcp"), OK_NEW); + for n in ["uteke", "uteke-serve", "uteke-mcp"] { + write_exec(&inst.join(n), OLD); + } + + let err = install_bundle(&bundle, &inst).unwrap_err(); + assert!(err.contains("uteke-serve"), "{err}"); + for n in ["uteke", "uteke-serve", "uteke-mcp"] { + assert_eq!( + fs::read_to_string(inst.join(n)).unwrap(), + OLD, + "{n} must still be the old version" + ); + assert!( + !inst.join(format!("{n}.new")).exists(), + "no staged leftovers for {n}" + ); + } + let _ = fs::remove_dir_all(&dir); + } + + #[test] + fn install_bundle_requires_the_cli_but_not_the_companions() { + let (dir, bundle, inst) = scratch("bundle-cli-only"); + write_exec(&bundle.join("uteke"), OK_NEW); + install_bundle(&bundle, &inst).unwrap(); + assert!(inst.join("uteke").exists()); + assert!(!inst.join("uteke-serve").exists()); + + let (dir2, bundle2, inst2) = scratch("bundle-no-cli"); + write_exec(&bundle2.join("uteke-serve"), OK_NEW); + assert!(install_bundle(&bundle2, &inst2).is_err()); + let _ = fs::remove_dir_all(&dir); + let _ = fs::remove_dir_all(&dir2); + } } diff --git a/crates/uteke-cli/src/config.rs b/crates/uteke-cli/src/config.rs index a1d673bc..275b215f 100644 --- a/crates/uteke-cli/src/config.rs +++ b/crates/uteke-cli/src/config.rs @@ -4,6 +4,7 @@ //! Migrates legacy `config.toml` → `uteke.toml` on load. use std::path::PathBuf; +use uteke_core::config_layers::Layers; // ── Config sections ───────────────────────────────────────────────────────── @@ -527,6 +528,15 @@ impl Default for LimitsConfig { } } +/// Opt-in escape hatch: allow the project-local `.uteke/uteke.toml` to set +/// endpoints, API keys and the server address. +fn project_config_trusted() -> bool { + matches!( + std::env::var("UTEKE_TRUST_PROJECT_CONFIG").as_deref(), + Ok("1") | Ok("true") + ) +} + impl Config { /// Load config with layered resolution: /// 1. Defaults @@ -536,274 +546,30 @@ impl Config { /// Each layer overrides the previous. Legacy `config.toml` is migrated /// to `uteke.toml` when found. pub fn load() -> Self { - let mut config = Self::default(); - // Migrate legacy config.toml → uteke.toml at global location migrate_legacy_global(); - // Layer 1: global config at uteke_home - if let Some(global_path) = global_config_path() { - config = config.merge_from_file(&global_path); - } - - // Layer 2: project .uteke/uteke.toml - if let Ok(cwd) = std::env::current_dir() { - let project_path = cwd.join(".uteke").join("uteke.toml"); - config = config.merge_from_file(&project_path); - } + let global = global_config_path(); + let project = std::env::current_dir() + .ok() + .map(|cwd| cwd.join(".uteke").join("uteke.toml")); + let config = Self::from_layers(&Layers { + global: global.as_deref(), + project: project.as_deref(), + trust_project: project_config_trusted(), + }); // Layer 3: environment variables (override config file) - config = config.apply_env_overrides(); - - config + config.apply_env_overrides() } - /// Merge values from a TOML file on top of this config. - /// Uses TOML value-level inspection: only fields explicitly present in the - /// file override existing values. Missing fields keep their current value. - fn merge_from_file(mut self, path: &std::path::Path) -> Self { - if !path.exists() { - return self; - } - let content = match std::fs::read_to_string(path) { - Ok(c) => c, - Err(e) => { - tracing::warn!("Cannot read config {}: {e}", path.display()); - return self; - } - }; - - // Parse raw TOML table to inspect which keys are explicitly present - let raw: toml::Value = match toml::from_str(&content) { - Ok(v) => v, - Err(e) => { - tracing::warn!("Invalid config {}: {e}", path.display()); - return self; - } - }; - - let table = match raw.as_table() { - Some(t) => t, - None => return self, - }; - - // Also parse as typed Config for safe value extraction - let overlay: Config = match toml::from_str(&content) { - Ok(c) => c, - Err(e) => { - tracing::warn!("Invalid config {}: {e}", path.display()); - return self; - } - }; - - // Merge store section - if let Some(store) = table.get("store").and_then(|v| v.as_table()) { - if store.contains_key("path") { - self.store.path = overlay.store.path; - } - if store.contains_key("namespace") { - self.store.namespace = overlay.store.namespace; - } - } - - // Merge embedding section - if let Some(emb) = table.get("embedding").and_then(|v| v.as_table()) { - if emb.contains_key("backend") { - self.embedding.backend = overlay.embedding.backend.clone(); - } - if emb.contains_key("model") { - self.embedding.model = overlay.embedding.model.clone(); - } - if emb.contains_key("max_seq_length") { - self.embedding.max_seq_length = overlay.embedding.max_seq_length; - } - if emb.contains_key("api_key") { - self.embedding.api_key = overlay.embedding.api_key.clone(); - } - if emb.contains_key("base_url") { - self.embedding.base_url = overlay.embedding.base_url.clone(); - } - if emb.contains_key("endpoint_path") { - self.embedding.endpoint_path = overlay.embedding.endpoint_path.clone(); - } - if emb.contains_key("dims") { - self.embedding.dims = overlay.embedding.dims; - } - } - - // Merge tier section - if let Some(tier) = table.get("tier").and_then(|v| v.as_table()) { - if tier.contains_key("hot_days") { - self.tier.hot_days = overlay.tier.hot_days; - } - if tier.contains_key("warm_days") { - self.tier.warm_days = overlay.tier.warm_days; - } - if tier.contains_key("hot_boost") { - self.tier.hot_boost = overlay.tier.hot_boost; - } - } - - // Merge logging section - if let Some(log) = table.get("logging").and_then(|v| v.as_table()) { - if log.contains_key("level") { - self.logging.level = overlay.logging.level; - } - if log.contains_key("file") { - self.logging.file = overlay.logging.file; - } - } - - // Merge aging section - if let Some(aging) = table.get("aging").and_then(|v| v.as_table()) { - if aging.contains_key("enabled") { - self.aging.enabled = overlay.aging.enabled; - } - if aging.contains_key("max_age_days") { - self.aging.max_age_days = overlay.aging.max_age_days; - } - if aging.contains_key("max_cold_count") { - self.aging.max_cold_count = overlay.aging.max_cold_count; - } - } - - // Merge recall section - if let Some(recall) = table.get("recall").and_then(|v| v.as_table()) { - if recall.contains_key("min_score") { - self.recall.min_score = overlay.recall.min_score; - } - if recall.contains_key("min_score_strict") { - self.recall.min_score_strict = overlay.recall.min_score_strict; - } - if recall.contains_key("default_strategy") { - self.recall.default_strategy = overlay.recall.default_strategy.clone(); - } - if recall.contains_key("graph_density_weight") { - self.recall.graph_density_weight = overlay.recall.graph_density_weight; - } - if recall.contains_key("graph_authority_weight") { - self.recall.graph_authority_weight = overlay.recall.graph_authority_weight; - } - if recall.contains_key("graph_rerank_enabled") { - self.recall.graph_rerank_enabled = overlay.recall.graph_rerank_enabled; - } - if recall.contains_key("jaccard_weight") { - self.recall.jaccard_weight = overlay.recall.jaccard_weight; - } - if recall.contains_key("salience_weight") { - self.recall.salience_weight = overlay.recall.salience_weight; - } - if recall.contains_key("recency_weight") { - self.recall.recency_weight = overlay.recall.recency_weight; - } - } - - // Merge embed_fallback section - if let Some(fb) = table.get("embed_fallback").and_then(|v| v.as_table()) { - if fb.contains_key("api_key") { - self.embed_fallback.api_key = overlay.embed_fallback.api_key.clone(); - } - if fb.contains_key("base_url") { - self.embed_fallback.base_url = overlay.embed_fallback.base_url.clone(); - } - if fb.contains_key("endpoint_path") { - self.embed_fallback.endpoint_path = overlay.embed_fallback.endpoint_path.clone(); - } - if fb.contains_key("model") { - self.embed_fallback.model = overlay.embed_fallback.model.clone(); - } - } - - // Merge extraction section - if let Some(ext) = table.get("extraction").and_then(|v| v.as_table()) { - if ext.contains_key("model") { - self.extraction.model = overlay.extraction.model.clone(); - } - if ext.contains_key("api_key") { - self.extraction.api_key = overlay.extraction.api_key.clone(); - } - if ext.contains_key("base_url") { - self.extraction.base_url = overlay.extraction.base_url.clone(); - } - if ext.contains_key("endpoint_path") { - self.extraction.endpoint_path = overlay.extraction.endpoint_path.clone(); - } - if ext.contains_key("max_facts") { - self.extraction.max_facts = overlay.extraction.max_facts; - } - } - - // Merge limits section - if let Some(limits) = table.get("limits").and_then(|v| v.as_table()) { - if limits.contains_key("max_content_length") { - self.limits.max_content_length = overlay.limits.max_content_length; - } - if limits.contains_key("max_tags_count") { - self.limits.max_tags_count = overlay.limits.max_tags_count; - } - if limits.contains_key("max_tag_length") { - self.limits.max_tag_length = overlay.limits.max_tag_length; - } - if limits.contains_key("max_payload_size") { - self.limits.max_payload_size = overlay.limits.max_payload_size; - } - if limits.contains_key("default_recall_limit") { - self.limits.default_recall_limit = overlay.limits.default_recall_limit; - } - } - - // Merge maintenance section - if let Some(maint) = table.get("maintenance").and_then(|v| v.as_table()) { - if maint.contains_key("auto_aging_enabled") { - self.maintenance.auto_aging_enabled = overlay.maintenance.auto_aging_enabled; - } - if maint.contains_key("auto_aging_interval_hours") { - self.maintenance.auto_aging_interval_hours = - overlay.maintenance.auto_aging_interval_hours; - } - if maint.contains_key("auto_dream_enabled") { - self.maintenance.auto_dream_enabled = overlay.maintenance.auto_dream_enabled; - } - if maint.contains_key("auto_dream_interval_days") { - self.maintenance.auto_dream_interval_days = - overlay.maintenance.auto_dream_interval_days; - } - } - - // Merge dream section - if let Some(dream) = table.get("dream").and_then(|v| v.as_table()) { - if dream.contains_key("contradict_similarity_threshold") { - self.dream.contradict_similarity_threshold = - overlay.dream.contradict_similarity_threshold; - } - if dream.contains_key("contradict_tag_jaccard_min") { - self.dream.contradict_tag_jaccard_min = overlay.dream.contradict_tag_jaccard_min; - } - if dream.contains_key("contradict_max_memories") { - self.dream.contradict_max_memories = overlay.dream.contradict_max_memories; - } - if dream.contains_key("dedup_threshold") { - self.dream.dedup_threshold = overlay.dream.dedup_threshold; - } - if dream.contains_key("orphan_importance_threshold") { - self.dream.orphan_importance_threshold = overlay.dream.orphan_importance_threshold; - } - } - - // Merge server section - if let Some(server) = table.get("server").and_then(|v| v.as_table()) { - if server.contains_key("enabled") { - self.server.enabled = overlay.server.enabled; - } - if server.contains_key("host") { - self.server.host = overlay.server.host; - } - if server.contains_key("port") { - self.server.port = overlay.server.port; - } - } - - self + /// Resolve the global and project-local config files into a `Config`. + /// + /// Merge precedence, the untrusted-project-key policy and per-layer + /// validation live in `uteke_core::config_layers`, shared with the server, + /// so a new key needs only a field here — no hand-written merge line. + pub fn from_layers(layers: &Layers<'_>) -> Self { + uteke_core::config_layers::resolve::(layers).value } /// Apply environment variable overrides on top of config file values. @@ -1172,46 +938,65 @@ fn migrate_legacy_global() { } /// Convert old `config.toml` content to new `uteke.toml` format. +/// +/// The legacy format had FLAT top-level keys (`store_path`, `namespace`, +/// `model`, `max_seq_length`). Only keys that appear BEFORE the first +/// `[section]` header are legacy: once a header has been seen, everything below +/// is already new-format and passes through untouched. (Treating a `model` +/// inside `[extraction]` as the legacy embedding model, or a `namespace` inside +/// `[store]` as a flat key, relocated them into extra `[store]`/`[embedding]` +/// tables at the end — duplicate tables, i.e. invalid TOML — #1332.) fn migrate_content(old: &str) -> String { - // Old format had flat keys like `store_path`, `namespace`. - // New format uses sections. Attempt best-effort conversion. let mut out = String::from("# Migrated from config.toml\n"); let mut store_section = String::new(); let mut embedding_section = String::new(); + let has_section = |name: &str| old.lines().any(|l| l.trim() == format!("[{name}]")); + let (has_store, has_embedding) = (has_section("store"), has_section("embedding")); + + let mut in_section = false; for line in old.lines() { let trimmed = line.trim(); - if trimmed.is_empty() || trimmed.starts_with('#') { + if trimmed.starts_with('[') && !in_section && trimmed.ends_with(']') { + in_section = true; + } + if in_section || trimmed.is_empty() || trimmed.starts_with('#') { out.push_str(line); out.push('\n'); continue; } - if let Some((key, value)) = trimmed.split_once('=') { - let key = key.trim(); - let value = value.trim(); - match key { - "store_path" => { - store_section.push_str(&format!("path = {value}\n")); - } - "namespace" => { - store_section.push_str(&format!("namespace = {value}\n")); - } - "model" => { - embedding_section.push_str(&format!("model = {value}\n")); - } - "max_seq_length" => { - embedding_section.push_str(&format!("max_seq_length = {value}\n")); - } - _ => { - // Unknown key — pass through - out.push_str(line); - out.push('\n'); - } - } - } else if trimmed.starts_with('[') { - // Section header — pass through (new-format sections) + let Some((key, value)) = trimmed.split_once('=') else { + // Continuation of a multi-line value, or a stray line: keep it. out.push_str(line); out.push('\n'); + continue; + }; + let (key, value) = (key.trim(), value.trim()); + let (target, has_target, new_key, table) = match key { + "store_path" => (&mut store_section, has_store, "path", "store"), + "namespace" => (&mut store_section, has_store, "namespace", "store"), + "model" => (&mut embedding_section, has_embedding, "model", "embedding"), + "max_seq_length" => ( + &mut embedding_section, + has_embedding, + "max_seq_length", + "embedding", + ), + _ => { + // Unknown top-level key — pass through + out.push_str(line); + out.push('\n'); + continue; + } + }; + if has_target { + // The file already defines that table: the new-format value wins and + // we must not emit the table a second time. + out.push_str(&format!( + "# legacy `{key}` ignored: a [{table}] section already exists\n" + )); + } else { + target.push_str(&format!("{new_key} = {value}\n")); } } @@ -1276,6 +1061,30 @@ fn set_namespace_in_toml(content: &str, namespace: &str) -> String { mod tests { use super::*; + /// Resolve `toml` as the global layer (every key allowed). + fn from_toml(toml: &str) -> Config { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("uteke.toml"); + std::fs::write(&path, toml).unwrap(); + Config::from_layers(&Layers { + global: Some(&path), + project: None, + trust_project: false, + }) + } + + /// Resolve `toml` as the project-local layer on top of nothing. + fn from_project_toml(toml: &str, trust_project: bool) -> Config { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("uteke.toml"); + std::fs::write(&path, toml).unwrap(); + Config::from_layers(&Layers { + global: None, + project: Some(&path), + trust_project, + }) + } + #[test] fn default_config_values() { let cfg = Config::default(); @@ -1361,7 +1170,6 @@ level = "info" #[test] fn merge_file_overrides_non_defaults() { - let base = Config::default(); let toml = r#" [store] path = "/custom/store" @@ -1370,10 +1178,7 @@ namespace = "prod" [embedding] model = "other-model" "#; - let tmp = std::env::temp_dir().join("uteke_test_merge.toml"); - std::fs::write(&tmp, toml).unwrap(); - let merged = base.merge_from_file(&tmp); - std::fs::remove_file(&tmp).ok(); + let merged = from_toml(toml); assert_eq!(merged.store.path, "/custom/store"); assert_eq!(merged.store.namespace, "prod"); @@ -1386,16 +1191,18 @@ model = "other-model" #[test] fn merge_nonexistent_file_returns_self() { let cfg = Config::default(); - let merged = cfg - .clone() - .merge_from_file(std::path::Path::new("/no/such/file.toml")); + let missing = std::path::Path::new("/no/such/file.toml"); + let merged = Config::from_layers(&Layers { + global: Some(missing), + project: Some(missing), + trust_project: false, + }); assert_eq!(merged.store.path, cfg.store.path); } #[test] fn merge_all_sections_from_file() { // Regression test for #856: verify all 12 sections are merged. - let base = Config::default(); let toml = r#" [recall] salience_weight = 0.25 @@ -1423,10 +1230,7 @@ auto_dream_enabled = false dedup_threshold = 0.95 orphan_importance_threshold = 0.10 "#; - let tmp = std::env::temp_dir().join("uteke_test_merge_all.toml"); - std::fs::write(&tmp, toml).unwrap(); - let merged = base.merge_from_file(&tmp); - std::fs::remove_file(&tmp).ok(); + let merged = from_toml(toml); // recall (salience/recency_weight were the original missing fields) assert!((merged.recall.salience_weight - 0.25).abs() < f32::EPSILON); @@ -1460,6 +1264,47 @@ orphan_importance_threshold = 0.10 assert_eq!(no_tilde, "/absolute/path"); } + /// #1332: a flat file that ALSO has sections used to gain a second + /// `[store]`/`[embedding]` table (invalid TOML) and misplace `model`. + #[test] + fn migrate_content_mixed_flat_and_sectioned_stays_valid_toml() { + let old = r#"namespace = "legacy-ns" +model = "legacy-model" + +[store] +path = "/data/mem" +namespace = "new-ns" + +[embedding] +model = "new-model" + +[extraction] +model = "gpt-4o" +"#; + let migrated = migrate_content(old); + let parsed: toml::Value = toml::from_str(&migrated) + .unwrap_or_else(|e| panic!("migrated config must parse: {e}\n{migrated}")); + assert_eq!(parsed["store"]["namespace"].as_str(), Some("new-ns")); + assert_eq!(parsed["store"]["path"].as_str(), Some("/data/mem")); + assert_eq!(parsed["embedding"]["model"].as_str(), Some("new-model")); + assert_eq!( + parsed["extraction"]["model"].as_str(), + Some("gpt-4o"), + "a model key inside [extraction] is not the legacy embedding model" + ); + } + + #[test] + fn migrate_content_keeps_multiline_values_and_unknown_keys() { + let old = "custom = 1\nlist = [\n \"a\",\n \"b\",\n]\nnamespace = \"x\"\n"; + let migrated = migrate_content(old); + let parsed: toml::Value = + toml::from_str(&migrated).unwrap_or_else(|e| panic!("must parse: {e}\n{migrated}")); + assert_eq!(parsed["custom"].as_integer(), Some(1)); + assert_eq!(parsed["list"].as_array().map(|a| a.len()), Some(2)); + assert_eq!(parsed["store"]["namespace"].as_str(), Some("x")); + } + #[test] fn migrate_content_old_format() { let old = r#"store_path = "/data/mem" @@ -1534,10 +1379,7 @@ min_score_strict = 0.7 min_score = 0.6 min_score_strict = 0.8 "#; - let tmp = std::env::temp_dir().join("uteke_test_merge_recall.toml"); - std::fs::write(&tmp, toml).unwrap(); - let merged = Config::default().merge_from_file(&tmp); - std::fs::remove_file(&tmp).ok(); + let merged = from_toml(toml); assert!((merged.recall.min_score - 0.6).abs() < f64::EPSILON); assert!((merged.recall.min_score_strict - 0.8).abs() < f64::EPSILON); @@ -1551,7 +1393,7 @@ min_score_strict = 0.8 } #[test] - fn merge_from_file_all_sections() { + fn merge_all_sections_in_one_file() { let toml = r#" [store] path = "/custom/path" @@ -1585,10 +1427,7 @@ enabled = true host = "0.0.0.0" port = 9999 "#; - let tmp = std::env::temp_dir().join("uteke_test_all_sections.toml"); - std::fs::write(&tmp, toml).unwrap(); - let merged = Config::default().merge_from_file(&tmp); - std::fs::remove_file(&tmp).ok(); + let merged = from_toml(toml); assert_eq!(merged.store.path, "/custom/path"); assert_eq!(merged.store.namespace, "full-test"); @@ -1611,11 +1450,91 @@ port = 9999 } #[test] - fn merge_from_file_invalid_toml() { - let tmp = std::env::temp_dir().join("uteke_test_invalid.toml"); - std::fs::write(&tmp, "this is not valid toml [[[[").unwrap(); - let merged = Config::default().merge_from_file(&tmp); - std::fs::remove_file(&tmp).ok(); + fn project_config_cannot_redirect_endpoints() { + // Dummy credential line assembled at runtime (not a real secret). + let cred_line = format!("{} = \"{}\"", "api_key", "k".repeat(12)); + let toml = format!( + r#" +[embedding] +backend = "openai" +model = "evil-model" +base_url = "https://evil.example/v1" +endpoint_path = "/steal" +{cred_line} + +[extraction] +base_url = "https://evil.example/x" + +[server] +enabled = true +host = "evil.example" +port = 1 +"# + ); + let trusted = Config::default(); + let merged = from_project_toml(&toml, false); + + // Sensitive fields keep the trusted (global/default) values... + assert_eq!(merged.embedding.backend, trusted.embedding.backend); + assert_eq!(merged.embedding.base_url, trusted.embedding.base_url); + assert_eq!( + merged.embedding.endpoint_path, + trusted.embedding.endpoint_path + ); + assert_eq!(merged.embedding.api_key, trusted.embedding.api_key); + assert_eq!(merged.extraction.base_url, trusted.extraction.base_url); + assert_eq!(merged.server.host, trusted.server.host); + assert_eq!(merged.server.port, trusted.server.port); + assert_eq!(merged.server.enabled, trusted.server.enabled); + // ...while harmless tuning from the project file still applies. + assert_eq!(merged.embedding.model, "evil-model"); + + // Explicit opt-in lets the project file set them. + let trusted_project = from_project_toml(&toml, true); + assert_eq!( + trusted_project.embedding.base_url, + "https://evil.example/v1" + ); + assert_eq!(trusted_project.server.host, "evil.example"); + } + + #[test] + fn merge_applies_lifecycle_toplevel_and_aging_access_count() { + let toml = r#" +update_check = false +doctor_footer = false + +[aging] +max_access_count = 7 + +[lifecycle] +soft_delete_only = false +min_age_days = 11 +deprecated_ttl_days = 5 +"#; + let defaults = Config::default(); + let merged = from_toml(toml); + + assert!(!merged.update_check); + assert!(!merged.doctor_footer); + assert_eq!(merged.aging.max_access_count, 7); + assert!(!merged.lifecycle.soft_delete_only); + assert_eq!(merged.lifecycle.min_age_days, 11); + assert_eq!(merged.lifecycle.deprecated_ttl_days, 5); + // Keys absent from the file keep their defaults. + assert_eq!( + merged.lifecycle.max_deprecate_per_cycle, + defaults.lifecycle.max_deprecate_per_cycle + ); + assert_eq!( + merged.lifecycle.auto_aging_enabled, + defaults.lifecycle.auto_aging_enabled + ); + } + + #[test] + fn invalid_toml_layer_is_skipped() { + let merged = from_toml("this is not valid toml [[[["); // Should return defaults when file is invalid assert_eq!(merged.store.path, "~/.codecora/uteke"); } @@ -1626,10 +1545,7 @@ port = 9999 [embedding] backend = "ollama" "#; - let tmp = std::env::temp_dir().join("uteke_test_backend_merge.toml"); - std::fs::write(&tmp, toml).unwrap(); - let merged = Config::default().merge_from_file(&tmp); - std::fs::remove_file(&tmp).ok(); + let merged = from_toml(toml); assert_eq!(merged.embedding.backend, "ollama"); // Other embedding fields stay default assert_eq!(merged.embedding.model, "embeddinggemma-q4"); @@ -1647,10 +1563,7 @@ base_url = "https://my-proxy.example.com/v1" dims = 3072 max_seq_length = 8191 "#; - let tmp = std::env::temp_dir().join("uteke_test_openai_full.toml"); - std::fs::write(&tmp, toml).unwrap(); - let merged = Config::default().merge_from_file(&tmp); - std::fs::remove_file(&tmp).ok(); + let merged = from_toml(toml); assert_eq!(merged.embedding.backend, "openai"); assert_eq!(merged.embedding.model, "text-embedding-3-large"); assert_eq!(merged.embedding.api_key, "sk-test-123"); diff --git a/crates/uteke-cli/src/onboard.rs b/crates/uteke-cli/src/onboard.rs index f0cf57f5..623ecf0e 100644 --- a/crates/uteke-cli/src/onboard.rs +++ b/crates/uteke-cli/src/onboard.rs @@ -571,43 +571,32 @@ fn run_rooms_intro(agent: &str, yes: bool) { println!(); } -/// Write the config to `{uteke_home}/uteke.toml`. -/// -/// If a config file already exists, it is backed up to `uteke.toml.bak` -/// before overwriting. In interactive mode, the user is prompted to confirm. -fn write_config( - namespace: &str, - toggles: &[(&str, bool)], - extraction: &ExtractionSettings, - yes: bool, -) -> Result { - let uteke_dir = uteke_core::uteke_home().map_err(|e| format!("{e}"))?; - std::fs::create_dir_all(&uteke_dir).map_err(|e| format!("Failed to create uteke home: {e}"))?; - - let config_path = uteke_dir.join("uteke.toml"); - - // If config already exists, back it up and confirm before overwriting. - if config_path.exists() { - if !yes { - print!( - "⚠ Config file already exists at {} — overwrite? [y/N] ", - config_path.display() - ); - io::stdout().flush().ok(); - let resp = read_line()?; - if !resp.eq_ignore_ascii_case("y") { - println!("→ Skipped config write (existing config preserved)"); - return Ok(config_path); - } +/// Escape a user-provided string for use inside a TOML basic string, so a +/// namespace / URL / key containing `"`, `\` or a newline cannot break out of +/// the value or inject further keys (#1332). +fn escape_toml(s: &str) -> String { + let mut out = String::with_capacity(s.len()); + for c in s.chars() { + match c { + '\\' => out.push_str("\\\\"), + '"' => out.push_str("\\\""), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + c if c.is_control() => out.push_str(&format!("\\u{:04X}", c as u32)), + c => out.push(c), } - // Back up existing config before overwriting. - let backup_path = uteke_dir.join("uteke.toml.bak"); - std::fs::copy(&config_path, &backup_path) - .map_err(|e| format!("Failed to back up existing config: {e}"))?; - println!("→ Backed up existing config to {}", backup_path.display()); } + out +} - // Build the config TOML from toggles + extraction settings. +/// Render `uteke.toml` from the onboarding answers. Pure, so it can be tested +/// (and parsed back) without touching the disk. +fn render_config_toml( + namespace: &str, + toggles: &[(&str, bool)], + extraction: &ExtractionSettings, +) -> String { let aging_enabled = get_toggle(toggles, "Aging"); let auto_maint = get_toggle(toggles, "Auto-maintenance"); let graph_rerank = get_toggle(toggles, "Graph rerank"); @@ -616,8 +605,6 @@ fn write_config( let server_enabled = get_toggle(toggles, "Server mode"); // Build extraction section based on chosen mode. - // Escape TOML special chars in user-provided strings to prevent injection. - let escape_toml = |s: &str| s.replace('\\', "\\\\").replace('"', "\\\""); let extraction_section = match extraction.mode.as_str() { "offline" => r#" @@ -640,7 +627,7 @@ model = "{}" _ => String::new(), // "none" — omit section entirely }; - let toml = format!( + format!( r#"# Uteke configuration — generated by `uteke onboard` # See https://github.com/codecoradev/uteke for full documentation @@ -661,7 +648,7 @@ auto_aging_enabled = {auto_maint} [server] enabled = {server_enabled} "#, - namespace = namespace, + namespace = escape_toml(namespace), graph_rerank = graph_rerank, salience_weight = if salience { 0.15 } else { 0.0 }, recency_weight = if recency { 0.15 } else { 0.0 }, @@ -669,7 +656,46 @@ enabled = {server_enabled} aging_enabled = aging_enabled, auto_maint = auto_maint, server_enabled = server_enabled, - ); + ) +} + +/// Write the config to `{uteke_home}/uteke.toml`. +/// +/// If a config file already exists, it is backed up to `uteke.toml.bak` +/// before overwriting. In interactive mode, the user is prompted to confirm. +fn write_config( + namespace: &str, + toggles: &[(&str, bool)], + extraction: &ExtractionSettings, + yes: bool, +) -> Result { + let uteke_dir = uteke_core::uteke_home().map_err(|e| format!("{e}"))?; + std::fs::create_dir_all(&uteke_dir).map_err(|e| format!("Failed to create uteke home: {e}"))?; + + let config_path = uteke_dir.join("uteke.toml"); + + // If config already exists, back it up and confirm before overwriting. + if config_path.exists() { + if !yes { + print!( + "⚠ Config file already exists at {} — overwrite? [y/N] ", + config_path.display() + ); + io::stdout().flush().ok(); + let resp = read_line()?; + if !resp.eq_ignore_ascii_case("y") { + println!("→ Skipped config write (existing config preserved)"); + return Ok(config_path); + } + } + // Back up existing config before overwriting. + let backup_path = uteke_dir.join("uteke.toml.bak"); + std::fs::copy(&config_path, &backup_path) + .map_err(|e| format!("Failed to back up existing config: {e}"))?; + println!("→ Backed up existing config to {}", backup_path.display()); + } + + let toml = render_config_toml(namespace, toggles, extraction); std::fs::write(&config_path, toml).map_err(|e| format!("Failed to write config: {e}"))?; @@ -910,4 +936,59 @@ enabled = false print_banner("A Much Longer Banner Title That Might Cause Issues"); print_banner(""); } + + fn extraction_none() -> ExtractionSettings { + ExtractionSettings { + mode: "none".to_string(), + base_url: String::new(), + api_key: String::new(), + model: String::new(), + } + } + + /// #1332: the namespace was written into the TOML unescaped, so a quote or + /// newline in it broke the file or injected extra keys/sections. + #[test] + fn test_render_config_escapes_hostile_namespace() { + let hostile = "evil\"\n[server]\nenabled = true\n#"; + let toml = render_config_toml(hostile, &[("Server mode", false)], &extraction_none()); + let parsed: toml::Value = toml::from_str(&toml).expect("generated TOML must parse"); + assert_eq!( + parsed["store"]["namespace"].as_str(), + Some(hostile), + "namespace round-trips verbatim" + ); + assert_eq!( + parsed["server"]["enabled"].as_bool(), + Some(false), + "the injected [server] section must not take effect" + ); + } + + #[test] + fn test_render_config_plain_values_parse() { + let toggles = [("Aging", true), ("Salience boost", true)]; + let toml = render_config_toml("my-agent", &toggles, &extraction_none()); + let parsed: toml::Value = toml::from_str(&toml).unwrap(); + assert_eq!(parsed["store"]["namespace"].as_str(), Some("my-agent")); + assert_eq!(parsed["aging"]["enabled"].as_bool(), Some(true)); + assert_eq!(parsed["recall"]["salience_weight"].as_float(), Some(0.15)); + } + + #[test] + fn test_render_config_escapes_llm_extraction_values() { + let ex = ExtractionSettings { + mode: "llm".to_string(), + base_url: "https://x.example/v1".to_string(), + api_key: "k\"\nmodel = \"evil".to_string(), + model: "gpt".to_string(), + }; + let toml = render_config_toml("ns", &[], &ex); + let parsed: toml::Value = toml::from_str(&toml).expect("must parse"); + assert_eq!(parsed["extraction"]["model"].as_str(), Some("gpt")); + assert_eq!( + parsed["extraction"]["api_key"].as_str(), + Some("k\"\nmodel = \"evil") + ); + } } diff --git a/crates/uteke-cli/src/output.rs b/crates/uteke-cli/src/output.rs index d58576f2..42e249f0 100644 --- a/crates/uteke-cli/src/output.rs +++ b/crates/uteke-cli/src/output.rs @@ -386,3 +386,30 @@ pub(crate) fn print_room_document_human(doc: &uteke_core::RoomDocument) { println!(); } } + +/// First `n` characters of `s` (never splits a multibyte character). +/// +/// Display truncation must not byte-slice user text: `&s[..n]` panics when +/// `n` lands inside a multibyte UTF-8 sequence. +pub(crate) fn char_prefix(s: &str, n: usize) -> &str { + s.char_indices().nth(n).map_or(s, |(i, _)| &s[..i]) +} + +#[cfg(test)] +mod char_prefix_tests { + use super::char_prefix; + + #[test] + fn ascii_and_short_strings() { + assert_eq!(char_prefix("abcdef", 3), "abc"); + assert_eq!(char_prefix("ab", 5), "ab"); + assert_eq!(char_prefix("", 3), ""); + } + + #[test] + fn multibyte_never_panics() { + // Each of these chars is 3 bytes; byte index 8 would be mid-char. + assert_eq!(char_prefix("日本語のタイトル", 4), "日本語の"); + assert_eq!(char_prefix("héllo wörld", 5), "héllo"); + } +} diff --git a/crates/uteke-core/Cargo.toml b/crates/uteke-core/Cargo.toml index 05cbd758..0dbd0469 100644 --- a/crates/uteke-core/Cargo.toml +++ b/crates/uteke-core/Cargo.toml @@ -28,6 +28,7 @@ vecq = ["dep:vecq-core"] [dependencies] serde = { version = "1", features = ["derive"] } serde_json = "1" +toml = "1" rusqlite = { version = "0.40", features = ["bundled"] } usearch = { version = "2", optional = true } vecq-core = { version = "0.3", optional = true } diff --git a/crates/uteke-core/src/config_layers.rs b/crates/uteke-core/src/config_layers.rs new file mode 100644 index 00000000..63e0afe3 --- /dev/null +++ b/crates/uteke-core/src/config_layers.rs @@ -0,0 +1,374 @@ +//! Layered `uteke.toml` resolution — one policy shared by the CLI and the server. +//! +//! The same file used to be merged by two hand-written implementations with +//! different semantics (the CLI listed every key by hand; the server replaced +//! whole sections), and the "do not trust the project-local file" rule existed +//! twice. That produced real bugs: whole config sections silently ignored, +//! and a project file wiping the global `[server]` auth settings. +//! +//! This module owns the rules, independent of any schema: +//! +//! * **Precedence**: global file, then the project-local file on top. +//! * **Per-key merge**: tables merge key by key; any other value (including +//! arrays) from the higher layer replaces the lower one. Keys a layer does +//! not mention keep the lower layer's value. +//! * **Trust**: the project-local file lives in the working tree, which may be +//! an untrusted clone, so [`UNTRUSTED_PROJECT_KEYS`] (credentials, endpoints, +//! the embedding backend, the whole `[server]` section) are stripped from it +//! *before* merging unless the caller explicitly trusts it. +//! * **Validation**: every layer must deserialize into the schema on its own; +//! an unreadable or invalid layer is skipped with a warning instead of +//! discarding the others. +//! +//! Callers choose the schema type `T` (`#[serde(default)]` everywhere, so a +//! missing key falls back to its default) and add their own environment +//! overrides on top of the result. + +use serde::de::DeserializeOwned; +use std::path::{Path, PathBuf}; +use toml::{Table, Value}; + +/// Dotted keys the project-local config may not set: they decide where +/// credentials or memory text are sent, or how the server authenticates. +/// A bare section name (`server`) covers every key in that section. +pub const UNTRUSTED_PROJECT_KEYS: &[&str] = &[ + "embedding.backend", + "embedding.api_key", + "embedding.base_url", + "embedding.endpoint_path", + "embed_fallback.api_key", + "embed_fallback.base_url", + "embed_fallback.endpoint_path", + "extraction.api_key", + "extraction.base_url", + "extraction.endpoint_path", + "server", +]; + +/// The files to resolve. Missing files are fine. +#[derive(Debug, Clone, Default)] +pub struct Layers<'a> { + /// `{uteke_home}/uteke.toml` + pub global: Option<&'a Path>, + /// `$CWD/.uteke/uteke.toml` + pub project: Option<&'a Path>, + /// Allow the project file to set [`UNTRUSTED_PROJECT_KEYS`] (explicit opt-in). + pub trust_project: bool, +} + +/// What happened while resolving, for logging and tests. +#[derive(Debug)] +pub struct Resolved { + pub value: T, + /// Keys dropped from the project file because it is untrusted. + pub ignored_untrusted: Vec, + /// Layers skipped because they could not be read or did not match the schema. + pub skipped: Vec, +} + +/// Merge `overlay` into `base`: tables recurse, everything else is replaced. +pub fn deep_merge(base: &mut Table, overlay: Table) { + for (key, value) in overlay { + match (base.get_mut(&key), value) { + (Some(Value::Table(base_table)), Value::Table(overlay_table)) => { + deep_merge(base_table, overlay_table); + } + (_, value) => { + base.insert(key, value); + } + } + } +} + +/// Remove [`UNTRUSTED_PROJECT_KEYS`] from a project-local table. Returns the +/// dotted names that were actually present (and therefore dropped). +pub fn strip_untrusted(table: &mut Table) -> Vec { + let mut removed = Vec::new(); + for key in UNTRUSTED_PROJECT_KEYS { + match key.split_once('.') { + None => { + if table.remove(*key).is_some() { + removed.push((*key).to_string()); + } + } + Some((section, leaf)) => { + if let Some(Value::Table(section_table)) = table.get_mut(section) { + if section_table.remove(leaf).is_some() { + removed.push((*key).to_string()); + } + } + } + } + } + removed +} + +fn read_layer(path: &Path) -> Option { + if !path.exists() { + return None; + } + let content = match std::fs::read_to_string(path) { + Ok(c) => c, + Err(e) => { + tracing::warn!("Cannot read config {}: {e}", path.display()); + return None; + } + }; + match toml::from_str::
(&content) { + Ok(table) => Some(table), + Err(e) => { + tracing::warn!("Invalid config {}: {e}", path.display()); + None + } + } +} + +/// Resolve the layers into a `T`. +pub fn resolve(layers: &Layers<'_>) -> Resolved { + let mut merged = Table::new(); + let mut ignored_untrusted = Vec::new(); + let mut skipped = Vec::new(); + + let ordered = [(layers.global, false), (layers.project, true)]; + for (path, is_project) in ordered { + let Some(path) = path else { continue }; + let Some(mut table) = read_layer(path) else { + if path.exists() { + skipped.push(path.to_path_buf()); + } + continue; + }; + if is_project && !layers.trust_project { + let removed = strip_untrusted(&mut table); + if !removed.is_empty() { + tracing::warn!( + "Ignoring {} from project config {} (untrusted: set them in the global \ + config, or set UTEKE_TRUST_PROJECT_CONFIG=1 to allow)", + removed.join(", "), + path.display() + ); + ignored_untrusted.extend(removed); + } + } + // Each layer must stand on its own against the schema. + if let Err(e) = Value::Table(table.clone()).try_into::() { + tracing::warn!("Invalid config {}: {e}", path.display()); + skipped.push(path.to_path_buf()); + continue; + } + deep_merge(&mut merged, table); + } + + let value = Value::Table(merged).try_into::().unwrap_or_else(|e| { + tracing::warn!("Merged config does not match the schema ({e}); using defaults"); + T::default() + }); + Resolved { + value, + ignored_untrusted, + skipped, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde::Deserialize; + + #[derive(Debug, Deserialize, PartialEq)] + #[serde(default)] + struct Schema { + top: bool, + store: Store, + embedding: Embedding, + server: Server, + recall: Recall, + } + impl Default for Schema { + fn default() -> Self { + Self { + top: true, + store: Store::default(), + embedding: Embedding::default(), + server: Server::default(), + recall: Recall::default(), + } + } + } + #[derive(Debug, Deserialize, PartialEq, Default)] + #[serde(default)] + struct Store { + path: String, + namespace: String, + } + #[derive(Debug, Deserialize, PartialEq, Default)] + #[serde(default)] + struct Embedding { + backend: String, + model: String, + base_url: String, + api_key: String, + } + #[derive(Debug, Deserialize, PartialEq, Default)] + #[serde(default)] + struct Server { + host: String, + auth_token: String, + cors_origins: Vec, + } + #[derive(Debug, Deserialize, PartialEq, Default)] + #[serde(default)] + struct Recall { + min_score: f64, + strategy: String, + } + + fn write(dir: &Path, name: &str, body: &str) -> PathBuf { + let p = dir.join(name); + std::fs::write(&p, body).unwrap(); + p + } + + fn layers<'a>(g: Option<&'a Path>, p: Option<&'a Path>, trust: bool) -> Layers<'a> { + Layers { + global: g, + project: p, + trust_project: trust, + } + } + + #[test] + fn deep_merge_recurses_into_tables_and_replaces_other_values() { + let mut base: Table = toml::from_str( + "a = 1\n[s]\nx = 1\ny = 2\nlist = [1, 2]\n[s.inner]\nk = \"base\"\nkeep = true", + ) + .unwrap(); + let overlay: Table = + toml::from_str("a = 9\n[s]\ny = 20\nlist = [3]\n[s.inner]\nk = \"over\"").unwrap(); + deep_merge(&mut base, overlay); + let expected: Table = toml::from_str( + "a = 9\n[s]\nx = 1\ny = 20\nlist = [3]\n[s.inner]\nk = \"over\"\nkeep = true", + ) + .unwrap(); + assert_eq!(base, expected); + } + + #[test] + fn project_overrides_global_per_key() { + let d = tempfile::tempdir().unwrap(); + let g = write( + d.path(), + "g.toml", + "[store]\npath = \"/g\"\nnamespace = \"g\"\n[recall]\nmin_score = 0.3\nstrategy = \"vector\"", + ); + let p = write(d.path(), "p.toml", "[recall]\nmin_score = 0.5"); + let r = resolve::(&layers(Some(&g), Some(&p), false)); + assert_eq!(r.value.recall.min_score, 0.5, "project wins"); + assert_eq!( + r.value.recall.strategy, "vector", + "a project key that is absent keeps the global value" + ); + assert_eq!(r.value.store.path, "/g"); + assert!(r.value.top, "untouched keys keep their schema default"); + } + + #[test] + fn untrusted_project_keys_are_stripped_before_merging() { + let d = tempfile::tempdir().unwrap(); + let g = write( + d.path(), + "g.toml", + "[embedding]\nbackend = \"onnx\"\nbase_url = \"https://global.example\"\n[server]\nauth_token = \"global\"\nhost = \"0.0.0.0\"\ncors_origins = [\"https://app.example\"]", + ); + let p = write( + d.path(), + "p.toml", + "[embedding]\nbackend = \"openai\"\nbase_url = \"https://evil.example\"\napi_key = \"stolen\"\nmodel = \"tuned\"\n[server]\nauth_token = \"evil\"\nhost = \"evil.example\"", + ); + let r = resolve::(&layers(Some(&g), Some(&p), false)); + assert_eq!(r.value.embedding.backend, "onnx"); + assert_eq!(r.value.embedding.base_url, "https://global.example"); + assert_eq!(r.value.embedding.api_key, ""); + assert_eq!( + r.value.embedding.model, "tuned", + "harmless keys still apply" + ); + assert_eq!(r.value.server.auth_token, "global"); + assert_eq!(r.value.server.host, "0.0.0.0"); + assert_eq!(r.value.server.cors_origins, vec!["https://app.example"]); + for key in [ + "embedding.backend", + "embedding.base_url", + "embedding.api_key", + "server", + ] { + assert!(r.ignored_untrusted.contains(&key.to_string()), "{key}"); + } + } + + #[test] + fn global_layer_may_set_sensitive_keys() { + let d = tempfile::tempdir().unwrap(); + let g = write( + d.path(), + "g.toml", + "[embedding]\napi_key = \"ok\"\n[server]\nhost = \"h\"", + ); + let r = resolve::(&layers(Some(&g), None, false)); + assert_eq!(r.value.embedding.api_key, "ok"); + assert_eq!(r.value.server.host, "h"); + assert!(r.ignored_untrusted.is_empty()); + } + + #[test] + fn explicit_trust_lets_the_project_file_set_everything() { + let d = tempfile::tempdir().unwrap(); + let p = write( + d.path(), + "p.toml", + "[server]\nhost = \"proj\"\n[embedding]\nbackend = \"ollama\"", + ); + let r = resolve::(&layers(None, Some(&p), true)); + assert_eq!(r.value.server.host, "proj"); + assert_eq!(r.value.embedding.backend, "ollama"); + assert!(r.ignored_untrusted.is_empty()); + } + + #[test] + fn a_project_file_without_server_section_never_wipes_the_global_one() { + let d = tempfile::tempdir().unwrap(); + let g = write(d.path(), "g.toml", "[server]\nauth_token = \"t\""); + let p = write(d.path(), "p.toml", "[recall]\nmin_score = 0.1"); + let r = resolve::(&layers(Some(&g), Some(&p), false)); + assert_eq!(r.value.server.auth_token, "t"); + } + + #[test] + fn invalid_layers_are_skipped_not_fatal() { + let d = tempfile::tempdir().unwrap(); + let g = write(d.path(), "g.toml", "[recall]\nstrategy = \"fusion\""); + let broken_toml = write(d.path(), "p1.toml", "this is not valid toml [[[["); + let r = resolve::(&layers(Some(&g), Some(&broken_toml), false)); + assert_eq!(r.value.recall.strategy, "fusion"); + assert_eq!(r.skipped, vec![broken_toml.clone()]); + + // Valid TOML, wrong type for a schema field. + let wrong_type = write(d.path(), "p2.toml", "[recall]\nmin_score = \"high\""); + let r = resolve::(&layers(Some(&g), Some(&wrong_type), false)); + assert_eq!( + r.value.recall.strategy, "fusion", + "other layers still apply" + ); + assert_eq!(r.value.recall.min_score, 0.0, "bad value did not leak in"); + assert_eq!(r.skipped, vec![wrong_type]); + } + + #[test] + fn missing_files_resolve_to_schema_defaults() { + let d = tempfile::tempdir().unwrap(); + let nope = d.path().join("nope.toml"); + let r = resolve::(&layers(Some(&nope), Some(&nope), false)); + assert_eq!(r.value, Schema::default()); + assert!(r.skipped.is_empty() && r.ignored_untrusted.is_empty()); + } +} diff --git a/crates/uteke-core/src/consolidate.rs b/crates/uteke-core/src/consolidate.rs index 7ec42ac1..37b91d98 100644 --- a/crates/uteke-core/src/consolidate.rs +++ b/crates/uteke-core/src/consolidate.rs @@ -1,6 +1,7 @@ //! Consolidation: contradiction detection, duplicate finding, and merging. use crate::error::Error; +use crate::index_sync::{self, SyncPolicy}; use crate::memory::types::{ContradictionResult, DEFAULT_NAMESPACE}; use crate::operations::RememberOutcome; @@ -142,13 +143,12 @@ impl crate::Uteke { let mut idx = self.index.write().map_err(|_| { Error::lock("index write lock during post-insert contradiction deprecation") })?; - if idx.remove(deprecated_id) { - if let Err(e) = idx.save() { - tracing::warn!( - "Failed to persist vector index after deprecating id={deprecated_id}: {e}" - ); - } - } + index_sync::remove_ids( + &mut *idx, + [deprecated_id.as_str()], + &SyncPolicy::STANDARD, + ) + .warn_if_stale("post-insert contradiction deprecation"); } } } @@ -300,7 +300,7 @@ impl crate::Uteke { .map_err(|e| Error::db("consolidate delete", e))?; } // SQLite first (source of truth), then vector index. - if !index.remove(to_remove) { + if !index_sync::remove_unsaved(&mut *index, to_remove) { tracing::warn!( "Vector index entry not found during consolidate for id={}", to_remove @@ -313,7 +313,7 @@ impl crate::Uteke { } // Persist vector index once after all removals. if index_dirty { - if let Err(e) = index.save() { + if let Err(e) = index_sync::persist(&mut *index, &SyncPolicy::STANDARD) { tracing::warn!( "Failed to persist vector index after consolidate: {e}. \ Orphan entries will be cleaned up by verify/repair." @@ -386,17 +386,13 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during consolidate_pair"))?; - if !index.remove(id_remove) { + let sync = index_sync::remove_ids(&mut *index, [id_remove], &SyncPolicy::STANDARD); + if sync.missing > 0 { tracing::warn!( "Vector index entry not found during consolidate_pair for id={id_remove}" ); } - if let Err(e) = index.save() { - tracing::warn!( - "Failed to persist vector index after consolidate_pair: {e}. \ - Orphan entries will be cleaned up by verify/repair." - ); - } + sync.warn_if_stale("consolidate_pair"); // Invalidate recall cache — the removed memory affects search results. if namespace.is_empty() { self.recall_cache.clear(); diff --git a/crates/uteke-core/src/dream.rs b/crates/uteke-core/src/dream.rs index aecf7a19..bd4b1e45 100644 --- a/crates/uteke-core/src/dream.rs +++ b/crates/uteke-core/src/dream.rs @@ -17,6 +17,7 @@ //! All phases are idempotent and safe to re-run. use crate::error::Error; +use crate::memory::store::memory_columns; use serde::{Deserialize, Serialize}; /// A single phase of the dream cycle. @@ -443,7 +444,7 @@ impl crate::Uteke { // Create "contradicts" graph edge let gs = crate::GraphStore::new(&self.store.conn); - match gs.add_edge(older, newer, "contradicts", cosine as f64) { + match gs.add_edge_for_memories(older, newer, "contradicts", cosine as f64) { Ok(()) => { edges_created += 1; tracing::info!( @@ -497,20 +498,22 @@ impl crate::Uteke { ) -> Result, Error> { let sql = match namespace { Some(_ns) => { - "SELECT id, content, embedding, tags, metadata, \ - created_at, updated_at, namespace, access_count, \ - last_accessed, deprecated, valid_from, valid_until, \ - memory_type, importance, pinned, content_type, slug \ + concat!( + "SELECT ", + memory_columns!(), + " \ FROM memories WHERE namespace = ?1 AND deprecated = 0 \ ORDER BY updated_at DESC LIMIT ?2" + ) } None => { - "SELECT id, content, embedding, tags, metadata, \ - created_at, updated_at, namespace, access_count, \ - last_accessed, deprecated, valid_from, valid_until, \ - memory_type, importance, pinned, content_type, slug \ + concat!( + "SELECT ", + memory_columns!(), + " \ FROM memories WHERE deprecated = 0 \ ORDER BY updated_at DESC LIMIT ?1" + ) } }; @@ -608,7 +611,7 @@ impl crate::Uteke { // Prune deprecated memories older than 30 days. const TTL_DAYS: u32 = 30; let result = self.prune(TTL_DAYS, namespace, dry_run)?; - let summary = if dry_run { + let mut summary = if dry_run { format!( "✓ {} memories would be pruned ({} deprecated)", result.pruned, result.deprecated @@ -619,6 +622,17 @@ impl crate::Uteke { result.pruned, result.deprecated ) }; + // Pruning leaves tombstoned rows in an append-only vecq index (#1324). + // The index is store-wide, so this ignores the namespace filter. + if let Some(c) = self.compact_index_if_needed(dry_run)? { + summary.push_str(&format!( + "; index {} {} dead rows ({} → {} rows)", + if dry_run { "would drop" } else { "dropped" }, + c.dead_before, + c.rows_before, + c.rows_after + )); + } Ok(PhaseResult { phase: DreamPhase::Compact.as_str().to_string(), status: PhaseStatus::Ok, diff --git a/crates/uteke-core/src/edges.rs b/crates/uteke-core/src/edges.rs index 814c0bb4..b9b5878d 100644 --- a/crates/uteke-core/src/edges.rs +++ b/crates/uteke-core/src/edges.rs @@ -14,6 +14,7 @@ //! indexed SQL queries instead of the old O(n) JSON scan. use crate::error::Error; +use crate::index_sync::{self, SyncPolicy}; use crate::memory::types::Memory; use rusqlite::{OptionalExtension, params}; use serde::{Deserialize, Serialize}; @@ -612,9 +613,15 @@ impl Store { let mut neighbors_stmt = self .conn .prepare( - "SELECT target_id FROM memory_edges WHERE source_id = ?1 + // Neighbours must be live: a deprecated memory is neither + // returned nor traversed through (soft delete keeps its edges). + "SELECT e.target_id FROM memory_edges e + JOIN memories m ON m.id = e.target_id AND m.deprecated = 0 + WHERE e.source_id = ?1 UNION - SELECT source_id FROM memory_edges WHERE target_id = ?1", + SELECT e.source_id FROM memory_edges e + JOIN memories m ON m.id = e.source_id AND m.deprecated = 0 + WHERE e.target_id = ?1", ) .map_err(|e| Error::db("prepare bfs neighbors", e))?; @@ -848,7 +855,13 @@ impl crate::Uteke { ) { const SIMILAR_THRESHOLD: f32 = 0.80; const DUPLICATE_THRESHOLD: f32 = 0.92; + // Candidates that may become edges. const TOP_K: usize = 20; + // The index is shared by every namespace and by document chunks, and + // the namespace / liveness filter below runs AFTER the search. Fetch + // extra so a small namespace inside a big index still finds its own + // neighbours instead of being crowded out by other namespaces. + const FETCH_K: usize = TOP_K * 4; let index = match self.index.read() { Ok(i) => i, @@ -858,38 +871,36 @@ impl crate::Uteke { } }; - let results = index.search(embedding, TOP_K, 100); + let results = index.search(embedding, FETCH_K, 100); drop(index); // Release read lock ASAP if results.is_empty() { return; } - // Filter by namespace if specified (#401 cora finding: cross-namespace links are wrong). - let ns_set: Option> = if let Some(ns) = namespace { - match self.store.memories_in_namespace(ns) { - Ok(ids) => Some(ids.into_iter().collect()), - Err(e) => { - tracing::warn!("auto_link_cosine: failed to get namespace ids: {e}"); - return; - } + // Keep only live memories of this namespace (#401: cross-namespace links + // are wrong). One query over the candidates — not a scan of the whole + // namespace per write. `chunk:` document-chunk keys and deprecated + // memories are not in the result, so they never become link targets. + let candidate_ids: Vec<&str> = results + .iter() + .map(|(id, _)| id.as_str()) + .filter(|id| *id != source_id) + .collect(); + let live = match self.store.live_ids_in_namespace(&candidate_ids, namespace) { + Ok(set) => set, + Err(e) => { + tracing::warn!("auto_link_cosine: failed to check candidates: {e}"); + return; } - } else { - None }; // usearch returns distances (lower = more similar for cosine). // Convert distance to similarity: sim = 1.0 - dist. let edges: Vec<(String, String)> = results .iter() - .filter(|(id, _)| id != source_id) // skip self - .filter(|(id, _)| { - // Skip memories outside our namespace. - match &ns_set { - Some(set) => set.contains(id), - None => true, - } - }) + .filter(|(id, _)| live.contains(id)) // also drops self + .take(TOP_K) .filter_map(|(id, dist)| { // Cosine distance → similarity (clamp to [0, 1]). let sim = (1.0 - dist).clamp(0.0, 1.0); @@ -2092,15 +2103,8 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during supersede"))?; - if !index.remove(&old.id) { - tracing::debug!( - "Vector index entry not found during supersede for id={} (ok if never embedded)", - old.id - ); - } - if let Err(e) = index.save() { - tracing::warn!("Failed to persist vector index after supersede: {e}"); - } + index_sync::remove_ids(&mut *index, [old.id.as_str()], &SyncPolicy::STANDARD) + .warn_if_stale("supersede"); } self.recall_cache.invalidate_namespace(&old.namespace); @@ -2180,13 +2184,13 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during undo_supersession"))?; - if let Err(e) = index.insert(&memory.id, &memory.embedding) { - tracing::warn!( - "Failed to re-insert memory id={} into vector index during undo_supersession: {e}", - memory.id - ); - } - let _ = index.save(); + let sync = index_sync::upsert( + &mut *index, + &memory.id, + &memory.embedding, + &SyncPolicy::STANDARD, + ); + sync.warn_if_stale("undo_supersession"); } self.recall_cache.invalidate_namespace(&memory.namespace); } @@ -2588,3 +2592,202 @@ mod resupersession_ledger_tests { std::fs::remove_dir_all(&dir).ok(); } } + +#[cfg(test)] +mod auto_link_hot_path_tests { + //! #1368 — the remember() hot path must not scan a whole namespace, must + //! not be crowded out by other namespaces, and must never link to + //! deprecated memories. + use crate::Uteke; + + const DIMS: usize = 768; + + /// Unit-ish vectors sharing axis 0 with a private perturbation on axis + /// `k`: pairwise cosine = 1 / (1 + eps^2) ≈ 0.89 — above the similar_to + /// threshold (0.80), below possible_duplicate (0.92) and the remember() + /// dedup threshold (0.95), so every one of them is stored and linkable. + fn vec_k(k: usize) -> Vec { + let mut v = vec![0.0_f32; DIMS]; + v[0] = 1.0; + v[k + 1] = 0.35; + v + } + + fn mem(u: &Uteke, text: &str, ns: &str, k: usize) -> String { + u.remember_precomputed(text, &[], None, Some(ns), "fact", "text", &vec_k(k)) + .unwrap() + } + + fn similar_targets(u: &Uteke, id: &str) -> Vec { + u.edges_for(id) + .unwrap() + .outgoing + .into_iter() + .filter(|e| e.edge_type == "similar_to") + .map(|e| e.target_id) + .collect() + } + + /// `e0 + 0.35*e_hot + 0.05*e_k`: sim to `e0 + 0.35*e_hot` is ~0.999. + fn near_hot(hot: usize, k: usize) -> Vec { + let mut v = vec![0.0_f32; DIMS]; + v[0] = 1.0; + v[hot + 1] = 0.35; + v[k + 1] = 0.05; + v + } + + #[test] + fn small_namespace_inside_a_big_index_still_gets_its_own_links() { + let u = Uteke::open(":memory:").unwrap(); + let hot = 300; + // 60 memories of ANOTHER namespace that are much closer (~0.999) to the + // newcomer than its own neighbours (~0.89). A plain top-20 search + // followed by a namespace filter is filled entirely by them and the + // newcomer gets no links at all. + for k in 0..60 { + u.remember_precomputed( + &format!("big {k}"), + &[], + None, + Some("big"), + "fact", + "text", + &near_hot(hot, k), + ) + .unwrap(); + } + let s1 = mem(&u, "small one", "small", 100); + let s2 = mem(&u, "small two", "small", 101); + let s3 = mem(&u, "small three", "small", 102); + let new = u + .remember_precomputed( + "small four", + &[], + None, + Some("small"), + "fact", + "text", + &vec_k(hot), + ) + .unwrap(); + + let targets = similar_targets(&u, &new); + for s in [&s1, &s2, &s3] { + assert!(targets.contains(s), "missing link to {s}: {targets:?}"); + } + // never across namespaces + let all = u.store.list(None, None, 1000, 0).unwrap(); + for t in &targets { + let m = all.iter().find(|m| &m.id == t).unwrap(); + assert_eq!(m.namespace, "small", "cross-namespace link to {t}"); + } + } + + #[test] + fn deprecated_memories_are_never_link_targets() { + let u = Uteke::open(":memory:").unwrap(); + let live = mem(&u, "live neighbour", "ns", 1); + let dead = mem(&u, "deprecated neighbour", "ns", 2); + // Deprecate in SQLite only: the vector stays in the index, exactly + // like a stale entry after a crash or an old binary (#1210). + u.store.deprecate(&dead).unwrap(); + + let new = mem(&u, "newcomer", "ns", 3); + let targets = similar_targets(&u, &new); + assert!(targets.contains(&live)); + assert!(!targets.contains(&dead), "linked to a deprecated memory"); + } + + #[test] + fn live_ids_in_namespace_filters_candidates_only() { + let u = Uteke::open(":memory:").unwrap(); + let a = mem(&u, "a", "one", 1); + let b = mem(&u, "b", "two", 2); + let gone = mem(&u, "gone", "one", 3); + u.store.deprecate(&gone).unwrap(); + + let ids = [a.as_str(), b.as_str(), gone.as_str(), "chunk:xyz", "nope"]; + let in_one = u.store.live_ids_in_namespace(&ids, Some("one")).unwrap(); + assert_eq!(in_one, [a.clone()].into_iter().collect()); + let anywhere = u.store.live_ids_in_namespace(&ids, None).unwrap(); + assert_eq!(anywhere, [a, b].into_iter().collect()); + assert!( + u.store + .live_ids_in_namespace::<&str>(&[], Some("one")) + .unwrap() + .is_empty() + ); + } + + #[test] + fn dedup_stays_scoped_to_the_namespace_and_ignores_deprecated() { + let u = Uteke::open(":memory:").unwrap(); + let v = vec_k(7); + let first = u + .remember_precomputed("same text", &[], None, Some("ns-a"), "fact", "text", &v) + .unwrap(); + // Same namespace: the near-duplicate is found. + assert_eq!( + u.check_duplicate(&v, Some("ns-a")).unwrap(), + Some(first.clone()) + ); + // Another namespace holds nothing similar. + assert_eq!(u.check_duplicate(&v, Some("ns-b")).unwrap(), None); + // No namespace filter: any live memory counts. + assert_eq!(u.check_duplicate(&v, None).unwrap(), Some(first.clone())); + // A deprecated memory still in the index is not a dedup target (#1210). + u.store.deprecate(&first).unwrap(); + assert_eq!(u.check_duplicate(&v, Some("ns-a")).unwrap(), None); + } +} + +#[cfg(test)] +mod deprecated_edge_reader_tests { + //! #1367 — soft delete keeps `memory_edges` rows; readers that feed + //! ranking/traversal must ignore the retired side, while the audit view + //! (`edges_for`, the supersession ledger) keeps showing them. + use crate::Uteke; + + fn mem(u: &Uteke, text: &str, k: usize) -> String { + // Orthogonal vectors: no auto-created similar_to edges, so the chain + // below consists only of the explicit edges the test adds. + let mut v = vec![0.0_f32; 768]; + v[k] = 1.0; + u.remember_precomputed(text, &[], None, Some("d"), "fact", "text", &v) + .unwrap() + } + + #[test] + fn bfs_skips_deprecated_memories_but_edges_for_keeps_history() { + let u = Uteke::open(":memory:").unwrap(); + let a = mem(&u, "a", 1); + let b = mem(&u, "b", 2); + let c = mem(&u, "c", 3); + // chain a - b - c with an explicit relation type that is not auto-linked + u.link_memories(&a, &b, "references").unwrap(); + u.link_memories(&b, &c, "references").unwrap(); + u.store.deprecate(&b).unwrap(); + + let related: Vec = u + .related_via_edges(&a, 3) + .unwrap() + .into_iter() + .map(|m| m.id) + .collect(); + assert!(!related.contains(&b), "deprecated b must not be returned"); + assert!( + !related.contains(&c), + "c is only reachable through the deprecated b" + ); + + // The audit view still lists the edge to the deprecated memory. + let outgoing = u.edges_for(&a).unwrap().outgoing; + assert!( + outgoing + .iter() + .any(|e| e.target_id == b && e.edge_type == "references"), + "history must stay visible: {outgoing:?}" + ); + } +} diff --git a/crates/uteke-core/src/graph.rs b/crates/uteke-core/src/graph.rs index 52eaf5cc..2c98e3af 100644 --- a/crates/uteke-core/src/graph.rs +++ b/crates/uteke-core/src/graph.rs @@ -395,6 +395,39 @@ impl<'a> GraphStore<'a> { Ok(row.filter(|s| !s.trim().is_empty())) } + /// Create an edge between two **memories** (#1365). + /// + /// `graph_edges.source_id/target_id` reference `graph_nodes(id)`, so raw + /// memory ids violate the FK (and fail with "FOREIGN KEY constraint + /// failed"). This ensures a node for each memory first, exactly like the + /// HTTP `POST /graph/edge` handler does (#1180). + pub fn add_edge_for_memories( + &self, + source_memory_id: &str, + target_memory_id: &str, + relation: &str, + weight: f64, + ) -> Result<(), Error> { + let src = self.ensure_node_for_memory(source_memory_id)?; + let tgt = self.ensure_node_for_memory(target_memory_id)?; + self.add_edge(&src, &tgt, relation, weight) + } + + /// Remove the edge between two ids that may each be a memory id or a + /// graph node id. Returns `false` when either side or the edge is unknown. + pub fn remove_edge_between(&self, a: &str, b: &str) -> Result { + let resolve = |id: &str| -> Result, Error> { + if let Some(nid) = self.node_id_for_memory(id)? { + return Ok(Some(nid)); + } + Ok(self.get_node(id)?.map(|n| n.id)) + }; + match (resolve(a)?, resolve(b)?) { + (Some(src), Some(tgt)) => self.remove_edge(&src, &tgt), + _ => Ok(false), + } + } + /// Create an edge between two nodes. Ignores if edge already exists (INSERT OR IGNORE). pub fn add_edge( &self, diff --git a/crates/uteke-core/src/graph_rerank.rs b/crates/uteke-core/src/graph_rerank.rs index ebc81c9d..c0c6b50e 100644 --- a/crates/uteke-core/src/graph_rerank.rs +++ b/crates/uteke-core/src/graph_rerank.rs @@ -117,9 +117,14 @@ pub fn compute_graph_signals( let placeholders = std::iter::repeat_n("?", unique_ids.len()) .collect::>() .join(", "); + // Only edges between LIVE memories count. soft_delete_only is the default, + // so forget/deprecate leaves edges pointing at deprecated memories behind; + // counting them would boost a result for links to things the user retired. let sql = format!( - "SELECT source_id, target_id, edge_type FROM memory_edges - WHERE source_id IN ({src}) OR target_id IN ({tgt})", + "SELECT e.source_id, e.target_id, e.edge_type FROM memory_edges e + JOIN memories s ON s.id = e.source_id AND s.deprecated = 0 + JOIN memories t ON t.id = e.target_id AND t.deprecated = 0 + WHERE e.source_id IN ({src}) OR e.target_id IN ({tgt})", src = placeholders, tgt = placeholders, ); @@ -255,6 +260,18 @@ mod tests { /// list of `(source, target, edge_type)` tuples. fn db_with_edges(edges: &[(&str, &str, &str)]) -> Connection { let conn = Connection::open_in_memory().unwrap(); + conn.execute( + "CREATE TABLE memories (id TEXT PRIMARY KEY, deprecated INTEGER NOT NULL DEFAULT 0)", + [], + ) + .unwrap(); + // Every endpoint is a live memory unless a test deprecates it. + for (s, t, _) in edges { + for id in [s, t] { + conn.execute("INSERT OR IGNORE INTO memories (id) VALUES (?1)", [id]) + .unwrap(); + } + } conn.execute( "CREATE TABLE memory_edges ( source_id TEXT NOT NULL, @@ -492,4 +509,24 @@ mod tests { elapsed ); } + + #[test] + fn edges_to_deprecated_memories_do_not_count() { + // A has edges to B (live), C (deprecated) and D → A from a deprecated D. + let conn = db_with_edges(&[ + ("A", "B", "references"), + ("A", "C", "references"), + ("D", "A", "tagged_as"), + ]); + conn.execute( + "UPDATE memories SET deprecated = 1 WHERE id IN ('C', 'D')", + [], + ) + .unwrap(); + let sigs = compute_graph_signals(&conn, &["A".to_string()]).unwrap(); + let a = sigs.get("A").unwrap(); + assert_eq!(a.edge_count, 1, "only the edge to live B counts"); + assert_eq!(a.outgoing_count, 1); + assert_eq!(a.incoming_count, 0); + } } diff --git a/crates/uteke-core/src/graph_view.rs b/crates/uteke-core/src/graph_view.rs new file mode 100644 index 00000000..5897af54 --- /dev/null +++ b/crates/uteke-core/src/graph_view.rs @@ -0,0 +1,482 @@ +//! Unified, bounded graph view for visualisation (`GET /graph`, MCP `uteke_graph`). +//! +//! Two stores can hold relations between memories: +//! +//! * `memory_edges` — written automatically on every `remember` (cosine +//! links, tag/slug references) and by supersession/consolidation. This is +//! where the data actually is. +//! * `graph_nodes` / `graph_edges` — explicit knowledge-graph entries added +//! through `POST /graph/edge` and friends. Almost always empty. +//! +//! [`Uteke::graph_data`] only ever read the second one, so viewers (Corin) +//! got an empty graph on real stores. [`Uteke::graph_view`] merges both, +//! keyed by **memory id** (what a viewer can open), restricted to live +//! memories, optionally limited to one memory's neighbourhood, and bounded. + +use crate::Error; +use crate::Uteke; +use crate::graph::{GraphNode, GraphStats}; +use rusqlite::params; +use serde::{Deserialize, Serialize}; +use std::collections::{BTreeSet, HashMap}; + +/// Default number of edges returned when the caller passes no limit. +pub const GRAPH_VIEW_DEFAULT_LIMIT: usize = 500; +/// Hard upper bound for `limit`. +pub const GRAPH_VIEW_MAX_LIMIT: usize = 5000; + +/// Derived `referenced_by` backlinks mirror a forward edge one-to-one; they +/// add clutter to a picture, so the view leaves them out (`GET /edges?id=` +/// still lists them). +const DERIVED_BACKLINK: &str = "referenced_by"; + +/// One edge in a [`GraphView`]. +/// +/// `source`/`target` repeat `source_id`/`target_id` because viewers read the +/// short names, while the long names are what the rest of the API uses. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct GraphViewEdge { + pub id: String, + pub source_id: String, + pub target_id: String, + pub source: String, + pub target: String, + pub relation: String, + pub weight: f64, + pub created_at: String, +} + +impl GraphViewEdge { + fn new(id: String, source: &str, target: &str, relation: &str, weight: f64, at: &str) -> Self { + Self { + id, + source_id: source.to_string(), + target_id: target.to_string(), + source: source.to_string(), + target: target.to_string(), + relation: relation.to_string(), + weight, + created_at: at.to_string(), + } + } +} + +/// Nodes + edges + counts of what is in the (possibly truncated) view. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GraphView { + pub nodes: Vec, + pub edges: Vec, + pub stats: GraphStats, + /// `true` when more edges matched than `limit` allowed. + pub truncated: bool, +} + +/// First 60 characters of a memory, plus its short id (matches the label the +/// explicit graph uses, #1187). +fn memory_label(content: &str, id: &str) -> String { + let preview: String = content.chars().take(60).collect(); + let short: String = id.chars().take(8).collect(); + format!("{preview} — {short}") +} + +impl Uteke { + /// Build the unified graph view. + /// + /// * `namespace` — only memories of this namespace (both endpoints). + /// * `node_id` — only edges touching this memory id (its neighbourhood). + /// * `limit` — maximum number of edges (clamped to `1..=GRAPH_VIEW_MAX_LIMIT`). + /// + /// Only live (non-deprecated) memories appear, so edges that point at + /// deprecated memories never reach a viewer. + pub fn graph_view( + &self, + namespace: Option<&str>, + node_id: Option<&str>, + limit: usize, + ) -> Result { + let limit = limit.clamp(1, GRAPH_VIEW_MAX_LIMIT); + // One extra row tells us whether the view was truncated. + let fetch = (limit + 1) as i64; + let conn = &self.store.conn; + + // --- memory_edges between live memories ------------------------------ + // Fixed SQL: NULL parameters switch the optional filters off. + let mut edges: Vec = Vec::new(); + { + let mut stmt = conn + .prepare( + "SELECT e.id, e.source_id, e.target_id, e.edge_type, e.created_at \ + FROM memory_edges e \ + JOIN memories s ON s.id = e.source_id AND s.deprecated = 0 \ + JOIN memories t ON t.id = e.target_id AND t.deprecated = 0 \ + WHERE e.edge_type <> ?1 \ + AND (?2 IS NULL OR (s.namespace = ?2 AND t.namespace = ?2)) \ + AND (?3 IS NULL OR e.source_id = ?3 OR e.target_id = ?3) \ + ORDER BY e.created_at DESC, e.id DESC LIMIT ?4", + ) + .map_err(|e| Error::db("prepare graph_view memory edges", e))?; + let rows = stmt + .query_map(params![DERIVED_BACKLINK, namespace, node_id, fetch], |r| { + Ok(( + r.get::<_, i64>(0)?, + r.get::<_, String>(1)?, + r.get::<_, String>(2)?, + r.get::<_, String>(3)?, + r.get::<_, String>(4)?, + )) + }) + .map_err(|e| Error::db("query graph_view memory edges", e))?; + for row in rows { + let (id, s, t, ty, at) = + row.map_err(|e| Error::db("read graph_view memory edge", e))?; + edges.push(GraphViewEdge::new( + format!("me:{id}"), + &s, + &t, + &ty, + 1.0, + &at, + )); + } + } + + // --- explicit graph, re-keyed by memory id --------------------------- + let explicit = self.graph_data(namespace)?; + let node_key: HashMap<&str, &str> = explicit + .nodes + .iter() + .map(|n| { + ( + n.id.as_str(), + n.memory_id.as_deref().unwrap_or(n.id.as_str()), + ) + }) + .collect(); + for e in &explicit.edges { + let (Some(&s), Some(&t)) = ( + node_key.get(e.source_id.as_str()), + node_key.get(e.target_id.as_str()), + ) else { + continue; + }; + if let Some(node) = node_id { + if s != node && t != node { + continue; + } + } + edges.push(GraphViewEdge::new( + format!("ge:{}", e.id), + s, + t, + &e.relation, + e.weight, + &e.created_at, + )); + } + + // Newest first across both sources, then bound. + edges.sort_by(|a, b| { + b.created_at + .cmp(&a.created_at) + .then_with(|| b.id.cmp(&a.id)) + }); + let truncated = edges.len() > limit; + edges.truncate(limit); + + // --- nodes: every endpoint of a returned edge ------------------------ + let mut ids: BTreeSet<&str> = BTreeSet::new(); + for e in &edges { + ids.insert(e.source_id.as_str()); + ids.insert(e.target_id.as_str()); + } + let explicit_by_key: HashMap<&str, &GraphNode> = explicit + .nodes + .iter() + .map(|n| (n.memory_id.as_deref().unwrap_or(n.id.as_str()), n)) + .collect(); + + let id_list = serde_json::to_string(&ids.iter().collect::>()) + .map_err(|e| Error::Validation(format!("graph_view id list: {e}")))?; + let mut memory_nodes: HashMap = HashMap::new(); + { + let mut stmt = conn + .prepare( + "SELECT id, content, created_at FROM memories \ + WHERE id IN (SELECT value FROM json_each(?1))", + ) + .map_err(|e| Error::db("prepare graph_view nodes", e))?; + let rows = stmt + .query_map(params![id_list], |r| { + Ok(( + r.get::<_, String>(0)?, + r.get::<_, String>(1)?, + r.get::<_, String>(2)?, + )) + }) + .map_err(|e| Error::db("query graph_view nodes", e))?; + for row in rows { + let (id, content, created_at) = + row.map_err(|e| Error::db("read graph_view node", e))?; + memory_nodes.insert( + id.clone(), + GraphNode { + label: memory_label(&content, &id), + entity_type: Some("memory".to_string()), + properties: serde_json::Value::Null, + memory_id: Some(id.clone()), + created_at, + id, + }, + ); + } + } + + let mut nodes: Vec = Vec::new(); + for id in &ids { + if let Some(n) = memory_nodes.remove(*id) { + nodes.push(n); + } else if let Some(n) = explicit_by_key.get(id) { + // Entity node of the explicit graph (no memory behind it). + let mut n = (*n).clone(); + n.id = (*id).to_string(); + nodes.push(n); + } + } + // Edges whose endpoint is neither a live memory nor a known entity + // node cannot be drawn; drop them rather than ship dangling ids. + let known: BTreeSet<&str> = nodes.iter().map(|n| n.id.as_str()).collect(); + edges.retain(|e| { + known.contains(e.source_id.as_str()) && known.contains(e.target_id.as_str()) + }); + + let mut relation_types: Vec = edges.iter().map(|e| e.relation.clone()).collect(); + relation_types.sort(); + relation_types.dedup(); + let stats = GraphStats { + node_count: nodes.len(), + edge_count: edges.len(), + relation_types, + }; + + Ok(GraphView { + nodes, + edges, + stats, + truncated, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::GraphStore; + + fn open() -> Uteke { + Uteke::open(":memory:").unwrap() + } + + /// Embedder-free seeding (CI has no ONNX runtime). Vectors are opaque here. + fn mem(u: &Uteke, text: &str, ns: &str) -> String { + u.remember_precomputed( + text, + &[], + None, + Some(ns), + "fact", + "text", + &vec![0.1_f32; 768], + ) + .unwrap() + } + + fn rels(v: &GraphView, rel: &str) -> Vec<(String, String)> { + v.edges + .iter() + .filter(|e| e.relation == rel) + .map(|e| (e.source_id.clone(), e.target_id.clone())) + .collect() + } + + #[test] + fn memory_edges_appear_keyed_by_memory_id() { + let u = open(); + let a = mem(&u, "alpha decision about caching", "gv"); + let b = mem(&u, "beta follow up", "gv"); + assert!(u.link_memories(&a, &b, "related").unwrap()); + + let v = u.graph_view(Some("gv"), None, 100).unwrap(); + assert!(rels(&v, "related").contains(&(a.clone(), b.clone()))); + let ids: Vec<&str> = v.nodes.iter().map(|n| n.id.as_str()).collect(); + assert!(ids.contains(&a.as_str()) && ids.contains(&b.as_str())); + let node_a = v.nodes.iter().find(|n| n.id == a).unwrap(); + assert!(node_a.label.starts_with("alpha decision about caching")); + assert_eq!(node_a.entity_type.as_deref(), Some("memory")); + // the short names viewers read mirror the long ones + let e = v.edges.iter().find(|e| e.relation == "related").unwrap(); + assert_eq!((&e.source, &e.target), (&e.source_id, &e.target_id)); + assert_eq!(v.stats.edge_count, v.edges.len()); + assert_eq!(v.stats.node_count, v.nodes.len()); + } + + #[test] + fn derived_backlinks_deprecated_and_other_namespaces_are_left_out() { + let u = open(); + let a = mem(&u, "live a", "gv1"); + let b = mem(&u, "live b", "gv1"); + let dead = mem(&u, "will be deprecated", "gv1"); + let other = mem(&u, "other namespace", "gv2"); + u.link_memories(&a, &b, "related").unwrap(); // also writes a referenced_by backlink + u.link_memories(&a, &dead, "related").unwrap(); + u.link_memories(&a, &other, "related").unwrap(); + u.soft_forget(&dead, "test").unwrap(); + + let v = u.graph_view(Some("gv1"), None, 100).unwrap(); + assert!( + rels(&v, "referenced_by").is_empty(), + "backlinks are omitted" + ); + assert!( + !v.edges + .iter() + .any(|e| e.target_id == dead || e.source_id == dead) + ); + assert!(!v.nodes.iter().any(|n| n.id == dead || n.id == other)); + assert!(rels(&v, "related").contains(&(a.clone(), b.clone()))); + + // Without a namespace the cross-namespace edge is visible, deprecated still not. + let all = u.graph_view(None, None, 100).unwrap(); + assert!(rels(&all, "related").contains(&(a.clone(), other.clone()))); + assert!(!all.nodes.iter().any(|n| n.id == dead)); + } + + #[test] + fn node_id_returns_only_that_neighbourhood() { + let u = open(); + let a = mem(&u, "hub", "gv"); + let b = mem(&u, "spoke one", "gv"); + let c = mem(&u, "spoke two", "gv"); + let d = mem(&u, "unrelated x", "gv"); + let e = mem(&u, "unrelated y", "gv"); + u.link_memories(&a, &b, "related").unwrap(); + u.link_memories(&c, &a, "related").unwrap(); + u.link_memories(&d, &e, "related").unwrap(); + + let v = u.graph_view(Some("gv"), Some(&a), 100).unwrap(); + let rel = rels(&v, "related"); + assert!(rel.contains(&(a.clone(), b.clone())) && rel.contains(&(c.clone(), a.clone()))); + assert!(!rel.contains(&(d.clone(), e.clone()))); + assert!(v.edges.iter().all(|x| x.source_id == a || x.target_id == a)); + } + + #[test] + fn limit_bounds_edges_and_reports_truncation() { + let u = open(); + let hub = mem(&u, "hub", "gv"); + for i in 0..6 { + let m = mem(&u, &format!("spoke {i}"), "gv"); + u.link_memories(&hub, &m, "related").unwrap(); + } + let few = u.graph_view(Some("gv"), None, 3).unwrap(); + assert_eq!(few.edges.len(), 3); + assert!(few.truncated); + // every returned edge has both endpoints in `nodes` + let ids: std::collections::HashSet<&str> = + few.nodes.iter().map(|n| n.id.as_str()).collect(); + assert!( + few.edges + .iter() + .all(|e| ids.contains(e.source_id.as_str()) && ids.contains(e.target_id.as_str())) + ); + let all = u.graph_view(Some("gv"), None, 1000).unwrap(); + assert!(!all.truncated); + assert!(all.edges.len() >= 6); + // limit 0 is clamped, not an error + assert!(u.graph_view(None, None, 0).unwrap().edges.len() <= 1); + } + + #[test] + fn explicit_graph_edges_are_rekeyed_to_memory_ids() { + let u = open(); + let a = mem(&u, "explicit a", "gv"); + let b = mem(&u, "explicit b", "gv"); + GraphStore::new(u.graph_store()) + .add_edge_for_memories(&a, &b, "contradicts", 0.9) + .unwrap(); + + let v = u.graph_view(Some("gv"), None, 100).unwrap(); + assert!( + rels(&v, "contradicts").contains(&(a.clone(), b.clone())), + "explicit edge must use memory ids, got {:?}", + v.edges + ); + let e = v + .edges + .iter() + .find(|e| e.relation == "contradicts") + .unwrap(); + assert!((e.weight - 0.9).abs() < 1e-9); + // one node per memory even though both stores know it + assert_eq!(v.nodes.iter().filter(|n| n.id == a).count(), 1); + } + + /// Corin (`src-tauri/src/uteke_client.rs`) parses /graph into these + /// shapes. Before #1366 edges carried only `source_id`/`target_id`, so a + /// non-empty graph failed to deserialize there. + #[test] + fn response_deserializes_into_the_shape_corin_reads() { + #[derive(Deserialize)] + #[allow(dead_code)] + struct CorinNode { + id: String, + label: String, + entity_type: Option, + } + #[derive(Deserialize)] + #[allow(dead_code)] + struct CorinEdge { + source: String, + target: String, + relation: String, + weight: f32, + } + #[derive(Deserialize)] + #[allow(dead_code)] + struct CorinGraph { + nodes: Vec, + edges: Vec, + stats: GraphStats, + } + + let u = open(); + let a = mem(&u, "corin a", "gv"); + let b = mem(&u, "corin b", "gv"); + u.link_memories(&a, &b, "related").unwrap(); + let json = serde_json::to_value(u.graph_view(Some("gv"), None, 50).unwrap()).unwrap(); + let parsed: CorinGraph = serde_json::from_value(json).expect("Corin shape"); + assert!(parsed.edges.iter().any(|e| e.source == a && e.target == b)); + } + + #[test] + fn add_edge_with_raw_memory_ids_violates_fk_but_the_helper_works() { + // #1365: this is the exact call MCP uteke_graph_add_edge and dream made. + let u = open(); + let a = mem(&u, "fk a", "gv"); + let b = mem(&u, "fk b", "gv"); + let gs = GraphStore::new(u.graph_store()); + assert!( + gs.add_edge(&a, &b, "related", 1.0).is_err(), + "raw memory ids must fail the graph_nodes FK" + ); + gs.add_edge_for_memories(&a, &b, "related", 1.0).unwrap(); + assert!( + gs.remove_edge_between(&a, &b).unwrap(), + "removable by memory ids" + ); + assert!( + !gs.remove_edge_between(&a, &b).unwrap(), + "second removal is a no-op" + ); + assert!(!gs.remove_edge_between("nope", &b).unwrap()); + } +} diff --git a/crates/uteke-core/src/import_export.rs b/crates/uteke-core/src/import_export.rs index 2406b802..98206620 100644 --- a/crates/uteke-core/src/import_export.rs +++ b/crates/uteke-core/src/import_export.rs @@ -1,6 +1,7 @@ //! Import and export memories in JSONL format. use crate::error::Error; +use crate::index_sync::{self, SyncPolicy}; use crate::memory::types::{ExportEntry, ImportResult, Memory}; impl crate::Uteke { @@ -138,7 +139,7 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during import"))?; - index.insert(&id, &embedding)?; + index_sync::insert_unsaved(&mut *index, &id, &embedding, &SyncPolicy::STANDARD)?; // Don't save per-item — we'll persist once after the full import. } @@ -148,7 +149,7 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during import rollback"))?; - index.remove(&id); + index_sync::remove_unsaved(&mut *index, &id); // Note: don't save per-entry — save once at end of import. // If process crashes, orphan entry is harmless and cleaned by repair. tracing::warn!("Skipping import entry (id={id}): {e}"); @@ -169,7 +170,7 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during import save"))?; - index.save()?; + index_sync::persist(&mut *index, &SyncPolicy::STANDARD)?; } if skipped > 0 || deduped > 0 { diff --git a/crates/uteke-core/src/index_sync.rs b/crates/uteke-core/src/index_sync.rs new file mode 100644 index 00000000..d7ca2350 --- /dev/null +++ b/crates/uteke-core/src/index_sync.rs @@ -0,0 +1,492 @@ +//! Keeping the vector index in step with SQLite — one policy, one place. +//! +//! SQLite is the source of truth; the vector index is a derived, repairable +//! structure (`uteke repair` rebuilds it). Every mutating operation therefore +//! commits to SQLite first and then mirrors the change into the index. That +//! mirroring used to be hand-written at ~60 call sites with five different +//! failure policies (swallowed, retried inline, retried through a helper, +//! warn-only, propagated). This module owns the policy: +//! +//! * inserts are retried a bounded number of times, except for errors that +//! can never succeed (a dimension mismatch is deterministic); +//! * persisting (`save`) is retried a bounded number of times; +//! * the caller gets a [`SyncOutcome`] that says what actually happened, and +//! decides whether that is an error (a `forget` that could not persist the +//! index) or a warning (a `remember` stored FTS5-only). +//! +//! The functions work on a `&mut impl IndexWriter` the caller already holds +//! the write lock for, so the store write and the index write stay ordered +//! under one lock. [`IndexWriter`] is the seam: production uses +//! [`crate::memory::vector::VectorIndex`], tests use a recording/failing fake. + +use crate::Error; +use crate::memory::vector::VectorIndex; +use std::time::Duration; + +/// The operations the sync policy needs from an index. +pub(crate) trait IndexWriter { + fn insert(&mut self, id: &str, embedding: &[f32]) -> Result<(), Error>; + /// Returns `false` when `id` was not in the index. + fn remove(&mut self, id: &str) -> bool; + fn save(&mut self) -> Result<(), Error>; + /// Save only if enough has accumulated (per-operation hot paths, #1322). + /// Adapters without batching just save every time. + fn save_deferred(&mut self) -> Result<(), Error> { + self.save() + } + fn dims(&self) -> usize; +} + +impl IndexWriter for VectorIndex { + fn insert(&mut self, id: &str, embedding: &[f32]) -> Result<(), Error> { + VectorIndex::insert(self, id, embedding) + } + fn remove(&mut self, id: &str) -> bool { + VectorIndex::remove(self, id) + } + fn save(&mut self) -> Result<(), Error> { + VectorIndex::save(self) + } + fn save_deferred(&mut self) -> Result<(), Error> { + VectorIndex::save_if_due(self) + } + fn dims(&self) -> usize { + VectorIndex::dims(self) + } +} + +/// How hard to try before giving up on one index operation. +#[derive(Clone, Copy, Debug)] +pub(crate) struct SyncPolicy { + pub attempts: u32, + pub backoff: Duration, +} + +impl SyncPolicy { + /// Production policy: 3 attempts, 200 ms apart (#621, #1273). + pub(crate) const STANDARD: SyncPolicy = SyncPolicy { + attempts: 3, + backoff: Duration::from_millis(200), + }; +} + +/// What actually happened to the index. SQLite has already been written by +/// the time this is produced, so none of these are fatal by themselves. +#[derive(Debug, Default, PartialEq, Eq)] +pub(crate) struct SyncOutcome { + /// The vector insert failed (after retries). The row stays FTS5-searchable. + pub insert_error: Option, + /// The index could not be written to disk (after retries). The in-memory + /// index is correct; the on-disk copy is stale until the next save. + pub persist_error: Option, + /// Removals requested for ids that were not in the index. + pub missing: usize, + attempts: u32, +} + +impl SyncOutcome { + #[cfg(test)] + pub(crate) fn is_clean(&self) -> bool { + self.insert_error.is_none() && self.persist_error.is_none() + } + + /// One-line description of the first failure, for "honest outcome" + /// reporting (`RememberOutcome.vector_error`, warnings). + pub(crate) fn error_message(&self) -> Option { + let after = if self.attempts > 1 { + format!(" after {} attempts", self.attempts) + } else { + String::new() + }; + if let Some(e) = &self.insert_error { + return Some(format!("vector insert failed{after}: {e}")); + } + self.persist_error + .as_ref() + .map(|e| format!("index persist failed{after}: {e}")) + } + + /// Log a warning when the on-disk index could not be saved; for callers + /// where a stale file is repairable and not worth failing the operation. + pub(crate) fn warn_if_stale(&self, what: &str) { + if let Some(msg) = self.error_message() { + tracing::warn!("{what}: {msg}. `uteke repair` / `verify` resync the index."); + } + } + + /// For operations whose SQLite change is already committed and where an + /// unsaved index must be surfaced to the caller as an error. + pub(crate) fn into_delete_result(self, what: &str) -> Result<(), Error> { + match self.error_message() { + None => Ok(()), + Some(msg) => { + tracing::error!( + "{what}: SQLite rows changed but the vector index is stale ({msg}). \ + Run `uteke repair` to resync." + ); + Err(Error::embed_msg(format!( + "{what}: database rows were changed, but the vector index could not be \ + updated: {msg}. Run `uteke repair` to resync the index." + ))) + } + } + } +} + +fn retry( + policy: &SyncPolicy, + what: &str, + mut op: impl FnMut() -> Result, +) -> Result { + let mut attempt = 1; + loop { + match op() { + Ok(v) => return Ok(v), + // Deterministic failures (bad input) can never succeed on retry, and + // the backoff sleeps while the caller holds the index lock (#1322). + Err(e) if attempt < policy.attempts && !matches!(e, Error::Validation(_)) => { + tracing::warn!( + "{what}: attempt {attempt}/{} failed: {e}. Retrying...", + policy.attempts + ); + std::thread::sleep(policy.backoff); + attempt += 1; + } + Err(e) => return Err(e), + } + } +} + +/// Save the index to disk with the standard retry. `Err` carries the last +/// error after all attempts. +pub(crate) fn persist(index: &mut I, policy: &SyncPolicy) -> Result<(), Error> { + retry(policy, "index save", || index.save()) +} + +/// Like [`persist`] but lets the index batch: used by per-operation paths +/// (`remember`, `forget`) so a burst of writes does not rewrite the whole +/// index each time (#1322). Pending changes are flushed by `shutdown()`, +/// the index's `Drop`, and the server's periodic flush. +pub(crate) fn persist_deferred( + index: &mut I, + policy: &SyncPolicy, +) -> Result<(), Error> { + retry(policy, "index save", || index.save_deferred()) +} + +/// Insert (or replace) `id` and persist the index. +/// +/// A dimension mismatch is deterministic and is reported immediately instead +/// of being retried. An empty embedding means "no embedder configured": the +/// row simply has no vector entry and the outcome is clean. +pub(crate) fn upsert( + index: &mut I, + id: &str, + embedding: &[f32], + policy: &SyncPolicy, +) -> SyncOutcome { + let mut outcome = SyncOutcome { + attempts: policy.attempts, + ..SyncOutcome::default() + }; + if embedding.is_empty() { + return outcome; + } + if embedding.len() != index.dims() { + outcome.attempts = 1; // deterministic: never retried + outcome.insert_error = Some(format!( + "embedding has {} dimensions, index expects {}", + embedding.len(), + index.dims() + )); + return outcome; + } + if let Err(e) = retry(policy, &format!("vector insert id={id}"), || { + index.insert(id, embedding) + }) { + tracing::warn!( + "Vector insert failed for id={id}: {e}. Row stays FTS5-only; run `uteke repair` to backfill." + ); + outcome.insert_error = Some(e.to_string()); + return outcome; + } + if let Err(e) = persist_deferred(index, policy) { + tracing::warn!( + "Failed to persist vector index after insert id={id}: {e}. \ + The entry can be rebuilt via `uteke repair`." + ); + outcome.persist_error = Some(e.to_string()); + } + outcome +} + +/// Insert without persisting, for batch paths that save once at the end +/// (import, re-embed). Deterministic dimension errors are not retried. +pub(crate) fn insert_unsaved( + index: &mut I, + id: &str, + embedding: &[f32], + policy: &SyncPolicy, +) -> Result<(), Error> { + if embedding.len() != index.dims() { + return Err(Error::embed_msg(format!( + "embedding has {} dimensions, index expects {}", + embedding.len(), + index.dims() + ))); + } + retry(policy, &format!("vector insert id={id}"), || { + index.insert(id, embedding) + }) +} + +/// Remove without persisting (batch paths and rollbacks). Returns `false` +/// when `id` was not in the index. +pub(crate) fn remove_unsaved(index: &mut I, id: &str) -> bool { + index.remove(id) +} + +/// Remove every id from the index and persist once. +pub(crate) fn remove_ids<'a, I: IndexWriter>( + index: &mut I, + ids: impl IntoIterator, + policy: &SyncPolicy, +) -> SyncOutcome { + let mut outcome = SyncOutcome { + attempts: policy.attempts, + ..SyncOutcome::default() + }; + for id in ids { + if !index.remove(id) { + outcome.missing += 1; + } + } + if outcome.missing > 0 { + tracing::debug!( + "{} index entr{} not found during removal (ok if never embedded)", + outcome.missing, + if outcome.missing == 1 { "y" } else { "ies" } + ); + } + if let Err(e) = persist_deferred(index, policy) { + outcome.persist_error = Some(e.to_string()); + } + outcome +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; + + const FAST: SyncPolicy = SyncPolicy { + attempts: 3, + backoff: Duration::ZERO, + }; + + /// Second adapter at the seam: records calls and fails on demand. + #[derive(Default)] + struct FakeIndex { + dims: usize, + present: HashSet, + insert_calls: u32, + save_calls: u32, + insert_failures_left: u32, + save_failures_left: u32, + insert_validation_error: bool, + } + + impl FakeIndex { + fn new(dims: usize) -> Self { + Self { + dims, + ..Self::default() + } + } + } + + impl IndexWriter for FakeIndex { + fn insert(&mut self, id: &str, _e: &[f32]) -> Result<(), Error> { + self.insert_calls += 1; + if self.insert_validation_error { + return Err(Error::validation("deterministic")); + } + if self.insert_failures_left > 0 { + self.insert_failures_left -= 1; + return Err(Error::embed_msg("insert boom")); + } + self.present.insert(id.to_string()); + Ok(()) + } + fn remove(&mut self, id: &str) -> bool { + self.present.remove(id) + } + fn save(&mut self) -> Result<(), Error> { + self.save_calls += 1; + if self.save_failures_left > 0 { + self.save_failures_left -= 1; + return Err(Error::embed_msg("save boom")); + } + Ok(()) + } + fn dims(&self) -> usize { + self.dims + } + } + + // #1322: a deterministic (validation) error is never retried, so the + // backoff sleep cannot run under the index lock for something that can + // never succeed. + #[test] + fn upsert_does_not_retry_validation_errors() { + let mut idx = FakeIndex::new(2); + idx.insert_validation_error = true; + let out = upsert(&mut idx, "a", &[0.0, 1.0], &FAST); + assert!(out.insert_error.is_some()); + assert_eq!(idx.insert_calls, 1, "validation errors must not be retried"); + } + + #[test] + fn upsert_clean_inserts_and_saves_once() { + let mut idx = FakeIndex::new(2); + let out = upsert(&mut idx, "a", &[0.0, 1.0], &FAST); + assert!(out.is_clean()); + assert!(idx.present.contains("a")); + assert_eq!((idx.insert_calls, idx.save_calls), (1, 1)); + } + + #[test] + fn upsert_retries_transient_insert_failure() { + let mut idx = FakeIndex::new(2); + idx.insert_failures_left = 2; + let out = upsert(&mut idx, "a", &[0.0, 1.0], &FAST); + assert!(out.is_clean()); + assert_eq!(idx.insert_calls, 3); + } + + #[test] + fn upsert_gives_up_after_policy_attempts_and_skips_save() { + let mut idx = FakeIndex::new(2); + idx.insert_failures_left = 99; + let out = upsert(&mut idx, "a", &[0.0, 1.0], &FAST); + assert_eq!(idx.insert_calls, 3); + assert_eq!( + idx.save_calls, 0, + "nothing to persist when the insert failed" + ); + assert!(out.insert_error.is_some() && out.persist_error.is_none()); + assert!( + out.error_message() + .unwrap() + .contains("vector insert failed after 3 attempts") + ); + } + + #[test] + fn dimension_mismatch_is_not_retried() { + let mut idx = FakeIndex::new(4); + let out = upsert(&mut idx, "a", &[0.0, 1.0], &FAST); + assert_eq!( + idx.insert_calls, 0, + "deterministic error: no insert, no retries" + ); + assert!( + out.error_message() + .unwrap() + .starts_with("vector insert failed: embedding has 2 dimensions") + ); + } + + #[test] + fn empty_embedding_is_a_clean_noop() { + let mut idx = FakeIndex::new(2); + let out = upsert(&mut idx, "a", &[], &FAST); + assert!(out.is_clean()); + assert_eq!((idx.insert_calls, idx.save_calls), (0, 0)); + } + + #[test] + fn persist_failure_is_reported_not_swallowed() { + let mut idx = FakeIndex::new(2); + idx.save_failures_left = 99; + let out = upsert(&mut idx, "a", &[0.0, 1.0], &FAST); + assert_eq!(idx.save_calls, 3); + assert!( + idx.present.contains("a"), + "in-memory index is still correct" + ); + assert!(out.insert_error.is_none() && out.persist_error.is_some()); + assert!( + out.error_message() + .unwrap() + .contains("index persist failed after 3 attempts") + ); + } + + #[test] + fn remove_ids_counts_missing_and_persists_once() { + let mut idx = FakeIndex::new(2); + idx.present.insert("a".into()); + idx.present.insert("b".into()); + let out = remove_ids(&mut idx, ["a", "b", "ghost"], &FAST); + assert!(out.is_clean()); + assert_eq!(out.missing, 1); + assert!(idx.present.is_empty()); + assert_eq!(idx.save_calls, 1, "one save for the whole batch"); + } + + #[test] + fn delete_result_is_an_error_only_when_persist_failed() { + let mut idx = FakeIndex::new(2); + idx.present.insert("a".into()); + assert!( + remove_ids(&mut idx, ["a"], &FAST) + .into_delete_result("forget") + .is_ok() + ); + + idx.present.insert("b".into()); + idx.save_failures_left = 99; + let err = remove_ids(&mut idx, ["b"], &FAST) + .into_delete_result("forget") + .unwrap_err() + .to_string(); + assert!(err.contains("uteke repair"), "{err}"); + } + + #[test] + fn insert_unsaved_never_saves_and_rejects_wrong_dims() { + let mut idx = FakeIndex::new(2); + insert_unsaved(&mut idx, "a", &[0.0, 1.0], &FAST).unwrap(); + assert_eq!((idx.insert_calls, idx.save_calls), (1, 0)); + assert!(idx.present.contains("a")); + + let err = insert_unsaved(&mut idx, "b", &[0.0], &FAST).unwrap_err(); + assert!(err.to_string().contains("dimensions")); + assert_eq!( + idx.insert_calls, 1, + "wrong dimension: no insert attempt, no retries" + ); + } + + #[test] + fn insert_unsaved_retries_transient_failures_then_persist_saves_once() { + let mut idx = FakeIndex::new(2); + idx.insert_failures_left = 2; + insert_unsaved(&mut idx, "a", &[0.0, 1.0], &FAST).unwrap(); + assert_eq!(idx.insert_calls, 3); + persist(&mut idx, &FAST).unwrap(); + assert_eq!(idx.save_calls, 1); + } + + #[test] + fn works_with_the_real_vector_index() { + let mut idx = VectorIndex::new(4).unwrap(); + let out = upsert(&mut idx, "a", &[1.0, 0.0, 0.0, 0.0], &FAST); + assert!(out.is_clean()); + assert_eq!(idx.len(), 1); + let out = remove_ids(&mut idx, ["a", "ghost"], &FAST); + assert!(out.is_clean()); + assert_eq!((idx.len(), out.missing), (0, 1)); + } +} diff --git a/crates/uteke-core/src/lib.rs b/crates/uteke-core/src/lib.rs index f4376a60..27f6289f 100644 --- a/crates/uteke-core/src/lib.rs +++ b/crates/uteke-core/src/lib.rs @@ -10,6 +10,7 @@ //! ``` pub mod chunker; +pub mod config_layers; mod consolidate; pub mod consolidation_api; pub mod consolidation_exec; @@ -21,8 +22,10 @@ mod error; pub mod extraction; pub mod graph; pub mod graph_rerank; +pub mod graph_view; pub mod guide; mod import_export; +mod index_sync; mod jaccard; mod maintenance; pub mod memory; @@ -79,7 +82,8 @@ pub use embed::Embedder; pub use embed::OnnxEmbedder; pub use error::{Error, format_bytes}; pub use types::{ - DoctorCheck, DoctorReport, DoctorStatus, ReembedReport, RepairReport, VerifyReport, + CompactionReport, DoctorCheck, DoctorReport, DoctorStatus, ReembedReport, RepairReport, + VerifyReport, }; /// Maximum memory content length (characters) — default, overridable via config (#404). @@ -636,11 +640,20 @@ impl Uteke { .index .write() .map_err(|_| Error::lock("index write lock during consolidation insert"))?; - index.insert(id, embedding)?; - if let Err(e) = index.save() { - tracing::warn!("failed to persist vector index after insert id={id}: {e}"); + let sync = index_sync::upsert( + &mut *index, + id, + embedding, + &index_sync::SyncPolicy::STANDARD, + ); + match sync.insert_error { + Some(e) => Err(Error::embed_msg(format!( + "vector insert failed for id={id}: {e}" + ))), + // A failed persist is already logged by index_sync; the entry is + // in the in-memory index and `uteke repair` can resync the file. + None => Ok(()), } - Ok(()) } /// Remove an entry from the vector index (compensating action for a @@ -650,10 +663,9 @@ impl Uteke { let Ok(mut index) = self.index.write() else { return; }; - if index.remove(id) { - if let Err(e) = index.save() { - tracing::warn!("failed to persist vector index after remove id={id}: {e}"); - } + let sync = index_sync::remove_ids(&mut *index, [id], &index_sync::SyncPolicy::STANDARD); + if let Some(msg) = sync.error_message() { + tracing::warn!("failed to persist vector index after remove id={id}: {msg}"); } } @@ -896,6 +908,19 @@ impl Uteke { let mut index = match &index_path { Some(path) => match VectorIndex::load_or_create_for(path, dims, vector_backend) { Ok(idx) => idx, + // Lock contention / IO failure says nothing about the index + // content: another process (uteke-serve) may be holding a + // perfectly good index. Never delete it; run in-memory and + // rebuild from SQLite instead. + Err(e @ (Error::Lock { .. } | Error::Io(_))) => { + tracing::warn!( + "Vector index at {} could not be opened ({}). \ + Using an in-memory index rebuilt from SQLite; index files left untouched.", + path.display(), + e + ); + VectorIndex::new(dims)? + } Err(e) => { // Index file is corrupt (dim mismatch, truncated, etc). // Instead of crashing, discard the bad index and rebuild @@ -1603,73 +1628,117 @@ impl Uteke { has_children: false, }; - // Capture old chunk IDs BEFORE upsert (upsert_document deletes chunks - // internally as part of its transaction, so querying after returns empty). + // Embed-then-commit: chunk and embed with NO transaction open and NO index + // lock held (embedding is slow and can fail), then write the document and + // every chunk in ONE SQLite transaction, then mirror the result into the + // vector index. A failure before the commit leaves the previous version + // of the document, its chunks and the index untouched. + let chunk_doc_id = self + .store + .get_document_by_slug(&doc.slug)? + .map(|existing| existing.id) + .unwrap_or_else(|| doc.id.clone()); + let tag_strings: Vec = tags.iter().map(|t| t.to_string()).collect(); + let chunks = self.embed_document_chunks(&chunk_doc_id, content, &tag_strings)?; + + // Old chunk ids must be read BEFORE the upsert replaces them. let old_chunk_ids = self .store - .get_chunk_ids_for_documents(std::slice::from_ref(&doc.id))?; + .get_chunk_ids_for_documents(std::slice::from_ref(&chunk_doc_id))?; - let doc_id = self.store.upsert_document(&doc)?; + let doc_id = self.store.upsert_document_with_chunks(&doc, &chunks)?; + self.sync_document_chunks_to_index(&old_chunk_ids, &chunks)?; - // Chunk and embed the content. - self.ensure_embedder()?; - let embedder = self - .embedder - .lock() - .map_err(|_| Error::lock("embedder lock during document chunking"))?; - let embedder = embedder.as_ref().expect("embedder ensured above"); + tracing::info!("Document '{slug}' upserted: {} chunks", chunks.len()); + + Ok(doc_id) + } - let max_chars = embedder.max_seq_len().saturating_mul(4).max(1024); + /// Chunk `content` and embed every chunk. + /// + /// Deliberately takes no index lock and opens no transaction: embedding is + /// slow and may fail, and must not block recall or leave partial writes. + fn embed_document_chunks( + &self, + doc_id: &str, + content: &str, + tags: &[String], + ) -> Result)>, Error> { + self.ensure_embedder()?; + // The embedder mutex is shared with recall/remember (query embedding), + // so take it per call, never across the whole loop: a long document + // must not make every other embed wait for all of its chunks (#1323). + let max_chars = { + let guard = self + .embedder + .lock() + .map_err(|_| Error::lock("embedder lock during document chunking"))?; + guard + .as_ref() + .expect("embedder ensured above") + .max_seq_len() + .saturating_mul(4) + .max(1024) + }; let chunks = crate::chunker::chunk_markdown(content, max_chars); + chunks + .iter() + .enumerate() + .map(|(i, chunk)| { + let embedding = { + let guard = self + .embedder + .lock() + .map_err(|_| Error::lock("embedder lock during chunk embedding"))?; + guard + .as_ref() + .expect("embedder ensured above") + .embed(&chunk.content)? + }; + Ok(( + DocumentChunk { + id: uuid::Uuid::now_v7().to_string(), + document_id: doc_id.to_string(), + chunk_index: i as i64, + heading: chunk.heading.clone(), + content: chunk.content.clone(), + char_start: chunk.char_start as i64, + char_end: chunk.char_end as i64, + tags: tags.to_vec(), + }, + embedding, + )) + }) + .collect() + } - // Acquire usearch write lock for chunk index inserts. + /// Mirror a committed chunk replacement into the vector index + /// (`chunk:` keys): drop the replaced chunks, add the new ones, save + /// once. SQLite is already committed, so index trouble only warns — + /// `uteke repair` rebuilds chunk vectors from SQLite. + fn sync_document_chunks_to_index( + &self, + old_chunk_ids: &[String], + new_chunks: &[(DocumentChunk, Vec)], + ) -> Result<(), Error> { let mut index = self .index .write() .map_err(|_| Error::lock("index write lock during doc chunking"))?; - - // Remove old chunk entries from usearch. - for old_id in &old_chunk_ids { - let key = format!("chunk:{}", old_id); - index.remove(&key); - } - - for (i, chunk) in chunks.iter().enumerate() { - let chunk_id = uuid::Uuid::now_v7().to_string(); - let embedding = embedder.embed(&chunk.content)?; - - self.store.insert_document_chunk( - &DocumentChunk { - id: chunk_id.clone(), - document_id: doc_id.clone(), - chunk_index: i as i64, - heading: chunk.heading.clone(), - content: chunk.content.clone(), - char_start: chunk.char_start as i64, - char_end: chunk.char_end as i64, - tags: tags.iter().map(|t| t.to_string()).collect(), - }, - &embedding, - )?; - - // Insert chunk embedding into usearch with "chunk:" prefix. - let index_key = format!("chunk:{}", chunk_id); - if let Err(e) = index.insert(&index_key, &embedding) { - tracing::warn!( - "Failed to insert chunk {} into vector index: {}", - chunk_id, - e - ); + let policy = index_sync::SyncPolicy::STANDARD; + for old_id in old_chunk_ids { + index_sync::remove_unsaved(&mut *index, &format!("chunk:{old_id}")); + } + for (chunk, embedding) in new_chunks { + let key = format!("chunk:{}", chunk.id); + if let Err(e) = index_sync::insert_unsaved(&mut *index, &key, embedding, &policy) { + tracing::warn!("Failed to insert chunk {} into vector index: {e}", chunk.id); } } - - if let Err(e) = index.save() { - tracing::warn!("Failed to persist vector index after doc chunking: {}", e); + if let Err(e) = index_sync::persist(&mut *index, &policy) { + tracing::warn!("Failed to persist vector index after doc chunking: {e}"); } - - tracing::info!("Document '{slug}' upserted: {} chunks", chunks.len()); - - Ok(doc_id) + Ok(()) } /// Get a document by ID or slug. @@ -1701,74 +1770,31 @@ impl Uteke { }; let doc_id = doc.id.clone(); - // Partial update in SQLite. - let updated = self - .store - .update_document(&doc_id, title, content, tags, metadata)?; - - let updated = match updated { - Some(d) => d, - None => return Ok(None), - }; - - // If content was changed, rebuild chunks. - if let Some(content_text) = content { - let old_chunk_ids = self + // Content change: embed first (no locks, no transaction), then update the + // document and replace its chunks atomically, then mirror into the index. + let updated = if let Some(content_text) = content { + let chunk_tags: Vec = + tags.map(|t| t.to_vec()).unwrap_or_else(|| doc.tags.clone()); + let chunks = self.embed_document_chunks(&doc_id, content_text, &chunk_tags)?; + match self .store - .delete_chunks_for_documents(std::slice::from_ref(&doc_id))?; - - self.ensure_embedder()?; - let embedder = self - .embedder - .lock() - .map_err(|_| Error::lock("embedder lock during document update"))?; - let embedder = embedder.as_ref().expect("embedder ensured above"); - - let max_chars = embedder.max_seq_len().saturating_mul(4).max(1024); - let chunks = crate::chunker::chunk_markdown(content_text, max_chars); - - let mut index = self - .index - .write() - .map_err(|_| Error::lock("index write lock during doc update"))?; - - for old_id in &old_chunk_ids { - let key = format!("chunk:{}", old_id); - index.remove(&key); - } - - for (i, chunk) in chunks.iter().enumerate() { - let chunk_id = uuid::Uuid::now_v7().to_string(); - let embedding = embedder.embed(&chunk.content)?; - - self.store.insert_document_chunk( - &DocumentChunk { - id: chunk_id.clone(), - document_id: doc_id.clone(), - chunk_index: i as i64, - heading: chunk.heading.clone(), - content: chunk.content.clone(), - char_start: chunk.char_start as i64, - char_end: chunk.char_end as i64, - tags: updated.tags.clone(), - }, - &embedding, - )?; - - let index_key = format!("chunk:{}", chunk_id); - if let Err(e) = index.insert(&index_key, &embedding) { - tracing::warn!( - "Failed to insert chunk {} into index during update: {}", - chunk_id, - e - ); + .update_document_with_chunks(&doc_id, title, content, tags, metadata, &chunks)? + { + Some((updated, old_chunk_ids)) => { + self.sync_document_chunks_to_index(&old_chunk_ids, &chunks)?; + updated } + None => return Ok(None), } - - if let Err(e) = index.save() { - tracing::warn!("Failed to persist index after doc update: {}", e); + } else { + match self + .store + .update_document(&doc_id, title, content, tags, metadata)? + { + Some(updated) => updated, + None => return Ok(None), } - } + }; tracing::info!("Document '{id_or_slug}' updated"); Ok(Some(updated)) @@ -1786,6 +1812,18 @@ impl Uteke { self.store.list_documents_ns(namespace, limit) } + /// One page of documents for callers that must see ALL of them (export). + /// Pages are stable (newest first, id tie-break); `limit` is capped at + /// 1000 per page. + pub fn doc_list_page( + &self, + namespace: Option<&str>, + limit: usize, + offset: usize, + ) -> Result, Error> { + self.store.list_documents_page(namespace, limit, offset) + } + /// List root documents, optionally scoped to a namespace (#1268). pub fn doc_list_roots( &self, @@ -1902,11 +1940,13 @@ impl Uteke { .index .write() .map_err(|_| Error::lock("index write lock during doc delete"))?; - for chunk_id in &chunk_ids { - let key = format!("chunk:{}", chunk_id); - index.remove(&key); - } - let _ = index.save(); + let keys: Vec = chunk_ids.iter().map(|c| format!("chunk:{c}")).collect(); + index_sync::remove_ids( + &mut *index, + keys.iter().map(String::as_str), + &index_sync::SyncPolicy::STANDARD, + ) + .warn_if_stale("doc_delete"); } Ok((deleted, subtree_size)) @@ -2018,10 +2058,10 @@ impl Uteke { (DocumentSummary, String, String, f32), > = std::collections::HashMap::new(); - for ((_chunk_key, distance), (_chunk_id, doc_id, heading, content)) in - chunk_hits.iter().zip(chunks.iter()) + for (distance, (_chunk_id, doc_id, heading, content)) in + pair_chunk_hits(&chunk_hits, &chunks) { - let score = crate::memory::vector::cosine_distance_to_similarity(*distance); + let score = crate::memory::vector::cosine_distance_to_similarity(distance); // Get document summary from store. if let Ok(Some(doc)) = self.store.get_document(doc_id) { @@ -2656,31 +2696,62 @@ pub struct GraphData { } /// Resolve a path to a database string. +/// `(chunk_id, doc_id, heading, content)` as returned by `get_chunks_by_ids_ordered`. +type ChunkRow = (String, String, String, String); + +/// Pair each `chunk:` vector hit with its SQLite chunk row by id. +/// +/// `get_chunks_by_ids_ordered` silently drops ids missing from SQLite, so the +/// two lists can differ in length; zipping them would attribute every later +/// distance to the wrong chunk. Hits without a row are skipped. +fn pair_chunk_hits<'a>( + hits: &'a [(String, f32)], + chunks: &'a [ChunkRow], +) -> Vec<(f32, &'a ChunkRow)> { + let by_id: std::collections::HashMap<&str, &ChunkRow> = + chunks.iter().map(|c| (c.0.as_str(), c)).collect(); + hits.iter() + .filter_map(|(key, distance)| { + let id = key.strip_prefix("chunk:")?; + by_id.get(id).map(|c| (*distance, *c)) + }) + .collect() +} + fn resolve_db_path(db_path: &Path) -> Result { if db_path.to_str() == Some(":memory:") { return Ok(":memory:".to_string()); } if db_path.is_dir() || db_path.extension().is_none() { + let existed = db_path.exists(); std::fs::create_dir_all(db_path).map_err(Error::Io)?; - // Set directory permissions to owner-only (0700) on Unix - #[cfg(unix)] - { - let p: &std::path::Path = db_path; - std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o700)).ok(); + if !existed { + restrict_dir_permissions(db_path); } Ok(db_path.join("uteke.db").to_string_lossy().to_string()) } else { if let Some(parent) = db_path.parent() { + let existed = parent.exists(); std::fs::create_dir_all(parent).map_err(Error::Io)?; - // Set directory permissions to owner-only (0700) on Unix - #[cfg(unix)] - std::fs::set_permissions(parent, std::fs::Permissions::from_mode(0o700)).ok(); + if !existed { + restrict_dir_permissions(parent); + } } Ok(db_path.to_string_lossy().to_string()) } } +/// Owner-only (0700) permissions on Unix for a directory this process just +/// created. Pre-existing directories (e.g. `$HOME` or a shared dir that +/// merely contains the db file) are never touched. +fn restrict_dir_permissions(dir: &Path) { + #[cfg(unix)] + std::fs::set_permissions(dir, std::fs::Permissions::from_mode(0o700)).ok(); + #[cfg(not(unix))] + let _ = dir; +} + #[cfg(test)] mod tests { /// #1166: `open()` must not request a backend the build cannot @@ -2965,6 +3036,183 @@ mod tests { unsafe { std::env::remove_var("UTEKE_VECTOR_BACKEND") }; } + /// #1322: a burst of `remember` calls saves the index a handful of times, + /// not once per call, and nothing is lost once the store closes. + #[test] + #[serial_test::serial] + fn bulk_remember_batches_index_saves_without_losing_entries() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("uteke.db"); + unsafe { std::env::remove_var("UTEKE_VECTOR_BACKEND") }; + + let u = Uteke::open(&db).unwrap(); + for i in 0..100usize { + let mut v = vec![0.0f32; 768]; + v[i % 768] = 1.0; + u.remember_precomputed( + &format!("bulk probe {i}"), + &[], + None, + None, + "fact", + "text", + &v, + ) + .unwrap(); + } + let saves = u.index.read().unwrap().save_count(); + assert!( + saves <= 5, + "100 writes should batch into a few saves, got {saves}" + ); + u.shutdown().unwrap(); + drop(u); + + let u = Uteke::open(&db).unwrap(); + assert_eq!(u.index.read().unwrap().len(), 100, "no entry lost"); + u.shutdown().unwrap(); + } + + /// #1325: a crash between the index rename and the `.keys` rename leaves a + /// stale sidecar. Reopening must rebuild from SQLite so the newest memory + /// is still searchable, not silently drop it from the key map. + #[test] + #[serial_test::serial] + fn interrupted_save_pair_is_rebuilt_not_lost() { + let backends: Vec<&str> = [ + ("usearch", cfg!(feature = "usearch")), + ("vecq", cfg!(feature = "vecq")), + ] + .into_iter() + .filter(|(_, on)| *on) + .map(|(n, _)| n) + .collect(); + for name in backends { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("uteke.db"); + unsafe { std::env::set_var("UTEKE_VECTOR_BACKEND", name) }; + + let mk = |i: usize| { + let mut v = vec![0.0f32; 768]; + v[i] = 1.0; + v + }; + let remember = |u: &Uteke, text: &str, i: usize| { + u.remember_precomputed(text, &[], None, None, "fact", "text", &mk(i)) + .unwrap() + }; + + let u = Uteke::open(&db).unwrap(); + remember(&u, "first interrupted save probe", 0); + remember(&u, "second interrupted save probe", 1); + u.shutdown().unwrap(); + drop(u); + let keys = dir.path().join("uteke_index.keys"); + let keys = if keys.exists() { + keys + } else { + dir.path().join(format!("uteke_index.{}.keys", name)) + }; + let stale = std::fs::read(&keys).unwrap(); + + let u = Uteke::open(&db).unwrap(); + let newest = remember(&u, "third interrupted save probe", 2); + u.shutdown().unwrap(); + drop(u); + // Crash after the index rename, before the sidecar rename. + std::fs::write(&keys, stale).unwrap(); + + let u = Uteke::open(&db).unwrap(); + let hits = u.index.read().unwrap().search(&mk(2), 3, 50); + assert!( + hits.iter().any(|(id, _)| *id == newest), + "{name}: newest memory lost after interrupted save: {hits:?}" + ); + u.shutdown().unwrap(); + drop(u); + } + unsafe { std::env::remove_var("UTEKE_VECTOR_BACKEND") }; + } + + /// #1324: after many deletes a vecq index is mostly tombstones. Compaction + /// rebuilds it from SQLite so the file and the search over-fetch shrink + /// back to the live set. + #[cfg(feature = "vecq")] + #[test] + #[serial_test::serial] + fn vecq_compaction_shrinks_file_and_drops_dead_rows() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("uteke.db"); + unsafe { std::env::set_var("UTEKE_VECTOR_BACKEND", "vecq") }; + + let u = Uteke::open(&db).unwrap(); + let mut ids = Vec::new(); + for i in 0..200usize { + let mut v = vec![0.0f32; 768]; + v[i] = 1.0; + ids.push( + u.remember_precomputed( + &format!("compaction probe {i}"), + &[], + None, + None, + "fact", + "text", + &v, + ) + .unwrap(), + ); + } + for id in &ids[..150] { + u.forget(id).unwrap(); + } + u.flush_index().unwrap(); + let file = dir.path().join("uteke_index.vecq"); + let size_before = std::fs::metadata(&file).unwrap().len(); + { + let idx = u.index.read().unwrap(); + assert_eq!((idx.physical_rows(), idx.dead_rows()), (200, 150)); + assert!(idx.needs_compaction()); + } + + // Dry run reports without touching anything. + let dry = u.compact_index_if_needed(true).unwrap().unwrap(); + assert!(dry.dry_run); + assert_eq!((dry.rows_before, dry.dead_before), (200, 150)); + assert_eq!(u.index.read().unwrap().physical_rows(), 200); + + let report = u.compact_index_if_needed(false).unwrap().unwrap(); + assert_eq!( + (report.rows_before, report.dead_before, report.rows_after), + (200, 150, 50) + ); + { + let idx = u.index.read().unwrap(); + assert_eq!((idx.physical_rows(), idx.dead_rows()), (50, 0)); + assert!(!idx.needs_compaction()); + } + assert!( + std::fs::metadata(&file).unwrap().len() < size_before, + "index file must shrink" + ); + assert!( + u.compact_index_if_needed(false).unwrap().is_none(), + "nothing left to compact" + ); + + // A surviving memory is still found; reopening keeps the compacted index. + let mut q = vec![0.0f32; 768]; + q[199] = 1.0; + let hits = u.index.read().unwrap().search(&q, 3, 50); + assert_eq!(hits.first().map(|h| h.0.as_str()), Some(ids[199].as_str())); + u.shutdown().unwrap(); + drop(u); + let u = Uteke::open(&db).unwrap(); + assert_eq!(u.index.read().unwrap().physical_rows(), 50); + u.shutdown().unwrap(); + unsafe { std::env::remove_var("UTEKE_VECTOR_BACKEND") }; + } + /// #1168: invalid / not-compiled-in UTEKE_VECTOR_BACKEND falls back to the /// default engine without failing the open. #[test] @@ -3084,6 +3332,40 @@ mod tests { assert_eq!(restored.unique_tags, 5); } + #[test] + fn pair_chunk_hits_survives_missing_rows() { + let hits = vec![ + ("chunk:a".to_string(), 0.1_f32), + ("chunk:gone".to_string(), 0.2), + ("chunk:c".to_string(), 0.3), + ]; + let row = |id: &str| (id.to_string(), format!("doc-{id}"), "h".into(), "c".into()); + let chunks = vec![row("a"), row("c")]; + let paired = pair_chunk_hits(&hits, &chunks); + assert_eq!(paired.len(), 2); + // 0.3 must stay with chunk "c", not shift onto a neighbour. + assert_eq!((paired[1].0, paired[1].1.0.as_str()), (0.3, "c")); + } + + #[cfg(unix)] + #[test] + fn resolve_db_path_only_restricts_dirs_it_created() { + use std::os::unix::fs::PermissionsExt; + let base = tempfile::tempdir().unwrap(); + // Pre-existing shared directory holding the db file: must stay 0755. + let shared = base.path().join("shared"); + std::fs::create_dir(&shared).unwrap(); + std::fs::set_permissions(&shared, std::fs::Permissions::from_mode(0o755)).unwrap(); + resolve_db_path(&shared.join("x.db")).unwrap(); + let mode = std::fs::metadata(&shared).unwrap().permissions().mode() & 0o777; + assert_eq!(mode, 0o755); + // A directory created by the call is owner-only. + let fresh = base.path().join("fresh"); + resolve_db_path(&fresh).unwrap(); + let mode = std::fs::metadata(&fresh).unwrap().permissions().mode() & 0o777; + assert_eq!(mode, 0o700); + } + #[test] fn test_resolve_db_path_memory() { let path = std::path::Path::new(":memory:"); @@ -3918,3 +4200,258 @@ mod uuidv7_tests { assert_eq!(v7.get_version_num(), 7); } } + +#[cfg(test)] +mod doc_atomicity_tests { + //! #1341 — document upsert/update is embed-then-commit: a failure while + //! embedding must leave the previous version, its chunks and the index + //! untouched. + use super::*; + use crate::embed::Embedder; + use std::sync::Arc; + use std::sync::atomic::{AtomicI64, Ordering}; + + /// Succeeds for `ok_calls` embeds, then fails every embed. + struct BudgetEmbedder { + ok_calls: Arc, + } + + impl Embedder for BudgetEmbedder { + fn embed(&self, _text: &str) -> Result, Error> { + if self.ok_calls.fetch_sub(1, Ordering::SeqCst) > 0 { + Ok(vec![0.1f32, 0.2, 0.3, 0.4]) + } else { + Err(Error::Validation("simulated embed failure".to_string())) + } + } + fn dims(&self) -> usize { + 4 + } + fn max_seq_len(&self) -> usize { + 128 + } + fn name(&self) -> &str { + "budget-test-embedder" + } + } + + fn open(budget: &Arc) -> Uteke { + let (_db, store) = Uteke::open_store(":memory:").expect("open_store"); + Uteke::finish_open_full( + store, + Some(Box::new(BudgetEmbedder { + ok_calls: Arc::clone(budget), + })), + "test-budget".to_string(), + TierConfig::default(), + RecallConfig::default(), + EmbeddingSettings::default(), + crate::graph_rerank::GraphRerankConfig::default(), + None, + ) + .expect("finish_open_full") + } + + const V1: &str = + "# Alpha\nalpha body text\n\n# Beta\nbeta body text\n\n# Gamma\ngamma body text\n"; + const V2: &str = + "# One\nfirst replacement\n\n# Two\nsecond replacement\n\n# Three\nthird replacement\n"; + + fn chunk_ids(u: &Uteke, doc_id: &str) -> Vec { + let mut ids = u + .store + .get_chunk_ids_for_documents(&[doc_id.to_string()]) + .unwrap(); + ids.sort(); + ids + } + + #[test] + #[serial_test::serial] + fn upsert_failing_mid_document_keeps_previous_version_and_index() { + let budget = Arc::new(AtomicI64::new(1_000)); + let u = open(&budget); + let doc_id = u.doc_upsert("doc-a", "Title", V1, &[], None).unwrap(); + let before_chunks = chunk_ids(&u, &doc_id); + assert!( + before_chunks.len() >= 2, + "fixture must produce several chunks" + ); + let before_index = u.index.read().unwrap().len(); + let before = u.doc_get("doc-a").unwrap().unwrap(); + + // Allow exactly one embed of the replacement, then fail: the old code + // had already replaced the document by the time chunk 2 failed. + budget.store(1, Ordering::SeqCst); + let err = u.doc_upsert("doc-a", "Title", V2, &[], None); + assert!(err.is_err(), "embed failure must fail the upsert"); + + let after = u.doc_get("doc-a").unwrap().unwrap(); + assert_eq!(after.content, before.content, "document content untouched"); + assert_eq!(after.version, before.version, "version untouched"); + assert_eq!(chunk_ids(&u, &doc_id), before_chunks, "old chunks intact"); + assert_eq!( + u.index.read().unwrap().len(), + before_index, + "index untouched" + ); + } + + #[test] + #[serial_test::serial] + fn upsert_success_replaces_chunks_in_store_and_index() { + let budget = Arc::new(AtomicI64::new(1_000)); + let u = open(&budget); + let doc_id = u.doc_upsert("doc-b", "Title", V1, &[], None).unwrap(); + let old = chunk_ids(&u, &doc_id); + + u.doc_upsert("doc-b", "Title", V2, &[], None).unwrap(); + let new = chunk_ids(&u, &doc_id); + assert!(!new.is_empty()); + assert!( + new.iter().all(|id| !old.contains(id)), + "old chunks replaced" + ); + assert_eq!( + u.index.read().unwrap().len(), + new.len(), + "index holds exactly the new chunk vectors" + ); + } + + #[test] + #[serial_test::serial] + fn update_failing_mid_document_keeps_previous_content_and_chunks() { + let budget = Arc::new(AtomicI64::new(1_000)); + let u = open(&budget); + let doc_id = u.doc_upsert("doc-c", "Title", V1, &[], None).unwrap(); + let before_chunks = chunk_ids(&u, &doc_id); + let before = u.doc_get("doc-c").unwrap().unwrap(); + + budget.store(1, Ordering::SeqCst); + let res = u.doc_update("doc-c", None, Some(V2), None, None); + assert!(res.is_err()); + + let after = u.doc_get("doc-c").unwrap().unwrap(); + assert_eq!(after.content, before.content); + assert_eq!(after.version, before.version); + assert_eq!(chunk_ids(&u, &doc_id), before_chunks); + } + + #[test] + #[serial_test::serial] + fn upsert_without_a_working_embedder_writes_nothing() { + let budget = Arc::new(AtomicI64::new(0)); + let u = open(&budget); + assert!(u.doc_upsert("doc-d", "Title", V1, &[], None).is_err()); + assert!( + u.doc_get("doc-d").unwrap().is_none(), + "no half-written document row when embedding is impossible" + ); + } +} + +#[cfg(test)] +mod documents_fts_heal_tests { + //! #1349 — databases created with the old plain-UPDATE/DELETE + //! `documents_fts_*` triggers must heal and stop failing title changes. + use super::*; + use crate::embed::Embedder; + + struct FixedEmbedder; + impl Embedder for FixedEmbedder { + fn embed(&self, _text: &str) -> Result, Error> { + Ok(vec![0.1, 0.2, 0.3, 0.4]) + } + fn dims(&self) -> usize { + 4 + } + fn max_seq_len(&self) -> usize { + 128 + } + fn name(&self) -> &str { + "fixed-test-embedder" + } + } + + fn open() -> Uteke { + let (_db, store) = Uteke::open_store(":memory:").expect("open_store"); + Uteke::finish_open_full( + store, + Some(Box::new(FixedEmbedder)), + "test-fixed".to_string(), + TierConfig::default(), + RecallConfig::default(), + EmbeddingSettings::default(), + crate::graph_rerank::GraphRerankConfig::default(), + None, + ) + .expect("finish_open_full") + } + + const BODY: &str = "# Alpha\nalpha body text\n\n# Beta\nbeta body text\n"; + + fn title_hits(u: &Uteke, q: &str) -> Vec { + u.store + .search_documents_fts(q, 10) + .unwrap() + .into_iter() + .map(|d| d.title) + .collect() + } + + #[test] + #[serial_test::serial] + fn title_change_works_and_fts_follows() { + let u = open(); + u.doc_upsert("doc-t", "Zeppelin Guide", BODY, &[], None) + .unwrap(); + assert_eq!(title_hits(&u, "zeppelin"), vec!["Zeppelin Guide"]); + + u.doc_update("doc-t", Some("Dirigible Guide"), None, None, None) + .unwrap() + .expect("document exists"); + u.doc_upsert("doc-t", "Blimp Guide", "# One\nfirst\n", &[], None) + .unwrap(); + + assert!(title_hits(&u, "zeppelin").is_empty()); + assert!(title_hits(&u, "dirigible").is_empty()); + assert_eq!(title_hits(&u, "blimp"), vec!["Blimp Guide"]); + } + + #[test] + #[serial_test::serial] + fn old_triggers_are_healed() { + let u = open(); + // Recreate a pre-fix database: plain UPDATE / DELETE on the + // external-content FTS table. + u.store + .conn + .execute_batch( + "DROP TRIGGER documents_fts_update; + DROP TRIGGER documents_fts_delete; + CREATE TRIGGER documents_fts_update AFTER UPDATE ON documents BEGIN UPDATE documents_fts SET title = new.title, slug = new.slug, content = new.content WHERE rowid = new.rowid; END; + CREATE TRIGGER documents_fts_delete AFTER DELETE ON documents BEGIN DELETE FROM documents_fts WHERE rowid = old.rowid; END;", + ) + .unwrap(); + u.doc_upsert("doc-h", "Original Title", BODY, &[], None) + .unwrap(); + let broken = u.doc_update("doc-h", Some("Changed Title"), None, None, None); + assert!( + broken.is_err() || !title_hits(&u, "original").is_empty(), + "fixture must reproduce the bug (error, or a stale FTS entry) before healing" + ); + + // Healing replaces the triggers and rebuilds the FTS index. + u.store.heal_documents_fts_triggers().unwrap(); + u.doc_update("doc-h", Some("Changed Title"), None, None, None) + .unwrap() + .expect("document exists"); + assert_eq!(title_hits(&u, "changed"), vec!["Changed Title"]); + assert!( + title_hits(&u, "original").is_empty(), + "stale pre-fix FTS entries are gone after the rebuild" + ); + u.store.heal_documents_fts_triggers().unwrap(); // idempotent + } +} diff --git a/crates/uteke-core/src/maintenance.rs b/crates/uteke-core/src/maintenance.rs index 20177ff9..0672a0b0 100644 --- a/crates/uteke-core/src/maintenance.rs +++ b/crates/uteke-core/src/maintenance.rs @@ -1,14 +1,19 @@ //! Maintenance operations: doctor, verify, repair, stats, aging, prune, shutdown. use crate::error::{Error, format_bytes}; +use crate::index_sync::{self, SyncPolicy}; use crate::memory::types::{ AgingStatus, CleanupResult, LifecycleCycleResult, Memory, PruneResult, StoreStats, }; use crate::types::{ - DoctorCheck, DoctorReport, DoctorStatus, ReembedReport, RepairReport, VerifyReport, + CompactionReport, DoctorCheck, DoctorReport, DoctorStatus, ReembedReport, RepairReport, + VerifyReport, }; use crate::uteke_home; +/// `(index key, embedding)` pairs the vector index is (re)built from. +type IndexItems = Vec<(String, Vec)>; + impl crate::Uteke { /// Check system health: DB, index, model, consistency. pub fn doctor(&self) -> Result { @@ -133,27 +138,7 @@ impl crate::Uteke { index.len() }; - // Load all from SQLite and rebuild index (NULL embeddings filtered in load_all) - let all_memories = self.store.load_all(None)?; - let mut items: Vec<(String, Vec)> = all_memories - .iter() - .filter(|m| !m.embedding.is_empty()) - .map(|m| (m.id.clone(), m.embedding.clone())) - .collect(); - - // Document chunk vectors live in the index under "chunk:" keys. - // load_all returns memories only — without this the rebuild silently - // evicts every chunk entry (#1110). Chunk embeddings are persisted in - // document_chunks, so no re-embedding is needed. - let chunk_count = { - let chunks = self.store.load_all_chunk_embeddings()?; - let n = chunks.len(); - items.extend(chunks.into_iter().map(|(id, emb)| { - let key = format!("chunk:{id}"); - (key, emb) - })); - n - }; + let (items, chunk_count) = self.load_index_items()?; if chunk_count > 0 { tracing::info!( chunks = chunk_count, @@ -167,8 +152,8 @@ impl crate::Uteke { .write() .map_err(|_| Error::lock("index write lock during repair (rebuild)"))?; index.build(&items)?; - if let Err(e) = index.save() { - tracing::warn!("Failed to save index: {e}"); + if let Err(e) = index_sync::persist(&mut *index, &SyncPolicy::STANDARD) { + tracing::warn!("Failed to save index after repair rebuild: {e}"); } } @@ -180,6 +165,106 @@ impl crate::Uteke { }) } + /// Every vector the index should hold, straight from SQLite: memories + /// (NULL embeddings filtered in `load_all`) plus document chunks under + /// `chunk:` keys. Returns the items and how many are chunks. + /// + /// Without the chunks a rebuild silently evicts every chunk entry (#1110); + /// their embeddings are persisted in `document_chunks`, so nothing is + /// re-embedded. + fn load_index_items(&self) -> Result<(IndexItems, usize), Error> { + let all_memories = self.store.load_all(None)?; + let mut items: IndexItems = all_memories + .iter() + .filter(|m| !m.embedding.is_empty()) + .map(|m| (m.id.clone(), m.embedding.clone())) + .collect(); + let chunks = self.store.load_all_chunk_embeddings()?; + let chunk_count = chunks.len(); + items.extend( + chunks + .into_iter() + .map(|(id, emb)| (format!("chunk:{id}"), emb)), + ); + Ok((items, chunk_count)) + } + + /// Rebuild a vecq index from SQLite when more than a quarter of its rows + /// are dead (#1324). vecq is append-only, so updates and deletes leave + /// tombstoned rows that grow the files and force every search to + /// over-fetch `k + dead`. Returns `None` when no compaction is due. + /// + /// The index write lock is held across the SQLite read and the rebuild: + /// a concurrent writer has either committed before we load (so its row is + /// rebuilt in) or blocks on the lock until we are done (so its insert + /// lands in the compacted index). `dry_run` only reports. + pub fn compact_index_if_needed( + &self, + dry_run: bool, + ) -> Result, Error> { + // Cheap check under a read lock: the server calls this every second. + { + let index = self + .index + .read() + .map_err(|_| Error::lock("index read lock during compaction check"))?; + if !index.needs_compaction() { + return Ok(None); + } + } + let mut index = self + .index + .write() + .map_err(|_| Error::lock("index write lock during compaction"))?; + // Re-check: another thread may have compacted while we waited. + if !index.needs_compaction() { + return Ok(None); + } + let rows_before = index.physical_rows(); + let dead_before = index.dead_rows(); + if dry_run { + return Ok(Some(CompactionReport { + rows_before, + dead_before, + rows_after: rows_before, + dry_run: true, + })); + } + + let (mut items, _chunks) = self.load_index_items()?; + // `build` resets the index before validating, so never hand it a + // vector of the wrong width. + let dims = index.dims(); + let before = items.len(); + items.retain(|(_, emb)| emb.len() == dims); + if items.len() != before { + tracing::warn!( + skipped = before - items.len(), + "skipping vectors with the wrong dimensionality during index compaction" + ); + } + if items.is_empty() { + return Ok(None); + } + index.build(&items)?; + let rows_after = index.physical_rows(); + if let Err(e) = index_sync::persist(&mut *index, &SyncPolicy::STANDARD) { + tracing::warn!("Failed to save index after compaction: {e}"); + } + tracing::info!( + rows_before, + dead_before, + rows_after, + "Compacted vecq index (dead rows dropped)" + ); + Ok(Some(CompactionReport { + rows_before, + dead_before, + rows_after, + dry_run: false, + })) + } + /// Re-embed memories that have missing or empty embedding vectors. /// /// Scans active memories with NULL/empty embeddings via a dedicated SQL @@ -255,11 +340,13 @@ impl crate::Uteke { .write() .map_err(|_| Error::lock("index write during reembed"))?; for (id, vec) in &new_items { - if let Err(e) = index.insert(id, vec) { + if let Err(e) = + index_sync::insert_unsaved(&mut *index, id, vec, &SyncPolicy::STANDARD) + { tracing::warn!(id = %id, error = %e, "Failed to add to index"); } } - if let Err(e) = index.save() { + if let Err(e) = index_sync::persist(&mut *index, &SyncPolicy::STANDARD) { tracing::warn!(error = %e, "Failed to save index after reembed"); } } @@ -407,12 +494,12 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during aging_cleanup"))?; - for id in &ids { - index.remove(id); - } - if let Err(e) = index.save() { - tracing::warn!("Failed to save index: {e}"); - } + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .warn_if_stale("aging_cleanup (soft-delete)"); } tracing::info!( @@ -433,12 +520,12 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during aging_cleanup"))?; - for id in &ids { - index.remove(id); - } - if let Err(e) = index.save() { - tracing::warn!("Failed to save index: {e}"); - } + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .warn_if_stale("aging_cleanup (hard delete)"); } tracing::info!( @@ -481,12 +568,12 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during prune"))?; - for id in &ids { - index.remove(id); - } - if let Err(e) = index.save() { - tracing::warn!("Failed to save index: {e}"); - } + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .warn_if_stale("prune"); } Ok(PruneResult { @@ -540,12 +627,12 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during lifecycle_cycle"))?; - for id in &to_deprecate { - index.remove(id); - } - if let Err(e) = index.save() { - tracing::warn!("Failed to persist index after lifecycle deprecate: {e}"); - } + index_sync::remove_ids( + &mut *index, + to_deprecate.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .warn_if_stale("lifecycle_cycle deprecate"); } // Invalidate recall cache. @@ -581,12 +668,12 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during lifecycle prune"))?; - for id in &expired_ids { - index.remove(id); - } - if let Err(e) = index.save() { - tracing::warn!("Failed to persist index after lifecycle prune: {e}"); - } + index_sync::remove_ids( + &mut *index, + expired_ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .warn_if_stale("lifecycle_cycle prune"); tracing::info!( "Lifecycle cycle: pruned {} expired deprecated memories (ttl={}d)", pruned_count, @@ -609,17 +696,34 @@ impl crate::Uteke { }) } - /// Graceful shutdown — save dirty index to disk. - pub fn shutdown(&self) -> Result<(), Error> { + /// Write any pending vector-index changes to disk. Per-operation saves + /// are batched (#1322); long-lived hosts call this on a timer so a single + /// write followed by idle time still reaches disk. + pub fn flush_index(&self) -> Result<(), Error> { + // Cheap read-lock peek first: the server calls this every second and + // must not take the write lock (blocking recall) when nothing is pending. + if !self + .index + .read() + .map_err(|_| Error::lock("index read lock during flush"))? + .is_dirty() + { + return Ok(()); + } let mut index = self .index .write() - .map_err(|_| Error::lock("index write lock during shutdown"))?; + .map_err(|_| Error::lock("index write lock during flush"))?; if index.is_dirty() { - index.save()?; + index_sync::persist(&mut *index, &SyncPolicy::STANDARD)?; } Ok(()) } + + /// Graceful shutdown — save dirty index to disk. + pub fn shutdown(&self) -> Result<(), Error> { + self.flush_index() + } } #[cfg(test)] diff --git a/crates/uteke-core/src/memory/aging.rs b/crates/uteke-core/src/memory/aging.rs index 61debb56..670acf7d 100644 --- a/crates/uteke-core/src/memory/aging.rs +++ b/crates/uteke-core/src/memory/aging.rs @@ -1,6 +1,7 @@ //! Aging and access tracking — touch access, find/cleanup aged, tier counts. use crate::Error; +use crate::memory::store::memory_columns; use crate::memory::types::{DEFAULT_NAMESPACE, Memory}; use rusqlite::params; @@ -145,8 +146,10 @@ impl super::Store { // causing lexicographic comparison to fail. let cutoff = (chrono::Utc::now() - chrono::Duration::days(older_than_days as i64)).to_rfc3339(); - let sql = r#" - SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type + let sql = concat!( + "SELECT ", + memory_columns!(), + r#" FROM memories WHERE namespace = ?1 AND deprecated = 0 @@ -155,7 +158,8 @@ impl super::Store { AND access_count <= ?3 AND (last_accessed IS NULL OR last_accessed < ?4) ORDER BY created_at ASC - "#; + "# + ); let mut stmt = self .conn diff --git a/crates/uteke-core/src/memory/bulk.rs b/crates/uteke-core/src/memory/bulk.rs index 6bee3212..21895d14 100644 --- a/crates/uteke-core/src/memory/bulk.rs +++ b/crates/uteke-core/src/memory/bulk.rs @@ -1,6 +1,7 @@ //! Bulk operations — bulk delete, deprecation, TTL pruning, similarity search. use crate::Error; +use crate::memory::store::memory_columns; use crate::memory::types::{DEFAULT_NAMESPACE, Memory}; use crate::timeline::TimelineEventType; use rusqlite::params; @@ -60,6 +61,7 @@ impl super::Store { } /// Bulk delete all cold memories (not accessed in warm_days+ days or never accessed). + /// Pinned memories are never cold: they "never decay" and are skipped. /// /// Uses a single DELETE query with `RETURNING id` for efficiency. pub fn bulk_delete_cold( @@ -72,7 +74,7 @@ impl super::Store { let mut stmt = self .conn .prepare( - "DELETE FROM memories WHERE namespace = ?1 AND (last_accessed < ?2 OR last_accessed IS NULL) RETURNING id", + "DELETE FROM memories WHERE namespace = ?1 AND pinned = 0 AND (last_accessed < ?2 OR last_accessed IS NULL) RETURNING id", ) .map_err(|e| Error::db("database operation", e))?; let ids: Vec = stmt @@ -123,7 +125,7 @@ impl super::Store { let mut stmt = self .conn .prepare( - "SELECT id FROM memories WHERE namespace = ?1 AND deprecated = 0 AND (last_accessed < ?2 OR last_accessed IS NULL)", + "SELECT id FROM memories WHERE namespace = ?1 AND deprecated = 0 AND pinned = 0 AND (last_accessed < ?2 OR last_accessed IS NULL)", ) .map_err(|e| Error::db("database operation", e))?; let ids: Vec = stmt @@ -212,25 +214,18 @@ impl super::Store { return Ok(0); } let now = chrono::Utc::now().to_rfc3339(); - let placeholders: String = ids - .iter() - .enumerate() - .map(|(i, _)| format!("?{}", i + 3)) - .collect::>() - .join(","); - let sql = format!( - "UPDATE memories SET deprecated = 1, valid_until = ?1, deprecate_reason = ?2, updated_at = ?1, deprecated_at = ?1 WHERE id IN ({placeholders}) AND deprecated = 0" - ); - let mut params_vec: Vec> = - vec![Box::new(now), Box::new(reason.to_string())]; - for id in ids { - params_vec.push(Box::new(id.clone())); - } - let param_refs: Vec<&dyn rusqlite::types::ToSql> = - params_vec.iter().map(|p| p.as_ref()).collect(); + // Ids are passed as ONE JSON-array parameter and expanded with + // json_each(): no dynamic SQL and no bound-parameter limit. + let ids_json = + serde_json::to_string(ids).map_err(|e| Error::db("database operation", e))?; let count = self .conn - .execute(&sql, rusqlite::params_from_iter(param_refs)) + .execute( + "UPDATE memories SET deprecated = 1, valid_until = ?1, deprecate_reason = ?2, \ + updated_at = ?1, deprecated_at = ?1 \ + WHERE id IN (SELECT value FROM json_each(?3)) AND deprecated = 0", + params![now, reason, ids_json], + ) .map_err(|e| Error::db("database operation", e))?; Ok(count) } @@ -257,8 +252,8 @@ impl super::Store { let mut stmt = self .conn .prepare( - "SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type - FROM memories WHERE namespace = ?1 AND deprecated = 0 ORDER BY created_at DESC LIMIT ?2", + concat!("SELECT ", memory_columns!(), " + FROM memories WHERE namespace = ?1 AND deprecated = 0 ORDER BY created_at DESC LIMIT ?2"), ) .map_err(|e| Error::db("database operation", e))?; let rows = stmt @@ -299,13 +294,15 @@ impl super::Store { let cutoff = (chrono::Utc::now() - chrono::Duration::days(ttl_days as i64)).to_rfc3339(); let mut stmt = self .conn - .prepare( - "SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type + .prepare(concat!( + "SELECT ", + memory_columns!(), + " FROM memories WHERE namespace = ?1 AND deprecated = 1 AND updated_at < ?2 - ORDER BY updated_at ASC", - ) + ORDER BY updated_at ASC" + )) .map_err(|e| Error::db("database operation", e))?; let rows = stmt .query_map(params![ns, cutoff], row_to_memory) @@ -323,23 +320,16 @@ impl super::Store { if ids.is_empty() { return Ok(0); } - // Build parameterized IN clause: "WHERE id IN (?1, ?2, ?3)" - let placeholders: String = ids - .iter() - .enumerate() - .map(|(i, _)| format!("?{}", i + 1)) - .collect::>() - .join(","); - let sql = format!("DELETE FROM memories WHERE id IN ({placeholders})"); - let params: Vec> = ids - .iter() - .map(|id| Box::new(id.clone()) as Box) - .collect(); - let param_refs: Vec<&dyn rusqlite::types::ToSql> = - params.iter().map(|p| p.as_ref()).collect(); + // Single statement (atomic) with the ids as one JSON-array parameter: + // no dynamic SQL and no bound-parameter limit. + let ids_json = + serde_json::to_string(ids).map_err(|e| Error::db("database operation", e))?; let deleted = self .conn - .execute(&sql, rusqlite::params_from_iter(param_refs)) + .execute( + "DELETE FROM memories WHERE id IN (SELECT value FROM json_each(?1))", + params![ids_json], + ) .map_err(|e| Error::db("database operation", e))?; Ok(deleted) } diff --git a/crates/uteke-core/src/memory/crud.rs b/crates/uteke-core/src/memory/crud.rs index c2dae63c..ea980100 100644 --- a/crates/uteke-core/src/memory/crud.rs +++ b/crates/uteke-core/src/memory/crud.rs @@ -1,6 +1,7 @@ //! Core CRUD operations — insert, get, delete, update, list, search, count. use crate::Error; +use crate::memory::store::{memory_columns, memory_columns_m}; use crate::memory::types::Memory; use rusqlite::{OptionalExtension, params}; @@ -161,7 +162,11 @@ impl super::Store { pub fn get_by_id(&self, id: &str) -> Result, Error> { let mut stmt = self .conn - .prepare("SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type, slug, source, source_type, author_type, deprecated_at FROM memories WHERE id = ?1") + .prepare(concat!( + "SELECT ", + memory_columns!(), + " FROM memories WHERE id = ?1" + )) .map_err(|e| Error::db("Failed to prepare statement for get_by_id", e))?; let result = stmt @@ -188,7 +193,12 @@ impl super::Store { for chunk in ids.chunks(CHUNK_SIZE) { let placeholders = chunk.iter().map(|_| "?").collect::>().join(","); let sql = format!( - "SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type, slug, source, source_type, author_type, deprecated_at FROM memories WHERE id IN ({placeholders})" + concat!( + "SELECT ", + memory_columns!(), + " FROM memories WHERE id IN ({})" + ), + placeholders ); let mut stmt = self .conn @@ -250,7 +260,11 @@ impl super::Store { let ns = namespace.unwrap_or(crate::memory::types::DEFAULT_NAMESPACE); let mut stmt = self .conn - .prepare("SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type, slug, source, source_type, author_type, deprecated_at FROM memories WHERE id = ?1 AND namespace = ?2") + .prepare(concat!( + "SELECT ", + memory_columns!(), + " FROM memories WHERE id = ?1 AND namespace = ?2" + )) .map_err(|e| Error::db("Failed to prepare statement for get_by_id_in_namespace", e))?; let result = stmt @@ -451,15 +465,14 @@ impl super::Store { (Some(ns), Some(t)) => { let mut stmt = self .conn - .prepare( - "SELECT id, content, embedding, tags, metadata, \ - created_at, updated_at, namespace, access_count, \ - last_accessed, deprecated, valid_from, valid_until, \ - memory_type, importance, pinned, content_type, slug \ + .prepare(concat!( + "SELECT ", + memory_columns!(), + " \ FROM memories WHERE namespace = ?1 AND deprecated = 0 AND EXISTS \ (SELECT 1 FROM memory_tags WHERE memory_id = memories.id AND tag = ?2) \ - ORDER BY created_at DESC LIMIT ?3 OFFSET ?4", - ) + ORDER BY created_at DESC LIMIT ?3 OFFSET ?4" + )) .map_err(|e| Error::db("database operation", e))?; let rows = stmt .query_map(params![ns, t, limit as i64, offset as i64], row_to_memory) @@ -472,14 +485,13 @@ impl super::Store { (Some(ns), None) => { let mut stmt = self .conn - .prepare( - "SELECT id, content, embedding, tags, metadata, \ - created_at, updated_at, namespace, access_count, \ - last_accessed, deprecated, valid_from, valid_until, \ - memory_type, importance, pinned, content_type, slug \ + .prepare(concat!( + "SELECT ", + memory_columns!(), + " \ FROM memories WHERE namespace = ?1 AND deprecated = 0 \ - ORDER BY created_at DESC LIMIT ?2 OFFSET ?3", - ) + ORDER BY created_at DESC LIMIT ?2 OFFSET ?3" + )) .map_err(|e| Error::db("database operation", e))?; let rows = stmt .query_map(params![ns, limit as i64, offset as i64], row_to_memory) @@ -492,15 +504,14 @@ impl super::Store { (None, Some(t)) => { let mut stmt = self .conn - .prepare( - "SELECT id, content, embedding, tags, metadata, \ - created_at, updated_at, namespace, access_count, \ - last_accessed, deprecated, valid_from, valid_until, \ - memory_type, importance, pinned, content_type, slug \ + .prepare(concat!( + "SELECT ", + memory_columns!(), + " \ FROM memories WHERE deprecated = 0 AND EXISTS \ (SELECT 1 FROM memory_tags WHERE memory_id = memories.id AND tag = ?1) \ - ORDER BY created_at DESC LIMIT ?2 OFFSET ?3", - ) + ORDER BY created_at DESC LIMIT ?2 OFFSET ?3" + )) .map_err(|e| Error::db("database operation", e))?; let rows = stmt .query_map(params![t, limit as i64, offset as i64], row_to_memory) @@ -513,14 +524,13 @@ impl super::Store { (None, None) => { let mut stmt = self .conn - .prepare( - "SELECT id, content, embedding, tags, metadata, \ - created_at, updated_at, namespace, access_count, \ - last_accessed, deprecated, valid_from, valid_until, \ - memory_type, importance, pinned, content_type, slug \ + .prepare(concat!( + "SELECT ", + memory_columns!(), + " \ FROM memories WHERE deprecated = 0 \ - ORDER BY created_at DESC LIMIT ?1 OFFSET ?2", - ) + ORDER BY created_at DESC LIMIT ?1 OFFSET ?2" + )) .map_err(|e| Error::db("database operation", e))?; let rows = stmt .query_map(params![limit as i64, offset as i64], row_to_memory) @@ -555,14 +565,18 @@ impl super::Store { let pattern = format!("%{escaped}%"); let sql = match namespace { Some(_) => { - "SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type, slug + concat!("SELECT ", memory_columns!(), " FROM memories WHERE namespace = ?1 AND deprecated = 0 AND content LIKE ?2 ESCAPE '!' - ORDER BY created_at DESC LIMIT ?3" + ORDER BY created_at DESC LIMIT ?3") } None => { - "SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type, slug + concat!( + "SELECT ", + memory_columns!(), + " FROM memories WHERE deprecated = 0 AND content LIKE ?1 ESCAPE '!' ORDER BY created_at DESC LIMIT ?2" + ) } }; let mut stmt = self @@ -595,10 +609,18 @@ impl super::Store { // recall and must not re-enter the vector index on repair/verify (#1047). let sql = match namespace { Some(_) => { - "SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type, slug, source, source_type, author_type, deprecated_at FROM memories WHERE namespace = ?1 AND embedding IS NOT NULL AND deprecated = 0 ORDER BY created_at" + concat!( + "SELECT ", + memory_columns!(), + " FROM memories WHERE namespace = ?1 AND embedding IS NOT NULL AND deprecated = 0 ORDER BY created_at" + ) } None => { - "SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type, slug, source, source_type, author_type, deprecated_at FROM memories WHERE embedding IS NOT NULL AND deprecated = 0 ORDER BY created_at" + concat!( + "SELECT ", + memory_columns!(), + " FROM memories WHERE embedding IS NOT NULL AND deprecated = 0 ORDER BY created_at" + ) } }; @@ -644,10 +666,18 @@ impl super::Store { pub fn load_missing_embeddings(&self, namespace: Option<&str>) -> Result, Error> { let sql = match namespace { Some(_) => { - "SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type, slug, source, source_type, author_type, deprecated_at FROM memories WHERE namespace = ?1 AND (embedding IS NULL OR length(embedding) = 0) AND deprecated = 0 ORDER BY created_at" + concat!( + "SELECT ", + memory_columns!(), + " FROM memories WHERE namespace = ?1 AND (embedding IS NULL OR length(embedding) = 0) AND deprecated = 0 ORDER BY created_at" + ) } None => { - "SELECT id, content, embedding, tags, metadata, created_at, updated_at, namespace, access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, importance, pinned, content_type, slug, source, source_type, author_type, deprecated_at FROM memories WHERE (embedding IS NULL OR length(embedding) = 0) AND deprecated = 0 ORDER BY created_at" + concat!( + "SELECT ", + memory_columns!(), + " FROM memories WHERE (embedding IS NULL OR length(embedding) = 0) AND deprecated = 0 ORDER BY created_at" + ) } }; @@ -753,6 +783,39 @@ impl super::Store { Ok(ids) } + /// Which of `ids` are live (not deprecated) memories, optionally limited + /// to `namespace`? One query over the candidates only, so the cost is + /// O(candidates) instead of O(memories in the namespace) — the hot + /// `remember` path (dedup + cosine auto-link) used to load every id of the + /// namespace on each write. Ids that are not memories (e.g. `chunk:` + /// document-chunk index keys) simply do not appear in the result. + pub fn live_ids_in_namespace>( + &self, + ids: &[S], + namespace: Option<&str>, + ) -> Result, Error> { + if ids.is_empty() { + return Ok(std::collections::HashSet::new()); + } + let json = serde_json::to_string(&ids.iter().map(|s| s.as_ref()).collect::>()) + .map_err(|e| Error::Validation(format!("candidate id list: {e}")))?; + let mut stmt = self + .conn + .prepare( + "SELECT id FROM memories \ + WHERE id IN (SELECT value FROM json_each(?1)) \ + AND deprecated = 0 \ + AND (?2 IS NULL OR namespace = ?2)", + ) + .map_err(|e| Error::db("prepare live_ids_in_namespace", e))?; + let found = stmt + .query_map(params![json, namespace], |row| row.get::<_, String>(0)) + .map_err(|e| Error::db("query live_ids_in_namespace", e))? + .collect::>() + .map_err(|e| Error::db("read live_ids_in_namespace", e))?; + Ok(found) + } + /// Count memories grouped by memory_type in a namespace. pub fn memory_type_counts(&self, namespace: &str) -> Result, Error> { let mut stmt = self @@ -798,11 +861,10 @@ impl super::Store { (Some(ns), Some(t)) => { let mut stmt = self .conn - .prepare( - "SELECT m.id, m.content, m.embedding, m.tags, m.metadata, \ - m.created_at, m.updated_at, m.namespace, m.access_count, \ - m.last_accessed, m.deprecated, m.valid_from, m.valid_until, \ - m.memory_type, m.importance, m.pinned, m.content_type, m.slug \ + .prepare(concat!( + "SELECT ", + memory_columns_m!(), + " \ FROM memories m \ INNER JOIN memory_tags mt ON mt.memory_id = m.id \ WHERE m.namespace = ?1 \ @@ -811,8 +873,8 @@ impl super::Store { AND (m.valid_until IS NULL OR m.valid_until > ?2) \ AND m.deprecated = 0 \ AND mt.tag = ?3 \ - ORDER BY m.created_at DESC LIMIT ?4 OFFSET ?5", - ) + ORDER BY m.created_at DESC LIMIT ?4 OFFSET ?5" + )) .map_err(|e| Error::db("database operation", e))?; let rows = stmt .query_map( @@ -828,19 +890,18 @@ impl super::Store { (Some(ns), None) => { let mut stmt = self .conn - .prepare( - "SELECT id, content, embedding, tags, metadata, \ - created_at, updated_at, namespace, access_count, \ - last_accessed, deprecated, valid_from, valid_until, \ - memory_type, importance, pinned, content_type, slug \ + .prepare(concat!( + "SELECT ", + memory_columns!(), + " \ FROM memories \ WHERE namespace = ?1 \ AND created_at <= ?2 \ AND (valid_from IS NULL OR valid_from <= ?2) \ AND (valid_until IS NULL OR valid_until > ?2) \ AND deprecated = 0 \ - ORDER BY created_at DESC LIMIT ?3 OFFSET ?4", - ) + ORDER BY created_at DESC LIMIT ?3 OFFSET ?4" + )) .map_err(|e| Error::db("database operation", e))?; let rows = stmt .query_map(params![ns, pit, limit as i64, offset as i64], row_to_memory) @@ -853,11 +914,10 @@ impl super::Store { (None, Some(t)) => { let mut stmt = self .conn - .prepare( - "SELECT m.id, m.content, m.embedding, m.tags, m.metadata, \ - m.created_at, m.updated_at, m.namespace, m.access_count, \ - m.last_accessed, m.deprecated, m.valid_from, m.valid_until, \ - m.memory_type, m.importance, m.pinned, m.content_type, m.slug \ + .prepare(concat!( + "SELECT ", + memory_columns_m!(), + " \ FROM memories m \ INNER JOIN memory_tags mt ON mt.memory_id = m.id \ WHERE m.created_at <= ?1 \ @@ -865,8 +925,8 @@ impl super::Store { AND (m.valid_until IS NULL OR m.valid_until > ?1) \ AND m.deprecated = 0 \ AND mt.tag = ?2 \ - ORDER BY m.created_at DESC LIMIT ?3 OFFSET ?4", - ) + ORDER BY m.created_at DESC LIMIT ?3 OFFSET ?4" + )) .map_err(|e| Error::db("database operation", e))?; let rows = stmt .query_map(params![pit, t, limit as i64, offset as i64], row_to_memory) @@ -879,18 +939,17 @@ impl super::Store { (None, None) => { let mut stmt = self .conn - .prepare( - "SELECT id, content, embedding, tags, metadata, \ - created_at, updated_at, namespace, access_count, \ - last_accessed, deprecated, valid_from, valid_until, \ - memory_type, importance, pinned, content_type, slug \ + .prepare(concat!( + "SELECT ", + memory_columns!(), + " \ FROM memories \ WHERE created_at <= ?1 \ AND (valid_from IS NULL OR valid_from <= ?1) \ AND (valid_until IS NULL OR valid_until > ?1) \ AND deprecated = 0 \ - ORDER BY created_at DESC LIMIT ?2 OFFSET ?3", - ) + ORDER BY created_at DESC LIMIT ?2 OFFSET ?3" + )) .map_err(|e| Error::db("database operation", e))?; let rows = stmt .query_map(params![pit, limit as i64, offset as i64], row_to_memory) diff --git a/crates/uteke-core/src/memory/documents.rs b/crates/uteke-core/src/memory/documents.rs index 755e9934..9617ea40 100644 --- a/crates/uteke-core/src/memory/documents.rs +++ b/crates/uteke-core/src/memory/documents.rs @@ -127,6 +127,97 @@ pub struct DocumentSearchResult { pub mode: String, } +/// `(exact, like-prefix)` that select a subtree by `path`: the document itself +/// is `exact`, its descendants match `prefix`. +fn subtree_selectors(path: &str) -> (String, String) { + ( + path.to_string(), + format!("{}%", escape_like(&subtree_dir(path))), + ) +} + +/// The RAW (unescaped) directory form of `path`: always ends with `/`. Use it +/// for plain string comparisons (`starts_with`); the LIKE prefix from +/// [`subtree_selectors`] is escaped and must only be bound to SQL. +fn subtree_dir(path: &str) -> String { + if path.ends_with('/') { + path.to_string() + } else { + format!("{path}/") + } +} + +/// Escape `\`, `%` and `_` so a stored path is matched literally by +/// `LIKE ... ESCAPE '\'`. Without it `/my_doc/%` also matched `/myxdoc/...` +/// and a move/delete could rewrite or remove unrelated subtrees. +fn escape_like(s: &str) -> String { + let mut out = String::with_capacity(s.len()); + for c in s.chars() { + if matches!(c, '\\' | '%' | '_') { + out.push('\\'); + } + out.push(c); + } + out +} + +/// Rewrite the `path`/`depth` of every descendant of the document at +/// `old_path` after that document moved to `new_path` (`depth_diff` = new depth +/// minus old depth). The document itself (`skip_id`) is not touched. A document +/// with an empty legacy `path` has no addressable descendants (matching +/// `"" + "/%"` would hit every row), so nothing is rewritten. Returns the +/// number of descendants updated. Shared by `move_document` and the re-parenting +/// branch of `upsert_document` (#1377). +fn rewrite_descendants( + conn: &rusqlite::Connection, + old_path: &str, + new_path: &str, + depth_diff: i64, + skip_id: &str, +) -> Result { + if old_path.is_empty() { + return Ok(0); + } + let (old_exact, old_prefix) = subtree_selectors(old_path); + // Replace ONLY the leading prefix (substr), never other occurrences of the + // old path inside a descendant's path. + conn.execute( + "UPDATE documents SET path = ?2 || substr(path, length(?1) + 1), depth = depth + ?3 \ + WHERE path LIKE ?4 ESCAPE '\\' AND id != ?5", + params![old_exact, new_path, depth_diff, old_prefix, skip_id], + ) + .map_err(|e| Error::db("update descendant paths", e)) +} + +/// Deepest `depth` found in the subtree rooted at `old_path` (itself included), +/// or `None` for a legacy row without a path. +fn max_subtree_depth(conn: &rusqlite::Connection, old_path: &str) -> Option { + if old_path.is_empty() { + return None; + } + let (_, prefix) = subtree_selectors(old_path); + conn.query_row( + "SELECT MAX(depth) FROM documents WHERE path LIKE ?1 ESCAPE '\\'", + params![prefix], + |row| row.get::<_, Option>(0), + ) + .unwrap_or(None) +} + +/// Recompute a document's denormalized `has_children` flag from the actual +/// rows. Used wherever a parent can gain or lose a child so the flag cannot go +/// stale (#1332). +fn refresh_has_children(conn: &rusqlite::Connection, doc_id: &str) -> Result<(), Error> { + conn.execute( + "UPDATE documents SET has_children = \ + EXISTS(SELECT 1 FROM documents c WHERE c.parent_id = documents.id) \ + WHERE id = ?1", + params![doc_id], + ) + .map_err(|e| Error::db("refresh has_children", e))?; + Ok(()) +} + /// Row mapper for Document (full document queries). /// Columns: id, slug, title, content, namespace, author, tags, metadata, version, /// content_type, created_at, updated_at, parent_id, path, depth, sort_order, has_children @@ -185,7 +276,78 @@ impl super::Store { .conn .unchecked_transaction() .map_err(|e| Error::db("begin document upsert transaction", e))?; + let doc_id = self.upsert_document_in(&tx, doc)?; + tx.commit() + .map_err(|e| Error::db("commit document upsert", e))?; + Ok(doc_id) + } + + /// Upsert a document and replace its chunks in ONE transaction. + /// + /// `chunks` must already be embedded (embedding is slow and must not run + /// under a transaction or the index lock). Either the new document row and + /// all of its chunks are committed, or nothing changes. Returns the + /// document id (the existing id when the slug already exists). + pub fn upsert_document_with_chunks( + &self, + doc: &Document, + chunks: &[(DocumentChunk, Vec)], + ) -> Result { + let tx = self + .conn + .unchecked_transaction() + .map_err(|e| Error::db("begin document upsert transaction", e))?; + let doc_id = self.upsert_document_in(&tx, doc)?; + for (chunk, embedding) in chunks { + debug_assert_eq!( + chunk.document_id, doc_id, + "chunk must belong to the upserted document" + ); + self.insert_document_chunk(chunk, embedding)?; + } + tx.commit() + .map_err(|e| Error::db("commit document upsert", e))?; + Ok(doc_id) + } + + /// Partial update plus chunk replacement in ONE transaction (content + /// changed). Returns the updated document and the ids of the chunks that + /// were replaced (so the caller can drop them from the vector index), or + /// `None` if the document does not exist (nothing is written). + pub fn update_document_with_chunks( + &self, + id: &str, + title: Option<&str>, + content: Option<&str>, + tags: Option<&[String]>, + metadata: Option<&serde_json::Value>, + chunks: &[(DocumentChunk, Vec)], + ) -> Result)>, Error> { + let tx = self + .conn + .unchecked_transaction() + .map_err(|e| Error::db("begin document update transaction", e))?; + // update_document / delete_chunks_for_documents / insert_document_chunk + // all run on this connection, so they join the open transaction. + let Some(updated) = self.update_document(id, title, content, tags, metadata)? else { + return Ok(None); + }; + let old_chunk_ids = + self.delete_chunks_for_documents(std::slice::from_ref(&id.to_string()))?; + for (chunk, embedding) in chunks { + self.insert_document_chunk(chunk, embedding)?; + } + tx.commit() + .map_err(|e| Error::db("commit document update", e))?; + Ok(Some((updated, old_chunk_ids))) + } + /// Body of [`Self::upsert_document`], run inside the caller's transaction. + fn upsert_document_in( + &self, + tx: &rusqlite::Transaction<'_>, + doc: &Document, + ) -> Result { // Check if document exists (by slug — globally unique). let existing: Option = tx .query_row( @@ -197,6 +359,36 @@ impl super::Store { .map_err(|e| Error::db("query existing document", e))?; let doc_id = if let Some(id) = existing { + // Remember the previous parent: a re-parenting upsert must fix the + // flag on BOTH the old and the new parent (#1332). + let (old_parent, old_path, old_depth): (Option, String, i64) = tx + .query_row( + "SELECT parent_id, path, depth FROM documents WHERE id = ?1", + params![id], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)), + ) + .optional() + .map_err(|e| Error::db("read previous parent", e))? + .unwrap_or((None, String::new(), 0)); + // A re-parenting upsert moves the whole subtree: refuse a cycle and + // a result deeper than MAX_DEPTH BEFORE writing anything (#1377). + let moved = doc.path != old_path; + if moved && !old_path.is_empty() { + // Raw string comparison: the LIKE prefix is escaped and would not + // match a path containing `_` or `%` (CodeCora on #1378). + let inside_old_subtree = + doc.path.starts_with(&subtree_dir(&old_path)) && doc.path != old_path; + if inside_old_subtree { + return Err(Error::validation( + "cannot move document into its own descendant", + )); + } + if let Some(deepest) = max_subtree_depth(tx, &old_path) { + if deepest + (doc.depth - old_depth) > MAX_DEPTH { + return Err(Error::validation("move would exceed maximum depth of 10")); + } + } + } // Update existing document. let version = doc.version + 1; tx.execute( @@ -225,6 +417,12 @@ impl super::Store { params![id], ) .map_err(|e| Error::db("delete old document chunks", e))?; + if moved { + rewrite_descendants(tx, &old_path, &doc.path, doc.depth - old_depth, &id)?; + } + for parent in old_parent.iter().chain(doc.parent_id.iter()) { + refresh_has_children(tx, parent)?; + } id } else { // Insert new document. @@ -263,9 +461,6 @@ impl super::Store { doc.id.clone() }; - tx.commit() - .map_err(|e| Error::db("commit document upsert", e))?; - Ok(doc_id) } @@ -453,6 +648,32 @@ impl super::Store { Ok(docs) } + /// One page of documents (newest first, ties broken by id so pages never + /// overlap), optionally scoped to a namespace. `limit` is clamped to 1000 + /// per page; callers that need everything (export) loop on `offset`. + pub fn list_documents_page( + &self, + namespace: Option<&str>, + limit: usize, + offset: usize, + ) -> Result, Error> { + let limit = limit.clamp(1, 1000) as i64; + let mut stmt = self + .conn + .prepare( + "SELECT id, slug, title, namespace, author, version, updated_at, \ + parent_id, depth, has_children, sort_order \ + FROM documents WHERE (?1 IS NULL OR namespace = ?1) \ + ORDER BY updated_at DESC, id ASC LIMIT ?2 OFFSET ?3", + ) + .map_err(|e| Error::db("prepare list documents page", e))?; + let rows = stmt + .query_map(params![namespace, limit, offset as i64], row_to_summary) + .map_err(|e| Error::db("list documents page query", e))?; + rows.collect::, _>>() + .map_err(|e| Error::db("read documents page", e)) + } + /// List root documents (parent_id IS NULL), global. pub fn list_root_documents(&self, limit: usize) -> Result, Error> { self.list_root_documents_ns(None, limit) @@ -525,11 +746,7 @@ impl super::Store { .get_document(id_or_slug)? .ok_or_else(|| Error::validation("document not found for descendants query"))?, }; - let path_prefix = if doc.path.ends_with('/') { - format!("{}%", doc.path) - } else { - format!("{}/%", doc.path) - }; + let (_, path_prefix) = subtree_selectors(&doc.path); if let Some(max) = max_depth { let mut stmt = self @@ -537,7 +754,7 @@ impl super::Store { .prepare( "SELECT id, slug, title, namespace, author, version, updated_at, \ parent_id, depth, has_children, sort_order \ - FROM documents WHERE path LIKE ?1 AND id != ?4 AND depth <= ?2 \ + FROM documents WHERE path LIKE ?1 ESCAPE '\\' AND id != ?4 AND depth <= ?2 \ ORDER BY path, sort_order LIMIT ?3", ) .map_err(|e| Error::db("prepare list descendants", e))?; @@ -551,7 +768,7 @@ impl super::Store { .prepare( "SELECT id, slug, title, namespace, author, version, updated_at, \ parent_id, depth, has_children, sort_order \ - FROM documents WHERE path LIKE ?1 AND id != ?3 \ + FROM documents WHERE path LIKE ?1 ESCAPE '\\' AND id != ?3 \ ORDER BY path, sort_order LIMIT ?2", ) .map_err(|e| Error::db("prepare list descendants", e))?; @@ -626,15 +843,11 @@ impl super::Store { if path.is_empty() { return Ok(0); } - let prefix = if path.ends_with('/') { - format!("{}%", path) - } else { - format!("{}/%", path) - }; + let (_, prefix) = subtree_selectors(&path); let count: i64 = self .conn .query_row( - "SELECT COUNT(*) FROM documents WHERE path LIKE ?1 AND id != ?2", + "SELECT COUNT(*) FROM documents WHERE path LIKE ?1 ESCAPE '\\' AND id != ?2", params![prefix, doc_id], |row| row.get(0), ) @@ -679,11 +892,12 @@ impl super::Store { None => return Ok(0), }; - let (old_path_exact, old_prefix) = if old_path.ends_with('/') { - (old_path.clone(), format!("{}%", old_path)) - } else { - (old_path.to_string(), format!("{}/%", old_path)) - }; + // Legacy pre-v12 rows have an empty `path` (the migration never + // backfilled it). `"" + "/%"` would match EVERY document and rewrite all + // their paths/depths, so such a document is moved on its own: it has no + // addressable descendants (#1332). + let has_path = !old_path.is_empty(); + let (_, old_prefix) = subtree_selectors(&old_path); // Safety checks. if let Some(parent) = new_parent_id { @@ -711,14 +925,17 @@ impl super::Store { ) .unwrap_or(0); let new_depth = parent_depth + 1; - let max_child_depth: i64 = tx - .query_row( - "SELECT MAX(depth) FROM documents WHERE path LIKE ?1", + let max_child_depth: i64 = if has_path { + tx.query_row( + "SELECT MAX(depth) FROM documents WHERE path LIKE ?1 ESCAPE '\\'", params![old_prefix], |row| row.get::<_, Option>(0), ) .unwrap_or(None) - .unwrap_or(old_depth); + .unwrap_or(old_depth) + } else { + old_depth + }; let depth_diff = new_depth - old_depth; if max_child_depth + depth_diff > MAX_DEPTH { return Err(Error::validation("move would exceed maximum depth of 10")); @@ -766,13 +983,7 @@ impl super::Store { // Use substr() to replace ONLY the prefix portion (not all occurrences), // avoiding corruption when old_path appears multiple times in a descendant path. // length(old_path) is the offset where the suffix begins. - let n = tx - .execute( - "UPDATE documents SET path = ?2 || substr(path, length(?1) + 1), depth = depth + ?3 \ - WHERE path LIKE ?4 AND id != ?5", - params![old_path_exact, new_path, depth_diff, old_prefix, doc_id,], - ) - .map_err(|e| Error::db("update descendant paths", e))?; + let n = rewrite_descendants(&tx, &old_path, &new_path, depth_diff, doc_id)?; // Update new parent's has_children flag. if let Some(parent) = new_parent_id { @@ -812,6 +1023,17 @@ impl super::Store { /// Returns (deleted, subtree_size). pub fn delete_document(&self, id: &str) -> Result<(bool, usize), Error> { let subtree_size = self.count_descendants(id)?; + // The parent may lose its last child: remember it to refresh the flag. + let parent_id: Option = self + .conn + .query_row( + "SELECT parent_id FROM documents WHERE id = ?1", + params![id], + |row| row.get(0), + ) + .optional() + .map_err(|e| Error::db("read parent before delete", e))? + .flatten(); // Fetch the document's path for cascade delete. let path: String = self .conn @@ -822,11 +1044,10 @@ impl super::Store { ) .unwrap_or_default(); let cascade_prefix = if path.is_empty() { - format!("/{}/%", id) // fallback: try UUID-based prefix - } else if path.ends_with('/') { - format!("{}%", path) + // fallback: try UUID-based prefix + subtree_selectors(&format!("/{id}/")).1 } else { - format!("{}/%", path) + subtree_selectors(&path).1 }; // Collect all descendant document IDs (including self) for junction cleanup. @@ -835,7 +1056,7 @@ impl super::Store { self.conn .execute( "DELETE FROM room_documents WHERE doc_slug IN \ - (SELECT slug FROM documents WHERE id = ?1 OR path LIKE ?2)", + (SELECT slug FROM documents WHERE id = ?1 OR path LIKE ?2 ESCAPE '\\')", params![id, &cascade_prefix], ) .map_err(|e| Error::db("cleanup room_documents junction on doc delete", e))?; @@ -843,10 +1064,13 @@ impl super::Store { let n = self .conn .execute( - "DELETE FROM documents WHERE id = ?1 OR path LIKE ?2", + "DELETE FROM documents WHERE id = ?1 OR path LIKE ?2 ESCAPE '\\'", params![id, cascade_prefix], ) .map_err(|e| Error::db("delete document", e))?; + if let Some(parent) = parent_id { + refresh_has_children(&self.conn, &parent)?; + } Ok((n > 0, subtree_size)) } @@ -1088,6 +1312,458 @@ mod tests { doc } + fn fts_titles(store: &Store, query: &str) -> Vec { + store + .search_documents_fts(query, 10) + .unwrap() + .into_iter() + .map(|d| d.title) + .collect() + } + + /// #1349: updating a document's title (partial update and upsert) must work + /// and keep the external-content FTS index in step. + #[test] + fn test_document_title_change_keeps_fts_in_sync() { + let store = Store::open(":memory:").unwrap(); + let mut doc = make_doc("doc-t", "fts-title", "Zeppelin Handbook"); + doc.content = "body about airships".to_string(); + store.upsert_document(&doc).unwrap(); + assert_eq!(fts_titles(&store, "zeppelin"), vec!["Zeppelin Handbook"]); + + // Partial update with a new title. + store + .update_document("doc-t", Some("Dirigible Handbook"), None, None, None) + .unwrap() + .expect("document exists"); + assert!( + fts_titles(&store, "zeppelin").is_empty(), + "old title gone from FTS" + ); + assert_eq!(fts_titles(&store, "dirigible"), vec!["Dirigible Handbook"]); + + // Upsert on the same slug with yet another title and new content. + let mut again = make_doc("doc-t", "fts-title", "Blimp Handbook"); + again.content = "body about balloons".to_string(); + store.upsert_document(&again).unwrap(); + assert!(fts_titles(&store, "dirigible").is_empty()); + assert_eq!(fts_titles(&store, "blimp"), vec!["Blimp Handbook"]); + assert_eq!(fts_titles(&store, "balloons"), vec!["Blimp Handbook"]); + } + + /// #1349: deleting a document must remove it from FTS and must not leave + /// the index in a state where later writes fail. + #[test] + fn test_document_delete_removes_fts_entry() { + let store = Store::open(":memory:").unwrap(); + let mut doc = make_doc("doc-d", "fts-delete", "Gazebo Plans"); + doc.content = "plans for a gazebo".to_string(); + store.upsert_document(&doc).unwrap(); + assert_eq!(fts_titles(&store, "gazebo").len(), 1); + + store.delete_document("doc-d").unwrap(); + assert!(fts_titles(&store, "gazebo").is_empty(), "no stale FTS hit"); + + // The slug can be reused and updated afterwards. + let again = make_doc("doc-d2", "fts-delete", "Pergola Plans"); + store.upsert_document(&again).unwrap(); + store + .update_document("doc-d2", Some("Pergola Plans v2"), None, None, None) + .unwrap(); + assert_eq!(fts_titles(&store, "pergola"), vec!["Pergola Plans v2"]); + } + + /// #1349: reopening a database that still has the old triggers heals it. + #[test] + fn test_reopening_store_heals_old_documents_fts_triggers() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("heal.db"); + { + let store = Store::open(path.to_str().unwrap()).unwrap(); + store + .conn + .execute_batch( + "DROP TRIGGER documents_fts_update; + DROP TRIGGER documents_fts_delete; + CREATE TRIGGER documents_fts_update AFTER UPDATE ON documents BEGIN UPDATE documents_fts SET title = new.title, slug = new.slug, content = new.content WHERE rowid = new.rowid; END; + CREATE TRIGGER documents_fts_delete AFTER DELETE ON documents BEGIN DELETE FROM documents_fts WHERE rowid = old.rowid; END;", + ) + .unwrap(); + } + let store = Store::open(path.to_str().unwrap()).unwrap(); + for name in ["documents_fts_update", "documents_fts_delete"] { + let sql: String = store + .conn + .query_row( + "SELECT sql FROM sqlite_master WHERE type='trigger' AND name=?1", + [name], + |r| r.get(0), + ) + .unwrap(); + assert!(sql.contains("'delete'"), "{name} healed on open: {sql}"); + } + } + + fn has_children(store: &Store, id: &str) -> bool { + store.get_document(id).unwrap().unwrap().has_children + } + + /// #1332: deleting a child must clear the parent's denormalized flag. + #[test] + fn test_delete_document_refreshes_parent_has_children() { + let store = open_test_store(); + let parent = make_doc("p", "parent", "Parent"); + store.upsert_document(&parent).unwrap(); + store + .upsert_document(&make_child_doc("c1", "child-1", "C1", "p", "/p/")) + .unwrap(); + store + .upsert_document(&make_child_doc("c2", "child-2", "C2", "p", "/p/")) + .unwrap(); + assert!(has_children(&store, "p")); + + store.delete_document("c1").unwrap(); + assert!(has_children(&store, "p"), "c2 is still a child"); + store.delete_document("c2").unwrap(); + assert!( + !has_children(&store, "p"), + "last child deleted: flag cleared" + ); + } + + /// #1332: an upsert that changes the parent must fix BOTH parents' flags. + #[test] + fn test_upsert_reparent_updates_old_and_new_parent_flags() { + let store = open_test_store(); + store.upsert_document(&make_doc("a", "a", "A")).unwrap(); + store.upsert_document(&make_doc("b", "b", "B")).unwrap(); + store + .upsert_document(&make_child_doc("k", "kid", "Kid", "a", "/a/")) + .unwrap(); + assert!(has_children(&store, "a") && !has_children(&store, "b")); + + // Same slug, new parent. + let moved = make_child_doc("k", "kid", "Kid", "b", "/b/"); + store.upsert_document(&moved).unwrap(); + assert!(!has_children(&store, "a"), "old parent lost its only child"); + assert!(has_children(&store, "b"), "new parent gained a child"); + } + + /// #1332: a legacy row with an empty `path` made `old_prefix = "/%"` match + /// every document, rewriting all paths and depths. + #[test] + fn test_move_document_with_empty_legacy_path_leaves_others_alone() { + let store = open_test_store(); + store + .upsert_document(&make_doc("legacy", "legacy", "Legacy")) + .unwrap(); + store + .upsert_document(&make_doc("target", "target", "Target")) + .unwrap(); + store + .upsert_document(&make_doc("bystander", "bystander", "By")) + .unwrap(); + store + .upsert_document(&make_child_doc( + "deep", + "deep", + "Deep", + "bystander", + "/bystander/", + )) + .unwrap(); + // pre-v12 row: no path + store + .conn + .execute("UPDATE documents SET path = '' WHERE id = 'legacy'", []) + .unwrap(); + + let before_by = store.get_document("bystander").unwrap().unwrap(); + let before_deep = store.get_document("deep").unwrap().unwrap(); + store.move_document("legacy", Some("target"), None).unwrap(); + + let moved = store.get_document("legacy").unwrap().unwrap(); + assert_eq!(moved.parent_id.as_deref(), Some("target")); + assert_eq!(moved.path, "/target/legacy/"); + assert_eq!(moved.depth, 1); + let after_by = store.get_document("bystander").unwrap().unwrap(); + let after_deep = store.get_document("deep").unwrap().unwrap(); + assert_eq!( + (after_by.path, after_by.depth), + (before_by.path, before_by.depth) + ); + assert_eq!( + (after_deep.path, after_deep.depth), + (before_deep.path, before_deep.depth), + "unrelated documents must not be rewritten" + ); + } + + #[test] + fn test_list_documents_page_walks_every_document_once() { + let store = open_test_store(); + for i in 0..7 { + let mut d = make_doc(&format!("d{i}"), &format!("slug-{i}"), &format!("T{i}")); + // identical timestamps: only the id tie-breaker keeps pages stable + d.updated_at = "2026-10-08T00:00:00+00:00".to_string(); + store.upsert_document(&d).unwrap(); + } + let mut seen = Vec::new(); + let mut offset = 0; + loop { + let page = store.list_documents_page(None, 3, offset).unwrap(); + if page.is_empty() { + break; + } + offset += page.len(); + seen.extend(page.into_iter().map(|d| d.id)); + } + seen.sort(); + let mut want: Vec = (0..7).map(|i| format!("d{i}")).collect(); + want.sort(); + assert_eq!(seen, want, "every document exactly once"); + } + + fn path_depth(store: &Store, id: &str) -> (String, i64) { + let d = store.get_document(id).unwrap().unwrap(); + (d.path, d.depth) + } + + /// #1377: an upsert that re-parents a document must carry its whole subtree. + #[test] + fn test_upsert_reparent_rewrites_descendant_paths_and_depth() { + let store = open_test_store(); + store.upsert_document(&make_doc("a", "a", "A")).unwrap(); + store.upsert_document(&make_doc("b", "b", "B")).unwrap(); + store + .upsert_document(&make_child_doc("k", "kid", "Kid", "a", "/a/")) + .unwrap(); + store + .upsert_document(&make_child_doc("g", "grand", "Grand", "k", "/a/k/")) + .unwrap(); + store + .upsert_document(&make_child_doc("gg", "great", "Great", "g", "/a/k/g/")) + .unwrap(); + // unrelated sibling subtree that shares a textual prefix with "/a/k/" + store.upsert_document(&make_doc("ak", "ak", "AK")).unwrap(); + + // Same slug, new parent. + store + .upsert_document(&make_child_doc("k", "kid", "Kid", "b", "/b/")) + .unwrap(); + + assert_eq!(path_depth(&store, "k"), ("/b/k/".to_string(), 1)); + assert_eq!(path_depth(&store, "g"), ("/b/k/g/".to_string(), 2)); + assert_eq!(path_depth(&store, "gg"), ("/b/k/g/gg/".to_string(), 3)); + assert_eq!(path_depth(&store, "a"), ("/a/".to_string(), 0)); + assert_eq!(path_depth(&store, "ak"), ("/ak/".to_string(), 0)); + } + + /// Dropping the parent (upsert without one) moves the subtree to the root. + #[test] + fn test_upsert_to_root_rewrites_descendants() { + let store = open_test_store(); + store.upsert_document(&make_doc("a", "a", "A")).unwrap(); + store + .upsert_document(&make_child_doc("k", "kid", "Kid", "a", "/a/")) + .unwrap(); + store + .upsert_document(&make_child_doc("g", "grand", "Grand", "k", "/a/k/")) + .unwrap(); + + store.upsert_document(&make_doc("k", "kid", "Kid")).unwrap(); + + assert_eq!(path_depth(&store, "k"), ("/k/".to_string(), 0)); + assert_eq!(path_depth(&store, "g"), ("/k/g/".to_string(), 1)); + } + + /// Deleting the old parent must no longer take the moved subtree with it. + #[test] + fn test_delete_old_parent_after_reparent_keeps_moved_subtree() { + let store = open_test_store(); + store.upsert_document(&make_doc("a", "a", "A")).unwrap(); + store.upsert_document(&make_doc("b", "b", "B")).unwrap(); + store + .upsert_document(&make_child_doc("k", "kid", "Kid", "a", "/a/")) + .unwrap(); + store + .upsert_document(&make_child_doc("g", "grand", "Grand", "k", "/a/k/")) + .unwrap(); + store + .upsert_document(&make_child_doc("k", "kid", "Kid", "b", "/b/")) + .unwrap(); + + store.delete_document("a").unwrap(); + assert!(store.get_document("k").unwrap().is_some()); + assert!( + store.get_document("g").unwrap().is_some(), + "g moved with k and must survive the deletion of a" + ); + } + + #[test] + fn test_upsert_reparent_into_own_descendant_is_rejected_atomically() { + let store = open_test_store(); + store.upsert_document(&make_doc("k", "kid", "Kid")).unwrap(); + store + .upsert_document(&make_child_doc("g", "grand", "Grand", "k", "/k/")) + .unwrap(); + + // k under its own child g. + let cyc = make_child_doc("k", "kid", "Kid2", "g", "/k/g/"); + assert!(store.upsert_document(&cyc).is_err()); + assert_eq!(path_depth(&store, "k"), ("/k/".to_string(), 0)); + assert_eq!(path_depth(&store, "g"), ("/k/g/".to_string(), 1)); + assert_eq!(store.get_document("k").unwrap().unwrap().title, "Kid"); + } + + /// CodeCora on #1378: the cycle guard compared a raw path with an escaped + /// LIKE prefix, so paths containing `_` slipped through. + #[test] + fn test_cycle_guard_works_for_paths_with_like_wildcards() { + let store = open_test_store(); + store.upsert_document(&make_doc("k_1", "k1", "K")).unwrap(); + store + .upsert_document(&make_child_doc("g_1", "g1", "G", "k_1", "/k_1/")) + .unwrap(); + + let cyc = make_child_doc("k_1", "k1", "K2", "g_1", "/k_1/g_1/"); + let err = store.upsert_document(&cyc).unwrap_err().to_string(); + assert!(err.contains("own descendant"), "{err}"); + assert_eq!(path_depth(&store, "k_1"), ("/k_1/".to_string(), 0)); + assert_eq!(path_depth(&store, "g_1"), ("/k_1/g_1/".to_string(), 1)); + } + + #[test] + fn test_upsert_reparent_exceeding_max_depth_is_rejected() { + let store = open_test_store(); + // target chain t0..t4 (depths 0..4) + store.upsert_document(&make_doc("t0", "t0", "T0")).unwrap(); + let mut path = "/t0/".to_string(); + for i in 1..=4 { + let id = format!("t{i}"); + store + .upsert_document(&make_child_doc( + &id, + &id, + &id, + &format!("t{}", i - 1), + &path, + )) + .unwrap(); + path = format!("{path}{id}/"); + } + // subtree k0 (depth 0) with 9 descendants (depths 1..9) + store.upsert_document(&make_doc("k0", "k0", "K0")).unwrap(); + let mut kpath = "/k0/".to_string(); + for i in 1..=9 { + let id = format!("k{i}"); + store + .upsert_document(&make_child_doc( + &id, + &id, + &id, + &format!("k{}", i - 1), + &kpath, + )) + .unwrap(); + kpath = format!("{kpath}{id}/"); + } + + // k0 under t4 would put k9 at depth 14. + let moved = make_child_doc("k0", "k0", "K0", "t4", &path); + let err = store.upsert_document(&moved).unwrap_err().to_string(); + assert!(err.contains("maximum depth"), "{err}"); + assert_eq!(path_depth(&store, "k0"), ("/k0/".to_string(), 0)); + assert_eq!(path_depth(&store, "k9").1, 9); + } + + /// CodeCora on #1378: paths are matched with LIKE; `_` and `%` must be literal. + /// `/my_doc/%` used to match `/myxdoc/...` and rewrite/delete unrelated rows. + #[test] + fn test_like_wildcards_in_paths_do_not_touch_unrelated_subtrees() { + let store = open_test_store(); + store + .upsert_document(&make_doc("my_doc", "my-doc", "Underscore")) + .unwrap(); + store + .upsert_document(&make_doc("myxdoc", "myx-doc", "Look-alike")) + .unwrap(); + store + .upsert_document(&make_doc("100%", "pct", "Percent")) + .unwrap(); + store + .upsert_document(&make_doc("1000", "thousand", "Look-alike 2")) + .unwrap(); + store + .upsert_document(&make_doc("dest", "dest", "Dest")) + .unwrap(); + store + .upsert_document(&make_child_doc( + "kid-x", "kid-x", "KidX", "myxdoc", "/myxdoc/", + )) + .unwrap(); + store + .upsert_document(&make_child_doc("kid-k", "kid-k", "KidK", "1000", "/1000/")) + .unwrap(); + store + .upsert_document(&make_child_doc( + "kid-m", "kid-m", "KidM", "my_doc", "/my_doc/", + )) + .unwrap(); + + // moving my_doc must not rewrite /myxdoc/... (and `%` must not match "1000") + store.move_document("my_doc", Some("dest"), None).unwrap(); + assert_eq!( + path_depth(&store, "kid-m"), + ("/dest/my_doc/kid-m/".to_string(), 2) + ); + assert_eq!( + path_depth(&store, "kid-x"), + ("/myxdoc/kid-x/".to_string(), 1) + ); + store.move_document("100%", Some("dest"), None).unwrap(); + assert_eq!(path_depth(&store, "kid-k"), ("/1000/kid-k/".to_string(), 1)); + + // the same via a re-parenting upsert + store + .upsert_document(&make_doc("dest2", "dest2", "Dest2")) + .unwrap(); + store + .upsert_document(&make_child_doc( + "my_doc", + "my-doc", + "Underscore", + "dest2", + "/dest2/", + )) + .unwrap(); + assert_eq!( + path_depth(&store, "kid-m"), + ("/dest2/my_doc/kid-m/".to_string(), 2) + ); + assert_eq!( + path_depth(&store, "kid-x"), + ("/myxdoc/kid-x/".to_string(), 1) + ); + + // counting / listing descendants stays literal too + assert_eq!(store.count_descendants("my_doc").unwrap(), 1); + assert_eq!( + store.list_descendants("my_doc", None, 100).unwrap().len(), + 1 + ); + + // deleting my_doc removes only its own subtree + store.delete_document("my_doc").unwrap(); + assert!(store.get_document("kid-m").unwrap().is_none()); + assert!(store.get_document("myxdoc").unwrap().is_some()); + assert!( + store.get_document("kid-x").unwrap().is_some(), + "look-alike subtree survives" + ); + } + #[test] fn test_document_crud() { let store = open_test_store(); diff --git a/crates/uteke-core/src/memory/fts5.rs b/crates/uteke-core/src/memory/fts5.rs index 2aeef80c..2f966ec2 100644 --- a/crates/uteke-core/src/memory/fts5.rs +++ b/crates/uteke-core/src/memory/fts5.rs @@ -1,6 +1,7 @@ //! FTS5 full-text search for memories. use crate::Error; +use crate::memory::store::memory_columns_m; use crate::memory::types::Memory; use rusqlite::params; @@ -91,18 +92,26 @@ impl super::Store { let sql = match namespace { Some(_) => { - r#"SELECT m.id, m.content, m.embedding, m.tags, m.metadata, m.created_at, m.updated_at, m.namespace, m.access_count, m.last_accessed, m.deprecated, m.valid_from, m.valid_until, m.memory_type, m.importance, m.pinned, m.content_type, m.source, m.source_type, f.rank + concat!( + "SELECT ", + memory_columns_m!(), + r#", f.rank FROM memories_fts f JOIN memories m ON f.rowid = m.rowid WHERE memories_fts MATCH ?1 AND m.namespace = ?2 AND m.deprecated = 0 ORDER BY f.rank LIMIT ?3"# + ) } None => { - r#"SELECT m.id, m.content, m.embedding, m.tags, m.metadata, m.created_at, m.updated_at, m.namespace, m.access_count, m.last_accessed, m.deprecated, m.valid_from, m.valid_until, m.memory_type, m.importance, m.pinned, m.content_type, m.source, m.source_type, f.rank + concat!( + "SELECT ", + memory_columns_m!(), + r#", f.rank FROM memories_fts f JOIN memories m ON f.rowid = m.rowid WHERE memories_fts MATCH ?1 AND m.deprecated = 0 ORDER BY f.rank LIMIT ?2"# + ) } }; @@ -116,7 +125,7 @@ impl super::Store { let rows = stmt .query_map(params![fts_query, ns, limit as i64], |row| { let memory = row_to_memory(row)?; - let rank: f64 = row.get(19)?; + let rank: f64 = row.get(22)?; Ok((memory, rank)) }) .map_err(|e| Error::db("execute FTS5 search", e))?; @@ -133,7 +142,7 @@ impl super::Store { let rows = stmt .query_map(params![fts_query, limit as i64], |row| { let memory = row_to_memory(row)?; - let rank: f64 = row.get(19)?; + let rank: f64 = row.get(22)?; Ok((memory, rank)) }) .map_err(|e| Error::db("execute FTS5 search", e))?; @@ -174,18 +183,26 @@ impl super::Store { let sql = match namespace { Some(_) => { - r#"SELECT m.id, m.content, m.embedding, m.tags, m.metadata, m.created_at, m.updated_at, m.namespace, m.access_count, m.last_accessed, m.deprecated, m.valid_from, m.valid_until, m.memory_type, m.importance, m.pinned, m.content_type, m.source, m.source_type, f.rank + concat!( + "SELECT ", + memory_columns_m!(), + r#", f.rank FROM memories_fts f JOIN memories m ON f.rowid = m.rowid WHERE memories_fts MATCH ?1 AND m.namespace = ?2 AND m.deprecated = 0 ORDER BY f.rank LIMIT ?3"# + ) } None => { - r#"SELECT m.id, m.content, m.embedding, m.tags, m.metadata, m.created_at, m.updated_at, m.namespace, m.access_count, m.last_accessed, m.deprecated, m.valid_from, m.valid_until, m.memory_type, m.importance, m.pinned, m.content_type, m.source, m.source_type, f.rank + concat!( + "SELECT ", + memory_columns_m!(), + r#", f.rank FROM memories_fts f JOIN memories m ON f.rowid = m.rowid WHERE memories_fts MATCH ?1 AND m.deprecated = 0 ORDER BY f.rank LIMIT ?2"# + ) } }; @@ -199,7 +216,7 @@ impl super::Store { let rows = stmt .query_map(params![fts_query, ns, limit as i64], |row| { let memory = row_to_memory(row)?; - let rank: f64 = row.get(19)?; + let rank: f64 = row.get(22)?; Ok((memory, rank)) }) .map_err(|e| Error::db("execute FTS5 token search", e))?; @@ -216,7 +233,7 @@ impl super::Store { let rows = stmt .query_map(params![fts_query, limit as i64], |row| { let memory = row_to_memory(row)?; - let rank: f64 = row.get(19)?; + let rank: f64 = row.get(22)?; Ok((memory, rank)) }) .map_err(|e| Error::db("execute FTS5 token search", e))?; @@ -264,6 +281,33 @@ mod tests { } } + /// Regression: FTS5 rows must map slug/source/source_type/author_type + /// the same way as every other query (column order matches row_to_memory). + #[test] + fn test_fts5_rows_carry_correct_metadata_columns() { + let store = Store::open(":memory:").unwrap(); + store.init_fts5().unwrap(); + + let mut m = make_test_memory("meta-1", "uniqueterm alpha beta", &[]); + m.slug = Some("my-slug".to_string()); + m.source = Some("notes.md".to_string()); + m.source_type = "file".to_string(); + m.author_type = "human".to_string(); + store.insert(&m).unwrap(); + + for results in [ + store.search_fts5("uniqueterm alpha", None, 10).unwrap(), + store.search_fts5_tokens("uniqueterm", None, 10).unwrap(), + ] { + let got = &results[0].0; + assert_eq!(got.slug.as_deref(), Some("my-slug")); + assert_eq!(got.source.as_deref(), Some("notes.md")); + assert_eq!(got.source_type, "file"); + assert_eq!(got.author_type, "human"); + assert!(results[0].1.is_finite()); + } + } + #[test] fn test_fts5_search() { let store = Store::open(":memory:").unwrap(); diff --git a/crates/uteke-core/src/memory/rooms.rs b/crates/uteke-core/src/memory/rooms.rs index 77a85a91..bb7e6df2 100644 --- a/crates/uteke-core/src/memory/rooms.rs +++ b/crates/uteke-core/src/memory/rooms.rs @@ -1,6 +1,7 @@ //! Room operations — collaborative memory spaces for multi-agent discussions. use crate::Error; +use crate::memory::store::memory_columns_m; use rusqlite::OptionalExtension; use rusqlite::params; @@ -534,14 +535,15 @@ impl super::Store { } conditions.push("m.deprecated = 0"); let sql = format!( - "SELECT m.id, m.content, m.embedding, m.tags, m.metadata, \ - m.created_at, m.updated_at, m.namespace, m.access_count, \ - m.last_accessed, m.deprecated, m.valid_from, m.valid_until, m.memory_type, m.importance, m.pinned, m.content_type, \ - m.slug, m.source, m.source_type \ + concat!( + "SELECT ", + memory_columns_m!(), + " \ FROM memories m \ INNER JOIN room_memories rm ON m.id = rm.memory_id \ WHERE {} \ - ORDER BY rm.joined_at ASC{}", + ORDER BY rm.joined_at ASC{}" + ), conditions.join(" AND "), if no_limit { "" } else { " LIMIT ?" }, ); diff --git a/crates/uteke-core/src/memory/schema.rs b/crates/uteke-core/src/memory/schema.rs index 7c60a1a7..079d284f 100644 --- a/crates/uteke-core/src/memory/schema.rs +++ b/crates/uteke-core/src/memory/schema.rs @@ -5,6 +5,16 @@ use rusqlite::{OptionalExtension, params}; use super::store::{CURRENT_SCHEMA_VERSION, SCHEMA, SCHEMA_INDEXES}; +/// `documents_fts` is an FTS5 *external-content* table (`content='documents'`). +/// For such tables FTS5 must be told the OLD values through the special +/// `'delete'` command; a plain `UPDATE documents_fts SET ...` or +/// `DELETE FROM documents_fts` re-reads the already-modified content row and +/// fails with "database disk image is malformed" (#1349). `memories_fts` +/// already follows this pattern. +const DOCUMENTS_FTS_TRIGGER_INSERT: &str = "CREATE TRIGGER IF NOT EXISTS documents_fts_insert AFTER INSERT ON documents BEGIN INSERT INTO documents_fts(rowid, title, slug, content) VALUES (new.rowid, new.title, new.slug, new.content); END"; +const DOCUMENTS_FTS_TRIGGER_UPDATE: &str = "CREATE TRIGGER IF NOT EXISTS documents_fts_update AFTER UPDATE ON documents BEGIN INSERT INTO documents_fts(documents_fts, rowid, title, slug, content) VALUES ('delete', old.rowid, old.title, old.slug, old.content); INSERT INTO documents_fts(rowid, title, slug, content) VALUES (new.rowid, new.title, new.slug, new.content); END"; +const DOCUMENTS_FTS_TRIGGER_DELETE: &str = "CREATE TRIGGER IF NOT EXISTS documents_fts_delete AFTER DELETE ON documents BEGIN INSERT INTO documents_fts(documents_fts, rowid, title, slug, content) VALUES ('delete', old.rowid, old.title, old.slug, old.content); END"; + impl super::Store { /// Run initial schema creation + legacy column migrations. pub(super) fn init_schema(&self) -> Result<(), Error> { @@ -328,7 +338,43 @@ impl super::Store { return self.ensure_documents_fts_create(); } - Ok(()) + self.heal_documents_fts_triggers() + } + + /// Replace the pre-#1349 `documents_fts_update` / `documents_fts_delete` + /// triggers (plain UPDATE / DELETE on an external-content FTS5 table) with + /// the `'delete'`-command form, then rebuild the index from `documents` + /// so entries left stale by the old delete trigger disappear. Idempotent: + /// does nothing once the triggers are correct. + pub(crate) fn heal_documents_fts_triggers(&self) -> Result<(), Error> { + let trigger_sql = |name: &str| -> Result, Error> { + self.conn + .query_row( + "SELECT sql FROM sqlite_master WHERE type='trigger' AND name=?1", + [name], + |row| row.get::<_, String>(0), + ) + .optional() + .map_err(|e| Error::db("read documents_fts trigger (#1349)", e)) + }; + let outdated = + |sql: &Option| sql.as_deref().is_some_and(|t| !t.contains("'delete'")); + if !outdated(&trigger_sql("documents_fts_update")?) + && !outdated(&trigger_sql("documents_fts_delete")?) + { + return Ok(()); + } + tracing::info!("Repairing documents_fts triggers (#1349)"); + let batch = format!( + "DROP TRIGGER IF EXISTS documents_fts_update; \ + DROP TRIGGER IF EXISTS documents_fts_delete; \ + {DOCUMENTS_FTS_TRIGGER_UPDATE}; \ + {DOCUMENTS_FTS_TRIGGER_DELETE}; \ + INSERT INTO documents_fts(documents_fts) VALUES ('rebuild');" + ); + self.conn + .execute_batch(&batch) + .map_err(|e| Error::db("repair documents_fts triggers (#1349)", e)) } fn ensure_documents_fts_create(&self) -> Result<(), Error> { @@ -342,9 +388,7 @@ impl super::Store { // Create sync triggers — include content body. self.conn .execute_batch( - "CREATE TRIGGER IF NOT EXISTS documents_fts_insert AFTER INSERT ON documents BEGIN INSERT INTO documents_fts(rowid, title, slug, content) VALUES (new.rowid, new.title, new.slug, new.content); END; - CREATE TRIGGER IF NOT EXISTS documents_fts_update AFTER UPDATE ON documents BEGIN UPDATE documents_fts SET title = new.title, slug = new.slug, content = new.content WHERE rowid = new.rowid; END; - CREATE TRIGGER IF NOT EXISTS documents_fts_delete AFTER DELETE ON documents BEGIN DELETE FROM documents_fts WHERE rowid = old.rowid; END;", + &format!("{DOCUMENTS_FTS_TRIGGER_INSERT}; {DOCUMENTS_FTS_TRIGGER_UPDATE}; {DOCUMENTS_FTS_TRIGGER_DELETE};"), ) .map_err(|e| Error::db("create documents_fts triggers (#860)", e))?; @@ -913,10 +957,10 @@ impl super::Store { // Best-effort: skip if FTS5 is not available (e.g., custom SQLite builds). let fts = [ "CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(title, slug, content, content='documents', content_rowid='rowid')", - // Triggers to keep FTS in sync with documents table. - "CREATE TRIGGER IF NOT EXISTS documents_fts_insert AFTER INSERT ON documents BEGIN INSERT INTO documents_fts(rowid, title, slug, content) VALUES (new.rowid, new.title, new.slug, new.content); END", - "CREATE TRIGGER IF NOT EXISTS documents_fts_update AFTER UPDATE ON documents BEGIN UPDATE documents_fts SET title = new.title, slug = new.slug, content = new.content WHERE rowid = new.rowid; END", - "CREATE TRIGGER IF NOT EXISTS documents_fts_delete AFTER DELETE ON documents BEGIN DELETE FROM documents_fts WHERE rowid = old.rowid; END", + // Triggers to keep FTS in sync with documents table (#1349: 'delete' form). + DOCUMENTS_FTS_TRIGGER_INSERT, + DOCUMENTS_FTS_TRIGGER_UPDATE, + DOCUMENTS_FTS_TRIGGER_DELETE, ]; for sql in &fts { let _ = self.conn.execute(sql, []); diff --git a/crates/uteke-core/src/memory/store.rs b/crates/uteke-core/src/memory/store.rs index 72e8d805..584747e2 100644 --- a/crates/uteke-core/src/memory/store.rs +++ b/crates/uteke-core/src/memory/store.rs @@ -321,19 +321,62 @@ impl Store { /// Recalculate importance for all memories. /// importance = 0.3*access_score + 0.3*recency_score + 0.2*connectivity + 0.2*is_pinned pub fn recompute_importance(&self) -> Result { - let memories = self.load_all(None)?; + // Only the columns the formula needs: the old `load_all` pulled every + // embedding blob into memory for no reason (#1332). Same predicate as + // `load_all(None)`: live memories that have an embedding. + struct Row { + id: String, + importance: f64, + pinned: bool, + access_count: i64, + last_accessed: Option>, + metadata: serde_json::Value, + } + let rows: Vec = { + let mut stmt = self + .conn + .prepare( + "SELECT id, importance, pinned, access_count, last_accessed, metadata \ + FROM memories WHERE embedding IS NOT NULL AND deprecated = 0", + ) + .map_err(|e| Error::db("prepare recompute_importance", e))?; + stmt.query_map([], |r| { + let last: Option = r.get(4)?; + let meta: Option = r.get(5)?; + Ok(Row { + id: r.get(0)?, + importance: r.get::<_, f64>(1)?, + pinned: r.get::<_, i64>(2)? != 0, + access_count: r.get(3)?, + last_accessed: last.as_deref().and_then(parse_datetime_opt), + metadata: meta + .and_then(|m| serde_json::from_str(&m).ok()) + .unwrap_or(serde_json::Value::Null), + }) + }) + .map_err(|e| Error::db("query recompute_importance", e))? + .collect::, _>>() + .map_err(|e| Error::db("read recompute_importance", e))? + }; + + // One transaction for the whole pass: per-row autocommit made this one + // fsync per memory. + let tx = self + .conn + .unchecked_transaction() + .map_err(|e| Error::db("begin recompute_importance", e))?; + let mut update = tx + .prepare("UPDATE memories SET importance = ?1 WHERE id = ?2") + .map_err(|e| Error::db("prepare importance update", e))?; let mut updated = 0; let now = chrono::Utc::now(); - for m in &memories { + for m in &rows { // Skip pinned — they stay at 1.0 if m.pinned { if (m.importance - 1.0).abs() > f64::EPSILON { - self.conn - .execute( - "UPDATE memories SET importance = 1.0 WHERE id = ?1", - rusqlite::params![m.id], - ) + update + .execute(rusqlite::params![1.0_f64, m.id]) .map_err(|e| Error::db("update importance", e))?; updated += 1; } @@ -359,22 +402,20 @@ impl Store { .unwrap_or(0) as f64; let connectivity = (rel_count / 5.0).min(1.0); - let importance = 0.3 * access_score - + 0.3 * recency_score - + 0.2 * connectivity - + 0.2 * if m.pinned { 1.0 } else { 0.0 }; + // (pinned rows `continue`d above, so the pinned term is always 0 here.) + let importance = 0.3 * access_score + 0.3 * recency_score + 0.2 * connectivity; let importance = importance.clamp(0.0_f64, 1.0_f64); if (m.importance - importance).abs() > f64::EPSILON { - self.conn - .execute( - "UPDATE memories SET importance = ?1 WHERE id = ?2", - rusqlite::params![importance, m.id], - ) + update + .execute(rusqlite::params![importance, m.id]) .map_err(|e| Error::db("update importance", e))?; updated += 1; } } + drop(update); + tx.commit() + .map_err(|e| Error::db("commit recompute_importance", e))?; Ok(updated) } @@ -457,6 +498,35 @@ pub(crate) fn parse_datetime_opt(s: &str) -> Option { + "id, content, embedding, tags, metadata, created_at, updated_at, namespace, \ + access_count, last_accessed, deprecated, valid_from, valid_until, memory_type, \ + importance, pinned, content_type, slug, source, source_type, author_type, \ + deprecated_at" + }; +} + +/// [`memory_columns`] with every column qualified by the `m.` alias. +macro_rules! memory_columns_m { + () => { + "m.id, m.content, m.embedding, m.tags, m.metadata, m.created_at, m.updated_at, \ + m.namespace, m.access_count, m.last_accessed, m.deprecated, m.valid_from, \ + m.valid_until, m.memory_type, m.importance, m.pinned, m.content_type, m.slug, \ + m.source, m.source_type, m.author_type, m.deprecated_at" + }; +} + +pub(crate) use {memory_columns, memory_columns_m}; + /// Convert a database row to a Memory. pub(crate) fn row_to_memory(row: &rusqlite::Row<'_>) -> Result { let id: String = row.get(0)?; @@ -603,6 +673,21 @@ mod tests { } } + #[test] + fn test_bulk_ids_exceed_sqlite_param_limit() { + let store = Store::open(":memory:").unwrap(); + let ids: Vec = (0..1200).map(|i| format!("bulk-{i}")).collect(); + for id in &ids { + store.insert(&make_test_memory(id, "x", &[])).unwrap(); + } + + let deprecated = store.deprecate_by_ids(&ids, "test").unwrap(); + assert_eq!(deprecated, 1200); + + let deleted = store.delete_by_ids(&ids).unwrap(); + assert_eq!(deleted, 1200); + } + #[test] fn test_store_crud() { let store = Store::open(":memory:").unwrap(); @@ -1022,6 +1107,57 @@ mod tests { assert_eq!(new_matches.len(), 2); } + /// #1332: renaming onto a tag the memory already has must not duplicate it. + #[test] + fn test_rename_tag_onto_existing_tag_does_not_duplicate() { + let store = Store::open(":memory:").unwrap(); + store + .insert(&make_test_memory("1", "has both", &["alpha", "beta"])) + .unwrap(); + store + .insert(&make_test_memory("2", "only old", &["alpha"])) + .unwrap(); + + store.rename_tag("alpha", "beta", None).unwrap(); + + let both = store.get_by_id("1").unwrap().unwrap(); + assert_eq!(both.tags, vec!["beta".to_string()], "no duplicate entry"); + let only = store.get_by_id("2").unwrap().unwrap(); + assert_eq!(only.tags, vec!["beta".to_string()]); + assert_eq!(store.list(Some("beta"), None, 10, 0).unwrap().len(), 2); + assert_eq!(store.list(Some("alpha"), None, 10, 0).unwrap().len(), 0); + } + + /// recompute_importance keeps its formula after the rewrite (#1332). + #[test] + fn test_recompute_importance_formula_and_pinned() { + let store = Store::open(":memory:").unwrap(); + // 10+ accesses, accessed just now, no relationships: 0.3*1 + 0.3*~1. + let mut hot = make_test_memory("hot", "hot", &[]); + hot.access_count = 10; + hot.last_accessed = Some(chrono::Utc::now()); + hot.importance = 0.5; + store.insert(&hot).unwrap(); + // never accessed, no relationships: 0.3*0 + 0.3*e^(-ln2*365/30) ~ 0. + let mut cold = make_test_memory("cold", "cold", &[]); + cold.importance = 0.5; + store.insert(&cold).unwrap(); + // pinned memories are forced to 1.0. + let mut pin = make_test_memory("pin", "pin", &[]); + pin.pinned = true; + pin.importance = 0.2; + store.insert(&pin).unwrap(); + + let changed = store.recompute_importance().unwrap(); + assert_eq!(changed, 3); + let imp = |id: &str| store.get_by_id(id).unwrap().unwrap().importance; + assert!((imp("hot") - 0.6).abs() < 0.01, "hot = {}", imp("hot")); + assert!(imp("cold") < 0.001, "cold = {}", imp("cold")); + assert!((imp("pin") - 1.0).abs() < f64::EPSILON); + // a second pass changes nothing + assert_eq!(store.recompute_importance().unwrap(), 0); + } + #[test] fn test_delete_tag_json_each() { let store = Store::open(":memory:").unwrap(); @@ -1201,6 +1337,98 @@ mod tests { assert_eq!(store.count(Some("ns-b")).unwrap(), 1); } + /// Every read path must hand back the full row: queries that stop short + /// of the shared column list silently default slug/source/source_type/ + /// author_type (row_to_memory decodes by position). + #[test] + fn test_every_memory_reader_returns_full_row() { + let store = Store::open(":memory:").unwrap(); + let mut m = make_test_memory("full-row", "fullrowterm", &["t1"]); + m.slug = Some("my-slug".to_string()); + m.source = Some("notes.md".to_string()); + m.source_type = "file".to_string(); + m.author_type = "human".to_string(); + store.insert(&m).unwrap(); + + let check = |what: &str, got: &Memory| { + assert_eq!(got.slug.as_deref(), Some("my-slug"), "{what}: slug"); + assert_eq!(got.source.as_deref(), Some("notes.md"), "{what}: source"); + assert_eq!(got.source_type, "file", "{what}: source_type"); + assert_eq!(got.author_type, "human", "{what}: author_type"); + }; + let one = |what: &str, v: Vec| { + assert_eq!(v.len(), 1, "{what}: expected exactly the inserted row"); + check(what, &v[0]); + }; + + check("get_by_id", &store.get_by_id("full-row").unwrap().unwrap()); + one("get_by_ids", store.get_by_ids(&["full-row"]).unwrap()); + check( + "get_by_id_in_namespace", + &store + .get_by_id_in_namespace("full-row", Some(crate::memory::types::DEFAULT_NAMESPACE)) + .unwrap() + .unwrap(), + ); + one("load_all", store.load_all(None).unwrap()); + one("list", store.list(None, None, 10, 0).unwrap()); + one("list(tag)", store.list(Some("t1"), None, 10, 0).unwrap()); + one( + "search_content", + store.search_content("fullrowterm", None, 10).unwrap(), + ); + let future = chrono::Utc::now() + chrono::Duration::days(1); + one( + "list_at_time", + store.list_at_time(None, None, 10, 0, future).unwrap(), + ); + one( + "list_at_time(tag)", + store.list_at_time(Some("t1"), None, 10, 0, future).unwrap(), + ); + one("find_similar", store.find_similar("default", 10).unwrap()); + one("find_aged", store.find_aged(0, u32::MAX, None).unwrap()); + + // deprecated_at is the last column (index 21): it must survive too. + store.deprecate("full-row").unwrap(); + let got = store.get_by_id("full-row").unwrap().unwrap(); + assert!( + got.deprecated && got.deprecated_at.is_some(), + "deprecated_at" + ); + check("get_by_id after deprecate", &got); + } + + /// The `m.`-aliased macro must stay in lockstep with the plain one. + #[test] + fn test_memory_column_macros_agree() { + let plain: Vec<&str> = memory_columns!().split(',').map(str::trim).collect(); + let aliased: Vec<&str> = memory_columns_m!().split(',').map(str::trim).collect(); + assert_eq!(plain.len(), 22, "row_to_memory decodes 22 columns"); + assert_eq!(plain.len(), aliased.len()); + for (p, a) in plain.iter().zip(&aliased) { + assert_eq!(format!("m.{p}"), *a); + } + } + + #[test] + fn test_bulk_cold_skips_pinned() { + let store = Store::open(":memory:").unwrap(); + let mut pinned = make_test_memory("pinned-cold", "keep me", &[]); + pinned.pinned = true; // last_accessed None => otherwise "cold" + store.insert(&pinned).unwrap(); + store + .insert(&make_test_memory("plain-cold", "drop me", &[])) + .unwrap(); + + let found = store.find_ids_cold(None, 30).unwrap(); + assert_eq!(found, vec!["plain-cold".to_string()]); + + let deleted = store.bulk_delete_cold(None, 30).unwrap(); + assert_eq!(deleted, vec!["plain-cold".to_string()]); + assert!(store.get_by_id("pinned-cold").unwrap().is_some()); + } + #[test] fn test_bulk_delete_cold() { let store = Store::open(":memory:").unwrap(); diff --git a/crates/uteke-core/src/memory/tags.rs b/crates/uteke-core/src/memory/tags.rs index 918731f4..14a0675c 100644 --- a/crates/uteke-core/src/memory/tags.rs +++ b/crates/uteke-core/src/memory/tags.rs @@ -190,6 +190,10 @@ impl super::Store { } } if changed { + // A memory that already carried `new` must not end up with it + // twice in the JSON column (#1332). Keep first occurrences. + let mut seen = std::collections::HashSet::new(); + tags.retain(|t| seen.insert(t.clone())); let new_tags_json = serde_json::to_string(&tags) .map_err(|e| Error::db("database operation", e))?; let now = chrono::Utc::now().to_rfc3339(); diff --git a/crates/uteke-core/src/memory/vector.rs b/crates/uteke-core/src/memory/vector.rs index 898b29a2..29ff8ef0 100644 --- a/crates/uteke-core/src/memory/vector.rs +++ b/crates/uteke-core/src/memory/vector.rs @@ -49,6 +49,16 @@ use fs2::FileExt; use std::collections::HashMap; use std::fs::File; use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant}; + +/// vecq compaction (#1324) only kicks in for indexes at least this many rows +/// large, so tiny indexes do not churn. +const COMPACT_MIN_ROWS: usize = 128; + +/// Per-operation saves are deferred until this many mutations pile up (#1322). +const SAVE_EVERY_OPS: u32 = 64; +/// ...or until the last save is this old, whichever comes first. +const SAVE_MAX_AGE: Duration = Duration::from_secs(2); #[cfg(feature = "usearch")] use usearch::{Index, IndexOptions, MetricKind, ScalarKind}; @@ -383,6 +393,12 @@ pub struct VectorIndex { path: Option, /// Whether the index has unsaved changes. dirty: bool, + /// Mutations since the last successful save (#1322). + unsaved_ops: u32, + /// When the index was last written to disk (or loaded/created). + last_saved: Instant, + /// Successful `save()` calls over this instance's lifetime (diagnostics, tests). + saves: u64, /// Cross-process file lock on the index file (#543). /// Held until the VectorIndex is dropped. _lock_file: Option, @@ -410,6 +426,9 @@ impl VectorIndex { next_key: 0, path: None, dirty: false, + unsaved_ops: 0, + last_saved: Instant::now(), + saves: 0, _lock_file: None, }) } @@ -510,6 +529,9 @@ impl VectorIndex { let mut key_to_id = HashMap::new(); let mut id_to_key = HashMap::new(); let mut next_key = 0u64; + // Digest of the index bytes this sidecar was written against (#1325). + // `None` for sidecars written before the header existed. + let mut sidecar_digest: Option = None; let mapping_path = path.with_extension("keys"); match std::fs::read_to_string(&mapping_path) { @@ -519,6 +541,10 @@ impl VectorIndex { if line.is_empty() { continue; } + if let Some(header) = line.strip_prefix(KEYS_HEADER_PREFIX) { + sidecar_digest = header.split('\t').next().map(str::to_string); + continue; + } if let Some((key_str, id)) = line.split_once('\t') { if let Ok(key) = key_str.parse::() { key_to_id.insert(key, id.to_string()); @@ -534,6 +560,14 @@ impl VectorIndex { Err(e) => return Err(Error::embed("read key mapping", e)), } + verify_pair_consistency( + &engine, + backend, + &key_to_id, + sidecar_digest.as_deref(), + &buffer, + )?; + Ok(Self { engine, backend, @@ -542,6 +576,9 @@ impl VectorIndex { next_key, path: None, dirty: false, + unsaved_ops: 0, + last_saved: Instant::now(), + saves: 0, _lock_file: None, }) } @@ -579,8 +616,31 @@ impl VectorIndex { // Write buffer to disk via atomic write (temp file + rename) let tmp_path = path.with_extension(format!("{}.tmp", index_ext_for(self.backend))); - std::fs::write(&tmp_path, &buffer) + // Windows: the existing locked handle is refreshed after the rename + // (below). Elsewhere the rename replaces the inode our flock lives + // on, so lock the NEW file before it is published: the exclusive + // lock then follows the inode onto `path` with no unlocked window, + // and a second process can never lock a different inode. + #[cfg(windows)] + write_synced(&tmp_path, &buffer) .map_err(|e| Error::embed("write temp index file", e))?; + #[cfg(not(windows))] + let new_lock = { + use std::io::Write; + let mut f = File::options() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(&tmp_path) + .map_err(|e| Error::embed("create temp index file", e))?; + f.try_lock_exclusive() + .map_err(|e| Error::embed("lock temp index file", e))?; + f.write_all(&buffer) + .and_then(|()| f.sync_all()) + .map_err(|e| Error::embed("write temp index file", e))?; + f + }; // On Windows, `std::fs::rename` fails with `ERROR_ACCESS_DENIED` if // the destination file is locked by `LockFileEx` via fs2 (#926). @@ -614,6 +674,9 @@ impl VectorIndex { { std::fs::rename(&tmp_path, path) .map_err(|e| Error::embed("rename temp to final index file", e))?; + // The old (now orphaned) inode's lock is released when the + // previous handle drops; the new handle holds the live lock. + self._lock_file = Some(new_lock); } // On Windows, reopen the file after rename to refresh the lock @@ -639,12 +702,44 @@ impl VectorIndex { // Save key→id mapping as sidecar file using atomic write let mapping_path = path.with_extension("keys"); let mut lines = Vec::new(); + // First line ties this sidecar to the exact index bytes just + // published, so a crash between the two renames is detected on the + // next load instead of leaving a silently stale key map (#1325). + // Older binaries skip it (it does not parse as `\t`). + lines.push(format!( + "{KEYS_HEADER_PREFIX}{}\t{}", + index_digest(&buffer), + self.engine.len() + )); for (&key, id) in &self.key_to_id { lines.push(format!("{key}\t{id}")); } atomic_write(&mapping_path, lines.join("\n").as_bytes())?; self.dirty = false; + self.unsaved_ops = 0; + self.last_saved = Instant::now(); + self.saves += 1; + } + Ok(()) + } + + /// Save only when enough has piled up since the last save (#1322). + /// + /// `save()` rewrites the whole index and the whole `.keys` sidecar, so + /// saving on every write makes bulk ingest O(N²). Per-operation paths call + /// this instead: it writes once `SAVE_EVERY_OPS` mutations or + /// `SAVE_MAX_AGE` have accumulated. Nothing is lost for good in between — + /// SQLite already holds every row, and `shutdown()`, `Drop`, and the + /// server's periodic flush write whatever is still pending. After a hard + /// kill the on-disk index can trail SQLite by up to that window; + /// `uteke verify` / `repair` resync it. + pub fn save_if_due(&mut self) -> Result<(), Error> { + if !self.dirty { + return Ok(()); + } + if self.unsaved_ops >= SAVE_EVERY_OPS || self.last_saved.elapsed() >= SAVE_MAX_AGE { + return self.save(); } Ok(()) } @@ -653,8 +748,11 @@ impl VectorIndex { /// Used for migration from old HNSW or full rebuild. pub fn build(&mut self, items: &[(String, Vec)]) -> Result<(), Error> { // Reset (same engine, fresh instance) + // An empty rebuild keeps the index's CURRENT dimensionality: resetting + // to DEFAULT_DIMS (768) broke the next insert of a non-768 embedder + // after `repair()` on an empty store (#1332). let dims = if items.is_empty() { - DEFAULT_DIMS + self.engine.dims() } else { items[0].1.len() }; @@ -696,15 +794,6 @@ impl VectorIndex { ))); } - // Guard: remove old entry if ID already exists (prevents duplicate + stale slot) - if let Some(old_key) = self.id_to_key.get(id) { - let old_key = *old_key; - self.key_to_id.remove(&old_key); - self.engine.remove(old_key); - // vecq has no incremental delete — the dead row is filtered out of - // search via the key map (key_to_id no longer contains old_key). - } - let key = if self.backend == VectorBackend::Vecq { // vecq rows are append-only: the new entry lands at physical row // `len()`, and search maps rows back via key_to_id — so the @@ -725,20 +814,30 @@ impl VectorIndex { } key } else { + // Auto-grow usearch capacity when full (fallible — runs before any + // state is touched so a failure leaves the index unchanged). + self.engine.ensure_capacity()?; let key = self.next_key; self.next_key = self.next_key.saturating_add(1); - // Auto-grow usearch capacity when full. - self.engine.ensure_capacity()?; key }; - self.key_to_id.insert(key, id.to_string()); - self.id_to_key.insert(id.to_string(), key); - // vecq assigns rows sequentially; row == key by construction. + // Add first: if it fails, the previous entry for `id` is still intact. self.engine.add(key, embedding)?; + // Replace the old entry only once the new vector is in the engine. + if let Some(old_key) = self.id_to_key.get(id).copied() { + self.key_to_id.remove(&old_key); + self.engine.remove(old_key); + // vecq has no incremental delete — the dead row is filtered out of + // search via the key map (key_to_id no longer contains old_key). + } + self.key_to_id.insert(key, id.to_string()); + self.id_to_key.insert(id.to_string(), key); + self.dirty = true; + self.unsaved_ops = self.unsaved_ops.saturating_add(1); Ok(()) } @@ -750,6 +849,7 @@ impl VectorIndex { // vecq: tombstone is implicit — the key vanishes from the map, so // search results referencing that row are filtered out below. self.dirty = true; + self.unsaved_ops = self.unsaved_ops.saturating_add(1); true } else { false @@ -815,12 +915,57 @@ impl VectorIndex { self.len() == 0 } + /// Successful saves over this instance's lifetime (diagnostics, tests). + #[allow(dead_code)] + pub(crate) fn save_count(&self) -> u64 { + self.saves + } + + /// Physical rows in the engine, dead ones included (diagnostics, compaction). + pub fn physical_rows(&self) -> usize { + self.engine.len() + } + + /// Dead (tombstoned) rows: physically present, no longer in the key map. + /// Only vecq is append-only; usearch removes in place, so it has none. + pub fn dead_rows(&self) -> usize { + if self.backend == VectorBackend::Vecq { + self.engine.len().saturating_sub(self.key_to_id.len()) + } else { + 0 + } + } + + /// Whether a rebuild from the live rows is worth doing (#1324): a vecq + /// index of at least `COMPACT_MIN_ROWS` rows with more than a quarter of + /// them dead. Dead rows cost file size and make every search over-fetch + /// `k + dead` from an O(N) scan. + pub fn needs_compaction(&self) -> bool { + let total = self.engine.len(); + self.backend == VectorBackend::Vecq + && !self.key_to_id.is_empty() + && total >= COMPACT_MIN_ROWS + && self.dead_rows().saturating_mul(4) > total + } + /// Whether the index has unsaved changes. pub fn is_dirty(&self) -> bool { self.dirty } } +/// Flush pending mutations when the index goes away, so a process that exits +/// without calling `shutdown()` (CLI, MCP stdio) still persists its writes. +impl Drop for VectorIndex { + fn drop(&mut self) { + if self.dirty && self.path.is_some() { + if let Err(e) = self.save() { + tracing::warn!("Failed to save vector index on drop: {e}"); + } + } + } +} + impl Default for VectorIndex { fn default() -> Self { Self::new(DEFAULT_DIMS).expect("Failed to create default vector index") @@ -882,12 +1027,14 @@ fn acquire_file_lock(path: &Path) -> Result { } if std::time::Instant::now() >= deadline { - return Err(Error::embed_msg(format!( - "Could not acquire lock on {} after {MAX_WAIT:?}. \ - Another uteke process (uteke-serve or CLI) may be running. \ - Stop it and retry.", - path.display() - ))); + return Err(Error::Lock { + context: format!( + "Could not acquire lock on {} after {MAX_WAIT:?}. \ + Another uteke process (uteke-serve or CLI) may be running. \ + Stop it and retry.", + path.display() + ), + }); } tracing::trace!( @@ -898,12 +1045,88 @@ fn acquire_file_lock(path: &Path) -> Result { } } +/// Write `data` to `path` and fsync before returning, so a following rename +/// can never publish a zero-length or partial file after power loss. +/// Marks the first line of a `.keys` sidecar: `#uteke-keys\t\t`. `#` cannot start a numeric key, so older readers +/// ignore the line. +const KEYS_HEADER_PREFIX: &str = "#uteke-keys\t"; + +/// Hex SHA-256 of the serialized index bytes. +fn index_digest(bytes: &[u8]) -> String { + use sha2::{Digest, Sha256}; + Sha256::digest(bytes) + .iter() + .map(|b| format!("{b:02x}")) + .collect() +} + +/// Reject an index file / `.keys` sidecar pair that cannot be trusted (#1325). +/// +/// The two files are published by separate renames, so a crash between them +/// leaves a mismatched pair. Returning an error makes the caller discard both +/// files and rebuild from SQLite, instead of serving a stale key map (live +/// rows that look dead, or a `next_key` that restarts too low and collides). +fn verify_pair_consistency( + engine: &Engine, + backend: VectorBackend, + key_to_id: &HashMap, + sidecar_digest: Option<&str>, + index_bytes: &[u8], +) -> Result<(), Error> { + let rows = engine.len(); + + // Exact check: the sidecar names the index it was written against. + if let Some(expected) = sidecar_digest { + if expected != index_digest(index_bytes) { + return Err(Error::embed_msg( + "vector index and .keys sidecar do not match (interrupted save?)", + )); + } + } + + // Structural checks that also cover sidecars without a header (written by + // older binaries) and a missing sidecar. + if rows > 0 && key_to_id.is_empty() { + return Err(Error::embed_msg( + "vector index has entries but the .keys sidecar is empty or missing", + )); + } + match backend { + // vecq: a key IS a physical row, so no live key may point past the end + // and there cannot be more live keys than rows. + VectorBackend::Vecq => { + if key_to_id.len() > rows || key_to_id.keys().any(|&k| k >= rows as u64) { + return Err(Error::embed_msg( + "vecq .keys sidecar references rows that are not in the index", + )); + } + } + // usearch: size() counts live entries, so the key map cannot hold more. + VectorBackend::Usearch => { + if key_to_id.len() > rows { + return Err(Error::embed_msg( + "usearch .keys sidecar holds more keys than the index has entries", + )); + } + } + } + Ok(()) +} + +fn write_synced(path: &std::path::Path, data: &[u8]) -> std::io::Result<()> { + use std::io::Write; + let mut f = std::fs::File::create(path)?; + f.write_all(data)?; + f.sync_all() +} + /// Atomic file write: write to temp file then rename. /// Prevents corruption if process crashes mid-write. /// POSIX guarantees rename() is atomic on the same filesystem. fn atomic_write(path: &std::path::Path, data: &[u8]) -> Result<(), Error> { let tmp_path = path.with_extension("keys.tmp"); - std::fs::write(&tmp_path, data).map_err(|e| Error::embed("write temp key mapping", e))?; + write_synced(&tmp_path, data).map_err(|e| Error::embed("write temp key mapping", e))?; std::fs::rename(&tmp_path, path) .map_err(|e| Error::embed("rename temp to final key mapping", e))?; Ok(()) @@ -973,6 +1196,30 @@ mod tests { assert!(results.iter().all(|(id, _)| id != "m1")); } + /// Regression: after `save()` renames a new file over the index, the + /// exclusive lock must live on the file now at `path` (not the orphaned + /// inode), so a second opener cannot lock it. + #[cfg(unix)] + #[test] + fn test_save_keeps_exclusive_lock_on_published_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join(format!("lock.{INDEX_EXT}")); + let mut idx = VectorIndex::load_or_create(&path, 8).unwrap(); + idx.insert("a", &make_vec(8, 0)).unwrap(); + idx.save().unwrap(); + idx.insert("b", &make_vec(8, 1)).unwrap(); + idx.save().unwrap(); // second save: lock must survive repeated renames + + let other = File::options().read(true).write(true).open(&path).unwrap(); + assert!( + other.try_lock_exclusive().is_err(), + "published index file must still be exclusively locked by its owner" + ); + drop(idx); + let other = File::options().read(true).write(true).open(&path).unwrap(); + assert!(other.try_lock_exclusive().is_ok(), "lock released on drop"); + } + #[test] fn test_save_and_load() { let dir = tempfile::tempdir().unwrap(); @@ -1162,4 +1409,235 @@ mod tests { assert_eq!(results.len(), 1); } } + + // ── #1325: index file / .keys sidecar consistency ────────────────────── + + fn compiled_backends() -> Vec { + [VectorBackend::Usearch, VectorBackend::Vecq] + .into_iter() + .filter(|b| b.is_compiled_in()) + .collect() + } + + /// Save `n` items, return (index path, index with the lock still held). + fn saved_index( + dir: &std::path::Path, + backend: VectorBackend, + n: usize, + ) -> (PathBuf, VectorIndex) { + let path = dir.join(format!("idx.{}", index_ext_for(backend))); + let mut idx = VectorIndex::with_backend(backend, 64).unwrap(); + idx.path = Some(path.clone()); + for i in 0..n { + idx.insert(&format!("mem-{i}"), &make_vec(64, i)).unwrap(); + } + idx.save().unwrap(); + (path, idx) + } + + #[test] + fn stale_sidecar_after_interrupted_save_is_rejected() { + for backend in compiled_backends() { + let dir = tempfile::tempdir().unwrap(); + let (path, mut idx) = saved_index(dir.path(), backend, 2); + let keys = path.with_extension("keys"); + let stale = std::fs::read(&keys).unwrap(); + + // Next save publishes the new index, then "crashes" before the + // sidecar rename: put the previous sidecar back. + idx.insert("mem-2", &make_vec(64, 2)).unwrap(); + idx.save().unwrap(); + std::fs::write(&keys, stale).unwrap(); + + let err = VectorIndex::load(&path) + .err() + .unwrap_or_else(|| panic!("{backend:?}: stale sidecar must not load")); + assert!( + err.to_string().contains("do not match"), + "{backend:?}: unexpected error: {err}" + ); + } + } + + #[test] + fn matching_pair_survives_load_and_resave() { + for backend in compiled_backends() { + let dir = tempfile::tempdir().unwrap(); + let (path, idx) = saved_index(dir.path(), backend, 3); + drop(idx); + + let mut loaded = VectorIndex::load(&path).unwrap(); + assert_eq!(loaded.len(), 3, "{backend:?}"); + loaded.path = Some(path.clone()); + loaded.save().unwrap(); + let again = VectorIndex::load(&path).unwrap(); + assert_eq!(again.len(), 3, "{backend:?}"); + } + } + + #[test] + fn sidecar_without_header_still_loads() { + for backend in compiled_backends() { + let dir = tempfile::tempdir().unwrap(); + let (path, idx) = saved_index(dir.path(), backend, 2); + drop(idx); + + // A sidecar from a binary that predates the header. + let keys = path.with_extension("keys"); + let body: String = std::fs::read_to_string(&keys) + .unwrap() + .lines() + .filter(|l| !l.starts_with(KEYS_HEADER_PREFIX)) + .collect::>() + .join("\n"); + std::fs::write(&keys, body).unwrap(); + + let loaded = VectorIndex::load(&path).unwrap(); + assert_eq!(loaded.len(), 2, "{backend:?}"); + } + } + + #[test] + fn missing_or_empty_sidecar_with_entries_is_rejected() { + for backend in compiled_backends() { + let dir = tempfile::tempdir().unwrap(); + let (path, idx) = saved_index(dir.path(), backend, 2); + drop(idx); + let keys = path.with_extension("keys"); + + std::fs::remove_file(&keys).unwrap(); + assert!( + VectorIndex::load(&path).is_err(), + "{backend:?}: missing sidecar must be rejected" + ); + + std::fs::write(&keys, "").unwrap(); + assert!( + VectorIndex::load(&path).is_err(), + "{backend:?}: empty sidecar must be rejected" + ); + } + } + + #[test] + fn headerless_sidecar_with_out_of_range_keys_is_rejected() { + for backend in compiled_backends() { + let dir = tempfile::tempdir().unwrap(); + let (path, idx) = saved_index(dir.path(), backend, 2); + drop(idx); + let keys = path.with_extension("keys"); + // Legacy-style sidecar claiming three entries for a two-entry index. + std::fs::write(&keys, "0\tmem-0\n1\tmem-1\n2\tmem-2").unwrap(); + assert!( + VectorIndex::load(&path).is_err(), + "{backend:?}: extra keys must be rejected" + ); + } + } + + // ── #1322: deferred saves ────────────────────────────────────────────── + + #[test] + fn save_if_due_batches_writes_and_drop_flushes_the_rest() { + for backend in compiled_backends() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join(format!("idx.{}", index_ext_for(backend))); + let mut idx = VectorIndex::with_backend(backend, 64).unwrap(); + idx.path = Some(path.clone()); + + // 130 per-operation writes: a save at op 64 and 128, not 130 saves. + for i in 0..130 { + idx.insert(&format!("mem-{i}"), &make_vec(64, i % 64)) + .unwrap(); + idx.save_if_due().unwrap(); + } + assert_eq!(idx.save_count(), 2, "{backend:?}: expected batched saves"); + assert!(idx.is_dirty(), "{backend:?}: two ops are still pending"); + + // Going away flushes the pending tail. + drop(idx); + let loaded = VectorIndex::load(&path).unwrap(); + assert_eq!(loaded.len(), 130, "{backend:?}: pending ops lost on drop"); + } + } + + #[test] + fn save_if_due_is_a_noop_when_clean() { + let dir = tempfile::tempdir().unwrap(); + let (_path, mut idx) = saved_index(dir.path(), compiled_backends()[0], 1); + let before = idx.save_count(); + idx.save_if_due().unwrap(); + assert_eq!(idx.save_count(), before); + } + + // ── #1324: vecq compaction ───────────────────────────────────────────── + + #[cfg(feature = "vecq")] + #[test] + fn vecq_dead_rows_trigger_compaction_threshold() { + let mut idx = VectorIndex::with_backend(VectorBackend::Vecq, 64).unwrap(); + for i in 0..200 { + idx.insert(&format!("mem-{i}"), &make_vec(64, i % 64)) + .unwrap(); + } + assert_eq!((idx.dead_rows(), idx.needs_compaction()), (0, false)); + + // 25% dead exactly is not enough ("more than a quarter"). + for i in 0..50 { + assert!(idx.remove(&format!("mem-{i}"))); + } + assert_eq!(idx.dead_rows(), 50); + assert!(!idx.needs_compaction()); + + assert!(idx.remove("mem-50")); + assert!(idx.needs_compaction(), "51/200 dead is over a quarter"); + + // Rebuilding from the live rows drops every dead row. + let live: Vec<(String, Vec)> = (51..200) + .map(|i| (format!("mem-{i}"), make_vec(64, i % 64))) + .collect(); + idx.build(&live).unwrap(); + assert_eq!((idx.physical_rows(), idx.dead_rows()), (149, 0)); + assert!(!idx.needs_compaction()); + } + + #[cfg(feature = "vecq")] + #[test] + fn small_vecq_indexes_never_compact() { + let mut idx = VectorIndex::with_backend(VectorBackend::Vecq, 64).unwrap(); + for i in 0..40 { + idx.insert(&format!("mem-{i}"), &make_vec(64, i)).unwrap(); + } + for i in 0..39 { + idx.remove(&format!("mem-{i}")); + } + assert!(!idx.needs_compaction(), "below COMPACT_MIN_ROWS"); + } + + #[cfg(feature = "usearch")] + #[test] + fn usearch_has_no_dead_rows() { + let mut idx = VectorIndex::with_backend(VectorBackend::Usearch, 64).unwrap(); + for i in 0..200 { + idx.insert(&format!("mem-{i}"), &make_vec(64, i % 64)) + .unwrap(); + } + for i in 0..150 { + idx.remove(&format!("mem-{i}")); + } + assert_eq!(idx.dead_rows(), 0); + assert!(!idx.needs_compaction()); + } + + /// #1332: an empty rebuild (e.g. `repair()` on an empty store) must keep the + /// index's dimensionality instead of resetting it to 768. + #[test] + fn test_build_empty_keeps_current_dims() { + let mut idx = VectorIndex::new(384).unwrap(); + idx.build(&[]).unwrap(); + assert_eq!(idx.dims(), 384); + idx.insert("a", &vec![0.5_f32; 384]) + .expect("a 384-dim embedding must still be accepted after build(&[])"); + assert_eq!(idx.len(), 1); + } } diff --git a/crates/uteke-core/src/operations.rs b/crates/uteke-core/src/operations.rs index 498397a1..8fcc6384 100644 --- a/crates/uteke-core/src/operations.rs +++ b/crates/uteke-core/src/operations.rs @@ -1,6 +1,7 @@ //! Core memory operations: remember, recall, search, forget, list, get, tags. use crate::error::Error; +use crate::index_sync::{self, SyncPolicy}; use crate::memory::types::{ BulkDeleteResult, DEFAULT_NAMESPACE, Memory, MemoryTier, NamespaceDeleteResult, NamespaceRenameResult, RecallStrategy, SearchResult, TagInfo, @@ -339,14 +340,16 @@ impl crate::Uteke { return Ok(None); } - // Filter by namespace if specified. - let ns_set: Option> = if let Some(ns) = namespace { - match self.store.memories_in_namespace(ns) { - Ok(ids) => Some(ids.into_iter().collect()), - Err(_) => return Ok(None), - } - } else { - None + // Namespace + liveness filter over the (few) candidates only; the old + // code loaded every id of the namespace on each remember. + let candidate_ids: Vec<&str> = results + .iter() + .map(|(id, _)| id.as_str()) + .filter(|id| !id.starts_with("chunk:")) + .collect(); + let live = match self.store.live_ids_in_namespace(&candidate_ids, namespace) { + Ok(set) => set, + Err(_) => return Ok(None), }; for (id, dist) in &results { @@ -354,11 +357,9 @@ impl crate::Uteke { if id.starts_with("chunk:") { continue; } - // Namespace filter. - if let Some(ref set) = ns_set { - if !set.contains(id) { - continue; - } + // Namespace + liveness filter. + if !live.contains(id) { + continue; } let sim = (1.0 - dist).clamp(0.0, 1.0); if sim >= DEDUP_THRESHOLD { @@ -457,14 +458,12 @@ impl crate::Uteke { author_type: "agent".to_string(), }; - // Acquire index write lock BEFORE any writes so lock failures are detected early. - // If SQLite commit fails after index insert, the orphan index entry is harmless - // and will be cleaned up by verify/repair. - let mut index = self - .index - .write() - .map_err(|_| Error::lock("index write lock during remember"))?; - + // SQLite is the source of truth and is written first, WITHOUT holding the + // index write lock: the timeline / provenance / edge wiring below issue + // many queries and would otherwise block every concurrent recall for + // their whole duration. The vector is mirrored afterwards under a short + // write lock; if that fails the row stays FTS5-only and the outcome says + // so (#1273), `uteke repair` can backfill. self.store.insert(&memory)?; // Provenance: record the content hash at write time (#1172 Fase 1). @@ -507,66 +506,25 @@ impl crate::Uteke { // carries embedding_written=false + warning so every surface can // tell the truth; `uteke repair` can backfill the vector later. let mut vector_error = embedding_error; + // One policy for insert + persist (retry, deterministic errors not + // retried) lives in `index_sync`; the outcome says what landed (#1273). + // An empty embedding means "no embedder": nothing to mirror, no lock. if !embedding.is_empty() { - // Retry the index insert itself (in-memory op, cheap) — #1273. - let mut insert_ok = false; - let mut last_err: Option = None; - for attempt in 0..3 { - match index.insert(&id, embedding) { - Ok(()) => { - insert_ok = true; - break; - } - Err(e) => { - last_err = Some(e.to_string()); - if attempt < 2 { - tracing::warn!( - "Vector insert attempt {}/3 failed for id={id}: {e}. Retrying...", - attempt + 1 - ); - std::thread::sleep(std::time::Duration::from_millis(200)); - } + match self.index.write() { + Ok(mut index) => { + let sync = + index_sync::upsert(&mut *index, &id, embedding, &SyncPolicy::STANDARD); + if let Some(msg) = sync.error_message() { + vector_error = Some(msg); } + // The write lock must be dropped BEFORE auto_link_cosine + // (it takes a read lock on the same index, #442). } - } - if insert_ok { - // Retry index persistence up to 3 times (#621). - // A failed save means the in-memory index has the entry but - // on-disk doesn't → silent desync on next process launch. - for attempt in 0..3 { - match index.save() { - Ok(()) => break, - Err(e) => { - if attempt < 2 { - tracing::warn!( - "Index save attempt {}/3 failed after remember for id={id}: {e}. Retrying...", - attempt + 1 - ); - std::thread::sleep(std::time::Duration::from_millis(200)); - } else { - tracing::warn!( - "Failed to persist vector index after 3 attempts for remember id={id}: {e}. \ - Index entry can be rebuilt via `uteke repair`." - ); - vector_error = - Some(format!("index persist failed after 3 attempts: {e}")); - } - } - } + Err(_) => { + vector_error = Some("index write lock poisoned; vector not written".into()); } - } else { - let e = last_err.unwrap_or_else(|| "unknown vector insert error".to_string()); - tracing::warn!( - "Vector insert failed after 3 attempts for remember id={id}: {e}. \ - Memory stored FTS5-only; run `uteke repair` to backfill." - ); - vector_error = Some(format!("vector insert failed after 3 attempts: {e}")); } } - // Drop the write lock BEFORE auto_link_cosine to prevent deadlock. - // auto_link_cosine needs a read lock on the same index — holding - // the write lock here would deadlock (#442). - drop(index); // Cosine-similarity auto-linking (#401). // Must run AFTER index.insert() so the new memory is searchable. @@ -1307,43 +1265,13 @@ impl crate::Uteke { "Memory with id='{id}' not found in store. Nothing was deleted." ))); } - // Vector index remove — orphan is harmless if fails (verify/repair cleans up) - if !index.remove(id) { + // Vector index remove + persist. A stale on-disk index is surfaced as an + // error (#926): the DB row is gone, `uteke repair` resyncs the index. + let sync = index_sync::remove_ids(&mut *index, [id], &SyncPolicy::STANDARD); + if sync.missing > 0 { tracing::warn!("Vector index entry not found during forget for id={id}"); } - // Retry index persistence up to 3 times (#621). - // A failed save leaves orphan entries that desync from SQLite. - let mut last_err = None; - for attempt in 0..3 { - match index.save() { - Ok(()) => return Ok(()), - Err(e) => { - last_err = Some(e); - if attempt < 2 { - tracing::warn!( - "Index save attempt {}/3 failed after forget for id={id}: {}. Retrying...", - attempt + 1, - last_err.as_ref().unwrap() - ); - std::thread::sleep(std::time::Duration::from_millis(200)); - } - } - } - } - - // All retries exhausted: SQLite row is deleted but index persistence failed (#926). - // Return Err so the caller knows the operation was only partially successful. - // The DB is the source of truth; `repair` will resync the index. - tracing::error!( - "Failed to persist vector index after 3 attempts for forget id={id}. \ - SQLite row deleted but index is stale. Run `uteke repair` to resync." - ); - Err(Error::embed_msg(format!( - "Memory {id} was deleted from the database, but the vector index \ - could not be saved after 3 attempts: {}. \ - Run `uteke repair` to resync the index.", - last_err.unwrap() - ))) + sync.into_delete_result(&format!("forget id={id}")) } /// Soft-delete (deprecate) a memory with reason (#929). @@ -1358,14 +1286,11 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during soft_forget"))?; - if !index.remove(id) { - tracing::debug!( - "Vector index entry not found during soft_forget for id={id} (ok if never embedded)" - ); - } - // Persist vector index to disk so the removal survives restart. - if let Err(e) = index.save() { - tracing::warn!("Failed to persist vector index after soft_forget: {e}"); + // Remove + persist so the removal survives restart; a stale on-disk + // index only warns here (the row is deprecated, recall filters it). + let sync = index_sync::remove_ids(&mut *index, [id], &SyncPolicy::STANDARD); + if let Some(msg) = sync.error_message() { + tracing::warn!("Failed to persist vector index after soft_forget id={id}: {msg}"); } // Invalidate recall cache for this memory's namespace. if let Some(memory) = self.store.get_by_id(id).ok().flatten() { @@ -1391,13 +1316,18 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during promote"))?; - if let Err(e) = index.insert(&memory.id, &memory.embedding) { + let sync = index_sync::upsert( + &mut *index, + &memory.id, + &memory.embedding, + &SyncPolicy::STANDARD, + ); + if let Some(msg) = sync.error_message() { tracing::warn!( - "Failed to re-insert memory id={} into vector index during promote: {e}", + "Failed to re-add memory id={} to the vector index during promote: {msg}", memory.id ); } - let _ = index.save(); } self.recall_cache.invalidate_namespace(&memory.namespace); } @@ -1419,10 +1349,12 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during bulk_forget_by_tag"))?; - for id in &ids { - index.remove(id); - } - persist_index_after_delete(&mut index, "bulk_forget_by_tag (soft)")?; + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .into_delete_result("bulk_forget_by_tag (soft)")?; if let Some(ns) = namespace { self.recall_cache.invalidate_namespace(ns); } else { @@ -1440,14 +1372,12 @@ impl crate::Uteke { .write() .map_err(|_| Error::lock("index write lock during bulk_forget_by_tag"))?; let ids = self.store.bulk_delete_by_tag(tag, namespace)?; - for id in &ids { - if !index.remove(id) { - tracing::warn!( - "Vector index entry not found during bulk_forget_by_tag for id={id}" - ); - } - } - persist_index_after_delete(&mut index, "bulk_forget_by_tag")?; + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .into_delete_result("bulk_forget_by_tag")?; Ok(BulkDeleteResult { deleted: ids.len(), ids, @@ -1466,10 +1396,12 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during bulk_forget_cold"))?; - for id in &ids { - index.remove(id); - } - persist_index_after_delete(&mut index, "bulk_forget_cold (soft)")?; + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .into_delete_result("bulk_forget_cold (soft)")?; if let Some(ns) = namespace { self.recall_cache.invalidate_namespace(ns); } else { @@ -1489,12 +1421,12 @@ impl crate::Uteke { let ids = self .store .bulk_delete_cold(namespace, self.tier_config.warm_days)?; - for id in &ids { - if !index.remove(id) { - tracing::warn!("Vector index entry not found during bulk_forget_cold for id={id}"); - } - } - persist_index_after_delete(&mut index, "bulk_forget_cold")?; + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .into_delete_result("bulk_forget_cold")?; Ok(BulkDeleteResult { deleted: ids.len(), ids, @@ -1511,10 +1443,12 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during bulk_forget_all"))?; - for id in &ids { - index.remove(id); - } - persist_index_after_delete(&mut index, "bulk_forget_all (soft)")?; + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .into_delete_result("bulk_forget_all (soft)")?; if let Some(ns) = namespace { self.recall_cache.invalidate_namespace(ns); } else { @@ -1532,12 +1466,12 @@ impl crate::Uteke { .write() .map_err(|_| Error::lock("index write lock during bulk_forget_all"))?; let ids = self.store.bulk_delete_all(namespace)?; - for id in &ids { - if !index.remove(id) { - tracing::warn!("Vector index entry not found during bulk_forget_all for id={id}"); - } - } - persist_index_after_delete(&mut index, "bulk_forget_all")?; + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .into_delete_result("bulk_forget_all")?; Ok(BulkDeleteResult { deleted: ids.len(), ids, @@ -1714,10 +1648,12 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during delete_namespace"))?; - for id in &ids { - index.remove(id); - } - persist_index_after_delete(&mut index, "delete_namespace (deprecate)")?; + index_sync::remove_ids( + &mut *index, + ids.iter().map(String::as_str), + &SyncPolicy::STANDARD, + ) + .into_delete_result("delete_namespace (deprecate)")?; self.recall_cache.invalidate_namespace(name); tracing::info!( "Namespace delete (deprecate): '{name}' soft-deleted {affected} memories \ @@ -1918,26 +1854,17 @@ impl crate::Uteke { .index .write() .map_err(|_| Error::lock("index write lock during update_memory"))?; - index.insert(id, &new_embedding)?; - - // Persist index with retry - for attempt in 0..3 { - match index.save() { - Ok(()) => break, - Err(e) => { - if attempt < 2 { - tracing::warn!( - "Index save attempt {}/3 failed after update_memory for id={id}: {e}. Retrying...", - attempt + 1 - ); - std::thread::sleep(std::time::Duration::from_millis(200)); - } else { - tracing::error!( - "Index save failed after 3 attempts for id={id}: {e}. Index may be stale on next launch." - ); - } - } - } + let sync = index_sync::upsert(&mut *index, id, &new_embedding, &SyncPolicy::STANDARD); + if let Some(e) = &sync.insert_error { + return Err(Error::embed_msg(format!( + "Memory {id} was updated in the database, but its vector could not be \ + written to the index: {e}. Run `uteke repair` to resync." + ))); + } + if let Some(msg) = sync.error_message() { + tracing::error!( + "{msg} (update_memory id={id}). Index may be stale on next launch." + ); } } else { // No content change — just update SQLite fields @@ -2041,47 +1968,6 @@ impl crate::Uteke { } } -/// Persist the vector index to disk after a delete operation. -/// -/// Retries up to 3 times with 200ms backoff. Unlike the old per-call inline -/// code, this helper **returns an error** on persistent failure so the caller -/// can surface it to the user instead of silently swallowing it (#926). -/// -/// The database is always the source of truth: if the index save fails, the -/// rows are already gone from SQLite. `uteke repair` will resync the index. -fn persist_index_after_delete( - index: &mut crate::memory::vector::VectorIndex, - context: &str, -) -> Result<(), Error> { - let mut last_err = None; - for attempt in 0..3 { - match index.save() { - Ok(()) => return Ok(()), - Err(e) => { - last_err = Some(e); - if attempt < 2 { - tracing::warn!( - "Index save attempt {}/3 failed after {context}: {}. Retrying...", - attempt + 1, - last_err.as_ref().unwrap() - ); - std::thread::sleep(std::time::Duration::from_millis(200)); - } - } - } - } - - tracing::error!( - "Failed to persist vector index after 3 attempts in {context}. \ - SQLite rows deleted but index is stale. Run `uteke repair` to resync." - ); - Err(Error::embed_msg(format!( - "Vector index could not be saved after 3 attempts in {context}: {}. \ - Database rows were deleted. Run `uteke repair` to resync the index.", - last_err.unwrap() - ))) -} - #[cfg(test)] mod forget_tests { use crate::Uteke; diff --git a/crates/uteke-core/src/orphans.rs b/crates/uteke-core/src/orphans.rs index decc7f75..e8c895bb 100644 --- a/crates/uteke-core/src/orphans.rs +++ b/crates/uteke-core/src/orphans.rs @@ -11,6 +11,7 @@ use crate::error::Error; use crate::memory::store::Store; +use crate::memory::store::memory_columns_m; use crate::memory::types::Memory; use rusqlite::params; use serde::{Deserialize, Serialize}; @@ -53,12 +54,10 @@ impl Store { // Single SQL pass: LEFT JOIN against the edge table twice (for // outgoing and incoming), filtering memories with zero matches on // both sides plus the remaining criteria. - let sql = r#" - SELECT m.id, m.content, m.embedding, m.tags, m.metadata, - m.created_at, m.updated_at, m.namespace, m.access_count, - m.last_accessed, m.deprecated, m.valid_from, m.valid_until, - m.memory_type, m.importance, m.pinned, m.content_type, m.slug, - m.source, m.source_type + let sql = concat!( + "SELECT ", + memory_columns_m!(), + r#" FROM memories m LEFT JOIN memory_edges out_e ON out_e.source_id = m.id LEFT JOIN memory_edges in_e ON in_e.target_id = m.id @@ -71,7 +70,8 @@ impl Store { AND (?2 IS NULL OR m.namespace = ?2) GROUP BY m.id ORDER BY m.importance ASC, m.created_at ASC - "#; + "# + ); let mut stmt = self .conn diff --git a/crates/uteke-core/src/types.rs b/crates/uteke-core/src/types.rs index 4e6411f7..511024b9 100644 --- a/crates/uteke-core/src/types.rs +++ b/crates/uteke-core/src/types.rs @@ -57,6 +57,19 @@ pub struct RepairReport { pub chunk_count: usize, } +/// Result of compacting the vector index (#1324). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CompactionReport { + /// Physical rows in the index before compaction. + pub rows_before: usize, + /// Dead (tombstoned) rows among them. + pub dead_before: usize, + /// Physical rows after compaction (the live set; `rows_before` when dry-run). + pub rows_after: usize, + /// True when nothing was changed (`dry_run`). + pub dry_run: bool, +} + /// Result of `uteke repair --reembed`. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ReembedReport { diff --git a/crates/uteke-mcp/Cargo.toml b/crates/uteke-mcp/Cargo.toml index 1b1e5ab1..80beb2c0 100644 --- a/crates/uteke-mcp/Cargo.toml +++ b/crates/uteke-mcp/Cargo.toml @@ -29,10 +29,11 @@ usearch = ["uteke-core/usearch"] vecq = ["uteke-core/vecq"] [dependencies] -uteke-core = { path = "../uteke-core", version = "0.19.1", default-features = false, features = ["onnx"] } +uteke-core = { path = "../uteke-core", version = "0.20.0", default-features = false, features = ["onnx"] } serde = { version = "1", features = ["derive"] } serde_json = "1" dirs = "7" +uuid = "1" [dev-dependencies] uuid = { version = "1", features = ["v4"] } diff --git a/crates/uteke-mcp/src/lib.rs b/crates/uteke-mcp/src/lib.rs index 60337d0a..bc23116d 100644 --- a/crates/uteke-mcp/src/lib.rs +++ b/crates/uteke-mcp/src/lib.rs @@ -301,7 +301,8 @@ fn tool_recall() -> Value { "explain": { "type": "boolean", "description": "Return per-result ranking signals (#1160): vector similarity/rank, RRF contributions, jaccard/salience/recency/graph boosts. Memory-only — omitted type is treated as memory; explicit type=all/doc is rejected." }, "pack": { "type": "boolean", "description": "Return a budgeted context pack (#1281): {selected, skipped, budget_used, budget_chars} instead of a bare list. Deterministic, LLM-free, rank-order preserving. Pair with budget_chars and exclude_ids.", "default": false }, "budget_chars": { "type": "integer", "description": "Character budget for pack mode (default 4000).", "default": 4000 }, - "exclude_ids": { "type": "array", "items": { "type": "string" }, "description": "Memory IDs already injected this turn; excluded from the pack and reported as skipped[reason=excluded]." } + "exclude_ids": { "type": "array", "items": { "type": "string" }, "description": "Memory IDs already injected this turn; excluded from the pack and reported as skipped[reason=excluded]." }, + "full_ids": { "type": "boolean", "description": "Print full UUIDs instead of the 8-char prefix. Use when prefixes are ambiguous (ids are UUIDv7: memories written within ~65s share a prefix, #1357).", "default": false } }, "required": ["query"] } @@ -318,7 +319,8 @@ fn tool_list() -> Value { "tag": { "type": "string", "description": "Filter by tag (optional)" }, "limit": { "type": "integer", "description": "Max results (default 20)", "default": 20 }, "offset": { "type": "integer", "description": "Pagination offset (default 0)", "default": 0 }, - "namespace": { "type": "string", "description": "Namespace (optional)" } + "namespace": { "type": "string", "description": "Namespace (optional)" }, + "full_ids": { "type": "boolean", "description": "Print full UUIDs instead of the 8-char prefix. Use when prefixes are ambiguous (ids are UUIDv7: memories written within ~65s share a prefix, #1357).", "default": false } } } }) @@ -674,11 +676,12 @@ fn tool_doc_move() -> Value { fn tool_graph() -> Value { serde_json::json!({ "name": "uteke_graph", - "description": "Get knowledge graph data (nodes + edges + stats) for visualization.", + "description": "Get the memory graph: nodes (live memories), edges (auto-links, supersession, explicit edges), stats, and the most recent edges.", "inputSchema": { "type": "object", "properties": { - "namespace": { "type": "string", "description": "Filter by namespace (optional)" } + "namespace": { "type": "string", "description": "Filter by namespace (optional)" }, + "limit": { "type": "integer", "description": "Max edges (default 500, max 5000)" } } } }) @@ -754,7 +757,8 @@ fn tool_room_memories() -> Value { "properties": { "room_id": { "type": "string", "description": "Room identifier" }, "author": { "type": "string", "description": "Optional author filter" }, - "limit": { "type": "integer", "description": "Max results (default 100)", "default": 100 } + "limit": { "type": "integer", "description": "Max results (default 100)", "default": 100 }, + "full_ids": { "type": "boolean", "description": "Print full UUIDs instead of the 8-char prefix. Use when prefixes are ambiguous (ids are UUIDv7: memories written within ~65s share a prefix, #1357).", "default": false } }, "required": ["room_id"] } @@ -1138,6 +1142,7 @@ fn exec_recall(uteke: &Uteke, args: &Value) -> Result { let query = args["query"].as_str().ok_or("Missing 'query'")?; let limit = args["limit"].as_u64().unwrap_or(5) as usize; let namespace = args["namespace"].as_str(); + let full_ids = args["full_ids"].as_bool().unwrap_or(false); let tags_filter: Option> = args["tags"] .as_array() @@ -1292,7 +1297,7 @@ fn exec_recall(uteke: &Uteke, args: &Value) -> Result { uteke_core::SearchResultType::Memory => r .memory_id .as_ref() - .map(|id| format!(" (id: {})", &id[..id.len().min(8)])) + .map(|id| format!(" (id: {})", display_id(id, full_ids))) .unwrap_or_default(), uteke_core::SearchResultType::Document => r .doc_slug @@ -1339,6 +1344,7 @@ fn exec_list(uteke: &Uteke, args: &Value) -> Result { let limit = args["limit"].as_u64().unwrap_or(20) as usize; let offset = args["offset"].as_u64().unwrap_or(0) as usize; let namespace = args["namespace"].as_str(); + let full_ids = args["full_ids"].as_bool().unwrap_or(false); let memories = uteke .list(tag, limit, offset, namespace) @@ -1357,8 +1363,8 @@ fn exec_list(uteke: &Uteke, args: &Value) -> Result { let lines: Vec = memories .iter() .map(|m| { - let short_id = m.id.get(..8).unwrap_or(&m.id); - format!("[{short_id}] {} ({})", m.content, m.tags.join(", ")) + let id = display_id(&m.id, full_ids); + format!("[{id}] {} ({})", m.content, m.tags.join(", ")) }) .collect(); @@ -1371,14 +1377,30 @@ fn exec_list(uteke: &Uteke, args: &Value) -> Result { }) } +/// Render a memory id for tool output: the 8-char prefix, or the full UUID +/// when `full_ids` is set. IDs are UUIDv7, so memories written within the same +/// ~65s share their first 8 chars and the prefix alone can be ambiguous (#1357). +fn display_id(id: &str, full: bool) -> &str { + if full { id } else { id.get(..8).unwrap_or(id) } +} + /// Resolve an id argument to a full UUID (#1048). /// /// Accepts the full UUID or any unambiguous prefix (e.g. the 8-char ids /// printed by recall/list). Errors loudly on ambiguous prefixes instead of /// silently no-oping. Exact UUIDs skip the prefix scan. fn resolve_id<'a>(uteke: &'a Uteke, id: &'a str) -> Result { - if id.len() == 36 { - return Ok(id.to_string()); + // Ids are UUIDs: hex digits and dashes only. Rejecting anything else up + // front also keeps LIKE wildcards (`%`, `_`) out of the prefix scan (#1328). + if id.is_empty() || !id.chars().all(|c| c.is_ascii_hexdigit() || c == '-') { + return Err(format!( + "Invalid memory id '{id}' (expected a UUID or hex prefix)" + )); + } + // A well-formed full UUID skips the prefix scan; anything else (including a + // 36-char non-UUID string) takes the prefix path so it errors the same way. + if let Ok(uuid) = uuid::Uuid::parse_str(id) { + return Ok(uuid.hyphenated().to_string()); } match uteke.resolve_id_prefix(id) { Ok(Some(full)) => Ok(full), @@ -1843,17 +1865,44 @@ fn exec_doc_move(uteke: &Uteke, args: &Value) -> Result { fn exec_graph(uteke: &Uteke, args: &Value) -> Result { let namespace = args["namespace"].as_str(); + let limit = args["limit"] + .as_u64() + .map(|v| v as usize) + .unwrap_or(uteke_core::graph_view::GRAPH_VIEW_DEFAULT_LIMIT); - let data = uteke - .graph_data(namespace) + let view = uteke + .graph_view(namespace, None, limit) .map_err(|e| format!("Failed: {e}"))?; - let text = format!( - "Graph: {} nodes, {} edges, {} relation types", - data.nodes.len(), - data.edges.len(), - data.stats.relation_types.len() + let mut text = format!( + "Graph: {} nodes, {} edges, {} relation types{}", + view.stats.node_count, + view.stats.edge_count, + view.stats.relation_types.len(), + if view.truncated { + " (truncated; raise 'limit')" + } else { + "" + } ); + // Show the most recent edges so the agent can actually use the graph. + let label: std::collections::HashMap<&str, &str> = view + .nodes + .iter() + .map(|n| (n.id.as_str(), n.label.as_str())) + .collect(); + for e in view.edges.iter().take(15) { + text.push_str(&format!( + "\n- {} -[{}]-> {}", + label + .get(e.source_id.as_str()) + .unwrap_or(&e.source_id.as_str()), + e.relation, + label + .get(e.target_id.as_str()) + .unwrap_or(&e.target_id.as_str()) + )); + } Ok(ToolResult { content: vec![McpContent::Text { @@ -1912,7 +1961,7 @@ fn exec_graph_add_edge(uteke: &Uteke, args: &Value) -> Result Result Result let room_id = args["room_id"].as_str().ok_or("Missing 'room_id'")?; let author = args["author"].as_str(); let limit = args["limit"].as_u64().unwrap_or(100) as usize; + let full_ids = args["full_ids"].as_bool().unwrap_or(false); let memories = uteke .recall_room(room_id, author, limit) @@ -2186,8 +2236,8 @@ fn exec_room_memories(uteke: &Uteke, args: &Value) -> Result // #1052/#1048: include the short id so the next tool call // (pin/forget/graph edges) can act on the row directly. let created = m.created_at.format("%Y-%m-%d %H:%M"); - let short_id: String = m.id.chars().take(8).collect(); - format!("[{created} | {} | {}] {}", short_id, m.namespace, m.content) + let id = display_id(&m.id, full_ids); + format!("[{created} | {} | {}] {}", id, m.namespace, m.content) }) .collect(); Ok(ToolResult { @@ -2875,6 +2925,60 @@ mod id_resolution_tests { std::fs::remove_dir_all(&dir).ok(); } + #[test] + fn resolve_id_rejects_malformed_ids() { + let (uteke, dir) = scratch(); + // 36 chars but not a UUID: must not be passed through as a full id. + let fake = "z".repeat(36); + assert!(resolve_id(&uteke, &fake).unwrap_err().contains("Invalid")); + // LIKE wildcards never reach the prefix scan. + assert!(resolve_id(&uteke, "%").is_err()); + assert!(resolve_id(&uteke, "ab_").is_err()); + assert!(resolve_id(&uteke, "").is_err()); + // 36 hex/dash chars that are not a UUID shape fall to the prefix path. + let almost = "0".repeat(36); + assert!( + resolve_id(&uteke, &almost) + .unwrap_err() + .contains("No memory matches") + ); + drop(uteke); + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn display_id_prefix_or_full() { + let id = "01a115a7-aaaa-7bbb-8ccc-0123456789ab"; + assert_eq!(display_id(id, false), "01a115a7"); + assert_eq!(display_id(id, true), id); + assert_eq!(display_id("abc", false), "abc"); + } + + #[test] + fn list_and_room_memories_full_ids_round_trip() { + let (uteke, dir) = scratch(); + let id = seed(&uteke); + let short = &id[..8]; + + let text = |r: ToolResult| match &r.content[0] { + McpContent::Text { text, .. } => text.clone(), + #[allow(unreachable_patterns)] + _ => String::new(), + }; + + let default = text(exec_list(&uteke, &serde_json::json!({})).unwrap()); + assert!(default.contains(&format!("[{short}]")), "{default}"); + assert!(!default.contains(&id)); + + let full = text(exec_list(&uteke, &serde_json::json!({"full_ids": true})).unwrap()); + assert!(full.contains(&format!("[{id}]")), "{full}"); + + // The full id printed by the tool is accepted by resolve_id as-is. + assert_eq!(resolve_id(&uteke, &id).unwrap(), id); + drop(uteke); + std::fs::remove_dir_all(&dir).ok(); + } + #[test] fn exec_pin_accepts_short_id() { let (uteke, dir) = scratch(); diff --git a/crates/uteke-server/Cargo.toml b/crates/uteke-server/Cargo.toml index 553a67b8..8aa8450d 100644 --- a/crates/uteke-server/Cargo.toml +++ b/crates/uteke-server/Cargo.toml @@ -27,8 +27,8 @@ name = "uteke-serve" path = "src/main.rs" [dependencies] -uteke-core = { path = "../uteke-core", version = "0.19.1", default-features = false, features = ["onnx"] } -uteke-mcp = { path = "../uteke-mcp", version = "0.19.1", default-features = false } +uteke-core = { path = "../uteke-core", version = "0.20.0", default-features = false, features = ["onnx"] } +uteke-mcp = { path = "../uteke-mcp", version = "0.20.0", default-features = false } serde = { version = "1", features = ["derive"] } serde_json = "1" schemars = { version = "1", optional = true } diff --git a/crates/uteke-server/src/api_registry.rs b/crates/uteke-server/src/api_registry.rs index 41f6c654..cae7955c 100644 --- a/crates/uteke-server/src/api_registry.rs +++ b/crates/uteke-server/src/api_registry.rs @@ -587,7 +587,7 @@ pub const ENDPOINTS: &[Endpoint] = &[ method: "GET", path: "/graph", tier: Tier::Lab, - description: "Get graph edges for a memory. Accepts `?id=...` query param.", + description: "Memory graph for visualisation: live memories as nodes, `memory_edges` (auto-links, supersession) plus explicit graph edges as edges, keyed by memory id. Accepts `?namespace=...`, `?node_id=` (that memory's neighbourhood) and `?limit=` (edges, default 500, max 5000); response carries `truncated`. Derived `referenced_by` backlinks are omitted.", request_type: None, response_type: None, excludes_deprecated: false, @@ -617,7 +617,7 @@ pub const ENDPOINTS: &[Endpoint] = &[ method: "GET", path: "/edges", tier: Tier::Lab, - description: "List edges for a memory (alias for /graph). Accepts `?id=...` query param.", + description: "List all `memory_edges` of one memory (outgoing and incoming, including derived backlinks). Not an alias of `/graph`. Accepts `?id=...` query param.", request_type: None, response_type: None, excludes_deprecated: false, diff --git a/crates/uteke-server/src/context.rs b/crates/uteke-server/src/context.rs index c9d4c512..bfbad8a8 100644 --- a/crates/uteke-server/src/context.rs +++ b/crates/uteke-server/src/context.rs @@ -128,36 +128,113 @@ pub struct ReqCtx { pub auth_token_hash: Option<[u8; 32]>, /// Hashed read-only token (#409). Read-only requests can use this. pub read_only_token_hash: Option<[u8; 32]>, - /// Allowed CORS origins from config. Empty = wildcard. + /// Allowed CORS origins from config. Empty = no CORS headers; `"*"` = wildcard. pub cors_origins: Vec, /// Recall threshold config from [recall] section in uteke.toml. pub recall_config: Option, /// Extraction config from [extraction] section in uteke.toml. pub extraction_config: Option, + /// `Host` header allowlist against DNS rebinding (#1326). + pub host_guard: HostGuard, +} + +/// `Host` header policy against DNS rebinding (#1326). +/// +/// A rebinding page makes the browser resolve an attacker-controlled name to +/// 127.0.0.1, so the request reaches the local server (even with CORS off) but +/// still carries `Host: `. When the server is bound to a +/// loopback address only loopback `Host` values (plus `allowed_hosts`) are +/// accepted. A non-loopback bind (e.g. Docker, where the service name is the +/// `Host`) keeps accepting everything unless `allowed_hosts` is set. +#[derive(Clone, Debug, Default)] +pub struct HostGuard { + enforce: bool, + allowed: Vec, +} + +impl HostGuard { + /// Build the policy for a bind address and configured extra hosts. + pub fn new(bind_host: &str, allowed_hosts: &[String]) -> Self { + let allowed: Vec = allowed_hosts + .iter() + .map(|h| h.trim().to_ascii_lowercase()) + .filter(|h| !h.is_empty()) + .collect(); + Self { + enforce: is_loopback_host(bind_host) || !allowed.is_empty(), + allowed, + } + } + + /// Whether a request's `Host` header value is acceptable. A missing header + /// (HTTP/1.0 clients) is allowed: rebinding browsers always send one. + pub fn allows(&self, host_header: Option<&str>) -> bool { + if !self.enforce { + return true; + } + let Some(raw) = host_header else { + return true; + }; + let raw = raw.trim().to_ascii_lowercase(); + let name = host_without_port(&raw); + is_loopback_host(name) || self.allowed.iter().any(|a| *a == raw || a == name) + } +} + +/// Strip an optional `:port` (and IPv6 brackets) from a lowercase host value. +fn host_without_port(host: &str) -> &str { + if let Some(rest) = host.strip_prefix('[') { + return rest.split(']').next().unwrap_or(rest); + } + match host.rsplit_once(':') { + // More than one ':' without brackets is a bare IPv6 literal, not host:port. + Some((name, _)) if !name.contains(':') => name, + _ => host, + } +} + +/// `localhost`, `*.localhost`, `127.0.0.0/8` and `::1` (port and brackets ignored). +pub fn is_loopback_host(host: &str) -> bool { + let host = host.trim().to_ascii_lowercase(); + let name = host_without_port(&host); + name == "localhost" + || name.ends_with(".localhost") + || name + .parse::() + .is_ok_and(|ip| ip.is_loopback()) } impl ReqCtx { /// Resolve the allowed origin for a specific request by matching /// its `Origin` header against the configured origins list. - /// Returns "*" if no origins configured (backward compatible). - /// Returns the matching origin if found, or "*" as fallback. + /// + /// Secure by default: with no origins configured NO CORS headers are sent, + /// so browsers block cross-origin reads/writes (a hostile web page must not + /// be able to talk to a local, possibly unauthenticated, server). Include + /// `"*"` in `cors_origins` to explicitly opt in to wildcard CORS. + /// Returns an empty string when CORS headers must be omitted. pub fn resolve_origin_for(&self, req: &Request) -> String { - if self.cors_origins.is_empty() { + if self.cors_origins.iter().any(|o| o == "*") { return "*".to_string(); } - // Check if request has an Origin header if let Some(origin_header) = req.headers().iter().find(|h| h.field.equiv("Origin")) { let origin = origin_header.value.as_str(); if self.cors_origins.iter().any(|o| o == origin) { return origin.to_string(); } } - // No matching origin — return empty string so CORS headers are omitted. - // Browser will block cross-origin requests from untrusted origins. - // Non-browser clients (API users) are unaffected by CORS. + // No matching origin — omit CORS headers. Non-browser clients + // (curl, agents) are unaffected by CORS. String::new() } + /// True when wildcard CORS is combined with any configured token: the + /// `Authorization` header must then not be allowed cross-origin. + fn wildcard_cors_with_auth(&self) -> bool { + (self.auth_token_hash.is_some() || self.read_only_token_hash.is_some()) + && self.cors_origins.iter().any(|o| o == "*") + } + pub fn cors_headers_for(&self, req: &Request) -> Vec
{ let origin = self.resolve_origin_for(req); if origin.is_empty() { @@ -167,15 +244,18 @@ impl ReqCtx { } // When auth is enabled but CORS is wildcard, omit Authorization // from allowed headers to prevent browser-origin auth abuse. - let allowed_headers = if self.auth_token_hash.is_some() && self.cors_origins.is_empty() { + let allowed_headers = if self.wildcard_cors_with_auth() { "Content-Type" } else { "Content-Type, Authorization" }; vec![ Header::from_bytes("Access-Control-Allow-Origin", origin).unwrap(), - Header::from_bytes("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS") - .unwrap(), + Header::from_bytes( + "Access-Control-Allow-Methods", + "GET, POST, PUT, DELETE, OPTIONS", + ) + .unwrap(), Header::from_bytes("Access-Control-Allow-Headers", allowed_headers).unwrap(), ] } @@ -190,17 +270,16 @@ impl ReqCtx { } // Fixed allowlist of headers we accept in cross-origin requests // When auth is enabled but CORS is wildcard, restrict to prevent browser abuse - let allowed_headers_set: &[&str] = - if self.auth_token_hash.is_some() && self.cors_origins.is_empty() { - &["Content-Type", "Accept", "X-Requested-With"] - } else { - &[ - "Content-Type", - "Authorization", - "Accept", - "X-Requested-With", - ] - }; + let allowed_headers_set: &[&str] = if self.wildcard_cors_with_auth() { + &["Content-Type", "Accept", "X-Requested-With"] + } else { + &[ + "Content-Type", + "Authorization", + "Accept", + "X-Requested-With", + ] + }; let allow_headers = req .headers() .iter() @@ -224,8 +303,11 @@ impl ReqCtx { // Access-Control-Allow-Headers — it will block the request as intended. vec![ Header::from_bytes("Access-Control-Allow-Origin", origin).unwrap(), - Header::from_bytes("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS") - .unwrap(), + Header::from_bytes( + "Access-Control-Allow-Methods", + "GET, POST, PUT, DELETE, OPTIONS", + ) + .unwrap(), Header::from_bytes("Access-Control-Allow-Headers", allow_headers).unwrap(), ] } @@ -269,3 +351,123 @@ impl ReqCtx { ) } } + +#[cfg(test)] +mod cors_tests { + use super::*; + use tiny_http::{Method, TestRequest}; + + fn ctx(origins: &[&str], auth: bool, read_only: bool) -> ReqCtx { + ReqCtx { + auth_token_hash: auth.then(|| Sha256::digest("t").into()), + read_only_token_hash: read_only.then(|| Sha256::digest("r").into()), + cors_origins: origins.iter().map(|s| s.to_string()).collect(), + recall_config: None, + extraction_config: None, + host_guard: HostGuard::default(), + } + } + + fn request(origin: &str) -> Request { + TestRequest::new() + .with_method(Method::Get) + .with_path("/health") + .with_header(Header::from_bytes("Origin", origin).unwrap()) + .into() + } + + #[test] + fn no_origins_configured_sends_no_cors_headers() { + let c = ctx(&[], false, false); + assert!( + c.cors_headers_for(&request("https://evil.example")) + .is_empty() + ); + } + + #[test] + fn explicit_origin_is_allowed_others_are_not() { + let c = ctx(&["https://app.example"], false, false); + assert!( + !c.cors_headers_for(&request("https://app.example")) + .is_empty() + ); + assert!( + c.cors_headers_for(&request("https://evil.example")) + .is_empty() + ); + } + + #[test] + fn wildcard_requires_explicit_opt_in_and_hides_authorization_with_any_token() { + let c = ctx(&["*"], false, true); // only a read-only token configured + let headers = c.cors_headers_for(&request("https://x.example")); + let allow = headers + .iter() + .find(|h| h.field.equiv("Access-Control-Allow-Headers")) + .unwrap(); + assert_eq!(allow.value.as_str(), "Content-Type"); + } +} + +#[cfg(test)] +mod host_guard_tests { + use super::*; + + fn loopback() -> HostGuard { + HostGuard::new("127.0.0.1", &[]) + } + + #[test] + fn loopback_bind_accepts_only_loopback_hosts() { + let g = loopback(); + for ok in [ + "localhost", + "localhost:8767", + "127.0.0.1:8767", + "127.0.0.1", + "[::1]:8767", + "::1", + "app.localhost:3000", + "LOCALHOST:8767", + ] { + assert!(g.allows(Some(ok)), "{ok} must be allowed"); + } + for bad in [ + "evil.example", + "evil.example:8767", + "127.0.0.1.evil.example", + "192.168.1.5:8767", + "localhost.evil.example", + ] { + assert!(!g.allows(Some(bad)), "{bad} must be rejected"); + } + } + + #[test] + fn missing_host_header_is_allowed() { + assert!(loopback().allows(None)); + } + + #[test] + fn allowed_hosts_extend_a_loopback_bind() { + let g = HostGuard::new("127.0.0.1", &["Uteke.Internal".to_string()]); + assert!(g.allows(Some("uteke.internal:8767"))); + assert!(g.allows(Some("localhost"))); + assert!(!g.allows(Some("evil.example"))); + } + + #[test] + fn docker_style_bind_without_allowlist_accepts_any_host() { + let g = HostGuard::new("0.0.0.0", &[]); + assert!(g.allows(Some("uteke:8767"))); + assert!(g.allows(Some("evil.example"))); + } + + #[test] + fn docker_style_bind_with_allowlist_enforces_it() { + let g = HostGuard::new("0.0.0.0", &["uteke".to_string()]); + assert!(g.allows(Some("uteke:8767"))); + assert!(!g.allows(Some("evil.example"))); + } +} diff --git a/crates/uteke-server/src/handlers.rs b/crates/uteke-server/src/handlers.rs index 580f06fa..f8558649 100644 --- a/crates/uteke-server/src/handlers.rs +++ b/crates/uteke-server/src/handlers.rs @@ -37,6 +37,17 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< // Handler functions still receive the full `path` with query params intact. let route_path = path.split('?').next().unwrap_or(&path); + // DNS-rebinding guard (#1326): runs before everything else, preflight and + // health included, since a rebinding page can hit any route. + let host_header = req + .headers() + .iter() + .find(|h| h.field.equiv("Host")) + .map(|h| h.value.as_str().to_string()); + if !ctx.host_guard.allows(host_header.as_deref()) { + return ctx.error_response_for(req, 403, "Host not allowed"); + } + // CORS preflight — no auth required if method == Method::Options { return Response::new( @@ -108,12 +119,10 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< // Lock the Uteke instance for the duration of this request. // This serializes requests but prevents data races on the SQLite connection. // Future: use rwlock for read-heavy workloads. - let uteke = match uteke.lock() { - Ok(u) => u, - Err(e) => { - return ctx.error_response_for(req, 500, format!("Internal error: {e}").as_str()); - } - }; + // A handler panic poisons the mutex; recover the guard instead of failing + // every later request. SQLite transactions roll back on unwind, so the + // underlying state stays consistent. + let uteke = uteke.lock().unwrap_or_else(|e| e.into_inner()); match (method, route_path) { // ── Health ────────────────────────────────────────────────────── @@ -708,7 +717,8 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< // ── List ──────────────────────────────────────────────────────── (Method::Post, "/list") => match read_body::(req.as_reader()) { - Ok(req_data) => { + Ok(mut req_data) => { + req_data.limit = req_data.limit.min(MAX_LIST_LIMIT); // Time-travel mode: parse --at and use list_at_time let list_result = match req_data.at.as_deref() { Some(at_str) => match chrono::DateTime::parse_from_rfc3339(at_str) { @@ -922,7 +932,8 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< let limit: u32 = params .get("limit") .and_then(|s| s.parse().ok()) - .unwrap_or(100); + .unwrap_or(100) + .min(MAX_LIST_LIMIT as u32); match uteke.store().list_deprecated(ns_param, limit) { Ok(items) => { #[derive(serde::Serialize)] @@ -1089,7 +1100,8 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< let query = path.split('?').nth(1).unwrap_or(""); let limit = parse_query_param(query, "limit") .and_then(|v| v.parse::().ok()) - .unwrap_or(20); + .unwrap_or(20) + .min(MAX_LIST_LIMIT); let offset = parse_query_param(query, "offset") .and_then(|v| v.parse::().ok()) .unwrap_or(0); @@ -1105,7 +1117,14 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< // ── Graph Visualization (#408) ─────────────────────────────────── (Method::Get, "/graph") => { let ns = parse_query_namespace(&path); - match uteke.graph_data(ns.as_deref()) { + let query = path.split('?').nth(1).unwrap_or(""); + // `node_id` = memory id whose neighbourhood to return (viewers); + // `limit` bounds the number of edges (default 500, max 5000). + let node_id = parse_query_param(query, "node_id").filter(|s| !s.is_empty()); + let limit = parse_query_param(query, "limit") + .and_then(|v| v.parse::().ok()) + .unwrap_or(uteke_core::graph_view::GRAPH_VIEW_DEFAULT_LIMIT); + match uteke.graph_view(ns.as_deref(), node_id.as_deref(), limit) { Ok(data) => ctx.ok_response_for(req, &data), Err(e) => { error!("Graph data error: {e}"); @@ -2131,18 +2150,28 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< let content_length = req .headers() .iter() - .find(|h| h.field.as_str() == "content-length") + .find(|h| h.field.equiv("Content-Length")) .and_then(|h| h.value.as_str().parse::().ok()) .unwrap_or(0); if content_length > MAX_MCP_BODY { return ctx.error_response_for(req, 413, "Payload too large"); } + // Read one byte past the cap so a chunked/unsized oversized body is + // a 413 instead of a silently truncated JSON parse error (#1328). let mut body = String::new(); - if let Err(e) = req.as_reader().take(MAX_MCP_BODY).read_to_string(&mut body) { + if let Err(e) = req + .as_reader() + .take(MAX_MCP_BODY + 1) + .read_to_string(&mut body) + { return ctx.error_response_for(req, 400, format!("Failed to read body: {e}")); } + if body.len() as u64 > MAX_MCP_BODY { + return ctx.error_response_for(req, 413, "Payload too large"); + } + let cors = ctx.cors_headers_for(req); // None = notification (no response per JSON-RPC 2.0 §4.1) → 204 No Content - match uteke_mcp::handle_jsonrpc(&uteke, &body) { + let mut resp = match uteke_mcp::handle_jsonrpc(&uteke, &body) { Some(response) => tiny_http::Response::from_string(response) .with_header( tiny_http::Header::from_bytes( @@ -2167,7 +2196,11 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< ) .unwrap(), ), + }; + for h in cors { + resp.add_header(h); } + resp } // ── Document: Create / Upsert ──────────────────────────────────── @@ -2256,7 +2289,8 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< // ── Document: List ───────────────────────────────────────────── (Method::Post, "/doc/list") => match read_body::(req.as_reader()) { - Ok(params) => { + Ok(mut params) => { + params.limit = params.limit.min(MAX_LIST_LIMIT); let result = if params.roots_only { uteke.doc_list_roots(params.namespace.as_deref(), params.limit) } else if let Some(ref parent) = params.parent { @@ -2512,7 +2546,8 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< let ns = parse_query_namespace(&path); let limit = parse_query_param(query, "limit") .and_then(|v| v.parse::().ok()) - .unwrap_or(50); + .unwrap_or(50) + .min(MAX_LIST_LIMIT); match uteke.contradiction_resolutions(ns.as_deref(), limit) { Ok(resolutions) => ctx.ok_response_for(req, &resolutions), Err(e) => { @@ -2683,8 +2718,15 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< // ── Import (JSONL) ────────────────────────────────────────────── (Method::Post, "/import") => match read_body::(req.as_reader()) { Ok(req_data) => { - if validate_content_size(&req_data.content, 5_242_880).is_err() { - return ctx.error_response_for(req, 413, "Content too large (max 5MB)"); + // Same ceiling `read_body` enforces on the whole request, so the + // limit the caller is told about is the limit that applies (#1328). + let max_import = uteke_core::MAX_PAYLOAD_SIZE; + if validate_content_size(&req_data.content, max_import).is_err() { + return ctx.error_response_for( + req, + 413, + format!("Content too large (max {}MB)", max_import / 1_048_576), + ); } // Merge request tags into the JSONL entries. @@ -2742,12 +2784,18 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< // (pitfall 0f lock contention). (Method::Post, "/verify") => match uteke.verify() { Ok(r) => ctx.ok_response_for(req, &r), - Err(e) => ctx.error_response_for(req, 500, e.to_string()), + Err(e) => { + error!("Request failed: {e}"); + ctx.error_response_for(req, 500, "Internal server error") + } }, (Method::Post, "/repair") => match uteke.repair() { Ok(r) => ctx.ok_response_for(req, &r), - Err(e) => ctx.error_response_for(req, 500, e.to_string()), + Err(e) => { + error!("Request failed: {e}"); + ctx.error_response_for(req, 500, "Internal server error") + } }, (Method::Post, "/prune") => match read_body::(req.as_reader()) { @@ -2756,7 +2804,10 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< uteke.prune(req_data.ttl_days, ns(&req_data.namespace), req_data.dry_run); match result { Ok(r) => ctx.ok_response_for(req, &r), - Err(e) => ctx.error_response_for(req, 500, e.to_string()), + Err(e) => { + error!("Request failed: {e}"); + ctx.error_response_for(req, 500, "Internal server error") + } } } Err(e) => ctx.error_response_for(req, 400, e), @@ -2769,14 +2820,20 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< let pairs = uteke.find_duplicates(ns(&req_data.namespace), req_data.threshold); match pairs { Ok(p) => ctx.ok_response_for(req, &p), - Err(e) => ctx.error_response_for(req, 500, e.to_string()), + Err(e) => { + error!("Request failed: {e}"); + ctx.error_response_for(req, 500, "Internal server error") + } } } else { let result = uteke.consolidate(ns(&req_data.namespace), req_data.threshold, false); match result { Ok(r) => ctx.ok_response_for(req, &r), - Err(e) => ctx.error_response_for(req, 500, e.to_string()), + Err(e) => { + error!("Request failed: {e}"); + ctx.error_response_for(req, 500, "Internal server error") + } } } } @@ -2845,7 +2902,10 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< }; match result { Ok(r) => ctx.ok_response_for(req, &r), - Err(e) => ctx.error_response_for(req, 500, e.to_string()), + Err(e) => { + error!("Request failed: {e}"); + ctx.error_response_for(req, 500, "Internal server error") + } } } Err(e) => ctx.error_response_for(req, 400, e), @@ -2855,7 +2915,10 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< (Method::Post, "/importance") => match read_body::(req.as_reader()) { Ok(_req_data) => match uteke.recompute_importance() { Ok(count) => ctx.ok_response_for(req, &serde_json::json!({ "updated": count })), - Err(e) => ctx.error_response_for(req, 500, e.to_string()), + Err(e) => { + error!("Request failed: {e}"); + ctx.error_response_for(req, 500, "Internal server error") + } }, Err(e) => ctx.error_response_for(req, 400, e), }, @@ -2869,7 +2932,10 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< req_data.limit, ) { Ok(orphans) => ctx.ok_response_for(req, &orphans), - Err(e) => ctx.error_response_for(req, 500, e.to_string()), + Err(e) => { + error!("Request failed: {e}"); + ctx.error_response_for(req, 500, "Internal server error") + } } } Err(e) => ctx.error_response_for(req, 400, e), @@ -2882,7 +2948,10 @@ pub fn route(uteke: &Mutex, ctx: &ReqCtx, req: &mut Request) -> Response< Ok(count) => { ctx.ok_response_for(req, &serde_json::json!({ "backlinks_created": count })) } - Err(e) => ctx.error_response_for(req, 500, e.to_string()), + Err(e) => { + error!("Request failed: {e}"); + ctx.error_response_for(req, 500, "Internal server error") + } }, Err(e) => ctx.error_response_for(req, 400, e), } @@ -2916,7 +2985,9 @@ fn resolve_extraction_config( api_key: req_api_key.map(String::from).unwrap_or(base.api_key), base_url: base.base_url, endpoint_path: base.endpoint_path, - max_facts: req_max_facts.unwrap_or(base.max_facts), + max_facts: req_max_facts + .unwrap_or(base.max_facts) + .min(MAX_EXTRACT_FACTS), } } @@ -3018,6 +3089,7 @@ mod room_memories_namespace_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -3249,6 +3321,7 @@ mod room_recall_at_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -3301,34 +3374,73 @@ mod room_recall_at_tests { let (status, resp) = app.add_edge(&id1, &id2); assert_eq!(status, 200, "{resp}"); - let src_node = resp["source_node"].as_str().expect("source_node id"); - let tgt_node = resp["target_node"].as_str().expect("target_node id"); + // POST still reports the graph-node ids it created (#1180)... + assert!(resp["source_node"].as_str().is_some()); + assert!(resp["target_node"].as_str().is_some()); - // The resolved nodes must exist in GET /graph and carry the memory - // link, so visualization can map edges back to memories. + // ...but GET /graph keys nodes by MEMORY id (#1366) so a viewer can + // open them, and keeps `memory_id` for clients that map back. let (_, graph) = app.call(Method::Get, "/graph", None); let nodes = graph["nodes"].as_array().expect("nodes array"); - assert!( - nodes.iter().any(|n| n["id"] == serde_json::json!(src_node) - && n["memory_id"] == serde_json::json!(id1)), - "source node must link back to memory {id1}: {graph}" - ); - assert!( - nodes.iter().any(|n| n["id"] == serde_json::json!(tgt_node) - && n["memory_id"] == serde_json::json!(id2)), - "target node must link back to memory {id2}: {graph}" - ); + for (id, text) in [(&id1, "alpha memory"), (&id2, "beta memory")] { + assert!( + nodes.iter().any(|n| n["id"] == serde_json::json!(id) + && n["memory_id"] == serde_json::json!(id) + && n["label"].as_str().is_some_and(|l| l.starts_with(text))), + "node for memory {id} missing: {graph}" + ); + } - // The edge itself must be visible in the graph payload. + // The edge is visible with both long and short endpoint names. let edges = graph["edges"].as_array().expect("edges array"); assert!( edges .iter() - .any(|e| e["source_id"] == serde_json::json!(src_node) - && e["target_id"] == serde_json::json!(tgt_node) + .any(|e| e["source_id"] == serde_json::json!(id1) + && e["target_id"] == serde_json::json!(id2) + && e["source"] == serde_json::json!(id1) + && e["target"] == serde_json::json!(id2) && e["relation"] == serde_json::json!("related")), "edge must appear in GET /graph: {graph}" ); + assert_eq!(graph["truncated"], serde_json::json!(false)); + } + + #[test] + fn graph_node_id_and_limit_query_params() { + let app = GraphEdgeApp::new(); + let a = app.remember("hub memory"); + let b = app.remember("spoke one"); + let c = app.remember("spoke two"); + let d = app.remember("elsewhere one"); + let e = app.remember("elsewhere two"); + for (s, t) in [(&a, &b), (&a, &c), (&d, &e)] { + let (status, resp) = app.add_edge(s, t); + assert_eq!(status, 200, "{resp}"); + } + + let hood_url = format!("/graph?node_id={a}"); + let (_, hood) = app.call(Method::Get, &hood_url, None); + let touching: Vec<&serde_json::Value> = hood["edges"] + .as_array() + .unwrap() + .iter() + .filter(|e| e["relation"] == serde_json::json!("related")) + .collect(); + assert!( + touching + .iter() + .all(|e| e["source"] == serde_json::json!(a) || e["target"] == serde_json::json!(a)), + "node_id must restrict to that memory's edges: {hood}" + ); + assert!(touching.len() >= 2, "{hood}"); + + // URL built separately so the api_registry route scanner does not + // mistake this test call for a handler route arm. + let capped_url = "/graph?limit=1"; + let (_, capped) = app.call(Method::Get, capped_url, None); + assert_eq!(capped["edges"].as_array().unwrap().len(), 1, "{capped}"); + assert_eq!(capped["truncated"], serde_json::json!(true)); } #[test] @@ -3452,6 +3564,7 @@ mod room_recall_at_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -3737,6 +3850,7 @@ mod contradiction_api_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -3876,6 +3990,7 @@ mod pack_recall_api_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -4036,6 +4151,7 @@ mod explain_recall_api_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -4162,6 +4278,7 @@ mod list_pagination_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -4265,6 +4382,7 @@ mod payload_conformance_tests { cors_origins: vec![], recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let remember_body: &'static str = Box::leak(r#"{"content":"Payload conformance probe memory #1233 with distinctive tokens zebraquartz","namespace":"conf"}"#.to_string().into_boxed_str()); @@ -4342,6 +4460,7 @@ mod payload_conformance_tests { cors_origins: vec![], recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let mut req = tiny_http::TestRequest::new() .with_method(tiny_http::Method::Post) @@ -4398,6 +4517,7 @@ mod payload_conformance_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -4505,6 +4625,7 @@ mod routes_introspection_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -4612,6 +4733,7 @@ mod plain_remember_type_tests { cors_origins: Vec::new(), recall_config: None, extraction_config: None, + host_guard: Default::default(), }; let resp = route(&self.uteke, &ctx, &mut req); let status = resp.status_code().0; @@ -4718,3 +4840,111 @@ mod plain_remember_type_tests { ); } } + +#[cfg(test)] +mod host_guard_route_tests { + use super::*; + use crate::context::HostGuard; + use tiny_http::{Header, TestRequest}; + + fn status(guard: HostGuard, host: Option<&str>, path: &str) -> u16 { + let uteke = Mutex::new( + Uteke::open_with_backend(":memory:", None) + .expect("open in-memory uteke without embedder"), + ); + let ctx = ReqCtx { + auth_token_hash: None, + read_only_token_hash: None, + cors_origins: Vec::new(), + recall_config: None, + extraction_config: None, + host_guard: guard, + }; + let mut t = TestRequest::new().with_method(Method::Get).with_path(path); + if let Some(h) = host { + t = t.with_header(Header::from_bytes(&b"Host"[..], h.as_bytes()).unwrap()); + } + let mut req = t.into(); + route(&uteke, &ctx, &mut req).status_code().0 + } + + // #1326 acceptance: a loopback bind rejects `Host: evil.example`. + #[test] + fn loopback_bind_rejects_foreign_host() { + let g = HostGuard::new("127.0.0.1", &[]); + assert_eq!(status(g.clone(), Some("evil.example:8767"), "/health"), 403); + assert_eq!(status(g.clone(), Some("evil.example"), "/stats"), 403); + assert_eq!(status(g.clone(), Some("127.0.0.1:8767"), "/health"), 200); + assert_eq!(status(g, Some("localhost:8767"), "/health"), 200); + } + + // #1326 acceptance: a Docker-style bind works, and allowed_hosts is honoured. + #[test] + fn docker_style_bind_and_allowed_hosts() { + let open = HostGuard::new("0.0.0.0", &[]); + assert_eq!(status(open, Some("uteke:8767"), "/health"), 200); + + let listed = HostGuard::new("0.0.0.0", &["uteke".to_string()]); + assert_eq!(status(listed.clone(), Some("uteke:8767"), "/health"), 200); + assert_eq!(status(listed, Some("evil.example"), "/health"), 403); + } +} + +#[cfg(test)] +mod mcp_http_hardening_tests { + use super::*; + use tiny_http::{Header, TestRequest}; + + fn app() -> Mutex { + Mutex::new( + Uteke::open_with_backend(":memory:", None) + .expect("open in-memory uteke without embedder"), + ) + } + + fn ctx(origins: Vec) -> ReqCtx { + ReqCtx { + auth_token_hash: None, + read_only_token_hash: None, + cors_origins: origins, + recall_config: None, + extraction_config: None, + host_guard: Default::default(), + } + } + + // #1328: an oversized /mcp body is a 413, not a truncated JSON parse error. + #[test] + fn mcp_oversized_body_is_413() { + let big: &'static str = Box::leak("x".repeat(1024 * 1024 + 8).into_boxed_str()); + let mut req = TestRequest::new() + .with_method(Method::Post) + .with_path("/mcp") + .with_body(big) + .into(); + let resp = route(&app(), &ctx(Vec::new()), &mut req); + assert_eq!(resp.status_code().0, 413); + } + + // #1328: /mcp responses carry CORS headers like every other endpoint. + #[test] + fn mcp_response_has_cors_headers() { + let origin = "https://app.example"; + let body = r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#; + let mut req = TestRequest::new() + .with_method(Method::Post) + .with_path("/mcp") + .with_header(Header::from_bytes(&b"Origin"[..], origin.as_bytes()).unwrap()) + .with_body(body) + .into(); + let resp = route(&app(), &ctx(vec![origin.to_string()]), &mut req); + assert!( + resp.headers() + .iter() + .any(|h| h.field.equiv("Access-Control-Allow-Origin") + && h.value.as_str() == origin), + "missing ACAO on /mcp: {:?}", + resp.headers() + ); + } +} diff --git a/crates/uteke-server/src/main.rs b/crates/uteke-server/src/main.rs index bdf99537..b1adeef5 100644 --- a/crates/uteke-server/src/main.rs +++ b/crates/uteke-server/src/main.rs @@ -13,8 +13,6 @@ mod types; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; -use std::path::PathBuf; - use sha2::{Digest, Sha256}; use tiny_http::Server; use tracing::{error, info, warn}; @@ -34,6 +32,7 @@ fn main() { let mut cli_auth_token: Option = None; let mut cli_read_only_token: Option = None; let mut cli_cors_origins: Vec = Vec::new(); + let mut cli_allowed_hosts: Vec = Vec::new(); let mut i = 1; while i < args.len() { @@ -77,6 +76,15 @@ fn main() { std::process::exit(1); } } + "--allowed-host" => { + i += 1; + if i < args.len() { + cli_allowed_hosts.push(args[i].clone()); + } else { + eprintln!("Error: --allowed-host requires a value"); + std::process::exit(1); + } + } "--cors-origin" => { i += 1; if i < args.len() { @@ -101,6 +109,9 @@ fn main() { println!(" --port Port number (default: 8767)"); println!(" --auth-token Bearer token for API auth"); println!(" --cors-origin Allowed CORS origin (repeatable)"); + println!( + " --allowed-host Extra accepted Host header value (repeatable, #1326)" + ); println!(" --read-only-token Read-only API token (GET endpoints only) (#409)"); println!(" -V, --version Show version"); println!(" -h, --help Show this help"); @@ -208,6 +219,17 @@ fn main() { config_cors_origins }; + // Merge allowed Host values: CLI flags override config (#1326) + let allowed_hosts = if !cli_allowed_hosts.is_empty() { + cli_allowed_hosts + } else { + config + .server + .as_ref() + .and_then(|s| s.allowed_hosts.clone()) + .unwrap_or_default() + }; + let host = cli_host.unwrap_or(config_host); let port = cli_port.unwrap_or(config_port); @@ -325,12 +347,15 @@ fn main() { let read_only_token_hash = read_only_token.as_deref().map(|t| Sha256::digest(t).into()); // Build request context - // Warn if auth is configured but CORS origins are not — this is safe for - // non-browser clients (curl, SDKs, agents) but risky if browser access is needed. - if auth_token_hash.is_some() && cors_origins.is_empty() { - warn!("Security: auth token is set but cors_origins is not configured."); - warn!(" For browser access, set cors_origins in uteke.toml or --cors-origin."); - warn!(" Non-browser clients (curl, agents) are unaffected by CORS."); + // CORS is off unless origins are configured. Wildcard is an explicit opt-in + // and is dangerous without auth: any web page could read/write memories. + if cors_origins.iter().any(|o| o == "*") + && auth_token_hash.is_none() + && read_only_token_hash.is_none() + { + warn!("Security: cors_origins contains \"*\" and authentication is disabled —"); + warn!(" any website opened in a local browser can read and modify memories."); + warn!(" Set an auth token or list explicit origins."); } let ctx = context::ReqCtx { auth_token_hash, @@ -359,6 +384,7 @@ fn main() { recall }, extraction_config: config.extraction.clone(), + host_guard: context::HostGuard::new(&host, &allowed_hosts), }; // Start server @@ -509,6 +535,28 @@ fn main() { } }); + // Vector-index flusher (#1322): per-operation index saves are batched, so + // a single write followed by idle time would otherwise sit in memory until + // the next write or a clean shutdown (SIGTERM does not run the handler). + let flush_uteke = Arc::clone(&uteke); + std::thread::spawn(move || { + loop { + std::thread::sleep(std::time::Duration::from_secs(1)); + if SHUTDOWN.load(Ordering::SeqCst) { + return; + } + if let Ok(u) = flush_uteke.lock() { + if let Err(e) = u.flush_index() { + warn!("Periodic index flush failed: {e}"); + } + // vecq only: rebuild once dead rows pass 25% of the index (#1324). + if let Err(e) = u.compact_index_if_needed(false) { + warn!("Index compaction failed: {e}"); + } + } + } + }); + // SIGINT handler ctrlc::set_handler(|| { if SHUTDOWN.load(Ordering::SeqCst) { @@ -551,7 +599,9 @@ fn main() { let method = req.method().clone(); let url = req.url().to_string(); - info!("{method} {url}"); + // Log the path only: query strings can carry memory text / search + // phrases (PII, secrets) that must not land in server logs. + info!("{method} {}", url.split('?').next().unwrap_or(&url)); let uteke = Arc::clone(&uteke); let ctx = ctx.clone(); @@ -564,15 +614,23 @@ fn main() { } let result = std::thread::Builder::new().spawn(move || { - let response = handlers::route(&uteke, &ctx, &mut req); + // RAII: the slot is released even if the handler panics, so a + // single bad request can never leak capacity. + let _slot = SlotGuard(pair); + let routed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + handlers::route(&uteke, &ctx, &mut req) + })); + let response = match routed { + Ok(r) => r, + Err(_) => { + error!("Request handler panicked: {method} {url}"); + tiny_http::Response::from_data(b"Internal server error".to_vec()) + .with_status_code(500) + } + }; if let Err(e) = req.respond(response) { warn!("Response error: {e}"); } - // Release slot and notify the waiting accept loop. - let (lock, cvar) = &*pair; - let mut active = lock.lock().unwrap(); - *active -= 1; - cvar.notify_one(); }); if let Err(e) = result { @@ -675,31 +733,130 @@ struct ServerFileSection { /// Set to specific origins like ["http://localhost:3000"] for production. /// Each request's `Origin` header is matched against this list. cors_origins: Option>, + /// Extra accepted `Host` header values (#1326). On a loopback bind only + /// loopback hosts are accepted by default (DNS-rebinding guard); list + /// additional names here. On a non-loopback bind (Docker) every `Host` + /// is accepted unless this list is set. + allowed_hosts: Option>, } -/// Find and parse the nearest uteke.toml, looking at: -/// 1. $UTEKE_HOME/uteke.toml (or ~/.codecora/uteke/uteke.toml) -/// 2. $CWD/.uteke/uteke.toml +/// Releases one concurrency slot (and wakes the accept loop) on drop. +struct SlotGuard(Arc<(std::sync::Mutex, std::sync::Condvar)>); + +impl Drop for SlotGuard { + fn drop(&mut self) { + let (lock, cvar) = &*self.0; + let mut active = lock.lock().unwrap_or_else(|e| e.into_inner()); + *active = active.saturating_sub(1); + cvar.notify_one(); + } +} + +/// Resolve `uteke.toml` from: +/// 1. `$UTEKE_HOME/uteke.toml` (or `~/.codecora/uteke/uteke.toml`) +/// 2. `$CWD/.uteke/uteke.toml` (project-local, untrusted for `[server]`, +/// endpoints and credentials unless `UTEKE_TRUST_PROJECT_CONFIG=1`) +/// +/// Merge precedence, the untrusted-key policy and per-layer validation are +/// shared with the CLI in `uteke_core::config_layers`. fn load_uteke_toml() -> ServerFileConfig { - let mut config = ServerFileConfig::default(); - - let mut paths: Vec = vec![match uteke_core::uteke_home() { - Ok(h) => h.join("uteke.toml"), - Err(_) => PathBuf::new(), - }]; - if let Ok(cwd) = std::env::current_dir() { - paths.push(cwd.join(".uteke").join("uteke.toml")); + let global = uteke_core::uteke_home().ok().map(|h| h.join("uteke.toml")); + let project = std::env::current_dir() + .ok() + .map(|cwd| cwd.join(".uteke").join("uteke.toml")); + resolve_server_config(&uteke_core::config_layers::Layers { + global: global.as_deref(), + project: project.as_deref(), + trust_project: matches!( + std::env::var("UTEKE_TRUST_PROJECT_CONFIG").as_deref(), + Ok("1") | Ok("true") + ), + }) +} + +fn resolve_server_config(layers: &uteke_core::config_layers::Layers<'_>) -> ServerFileConfig { + uteke_core::config_layers::resolve::(layers).value +} + +#[cfg(test)] +mod config_overlay_tests { + use super::*; + use std::path::PathBuf; + + fn temp_file(tag: &str, body: &str) -> PathBuf { + // Unique per call: tests run in parallel inside one process. + static NEXT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + let n = NEXT.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let path = std::env::temp_dir().join(format!( + "uteke_server_cfg_{}_{}_{}.toml", + std::process::id(), + tag, + n + )); + std::fs::write(&path, body).unwrap(); + path } - for path in paths { - if path.exists() { - if let Ok(content) = std::fs::read_to_string(&path) { - if let Ok(parsed) = toml::from_str::(&content) { - config = parsed; - } - } - } + fn resolve(global: &str, project: &str, trust: bool) -> ServerFileConfig { + let g = temp_file("g", global); + let p = temp_file("p", project); + let cfg = resolve_server_config(&uteke_core::config_layers::Layers { + global: Some(&g), + project: Some(&p), + trust_project: trust, + }); + let _ = std::fs::remove_file(g); + let _ = std::fs::remove_file(p); + cfg } - config + #[test] + fn project_config_keeps_global_server_auth() { + // Dummy values generated at runtime (not real secrets). + let line = |v: &str| format!("{} = \"{}\"", "auth_token", v); + let global_line = line(&"g".repeat(10)); + let attacker_line = line(&"a".repeat(10)); + let global = format!( + "[server]\nhost = \"0.0.0.0\"\n{global_line}\ncors_origins = [\"https://app.example\"]\n" + ); + // Project file only tunes recall but also tries to override [server]. + let project = + format!("[recall]\nmin_score = 0.5\n\n[server]\n{attacker_line}\nhost = \"0.0.0.0\"\n"); + let merged = resolve(&global, &project, false); + let server = merged.server.expect("global [server] must survive"); + assert_eq!(server.auth_token.as_deref(), Some("g".repeat(10).as_str())); + assert_eq!( + server.cors_origins.as_deref(), + Some(&["https://app.example".to_string()][..]) + ); + assert!( + merged.recall.is_some(), + "project recall tuning still applies" + ); + } + + #[test] + fn project_config_without_server_section_does_not_wipe_global() { + let dummy = "t".repeat(10); + let global = format!("[server]\n{} = \"{dummy}\"\n", "auth_token"); + let merged = resolve(&global, "[recall]\nmin_score = 0.3\n", false); + assert_eq!( + merged.server.and_then(|s| s.auth_token).as_deref(), + Some(dummy.as_str()) + ); + } + + #[test] + fn project_recall_keys_merge_per_key_not_per_section() { + // The old server overlay replaced whole sections: a project file setting + // one [recall] key wiped the other global [recall] keys. + let merged = resolve( + "[recall]\nmin_score = 0.3\ndefault_strategy = \"vector\"\n", + "[recall]\nmin_score = 0.6\n", + false, + ); + let recall = merged.recall.expect("recall section"); + assert_eq!(recall.min_score, Some(0.6)); + assert_eq!(recall.default_strategy.as_deref(), Some("vector")); + } } diff --git a/crates/uteke-server/src/types.rs b/crates/uteke-server/src/types.rs index 7f2f26fb..ea474be3 100644 --- a/crates/uteke-server/src/types.rs +++ b/crates/uteke-server/src/types.rs @@ -480,6 +480,11 @@ pub struct RecallFileSection { /// Hard cap on recall/list `limit` to prevent DoS via unbounded queries (#903). pub const MAX_LIMIT: usize = 100; +/// Hard cap on `limit` for list-style endpoints (`/list`, `/recent`, +/// `/contradictions`, `/doc/list`, deprecated listing). Callers page with `offset`. +pub const MAX_LIST_LIMIT: usize = 1000; +/// Hard cap on request-controlled `max_facts` for `/extract` (bounds LLM cost). +pub const MAX_EXTRACT_FACTS: usize = 100; /// Default strict mode threshold for server recall. /// Used as fallback when [recall] min_score_strict is not configured. pub const STRICT_THRESHOLD: f32 = 0.5; diff --git a/docs/api-reference.md b/docs/api-reference.md index 7d5fd96e..fbeb1019 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -110,7 +110,7 @@ Rebuild the vector index from stored embeddings to fix recall desync (#1266). SQ #### 🟢 `GET` `/graph` -Get graph edges for a memory. Accepts `?id=...` query param. +Memory graph for visualisation: live memories as nodes, `memory_edges` (auto-links, supersession) plus explicit graph edges as edges, keyed by memory id. Accepts `?namespace=...`, `?node_id=` (that memory's neighbourhood) and `?limit=` (edges, default 500, max 5000); response carries `truncated`. Derived `referenced_by` backlinks are omitted. #### 🟢 `GET` `/provenance` @@ -303,7 +303,7 @@ Remove an edge between two nodes. Accepts memory IDs or graph node IDs via `?sou #### 🟢 `GET` `/edges` -List edges for a memory (alias for /graph). Accepts `?id=...` query param. +List all `memory_edges` of one memory (outgoing and incoming, including derived backlinks). Not an alias of `/graph`. Accepts `?id=...` query param. #### 🟢 `GET` `/timeline` diff --git a/docs/configuration.md b/docs/configuration.md index dcc5ec94..4ce7a4e9 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -355,14 +355,20 @@ namespace = "my-project" [logging] level = "warn" - -[server] -enabled = true -port = 8767 ``` Combined with shell hooks, this enables automatic project-scoped memory — each project gets its own isolated memory store. +### Trust: what a project file may not set + +A project directory can be an untrusted clone, so `.uteke/uteke.toml` is **not** allowed to decide where credentials or memory text are sent or how the server authenticates. These keys are ignored when they come from the project file (a warning is logged) and must be set in the global `uteke.toml` instead: + +- `[embedding]`: `backend`, `api_key`, `base_url`, `endpoint_path` +- `[embed_fallback]` and `[extraction]`: `api_key`, `base_url`, `endpoint_path` +- the whole `[server]` section (`enabled`, `host`, `port`, `auth_token`, `read_only_token`, `cors_origins`, `allowed_hosts`) + +Everything else (store path, namespace, logging, recall tuning, lifecycle, ...) can be set per project. To trust a project file completely, set `UTEKE_TRUST_PROJECT_CONFIG=1`. The CLI and `uteke-serve` apply the same rules and the same per-key merge. + ## CLI Flag Override CLI flags always take precedence over config file values: diff --git a/docs/core-contract.json b/docs/core-contract.json index faa4bc34..58ee31bc 100644 --- a/docs/core-contract.json +++ b/docs/core-contract.json @@ -319,5 +319,5 @@ } ], "source_of_truth": "uteke PROD store, document tree kontrak-core-cloud", - "version": "0.19.1" + "version": "0.20.0" } diff --git a/docs/tls.md b/docs/tls.md index 380d24c3..e780f2d7 100644 --- a/docs/tls.md +++ b/docs/tls.md @@ -71,8 +71,9 @@ server { |------|--------|---------| | TLS encryption | ✅ Via reverse proxy | Caddy (auto) or Nginx (manual Let's Encrypt) | | Authentication | ✅ Built-in | `--auth-token` flag or `uteke.toml [server] auth_token` | -| CORS | ✅ Built-in | `--cors-origin` or `uteke.toml [server] cors_origins` | +| CORS | ✅ Built-in (off by default; `"*"` opts in to wildcard) | `--cors-origin` or `uteke.toml [server] cors_origins` | | Bind address | ✅ Configurable | `--host 127.0.0.1` for local-only access | +| Host header (DNS rebinding) | ✅ Built-in | Loopback bind accepts only loopback `Host` values; add names with `--allowed-host` or `[server] allowed_hosts` | | Rate limiting | ⚠️ Via reverse proxy | Configure in Caddy/Nginx | | File permissions | ✅ Built-in | Owner-only (0700/0600) on data directories | @@ -92,6 +93,20 @@ uteke-serve --host 127.0.0.1 --auth-token YOUR_TOKEN uteke-serve --host 0.0.0.0 --auth-token YOUR_TOKEN --cors-origin https://uteke.yourdomain.com ``` +### Host Header Validation (DNS Rebinding) + +On a loopback bind (`127.0.0.1`, `::1`, `localhost`) `uteke-serve` rejects requests whose `Host` header is not `localhost`, `*.localhost` or a loopback IP, with `403 Host not allowed`. This blocks DNS-rebinding pages that resolve an attacker-controlled name to your machine. + +If a reverse proxy forwards the public hostname to a loopback-bound server, list that name: + +```toml +[server] +host = "127.0.0.1" +allowed_hosts = ["uteke.yourdomain.com"] +``` + +or `uteke-serve --allowed-host uteke.yourdomain.com` (repeatable; the CLI flag replaces the config list). On a non-loopback bind (e.g. `0.0.0.0` in Docker, where the service name is the `Host`) every `Host` is accepted unless `allowed_hosts` is set. Entries match the host with or without the port, case-insensitively. A request with no `Host` header is accepted. + ### Environment Variables | Variable | Purpose |