diff --git a/surfsense_backend/app/capabilities/billing.py b/surfsense_backend/app/capabilities/billing.py index 20d1059bf..636920a9f 100644 --- a/surfsense_backend/app/capabilities/billing.py +++ b/surfsense_backend/app/capabilities/billing.py @@ -12,8 +12,10 @@ from app.capabilities.types import ( BillingUnit, CapabilityContext, ) +from app.config import config from app.services.token_tracking_service import record_token_usage from app.services.web_crawl_credit_service import WebCrawlCreditService +from app.utils.captcha import captcha_enabled if TYPE_CHECKING: from uuid import UUID @@ -36,29 +38,49 @@ async def gate_capability( async def _gate_web_crawl(ctx: CapabilityContext, estimated_successes: int) -> None: + """Reserve the worst-case cost: crawl successes + worst-case captcha attempts. + + Captcha budget is only reserved when solving is actually enabled — with + solving off, attempts can never happen, so reserving would wrongly block a + run for captcha that will never be attempted. Mirrors the indexer path (3d). + """ service = WebCrawlCreditService(ctx.session) - if not service.billing_enabled(): + crawl_on = service.billing_enabled() + captcha_on = service.captcha_billing_enabled() and captcha_enabled() + if not crawl_on and not captcha_on: return owner_user_id = await _resolve_workspace_owner(ctx.session, ctx.workspace_id) if owner_user_id is None: return - await service.check_credits(owner_user_id, estimated_successes) + + required_micros = 0 + if crawl_on: + required_micros += service.successes_to_micros(estimated_successes) + if captcha_on: + worst_case_attempts = estimated_successes * config.CAPTCHA_MAX_ATTEMPTS_PER_URL + required_micros += service.captcha_solves_to_micros(worst_case_attempts) + await service.check_balance(owner_user_id, required_micros) async def charge_capability( output: BillableOutput, unit: BillingUnit | None, ctx: CapabilityContext ) -> None: - """Bill the workspace owner for this result's billable successes (03c). ``None`` = free.""" + """Bill the workspace owner for this result's billable successes (03c). ``None`` = free. + + For crawl-backed verbs this also bills any captcha *attempts* (Phase 3d) as a + separate per-attempt unit — the solver charges per attempt even when the crawl + ultimately failed, so it can't ride the per-success crawl meter. + """ if unit is None: return - units = output.billable_units - if units <= 0: - return if unit is BillingUnit.WEB_CRAWL: - await _charge_web_crawl(ctx, units) + await _charge_web_crawl(ctx, output.billable_units) + await _charge_captcha(ctx, getattr(output, "captcha_attempts", 0)) async def _charge_web_crawl(ctx: CapabilityContext, successes: int) -> None: + if successes <= 0: + return service = WebCrawlCreditService(ctx.session) if not service.billing_enabled(): return @@ -77,6 +99,27 @@ async def _charge_web_crawl(ctx: CapabilityContext, successes: int) -> None: await service.charge_credits(owner_user_id, successes) +async def _charge_captcha(ctx: CapabilityContext, attempts: int) -> None: + if attempts <= 0: + return + service = WebCrawlCreditService(ctx.session) + if not service.captcha_billing_enabled(): + return + owner_user_id = await _resolve_workspace_owner(ctx.session, ctx.workspace_id) + if owner_user_id is None: + return + # Stage the audit row before charge_captcha's commit flushes both. + await record_token_usage( + ctx.session, + usage_type="web_crawl_captcha", + workspace_id=ctx.workspace_id, + user_id=owner_user_id, + cost_micros=service.captcha_solves_to_micros(attempts), + call_details={"attempts": attempts}, + ) + await service.charge_captcha(owner_user_id, attempts) + + async def _resolve_workspace_owner( session: AsyncSession, workspace_id: int ) -> UUID | None: diff --git a/surfsense_backend/app/capabilities/web/discover/schemas.py b/surfsense_backend/app/capabilities/web/discover/schemas.py index ae74d2c56..229454e7e 100644 --- a/surfsense_backend/app/capabilities/web/discover/schemas.py +++ b/surfsense_backend/app/capabilities/web/discover/schemas.py @@ -10,7 +10,7 @@ class DiscoverInput(BaseModel): description="What to search the web for, phrased in natural language." ) top_k: int = Field( - default=10, description="Maximum number of results to return (1-50)." + default=10, ge=1, le=50, description="Maximum number of results to return (1-50)." ) diff --git a/surfsense_backend/app/capabilities/web/scrape/executor.py b/surfsense_backend/app/capabilities/web/scrape/executor.py index 46f067b56..05459ee6a 100644 --- a/surfsense_backend/app/capabilities/web/scrape/executor.py +++ b/surfsense_backend/app/capabilities/web/scrape/executor.py @@ -16,18 +16,29 @@ def build_scrape_executor(engine: WebCrawlerConnector | None = None) -> Executor engine = engine or WebCrawlerConnector() async def execute(payload: ScrapeInput) -> ScrapeOutput: - rows = [_to_row(url, await engine.crawl_url(url)) for url in payload.urls] - return ScrapeOutput(rows=rows) + outcomes = [await engine.crawl_url(url) for url in payload.urls] + rows = [ + _to_row(url, outcome, payload.max_length) + for url, outcome in zip(payload.urls, outcomes, strict=True) + ] + return ScrapeOutput( + rows=rows, + captcha_attempts=sum(o.captcha_attempts for o in outcomes), + captcha_solved=sum(1 for o in outcomes if o.captcha_solved), + ) return execute -def _to_row(url: str, outcome: CrawlOutcome) -> ScrapeRow: +def _to_row(url: str, outcome: CrawlOutcome, max_length: int) -> ScrapeRow: if outcome.status is CrawlOutcomeStatus.SUCCESS and outcome.result: + content = outcome.result.get("content") + if content is not None: + content = content[:max_length] return ScrapeRow( url=url, status="success", - content=outcome.result.get("content"), + content=content, metadata=outcome.result.get("metadata"), ) status = "empty" if outcome.status is CrawlOutcomeStatus.EMPTY else "failed" diff --git a/surfsense_backend/app/capabilities/web/scrape/schemas.py b/surfsense_backend/app/capabilities/web/scrape/schemas.py index 415a64554..04faaa516 100644 --- a/surfsense_backend/app/capabilities/web/scrape/schemas.py +++ b/surfsense_backend/app/capabilities/web/scrape/schemas.py @@ -56,6 +56,10 @@ class ScrapeOutput(BaseModel): rows: list[ScrapeRow] = Field( description="One result per requested URL, in the same order." ) + # Billing-only telemetry (Phase 3d), excluded from the client-facing result: + # captcha solves are metered per attempt, independent of crawl success. + captcha_attempts: int = Field(default=0, exclude=True) + captcha_solved: int = Field(default=0, exclude=True) @property def billable_units(self) -> int: diff --git a/surfsense_backend/tests/unit/capabilities/test_billing.py b/surfsense_backend/tests/unit/capabilities/test_billing.py index 45e65c257..1f70afc56 100644 --- a/surfsense_backend/tests/unit/capabilities/test_billing.py +++ b/surfsense_backend/tests/unit/capabilities/test_billing.py @@ -92,6 +92,53 @@ async def test_charges_workspace_owner_per_successful_crawl(monkeypatch, record_ assert kwargs["cost_micros"] == 2000 +def _output_with_captcha(*statuses: str, attempts: int, solved: int) -> ScrapeOutput: + out = _output(*statuses) + out.captcha_attempts = attempts + out.captcha_solved = solved + return out + + +async def test_charges_workspace_owner_per_captcha_attempt_even_when_crawl_failed( + monkeypatch, record_usage +): + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True) + monkeypatch.setattr(config, "WEB_CRAWL_CAPTCHA_BILLING_ENABLED", True) + monkeypatch.setattr(config, "WEB_CRAWL_CAPTCHA_MICROS_PER_SOLVE", 3000) + session, user = _make_session(_OWNER, balance_micros=100_000) + + # Crawl failed (no billable success) but the solver ran twice — attempts bill. + await charge_capability( + _output_with_captcha("failed", attempts=2, solved=1), + BillingUnit.WEB_CRAWL, + _ctx(session), + ) + + assert user.credit_micros_balance == 100_000 - 2 * 3000 + record_usage.assert_awaited_once() + kwargs = record_usage.await_args.kwargs + assert kwargs["usage_type"] == "web_crawl_captcha" + assert kwargs["user_id"] == _OWNER + assert kwargs["cost_micros"] == 6000 + + +async def test_captcha_billing_disabled_does_not_charge_attempts( + monkeypatch, record_usage +): + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True) + monkeypatch.setattr(config, "WEB_CRAWL_CAPTCHA_BILLING_ENABLED", False) + session, user = _make_session(_OWNER, balance_micros=100_000) + + await charge_capability( + _output_with_captcha("failed", attempts=2, solved=1), + BillingUnit.WEB_CRAWL, + _ctx(session), + ) + + record_usage.assert_not_awaited() + assert user.credit_micros_balance == 100_000 + + async def test_no_successful_rows_is_free(monkeypatch, record_usage): monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True) session, user = _make_session(_OWNER, balance_micros=100_000) @@ -176,6 +223,32 @@ async def test_gate_is_noop_when_disabled(monkeypatch): ) +async def test_gate_reserves_worst_case_captcha_when_solving_enabled(monkeypatch): + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False) + monkeypatch.setattr(config, "WEB_CRAWL_CAPTCHA_BILLING_ENABLED", True) + monkeypatch.setattr(config, "WEB_CRAWL_CAPTCHA_MICROS_PER_SOLVE", 3000) + monkeypatch.setattr(config, "CAPTCHA_MAX_ATTEMPTS_PER_URL", 3) + monkeypatch.setattr(billing, "captcha_enabled", lambda: True) + session = _gate_session(_OWNER, balance_micros=5000) # < 1 url * 3 * 3000 + + with pytest.raises(InsufficientCreditsError): + await gate_capability( + ScrapeInput(urls=["https://a.com"]), BillingUnit.WEB_CRAWL, _ctx(session) + ) + + +async def test_gate_does_not_reserve_captcha_when_solving_disabled(monkeypatch): + monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False) + monkeypatch.setattr(config, "WEB_CRAWL_CAPTCHA_BILLING_ENABLED", True) + monkeypatch.setattr(billing, "captcha_enabled", lambda: False) + session = _gate_session(_OWNER, balance_micros=0) + + # Solving off → attempts can never happen → nothing to reserve → passes. + await gate_capability( + ScrapeInput(urls=["https://a.com"]), BillingUnit.WEB_CRAWL, _ctx(session) + ) + + async def test_gate_is_noop_for_free_verb(monkeypatch): monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True) session = _gate_session(_OWNER, balance_micros=0) diff --git a/surfsense_backend/tests/unit/capabilities/web/discover/test_schemas.py b/surfsense_backend/tests/unit/capabilities/web/discover/test_schemas.py new file mode 100644 index 000000000..f098014d8 --- /dev/null +++ b/surfsense_backend/tests/unit/capabilities/web/discover/test_schemas.py @@ -0,0 +1,22 @@ +"""``web.discover`` I/O contract bounds.""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from app.capabilities.web.discover.schemas import DiscoverInput + +pytestmark = pytest.mark.unit + + +def test_top_k_defaults_and_accepts_the_valid_range(): + assert DiscoverInput(query="q").top_k == 10 + assert DiscoverInput(query="q", top_k=1).top_k == 1 + assert DiscoverInput(query="q", top_k=50).top_k == 50 + + +@pytest.mark.parametrize("bad", [0, -1, 51, 100]) +def test_top_k_outside_1_to_50_is_rejected(bad): + with pytest.raises(ValidationError): + DiscoverInput(query="q", top_k=bad) diff --git a/surfsense_backend/tests/unit/capabilities/web/scrape/test_executor.py b/surfsense_backend/tests/unit/capabilities/web/scrape/test_executor.py index 9118b6944..b37a0346c 100644 --- a/surfsense_backend/tests/unit/capabilities/web/scrape/test_executor.py +++ b/surfsense_backend/tests/unit/capabilities/web/scrape/test_executor.py @@ -71,6 +71,53 @@ async def test_scrape_returns_one_row_per_url_in_input_order(): assert [row.content for row in out.rows] == ["A", "B", "C"] +async def test_content_longer_than_max_length_is_truncated(): + url = "https://long.com" + crawler = _FakeCrawler({url: _success("A" * 100, {"title": "Long"})}) + execute = build_scrape_executor(engine=crawler) + + out = await execute(ScrapeInput(urls=[url], max_length=10)) + + assert out.rows[0].content == "A" * 10 + + +async def test_content_within_max_length_is_untouched(): + url = "https://short.com" + crawler = _FakeCrawler({url: _success("hello", {"title": "Short"})}) + execute = build_scrape_executor(engine=crawler) + + out = await execute(ScrapeInput(urls=[url], max_length=10)) + + assert out.rows[0].content == "hello" + + +async def test_scrape_surfaces_total_captcha_attempts_for_billing(): + ok, blocked = "https://ok.com", "https://blocked.com" + crawler = _FakeCrawler( + { + ok: CrawlOutcome( + status=CrawlOutcomeStatus.SUCCESS, + result={"content": "OK", "metadata": {}}, + captcha_attempts=2, + captcha_solved=True, + ), + blocked: CrawlOutcome( + status=CrawlOutcomeStatus.FAILED, + error="blocked", + captcha_attempts=1, + captcha_solved=False, + ), + } + ) + execute = build_scrape_executor(engine=crawler) + + out = await execute(ScrapeInput(urls=[ok, blocked])) + + # Attempts bill even when the crawl ultimately failed (Phase 3d). + assert out.captcha_attempts == 3 + assert out.captcha_solved == 1 + + async def test_partial_failure_keeps_the_batch_and_labels_each_url(): ok, empty, failed = "https://ok.com", "https://empty.com", "https://failed.com" crawler = _FakeCrawler(