"""Session-scoped scratchpad for capturing working notes during task execution — backed by PostgreSQL KV store."""
from __future__ import annotations
from navi.tools._internal.base import Tool, ToolContext, ToolResult, current_session_id, current_user_id
# Global KV store reference — injected at startup by registry.py
_kv_store = None
def set_kv_store(kv) -> None:
"""Inject the shared KvStore instance (called once at startup)."""
global _kv_store
_kv_store = kv
def _sid(explicit: str | None = None) -> str:
return explicit or current_session_id.get() or "__default__"
def _uid(explicit: str | None = None) -> str | None:
return explicit if explicit is not None else current_user_id.get(None)
async def get_section(session_id: str, section: str, user_id: str | None = None) -> str:
"""Read one scratchpad section for the given session. Returns '' if absent."""
if _kv_store is None:
return ""
try:
val = await _kv_store.get(_uid(user_id), session_id, "scratchpad", section)
return val or ""
except Exception:
return ""
class ScratchpadTool(Tool):
name = "scratchpad"
description = (
"Working memory for the current session — for facts discovered mid-task, not for progress tracking. "
"Sections: goal (one-sentence objective), findings (key facts from tool results), "
"artifacts (file paths, URLs), errors (failures and what was tried). "
"Write 'goal' at the start of any multi-step task. "
"Read before composing a final answer — findings may contain facts needed for the response."
)
parameters = {
"type": "object",
"properties": {
"op": {
"type": "string",
"enum": ["write", "append", "read", "clear"],
"description": (
"write — overwrite a section with new content (requires 'content'). "
"append — add text to the end of an existing section (requires 'content'). "
"read — return one section (requires 'section') or all sections (omit 'section'). "
"clear — erase one section (requires 'section') or the entire pad (omit 'section')."
),
},
"section": {
"type": "string",
"description": (
"Which section to target. Standard names: "
"goal, findings, artifacts, errors, main. "
"Defaults to 'main' when omitted on write/append. "
"Omit entirely on read/clear to target all sections."
),
},
"content": {
"type": "string",
"description": (
"Text to write or append. REQUIRED for 'write' and 'append' — "
"the call will fail without it."
),
},
},
"required": ["op"],
}
def __init__(self, kv_store=None) -> None:
if kv_store is not None:
set_kv_store(kv_store)
async def execute(self, params: dict, ctx: ToolContext | None = None) -> ToolResult:
sid = _sid(ctx.session_id if ctx else None)
op = params.get("op")
section: str | None = params.get("section") or None
content: str = params.get("content", "")
if op == "write":
if not content:
return ToolResult(success=False, output="", error="'content' is required for 'write'")
key = section or "main"
if _kv_store is not None:
await _kv_store.set(_uid(ctx.user_id if ctx else None), sid, "scratchpad", key, content)
return ToolResult(success=True, output=f"[{key}] written ({len(content)} chars).")
if op == "append":
if not content:
return ToolResult(success=False, output="", error="'content' is required for 'append'")
key = section or "main"
if _kv_store is not None:
existing = await _kv_store.get(_uid(ctx.user_id if ctx else None), sid, "scratchpad", key) or ""
new = (existing + "\n" + content).lstrip("\n") if existing else content
await _kv_store.set(_uid(ctx.user_id if ctx else None), sid, "scratchpad", key, new)
return ToolResult(success=True, output=f"[{key}] updated ({len(new)} chars total).")
return ToolResult(success=True, output=f"[{key}] updated.")
if op == "read":
if _kv_store is None:
return ToolResult(success=True, output="Scratchpad is empty.")
if section is not None:
text = await _kv_store.get(_uid(ctx.user_id if ctx else None), sid, "scratchpad", section)
if not text:
return ToolResult(success=True, output=f"[{section}] is empty.")
return ToolResult(success=True, output=f"[{section}]:\n{text}")
# No section → read all
all_data = await _kv_store.get_all(_uid(ctx.user_id if ctx else None), sid, "scratchpad")
if not all_data:
return ToolResult(success=True, output="Scratchpad is empty.")
parts = [f"[{k}]:\n{v}" for k, v in all_data.items()]
return ToolResult(success=True, output="\n\n".join(parts))
if op == "clear":
if _kv_store is None:
return ToolResult(success=True, output="Scratchpad cleared.")
if section is not None:
existing = await _kv_store.get(_uid(ctx.user_id if ctx else None), sid, "scratchpad", section)
await _kv_store.delete(_uid(ctx.user_id if ctx else None), sid, "scratchpad", section)
return ToolResult(
success=True,
output=f"[{section}] cleared." if existing else f"[{section}] was already empty.",
)
await _kv_store.clear_scope(_uid(ctx.user_id if ctx else None), sid, "scratchpad")
return ToolResult(success=True, output="Scratchpad cleared.")
return ToolResult(success=False, output="", error=f"Unknown op: {op!r}")