Skip to content

fix(storage): route the last two id sanitisers through the single source + scope every record-id fetch to the brain #63

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
RobertSigmundsson wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(storage): route the last two id sanitisers through the single source + scope every record-id fetch to the brain #63

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 3 additions & 13 deletions src/surreal_memory/storage/surrealdb/migrations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -123,17 +124,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)
Expand Down Expand Up @@ -325,8 +315,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", "")))

@RobertSigmundsson RobertSigmundsson Jul 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Second follow-up: the old _sanitize_id was dash-only, so it both duplicated the fold and under-folded. _to_surreal_id here also quietly corrects endpoints — neuron records are stored folded, so a legacy source_id/target_id with a . (or any non-- char) previously produced a neuron:{id} endpoint that didn't match the stored record. uuid4 / hash ids are unaffected.

tgt = _to_surreal_id(str(row.get("target_id", "")))
relation: dict[str, Any] = {
"id": RecordID("synapse", sid),
"in": RecordID("neuron", src),
Expand Down
91 changes: 67 additions & 24 deletions src/surreal_memory/storage/surrealdb/store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -636,12 +636,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",

@RobertSigmundsson RobertSigmundsson Jul 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You flagged the two batch methods; the same bare-select read-by-id class was also open in get_neuron, get_neuron_state, get_synapse and get_fiber, so I scoped all of them here rather than leave a half-guarantee — a caller could otherwise still read a foreign brain's record via the single-record path. brain_id is a bound param from _get_brain_id() (→ _safe_brain_id, fail-closed), and the record stays pinned in FROM, so it's still a direct fetch, not a scan. Routing through _query also gives the single getters the reconnect-on-dropped-transport retry they lacked. (device was already brain-keyed in its record id.)

brain_id=brain_id,
)
if rows:
return _row_to_neuron(rows[0])
except Exception:
pass
return None
Expand All @@ -660,17 +665,29 @@ async def get_neurons_batch(self, neuron_ids: list[str]) -> dict[str, Neuron]:
"""
if not neuron_ids:
return {}
# 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.
# Resolve the brain id ONCE, before any concurrent fetch starts, so a
# hostile/unset brain context fails closed (_safe_brain_id raises) rather
# than firing 16 unscoped selects.
brain_id = self._get_brain_id()
semaphore = asyncio.Semaphore(_BATCH_FETCH_CONCURRENCY)

async def _fetch_one(nid: str) -> tuple[str, Neuron | None]:
sid = _to_surreal_id(nid)
async with semaphore:
try:
result = await self._conn.select(f"neuron:{sid}")
rows = await self._query(
f"SELECT * FROM neuron:{sid} WHERE brain_id = $brain_id",
brain_id=brain_id,
)
except Exception:
return nid, None
if result and isinstance(result, list) and len(result) > 0:
return nid, _row_to_neuron(result[0])
if rows:
return nid, _row_to_neuron(rows[0])
return nid, None

pairs = await asyncio.gather(*(_fetch_one(nid) for nid in neuron_ids))
Expand Down Expand Up @@ -790,12 +807,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",

@RobertSigmundsson RobertSigmundsson Jul 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The ids in FROM already pass _to_surreal_id + an isalnum()/_ filter, so the WHERE only adds the brain check on records already pinned by id — no change to the 2.7.2 dashboard-perf profile (still a pinned-record fetch, not a FROM neuron scan).

brain_id=brain_id,
)
out.extend(_row_to_neuron(r) for r in rows)
return out

Expand Down Expand Up @@ -939,12 +963,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
Expand Down Expand Up @@ -1080,12 +1107,15 @@ async def add_synapses_batch(
return len(synapses)

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
Expand All @@ -1098,18 +1128,28 @@ async def get_synapses_batch(self, synapse_ids: list[str]) -> dict[str, Synapse]
"""
if not synapse_ids:
return {}
conn = self._ensure_conn()
# Scope every fetch to the current brain: like get_neurons_batch these are
# direct record-id selects, so without a brain_id filter a caller could
# read a synapse owned by another brain just by knowing (or guessing) its
# id. The FROM stays a pinned record (synapse:{sid}) — a direct fetch, not
# a table scan — and the WHERE brain_id only rejects a cross-brain hit.
# Resolve the brain id ONCE, before any concurrent fetch starts, so a
# hostile/unset brain context fails closed (_safe_brain_id raises).
brain_id = self._get_brain_id()
semaphore = asyncio.Semaphore(_BATCH_FETCH_CONCURRENCY)

async def _fetch_one(syn_id: str) -> tuple[str, Synapse | None]:
sid = _to_surreal_id(syn_id)
async with semaphore:
try:
result = await conn.select(f"synapse:{sid}")
rows = await self._query(
f"SELECT * FROM synapse:{sid} WHERE brain_id = $brain_id",
brain_id=brain_id,
)
except Exception:
return syn_id, None
if result:
return syn_id, _row_to_synapse(result[0] if isinstance(result, list) else result)
if rows:
return syn_id, _row_to_synapse(rows[0])
return syn_id, None

pairs = await asyncio.gather(*(_fetch_one(sid) for sid in synapse_ids))
Expand Down Expand Up @@ -1415,12 +1455,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
Expand Down
3 changes: 2 additions & 1 deletion src/surreal_memory/storage/surrealdb/tool_events.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down Expand Up @@ -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", ""),
Expand Down
5 changes: 4 additions & 1 deletion tests/unit/test_dashboard_perf_queries.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,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()
Expand Down
Loading
Loading