From 19794c430d3130113fa924c8fb391a9d1c9d72ee Mon Sep 17 00:00:00 2001 From: jaylfc Date: Sun, 19 Jul 2026 23:57:32 +0100 Subject: [PATCH 01/16] feat(config): collections.allowed_roots setting (env override, default off) The allowed-roots list gates the whole collections feature: a collection source_path must resolve inside one of these directories or create/index are refused. Default is empty (feature off) so nothing changes for existing installs. TAOSMD_COLLECTIONS_ALLOWED_ROOTS overrides the config file, matching the other server settings. --- taosmd/config.py | 72 +++++++++++++++++++++++++ tests/test_config_collections_roots.py | 75 ++++++++++++++++++++++++++ 2 files changed, 147 insertions(+) create mode 100644 tests/test_config_collections_roots.py diff --git a/taosmd/config.py b/taosmd/config.py index 2262880b..f7abac73 100644 --- a/taosmd/config.py +++ b/taosmd/config.py @@ -50,6 +50,11 @@ _GENERATOR_PROFILE_KEY = "generator_profile" # Whether A2A registry auth runs in enforce mode (True) or verify-and-warn mode (False). _A2A_AUTH_ENFORCE_KEY = "a2a_auth_enforce" +# Section under which collections settings live. ``allowed_roots`` is the +# safety line of the collections contract: source paths must resolve inside +# one of these directories. Empty (the default) means collections are off. +_COLLECTIONS_SECTION_KEY = "collections" +_COLLECTIONS_ALLOWED_ROOTS_KEY = "allowed_roots" MANAGED_BY_STANDALONE = "standalone" MANAGED_BY_TAOS = "taos" @@ -597,6 +602,71 @@ def set_a2a_auth_enforce(value: bool, data_dir=None) -> None: _write(data, data_dir) +# --------------------------------------------------------------------------- +# Collections: allowed roots +# --------------------------------------------------------------------------- + +def get_collections_allowed_roots(data_dir=None) -> list[str]: + """Return the configured collections allowed-roots list (default empty). + + The allowed-roots list is the single most important safety line in the + collections contract: a collection's ``source_path`` must resolve inside + one of these directories or collection creation and indexing are refused. + The default is an EMPTY list, which means the collections feature is + effectively off until an operator opts in. + + Resolution order (first non-empty wins): + + 1. ``TAOSMD_COLLECTIONS_ALLOWED_ROOTS`` env var (``os.pathsep``-separated) + 2. ``collections.allowed_roots`` list in ``~/.taosmd/config.json`` + 3. ``[]`` (collections off) + + Blank entries are dropped; a corrupt config section reads as empty. + """ + env = os.environ.get("TAOSMD_COLLECTIONS_ALLOWED_ROOTS") + if env is not None and env.strip(): + return [p.strip() for p in env.split(os.pathsep) if p.strip()] + section = _read(data_dir).get(_COLLECTIONS_SECTION_KEY) + if not isinstance(section, dict): + return [] + roots = section.get(_COLLECTIONS_ALLOWED_ROOTS_KEY) + if not isinstance(roots, list): + return [] + return [r.strip() for r in roots if isinstance(r, str) and r.strip()] + + +def set_collections_allowed_roots(roots, clear: bool = False, data_dir=None) -> None: + """Persist the collections allowed-roots list (or clear it). + + Args: + roots: List of directory paths. Blank entries are dropped. Ignored + when ``clear`` is True. + clear: when True, remove the setting (collections revert to off). + + Raises: + ValueError: when ``clear`` is False and ``roots`` is not a list of + strings. + """ + data = _read(data_dir) + section = data.get(_COLLECTIONS_SECTION_KEY) + if not isinstance(section, dict): + section = {} + if clear: + section.pop(_COLLECTIONS_ALLOWED_ROOTS_KEY, None) + else: + if not isinstance(roots, list): + raise ValueError("roots must be a list of strings (or pass clear=True)") + cleaned = [] + for r in roots: + if not isinstance(r, str): + raise ValueError(f"allowed root must be a string, got {type(r).__name__}") + if r.strip(): + cleaned.append(r.strip()) + section[_COLLECTIONS_ALLOWED_ROOTS_KEY] = cleaned + data[_COLLECTIONS_SECTION_KEY] = section + _write(data, data_dir) + + __all__ = [ "get_memory_model", "set_memory_model", @@ -623,4 +693,6 @@ def set_a2a_auth_enforce(value: bool, data_dir=None) -> None: "MANAGED_BY_TAOS", "get_generator_profile", "set_generator_profile", + "get_collections_allowed_roots", + "set_collections_allowed_roots", ] diff --git a/tests/test_config_collections_roots.py b/tests/test_config_collections_roots.py new file mode 100644 index 00000000..2ab89358 --- /dev/null +++ b/tests/test_config_collections_roots.py @@ -0,0 +1,75 @@ +"""Tests for the ``collections.allowed_roots`` config key. + +The allowed-roots list is the single most important safety line in the +collections contract: ``source_path`` must resolve inside one of these +directories or collection creation and indexing are refused. The default is +an EMPTY list, which means the collections feature is effectively off until +an operator opts in. + +Resolution mirrors the other server settings: the +``TAOSMD_COLLECTIONS_ALLOWED_ROOTS`` env var (``os.pathsep``-separated) wins +over the ``collections.allowed_roots`` config-file key. +""" +from __future__ import annotations + +import os + +import pytest + +from taosmd import config + + +@pytest.fixture +def data_dir(tmp_path, monkeypatch): + monkeypatch.delenv("TAOSMD_COLLECTIONS_ALLOWED_ROOTS", raising=False) + return str(tmp_path) + + +def test_unset_is_empty_list(data_dir): + assert config.get_collections_allowed_roots(data_dir) == [] + + +def test_set_then_get_round_trip(data_dir): + config.set_collections_allowed_roots(["/srv/docs", "/srv/repos"], data_dir=data_dir) + assert config.get_collections_allowed_roots(data_dir) == ["/srv/docs", "/srv/repos"] + + +def test_env_overrides_config_file(data_dir, monkeypatch): + config.set_collections_allowed_roots(["/from-file"], data_dir=data_dir) + monkeypatch.setenv( + "TAOSMD_COLLECTIONS_ALLOWED_ROOTS", os.pathsep.join(["/a", "/b"]) + ) + assert config.get_collections_allowed_roots(data_dir) == ["/a", "/b"] + + +def test_clear_returns_empty(data_dir): + config.set_collections_allowed_roots(["/srv/docs"], data_dir=data_dir) + config.set_collections_allowed_roots([], clear=True, data_dir=data_dir) + assert config.get_collections_allowed_roots(data_dir) == [] + + +def test_set_rejects_non_list(data_dir): + with pytest.raises(ValueError): + config.set_collections_allowed_roots("/srv/docs", data_dir=data_dir) + + +def test_set_rejects_non_string_entries(data_dir): + with pytest.raises(ValueError): + config.set_collections_allowed_roots(["/ok", 42], data_dir=data_dir) + + +def test_blank_entries_are_dropped(data_dir): + config.set_collections_allowed_roots(["/srv/docs", "", " "], data_dir=data_dir) + assert config.get_collections_allowed_roots(data_dir) == ["/srv/docs"] + + +def test_corrupt_section_is_empty(data_dir, tmp_path): + (tmp_path / "config.json").write_text('{"collections": "not-a-dict"}') + assert config.get_collections_allowed_roots(data_dir) == [] + + +def test_independent_of_other_keys(data_dir): + config.set_collections_allowed_roots(["/srv/docs"], data_dir=data_dir) + assert config.get_server_token(data_dir) is None + config.set_server_token("tok", data_dir=data_dir) + assert config.get_collections_allowed_roots(data_dir) == ["/srv/docs"] From 2205c57a4a86650a1df7d86ce193a6546c239f6c Mon Sep 17 00:00:00 2001 From: jaylfc Date: Sun, 19 Jul 2026 23:59:13 +0100 Subject: [PATCH 02/16] feat(collections): storage layer (rows, typed links, grants, archive) First-class collection rows in collections.db with the lifecycle created|indexing|ready|error|archived, plus three side tables: typed project links {taos|git, id} (metadata only, no transitivity), grants (canonical_id, scope='collection', collection_id) UNIQUE together, and per-file content hashes for incremental re-index. source_path is validated against collections.allowed_roots via resolve_within at create time. Delete is an archive alias: reversible status change, nothing destroyed. Collection ids (col- + 12 hex) match the agent-name grammar so content rows can live under the existing per-agent scoping. --- taosmd/collections.py | 401 ++++++++++++++++++++++++++++++++ tests/test_collections_store.py | 224 ++++++++++++++++++ 2 files changed, 625 insertions(+) create mode 100644 taosmd/collections.py create mode 100644 tests/test_collections_store.py diff --git a/taosmd/collections.py b/taosmd/collections.py new file mode 100644 index 00000000..daa7c048 --- /dev/null +++ b/taosmd/collections.py @@ -0,0 +1,401 @@ +"""Collections: named, typed containers for content indexed from a folder. + +A collection is a first-class row (not a metadata tag): it names one source +folder, tracks its indexing lifecycle (created -> indexing -> ready | error), +and composes with the existing scoping surfaces through two side tables: + +- **links** attach a collection to one or more projects for discovery. A link + row is typed ``{type: "taos" | "git", id}`` because the taOS project id + (``prj-xxx``) and the taOSmd git-fingerprint project id (12 hex) are + different namespaces; a collection for a repo that taOS also manages + typically carries one of each. Links are metadata only and never grant + query access (no transitivity). +- **grants** give a named agent query access to a collection. A grant row is + ``(canonical_id, scope='collection', collection_id)``, unique together, + and is enforced at search time: collection hits are only returned to + agents holding a grant. + +Safety: a collection's ``source_path`` must resolve inside one of the +directories in the ``collections.allowed_roots`` config list (default EMPTY, +which turns the feature off), checked at create time and again at every +index. Symlink escapes are rejected by ``loaders._safety.resolve_within``. + +Zero-loss: DELETE is an alias for archive (``status='archived'``, +reversible); re-indexing supersedes replaced rows via the existing +``valid_to`` machinery and never deletes; only the wipe surface destroys. + +Content rows live in the normal vector/archive stores under the agent name +```` (collection ids match the agent-name grammar by +construction), so the existing per-agent search scoping doubles as the +collection scoping mechanism with no new query machinery. +""" + +from __future__ import annotations + +import fnmatch +import hashlib +import json +import logging +import os +import secrets +import sqlite3 +import time +from pathlib import Path + +from . import config as _config +from .loaders import check_size, resolve_within +from .loaders.registry import REGISTRY as _LOADER_REGISTRY, _path_to_extension + +logger = logging.getLogger(__name__) + +KINDS = ("docs", "codebase", "mixed") +STATUSES = ("created", "indexing", "ready", "error", "archived") +LINK_TYPES = ("taos", "git") +GRANT_SCOPE = "collection" + +#: Per-file size cap for the folder walker. Deliberately far below the +#: loaders' generous 100 MB default: a docs collection should never contain +#: a file this large, and the cap stops a stray artifact from bloating the +#: index. Overridable per ingest_folder call. +DEFAULT_MAX_FILE_BYTES = 10 * 1024 * 1024 + +#: Directory names never descended into, regardless of gitignore rules. +_SKIP_DIRS = frozenset({ + ".git", ".hg", ".svn", "node_modules", "__pycache__", ".venv", "venv", + ".tox", ".mypy_cache", ".ruff_cache", ".pytest_cache", +}) + +#: Extensions treated as binary and skipped without opening the file. +_BINARY_EXTS = frozenset({ + "png", "jpg", "jpeg", "gif", "webp", "ico", "bmp", "pdf", "zip", "gz", + "bz2", "xz", "tar", "7z", "whl", "so", "dylib", "dll", "exe", "bin", + "onnx", "gguf", "pt", "safetensors", "db", "sqlite", "sqlite3", "woff", + "woff2", "ttf", "eot", "otf", "mp3", "mp4", "mov", "avi", "wav", "flac", + "pyc", "pyo", "class", "jar", "o", "a", "ds_store", +}) + + +class CollectionNotFoundError(KeyError): + """Raised when referencing a collection id that does not exist.""" + + +def _new_collection_id() -> str: + """``col-`` + 12 lowercase hex. Matches the agent-name grammar + (``^[a-z][a-z0-9_-]{0,62}$``) so the id can double as the agent name + under which the collection's content rows are stored.""" + return f"col-{secrets.token_hex(6)}" + + +class CollectionStore: + """SQLite-backed store for collection rows, links, grants, and the + per-file hash state that backs incremental re-index. + + Synchronous by design (like :class:`taosmd.agents.AgentRegistry`); all + callers route through the single service loop so the thread-affine + connection is used from exactly one thread. + """ + + def __init__(self, data_dir) -> None: + self._data_dir = os.fspath(data_dir) + path = Path(self._data_dir) + path.mkdir(parents=True, exist_ok=True) + self._conn = sqlite3.connect(str(path / "collections.db")) + self._conn.row_factory = sqlite3.Row + self._init_schema() + + def _init_schema(self) -> None: + self._conn.executescript( + """ + CREATE TABLE IF NOT EXISTS collections ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + kind TEXT NOT NULL, + source_path TEXT NOT NULL, + embedder TEXT, + status TEXT NOT NULL, + created_at REAL NOT NULL, + last_indexed REAL, + stats_json TEXT NOT NULL DEFAULT '{}', + error TEXT + ); + CREATE TABLE IF NOT EXISTS collection_links ( + collection_id TEXT NOT NULL, + type TEXT NOT NULL, + ext_id TEXT NOT NULL, + created_at REAL NOT NULL, + UNIQUE(collection_id, type, ext_id) + ); + CREATE TABLE IF NOT EXISTS collection_grants ( + canonical_id TEXT NOT NULL, + scope TEXT NOT NULL DEFAULT 'collection', + collection_id TEXT NOT NULL, + created_at REAL NOT NULL, + UNIQUE(canonical_id, scope, collection_id) + ); + CREATE TABLE IF NOT EXISTS collection_files ( + collection_id TEXT NOT NULL, + file_path TEXT NOT NULL, + content_hash TEXT NOT NULL, + updated_at REAL NOT NULL, + UNIQUE(collection_id, file_path) + ); + """ + ) + self._conn.commit() + + def close(self) -> None: + self._conn.close() + + # ----- validation ------------------------------------------------------ + + def resolve_source_path(self, source_path: str) -> Path: + """Validate ``source_path`` against the allowed roots; return it resolved. + + Raises ``ValueError`` when no allowed roots are configured (feature + off), when the path escapes every root (including via symlink), or + when it is not an existing directory. Called at create time AND again + at every index, so a root removed from config retroactively disables + indexing of collections created under it. + """ + roots = _config.get_collections_allowed_roots(self._data_dir) + if not roots: + raise ValueError( + "collections are disabled: no collections.allowed_roots configured " + "(set the config key or TAOSMD_COLLECTIONS_ALLOWED_ROOTS)" + ) + resolved = None + for root in roots: + try: + resolved = resolve_within(source_path, root) + break + except ValueError: + continue + if resolved is None: + raise ValueError( + f"source_path {source_path!r} is outside every configured allowed root" + ) + if not resolved.is_dir(): + raise ValueError(f"source_path {source_path!r} is not a directory") + return resolved + + # ----- CRUD ------------------------------------------------------------ + + def create( + self, + *, + name: str, + kind: str, + source_path: str, + embedder: str | None = None, + ) -> dict: + if not isinstance(name, str) or not name.strip(): + raise ValueError("name must be a non-empty string") + if kind not in KINDS: + raise ValueError(f"kind must be one of {'|'.join(KINDS)}, got {kind!r}") + if embedder is not None and (not isinstance(embedder, str) or not embedder.strip()): + raise ValueError("embedder must be a non-empty string when provided") + resolved = self.resolve_source_path(source_path) + cid = _new_collection_id() + self._conn.execute( + "INSERT INTO collections (id, name, kind, source_path, embedder, status, " + "created_at, stats_json) VALUES (?, ?, ?, ?, ?, 'created', ?, '{}')", + (cid, name.strip(), kind, str(resolved), + embedder.strip() if embedder else None, time.time()), + ) + self._conn.commit() + return self.get(cid) + + def _row(self, collection_id: str) -> sqlite3.Row: + row = self._conn.execute( + "SELECT * FROM collections WHERE id = ?", (collection_id,) + ).fetchone() + if row is None: + raise CollectionNotFoundError(f"collection {collection_id!r} not found") + return row + + def get(self, collection_id: str) -> dict: + row = self._row(collection_id) + try: + stats = json.loads(row["stats_json"]) + except (json.JSONDecodeError, TypeError): + stats = {} + links = [ + {"type": r["type"], "id": r["ext_id"]} + for r in self._conn.execute( + "SELECT type, ext_id FROM collection_links WHERE collection_id = ? " + "ORDER BY created_at", + (collection_id,), + ) + ] + grants = [ + r["canonical_id"] + for r in self._conn.execute( + "SELECT canonical_id FROM collection_grants " + "WHERE collection_id = ? AND scope = ? ORDER BY created_at", + (collection_id, GRANT_SCOPE), + ) + ] + return { + "id": row["id"], + "name": row["name"], + "kind": row["kind"], + "source_path": row["source_path"], + "embedder": row["embedder"], + "status": row["status"], + "created_at": row["created_at"], + "last_indexed": row["last_indexed"], + "stats": stats if isinstance(stats, dict) else {}, + "error": row["error"], + "links": links, + "grants": grants, + } + + def list( + self, *, project: str | None = None, include_archived: bool = False + ) -> list[dict]: + """All collections, oldest first. ``project`` filters to collections + holding a link whose ``ext_id`` matches, regardless of link type + (taOS ``prj-*`` ids and git fingerprints are different namespaces, + so a raw id is unambiguous in practice).""" + rows = self._conn.execute( + "SELECT id, status FROM collections ORDER BY created_at" + ).fetchall() + out = [] + for row in rows: + if not include_archived and row["status"] == "archived": + continue + col = self.get(row["id"]) + if project is not None and not any( + link["id"] == project for link in col["links"] + ): + continue + out.append(col) + return out + + def archive(self, collection_id: str) -> dict: + """Archive a collection (the DELETE alias). Reversible: the row and + every content row stay on disk; search-time grant resolution skips + archived collections, which hides the content from query.""" + self._row(collection_id) # 404 check + self._conn.execute( + "UPDATE collections SET status = 'archived' WHERE id = ?", + (collection_id,), + ) + self._conn.commit() + return self.get(collection_id) + + def set_status( + self, + collection_id: str, + status: str, + *, + error: str | None = None, + last_indexed: float | None = None, + ) -> None: + if status not in STATUSES: + raise ValueError(f"status must be one of {'|'.join(STATUSES)}, got {status!r}") + self._row(collection_id) + self._conn.execute( + "UPDATE collections SET status = ?, error = ?, " + "last_indexed = COALESCE(?, last_indexed) WHERE id = ?", + (status, error if status == "error" else None, last_indexed, collection_id), + ) + self._conn.commit() + + def set_stats(self, collection_id: str, stats: dict) -> None: + self._row(collection_id) + self._conn.execute( + "UPDATE collections SET stats_json = ? WHERE id = ?", + (json.dumps(stats), collection_id), + ) + self._conn.commit() + + # ----- links ----------------------------------------------------------- + + def link(self, collection_id: str, link_type: str, ext_id: str) -> dict: + if link_type not in LINK_TYPES: + raise ValueError( + f"link type must be one of {'|'.join(LINK_TYPES)}, got {link_type!r}" + ) + if not isinstance(ext_id, str) or not ext_id.strip(): + raise ValueError("link id must be a non-empty string") + self._row(collection_id) + self._conn.execute( + "INSERT OR IGNORE INTO collection_links " + "(collection_id, type, ext_id, created_at) VALUES (?, ?, ?, ?)", + (collection_id, link_type, ext_id.strip(), time.time()), + ) + self._conn.commit() + return self.get(collection_id) + + def unlink(self, collection_id: str, link_type: str, ext_id: str) -> dict: + self._row(collection_id) + self._conn.execute( + "DELETE FROM collection_links " + "WHERE collection_id = ? AND type = ? AND ext_id = ?", + (collection_id, link_type, ext_id), + ) + self._conn.commit() + return self.get(collection_id) + + # ----- grants ---------------------------------------------------------- + + def grant(self, collection_id: str, canonical_id: str) -> dict: + if not isinstance(canonical_id, str) or not canonical_id.strip(): + raise ValueError("agent (canonical_id) must be a non-empty string") + self._row(collection_id) + self._conn.execute( + "INSERT OR IGNORE INTO collection_grants " + "(canonical_id, scope, collection_id, created_at) VALUES (?, ?, ?, ?)", + (canonical_id.strip(), GRANT_SCOPE, collection_id, time.time()), + ) + self._conn.commit() + return self.get(collection_id) + + def revoke(self, collection_id: str, canonical_id: str) -> dict: + self._row(collection_id) + self._conn.execute( + "DELETE FROM collection_grants " + "WHERE canonical_id = ? AND scope = ? AND collection_id = ?", + (canonical_id, GRANT_SCOPE, collection_id), + ) + self._conn.commit() + return self.get(collection_id) + + def has_grant(self, canonical_id: str, collection_id: str) -> bool: + row = self._conn.execute( + "SELECT 1 FROM collection_grants " + "WHERE canonical_id = ? AND scope = ? AND collection_id = ?", + (canonical_id, GRANT_SCOPE, collection_id), + ).fetchone() + return row is not None + + # ----- per-file hash state (incremental re-index) ---------------------- + + def file_states(self, collection_id: str) -> dict[str, str]: + return { + r["file_path"]: r["content_hash"] + for r in self._conn.execute( + "SELECT file_path, content_hash FROM collection_files " + "WHERE collection_id = ?", + (collection_id,), + ) + } + + def set_file_state(self, collection_id: str, file_path: str, content_hash: str) -> None: + self._conn.execute( + "INSERT INTO collection_files (collection_id, file_path, content_hash, updated_at) " + "VALUES (?, ?, ?, ?) " + "ON CONFLICT(collection_id, file_path) " + "DO UPDATE SET content_hash = excluded.content_hash, " + "updated_at = excluded.updated_at", + (collection_id, file_path, content_hash, time.time()), + ) + self._conn.commit() + + def remove_file_state(self, collection_id: str, file_path: str) -> None: + self._conn.execute( + "DELETE FROM collection_files WHERE collection_id = ? AND file_path = ?", + (collection_id, file_path), + ) + self._conn.commit() diff --git a/tests/test_collections_store.py b/tests/test_collections_store.py new file mode 100644 index 00000000..8a226c24 --- /dev/null +++ b/tests/test_collections_store.py @@ -0,0 +1,224 @@ +"""Tests for the collections store: CRUD, links, grants, archive. + +Collections are first-class rows (not metadata tags): a named, typed +container for content indexed from a folder. The store enforces: + +- ``kind`` in {docs, codebase, mixed} +- ``source_path`` must resolve inside a configured allowed root + (``collections.allowed_roots``); empty roots means creation is refused +- links are typed rows {type: taos|git, id}, unique per collection +- grants are (canonical_id, scope='collection', collection_id) unique rows +- delete is archive (status='archived', reversible), never destruction +""" +from __future__ import annotations + +import pytest + +from taosmd import config +from taosmd.collections import ( + CollectionNotFoundError, + CollectionStore, +) + + +@pytest.fixture +def data_dir(tmp_path, monkeypatch): + monkeypatch.delenv("TAOSMD_COLLECTIONS_ALLOWED_ROOTS", raising=False) + d = tmp_path / "taosmd-data" + d.mkdir() + return str(d) + + +@pytest.fixture +def source_dir(tmp_path): + d = tmp_path / "docs-src" + d.mkdir() + (d / "readme.md").write_text("# Hello\n\nSome docs.") + return str(d) + + +@pytest.fixture +def store(data_dir, source_dir): + config.set_collections_allowed_roots([source_dir], data_dir=data_dir) + return CollectionStore(data_dir) + + +# --------------------------------------------------------------------------- +# create / get / list +# --------------------------------------------------------------------------- + +def test_create_returns_row_with_defaults(store, source_dir): + col = store.create(name="repo docs", kind="docs", source_path=source_dir) + assert col["id"].startswith("col-") + assert len(col["id"]) == len("col-") + 12 + assert col["name"] == "repo docs" + assert col["kind"] == "docs" + assert col["status"] == "created" + assert col["embedder"] is None # default = global embedder + assert col["last_indexed"] is None + assert col["stats"] == {} + assert col["links"] == [] + assert col["grants"] == [] + + +def test_create_rejects_bad_kind(store, source_dir): + with pytest.raises(ValueError): + store.create(name="x", kind="movies", source_path=source_dir) + + +def test_create_requires_allowed_root(data_dir, source_dir): + # No allowed_roots configured: collections are off. + s = CollectionStore(data_dir) + with pytest.raises(ValueError, match="allowed_roots"): + s.create(name="x", kind="docs", source_path=source_dir) + + +def test_create_rejects_path_outside_roots(store, tmp_path): + outside = tmp_path / "elsewhere" + outside.mkdir() + with pytest.raises(ValueError): + store.create(name="x", kind="docs", source_path=str(outside)) + + +def test_create_rejects_missing_dir(store, source_dir): + with pytest.raises(ValueError): + store.create(name="x", kind="docs", source_path=source_dir + "/nope") + + +def test_create_stores_embedder(store, source_dir): + col = store.create( + name="x", kind="docs", source_path=source_dir, embedder="arctic-embed-s" + ) + assert store.get(col["id"])["embedder"] == "arctic-embed-s" + + +def test_get_unknown_raises(store): + with pytest.raises(CollectionNotFoundError): + store.get("col-000000000000") + + +def test_list_and_get_round_trip(store, source_dir): + a = store.create(name="a", kind="docs", source_path=source_dir) + b = store.create(name="b", kind="mixed", source_path=source_dir) + ids = {c["id"] for c in store.list()} + assert ids == {a["id"], b["id"]} + assert store.get(a["id"])["name"] == "a" + + +# --------------------------------------------------------------------------- +# links +# --------------------------------------------------------------------------- + +def test_link_unlink_round_trip(store, source_dir): + col = store.create(name="a", kind="docs", source_path=source_dir) + store.link(col["id"], "taos", "prj-123") + store.link(col["id"], "git", "abc123def456") + links = store.get(col["id"])["links"] + assert {"type": "taos", "id": "prj-123"} in links + assert {"type": "git", "id": "abc123def456"} in links + store.unlink(col["id"], "taos", "prj-123") + links = store.get(col["id"])["links"] + assert {"type": "taos", "id": "prj-123"} not in links + assert {"type": "git", "id": "abc123def456"} in links + + +def test_link_is_idempotent(store, source_dir): + col = store.create(name="a", kind="docs", source_path=source_dir) + store.link(col["id"], "taos", "prj-123") + store.link(col["id"], "taos", "prj-123") + assert len(store.get(col["id"])["links"]) == 1 + + +def test_link_rejects_bad_type(store, source_dir): + col = store.create(name="a", kind="docs", source_path=source_dir) + with pytest.raises(ValueError): + store.link(col["id"], "jira", "PROJ-1") + + +def test_list_project_filter_matches_either_link_type(store, source_dir): + a = store.create(name="a", kind="docs", source_path=source_dir) + b = store.create(name="b", kind="docs", source_path=source_dir) + store.link(a["id"], "taos", "prj-123") + store.link(b["id"], "git", "abc123def456") + assert [c["id"] for c in store.list(project="prj-123")] == [a["id"]] + assert [c["id"] for c in store.list(project="abc123def456")] == [b["id"]] + assert store.list(project="nothing") == [] + + +# --------------------------------------------------------------------------- +# grants +# --------------------------------------------------------------------------- + +def test_grant_revoke_round_trip(store, source_dir): + col = store.create(name="a", kind="docs", source_path=source_dir) + assert not store.has_grant("agent-a", col["id"]) + store.grant(col["id"], "agent-a") + assert store.has_grant("agent-a", col["id"]) + assert "agent-a" in store.get(col["id"])["grants"] + store.revoke(col["id"], "agent-a") + assert not store.has_grant("agent-a", col["id"]) + + +def test_grant_is_idempotent(store, source_dir): + col = store.create(name="a", kind="docs", source_path=source_dir) + store.grant(col["id"], "agent-a") + store.grant(col["id"], "agent-a") + assert store.get(col["id"])["grants"] == ["agent-a"] + + +def test_grant_unknown_collection_raises(store): + with pytest.raises(CollectionNotFoundError): + store.grant("col-000000000000", "agent-a") + + +# --------------------------------------------------------------------------- +# archive (delete alias) + status +# --------------------------------------------------------------------------- + +def test_archive_is_reversible_status_change(store, source_dir): + col = store.create(name="a", kind="docs", source_path=source_dir) + out = store.archive(col["id"]) + assert out["status"] == "archived" + # Row still exists: nothing destroyed. + assert store.get(col["id"])["status"] == "archived" + # Archived collections are hidden from the default listing… + assert store.list() == [] + # …but visible when explicitly asked for. + assert [c["id"] for c in store.list(include_archived=True)] == [col["id"]] + + +def test_set_status_and_stats(store, source_dir): + col = store.create(name="a", kind="docs", source_path=source_dir) + store.set_status(col["id"], "indexing") + assert store.get(col["id"])["status"] == "indexing" + store.set_status(col["id"], "error", error="boom") + got = store.get(col["id"]) + assert got["status"] == "error" + assert got["error"] == "boom" + store.set_stats(col["id"], {"files_indexed": 3}) + store.set_status(col["id"], "ready", last_indexed=123.0) + got = store.get(col["id"]) + assert got["stats"] == {"files_indexed": 3} + assert got["last_indexed"] == 123.0 + assert got["error"] is None # cleared on non-error status + + +def test_set_status_rejects_unknown(store, source_dir): + col = store.create(name="a", kind="docs", source_path=source_dir) + with pytest.raises(ValueError): + store.set_status(col["id"], "exploded") + + +# --------------------------------------------------------------------------- +# file state (incremental re-index support) +# --------------------------------------------------------------------------- + +def test_file_state_round_trip(store, source_dir): + col = store.create(name="a", kind="docs", source_path=source_dir) + assert store.file_states(col["id"]) == {} + store.set_file_state(col["id"], "readme.md", "hash1") + assert store.file_states(col["id"]) == {"readme.md": "hash1"} + store.set_file_state(col["id"], "readme.md", "hash2") + assert store.file_states(col["id"]) == {"readme.md": "hash2"} + store.remove_file_state(col["id"], "readme.md") + assert store.file_states(col["id"]) == {} From ddd8bb861a72010466c403d6284023b9d3f6ad6a Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 00:03:14 +0100 Subject: [PATCH 03/16] feat(collections): gitignore-aware walker, chunker, folder ingest, search scoping collect_files walks a collection source with simplified stdlib gitignore rules (nested files, negation, dir-only, anchored patterns), skips VCS/dependency/hidden dirs, binary files (extension + null-byte sniff), oversized files, symlink escapes (resolve_within), and anything no registered loader explicitly claims - wiring the previously-unwired loader registry into a real ingest path. ingest_folder chunks each doc (zero-dep paragraph packer, upgrade-path comment for structure-aware chunking) and routes chunks through api.ingest_batch under the collection id as agent namespace with per-chunk content-hash ids, so re-index dedups unchanged files for free. Changed and deleted files get their old rows soft-superseded (valid_to + collection-reindex marker), never deleted; the archive keeps every version. Status runs created -> indexing -> ready|error with stats and last_indexed on the row; allowed_roots is re-checked at every index. api.search gains collections/collections_only: granted collections join search_agents (grants enforced per requesting agent, archived or ungranted collections contribute nothing), and collection hits carry collection_id/file_path/source metadata. --- taosmd/api.py | 44 ++++ taosmd/collections.py | 392 +++++++++++++++++++++++++++++++ tests/test_collections_ingest.py | 335 ++++++++++++++++++++++++++ 3 files changed, 771 insertions(+) create mode 100644 tests/test_collections_ingest.py diff --git a/taosmd/api.py b/taosmd/api.py index 265e484d..bcf3e7ac 100644 --- a/taosmd/api.py +++ b/taosmd/api.py @@ -477,6 +477,8 @@ async def search( limit: int = 5, mode: str | None = None, prefer_verified: str | None = None, + collections: list[str] | None = None, + collections_only: bool = False, data_dir=None, ) -> list[dict]: """Search the librarian's shelves for passages relevant to ``query``. @@ -500,6 +502,15 @@ async def search( recipe resolution entirely and returns BM25-only hits (the #25 user-memory contract: keyword search-as-you-type, sub-300ms). Default ``None`` is the full recipe-driven retrieval path. + collections: Optional list of collection ids whose indexed content + should be searched alongside conversation memory. Grants are + enforced per requesting agent: a collection the agent holds no + grant for (or that is archived) contributes no hits. Collection + hits carry ``collection_id``, ``file_path``, and ``source`` in + their metadata. + collections_only: When True (with ``collections``), restrict the + search to the granted collections and exclude conversation + memory. Returns ``[]`` when no requested collection is granted. data_dir: Optional taosmd data dir (see :func:`ingest`). """ if not agent: @@ -535,6 +546,39 @@ async def search( if name != agent: search_agents.append(name) + # Collection scoping: content rows live under the collection id as their + # agent namespace, so granting search access is just extending the + # search_agents list. Grants are enforced here (per requesting agent); + # unknown, archived, or ungranted collections are silently skipped so a + # caller cannot probe for their existence. + if collections: + from taosmd.collections import CollectionNotFoundError, CollectionStore # noqa: PLC0415 + cstore = CollectionStore(_resolve_data_dir(data_dir)) + try: + granted: list[str] = [] + for cid in collections: + if not isinstance(cid, str) or not cid: + continue + try: + col = cstore.get(cid) + except CollectionNotFoundError: + continue + if col["status"] == "archived": + continue + if not cstore.has_grant(agent, cid): + continue + granted.append(cid) + finally: + cstore.close() + if collections_only: + if not granted: + return [] + search_agents = granted + else: + search_agents.extend(granted) + elif collections_only: + return [] + if mode == "bm25": # BM25-only path: no embed call, no recipe/reranker resolution. Hits # come straight off the vector store's BM25 index in the same diff --git a/taosmd/collections.py b/taosmd/collections.py index daa7c048..72538154 100644 --- a/taosmd/collections.py +++ b/taosmd/collections.py @@ -399,3 +399,395 @@ def remove_file_state(self, collection_id: str, file_path: str) -> None: (collection_id, file_path), ) self._conn.commit() + + +# --------------------------------------------------------------------------- +# Folder walker (gitignore-aware, stdlib only) +# --------------------------------------------------------------------------- + +def _parse_gitignore(path: Path) -> list[tuple[str, bool, bool]]: + """Parse one .gitignore into ``(pattern, negate, dir_only)`` rules. + + Simplified gitwildmatch: comments and blanks dropped, ``!`` negation, + trailing ``/`` marks dir-only, leading ``/`` anchors to the .gitignore's + own directory, patterns without ``/`` match the basename at any depth, + patterns with ``/`` are fnmatch-ed against the path relative to the + .gitignore's directory. + # upgrade-path: full gitwildmatch (``**`` semantics, escaped chars) if the + # simplified rules ever mis-walk a real repo. + """ + rules: list[tuple[str, bool, bool]] = [] + try: + lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + except OSError: + return rules + for line in lines: + line = line.rstrip() + if not line or line.lstrip().startswith("#"): + continue + negate = line.startswith("!") + if negate: + line = line[1:] + dir_only = line.endswith("/") + pattern = line.rstrip("/") + if pattern: + rules.append((pattern, negate, dir_only)) + return rules + + +def _match_one(pattern: str, rel: str, name: str) -> bool: + if pattern.startswith("/"): + return fnmatch.fnmatch(rel, pattern[1:]) + if "/" in pattern: + return fnmatch.fnmatch(rel, pattern) + return fnmatch.fnmatch(name, pattern) + + +def _is_ignored( + rel_posix: str, + name: str, + is_dir: bool, + rule_sets: list[tuple[str, list[tuple[str, bool, bool]]]], +) -> bool: + """Apply the collected .gitignore rule sets to one path. + + ``rule_sets`` is ``[(prefix, rules)]`` where ``prefix`` is the rule + file's directory relative to the source root (``""`` for the root). + Later-matching rules win, mirroring git's last-match-wins semantics. + """ + ignored = False + for prefix, rules in rule_sets: + if prefix: + if not rel_posix.startswith(prefix + "/"): + continue + sub = rel_posix[len(prefix) + 1:] + else: + sub = rel_posix + for pattern, negate, dir_only in rules: + if dir_only and not is_dir: + continue + if _match_one(pattern, sub, name): + ignored = not negate + return ignored + + +def _loader_for(path: Path): + """Return an instance of the loader that explicitly claims ``path``, + or ``None`` when no registered loader does. + + This deliberately does NOT use ``pick_loader``'s catch-all fallback: + the walker only ingests files a loader positively claims (DocLoader + md/txt/markdown/rst, ChatLoader *.chat.json, ...), so arbitrary source + files are skipped rather than mangled through the doc path. + """ + import mimetypes # noqa: PLC0415 + + ext = _path_to_extension(path) + mime, _ = mimetypes.guess_type(str(path)) + for loader_cls in _LOADER_REGISTRY: + if loader_cls.can_handle(extension=ext, mime_type=mime or ""): + return loader_cls() + return None + + +def collect_files( + source_root: Path | str, + *, + max_file_bytes: int = DEFAULT_MAX_FILE_BYTES, +) -> tuple[list[tuple[Path, str]], dict]: + """Walk ``source_root`` and return ``([(abs_path, rel_posix)], skips)``. + + Gitignore-aware (root and nested .gitignore files), skips VCS/dependency + directories and hidden directories, binary files (by extension and by a + null-byte sniff), files over ``max_file_bytes``, symlinks that escape the + root, and files no registered loader claims. ``skips`` counts each skip + reason so ingest stats can surface them. + """ + root = Path(source_root).resolve() + skips = { + "skipped_ignored": 0, + "skipped_binary": 0, + "skipped_size": 0, + "skipped_symlink": 0, + "skipped_unclaimed": 0, + } + files: list[tuple[Path, str]] = [] + rule_sets: list[tuple[str, list[tuple[str, bool, bool]]]] = [] + + for dirpath, dirnames, filenames in os.walk(root, followlinks=False): + dpath = Path(dirpath) + rel_dir = "" if dpath == root else dpath.relative_to(root).as_posix() + + gi = dpath / ".gitignore" + if gi.is_file(): + rules = _parse_gitignore(gi) + if rules: + rule_sets.append((rel_dir, rules)) + + keep_dirs = [] + for d in sorted(dirnames): + rel_d = f"{rel_dir}/{d}" if rel_dir else d + if d in _SKIP_DIRS or d.startswith("."): + continue + if _is_ignored(rel_d, d, True, rule_sets): + skips["skipped_ignored"] += 1 + continue + keep_dirs.append(d) + dirnames[:] = keep_dirs + + for fname in sorted(filenames): + fpath = dpath / fname + rel_f = f"{rel_dir}/{fname}" if rel_dir else fname + if _is_ignored(rel_f, fname, False, rule_sets): + skips["skipped_ignored"] += 1 + continue + try: + resolve_within(fpath, root) + except ValueError: + skips["skipped_symlink"] += 1 + continue + ext = fname.rsplit(".", 1)[-1].lower() if "." in fname else "" + if ext in _BINARY_EXTS: + skips["skipped_binary"] += 1 + continue + try: + check_size(fpath, max_file_bytes) + except ValueError: + skips["skipped_size"] += 1 + continue + except OSError: + skips["skipped_symlink"] += 1 + continue + if _loader_for(fpath) is None: + skips["skipped_unclaimed"] += 1 + continue + try: + with open(fpath, "rb") as fh: + head = fh.read(1024) + except OSError: + skips["skipped_symlink"] += 1 + continue + if b"\x00" in head: + skips["skipped_binary"] += 1 + continue + files.append((fpath, rel_f)) + return files, skips + + +# --------------------------------------------------------------------------- +# Chunker (zero-dep) +# --------------------------------------------------------------------------- + +def chunk_text(text: str, max_chars: int = 2000) -> list[str]: + """Split ``text`` into chunks of at most ``max_chars`` characters. + + Greedy paragraph packing: paragraphs (blank-line separated) are packed + into chunks until the cap; a paragraph longer than the cap is hard-split + at the nearest space. Zero-loss: concatenating the chunks preserves every + paragraph's content. + # upgrade-path: heading/structure-aware chunking (keep md sections whole) + # and token-based budgets once the Phase-1 eval says boundaries matter. + """ + import re # noqa: PLC0415 + + text = text.strip() + if not text: + return [] + paras = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()] + chunks: list[str] = [] + cur = "" + for para in paras: + while len(para) > max_chars: + cut = para.rfind(" ", 1, max_chars) + if cut <= 0: + cut = max_chars + if cur: + chunks.append(cur) + cur = "" + chunks.append(para[:cut].strip()) + para = para[cut:].strip() + if not para: + continue + if not cur: + cur = para + elif len(cur) + 2 + len(para) <= max_chars: + cur = f"{cur}\n\n{para}" + else: + chunks.append(cur) + cur = para + if cur: + chunks.append(cur) + return chunks + + +# --------------------------------------------------------------------------- +# Ingest pipeline +# --------------------------------------------------------------------------- + +def _supersede_collection_rows(vmem, collection_id: str, file_path: str) -> int: + """Soft-supersede the active vector rows of one collection file. + + Zero-loss: rows are stamped ``valid_to`` (the existing supersede + machinery) with a ``hidden_by: collection-reindex:`` marker in their + metadata; nothing is deleted and the archive rows are untouched. Used on + re-index for changed and deleted files, mirroring the shelf-archive + pattern in :mod:`taosmd.admin`. + """ + ts = time.time() + marker = f"collection-reindex:{ts}" + rows = vmem._conn.execute( + "SELECT id, metadata_json FROM vector_memory WHERE valid_to IS NULL" + ).fetchall() + superseded = 0 + for row in rows: + try: + meta = json.loads(row["metadata_json"]) + except (json.JSONDecodeError, TypeError): + continue + user_md = meta.get("metadata") if isinstance(meta, dict) else None + if not isinstance(user_md, dict): + continue + if user_md.get("collection_id") != collection_id: + continue + if user_md.get("file_path") != file_path: + continue + meta["hidden_by"] = marker + vmem._conn.execute( + "UPDATE vector_memory SET valid_to = ?, metadata_json = ? " + "WHERE id = ? AND valid_to IS NULL", + (ts, json.dumps(meta), row["id"]), + ) + superseded += 1 + if superseded: + vmem._conn.commit() + vmem._bm25_dirty = True + return superseded + + +async def ingest_folder( + collection_id: str, + *, + data_dir=None, + max_file_bytes: int = DEFAULT_MAX_FILE_BYTES, + chunk_chars: int = 2000, +) -> dict: + """Walk a collection's source folder and index its documents. + + Incremental by content hash: files whose hash matches the stored state + are skipped; changed files get their old rows superseded (never deleted) + and their new chunks ingested; files that disappeared from the source + are superseded too. Chunks route through :func:`taosmd.api.ingest_batch` + under the collection id as the agent namespace, so the batch dedup and + metadata preservation come for free and every chunk lands in the + zero-loss archive. + + Sets the collection status to ``indexing`` for the duration, then + ``ready`` (stats updated, ``last_indexed`` stamped) or ``error`` (the + failure recorded on the row). Raises on validation errors so a caller + driving it synchronously still sees them. + """ + from . import api as _api # noqa: PLC0415 - avoid import cycle at module load + + resolved_dir = _api._resolve_data_dir(data_dir) + store = CollectionStore(resolved_dir) + col = store.get(collection_id) + if col["status"] == "archived": + raise ValueError(f"collection {collection_id!r} is archived; unarchive before indexing") + if col["embedder"]: + # Per-collection embedder is stored and returned now (the mechanism); + # Phase 1 indexes with the global default regardless. + logger.info( + "collection %s requests embedder %r; Phase 1 indexes with the " + "global default embedder", collection_id, col["embedder"], + ) + store.set_status(collection_id, "indexing") + try: + source_root = store.resolve_source_path(col["source_path"]) + prior = store.file_states(collection_id) + files, skips = collect_files(source_root, max_file_bytes=max_file_bytes) + + stores = await _api._ensure_stores(data_dir) + vmem = stores["vector"] + + items: list[dict] = [] + indexed: list[tuple[str, str]] = [] + changed: list[str] = [] + unchanged = 0 + errors: list[str] = [] + seen: set[str] = set() + + for abs_path, rel in files: + seen.add(rel) + loader = _loader_for(abs_path) + if loader is None: # pragma: no cover - collect_files already filtered + continue + try: + blob = await loader.load( + abs_path, max_bytes=max_file_bytes, base_dir=source_root + ) + except Exception as exc: # noqa: BLE001 - per-file failures are non-fatal + errors.append(f"{rel}: {type(exc).__name__}: {exc}") + continue + text = blob.raw_text or getattr(blob, "content", "") or "" + if not text.strip(): + continue + file_hash = hashlib.sha256(text.encode("utf-8")).hexdigest() + if prior.get(rel) == file_hash: + unchanged += 1 + continue + if rel in prior: + changed.append(rel) + for i, chunk in enumerate(chunk_text(text, max_chars=chunk_chars)): + chunk_id = hashlib.sha256( + f"{collection_id}:{rel}:{file_hash}:{i}".encode("utf-8") + ).hexdigest() + items.append({ + "text": chunk, + "id": chunk_id, + "metadata": { + "collection_id": collection_id, + "file_path": rel, + "source": "collection", + "chunk_index": i, + "file_hash": file_hash, + }, + }) + indexed.append((rel, file_hash)) + + deleted = sorted(set(prior) - seen) + + chunks_superseded = 0 + for rel in [*changed, *deleted]: + chunks_superseded += _supersede_collection_rows(vmem, collection_id, rel) + + if items: + result = await _api.ingest_batch(items, agent=collection_id, data_dir=data_dir) + else: + result = {"ingested": 0, "skipped": 0} + + for rel, file_hash in indexed: + store.set_file_state(collection_id, rel, file_hash) + for rel in deleted: + store.remove_file_state(collection_id, rel) + + now = time.time() + stats = { + "files_indexed": len(indexed), + "files_unchanged": unchanged, + "files_deleted": len(deleted), + "files_total": len(store.file_states(collection_id)), + "chunks_ingested": result.get("ingested", 0), + "chunks_skipped": result.get("skipped", 0), + "chunks_superseded": chunks_superseded, + "errors": errors[:20], + **skips, + } + if result.get("vector_failures"): + stats["vector_failures"] = result["vector_failures"] + stats["degraded"] = True + store.set_stats(collection_id, stats) + store.set_status(collection_id, "ready", last_indexed=now) + return stats + except Exception as exc: + store.set_status(collection_id, "error", error=f"{type(exc).__name__}: {exc}") + raise diff --git a/tests/test_collections_ingest.py b/tests/test_collections_ingest.py new file mode 100644 index 00000000..f66b762d --- /dev/null +++ b/tests/test_collections_ingest.py @@ -0,0 +1,335 @@ +"""Tests for the collections folder walker, chunker, and ingest pipeline. + +The walker is gitignore-aware (simplified stdlib fnmatch rules), skips +version-control/binary/oversized files, and only ingests files an existing +loader explicitly claims (DocLoader md/txt/markdown/rst plus the other +registered loaders). ingest_folder routes chunks through api.ingest_batch +under the collection's own agent namespace, so re-index dedups on chunk +content hashes and changed/deleted files are superseded, never destroyed. +""" +from __future__ import annotations + +import asyncio +import os +import time + +import pytest + +from taosmd import api as taosmd_api +from taosmd import config +from taosmd.collections import ( + CollectionStore, + chunk_text, + collect_files, + ingest_folder, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def source_dir(tmp_path): + d = tmp_path / "src" + d.mkdir() + (d / "readme.md").write_text("# Widget\n\nThe widget frobnicates the sprocket.") + (d / "guide.txt").write_text("Install with pip. Configure the flux capacitor.") + (d / "notes.rst").write_text("Deployment notes\n----------------\n\nUse systemd.") + sub = d / "deep" + sub.mkdir() + (sub / "inner.md").write_text("# Inner\n\nNested documentation file.") + return d + + +@pytest.fixture +def data_dir(tmp_path, monkeypatch, source_dir): + d = tmp_path / "taosmd-data" + d.mkdir() + monkeypatch.setenv("TAOSMD_DATA_DIR", str(d)) + monkeypatch.delenv("TAOSMD_COLLECTIONS_ALLOWED_ROOTS", raising=False) + monkeypatch.setattr(taosmd_api, "_stores_cache", {}) + config.set_collections_allowed_roots([str(source_dir.parent)], data_dir=str(d)) + return str(d) + + +def _patch_embedder(data_dir) -> None: + """Deterministic 8-dim hash embedder, same pattern as test_http_server.""" + stores = asyncio.run(taosmd_api._ensure_stores(data_dir)) + vmem = stores["vector"] + + async def _fake_embed(text: str, task: str = "search_document") -> list[float]: + h = hash(text) & 0xFFFFFFFF + return [((h >> (i * 4)) & 0xFF) / 255.0 for i in range(8)] + + vmem.embed = _fake_embed # type: ignore[assignment] + + +# --------------------------------------------------------------------------- +# Walker +# --------------------------------------------------------------------------- + +def test_collect_files_finds_claimed_docs(source_dir): + files, skips = collect_files(source_dir) + rels = sorted(rel for _, rel in files) + assert rels == ["deep/inner.md", "guide.txt", "notes.rst", "readme.md"] + + +def test_collect_files_skips_unclaimed_extensions(source_dir): + (source_dir / "main.py").write_text("print('hi')") + (source_dir / "data.json").write_text("{}") + files, skips = collect_files(source_dir) + rels = {rel for _, rel in files} + assert "main.py" not in rels + assert "data.json" not in rels + assert skips["skipped_unclaimed"] == 2 + + +def test_collect_files_skips_vcs_and_dep_dirs(source_dir): + git = source_dir / ".git" + git.mkdir() + (git / "config.txt").write_text("secret") + nm = source_dir / "node_modules" + nm.mkdir() + (nm / "pkg.md").write_text("# dep readme") + files, _ = collect_files(source_dir) + rels = {rel for _, rel in files} + assert not any(r.startswith(".git") or r.startswith("node_modules") for r in rels) + + +def test_collect_files_respects_gitignore(source_dir): + (source_dir / ".gitignore").write_text("*.log\nbuild/\n/secret.md\n") + (source_dir / "debug.log").write_text("log log log") + (source_dir / "secret.md").write_text("# do not index") + build = source_dir / "build" + build.mkdir() + (build / "out.md").write_text("# generated") + # A nested .gitignore applies within its own directory. + (source_dir / "deep" / ".gitignore").write_text("inner.md\n") + files, skips = collect_files(source_dir) + rels = {rel for _, rel in files} + assert "debug.log" not in rels + assert "secret.md" not in rels + assert "build/out.md" not in rels + assert "deep/inner.md" not in rels + assert "readme.md" in rels + assert skips["skipped_ignored"] >= 3 + + +def test_collect_files_skips_binary_and_oversized(source_dir): + (source_dir / "logo.png").write_bytes(b"\x89PNG\r\n" + b"\x00" * 32) + # A claimed extension whose content is binary is sniffed out. + (source_dir / "fake.md").write_bytes(b"\x00\x01\x02 binary sneaking as md") + (source_dir / "huge.md").write_text("x" * 4096) + files, skips = collect_files(source_dir, max_file_bytes=1024) + rels = {rel for _, rel in files} + assert "logo.png" not in rels + assert "fake.md" not in rels + assert "huge.md" not in rels + assert skips["skipped_binary"] >= 2 + assert skips["skipped_size"] == 1 + + +def test_collect_files_skips_symlink_escape(source_dir, tmp_path): + outside = tmp_path / "outside.md" + outside.write_text("# outside the root") + os.symlink(outside, source_dir / "escape.md") + files, skips = collect_files(source_dir) + rels = {rel for _, rel in files} + assert "escape.md" not in rels + + +# --------------------------------------------------------------------------- +# Chunker +# --------------------------------------------------------------------------- + +def test_chunk_text_short_is_single_chunk(): + assert chunk_text("hello world") == ["hello world"] + + +def test_chunk_text_packs_paragraphs(): + paras = [f"Paragraph {i} " + "word " * 30 for i in range(10)] + text = "\n\n".join(paras) + chunks = chunk_text(text, max_chars=500) + assert len(chunks) > 1 + assert all(len(c) <= 500 for c in chunks) + # Nothing lost: every paragraph's marker appears in exactly one chunk. + joined = "\n\n".join(chunks) + for i in range(10): + assert f"Paragraph {i} " in joined + + +def test_chunk_text_hard_splits_long_paragraph(): + text = "word " * 1000 # one giant paragraph, no blank lines + chunks = chunk_text(text, max_chars=800) + assert len(chunks) > 1 + assert all(len(c) <= 800 for c in chunks) + + +def test_chunk_text_empty(): + assert chunk_text("") == [] + assert chunk_text(" \n\n ") == [] + + +# --------------------------------------------------------------------------- +# ingest_folder +# --------------------------------------------------------------------------- + +def _make_collection(data_dir, source_dir): + store = CollectionStore(data_dir) + return store, store.create(name="docs", kind="docs", source_path=str(source_dir)) + + +def test_ingest_folder_indexes_docs(data_dir, source_dir): + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + stats = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + assert stats["files_indexed"] == 4 + assert stats["chunks_ingested"] >= 4 + got = store.get(col["id"]) + assert got["status"] == "ready" + assert got["last_indexed"] is not None + assert got["stats"]["files_indexed"] == 4 + # Chunk rows carry collection metadata for provenance. + hits = asyncio.run( + taosmd_api.search( + "widget frobnicates", agent="dev", mode="bm25", + collections=[col["id"]], collections_only=True, data_dir=data_dir, + ) + ) + # No grant yet: nothing visible. + assert hits == [] + store.grant(col["id"], "dev") + hits = asyncio.run( + taosmd_api.search( + "widget frobnicates", agent="dev", mode="bm25", + collections=[col["id"]], collections_only=True, data_dir=data_dir, + ) + ) + assert hits + top = hits[0] + assert top["metadata"]["collection_id"] == col["id"] + assert top["metadata"]["file_path"] == "readme.md" + assert top["metadata"]["source"] == "collection" + + +def test_reindex_unchanged_skips_everything(data_dir, source_dir): + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + stats2 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + assert stats2["files_indexed"] == 0 + assert stats2["files_unchanged"] == 4 + assert stats2["chunks_ingested"] == 0 + + +def test_reindex_changed_file_supersedes_old_rows(data_dir, source_dir): + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + store.grant(col["id"], "dev") + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + (source_dir / "readme.md").write_text("# Widget\n\nThe widget now defenestrates.") + stats2 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + assert stats2["files_indexed"] == 1 + assert stats2["files_unchanged"] == 3 + assert stats2["chunks_superseded"] >= 1 + + def _search(q): + return asyncio.run( + taosmd_api.search( + q, agent="dev", mode="bm25", + collections=[col["id"]], collections_only=True, data_dir=data_dir, + ) + ) + + # The old content is out of active recall; the new content is in. + assert not any("frobnicates" in h["text"] for h in _search("frobnicates sprocket")) + assert any("defenestrates" in h["text"] for h in _search("defenestrates")) + # Zero-loss: the superseded row still exists physically (valid_to set). + stores = asyncio.run(taosmd_api._ensure_stores(data_dir)) + row = stores["vector"]._conn.execute( + "SELECT COUNT(*) AS n FROM vector_memory WHERE valid_to IS NOT NULL" + ).fetchone() + assert row["n"] >= 1 + + +def test_reindex_deleted_file_supersedes_rows(data_dir, source_dir): + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + store.grant(col["id"], "dev") + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + (source_dir / "guide.txt").unlink() + stats2 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + assert stats2["files_deleted"] == 1 + hits = asyncio.run( + taosmd_api.search( + "flux capacitor", agent="dev", mode="bm25", + collections=[col["id"]], collections_only=True, data_dir=data_dir, + ) + ) + assert not any("flux capacitor" in h["text"] for h in hits) + assert "guide.txt" not in store.file_states(col["id"]) + + +def test_ingest_folder_archived_collection_refused(data_dir, source_dir): + store, col = _make_collection(data_dir, source_dir) + store.archive(col["id"]) + with pytest.raises(ValueError): + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + + +def test_ingest_folder_root_removed_from_config_refused(data_dir, source_dir): + store, col = _make_collection(data_dir, source_dir) + config.set_collections_allowed_roots([], clear=True, data_dir=data_dir) + with pytest.raises(ValueError): + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + assert store.get(col["id"])["status"] == "error" + + +# --------------------------------------------------------------------------- +# Search integration +# --------------------------------------------------------------------------- + +def test_search_without_collections_excludes_collection_rows(data_dir, source_dir): + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + store.grant(col["id"], "dev") + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + hits = asyncio.run( + taosmd_api.search("widget frobnicates", agent="dev", mode="bm25", data_dir=data_dir) + ) + assert not any(h["metadata"].get("collection_id") for h in hits) + + +def test_search_with_collections_merges_own_memory(data_dir, source_dir): + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + store.grant(col["id"], "dev") + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + asyncio.run( + taosmd_api.ingest("The widget owner is Jay.", agent="dev", data_dir=data_dir) + ) + hits = asyncio.run( + taosmd_api.search( + "widget", agent="dev", mode="bm25", limit=10, + collections=[col["id"]], data_dir=data_dir, + ) + ) + sources = {h["metadata"].get("collection_id") for h in hits} + assert col["id"] in sources # collection hit present + assert None in sources # own conversational memory present too + + +def test_search_archived_collection_hidden(data_dir, source_dir): + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + store.grant(col["id"], "dev") + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + store.archive(col["id"]) + hits = asyncio.run( + taosmd_api.search( + "widget frobnicates", agent="dev", mode="bm25", + collections=[col["id"]], collections_only=True, data_dir=data_dir, + ) + ) + assert hits == [] From 485ef095513aca1b32c23f5c44be456974dfe880 Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 00:08:06 +0100 Subject: [PATCH 04/16] feat(collections): service wrappers + HTTP surface (admin create/index, async 202 poll) Service layer gains local-only collections_* wrappers (like the shelf admin wrappers: the server that owns the indexed filesystem runs the ops). HTTP adds the full contract from the design spec: - POST /collections and POST /collections/{id}/index are admin-token gated (fail closed), joining the shelves routes in _is_admin_route - DELETE /collections/{id} (admin) archives reversibly; do_DELETE is new - data plane: GET /collections[?project=], GET /collections/{id}, POST link/unlink (typed {taos|git, id}), POST grants, DELETE /collections/{id}/grants/{agent} - indexing is async: 202 + {status: indexing, job}, poll the GET until ready|error; the walk runs via a new _ServiceLoop.spawn so all store access stays on the single service-loop thread and HTTP never blocks - GET/POST /search gain collection/collections/collections_only params with per-agent grant enforcement Endpoint docstring table and serve() startup summary updated. --- taosmd/http_server.py | 249 ++++++++++++++++++++++- taosmd/service.py | 150 +++++++++++++- tests/test_collections_http.py | 356 +++++++++++++++++++++++++++++++++ 3 files changed, 747 insertions(+), 8 deletions(-) create mode 100644 tests/test_collections_http.py diff --git a/taosmd/http_server.py b/taosmd/http_server.py index 8db250b9..aae42870 100644 --- a/taosmd/http_server.py +++ b/taosmd/http_server.py @@ -76,8 +76,12 @@ Adds ``"vector_failures": int`` and ``"degraded": true`` to the result when embedding failed for one or more items (archived, repairable via reconcile), same as single ingest. -``POST /search`` ``{"query", "agent", "limit"?, "project"?, "also_include"?, "mode"?}`` -> ``{"hits": [...]}`` -``GET /search?q=&agent=&limit=&project=&also_include=a,b&mode=bm25`` -> ``{"hits": [...]}`` +``POST /search`` ``{"query", "agent", "limit"?, "project"?, "also_include"?, "mode"?, "collection"?, "collections"?: [...], "collections_only"?: bool}`` -> ``{"hits": [...]}`` + ``collection``/``collections`` add granted collections' indexed + content to the search (grants enforced per requesting agent; + collection hits carry ``collection_id``/``file_path``/``source`` + metadata); ``collections_only`` restricts to them. +``GET /search?q=&agent=&limit=&project=&also_include=a,b&mode=bm25&collection=&collections=a,b&collections_only=true`` -> ``{"hits": [...]}`` ``GET /projects`` -> ``{"projects": [...]}`` ``GET /shelves?project=`` -> ``{"shelves": [...]}`` ``GET /pending?agent=`` -> ``{"pending": [...]}`` @@ -95,7 +99,25 @@ ``POST /tasks/{id}/edges`` ``{"to_id", "type", "created_by"}`` -> edge record ``POST /tasks/{id}/edges/remove`` ``{"to_id", "type"}`` -> edge record with removed_ts +Collections (data plane; create/index/delete are admin, see below) +``GET /collections [?project=]`` -> ``{"collections": [...]}`` (project matches links of either type) +``GET /collections/{id}`` -> ``{"collection": {...}}`` with status/stats/links/grants +``POST /collections/{id}/link`` ``{"type": "taos"|"git", "id"}`` -> ``{"collection": {...}}`` +``POST /collections/{id}/unlink`` same body; metadata only, never touches content +``POST /collections/{id}/grants`` ``{"agent"}`` -> grant query access -> ``{"collection": {...}}`` +``DELETE /collections/{id}/grants/{agent}`` -> revoke -> ``{"collection": {...}}`` + Admin endpoints (all require a configured server token; 403 if none is set) +``POST /collections`` ``{"name", "kind": docs|codebase|mixed, "source_path", "embedder"?}`` + -> ``{"collection": {...}, "created": true}`` + source_path must resolve inside a configured + collections.allowed_roots directory (400 otherwise; + empty roots = collections off) +``POST /collections/{id}/index`` -> 202 ``{"status": "indexing", "job": }`` + async; poll GET /collections/{id} until status is + "ready" or "error"; stats update on completion +``DELETE /collections/{id}`` -> archive (reversible; content hidden from query, + nothing destroyed; destruction only via wipe) ``POST /shelves`` ``{"shelf_id", "project_id"?, "display_name"?}`` -> ``{"shelf": {...}, "created": bool}`` ``POST /shelves/{id}/archive`` ``?expect_empty=true`` -> ``{"archived": true, "rows_hidden": int}`` ``POST /shelves/{id}/unarchive`` -> ``{"archived": false, "rows_restored": int}`` @@ -507,6 +529,19 @@ def run(self, coro): raise return future.result() + def spawn(self, coro) -> None: + """Fire-and-forget: schedule ``coro`` on the loop without blocking. + + Used for background jobs (collection indexing) that outlive the HTTP + request. The coroutine must handle its own exceptions (the returned + future is intentionally not awaited). + """ + try: + asyncio.run_coroutine_threadsafe(coro, self._loop) + except RuntimeError: + coro.close() + raise + def close(self) -> None: self._loop.call_soon_threadsafe(self._loop.stop) self._thread.join(timeout=5) @@ -667,10 +702,20 @@ def _is_admin_route(method: str, path: str) -> bool: that gating the admin surface never locks the data plane. ``path`` has already had any trailing slash stripped by the caller. """ + if method == "DELETE" and path.startswith("/collections/"): + # DELETE /collections/{id} (archive) is admin; the grants + # revoke sub-path (DELETE /collections/{id}/grants/{agent}) + # stays on the data plane. + rest = path[len("/collections/"):] + return bool(rest) and "/" not in rest if method != "POST": return False if path == "/shelves" or path.startswith("/shelves/"): return True + if path == "/collections": + return True + if path.startswith("/collections/") and path.endswith("/index"): + return True return path in ( "/a2a/admin/delete-channel", "/a2a/admin/rename-channel", @@ -805,6 +850,9 @@ def do_HEAD(self) -> None: # noqa: N802 def do_POST(self) -> None: # noqa: N802 self._dispatch("POST") + def do_DELETE(self) -> None: # noqa: N802 + self._dispatch("DELETE") + def _dispatch(self, method: str) -> None: parts = urlsplit(self.path) path = parts.path.rstrip("/") or "/" @@ -919,6 +967,41 @@ def _dispatch(self, method: str) -> None: self._handle_admin_shelf_unarchive(shelf_id) else: self._send_json(404, {"error": f"unknown shelf action: {rest}"}) + # ----- collections ---------------------------------------- + elif method == "GET" and path == "/collections": + self._handle_collections_list(query) + elif method == "GET" and path.startswith("/collections/"): + cid = path[len("/collections/"):] + if not cid or "/" in cid: + self._send_json(404, {"error": "collection id required"}) + else: + self._handle_collections_get(cid) + elif method == "POST" and path == "/collections": + self._handle_collections_create() + elif method == "POST" and path.startswith("/collections/"): + rest = path[len("/collections/"):] + if rest.endswith("/index"): + self._handle_collections_index(rest[: -len("/index")]) + elif rest.endswith("/link"): + self._handle_collections_link(rest[: -len("/link")], unlink=False) + elif rest.endswith("/unlink"): + self._handle_collections_link(rest[: -len("/unlink")], unlink=True) + elif rest.endswith("/grants"): + self._handle_collections_grant(rest[: -len("/grants")]) + else: + self._send_json(404, {"error": f"unknown collection action: {rest}"}) + elif method == "DELETE" and path.startswith("/collections/"): + rest = path[len("/collections/"):] + if "/grants/" in rest: + cid, _, agent = rest.partition("/grants/") + if not cid or not agent: + self._send_json(404, {"error": "collection id and agent required"}) + else: + self._handle_collections_revoke(cid, agent) + elif rest and "/" not in rest: + self._handle_collections_delete(rest) + else: + self._send_json(404, {"error": f"unknown collection action: {rest}"}) # ----- admin surface: A2A channel admin ------------------- elif method == "POST" and path == "/a2a/admin/delete-channel": self._handle_admin_a2a_delete_channel() @@ -991,10 +1074,14 @@ def _handle_search_post(self) -> None: project = body.get("project") also_include = body.get("also_include") mode = body.get("mode") + collections = body.get("collections") + collection = body.get("collection") + collections_only = body.get("collections_only", False) project, ok = self._apply_token_binding(agent, project) if not ok: return - self._do_search(query, agent, limit, project, also_include, mode) + self._do_search(query, agent, limit, project, also_include, mode, + collections, collection, collections_only) def _handle_search_get(self, qs: dict) -> None: query = (qs.get("q") or qs.get("query") or [None])[0] @@ -1005,12 +1092,21 @@ def _handle_search_get(self, qs: dict) -> None: ai_raw = (qs.get("also_include") or [None])[0] also_include = [s for s in ai_raw.split(",") if s] if ai_raw else None mode = (qs.get("mode") or [None])[0] + # Collection scoping: ?collection= or ?collections=a,b, plus + # ?collections_only=true to exclude conversation memory. + col_raw = (qs.get("collections") or [None])[0] + collections = [s for s in col_raw.split(",") if s] if col_raw else None + collection = (qs.get("collection") or [None])[0] + co_raw = (qs.get("collections_only") or [None])[0] + collections_only = co_raw is not None and co_raw.lower() == "true" project, ok = self._apply_token_binding(agent, project) if not ok: return - self._do_search(query, agent, limit, project, also_include, mode) + self._do_search(query, agent, limit, project, also_include, mode, + collections, collection, collections_only) - def _do_search(self, query, agent, limit, project=None, also_include=None, mode=None) -> None: + def _do_search(self, query, agent, limit, project=None, also_include=None, mode=None, + collections=None, collection=None, collections_only=False) -> None: if not isinstance(query, str) or not query: raise _BadRequest("'query' (non-empty string) is required") if not isinstance(agent, str) or not agent: @@ -1027,6 +1123,15 @@ def _do_search(self, query, agent, limit, project=None, also_include=None, mode= raise _BadRequest("'limit' must be an integer") from exc if mode is not None and not isinstance(mode, str): raise _BadRequest("'mode' must be a string when provided") + if collection is not None and not isinstance(collection, str): + raise _BadRequest("'collection' must be a string when provided") + if collections is not None and not ( + isinstance(collections, list) and all(isinstance(s, str) for s in collections) + ): + raise _BadRequest("'collections' must be a list of strings when provided") + all_collections = list(collections or []) + if collection and collection not in all_collections: + all_collections.append(collection) opts: dict = {} if project: opts["project"] = project @@ -1034,6 +1139,10 @@ def _do_search(self, query, agent, limit, project=None, also_include=None, mode= opts["also_include"] = also_include if mode: opts["mode"] = mode + if all_collections: + opts["collections"] = all_collections + if collections_only: + opts["collections_only"] = True hits = runner.run( service.search(query, agent=agent, data_dir=data_dir, limit=limit_i, **opts) ) @@ -1686,6 +1795,127 @@ def _handle_admin_shelf_unarchive(self, shelf_id: str) -> None: return self._send_json(200, result) + # ----- collections ------------------------------------------------ + + def _handle_collections_create(self) -> None: + if not self._check_admin_token(): + return + body = self._read_json_body() + name = body.get("name") + kind = body.get("kind") + source_path = body.get("source_path") + embedder = body.get("embedder") + if not isinstance(name, str) or not name: + raise _BadRequest("'name' (non-empty string) is required") + if not isinstance(kind, str) or not kind: + raise _BadRequest("'kind' (non-empty string) is required") + if not isinstance(source_path, str) or not source_path: + raise _BadRequest("'source_path' (non-empty string) is required") + if embedder is not None and not isinstance(embedder, str): + raise _BadRequest("'embedder' must be a string when provided") + result = runner.run( + service.collections_create( + name=name, kind=kind, source_path=source_path, + embedder=embedder, data_dir=data_dir, + ) + ) + self._send_json(200, {"collection": result, "created": True}) + + def _handle_collections_list(self, qs: dict) -> None: + project = (qs.get("project") or [None])[0] + cols = runner.run( + service.collections_list(project=project, data_dir=data_dir) + ) + self._send_json(200, {"collections": cols}) + + def _handle_collections_get(self, collection_id: str) -> None: + from .collections import CollectionNotFoundError # noqa: PLC0415 + try: + col = runner.run( + service.collections_get(collection_id, data_dir=data_dir) + ) + except CollectionNotFoundError as exc: + self._send_json(404, {"error": str(exc)}) + return + self._send_json(200, {"collection": col}) + + def _handle_collections_index(self, collection_id: str) -> None: + if not self._check_admin_token(): + return + from .collections import CollectionNotFoundError # noqa: PLC0415 + try: + receipt = runner.run( + service.collections_index_start(collection_id, data_dir=data_dir) + ) + except CollectionNotFoundError as exc: + self._send_json(404, {"error": str(exc)}) + return + # Async by contract: 202 now, poll GET /collections/{id} until the + # status is ready|error. The walk runs on the service loop so all + # store access stays in the single-threaded context. + runner.spawn( + service.collections_index_background(collection_id, data_dir=data_dir) + ) + self._send_json(202, receipt) + + def _handle_collections_link(self, collection_id: str, *, unlink: bool) -> None: + from .collections import CollectionNotFoundError # noqa: PLC0415 + body = self._read_json_body() + link_type = body.get("type") + ext_id = body.get("id") + if not isinstance(link_type, str) or not link_type: + raise _BadRequest("'type' (\"taos\" or \"git\") is required") + if not isinstance(ext_id, str) or not ext_id: + raise _BadRequest("'id' (non-empty string) is required") + fn = service.collections_unlink if unlink else service.collections_link + try: + col = runner.run( + fn(collection_id, link_type, ext_id, data_dir=data_dir) + ) + except CollectionNotFoundError as exc: + self._send_json(404, {"error": str(exc)}) + return + self._send_json(200, {"collection": col}) + + def _handle_collections_grant(self, collection_id: str) -> None: + from .collections import CollectionNotFoundError # noqa: PLC0415 + body = self._read_json_body() + agent = body.get("agent") + if not isinstance(agent, str) or not agent: + raise _BadRequest("'agent' (non-empty string) is required") + try: + col = runner.run( + service.collections_grant(collection_id, agent, data_dir=data_dir) + ) + except CollectionNotFoundError as exc: + self._send_json(404, {"error": str(exc)}) + return + self._send_json(200, {"collection": col}) + + def _handle_collections_revoke(self, collection_id: str, agent: str) -> None: + from .collections import CollectionNotFoundError # noqa: PLC0415 + try: + col = runner.run( + service.collections_revoke(collection_id, agent, data_dir=data_dir) + ) + except CollectionNotFoundError as exc: + self._send_json(404, {"error": str(exc)}) + return + self._send_json(200, {"collection": col}) + + def _handle_collections_delete(self, collection_id: str) -> None: + if not self._check_admin_token(): + return + from .collections import CollectionNotFoundError # noqa: PLC0415 + try: + col = runner.run( + service.collections_archive(collection_id, data_dir=data_dir) + ) + except CollectionNotFoundError as exc: + self._send_json(404, {"error": str(exc)}) + return + self._send_json(200, {"collection": col}) + # ----- admin: A2A channel admin ---------------------------------- def _handle_admin_a2a_delete_channel(self) -> None: @@ -1785,8 +2015,13 @@ def serve(host: str = DEFAULT_HOST, port: int = DEFAULT_PORT, data_dir=None) -> "POST /a2a/send, GET /a2a/messages, GET /a2a/stream, " "GET /a2a/channels, GET /a2a/members, " "POST /tasks, GET /tasks, GET /tasks/ready, GET /tasks/prime, " - "POST /tasks/{id}, POST /tasks/{id}/edges, POST /tasks/{id}/edges/remove") - print("Admin (admin token required): POST /shelves, POST /shelves/{id}/archive, " + "POST /tasks/{id}, POST /tasks/{id}/edges, POST /tasks/{id}/edges/remove, " + "GET /collections, GET /collections/{id}, POST /collections/{id}/link, " + "POST /collections/{id}/unlink, POST /collections/{id}/grants, " + "DELETE /collections/{id}/grants/{agent}") + print("Admin (admin token required): POST /collections, POST /collections/{id}/index, " + "DELETE /collections/{id}, " + "POST /shelves, POST /shelves/{id}/archive, " "POST /shelves/{id}/unarchive, " "POST /a2a/admin/delete-channel, POST /a2a/admin/rename-channel, " "POST /a2a/admin/supersede-message") diff --git a/taosmd/service.py b/taosmd/service.py index dad5f98e..2de88a44 100644 --- a/taosmd/service.py +++ b/taosmd/service.py @@ -26,11 +26,14 @@ from __future__ import annotations import json +import logging from . import api as _api from . import config as _config from .archive import EVENT_A2A +logger = logging.getLogger(__name__) + # Cache of RemoteClient instances keyed by (base_url, token) so we don't # create a fresh object on every call. Access from async coroutines is safe # because Python dict operations are GIL-protected. @@ -851,10 +854,155 @@ async def admin_a2a_supersede_message(msg_id: int, *, data_dir=None) -> dict: return await a2a_admin_supersede_message(msg_id, data_dir=data_dir, stores=stores) +# --------------------------------------------------------------------------- +# Collections service wrappers +# --------------------------------------------------------------------------- +# +# Like the shelf admin wrappers these are local-only (no remote forwarding): +# the server that owns the filesystem being indexed is the server that runs +# the collection ops. All wrappers open the store against the resolved data +# dir and close it, so no sqlite connection outlives a call. + +def _collection_store(data_dir): + from .collections import CollectionStore # noqa: PLC0415 + return CollectionStore(_api._resolve_data_dir(data_dir)) + + +async def collections_create( + *, + name: str, + kind: str, + source_path: str, + embedder: str | None = None, + data_dir=None, +) -> dict: + """Create a collection row (admin operation). Returns the collection.""" + store = _collection_store(data_dir) + try: + return store.create( + name=name, kind=kind, source_path=source_path, embedder=embedder, + ) + finally: + store.close() + + +async def collections_list(*, project: str | None = None, data_dir=None) -> list[dict]: + """List collections, optionally filtered to one project's links.""" + store = _collection_store(data_dir) + try: + return store.list(project=project) + finally: + store.close() + + +async def collections_get(collection_id: str, *, data_dir=None) -> dict: + """Return one collection with full stats, links, and grants.""" + store = _collection_store(data_dir) + try: + return store.get(collection_id) + finally: + store.close() + + +async def collections_index_start(collection_id: str, *, data_dir=None) -> dict: + """Validate and mark a collection ``indexing``; the walk runs separately. + + Raises ``CollectionNotFoundError`` (404) for an unknown id and + ``ValueError`` (400) for an archived collection or a source path that no + longer resolves inside an allowed root, so callers get a synchronous + error before the background job is spawned. + """ + store = _collection_store(data_dir) + try: + col = store.get(collection_id) + if col["status"] == "archived": + raise ValueError( + f"collection {collection_id!r} is archived; unarchive before indexing" + ) + store.resolve_source_path(col["source_path"]) + store.set_status(collection_id, "indexing") + finally: + store.close() + return {"status": "indexing", "job": collection_id} + + +async def collections_index_run(collection_id: str, *, data_dir=None) -> dict: + """Run the folder walk + ingest for one collection (blocking variant).""" + from .collections import ingest_folder # noqa: PLC0415 + return await ingest_folder(collection_id, data_dir=data_dir) + + +async def collections_index_background(collection_id: str, *, data_dir=None) -> None: + """Background wrapper for the HTTP 202 path: never raises, only logs. + + ``ingest_folder`` records failures on the collection row + (status='error' + message) so pollers see them; this wrapper keeps the + fire-and-forget future from warning about an unobserved exception. + """ + try: + await collections_index_run(collection_id, data_dir=data_dir) + except Exception: # noqa: BLE001 - surfaced via the collection row + logger.exception("collections: background index failed for %s", collection_id) + + +async def collections_link( + collection_id: str, link_type: str, ext_id: str, *, data_dir=None +) -> dict: + """Attach a typed project link ({taos|git, id}). Metadata only.""" + store = _collection_store(data_dir) + try: + return store.link(collection_id, link_type, ext_id) + finally: + store.close() + + +async def collections_unlink( + collection_id: str, link_type: str, ext_id: str, *, data_dir=None +) -> dict: + """Remove a typed project link. Metadata only, content untouched.""" + store = _collection_store(data_dir) + try: + return store.unlink(collection_id, link_type, ext_id) + finally: + store.close() + + +async def collections_grant(collection_id: str, agent: str, *, data_dir=None) -> dict: + """Grant ``agent`` query access to the collection.""" + store = _collection_store(data_dir) + try: + return store.grant(collection_id, agent) + finally: + store.close() + + +async def collections_revoke(collection_id: str, agent: str, *, data_dir=None) -> dict: + """Revoke ``agent``'s query access.""" + store = _collection_store(data_dir) + try: + return store.revoke(collection_id, agent) + finally: + store.close() + + +async def collections_archive(collection_id: str, *, data_dir=None) -> dict: + """Archive a collection (the DELETE alias). Reversible; nothing destroyed.""" + store = _collection_store(data_dir) + try: + return store.archive(collection_id) + finally: + store.close() + + __all__ = ["ingest", "search", "pending_list", "pending_resolve", "reconcile", "stats", "supersede", "a2a_send", "a2a_feed", "a2a_channels", "a2a_members", "task_create", "task_list", "task_ready", "task_prime", "task_update", "task_add_edge", "task_remove_edge", "task_projects", "admin_shelf_create", "admin_shelf_archive", "admin_shelf_unarchive", "admin_a2a_delete_channel", "admin_a2a_rename_channel", - "admin_a2a_supersede_message"] + "admin_a2a_supersede_message", + "collections_create", "collections_list", "collections_get", + "collections_index_start", "collections_index_run", + "collections_index_background", "collections_link", + "collections_unlink", "collections_grant", "collections_revoke", + "collections_archive"] diff --git a/tests/test_collections_http.py b/tests/test_collections_http.py new file mode 100644 index 00000000..b660283a --- /dev/null +++ b/tests/test_collections_http.py @@ -0,0 +1,356 @@ +"""HTTP surface tests for collections. + +Contract (docs/specs/codebase-indexing-collections-design.md + the Jul 19 +decisions): POST /collections and POST /collections/{id}/index are +admin-token-gated (fail closed); list/get/link/unlink/grants and search are +data-plane. DELETE /collections/{id} archives (admin). Indexing is async: +202 then poll GET /collections/{id} until status is ready|error. Search +gains collection/collections params with per-agent grant enforcement. +""" +from __future__ import annotations + +import json +import threading +import time +import urllib.error +import urllib.request + +import pytest + +from taosmd import api as taosmd_api +from taosmd import config as taosmd_config +from taosmd import http_server + +_TOKEN = "test-admin-token-abc123" + + +def _patch_embedder(stores: dict) -> None: + vmem = stores["vector"] + + async def _fake_embed(text: str, task: str = "search_document") -> list[float]: + h = hash(text) & 0xFFFFFFFF + return [((h >> (i * 4)) & 0xFF) / 255.0 for i in range(8)] + + vmem.embed = _fake_embed # type: ignore[assignment] + + +def _req(method: str, url: str, payload=None, token: str | None = None): + headers = {} + data = None + if payload is not None: + data = json.dumps(payload).encode() + headers["Content-Type"] = "application/json" + if token: + headers["Authorization"] = f"Bearer {token}" + req = urllib.request.Request(url, data=data, headers=headers, method=method) + try: + with urllib.request.urlopen(req, timeout=10) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError as exc: + return exc.code, json.loads(exc.read().decode()) + + +@pytest.fixture +def source_dir(tmp_path): + d = tmp_path / "src" + d.mkdir() + (d / "readme.md").write_text("# Widget\n\nThe widget frobnicates the sprocket.") + (d / "guide.txt").write_text("Install with pip and enjoy.") + return d + + +@pytest.fixture +def live_server(tmp_path, monkeypatch, source_dir): + data_dir = tmp_path / "taosmd-data" + data_dir.mkdir() + monkeypatch.setattr(taosmd_api, "_stores_cache", {}) + monkeypatch.setenv("TAOSMD_TOKEN", _TOKEN) + monkeypatch.setenv("TAOSMD_DATA_DIR", str(data_dir)) + monkeypatch.delenv("TAOSMD_COLLECTIONS_ALLOWED_ROOTS", raising=False) + taosmd_config.set_collections_allowed_roots( + [str(source_dir)], data_dir=str(data_dir) + ) + + httpd = http_server.make_server("127.0.0.1", 0, data_dir=str(data_dir)) + stores = httpd.service_loop.run(taosmd_api._ensure_stores(str(data_dir))) + _patch_embedder(stores) + + host, port = httpd.server_address[:2] + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + try: + yield f"http://{host}:{port}", str(data_dir), source_dir + finally: + httpd.shutdown() + httpd.server_close() + thread.join(timeout=5) + for s in list(taosmd_api._stores_cache.values()): + for store in (s.get("archive"), s.get("vector"), s.get("kg")): + if store and hasattr(store, "close"): + try: + httpd.service_loop.run(store.close()) + except Exception: + pass + httpd.service_loop.close() + + +@pytest.fixture +def live_server_no_token(tmp_path, monkeypatch, source_dir): + data_dir = tmp_path / "taosmd-data-nt" + data_dir.mkdir() + monkeypatch.setattr(taosmd_api, "_stores_cache", {}) + monkeypatch.delenv("TAOSMD_TOKEN", raising=False) + monkeypatch.delenv("TAOSMD_ADMIN_TOKEN", raising=False) + httpd = http_server.make_server("127.0.0.1", 0, data_dir=str(data_dir)) + host, port = httpd.server_address[:2] + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + try: + yield f"http://{host}:{port}" + finally: + httpd.shutdown() + httpd.server_close() + thread.join(timeout=5) + httpd.service_loop.close() + + +def _create(base, source_dir, name="repo docs"): + status, body = _req( + "POST", f"{base}/collections", + {"name": name, "kind": "docs", "source_path": str(source_dir)}, + token=_TOKEN, + ) + assert status == 200, body + return body["collection"] + + +def _wait_ready(base, cid, timeout=10.0): + deadline = time.time() + timeout + while time.time() < deadline: + status, body = _req("GET", f"{base}/collections/{cid}", token=_TOKEN) + assert status == 200 + if body["collection"]["status"] in ("ready", "error"): + return body["collection"] + time.sleep(0.05) + raise AssertionError("collection never left 'indexing'") + + +# --------------------------------------------------------------------------- +# Admin gating +# --------------------------------------------------------------------------- + +def test_create_fails_closed_without_any_token(live_server_no_token): + status, body = _req( + "POST", f"{live_server_no_token}/collections", + {"name": "x", "kind": "docs", "source_path": "/tmp"}, + ) + assert status == 403 + + +def test_create_requires_admin_token(live_server): + base, _, source_dir = live_server + status, _ = _req( + "POST", f"{base}/collections", + {"name": "x", "kind": "docs", "source_path": str(source_dir)}, + ) + assert status == 401 + status, _ = _req( + "POST", f"{base}/collections", + {"name": "x", "kind": "docs", "source_path": str(source_dir)}, + token="wrong-token", + ) + assert status == 401 + + +def test_index_requires_admin_token(live_server): + base, _, source_dir = live_server + col = _create(base, source_dir) + status, _ = _req("POST", f"{base}/collections/{col['id']}/index") + assert status == 401 + + +def test_delete_requires_admin_token(live_server): + base, _, source_dir = live_server + col = _create(base, source_dir) + status, _ = _req("DELETE", f"{base}/collections/{col['id']}") + assert status == 401 + + +# --------------------------------------------------------------------------- +# Create / list / get +# --------------------------------------------------------------------------- + +def test_create_and_get(live_server): + base, _, source_dir = live_server + col = _create(base, source_dir) + assert col["status"] == "created" + assert col["kind"] == "docs" + status, body = _req("GET", f"{base}/collections/{col['id']}", token=_TOKEN) + assert status == 200 + assert body["collection"]["name"] == "repo docs" + assert body["collection"]["embedder"] is None + + +def test_create_rejects_disallowed_path(live_server, tmp_path): + base, _, _ = live_server + outside = tmp_path / "definitely-elsewhere" + outside.mkdir() + status, body = _req( + "POST", f"{base}/collections", + {"name": "x", "kind": "docs", "source_path": str(outside)}, + token=_TOKEN, + ) + assert status == 400 + assert "allowed root" in body["error"] + + +def test_get_unknown_404(live_server): + base, _, _ = live_server + status, _ = _req("GET", f"{base}/collections/col-000000000000", token=_TOKEN) + assert status == 404 + + +def test_list_and_project_filter(live_server): + base, _, source_dir = live_server + a = _create(base, source_dir, name="a") + b = _create(base, source_dir, name="b") + status, body = _req( + "POST", f"{base}/collections/{a['id']}/link", + {"type": "taos", "id": "prj-123"}, token=_TOKEN, + ) + assert status == 200 + status, body = _req("GET", f"{base}/collections", token=_TOKEN) + assert status == 200 + assert {c["id"] for c in body["collections"]} == {a["id"], b["id"]} + status, body = _req("GET", f"{base}/collections?project=prj-123", token=_TOKEN) + assert [c["id"] for c in body["collections"]] == [a["id"]] + + +# --------------------------------------------------------------------------- +# Link / unlink / grants +# --------------------------------------------------------------------------- + +def test_link_unlink(live_server): + base, _, source_dir = live_server + col = _create(base, source_dir) + status, body = _req( + "POST", f"{base}/collections/{col['id']}/link", + {"type": "git", "id": "abc123def456"}, token=_TOKEN, + ) + assert status == 200 + assert {"type": "git", "id": "abc123def456"} in body["collection"]["links"] + status, body = _req( + "POST", f"{base}/collections/{col['id']}/unlink", + {"type": "git", "id": "abc123def456"}, token=_TOKEN, + ) + assert status == 200 + assert body["collection"]["links"] == [] + + +def test_link_bad_type_400(live_server): + base, _, source_dir = live_server + col = _create(base, source_dir) + status, _ = _req( + "POST", f"{base}/collections/{col['id']}/link", + {"type": "jira", "id": "X-1"}, token=_TOKEN, + ) + assert status == 400 + + +def test_grants_add_and_revoke(live_server): + base, _, source_dir = live_server + col = _create(base, source_dir) + status, body = _req( + "POST", f"{base}/collections/{col['id']}/grants", + {"agent": "dev"}, token=_TOKEN, + ) + assert status == 200 + assert "dev" in body["collection"]["grants"] + status, body = _req( + "DELETE", f"{base}/collections/{col['id']}/grants/dev", token=_TOKEN, + ) + assert status == 200 + assert body["collection"]["grants"] == [] + + +# --------------------------------------------------------------------------- +# Index (async, 202 + poll) and search integration +# --------------------------------------------------------------------------- + +def test_index_flow_and_search_grant_enforcement(live_server): + base, _, source_dir = live_server + col = _create(base, source_dir) + status, body = _req( + "POST", f"{base}/collections/{col['id']}/index", token=_TOKEN, + ) + assert status == 202 + assert body["status"] == "indexing" + assert body["job"] == col["id"] + + ready = _wait_ready(base, col["id"]) + assert ready["status"] == "ready" + assert ready["stats"]["files_indexed"] == 2 + assert ready["last_indexed"] is not None + + # No grant: the collection contributes nothing. + status, body = _req( + "POST", f"{base}/search", + {"query": "widget frobnicates", "agent": "dev", "mode": "bm25", + "collections": [col["id"]], "collections_only": True}, + token=_TOKEN, + ) + assert status == 200 + assert body["hits"] == [] + + status, _ = _req( + "POST", f"{base}/collections/{col['id']}/grants", + {"agent": "dev"}, token=_TOKEN, + ) + assert status == 200 + status, body = _req( + "POST", f"{base}/search", + {"query": "widget frobnicates", "agent": "dev", "mode": "bm25", + "collections": [col["id"]], "collections_only": True}, + token=_TOKEN, + ) + assert status == 200 + assert body["hits"] + top = body["hits"][0] + assert top["metadata"]["file_path"] == "readme.md" + assert top["metadata"]["collection_id"] == col["id"] + + # Singular form works too (GET-style single collection). + status, body = _req( + "POST", f"{base}/search", + {"query": "widget frobnicates", "agent": "dev", "mode": "bm25", + "collection": col["id"], "collections_only": True}, + token=_TOKEN, + ) + assert status == 200 + assert body["hits"] + + +def test_index_unknown_404(live_server): + base, _, _ = live_server + status, _ = _req( + "POST", f"{base}/collections/col-000000000000/index", token=_TOKEN, + ) + assert status == 404 + + +# --------------------------------------------------------------------------- +# Delete (archive) +# --------------------------------------------------------------------------- + +def test_delete_archives_reversibly(live_server): + base, _, source_dir = live_server + col = _create(base, source_dir) + status, body = _req("DELETE", f"{base}/collections/{col['id']}", token=_TOKEN) + assert status == 200 + assert body["collection"]["status"] == "archived" + # Hidden from the default listing, still fetchable by id (nothing destroyed). + status, body = _req("GET", f"{base}/collections", token=_TOKEN) + assert body["collections"] == [] + status, body = _req("GET", f"{base}/collections/{col['id']}", token=_TOKEN) + assert status == 200 + assert body["collection"]["status"] == "archived" From 394f0d9fe77ade56f7648a8f2405e7c7cd4890bb Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 00:11:11 +0100 Subject: [PATCH 05/16] feat(collections): CLI subcommands + MCP surface, uniform hit metadata CLI: taosmd collections list|create|index|link|unlink|grant|revoke in the shelves/tasks subcommand style; index runs synchronously and prints the final stats (the async 202 path is HTTP-only). MCP: memory_list_collections tool and a collection parameter on memory_search (grant-gated, same as the HTTP surface). api._format_hit now unwraps to the innermost user metadata so the full retrieval path exposes the same metadata contract as the BM25 path for batch-ingested rows - collection hits carry file_path/source/ collection_id on both paths. --- taosmd/api.py | 11 ++- taosmd/cli.py | 144 ++++++++++++++++++++++++++++++++++ taosmd/mcp_server.py | 20 +++++ tests/test_collections_cli.py | 104 ++++++++++++++++++++++++ tests/test_collections_mcp.py | 93 ++++++++++++++++++++++ 5 files changed, 370 insertions(+), 2 deletions(-) create mode 100644 tests/test_collections_cli.py create mode 100644 tests/test_collections_mcp.py diff --git a/taosmd/api.py b/taosmd/api.py index bcf3e7ac..5e3d0f2b 100644 --- a/taosmd/api.py +++ b/taosmd/api.py @@ -425,8 +425,15 @@ def _format_hit(hit: dict) -> dict: underlying row's ``created_at`` / ``timestamp`` field. """ md = hit.get("metadata", {}) or {} - inner = md.get("metadata") if isinstance(md, dict) else None - user_md = inner if isinstance(inner, dict) else md + # Unwrap to the innermost user metadata. Depending on the path a hit's + # metadata is nested differently: the BM25 path passes the row metadata + # (one ``metadata`` level for batch rows), while the full retrieval path + # wraps the row metadata in a result envelope (two levels). Descending + # until there is no further ``metadata`` dict gives both paths the same + # user-metadata contract (e.g. collection hits expose ``file_path``). + user_md = md + while isinstance(user_md, dict) and isinstance(user_md.get("metadata"), dict): + user_md = user_md["metadata"] confidence = ( md.get("similarity") diff --git a/taosmd/cli.py b/taosmd/cli.py index a78cb02e..6360c615 100644 --- a/taosmd/cli.py +++ b/taosmd/cli.py @@ -991,6 +991,106 @@ def _projects_cmd(args: argparse.Namespace) -> int: return 0 +def _collections_cmd(args: argparse.Namespace) -> int: + """Handle ``taosmd collections``: collection lifecycle + access control.""" + import asyncio # noqa: PLC0415 + + from . import service # noqa: PLC0415 + from .collections import CollectionNotFoundError # noqa: PLC0415 + + data_dir = args.data_dir + + def _fmt(col: dict) -> str: + stats = col.get("stats") or {} + chunks = stats.get("chunks_ingested", 0) + return ( + f"{col['id']} {col['status']:<9} kind={col['kind']:<8} " + f"name={col['name']!r} files={stats.get('files_total', 0)} " + f"chunks={chunks} grants={len(col.get('grants') or [])}" + ) + + try: + if args.collections_cmd == "list": + cols = asyncio.run( + service.collections_list(project=args.project, data_dir=data_dir) + ) + if not cols: + print("No collections. Create one with `taosmd collections create`.") + return 0 + for col in cols: + print(_fmt(col)) + return 0 + + if args.collections_cmd == "create": + col = asyncio.run( + service.collections_create( + name=args.name, kind=args.kind, source_path=args.source, + embedder=args.embedder, data_dir=data_dir, + ) + ) + print(_fmt(col)) + return 0 + + if args.collections_cmd == "index": + # The CLI runs the walk synchronously (unlike the HTTP 202 path) + # so the caller sees the final stats on exit. + asyncio.run( + service.collections_index_start(args.collection_id, data_dir=data_dir) + ) + stats = asyncio.run( + service.collections_index_run(args.collection_id, data_dir=data_dir) + ) + print( + f"{args.collection_id}: files_indexed={stats['files_indexed']} " + f"unchanged={stats['files_unchanged']} deleted={stats['files_deleted']} " + f"chunks_ingested={stats['chunks_ingested']} " + f"superseded={stats['chunks_superseded']} errors={len(stats['errors'])}" + ) + if stats.get("degraded"): + print( + "warning: embedder unavailable for some chunks; content is " + "archived but not searchable until reconcile", + file=sys.stderr, + ) + return 0 + + if args.collections_cmd in ("link", "unlink"): + fn = ( + service.collections_link + if args.collections_cmd == "link" + else service.collections_unlink + ) + col = asyncio.run( + fn(args.collection_id, args.link_type, args.ext_id, data_dir=data_dir) + ) + links = ", ".join(f"{ln['type']}:{ln['id']}" for ln in col["links"]) or "(none)" + print(f"{col['id']} links: {links}") + return 0 + + if args.collections_cmd == "grant": + col = asyncio.run( + service.collections_grant(args.collection_id, args.agent, data_dir=data_dir) + ) + print(f"{col['id']} grants: {', '.join(col['grants']) or '(none)'}") + return 0 + + if args.collections_cmd == "revoke": + col = asyncio.run( + service.collections_revoke(args.collection_id, args.agent, data_dir=data_dir) + ) + print(f"{col['id']} grants: {', '.join(col['grants']) or '(none)'}") + return 0 + except CollectionNotFoundError as exc: + # KeyError str() wraps the message in quotes; unwrap for readability. + print(f"error: {exc.args[0] if exc.args else exc}", file=sys.stderr) + return 2 + except ValueError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + + return 1 + + def _reconcile_cmd(args: argparse.Namespace) -> int: """Handle ``taosmd reconcile``: compare archive to vector store and repair gaps.""" import asyncio # noqa: PLC0415 @@ -1597,6 +1697,47 @@ def _build_parser() -> argparse.ArgumentParser: help="Project id (from `taosmd projects` or taosmd.get_project_id())", ) + # ----- collections subcommand group --------------------------------- + collections_p = sub.add_parser( + "collections", + help="Manage collections: indexed folders agents can query (grant-gated)", + ) + collections_sub = collections_p.add_subparsers(dest="collections_cmd", required=True) + c_list_p = collections_sub.add_parser("list", help="List collections") + c_list_p.add_argument("--project", help="Only collections linked to this project id") + c_create_p = collections_sub.add_parser( + "create", help="Create a collection (source must be inside an allowed root)" + ) + c_create_p.add_argument("--name", required=True, help="Human-readable name") + c_create_p.add_argument( + "--kind", required=True, choices=["docs", "codebase", "mixed"], + help="Collection kind (Phase 1 indexes docs-shaped files only)", + ) + c_create_p.add_argument("--source", required=True, help="Source folder path") + c_create_p.add_argument( + "--embedder", help="Per-collection embedder id (stored; Phase 1 indexes with the global default)" + ) + c_index_p = collections_sub.add_parser( + "index", help="Index (or re-index) a collection's source folder" + ) + c_index_p.add_argument("collection_id", help="Collection id (col-...)") + c_link_p = collections_sub.add_parser("link", help="Link a collection to a project") + c_link_p.add_argument("collection_id") + c_link_p.add_argument("--type", required=True, choices=["taos", "git"], dest="link_type") + c_link_p.add_argument("--id", required=True, dest="ext_id", help="Project id") + c_unlink_p = collections_sub.add_parser("unlink", help="Remove a project link") + c_unlink_p.add_argument("collection_id") + c_unlink_p.add_argument("--type", required=True, choices=["taos", "git"], dest="link_type") + c_unlink_p.add_argument("--id", required=True, dest="ext_id", help="Project id") + c_grant_p = collections_sub.add_parser( + "grant", help="Grant an agent query access to a collection" + ) + c_grant_p.add_argument("collection_id") + c_grant_p.add_argument("agent") + c_revoke_p = collections_sub.add_parser("revoke", help="Revoke an agent's access") + c_revoke_p.add_argument("collection_id") + c_revoke_p.add_argument("agent") + # ----- tasks subcommand group --------------------------------------- tasks_p = sub.add_parser( "tasks", @@ -1773,6 +1914,9 @@ def main(argv: list[str] | None = None) -> int: if args.cmd in ("projects", "shelves"): return _projects_cmd(args) + if args.cmd == "collections": + return _collections_cmd(args) + if args.cmd == "claims": return _claims_cmd(args) diff --git a/taosmd/mcp_server.py b/taosmd/mcp_server.py index 1f3239e6..d1cea67e 100644 --- a/taosmd/mcp_server.py +++ b/taosmd/mcp_server.py @@ -109,6 +109,7 @@ async def memory_search( limit: int = 5, project: str | None = None, also_include: list[str] | None = None, + collection: str | None = None, ) -> list[dict]: """Search an agent's memory for passages relevant to ``query``. @@ -117,12 +118,17 @@ async def memory_search( to ``agent``. Optional ``project`` scopes the search to one project; ``also_include`` (a list of agent names, only honoured with ``project``) adds those agents' memories within the project (cross-agent reads). + Optional ``collection`` (an id from ``memory_list_collections``) adds + that collection's indexed content when the agent holds a grant for + it; collection hits carry ``file_path`` metadata. """ opts: dict = {} if project: opts["project"] = project if also_include: opts["also_include"] = also_include + if collection: + opts["collections"] = [collection] return await _dispatch( service.search(query, agent=agent, data_dir=data_dir, limit=limit, **opts) ) @@ -137,6 +143,20 @@ async def memory_list_projects() -> list[dict]: """ return await _dispatch(service.list_projects(data_dir=data_dir)) + @mcp.tool() + async def memory_list_collections(project: str | None = None) -> list[dict]: + """List collections: indexed folders agents can query when granted. + + Returns each collection's ``id``, ``name``, ``kind``, ``status``, + ``stats``, ``links``, and ``grants``. Optional ``project`` filters to + collections linked to that project id (taOS or git fingerprint). Pass + a collection ``id`` to ``memory_search``'s ``collection`` parameter + to search its content (requires a grant for the calling agent). + """ + return await _dispatch( + service.collections_list(project=project, data_dir=data_dir) + ) + @mcp.tool() async def memory_list_shelves(project: str) -> list[dict]: """List the agent shelves that have memories within ``project``. diff --git a/tests/test_collections_cli.py b/tests/test_collections_cli.py new file mode 100644 index 00000000..1146a962 --- /dev/null +++ b/tests/test_collections_cli.py @@ -0,0 +1,104 @@ +"""CLI tests for ``taosmd collections`` subcommands. + +Mirrors the shelves/projects CLI conventions: thin argparse layer over the +service wrappers, human-readable one-line-per-row output, exit code 2 with +an error on stderr for validation failures. +""" +from __future__ import annotations + +import pytest + +from taosmd import config +from taosmd.cli import main +from taosmd.collections import CollectionStore + + +@pytest.fixture +def source_dir(tmp_path): + d = tmp_path / "src" + d.mkdir() + (d / "readme.md").write_text("# Hello\n\nDocs body.") + return d + + +@pytest.fixture +def data_dir(tmp_path, monkeypatch, source_dir): + d = tmp_path / "taosmd-data" + d.mkdir() + monkeypatch.setenv("TAOSMD_DATA_DIR", str(d)) + monkeypatch.delenv("TAOSMD_COLLECTIONS_ALLOWED_ROOTS", raising=False) + from taosmd import api as taosmd_api + monkeypatch.setattr(taosmd_api, "_stores_cache", {}) + config.set_collections_allowed_roots([str(source_dir)], data_dir=str(d)) + return str(d) + + +def test_create_and_list(data_dir, source_dir, capsys): + rc = main([ + "collections", "create", "--name", "repo docs", "--kind", "docs", + "--source", str(source_dir), + ]) + assert rc == 0 + out = capsys.readouterr().out + assert "col-" in out + + rc = main(["collections", "list"]) + assert rc == 0 + out = capsys.readouterr().out + assert "repo docs" in out + assert "created" in out + + +def test_create_outside_roots_fails(data_dir, tmp_path, capsys): + outside = tmp_path / "outside" + outside.mkdir() + rc = main([ + "collections", "create", "--name", "x", "--kind", "docs", + "--source", str(outside), + ]) + assert rc == 2 + assert "allowed root" in capsys.readouterr().err + + +def test_link_unlink_grant_revoke(data_dir, source_dir, capsys): + store = CollectionStore(data_dir) + col = store.create(name="a", kind="docs", source_path=str(source_dir)) + cid = col["id"] + + assert main(["collections", "link", cid, "--type", "taos", "--id", "prj-1"]) == 0 + assert "prj-1" in store.get(cid)["links"][0]["id"] + assert main(["collections", "unlink", cid, "--type", "taos", "--id", "prj-1"]) == 0 + assert store.get(cid)["links"] == [] + + assert main(["collections", "grant", cid, "dev"]) == 0 + assert store.get(cid)["grants"] == ["dev"] + assert main(["collections", "revoke", cid, "dev"]) == 0 + assert store.get(cid)["grants"] == [] + capsys.readouterr() + + +def test_index_runs_synchronously(data_dir, source_dir, capsys, monkeypatch): + from taosmd import api as taosmd_api + import asyncio + + stores = asyncio.run(taosmd_api._ensure_stores(data_dir)) + + async def _fake_embed(text, task="search_document"): + h = hash(text) & 0xFFFFFFFF + return [((h >> (i * 4)) & 0xFF) / 255.0 for i in range(8)] + + stores["vector"].embed = _fake_embed + + store = CollectionStore(data_dir) + col = store.create(name="a", kind="docs", source_path=str(source_dir)) + rc = main(["collections", "index", col["id"]]) + assert rc == 0 + out = capsys.readouterr().out + assert "files_indexed=1" in out + assert store.get(col["id"])["status"] == "ready" + + +def test_unknown_collection_exits_2(data_dir, capsys): + rc = main(["collections", "grant", "col-000000000000", "dev"]) + assert rc == 2 + assert "not found" in capsys.readouterr().err diff --git a/tests/test_collections_mcp.py b/tests/test_collections_mcp.py new file mode 100644 index 00000000..05857b5a --- /dev/null +++ b/tests/test_collections_mcp.py @@ -0,0 +1,93 @@ +"""MCP surface tests for collections: memory_list_collections + the +``collection`` parameter on memory_search. Skips when the optional ``mcp`` +SDK is absent, matching tests/test_mcp_server.py.""" +from __future__ import annotations + +import asyncio + +import pytest + +pytest.importorskip("mcp") + +from taosmd import api as taosmd_api +from taosmd import config as taosmd_config +from taosmd import mcp_server +from taosmd.collections import CollectionStore, ingest_folder + + +def _patch_embedder(stores: dict) -> None: + vmem = stores["vector"] + + async def _fake_embed(text: str, task: str = "search_document") -> list[float]: + h = hash(text) & 0xFFFFFFFF + return [((h >> (i * 4)) & 0xFF) / 255.0 for i in range(8)] + + vmem.embed = _fake_embed # type: ignore[assignment] + + +@pytest.fixture +def env(tmp_path, monkeypatch): + source = tmp_path / "src" + source.mkdir() + (source / "readme.md").write_text("# Widget\n\nThe widget frobnicates the sprocket.") + data_dir = tmp_path / "taosmd-data" + data_dir.mkdir() + monkeypatch.setenv("TAOSMD_DATA_DIR", str(data_dir)) + monkeypatch.delenv("TAOSMD_COLLECTIONS_ALLOWED_ROOTS", raising=False) + monkeypatch.setattr(taosmd_api, "_stores_cache", {}) + taosmd_config.set_collections_allowed_roots([str(source)], data_dir=str(data_dir)) + + mcp = mcp_server.build_server(data_dir=str(data_dir)) + loop = mcp._taosmd_service_loop + stores = loop.run(taosmd_api._ensure_stores(str(data_dir))) + _patch_embedder(stores) + try: + yield mcp, loop, str(data_dir), source + finally: + for s in list(taosmd_api._stores_cache.values()): + for store in (s.get("archive"), s.get("vector"), s.get("kg")): + if store and hasattr(store, "close"): + try: + loop.run(store.close()) + except Exception: + pass + loop.close() + + +async def _call(mcp, name: str, args: dict): + result = await mcp.call_tool(name, args) + if isinstance(result, tuple): + structured = result[1] + if isinstance(structured, dict) and "result" in structured: + return structured["result"] + import json + return json.loads(result[0].text) + + +def test_tools_registered(env): + mcp, *_ = env + names = {t.name for t in asyncio.run(mcp.list_tools())} + assert "memory_list_collections" in names + + +def test_list_collections_and_scoped_search(env): + mcp, loop, data_dir, source = env + store = CollectionStore(data_dir) + col = store.create(name="docs", kind="docs", source_path=str(source)) + store.grant(col["id"], "dev") + loop.run(ingest_folder(col["id"], data_dir=data_dir)) + + cols = asyncio.run(_call(mcp, "memory_list_collections", {})) + assert [c["id"] for c in cols] == [col["id"]] + assert cols[0]["status"] == "ready" + + hits = asyncio.run(_call(mcp, "memory_search", { + "query": "widget frobnicates", "agent": "dev", "collection": col["id"], + })) + assert any(h["metadata"].get("file_path") == "readme.md" for h in hits) + + # No grant, no hits from the collection. + hits = asyncio.run(_call(mcp, "memory_search", { + "query": "widget frobnicates", "agent": "stranger", "collection": col["id"], + })) + assert not any(h["metadata"].get("collection_id") for h in hits) From d1a08275ca9d32d31f78c6196272af24be94a714 Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 00:12:14 +0100 Subject: [PATCH 06/16] docs(collections): user page, changelog entry, spec decisions appendix docs/collections.md covers enabling allowed_roots, the lifecycle, query semantics with grants, the HTTP surface split (admin vs data plane), and the zero-loss guarantees. The spec gains section 9 recording the Jul 19 decisions: no tree-sitter/new deps, admin gating on create+index, ship Phase 1 alone, per-collection embedder mechanism now (global default used), taOS owns the panel, plus the typed-links/grants verdict. --- CHANGELOG.md | 2 + docs/collections.md | 107 ++++++++++++++++++ .../codebase-indexing-collections-design.md | 14 +++ 3 files changed, 123 insertions(+) create mode 100644 docs/collections.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a3787bc..0997bd77 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Unreleased +Collections, Phase 1 (docs MVP, per `docs/specs/codebase-indexing-collections-design.md`): named containers of content indexed from a folder, queryable by granted agents alongside conversation memory. A collection is a first-class row (`created -> indexing -> ready | error`, plus reversible `archived`) with typed project links (`{type: taos|git, id}`, metadata only, never access-granting) and per-agent grants (`(canonical_id, scope='collection', collection_id)` unique rows, enforced at search time). Indexing wires the previously-unwired loader framework into a real ingest path: a gitignore-aware walker (stdlib rules; VCS/dependency/hidden dirs, binaries, oversized files, and symlink escapes skipped) feeds files a registered loader claims through a zero-dep paragraph chunker into `ingest_batch` under the collection's own agent namespace, with per-chunk content-hash ids so re-index dedups unchanged files; changed and deleted files have their old rows superseded (`valid_to` + marker), never deleted. The feature is off by default: the new `collections.allowed_roots` config list (or `TAOSMD_COLLECTIONS_ALLOWED_ROOTS`) must name the directories collections may index, and `source_path` is containment-checked (`resolve_within`) at create and at every index. Surfaces: HTTP (`POST /collections` and `POST /collections/{id}/index` admin-gated with async 202+poll indexing, `DELETE /collections/{id}` archives; list/get/link/unlink/grants on the data plane; `collection`/`collections`/`collections_only` on search), CLI (`taosmd collections list|create|index|link|unlink|grant|revoke`), and MCP (`memory_list_collections`, `collection` on `memory_search`). Collection hits carry `collection_id`/`file_path`/`source` metadata. A per-collection `embedder` field is stored and returned now (the mechanism for the code-embedder bake-off); Phase 1 always indexes with the global default. `benchmarks/collections_eval.py` pre-registers the file-level Recall@5 eval over the repo's own docs. + Admin token separation (#154, phase 1). Admin operations are now gated by a dedicated `admin_token`, distinct from the data-plane `server_token`. Previously the server token gated every data and A2A endpoint AND the admin surface, so on a token-less deployment the only way to authorize an admin op was to set a server token, which locked out every agent on the data plane for the duration of the admin window (this hit the Pi bus in production for about three minutes during a channel cleanup). Now the admin write routes (`POST /shelves`, `POST /shelves/{id}/archive|unarchive`, `POST /a2a/admin/delete-channel|rename-channel|supersede-message`) are exempt from the data-plane token gate and enforce the admin token themselves. Resolution prefers `admin_token` and falls back to `server_token`: existing token-secured installs keep working unchanged; setting only `admin_token` gates admin while leaving data and A2A endpoints open; with both set the data plane is gated by `server_token` and admin by `admin_token`, so a caller holding only the server token cannot run admin ops; with neither set the admin surface still fails closed (403). Configure via `admin_token` in config, the `TAOSMD_ADMIN_TOKEN` env var, or `taosmd config set-admin-token`. Phase 2 (isolating admin operations from the single service loop so a slow admin op cannot stall data reads/writes) is not part of this change and is tracked separately. Scoping fix: `reindex` now carries the `project` scope and provenance through when it rebuilds an agent's vector rows from the archive. The rebuild loop stamped only `{"agent": agent}` plus role/timestamp, dropping the `project` tag, the `archive_span_id`, and the nested user metadata (`source_id`/`forget_after`). Losing the project tag was a cross-project scope leak: a project-scoped row came back project-untagged but still agent-tagged, so after a reindex it surfaced in a DIFFERENT project's search for the same agent. `reindex` now mirrors `reconcile`'s metadata reconstruction exactly, so a reindexed row is indistinguishable from a reconcile-repaired one in scope and provenance: it stays visible to its own project's search, stays invisible to other projects, keeps its archive-span linkage for the claims gate, and keeps its `source_id` so a re-POST of the same batch still dedupes. Zero-loss was never at risk (the archive is untouched); this is a scoping-correctness repair. diff --git a/docs/collections.md b/docs/collections.md new file mode 100644 index 00000000..58bcd12b --- /dev/null +++ b/docs/collections.md @@ -0,0 +1,107 @@ +# Collections: feed folders to your agents + +A collection is a named container of content indexed from one folder. Point +it at a repo's `docs/` (or any documentation folder inside an allowed root), +index it, grant agents access, and they can query it alongside their +conversation memory. Phase 1 indexes docs-shaped files (md, txt, markdown, +rst, plus anything another registered loader claims, such as `*.chat.json`); +the code path is Phase 2. + +## Enable the feature (off by default) + +Collections read the server's filesystem, so they are disabled until an +operator opts in by listing the directories collections may be created +under: + +```json +// ~/.taosmd/config.json +{ "collections": { "allowed_roots": ["/srv/docs", "/home/jay/repos"] } } +``` + +Or via the environment: `TAOSMD_COLLECTIONS_ALLOWED_ROOTS=/srv/docs:/home/jay/repos`. +A collection's `source_path` must resolve inside one of these roots +(symlink escapes are rejected), checked at create time and again at every +index. An empty list means collections are off. + +## Lifecycle + +```bash +# create (admin operation over HTTP; local CLI works directly) +taosmd collections create --name "repo docs" --kind docs --source /srv/docs/myrepo + +# index (walks the folder, chunks, embeds; re-run any time) +taosmd collections index col-ab12cd34ef56 + +# give an agent query access +taosmd collections grant col-ab12cd34ef56 my-agent + +# attach to a project for discovery (taOS prj-* id or git fingerprint) +taosmd collections link col-ab12cd34ef56 --type git --id abc123def456 + +taosmd collections list +``` + +Statuses: `created -> indexing -> ready | error` (and `archived` after a +delete). Re-indexing is incremental: unchanged files are skipped by content +hash, changed and deleted files have their old rows superseded (hidden from +recall, never destroyed; the archive keeps every version). + +The walker respects `.gitignore` files (simplified rules), skips VCS and +dependency directories, hidden directories, binaries, and oversized files. + +## Querying + +Search merges granted collections in with conversation memory: + +```bash +curl -s localhost:7900/search -d '{ + "query": "how do I configure the widget?", + "agent": "my-agent", + "collections": ["col-ab12cd34ef56"], + "limit": 5 +}' +``` + +- `collections` (list) or `collection` (single id) adds collection content. +- `collections_only: true` restricts the search to the collections. +- Grants are enforced per requesting agent: a collection the agent holds no + grant for contributes nothing (and its existence is not revealed). +- Collection hits carry `collection_id`, `file_path` (relative to the + source root), and `source: "collection"` in their metadata. + +Over MCP: `memory_list_collections` lists them; `memory_search` takes a +`collection` parameter. + +## HTTP surface + +Data plane (bearer token when one is configured): + +``` +GET /collections [?project=] +GET /collections/{id} +POST /collections/{id}/link {"type": "taos"|"git", "id": "..."} +POST /collections/{id}/unlink same body +POST /collections/{id}/grants {"agent": "..."} +DELETE /collections/{id}/grants/{agent} +POST /search with collection/collections/collections_only +``` + +Admin (dedicated admin token, fail-closed): + +``` +POST /collections {"name", "kind", "source_path", "embedder"?} +POST /collections/{id}/index -> 202; poll GET /collections/{id} +DELETE /collections/{id} -> archive (reversible) +``` + +The optional `embedder` field is stored and returned per collection (the +per-collection embedder mechanism); Phase 1 always indexes with the global +default embedder. + +## Zero-loss guarantees + +Delete archives, it never destroys: the collection row, its vector rows, +and its archive entries all stay on disk; archived collections simply stop +contributing to search. Re-index supersedes replaced rows with the same +`valid_to` machinery corrections use. Destruction remains exclusive to the +wipe surface. diff --git a/docs/specs/codebase-indexing-collections-design.md b/docs/specs/codebase-indexing-collections-design.md index b44b6c0c..f001226d 100644 --- a/docs/specs/codebase-indexing-collections-design.md +++ b/docs/specs/codebase-indexing-collections-design.md @@ -134,3 +134,17 @@ Retrieval quality on a real repo, judged-free, before any default surface ships. 3. **Ship Phase 1 alone first?** Docs collections are useful on their own (index a project's docs folder, grant it to the agents). Shipping it before the code path exists gets real usage feedback on the container model early. My lean is yes. 4. **Code-embedder bake-off timing.** Follow-up experiment after Phase 2 lands, or run it in parallel with Phase 2 so the per-collection embedder mechanism is designed in from the start? 5. **taOS UI ownership.** Section 4 puts the Collections panel in the taOS Projects app and keeps only a minimal view in the standalone dashboard. If you want the standalone dashboard to be feature-complete for non-taOS users, that roughly doubles the UI work and should be scoped now, not discovered later. + +## 9. Decisions (2026-07-19) + +Settled with Jay and @taOS-dev; these bind the Phase 1 build. + +1. **tree-sitter: no.** No new dependencies anywhere in this workstream. Phase 1 is docs-only (`DocLoader` md/txt/markdown/rst plus whatever the other registered loaders claim); chunking is the zero-dep paragraph packer with an `# upgrade-path:` comment. The Phase 2 code splitter starts zero-dep too, and only the pre-registered eval can argue for the dep later. +2. **Admin gating: yes, on create and index.** `POST /collections` and `POST /collections/{id}/index` require the admin token (fail closed, same as shelves). `DELETE /collections/{id}` (archive) is admin too. The data plane (list, get, link, unlink, grants, query) uses the normal server token plus grants. +3. **Ship Phase 1 alone: yes.** Docs collections ship first to get real usage on the container model before the code path exists. +4. **Per-collection embedder: mechanism now, minimal.** Collections carry an `embedder` field (default = global embedder), stored at create, returned in GET, and read by the index path. Phase 1 always indexes with the global default; the field exists so the code-embedder bake-off can use it without a schema change. +5. **taOS owns the Collections panel.** The standalone dashboard gets at most a minimal read view later; nothing UI ships in Phase 1. + +**Links and grants verdict.** Link rows are typed `{type: "taos" | "git", id}`, multiple per collection, metadata only, and strictly non-transitive: a project link never grants query access. Grants are their own table, `(canonical_id, scope, collection_id)` UNIQUE together with `scope='collection'`, stored in taosmd, enforced on search per requesting agent. + +**Safety and zero-loss (confirmed).** `collections.allowed_roots` defaults to EMPTY (feature off); `source_path` must `resolve_within` an allowed root at create and at every index, symlink escapes rejected, `check_size` per file. DELETE archives (reversible status change); re-index supersedes changed-file rows via the existing supersede machinery; destruction stays exclusive to wipe. From 42ab4ad26df0943b38c7e3866e0096180dfdc101 Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 00:14:02 +0100 Subject: [PATCH 07/16] bench(collections): pre-registered file-level Recall@5 eval over repo docs 20 questions with gold file paths (written before the index was built) over the repo's own docs/ folder, per section 5 of the design spec. Metric is judge-free file-level Recall@5. Two arms: semantic (real ONNX embedder; the arm the kill bar judges, run on the bench host) and a lexical fallback for hosts with no local ONNX model, which stubs the embedder so rows land and retrieves through the engine's BM25-only mode, exercising the full walk/chunk/ingest/grant/scope path. auto picks per host. Lexical smoke on this host: 47 files / 493 chunks indexed in 3.1s, Recall@5 = 19/20 = 0.950 (one miss: the LoCoMo-scorecards question, cross-file term collision on 'LoCoMo'). The question set is a tiny pre-registration file, exempted from the benchmarks/data size gitignore like the README pointer. --- .gitignore | 2 + benchmarks/collections_eval.py | 128 ++++++++++++++++++ .../data/collections_eval_questions.json | 25 ++++ 3 files changed, 155 insertions(+) create mode 100644 benchmarks/collections_eval.py create mode 100644 benchmarks/data/collections_eval_questions.json diff --git a/.gitignore b/.gitignore index 8dc47b7c..cc3f8685 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,8 @@ data/ !benchmarks/data/ benchmarks/data/* !benchmarks/data/README.md +# tiny pre-registered eval sets are tracked (not datasets) +!benchmarks/data/collections_eval_questions.json *.rknn *.rkllm models/minilm-onnx/model.onnx diff --git a/benchmarks/collections_eval.py b/benchmarks/collections_eval.py new file mode 100644 index 00000000..a0579188 --- /dev/null +++ b/benchmarks/collections_eval.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +"""Phase 1 collections eval: file-level Recall@5 over the repo's own docs/. + +Pre-registered per the design spec (section 5): index the taosmd repo's +docs/ folder into a temporary collection and answer the 20-question set in +``benchmarks/data/collections_eval_questions.json`` (written before the +index was built, gold labels are file paths). A question scores 1 when any +of the top-5 results' ``file_path`` metadata matches a gold file. No LLM +judge anywhere in the loop. + +Modes +----- +- ``semantic``: the real path; requires a local ONNX embedding model. This + is the number that counts against the section-5 kill bar; run it on the + bench host where the ONNX embedder is installed. +- ``lexical``: the fallback smoke for hosts with no local ONNX model. The + embedder is stubbed with a deterministic hash vector purely so rows land + in the vector store; retrieval then uses the engine's BM25-only mode + (``mode="bm25"``), which never touches the stubbed vectors. This + exercises the full production walk/chunk/ingest/grant/scope path with a + lexical ranker, and validates the walker + container plumbing + independently of the embedder (the spec's stated purpose for the Phase 1 + eval). +- ``auto`` (default): semantic when an ONNX model is present, else lexical. + +Usage:: + + python3 benchmarks/collections_eval.py [--mode auto|semantic|lexical] + [--docs-dir docs] [--k 5] [--questions benchmarks/data/...json] +""" +from __future__ import annotations + +import argparse +import asyncio +import json +import sys +import tempfile +import time +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT)) + +AGENT = "eval-agent" + + +async def _hash_embed(text: str, task: str = "search_document") -> list[float]: + """Deterministic stub (same as the test suite's). BM25 mode ignores it.""" + h = hash(text) & 0xFFFFFFFF + return [((h >> (i * 4)) & 0xFF) / 255.0 for i in range(8)] + + +async def run(docs_dir: Path, questions_path: Path, mode: str, k: int) -> int: + from taosmd import api, config + from taosmd.collections import CollectionStore, ingest_folder + + spec = json.loads(questions_path.read_text()) + questions = spec["questions"] + + with tempfile.TemporaryDirectory(prefix="taosmd-collections-eval-") as tmp: + data_dir = str(Path(tmp) / "data") + Path(data_dir).mkdir() + config.set_collections_allowed_roots([str(docs_dir)], data_dir=data_dir) + + stores = await api._ensure_stores(data_dir) + if mode == "auto": + has_onnx = api._resolve_onnx_path(data_dir) is not None + mode = "semantic" if has_onnx else "lexical" + if mode == "lexical": + print("auto: no local ONNX embedding model found; using the " + "lexical (BM25) fallback. Run the semantic arm on the " + "bench host for the number that counts.") + if mode == "lexical": + stores["vector"].embed = _hash_embed # rows must land; BM25 ignores vectors + search_opts = {"mode": "bm25"} if mode == "lexical" else {} + + store = CollectionStore(data_dir) + col = store.create(name="repo-docs", kind="docs", source_path=str(docs_dir)) + store.grant(col["id"], AGENT) + + t0 = time.time() + stats = await ingest_folder(col["id"], data_dir=data_dir) + dt = time.time() - t0 + print(f"indexed {stats['files_indexed']} files / " + f"{stats['chunks_ingested']} chunks in {dt:.1f}s " + f"(skipped: unclaimed={stats['skipped_unclaimed']} " + f"ignored={stats['skipped_ignored']} binary={stats['skipped_binary']})") + if stats.get("degraded"): + print("ERROR: embedder unavailable; aborting", file=sys.stderr) + return 2 + + hits_n = 0 + for q in questions: + results = await api.search( + q["question"], agent=AGENT, limit=k, + collections=[col["id"]], collections_only=True, + data_dir=data_dir, **search_opts, + ) + got_files = [h["metadata"].get("file_path") for h in results] + hit = any(f in q["answer_files"] for f in got_files) + hits_n += hit + mark = "HIT " if hit else "MISS" + print(f" [{mark}] {q['question'][:70]:<70} -> {got_files[:3]}") + + recall = hits_n / len(questions) + print(f"\nmode={mode} file-level Recall@{k}: {hits_n}/{len(questions)} = {recall:.3f}") + if mode == "lexical": + print("note: lexical fallback smoke. The pre-registered kill bar " + "(Recall@5 >= 0.8) is judged on the semantic arm with the " + "ONNX embedder on the bench host.") + store.close() + return 0 + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--docs-dir", default=str(REPO_ROOT / "docs")) + p.add_argument("--questions", + default=str(REPO_ROOT / "benchmarks" / "data" / "collections_eval_questions.json")) + p.add_argument("--mode", choices=["auto", "semantic", "lexical"], default="auto") + p.add_argument("--k", type=int, default=5) + args = p.parse_args() + return asyncio.run(run(Path(args.docs_dir).resolve(), + Path(args.questions), args.mode, args.k)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/data/collections_eval_questions.json b/benchmarks/data/collections_eval_questions.json new file mode 100644 index 00000000..74c9cdaf --- /dev/null +++ b/benchmarks/data/collections_eval_questions.json @@ -0,0 +1,25 @@ +{ + "description": "Phase 1 collections eval: 20 questions over the taosmd repo's own docs/ folder, written before the index was built. Gold labels are file paths relative to docs/. Metric: file-level Recall@5 (a question scores 1 if any of the top-5 results' file_path is in answer_files). No LLM judge.", + "questions": [ + {"question": "How do I pause a long benchmark run, reboot the machine, and resume it later without repeating work?", "answer_files": ["bench-pause-resume.md"]}, + {"question": "Can several agents share a single taosmd install while keeping their memories isolated?", "answer_files": ["multi-agent.md"]}, + {"question": "How do I keep the taosmd HTTP server running in the background so it survives logout and starts at login?", "answer_files": ["serve-service.md"]}, + {"question": "Which memory controls should the taOS settings UI expose and what scope does each control have?", "answer_files": ["INTEGRATION-memory-config.md"]}, + {"question": "How do agent trace events and memory operations map onto OpenTelemetry GenAI spans and attributes?", "answer_files": ["otel-genai-mapping.md"]}, + {"question": "What prompt do I paste into a tool-capable agent to install taosmd for me?", "answer_files": ["INSTALL-AGENT-PROMPT.md"]}, + {"question": "What Recall@5 and judge scores has taosmd published on LongMemEval?", "answer_files": ["benchmarks.md", "research-report.md"]}, + {"question": "Where are negative experiment results and methodology notes recorded?", "answer_files": ["research-report.md"]}, + {"question": "What is the design for feeding codebases and folders into named collections agents can query?", "answer_files": ["specs/codebase-indexing-collections-design.md", "collections.md"]}, + {"question": "What is the proposal for routing retrieval by semantic category?", "answer_files": ["specs/category-routed-retrieval-design.md"]}, + {"question": "How does the nightly session catalog pipeline enrich and crystallize conversations?", "answer_files": ["specs/session-catalog-pipeline.md", "plans/session-catalog-pipeline.md"]}, + {"question": "What is the librarian and how does it decide what to enrich at ingest time?", "answer_files": ["specs/2026-04-15-librarian-design.md"]}, + {"question": "What was the plan for running retrieval sources in parallel and reranking their results?", "answer_files": ["specs/parallel-retrieval-reranking.md", "plans/parallel-retrieval-reranking.md"]}, + {"question": "What is the spec for the memory management app?", "answer_files": ["specs/memory-management-app.md", "plans/memory-management-app.md"]}, + {"question": "How are tasks, edges, and the ready queue designed in the task graph?", "answer_files": ["superpowers/specs/2026-06-10-task-graph-design.md"]}, + {"question": "What is the long-term vision for the memory cockpit dashboard?", "answer_files": ["superpowers/specs/2026-06-21-memory-cockpit-vision.md"]}, + {"question": "How does the smart installer pick what to install for a machine?", "answer_files": ["superpowers/specs/2026-06-16-smart-installer-design.md", "superpowers/plans/2026-06-17-smart-installer.md"]}, + {"question": "What were the LoCoMo scorecard results per conversation?", "answer_files": ["specs/2026-04-19-locomo-scorecards.md"]}, + {"question": "What went wrong with the cross-encoder model path and how was it fixed?", "answer_files": ["agent-jobs/job-002-cross-encoder-path-fix.md"]}, + {"question": "How do grants control which agents may query an indexed folder collection?", "answer_files": ["collections.md", "specs/codebase-indexing-collections-design.md"]} + ] +} From 0a24e8c7197884dff9d98a7b48beb664076e6578 Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 05:19:41 +0100 Subject: [PATCH 08/16] fix(collections): preserve claims-gate metadata through _format_hit unwrap The unwrap-to-innermost loop gave every path the same user-metadata contract but stripped the provenance envelope on the way down: batch rows are double-wrapped on the semantic path (retrieval envelope -> row meta -> user metadata), so archive_span_id vanished from the formatted hit and the prefer_verified claims gate went blind for those rows (a contradicted-claim row survived recall). Capture archive_span_id/agent/project while descending and re-attach them to a copy of the innermost metadata (user keys never clobbered), so the gate keeps its inputs and collection hits still expose file_path. Regression tests cover the semantic and bm25 paths plus the contradicted-row drop. --- taosmd/api.py | 15 +++++++ tests/test_api.py | 105 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/taosmd/api.py b/taosmd/api.py index 5e3d0f2b..4ad44ea7 100644 --- a/taosmd/api.py +++ b/taosmd/api.py @@ -431,9 +431,24 @@ def _format_hit(hit: dict) -> dict: # wraps the row metadata in a result envelope (two levels). Descending # until there is no further ``metadata`` dict gives both paths the same # user-metadata contract (e.g. collection hits expose ``file_path``). + # + # Provenance fields live on the envelope levels, not in the user metadata, + # so they are captured on the way down and re-attached to the formatted + # hit: the prefer_verified claims gate reads ``archive_span_id`` off the + # formatted metadata, and stripping it would blind the gate for batch + # rows (a contradicted-claim row would survive recall). Innermost + # envelope wins; an explicit user key of the same name is never clobbered. + preserved: dict = {} user_md = md while isinstance(user_md, dict) and isinstance(user_md.get("metadata"), dict): + for key in ("archive_span_id", "agent", "project"): + if key in user_md: + preserved[key] = user_md[key] user_md = user_md["metadata"] + if isinstance(user_md, dict): + user_md = dict(user_md) # copy: never mutate the stored row metadata + for key, value in preserved.items(): + user_md.setdefault(key, value) confidence = ( md.get("similarity") diff --git a/tests/test_api.py b/tests/test_api.py index e693cf7b..fba710f9 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -392,6 +392,111 @@ def test_bm25_python_rank_orders_by_relevance(): assert ranked[-1][1] == 0.0, "no-overlap doc should score zero" +# --------------------------------------------------------------------------- +# Claims-gate provenance must survive _format_hit's metadata unwrap +# --------------------------------------------------------------------------- + +def _batch_row_span(stores, text: str) -> int: + """Read a batch row's archive_span_id straight off the raw vector row.""" + rows = stores["vector"]._conn.execute( + "SELECT text, metadata_json FROM vector_memory" + ).fetchall() + for r in rows: + if r["text"] == text: + meta = json.loads(r["metadata_json"]) + span = meta.get("archive_span_id") + assert isinstance(span, int), "batch row lost its archive_span_id at write time" + return span + raise AssertionError(f"no vector row stored for {text!r}") + + +def test_semantic_search_batch_row_keeps_archive_span_id(isolated_data_dir): + """Regression: _format_hit's unwrap-to-innermost loop must not strip the + provenance envelope. Batch rows are double-wrapped on the semantic path + (retrieval envelope -> row meta -> user metadata); descending all the way + down dropped archive_span_id, blinding the prefer_verified claims gate. + The formatted hit must expose BOTH the user metadata (e.g. file_path for + collection hits) and the archive span the gate reads.""" + stores = _setup_stores(isolated_data_dir) + asyncio.run(taosmd.ingest_batch( + [{"text": "The rack in bay four is painted teal.", + "id": "hash-rack", + "metadata": {"file_path": "docs/rack.md"}}], + agent="batch-agent", data_dir=str(isolated_data_dir), + )) + span = _batch_row_span(stores, "The rack in bay four is painted teal.") + + hits = asyncio.run(taosmd.search( + "The rack in bay four is painted teal.", + agent="batch-agent", + prefer_verified="off", + data_dir=str(isolated_data_dir), + )) + assert hits, "expected a semantic hit for an exact-content query" + top = hits[0] + assert top["metadata"].get("file_path") == "docs/rack.md" # client-facing win stays + assert top["metadata"].get("archive_span_id") == span # gate input preserved + + +def test_bm25_search_batch_row_keeps_archive_span_id(isolated_data_dir): + """Same contract on the BM25 path: row meta is one level shallower there, + but the formatted hit must still carry the provenance span.""" + stores = _setup_stores(isolated_data_dir) + asyncio.run(taosmd.ingest_batch( + [{"text": "The quarterly review moved to Friday morning.", + "id": "hash-review", + "metadata": {"file_path": "notes/review.md"}}], + agent="batch-agent", data_dir=str(isolated_data_dir), + )) + span = _batch_row_span(stores, "The quarterly review moved to Friday morning.") + hits = asyncio.run(taosmd.search( + "quarterly review Friday", + agent="batch-agent", + mode="bm25", + prefer_verified="off", + data_dir=str(isolated_data_dir), + )) + assert hits + assert hits[0]["metadata"].get("file_path") == "notes/review.md" + assert hits[0]["metadata"].get("archive_span_id") == span + + +def test_prefer_verified_drops_contradicted_batch_row(isolated_data_dir): + """Regression: with a batch row's backing claim contradicted, the + prefer_verified gate must drop the row from semantic recall. This is the + master behaviour the metadata unwrap regressed (gate went blind because + the formatted hit no longer carried archive_span_id).""" + stores = _setup_stores(isolated_data_dir) + asyncio.run(taosmd.ingest_batch( + [{"text": "The rack in bay four is painted teal.", + "id": "hash-rack", + "metadata": {"file_path": "docs/rack.md"}}], + agent="batch-agent", data_dir=str(isolated_data_dir), + )) + span = _batch_row_span(stores, "The rack in bay four is painted teal.") + + # Sanity: without the gate the row is recalled. + ungated = asyncio.run(taosmd.search( + "The rack in bay four is painted teal.", + agent="batch-agent", prefer_verified="off", + data_dir=str(isolated_data_dir), + )) + assert any("teal" in h["text"] for h in ungated) + + # Contradict the claim backing that span; the gate must now drop the row. + cs = stores["claims"] + cid = asyncio.run(cs.add_claim("rack colour", [span], source_extractor="test")) + asyncio.run(cs.set_status(cid, "contradicted", verifier_model="m", now=1.0)) + gated = asyncio.run(taosmd.search( + "The rack in bay four is painted teal.", + agent="batch-agent", prefer_verified="prefer_verified", + data_dir=str(isolated_data_dir), + )) + assert not any("teal" in h["text"] for h in gated), ( + "contradicted-claim row survived: the claims gate is blind to batch rows" + ) + + def test_search_prefer_verified_resolves_from_config(): """Provable memory ships on by default: search()'s prefer_verified param is a sentinel (None) that resolves from the persisted controls at call time, and From a30ea3b1f70e165ff1156f405d18e38e6536b84e Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 05:20:19 +0100 Subject: [PATCH 09/16] fix(collections): run the index walk and chunking off the service loop The 202+poll index contract still blocked the single service loop: collect_files did the full os.walk, per-file size checks, and the 1KB null-byte sniff (opening every candidate file) synchronously on the loop thread, so /search and /ingest stalled for the duration of an index. Run the walk and the per-file hash+chunk step through asyncio.to_thread; the store writes stay on the loop where the thread-affine sqlite connections live. Test pins the walk to a worker thread and that the pipeline still lands created -> indexing -> ready. --- taosmd/collections.py | 32 ++++++++++++++++++++++++---- tests/test_collections_ingest.py | 36 ++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 4 deletions(-) diff --git a/taosmd/collections.py b/taosmd/collections.py index 72538154..6561aa05 100644 --- a/taosmd/collections.py +++ b/taosmd/collections.py @@ -32,6 +32,7 @@ from __future__ import annotations +import asyncio import fnmatch import hashlib import json @@ -664,6 +665,20 @@ def _supersede_collection_rows(vmem, collection_id: str, file_path: str) -> int: return superseded +def _hash_and_chunk( + text: str, chunk_chars: int, prior_hash: str | None +) -> tuple[str, list[str] | None]: + """CPU-bound half of the per-file ingest step, run off the event loop. + + Returns ``(file_hash, chunks)``; ``chunks`` is ``None`` when the hash + matches ``prior_hash`` (unchanged file, nothing to chunk). + """ + file_hash = hashlib.sha256(text.encode("utf-8")).hexdigest() + if prior_hash == file_hash: + return file_hash, None + return file_hash, chunk_text(text, max_chars=chunk_chars) + + async def ingest_folder( collection_id: str, *, @@ -704,7 +719,13 @@ async def ingest_folder( try: source_root = store.resolve_source_path(col["source_path"]) prior = store.file_states(collection_id) - files, skips = collect_files(source_root, max_file_bytes=max_file_bytes) + # The walk (os.walk + per-file stat + null-byte sniff over every + # candidate) is blocking filesystem work; run it on a worker thread so + # the 202+poll contract holds and the single service loop keeps + # serving /search and /ingest while a collection indexes. + files, skips = await asyncio.to_thread( + collect_files, source_root, max_file_bytes=max_file_bytes, + ) stores = await _api._ensure_stores(data_dir) vmem = stores["vector"] @@ -731,13 +752,16 @@ async def ingest_folder( text = blob.raw_text or getattr(blob, "content", "") or "" if not text.strip(): continue - file_hash = hashlib.sha256(text.encode("utf-8")).hexdigest() - if prior.get(rel) == file_hash: + # Hashing + chunking are CPU-bound; off-loop like the walk above. + file_hash, chunks = await asyncio.to_thread( + _hash_and_chunk, text, chunk_chars, prior.get(rel) + ) + if chunks is None: unchanged += 1 continue if rel in prior: changed.append(rel) - for i, chunk in enumerate(chunk_text(text, max_chars=chunk_chars)): + for i, chunk in enumerate(chunks): chunk_id = hashlib.sha256( f"{collection_id}:{rel}:{file_hash}:{i}".encode("utf-8") ).hexdigest() diff --git a/tests/test_collections_ingest.py b/tests/test_collections_ingest.py index f66b762d..28e0ed55 100644 --- a/tests/test_collections_ingest.py +++ b/tests/test_collections_ingest.py @@ -271,6 +271,42 @@ def test_reindex_deleted_file_supersedes_rows(data_dir, source_dir): assert "guide.txt" not in store.file_states(col["id"]) +def test_ingest_folder_walks_off_the_event_loop(data_dir, source_dir, monkeypatch): + """The 202+poll contract promises a responsive server during an index: + the filesystem walk (os.walk + per-file stat + null-byte sniff) must run + on a worker thread, not on the single service loop thread where it would + block /search and /ingest for the duration.""" + import threading + + from taosmd import collections as collections_mod + + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + + seen: dict = {} + real_collect = collections_mod.collect_files + + def _spy(*args, **kwargs): + seen["thread"] = threading.current_thread() + return real_collect(*args, **kwargs) + + monkeypatch.setattr(collections_mod, "collect_files", _spy) + + async def _run(): + seen["loop_thread"] = threading.current_thread() + return await ingest_folder(col["id"], data_dir=data_dir) + + stats = asyncio.run(_run()) + # The pipeline still completes correctly end to end. + assert stats["files_indexed"] == 4 + assert store.get(col["id"])["status"] == "ready" + # And the walk ran off the loop thread. + assert "thread" in seen + assert seen["thread"] is not seen["loop_thread"], ( + "collect_files ran on the event loop thread; the async index blocks the server" + ) + + def test_ingest_folder_archived_collection_refused(data_dir, source_dir): store, col = _make_collection(data_dir, source_dir) store.archive(col["id"]) From 93574a37f3a3a294b5971d8ff78502bdf510208a Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 05:20:49 +0100 Subject: [PATCH 10/16] fix(collections): reject concurrent index starts with 409 collections_index_start set status=indexing without checking it, so a second POST /collections/{id}/index while one index was running would double-walk the same tree into the batch dedup. Raise CollectionBusyError when the status is already indexing; the endpoint maps it to 409 with a poll hint. ready and error both re-arm the start, so retries after completion or failure keep working. --- taosmd/collections.py | 5 +++++ taosmd/http_server.py | 12 ++++++++++-- taosmd/service.py | 9 ++++++++- tests/test_collections_http.py | 26 ++++++++++++++++++++++++++ tests/test_collections_ingest.py | 24 ++++++++++++++++++++++++ 5 files changed, 73 insertions(+), 3 deletions(-) diff --git a/taosmd/collections.py b/taosmd/collections.py index 6561aa05..9084f9d4 100644 --- a/taosmd/collections.py +++ b/taosmd/collections.py @@ -80,6 +80,11 @@ class CollectionNotFoundError(KeyError): """Raised when referencing a collection id that does not exist.""" +class CollectionBusyError(RuntimeError): + """Raised when starting an index on a collection that is already + ``indexing`` (mapped to HTTP 409 by the endpoint).""" + + def _new_collection_id() -> str: """``col-`` + 12 lowercase hex. Matches the agent-name grammar (``^[a-z][a-z0-9_-]{0,62}$``) so the id can double as the agent name diff --git a/taosmd/http_server.py b/taosmd/http_server.py index aae42870..a0757734 100644 --- a/taosmd/http_server.py +++ b/taosmd/http_server.py @@ -115,7 +115,8 @@ empty roots = collections off) ``POST /collections/{id}/index`` -> 202 ``{"status": "indexing", "job": }`` async; poll GET /collections/{id} until status is - "ready" or "error"; stats update on completion + "ready" or "error"; stats update on completion; + 409 while an index is already running ``DELETE /collections/{id}`` -> archive (reversible; content hidden from query, nothing destroyed; destruction only via wipe) ``POST /shelves`` ``{"shelf_id", "project_id"?, "display_name"?}`` -> ``{"shelf": {...}, "created": bool}`` @@ -1842,7 +1843,10 @@ def _handle_collections_get(self, collection_id: str) -> None: def _handle_collections_index(self, collection_id: str) -> None: if not self._check_admin_token(): return - from .collections import CollectionNotFoundError # noqa: PLC0415 + from .collections import ( # noqa: PLC0415 + CollectionBusyError, + CollectionNotFoundError, + ) try: receipt = runner.run( service.collections_index_start(collection_id, data_dir=data_dir) @@ -1850,6 +1854,10 @@ def _handle_collections_index(self, collection_id: str) -> None: except CollectionNotFoundError as exc: self._send_json(404, {"error": str(exc)}) return + except CollectionBusyError as exc: + # Concurrent-index guard: one index per collection at a time. + self._send_json(409, {"error": str(exc)}) + return # Async by contract: 202 now, poll GET /collections/{id} until the # status is ready|error. The walk runs on the service loop so all # store access stays in the single-threaded context. diff --git a/taosmd/service.py b/taosmd/service.py index 2de88a44..9512eb8d 100644 --- a/taosmd/service.py +++ b/taosmd/service.py @@ -907,11 +907,13 @@ async def collections_get(collection_id: str, *, data_dir=None) -> dict: async def collections_index_start(collection_id: str, *, data_dir=None) -> dict: """Validate and mark a collection ``indexing``; the walk runs separately. - Raises ``CollectionNotFoundError`` (404) for an unknown id and + Raises ``CollectionNotFoundError`` (404) for an unknown id, + ``CollectionBusyError`` (409) when an index is already running, and ``ValueError`` (400) for an archived collection or a source path that no longer resolves inside an allowed root, so callers get a synchronous error before the background job is spawned. """ + from .collections import CollectionBusyError # noqa: PLC0415 store = _collection_store(data_dir) try: col = store.get(collection_id) @@ -919,6 +921,11 @@ async def collections_index_start(collection_id: str, *, data_dir=None) -> dict: raise ValueError( f"collection {collection_id!r} is archived; unarchive before indexing" ) + if col["status"] == "indexing": + raise CollectionBusyError( + f"collection {collection_id!r} is already indexing; " + f"poll GET /collections/{collection_id} until it settles" + ) store.resolve_source_path(col["source_path"]) store.set_status(collection_id, "indexing") finally: diff --git a/tests/test_collections_http.py b/tests/test_collections_http.py index b660283a..6329b0e2 100644 --- a/tests/test_collections_http.py +++ b/tests/test_collections_http.py @@ -330,6 +330,32 @@ def test_index_flow_and_search_grant_enforcement(live_server): assert body["hits"] +def test_index_409_while_already_indexing(live_server): + """Concurrent-index guard: a second index start while the collection is + already ``indexing`` is rejected with 409; ready re-arms it.""" + from taosmd.collections import CollectionStore + + base, data_dir, source_dir = live_server + col = _create(base, source_dir) + store = CollectionStore(data_dir) + try: + store.set_status(col["id"], "indexing") + status, body = _req( + "POST", f"{base}/collections/{col['id']}/index", token=_TOKEN, + ) + assert status == 409 + assert "indexing" in body["error"] + # Back to a settled state: indexing is allowed again. + store.set_status(col["id"], "ready") + finally: + store.close() + status, _ = _req( + "POST", f"{base}/collections/{col['id']}/index", token=_TOKEN, + ) + assert status == 202 + _wait_ready(base, col["id"]) + + def test_index_unknown_404(live_server): base, _, _ = live_server status, _ = _req( diff --git a/tests/test_collections_ingest.py b/tests/test_collections_ingest.py index 28e0ed55..2953114a 100644 --- a/tests/test_collections_ingest.py +++ b/tests/test_collections_ingest.py @@ -322,6 +322,30 @@ def test_ingest_folder_root_removed_from_config_refused(data_dir, source_dir): assert store.get(col["id"])["status"] == "error" +def test_index_start_rejects_concurrent_index(data_dir, source_dir): + """collections_index_start must refuse to double-start: a second start + while the collection is already ``indexing`` raises CollectionBusyError + (409 over HTTP); ready and error states re-arm it.""" + from taosmd import service + from taosmd.collections import CollectionBusyError + + store, col = _make_collection(data_dir, source_dir) + receipt = asyncio.run(service.collections_index_start(col["id"], data_dir=data_dir)) + assert receipt["status"] == "indexing" + with pytest.raises(CollectionBusyError): + asyncio.run(service.collections_index_start(col["id"], data_dir=data_dir)) + + # A finished index (ready) re-arms the start. + store.set_status(col["id"], "ready") + receipt = asyncio.run(service.collections_index_start(col["id"], data_dir=data_dir)) + assert receipt["status"] == "indexing" + + # So does a failed one (error): retries stay possible. + store.set_status(col["id"], "error", error="boom") + receipt = asyncio.run(service.collections_index_start(col["id"], data_dir=data_dir)) + assert receipt["status"] == "indexing" + + # --------------------------------------------------------------------------- # Search integration # --------------------------------------------------------------------------- From fabfe4c99c4c48f5fa61db81750a83e4fcb9b273 Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 05:21:31 +0100 Subject: [PATCH 11/16] feat(collections): cap the walker at 20000 files per tree The walker had a per-file size cap but no tree cap, so a collection pointed at a huge tree (a home dir, a monorepo root) would grind through every file. collect_files now raises past DEFAULT_MAX_FILES (20000, upgrade-path comment for making it configurable) and the index errors cleanly with the offending path and a pointer to narrow the source folder; ingest_folder passes the cap through for callers. --- taosmd/collections.py | 23 ++++++++++++++++++++++- tests/test_collections_ingest.py | 20 ++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/taosmd/collections.py b/taosmd/collections.py index 9084f9d4..038b8e63 100644 --- a/taosmd/collections.py +++ b/taosmd/collections.py @@ -60,6 +60,14 @@ #: index. Overridable per ingest_folder call. DEFAULT_MAX_FILE_BYTES = 10 * 1024 * 1024 +#: Tree-wide file-count cap for one collection walk. A docs folder should be +#: nowhere near this; hitting it means the collection points at something too +#: big (a whole home dir, a monorepo root) and the index errors cleanly with +#: a message instead of grinding through the tree. +# upgrade-path: make configurable (collections.max_files) when a legitimate +# corpus needs more than this. +DEFAULT_MAX_FILES = 20000 + #: Directory names never descended into, regardless of gitignore rules. _SKIP_DIRS = frozenset({ ".git", ".hg", ".svn", "node_modules", "__pycache__", ".venv", "venv", @@ -500,6 +508,7 @@ def collect_files( source_root: Path | str, *, max_file_bytes: int = DEFAULT_MAX_FILE_BYTES, + max_files: int = DEFAULT_MAX_FILES, ) -> tuple[list[tuple[Path, str]], dict]: """Walk ``source_root`` and return ``([(abs_path, rel_posix)], skips)``. @@ -508,6 +517,10 @@ def collect_files( null-byte sniff), files over ``max_file_bytes``, symlinks that escape the root, and files no registered loader claims. ``skips`` counts each skip reason so ingest stats can surface them. + + Raises ``ValueError`` when the tree holds more than ``max_files`` + ingestable files: the walk stops instead of grinding through a tree far + bigger than any docs collection should be, and the index errors cleanly. """ root = Path(source_root).resolve() skips = { @@ -577,6 +590,12 @@ def collect_files( skips["skipped_binary"] += 1 continue files.append((fpath, rel_f)) + if len(files) > max_files: + raise ValueError( + f"collection walk exceeded the {max_files}-file cap at " + f"{rel_f!r}; point the collection at a smaller folder " + f"(or tighten .gitignore rules)" + ) return files, skips @@ -689,6 +708,7 @@ async def ingest_folder( *, data_dir=None, max_file_bytes: int = DEFAULT_MAX_FILE_BYTES, + max_files: int = DEFAULT_MAX_FILES, chunk_chars: int = 2000, ) -> dict: """Walk a collection's source folder and index its documents. @@ -729,7 +749,8 @@ async def ingest_folder( # the 202+poll contract holds and the single service loop keeps # serving /search and /ingest while a collection indexes. files, skips = await asyncio.to_thread( - collect_files, source_root, max_file_bytes=max_file_bytes, + collect_files, source_root, + max_file_bytes=max_file_bytes, max_files=max_files, ) stores = await _api._ensure_stores(data_dir) diff --git a/tests/test_collections_ingest.py b/tests/test_collections_ingest.py index 2953114a..0f28f1f5 100644 --- a/tests/test_collections_ingest.py +++ b/tests/test_collections_ingest.py @@ -130,6 +130,26 @@ def test_collect_files_skips_binary_and_oversized(source_dir): assert skips["skipped_size"] == 1 +def test_collect_files_rejects_trees_over_the_file_cap(source_dir): + """The walker has a per-file size cap but also needs a tree cap: pointing + a collection at a huge tree must fail fast with a clear message instead + of grinding through it.""" + with pytest.raises(ValueError, match="file cap"): + collect_files(source_dir, max_files=2) + # At or under the cap is fine (the fixture tree has 4 claimed files). + files, _ = collect_files(source_dir, max_files=4) + assert len(files) == 4 + + +def test_ingest_folder_errors_cleanly_over_file_cap(data_dir, source_dir): + store, col = _make_collection(data_dir, source_dir) + with pytest.raises(ValueError, match="file cap"): + asyncio.run(ingest_folder(col["id"], data_dir=data_dir, max_files=2)) + got = store.get(col["id"]) + assert got["status"] == "error" + assert "file cap" in (got["error"] or "") + + def test_collect_files_skips_symlink_escape(source_dir, tmp_path): outside = tmp_path / "outside.md" outside.write_text("# outside the root") From 5edf1750e6fdc228f2b7a1d6e78f74e717cf8c8c Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 05:21:31 +0100 Subject: [PATCH 12/16] docs(collections): state the grants trust model, walker caps, index 409 Make the Jul 19 data-plane-grants decision explicit rather than a surprise: grants scope collections per agent inside the server-token boundary, any holder of the server token can manage grants, and collection access is exactly as strong as that token. Only create/index/delete sit behind the admin token because they touch the server filesystem. Also document the 20000-file walk cap and the 409 on concurrent index starts. --- docs/collections.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/docs/collections.md b/docs/collections.md index 58bcd12b..2676c3f4 100644 --- a/docs/collections.md +++ b/docs/collections.md @@ -48,6 +48,13 @@ recall, never destroyed; the archive keeps every version). The walker respects `.gitignore` files (simplified rules), skips VCS and dependency directories, hidden directories, binaries, and oversized files. +A tree with more than 20000 ingestable files errors the index with a clear +message instead of grinding through it; point the collection at a smaller +folder (a `docs/` directory, not a monorepo root). + +One index runs per collection at a time: starting an index while one is +already running returns `409` (poll `GET /collections/{id}` until the +status settles at `ready` or `error`, then retry). ## Querying @@ -69,6 +76,15 @@ curl -s localhost:7900/search -d '{ - Collection hits carry `collection_id`, `file_path` (relative to the source root), and `source: "collection"` in their metadata. +Trust model: grants protect collections from ungranted *agents*, not from +holders of the server token. Grant and revoke live on the data plane, so +anyone presenting the server's bearer token can manage grants (and could +grant themselves access); collection access is therefore exactly as strong +as the server token. Treat the token as the security boundary and grants +as the per-agent scoping mechanism inside it. Only create/index/delete sit +behind the separate admin token, because those touch the server's +filesystem. + Over MCP: `memory_list_collections` lists them; `memory_search` takes a `collection` parameter. From 5a9c0b6fbedc018550ae6ce711c276e7796f2e69 Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 09:41:12 +0100 Subject: [PATCH 13/16] fix(collections): supersede rows when a previously-indexed file becomes empty A file whose content was emptied still exists on disk, so it stayed in the walker's seen set and never reached the deleted path: its old vector rows were never superseded and its hash state was never cleared, leaving stale content searchable indefinitely against the incremental/zero-loss contract in the function's own docstring. Drop such a file from seen so the existing supersede path handles it (rows stamped valid_to with the hidden_by marker, never hard-deleted) and its file state is cleared. A file that was already blank is absent from the prior state too, so it stays a no-op and does not churn on re-index. --- taosmd/collections.py | 12 ++++++-- tests/test_collections_ingest.py | 53 ++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/taosmd/collections.py b/taosmd/collections.py index 038b8e63..46d1b3b4 100644 --- a/taosmd/collections.py +++ b/taosmd/collections.py @@ -715,8 +715,8 @@ async def ingest_folder( Incremental by content hash: files whose hash matches the stored state are skipped; changed files get their old rows superseded (never deleted) - and their new chunks ingested; files that disappeared from the source - are superseded too. Chunks route through :func:`taosmd.api.ingest_batch` + and their new chunks ingested; files that disappeared from the source, + or whose content was emptied, are superseded too. Chunks route through :func:`taosmd.api.ingest_batch` under the collection id as the agent namespace, so the batch dedup and metadata preservation come for free and every chunk lands in the zero-loss archive. @@ -777,6 +777,14 @@ async def ingest_folder( continue text = blob.raw_text or getattr(blob, "content", "") or "" if not text.strip(): + # A file with no content contributes nothing, but if we had + # indexed it before, leaving it in ``seen`` would strand its + # old rows in active recall forever (the file still exists, + # so the deleted set never catches it). Drop it from ``seen`` + # and let the deleted path supersede its rows and clear its + # hash state. A file that was already blank is absent from + # ``prior`` too, so this stays a no-op for it. + seen.discard(rel) continue # Hashing + chunking are CPU-bound; off-loop like the walk above. file_hash, chunks = await asyncio.to_thread( diff --git a/tests/test_collections_ingest.py b/tests/test_collections_ingest.py index 0f28f1f5..33b40f05 100644 --- a/tests/test_collections_ingest.py +++ b/tests/test_collections_ingest.py @@ -291,6 +291,59 @@ def test_reindex_deleted_file_supersedes_rows(data_dir, source_dir): assert "guide.txt" not in store.file_states(col["id"]) +def test_reindex_emptied_file_supersedes_rows(data_dir, source_dir): + """A previously-indexed file whose content is emptied must lose its rows. + + Zero-loss cuts both ways: the old text has to leave active recall (it no + longer exists in the source) while the physical rows survive as + superseded history. Skipping blank files outright left the stale content + searchable forever and the hash state pointing at content that is gone. + """ + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + store.grant(col["id"], "dev") + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + + def _search(q): + return asyncio.run( + taosmd_api.search( + q, agent="dev", mode="bm25", + collections=[col["id"]], collections_only=True, data_dir=data_dir, + ) + ) + + assert any("frobnicates" in h["text"] for h in _search("frobnicates sprocket")) + + # Truncate to whitespace: the file still exists, but has no content. + (source_dir / "readme.md").write_text(" \n\n \n") + stats2 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + assert stats2["chunks_superseded"] >= 1 + + # The stale content is out of active recall. + assert not any("frobnicates" in h["text"] for h in _search("frobnicates sprocket")) + # And its hash state is cleared, so a later refill re-indexes cleanly. + assert "readme.md" not in store.file_states(col["id"]) + # Zero-loss: superseded, not hard-deleted. + stores = asyncio.run(taosmd_api._ensure_stores(data_dir)) + rows = stores["vector"]._conn.execute( + "SELECT metadata_json FROM vector_memory WHERE valid_to IS NOT NULL" + ).fetchall() + assert rows + assert any("collection-reindex:" in (r["metadata_json"] or "") for r in rows) + + +def test_reindex_already_empty_file_is_a_no_op(data_dir, source_dir): + """A file that was blank and is still blank must not churn on re-index.""" + _patch_embedder(data_dir) + (source_dir / "blank.md").write_text("\n\n \n") + store, col = _make_collection(data_dir, source_dir) + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + stats2 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + assert stats2["chunks_superseded"] == 0 + assert stats2["files_deleted"] == 0 + assert stats2["files_indexed"] == 0 + + def test_ingest_folder_walks_off_the_event_loop(data_dir, source_dir, monkeypatch): """The 202+poll contract promises a responsive server during an index: the filesystem walk (os.walk + per-file stat + null-byte sniff) must run From 5621882252fa263450aca8588d69c862a28e0d58 Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 09:41:25 +0100 Subject: [PATCH 14/16] fix(collections): close the collection store connection in ingest_folder ingest_folder opened its own CollectionStore and never closed it on any path, so a server re-indexing on a timer leaked one sqlite connection (and file handle) per run. The service wrappers all close theirs in a finally; this path was the outlier. Wrap the whole body in try/finally so the connection closes on success and on the error path that records status=error. --- taosmd/collections.py | 240 ++++++++++++++++--------------- tests/test_collections_ingest.py | 57 ++++++++ 2 files changed, 181 insertions(+), 116 deletions(-) diff --git a/taosmd/collections.py b/taosmd/collections.py index 46d1b3b4..37d092b8 100644 --- a/taosmd/collections.py +++ b/taosmd/collections.py @@ -730,122 +730,130 @@ async def ingest_folder( resolved_dir = _api._resolve_data_dir(data_dir) store = CollectionStore(resolved_dir) - col = store.get(collection_id) - if col["status"] == "archived": - raise ValueError(f"collection {collection_id!r} is archived; unarchive before indexing") - if col["embedder"]: - # Per-collection embedder is stored and returned now (the mechanism); - # Phase 1 indexes with the global default regardless. - logger.info( - "collection %s requests embedder %r; Phase 1 indexes with the " - "global default embedder", collection_id, col["embedder"], - ) - store.set_status(collection_id, "indexing") + # One sqlite connection per index run: without the finally below, a + # server re-indexing on a timer leaks a file handle every pass. The + # service wrappers already close theirs; this path was the outlier. try: - source_root = store.resolve_source_path(col["source_path"]) - prior = store.file_states(collection_id) - # The walk (os.walk + per-file stat + null-byte sniff over every - # candidate) is blocking filesystem work; run it on a worker thread so - # the 202+poll contract holds and the single service loop keeps - # serving /search and /ingest while a collection indexes. - files, skips = await asyncio.to_thread( - collect_files, source_root, - max_file_bytes=max_file_bytes, max_files=max_files, - ) - - stores = await _api._ensure_stores(data_dir) - vmem = stores["vector"] - - items: list[dict] = [] - indexed: list[tuple[str, str]] = [] - changed: list[str] = [] - unchanged = 0 - errors: list[str] = [] - seen: set[str] = set() + col = store.get(collection_id) + if col["status"] == "archived": + raise ValueError( + f"collection {collection_id!r} is archived; unarchive before indexing" + ) + if col["embedder"]: + # Per-collection embedder is stored and returned now (the mechanism); + # Phase 1 indexes with the global default regardless. + logger.info( + "collection %s requests embedder %r; Phase 1 indexes with the " + "global default embedder", collection_id, col["embedder"], + ) + store.set_status(collection_id, "indexing") + try: + source_root = store.resolve_source_path(col["source_path"]) + prior = store.file_states(collection_id) + # The walk (os.walk + per-file stat + null-byte sniff over every + # candidate) is blocking filesystem work; run it on a worker thread so + # the 202+poll contract holds and the single service loop keeps + # serving /search and /ingest while a collection indexes. + files, skips = await asyncio.to_thread( + collect_files, source_root, + max_file_bytes=max_file_bytes, max_files=max_files, + ) - for abs_path, rel in files: - seen.add(rel) - loader = _loader_for(abs_path) - if loader is None: # pragma: no cover - collect_files already filtered - continue - try: - blob = await loader.load( - abs_path, max_bytes=max_file_bytes, base_dir=source_root + stores = await _api._ensure_stores(data_dir) + vmem = stores["vector"] + + items: list[dict] = [] + indexed: list[tuple[str, str]] = [] + changed: list[str] = [] + unchanged = 0 + errors: list[str] = [] + seen: set[str] = set() + + for abs_path, rel in files: + seen.add(rel) + loader = _loader_for(abs_path) + if loader is None: # pragma: no cover - collect_files already filtered + continue + try: + blob = await loader.load( + abs_path, max_bytes=max_file_bytes, base_dir=source_root + ) + except Exception as exc: # noqa: BLE001 - per-file failures are non-fatal + errors.append(f"{rel}: {type(exc).__name__}: {exc}") + continue + text = blob.raw_text or getattr(blob, "content", "") or "" + if not text.strip(): + # A file with no content contributes nothing, but if we had + # indexed it before, leaving it in ``seen`` would strand its + # old rows in active recall forever (the file still exists, + # so the deleted set never catches it). Drop it from ``seen`` + # and let the deleted path supersede its rows and clear its + # hash state. A file that was already blank is absent from + # ``prior`` too, so this stays a no-op for it. + seen.discard(rel) + continue + # Hashing + chunking are CPU-bound; off-loop like the walk above. + file_hash, chunks = await asyncio.to_thread( + _hash_and_chunk, text, chunk_chars, prior.get(rel) ) - except Exception as exc: # noqa: BLE001 - per-file failures are non-fatal - errors.append(f"{rel}: {type(exc).__name__}: {exc}") - continue - text = blob.raw_text or getattr(blob, "content", "") or "" - if not text.strip(): - # A file with no content contributes nothing, but if we had - # indexed it before, leaving it in ``seen`` would strand its - # old rows in active recall forever (the file still exists, - # so the deleted set never catches it). Drop it from ``seen`` - # and let the deleted path supersede its rows and clear its - # hash state. A file that was already blank is absent from - # ``prior`` too, so this stays a no-op for it. - seen.discard(rel) - continue - # Hashing + chunking are CPU-bound; off-loop like the walk above. - file_hash, chunks = await asyncio.to_thread( - _hash_and_chunk, text, chunk_chars, prior.get(rel) - ) - if chunks is None: - unchanged += 1 - continue - if rel in prior: - changed.append(rel) - for i, chunk in enumerate(chunks): - chunk_id = hashlib.sha256( - f"{collection_id}:{rel}:{file_hash}:{i}".encode("utf-8") - ).hexdigest() - items.append({ - "text": chunk, - "id": chunk_id, - "metadata": { - "collection_id": collection_id, - "file_path": rel, - "source": "collection", - "chunk_index": i, - "file_hash": file_hash, - }, - }) - indexed.append((rel, file_hash)) - - deleted = sorted(set(prior) - seen) - - chunks_superseded = 0 - for rel in [*changed, *deleted]: - chunks_superseded += _supersede_collection_rows(vmem, collection_id, rel) - - if items: - result = await _api.ingest_batch(items, agent=collection_id, data_dir=data_dir) - else: - result = {"ingested": 0, "skipped": 0} - - for rel, file_hash in indexed: - store.set_file_state(collection_id, rel, file_hash) - for rel in deleted: - store.remove_file_state(collection_id, rel) - - now = time.time() - stats = { - "files_indexed": len(indexed), - "files_unchanged": unchanged, - "files_deleted": len(deleted), - "files_total": len(store.file_states(collection_id)), - "chunks_ingested": result.get("ingested", 0), - "chunks_skipped": result.get("skipped", 0), - "chunks_superseded": chunks_superseded, - "errors": errors[:20], - **skips, - } - if result.get("vector_failures"): - stats["vector_failures"] = result["vector_failures"] - stats["degraded"] = True - store.set_stats(collection_id, stats) - store.set_status(collection_id, "ready", last_indexed=now) - return stats - except Exception as exc: - store.set_status(collection_id, "error", error=f"{type(exc).__name__}: {exc}") - raise + if chunks is None: + unchanged += 1 + continue + if rel in prior: + changed.append(rel) + for i, chunk in enumerate(chunks): + chunk_id = hashlib.sha256( + f"{collection_id}:{rel}:{file_hash}:{i}".encode("utf-8") + ).hexdigest() + items.append({ + "text": chunk, + "id": chunk_id, + "metadata": { + "collection_id": collection_id, + "file_path": rel, + "source": "collection", + "chunk_index": i, + "file_hash": file_hash, + }, + }) + indexed.append((rel, file_hash)) + + deleted = sorted(set(prior) - seen) + + chunks_superseded = 0 + for rel in [*changed, *deleted]: + chunks_superseded += _supersede_collection_rows(vmem, collection_id, rel) + + if items: + result = await _api.ingest_batch(items, agent=collection_id, data_dir=data_dir) + else: + result = {"ingested": 0, "skipped": 0} + + for rel, file_hash in indexed: + store.set_file_state(collection_id, rel, file_hash) + for rel in deleted: + store.remove_file_state(collection_id, rel) + + now = time.time() + stats = { + "files_indexed": len(indexed), + "files_unchanged": unchanged, + "files_deleted": len(deleted), + "files_total": len(store.file_states(collection_id)), + "chunks_ingested": result.get("ingested", 0), + "chunks_skipped": result.get("skipped", 0), + "chunks_superseded": chunks_superseded, + "errors": errors[:20], + **skips, + } + if result.get("vector_failures"): + stats["vector_failures"] = result["vector_failures"] + stats["degraded"] = True + store.set_stats(collection_id, stats) + store.set_status(collection_id, "ready", last_indexed=now) + return stats + except Exception as exc: + store.set_status(collection_id, "error", error=f"{type(exc).__name__}: {exc}") + raise + finally: + store.close() diff --git a/tests/test_collections_ingest.py b/tests/test_collections_ingest.py index 33b40f05..e04eb4c2 100644 --- a/tests/test_collections_ingest.py +++ b/tests/test_collections_ingest.py @@ -344,6 +344,63 @@ def test_reindex_already_empty_file_is_a_no_op(data_dir, source_dir): assert stats2["files_indexed"] == 0 +def _closing_store_spy(monkeypatch): + """Patch CollectionStore so every instance records whether it was closed.""" + from taosmd import collections as collections_mod + + real = collections_mod.CollectionStore + made: list = [] + + class _Spy(real): # type: ignore[misc, valid-type] + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.closed = False + made.append(self) + + def close(self) -> None: + super().close() + self.closed = True + + monkeypatch.setattr(collections_mod, "CollectionStore", _Spy) + return made + + +def test_ingest_folder_closes_the_store_on_success(data_dir, source_dir, monkeypatch): + """ingest_folder opens its own CollectionStore; a server that re-indexes + on a timer would leak one sqlite connection per run if it never closed.""" + _patch_embedder(data_dir) + _store, col = _make_collection(data_dir, source_dir) + made = _closing_store_spy(monkeypatch) + + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + + assert made, "ingest_folder did not open a CollectionStore" + assert all(s.closed for s in made) + for s in made: + with pytest.raises(Exception): + s._conn.execute("SELECT 1") + + +def test_ingest_folder_closes_the_store_on_error(data_dir, source_dir, monkeypatch): + """The failure path (status=error) must close the connection too.""" + from taosmd import collections as collections_mod + + _patch_embedder(data_dir) + _store, col = _make_collection(data_dir, source_dir) + made = _closing_store_spy(monkeypatch) + + def _boom(*args, **kwargs): + raise RuntimeError("walk exploded") + + monkeypatch.setattr(collections_mod, "collect_files", _boom) + + with pytest.raises(RuntimeError, match="walk exploded"): + asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + + assert made, "ingest_folder did not open a CollectionStore" + assert all(s.closed for s in made) + + def test_ingest_folder_walks_off_the_event_loop(data_dir, source_dir, monkeypatch): """The 202+poll contract promises a responsive server during an index: the filesystem walk (os.walk + per-file stat + null-byte sniff) must run From b63277ebd156cf9530b01441054056cc9dc0ff6e Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 09:41:25 +0100 Subject: [PATCH 15/16] docs(collections): language specifiers on fenced blocks The two HTTP surface blocks had bare fences (markdownlint MD040). --- docs/collections.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/collections.md b/docs/collections.md index 2676c3f4..f4a2cf56 100644 --- a/docs/collections.md +++ b/docs/collections.md @@ -92,7 +92,7 @@ Over MCP: `memory_list_collections` lists them; `memory_search` takes a Data plane (bearer token when one is configured): -``` +```text GET /collections [?project=] GET /collections/{id} POST /collections/{id}/link {"type": "taos"|"git", "id": "..."} @@ -104,7 +104,7 @@ POST /search with collection/collections/collections_only Admin (dedicated admin token, fail-closed): -``` +```text POST /collections {"name", "kind", "source_path", "embedder"?} POST /collections/{id}/index -> 202; poll GET /collections/{id} DELETE /collections/{id} -> archive (reversible) From e226d2980ed0d81a1d58ecc8c490674e03f74b7d Mon Sep 17 00:00:00 2001 From: jaylfc Date: Mon, 20 Jul 2026 09:56:39 +0100 Subject: [PATCH 16/16] feat(collections): report emptied files under their own stats counter A previously-indexed file whose content becomes empty was routed through the deleted path and counted in files_deleted, conflating two things a UI needs to show apart: a file that vanished from disk versus a file that still exists but has no content left. Split the reporting with a distinct files_emptied counter. The underlying behaviour is unchanged: both cases still supersede their old rows and clear their file_states hash, so the zero-loss retirement holds. The key is always present (0 when none) so the stats shape stays stable. --- docs/collections.md | 23 ++++++- .../codebase-indexing-collections-design.md | 7 ++- taosmd/cli.py | 1 + taosmd/collections.py | 25 ++++++-- taosmd/http_server.py | 5 ++ tests/test_collections_ingest.py | 61 +++++++++++++++++++ 6 files changed, 115 insertions(+), 7 deletions(-) diff --git a/docs/collections.md b/docs/collections.md index f4a2cf56..b665ca9d 100644 --- a/docs/collections.md +++ b/docs/collections.md @@ -44,7 +44,28 @@ taosmd collections list Statuses: `created -> indexing -> ready | error` (and `archived` after a delete). Re-indexing is incremental: unchanged files are skipped by content hash, changed and deleted files have their old rows superseded (hidden from -recall, never destroyed; the archive keeps every version). +recall, never destroyed; the archive keeps every version). A file that still +exists but whose content has been emptied is retired the same way, and +reported separately as `files_emptied` rather than `files_deleted`, so a +blanked file is never mistaken for one that vanished from disk. + +Index stats (on the collection row, and in the `GET /collections/{id}` +response) always carry the same keys: + +```text +files_indexed files whose content was chunked and ingested this pass +files_unchanged files skipped because their content hash matched +files_deleted previously-indexed files no longer present on disk +files_emptied previously-indexed files still present but now empty +files_total files currently tracked in the collection's hash state +chunks_ingested chunks written this pass +chunks_skipped chunks the batch deduped away +chunks_superseded active rows retired this pass (never destroyed) +errors up to 20 per-file failure strings +``` + +`files_deleted` and `files_emptied` count disjoint sets, and both are +present (as `0`) when nothing was retired. The walker respects `.gitignore` files (simplified rules), skips VCS and dependency directories, hidden directories, binaries, and oversized files. diff --git a/docs/specs/codebase-indexing-collections-design.md b/docs/specs/codebase-indexing-collections-design.md index f001226d..6c12e968 100644 --- a/docs/specs/codebase-indexing-collections-design.md +++ b/docs/specs/codebase-indexing-collections-design.md @@ -69,7 +69,12 @@ POST /collections (admin) {"name", "kind", "source_path GET /collections [?project=] -> {"collections": [ {id, name, kind, source_path, project_id, status, stats, last_indexed, links, grants} ]} GET /collections/{id} -> {"collection": {...}} with full stats - (file_count, chunk_count, last_indexed, errors: [...]) + (files_indexed, files_unchanged, files_deleted, + files_emptied, files_total, chunks_ingested, + chunks_skipped, chunks_superseded, errors: [...]) + files_deleted counts files gone from disk; + files_emptied counts files still present whose + content is now empty. Disjoint, always present. POST /collections/{id}/index (admin) -> {"status": "indexing", "job": ""} async; poll GET /collections/{id} until status is "ready" or "error"; stats update live diff --git a/taosmd/cli.py b/taosmd/cli.py index 6360c615..7bd4fc5b 100644 --- a/taosmd/cli.py +++ b/taosmd/cli.py @@ -1043,6 +1043,7 @@ def _fmt(col: dict) -> str: print( f"{args.collection_id}: files_indexed={stats['files_indexed']} " f"unchanged={stats['files_unchanged']} deleted={stats['files_deleted']} " + f"emptied={stats['files_emptied']} " f"chunks_ingested={stats['chunks_ingested']} " f"superseded={stats['chunks_superseded']} errors={len(stats['errors'])}" ) diff --git a/taosmd/collections.py b/taosmd/collections.py index 37d092b8..85e9047b 100644 --- a/taosmd/collections.py +++ b/taosmd/collections.py @@ -716,7 +716,10 @@ async def ingest_folder( Incremental by content hash: files whose hash matches the stored state are skipped; changed files get their old rows superseded (never deleted) and their new chunks ingested; files that disappeared from the source, - or whose content was emptied, are superseded too. Chunks route through :func:`taosmd.api.ingest_batch` + or whose content was emptied, are superseded too (and reported apart, as + ``files_deleted`` and ``files_emptied``, because a vanished file and a + blanked one call for different responses). + Chunks route through :func:`taosmd.api.ingest_batch` under the collection id as the agent namespace, so the batch dedup and metadata preservation come for free and every chunk lands in the zero-loss archive. @@ -765,6 +768,7 @@ async def ingest_folder( items: list[dict] = [] indexed: list[tuple[str, str]] = [] changed: list[str] = [] + emptied: set[str] = set() unchanged = 0 errors: list[str] = [] seen: set[str] = set() @@ -791,6 +795,11 @@ async def ingest_folder( # hash state. A file that was already blank is absent from # ``prior`` too, so this stays a no-op for it. seen.discard(rel) + if rel in prior: + # Same retirement, different report: a file that still + # exists but has been emptied is not a file that + # vanished, and clients show the two separately. + emptied.add(rel) continue # Hashing + chunking are CPU-bound; off-loop like the walk above. file_hash, chunks = await asyncio.to_thread( @@ -818,10 +827,15 @@ async def ingest_folder( }) indexed.append((rel, file_hash)) - deleted = sorted(set(prior) - seen) + # Everything that has to be retired this pass: files gone from + # disk plus files still on disk whose content is now empty. They + # are superseded and have their hash state cleared identically; + # only the stats split them apart. + retired = sorted(set(prior) - seen) + removed = [rel for rel in retired if rel not in emptied] chunks_superseded = 0 - for rel in [*changed, *deleted]: + for rel in [*changed, *retired]: chunks_superseded += _supersede_collection_rows(vmem, collection_id, rel) if items: @@ -831,14 +845,15 @@ async def ingest_folder( for rel, file_hash in indexed: store.set_file_state(collection_id, rel, file_hash) - for rel in deleted: + for rel in retired: store.remove_file_state(collection_id, rel) now = time.time() stats = { "files_indexed": len(indexed), "files_unchanged": unchanged, - "files_deleted": len(deleted), + "files_deleted": len(removed), + "files_emptied": len(emptied), "files_total": len(store.file_states(collection_id)), "chunks_ingested": result.get("ingested", 0), "chunks_skipped": result.get("skipped", 0), diff --git a/taosmd/http_server.py b/taosmd/http_server.py index a0757734..12881d9b 100644 --- a/taosmd/http_server.py +++ b/taosmd/http_server.py @@ -102,6 +102,11 @@ Collections (data plane; create/index/delete are admin, see below) ``GET /collections [?project=]`` -> ``{"collections": [...]}`` (project matches links of either type) ``GET /collections/{id}`` -> ``{"collection": {...}}`` with status/stats/links/grants + stats: files_indexed/files_unchanged/files_deleted/ + files_emptied/files_total, chunks_ingested/skipped/ + superseded, errors. files_deleted is files gone from + disk, files_emptied is files still there but now empty + (disjoint; both always present, 0 when none). ``POST /collections/{id}/link`` ``{"type": "taos"|"git", "id"}`` -> ``{"collection": {...}}`` ``POST /collections/{id}/unlink`` same body; metadata only, never touches content ``POST /collections/{id}/grants`` ``{"agent"}`` -> grant query access -> ``{"collection": {...}}`` diff --git a/tests/test_collections_ingest.py b/tests/test_collections_ingest.py index e04eb4c2..867fa98a 100644 --- a/tests/test_collections_ingest.py +++ b/tests/test_collections_ingest.py @@ -10,6 +10,7 @@ from __future__ import annotations import asyncio +import json import os import time @@ -281,6 +282,8 @@ def test_reindex_deleted_file_supersedes_rows(data_dir, source_dir): (source_dir / "guide.txt").unlink() stats2 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) assert stats2["files_deleted"] == 1 + # A file gone from disk is not an emptied file. + assert stats2["files_emptied"] == 0 hits = asyncio.run( taosmd_api.search( "flux capacitor", agent="dev", mode="bm25", @@ -318,6 +321,9 @@ def _search(q): (source_dir / "readme.md").write_text(" \n\n \n") stats2 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) assert stats2["chunks_superseded"] >= 1 + # An emptied file is not a deleted file: it reports under its own counter. + assert stats2["files_emptied"] == 1 + assert stats2["files_deleted"] == 0 # The stale content is out of active recall. assert not any("frobnicates" in h["text"] for h in _search("frobnicates sprocket")) @@ -341,9 +347,64 @@ def test_reindex_already_empty_file_is_a_no_op(data_dir, source_dir): stats2 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) assert stats2["chunks_superseded"] == 0 assert stats2["files_deleted"] == 0 + assert stats2["files_emptied"] == 0 assert stats2["files_indexed"] == 0 +def test_emptied_and_deleted_files_report_under_separate_counters(data_dir, source_dir): + """``files_emptied`` and ``files_deleted`` describe two different events. + + A UI showing index stats needs to tell "the file is gone from disk" apart + from "the file is still there but has no content left" - they call for + different operator responses. Both retire their old rows the same way + (superseded, never destroyed), so the split is purely about reporting. + """ + _patch_embedder(data_dir) + store, col = _make_collection(data_dir, source_dir) + store.grant(col["id"], "dev") + stats1 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + # The counter is always present, so clients can rely on the shape. + assert stats1["files_emptied"] == 0 + + # One file emptied in place, one removed from disk, in the same pass. + (source_dir / "readme.md").write_text(" \n\n \n") + (source_dir / "guide.txt").unlink() + stats2 = asyncio.run(ingest_folder(col["id"], data_dir=data_dir)) + assert stats2["files_emptied"] == 1 + assert stats2["files_deleted"] == 1 + + def _search(q): + return asyncio.run( + taosmd_api.search( + q, agent="dev", mode="bm25", + collections=[col["id"]], collections_only=True, data_dir=data_dir, + ) + ) + + # Both leave active recall, and both clear their hash state. + assert not any("frobnicates" in h["text"] for h in _search("frobnicates sprocket")) + assert not any("flux capacitor" in h["text"] for h in _search("flux capacitor")) + states = store.file_states(col["id"]) + assert "readme.md" not in states + assert "guide.txt" not in states + + # Zero-loss: the emptied file's rows survive as superseded history. + stores = asyncio.run(taosmd_api._ensure_stores(data_dir)) + rows = stores["vector"]._conn.execute( + "SELECT metadata_json, valid_to FROM vector_memory WHERE valid_to IS NOT NULL" + ).fetchall() + emptied_rows = [ + json.loads(r["metadata_json"]) + for r in rows + if "readme.md" in (r["metadata_json"] or "") + ] + assert emptied_rows + assert all( + str(meta.get("hidden_by", "")).startswith("collection-reindex:") + for meta in emptied_rows + ) + + def _closing_store_spy(monkeypatch): """Patch CollectionStore so every instance records whether it was closed.""" from taosmd import collections as collections_mod