From b6e378b070874e0da455aaabf104eab646d1f38e Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Sat, 4 Jul 2026 22:55:55 -0700 Subject: [PATCH] feat(runs): introduce Run and ToolOutputSpill models for logging scraper invocations - Added `Run` model to track scraper invocations, including metadata such as status, input, and output. - Implemented `ToolOutputSpill` model to store context-editing spills separately from user-facing logs. - Updated middleware to handle spill placeholders and integrate with the new models. - Enhanced REST API to record runs and expose run history through new endpoints. - Adjusted tests to validate the new run logging functionality and ensure proper integration with existing capabilities. --- .../171_add_runs_and_tool_output_spills.py | 80 +++++++ .../middleware/context_editing/builder.py | 5 +- .../middleware/context_editing/middleware.py | 175 ++++++++-------- .../main_agent/middleware/stack.py | 7 +- .../main_agent/runtime/factory.py | 9 + .../subagents/shared/run_reader.py | 175 ++++++++++++++++ .../app/capabilities/core/access/agent.py | 124 ++++++++++- .../app/capabilities/core/access/rest.py | 196 +++++++++++++++++- .../app/capabilities/core/runs.py | 191 +++++++++++++++++ surfsense_backend/app/db.py | 75 +++++++ .../agents/new_chat/test_context_editing.py | 12 +- .../capabilities/access/test_agent_tools.py | 23 +- .../capabilities/access/test_rest_router.py | 114 +++++++++- .../unit/capabilities/test_run_truncation.py | 191 +++++++++++++++++ 14 files changed, 1247 insertions(+), 130 deletions(-) create mode 100644 surfsense_backend/alembic/versions/171_add_runs_and_tool_output_spills.py create mode 100644 surfsense_backend/app/agents/chat/multi_agent_chat/subagents/shared/run_reader.py create mode 100644 surfsense_backend/app/capabilities/core/runs.py create mode 100644 surfsense_backend/tests/unit/capabilities/test_run_truncation.py diff --git a/surfsense_backend/alembic/versions/171_add_runs_and_tool_output_spills.py b/surfsense_backend/alembic/versions/171_add_runs_and_tool_output_spills.py new file mode 100644 index 000000000..34c3628eb --- /dev/null +++ b/surfsense_backend/alembic/versions/171_add_runs_and_tool_output_spills.py @@ -0,0 +1,80 @@ +"""Add runs and tool_output_spills tables. + +``runs`` backs the user-facing Scraper-API logs and the agent's tool-boundary +truncation (full scraper output stored here, model sees a capped preview + id). +``tool_output_spills`` is the internal scratch store for main-agent context +spills, kept separate so the customer log stays clean. + +Revision ID: 171 +Revises: 170 +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "171" +down_revision: str | None = "170" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.execute( + """ + CREATE TABLE IF NOT EXISTS runs ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + workspace_id INTEGER NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, + user_id UUID REFERENCES "user"(id) ON DELETE CASCADE, + thread_id VARCHAR(255), + capability VARCHAR(100) NOT NULL, + origin VARCHAR(16) NOT NULL, + status VARCHAR(16) NOT NULL, + error TEXT, + input JSONB, + output_text TEXT, + item_count INTEGER NOT NULL DEFAULT 0, + char_count INTEGER NOT NULL DEFAULT 0, + duration_ms INTEGER, + cost_micros BIGINT, + created_at TIMESTAMP WITH TIME ZONE NOT NULL + ); + """ + ) + op.execute( + "CREATE INDEX IF NOT EXISTS ix_runs_workspace_id ON runs (workspace_id)" + ) + op.execute("CREATE INDEX IF NOT EXISTS ix_runs_user_id ON runs (user_id)") + op.execute("CREATE INDEX IF NOT EXISTS ix_runs_capability ON runs (capability)") + op.execute("CREATE INDEX IF NOT EXISTS ix_runs_created_at ON runs (created_at)") + op.execute( + "CREATE INDEX IF NOT EXISTS ix_runs_workspace_created " + "ON runs (workspace_id, created_at)" + ) + + op.execute( + """ + CREATE TABLE IF NOT EXISTS tool_output_spills ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + workspace_id INTEGER REFERENCES workspaces(id) ON DELETE CASCADE, + thread_id VARCHAR(255), + tool_name VARCHAR(255), + content TEXT NOT NULL, + char_count INTEGER NOT NULL DEFAULT 0, + created_at TIMESTAMP WITH TIME ZONE NOT NULL + ); + """ + ) + op.execute( + "CREATE INDEX IF NOT EXISTS ix_tool_output_spills_workspace_id " + "ON tool_output_spills (workspace_id)" + ) + op.execute( + "CREATE INDEX IF NOT EXISTS ix_tool_output_spills_created_at " + "ON tool_output_spills (created_at)" + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS tool_output_spills") + op.execute("DROP TABLE IF EXISTS runs") diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/builder.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/builder.py index 1d7a2f47f..784cb36f7 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/builder.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/builder.py @@ -3,7 +3,6 @@ from __future__ import annotations from collections.abc import Sequence -from typing import Any from langchain_core.tools import BaseTool @@ -25,7 +24,7 @@ def build_context_editing_mw( flags: AgentFeatureFlags, max_input_tokens: int | None, tools: Sequence[BaseTool], - backend_resolver: Any, + workspace_id: int | None = None, ) -> SpillingContextEditingMiddleware | None: if not enabled(flags, "enable_context_editing") or not max_input_tokens: return None @@ -46,5 +45,5 @@ def build_context_editing_mw( ) return SpillingContextEditingMiddleware( edits=[spill_edit, clear_edit], - backend_resolver=backend_resolver, + workspace_id=workspace_id, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/middleware.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/middleware.py index 39bc57c8b..32ace056a 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/middleware.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/middleware.py @@ -4,25 +4,31 @@ SpillToBackendEdit + SpillingContextEditingMiddleware. LangChain's :class:`ClearToolUsesEdit` discards old ``ToolMessage.content`` when the context-editing budget triggers, replacing the body with a fixed placeholder. That's lossy: anything the agent might want to revisit is -gone. The spill-to-disk pattern (originally from OpenCode's +gone. The spill pattern (originally from OpenCode's ``opencode/packages/opencode/src/tool/truncate.ts``) keeps the prune -behavior but writes the full original payload to the runtime backend -under ``/tool_outputs/{thread_id}/{message_id}.txt`` first. The -placeholder is then upgraded to point at the spill path so the agent -(or a subagent) can read it back on demand. +behavior but persists the full original payload first — to the +``tool_output_spills`` table — and upgrades the placeholder to a +``spill_`` reference the agent can read back with the shared +``read_run``/``search_run`` tools on demand. + +Why the DB and not the runtime filesystem backend: the previous version +wrote via ``backend.aupload_files``, which is a no-op on cloud +(``KBPostgresBackend`` raises ``NotImplementedError``) and mismatched on +desktop (paths must be ``/{mount_id}/...``), so spills were unrecoverable +in production. A table works on every deployment and needs no sandbox. Why this is a middleware subclass instead of a plain ``ContextEdit``: -``ContextEdit.apply`` is sync, but writing to the backend is async. We -capture the spill payloads inside ``apply`` and flush them via -``await backend.aupload_files(...)`` from ``awrap_model_call`` *before* -delegating to the handler, so the explore subagent can always read what -the placeholder advertises. +``ContextEdit.apply`` is sync, but the DB write is async. We generate the +spill id and capture the payload inside ``apply`` (so the placeholder can +reference the final id immediately) and flush the rows from +``awrap_model_call`` *before* delegating to the handler. """ from __future__ import annotations import logging import threading +import uuid from collections.abc import Awaitable, Callable, Sequence from copy import deepcopy from dataclasses import dataclass, field @@ -44,7 +50,6 @@ from langchain_core.messages.utils import count_tokens_approximately from langgraph.config import get_config if TYPE_CHECKING: - from deepagents.backends.protocol import BackendProtocol from langchain.agents.middleware.types import ( ModelRequest, ModelResponse, @@ -52,24 +57,27 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -DEFAULT_SPILL_PREFIX = "/tool_outputs" +# Namespace for deterministic spill ids: the same (thread, message) always maps +# to the same row, so re-running the edit on later model calls (edits apply to a +# per-call copy of the messages, never to persisted state) re-references the +# existing spill instead of inserting a duplicate every turn. +_SPILL_NAMESPACE = uuid.UUID("6ba7b810-9dad-11d1-80b4-00c04fd430c8") -def _build_spill_placeholder(spill_path: str) -> str: +def _spill_id_for(thread_id: str | None, message_key: str) -> uuid.UUID: + return uuid.uuid5(_SPILL_NAMESPACE, f"{thread_id or 'no_thread'}:{message_key}") + + +def _build_spill_placeholder(spill_id: uuid.UUID) -> str: """Build the user-facing placeholder text shown to the model.""" return ( - f"[cleared — full output at {spill_path}; ask the explore subagent to read it]" + f"[cleared — full output stored as spill_{spill_id}; " + "use read_run/search_run to read it]" ) -def _get_thread_id_or_session() -> str: - """Best-effort thread_id discovery for the spill path. - - Falls back to a process-stable string if no LangGraph config is - available (e.g. unit tests). The exact value doesn't matter as long - as it's stable within one stream so the placeholder paths line up - with the actual upload path. - """ +def _get_thread_id() -> str | None: + """Best-effort ``configurable.thread_id`` for the spill row (``None`` if absent).""" try: config = get_config() thread_id = config.get("configurable", {}).get("thread_id") @@ -77,17 +85,18 @@ def _get_thread_id_or_session() -> str: return str(thread_id) except RuntimeError: pass - return "no_thread" + return None @dataclass(slots=True) class SpillToBackendEdit(ContextEdit): - """Capture-and-replace context edit that spills full tool output to the backend. + """Capture-and-replace context edit that spills full tool output to the DB. Behaves like :class:`ClearToolUsesEdit` (same trigger / keep / exclude semantics) **and** records the original ``ToolMessage.content`` in - :attr:`pending_spills` so the wrapping middleware can flush them - before the model call. + :attr:`pending_spills` so the wrapping middleware can flush the rows to + ``tool_output_spills`` before the model call. The spill id is generated up + front so the placeholder can reference it immediately. Args: trigger: Token threshold above which the edit fires. @@ -97,8 +106,6 @@ class SpillToBackendEdit(ContextEdit): exclude_tools: Names of tools whose output is NOT spilled. clear_tool_inputs: Also clear the originating ``AIMessage.tool_calls`` args when their pair is cleared. - path_prefix: Path under the backend where spills are written. - Default ``"/tool_outputs"``. """ trigger: int = 100_000 @@ -106,12 +113,16 @@ class SpillToBackendEdit(ContextEdit): keep: int = 3 clear_tool_inputs: bool = False exclude_tools: Sequence[str] = () - path_prefix: str = DEFAULT_SPILL_PREFIX - pending_spills: list[tuple[str, bytes]] = field(default_factory=list) + # (spill_id, content_bytes, tool_name, thread_id) + pending_spills: list[tuple[uuid.UUID, bytes, str | None, str | None]] = field( + default_factory=list + ) _lock: threading.Lock = field(default_factory=threading.Lock) - def drain_pending(self) -> list[tuple[str, bytes]]: + def drain_pending( + self, + ) -> list[tuple[uuid.UUID, bytes, str | None, str | None]]: """Return and clear the pending-spill list atomically.""" with self._lock: out = list(self.pending_spills) @@ -139,7 +150,7 @@ class SpillToBackendEdit(ContextEdit): if self.keep: candidates = candidates[: -self.keep] - thread_id = _get_thread_id_or_session() + thread_id = _get_thread_id() excluded_tools = set(self.exclude_tools) for idx, tool_message in candidates: @@ -168,24 +179,23 @@ class SpillToBackendEdit(ContextEdit): if tool_name in excluded_tools: continue - message_id = tool_message.id or tool_message.tool_call_id or "unknown" - spill_path = f"{self.path_prefix}/{thread_id}/{message_id}.txt" - + message_key = tool_message.id or tool_message.tool_call_id or "unknown" + spill_id = _spill_id_for(thread_id, message_key) original = tool_message.content payload = self._encode_payload(original) with self._lock: - self.pending_spills.append((spill_path, payload)) + self.pending_spills.append((spill_id, payload, tool_name, thread_id)) messages[idx] = tool_message.model_copy( update={ "artifact": None, - "content": _build_spill_placeholder(spill_path), + "content": _build_spill_placeholder(spill_id), "response_metadata": { **tool_message.response_metadata, "context_editing": { "cleared": True, - "strategy": "spill_to_backend", - "spill_path": spill_path, + "strategy": "spill_to_db", + "spill_id": str(spill_id), }, }, } @@ -243,52 +253,30 @@ class SpillToBackendEdit(ContextEdit): ) -BackendResolver = "Callable[[Any], BackendProtocol] | BackendProtocol" - - class SpillingContextEditingMiddleware(ContextEditingMiddleware): """:class:`ContextEditingMiddleware` that flushes :class:`SpillToBackendEdit` writes. - Runs the configured edits as the parent does, then flushes any - pending spills via the supplied backend resolver before delegating - to the model handler. Spill failures are logged but never abort the - model call — the placeholder text is already in the message, so the - worst case is the agent gets a placeholder it cannot follow up on. + Runs the configured edits as the parent does, then persists any pending + spills to ``tool_output_spills`` before delegating to the model handler. + Spill failures are logged but never abort the model call — the placeholder + text is already in the message, so the worst case is the agent gets a + placeholder it cannot follow up on. """ def __init__( self, *, edits: Sequence[ContextEdit], - backend_resolver: BackendResolver | None = None, + workspace_id: int | None = None, token_count_method: str = "approximate", ) -> None: super().__init__(edits=list(edits), token_count_method=token_count_method) # type: ignore[arg-type] - self._backend_resolver = backend_resolver + self._workspace_id = workspace_id - def _resolve_backend(self, request: ModelRequest) -> BackendProtocol | None: - if self._backend_resolver is None: - return None - if callable(self._backend_resolver): - try: - from langchain.tools import ToolRuntime - - tool_runtime = ToolRuntime( - state=getattr(request, "state", {}), - context=getattr(request.runtime, "context", None), - stream_writer=getattr(request.runtime, "stream_writer", None), - store=getattr(request.runtime, "store", None), - config=getattr(request.runtime, "config", None) or {}, - tool_call_id=None, - ) - return self._backend_resolver(tool_runtime) - except Exception: - logger.exception("Failed to resolve spill backend") - return None - return self._backend_resolver # type: ignore[return-value] - - def _collect_pending(self) -> list[tuple[str, bytes]]: - out: list[tuple[str, bytes]] = [] + def _collect_pending( + self, + ) -> list[tuple[uuid.UUID, bytes, str | None, str | None]]: + out: list[tuple[uuid.UUID, bytes, str | None, str | None]] = [] for edit in self.edits: if isinstance(edit, SpillToBackendEdit): out.extend(edit.drain_pending()) @@ -321,28 +309,37 @@ class SpillingContextEditingMiddleware(ContextEditingMiddleware): pending = self._collect_pending() if pending: - backend = self._resolve_backend(request) - if backend is not None: - try: - await backend.aupload_files(pending) - except Exception: - logger.exception( - "Spill-to-backend upload failed (%d files); placeholders " - "remain in messages but content is unrecoverable", - len(pending), - ) - else: - logger.warning( - "SpillToBackendEdit produced %d pending spills but no backend " - "resolver was configured; content is unrecoverable", - len(pending), - ) + await self._flush_spills(pending) return await handler(request.override(messages=edited_messages)) + async def _flush_spills( + self, pending: list[tuple[uuid.UUID, bytes, str | None, str | None]] + ) -> None: + """Persist spilled tool outputs to the DB (best-effort).""" + from app.capabilities.core.runs import record_spill + from app.db import async_session_maker + + try: + async with async_session_maker() as session: + for spill_id, payload, tool_name, thread_id in pending: + await record_spill( + session, + spill_id=spill_id, + content=payload.decode("utf-8", errors="replace"), + workspace_id=self._workspace_id, + thread_id=thread_id, + tool_name=tool_name, + ) + except Exception: + logger.exception( + "Spill-to-DB flush failed (%d rows); placeholders remain in " + "messages but content is unrecoverable", + len(pending), + ) + __all__ = [ - "DEFAULT_SPILL_PREFIX", "ClearToolUsesEdit", "SpillToBackendEdit", "SpillingContextEditingMiddleware", diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/stack.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/stack.py index 143cceaef..7da9a7ccb 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/stack.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/stack.py @@ -7,6 +7,11 @@ retrieval are owned by the ``knowledge_base`` subagent and reached via the the main agent only sees the specialist's grounded summary. The stack here computes the workspace tree, commits any subagent-side staged writes at end of turn (cloud mode), and wires the supporting middleware. + +One deliberate, read-only exception to the pure-router stance: the main agent +also carries the ``read_run``/``search_run`` tools (added in ``runtime/factory``) +so it can follow the context-editing spill placeholder — evicted tool output is +persisted to ``tool_output_spills`` and the placeholder advertises those tools. """ from __future__ import annotations @@ -260,7 +265,7 @@ def build_main_agent_deepagent_middleware( flags=flags, max_input_tokens=max_input_tokens, tools=tools, - backend_resolver=backend_resolver, + workspace_id=workspace_id, ), build_compaction_mw(llm), build_noop_injection_mw(flags), diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/runtime/factory.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/runtime/factory.py index 331ff363f..9f1ed344c 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/runtime/factory.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/runtime/factory.py @@ -224,6 +224,15 @@ async def create_multi_agent_chat_deep_agent( additional_tools=list(additional_tools) if additional_tools else None, ) + # Read-only exception to the "main agent is a pure router" stance: the + # context-editing spill placeholder points at read_run/search_run, so the + # main agent needs those tools to follow it. See middleware/stack.py. + from app.agents.chat.multi_agent_chat.subagents.shared.run_reader import ( + build_run_reader_tools, + ) + + tools = [*list(tools), *build_run_reader_tools(workspace_id=workspace_id)] + _flags: AgentFeatureFlags = get_flags() if _flags.enable_tool_call_repair and INVALID_TOOL_NAME not in { t.name for t in tools diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/shared/run_reader.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/shared/run_reader.py new file mode 100644 index 000000000..3a37d4979 --- /dev/null +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/shared/run_reader.py @@ -0,0 +1,175 @@ +"""``read_run`` / ``search_run``: page and grep a stored run or spill by line. + +Scraper capability outputs and evicted context spills are stored full in Postgres +(``runs`` / ``tool_output_spills``); the model only ever sees a capped preview plus +a reference like ``run_`` or ``spill_``. These two tools let the agent +retrieve the rest on demand — line-based paging and pattern search — without ever +loading the whole payload into context. Every lookup is scoped to the caller's +workspace (the trust boundary). +""" + +from __future__ import annotations + +import re +from typing import Annotated +from uuid import UUID + +from langchain_core.tools import BaseTool, StructuredTool +from sqlalchemy import select + +from app.capabilities.core.runs import RUN_OUTPUT_CHAR_CAP +from app.db import Run, ToolOutputSpill, shielded_async_session + +_MAX_LIMIT = 100 +_MAX_PATTERN_LEN = 200 +"""ReDoS guard: reject/simplify absurdly long model-supplied patterns.""" + + +def _parse_ref(ref: str) -> tuple[str, UUID] | None: + """Split ``run_`` / ``spill_`` into ``(kind, uuid)``; ``None`` if malformed.""" + ref = (ref or "").strip() + for prefix, kind in (("run_", "run"), ("spill_", "spill")): + if ref.startswith(prefix): + try: + return kind, UUID(ref[len(prefix) :]) + except ValueError: + return None + return None + + +async def _load_body(ref: str, workspace_id: int) -> tuple[str, str] | str: + """Return ``(body_text, kind)`` for a ref, or an error string. + + Workspace-scoped: a ref belonging to another workspace reads as not found. + """ + parsed = _parse_ref(ref) + if parsed is None: + return ( + f"Error: '{ref}' is not a valid run reference. " + "Expected 'run_' or 'spill_'." + ) + kind, ref_id = parsed + async with shielded_async_session() as session: + if kind == "run": + row = ( + await session.execute( + select(Run.output_text).where( + Run.id == ref_id, Run.workspace_id == workspace_id + ) + ) + ).scalar_one_or_none() + else: + row = ( + await session.execute( + select(ToolOutputSpill.content).where( + ToolOutputSpill.id == ref_id, + ToolOutputSpill.workspace_id == workspace_id, + ) + ) + ).scalar_one_or_none() + if row is None: + return f"Error: {ref} not found in this workspace." + return (row or ""), kind + + +def _cap(body: str) -> str: + """Clip a response to the shared char cap with an explicit truncation note.""" + if len(body) <= RUN_OUTPUT_CHAR_CAP: + return body + return ( + body[:RUN_OUTPUT_CHAR_CAP] + + f"\n\n...[response truncated at {RUN_OUTPUT_CHAR_CAP} chars; " + "narrow with a larger offset or use search_run]..." + ) + + +def build_run_reader_tools(*, workspace_id: int) -> list[BaseTool]: + """Build the ``read_run`` / ``search_run`` tools bound to one workspace.""" + + async def _read_run( + ref: Annotated[ + str, "The run reference to read, e.g. 'run_' or 'spill_'." + ], + offset: Annotated[int, "0-based line (item) index to start from."] = 0, + limit: Annotated[int, "Max lines (items) to return (default 20)."] = 20, + ) -> str: + loaded = await _load_body(ref, workspace_id) + if isinstance(loaded, str): + return loaded + body, _kind = loaded + lines = body.split("\n") + start = max(0, offset) + count = min(max(1, limit), _MAX_LIMIT) + window = lines[start : start + count] + if not window: + return ( + f"No lines at offset {start} (total {len(lines)} lines in {ref})." + ) + header = ( + f"Showing lines {start}-{start + len(window) - 1} of {len(lines)} in {ref}:\n" + ) + return _cap(header + "\n".join(window)) + + async def _search_run( + ref: Annotated[str, "The run reference to search, e.g. 'run_'."], + pattern: Annotated[str, "Substring or regular expression to match per line."], + max_matches: Annotated[int, "Max matching lines to return (default 20)."] = 20, + ) -> str: + loaded = await _load_body(ref, workspace_id) + if isinstance(loaded, str): + return loaded + body, _kind = loaded + pattern = (pattern or "").strip() + if not pattern: + return "Error: provide a non-empty search pattern." + + matcher = _build_matcher(pattern) + limit = min(max(1, max_matches), _MAX_LIMIT) + matches: list[str] = [] + total = 0 + for idx, line in enumerate(body.split("\n")): + if matcher(line): + total += 1 + if len(matches) < limit: + matches.append(f"[{idx}] {line}") + if not matches: + return f"No lines in {ref} matched {pattern!r}." + header = ( + f"Found {total} matching line(s) in {ref}" + f"{f' (showing first {limit})' if total > limit else ''}:\n" + ) + return _cap(header + "\n".join(matches)) + + return [ + StructuredTool.from_function( + name="read_run", + description=( + "Read a stored scraper run or spilled tool output by line, in pages. " + "Use the reference from a truncated tool result (e.g. 'run_'). " + "Each line is one result item (JSON). Page with offset/limit; prefer " + "search_run when hunting for something specific." + ), + coroutine=_read_run, + ), + StructuredTool.from_function( + name="search_run", + description=( + "Search a stored scraper run or spilled tool output for lines matching " + "a substring or regular expression. Returns matching items with their " + "line index. Cheaper than reading the whole run when you know what you " + "are looking for." + ), + coroutine=_search_run, + ), + ] + + +def _build_matcher(pattern: str): + """Compile a line matcher; fall back to substring on bad/oversized regex (ReDoS guard).""" + if len(pattern) > _MAX_PATTERN_LEN: + return lambda line: pattern in line + try: + compiled = re.compile(pattern, re.IGNORECASE) + except re.error: + return lambda line: pattern.lower() in line.lower() + return lambda line: compiled.search(line) is not None diff --git a/surfsense_backend/app/capabilities/core/access/agent.py b/surfsense_backend/app/capabilities/core/access/agent.py index 003f42ae2..7ca446edb 100644 --- a/surfsense_backend/app/capabilities/core/access/agent.py +++ b/surfsense_backend/app/capabilities/core/access/agent.py @@ -1,16 +1,26 @@ """Generate the agent door from the capability registry (05). One LangChain tool per verb; each runs the same thin adapter as the REST door -(``access/rest.py``): meter-gate -> executor -> charge. The tool returns the -verb's serialized output so the model can reason over it; UI cards are the SSE -emission handler's job, not this generator's. +(``access/rest.py``): meter-gate -> executor -> charge. Every run is recorded to +the ``runs`` table (best-effort). Outputs that fit under ``RUN_OUTPUT_CHAR_CAP`` +are returned inline; larger ones are stored and the model gets a char-budgeted +preview plus a ``run_`` reference it can page with ``read_run``/``search_run``. +Those two read tools are appended to the tool list so every capability-using +subagent can follow a truncation reference without extra wiring. """ from __future__ import annotations +import time + from langchain_core.tools import BaseTool, StructuredTool from app.capabilities.core.billing import charge_capability, gate_capability +from app.capabilities.core.runs import ( + RUN_OUTPUT_CHAR_CAP, + record_run, + serialize_output, +) from app.capabilities.core.store import all_capabilities from app.capabilities.core.types import Capability, CapabilityContext from app.db import async_session_maker @@ -22,31 +32,129 @@ def build_capability_tools( workspace_id: int, capabilities: list[Capability] | None = None, ) -> list[BaseTool]: - """Emit one tool per verb (defaults to the whole registry).""" + """Emit one tool per verb (defaults to the whole registry), plus the run readers.""" caps = capabilities if capabilities is not None else all_capabilities() - return [_capability_tool(cap, workspace_id) for cap in caps] + tools = [_capability_tool(cap, workspace_id) for cap in caps] + # Deferred import: the reader lives in the agents package (which imports from + # here), so importing it lazily avoids an import-time cycle. + from app.agents.chat.multi_agent_chat.subagents.shared.run_reader import ( + build_run_reader_tools, + ) + + tools.extend(build_run_reader_tools(workspace_id=workspace_id)) + return tools + + +def _current_thread_id() -> str | None: + """Best-effort ``configurable.thread_id`` from the active LangGraph config.""" + try: + from langgraph.config import get_config + + cfg = get_config() + tid = (cfg.get("configurable") or {}).get("thread_id") + return str(tid) if tid is not None else None + except Exception: + return None def _capability_tool(capability: Capability, workspace_id: int) -> BaseTool: input_model = capability.input_schema unit = capability.billing_unit executor = capability.executor + name = capability.name async def _run(**kwargs: object) -> dict | str: payload = input_model(**kwargs) + input_dump = payload.model_dump(exclude_none=True) + thread_id = _current_thread_id() + async with async_session_maker() as session: ctx = CapabilityContext(session=session, workspace_id=workspace_id) try: await gate_capability(payload, unit, ctx) except InsufficientCreditsError as exc: return str(exc) - output = await executor(payload) + + started = time.perf_counter() + try: + output = await executor(payload) + except Exception as exc: + duration_ms = int((time.perf_counter() - started) * 1000) + async with async_session_maker() as rec_session: + await record_run( + rec_session, + workspace_id=workspace_id, + capability=name, + origin="agent", + status="error", + input=input_dump, + error=str(exc), + thread_id=thread_id, + duration_ms=duration_ms, + ) + raise + + duration_ms = int((time.perf_counter() - started) * 1000) await charge_capability(output, unit, ctx) - return output.model_dump() + + serialized = serialize_output(output) + async with async_session_maker() as rec_session: + run_id = await record_run( + rec_session, + workspace_id=workspace_id, + capability=name, + origin="agent", + status="success", + serialized=serialized, + input=input_dump, + thread_id=thread_id, + duration_ms=duration_ms, + ) + + if serialized.char_count <= RUN_OUTPUT_CHAR_CAP: + dump = output.model_dump(exclude_none=True) + if run_id is not None: + dump["run_id"] = f"run_{run_id}" + return dump + + return _build_preview(serialized, run_id) return StructuredTool.from_function( coroutine=_run, - name=capability.name.replace(".", "_"), + name=name.replace(".", "_"), description=capability.description, args_schema=input_model, ) + + +def _build_preview(serialized, run_id: str | None) -> str: + """Char-budgeted preview: whole JSONL items until the cap is spent.""" + lines = serialized.text.split("\n") + preview_lines: list[str] = [] + used = 0 + for line in lines: + cost = len(line) + 1 + if used + cost > RUN_OUTPUT_CHAR_CAP: + break + preview_lines.append(line) + used += cost + + if not preview_lines and lines: + # A single item larger than the cap: show a clipped head so the model + # still sees the shape and can page/search for the rest. + preview_lines = [lines[0][:RUN_OUTPUT_CHAR_CAP]] + + shown = len(preview_lines) + preview = "\n".join(preview_lines) + + if run_id is None: + return ( + f"{preview}\n\n...Showing {shown} of {serialized.item_count} items " + f"({serialized.char_count} chars). Full output unavailable (storage error)." + ) + return ( + f"{preview}\n\n...Showing {shown} of {serialized.item_count} items " + f"({serialized.char_count} chars). Full run stored as run_{run_id}. Use " + f"read_run('run_{run_id}', offset, limit) or search_run('run_{run_id}', " + "pattern) to inspect the rest." + ) diff --git a/surfsense_backend/app/capabilities/core/access/rest.py b/surfsense_backend/app/capabilities/core/access/rest.py index 70bcf1b7f..c6f50efcd 100644 --- a/surfsense_backend/app/capabilities/core/access/rest.py +++ b/surfsense_backend/app/capabilities/core/access/rest.py @@ -1,32 +1,77 @@ """Generate the REST door from the capability registry (05). -One typed ``POST`` per verb; each runs the same thin adapter: -authn -> workspace authz -> meter-gate -> executor -> charge -> typed output. +One typed ``POST`` per verb under ``/workspaces/{id}/scrapers/{platform}/{verb}``; +each runs the same thin adapter: authn -> workspace authz -> meter-gate -> executor +-> charge -> typed output. Every request is recorded to the ``runs`` table +(best-effort) and its id returned via the ``X-Run-Id`` header. Two ``GET`` routes +expose the run history that backs the Scraper-API logs UI. """ -from fastapi import APIRouter, Depends, HTTPException, status +import time +import uuid +from datetime import datetime + +from fastapi import APIRouter, Depends, HTTPException, Query, Response, status +from pydantic import BaseModel +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.auth.context import AuthContext from app.capabilities.core.access.rate_limit import enforce_capability_rate_limit from app.capabilities.core.billing import charge_capability, gate_capability +from app.capabilities.core.runs import record_run, serialize_output from app.capabilities.core.store import all_capabilities from app.capabilities.core.types import Capability, CapabilityContext -from app.db import get_async_session +from app.db import Run, async_session_maker, get_async_session from app.exceptions import ExternalServiceError, SurfSenseError from app.services.web_crawl_credit_service import InsufficientCreditsError from app.users import get_auth_context from app.utils.rbac import check_workspace_access +class RunSummary(BaseModel): + """Metadata row for the runs list (output body omitted).""" + + id: str + capability: str + origin: str + status: str + item_count: int + char_count: int + duration_ms: int | None + cost_micros: int | None + error: str | None + created_at: datetime + + +class RunDetail(RunSummary): + """Full run including input and stored output.""" + + thread_id: str | None + input: dict | None + output_text: str | None + + +def _origin_for(auth: AuthContext) -> str: + """Session callers are the in-app UI; PAT/system callers are the public API.""" + return "ui" if getattr(auth, "method", None) == "session" else "api" + + +async def _record_rest_run(**kwargs) -> str | None: + """Record a run on a dedicated session so it survives a failed request txn.""" + async with async_session_maker() as session: + return await record_run(session, **kwargs) + + def build_capabilities_router( capabilities: list[Capability] | None = None, ) -> APIRouter: - """Emit one typed route per verb (defaults to the whole registry).""" - router = APIRouter(tags=["capabilities"]) + """Emit one typed route per verb (defaults to the whole registry) + run history.""" + router = APIRouter(tags=["scrapers"]) caps = capabilities if capabilities is not None else all_capabilities() for capability in caps: _register_verb(router, capability) + _register_run_history(router) return router @@ -35,10 +80,13 @@ def _register_verb(router: APIRouter, capability: Capability) -> None: output_model = capability.output_schema unit = capability.billing_unit executor = capability.executor + name = capability.name + platform, _, verb = name.partition(".") async def endpoint( workspace_id: int, payload: input_model, + response: Response, session: AsyncSession = Depends(get_async_session), auth: AuthContext = Depends(get_auth_context), ): @@ -57,24 +105,152 @@ def _register_verb(router: APIRouter, capability: Capability) -> None: }, ) from exc + input_dump = payload.model_dump(exclude_none=True) + user_id = getattr(auth.user, "id", None) + origin = _origin_for(auth) + started = time.perf_counter() try: output = await executor(payload) - except (SurfSenseError, HTTPException): + except (SurfSenseError, HTTPException) as exc: + await _record_rest_run( + workspace_id=workspace_id, + capability=name, + origin=origin, + status="error", + input=input_dump, + user_id=user_id, + error=str(exc), + duration_ms=int((time.perf_counter() - started) * 1000), + ) raise except Exception as exc: + await _record_rest_run( + workspace_id=workspace_id, + capability=name, + origin=origin, + status="error", + input=input_dump, + user_id=user_id, + error=str(exc), + duration_ms=int((time.perf_counter() - started) * 1000), + ) raise ExternalServiceError( - f"The '{capability.name}' capability failed due to an upstream error.", + f"The '{name}' capability failed due to an upstream error.", code="CAPABILITY_UPSTREAM_ERROR", ) from exc + duration_ms = int((time.perf_counter() - started) * 1000) await charge_capability(output, unit, ctx) + + serialized = serialize_output(output) + run_id = await _record_rest_run( + workspace_id=workspace_id, + capability=name, + origin=origin, + status="success", + serialized=serialized, + input=input_dump, + user_id=user_id, + duration_ms=duration_ms, + ) + if run_id is not None: + response.headers["X-Run-Id"] = f"run_{run_id}" return output router.add_api_route( - f"/workspaces/{{workspace_id}}/capabilities/{capability.name}", + f"/workspaces/{{workspace_id}}/scrapers/{platform}/{verb}", endpoint, methods=["POST"], response_model=output_model, - name=f"capability:{capability.name}", + name=f"scraper:{name}", dependencies=[Depends(enforce_capability_rate_limit)], ) + + +def _register_run_history(router: APIRouter) -> None: + """Register the ``GET`` list + detail routes that back the logs UI.""" + + async def list_runs( + workspace_id: int, + session: AsyncSession = Depends(get_async_session), + auth: AuthContext = Depends(get_auth_context), + limit: int = Query(default=50, ge=1, le=100), + offset: int = Query(default=0, ge=0), + capability: str | None = Query(default=None), + run_status: str | None = Query(default=None, alias="status"), + ) -> list[RunSummary]: + await check_workspace_access(session, auth, workspace_id) + stmt = select(Run).where(Run.workspace_id == workspace_id) + if capability: + stmt = stmt.where(Run.capability == capability) + if run_status: + stmt = stmt.where(Run.status == run_status) + stmt = stmt.order_by(Run.created_at.desc()).limit(limit).offset(offset) + rows = (await session.execute(stmt)).scalars().all() + return [_to_summary(row) for row in rows] + + async def get_run( + workspace_id: int, + run_id: str, + session: AsyncSession = Depends(get_async_session), + auth: AuthContext = Depends(get_auth_context), + ) -> RunDetail: + await check_workspace_access(session, auth, workspace_id) + raw_id = run_id[len("run_") :] if run_id.startswith("run_") else run_id + try: + parsed_id = uuid.UUID(raw_id) + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="Run not found." + ) from exc + row = ( + await session.execute( + select(Run).where( + Run.id == parsed_id, Run.workspace_id == workspace_id + ) + ) + ).scalar_one_or_none() + if row is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="Run not found." + ) + return _to_detail(row) + + router.add_api_route( + "/workspaces/{workspace_id}/scrapers/runs", + list_runs, + methods=["GET"], + response_model=list[RunSummary], + name="scraper:list_runs", + ) + router.add_api_route( + "/workspaces/{workspace_id}/scrapers/runs/{run_id}", + get_run, + methods=["GET"], + response_model=RunDetail, + name="scraper:get_run", + ) + + +def _to_summary(row: Run) -> RunSummary: + return RunSummary( + id=f"run_{row.id}", + capability=row.capability, + origin=row.origin, + status=row.status, + item_count=row.item_count, + char_count=row.char_count, + duration_ms=row.duration_ms, + cost_micros=row.cost_micros, + error=row.error, + created_at=row.created_at, + ) + + +def _to_detail(row: Run) -> RunDetail: + return RunDetail( + **_to_summary(row).model_dump(), + thread_id=row.thread_id, + input=row.input, + output_text=row.output_text, + ) diff --git a/surfsense_backend/app/capabilities/core/runs.py b/surfsense_backend/app/capabilities/core/runs.py new file mode 100644 index 000000000..7db9d52c7 --- /dev/null +++ b/surfsense_backend/app/capabilities/core/runs.py @@ -0,0 +1,191 @@ +"""Shared run recorder: persist one ``runs`` row per capability invocation. + +Both doors (the agent tool adapter and the REST endpoint) call :func:`record_run` +so agent and API runs land identically. Output is serialized to JSONL (one item +per line, ``exclude_none``) so the ``read_run``/``search_run`` tools can page and +grep by line. Recording is best-effort: a failure here never fails an otherwise +successful scrape — the caller degrades gracefully on a ``None`` return. +""" + +from __future__ import annotations + +import json +import logging +import random +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING, Any + +from pydantic import BaseModel +from sqlalchemy import text + +from app.db import Run, ToolOutputSpill + +if TYPE_CHECKING: + from sqlalchemy.ext.asyncio import AsyncSession + +logger = logging.getLogger(__name__) + +RUN_OUTPUT_CHAR_CAP = 40_000 +"""~10k tokens. Capability outputs at/under this are returned inline; larger ones +are stored and previewed. Read-tool responses pass through the same cap.""" + +RUNS_RETENTION_DAYS = 30 +SPILLS_RETENTION_DAYS = 7 +_CLEANUP_BATCH = 200 +_CLEANUP_SAMPLE_RATE = 0.01 +"""ponytail: opportunistic bounded cleanup fired on ~1% of inserts. A dedicated +cron/scheduler is the upgrade path if row volume ever outpaces this.""" + + +@dataclass(frozen=True) +class SerializedOutput: + """A capability output flattened to JSONL, computed once and reused.""" + + text: str + item_count: int + char_count: int + + +def serialize_output(output: BaseModel) -> SerializedOutput: + """Flatten a capability output to JSONL (one item per line, ``exclude_none``). + + Capability outputs wrap their results in an ``items`` list; each item becomes + one JSON line so line-based paging/grep works. A non-``items`` output is dumped + as a single line. + """ + items = getattr(output, "items", None) + if isinstance(items, list): + lines = [ + json.dumps(_dump(item), default=str, ensure_ascii=False) for item in items + ] + else: + lines = [ + json.dumps( + output.model_dump(exclude_none=True), default=str, ensure_ascii=False + ) + ] + + body = "\n".join(lines) + return SerializedOutput(text=body, item_count=len(lines), char_count=len(body)) + + +def _dump(item: Any) -> Any: + if isinstance(item, BaseModel): + return item.model_dump(exclude_none=True) + return item + + +async def record_run( + session: AsyncSession, + *, + workspace_id: int, + capability: str, + origin: str, + status: str, + serialized: SerializedOutput | None = None, + input: dict | None = None, + user_id: Any | None = None, + thread_id: str | None = None, + error: str | None = None, + duration_ms: int | None = None, + cost_micros: int | None = None, +) -> str | None: + """Persist a run row and return its id, or ``None`` on failure (best-effort). + + Both doors pass a dedicated session (from ``async_session_maker``), so this + function owns the commit — recording never entangles the request transaction + and survives an executor error that leaves the request session unusable. + """ + try: + run = Run( + workspace_id=workspace_id, + user_id=user_id, + thread_id=thread_id, + capability=capability, + origin=origin, + status=status, + error=error, + input=input, + output_text=serialized.text if serialized else None, + item_count=serialized.item_count if serialized else 0, + char_count=serialized.char_count if serialized else 0, + duration_ms=duration_ms, + cost_micros=cost_micros, + ) + session.add(run) + await session.flush() + run_id = str(run.id) + await _maybe_cleanup(session, "runs", RUNS_RETENTION_DAYS) + await session.commit() + return run_id + except Exception: + logger.exception("record_run failed for capability=%s", capability) + try: + await session.rollback() + except Exception: + logger.exception("record_run rollback failed") + return None + + +async def record_spill( + session: AsyncSession, + *, + content: str, + spill_id: Any | None = None, + workspace_id: int | None = None, + thread_id: str | None = None, + tool_name: str | None = None, +) -> str | None: + """Persist a context-editing spill row and return its id, or ``None``. + + ``spill_id`` may be supplied so the caller's placeholder can reference the id + before the row is flushed (the context-editing middleware needs this). The + write is idempotent on that id: context edits re-apply on every model call + (they operate on a per-call copy of the messages), so the same spill arrives + repeatedly — an existing row is left as-is and its id returned. + """ + try: + kwargs: dict[str, Any] = {} + if spill_id is not None: + kwargs["id"] = spill_id + existing = await session.get(ToolOutputSpill, spill_id) + if existing is not None: + return str(spill_id) + spill = ToolOutputSpill( + workspace_id=workspace_id, + thread_id=thread_id, + tool_name=tool_name, + content=content, + char_count=len(content), + **kwargs, + ) + session.add(spill) + await session.flush() + spill_id = str(spill.id) + await _maybe_cleanup(session, "tool_output_spills", SPILLS_RETENTION_DAYS) + await session.commit() + return spill_id + except Exception: + logger.exception("record_spill failed") + try: + await session.rollback() + except Exception: + logger.exception("record_spill rollback failed") + return None + + +async def _maybe_cleanup(session: AsyncSession, table: str, retention_days: int) -> None: + """Delete a bounded batch of expired rows on ~1% of inserts.""" + if random.random() >= _CLEANUP_SAMPLE_RATE: + return + cutoff = datetime.now(UTC) - timedelta(days=retention_days) + # ponytail: LIMIT-bounded so a long backlog never lands as one giant delete. + # `table` is one of two hardcoded literals below — never user input. + await session.execute( + text( + f"DELETE FROM {table} WHERE id IN " + f"(SELECT id FROM {table} WHERE created_at < :cutoff LIMIT :batch)" + ), + {"cutoff": cutoff, "batch": _CLEANUP_BATCH}, + ) diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index fd64b43ea..dd13b1274 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -2779,6 +2779,81 @@ class PersonalAccessToken(BaseModel, TimestampMixin): return not self.is_expired +class Run(Base, TimestampMixin): + """One row per scraper-capability invocation, from either the agent door + or the REST/API-key door. + + Backs the user-facing Scraper-API logs and the agent's tool-boundary + truncation: the full output lives here while the model sees only a capped + preview plus this row's id. ``output_text`` is stored as JSONL (one item + per line, ``exclude_none``) so ``read_run``/``search_run`` can page and grep + by line without parsing the whole payload. Retained ~30 days via opportunistic + bounded cleanup on insert. + + ``cost_micros`` ships nullable and unpopulated in this pass; the planned + per-verb pricing rework will fill it. ``thread_id`` is a free-form string + (subagent ids look like ``"2099::task:call_x"``), so it is intentionally not + a foreign key. + """ + + __tablename__ = "runs" + __allow_unmapped__ = True + + __table_args__ = ( + Index("ix_runs_workspace_created", "workspace_id", "created_at"), + ) + + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + workspace_id = Column( + Integer, + ForeignKey("workspaces.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + user_id = Column( + UUID(as_uuid=True), + ForeignKey("user.id", ondelete="CASCADE"), + nullable=True, + index=True, + ) + thread_id = Column(String(255), nullable=True) + capability = Column(String(100), nullable=False, index=True) + origin = Column(String(16), nullable=False) + status = Column(String(16), nullable=False) + error = Column(Text, nullable=True) + input = Column(JSONB, nullable=True) + output_text = Column(Text, nullable=True) + item_count = Column(Integer, nullable=False, default=0) + char_count = Column(Integer, nullable=False, default=0) + duration_ms = Column(Integer, nullable=True) + cost_micros = Column(BigInteger, nullable=True) + + +class ToolOutputSpill(Base, TimestampMixin): + """Internal scratch store for main-agent context-editing spills. + + Kept separate from ``runs`` so customer-facing scraper logs stay clean. + The full ``ToolMessage`` content that context editing evicts is written here + and the message body is replaced with a ``spill_{id}`` placeholder the agent + can read back via ``read_run``/``search_run``. Retained ~7 days. + """ + + __tablename__ = "tool_output_spills" + __allow_unmapped__ = True + + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + workspace_id = Column( + Integer, + ForeignKey("workspaces.id", ondelete="CASCADE"), + nullable=True, + index=True, + ) + thread_id = Column(String(255), nullable=True) + tool_name = Column(String(255), nullable=True) + content = Column(Text, nullable=False) + char_count = Column(Integer, nullable=False, default=0) + + # Register model packages that live outside this file so their classes # are present in Base.metadata before configure_mappers() resolves any # string-based relationship() references. diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py b/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py index 9632fd14d..f80e38154 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py @@ -59,7 +59,7 @@ class TestSpillEdit: assert edit.pending_spills == [] def test_above_trigger_clears_and_records(self) -> None: - edit = SpillToBackendEdit(trigger=100, keep=1, path_prefix="/tool_outputs") + edit = SpillToBackendEdit(trigger=100, keep=1) msgs = _build_history(4) edit.apply(msgs, count_tokens=_approx_count) @@ -102,7 +102,9 @@ class TestSpillEdit: assert edit.drain_pending() == [] def test_placeholder_format(self) -> None: - path = "/tool_outputs/thread-1/tool-msg-0.txt" - text = _build_spill_placeholder(path) - assert path in text - assert "explore" in text # mentions the recovery agent + import uuid + + spill_id = uuid.uuid4() + text = _build_spill_placeholder(spill_id) + assert f"spill_{spill_id}" in text + assert "read_run" in text # points at the recovery tools diff --git a/surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py b/surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py index ac3db487d..b2173292f 100644 --- a/surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py +++ b/surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py @@ -64,7 +64,12 @@ def isolate(monkeypatch): return SimpleNamespace(module=mod, charge=charge, gate=gate) -async def test_registry_becomes_one_tool_per_verb(isolate): +def _verb_tool(tools, name: str): + """Pick one capability tool out of the list (readers are appended after).""" + return next(t for t in tools if t.name == name) + + +async def test_registry_becomes_one_tool_per_verb_plus_readers(isolate): caps = [ _capability(name="web.scrape", output=_EchoOutput(echoed="a")), _capability(name="web.discover", output=_EchoOutput(echoed="b"), unit=None), @@ -73,7 +78,8 @@ async def test_registry_becomes_one_tool_per_verb(isolate): tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=caps) by_name = {t.name: t for t in tools} - assert set(by_name) == {"web_scrape", "web_discover"} + # One tool per verb, plus the two shared run-reader tools. + assert set(by_name) == {"web_scrape", "web_discover", "read_run", "search_run"} assert by_name["web_scrape"].description == "web.scrape does a thing." assert by_name["web_scrape"].args_schema is _EchoInput @@ -81,17 +87,20 @@ async def test_registry_becomes_one_tool_per_verb(isolate): async def test_input_field_docs_reach_the_model(isolate): """Per-field descriptions must surface in the tool's args schema (LLM context).""" cap = _capability(name="web.scrape", output=_EchoOutput(echoed="a")) - [tool] = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap]) + tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap]) + tool = _verb_tool(tools, "web_scrape") assert tool.args["text"]["description"] == "The text to echo back." async def test_tool_runs_executor_and_returns_serialized_output(isolate): cap = _capability(name="web.scrape", output=_EchoOutput(echoed="hi there")) - [tool] = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap]) + tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap]) + tool = _verb_tool(tools, "web_scrape") result = await tool.ainvoke({"text": "ping"}) + # Fake session makes record_run fail -> no run_id key, plain serialized output. assert result == {"echoed": "hi there"} assert cap.executor.seen.text == "ping" @@ -99,7 +108,8 @@ async def test_tool_runs_executor_and_returns_serialized_output(isolate): async def test_tool_charges_owner(isolate): output = _EchoOutput(echoed="hi") cap = _capability(name="web.scrape", output=output) - [tool] = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap]) + tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap]) + tool = _verb_tool(tools, "web_scrape") await tool.ainvoke({"text": "ping"}) @@ -117,7 +127,8 @@ async def test_over_budget_returns_friendly_message(isolate): balance_micros=0, required_micros=1_000_000, ) - [tool] = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap]) + tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap]) + tool = _verb_tool(tools, "web_scrape") result = await tool.ainvoke({"text": "ping"}) diff --git a/surfsense_backend/tests/unit/capabilities/access/test_rest_router.py b/surfsense_backend/tests/unit/capabilities/access/test_rest_router.py index 63d46f0be..f5f99b0de 100644 --- a/surfsense_backend/tests/unit/capabilities/access/test_rest_router.py +++ b/surfsense_backend/tests/unit/capabilities/access/test_rest_router.py @@ -40,6 +40,7 @@ def _build_app(capabilities, monkeypatch) -> FastAPI: from app.capabilities.core.access.rate_limit import enforce_capability_rate_limit monkeypatch.setattr(rest, "check_workspace_access", _noop_async, raising=True) + monkeypatch.setattr(rest, "_record_rest_run", _fake_record, raising=True) app = FastAPI() app.include_router(rest.build_capabilities_router(capabilities), prefix="/api/v1") @@ -57,6 +58,11 @@ async def _noop_async(*args, **kwargs) -> None: return None +async def _fake_record(**kwargs) -> str: + """Stand-in for the DB-backed recorder so unit tests never touch a database.""" + return "test-run-id" + + async def _allow() -> None: return None @@ -70,11 +76,12 @@ async def test_verb_is_exposed_as_typed_post_route(monkeypatch): app = _build_app([_ECHO], monkeypatch) async with _client(app) as client: resp = await client.post( - "/api/v1/workspaces/7/capabilities/test.echo", + "/api/v1/workspaces/7/scrapers/test/echo", json={"value": "hi"}, ) assert resp.status_code == 200 assert resp.json() == {"echo": "hi"} + assert resp.headers["X-Run-Id"] == "run_test-run-id" @pytest.mark.asyncio @@ -82,7 +89,7 @@ async def test_input_is_validated_against_the_verb_schema(monkeypatch): app = _build_app([_ECHO], monkeypatch) async with _client(app) as client: resp = await client.post( - "/api/v1/workspaces/7/capabilities/test.echo", + "/api/v1/workspaces/7/scrapers/test/echo", json={"wrong": "field"}, ) assert resp.status_code == 422 @@ -95,7 +102,7 @@ def test_registered_verbs_appear_on_rest(): router = rest.build_capabilities_router() paths = {route.path for route in router.routes} - assert "/workspaces/{workspace_id}/capabilities/web.crawl" in paths + assert "/workspaces/{workspace_id}/scrapers/web/crawl" in paths @pytest.mark.asyncio @@ -113,7 +120,7 @@ async def test_over_budget_is_blocked_before_the_executor(monkeypatch): app = _build_app([_ECHO], monkeypatch) async with _client(app) as client: resp = await client.post( - "/api/v1/workspaces/7/capabilities/test.echo", + "/api/v1/workspaces/7/scrapers/test/echo", json={"value": "hi"}, ) assert resp.status_code == 402 @@ -142,7 +149,7 @@ async def test_rate_limit_blocks_the_workspace(monkeypatch): async with _client(app) as client: resp = await client.post( - "/api/v1/workspaces/7/capabilities/test.echo", + "/api/v1/workspaces/7/scrapers/test/echo", json={"value": "hi"}, ) assert resp.status_code == 429 @@ -183,7 +190,7 @@ async def test_executor_fault_becomes_502(monkeypatch): _register_surfsense_handler(app) async with _client(app) as client: resp = await client.post( - "/api/v1/workspaces/7/capabilities/test.boom", + "/api/v1/workspaces/7/scrapers/test/boom", json={"value": "hi"}, ) assert resp.status_code == 502 @@ -211,13 +218,104 @@ async def test_surfsense_error_passes_through(monkeypatch): _register_surfsense_handler(app) async with _client(app) as client: resp = await client.post( - "/api/v1/workspaces/7/capabilities/test.forbidden", + "/api/v1/workspaces/7/scrapers/test/forbidden", json={"value": "hi"}, ) assert resp.status_code == 403 assert resp.json()["code"] == "GOOGLE_SIGNIN_REQUIRED" +def _fake_run_row(**overrides): + from datetime import UTC, datetime + from uuid import uuid4 + + defaults = { + "id": uuid4(), + "capability": "test.echo", + "origin": "api", + "status": "success", + "item_count": 2, + "char_count": 42, + "duration_ms": 10, + "cost_micros": None, + "error": None, + "created_at": datetime.now(UTC), + "thread_id": None, + "input": {"value": "hi"}, + "output_text": '{"echo": "hi"}', + } + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +def _build_app_with_rows(monkeypatch, rows): + """App whose fake session answers select() with the given Run-like rows.""" + from app.capabilities.core.access import rest + + monkeypatch.setattr(rest, "check_workspace_access", _noop_async, raising=True) + + class _Result: + def scalars(self): + return self + + def all(self): + return rows + + def scalar_one_or_none(self): + return rows[0] if rows else None + + class _Session: + async def execute(self, stmt): + return _Result() + + app = FastAPI() + app.include_router(rest.build_capabilities_router([]), prefix="/api/v1") + app.dependency_overrides[get_auth_context] = lambda: SimpleNamespace(user=None) + + async def _session(): + yield _Session() + + app.dependency_overrides[get_async_session] = _session + return app + + +@pytest.mark.asyncio +async def test_runs_list_returns_metadata_without_output(monkeypatch): + row = _fake_run_row() + app = _build_app_with_rows(monkeypatch, [row]) + async with _client(app) as client: + resp = await client.get("/api/v1/workspaces/7/scrapers/runs") + assert resp.status_code == 200 + [item] = resp.json() + assert item["id"] == f"run_{row.id}" + assert item["capability"] == "test.echo" + assert "output_text" not in item # list is metadata-only + + +@pytest.mark.asyncio +async def test_run_detail_includes_output(monkeypatch): + row = _fake_run_row() + app = _build_app_with_rows(monkeypatch, [row]) + async with _client(app) as client: + resp = await client.get(f"/api/v1/workspaces/7/scrapers/runs/run_{row.id}") + assert resp.status_code == 200 + body = resp.json() + assert body["output_text"] == '{"echo": "hi"}' + assert body["input"] == {"value": "hi"} + + +@pytest.mark.asyncio +async def test_run_detail_404s(monkeypatch): + app = _build_app_with_rows(monkeypatch, []) + async with _client(app) as client: + missing = await client.get( + "/api/v1/workspaces/7/scrapers/runs/run_00000000-0000-0000-0000-000000000000" + ) + malformed = await client.get("/api/v1/workspaces/7/scrapers/runs/garbage") + assert missing.status_code == 404 + assert malformed.status_code == 404 # bad UUID must not become a 500 + + @pytest.mark.asyncio async def test_success_charges_once(monkeypatch): from unittest.mock import AsyncMock @@ -230,7 +328,7 @@ async def test_success_charges_once(monkeypatch): app = _build_app([_ECHO], monkeypatch) async with _client(app) as client: resp = await client.post( - "/api/v1/workspaces/7/capabilities/test.echo", + "/api/v1/workspaces/7/scrapers/test/echo", json={"value": "hi"}, ) assert resp.status_code == 200 diff --git a/surfsense_backend/tests/unit/capabilities/test_run_truncation.py b/surfsense_backend/tests/unit/capabilities/test_run_truncation.py new file mode 100644 index 000000000..5a7d58f78 --- /dev/null +++ b/surfsense_backend/tests/unit/capabilities/test_run_truncation.py @@ -0,0 +1,191 @@ +"""Tool-boundary truncation + run read-tool behavior (no DB). + +Covers the pure pieces of the DB-backed run log: JSONL serialization, the +char-budgeted preview (including the single-oversized-item case and the +storage-failure degrade), and the ``read_run``/``search_run`` tools' paging, +search, ReDoS fallback, and workspace scoping — all with a fake session so the +unit suite never touches a database. +""" + +from __future__ import annotations + +import contextlib +import json + +import pytest +from pydantic import BaseModel + +from app.agents.chat.multi_agent_chat.subagents.shared import run_reader +from app.capabilities.core.access.agent import _build_preview +from app.capabilities.core.runs import ( + RUN_OUTPUT_CHAR_CAP, + SerializedOutput, + serialize_output, +) + +pytestmark = pytest.mark.unit + + +class _Item(BaseModel): + id: int + name: str + note: str | None = None + + +class _Output(BaseModel): + items: list[_Item] + + +class _Scalar(BaseModel): + value: str + + +def test_serialize_output_is_jsonl_and_excludes_none(): + out = _Output(items=[_Item(id=1, name="a"), _Item(id=2, name="b", note="x")]) + result = serialize_output(out) + + lines = result.text.split("\n") + assert result.item_count == 2 + assert len(lines) == 2 + # exclude_none: the first item has no "note" key + assert "note" not in json.loads(lines[0]) + assert json.loads(lines[1])["note"] == "x" + assert result.char_count == len(result.text) + + +def test_serialize_output_without_items_is_single_line(): + result = serialize_output(_Scalar(value="hi")) + assert result.item_count == 1 + assert json.loads(result.text) == {"value": "hi"} + + +def test_preview_is_char_budgeted_and_references_run(): + # Many small items whose total blows the cap. + per_item = "y" * 500 + items = [f'{{"i": {i}, "v": "{per_item}"}}' for i in range(500)] + body = "\n".join(items) + serialized = SerializedOutput(text=body, item_count=len(items), char_count=len(body)) + + preview = _build_preview(serialized, run_id="abc") + + assert len(preview) < serialized.char_count + assert "run_abc" in preview + assert "read_run" in preview + # Only a prefix of items is shown. + assert preview.count('"i":') < len(items) + + +def test_preview_handles_single_oversized_item(): + huge = "z" * (RUN_OUTPUT_CHAR_CAP * 2) + serialized = SerializedOutput(text=huge, item_count=1, char_count=len(huge)) + + preview = _build_preview(serialized, run_id="big") + + # Still returns a clipped head rather than nothing. + assert "z" in preview + assert "run_big" in preview + assert len(preview) < serialized.char_count + + +def test_preview_degrades_when_storage_failed(): + body = "\n".join(f'{{"i": {i}}}' for i in range(200)) + serialized = SerializedOutput(text=body, item_count=200, char_count=len(body)) + + preview = _build_preview(serialized, run_id=None) + + assert "storage error" in preview + assert "run_" not in preview + + +# --- read tools ----------------------------------------------------------- + + +class _FakeResult: + def __init__(self, value): + self._value = value + + def scalar_one_or_none(self): + return self._value + + +class _FakeSession: + def __init__(self, value, calls): + self._value = value + self._calls = calls + + async def execute(self, stmt): + self._calls.append(str(stmt)) + return _FakeResult(self._value) + + +def _patch_session(monkeypatch, value, calls): + @contextlib.asynccontextmanager + async def _maker(): + yield _FakeSession(value, calls) + + monkeypatch.setattr(run_reader, "shielded_async_session", _maker) + + +def _tools(): + read_run, search_run = run_reader.build_run_reader_tools(workspace_id=7) + return read_run, search_run + + +_BODY = "\n".join(f'{{"i": {i}, "name": "item_{i}"}}' for i in range(10)) + + +@pytest.mark.asyncio +async def test_read_run_paginates(monkeypatch): + calls: list[str] = [] + _patch_session(monkeypatch, _BODY, calls) + read_run, _ = _tools() + + out = await read_run.ainvoke( + {"ref": "run_" + "0" * 8 + "-0000-0000-0000-000000000000", + "offset": 2, "limit": 3} + ) + assert "item_2" in out and "item_3" in out and "item_4" in out + assert "item_0" not in out and "item_5" not in out + # Scoped by workspace_id in the query. + assert "workspace_id" in calls[0] + + +@pytest.mark.asyncio +async def test_read_run_rejects_bad_ref(monkeypatch): + _patch_session(monkeypatch, _BODY, []) + read_run, _ = _tools() + out = await read_run.ainvoke({"ref": "not-a-ref"}) + assert "not a valid run reference" in out + + +@pytest.mark.asyncio +async def test_read_run_not_found(monkeypatch): + _patch_session(monkeypatch, None, []) + read_run, _ = _tools() + out = await read_run.ainvoke( + {"ref": "run_" + "0" * 8 + "-0000-0000-0000-000000000000"} + ) + assert "not found" in out + + +@pytest.mark.asyncio +async def test_search_run_matches(monkeypatch): + _patch_session(monkeypatch, _BODY, []) + _, search_run = _tools() + out = await search_run.ainvoke( + {"ref": "spill_" + "0" * 8 + "-0000-0000-0000-000000000000", + "pattern": "item_7"} + ) + assert "item_7" in out + assert "item_1" not in out.split("item_7")[0] + + +@pytest.mark.asyncio +async def test_search_run_falls_back_on_bad_regex(monkeypatch): + _patch_session(monkeypatch, _BODY, []) + _, search_run = _tools() + # "(" is an invalid regex -> substring fallback, must not raise. + out = await search_run.ainvoke( + {"ref": "run_" + "0" * 8 + "-0000-0000-0000-000000000000", "pattern": "("} + ) + assert "matched" in out or "No lines" in out