mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-10 22:32:16 +02:00
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.
This commit is contained in:
parent
ab747e7a49
commit
b6e378b070
14 changed files with 1247 additions and 130 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"})
|
||||
|
||||
|
|
|
|||
|
|
@ -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
191
surfsense_backend/tests/unit/capabilities/test_run_truncation.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue