diff --git a/server/api/admin.py b/server/api/admin.py index 8b37722..07502d4 100644 --- a/server/api/admin.py +++ b/server/api/admin.py @@ -1231,12 +1231,13 @@ async def memory_provenance( raise HTTPException(status_code=400, detail="Invalid memory_id format") async with engine_module.get_session_factory()() as session: - mem_result = await session.execute( - select(MemoryRow).where( - MemoryRow.id == mem_uuid, - MemoryRow.subject_id == subject_id, - ) + mem_stmt = select(MemoryRow).where( + MemoryRow.id == mem_uuid, + MemoryRow.subject_id == subject_id, ) + if tenant_id is not None: + mem_stmt = mem_stmt.where(MemoryRow.tenant_id == tenant_id) + mem_result = await session.execute(mem_stmt) mem = mem_result.scalar_one_or_none() if mem is None: raise HTTPException(status_code=404, detail="Memory not found") @@ -1246,22 +1247,24 @@ async def memory_provenance( # Fetch the source episodes episodes: list[EpisodeRow] = [] if source_ep_ids: - ep_result = await session.execute( - select(EpisodeRow).where(EpisodeRow.id.in_(source_ep_ids)) - ) + ep_stmt = select(EpisodeRow).where(EpisodeRow.id.in_(source_ep_ids)) + if tenant_id is not None: + ep_stmt = ep_stmt.where(EpisodeRow.tenant_id == tenant_id) + ep_result = await session.execute(ep_stmt) episodes = list(ep_result.scalars().all()) # Sibling memories: share ≥1 source episode, are not this memory sibling_rows: list[MemoryRow] = [] if source_ep_ids: for ep_id in source_ep_ids: - sib_result = await session.execute( - select(MemoryRow).where( - MemoryRow.subject_id == subject_id, - MemoryRow.id != mem_uuid, - ep_id == any_(MemoryRow.source_episode_ids), - ) + sib_stmt = select(MemoryRow).where( + MemoryRow.subject_id == subject_id, + MemoryRow.id != mem_uuid, + ep_id == any_(MemoryRow.source_episode_ids), ) + if tenant_id is not None: + sib_stmt = sib_stmt.where(MemoryRow.tenant_id == tenant_id) + sib_result = await session.execute(sib_stmt) sibling_rows.extend(sib_result.scalars().all()) # Deduplicate by id (a sibling may share multiple episodes) seen: set[uuid_module.UUID] = set() diff --git a/tests/test_admin_memory_provenance_tenant_scope.py b/tests/test_admin_memory_provenance_tenant_scope.py new file mode 100644 index 0000000..50aa9b5 --- /dev/null +++ b/tests/test_admin_memory_provenance_tenant_scope.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import uuid +from datetime import datetime, timezone +from types import SimpleNamespace + +import pytest +from sqlalchemy.dialects import postgresql + +from server.api.admin import memory_provenance + + +class _FakeResult: + def __init__(self, rows): + self.rows = rows + + def scalar_one_or_none(self): + return self.rows[0] if self.rows else None + + def scalars(self): + return self + + def all(self): + return self.rows + + +class _FakeSession: + def __init__(self, results): + self.results = list(results) + self.statements = [] + + async def execute(self, statement): + self.statements.append(statement) + return self.results.pop(0) + + +class _FakeSessionContext: + def __init__(self, session): + self.session = session + + async def __aenter__(self): + return self.session + + async def __aexit__(self, exc_type, exc, tb): + return False + + +def _compiled(statement): + return statement.compile(dialect=postgresql.dialect()) + + +def _memory_row(*, memory_id, subject_id, tenant_id, source_episode_ids): + return SimpleNamespace( + id=memory_id, + subject_id=subject_id, + tenant_id=tenant_id, + kind="profile_fact", + content=f"{tenant_id} memory", + summary=f"{tenant_id} memory", + confidence=1.0, + status="active", + created_at=datetime.now(timezone.utc), + source_episode_ids=source_episode_ids, + ) + + +def _episode_row(*, episode_id, subject_id, tenant_id): + return SimpleNamespace( + id=episode_id, + subject_id=subject_id, + tenant_id=tenant_id, + source="test", + type="message", + payload={"text": f"{tenant_id} episode"}, + created_at=datetime.now(timezone.utc), + ) + + +@pytest.mark.anyio +async def test_memory_provenance_scopes_memory_episodes_and_siblings_by_tenant( + monkeypatch, +): + subject_id = "shared-subject" + tenant_id = "tenant-a" + source_episode_id = uuid.uuid4() + memory_id = uuid.uuid4() + sibling_id = uuid.uuid4() + + memory = _memory_row( + memory_id=memory_id, + subject_id=subject_id, + tenant_id=tenant_id, + source_episode_ids=[source_episode_id], + ) + source_episode = _episode_row( + episode_id=source_episode_id, + subject_id=subject_id, + tenant_id=tenant_id, + ) + sibling = _memory_row( + memory_id=sibling_id, + subject_id=subject_id, + tenant_id=tenant_id, + source_episode_ids=[source_episode_id], + ) + session = _FakeSession( + [ + _FakeResult([memory]), + _FakeResult([source_episode]), + _FakeResult([sibling]), + ] + ) + + monkeypatch.setattr( + "server.db.engine.get_session_factory", + lambda: lambda: _FakeSessionContext(session), + ) + + response = await memory_provenance( + subject_id, + str(memory_id), + tenant_id=tenant_id, + ) + + assert response.memory.id == str(memory_id) + assert [ep.id for ep in response.source_episodes] == [str(source_episode_id)] + assert [mem.id for mem in response.sibling_memories] == [str(sibling_id)] + + memory_stmt, episode_stmt, sibling_stmt = session.statements + for statement in (memory_stmt, episode_stmt, sibling_stmt): + compiled = _compiled(statement) + assert "tenant_id =" in str(compiled) + assert tenant_id in compiled.params.values() +