feat(R03): unify model routing and centralise the provider catalogue

EPIC R03 (Team Duy) — Model Providers & Routing. All six tasks done.

R03-T02 — Provider catalogue
  domain/models/provider_descriptor.py     ProviderDescriptor (frozen), WireProtocol, AuthKind
  infrastructure/providers/provider_registry.py
                                           thread-safe registry: id/alias lookup, dynamic
                                           lookup by model id, adapter selection by protocol
  providers/factory.py                     drops its own _REGISTRY table and delegates to the
                                           registry, still raising ProviderError for callers

R03-T03 — RoutingApplicationService (pure Python, 4 modes)
  application/model_routing/routing_models.py
                                           RoutingMode (off/auto/manual/fallback),
                                           RoutingRequest (immutable snapshot), RouteEvaluation,
                                           RoutingOutcome
  application/model_routing/routing_application_service.py
                                           the single decision flow, reached through two narrow
                                           ports plus a caller-supplied confirm callback, so no
                                           Qt import is needed
  application/model_routing/core_routing_adapter.py
                                           binds the ports to core/routing and AppContext

  Fallback is a new resilience mode: keep the selected model while it can serve the turn,
  re-route only when it cannot. Wired end to end through config.py, state.py,
  ui/routing_toggle.py and i18n.py (EN/JA/VI).

R03-T04 / T05 — Remove the duplicated routing flow
  ui/chat_panel.py (#L638), ui/co4e_tab.py, ui/folder_tab.py each drop ~35 lines of copied
  logic and call the shared service; the widgets now only build a RoutingRequest, host the
  Manual-mode modal and render the outcome.

R03-T06 — Token usage as an event
  infrastructure/telemetry/usage_sink.py   UsageEvent + UsageEventSink protocol, with tracker,
                                           in-memory and composite sinks
  providers/openai_compat.py, providers/anthropic.py
                                           publish a UsageEvent instead of writing to the
                                           usage tracker themselves
  core/usage_tracker.py                    adds current_context() so a sink can borrow and
                                           restore a thread's attribution

R03-T01 — Contract tests
  tests/contracts/test_providers.py parametrises over every provider in the registry: chat()
  signature, canonical assistant message, normalised tool calls, response closed, tool schema
  translation, ProviderError, list_models/test_connection, one UsageEvent per turn.

Test infrastructure fix (required to verify any of the above): tests/conftest.py used to put
the repository's PARENT directory on sys.path, so `import cowork_local.*` resolved against
whichever sibling folder happened to carry that name — on a dev machine, an unrelated older
checkout. The suite reported green while exercising different code. The conftest now binds
this checkout to the cowork_local name in sys.modules.

Verification
  pytest tests/                    236 passed in ~1.8s (102 before this change)
  scripts/check_imports.py         PASS, 0 forbidden imports in domain/ and application/
  new production files             largest is 288 lines, all under the 400 LOC ceiling
  new tests                        134 (50 contract, 70 unit, 14 integration), all offline

scripts/run_quality_gate.py does not exist yet (R10-T02), so DoD item 7 was covered by
check_imports.py plus the full suite.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-08-22 19:36:20 +09:00
co-authored by Claude Opus 5
parent 10739f19aa
commit f61c5474b0
30 changed files with 3458 additions and 166 deletions
+220
View File
@@ -0,0 +1,220 @@
"""Unit tests for the adapters that bridge the routing engine to the app service.
The integration suite covers the happy path over the real engine; this file pins
the translation edge cases that are hard to provoke there — malformed task
types, a missing ranking, and the service-caching contract.
"""
from __future__ import annotations
import pytest
from cowork_local.application.model_routing import (
AppContextModeResolver,
CoreRoutingEngine,
RoutingApplicationService,
RoutingMode,
RoutingRequest,
)
from cowork_local.application.model_routing.core_routing_adapter import (
build_routing_application_service,
)
from cowork_local.core.routing.models import SwitchDecision, SwitchMode, TaskType
class FakeRanking:
"""Just enough of ``selector.Ranking`` for the adapter's usability check."""
def __init__(self, scores) -> None:
self._scores = dict(scores)
def score_of(self, key: str) -> float:
return self._scores.get(key, 0.0)
class FakeRouteResult:
"""Stands in for ``core.routing.service.RouteResult``."""
def __init__(self, decision, task_type=TaskType.CODING, ranking=None, target=None) -> None:
self.decision = decision
self.task_type = task_type
self.ranking = ranking
self._target = target
@property
def should_switch(self) -> bool:
return self.decision.should_switch
def target(self):
return self._target
class FakeRoutingService:
"""Records the arguments the adapter forwards to the engine."""
def __init__(self, result: FakeRouteResult) -> None:
self.result = result
self.calls: list = []
def route(self, surface, prompt, current_provider, current_model, **kwargs):
self.calls.append({"surface": surface, "prompt": prompt,
"current_provider": current_provider,
"current_model": current_model, **kwargs})
return self.result
def make_decision(**overrides) -> SwitchDecision:
fields = dict(
should_switch=True,
from_model="anthropic/weak-model",
to_model="anthropic/strong-model",
score_gain=0.3,
reason="coding fit 0.9 > current 0.6",
mode=SwitchMode.AUTO,
task_type="coding",
)
fields.update(overrides)
return SwitchDecision(**fields)
def make_request(**overrides) -> RoutingRequest:
fields = dict(surface="cowork", prompt="Fix this bug",
current_provider="anthropic", current_model="weak-model")
fields.update(overrides)
return RoutingRequest(**fields)
# --------------------------------------------------------------------------- #
# CoreRoutingEngine translation
# --------------------------------------------------------------------------- #
def test_engine_flattens_the_route_result() -> None:
"""No ``core.routing`` type may leak past the adapter — the application
service and the widgets only ever see plain fields."""
service = FakeRoutingService(FakeRouteResult(
make_decision(),
ranking=FakeRanking({"anthropic/weak-model": 0.6}),
target=("anthropic", "strong-model"),
))
evaluation = CoreRoutingEngine(service).evaluate(make_request(), RoutingMode.AUTO)
assert evaluation.task_type == "coding" # str, not TaskType
assert evaluation.should_switch is True
assert evaluation.target_provider == "anthropic"
assert evaluation.target_model == "strong-model"
assert evaluation.score_gain == pytest.approx(0.3)
assert evaluation.current_is_usable is True
def test_engine_forwards_the_mode_as_a_plain_string() -> None:
"""``RoutingService.route`` takes the mode as a string; handing it an enum
would silently fall through to its "unknown mode -> off" branch."""
service = FakeRoutingService(FakeRouteResult(make_decision(should_switch=False)))
CoreRoutingEngine(service).evaluate(make_request(), RoutingMode.AUTO)
assert service.calls[0]["mode_override"] == "auto"
def test_engine_reports_an_unranked_model_as_unusable() -> None:
"""This is the signal Fallback acts on: absent from the ranking means the
selector already rejected it (unavailable / no probe / failed probe)."""
service = FakeRoutingService(FakeRouteResult(
make_decision(),
ranking=FakeRanking({"anthropic/strong-model": 0.9}), # current is absent
target=("anthropic", "strong-model"),
))
evaluation = CoreRoutingEngine(service).evaluate(make_request(), RoutingMode.AUTO)
assert evaluation.current_is_usable is False
def test_engine_assumes_usable_without_a_ranking() -> None:
"""No ranking (routing off, or the engine's own error path) is absence of
evidence — it must not trigger a surprise Fallback switch."""
service = FakeRoutingService(FakeRouteResult(make_decision(), ranking=None))
evaluation = CoreRoutingEngine(service).evaluate(make_request(), RoutingMode.AUTO)
assert evaluation.current_is_usable is True
def test_engine_assumes_usable_when_the_ranking_misbehaves() -> None:
"""A broken ranking object must not fail the turn."""
class BrokenRanking:
def score_of(self, key):
raise RuntimeError("corrupt ranking")
service = FakeRoutingService(FakeRouteResult(make_decision(), ranking=BrokenRanking()))
evaluation = CoreRoutingEngine(service).evaluate(make_request(), RoutingMode.AUTO)
assert evaluation.current_is_usable is True
@pytest.mark.parametrize(
"raw, expected",
[("coding", TaskType.CODING), ("QA", TaskType.QA), (None, None), ("nonsense", None)],
)
def test_task_type_strings_are_coerced_or_dropped(raw, expected) -> None:
"""A pinned task type is honoured; an unknown one falls back to letting the
engine classify the prompt rather than raising mid-turn."""
service = FakeRoutingService(FakeRouteResult(make_decision(should_switch=False)))
CoreRoutingEngine(service).evaluate(make_request(task_type=raw), RoutingMode.AUTO)
assert service.calls[0]["task_type"] == expected
def test_required_capabilities_are_passed_as_a_list_or_none() -> None:
"""``rank_models`` filters on a list; an empty tuple must become None so it
is treated as "no filter" rather than "require nothing, but filter"."""
service = FakeRoutingService(FakeRouteResult(make_decision(should_switch=False)))
engine = CoreRoutingEngine(service)
engine.evaluate(make_request(required_capabilities=("vision",)), RoutingMode.AUTO)
engine.evaluate(make_request(), RoutingMode.AUTO)
assert service.calls[0]["required_capabilities"] == ["vision"]
assert service.calls[1]["required_capabilities"] is None
# --------------------------------------------------------------------------- #
# Mode resolver + wiring
# --------------------------------------------------------------------------- #
def test_mode_resolver_reads_the_per_workspace_mode() -> None:
"""Per-workspace routing keeps working now that the lookup left the widgets."""
class StubCtx:
def project_routing_mode(self, surface):
return "fallback" if surface == "co4e" else "off"
resolver = AppContextModeResolver(StubCtx())
assert resolver.mode_for("co4e") is RoutingMode.FALLBACK
assert resolver.mode_for("cowork") is RoutingMode.OFF
def test_service_is_built_once_and_cached_on_the_context() -> None:
"""Every surface must share one instance, so future per-surface state (a
cool-down, a switch history) is shared rather than duplicated per widget."""
class StubCtx:
def __init__(self):
self.routing_calls = 0
self.config = type("Cfg", (), {"routing": {"confirm_timeout_sec": 45}})()
def routing(self):
self.routing_calls += 1
return FakeRoutingService(FakeRouteResult(make_decision(should_switch=False)))
def project_routing_mode(self, surface):
return "off"
ctx = StubCtx()
first = build_routing_application_service(ctx)
second = build_routing_application_service(ctx)
assert first is second
assert ctx.routing_calls == 1
assert isinstance(first, RoutingApplicationService)
# The confirm timeout is read from config at call time, not frozen at build.
assert first.confirm_timeout() == pytest.approx(45.0)
+204
View File
@@ -0,0 +1,204 @@
"""R03-T02 — unit tests for ProviderDescriptor and the central ProviderRegistry.
Covers what the rest of the app now relies on the catalogue for: resolving ids
and aliases, resolving a bare model id back to its provider, filling in default
models, and refusing to let a duplicate registration silently hijack a built-in.
"""
from __future__ import annotations
import pytest
from cowork_local.domain.models.provider_descriptor import (
AuthKind,
ProviderDescriptor,
WireProtocol,
)
from cowork_local.infrastructure.providers.provider_registry import (
BUILTIN_DESCRIPTORS,
ProviderNotFoundError,
ProviderRegistry,
)
def make_descriptor(**overrides) -> ProviderDescriptor:
"""A minimal valid descriptor; tests override just the field under test."""
fields = dict(
provider_id="demo",
display_name="Demo provider",
wire_protocol=WireProtocol.OPENAI_COMPAT,
default_model="demo-small",
models=("demo-small", "demo-large"),
)
fields.update(overrides)
return ProviderDescriptor(**fields)
# --------------------------------------------------------------------------- #
# ProviderDescriptor
# --------------------------------------------------------------------------- #
def test_descriptor_rejects_an_empty_id() -> None:
"""An id-less descriptor could never be looked up, so it must not exist."""
with pytest.raises(ValueError):
make_descriptor(provider_id="")
def test_descriptor_rejects_a_non_enum_protocol() -> None:
"""The protocol drives adapter selection; a stray string would silently
fall through to "no adapter" at build time instead of failing here."""
with pytest.raises(TypeError):
make_descriptor(wire_protocol="openai_compat")
def test_descriptor_is_immutable() -> None:
"""Descriptors are shared process-wide; a mutation would be visible to every
other reader mid-iteration."""
descriptor = make_descriptor()
with pytest.raises(Exception):
descriptor.default_model = "hacked" # type: ignore[misc]
def test_id_matching_ignores_case_and_honours_aliases() -> None:
"""Provider ids come from hand-edited config files and old app versions."""
descriptor = make_descriptor(aliases=("legacy-demo",))
assert descriptor.matches("DEMO")
assert descriptor.matches(" legacy-demo ")
assert not descriptor.matches("other")
def test_capabilities_use_the_routing_vocabulary() -> None:
"""The set must be feedable straight into the routing selector's filter."""
descriptor = make_descriptor(supports_vision=True, supports_tools=True,
supports_streaming=False)
assert descriptor.capabilities == frozenset({"vision", "tools"})
assert descriptor.has_capability("vision")
assert not descriptor.has_capability("streaming")
def test_average_cost_is_none_when_a_price_is_unknown() -> None:
"""Unknown prices stay unknown — a guessed number would silently skew the
routing scorer's cost term."""
assert make_descriptor(cost_per_1k_input=0.5).avg_cost_per_1k is None
priced = make_descriptor(cost_per_1k_input=1.0, cost_per_1k_output=3.0)
# Same 1:3 input:output weighting as ModelMetadata.avg_cost_per_1k.
assert priced.avg_cost_per_1k == pytest.approx((1.0 + 9.0) / 4.0)
def test_resolve_model_prefers_the_caller_then_the_default() -> None:
"""One place implements the "picked model or provider default" fallback that
every chat surface used to re-implement inline."""
descriptor = make_descriptor()
assert descriptor.resolve_model("demo-large") == "demo-large"
assert descriptor.resolve_model("") == "demo-small"
assert descriptor.resolve_model(" ") == "demo-small"
def test_with_models_repoints_a_default_that_vanished() -> None:
"""After discovery, the default must still name a model that exists."""
descriptor = make_descriptor()
updated = descriptor.with_models(["demo-v2", "demo-v2", "demo-v3"])
assert updated.models == ("demo-v2", "demo-v3") # de-duplicated, order kept
assert updated.default_model == "demo-v2"
assert descriptor.models == ("demo-small", "demo-large"), "original was mutated"
def test_with_models_keeps_a_default_that_survived() -> None:
"""Discovery must not reshuffle a user's working selection."""
updated = make_descriptor().with_models(["demo-large", "demo-small"])
assert updated.default_model == "demo-small"
# --------------------------------------------------------------------------- #
# ProviderRegistry
# --------------------------------------------------------------------------- #
def test_registry_resolves_ids_aliases_and_reports_unknowns() -> None:
"""Lookup must be forgiving about form, but loud about genuinely unknown
providers — a typo should fail at the call site, not as a None later."""
registry = ProviderRegistry([make_descriptor(aliases=("legacy-demo",))])
assert registry.get("demo").provider_id == "demo"
assert registry.get("legacy-demo").provider_id == "demo"
assert registry.find("missing") is None
assert "demo" in registry
with pytest.raises(ProviderNotFoundError):
registry.get("missing")
def test_registry_refuses_to_overwrite_silently_but_replace_works() -> None:
"""A second registration of the same id is almost always a bug; updating a
descriptor is a deliberate act with its own method."""
registry = ProviderRegistry([make_descriptor()])
with pytest.raises(ValueError):
registry.register(make_descriptor(display_name="Impostor"))
registry.replace(make_descriptor(display_name="Renamed"))
assert registry.get("demo").display_name == "Renamed"
assert len(registry) == 1
def test_registry_re_registering_an_identical_descriptor_is_a_no_op() -> None:
"""Idempotent registration keeps repeated bootstrap calls harmless."""
registry = ProviderRegistry([make_descriptor()])
registry.register(make_descriptor())
assert len(registry) == 1
def test_find_by_model_resolves_a_bare_model_id() -> None:
"""Routing decisions and saved conversations sometimes carry only a model
name; the registry is what turns that back into a provider."""
registry = ProviderRegistry([make_descriptor()])
assert registry.find_by_model("demo-large").provider_id == "demo"
# A gateway model we cannot enumerate offline is a miss, not an error — the
# caller falls back to the configured active provider.
assert registry.find_by_model("unknown-model") is None
assert registry.find_by_model("") is None
def test_builtin_catalogue_covers_every_configured_provider() -> None:
"""The catalogue and DEFAULT_CONFIG must not drift: a provider users can
configure but the registry cannot build is a dead Settings entry."""
from cowork_local.config import DEFAULT_CONFIG
registry = ProviderRegistry(BUILTIN_DESCRIPTORS)
for provider_id in DEFAULT_CONFIG["providers"]:
assert registry.find(provider_id) is not None, f"{provider_id} missing from registry"
def test_build_fills_in_the_default_model() -> None:
"""A half-written config must still produce a usable provider rather than an
empty model id that only fails once the request reaches the gateway."""
registry = ProviderRegistry(BUILTIN_DESCRIPTORS)
provider = registry.build("anthropic", {"api_key": "k"})
assert provider.model == registry.get("anthropic").default_model
def test_build_respects_an_explicit_model() -> None:
"""Per-tab model selection must win over the catalogue default."""
registry = ProviderRegistry(BUILTIN_DESCRIPTORS)
provider = registry.build("anthropic", {"api_key": "k", "model": "claude-opus-4-8"})
assert provider.model == "claude-opus-4-8"
def test_factory_still_raises_provider_error_for_unknown_ids() -> None:
"""Existing call sites catch ProviderError; routing lookups through the
registry must not change the exception type they see."""
from cowork_local.providers import build_provider
from cowork_local.providers.base import ProviderError
with pytest.raises(ProviderError):
build_provider("definitely-not-a-provider", {})
@@ -0,0 +1,384 @@
"""R03-T03 — unit tests for the unified routing decision rules.
The point of moving these rules out of the three chat widgets is that they can
now be exercised without Qt, without the assessment store and without a network:
the service talks to two narrow ports, so every mode is driven here by ~10-line
fakes. Each test names the behaviour a chat surface depends on.
"""
from __future__ import annotations
import pytest
from cowork_local.application.model_routing import (
RouteEvaluation,
RoutingApplicationService,
RoutingMode,
RoutingOutcome,
RoutingRequest,
)
class FakeDecisionPort:
"""A routing engine that returns a canned verdict and records its input."""
def __init__(self, evaluation: RouteEvaluation) -> None:
self.evaluation = evaluation
self.calls: list = []
def evaluate(self, request: RoutingRequest, mode: RoutingMode) -> RouteEvaluation:
self.calls.append((request, mode))
return self.evaluation
class ExplodingDecisionPort:
"""An engine that fails — proves routing degrades instead of breaking a turn."""
def evaluate(self, request: RoutingRequest, mode: RoutingMode) -> RouteEvaluation:
raise RuntimeError("assessment store is corrupt")
class FakeModeResolver:
"""Per-surface mode lookup, standing in for the workspace settings."""
def __init__(self, mode) -> None:
self.mode = mode
self.surfaces: list = []
def mode_for(self, surface: str):
self.surfaces.append(surface)
return self.mode
def make_request(**overrides) -> RoutingRequest:
"""A representative turn: Cowork chat, currently on a cheap OpenAI model."""
fields = dict(
surface="cowork",
prompt="Refactor this function",
current_provider="codex",
current_model="gpt-4o-mini",
)
fields.update(overrides)
return RoutingRequest(**fields)
def switch_evaluation(**overrides) -> RouteEvaluation:
"""An engine verdict that proposes a switch to a better coding model."""
fields = dict(
task_type="coding",
should_switch=True,
target_provider="anthropic",
target_model="claude-sonnet-4-6",
score_gain=0.21,
reason="coding fit 0.88 > current 0.67",
decision=object(),
)
fields.update(overrides)
return RouteEvaluation(**fields)
# --------------------------------------------------------------------------- #
# Off
# --------------------------------------------------------------------------- #
def test_off_mode_never_consults_the_engine() -> None:
"""Off must be free: no ranking, no store read, no decision at all."""
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.OFF))
outcome = service.resolve(make_request())
assert outcome.switched is False
assert outcome.provider is None and outcome.model is None
assert port.calls == [], "Off mode must not call the routing engine"
def test_missing_mode_resolver_defaults_to_off() -> None:
"""Routing stays opt-in: with no way to read the mode, never switch."""
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port)
outcome = service.resolve(make_request())
assert outcome.mode is RoutingMode.OFF
assert outcome.switched is False
def test_empty_prompt_is_not_routed() -> None:
"""An empty message carries no signal to classify, so the engine is skipped."""
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.AUTO))
outcome = service.resolve(make_request(prompt=" "))
assert outcome.switched is False
assert port.calls == []
# --------------------------------------------------------------------------- #
# Auto
# --------------------------------------------------------------------------- #
def test_auto_mode_switches_silently() -> None:
"""Auto applies the engine's verdict without asking the user."""
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.AUTO))
outcome = service.resolve(make_request())
assert outcome.switched is True
assert outcome.provider == "anthropic"
assert outcome.model == "claude-sonnet-4-6"
assert outcome.task_type == "coding"
assert outcome.score_gain == pytest.approx(0.21)
assert outcome.should_notify is True
def test_auto_mode_keeps_current_when_nothing_is_better() -> None:
"""No proposed switch means the surface's own selection is untouched."""
port = FakeDecisionPort(switch_evaluation(
should_switch=False, reason="current model is already best-fit"))
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.AUTO))
outcome = service.resolve(make_request())
assert outcome.switched is False
assert outcome.provider is None
assert "already best-fit" in outcome.reason
def test_switch_without_a_target_is_ignored() -> None:
"""A verdict that says "switch" but names nothing is not actionable — a
surface must never be handed an empty model id."""
port = FakeDecisionPort(switch_evaluation(target_provider=None, target_model=None))
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.AUTO))
outcome = service.resolve(make_request())
assert outcome.switched is False
def test_same_provider_switch_keeps_the_current_provider() -> None:
"""A model-only switch must not blank out the provider the surface uses."""
port = FakeDecisionPort(switch_evaluation(target_provider=None, target_model="o3"))
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.AUTO))
outcome = service.resolve(make_request())
assert outcome.switched is True
assert outcome.provider == "codex" # unchanged, from the request
assert outcome.model == "o3"
# --------------------------------------------------------------------------- #
# Manual
# --------------------------------------------------------------------------- #
def test_manual_mode_switches_only_after_approval() -> None:
"""Manual's contract: ask first, then apply exactly what was approved."""
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(
port, FakeModeResolver(RoutingMode.MANUAL),
confirm_timeout_sec=lambda: 30.0,
)
asked: list = []
def confirm(decision, timeout):
asked.append((decision, timeout))
return True
outcome = service.resolve(make_request(), confirm=confirm)
assert outcome.switched is True
assert len(asked) == 1
# The configured timeout must reach the dialog, not a hard-coded default.
assert asked[0][1] == pytest.approx(30.0)
def test_manual_mode_decline_is_reported_distinctly() -> None:
""""The user said no" must be distinguishable from "nothing better found",
so a surface can stay quiet in one case and explain itself in the other."""
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.MANUAL))
outcome = service.resolve(make_request(), confirm=lambda decision, timeout: False)
assert outcome.switched is False
assert outcome.declined is True
def test_manual_mode_without_a_callback_never_switches() -> None:
"""Silently switching in Manual mode would violate the mode's promise."""
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.MANUAL))
outcome = service.resolve(make_request(), confirm=None)
assert outcome.switched is False
def test_manual_mode_treats_a_broken_dialog_as_a_decline() -> None:
"""A crashing confirm dialog must not auto-approve a model change."""
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.MANUAL))
def confirm(decision, timeout):
raise RuntimeError("dialog blew up")
outcome = service.resolve(make_request(), confirm=confirm)
assert outcome.switched is False
assert outcome.declined is True
# --------------------------------------------------------------------------- #
# Fallback
# --------------------------------------------------------------------------- #
def test_fallback_keeps_a_healthy_model_even_when_a_better_one_exists() -> None:
"""Fallback is a resilience mode, not an optimiser: a usable pinned model
wins over a higher-scoring candidate."""
port = FakeDecisionPort(switch_evaluation(current_is_usable=True))
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.FALLBACK))
outcome = service.resolve(make_request())
assert outcome.switched is False
assert "healthy" in outcome.reason
def test_fallback_switches_when_the_current_model_cannot_serve_the_turn() -> None:
"""The one case Fallback exists for: rescue an unusable selection."""
port = FakeDecisionPort(switch_evaluation(current_is_usable=False))
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.FALLBACK))
outcome = service.resolve(make_request())
assert outcome.switched is True
assert outcome.model == "claude-sonnet-4-6"
def test_fallback_asks_the_engine_with_auto_semantics() -> None:
"""The engine only understands off/auto/manual, so Fallback must reach it as
Auto — otherwise the engine would reject the unknown mode and rank nothing."""
port = FakeDecisionPort(switch_evaluation(current_is_usable=False))
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.FALLBACK))
service.resolve(make_request())
assert port.calls[0][1] is RoutingMode.AUTO
def test_fallback_never_confirms_with_the_user() -> None:
"""Rescuing an unusable model is not a proposal — it happens silently."""
port = FakeDecisionPort(switch_evaluation(current_is_usable=False))
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.FALLBACK))
asked: list = []
outcome = service.resolve(
make_request(), confirm=lambda decision, timeout: asked.append(1) or True)
assert outcome.switched is True
assert asked == []
def test_fallback_with_no_replacement_keeps_current() -> None:
"""Nothing to fall back to means keep going with what we have and let the
provider surface the real error, rather than blanking the model."""
port = FakeDecisionPort(switch_evaluation(
current_is_usable=False, target_provider=None, target_model=None))
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.FALLBACK))
outcome = service.resolve(make_request())
assert outcome.switched is False
# --------------------------------------------------------------------------- #
# Robustness & plumbing
# --------------------------------------------------------------------------- #
def test_engine_failure_degrades_to_keep_current() -> None:
"""A broken assessment store must never stop a user sending a message."""
service = RoutingApplicationService(
ExplodingDecisionPort(), FakeModeResolver(RoutingMode.AUTO))
outcome = service.resolve(make_request())
assert isinstance(outcome, RoutingOutcome)
assert outcome.switched is False
assert "error" in outcome.reason
def test_mode_resolver_failure_degrades_to_off() -> None:
"""An unreadable workspace config must not enable routing by accident."""
class BrokenResolver:
def mode_for(self, surface):
raise OSError("workspace file unreadable")
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port, BrokenResolver())
outcome = service.resolve(make_request())
assert outcome.mode is RoutingMode.OFF
assert port.calls == []
def test_explicit_request_mode_overrides_the_resolver() -> None:
"""A surface may pin the mode for one turn (tests, replay, admin actions)."""
resolver = FakeModeResolver(RoutingMode.OFF)
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port, resolver)
outcome = service.resolve(make_request(mode=RoutingMode.AUTO))
assert outcome.switched is True
assert resolver.surfaces == [], "an explicit mode must skip the resolver"
def test_request_is_forwarded_to_the_engine_unchanged() -> None:
"""Surface, prompt and pinned task type must survive the hand-off — AI-Edit
relies on its "coding" pin reaching the engine."""
port = FakeDecisionPort(switch_evaluation())
service = RoutingApplicationService(port, FakeModeResolver(RoutingMode.AUTO))
request = make_request(surface="ai_edit", task_type="coding",
required_capabilities=("vision",))
service.resolve(request)
forwarded = port.calls[0][0]
assert forwarded is request
assert forwarded.surface == "ai_edit"
assert forwarded.task_type == "coding"
assert forwarded.required_capabilities == ("vision",)
@pytest.mark.parametrize(
"raw, expected",
[
("auto", RoutingMode.AUTO),
("MANUAL", RoutingMode.MANUAL),
(" fallback ", RoutingMode.FALLBACK),
("nonsense", RoutingMode.OFF),
("", RoutingMode.OFF),
(None, RoutingMode.OFF),
],
)
def test_mode_parsing_is_forgiving(raw, expected) -> None:
"""Config values are hand-edited; an unknown one must degrade, not raise."""
assert RoutingMode.parse(raw) is expected
def test_confirm_timeout_falls_back_to_the_default_when_unusable() -> None:
"""A corrupted timeout must not produce a zero-second dialog that declines
every switch before the user can read it."""
service = RoutingApplicationService(
FakeDecisionPort(switch_evaluation()),
FakeModeResolver(RoutingMode.MANUAL),
confirm_timeout_sec=lambda: 0.0,
)
assert service.confirm_timeout() == RoutingApplicationService.DEFAULT_CONFIRM_TIMEOUT_SEC
def test_routing_request_is_immutable() -> None:
"""The snapshot must not change under a turn that is already in flight."""
request = make_request()
with pytest.raises(Exception):
request.prompt = "something else" # type: ignore[misc]
+184
View File
@@ -0,0 +1,184 @@
"""R03-T06 — unit tests for the token-usage telemetry seam.
The seam exists so provider adapters stop owning telemetry policy. These tests
pin the two properties that makes that safe: events reach every subscriber, and
no telemetry failure can ever propagate back into the turn that produced it.
"""
from __future__ import annotations
import pytest
from cowork_local.infrastructure.telemetry import usage_sink
from cowork_local.infrastructure.telemetry.usage_sink import (
CompositeUsageSink,
InMemoryUsageSink,
UsageEvent,
UsageTrackerSink,
)
@pytest.fixture(autouse=True)
def isolated_sink(monkeypatch):
"""Give every test its own process-wide sink.
Autouse because a leaked sink would let one test's subscriber observe the
next test's events — and, worse, let a test write to the developer's real
usage files through the default tracker sink.
"""
monkeypatch.setattr(usage_sink, "_sink", None)
yield
monkeypatch.setattr(usage_sink, "_sink", None)
def make_event(**overrides) -> UsageEvent:
fields = dict(provider="anthropic", model="claude-sonnet-4-6",
input_tokens=100, output_tokens=40, cached_tokens=10)
fields.update(overrides)
return UsageEvent(**fields)
# --------------------------------------------------------------------------- #
# UsageEvent
# --------------------------------------------------------------------------- #
def test_event_is_immutable() -> None:
"""A subscriber must not be able to edit the event the next one receives."""
event = make_event()
with pytest.raises(Exception):
event.input_tokens = 0 # type: ignore[misc]
def test_total_tokens_does_not_double_count_cache_reads() -> None:
"""Every gateway we support already reports cached tokens inside the input
count, so adding them again would inflate the dashboard."""
assert make_event().total_tokens == 140
def test_to_dict_uses_the_stored_row_keys() -> None:
"""Matching the tracker's short keys lets a caller diff an event against a
persisted row without a translation table."""
row = make_event(source="cowork", label="Refactor chat").to_dict()
assert row["in"] == 100 and row["out"] == 40 and row["cache"] == 10
assert row["source"] == "cowork" and row["label"] == "Refactor chat"
assert row["estimated"] is False
# --------------------------------------------------------------------------- #
# Fan-out
# --------------------------------------------------------------------------- #
def test_publish_reaches_every_subscriber() -> None:
"""The whole point of the seam: extra consumers attach without patching
provider code."""
first, second = InMemoryUsageSink(), InMemoryUsageSink()
usage_sink.set_usage_sink(CompositeUsageSink([first, second]))
usage_sink.publish(make_event())
assert len(first.snapshot()) == 1
assert len(second.snapshot()) == 1
def test_one_failing_subscriber_does_not_starve_the_others() -> None:
"""A buggy consumer must not silently disable the Dashboard."""
class Exploding:
def emit(self, event):
raise RuntimeError("subscriber is broken")
healthy = InMemoryUsageSink()
usage_sink.set_usage_sink(CompositeUsageSink([Exploding(), healthy]))
usage_sink.publish(make_event())
assert len(healthy.snapshot()) == 1
def test_subscribe_and_unsubscribe_round_trip() -> None:
"""Teardown code calls unsubscribe unconditionally, so removing a sink that
was never added must be harmless."""
extra = InMemoryUsageSink()
usage_sink.subscribe(extra)
usage_sink.publish(make_event())
usage_sink.unsubscribe(extra)
usage_sink.unsubscribe(extra) # second removal is a no-op
usage_sink.publish(make_event(model="claude-opus-4-8"))
assert [e.model for e in extra.snapshot()] == ["claude-sonnet-4-6"]
def test_default_sink_is_the_usage_tracker() -> None:
"""Out of the box the seam must preserve the existing Dashboard pipeline."""
sinks = usage_sink.get_usage_sink().sinks()
assert any(isinstance(s, UsageTrackerSink) for s in sinks)
def test_in_memory_sink_totals_and_clears() -> None:
"""Test-double conveniences the contract suite relies on."""
sink = InMemoryUsageSink()
sink.emit(make_event())
sink.emit(make_event(input_tokens=1, output_tokens=1, cached_tokens=0))
assert sink.total_tokens == 142
sink.clear()
assert sink.snapshot() == []
# --------------------------------------------------------------------------- #
# UsageTrackerSink forwarding
# --------------------------------------------------------------------------- #
def test_tracker_sink_forwards_the_counts() -> None:
"""The adapter must hand the tracker exactly what the provider measured."""
recorded: list = []
def fake_record(provider, model, tokens_in, tokens_out, cached, estimated=False):
recorded.append((provider, model, tokens_in, tokens_out, cached, estimated))
UsageTrackerSink(recorder=fake_record).emit(make_event(estimated=True))
assert recorded == [("anthropic", "claude-sonnet-4-6", 100, 40, 10, True)]
def test_tracker_sink_restores_the_thread_context_it_borrowed() -> None:
"""An event carrying its own attribution must relabel ONE row, not every
later turn that happens to run on the same worker thread."""
from cowork_local.core import usage_tracker as tracker
tracker.set_context("cowork", "original chat")
seen: list = []
UsageTrackerSink(recorder=lambda *a, **k: seen.append(tracker.current_context())).emit(
make_event(source="co4e", label="flow run"))
assert seen == [("co4e", "flow run")], "event attribution was not applied"
assert tracker.current_context() == ("cowork", "original chat")
def test_tracker_sink_swallows_recorder_failures() -> None:
"""Telemetry is never allowed to abort an otherwise successful turn."""
def boom(*_args, **_kwargs):
raise OSError("usage directory is read-only")
UsageTrackerSink(recorder=boom).emit(make_event()) # must not raise
def test_publish_never_raises_even_with_a_broken_sink() -> None:
"""Last line of defence: providers call publish() inside their stream loop."""
class Hostile:
def emit(self, event):
raise RuntimeError("nope")
def sinks(self):
raise RuntimeError("nope")
usage_sink.set_usage_sink(Hostile())
usage_sink.publish(make_event()) # must not raise
def test_estimate_tokens_matches_the_tracker_heuristic() -> None:
"""Re-exported so adapters need one telemetry import; it must not drift."""
from cowork_local.core import usage_tracker as tracker
for text in ("", "a", "hello world", "x" * 4001):
assert usage_sink.estimate_tokens(text) == tracker.estimate_tokens(text)