From f473fed7087e17e32641a4e6d867c6deabc18e99 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Fri, 24 Jul 2026 19:33:45 +0800 Subject: [PATCH] core/memory: persist session context usage and reminder band --- src/plyngent/memory/database/schema.py | 8 ++++ src/plyngent/memory/database/store.py | 51 ++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/src/plyngent/memory/database/schema.py b/src/plyngent/memory/database/schema.py index 50547f4..ddf5550 100644 --- a/src/plyngent/memory/database/schema.py +++ b/src/plyngent/memory/database/schema.py @@ -34,6 +34,14 @@ class Session(PlyngentBase): model: Mapped[str | None] = mapped_column(String(256), nullable=True) # Todo/task stack JSON for multi-step sub-tasks (optional). todo_stack: Mapped[dict[str, object] | None] = mapped_column(JSON(), nullable=True) + # Last model request context size (API prompt_tokens preferred). + last_prompt_tokens: Mapped[int | None] = mapped_column(nullable=True) + peak_prompt_tokens: Mapped[int | None] = mapped_column(nullable=True) + last_completion_tokens: Mapped[int | None] = mapped_column(nullable=True) + # "api" | "estimate" | None when never recorded. + usage_source: Mapped[str | None] = mapped_column(String(16), nullable=True) + # Highest directive-reminder band already injected (append-only checkpoints). + reminder_last_band: Mapped[int | None] = mapped_column(nullable=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 a7bac5b..75b77e1 100644 --- a/src/plyngent/memory/database/store.py +++ b/src/plyngent/memory/database/store.py @@ -69,6 +69,7 @@ class MemoryStore: await conn.run_sync(_migrate_session_workspace) await conn.run_sync(_migrate_session_llm) await conn.run_sync(_migrate_session_todo_stack) + await conn.run_sync(_migrate_session_context_usage) async def close(self) -> None: """Dispose the underlying engine.""" @@ -234,6 +235,37 @@ class MemoryStore: await session.refresh(row) return row + async def update_session_context_usage( + self, + sid: int, + *, + last_prompt_tokens: int | None = None, + peak_prompt_tokens: int | None = None, + last_completion_tokens: int | None = None, + usage_source: str | None = None, + reminder_last_band: int | None = None, + ) -> Session: + """Update remembered context size / directive-reminder band (omit to leave).""" + 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) + if last_prompt_tokens is not None: + row.last_prompt_tokens = int(last_prompt_tokens) + if peak_prompt_tokens is not None: + row.peak_prompt_tokens = int(peak_prompt_tokens) + if last_completion_tokens is not None: + row.last_completion_tokens = int(last_completion_tokens) + if usage_source is not None: + row.usage_source = usage_source + if reminder_last_band is not None: + row.reminder_last_band = int(reminder_last_band) + row.updated_at = datetime.now(UTC) + await session.commit() + await session.refresh(row) + return row + async def rename_session(self, sid: int, name: str) -> Session: """Rename a session (max 64 characters, non-empty after strip).""" cleaned = name.strip() @@ -345,3 +377,22 @@ def _migrate_session_todo_stack(sync_conn: object) -> None: columns = _session_columns(sync_conn) if "todo_stack" not in columns: _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN todo_stack JSON")) + + +def _migrate_session_context_usage(sync_conn: object) -> None: + """Add session context usage + directive reminder band columns.""" + from sqlalchemy.engine import Connection + + if not isinstance(sync_conn, Connection): + return + columns = _session_columns(sync_conn) + if "last_prompt_tokens" not in columns: + _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN last_prompt_tokens INTEGER")) + if "peak_prompt_tokens" not in columns: + _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN peak_prompt_tokens INTEGER")) + if "last_completion_tokens" not in columns: + _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN last_completion_tokens INTEGER")) + if "usage_source" not in columns: + _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN usage_source VARCHAR(16)")) + if "reminder_last_band" not in columns: + _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN reminder_last_band INTEGER"))