diff --git a/src/plyngent/memory/database/schema.py b/src/plyngent/memory/database/schema.py index 7c2b42c..a315ded 100644 --- a/src/plyngent/memory/database/schema.py +++ b/src/plyngent/memory/database/schema.py @@ -27,6 +27,8 @@ class Session(PlyngentBase): sid: Mapped[int] = mapped_column(primary_key=True) uid: Mapped[int] = mapped_column(ForeignKey(User.uid, ondelete="CASCADE"), index=True) name: Mapped[str] = mapped_column(String(64)) + # Absolute workspace path this chat is bound to (tools root); null = legacy unbound. + workspace: Mapped[str | None] = mapped_column(String(1024), nullable=True, index=True) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now()) updated_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), onupdate=func.now() diff --git a/src/plyngent/memory/database/store.py b/src/plyngent/memory/database/store.py index 2603601..f7b70d2 100644 --- a/src/plyngent/memory/database/store.py +++ b/src/plyngent/memory/database/store.py @@ -1,9 +1,10 @@ from __future__ import annotations +from pathlib import Path from typing import TYPE_CHECKING, Self import msgspec -from sqlalchemy import select +from sqlalchemy import select, text from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker from plyngent.config.models import DatabaseConfig @@ -21,6 +22,16 @@ DEFAULT_USER_EMAIL = "local@localhost" DEFAULT_USER_PASSWORD_HASH = "" +def normalize_workspace(path: str | Path | None) -> str | None: + """Return a stable absolute workspace path string, or None if unset.""" + if path is None: + return None + text_path = str(path).strip() + if not text_path: + return None + return str(Path(text_path).expanduser().resolve()) + + class MemoryStore: """Async persistence for users, chat sessions, and messages.""" @@ -50,9 +61,10 @@ class MemoryStore: return store async def create_schema(self) -> None: - """Create all tables if they do not exist.""" + """Create all tables if they do not exist; apply lightweight migrations.""" async with self._engine.begin() as conn: _ = await conn.run_sync(PlyngentBase.metadata.create_all) + await conn.run_sync(_migrate_session_workspace) async def close(self) -> None: """Dispose the underlying engine.""" @@ -80,13 +92,23 @@ class MemoryStore: result = await session.execute(select(User).where(User.name == name)) return result.scalar_one_or_none() - async def create_session(self, *, uid: int | None = None, name: str = "default") -> Session: - """Create a chat session for ``uid`` (default local user when omitted).""" + async def create_session( + self, + *, + uid: int | None = None, + name: str = "default", + workspace: str | Path | None = None, + ) -> Session: + """Create a chat session for ``uid`` (default local user when omitted). + + ``workspace`` is stored as a resolved absolute path when provided. + """ if uid is None: user = await self.ensure_default_user() uid = user.uid + ws = normalize_workspace(workspace) async with self._session_factory() as session: - row = Session(uid=uid, name=name) + row = Session(uid=uid, name=name, workspace=ws) session.add(row) await session.commit() await session.refresh(row) @@ -96,14 +118,22 @@ class MemoryStore: async with self._session_factory() as session: return await session.get(Session, sid) - async def list_sessions(self, *, uid: int | None = None) -> Sequence[Session]: + async def list_sessions( + self, + *, + uid: int | None = None, + workspace: str | Path | None = None, + ) -> Sequence[Session]: + """List sessions; optionally filter by user and/or bound workspace path.""" + ws = normalize_workspace(workspace) async with self._session_factory() as session: stmt = select(Session).order_by(Session.sid) if uid is not None: stmt = stmt.where(Session.uid == uid) + if ws is not None: + stmt = stmt.where(Session.workspace == ws) result = await session.execute(stmt) return result.scalars().all() - async def append_message(self, sid: int, message: AnyChatMessage) -> Message: """Append a chat message to a session with the next sequence number.""" data = msgspec.to_builtins(message) @@ -135,3 +165,16 @@ class MemoryStore: async with self._session_factory() as session: result = await session.execute(select(Message).where(Message.sid == sid).order_by(Message.seq)) return result.scalars().all() + + +def _migrate_session_workspace(sync_conn: object) -> None: + """Add ``session.workspace`` on existing SQLite DBs created before the column existed.""" + from sqlalchemy.engine import Connection + + if not isinstance(sync_conn, Connection): + return + rows = sync_conn.execute(text("PRAGMA table_info(session)")).fetchall() + columns = {str(row[1]) for row in rows} + if "workspace" in columns: + return + _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN workspace VARCHAR(1024)")) diff --git a/tests/test_memory/test_store.py b/tests/test_memory/test_store.py index 9d5e0bc..4d1797e 100644 --- a/tests/test_memory/test_store.py +++ b/tests/test_memory/test_store.py @@ -59,3 +59,63 @@ async def test_list_sessions(store: MemoryStore) -> None: ids = {s.sid for s in sessions} assert s1.sid in ids assert s2.sid in ids + + +async def test_session_workspace_binding(store: MemoryStore, tmp_path: object) -> None: + from pathlib import Path + + assert isinstance(tmp_path, Path) + a = tmp_path / "proj-a" + b = tmp_path / "proj-b" + a.mkdir() + b.mkdir() + sa = await store.create_session(name="a", workspace=a) + sb = await store.create_session(name="b", workspace=b) + assert sa.workspace == str(a.resolve()) + listed_a = await store.list_sessions(workspace=a) + assert {s.sid for s in listed_a} == {sa.sid} + listed_b = await store.list_sessions(workspace=b) + assert {s.sid for s in listed_b} == {sb.sid} + + +async def test_workspace_column_migration(tmp_path: object) -> None: + """Existing DBs without session.workspace get the column via ALTER.""" + from pathlib import Path + + from sqlalchemy import text + from sqlalchemy.ext.asyncio import create_async_engine + + assert isinstance(tmp_path, Path) + db_path = tmp_path / "legacy.db" + engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}") + async with engine.begin() as conn: + _ = await conn.execute( + text( + "CREATE TABLE user (" + "uid INTEGER PRIMARY KEY, name VARCHAR(48) UNIQUE, " + "email VARCHAR(255) UNIQUE, password_hash VARCHAR(256), " + "created_at DATETIME)" + ) + ) + _ = await conn.execute( + text( + "CREATE TABLE session (" + "sid INTEGER PRIMARY KEY, uid INTEGER, name VARCHAR(64), " + "created_at DATETIME, updated_at DATETIME)" + ) + ) + _ = await conn.execute( + text( + "CREATE TABLE message (" + "mid INTEGER PRIMARY KEY, sid INTEGER, seq INTEGER, " + "data JSON, created_at DATETIME, updated_at DATETIME)" + ) + ) + await engine.dispose() + + store = await MemoryStore.open(DatabaseConfig(url=str(db_path))) + session = await store.create_session(name="migrated", workspace=tmp_path) + from plyngent.memory.database.store import normalize_workspace + + assert session.workspace == normalize_workspace(tmp_path) + await store.close()