296 lines
15 KiB
Python
296 lines
15 KiB
Python
"""Actual provider counters survive invalid actions and concurrent task writes."""
|
|
from __future__ import annotations
|
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import json
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from pydantic import BaseModel
|
|
|
|
from backend.browser import llm_factory
|
|
from backend.prospecting import analysis, dialogue, usage
|
|
from backend.prospecting.contracts import ProspectingBrief
|
|
|
|
|
|
@pytest.fixture
|
|
def ledger(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(usage, "task_doc_dir", lambda tid: tmp_path / tid)
|
|
usage._storage_errors.clear()
|
|
return tmp_path
|
|
|
|
|
|
def one_call(task_id, *, phase="browser", status="succeeded", raw=None):
|
|
with usage.usage_scope(task_id, phase):
|
|
call = usage.start_call("deepseek-flash")
|
|
usage.record_response(call, raw)
|
|
usage.finish_call(call, status)
|
|
return call
|
|
|
|
|
|
def test_new_task_unknown_and_zero_usage_are_distinct(ledger):
|
|
usage.initialize_task("new-task")
|
|
assert usage.snapshot("new-task")["input_tokens"] is None
|
|
one_call("new-task", raw={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
|
|
one_call("new-task")
|
|
state = usage.snapshot("new-task")
|
|
assert state["requests"] == 2 and state["input_tokens"] == state["output_tokens"] == 0
|
|
assert state["unknown_usage_requests"] == state["incomplete_usage_requests"] == 1
|
|
assert state["legacy_incomplete"] is False
|
|
assert state["cache_hit_tokens"] is None
|
|
|
|
|
|
def test_old_task_cannot_be_relabelled_as_fully_instrumented(ledger):
|
|
one_call("old-task", raw={"prompt_tokens": 8, "completion_tokens": 3, "total_tokens": 11})
|
|
usage.initialize_task("old-task")
|
|
state = usage.snapshot("old-task")
|
|
assert state["legacy_incomplete"] is True and state["requests"] == 1
|
|
assert state["coverage"] == "recorded_calls_only"
|
|
assert usage.snapshot("no-records")["total_tokens"] is None
|
|
|
|
|
|
def test_response_is_idempotent_and_private_fields_never_persist(ledger):
|
|
usage.initialize_task("task")
|
|
with usage.usage_scope("task", "browser"):
|
|
call = usage.start_call("deepseek-flash")
|
|
raw = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13,
|
|
"prompt_cache_hit_tokens": 8, "prompt_cache_miss_tokens": 2,
|
|
"api_key": "NEVER-PERSIST", "prompt": "PRIVATE-PROMPT", "screenshot": "PRIVATE-IMAGE"}
|
|
usage.record_response(call, raw, model="deepseek-flash-actual")
|
|
usage.record_response(call, {"prompt_tokens": 999})
|
|
usage.finish_call(call, "succeeded")
|
|
usage.finish_call(call, "schema_error")
|
|
state = usage.snapshot("task")
|
|
assert state["requests"] == 1 and state["input_tokens"] == 10 and state["output_tokens"] == 3
|
|
assert state["cache_hit_tokens"] == 8 and state["cache_miss_tokens"] == 2
|
|
assert state["failed_requests"] == 0 and state["models"] == ["deepseek-flash-actual"]
|
|
text = (ledger / "task" / "model-usage.jsonl").read_text()
|
|
assert all(value not in text for value in ("NEVER-PERSIST", "PRIVATE-PROMPT", "PRIVATE-IMAGE"))
|
|
|
|
|
|
def test_nested_phase_inherits_task_and_scope_is_restored(ledger):
|
|
with usage.usage_scope("task", "browser"):
|
|
first = usage.start_call("deepseek-flash")
|
|
with usage.usage_scope(phase="compression"):
|
|
compressed = usage.start_call("deepseek-flash")
|
|
usage.record_response(compressed, {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12})
|
|
usage.finish_call(compressed, "succeeded")
|
|
assert usage.current_scope() == ("task", "browser")
|
|
usage.finish_call(first, "connection_error")
|
|
assert usage.current_scope() == (None, "unknown") and usage.start_call("deepseek-flash") is None
|
|
state = usage.snapshot("task")
|
|
assert state["requests"] == 2 and {row["phase"] for row in state["phases"]} == {"browser", "compression"}
|
|
assert state["unknown_usage_requests"] == 1
|
|
|
|
|
|
def test_deepseek_cache_top_level_precedes_nested_and_partial_stays_unknown():
|
|
state = usage.normalize_usage(SimpleNamespace(prompt_tokens=10, completion_tokens=3, total_tokens=13,
|
|
prompt_cache_hit_tokens=8, prompt_cache_miss_tokens=2,
|
|
prompt_tokens_details=SimpleNamespace(cached_tokens=0)))
|
|
assert state["cache_hit_tokens"] == 8 and state["cache_miss_tokens"] == 2
|
|
state = usage.normalize_usage({"prompt_tokens": 10, "prompt_tokens_details": {"cached_tokens": 7}})
|
|
assert state["cache_hit_tokens"] == 7 and state["cache_miss_tokens"] == 3
|
|
assert state["output_tokens"] is state["total_tokens"] is None
|
|
assert usage.normalize_usage({"prompt_tokens": True, "completion_tokens": -1})["input_tokens"] is None
|
|
|
|
|
|
def test_threaded_appends_do_not_lose_calls(ledger):
|
|
usage.initialize_task("shared")
|
|
with ThreadPoolExecutor(max_workers=8) as pool:
|
|
list(pool.map(lambda i: one_call("shared", raw={"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}), range(32)))
|
|
state = usage.snapshot("shared")
|
|
assert state["requests"] == 32 and state["total_tokens"] == 192
|
|
assert state["pending_requests"] == 0 and state["ledger_error"] is False
|
|
|
|
|
|
def test_process_appends_are_locked(ledger):
|
|
usage.initialize_task("shared-process")
|
|
script = """from pathlib import Path
|
|
import sys
|
|
from backend.prospecting import usage
|
|
usage.task_doc_dir=lambda tid:Path(sys.argv[1])/tid
|
|
for _ in range(4):
|
|
with usage.usage_scope('shared-process','analysis'):
|
|
c=usage.start_call('deepseek-flash')
|
|
usage.record_response(c,{'prompt_tokens':2,'completion_tokens':1,'total_tokens':3})
|
|
usage.finish_call(c,'succeeded')
|
|
"""
|
|
processes = [subprocess.Popen([sys.executable, "-c", script, str(ledger)], stdout=subprocess.PIPE, stderr=subprocess.PIPE) for _ in range(3)]
|
|
for process in processes:
|
|
stdout, stderr = process.communicate(timeout=15)
|
|
assert process.returncode == 0, (stdout.decode(), stderr.decode())
|
|
state = usage.snapshot("shared-process")
|
|
assert state["requests"] == 12 and state["total_tokens"] == 36 and not state["ledger_error"]
|
|
|
|
|
|
def test_truncated_ledger_preserves_known_data_and_marks_gap(ledger):
|
|
usage.initialize_task("damaged")
|
|
path = ledger / "damaged" / "model-usage.jsonl"
|
|
with path.open("ab") as stream:
|
|
stream.write(b'{"interrupted":')
|
|
one_call("damaged", raw={"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3})
|
|
state = usage.snapshot("damaged")
|
|
assert state["requests"] == 1 and state["total_tokens"] == 3 and state["ledger_error"]
|
|
|
|
|
|
def test_storage_failure_cannot_turn_a_model_response_into_paid_retry(ledger):
|
|
(ledger / "bad").write_text("not a directory")
|
|
one_call("bad", raw={"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3})
|
|
state = usage.snapshot("bad")
|
|
assert state["ledger_error"] and state["total_tokens"] is None
|
|
|
|
|
|
@pytest.mark.parametrize("field,bad_value", [
|
|
("at", {"secret": "PRIVATE-NESTED-METADATA"}),
|
|
("at", "PRIVATE-ARBITRARY-TEXT"),
|
|
("at", "2026-10-29T10:00:00"),
|
|
("kind", {"secret": "PRIVATE-KIND"}),
|
|
("kind", "not-an-event"),
|
|
("call_id", "wrong-call"),
|
|
("event_id", "f" * 32 + ":response"),
|
|
("phase", {"secret": "PRIVATE-PHASE"}),
|
|
("model", {"secret": "PRIVATE-MODEL"}),
|
|
("prompt", "PRIVATE-UNEXPECTED-CONTENT"),
|
|
("finish_reason", "PRIVATE-UNKNOWN-FINISH"),
|
|
("usage", {"input_tokens": "PRIVATE-NON-NUMERIC"}),
|
|
])
|
|
def test_valid_json_bad_metadata_is_skipped_without_leak(ledger, field, bad_value):
|
|
usage.initialize_task("bad-metadata")
|
|
one_call("bad-metadata", raw={"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3})
|
|
path = ledger / "bad-metadata" / "model-usage.jsonl"
|
|
rows = [json.loads(line) for line in path.read_text().splitlines()]
|
|
bad = next(row for row in rows if row.get("kind") == "response")
|
|
bad[field] = bad_value
|
|
with path.open("a") as stream:
|
|
stream.write(json.dumps(bad) + "\n")
|
|
state = usage.snapshot("bad-metadata")
|
|
assert state["requests"] == 1 and state["total_tokens"] == 3
|
|
assert state["ledger_error"] and "PRIVATE" not in json.dumps(state)
|
|
|
|
|
|
def test_incremental_analysis_has_its_own_phase(ledger):
|
|
one_call("incremental", phase="analysis_incremental", raw={"prompt_tokens": 6, "completion_tokens": 1, "total_tokens": 7})
|
|
state = usage.snapshot("incremental")
|
|
assert state["phases"][0]["phase"] == "analysis_incremental"
|
|
|
|
|
|
class BrowserAction(BaseModel):
|
|
action: str
|
|
|
|
|
|
async def test_browser_schema_failure_still_records_provider_tokens(ledger, monkeypatch):
|
|
import openai
|
|
class Client:
|
|
def __init__(self, **kwargs):
|
|
self.chat = SimpleNamespace(completions=SimpleNamespace(create=self.create))
|
|
async def __aenter__(self):
|
|
return self
|
|
async def __aexit__(self, *args):
|
|
return False
|
|
async def create(self, **kwargs):
|
|
return SimpleNamespace(model="deepseek-flash", choices=[SimpleNamespace(message=SimpleNamespace(content=None, tool_calls=[
|
|
SimpleNamespace(type="function", function=SimpleNamespace(name="return_structured_output", arguments='{"unexpected":1}'))
|
|
]), finish_reason="tool_calls")],
|
|
usage=SimpleNamespace(prompt_tokens=20, completion_tokens=4, total_tokens=24,
|
|
prompt_cache_hit_tokens=17, prompt_cache_miss_tokens=3))
|
|
monkeypatch.setattr(openai, "AsyncOpenAI", Client)
|
|
from browser_use.llm.exceptions import ModelProviderError
|
|
from browser_use.llm.messages import UserMessage
|
|
with usage.usage_scope("invalid-action", "browser"):
|
|
with pytest.raises(ModelProviderError, match="不符合"):
|
|
await llm_factory.DeepSeekBrowserChat("deepseek-flash", "PRIVATE-TEST-KEY").ainvoke([UserMessage(content="PRIVATE-PROMPT")], BrowserAction)
|
|
state = usage.snapshot("invalid-action")
|
|
assert state["total_tokens"] == 24 and state["cache_hit_tokens"] == 17
|
|
assert state["parse_failed_requests"] == state["failed_requests"] == 1
|
|
assert "PRIVATE" not in (ledger / "invalid-action" / "model-usage.jsonl").read_text()
|
|
|
|
|
|
async def test_invalid_analysis_keeps_usage_and_parent_context(ledger, monkeypatch):
|
|
import httpx
|
|
calls = []
|
|
body = {"model": "deepseek-flash", "choices": [{"message": {"content": '{"results":[],"summary":"none"}'}}],
|
|
"usage": {"prompt_tokens": 14, "completion_tokens": 2, "total_tokens": 16}}
|
|
class Client:
|
|
def __init__(self, **kwargs):
|
|
pass
|
|
async def __aenter__(self):
|
|
return self
|
|
async def __aexit__(self, *args):
|
|
return False
|
|
async def post(self, url, **kwargs):
|
|
calls.append(kwargs)
|
|
return httpx.Response(200, json=body)
|
|
monkeypatch.setattr(analysis.httpx, "AsyncClient", Client)
|
|
monkeypatch.setattr(analysis, "get_llm_config", lambda: {"api_key": "PRIVATE-KEY", "model": "deepseek-flash", "base_url": "https://fixture", "provider": "deepseek"})
|
|
note = {"source_id": "note", "source_type": "note", "source_url": "https://fixture/n", "content": "parent context"}
|
|
reply = {"source_id": "reply", "source_type": "reply", "source_url": "https://fixture/n", "parent_source_id": "note", "content": "own words"}
|
|
with usage.usage_scope("bad-analysis", "review"):
|
|
with pytest.raises(ValueError, match="来源"):
|
|
await analysis.analyze_sources(ProspectingBrief(intent="research"), [reply], context_sources=[note, reply])
|
|
assert '"source_id": "note"' in calls[0]["json"]["messages"][0]["content"]
|
|
state = usage.snapshot("bad-analysis")
|
|
assert state["requests"] == 1 and state["total_tokens"] == 16 and state["parse_failed_requests"] == 1
|
|
assert state["phases"][0]["phase"] == "review"
|
|
|
|
|
|
async def test_intent_validation_failure_is_metered(ledger, monkeypatch):
|
|
import httpx
|
|
class Client:
|
|
def __init__(self, **kwargs):
|
|
pass
|
|
async def __aenter__(self):
|
|
return self
|
|
async def __aexit__(self, *args):
|
|
return False
|
|
async def post(self, url, **kwargs):
|
|
return httpx.Response(200, json={"choices": [{"message": {"content": "not JSON"}}],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}})
|
|
monkeypatch.setattr(dialogue.httpx, "AsyncClient", Client)
|
|
monkeypatch.setattr(dialogue, "get_llm_config", lambda: {"api_key": "PRIVATE-KEY", "model": "deepseek-flash", "base_url": "https://fixture", "provider": "deepseek"})
|
|
with usage.usage_scope("bad-intent", "intent"):
|
|
with pytest.raises(ValueError, match="可用计划"):
|
|
await dialogue.interpret_message("找客户", ProspectingBrief(), [])
|
|
state = usage.snapshot("bad-intent")
|
|
assert state["requests"] == 1 and state["total_tokens"] == 11 and state["parse_failed_requests"] == 1
|
|
|
|
|
|
async def test_legacy_wrapper_preserves_request_and_only_counts_once(ledger):
|
|
class Wrapped:
|
|
provider, name = "fixture", "fixture-model"
|
|
def __init__(self):
|
|
self.calls = []
|
|
def __repr__(self):
|
|
return "PRIVATE-KEY"
|
|
async def ainvoke(self, messages, output_format=None, **kwargs):
|
|
self.calls.append((messages, output_format, kwargs))
|
|
return SimpleNamespace(completion="same response", usage=SimpleNamespace(prompt_tokens=9, completion_tokens=2,
|
|
total_tokens=11, prompt_cached_tokens=7))
|
|
wrapped = Wrapped()
|
|
model = llm_factory.ObservedChatModel(wrapped, "fixture-model")
|
|
with usage.usage_scope("legacy", "compression"):
|
|
result = await model.ainvoke(["unchanged"], BrowserAction, temperature=0)
|
|
assert result.completion == "same response" and wrapped.calls == [(["unchanged"], BrowserAction, {"temperature": 0})]
|
|
state = usage.snapshot("legacy")
|
|
assert state["requests"] == 1 and state["total_tokens"] == 11 and state["cache_hit_tokens"] == 7
|
|
assert "PRIVATE-KEY" not in repr(model)
|
|
|
|
|
|
async def test_legacy_failed_adapter_records_unknown_without_retry(ledger):
|
|
class Wrapped:
|
|
provider, name = "fixture", "fixture-model"
|
|
calls = 0
|
|
async def ainvoke(self, *args, **kwargs):
|
|
self.calls += 1
|
|
raise ValueError("PRIVATE response content must never persist")
|
|
wrapped = Wrapped()
|
|
with usage.usage_scope("legacy-failed", "browser"):
|
|
with pytest.raises(ValueError):
|
|
await llm_factory.ObservedChatModel(wrapped, "fixture-model").ainvoke([])
|
|
state = usage.snapshot("legacy-failed")
|
|
assert wrapped.calls == state["requests"] == state["unknown_usage_requests"] == 1
|
|
assert state["input_tokens"] is None and state["output_tokens"] is None
|
|
assert "PRIVATE" not in (ledger / "legacy-failed" / "model-usage.jsonl").read_text()
|