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

feat: api playground, other tweaks

Browse source
This commit is contained in:
DESKTOP-RTLN3BA\$punk 2026-07-06 16:45:04 -07:00
parent 7bf559cd5b
commit 50f2d095aa
104 changed files with 5482 additions and 244 deletions
Download patch file Download diff file Expand all files Collapse all files

8
surfsense_backend/.env.example
View file

@ -282,10 +282,10 @@ MICROS_PER_PAGE=1000
# PLATFORM_SCRAPE_BILLING_ENABLED=FALSE
# REDDIT_SCRAPE_MICROS_PER_ITEM=3500
# GOOGLE_SEARCH_MICROS_PER_SERP=5500
# GOOGLE_MAPS_MICROS_PER_PLACE=5000
# GOOGLE_MAPS_MICROS_PER_REVIEW=2000
# YOUTUBE_MICROS_PER_VIDEO=3500
# YOUTUBE_MICROS_PER_COMMENT=3500
# GOOGLE_MAPS_MICROS_PER_PLACE=3500
# GOOGLE_MAPS_MICROS_PER_REVIEW=1500
# YOUTUBE_MICROS_PER_VIDEO=2500
# YOUTUBE_MICROS_PER_COMMENT=1500
# Low-balance warning threshold (micro-USD), surfaced to the UI. Default $0.50.
CREDIT_LOW_BALANCE_WARNING_MICROS=500000

158
surfsense_backend/alembic/env.py
View file

@ -1,9 +1,11 @@
import asyncio
import logging
import os
import sys
from logging.config import fileConfig
import sqlalchemy as sa
from alembic.script import ScriptDirectory
from sqlalchemy import pool
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import async_engine_from_config
@ -40,6 +42,155 @@ target_metadata = Base.metadata
MIGRATION_ADVISORY_LOCK_NAMESPACE = "surfsense"
MIGRATION_ADVISORY_LOCK_NAME = "alembic_migrations"
# Migration 170 renamed searchspaces -> workspaces, so a ``workspaces`` table
# can only exist once the schema is at revision >= 170. If it exists while the
# recorded revision is missing or still pre-170, the schema did not come from
# this migration history at all -- it was created by the startup bootstrap
# (``Base.metadata.create_all`` in ``app.db.create_db_and_tables``), which
# always builds the *current* model shape. Replaying history against such a
# schema fails (e.g. migration 5's ``ALTER COLUMN ... TYPE`` is rejected
# because the column already sits in zero_publication's column list), so the
# schema is adopted by stamping head instead.
BOOTSTRAP_MARKER_TABLE = "workspaces"
RENAME_REVISION = "170"
def _stamp_head(connection: Connection, script: ScriptDirectory) -> None:
context.get_context().stamp(script, script.get_current_head())
if connection.in_transaction():
# The outer begin_transaction() is a no-op under
# transaction_per_migration, so commit explicitly.
connection.commit()
def _fast_forward_fresh_db(connection: Connection) -> bool:
"""Build a fresh (empty) DB at head via create_all instead of replaying.
Historical migrations were written against the pre-workspace-rename
schema (``searchspaces``, ``search_space_id``), while migration 0's
``create_all`` builds the *current* models -- so replaying the chain on a
fresh DB crashes as soon as a migration touches a renamed object (first
at migration 18). A fresh DB needs no history: create the head-shape
schema directly, mirror migration 0's indexes, create the Zero
publication, and stamp head. Replay remains only for legacy DBs that
genuinely contain the old objects.
ponytail: seed-data migrations (114/128 default prompts) are skipped on
this path, same as always for create_all-bootstrapped DBs; the app copes
with missing seeds. If seeds ever become mandatory, add a runtime seeding
step rather than resurrecting the replay.
"""
for table in ("documents", "searchspaces", BOOTSTRAP_MARKER_TABLE):
if connection.execute(
sa.text("SELECT to_regclass(:t)"), {"t": table}
).scalar():
return False
if connection.execute(sa.text("SELECT to_regclass('alembic_version')")).scalar():
current = connection.execute(
sa.text("SELECT version_num FROM alembic_version")
).scalar()
if current:
return False
logging.getLogger("alembic.env").info(
"Fresh database detected: creating head-shape schema via create_all "
"and stamping head instead of replaying migration history."
)
connection.execute(sa.text("CREATE EXTENSION IF NOT EXISTS vector"))
connection.execute(sa.text("CREATE EXTENSION IF NOT EXISTS pg_trgm"))
Base.metadata.create_all(bind=connection)
# Same core indexes migration 0 created (runtime setup_indexes() adds the
# rest concurrently on app boot).
connection.execute(
sa.text(
"CREATE INDEX IF NOT EXISTS document_vector_index ON documents "
"USING hnsw (embedding public.vector_cosine_ops)"
)
)
connection.execute(
sa.text(
"CREATE INDEX IF NOT EXISTS document_search_index ON documents "
"USING gin (to_tsvector('english', content))"
)
)
connection.execute(
sa.text(
"CREATE INDEX IF NOT EXISTS chucks_vector_index ON chunks "
"USING hnsw (embedding public.vector_cosine_ops)"
)
)
connection.execute(
sa.text(
"CREATE INDEX IF NOT EXISTS chucks_search_index ON chunks "
"USING gin (to_tsvector('english', content))"
)
)
from app.zero_publication import ensure_publication
ensure_publication(connection)
_stamp_head(connection, ScriptDirectory.from_config(config))
return True
def _adopt_bootstrapped_schema(connection: Connection) -> bool:
"""Stamp head instead of replaying history on a create_all-created DB.
Returns True when the schema was adopted (migrations must then be
skipped for this run).
ponytail: assumes the bootstrapped schema matches the checked-out models
(true whenever the backend booted on this checkout, since create_all runs
on every startup). If the checkout moved ahead without a backend boot,
column-level drift from the skipped migrations is possible; the upgrade
path is re-bootstrapping (boot the backend once) before stamping.
"""
marker = connection.execute(
sa.text("SELECT to_regclass(:t)"), {"t": BOOTSTRAP_MARKER_TABLE}
).scalar()
if marker is None:
return False
# Guard against a legacy-shape DB that merely had missing tables filled in
# by a later create_all: adoption requires the core tables to be in the
# current (post-rename) shape too, not just the marker table to exist.
documents_renamed = connection.execute(
sa.text(
"SELECT 1 FROM information_schema.columns "
"WHERE table_schema = current_schema() "
"AND table_name = 'documents' AND column_name = 'workspace_id'"
)
).scalar()
if not documents_renamed:
return False
current = None
if connection.execute(sa.text("SELECT to_regclass('alembic_version')")).scalar():
current = connection.execute(
sa.text("SELECT version_num FROM alembic_version")
).scalar()
script = ScriptDirectory.from_config(config)
pre_rename_revisions = {
rev.revision for rev in script.iterate_revisions(RENAME_REVISION, "base")
} - {RENAME_REVISION}
if current is not None and current not in pre_rename_revisions:
# Genuinely migration-managed at >= 170; run migrations normally.
return False
logging.getLogger("alembic.env").info(
"Adopting bootstrap-created schema (%r exists, recorded revision %r "
"predates the workspace rename): stamping %s instead of replaying "
"migration history.",
BOOTSTRAP_MARKER_TABLE,
current,
script.get_current_head(),
)
_stamp_head(connection, script)
return True
# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
@ -86,8 +237,11 @@ def do_run_migrations(connection: Connection) -> None:
lock_params,
)
try:
with context.begin_transaction():
context.run_migrations()
if not _fast_forward_fresh_db(connection) and not _adopt_bootstrapped_schema(
connection
):
with context.begin_transaction():
context.run_migrations()
finally:
connection.execute(
sa.text("SELECT pg_advisory_unlock(hashtext(:namespace), hashtext(:name))"),

3
surfsense_backend/alembic/versions/172_remove_ai_file_sort.py
View file

@ -6,9 +6,10 @@ Revises: 171
from collections.abc import Sequence
from alembic import op
import sqlalchemy as sa
from alembic import op
revision: str = "172"
down_revision: str | None = "171"
branch_labels: str | Sequence[str] | None = None

26
surfsense_backend/alembic/versions/173_add_runs_progress.py Normal file
View file

@ -0,0 +1,26 @@
"""Add runs.progress for streamed scraper run progress.
Stores the coarse (throttled) progress log captured during a run; the live
fine-grained stream stays ephemeral (in-process bus / SSE only). Nullable so
existing rows and the sync/agent doors that don't report progress are unaffected.
Revision ID: 173
Revises: 172
"""
from collections.abc import Sequence
from alembic import op
revision: str = "173"
down_revision: str | None = "172"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
op.execute("ALTER TABLE runs ADD COLUMN IF NOT EXISTS progress JSONB")
def downgrade() -> None:
op.execute("ALTER TABLE runs DROP COLUMN IF EXISTS progress")

2
surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/index.py
View file

@ -22,8 +22,8 @@ from langchain_core.tools import BaseTool
from app.agents.chat.multi_agent_chat.shared.permissions import Rule, Ruleset
from .calendar.index import load_tools as load_calendar_tools
from .gmail.index import load_tools as load_gmail_tools
from .get_connected_accounts import create_get_connected_accounts_tool
from .gmail.index import load_tools as load_gmail_tools
NAME = "mcp_discovery"

1
surfsense_backend/app/agents/chat/multi_agent_chat/subagents/registry.py
View file

@ -39,6 +39,7 @@ from app.agents.chat.multi_agent_chat.subagents.builtins.web_crawler.agent impor
from app.agents.chat.multi_agent_chat.subagents.builtins.youtube.agent import (
build_subagent as build_youtube_subagent,
)
# File connectors stay native — they enrich the knowledge base. Every other
# connector (Slack/Jira/Linear/ClickUp/Airtable/Notion/Confluence/Gmail/
# Calendar) migrated to hosted MCP under ``mcp_discovery``; Discord/Teams/Luma

20
surfsense_backend/app/app.py
View file

@ -613,6 +613,25 @@ async def _warm_embedding_model() -> None:
)
async def _sweep_stale_scraper_runs() -> None:
"""Fail scraper runs left ``running`` by a previous process (single-process).
The async scraper door tracks in-flight runs as ``running``; a restart kills
those background tasks, so any such row at boot is dead. Non-fatal.
"""
logger = logging.getLogger(__name__)
try:
from app.capabilities.core.runs import fail_stale_running_runs
from app.db import async_session_maker
async with async_session_maker() as session:
swept = await fail_stale_running_runs(session)
if swept:
logger.info("[startup] Marked %d stale running scraper run(s) as error", swept)
except Exception:
logger.warning("[startup] Stale scraper-run sweep failed (non-fatal)", exc_info=True)
@asynccontextmanager
async def lifespan(app: FastAPI):
# Tune GC: lower gen-2 threshold so long-lived garbage is collected
@ -623,6 +642,7 @@ async def lifespan(app: FastAPI):
_enable_slow_callback_logging(threshold_sec=0.5)
init_otel(app)
await create_db_and_tables()
await _sweep_stale_scraper_runs()
await setup_checkpointer_tables()
initialize_openrouter_integration()
_start_openrouter_background_refresh()

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

@ -16,6 +16,7 @@ import time
from langchain_core.tools import BaseTool, StructuredTool
from app.capabilities.core.billing import charge_capability, gate_capability
from app.capabilities.core.progress import progress_scope
from app.capabilities.core.runs import (
RUN_OUTPUT_CHAR_CAP,
record_run,
@ -68,49 +69,55 @@ def _capability_tool(capability: Capability, workspace_id: int) -> BaseTool:
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)
# A buffer-only reporter: coarse progress lands in ``runs.progress`` and,
# because we're inside a LangGraph tool call, ``emit_progress`` also fires
# ``scraper_progress`` custom events that surface on the chat thinking step.
with progress_scope() as reporter:
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)
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,
progress=reporter.coarse,
)
raise
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
cost_micros = await charge_capability(output, unit, ctx)
duration_ms = int((time.perf_counter() - started) * 1000)
cost_micros = await charge_capability(output, unit, ctx)
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,
cost_micros=cost_micros,
)
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,
cost_micros=cost_micros,
progress=reporter.coarse,
)
if serialized.char_count <= RUN_OUTPUT_CHAR_CAP:
dump = output.model_dump(exclude_none=True)

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

@ -3,15 +3,25 @@
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.
(best-effort) and its id returned via the ``X-Run-Id`` header.
Runs can also be started in **async mode** (``?mode=async``): the POST inserts a
``running`` row, spawns the scrape as a background task, and returns ``202`` with
the run id. The client then tails ``GET .../runs/{run_id}/events`` (SSE) for live
progress and a terminal ``run.finished`` event, or cancels via
``POST .../runs/{run_id}/cancel``. Two ``GET`` routes expose the run history that
backs the Scraper-API logs UI.
"""
import asyncio
import json
import logging
import time
import uuid
from datetime import datetime
from fastapi import APIRouter, Depends, HTTPException, Query, Response, status
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
@ -19,7 +29,14 @@ 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.events import run_event_bus
from app.capabilities.core.progress import progress_scope
from app.capabilities.core.runs import (
create_pending_run,
finalize_run,
record_run,
serialize_output,
)
from app.capabilities.core.store import all_capabilities
from app.capabilities.core.types import Capability, CapabilityContext
from app.db import Run, async_session_maker, get_async_session
@ -28,9 +45,27 @@ from app.services.web_crawl_credit_service import InsufficientCreditsError
from app.users import get_auth_context
from app.utils.rbac import check_workspace_access
logger = logging.getLogger(__name__)
_HEARTBEAT_SEC = 10
_SSE_HEADERS = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
}
class CapabilitySummary(BaseModel):
"""A verb's identity + input/output JSON schemas, for the playground UI."""
name: str
description: str
input_schema: dict
output_schema: dict
class RunSummary(BaseModel):
"""Metadata row for the runs list (output body omitted)."""
"""Metadata row for the runs list (output body + progress log omitted)."""
id: str
capability: str
@ -45,11 +80,12 @@ class RunSummary(BaseModel):
class RunDetail(RunSummary):
"""Full run including input and stored output."""
"""Full run including input, stored output, and the coarse progress log."""
thread_id: str | None
input: dict | None
output_text: str | None
progress: list[dict] | None
def _origin_for(auth: AuthContext) -> str:
@ -57,6 +93,14 @@ def _origin_for(auth: AuthContext) -> str:
return "ui" if getattr(auth, "method", None) == "session" else "api"
def _now_ms() -> int:
return int(time.time() * 1000)
def _sse(event: dict) -> str:
return f"data: {json.dumps(event, default=str)}\n\n"
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:
@ -71,10 +115,164 @@ def build_capabilities_router(
caps = capabilities if capabilities is not None else all_capabilities()
for capability in caps:
_register_verb(router, capability)
_register_capabilities_list(router, caps)
_register_run_history(router)
return router
def _register_capabilities_list(
router: APIRouter, capabilities: list[Capability]
) -> None:
"""Register the ``GET`` that lists verbs + their input schemas for the UI."""
summaries = [
CapabilitySummary(
name=capability.name,
description=capability.description,
input_schema=capability.input_schema.model_json_schema(),
output_schema=capability.output_schema.model_json_schema(),
)
for capability in capabilities
]
async def list_capabilities(
workspace_id: int,
session: AsyncSession = Depends(get_async_session),
auth: AuthContext = Depends(get_auth_context),
) -> list[CapabilitySummary]:
await check_workspace_access(session, auth, workspace_id)
return summaries
router.add_api_route(
"/workspaces/{workspace_id}/scrapers/capabilities",
list_capabilities,
methods=["GET"],
response_model=list[CapabilitySummary],
name="scraper:list_capabilities",
)
async def _execute_async_run(
*,
run_id: str,
workspace_id: int,
capability: str,
unit,
executor,
payload,
) -> None:
"""Run a scrape in the background: stream progress, charge, finalize the row.
Owns its own DB sessions (the request session is long gone). Cancellation is
finalized by the cancel endpoint, so here we simply let ``CancelledError``
propagate. Every other failure finalizes the row as ``error`` and emits a
terminal event so subscribers unblock.
"""
prefixed = f"run_{run_id}"
started = time.perf_counter()
with progress_scope(run_id=run_id, bus=run_event_bus) as reporter:
run_event_bus.publish(
run_id,
{
"type": "run.started",
"run_id": prefixed,
"capability": capability,
"ts": _now_ms(),
},
)
try:
output = await executor(payload)
except asyncio.CancelledError:
raise
except (SurfSenseError, HTTPException) as exc:
await _finalize_async(
run_id, status="error", error=str(exc), started=started,
progress=reporter.coarse,
)
_publish_finished(run_id, "error", error=str(exc))
return
except Exception:
logger.exception("async run %s failed with an upstream error", run_id)
await _finalize_async(
run_id,
status="error",
error=f"The '{capability}' capability failed due to an upstream error.",
started=started,
progress=reporter.coarse,
)
_publish_finished(run_id, "error", error="upstream error")
return
duration_ms = int((time.perf_counter() - started) * 1000)
cost_micros: int | None
try:
async with async_session_maker() as session:
ctx = CapabilityContext(session=session, workspace_id=workspace_id)
cost_micros = await charge_capability(output, unit, ctx)
except Exception:
logger.exception("charge failed for async run %s", run_id)
cost_micros = None
serialized = serialize_output(output)
await _finalize_async(
run_id,
status="success",
serialized=serialized,
started=started,
duration_ms=duration_ms,
cost_micros=cost_micros,
progress=reporter.coarse,
)
_publish_finished(run_id, "success", item_count=serialized.item_count)
async def _finalize_async(
run_id: str,
*,
status: str,
serialized=None,
error: str | None = None,
started: float | None = None,
duration_ms: int | None = None,
cost_micros: int | None = None,
progress: list[dict] | None = None,
) -> None:
if duration_ms is None and started is not None:
duration_ms = int((time.perf_counter() - started) * 1000)
async with async_session_maker() as session:
await finalize_run(
session,
run_id=run_id,
status=status,
serialized=serialized,
error=error,
duration_ms=duration_ms,
cost_micros=cost_micros,
progress=progress,
)
def _publish_finished(run_id: str, status: str, **extra) -> None:
"""Emit the terminal event to subscribers, then drop the run's bus state."""
event = {
"type": "run.finished",
"run_id": f"run_{run_id}",
"status": status,
"ts": _now_ms(),
}
event.update(extra)
run_event_bus.publish(run_id, event)
run_event_bus.close(run_id)
def _log_task_result(run_id: str, task: asyncio.Task) -> None:
if task.cancelled():
return
exc = task.exception()
if exc is not None:
logger.error("async run %s task crashed: %r", run_id, exc)
def _register_verb(router: APIRouter, capability: Capability) -> None:
input_model = capability.input_schema
output_model = capability.output_schema
@ -89,6 +287,7 @@ def _register_verb(router: APIRouter, capability: Capability) -> None:
response: Response,
session: AsyncSession = Depends(get_async_session),
auth: AuthContext = Depends(get_auth_context),
mode: str = Query(default="sync", pattern="^(sync|async)$"),
):
await check_workspace_access(session, auth, workspace_id)
ctx = CapabilityContext(session=session, workspace_id=workspace_id)
@ -108,52 +307,89 @@ def _register_verb(router: APIRouter, capability: Capability) -> None:
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) as exc:
await _record_rest_run(
if mode == "async":
run_id = await create_pending_run(
session,
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(
if run_id is None:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Could not start run.",
)
task = asyncio.create_task(
_execute_async_run(
run_id=run_id,
workspace_id=workspace_id,
capability=name,
unit=unit,
executor=executor,
payload=payload,
)
)
run_event_bus.register_task(run_id, task)
task.add_done_callback(lambda t: _log_task_result(run_id, t))
return JSONResponse(
status_code=status.HTTP_202_ACCEPTED,
content={"run_id": f"run_{run_id}", "status": "running"},
)
# Sync mode: block until done, persisting the coarse progress log.
with progress_scope() as reporter:
started = time.perf_counter()
try:
output = await executor(payload)
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),
progress=reporter.coarse,
)
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),
progress=reporter.coarse,
)
raise ExternalServiceError(
f"The '{name}' capability failed due to an upstream error.",
code="CAPABILITY_UPSTREAM_ERROR",
) from exc
duration_ms = int((time.perf_counter() - started) * 1000)
cost_micros = 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="error",
status="success",
serialized=serialized,
input=input_dump,
user_id=user_id,
error=str(exc),
duration_ms=int((time.perf_counter() - started) * 1000),
duration_ms=duration_ms,
cost_micros=cost_micros,
progress=reporter.coarse,
)
raise ExternalServiceError(
f"The '{name}' capability failed due to an upstream error.",
code="CAPABILITY_UPSTREAM_ERROR",
) from exc
duration_ms = int((time.perf_counter() - started) * 1000)
cost_micros = 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,
cost_micros=cost_micros,
)
if run_id is not None:
response.headers["X-Run-Id"] = f"run_{run_id}"
return output
@ -168,8 +404,33 @@ def _register_verb(router: APIRouter, capability: Capability) -> None:
)
def _parse_run_uuid(run_id: str) -> uuid.UUID:
raw = run_id[len("run_") :] if run_id.startswith("run_") else run_id
try:
return uuid.UUID(raw)
except ValueError as exc:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND, detail="Run not found."
) from exc
async def _load_run(
session: AsyncSession, workspace_id: int, parsed_id: uuid.UUID
) -> Run:
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 row
def _register_run_history(router: APIRouter) -> None:
"""Register the ``GET`` list + detail routes that back the logs UI."""
"""Register the run list/detail + live-events + cancel routes."""
async def list_runs(
workspace_id: int,
@ -197,26 +458,106 @@ def _register_run_history(router: APIRouter) -> None:
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."
)
parsed_id = _parse_run_uuid(run_id)
row = await _load_run(session, workspace_id, parsed_id)
return _to_detail(row)
async def stream_run_events(
workspace_id: int,
run_id: str,
session: AsyncSession = Depends(get_async_session),
auth: AuthContext = Depends(get_auth_context),
):
"""SSE tail of a run's progress: replay buffered events, then live.
Note: the request ``session`` must not be used inside the generator (it
is torn down once the response starts streaming) — the generator opens
its own session for the terminal snapshot.
"""
await check_workspace_access(session, auth, workspace_id)
parsed_id = _parse_run_uuid(run_id)
await _load_run(session, workspace_id, parsed_id) # authz + 404
raw = str(parsed_id)
async def gen():
queue = run_event_bus.subscribe(raw)
try:
replayed = list(run_event_bus.replay(raw))
for event in replayed:
yield _sse(event)
if any(e.get("type") == "run.finished" for e in replayed):
return
if run_event_bus.get_task(raw) is None:
# Not actively streaming (finished before we attached, or a
# sync/agent run) — snapshot the terminal state and close.
async with async_session_maker() as snap_session:
row = (
await snap_session.execute(
select(Run).where(
Run.id == parsed_id,
Run.workspace_id == workspace_id,
)
)
).scalar_one_or_none()
yield _sse(
{
"type": "run.finished",
"run_id": f"run_{raw}",
"status": row.status if row else "error",
"item_count": row.item_count if row else 0,
"error": row.error if row else None,
"ts": _now_ms(),
}
)
return
while True:
try:
event = await asyncio.wait_for(
queue.get(), timeout=_HEARTBEAT_SEC
)
except TimeoutError:
yield _sse({"type": "run.heartbeat", "ts": _now_ms()})
continue
yield _sse(event)
if event.get("type") == "run.finished":
return
finally:
run_event_bus.unsubscribe(raw, queue)
return StreamingResponse(
gen(), media_type="text/event-stream", headers=_SSE_HEADERS
)
async def cancel_run(
workspace_id: int,
run_id: str,
session: AsyncSession = Depends(get_async_session),
auth: AuthContext = Depends(get_auth_context),
):
await check_workspace_access(session, auth, workspace_id)
parsed_id = _parse_run_uuid(run_id)
row = await _load_run(session, workspace_id, parsed_id)
if row.status != "running":
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail="Run is not in progress.",
)
raw = str(parsed_id)
task = run_event_bus.get_task(raw)
if task is not None and not task.done():
task.cancel()
# No output produced -> nothing charged. ponytail: any pre-cancel captcha
# attempts go unbilled; upgrade path is charging from progress counters.
async with async_session_maker() as cancel_session:
await finalize_run(
cancel_session,
run_id=raw,
status="cancelled",
error="Cancelled by user",
)
_publish_finished(raw, "cancelled")
return JSONResponse(content={"run_id": f"run_{raw}", "status": "cancelled"})
router.add_api_route(
"/workspaces/{workspace_id}/scrapers/runs",
list_runs,
@ -231,6 +572,18 @@ def _register_run_history(router: APIRouter) -> None:
response_model=RunDetail,
name="scraper:get_run",
)
router.add_api_route(
"/workspaces/{workspace_id}/scrapers/runs/{run_id}/events",
stream_run_events,
methods=["GET"],
name="scraper:stream_run_events",
)
router.add_api_route(
"/workspaces/{workspace_id}/scrapers/runs/{run_id}/cancel",
cancel_run,
methods=["POST"],
name="scraper:cancel_run",
)
def _to_summary(row: Run) -> RunSummary:
@ -254,4 +607,5 @@ def _to_detail(row: Run) -> RunDetail:
thread_id=row.thread_id,
input=row.input,
output_text=row.output_text,
progress=row.progress,
)

97
surfsense_backend/app/capabilities/core/events.py Normal file
View file

@ -0,0 +1,97 @@
"""In-process pub/sub bus for streaming a scraper run's progress over SSE.
One process, one bus (the module-level :data:`run_event_bus`). Per ``run_id`` it
holds three things:
* a set of subscriber ``asyncio.Queue`` s — one per open SSE connection;
* a bounded ring buffer of recent events so a late/reconnecting subscriber can
replay what it missed before tailing live;
* the background ``asyncio.Task`` running the scrape, so the cancel endpoint can
reach it.
State for a run is dropped when it reaches a terminal event (``run.finished``);
a client that connects *after* that reads the final snapshot from the ``runs``
row instead. ``ponytail:`` single-process only — a multi-worker deployment needs
Redis pub/sub (or Postgres LISTEN/NOTIFY) behind this same interface.
"""
from __future__ import annotations
import asyncio
import logging
from collections import deque
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Iterable
logger = logging.getLogger(__name__)
_BUFFER_SIZE = 500
_SUBSCRIBER_QUEUE_SIZE = 1000
class RunEventBus:
"""Fan-out of per-run progress events to live SSE subscribers."""
def __init__(self, *, buffer_size: int = _BUFFER_SIZE) -> None:
self._subscribers: dict[str, set[asyncio.Queue[dict[str, Any]]]] = {}
self._buffers: dict[str, deque[dict[str, Any]]] = {}
self._tasks: dict[str, asyncio.Task[Any]] = {}
self._buffer_size = buffer_size
# -- task registry (for cancellation) --------------------------------
def register_task(self, run_id: str, task: asyncio.Task[Any]) -> None:
self._tasks[run_id] = task
def get_task(self, run_id: str) -> asyncio.Task[Any] | None:
return self._tasks.get(run_id)
# -- publish / subscribe ---------------------------------------------
def publish(self, run_id: str, event: dict[str, Any]) -> None:
"""Buffer an event and fan it out to every live subscriber.
Called from within the event loop by the running scrape. Full
subscriber queues drop the event rather than block the scrape;
the replay buffer still preserves recent history.
"""
buffer = self._buffers.setdefault(run_id, deque(maxlen=self._buffer_size))
buffer.append(event)
for queue in list(self._subscribers.get(run_id, ())):
try:
queue.put_nowait(event)
except asyncio.QueueFull:
logger.debug("run %s: subscriber queue full, dropping event", run_id)
def subscribe(self, run_id: str) -> asyncio.Queue[dict[str, Any]]:
queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(
maxsize=_SUBSCRIBER_QUEUE_SIZE
)
self._subscribers.setdefault(run_id, set()).add(queue)
return queue
def unsubscribe(self, run_id: str, queue: asyncio.Queue[dict[str, Any]]) -> None:
subscribers = self._subscribers.get(run_id)
if subscribers is None:
return
subscribers.discard(queue)
if not subscribers:
self._subscribers.pop(run_id, None)
def replay(self, run_id: str) -> Iterable[dict[str, Any]]:
return list(self._buffers.get(run_id, ()))
def close(self, run_id: str) -> None:
"""Drop all state for a finished run.
Terminal events are published *before* this, so any live subscriber has
already received ``run.finished`` in its queue and will break its loop.
"""
self._buffers.pop(run_id, None)
self._subscribers.pop(run_id, None)
self._tasks.pop(run_id, None)
run_event_bus = RunEventBus()

163
surfsense_backend/app/capabilities/core/progress.py Normal file
View file

@ -0,0 +1,163 @@
"""Live progress plumbing shared by every capability run.
A single module-level :func:`emit_progress` lets the proprietary scraper code
report what it is doing without knowing *who* is listening. The active listener
is held in a :class:`contextvars.ContextVar`, so:
* the REST async door sets a reporter bound to the run's bus channel — every
event is streamed live over SSE and coarse ones are buffered for persistence;
* the REST sync door and the agent door set a buffer-only reporter — coarse
events still land in ``runs.progress`` and (in a chat/graph context) surface
as ``scraper_progress`` custom events on the active thinking step;
* outside any run (unit tests calling a scraper directly) the var is unset and
:func:`emit_progress` is a **no-op**, so scraper code can call it freely.
The event shape is one flexible dict, not a class hierarchy::
{"type": "run.progress", "ts": <epoch ms>, "phase": str,
"message"?: str, "current"?: int, "total"?: int, "unit"?: str,
"detail"?: {...}}
"""
from __future__ import annotations
import contextlib
import logging
import time
from contextvars import ContextVar
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Iterator
from app.capabilities.core.events import RunEventBus
logger = logging.getLogger(__name__)
# Coarse counters are persisted/chat-surfaced at most once per this window per
# (phase, unit) key; the raw fine-grained stream is unthrottled on the bus.
_COUNTER_THROTTLE_MS = 1000
# ponytail: bound the persisted coarse log so a runaway scrape can't write an
# unbounded JSONB blob. Upgrade path: page the log into its own table.
_MAX_COARSE_EVENTS = 1000
class ProgressReporter:
"""Sink for one run's progress events.
Publishes every event to the bus (live SSE tail) and keeps a throttled,
bounded list of *coarse* events for persistence + chat surfacing.
"""
def __init__(
self, *, run_id: str | None = None, bus: RunEventBus | None = None
) -> None:
self.run_id = run_id
self.bus = bus
self.coarse: list[dict[str, Any]] = []
self._last_phase: str | None = None
self._last_counter_ts: dict[tuple[str, str | None], int] = {}
def handle(self, event: dict[str, Any]) -> None:
# Live tail: every event, unthrottled (this is the "stream everything").
if self.bus is not None and self.run_id is not None:
self.bus.publish(self.run_id, event)
if not self._is_coarse(event):
return
if len(self.coarse) < _MAX_COARSE_EVENTS:
self.coarse.append(event)
# Chat surface: coarse events become thinking-step items when a LangGraph
# run context is active (agent door). No-op elsewhere.
_dispatch_chat_event(event)
def _is_coarse(self, event: dict[str, Any]) -> bool:
phase = event.get("phase", "")
if phase != self._last_phase:
self._last_phase = phase
self._last_counter_ts.clear()
return True
if "current" not in event:
return True
key = (phase, event.get("unit"))
now = int(event.get("ts", 0))
last = self._last_counter_ts.get(key, 0)
if now - last >= _COUNTER_THROTTLE_MS:
self._last_counter_ts[key] = now
return True
return False
_active_reporter: ContextVar[ProgressReporter | None] = ContextVar(
"active_progress_reporter", default=None
)
def emit_progress(
phase: str,
message: str | None = None,
*,
current: int | None = None,
total: int | None = None,
unit: str | None = None,
**detail: Any,
) -> None:
"""Report a progress event for the active run; a no-op when none is active.
Safe to call from anywhere in the scraper code path — never raises into the
caller.
"""
reporter = _active_reporter.get()
if reporter is None:
return
try:
event: dict[str, Any] = {
"type": "run.progress",
"ts": int(time.time() * 1000),
"phase": phase,
}
if message is not None:
event["message"] = message
if current is not None:
event["current"] = current
if total is not None:
event["total"] = total
if unit is not None:
event["unit"] = unit
if detail:
event["detail"] = detail
reporter.handle(event)
except Exception: # pragma: no cover - progress must never break a scrape
logger.debug("emit_progress failed; suppressed", exc_info=True)
@contextlib.contextmanager
def progress_scope(
*, run_id: str | None = None, bus: RunEventBus | None = None
) -> Iterator[ProgressReporter]:
"""Bind a :class:`ProgressReporter` for the duration of a run.
The contextvar set here propagates to every coroutine awaited inside the
``with`` block (same task), so deep scraper code sees it without wiring.
"""
reporter = ProgressReporter(run_id=run_id, bus=bus)
token = _active_reporter.set(reporter)
try:
yield reporter
finally:
_active_reporter.reset(token)
def _dispatch_chat_event(event: dict[str, Any]) -> None:
"""Best-effort forward to the chat SSE relay via a LangGraph custom event.
Only meaningful inside an agent tool call (there is a run context); raises
otherwise, which we swallow — mirrors ``retry_after``'s guarded dispatch.
"""
try:
from langchain_core.callbacks import dispatch_custom_event
dispatch_custom_event("scraper_progress", event)
except Exception:
logger.debug("scraper_progress dispatch skipped (no run context)")

116
surfsense_backend/app/capabilities/core/runs.py
View file

@ -90,6 +90,7 @@ async def record_run(
error: str | None = None,
duration_ms: int | None = None,
cost_micros: int | None = None,
progress: list[dict[str, Any]] | None = None,
) -> str | None:
"""Persist a run row and return its id, or ``None`` on failure (best-effort).
@ -112,6 +113,7 @@ async def record_run(
char_count=serialized.char_count if serialized else 0,
duration_ms=duration_ms,
cost_micros=cost_micros,
progress=progress or None,
)
session.add(run)
await session.flush()
@ -128,6 +130,120 @@ async def record_run(
return None
async def create_pending_run(
session: AsyncSession,
*,
workspace_id: int,
capability: str,
origin: str,
input: dict | None = None,
user_id: Any | None = None,
thread_id: str | None = None,
) -> str | None:
"""Insert a ``running`` run row up front and return its id (best-effort).
The async door needs a durable id before it spawns the background scrape so
the row is visible in history and streamable while it runs; :func:`finalize_run`
later flips it to a terminal status.
"""
try:
run = Run(
workspace_id=workspace_id,
user_id=user_id,
thread_id=thread_id,
capability=capability,
origin=origin,
status="running",
input=input,
)
session.add(run)
await session.flush()
run_id = str(run.id)
await session.commit()
return run_id
except Exception:
logger.exception("create_pending_run failed for capability=%s", capability)
try:
await session.rollback()
except Exception:
logger.exception("create_pending_run rollback failed")
return None
async def finalize_run(
session: AsyncSession,
*,
run_id: str,
status: str,
serialized: SerializedOutput | None = None,
error: str | None = None,
duration_ms: int | None = None,
cost_micros: int | None = None,
progress: list[dict[str, Any]] | None = None,
) -> bool:
"""Flip a pending run to a terminal status with its output/metrics.
Returns ``True`` on success. Best-effort like the other recorders: a failure
here is logged and swallowed rather than crashing the background task.
"""
import uuid as _uuid
try:
raw = run_id[len("run_") :] if run_id.startswith("run_") else run_id
run = await session.get(Run, _uuid.UUID(raw))
if run is None:
logger.warning("finalize_run: run %s not found", run_id)
return False
run.status = status
run.error = error
if serialized is not None:
run.output_text = serialized.text
run.item_count = serialized.item_count
run.char_count = serialized.char_count
if duration_ms is not None:
run.duration_ms = duration_ms
if cost_micros is not None:
run.cost_micros = cost_micros
if progress:
run.progress = progress
await _maybe_cleanup(session, "runs", RUNS_RETENTION_DAYS)
await session.commit()
return True
except Exception:
logger.exception("finalize_run failed for run=%s", run_id)
try:
await session.rollback()
except Exception:
logger.exception("finalize_run rollback failed")
return False
async def fail_stale_running_runs(session: AsyncSession) -> int:
"""Mark every leftover ``running`` run as ``error`` — called once at startup.
Single process: any row still ``running`` at boot belongs to a scrape that
died with the previous process, so it can never complete. Without this sweep
such rows would stay ``running`` forever.
"""
try:
result = await session.execute(
text(
"UPDATE runs SET status = 'error', "
"error = 'Interrupted by server restart' "
"WHERE status = 'running'"
)
)
await session.commit()
return result.rowcount or 0
except Exception:
logger.exception("fail_stale_running_runs failed")
try:
await session.rollback()
except Exception:
logger.exception("fail_stale_running_runs rollback failed")
return 0
async def record_spill(
session: AsyncSession,
*,

3
surfsense_backend/app/capabilities/google_maps/reviews/executor.py
View file

@ -5,6 +5,7 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable
from app.capabilities.core import Executor
from app.capabilities.core.progress import emit_progress
from app.capabilities.google_maps.reviews.schemas import ReviewsInput, ReviewsOutput
from app.exceptions import ForbiddenError
from app.proprietary.platforms.google_maps import (
@ -29,12 +30,14 @@ def build_reviews_executor(scrape_fn: ReviewsFn | None = None) -> Executor:
reviewsStartDate=payload.start_date,
language=payload.language,
)
emit_progress("starting", "Fetching Google Maps reviews", total=payload.max_reviews, unit="review")
try:
items = await scrape_fn(actor_input)
except SignInRequiredError as exc:
raise ForbiddenError(
f"Google sign in required: {exc}", code="GOOGLE_SIGNIN_REQUIRED"
) from exc
emit_progress("done", f"Scraped {len(items)} review(s)", current=len(items), unit="review")
return ReviewsOutput(items=items)
return execute

3
surfsense_backend/app/capabilities/google_maps/scrape/executor.py
View file

@ -5,6 +5,7 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable
from app.capabilities.core import Executor
from app.capabilities.core.progress import emit_progress
from app.capabilities.google_maps.scrape.schemas import ScrapeInput, ScrapeOutput
from app.exceptions import ForbiddenError
from app.proprietary.platforms.google_maps import (
@ -32,12 +33,14 @@ def build_scrape_executor(scrape_fn: ScrapeFn | None = None) -> Executor:
maxReviews=payload.max_reviews,
maxImages=payload.max_images,
)
emit_progress("starting", "Searching Google Maps", total=payload.max_places, unit="place")
try:
items = await scrape_fn(actor_input)
except SignInRequiredError as exc:
raise ForbiddenError(
f"Google sign in required: {exc}", code="GOOGLE_SIGNIN_REQUIRED"
) from exc
emit_progress("done", f"Scraped {len(items)} place(s)", current=len(items), unit="place")
return ScrapeOutput(items=items)
return execute

8
surfsense_backend/app/capabilities/google_search/scrape/executor.py
View file

@ -5,6 +5,7 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable
from app.capabilities.core import Executor
from app.capabilities.core.progress import emit_progress
from app.capabilities.google_search.scrape.schemas import (
MAX_PAGES_PER_QUERY,
MAX_SEARCH_QUERIES,
@ -35,7 +36,14 @@ def build_scrape_executor(scrape_fn: ScrapeFn | None = None) -> Executor:
languageCode=payload.language_code,
site=payload.site,
)
emit_progress(
"starting",
f"Searching {len(payload.queries)} query(ies)",
total=_MAX_SERP_ITEMS,
unit="page",
)
items = await scrape_fn(actor_input, limit=_MAX_SERP_ITEMS)
emit_progress("done", f"Scraped {len(items)} SERP page(s)", current=len(items), unit="page")
return ScrapeOutput(items=items)
return execute

3
surfsense_backend/app/capabilities/reddit/scrape/executor.py
View file

@ -5,6 +5,7 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable
from app.capabilities.core import Executor
from app.capabilities.core.progress import emit_progress
from app.capabilities.reddit.scrape.schemas import ScrapeInput, ScrapeOutput
from app.exceptions import ForbiddenError
from app.proprietary.platforms.reddit import (
@ -35,6 +36,7 @@ def build_scrape_executor(scrape_fn: ScrapeFn | None = None) -> Executor:
postDateLimit=payload.post_date_limit,
commentDateLimit=payload.comment_date_limit,
)
emit_progress("starting", "Resolving Reddit targets", total=payload.max_items, unit="item")
try:
items = await scrape_fn(actor_input, limit=payload.max_items)
except RedditAccessBlockedError as exc:
@ -44,6 +46,7 @@ def build_scrape_executor(scrape_fn: ScrapeFn | None = None) -> Executor:
f"Reddit refused anonymous access: {exc}",
code="REDDIT_ACCESS_BLOCKED",
) from exc
emit_progress("done", f"Scraped {len(items)} item(s)", current=len(items), unit="item")
return ScrapeOutput(items=items)
return execute

14
surfsense_backend/app/capabilities/web/crawl/executor.py
View file

@ -9,6 +9,7 @@ the captcha telemetry the billing seam reads).
from __future__ import annotations
from app.capabilities.core import Executor
from app.capabilities.core.progress import emit_progress
from app.capabilities.web.crawl.schemas import (
ContactRef,
Contacts,
@ -38,6 +39,12 @@ def build_crawl_executor(engine: WebCrawlerConnector | None = None) -> Executor:
crawler = engine or WebCrawlerConnector()
async def execute(payload: CrawlInput) -> CrawlOutput:
emit_progress(
"starting",
f"Crawling {len(payload.startUrls)} seed URL(s)",
total=payload.maxCrawlPages,
unit="page",
)
pages = await crawl_site(
crawler,
payload.startUrls,
@ -46,7 +53,14 @@ def build_crawl_executor(engine: WebCrawlerConnector | None = None) -> Executor:
include_patterns=payload.includeUrlPatterns,
exclude_patterns=payload.excludeUrlPatterns,
)
emit_progress(
"processing",
f"Processing {len(pages)} crawled page(s)",
current=len(pages),
unit="page",
)
items = [_to_item(page, payload.maxLength) for page in pages]
emit_progress("done", f"Crawled {len(items)} page(s)", current=len(items), unit="page")
return CrawlOutput(
items=items,
contacts=_aggregate_contacts(items),

3
surfsense_backend/app/capabilities/youtube/comments/executor.py
View file

@ -5,6 +5,7 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable
from app.capabilities.core import Executor
from app.capabilities.core.progress import emit_progress
from app.capabilities.youtube.comments.schemas import CommentsInput, CommentsOutput
from app.proprietary.platforms.youtube import (
YouTubeCommentsInput,
@ -24,7 +25,9 @@ def build_comments_executor(scrape_fn: CommentsFn | None = None) -> Executor:
maxComments=payload.max_comments,
sortCommentsBy=payload.sort_by,
)
emit_progress("starting", "Fetching YouTube comments", total=payload.max_comments, unit="comment")
items = await scrape_fn(actor_input)
emit_progress("done", f"Scraped {len(items)} comment(s)", current=len(items), unit="comment")
return CommentsOutput(items=items)
return execute

3
surfsense_backend/app/capabilities/youtube/scrape/executor.py
View file

@ -5,6 +5,7 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable
from app.capabilities.core import Executor
from app.capabilities.core.progress import emit_progress
from app.capabilities.youtube.scrape.schemas import ScrapeInput, ScrapeOutput
from app.proprietary.platforms.youtube import (
YouTubeScrapeInput,
@ -30,7 +31,9 @@ def build_scrape_executor(scrape_fn: ScrapeFn | None = None) -> Executor:
downloadSubtitles=payload.download_subtitles,
subtitlesLanguage=payload.subtitles_language,
)
emit_progress("starting", "Resolving YouTube targets", total=payload.max_results, unit="video")
items = await scrape_fn(actor_input)
emit_progress("done", f"Scraped {len(items)} video(s)", current=len(items), unit="video")
return ScrapeOutput(items=items)
return execute

8
surfsense_backend/app/config/__init__.py
View file

@ -705,16 +705,16 @@ class Config:
os.getenv("GOOGLE_SEARCH_MICROS_PER_SERP", "5500")
)
GOOGLE_MAPS_MICROS_PER_PLACE = int(
os.getenv("GOOGLE_MAPS_MICROS_PER_PLACE", "5000")
os.getenv("GOOGLE_MAPS_MICROS_PER_PLACE", "3500")
)
GOOGLE_MAPS_MICROS_PER_REVIEW = int(
os.getenv("GOOGLE_MAPS_MICROS_PER_REVIEW", "2000")
os.getenv("GOOGLE_MAPS_MICROS_PER_REVIEW", "1500")
)
YOUTUBE_MICROS_PER_VIDEO = int(os.getenv("YOUTUBE_MICROS_PER_VIDEO", "3500"))
YOUTUBE_MICROS_PER_VIDEO = int(os.getenv("YOUTUBE_MICROS_PER_VIDEO", "2500"))
# Kept separate from the video rate so comments can be re-tuned toward the
# cheaper per-comment market ($0.40-2.00/1k) without touching video pricing.
YOUTUBE_MICROS_PER_COMMENT = int(
os.getenv("YOUTUBE_MICROS_PER_COMMENT", "3500")
os.getenv("YOUTUBE_MICROS_PER_COMMENT", "1500")
)
# Low-balance WARNING threshold (micro-USD). Surfaced by the quota service

3
surfsense_backend/app/db.py
View file

@ -2823,6 +2823,9 @@ class Run(Base, TimestampMixin):
char_count = Column(Integer, nullable=False, default=0)
duration_ms = Column(Integer, nullable=True)
cost_micros = Column(BigInteger, nullable=True)
# Coarse progress log (list of throttled events) captured during the run;
# the live fine-grained stream is ephemeral (bus/SSE only).
progress = Column(JSONB, nullable=True)
class ToolOutputSpill(Base, TimestampMixin):

23
surfsense_backend/app/proprietary/platforms/google_maps/fetch.py
View file

@ -582,14 +582,21 @@ async def fetch_rpc_json(
if looks_like_signin_wall(text):
logger.warning("Maps RPC GET %s hit a sign-in/consent wall", url)
return None
body = strip_xssi(text)
# The proxy/HTML fetcher may wrap the JSON in <html><body><p>…</p>…,
# so extract just the balanced top-level array/object.
start = next((i for i, c in enumerate(body) if c in "[{"), -1)
if start == -1:
return None
blob = brace_match_json(body, start)
return json.loads(blob) if blob else None
# brace_match_json is a pure-Python scan and review payloads reach
# ~1MB; decode off-loop so it can't stall concurrent requests.
return await asyncio.to_thread(_decode_rpc_body, text)
except Exception as e:
logger.warning("Maps RPC GET %s failed: %s", url, e)
return None
def _decode_rpc_body(text: str) -> Any | None:
"""Strip the XSSI guard and decode the balanced top-level JSON blob."""
body = strip_xssi(text)
# The proxy/HTML fetcher may wrap the JSON in <html><body><p>…</p>…,
# so extract just the balanced top-level array/object.
start = next((i for i, c in enumerate(body) if c in "[{"), -1)
if start == -1:
return None
blob = brace_match_json(body, start)
return json.loads(blob) if blob else None

3
surfsense_backend/app/proprietary/platforms/google_maps/reviews.py
View file

@ -202,9 +202,12 @@ async def scrape_reviews(
input_model: GoogleMapsReviewsInput, *, limit: int | None = None
) -> list[dict[str, Any]]:
"""Collect :func:`iter_reviews` into a list, honoring an optional guard."""
from app.capabilities.core.progress import emit_progress
results: list[dict[str, Any]] = []
async for item in iter_reviews(input_model):
results.append(item)
emit_progress("scraping", current=len(results), total=limit, unit="review")
if limit is not None and len(results) >= limit:
break
return results

3
surfsense_backend/app/proprietary/platforms/google_maps/scraper.py
View file

@ -455,9 +455,12 @@ async def scrape_places(
``limit`` is a request-time policy guard (used by the route), NOT a ceiling
in the streaming core.
"""
from app.capabilities.core.progress import emit_progress
results: list[dict[str, Any]] = []
async for item in iter_places(input_model):
results.append(item)
emit_progress("scraping", current=len(results), total=limit, unit="place")
if limit is not None and len(results) >= limit:
break
return results

63
surfsense_backend/app/proprietary/platforms/google_search/fetch.py
View file

@ -33,6 +33,8 @@ from __future__ import annotations
import asyncio
import logging
import random
import sys
import threading
import time
from urllib.parse import urlsplit, urlunsplit
@ -228,9 +230,46 @@ _MOBILE_UA = "Mozilla/5.0 (Android 14; Mobile; rv:132.0) Gecko/132.0 Firefox/132
_MOBILE_VIEWPORT = {"width": 412, "height": 915}
# patchright launches Chromium via asyncio.create_subprocess_exec, which the
# server's main loop cannot do on Windows (main.py pins a SelectorEventLoop
# for psycopg; Selector loops raise NotImplementedError on subprocess_exec).
# All browser work therefore runs on ONE dedicated background loop that is
# explicitly subprocess-capable; callers await it across threads. This also
# keeps the persistent AsyncStealthySession (and its async page_action) intact
# — the sync-fetcher-in-a-thread pattern the other scrapers use would tear
# down the browser on every fetch.
_browser_loop: asyncio.AbstractEventLoop | None = None
_browser_loop_guard = threading.Lock()
def _get_browser_loop() -> asyncio.AbstractEventLoop:
"""The lazily-started, process-wide event loop the browser lives on."""
global _browser_loop
with _browser_loop_guard:
if _browser_loop is None:
loop = (
asyncio.ProactorEventLoop()
if sys.platform == "win32"
else asyncio.new_event_loop()
)
threading.Thread(
target=loop.run_forever, name="google-search-browser", daemon=True
).start()
_browser_loop = loop
return _browser_loop
async def _in_browser_loop(coro):
"""Run ``coro`` on the browser loop and await its result from this loop."""
return await asyncio.wrap_future(
asyncio.run_coroutine_threadsafe(coro, _get_browser_loop())
)
# One live browser per layout (desktop / mobile — the UA and viewport are
# session-level context options). Launching Chromium costs ~5 s, so it's paid
# once and every fetch just opens a fresh context on its vetted sticky proxy.
# Only ever touched from coroutines running on the browser loop.
_sessions: dict[bool, AsyncStealthySession] = {}
_session_lock = asyncio.Lock()
@ -263,8 +302,7 @@ async def _get_session(mobile: bool) -> AsyncStealthySession:
return session
async def _drop_session(mobile: bool) -> None:
"""Close and forget a session whose browser is presumed broken."""
async def _drop_session_on_loop(mobile: bool) -> None:
async with _session_lock:
session = _sessions.pop(mobile, None)
if session is not None:
@ -274,18 +312,32 @@ async def _drop_session(mobile: bool) -> None:
pass
async def _drop_session(mobile: bool) -> None:
"""Close and forget a session whose browser is presumed broken."""
await _in_browser_loop(_drop_session_on_loop(mobile))
async def close_sessions() -> None:
"""Shut down the shared browsers (for tests/scripts wanting a clean exit)."""
for mobile in (False, True):
await _drop_session(mobile)
async def _render(url: str, proxy: str | None, mobile: bool = False):
"""Headless render of a SERP on the shared browser (fresh proxy context)."""
async def _render_on_loop(url: str, proxy: str | None, mobile: bool):
session = await _get_session(mobile)
return await session.fetch(url, proxy=proxy)
async def _render(url: str, proxy: str | None, mobile: bool = False):
"""Headless render of a SERP on the shared browser (fresh proxy context).
The actual browser work is marshalled onto the dedicated subprocess-capable
loop (see :func:`_get_browser_loop`); this coroutine just awaits it from
the caller's loop.
"""
return await _in_browser_loop(_render_on_loop(url, proxy, mobile))
async def fetch_serp_html(url: str, *, mobile: bool = False) -> str | None:
"""Return fully-rendered SERP HTML for ``url``, or ``None`` if unobtainable.
@ -324,7 +376,8 @@ async def fetch_serp_html(url: str, *, mobile: bool = False) -> str | None:
except Exception as e:
# Renders on a walled IP still return HTML; an exception means the
# browser side is broken, so relaunch it rather than limp along.
logger.warning("[google_search] render failed: %s", e)
# repr(), not str(): e.g. NotImplementedError stringifies to "".
logger.warning("[google_search] render failed: %r", e)
_last_good_proxy = None
await _drop_session(mobile)
continue

11
surfsense_backend/app/proprietary/platforms/google_search/scraper.py
View file

@ -10,6 +10,7 @@ the YouTube and Maps flows were.
from __future__ import annotations
import asyncio
import logging
from collections.abc import AsyncIterator
from typing import Any
@ -74,7 +75,10 @@ async def _serp_page_flow(
if html is None:
logger.warning("[google_search] no SERP HTML for %s", url)
break
item = parse_serp(html, include_icons=input_model.includeIcons)
# Rendered SERPs are ~1MB; parse off-loop to keep the server responsive.
item = await asyncio.to_thread(
parse_serp, html, include_icons=input_model.includeIcons
)
if input_model.saveHtml:
item.html = html
if not input_model.focusOnPaidAds or item.paidResults or item.paidProducts:
@ -137,7 +141,7 @@ async def _ai_mode_flow(
if html is None:
logger.warning("[google_search] no AI Mode HTML for %r", term)
return
result = parse_ai_mode(html, query=term, url=url)
result = await asyncio.to_thread(parse_ai_mode, html, query=term, url=url)
if result is None:
logger.info("[google_search] AI Mode answer missing for %r", term)
return
@ -177,9 +181,12 @@ async def scrape_serps(
``limit`` is a request-time policy guard (used by the route), NOT a
ceiling in the streaming core.
"""
from app.capabilities.core.progress import emit_progress
results: list[dict[str, Any]] = []
async for item in iter_serps(input_model):
results.append(item)
emit_progress("scraping", current=len(results), total=limit, unit="page")
if limit is not None and len(results) >= limit:
break
return results

5
surfsense_backend/app/proprietary/platforms/reddit/scraper.py
View file

@ -444,9 +444,14 @@ async def scrape_reddit(
``limit`` is a request-time policy guard, NOT a ceiling in the streaming
core.
"""
from app.capabilities.core.progress import emit_progress
results: list[dict[str, Any]] = []
async for item in iter_reddit(input_model):
results.append(item)
emit_progress(
"scraping", current=len(results), total=limit, unit="item"
)
if limit is not None and len(results) >= limit:
break
return results

9
surfsense_backend/app/proprietary/platforms/youtube/comments.py
View file

@ -80,14 +80,16 @@ async def _comments_for_video(
html = await fetch_html(page_url)
if not html:
return
initial = extract_yt_initial_data(html)
# Both helpers scan 1-2MB of watch-page HTML in pure Python; keep that
# off the event loop.
initial = await asyncio.to_thread(extract_yt_initial_data, html)
if not initial:
return
section = comment_section_token(initial)
if not section: # comments disabled or absent
return
meta = parse_video_page(html) or {}
meta = await asyncio.to_thread(parse_video_page, html) or {}
base = {
"videoId": video_id,
"pageUrl": page_url,
@ -180,9 +182,12 @@ async def scrape_comments(
input_model: YouTubeCommentsInput, *, limit: int | None = None
) -> list[dict]:
"""Collect :func:`iter_comments` into a list, honoring an optional guard."""
from app.capabilities.core.progress import emit_progress
results: list[dict] = []
async for comment in iter_comments(input_model):
results.append(comment)
emit_progress("scraping", current=len(results), total=limit, unit="comment")
if limit is not None and len(results) >= limit:
break
return results

17
surfsense_backend/app/proprietary/platforms/youtube/scraper.py
View file

@ -174,7 +174,9 @@ async def _video_flow(
html = await fetch_html(url)
if not html:
return
partial = parse_video_page(html)
# Watch-page HTML is 1-2MB and the embedded-JSON scan is pure Python; run
# it off-loop so one video parse can't stall other requests.
partial = await asyncio.to_thread(parse_video_page, html)
if not partial:
return
yield await _finalize(
@ -294,7 +296,7 @@ async def _channel_tab_flow(
seed_html = await fetch_html(from_url)
if not seed_html:
return
initial = extract_yt_initial_data(seed_html)
initial = await asyncio.to_thread(extract_yt_initial_data, seed_html)
if not initial:
return
@ -355,7 +357,11 @@ async def _channel_flow(
(identity, banner, and the About panel's deep fields) stamped on every item.
"""
videos_seed = await fetch_html(_channel_tab_url(handle, "videos"))
initial = extract_yt_initial_data(videos_seed) if videos_seed else None
initial = (
await asyncio.to_thread(extract_yt_initial_data, videos_seed)
if videos_seed
else None
)
channel_meta: dict[str, Any] = {}
if initial:
channel_meta = parse_channel_metadata(initial)
@ -457,7 +463,7 @@ async def _hashtag_flow(
html = await fetch_html(url)
if not html:
return
data = extract_yt_initial_data(html)
data = await asyncio.to_thread(extract_yt_initial_data, html)
order = 0
while data:
items, token = parse_search_response(data)
@ -552,9 +558,12 @@ async def scrape_youtube(
``limit`` is a request-time policy guard (used by the route), NOT a ceiling
in the streaming core.
"""
from app.capabilities.core.progress import emit_progress
results: list[dict[str, Any]] = []
async for item in iter_youtube(input_model):
results.append(item)
emit_progress("scraping", current=len(results), total=limit, unit="video")
if limit is not None and len(results) >= limit:
break
return results

31
surfsense_backend/app/proprietary/platforms/youtube/subtitles.py
View file

@ -12,7 +12,7 @@ import logging
from fake_useragent import UserAgent
from requests import Session
from youtube_transcript_api import YouTubeTranscriptApi
from youtube_transcript_api import RequestBlocked, YouTubeTranscriptApi
from youtube_transcript_api.formatters import (
SRTFormatter,
TextFormatter,
@ -31,6 +31,11 @@ _FORMATTERS = {
"plaintext": TextFormatter,
}
# Blocked-IP retries, mirroring innertube's rotate-on-block. Each retry builds a
# fresh Session (new TCP connection), which draws a new exit IP on a rotating
# proxy gateway. Without a proxy we don't retry - the egress IP can't change.
_MAX_ROTATIONS = 3
def _build_client() -> YouTubeTranscriptApi:
http_client = Session()
@ -59,10 +64,26 @@ def _select_transcript(transcript_list, language: str, prefer_generated: bool):
def _fetch_subtitles_sync(
video_id: str, language: str, fmt: str, prefer_generated: bool
):
api = _build_client()
transcript_list = api.list(video_id)
transcript = _select_transcript(transcript_list, language, prefer_generated)
fetched = transcript.fetch()
attempts = (_MAX_ROTATIONS + 1) if get_requests_proxies() else 1
for attempt in range(attempts):
api = _build_client()
try:
transcript_list = api.list(video_id)
transcript = _select_transcript(
transcript_list, language, prefer_generated
)
fetched = transcript.fetch()
break
except RequestBlocked: # covers IpBlocked; rotate to a fresh exit IP
if attempt == attempts - 1:
raise
logger.info(
"Transcript request blocked for %s; retrying on a fresh proxy "
"connection (attempt %d/%d)",
video_id,
attempt + 2,
attempts,
)
if fmt == "xml":
# No XML formatter in the library; emit the raw snippet data as text.

6
surfsense_backend/app/proprietary/web_crawler/connector.py
View file

@ -514,7 +514,11 @@ class WebCrawlerConnector:
)
return None
result = self._build_result(
# Trafilatura extraction is CPU-bound (100ms+ on large pages); run it
# off-loop so concurrent requests aren't stalled. The browser tiers get
# this for free by calling _build_result inside their worker threads.
result = await asyncio.to_thread(
self._build_result,
page.html_content,
url,
"scrapling-static",

20
surfsense_backend/app/proprietary/web_crawler/site_crawler.py
View file

@ -25,6 +25,7 @@ from typing import TYPE_CHECKING, Any
from scrapling.spiders import CrawlerEngine, LinkExtractor, Response, Spider
from app.capabilities.core.progress import emit_progress
from app.proprietary.web_crawler.connector import (
CrawlOutcome,
CrawlOutcomeStatus,
@ -88,7 +89,16 @@ class _ConnectorSession:
self._is_alive = False
async def fetch(self, url: str, **_kwargs: Any) -> Response:
emit_progress("fetching", f"Fetching {url}", unit="page", url=url)
outcome = await self._connector.crawl_url(url)
if outcome.captcha_attempts:
emit_progress(
"captcha",
f"Captcha {'solved' if outcome.captcha_solved else 'attempted'} on {url}",
url=url,
attempts=outcome.captcha_attempts,
solved=outcome.captcha_solved,
)
result = outcome.result or {}
content = result.get("content")
# Selector chokes on empty content; a fetch that raises would be counted
@ -156,6 +166,16 @@ class _SiteSpider(Spider):
if len(self.pages) < self._max_pages:
self.pages.append(_to_page(req_url, outcome, depth, referrer))
emit_progress(
"crawled",
f"Crawled {req_url}",
current=len(self.pages),
total=self._max_pages,
unit="page",
url=req_url,
depth=depth,
status=outcome.status.value,
)
# Cap reached: stop the engine so queued-but-unfetched links are abandoned
# (never fetched, never billed), matching the old BFS's per-fetch guard.

2
surfsense_backend/app/services/platform_scrape_credit_service.py
View file

@ -21,11 +21,11 @@ from uuid import UUID
from sqlalchemy.ext.asyncio import AsyncSession
from app.config import config
from app.services import wallet_credit
# One "out of credit" type across every per-unit biller; the capability doors
# already catch exactly this one.
from app.services.etl_credit_service import InsufficientCreditsError
from app.services import wallet_credit
__all__ = ["InsufficientCreditsError", "PlatformScrapeCreditService"]

8
surfsense_backend/app/services/web_crawl_credit_service.py
View file

@ -24,14 +24,14 @@ from sqlalchemy.ext.asyncio import AsyncSession
from app.config import config
# Reuse the ETL service's error type so callers (and tests) have one exception
# to catch for "out of credit" across every per-unit wallet biller.
from app.services.etl_credit_service import InsufficientCreditsError
# Wallet math (spendable / check / debit) is shared with the platform-scrape
# biller — see app/services/wallet_credit.py.
from app.services import wallet_credit
# Reuse the ETL service's error type so callers (and tests) have one exception
# to catch for "out of credit" across every per-unit wallet biller.
from app.services.etl_credit_service import InsufficientCreditsError
__all__ = ["InsufficientCreditsError", "WebCrawlCreditService"]

15
surfsense_backend/app/tasks/chat/streaming/handlers/custom_event_dispatch.py
View file

@ -10,6 +10,7 @@ from app.tasks.chat.streaming.handlers.custom_events import (
handle_action_log_updated,
handle_document_created,
handle_report_progress,
handle_scraper_progress,
)
from app.tasks.chat.streaming.relay.state import AgentEventRelayState
@ -39,6 +40,20 @@ def iter_custom_event_frames(
yield frame
return
if name == "scraper_progress":
frame, state.last_active_step_items = handle_scraper_progress(
data,
last_active_step_id=state.last_active_step_id,
last_active_step_title=state.last_active_step_title,
last_active_step_items=state.last_active_step_items,
streaming_service=streaming_service,
content_builder=content_builder,
thinking_metadata=state.span_metadata_if_active(),
)
if frame:
yield frame
return
if name == "document_created":
frame = handle_document_created(data, streaming_service=streaming_service)
if frame:

47
surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py
View file

@ -54,6 +54,53 @@ def handle_report_progress(
return frame, new_items
def _scraper_progress_label(data: dict[str, Any]) -> str:
"""Build a one-line human status from a ``scraper_progress`` event."""
message = data.get("message")
phase = data.get("phase", "")
current = data.get("current")
total = data.get("total")
label = message or (phase.replace("_", " ").capitalize() if phase else "Working")
if current is not None:
counter = f"{current}/{total}" if total else str(current)
label = f"{label} ({counter})"
return label
def handle_scraper_progress(
data: dict[str, Any],
*,
last_active_step_id: str | None,
last_active_step_title: str,
last_active_step_items: list[str],
streaming_service: Any,
content_builder: Any | None,
thinking_metadata: dict[str, Any] | None = None,
) -> tuple[str | None, list[str]]:
"""Surface a scraper's live progress as an evolving thinking-step item.
Scraper capability tool calls own a fresh thinking step (see ``tool_start``),
so we show a single latest-status line rather than accumulating every event.
Returns (frame or None, items after update).
"""
if not last_active_step_id:
return None, last_active_step_items
label = _scraper_progress_label(data)
if not label:
return None, last_active_step_items
new_items = [label]
frame = emit_thinking_step_frame(
streaming_service=streaming_service,
content_builder=content_builder,
step_id=last_active_step_id,
title=last_active_step_title,
status="in_progress",
items=new_items,
metadata=thinking_metadata,
)
return frame, new_items
def handle_document_created(
data: dict[str, Any], *, streaming_service: Any
) -> str | None:

12
surfsense_backend/main.py
View file

@ -6,10 +6,6 @@ import sys
import uvicorn
from dotenv import load_dotenv
# Fix for Windows: psycopg requires SelectorEventLoop, not ProactorEventLoop
if sys.platform == "win32":
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
from app.config.uvicorn import load_uvicorn_config
_old_log_record_factory = logging.getLogRecordFactory()
@ -47,6 +43,14 @@ if __name__ == "__main__":
server = uvicorn.Server(config)
if sys.platform == "win32":
# Windows needs a split-loop setup: psycopg's async driver requires a
# SelectorEventLoop (no add_reader on Proactor), so the SERVER loop is
# forced to Selector via loop_factory. The global policy is left at
# its default (Proactor) so worker threads that call
# asyncio.new_event_loop() — playwright/patchright sync API used by
# the scrapers, unstructured, etc. — get a loop that CAN spawn
# subprocesses. Do not set WindowsSelectorEventLoopPolicy globally:
# that breaks every browser-based scraper with NotImplementedError.
asyncio.run(server.serve(), loop_factory=asyncio.SelectorEventLoop)
else:
server.run()

142
surfsense_backend/scripts/check_migration_flow.py Normal file
View file

@ -0,0 +1,142 @@
"""Self-check for the alembic fast-forward/adoption flow in alembic/env.py.
Verifies ``alembic upgrade head`` succeeds on the three DB states it must
handle without replaying pre-workspace-rename history against a
workspace-shape schema:
1. fresh -- empty database (fast-forward: create_all + stamp head)
2. bootstrap -- create_all-created schema + zero_publication, no alembic
history (adoption: stamp head)
3. midcrash -- bootstrap schema whose alembic_version is stuck at a
pre-rename revision from a failed replay (adoption: stamp head)
Run: python scripts/check_migration_flow.py
"""
import asyncio
import os
import subprocess
import sys
from pathlib import Path
BACKEND_DIR = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(BACKEND_DIR))
ADMIN_URL = os.getenv(
"ADMIN_DATABASE_URL", "postgresql://postgres:postgres@localhost:5432/postgres"
)
SCRATCH_DB = "surfsense_check_migration_flow"
SCRATCH_URL = ADMIN_URL.rsplit("/", 1)[0] + f"/{SCRATCH_DB}"
SCRATCH_URL_ASYNC = SCRATCH_URL.replace("postgresql://", "postgresql+asyncpg://")
async def recreate_scratch_db() -> None:
import asyncpg
admin = await asyncpg.connect(ADMIN_URL)
await admin.execute(f'DROP DATABASE IF EXISTS "{SCRATCH_DB}" WITH (FORCE)')
await admin.execute(f'CREATE DATABASE "{SCRATCH_DB}"')
await admin.close()
async def drop_scratch_db() -> None:
import asyncpg
admin = await asyncpg.connect(ADMIN_URL)
await admin.execute(f'DROP DATABASE IF EXISTS "{SCRATCH_DB}" WITH (FORCE)')
await admin.close()
def run_alembic_upgrade() -> None:
env = dict(os.environ, DATABASE_URL=SCRATCH_URL_ASYNC)
subprocess.run(
[sys.executable, "-m", "alembic", "upgrade", "head"],
cwd=BACKEND_DIR,
env=env,
check=True,
)
async def bootstrap_schema() -> None:
"""Mimic app startup bootstrap: create_all + ensure_publication, no stamp."""
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine
from app.db import Base
from app.zero_publication import ensure_publication
engine = create_async_engine(SCRATCH_URL_ASYNC)
async with engine.begin() as conn:
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm"))
await conn.run_sync(Base.metadata.create_all)
await conn.run_sync(ensure_publication)
await engine.dispose()
async def set_version(version: str | None) -> None:
import asyncpg
conn = await asyncpg.connect(SCRATCH_URL)
if version is None:
await conn.execute("DROP TABLE IF EXISTS alembic_version")
else:
await conn.execute(
"CREATE TABLE IF NOT EXISTS alembic_version ("
"version_num VARCHAR(32) NOT NULL PRIMARY KEY)"
)
await conn.execute("DELETE FROM alembic_version")
await conn.execute(
"INSERT INTO alembic_version (version_num) VALUES ($1)", version
)
await conn.close()
async def assert_at_head() -> None:
import asyncpg
from alembic.script import ScriptDirectory
head = ScriptDirectory(str(BACKEND_DIR / "alembic")).get_current_head()
conn = await asyncpg.connect(SCRATCH_URL)
version = await conn.fetchval("SELECT version_num FROM alembic_version")
workspaces = await conn.fetchval("SELECT to_regclass('workspaces')")
publication = await conn.fetchval(
"SELECT 1 FROM pg_publication WHERE pubname = 'zero_publication'"
)
await conn.close()
assert version == head, f"expected version {head}, got {version}"
assert workspaces, "workspaces table missing"
assert publication, "zero_publication missing"
async def main() -> None:
try:
# 1. Fresh empty DB -> fast-forward.
await recreate_scratch_db()
run_alembic_upgrade()
await assert_at_head()
print("OK: fresh DB fast-forwards to head")
# 2. Bootstrap-created schema, no alembic history -> adoption.
await recreate_scratch_db()
await bootstrap_schema()
await set_version(None)
run_alembic_upgrade()
await assert_at_head()
print("OK: bootstrap-created schema adopted (stamped head)")
# 3. Bootstrap schema stuck at a pre-rename revision -> adoption.
await set_version("4")
run_alembic_upgrade()
await assert_at_head()
print("OK: pre-rename stuck revision adopted (stamped head)")
# Re-run must be a clean no-op.
run_alembic_upgrade()
await assert_at_head()
print("OK: repeat upgrade is a no-op")
finally:
await drop_scratch_db()
asyncio.run(main())

5
surfsense_backend/tests/unit/agents/multi_agent_chat/test_web_search_removed.py
View file

@ -17,7 +17,10 @@ from app.agents.chat.multi_agent_chat.main_agent.tools.index import (
from app.agents.chat.multi_agent_chat.subagents.registry import (
SUBAGENT_BUILDERS_BY_NAME,
)
from app.utils.validators import DEPRECATED_CONNECTOR_TYPES, raise_if_connector_deprecated
from app.utils.validators import (
DEPRECATED_CONNECTOR_TYPES,
raise_if_connector_deprecated,
)
pytestmark = pytest.mark.unit

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

@ -105,6 +105,24 @@ def test_registered_verbs_appear_on_rest():
assert "/workspaces/{workspace_id}/scrapers/web/crawl" in paths
@pytest.mark.asyncio
async def test_capabilities_endpoint_lists_verbs_with_input_schema(monkeypatch):
"""The playground reads verb identity + input JSON schema from one GET."""
app = _build_app([_ECHO], monkeypatch)
async with _client(app) as client:
resp = await client.get("/api/v1/workspaces/7/scrapers/capabilities")
assert resp.status_code == 200
body = resp.json()
assert len(body) == 1
entry = body[0]
assert entry["name"] == "test.echo"
assert entry["description"] == "Echo the input back for tests."
# The schemas are the pydantic models' JSON schemas: the form renders the
# input schema, the API reference docs render both.
assert "value" in entry["input_schema"]["properties"]
assert "properties" in entry["output_schema"]
@pytest.mark.asyncio
async def test_over_budget_is_blocked_before_the_executor(monkeypatch):
from app.capabilities.core.access import rest
@ -243,6 +261,7 @@ def _fake_run_row(**overrides):
"thread_id": None,
"input": {"value": "hi"},
"output_text": '{"echo": "hi"}',
"progress": None,
}
defaults.update(overrides)
return SimpleNamespace(**defaults)
@ -337,3 +356,111 @@ async def test_success_charges_once(monkeypatch):
assert isinstance(output, _EchoOutput)
assert unit is None
assert ctx.workspace_id == 7
@pytest.mark.asyncio
async def test_async_mode_returns_202_and_pending_run(monkeypatch):
"""``?mode=async`` inserts a pending run and returns its id without blocking."""
from unittest.mock import AsyncMock
from app.capabilities.core.access import rest
monkeypatch.setattr(
rest, "create_pending_run", AsyncMock(return_value="async-id"), raising=True
)
# Don't actually run the scrape in the background during this unit test.
monkeypatch.setattr(rest, "_execute_async_run", AsyncMock(), raising=True)
app = _build_app([_ECHO], monkeypatch)
async with _client(app) as client:
resp = await client.post(
"/api/v1/workspaces/7/scrapers/test/echo?mode=async",
json={"value": "hi"},
)
assert resp.status_code == 202
assert resp.json() == {"run_id": "run_async-id", "status": "running"}
@pytest.mark.asyncio
async def test_run_events_replays_buffer_then_finishes(monkeypatch):
"""The SSE endpoint replays buffered events and closes on ``run.finished``."""
from app.capabilities.core.events import run_event_bus
row = _fake_run_row(status="running")
raw = str(row.id)
run_event_bus.publish(raw, {"type": "run.progress", "phase": "scraping", "current": 1})
run_event_bus.publish(raw, {"type": "run.finished", "status": "success", "item_count": 2})
app = _build_app_with_rows(monkeypatch, [row])
try:
async with _client(app) as client:
resp = await client.get(
f"/api/v1/workspaces/7/scrapers/runs/run_{raw}/events"
)
assert resp.status_code == 200
body = resp.text
assert '"type": "run.progress"' in body
assert '"type": "run.finished"' in body
assert '"status": "success"' in body
finally:
run_event_bus.close(raw)
@pytest.mark.asyncio
async def test_cancel_finalizes_running_run(monkeypatch):
"""Cancel signals the task, finalizes as ``cancelled``, and emits a terminal."""
import asyncio
from unittest.mock import AsyncMock
from app.capabilities.core.access import rest
from app.capabilities.core.events import run_event_bus
finalize = AsyncMock(return_value=True)
monkeypatch.setattr(rest, "finalize_run", finalize, raising=True)
row = _fake_run_row(status="running")
raw = str(row.id)
task = asyncio.create_task(asyncio.sleep(60))
run_event_bus.register_task(raw, task)
app = _build_app_with_rows(monkeypatch, [row])
try:
async with _client(app) as client:
resp = await client.post(
f"/api/v1/workspaces/7/scrapers/runs/run_{raw}/cancel"
)
assert resp.status_code == 200
assert resp.json() == {"run_id": f"run_{raw}", "status": "cancelled"}
finalize.assert_awaited_once()
assert finalize.await_args.kwargs["status"] == "cancelled"
finally:
task.cancel()
await asyncio.gather(task, return_exceptions=True)
run_event_bus.close(raw)
@pytest.mark.asyncio
async def test_cancel_conflicts_when_not_running(monkeypatch):
"""Cancelling a terminal run is a 409, not a silent overwrite."""
row = _fake_run_row(status="success")
app = _build_app_with_rows(monkeypatch, [row])
async with _client(app) as client:
resp = await client.post(
f"/api/v1/workspaces/7/scrapers/runs/run_{row.id}/cancel"
)
assert resp.status_code == 409
def test_emit_progress_is_a_noop_without_context():
"""Scraper code can call ``emit_progress`` freely; unset context = no-op."""
from app.capabilities.core.progress import emit_progress, progress_scope
# No active reporter -> returns without raising, records nothing.
emit_progress("phase", "message", current=1, total=2, unit="item")
# Inside a scope, coarse events are buffered for persistence.
with progress_scope() as reporter:
emit_progress("starting", "go")
emit_progress("done", current=5, unit="item")
assert len(reporter.coarse) == 2
assert reporter.coarse[0]["phase"] == "starting"

36
surfsense_backend/tests/unit/platforms/google_search/test_browser_loop.py Normal file
View file

@ -0,0 +1,36 @@
"""The Windows regression this guards: main.py runs the server on a
SelectorEventLoop (psycopg needs it), and Selector loops cannot spawn
subprocesses — so patchright's Chromium launch died with NotImplementedError
on every render. fetch.py now marshals all browser work onto a dedicated
subprocess-capable loop; this check proves that marshalling works from a
Selector main loop without paying for a real browser launch.
"""
import asyncio
import sys
import pytest
from app.proprietary.platforms.google_search import fetch
def test_browser_loop_can_spawn_subprocess_from_selector_loop():
async def spawn():
proc = await asyncio.create_subprocess_exec(
sys.executable,
"-c",
"print('ok')",
stdout=asyncio.subprocess.PIPE,
)
out, _ = await proc.communicate()
return out
async def main():
if sys.platform == "win32":
# The original bug: the server's own loop cannot do this.
with pytest.raises(NotImplementedError):
await spawn()
# The fix: the same work marshalled onto the browser loop succeeds.
assert b"ok" in await fetch._in_browser_loop(spawn())
asyncio.run(main(), loop_factory=asyncio.SelectorEventLoop)

77
surfsense_backend/tests/unit/platforms/youtube/test_subtitles_retry.py Normal file
View file

@ -0,0 +1,77 @@
"""Offline tests for the subtitles blocked-IP retry (rotate-on-block, no network.
Stubs ``_build_client``/``get_requests_proxies`` so the RequestBlocked retry
path is exercised deterministically: retries only when a proxy is configured,
each attempt gets a fresh client, and the final block is re-raised.
"""
from __future__ import annotations
import pytest
from youtube_transcript_api import RequestBlocked
from app.proprietary.platforms.youtube import subtitles
class _FakeTranscript:
is_generated = True
language_code = "en"
def fetch(self):
return [] # formatters iterate snippets; empty is fine here
class _FakeTranscriptList:
def find_transcript(self, codes):
return _FakeTranscript()
class _FakeApi:
"""One 'IP': blocks if ``blocked``, else returns a transcript list."""
def __init__(self, blocked: bool) -> None:
self.blocked = blocked
def list(self, video_id: str):
if self.blocked:
raise RequestBlocked(video_id)
return _FakeTranscriptList()
def _install(monkeypatch, outcomes: list[bool], proxied: bool) -> list[_FakeApi]:
"""Each ``_build_client`` call pops the next outcome (True = blocked)."""
built: list[_FakeApi] = []
def fake_build():
api = _FakeApi(outcomes[len(built)])
built.append(api)
return api
monkeypatch.setattr(subtitles, "_build_client", fake_build)
monkeypatch.setattr(
subtitles,
"get_requests_proxies",
lambda: {"http": "http://p", "https": "http://p"} if proxied else None,
)
return built
def test_blocked_then_success_retries_with_fresh_client(monkeypatch):
built = _install(monkeypatch, [True, True, False], proxied=True)
track = subtitles._fetch_subtitles_sync("vid", "en", "plaintext", False)
assert track.language == "en"
assert len(built) == 3 # two blocked attempts + one success, each a new client
def test_all_attempts_blocked_reraises(monkeypatch):
built = _install(monkeypatch, [True] * 10, proxied=True)
with pytest.raises(RequestBlocked):
subtitles._fetch_subtitles_sync("vid", "en", "plaintext", False)
assert len(built) == subtitles._MAX_ROTATIONS + 1
def test_no_proxy_means_single_attempt(monkeypatch):
built = _install(monkeypatch, [True] * 10, proxied=False)
with pytest.raises(RequestBlocked):
subtitles._fetch_subtitles_sync("vid", "en", "plaintext", False)
assert len(built) == 1 # same egress IP; retrying would be futile