mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-08 22:22:17 +02:00
80927a2872
Add per-item, per-platform billing for the platform-native connectors (Reddit, Google Search, Google Maps places/reviews, YouTube videos/comments) through the capability gate/charge seam. Rates are config-driven with a shared wallet-credit module (wallet_credit) and a dedicated PlatformScrapeCreditService; agent and REST capability runs now record cost_micros. Google Maps scrape dual-meters places and attached reviews. Remove the main-agent scrape_webpage tool now that the web.crawl capability covers single-page (maxCrawlDepth=0) and site crawling. The main agent now reaches crawling via task(web_crawler, ...). Update prompts, tool catalog, receipts, skills, proprietary docs, and tests; drop the obsolete chat-turn crawl fold path. Co-authored-by: Cursor <cursoragent@cursor.com>
427 lines
15 KiB
Python
427 lines
15 KiB
Python
"""Billing charges the workspace owner once per billable success at the executor (03c).
|
|
|
|
Boundaries mocked: the DB session and the audit helper. NOT mocked: the real
|
|
WebCrawlCreditService debit math and the owner-billed decision.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
from uuid import UUID
|
|
|
|
import pytest
|
|
|
|
import app.capabilities.core.billing as billing
|
|
from app.capabilities.core.billing import charge_capability, gate_capability
|
|
from app.capabilities.core.types import BillingUnit, CapabilityContext
|
|
from app.capabilities.web.crawl.schemas import CrawlInput, CrawlItem, CrawlOutput
|
|
from app.config import config
|
|
from app.services.web_crawl_credit_service import InsufficientCreditsError
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
_WORKSPACE_ID = 1
|
|
_OWNER = 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):
|
|
"""Mock session serving owner-resolution and the charge_credits debit."""
|
|
fake_user = _FakeUser(balance_micros)
|
|
session = AsyncMock()
|
|
session.add = MagicMock()
|
|
|
|
def _make_result(*_args, **_kwargs):
|
|
result = MagicMock()
|
|
result.scalar_one_or_none.return_value = owner_id # owner resolution
|
|
result.unique.return_value.scalar_one_or_none.return_value = fake_user # debit
|
|
return result
|
|
|
|
session.execute = AsyncMock(side_effect=_make_result)
|
|
return session, fake_user
|
|
|
|
|
|
def _output(*statuses: str) -> CrawlOutput:
|
|
return CrawlOutput(
|
|
items=[
|
|
CrawlItem(url=f"https://{i}.com", status=status)
|
|
for i, status in enumerate(statuses)
|
|
]
|
|
)
|
|
|
|
|
|
def _ctx(session) -> CapabilityContext:
|
|
return CapabilityContext(session=session, workspace_id=_WORKSPACE_ID)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _stub_auto_reload(monkeypatch):
|
|
import app.services.auto_reload_service as ar
|
|
|
|
monkeypatch.setattr(ar, "maybe_trigger_auto_reload", AsyncMock())
|
|
|
|
|
|
@pytest.fixture
|
|
def record_usage(monkeypatch):
|
|
rec = AsyncMock(return_value=MagicMock())
|
|
monkeypatch.setattr(billing, "record_token_usage", rec)
|
|
return rec
|
|
|
|
|
|
async def test_charges_workspace_owner_per_successful_crawl(monkeypatch, record_usage):
|
|
monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True)
|
|
monkeypatch.setattr(config, "WEB_CRAWL_MICROS_PER_SUCCESS", 1000)
|
|
session, user = _make_session(_OWNER, balance_micros=100_000)
|
|
|
|
await charge_capability(
|
|
_output("success", "empty", "success"), BillingUnit.WEB_CRAWL, _ctx(session)
|
|
)
|
|
|
|
# Owner debited 2 * 1000; one web_crawl audit row billed to the OWNER.
|
|
assert user.credit_micros_balance == 100_000 - 2000
|
|
record_usage.assert_awaited_once()
|
|
kwargs = record_usage.await_args.kwargs
|
|
assert kwargs["usage_type"] == "web_crawl"
|
|
assert kwargs["user_id"] == _OWNER
|
|
assert kwargs["workspace_id"] == _WORKSPACE_ID
|
|
assert kwargs["cost_micros"] == 2000
|
|
|
|
|
|
def _output_with_captcha(*statuses: str, attempts: int, solved: int) -> CrawlOutput:
|
|
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)
|
|
|
|
await charge_capability(
|
|
_output("empty", "failed"), BillingUnit.WEB_CRAWL, _ctx(session)
|
|
)
|
|
|
|
record_usage.assert_not_awaited()
|
|
assert user.credit_micros_balance == 100_000
|
|
|
|
|
|
async def test_disabled_is_noop(monkeypatch, record_usage):
|
|
monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False)
|
|
session, user = _make_session(_OWNER, balance_micros=100_000)
|
|
|
|
await charge_capability(
|
|
_output("success", "success"), BillingUnit.WEB_CRAWL, _ctx(session)
|
|
)
|
|
|
|
record_usage.assert_not_awaited()
|
|
session.execute.assert_not_called()
|
|
assert user.credit_micros_balance == 100_000
|
|
|
|
|
|
async def test_free_verb_without_a_unit_is_noop(monkeypatch, record_usage):
|
|
monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True)
|
|
session, user = _make_session(_OWNER, balance_micros=100_000)
|
|
|
|
await charge_capability(_output("success", "success"), None, _ctx(session))
|
|
|
|
record_usage.assert_not_awaited()
|
|
session.execute.assert_not_called()
|
|
assert user.credit_micros_balance == 100_000
|
|
|
|
|
|
def _gate_session(owner_id, balance_micros):
|
|
"""Mock session serving owner-resolution and the spendable-balance read."""
|
|
session = AsyncMock()
|
|
|
|
def _make_result(*_args, **_kwargs):
|
|
result = MagicMock()
|
|
result.scalar_one_or_none.return_value = owner_id # owner resolution
|
|
result.first.return_value = (balance_micros, 0) # balance, reserved
|
|
return result
|
|
|
|
session.execute = AsyncMock(side_effect=_make_result)
|
|
return session
|
|
|
|
|
|
async def test_gate_blocks_when_worst_case_exceeds_balance(monkeypatch):
|
|
monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True)
|
|
monkeypatch.setattr(config, "WEB_CRAWL_MICROS_PER_SUCCESS", 1000)
|
|
session = _gate_session(_OWNER, balance_micros=1500) # affords 1 crawl, not 2
|
|
|
|
with pytest.raises(InsufficientCreditsError):
|
|
await gate_capability(
|
|
CrawlInput(startUrls=["https://a.com", "https://b.com"]),
|
|
BillingUnit.WEB_CRAWL,
|
|
_ctx(session),
|
|
)
|
|
|
|
|
|
async def test_gate_passes_when_balance_covers_worst_case(monkeypatch):
|
|
monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True)
|
|
monkeypatch.setattr(config, "WEB_CRAWL_MICROS_PER_SUCCESS", 1000)
|
|
session = _gate_session(_OWNER, balance_micros=100_000)
|
|
|
|
await gate_capability(
|
|
CrawlInput(startUrls=["https://a.com", "https://b.com"]),
|
|
BillingUnit.WEB_CRAWL,
|
|
_ctx(session),
|
|
)
|
|
|
|
|
|
async def test_gate_is_noop_when_disabled(monkeypatch):
|
|
monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False)
|
|
session = _gate_session(_OWNER, balance_micros=0)
|
|
|
|
await gate_capability(
|
|
CrawlInput(startUrls=["https://a.com"]), BillingUnit.WEB_CRAWL, _ctx(session)
|
|
)
|
|
|
|
|
|
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(
|
|
CrawlInput(startUrls=["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(
|
|
CrawlInput(startUrls=["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)
|
|
|
|
await gate_capability(CrawlInput(startUrls=["https://a.com"]), None, _ctx(session))
|
|
|
|
session.execute.assert_not_called()
|
|
|
|
|
|
# ===================================================================
|
|
# Platform scraper per-item billing (Reddit / Search / Maps / YouTube)
|
|
# ===================================================================
|
|
|
|
|
|
class _FakePlatformOutput:
|
|
"""Stand-in for a verb output: only the billing-read properties matter."""
|
|
|
|
def __init__(self, items: int, attached_review_count: int = 0):
|
|
self._items = items
|
|
self._reviews = attached_review_count
|
|
|
|
@property
|
|
def billable_units(self) -> int:
|
|
return self._items
|
|
|
|
@property
|
|
def attached_review_count(self) -> int:
|
|
return self._reviews
|
|
|
|
|
|
class _FakePlatformInput:
|
|
"""Stand-in for a verb input reporting its worst-case unit counts."""
|
|
|
|
def __init__(self, estimated_units: int, estimated_review_units: int = 0):
|
|
self._units = estimated_units
|
|
self._review_units = estimated_review_units
|
|
|
|
@property
|
|
def estimated_units(self) -> int:
|
|
return self._units
|
|
|
|
@property
|
|
def estimated_review_units(self) -> int:
|
|
return self._review_units
|
|
|
|
|
|
@pytest.fixture
|
|
def _enable_platform_billing(monkeypatch):
|
|
monkeypatch.setattr(config, "PLATFORM_SCRAPE_BILLING_ENABLED", True)
|
|
|
|
|
|
async def test_platform_charges_owner_per_item(
|
|
monkeypatch, record_usage, _enable_platform_billing
|
|
):
|
|
monkeypatch.setattr(config, "REDDIT_SCRAPE_MICROS_PER_ITEM", 3500)
|
|
session, user = _make_session(_OWNER, balance_micros=1_000_000)
|
|
|
|
charged = await charge_capability(
|
|
_FakePlatformOutput(3), BillingUnit.REDDIT_ITEM, _ctx(session)
|
|
)
|
|
|
|
assert charged == 3 * 3500
|
|
assert user.credit_micros_balance == 1_000_000 - 3 * 3500
|
|
record_usage.assert_awaited_once()
|
|
kwargs = record_usage.await_args.kwargs
|
|
assert kwargs["usage_type"] == "reddit_item"
|
|
assert kwargs["user_id"] == _OWNER
|
|
assert kwargs["workspace_id"] == _WORKSPACE_ID
|
|
assert kwargs["cost_micros"] == 3 * 3500
|
|
|
|
|
|
async def test_platform_maps_scrape_dual_meters_places_and_reviews(
|
|
monkeypatch, record_usage, _enable_platform_billing
|
|
):
|
|
monkeypatch.setattr(config, "GOOGLE_MAPS_MICROS_PER_PLACE", 5000)
|
|
monkeypatch.setattr(config, "GOOGLE_MAPS_MICROS_PER_REVIEW", 2000)
|
|
session, user = _make_session(_OWNER, balance_micros=1_000_000)
|
|
|
|
# 2 places + 10 attached reviews -> 2*5000 + 10*2000 = 30000.
|
|
charged = await charge_capability(
|
|
_FakePlatformOutput(2, attached_review_count=10),
|
|
BillingUnit.GOOGLE_MAPS_PLACE,
|
|
_ctx(session),
|
|
)
|
|
|
|
assert charged == 2 * 5000 + 10 * 2000
|
|
assert user.credit_micros_balance == 1_000_000 - 30_000
|
|
assert record_usage.await_count == 2
|
|
usage_types = {c.kwargs["usage_type"] for c in record_usage.await_args_list}
|
|
assert usage_types == {"google_maps_place", "google_maps_review"}
|
|
|
|
|
|
async def test_platform_charge_disabled_is_noop(monkeypatch, record_usage):
|
|
monkeypatch.setattr(config, "PLATFORM_SCRAPE_BILLING_ENABLED", False)
|
|
monkeypatch.setattr(config, "REDDIT_SCRAPE_MICROS_PER_ITEM", 3500)
|
|
session, user = _make_session(_OWNER, balance_micros=1_000_000)
|
|
|
|
charged = await charge_capability(
|
|
_FakePlatformOutput(3), BillingUnit.REDDIT_ITEM, _ctx(session)
|
|
)
|
|
|
|
assert charged == 0
|
|
record_usage.assert_not_awaited()
|
|
session.execute.assert_not_called()
|
|
assert user.credit_micros_balance == 1_000_000
|
|
|
|
|
|
async def test_platform_no_items_is_free(
|
|
monkeypatch, record_usage, _enable_platform_billing
|
|
):
|
|
monkeypatch.setattr(config, "YOUTUBE_MICROS_PER_COMMENT", 3500)
|
|
session, user = _make_session(_OWNER, balance_micros=1_000_000)
|
|
|
|
charged = await charge_capability(
|
|
_FakePlatformOutput(0), BillingUnit.YOUTUBE_COMMENT, _ctx(session)
|
|
)
|
|
|
|
assert charged == 0
|
|
record_usage.assert_not_awaited()
|
|
assert user.credit_micros_balance == 1_000_000
|
|
|
|
|
|
async def test_platform_gate_blocks_when_worst_case_exceeds_balance(
|
|
monkeypatch, _enable_platform_billing
|
|
):
|
|
monkeypatch.setattr(config, "GOOGLE_SEARCH_MICROS_PER_SERP", 5500)
|
|
session = _gate_session(_OWNER, balance_micros=6000) # affords 1 SERP, not 2
|
|
|
|
with pytest.raises(InsufficientCreditsError):
|
|
await gate_capability(
|
|
_FakePlatformInput(estimated_units=2),
|
|
BillingUnit.GOOGLE_SEARCH_SERP,
|
|
_ctx(session),
|
|
)
|
|
|
|
|
|
async def test_platform_gate_maps_reserves_places_plus_reviews(
|
|
monkeypatch, _enable_platform_billing
|
|
):
|
|
monkeypatch.setattr(config, "GOOGLE_MAPS_MICROS_PER_PLACE", 5000)
|
|
monkeypatch.setattr(config, "GOOGLE_MAPS_MICROS_PER_REVIEW", 2000)
|
|
# 1 place (5000) + 10 worst-case reviews (20000) = 25000 required.
|
|
session = _gate_session(_OWNER, balance_micros=20_000)
|
|
|
|
with pytest.raises(InsufficientCreditsError):
|
|
await gate_capability(
|
|
_FakePlatformInput(estimated_units=1, estimated_review_units=10),
|
|
BillingUnit.GOOGLE_MAPS_PLACE,
|
|
_ctx(session),
|
|
)
|
|
|
|
|
|
async def test_platform_gate_passes_when_affordable(
|
|
monkeypatch, _enable_platform_billing
|
|
):
|
|
monkeypatch.setattr(config, "GOOGLE_SEARCH_MICROS_PER_SERP", 5500)
|
|
session = _gate_session(_OWNER, balance_micros=1_000_000)
|
|
|
|
await gate_capability(
|
|
_FakePlatformInput(estimated_units=2),
|
|
BillingUnit.GOOGLE_SEARCH_SERP,
|
|
_ctx(session),
|
|
)
|
|
|
|
|
|
async def test_platform_gate_disabled_is_noop(monkeypatch):
|
|
monkeypatch.setattr(config, "PLATFORM_SCRAPE_BILLING_ENABLED", False)
|
|
session = _gate_session(_OWNER, balance_micros=0)
|
|
|
|
await gate_capability(
|
|
_FakePlatformInput(estimated_units=1000),
|
|
BillingUnit.REDDIT_ITEM,
|
|
_ctx(session),
|
|
)
|
|
|
|
session.execute.assert_not_called()
|