feat(R03): unify provider catalogue, routing decisions and usage telemetry
EPIC R03 (Team Duy) - one provider catalogue, one routing flow, one usage seam.
R03-T01 tests/contracts/test_providers.py
29 contract tests every provider must satisfy: canonical assistant message,
streamed text == returned content, reasoning never joins the answer, parsed
tool arguments, ProviderError for every failure. Real adapters exercised
offline by stubbing Provider._request.
R03-T02 domain/models/provider_descriptor.py
infrastructure/providers/provider_registry.py
Provider facts declared once (was split across providers/factory.py,
DEFAULT_CONFIG and PROVIDER_LABELS). ProviderRegistry.build() also stamps the
descriptor id onto the instance, so ollama/github_copilot/codex usage is no
longer all attributed to "openai_compat", and never mutates the caller config.
R03-T03 application/model_routing/routing_application_service.py
Pure-Python routing policy with four modes: Off, Auto, Manual and the new
Fallback (switch only AFTER the current model fails). Depends on a RoutingPort
protocol; production wires the existing core.routing engine underneath.
R03-T04/T05 ui/chat_panel.py, ui/co4e_tab.py, ui/folder_tab.py
Three near-identical routing copies (~40 lines each) replaced by a call to
ctx.routing_application() plus a confirm callback. Mode vocabulary now lives
in one place (normalize_mode/is_valid_mode) instead of four literal tuples.
R03-T06 infrastructure/telemetry/usage_sink.py
Token usage extracted from both providers into UsageEvent + UsageEventSink.
Estimation pinned against core.usage_tracker so no recorded number changes.
Also fixes a deadlock introduced while wiring AppContext: routing_application()
held _routing_lock and called routing(), which takes the same non-reentrant lock.
Suite: 186 passed, 1.22s. check_imports: PASS. All new files < 400 LOC.
2 pre-existing failures remain in test_config_security.py (EPIC R02/Team Nam).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,346 @@
|
||||
"""Unit tests for :mod:`application.model_routing` (R03-T03).
|
||||
|
||||
These run against a hand-written fake router rather than ``core.routing``: the
|
||||
point of the service is the DECISION policy around the engine (mode handling,
|
||||
the manual confirm, never-raise behaviour, failure fallback), and mixing in the
|
||||
real scorer would test the wrong thing and drag the suite over its time budget.
|
||||
|
||||
No Qt, no config, no network - the whole file runs in milliseconds, which is the
|
||||
concrete payoff of moving this logic out of ``ui/chat_panel.py``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List, Optional, Tuple
|
||||
|
||||
import pytest
|
||||
|
||||
from cowork_local.application.model_routing import (
|
||||
RoutingApplicationService,
|
||||
RoutingDecision,
|
||||
RoutingMode,
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test doubles shaped like core.routing's RouteResult / SwitchDecision
|
||||
# --------------------------------------------------------------------------- #
|
||||
@dataclass
|
||||
class _TaskType:
|
||||
value: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Decision:
|
||||
score_gain: float = 0.0
|
||||
reason: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RouteResult:
|
||||
should_switch: bool
|
||||
to: Optional[Tuple[str, str]] = None
|
||||
task_type: Any = None
|
||||
decision: Any = None
|
||||
|
||||
def target(self) -> Optional[Tuple[str, str]]:
|
||||
return self.to
|
||||
|
||||
|
||||
class _FakeRouter:
|
||||
"""Records every route() call and replays a canned result."""
|
||||
|
||||
def __init__(self, result: Any = None, raises: bool = False) -> None:
|
||||
self._result = result or _RouteResult(should_switch=False, decision=_Decision())
|
||||
self._raises = raises
|
||||
self.calls: List[dict] = []
|
||||
|
||||
def route(self, surface, prompt, current_provider, current_model, **kwargs):
|
||||
self.calls.append({"surface": surface, "prompt": prompt,
|
||||
"provider": current_provider, "model": current_model, **kwargs})
|
||||
if self._raises:
|
||||
raise RuntimeError("assessment store is corrupt")
|
||||
return self._result
|
||||
|
||||
|
||||
def _switch_to(provider: str, model: str, gain: float = 0.2, task: str = "coding") -> _RouteResult:
|
||||
return _RouteResult(
|
||||
should_switch=True, to=(provider, model), task_type=_TaskType(task),
|
||||
decision=_Decision(score_gain=gain, reason=f"{task} fit beats current by {gain}"),
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Mode parsing
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.parametrize("raw,expected", [
|
||||
("off", RoutingMode.OFF),
|
||||
("AUTO", RoutingMode.AUTO),
|
||||
(" manual ", RoutingMode.MANUAL),
|
||||
("fallback", RoutingMode.FALLBACK),
|
||||
])
|
||||
def test_parse_accepts_the_config_spellings(raw, expected):
|
||||
assert RoutingMode.parse(raw) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["", None, "nonsense", 0])
|
||||
def test_parse_degrades_unknown_values_to_off(raw):
|
||||
"""A corrupt setting must leave the user's own model alone rather than
|
||||
silently moving their work onto another model."""
|
||||
assert RoutingMode.parse(raw) is RoutingMode.OFF
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# OFF
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_off_never_consults_the_engine():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "hi", "openai_compat", "gpt-4o-mini",
|
||||
mode="off")
|
||||
|
||||
assert router.calls == [] # not even scored: OFF costs nothing
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
def test_blank_prompt_is_never_routed():
|
||||
"""An empty message carries no signal to classify; all three legacy copies
|
||||
guarded this and the guard has to survive the move."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", " ", "openai_compat", "m", mode="auto")
|
||||
|
||||
assert router.calls == []
|
||||
assert decision.switched is False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# AUTO
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_auto_switches_silently_and_reports_the_target():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude-sonnet-4-6", gain=0.31))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "write a function", "openai_compat", "gpt-4o-mini",
|
||||
mode="auto")
|
||||
|
||||
assert decision.switched is True
|
||||
assert decision.target() == ("anthropic", "claude-sonnet-4-6")
|
||||
assert decision.task_type == "coding"
|
||||
assert decision.score_gain == pytest.approx(0.31)
|
||||
assert decision.should_notify is True
|
||||
|
||||
|
||||
def test_auto_keeps_the_current_model_when_no_candidate_wins():
|
||||
router = _FakeRouter(_RouteResult(should_switch=False, decision=_Decision(reason="no gain")))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "hello", "openai_compat", "gpt-4o-mini", mode="auto")
|
||||
|
||||
assert decision.switched is False
|
||||
# The decision still names a model to run on, so the call site never has to
|
||||
# re-derive the fallback itself - the exact drift the three copies suffered.
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
assert decision.should_notify is False
|
||||
|
||||
|
||||
def test_auto_never_asks_for_confirmation():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
asked: List[RoutingDecision] = []
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
service.route_turn("cowork", "q", "openai_compat", "m", mode="auto",
|
||||
confirm=lambda d: asked.append(d) or True)
|
||||
|
||||
assert asked == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# MANUAL
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_manual_switches_only_after_the_user_approves():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
seen: List[RoutingDecision] = []
|
||||
|
||||
def confirm(proposal: RoutingDecision) -> bool:
|
||||
seen.append(proposal)
|
||||
return True
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "m",
|
||||
mode="manual", confirm=confirm)
|
||||
|
||||
assert decision.switched is True
|
||||
assert decision.target() == ("anthropic", "claude")
|
||||
# The dialog is handed the full proposal so it can explain the trade-off.
|
||||
assert seen[0].model == "claude"
|
||||
assert seen[0].score_gain > 0
|
||||
|
||||
|
||||
def test_manual_keeps_the_current_model_when_declined():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "gpt-4o-mini",
|
||||
mode="manual", confirm=lambda d: False)
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.declined is True
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
def test_manual_without_a_confirm_callback_does_not_switch():
|
||||
"""A headless caller (scheduler) has nobody to ask, so Manual must behave as
|
||||
"not approved" rather than as "approved by default"."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "m", mode="manual")
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.declined is True
|
||||
|
||||
|
||||
def test_a_confirm_dialog_that_raises_counts_as_declined():
|
||||
"""If the modal blows up (window closing mid-turn) the safe reading is that
|
||||
the user did NOT consent to running on another model."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
def confirm(_proposal):
|
||||
raise RuntimeError("dialog destroyed")
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "m",
|
||||
mode="manual", confirm=confirm)
|
||||
|
||||
assert decision.switched is False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# FALLBACK
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_fallback_does_not_switch_up_front():
|
||||
"""The whole point of the mode: honour the user's model choice until it
|
||||
actually fails."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "gpt-4o-mini",
|
||||
mode="fallback")
|
||||
|
||||
assert router.calls == []
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
def test_fallback_switches_after_a_failure():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude", gain=0.4))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.fallback_after_failure("cowork", "q", "openai_compat", "gpt-4o-mini",
|
||||
mode="fallback")
|
||||
|
||||
assert decision is not None
|
||||
assert decision.switched is True
|
||||
assert decision.target() == ("anthropic", "claude")
|
||||
assert "failed" in decision.reason
|
||||
|
||||
|
||||
def test_fallback_never_returns_the_model_that_just_failed():
|
||||
"""Retrying the model that just went down would spin on the outage."""
|
||||
router = _FakeRouter(_switch_to("openai_compat", "gpt-4o-mini"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
assert service.fallback_after_failure(
|
||||
"cowork", "q", "openai_compat", "gpt-4o-mini", mode="fallback") is None
|
||||
|
||||
|
||||
def test_fallback_returns_none_when_there_is_no_alternative():
|
||||
router = _FakeRouter(_RouteResult(should_switch=False, decision=_Decision()))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
assert service.fallback_after_failure("cowork", "q", "openai_compat", "m",
|
||||
mode="auto") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["off", "manual"])
|
||||
def test_off_and_manual_do_not_auto_recover_from_a_failure(mode):
|
||||
"""Both modes exist to keep the user in control of which model runs their
|
||||
work; moving it on failure would break that promise silently."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
assert service.fallback_after_failure("cowork", "q", "openai_compat", "m",
|
||||
mode=mode) is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Robustness - routing must never break a chat turn
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_engine_failure_degrades_to_keeping_the_current_model():
|
||||
service = RoutingApplicationService(_FakeRouter(raises=True))
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "gpt-4o-mini", mode="auto")
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
def test_engine_failure_during_fallback_returns_none():
|
||||
"""A broken router must not mask the original provider error with its own."""
|
||||
service = RoutingApplicationService(_FakeRouter(raises=True))
|
||||
|
||||
assert service.fallback_after_failure("cowork", "q", "p", "m", mode="auto") is None
|
||||
|
||||
|
||||
def test_a_malformed_route_result_is_treated_as_no_switch():
|
||||
"""The engine is a legacy module still under refactor; a missing attribute
|
||||
must degrade, not raise into the middle of a turn."""
|
||||
class _Garbage:
|
||||
should_switch = True # claims a switch but exposes no target()
|
||||
|
||||
service = RoutingApplicationService(_FakeRouter(_Garbage()))
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "m", mode="auto")
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "m")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Per-surface mode lookup
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_mode_is_read_per_surface_when_not_passed_explicitly():
|
||||
"""Each screen has its own Off/Auto/Manual toggle, and workspaces override
|
||||
it - so the surface, not a global setting, decides."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
modes = {"cowork": "auto", "ai_edit": "off"}
|
||||
service = RoutingApplicationService(router, mode_reader=modes.get)
|
||||
|
||||
assert service.route_turn("cowork", "q", "p", "m").switched is True
|
||||
assert service.route_turn("ai_edit", "q", "p", "m").switched is False
|
||||
|
||||
|
||||
def test_a_failing_mode_reader_falls_back_to_off():
|
||||
def broken(_surface):
|
||||
raise KeyError("config not loaded yet")
|
||||
|
||||
service = RoutingApplicationService(_FakeRouter(_switch_to("a", "b")),
|
||||
mode_reader=broken)
|
||||
|
||||
assert service.route_turn("cowork", "q", "p", "m").switched is False
|
||||
|
||||
|
||||
def test_required_capabilities_are_passed_through_to_the_engine():
|
||||
"""An image turn must only be routed to a vision-capable model; the filter
|
||||
has to reach the scorer or the constraint is silently dropped."""
|
||||
router = _FakeRouter()
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
service.route_turn("cowork", "describe this", "p", "m", mode="auto",
|
||||
required_capabilities=["vision"])
|
||||
|
||||
assert router.calls[0]["required_capabilities"] == ["vision"]
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Integration test for the AppContext routing wiring (R03-T04 / R03-T05).
|
||||
|
||||
The three chat surfaces now call ``ctx.routing_application()`` instead of each
|
||||
carrying their own copy of the routing algorithm. The unit tests cover the
|
||||
policy; this file covers the WIRING, which unit tests with a fake router cannot
|
||||
see:
|
||||
|
||||
* the service is built and memoised on the context
|
||||
* it reads the per-workspace mode through ``project_routing_mode``
|
||||
* the legacy ``core.routing.RoutingService`` is what sits underneath it
|
||||
* ``fallback`` survives a round trip through the per-workspace mode store
|
||||
|
||||
Still Qt-free: ``AppContext`` itself imports no widgets, and the config is
|
||||
written into a tmp dir so nothing touches ``~/.cowork_local``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from cowork_local.application.model_routing import (
|
||||
RoutingApplicationService,
|
||||
RoutingMode,
|
||||
)
|
||||
from cowork_local.config import AppConfig
|
||||
from cowork_local.state import AppContext
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ctx(tmp_path: Path) -> AppContext:
|
||||
"""An AppContext backed by a throwaway config file."""
|
||||
return AppContext(AppConfig.load(tmp_path / "config.json"))
|
||||
|
||||
|
||||
def test_routing_application_is_built_and_memoised(ctx):
|
||||
"""One instance per app: the pending-switch registry underneath it must be
|
||||
shared by every surface, so a second call has to return the same object."""
|
||||
first = ctx.routing_application()
|
||||
|
||||
assert isinstance(first, RoutingApplicationService)
|
||||
assert ctx.routing_application() is first
|
||||
|
||||
|
||||
def test_the_legacy_engine_sits_underneath_the_new_service():
|
||||
"""Strangler-fig check (ADR-001 section 4): the scoring engine is reused, not
|
||||
reimplemented. If this ever stops holding, the assessment scores the
|
||||
scheduler probes would no longer be the ones routing decisions use."""
|
||||
from cowork_local.core.routing.service import RoutingService
|
||||
|
||||
config = AppConfig.load(Path("does-not-exist.json"))
|
||||
context = AppContext(config)
|
||||
|
||||
service = context.routing_application()
|
||||
|
||||
assert isinstance(service._router, RoutingService)
|
||||
assert service._router is context.routing()
|
||||
|
||||
|
||||
def test_mode_is_read_through_the_per_workspace_lookup(ctx, monkeypatch):
|
||||
seen = []
|
||||
|
||||
def fake_mode(surface: str) -> str:
|
||||
seen.append(surface)
|
||||
return "off"
|
||||
|
||||
monkeypatch.setattr(ctx, "project_routing_mode", fake_mode)
|
||||
# Built after the patch so the service captures the patched reader.
|
||||
service = RoutingApplicationService(ctx.routing(), mode_reader=ctx.project_routing_mode)
|
||||
|
||||
decision = service.route_turn("co4e", "hello", "openai_compat", "gpt-4o-mini")
|
||||
|
||||
assert seen == ["co4e"]
|
||||
assert decision.switched is False
|
||||
|
||||
|
||||
def test_routing_off_by_default_leaves_the_selected_model_alone(ctx):
|
||||
"""Default config has routing off on every surface, so a fresh install must
|
||||
never move a turn to another model."""
|
||||
decision = ctx.routing_application().route_turn(
|
||||
"cowork", "write a function", "openai_compat", "gpt-4o-mini")
|
||||
|
||||
assert decision.mode is RoutingMode.OFF
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["off", "auto", "manual", "fallback"])
|
||||
def test_every_mode_survives_a_round_trip_through_the_config(ctx, mode):
|
||||
"""``fallback`` is new (R03-T03); the per-surface store used to whitelist
|
||||
only three values and would have silently downgraded it to "off"."""
|
||||
ctx.set_project_routing_mode("cowork", mode)
|
||||
|
||||
assert ctx.project_routing_mode("cowork") == mode
|
||||
|
||||
|
||||
def test_an_unknown_mode_still_falls_back_to_off(ctx):
|
||||
ctx.set_project_routing_mode("cowork", "turbo")
|
||||
|
||||
assert ctx.project_routing_mode("cowork") == "off"
|
||||
|
||||
|
||||
def test_a_real_route_call_never_raises_without_any_assessments(ctx):
|
||||
"""The store is empty on a fresh install. Routing must degrade to "keep the
|
||||
current model" rather than raise into the middle of the first message."""
|
||||
ctx.set_project_routing_mode("cowork", "auto")
|
||||
|
||||
decision = ctx.routing_application().route_turn(
|
||||
"cowork", "hello there", "openai_compat", "gpt-4o-mini")
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
@@ -0,0 +1,249 @@
|
||||
"""Unit tests for :mod:`infrastructure.telemetry.usage_sink` (R03-T06).
|
||||
|
||||
Two things are being protected here:
|
||||
|
||||
1. The **numbers do not change**. Extracting usage recording out of the two
|
||||
providers is only safe if the events built from each wire format carry
|
||||
exactly what ``core.usage_tracker.record`` used to receive - a silent change
|
||||
would corrupt the Dashboard's cost history.
|
||||
2. The **sink can never break a turn**. Telemetry is observability; a broken
|
||||
store must be swallowed (and logged), never raised into a chat turn.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import pytest
|
||||
|
||||
from cowork_local.infrastructure.telemetry import usage_sink as telemetry
|
||||
from cowork_local.providers.anthropic import AnthropicProvider
|
||||
from cowork_local.providers.base import Provider
|
||||
from cowork_local.providers.openai_compat import OpenAICompatProvider
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Event construction - one per wire format
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_openai_usage_block_maps_onto_the_canonical_event():
|
||||
event = telemetry.openai_usage_event("openai_compat", "gpt-4o-mini", {
|
||||
"prompt_tokens": 120,
|
||||
"completion_tokens": 45,
|
||||
"prompt_tokens_details": {"cached_tokens": 100},
|
||||
})
|
||||
|
||||
assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (120, 45, 100)
|
||||
assert event.estimated is False
|
||||
# Cached tokens are a SUBSET of input, so adding them would double-count.
|
||||
assert event.total_tokens == 165
|
||||
|
||||
|
||||
def test_anthropic_usage_accumulator_maps_onto_the_canonical_event():
|
||||
"""Anthropic reports input on message_start and output on message_delta, so
|
||||
providers/anthropic.py accumulates them into in/out/cache keys."""
|
||||
event = telemetry.anthropic_usage_event("anthropic", "claude", {
|
||||
"in": 200, "out": 80, "cache": 150,
|
||||
})
|
||||
|
||||
assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (200, 80, 150)
|
||||
assert event.estimated is False
|
||||
|
||||
|
||||
def test_a_missing_usage_block_produces_an_estimated_event():
|
||||
event = telemetry.estimated_event("ollama", "llama3.1", "x" * 400, "y" * 40)
|
||||
|
||||
assert event.estimated is True
|
||||
assert event.input_tokens == 100 # ~4 characters per token
|
||||
assert event.output_tokens == 10
|
||||
assert event.cached_tokens == 0
|
||||
|
||||
|
||||
def test_estimation_matches_the_legacy_tracker_formula():
|
||||
"""The extraction must not shift a single recorded number, so the estimator
|
||||
is pinned against the one it replaced."""
|
||||
from cowork_local.core import usage_tracker
|
||||
|
||||
for text in ("", "short", "x" * 4001, "unicode - tiếng Việt"):
|
||||
assert telemetry.estimate_tokens(text) == usage_tracker.estimate_tokens(text)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Sinks
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_recording_sink_collects_events_for_assertions():
|
||||
sink = telemetry.RecordingUsageSink()
|
||||
|
||||
sink.record(telemetry.UsageEvent("p", "m", input_tokens=10, output_tokens=5))
|
||||
sink.record(telemetry.UsageEvent("p", "m", input_tokens=1, output_tokens=1))
|
||||
|
||||
assert len(sink.events) == 2
|
||||
assert sink.total_tokens == 17
|
||||
|
||||
|
||||
def test_null_sink_discards_without_error():
|
||||
telemetry.NullUsageSink().record(telemetry.UsageEvent("p", "m"))
|
||||
|
||||
|
||||
def test_tracker_sink_forwards_every_field_positionally():
|
||||
"""``core.usage_tracker.record`` takes positional counts plus an ``estimated``
|
||||
keyword; the adapter has to preserve that exact call shape."""
|
||||
seen: Dict[str, Any] = {}
|
||||
|
||||
class _Tracker:
|
||||
@staticmethod
|
||||
def record(provider, model, input_tokens, output_tokens, cached_tokens,
|
||||
estimated=False):
|
||||
# Fields captured explicitly rather than via locals(), which would
|
||||
# also drag in the closed-over `seen` binding itself.
|
||||
seen.update({"provider": provider, "model": model,
|
||||
"input_tokens": input_tokens, "output_tokens": output_tokens,
|
||||
"cached_tokens": cached_tokens, "estimated": estimated})
|
||||
|
||||
telemetry.UsageTrackerSink(tracker=_Tracker()).record(
|
||||
telemetry.UsageEvent("anthropic", "claude", 7, 3, 2, estimated=True))
|
||||
|
||||
assert seen == {"provider": "anthropic", "model": "claude", "input_tokens": 7,
|
||||
"output_tokens": 3, "cached_tokens": 2, "estimated": True}
|
||||
|
||||
|
||||
def test_a_failing_tracker_never_raises_into_the_turn():
|
||||
class _Broken:
|
||||
@staticmethod
|
||||
def record(*_args, **_kwargs):
|
||||
raise OSError("usage store is read-only")
|
||||
|
||||
# Must not raise - the turn that produced this event has already succeeded.
|
||||
telemetry.UsageTrackerSink(tracker=_Broken()).record(telemetry.UsageEvent("p", "m"))
|
||||
|
||||
|
||||
def test_set_default_sink_returns_the_previous_one_for_restoration():
|
||||
replacement = telemetry.RecordingUsageSink()
|
||||
|
||||
previous = telemetry.set_default_sink(replacement)
|
||||
try:
|
||||
assert telemetry.default_sink is replacement
|
||||
finally:
|
||||
telemetry.set_default_sink(previous)
|
||||
assert telemetry.default_sink is previous
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Provider integration - the seam actually being used
|
||||
# --------------------------------------------------------------------------- #
|
||||
class _StubResponse:
|
||||
"""The few members the provider streaming loop touches."""
|
||||
|
||||
def __init__(self, lines: List[str]) -> None:
|
||||
self._lines = lines
|
||||
self.status_code = 200
|
||||
self.headers: Dict[str, str] = {}
|
||||
self.encoding = "utf-8"
|
||||
self.text = ""
|
||||
|
||||
def iter_lines(self, decode_unicode: bool = False):
|
||||
yield from self._lines
|
||||
|
||||
def close(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sink(monkeypatch):
|
||||
"""A per-instance recording sink, so nothing touches the real usage store."""
|
||||
return telemetry.RecordingUsageSink()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def canned(monkeypatch):
|
||||
def _install(lines: List[str]):
|
||||
monkeypatch.setattr(Provider, "_request",
|
||||
lambda self, method, url, **kw: _StubResponse(lines))
|
||||
return _install
|
||||
|
||||
|
||||
def test_openai_provider_reports_server_counts_to_its_sink(canned, sink):
|
||||
canned([
|
||||
'data: ' + json.dumps({"choices": [{"delta": {"content": "hi"}}],
|
||||
"usage": {"prompt_tokens": 11, "completion_tokens": 2,
|
||||
"prompt_tokens_details": {"cached_tokens": 4}}}),
|
||||
"data: [DONE]",
|
||||
])
|
||||
provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
|
||||
"api_key": "k", "model": "gpt-4o-mini"})
|
||||
provider.usage_sink = sink
|
||||
|
||||
provider.chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
assert len(sink.events) == 1
|
||||
event = sink.events[0]
|
||||
assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (11, 2, 4)
|
||||
assert event.estimated is False
|
||||
assert event.model == "gpt-4o-mini"
|
||||
|
||||
|
||||
def test_openai_provider_estimates_when_the_gateway_sends_no_usage(canned, sink):
|
||||
canned(['data: ' + json.dumps({"choices": [{"delta": {"content": "hello"}}]}),
|
||||
"data: [DONE]"])
|
||||
provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
|
||||
"api_key": "k", "model": "m"})
|
||||
provider.usage_sink = sink
|
||||
|
||||
provider.chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
assert sink.events[0].estimated is True
|
||||
assert sink.events[0].output_tokens >= 1
|
||||
|
||||
|
||||
def test_anthropic_provider_reports_stream_counts_to_its_sink(canned, sink):
|
||||
canned(['data: ' + json.dumps(p) for p in (
|
||||
{"type": "message_start", "message": {"usage": {"input_tokens": 30,
|
||||
"cache_read_input_tokens": 10}}},
|
||||
{"type": "content_block_delta", "index": 0,
|
||||
"delta": {"type": "text_delta", "text": "ok"}},
|
||||
{"type": "message_delta", "usage": {"output_tokens": 5}},
|
||||
{"type": "message_stop"},
|
||||
)])
|
||||
provider = AnthropicProvider({"base_url": "https://x.invalid", "api_key": "k",
|
||||
"model": "claude"})
|
||||
provider.usage_sink = sink
|
||||
|
||||
provider.chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
event = sink.events[0]
|
||||
assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (30, 5, 10)
|
||||
assert event.estimated is False
|
||||
|
||||
|
||||
def test_a_provider_without_an_explicit_sink_uses_the_process_default(canned):
|
||||
"""Existing call sites set no sink, so the default has to keep working -
|
||||
that is what makes this extraction a no-op for production behaviour."""
|
||||
canned(['data: ' + json.dumps({"choices": [{"delta": {"content": "x"}}]}),
|
||||
"data: [DONE]"])
|
||||
recorder = telemetry.RecordingUsageSink()
|
||||
previous = telemetry.set_default_sink(recorder)
|
||||
try:
|
||||
OpenAICompatProvider({"base_url": "https://x.invalid/v1", "api_key": "k",
|
||||
"model": "m"}).chat([{"role": "user", "content": "hi"}])
|
||||
finally:
|
||||
telemetry.set_default_sink(previous)
|
||||
|
||||
assert len(recorder.events) == 1
|
||||
|
||||
|
||||
def test_a_sink_that_raises_does_not_fail_the_turn(canned):
|
||||
"""The answer has already been produced by the time usage is recorded;
|
||||
losing the telemetry is strictly better than losing the answer."""
|
||||
canned(['data: ' + json.dumps({"choices": [{"delta": {"content": "x"}}]}),
|
||||
"data: [DONE]"])
|
||||
|
||||
class _Exploding:
|
||||
def record(self, _event):
|
||||
raise RuntimeError("sink is down")
|
||||
|
||||
provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
|
||||
"api_key": "k", "model": "m"})
|
||||
provider.usage_sink = _Exploding()
|
||||
|
||||
result = provider.chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
assert result["content"] == "x"
|
||||
Reference in New Issue
Block a user