diff --git a/src/surreal_memory/storage/surrealdb/migrations.py b/src/surreal_memory/storage/surrealdb/migrations.py index 676759e9..99a13dc8 100644 --- a/src/surreal_memory/storage/surrealdb/migrations.py +++ b/src/surreal_memory/storage/surrealdb/migrations.py @@ -35,6 +35,7 @@ import uuid from typing import Any +from surreal_memory.storage.surrealdb._ids import _to_surreal_id from surreal_memory.storage.surrealdb.connection import ( MIN_SERVER_VERSION, parse_server_version, @@ -122,17 +123,6 @@ def _already_exists(exc: Exception) -> bool: return "already exists" in str(exc).lower() -def _sanitize_id(raw: str) -> str: - """Strip a table prefix and normalise dashes to underscores. - - Mirrors ``store._to_surreal_id`` so migrated endpoint ids match what the - query layer produces (``neuron:abc-123`` / ``abc-123`` -> ``abc_123``). - """ - if ":" in raw: - raw = raw.rsplit(":", 1)[1] - return raw.replace("-", "_") - - def _record_part(rid: Any) -> str: """Extract the identifier part of a RecordID (or a ``table:id`` string).""" part = getattr(rid, "id", None) @@ -321,8 +311,8 @@ def _to_relation_row(row: dict[str, Any]) -> dict[str, Any]: from surrealdb import RecordID # lazy sid = _record_part(row["id"]) - src = _sanitize_id(str(row.get("source_id", ""))) - tgt = _sanitize_id(str(row.get("target_id", ""))) + src = _to_surreal_id(str(row.get("source_id", ""))) + tgt = _to_surreal_id(str(row.get("target_id", ""))) relation: dict[str, Any] = { "id": RecordID("synapse", sid), "in": RecordID("neuron", src), diff --git a/src/surreal_memory/storage/surrealdb/store.py b/src/surreal_memory/storage/surrealdb/store.py index ca7f6235..749c6eaf 100755 --- a/src/surreal_memory/storage/surrealdb/store.py +++ b/src/surreal_memory/storage/surrealdb/store.py @@ -620,12 +620,17 @@ async def add_neuron(self, neuron: Neuron) -> str: return neuron.id async def get_neuron(self, neuron_id: str) -> Neuron | None: - conn = self._ensure_conn() + # Scope to the current brain: a bare record select would let a caller + # read another brain's neuron by id. The record is still pinned in FROM. + brain_id = self._get_brain_id() sid = _to_surreal_id(neuron_id) try: - result = await conn.select(f"neuron:{sid}") - if result and isinstance(result, list) and len(result) > 0: - return _row_to_neuron(result[0]) + rows = await self._query( + f"SELECT * FROM neuron:{sid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) + if rows: + return _row_to_neuron(rows[0]) except Exception: pass return None @@ -633,16 +638,23 @@ async def get_neuron(self, neuron_id: str) -> Neuron | None: async def get_neurons_batch(self, neuron_ids: list[str]) -> dict[str, Neuron]: if not neuron_ids: return {} - # Use direct select for each ID (more reliable than IN query with params) + # Scope every fetch to the current brain. These are direct record-id + # selects, so without a brain_id filter a caller could read a neuron + # owned by another brain just by knowing (or guessing) its id. The FROM + # is still a pinned record (neuron:{sid}), so this stays a direct fetch, + # not a table scan — the brain_id check only rejects a cross-brain hit. + brain_id = self._get_brain_id() results: dict[str, Neuron] = {} for nid in neuron_ids: sid = _to_surreal_id(nid) try: - result = await self._conn.select(f"neuron:{sid}") - if result and isinstance(result, list) and len(result) > 0: - neuron = _row_to_neuron(result[0]) - # Use the converted ID as key - results[nid] = neuron + rows = await self._query( + f"SELECT * FROM neuron:{sid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) + if rows: + # Use the original ID as key + results[nid] = _row_to_neuron(rows[0]) except Exception: pass return results @@ -718,12 +730,19 @@ async def find_neurons_by_ids( ] if not safe: return [] + # Scope to the current brain: the FROM is a pinned record-id list (still + # a direct fetch, not a table scan), and the WHERE brain_id filter keeps + # a caller from reading another brain's neuron by id. + brain_id = self._get_brain_id() projection = "SELECT *" if include_embedding else "SELECT * OMIT embedding_vec" out: list[Neuron] = [] # Chunk to keep the FROM record-id list a sane query size. for i in range(0, len(safe), 1000): things = ", ".join(f"neuron:{s}" for s in safe[i : i + 1000]) - rows = await self._query(f"{projection} FROM {things}") + rows = await self._query( + f"{projection} FROM {things} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) out.extend(_row_to_neuron(r) for r in rows) return out @@ -823,12 +842,15 @@ async def has_neuron_by_content_hash(self, content_hash: int) -> bool: # ================================================================ async def get_neuron_state(self, neuron_id: str) -> NeuronState | None: - conn = self._ensure_conn() + brain_id = self._get_brain_id() sid = _to_surreal_id(neuron_id) try: - result = await conn.select(f"neuron_state:state_{sid}") - if result: - return _row_to_neuron_state(result[0] if isinstance(result, list) else result) + rows = await self._query( + f"SELECT * FROM neuron_state:state_{sid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) + if rows: + return _row_to_neuron_state(rows[0]) except Exception: pass return None @@ -887,12 +909,15 @@ async def add_synapse(self, synapse: Synapse) -> str: return synapse.id async def get_synapse(self, synapse_id: str) -> Synapse | None: - conn = self._ensure_conn() + brain_id = self._get_brain_id() sid = _to_surreal_id(synapse_id) try: - result = await conn.select(f"synapse:{sid}") - if result: - return _row_to_synapse(result[0] if isinstance(result, list) else result) + rows = await self._query( + f"SELECT * FROM synapse:{sid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) + if rows: + return _row_to_synapse(rows[0]) except Exception: pass return None @@ -1185,12 +1210,15 @@ async def add_fiber(self, fiber: Fiber) -> str: return fiber.id async def get_fiber(self, fiber_id: str) -> Fiber | None: - conn = self._ensure_conn() + brain_id = self._get_brain_id() fid = _to_surreal_id(fiber_id) try: - result = await conn.select(f"fiber:{fid}") - if result: - return _row_to_fiber(result[0] if isinstance(result, list) else result) + rows = await self._query( + f"SELECT * FROM fiber:{fid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) + if rows: + return _row_to_fiber(rows[0]) except Exception: pass return None diff --git a/src/surreal_memory/storage/surrealdb/tool_events.py b/src/surreal_memory/storage/surrealdb/tool_events.py index 73f9388c..33d21216 100644 --- a/src/surreal_memory/storage/surrealdb/tool_events.py +++ b/src/surreal_memory/storage/surrealdb/tool_events.py @@ -15,6 +15,7 @@ from datetime import datetime, timedelta from typing import Any +from surreal_memory.storage.surrealdb._ids import _to_surreal_id from surreal_memory.utils.timeutils import utcnow # Cap per brain to prevent unbounded growth (mirrors the SQLite mixin). @@ -66,7 +67,7 @@ async def insert_tool_events( await conn.insert( "tool_events", { - "id": event_id.replace("-", "_"), + "id": _to_surreal_id(event_id), "event_id": event_id, "brain_id": brain_id, "tool_name": ev.get("tool_name", ""), diff --git a/tests/unit/test_dashboard_perf_queries.py b/tests/unit/test_dashboard_perf_queries.py index cd6cd5fa..eb47d1b7 100644 --- a/tests/unit/test_dashboard_perf_queries.py +++ b/tests/unit/test_dashboard_perf_queries.py @@ -98,7 +98,10 @@ async def test_fetches_records_directly_with_omit(self): await st.find_neurons_by_ids(["abc-123", "def-456"]) (sql,) = _queries(conn) assert sql.startswith("SELECT * OMIT embedding_vec FROM neuron:abc_123, neuron:def_456") - assert "WHERE" not in sql # record fetch, not a table scan + # Still a pinned-record fetch, not a `FROM neuron` table scan… + assert "FROM neuron WHERE" not in sql + # …but scoped to the current brain (no cross-brain id read). + assert sql.endswith("WHERE brain_id = $brain_id") async def test_empty_and_unsafe_ids_are_dropped(self): st, conn = _store_with_mock_conn() diff --git a/tests/unit/test_id_consolidation_and_brain_scope.py b/tests/unit/test_id_consolidation_and_brain_scope.py new file mode 100644 index 00000000..5dd340df --- /dev/null +++ b/tests/unit/test_id_consolidation_and_brain_scope.py @@ -0,0 +1,149 @@ +"""Single-source id sanitisation + brain-scoped batch fetch. + +Two follow-ups on the id-sanitiser consolidation: + + * two weak sanitisers still bypassed the ``_ids.py`` single choke point — + ``migrations._sanitize_id`` (dash-only ``.replace``) and an inline + ``.replace("-", "_")`` in ``tool_events`` — so both now route through + ``_to_surreal_id``; + * ``get_neurons_batch`` / ``find_neurons_by_ids`` issued direct record-id + selects with no ``WHERE brain_id``, so a caller could read another brain's + neuron just by knowing its id — both now scope to the current brain. + +The surrealdb SDK is stubbed (repo convention); the connection is an AsyncMock. +""" + +from __future__ import annotations + +import sys +from unittest.mock import AsyncMock, MagicMock + +if "surrealdb" not in sys.modules: + sys.modules["surrealdb"] = MagicMock() + sys.modules["surrealdb.errors"] = MagicMock() + +import surreal_memory.storage.surrealdb.tool_events as tool_events_mod +from surreal_memory.storage.surrealdb import migrations +from surreal_memory.storage.surrealdb._ids import _to_surreal_id +from surreal_memory.storage.surrealdb.store import SurrealDBStorage + + +def _store_with_mock_conn() -> tuple[SurrealDBStorage, AsyncMock]: + st = SurrealDBStorage(url="http://localhost:8001") + conn = AsyncMock() + conn.query = AsyncMock(return_value=[]) + conn.insert = AsyncMock(return_value=None) + st._conn = conn + st._current_brain_id = "b1" + return st, conn + + +# --------------------------------------------------------------------------- # +# brain-scoped batch fetch +# --------------------------------------------------------------------------- # +class TestBrainScopedBatchFetch: + async def test_find_neurons_by_ids_scopes_to_current_brain(self): + st, conn = _store_with_mock_conn() + await st.find_neurons_by_ids(["abc-123", "def-456"]) + sql, params = conn.query.call_args.args + # Still a pinned-record fetch (ids in FROM), never a `FROM neuron` scan… + assert sql.startswith("SELECT * OMIT embedding_vec FROM neuron:abc_123, neuron:def_456") + # …but now scoped, so a foreign-brain id cannot be read back. + assert sql.endswith("WHERE brain_id = $brain_id") + assert params == {"brain_id": "b1"} + + async def test_get_neurons_batch_scopes_to_current_brain(self): + st, conn = _store_with_mock_conn() + await st.get_neurons_batch(["abc-123"]) + sql, params = conn.query.call_args.args + assert sql == "SELECT * FROM neuron:abc_123 WHERE brain_id = $brain_id" + assert params == {"brain_id": "b1"} + + async def test_get_neurons_batch_uses_safe_brain_id(self): + # An unsafe brain context must fail closed (via _safe_brain_id) rather + # than inline a hostile id — same guarantee as the rest of the store. + st, conn = _store_with_mock_conn() + st._current_brain_id = 'bad"; DELETE neuron' + import pytest + + with pytest.raises(ValueError): + await st.get_neurons_batch(["abc-123"]) + + +# --------------------------------------------------------------------------- # +# single-source sanitisation +# --------------------------------------------------------------------------- # +class TestSanitizerConsolidation: + def test_migrations_has_no_local_sanitiser(self): + # The dash-only duplicate is gone; the single source is _ids._to_surreal_id. + assert not hasattr(migrations, "_sanitize_id") + + def test_relation_row_folds_via_single_source(self, monkeypatch): + # A dotted endpoint id must fold '.' -> '_' (full [A-Za-z0-9_] fold), + # which the removed dash-only sanitiser did NOT do. This both proves the + # single source is used and fixes a latent bug: neuron records are stored + # folded, so a dotted legacy id used to migrate to a non-matching edge. + rid_calls: list[tuple[str, str]] = [] + + def _fake_rid(table: str, ident: object) -> str: + rid_calls.append((table, str(ident))) + return f"{table}:{ident}" + + monkeypatch.setattr(sys.modules["surrealdb"], "RecordID", _fake_rid) + migrations._to_relation_row( + { + "id": "synapse:s1", + "source_id": "a.b-c", + "target_id": "x-y", + "brain_id": "default", + "type": "assoc", + "weight": 1.0, + } + ) + assert ("neuron", _to_surreal_id("a.b-c")) in rid_calls # -> a_b_c + assert ("neuron", "a_b_c") in rid_calls + assert ("neuron", "x_y") in rid_calls + + async def test_tool_events_id_routes_through_single_source(self, monkeypatch): + st, conn = _store_with_mock_conn() + monkeypatch.setattr(tool_events_mod, "_to_surreal_id", lambda s: "SENTINEL") + await st.insert_tool_events( + "b1", [{"tool_name": "t", "created_at": "2026-01-01T00:00:00+00:00"}] + ) + table, doc = conn.insert.call_args.args + assert table == "tool_events" + # The record id came from the single-source folder, not a raw .replace. + assert doc["id"] == "SENTINEL" + + +# --------------------------------------------------------------------------- # +# single-record fetches close the same read-by-id class +# --------------------------------------------------------------------------- # +class TestSingleRecordBrainScope: + async def test_get_neuron_scopes_to_brain(self): + st, conn = _store_with_mock_conn() + await st.get_neuron("abc-123") + sql, params = conn.query.call_args.args + assert sql == "SELECT * FROM neuron:abc_123 WHERE brain_id = $brain_id" + assert params == {"brain_id": "b1"} + + async def test_get_neuron_state_scopes_to_brain(self): + st, conn = _store_with_mock_conn() + await st.get_neuron_state("abc-123") + sql, params = conn.query.call_args.args + assert sql == "SELECT * FROM neuron_state:state_abc_123 WHERE brain_id = $brain_id" + assert params == {"brain_id": "b1"} + + async def test_get_synapse_scopes_to_brain(self): + st, conn = _store_with_mock_conn() + await st.get_synapse("s-1") + sql, params = conn.query.call_args.args + assert sql == "SELECT * FROM synapse:s_1 WHERE brain_id = $brain_id" + assert params == {"brain_id": "b1"} + + async def test_get_fiber_scopes_to_brain(self): + st, conn = _store_with_mock_conn() + await st.get_fiber("f-1") + sql, params = conn.query.call_args.args + assert sql == "SELECT * FROM fiber:f_1 WHERE brain_id = $brain_id" + assert params == {"brain_id": "b1"}