From 17bdb06825c308f08b6cca02c585f2584fe2c947 Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Mon, 29 Jun 2026 21:40:53 -0700 Subject: [PATCH] feat(03c): bill credits per successful web crawl Add a config-driven crawl meter (1 USD / 1000 successes by default, retunable via WEB_CRAWL_MICROS_PER_SUCCESS) mirroring EtlCreditService: - WebCrawlCreditService (gate -> pre-check -> post-charge) on the unified wallet - webcrawler indexer: pre-flight credit block + post-charge of the workspace owner per success, with a web_crawl TokenUsage audit row - chat scrape_webpage (main + research): fold each successful scrape into the active turn's token accumulator so it settles with the chat turn - WEB_CRAWL_CREDIT_BILLING_ENABLED / WEB_CRAWL_MICROS_PER_SUCCESS config + env - unit tests (service) + integration tests (indexer wiring) - scripts/e2e_phase3_crawl_billing.py: manual functional e2e (3a/3b/3c) Co-authored-by: Cursor --- surfsense_backend/.env.example | 9 + .../main_agent/tools/scrape_webpage.py | 28 ++ .../builtins/research/tools/scrape_webpage.py | 28 ++ surfsense_backend/app/config/__init__.py | 16 + .../app/services/web_crawl_credit_service.py | 143 ++++++++ .../connector_indexers/webcrawler_indexer.py | 68 ++++ .../scripts/e2e_phase3_crawl_billing.py | 344 ++++++++++++++++++ .../test_web_crawl_billing.py | 190 ++++++++++ .../services/test_web_crawl_credit_service.py | 215 +++++++++++ 9 files changed, 1041 insertions(+) create mode 100644 surfsense_backend/app/services/web_crawl_credit_service.py create mode 100644 surfsense_backend/scripts/e2e_phase3_crawl_billing.py create mode 100644 surfsense_backend/tests/unit/connector_indexers/test_web_crawl_billing.py create mode 100644 surfsense_backend/tests/unit/services/test_web_crawl_credit_service.py diff --git a/surfsense_backend/.env.example b/surfsense_backend/.env.example index 22851dd6d..07dd37dc3 100644 --- a/surfsense_backend/.env.example +++ b/surfsense_backend/.env.example @@ -257,6 +257,15 @@ DEFAULT_CREDIT_MICROS_BALANCE=5000000 ETL_CREDIT_BILLING_ENABLED=FALSE MICROS_PER_PAGE=1000 +# Debit the credit wallet per *successful* web crawl. Default FALSE keeps +# crawling effectively free for self-hosted/OSS installs; hosted sets TRUE. +# Price is fully config-driven (the only source of truth, no hardcoded rate): +# WEB_CRAWL_MICROS_PER_SUCCESS = round(USD_per_1000_crawls * 1_000) +# 1000 == $1 / 1000 crawls (default) | 2000 == $2/1000 | 500 == $0.50/1000 +# Retune anytime with just an env change + restart (no code/migration). +# WEB_CRAWL_CREDIT_BILLING_ENABLED=FALSE +# WEB_CRAWL_MICROS_PER_SUCCESS=1000 + # Low-balance warning threshold (micro-USD), surfaced to the UI. Default $0.50. CREDIT_LOW_BALANCE_WARNING_MICROS=500000 diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/scrape_webpage.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/scrape_webpage.py index e2eb85a3f..ad082818b 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/scrape_webpage.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/scrape_webpage.py @@ -28,6 +28,33 @@ from app.utils.proxy import get_proxy_url, get_requests_proxies logger = logging.getLogger(__name__) +def _bill_successful_scrape() -> None: + """Fold one successful crawl into the current chat turn's bill (Phase 3c). + + The cost rides the turn accumulator and settles at the premium + ``finalize_credit`` step — no separate wallet hit. Free / BYOK / anonymous + turns (which never reserve/finalize) still record the line in the + breakdown but are never debited. No-op when crawl billing is disabled or + there is no active turn (e.g. non-chat callers). + """ + from app.services.token_tracking_service import get_current_accumulator + from app.services.web_crawl_credit_service import WebCrawlCreditService + + if not WebCrawlCreditService.billing_enabled(): + return + acc = get_current_accumulator() + if acc is None: + return + acc.add( + model="web_crawl", + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + cost_micros=WebCrawlCreditService.successes_to_micros(1), + call_kind="web_crawl", + ) + + def extract_domain(url: str) -> str: """Extract the domain from a URL.""" try: @@ -243,6 +270,7 @@ def create_scrape_webpage_tool(): } result = outcome.result + _bill_successful_scrape() content = result.get("content", "") metadata = result.get("metadata", {}) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/research/tools/scrape_webpage.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/research/tools/scrape_webpage.py index a3812bb2c..e449cae98 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/research/tools/scrape_webpage.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/research/tools/scrape_webpage.py @@ -22,6 +22,33 @@ from app.utils.proxy import get_proxy_url, get_requests_proxies logger = logging.getLogger(__name__) +def _bill_successful_scrape() -> None: + """Fold one successful crawl into the current chat turn's bill (Phase 3c). + + The cost rides the turn accumulator and settles at the premium + ``finalize_credit`` step — no separate wallet hit. Free / BYOK / anonymous + turns (which never reserve/finalize) still record the line in the + breakdown but are never debited. No-op when crawl billing is disabled or + there is no active turn (e.g. non-chat callers). + """ + from app.services.token_tracking_service import get_current_accumulator + from app.services.web_crawl_credit_service import WebCrawlCreditService + + if not WebCrawlCreditService.billing_enabled(): + return + acc = get_current_accumulator() + if acc is None: + return + acc.add( + model="web_crawl", + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + cost_micros=WebCrawlCreditService.successes_to_micros(1), + call_kind="web_crawl", + ) + + def extract_domain(url: str) -> str: """Extract the domain from a URL.""" try: @@ -237,6 +264,7 @@ def create_scrape_webpage_tool(): } result = outcome.result + _bill_successful_scrape() content = result.get("content", "") metadata = result.get("metadata", {}) diff --git a/surfsense_backend/app/config/__init__.py b/surfsense_backend/app/config/__init__.py index 8673ded7d..e030b9ee7 100644 --- a/surfsense_backend/app/config/__init__.py +++ b/surfsense_backend/app/config/__init__.py @@ -654,6 +654,22 @@ class Config: ) MICROS_PER_PAGE = int(os.getenv("MICROS_PER_PAGE", "1000")) + # Web-crawl billing debits the credit wallet per *successful* crawl request + # (CrawlOutcomeStatus.SUCCESS). Off by default so self-hosted / OSS installs + # keep crawling effectively-free; hosted deployments set this TRUE. + # + # The price is fully config-driven — there is no hardcoded rate anywhere. + # ``WEB_CRAWL_MICROS_PER_SUCCESS`` is the single source of truth; retune it + # to any rate with just an env change + restart (no code/migration): + # WEB_CRAWL_MICROS_PER_SUCCESS = round(USD_per_1000_crawls * 1_000) + # $1/1000 -> 1000 (default) | $2/1000 -> 2000 | $0.50/1000 -> 500 + WEB_CRAWL_CREDIT_BILLING_ENABLED = ( + os.getenv("WEB_CRAWL_CREDIT_BILLING_ENABLED", "FALSE").upper() == "TRUE" + ) + WEB_CRAWL_MICROS_PER_SUCCESS = int( + os.getenv("WEB_CRAWL_MICROS_PER_SUCCESS", "1000") + ) + # Low-balance WARNING threshold (micro-USD). Surfaced by the quota service # so the UI can nudge the user to top up / enable auto-reload. $0.50. CREDIT_LOW_BALANCE_WARNING_MICROS = int( diff --git a/surfsense_backend/app/services/web_crawl_credit_service.py b/surfsense_backend/app/services/web_crawl_credit_service.py new file mode 100644 index 000000000..8d7539a99 --- /dev/null +++ b/surfsense_backend/app/services/web_crawl_credit_service.py @@ -0,0 +1,143 @@ +"""Service for charging the unified credit wallet per successful web crawl. + +Deliberately mirrors :class:`app.services.etl_credit_service.EtlCreditService`: +a simple **gate -> pre-check -> post-charge** model (no reserve/finalize), +because a crawl has no LLM token accumulator to settle against. The billable +unit is one *successful* crawl (``CrawlOutcomeStatus.SUCCESS``). + +The price is **fully config-driven** — there is no hardcoded rate anywhere. +``config.WEB_CRAWL_MICROS_PER_SUCCESS`` is the single source of truth (default +``1000`` micro-USD == $1 / 1000 crawls); retune it via env + restart, no code +change. When ``config.WEB_CRAWL_CREDIT_BILLING_ENABLED`` is False (the default +for self-hosted / OSS installs) every check/charge is a no-op, preserving the +prior effectively-free crawl behaviour. + +``billing_enabled()`` and ``successes_to_micros()`` are exposed as static +helpers so the chat ``scrape_webpage`` tools can share the flag/price math: +they fold a single success into the current chat turn's existing bill (via the +turn accumulator) instead of debiting the wallet directly. +""" + +from uuid import UUID + +from sqlalchemy import select +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 + +__all__ = ["InsufficientCreditsError", "WebCrawlCreditService"] + + +class WebCrawlCreditService: + """Checks and charges the credit wallet for successful web crawls.""" + + def __init__(self, session: AsyncSession): + self.session = session + + @staticmethod + def billing_enabled() -> bool: + return config.WEB_CRAWL_CREDIT_BILLING_ENABLED + + @staticmethod + def successes_to_micros(successes: int) -> int: + """Convert a successful-crawl count to USD micro-credits. + + Reads ``config.WEB_CRAWL_MICROS_PER_SUCCESS`` — the single, env-tunable + source of truth for crawl price. + """ + return int(successes) * config.WEB_CRAWL_MICROS_PER_SUCCESS + + async def get_available_micros(self, user_id: str | UUID) -> int | None: + """Return spendable credit in micro-USD (``balance - reserved``). + + Returns ``None`` when crawl billing is disabled, which callers treat as + "unlimited" (no blocking, no charge). + """ + if not config.WEB_CRAWL_CREDIT_BILLING_ENABLED: + return None + + from app.db import User + + result = await self.session.execute( + select(User.credit_micros_balance, User.credit_micros_reserved).where( + User.id == user_id + ) + ) + row = result.first() + if not row: + raise ValueError(f"User with ID {user_id} not found") + + balance, reserved = row + return balance - reserved + + async def check_credits( + self, user_id: str | UUID, estimated_successes: int = 1 + ) -> None: + """Raise :class:`InsufficientCreditsError` if the user can't afford + ``estimated_successes`` crawls. + + No-op when crawl billing is disabled. ``estimated_successes`` is a safe + upper bound (``len(urls)``) — actual successes are always <=, so a + passing pre-flight guarantees the wallet can never go negative. + """ + if not config.WEB_CRAWL_CREDIT_BILLING_ENABLED: + return + + required = self.successes_to_micros(estimated_successes) + available = await self.get_available_micros(user_id) + if available is None: + return + + if required > available: + raise InsufficientCreditsError( + message=( + "This crawl would exceed your available credit. " + f"Available: ${available / 1_000_000:.2f}. " + f"Up to {estimated_successes} URL(s) cost about " + f"${required / 1_000_000:.2f}. Add more credits to continue." + ), + balance_micros=available, + required_micros=required, + ) + + async def charge_credits( + self, user_id: str | UUID, successes: int + ) -> int | None: + """Debit the wallet for ``successes`` successful crawls. + + Commits the session (mirroring ``EtlCreditService.charge_credits``), + which also flushes any audit row staged by the caller before this call. + No-op when billing is disabled or ``successes <= 0``. + + Returns the new balance in micros, or ``None`` when nothing was charged. + """ + if not config.WEB_CRAWL_CREDIT_BILLING_ENABLED: + return None + if successes <= 0: + return None + + from app.db import User + + result = await self.session.execute(select(User).where(User.id == user_id)) + user = result.unique().scalar_one_or_none() + if not user: + raise ValueError(f"User with ID {user_id} not found") + + cost = self.successes_to_micros(successes) + user.credit_micros_balance -= cost + await self.session.commit() + await self.session.refresh(user) + + # Best-effort: fire an auto-reload check if the balance dropped low. + try: + from app.services.auto_reload_service import maybe_trigger_auto_reload + + await maybe_trigger_auto_reload(user_id) + except Exception: + pass + + return user.credit_micros_balance diff --git a/surfsense_backend/app/tasks/connector_indexers/webcrawler_indexer.py b/surfsense_backend/app/tasks/connector_indexers/webcrawler_indexer.py index b26e77dfa..8ecba2fc7 100644 --- a/surfsense_backend/app/tasks/connector_indexers/webcrawler_indexer.py +++ b/surfsense_backend/app/tasks/connector_indexers/webcrawler_indexer.py @@ -18,7 +18,10 @@ from app.proprietary.web_crawler import ( WebCrawlerConnector, ) from app.db import Document, DocumentStatus, DocumentType, SearchSourceConnectorType +from app.services.etl_credit_service import InsufficientCreditsError from app.services.task_logging_service import TaskLoggingService +from app.services.token_tracking_service import record_token_usage +from app.services.web_crawl_credit_service import WebCrawlCreditService from app.utils.document_converters import ( create_document_chunks, embed_text, @@ -44,6 +47,23 @@ HeartbeatCallbackType = Callable[[int], Awaitable[None]] HEARTBEAT_INTERVAL_SECONDS = 30 +async def _resolve_workspace_owner(session: AsyncSession, workspace_id: int): + """Return the ``user_id`` that owns ``workspace_id`` (the crawl payer). + + Phase 3c bills the **workspace owner**, not the triggering user (which for + periodic/scheduled runs may differ). Returns ``None`` if the workspace is + gone, letting the caller fall back to the triggering user. + """ + from sqlalchemy import select + + from app.db import Workspace + + result = await session.execute( + select(Workspace.user_id).where(Workspace.id == workspace_id) + ) + return result.scalar_one_or_none() + + async def index_crawled_urls( session: AsyncSession, connector_id: int, @@ -154,6 +174,32 @@ async def index_crawled_urls( ) return 0, "No URLs provided for indexing" + # ======================================================================= + # Phase 3c: crawl-credit pre-flight. Bill the workspace OWNER. ``len(urls)`` + # is a safe upper bound on successes, so a passing check guarantees the + # wallet can't go negative from this run. No-op when billing is disabled + # (self-hosted/OSS) — we also skip the owner lookup in that case. + # ======================================================================= + credit_service = WebCrawlCreditService(session) + owner_user_id = user_id + if credit_service.billing_enabled(): + owner_user_id = ( + await _resolve_workspace_owner(session, workspace_id) or user_id + ) + try: + await credit_service.check_credits(owner_user_id, len(urls)) + except InsufficientCreditsError as credit_error: + await task_logger.log_task_failure( + log_entry, + f"Insufficient crawl credit for connector {connector_id}", + str(credit_error), + {"error_type": "InsufficientCreditsError"}, + ) + logger.warning( + f"Skipping crawl for connector {connector_id}: {credit_error!s}" + ) + return 0, f"Insufficient crawl credit: {credit_error!s}" + await task_logger.log_task_progress( log_entry, f"Starting to process {len(urls)} URLs", @@ -458,6 +504,28 @@ async def index_crawled_urls( else: raise + # ======================================================================= + # Phase 3c: charge the workspace owner for successful crawls. Audit row + # is staged first (add only); charge_credits' commit flushes both the + # ``web_crawl`` TokenUsage row and the balance debit in one transaction. + # Skipped when billing is disabled or nothing crawled successfully. + # ======================================================================= + if credit_service.billing_enabled() and crawls_succeeded > 0: + await record_token_usage( + session, + usage_type="web_crawl", + workspace_id=workspace_id, + user_id=owner_user_id, + cost_micros=credit_service.successes_to_micros(crawls_succeeded), + call_details={ + "urls": len(urls), + "successes": crawls_succeeded, + "connector_id": connector_id, + }, + message_id=None, + ) + await credit_service.charge_credits(owner_user_id, crawls_succeeded) + # Build warning message if there were issues warning_parts = [] if duplicate_content_count > 0: diff --git a/surfsense_backend/scripts/e2e_phase3_crawl_billing.py b/surfsense_backend/scripts/e2e_phase3_crawl_billing.py new file mode 100644 index 000000000..0d51b238b --- /dev/null +++ b/surfsense_backend/scripts/e2e_phase3_crawl_billing.py @@ -0,0 +1,344 @@ +"""Manual functional e2e for Phase 3 crawler + billing (3a / 3b / 3c). + +Run from the backend directory: + cd surfsense_backend + uv run python scripts/e2e_phase3_crawl_billing.py + # or: .\\.venv\\Scripts\\python.exe scripts/e2e_phase3_crawl_billing.py + +What it exercises (everything REAL — live network, live proxy, live DB reads): + + Stage 1 (3a + 3b) — direct fetch + proxy egress-IP proof + crawl_url ladder. + Stage 2 (3c chat surface) — the scrape_webpage tool folds one successful + crawl into the current chat turn's accumulator (billed at finalize). + Stage 3 (3c connector surface) — index_crawled_urls debits the WORKSPACE + OWNER per successful crawl and writes one `web_crawl` TokenUsage row. + +SAFETY: Stage 3 creates a scratch user/workspace/connector inside an outer +transaction that is ALWAYS rolled back (``join_transaction_mode= +"create_savepoint"``), so nothing persists to your database. The only real +side effect is a handful of HTTP requests (small proxy spend). + +This is NOT a pytest test (it needs a live stack + proxy creds + network). It +is the manual functional counterpart to the unit suites; the undetectability / +anti-bot scorecard is a separate deliverable (03f), after 03d/03e. +""" + +import asyncio +import sys +import uuid +from pathlib import Path +from urllib.parse import urlsplit + +from dotenv import load_dotenv + +# --- bootstrap: load .env and put the backend root on sys.path before app.* --- +_BACKEND_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_BACKEND_ROOT)) +for _candidate in (_BACKEND_ROOT / ".env", _BACKEND_ROOT.parent / ".env"): + if _candidate.exists(): + load_dotenv(_candidate) + break + +from sqlalchemy import select, text # noqa: E402 +from sqlalchemy.ext.asyncio import AsyncSession # noqa: E402 + +from app.config import config # noqa: E402 + +# Content-rich, generally crawl-friendly targets (real extraction expected). +_ARTICLE_URLS = [ + "https://en.wikipedia.org/wiki/Competitive_intelligence", + "https://en.wikipedia.org/wiki/Market_research", +] +_IP_ECHO = "https://api.ipify.org?format=json" + + +def _mask(url: str | None) -> str: + if not url: + return "" + p = urlsplit(url) + host = p.hostname or "?" + port = f":{p.port}" if p.port else "" + creds = "***@" if p.username else "" + return f"{p.scheme}://{creds}{host}{port}" + + +def _hr(title: str) -> None: + print(f"\n{'=' * 70}\n{title}\n{'=' * 70}") + + +def _check(label: str, ok: bool, detail: str = "") -> bool: + print(f" [{'PASS' if ok else 'FAIL'}] {label}{f' — {detail}' if detail else ''}") + return ok + + +# =========================================================================== +# Stage 1 — crawl core (3a) + proxy routing (3b) +# =========================================================================== +async def stage1_crawl_and_proxy() -> bool: + _hr("STAGE 1 — crawl_url ladder (3a) + proxy egress (3b)") + from scrapling.fetchers import AsyncFetcher + + from app.proprietary.web_crawler import CrawlOutcomeStatus, WebCrawlerConnector + from app.utils.proxy import get_active_provider, get_proxy_url, is_pool_backed + + ok = True + provider = get_active_provider() + proxy_url = get_proxy_url() + print(f" active proxy provider : {provider.name}") + print(f" proxy url : {_mask(proxy_url)}") + print(f" pool-backed (rotates) : {is_pool_backed()}") + + # Proxy egress-IP proof: direct IP vs proxied IP should differ. + direct_ip = proxied_ip = None + try: + direct = await AsyncFetcher.get(_IP_ECHO, impersonate="chrome", timeout=30) + direct_ip = direct.json().get("ip") + except Exception as exc: # noqa: BLE001 + print(f" [INFO] direct IP fetch failed: {exc}") + if proxy_url: + try: + via = await AsyncFetcher.get( + _IP_ECHO, impersonate="chrome", proxy=proxy_url, timeout=45 + ) + proxied_ip = via.json().get("ip") + except Exception as exc: # noqa: BLE001 + print(f" [INFO] proxied IP fetch failed: {exc}") + print(f" egress IP (direct) : {direct_ip}") + print(f" egress IP (via proxy) : {proxied_ip}") + if proxy_url: + ok &= _check( + "proxy changes egress IP", + bool(proxied_ip) and proxied_ip != direct_ip, + f"{direct_ip} -> {proxied_ip}", + ) + else: + print(" [INFO] no proxy configured — skipping egress-IP assertion") + + # crawl_url end-to-end on a content-rich page. + crawler = WebCrawlerConnector() + outcome = await crawler.crawl_url(_ARTICLE_URLS[0]) + content = (outcome.result or {}).get("content", "") if outcome.result else "" + tier = (outcome.result or {}).get("crawler_type", "?") if outcome.result else "?" + ok &= _check( + "crawl_url returns SUCCESS with content", + outcome.status is CrawlOutcomeStatus.SUCCESS and len(content) > 200, + f"status={outcome.status.value} tier={tier} chars={len(content)}", + ) + return ok + + +# =========================================================================== +# Stage 2 — chat scrape folds cost into the turn accumulator (3c surface 2) +# =========================================================================== +async def stage2_chat_fold() -> bool: + _hr("STAGE 2 — chat scrape_webpage folds crawl cost into turn (3c)") + config.WEB_CRAWL_CREDIT_BILLING_ENABLED = True + price = config.WEB_CRAWL_MICROS_PER_SUCCESS + + from app.agents.chat.multi_agent_chat.main_agent.tools.scrape_webpage import ( + create_scrape_webpage_tool, + ) + from app.services.token_tracking_service import start_turn + + acc = start_turn() + tool = create_scrape_webpage_tool() + result = await tool.ainvoke({"url": _ARTICLE_URLS[0]}) + crawled_ok = "error" not in result and bool(result.get("content")) + print(f" scrape error : {result.get('error', '')}") + print(f" turn cost_micros : {acc.total_cost_micros}") + print(f" call kinds : {[c.call_kind for c in acc.calls]}") + if not crawled_ok: + print(" [INFO] crawl did not succeed (site/proxy) — cannot assert fold") + return False + return _check( + "one web_crawl line folded at configured price", + acc.total_cost_micros == price + and any(c.call_kind == "web_crawl" for c in acc.calls), + f"expected={price} got={acc.total_cost_micros}", + ) + + +# =========================================================================== +# Stage 3 — connector indexer bills the workspace owner (3c surface 1) +# =========================================================================== +async def stage3_indexer_billing() -> bool: + _hr("STAGE 3 — index_crawled_urls bills workspace owner (3c) [rolled back]") + config.WEB_CRAWL_CREDIT_BILLING_ENABLED = True + price = config.WEB_CRAWL_MICROS_PER_SUCCESS + start_balance = 10_000_000 # $10 — plenty for the pre-flight check + + from app.db import ( + Base, + SearchSourceConnector, + SearchSourceConnectorType, + TokenUsage, + User, + Workspace, + engine, + ) + from app.tasks.connector_indexers.webcrawler_indexer import index_crawled_urls + + # Self-bootstrap: if the configured DB has no schema (e.g. an empty + # surfsense_test), create it like the integration harness. No-op against an + # already-migrated DB. DDL is outside the rolled-back txn, so tables persist + # but the scratch rows below do not. + async with engine.connect() as probe: + has_schema = ( + await probe.execute(text("select to_regclass('public.\"user\"')")) + ).scalar() is not None + if not has_schema: + print(" [INFO] empty database — creating schema (vector ext + create_all)") + async with engine.begin() as ddl: + await ddl.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + await ddl.run_sync(Base.metadata.create_all) + + async with engine.connect() as conn: + outer = await conn.begin() + try: + async with AsyncSession( + bind=conn, + expire_on_commit=False, + join_transaction_mode="create_savepoint", + ) as session: + owner = User( + id=uuid.uuid4(), + email=f"e2e-phase3+{uuid.uuid4().hex[:8]}@surfsense.test", + hashed_password="not-a-real-hash", + is_active=True, + is_superuser=False, + is_verified=True, + credit_micros_balance=start_balance, + ) + session.add(owner) + await session.flush() + + # A DISTINCT triggering user (real row — documents.created_by_id + # FKs to it) to prove the OWNER, not the trigger, gets billed. + trigger = User( + id=uuid.uuid4(), + email=f"e2e-trigger+{uuid.uuid4().hex[:8]}@surfsense.test", + hashed_password="not-a-real-hash", + is_active=True, + is_superuser=False, + is_verified=True, + credit_micros_balance=start_balance, + ) + session.add(trigger) + await session.flush() + + ws = Workspace(name="E2E Phase3 Scratch", user_id=owner.id) + session.add(ws) + await session.flush() + + connector = SearchSourceConnector( + name="E2E WebCrawler Scratch", + connector_type=SearchSourceConnectorType.WEBCRAWLER_CONNECTOR, + config={"INITIAL_URLS": _ARTICLE_URLS}, + is_indexable=True, + workspace_id=ws.id, + user_id=owner.id, + ) + session.add(connector) + await session.flush() + + # Snapshot plain values before the indexer's commits so post-run + # reads never lazy-load an expired ORM attribute. + owner_id = owner.id + trigger_id = trigger.id + ws_id = ws.id + connector_id = connector.id + print(f" owner user id : {owner_id}") + print(f" triggering user id : {trigger_id}") + print(f" urls : {len(_ARTICLE_URLS)}") + + total, warning = await index_crawled_urls( + session, connector_id, ws_id, str(trigger_id) + ) + await session.refresh(owner) + await session.refresh(trigger) + debit = start_balance - owner.credit_micros_balance + trigger_debit = start_balance - trigger.credit_micros_balance + + rows = ( + ( + await session.execute( + select(TokenUsage).where( + TokenUsage.usage_type == "web_crawl", + TokenUsage.workspace_id == ws_id, + ) + ) + ) + .scalars() + .all() + ) + + print(f" index result : processed={total} warning={warning}") + print(f" owner debit (micros) : {debit}") + print(f" trigger debit (micros): {trigger_debit}") + print(f" web_crawl usage rows : {len(rows)}") + + ok = True + if not rows: + print( + " [INFO] 0 successful crawls (site/proxy blocked) — " + "nothing billed; cannot assert debit" + ) + return False + + row = rows[0] + successes = (row.call_details or {}).get("successes") + print( + f" audit row : cost_micros={row.cost_micros} " + f"successes={successes} user_id={row.user_id}" + ) + ok &= _check("exactly one web_crawl audit row", len(rows) == 1) + ok &= _check( + "audit billed to OWNER (not trigger user)", + str(row.user_id) == str(owner_id), + ) + ok &= _check( + "triggering user NOT debited", + trigger_debit == 0, + f"trigger_debit={trigger_debit}", + ) + ok &= _check( + "cost == successes * configured price", + row.cost_micros == successes * price, + f"{row.cost_micros} == {successes} * {price}", + ) + ok &= _check( + "wallet debit matches audit cost", + debit == row.cost_micros, + f"debit={debit} cost={row.cost_micros}", + ) + return ok + finally: + await outer.rollback() + print(" [INFO] transaction rolled back — no scratch rows persisted") + + +async def main() -> int: + print("Phase 3 functional e2e (3a/3b/3c) — live network + proxy, DB rolled back") + results: dict[str, bool] = {} + for name, coro in ( + ("Stage 1 crawl+proxy", stage1_crawl_and_proxy), + ("Stage 2 chat fold", stage2_chat_fold), + ("Stage 3 indexer billing", stage3_indexer_billing), + ): + try: + results[name] = await coro() + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" [ERROR] {name} raised: {exc}") + results[name] = False + + _hr("SUMMARY") + for name, ok in results.items(): + print(f" {'PASS' if ok else 'FAIL/SKIP'} — {name}") + return 0 if all(results.values()) else 1 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(main())) diff --git a/surfsense_backend/tests/unit/connector_indexers/test_web_crawl_billing.py b/surfsense_backend/tests/unit/connector_indexers/test_web_crawl_billing.py new file mode 100644 index 000000000..6b6c4037a --- /dev/null +++ b/surfsense_backend/tests/unit/connector_indexers/test_web_crawl_billing.py @@ -0,0 +1,190 @@ +"""Phase 3c crawl-billing wiring in the webcrawler indexer. + +System boundaries mocked: DB session, TaskLoggingService, the crawler, and the +document-conversion helpers. NOT mocked: WebCrawlCreditService (our own code) +and the indexer's billing decisions (owner resolution, pre-flight gate, +audit-then-charge, success counting). +""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID + +import pytest + +import app.tasks.connector_indexers.webcrawler_indexer as _mod +from app.config import config +from app.proprietary.web_crawler import CrawlOutcomeStatus + +pytestmark = pytest.mark.unit + +_CONNECTOR_ID = 7 +_WORKSPACE_ID = 1 +_TRIGGER_USER = "00000000-0000-0000-0000-0000000000aa" +_OWNER_USER = UUID("00000000-0000-0000-0000-0000000000bb") + + +class _FakeUser: + def __init__(self, balance_micros: int, reserved_micros: int = 0): + self.credit_micros_balance = balance_micros + self.credit_micros_reserved = reserved_micros + + +def _make_session(owner_id, balance_micros, reserved_micros=0): + """Mock session serving owner-resolution, get_available_micros and charge.""" + fake_user = _FakeUser(balance_micros, reserved_micros) + session = AsyncMock() + session.add = MagicMock() + session.no_autoflush = MagicMock() + + def _make_result(*_args, **_kwargs): + result = MagicMock() + # _resolve_workspace_owner → select(Workspace.user_id).scalar_one_or_none() + result.scalar_one_or_none.return_value = owner_id + # get_available_micros → .first() → (balance, reserved) + result.first.return_value = ( + fake_user.credit_micros_balance, + fake_user.credit_micros_reserved, + ) + # charge_credits → .unique().scalar_one_or_none() → User + result.unique.return_value.scalar_one_or_none.return_value = fake_user + return result + + session.execute = AsyncMock(side_effect=_make_result) + return session, fake_user + + +def _outcome(success: bool, content: str = "Hello content"): + o = MagicMock() + if success: + o.status = CrawlOutcomeStatus.SUCCESS + o.result = { + "content": content, + "metadata": {"title": "Title"}, + "crawler_type": "scrapling-static", + } + o.error = None + else: + o.status = CrawlOutcomeStatus.FAILED + o.result = None + o.error = "blocked" + return o + + +@pytest.fixture +def indexer_env(monkeypatch): + """Patch every system boundary the indexer touches; return the handles.""" + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True) + monkeypatch.setattr(config, "WEB_CRAWL_MICROS_PER_SUCCESS", 1000) + + # Task logger: async no-op methods, log_task_start returns a sentinel. + task_logger = MagicMock() + task_logger.log_task_start = AsyncMock(return_value=MagicMock()) + task_logger.log_task_progress = AsyncMock() + task_logger.log_task_failure = AsyncMock() + task_logger.log_task_success = AsyncMock() + monkeypatch.setattr(_mod, "TaskLoggingService", MagicMock(return_value=task_logger)) + + # Connector + URL parsing. + connector = MagicMock() + connector.name = "wc" + monkeypatch.setattr(_mod, "get_connector_by_id", AsyncMock(return_value=connector)) + monkeypatch.setattr(_mod, "parse_webcrawler_urls", lambda raw: list(raw)) + + # Crawler. + crawler = MagicMock() + crawler.crawl_url = AsyncMock() + crawler.format_to_structured_document = MagicMock(return_value="doc") + monkeypatch.setattr(_mod, "WebCrawlerConnector", MagicMock(return_value=crawler)) + + # Document-conversion + persistence helpers (system boundary). + monkeypatch.setattr( + _mod, "check_document_by_unique_identifier", AsyncMock(return_value=None) + ) + monkeypatch.setattr( + _mod, "check_duplicate_document_by_hash", AsyncMock(return_value=None) + ) + monkeypatch.setattr(_mod, "embed_text", lambda *_a, **_k: [0.1, 0.2]) + monkeypatch.setattr(_mod, "create_document_chunks", AsyncMock(return_value=[])) + monkeypatch.setattr(_mod, "safe_set_chunks", AsyncMock()) + monkeypatch.setattr(_mod, "update_connector_last_indexed", AsyncMock()) + + # Audit helper — assert it's the owner who is billed. + record_usage = AsyncMock(return_value=MagicMock()) + monkeypatch.setattr(_mod, "record_token_usage", record_usage) + + return { + "connector": connector, + "crawler": crawler, + "record_usage": record_usage, + "task_logger": task_logger, + } + + +async def _run(env, session, urls): + env["connector"].config = {"INITIAL_URLS": urls} + return await _mod.index_crawled_urls( + session, + _CONNECTOR_ID, + _WORKSPACE_ID, + _TRIGGER_USER, + ) + + +async def test_charges_owner_for_successful_crawls(indexer_env): + session, user = _make_session(_OWNER_USER, balance_micros=100_000) + indexer_env["crawler"].crawl_url.side_effect = [_outcome(True), _outcome(True)] + + total, warning = await _run(indexer_env, session, ["https://a.com", "https://b.com"]) + + assert total == 2 + assert warning is None + # Owner debited 2 * 1000. + assert user.credit_micros_balance == 100_000 - 2000 + # One audit row, billed to the OWNER (not the triggering user), correct cost. + indexer_env["record_usage"].assert_awaited_once() + kwargs = indexer_env["record_usage"].await_args.kwargs + assert kwargs["usage_type"] == "web_crawl" + assert kwargs["user_id"] == _OWNER_USER + assert kwargs["workspace_id"] == _WORKSPACE_ID + assert kwargs["cost_micros"] == 2000 + + +async def test_preflight_blocks_run_when_insufficient(indexer_env): + # balance 1000 < 3 URLs * 1000 required. + session, user = _make_session(_OWNER_USER, balance_micros=1000) + + total, warning = await _run( + indexer_env, session, ["https://a.com", "https://b.com", "https://c.com"] + ) + + assert total == 0 + assert "insufficient crawl credit" in warning.lower() + # Never crawled, never billed, balance untouched. + indexer_env["crawler"].crawl_url.assert_not_awaited() + indexer_env["record_usage"].assert_not_awaited() + assert user.credit_micros_balance == 1000 + + +async def test_failed_crawls_are_free(indexer_env): + session, user = _make_session(_OWNER_USER, balance_micros=100_000) + indexer_env["crawler"].crawl_url.side_effect = [_outcome(False)] + + total, _warning = await _run(indexer_env, session, ["https://a.com"]) + + assert total == 0 + # crawls_succeeded == 0 → no audit, no charge. + indexer_env["record_usage"].assert_not_awaited() + assert user.credit_micros_balance == 100_000 + + +async def test_disabled_skips_billing_but_still_crawls(indexer_env, monkeypatch): + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False) + session, user = _make_session(_OWNER_USER, balance_micros=100_000) + indexer_env["crawler"].crawl_url.side_effect = [_outcome(True)] + + total, _warning = await _run(indexer_env, session, ["https://a.com"]) + + assert total == 1 + indexer_env["crawler"].crawl_url.assert_awaited_once() + indexer_env["record_usage"].assert_not_awaited() + assert user.credit_micros_balance == 100_000 diff --git a/surfsense_backend/tests/unit/services/test_web_crawl_credit_service.py b/surfsense_backend/tests/unit/services/test_web_crawl_credit_service.py new file mode 100644 index 000000000..0e79836ac --- /dev/null +++ b/surfsense_backend/tests/unit/services/test_web_crawl_credit_service.py @@ -0,0 +1,215 @@ +"""Unit tests for WebCrawlCreditService and the chat-scrape fold helper (Phase 3c). + +Covers: + A) successes_to_micros — config-driven price (the single source of truth), + including retune behaviour. + B) billing_enabled gate. + C) check_credits — sufficient / insufficient / disabled no-op. + D) charge_credits — debit / disabled no-op / zero no-op. + E) Chat-scrape fold helper — folds one success into the turn accumulator only + when billing is enabled and a turn is active. + +The wallet logic runs against the real service with a mock DB session at the +system boundary (mirrors tests/unit/connector_indexers/test_etl_credits.py). +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from app.config import config +from app.services.etl_credit_service import InsufficientCreditsError +from app.services.web_crawl_credit_service import WebCrawlCreditService + +pytestmark = pytest.mark.unit + +_USER_ID = "00000000-0000-0000-0000-000000000001" + + +@pytest.fixture(autouse=True) +def _enable_crawl_billing(monkeypatch): + """Force crawl billing ON; default is off (self-hosted) which no-ops everything.""" + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True) + monkeypatch.setattr(config, "WEB_CRAWL_MICROS_PER_SUCCESS", 1000) + + +@pytest.fixture(autouse=True) +def _stub_auto_reload(monkeypatch): + """charge_credits fires a best-effort auto-reload check; stub it out.""" + import app.services.auto_reload_service as _ar + + monkeypatch.setattr(_ar, "maybe_trigger_auto_reload", AsyncMock()) + + +class _FakeUser: + def __init__(self, balance_micros: int = 0, reserved_micros: int = 0): + self.credit_micros_balance = balance_micros + self.credit_micros_reserved = reserved_micros + + +def _make_session(balance_micros: int = 100_000, reserved_micros: int = 0): + """Mock DB session compatible with get_available_micros + charge_credits.""" + fake_user = _FakeUser(balance_micros, reserved_micros) + session = AsyncMock() + + def _make_result(*_args, **_kwargs): + result = MagicMock() + result.first.return_value = ( + fake_user.credit_micros_balance, + fake_user.credit_micros_reserved, + ) + result.unique.return_value.scalar_one_or_none.return_value = fake_user + return result + + session.execute = AsyncMock(side_effect=_make_result) + return session, fake_user + + +# =================================================================== +# A) successes_to_micros — config-driven price +# =================================================================== + + +class TestSuccessesToMicros: + def test_default_is_one_dollar_per_thousand(self, monkeypatch): + monkeypatch.setattr(config, "WEB_CRAWL_MICROS_PER_SUCCESS", 1000) + assert WebCrawlCreditService.successes_to_micros(1) == 1000 + assert WebCrawlCreditService.successes_to_micros(1000) == 1_000_000 # $1 + + def test_zero_successes_cost_nothing(self): + assert WebCrawlCreditService.successes_to_micros(0) == 0 + + @pytest.mark.parametrize( + ("per_success", "successes", "expected"), + [ + (2000, 1000, 2_000_000), # $2 / 1000 + (500, 1000, 500_000), # $0.50 / 1000 + (5000, 10, 50_000), # $5 / 1000 + ], + ) + def test_retune_is_config_driven( + self, monkeypatch, per_success, successes, expected + ): + """No hardcoded rate: changing the config knob changes the price.""" + monkeypatch.setattr(config, "WEB_CRAWL_MICROS_PER_SUCCESS", per_success) + assert WebCrawlCreditService.successes_to_micros(successes) == expected + + +# =================================================================== +# B) billing_enabled gate +# =================================================================== + + +class TestBillingEnabled: + def test_reads_config_flag(self, monkeypatch): + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True) + assert WebCrawlCreditService.billing_enabled() is True + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False) + assert WebCrawlCreditService.billing_enabled() is False + + +# =================================================================== +# C) check_credits +# =================================================================== + + +class TestCheckCredits: + async def test_sufficient_credit_passes(self): + session, _user = _make_session(balance_micros=10_000) + svc = WebCrawlCreditService(session) + # 5 URLs * 1000 = 5000 <= 10000 → no raise + await svc.check_credits(_USER_ID, estimated_successes=5) + + async def test_insufficient_credit_raises(self): + session, _user = _make_session(balance_micros=3_000) + svc = WebCrawlCreditService(session) + # 5 URLs * 1000 = 5000 > 3000 → raise + with pytest.raises(InsufficientCreditsError): + await svc.check_credits(_USER_ID, estimated_successes=5) + + async def test_reserved_credit_reduces_available(self): + session, _user = _make_session(balance_micros=10_000, reserved_micros=8_000) + svc = WebCrawlCreditService(session) + # available = 2000; 3 * 1000 = 3000 > 2000 → raise + with pytest.raises(InsufficientCreditsError): + await svc.check_credits(_USER_ID, estimated_successes=3) + + async def test_disabled_is_noop(self, monkeypatch): + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False) + session, _user = _make_session(balance_micros=0) + svc = WebCrawlCreditService(session) + await svc.check_credits(_USER_ID, estimated_successes=10_000) + session.execute.assert_not_called() + + +# =================================================================== +# D) charge_credits +# =================================================================== + + +class TestChargeCredits: + async def test_debits_per_success(self): + session, user = _make_session(balance_micros=100_000) + svc = WebCrawlCreditService(session) + new_balance = await svc.charge_credits(_USER_ID, successes=3) + assert user.credit_micros_balance == 100_000 - 3 * 1000 + assert new_balance == 97_000 + session.commit.assert_awaited() + + async def test_zero_successes_is_noop(self): + session, user = _make_session(balance_micros=100_000) + svc = WebCrawlCreditService(session) + result = await svc.charge_credits(_USER_ID, successes=0) + assert result is None + assert user.credit_micros_balance == 100_000 + session.commit.assert_not_awaited() + + async def test_disabled_is_noop(self, monkeypatch): + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False) + session, user = _make_session(balance_micros=100_000) + svc = WebCrawlCreditService(session) + result = await svc.charge_credits(_USER_ID, successes=5) + assert result is None + assert user.credit_micros_balance == 100_000 + session.execute.assert_not_called() + + +# =================================================================== +# E) Chat-scrape fold helper +# =================================================================== + + +class TestChatScrapeFold: + """The scrape_webpage tools fold one success into the live turn accumulator.""" + + def _import_helper(self): + from app.agents.chat.multi_agent_chat.main_agent.tools.scrape_webpage import ( + _bill_successful_scrape, + ) + + return _bill_successful_scrape + + def test_folds_into_active_turn_when_enabled(self): + from app.services.token_tracking_service import start_turn + + acc = start_turn() + self._import_helper()() + assert acc.total_cost_micros == 1000 + assert len(acc.calls) == 1 + assert acc.calls[0].call_kind == "web_crawl" + + def test_noop_when_billing_disabled(self, monkeypatch): + from app.services.token_tracking_service import start_turn + + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False) + acc = start_turn() + self._import_helper()() + assert acc.calls == [] + assert acc.total_cost_micros == 0 + + def test_noop_when_no_active_turn(self): + import contextvars + + # Run in a fresh context so the turn ContextVar resolves to its default + # (None) regardless of other tests — must not raise. + contextvars.Context().run(self._import_helper())