apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-07-10 22:32:16 +02:00
Code Issues Projects Releases Packages Wiki Activity

feat(runs): introduce Run and ToolOutputSpill models for logging scraper invocations

Browse source 
- 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.
This commit is contained in:
DESKTOP-RTLN3BA\$punk 2026-07-04 22:55:55 -07:00
parent ab747e7a49
commit b6e378b070
14 changed files with 1247 additions and 130 deletions
Download patch file Download diff file Expand all files Collapse all files

80
surfsense_backend/alembic/versions/171_add_runs_and_tool_output_spills.py Normal file
View file

@ -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")

5
surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/builder.py
View file

@ -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,
)

175
surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/context_editing/middleware.py
View file

@ -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_<uuid>`` 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",

7
surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/stack.py
View file

@ -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),

9
surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/runtime/factory.py
View file

@ -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

175
surfsense_backend/app/agents/chat/multi_agent_chat/subagents/shared/run_reader.py Normal file
View file

@ -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_<uuid>`` or ``spill_<uuid>``. 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_<uuid>`` / ``spill_<uuid>`` 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_<uuid>' or 'spill_<uuid>'."
)
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_<uuid>' or 'spill_<uuid>'."
],
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_<uuid>'."],
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_<uuid>'). "
"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

124
surfsense_backend/app/capabilities/core/access/agent.py
View file

@ -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_<id>`` 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."
)

196
surfsense_backend/app/capabilities/core/access/rest.py
View file

@ -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,
)

191
surfsense_backend/app/capabilities/core/runs.py Normal file
View file

@ -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},
)

75
surfsense_backend/app/db.py
View file

@ -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.

12
surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py
View file

@ -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

23
surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py
View file

@ -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"})

114
surfsense_backend/tests/unit/capabilities/access/test_rest_router.py
View file

@ -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

191
surfsense_backend/tests/unit/capabilities/test_run_truncation.py Normal file
View file

@ -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