340 lines
14 KiB
Python
340 lines
14 KiB
Python
"""Task-local provider usage, recorded before output validation.
|
|
|
|
Only token counters and bounded request metadata are persisted. This is an
|
|
observation ledger, not a balance, a spending limit, or a provider invoice.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import threading
|
|
import uuid
|
|
from zoneinfo import ZoneInfo
|
|
|
|
from backend.orchestrator.progress_doc import task_doc_dir
|
|
|
|
log = logging.getLogger(__name__)
|
|
_scope: ContextVar[tuple[str | None, str]] = ContextVar("provider_usage_scope", default=(None, "unknown"))
|
|
_locks: dict[str, threading.RLock] = {}
|
|
_locks_guard = threading.Lock()
|
|
_storage_errors: set[str] = set()
|
|
_PHASES = {"intent", "browser", "compression", "analysis", "analysis_incremental", "review", "unknown"}
|
|
_STATUSES = {"succeeded", "schema_error", "truncated", "invalid_response", "http_error", "connection_error", "cancelled"}
|
|
_TOKEN_FIELDS = ("input_tokens", "output_tokens", "total_tokens", "cache_hit_tokens", "cache_miss_tokens")
|
|
_FINISH_REASONS = {"stop", "length", "tool_calls", "content_filter", "end_turn", "max_tokens"}
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(ZoneInfo("Asia/Shanghai")).isoformat(timespec="milliseconds")
|
|
|
|
|
|
def _phase(value: str | None) -> str:
|
|
return value if isinstance(value, str) and value in _PHASES else "unknown"
|
|
|
|
|
|
def _model(value) -> str:
|
|
return value if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9._:/-]{1,100}", value) else "unknown"
|
|
|
|
|
|
def _directory(task_id: str) -> Path:
|
|
if not isinstance(task_id, str) or not re.fullmatch(r"[A-Za-z0-9_-]{1,128}", task_id):
|
|
raise ValueError("任务用量记录标识无效")
|
|
directory = task_doc_dir(task_id)
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
return directory
|
|
|
|
|
|
@contextmanager
|
|
def _locked(directory: Path):
|
|
key = str(directory.resolve())
|
|
with _locks_guard:
|
|
lock = _locks.setdefault(key, threading.RLock())
|
|
with lock:
|
|
with (directory / "model-usage.lock").open("a+b") as stream:
|
|
if os.name == "nt":
|
|
import msvcrt
|
|
stream.seek(0, 2)
|
|
if stream.tell() == 0:
|
|
stream.write(b"\0")
|
|
stream.flush()
|
|
stream.seek(0)
|
|
msvcrt.locking(stream.fileno(), msvcrt.LK_LOCK, 1)
|
|
else:
|
|
import fcntl
|
|
fcntl.flock(stream.fileno(), fcntl.LOCK_EX)
|
|
try:
|
|
yield
|
|
finally:
|
|
if os.name == "nt":
|
|
stream.seek(0)
|
|
msvcrt.locking(stream.fileno(), msvcrt.LK_UNLCK, 1)
|
|
else:
|
|
fcntl.flock(stream.fileno(), fcntl.LOCK_UN)
|
|
|
|
|
|
def _clean_event(raw) -> dict:
|
|
"""Reject valid-JSON corruption as well as malformed JSON.
|
|
|
|
Reconstruct the whitelisted event instead of handing any on-disk metadata
|
|
through to an API response, including a forged timestamp or nested value.
|
|
"""
|
|
if not isinstance(raw, dict) or raw.get("kind") not in {"task_start", "started", "response", "finished"}:
|
|
raise ValueError("invalid usage event")
|
|
at = raw.get("at")
|
|
if not isinstance(at, str) or len(at) > 48:
|
|
raise ValueError("invalid usage event time")
|
|
timestamp = datetime.fromisoformat(at)
|
|
if timestamp.tzinfo is None or timestamp.utcoffset() is None:
|
|
raise ValueError("usage time must have a timezone")
|
|
clean = {"event_id": raw.get("event_id"), "kind": raw["kind"], "at": timestamp.isoformat(timespec="milliseconds")}
|
|
if raw["kind"] == "task_start":
|
|
if raw.get("event_id") != "task-start" or set(raw) != {"event_id", "kind", "at"}:
|
|
raise ValueError("invalid task usage origin")
|
|
return clean
|
|
cid = raw.get("call_id")
|
|
if not isinstance(cid, str) or not re.fullmatch(r"[a-f0-9]{32}", cid) or raw.get("event_id") != cid + ":" + raw["kind"]:
|
|
raise ValueError("invalid usage call identity")
|
|
phase, model = raw.get("phase"), raw.get("model")
|
|
if not isinstance(phase, str) or phase not in _PHASES or _model(model) != model:
|
|
raise ValueError("invalid usage call metadata")
|
|
clean.update(call_id=cid, phase=phase, model=model)
|
|
allowed = {"event_id", "kind", "at", "call_id", "phase", "model"}
|
|
if raw["kind"] == "response":
|
|
counters = raw.get("usage")
|
|
reason = raw.get("finish_reason")
|
|
if not isinstance(counters, dict) or set(counters) != set(_TOKEN_FIELDS):
|
|
raise ValueError("invalid usage counters")
|
|
if any(value is not None and _counter(value) is None for value in counters.values()):
|
|
raise ValueError("invalid usage counter value")
|
|
if reason is not None and (not isinstance(reason, str) or reason not in _FINISH_REASONS):
|
|
raise ValueError("invalid usage finish reason")
|
|
clean.update(usage={key: counters[key] for key in _TOKEN_FIELDS}, finish_reason=reason)
|
|
allowed.update({"usage", "finish_reason"})
|
|
elif raw["kind"] == "finished":
|
|
status, code = raw.get("status"), raw.get("http_status")
|
|
if not isinstance(status, str) or status not in _STATUSES:
|
|
raise ValueError("invalid usage completion status")
|
|
if code is not None and (type(code) is not int or not 100 <= code <= 599):
|
|
raise ValueError("invalid usage HTTP status")
|
|
clean.update(status=status, http_status=code)
|
|
allowed.update({"status", "http_status"})
|
|
if set(raw) != allowed:
|
|
raise ValueError("unexpected usage event fields")
|
|
return clean
|
|
|
|
|
|
def _read_events(directory: Path) -> tuple[list[dict], bool]:
|
|
path = directory / "model-usage.jsonl"
|
|
if not path.exists():
|
|
return [], False
|
|
events, damaged = [], False
|
|
with path.open(encoding="utf-8", errors="replace") as stream:
|
|
for line in stream:
|
|
if not line.strip():
|
|
continue
|
|
try:
|
|
event = _clean_event(json.loads(line))
|
|
events.append(event)
|
|
except (ValueError, TypeError):
|
|
damaged = True
|
|
return events, damaged
|
|
|
|
|
|
def _write_event(directory: Path, event: dict) -> None:
|
|
path = directory / "model-usage.jsonl"
|
|
# A process can die during one append. Preserve that fragment and separate
|
|
# future records, rather than joining two invalid lines.
|
|
with path.open("a+b") as stream:
|
|
stream.seek(0, 2)
|
|
size = stream.tell()
|
|
if size:
|
|
stream.seek(-1, 2)
|
|
if stream.read(1) != b"\n":
|
|
stream.seek(0, 2)
|
|
stream.write(b"\n")
|
|
stream.seek(0, 2)
|
|
stream.write((json.dumps(event, ensure_ascii=False, separators=(",", ":")) + "\n").encode())
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
|
|
|
|
def _append(task_id: str, event: dict) -> None:
|
|
"""Idempotent append under thread and process locks, without task rewrites."""
|
|
try:
|
|
directory = _directory(task_id)
|
|
with _locked(directory):
|
|
events, damaged = _read_events(directory)
|
|
if damaged:
|
|
_storage_errors.add(task_id)
|
|
if any(row["event_id"] == event["event_id"] for row in events):
|
|
return
|
|
_write_event(directory, event)
|
|
except (OSError, ValueError):
|
|
# Metering cannot trigger a paid retry or turn an already received
|
|
# browser action into a fresh request. The UI explicitly marks gaps.
|
|
_storage_errors.add(task_id)
|
|
log.warning("模型用量记录未能完整保存；该任务用量将标为不完整")
|
|
|
|
|
|
@contextmanager
|
|
def usage_scope(task_id: str | None = None, phase: str | None = None):
|
|
"""Bind a task, or override just phase while inheriting the current task."""
|
|
previous_task, previous_phase = _scope.get()
|
|
token = _scope.set((previous_task if task_id is None else task_id,
|
|
previous_phase if phase is None else _phase(phase)))
|
|
try:
|
|
yield
|
|
finally:
|
|
_scope.reset(token)
|
|
|
|
|
|
def current_scope() -> tuple[str | None, str]:
|
|
return _scope.get()
|
|
|
|
|
|
def initialize_task(task_id: str) -> None:
|
|
"""Call only for a newly created task, before its first model request."""
|
|
try:
|
|
directory = _directory(task_id)
|
|
with _locked(directory):
|
|
events, damaged = _read_events(directory)
|
|
if damaged or events:
|
|
# Initializing an old task must not claim its old calls were
|
|
# observed. Existing records and gaps remain authoritative.
|
|
return
|
|
_write_event(directory, {"event_id": "task-start", "kind": "task_start", "at": _now()})
|
|
except (OSError, ValueError):
|
|
_storage_errors.add(task_id)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UsageCall:
|
|
task_id: str
|
|
call_id: str
|
|
phase: str
|
|
model: str
|
|
|
|
|
|
def start_call(model: str, *, phase: str | None = None) -> UsageCall | None:
|
|
task_id, active_phase = _scope.get()
|
|
if task_id is None:
|
|
return None
|
|
call = UsageCall(task_id, uuid.uuid4().hex, _phase(phase or active_phase), _model(model))
|
|
_append(task_id, {"event_id": call.call_id + ":started", "kind": "started", "call_id": call.call_id,
|
|
"phase": call.phase, "model": call.model, "at": _now()})
|
|
return call
|
|
|
|
|
|
def _get(value, key):
|
|
return value.get(key) if isinstance(value, dict) else getattr(value, key, None)
|
|
|
|
|
|
def _counter(value) -> int | None:
|
|
return value if type(value) is int and value >= 0 else None
|
|
|
|
|
|
def normalize_usage(raw) -> dict[str, int | None]:
|
|
"""Keep only actual counters; absent information stays unknown."""
|
|
prompt = _counter(_get(raw, "prompt_tokens"))
|
|
output = _counter(_get(raw, "completion_tokens"))
|
|
if prompt is None:
|
|
prompt = _counter(_get(raw, "input_tokens"))
|
|
if output is None:
|
|
output = _counter(_get(raw, "output_tokens"))
|
|
cached = _counter(_get(raw, "prompt_cache_hit_tokens"))
|
|
if cached is None:
|
|
cached = _counter(_get(raw, "prompt_cached_tokens"))
|
|
if cached is None:
|
|
details = _get(raw, "prompt_tokens_details") or _get(raw, "input_tokens_details")
|
|
cached = _counter(_get(details, "cached_tokens"))
|
|
missed = _counter(_get(raw, "prompt_cache_miss_tokens"))
|
|
if cached is not None and prompt is not None and cached > prompt:
|
|
cached = None
|
|
if missed is not None and prompt is not None and missed > prompt:
|
|
missed = None
|
|
# Cached + uncached is an exact partition of the supplied input count.
|
|
if missed is None and prompt is not None and cached is not None:
|
|
missed = prompt - cached
|
|
return {"input_tokens": prompt, "output_tokens": output,
|
|
"total_tokens": _counter(_get(raw, "total_tokens")),
|
|
"cache_hit_tokens": cached, "cache_miss_tokens": missed}
|
|
|
|
|
|
def record_response(call: UsageCall | None, usage, *, model: str | None = None, finish_reason: str | None = None) -> None:
|
|
if call is None:
|
|
return
|
|
reason = finish_reason if isinstance(finish_reason, str) and finish_reason in _FINISH_REASONS else None
|
|
_append(call.task_id, {"event_id": call.call_id + ":response", "kind": "response", "call_id": call.call_id,
|
|
"phase": call.phase, "model": _model(model or call.model), "at": _now(),
|
|
"finish_reason": reason, "usage": normalize_usage(usage)})
|
|
|
|
|
|
def finish_call(call: UsageCall | None, status: str, *, http_status: int | None = None) -> None:
|
|
if call is None:
|
|
return
|
|
_append(call.task_id, {"event_id": call.call_id + ":finished", "kind": "finished", "call_id": call.call_id,
|
|
"phase": call.phase, "model": call.model, "at": _now(),
|
|
"status": status if status in _STATUSES else "invalid_response",
|
|
"http_status": http_status if type(http_status) is int and 100 <= http_status <= 599 else None})
|
|
|
|
|
|
def _aggregate(rows: list[dict]) -> dict:
|
|
stats = {"requests": len(rows), **{key: None for key in _TOKEN_FIELDS},
|
|
"unknown_usage_requests": 0, "incomplete_usage_requests": 0,
|
|
"parse_failed_requests": 0, "failed_requests": 0, "pending_requests": 0}
|
|
for row in rows:
|
|
usage = row.get("usage") or {}
|
|
for key in _TOKEN_FIELDS:
|
|
value = _counter(usage.get(key))
|
|
if value is not None:
|
|
stats[key] = (stats[key] or 0) + value
|
|
if not any(_counter(usage.get(key)) is not None for key in _TOKEN_FIELDS):
|
|
stats["unknown_usage_requests"] += 1
|
|
if any(_counter(usage.get(key)) is None for key in ("input_tokens", "output_tokens", "total_tokens")):
|
|
stats["incomplete_usage_requests"] += 1
|
|
status = row.get("status") if isinstance(row.get("status"), str) else None
|
|
if status in {"schema_error", "invalid_response", "truncated"}:
|
|
stats["parse_failed_requests"] += 1
|
|
if status and status != "succeeded":
|
|
stats["failed_requests"] += 1
|
|
if not status:
|
|
stats["pending_requests"] += 1
|
|
return stats
|
|
|
|
|
|
def snapshot(task_id: str) -> dict:
|
|
"""Sum known counters without presenting missing history as zero cost."""
|
|
events, damaged = [], False
|
|
try:
|
|
directory = _directory(task_id)
|
|
with _locked(directory):
|
|
events, damaged = _read_events(directory)
|
|
except (OSError, ValueError):
|
|
damaged = True
|
|
rows: dict[str, dict] = {}
|
|
for event in events:
|
|
cid = event.get("call_id")
|
|
if not isinstance(cid, str) or not re.fullmatch(r"[a-f0-9]{32}", cid):
|
|
continue
|
|
row = rows.setdefault(cid, {"phase": _phase(event.get("phase")), "model": _model(event.get("model"))})
|
|
if event.get("kind") == "response":
|
|
row["model"] = _model(event.get("model"))
|
|
row["usage"] = event.get("usage") if isinstance(event.get("usage"), dict) else {}
|
|
elif event.get("kind") == "finished":
|
|
row["status"] = event.get("status")
|
|
values = list(rows.values())
|
|
phases = sorted({row["phase"] for row in values})
|
|
return {"version": 1, **_aggregate(values), "coverage": "recorded_calls_only",
|
|
"legacy_incomplete": not bool(events and events[0].get("kind") == "task_start"),
|
|
"ledger_error": damaged or task_id in _storage_errors,
|
|
"recorded_since": next((row.get("at") for row in events if row.get("kind") == "started"), None),
|
|
"models": sorted({row["model"] for row in values}),
|
|
"phases": [{"phase": phase, **_aggregate([row for row in values if row["phase"] == phase])} for phase in phases]}
|