"""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. " "Use this to save intermediate findings (file paths, URLs, error details) so they are not lost across tool calls. " "Read before composing a final answer — saved findings may contain facts needed for the response.\n\n" "JSON schema:\n" " action: 'write' | 'append' | 'read' | 'clear'\n" " section: string — which section to target (goal, findings, artifacts, errors, main). Defaults to 'main'.\n" " content: string — text to write or append. REQUIRED for write and append.\n\n" "Examples (copy this structure exactly):\n" " {\"action\": \"write\", \"section\": \"findings\", \"content\": \"The config file is at /etc/app/config.yml\"}\n" " {\"action\": \"append\", \"section\": \"findings\", \"content\": \"Port is 8080\"}\n" " {\"action\": \"read\", \"section\": \"findings\"}\n" " {\"action\": \"clear\", \"section\": \"errors\"}\n\n" "Common mistakes to avoid:\n" " - Do NOT omit 'content' on write/append — the call will fail.\n" " - Do NOT pass the text as the action value (e.g. {\"action\": \"your text here\"}). " " The action MUST be exactly one of: write, append, read, clear.\n" " - Do NOT put the text inside 'section' — section is the category name, content is the text." ) parameters = { "type": "object", "properties": { "action": { "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 if this is missing or empty." ), }, }, "required": ["action"], } 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) # Support both 'action' (new) and 'op' (legacy) for backward compatibility action = params.get("action") or params.get("op") section: str | None = params.get("section") or None content: str = params.get("content", "") if not action: return ToolResult( success=False, output="", error=( "Missing required parameter 'action'. " "Expected one of: write, append, read, clear. " "Example: {\"action\": \"write\", \"section\": \"findings\", \"content\": \"...\"}" ), ) if action == "write": if not content: return ToolResult( success=False, output="", error=( "'content' is required for 'write'. " "You must provide the text to save in the 'content' field. " "Example: {\"action\": \"write\", \"section\": \"findings\", \"content\": \"The result is 42\"}" ), ) 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 action == "append": if not content: return ToolResult( success=False, output="", error=( "'content' is required for 'append'. " "You must provide the text to append in the 'content' field. " "Example: {\"action\": \"append\", \"section\": \"findings\", \"content\": \"Additional note\"}" ), ) 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 action == "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 action == "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 action: {action!r}. " f"Expected one of: write, append, read, clear. " f"Example: {{\"action\": \"write\", \"section\": \"findings\", \"content\": \"...\"}}" ), )