mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-08 16:37:55 +08:00
memory: persist session access grants
This commit is contained in:
@@ -34,6 +34,8 @@ class Session(PlyngentBase):
|
|||||||
model: Mapped[str | None] = mapped_column(String(256), nullable=True)
|
model: Mapped[str | None] = mapped_column(String(256), nullable=True)
|
||||||
# Todo/task stack JSON for multi-step sub-tasks (optional).
|
# Todo/task stack JSON for multi-step sub-tasks (optional).
|
||||||
todo_stack: Mapped[dict[str, object] | None] = mapped_column(JSON(), nullable=True)
|
todo_stack: Mapped[dict[str, object] | None] = mapped_column(JSON(), nullable=True)
|
||||||
|
# Directory-access grants approved for this session (resolved path → mode).
|
||||||
|
access_grants: Mapped[dict[str, str] | None] = mapped_column(JSON(), nullable=True)
|
||||||
# Last model request context size (API prompt_tokens preferred).
|
# Last model request context size (API prompt_tokens preferred).
|
||||||
last_prompt_tokens: Mapped[int | None] = mapped_column(nullable=True)
|
last_prompt_tokens: Mapped[int | None] = mapped_column(nullable=True)
|
||||||
peak_prompt_tokens: Mapped[int | None] = mapped_column(nullable=True)
|
peak_prompt_tokens: Mapped[int | None] = mapped_column(nullable=True)
|
||||||
|
|||||||
@@ -241,6 +241,31 @@ class MemoryStore:
|
|||||||
await session.refresh(row)
|
await session.refresh(row)
|
||||||
return row
|
return row
|
||||||
|
|
||||||
|
async def get_session_access_grants(self, sid: int) -> dict[str, str]:
|
||||||
|
"""Return stored directory-access grants (resolved path → mode)."""
|
||||||
|
async with self._session_factory() as session:
|
||||||
|
row = await session.get(Session, sid)
|
||||||
|
if row is None:
|
||||||
|
msg = f"session not found: {sid}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
raw = row.access_grants
|
||||||
|
if raw is None:
|
||||||
|
return {}
|
||||||
|
return {str(k): str(v) for k, v in cast("dict[object, object]", raw).items()}
|
||||||
|
|
||||||
|
async def update_session_access_grants(self, sid: int, grants: Mapping[str, str]) -> Session:
|
||||||
|
"""Persist directory-access grants for a session (empty clears)."""
|
||||||
|
async with self._session_factory() as session:
|
||||||
|
row = await session.get(Session, sid)
|
||||||
|
if row is None:
|
||||||
|
msg = f"session not found: {sid}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
row.access_grants = dict(grants)
|
||||||
|
row.updated_at = datetime.now(UTC)
|
||||||
|
await session.commit()
|
||||||
|
await session.refresh(row)
|
||||||
|
return row
|
||||||
|
|
||||||
async def update_session_context_usage(
|
async def update_session_context_usage(
|
||||||
self,
|
self,
|
||||||
sid: int,
|
sid: int,
|
||||||
@@ -400,18 +425,30 @@ def _migrate_session_v3(sync_conn: object, columns: set[str]) -> None:
|
|||||||
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN reminder_last_band INTEGER"))
|
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN reminder_last_band INTEGER"))
|
||||||
|
|
||||||
|
|
||||||
|
def _migrate_session_v4(sync_conn: object, columns: set[str]) -> None:
|
||||||
|
"""Post-v0.5.0: add ``access_grants`` JSON for session directory-access grants."""
|
||||||
|
from sqlalchemy.engine import Connection
|
||||||
|
|
||||||
|
if not isinstance(sync_conn, Connection):
|
||||||
|
return
|
||||||
|
if "access_grants" not in columns:
|
||||||
|
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN access_grants JSON"))
|
||||||
|
|
||||||
|
|
||||||
# Latest schema version this build knows. Bump (and append a step to
|
# Latest schema version this build knows. Bump (and append a step to
|
||||||
# _SCHEMA_MIGRATIONS) when a new release adds or changes columns.
|
# _SCHEMA_MIGRATIONS) when a new release adds or changes columns.
|
||||||
_SCHEMA_VERSION = 3
|
_SCHEMA_VERSION = 4
|
||||||
|
|
||||||
# Ordered (version, step) pairs. Each step receives the current ``session``
|
# Ordered (version, step) pairs. Each step receives the current ``session``
|
||||||
# column set (read once) and adds only its own missing columns, so it is
|
# column set (read once) and adds only its own missing columns, so it is
|
||||||
# idempotent for DBs that were upgraded incrementally without a version stamp.
|
# idempotent for DBs that were upgraded incrementally without a version stamp.
|
||||||
# Version 1 shipped in v0.1.0, version 2 in v0.1.2, version 3 post-v0.2.0.
|
# Version 1 shipped in v0.1.0, version 2 in v0.1.2, version 3 post-v0.2.0,
|
||||||
|
# version 4 post-v0.5.0.
|
||||||
_SCHEMA_MIGRATIONS: tuple[tuple[int, MigrationStep], ...] = (
|
_SCHEMA_MIGRATIONS: tuple[tuple[int, MigrationStep], ...] = (
|
||||||
(1, _migrate_session_v1),
|
(1, _migrate_session_v1),
|
||||||
(2, _migrate_session_v2),
|
(2, _migrate_session_v2),
|
||||||
(3, _migrate_session_v3),
|
(3, _migrate_session_v3),
|
||||||
|
(4, _migrate_session_v4),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -34,6 +34,21 @@ async def test_create_session_uses_default_user(store: MemoryStore) -> None:
|
|||||||
assert session.uid == user.uid
|
assert session.uid == user.uid
|
||||||
|
|
||||||
|
|
||||||
|
async def test_session_access_grants_roundtrip(store: MemoryStore) -> None:
|
||||||
|
session = await store.create_session(name="grants")
|
||||||
|
assert await store.get_session_access_grants(session.sid) == {}
|
||||||
|
_ = await store.update_session_access_grants(
|
||||||
|
session.sid,
|
||||||
|
{"/data/datasets": "read", "/tmp/build": "exec"},
|
||||||
|
)
|
||||||
|
assert await store.get_session_access_grants(session.sid) == {
|
||||||
|
"/data/datasets": "read",
|
||||||
|
"/tmp/build": "exec",
|
||||||
|
}
|
||||||
|
_ = await store.update_session_access_grants(session.sid, {})
|
||||||
|
assert await store.get_session_access_grants(session.sid) == {}
|
||||||
|
|
||||||
|
|
||||||
async def test_session_context_usage_roundtrip(store: MemoryStore) -> None:
|
async def test_session_context_usage_roundtrip(store: MemoryStore) -> None:
|
||||||
session = await store.create_session(name="usage")
|
session = await store.create_session(name="usage")
|
||||||
row = await store.update_session_context_usage(
|
row = await store.update_session_context_usage(
|
||||||
@@ -293,9 +308,9 @@ async def test_schema_migrations_stamp_version(tmp_path: object) -> None:
|
|||||||
db_path = tmp_path / "fresh.db"
|
db_path = tmp_path / "fresh.db"
|
||||||
store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
|
store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
|
||||||
await store.close()
|
await store.close()
|
||||||
assert await _user_version(db_path) == 3
|
assert await _user_version(db_path) == 4
|
||||||
columns = await _session_columns(db_path)
|
columns = await _session_columns(db_path)
|
||||||
assert {"workspace", "provider_name", "model", "todo_stack", "last_prompt_tokens"} <= columns
|
assert {"workspace", "provider_name", "model", "todo_stack", "access_grants", "last_prompt_tokens"} <= columns
|
||||||
|
|
||||||
|
|
||||||
async def test_legacy_db_migrated_to_latest(tmp_path: object) -> None:
|
async def test_legacy_db_migrated_to_latest(tmp_path: object) -> None:
|
||||||
@@ -311,9 +326,9 @@ async def test_legacy_db_migrated_to_latest(tmp_path: object) -> None:
|
|||||||
session = await store.create_session(name="migrated")
|
session = await store.create_session(name="migrated")
|
||||||
await store.close()
|
await store.close()
|
||||||
|
|
||||||
assert await _user_version(db_path) == 3
|
assert await _user_version(db_path) == 4
|
||||||
columns = await _session_columns(db_path)
|
columns = await _session_columns(db_path)
|
||||||
assert {"workspace", "provider_name", "model", "todo_stack", "last_prompt_tokens"} <= columns
|
assert {"workspace", "provider_name", "model", "todo_stack", "access_grants", "last_prompt_tokens"} <= columns
|
||||||
assert session.sid is not None
|
assert session.sid is not None
|
||||||
|
|
||||||
|
|
||||||
@@ -335,10 +350,10 @@ async def test_partially_migrated_legacy_db_noops_remaining_steps(tmp_path: obje
|
|||||||
session = await store.create_session(name="partial", workspace=tmp_path)
|
session = await store.create_session(name="partial", workspace=tmp_path)
|
||||||
await store.close()
|
await store.close()
|
||||||
|
|
||||||
assert await _user_version(db_path) == 3
|
assert await _user_version(db_path) == 4
|
||||||
columns = await _session_columns(db_path)
|
columns = await _session_columns(db_path)
|
||||||
assert "workspace" in columns
|
assert "workspace" in columns
|
||||||
assert {"provider_name", "model", "todo_stack", "last_prompt_tokens"} <= columns
|
assert {"provider_name", "model", "todo_stack", "access_grants", "last_prompt_tokens"} <= columns
|
||||||
assert session.workspace is not None
|
assert session.workspace is not None
|
||||||
|
|
||||||
|
|
||||||
@@ -350,8 +365,8 @@ async def test_stamped_db_skips_migrations_on_reopen(tmp_path: object) -> None:
|
|||||||
db_path = tmp_path / "stamped.db"
|
db_path = tmp_path / "stamped.db"
|
||||||
store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
|
store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
|
||||||
await store.close()
|
await store.close()
|
||||||
assert await _user_version(db_path) == 3
|
assert await _user_version(db_path) == 4
|
||||||
|
|
||||||
store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
|
store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
|
||||||
await store.close()
|
await store.close()
|
||||||
assert await _user_version(db_path) == 3
|
assert await _user_version(db_path) == 4
|
||||||
|
|||||||
Reference in New Issue
Block a user