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
+69 -4
View File
@@ -1,10 +1,75 @@
"""Make the repository package importable when pytest runs from the repo root."""
"""Make THIS checkout importable as the ``cowork_local`` package during tests.
Why this is not just a ``sys.path`` insert
------------------------------------------
Test modules import the app in two different styles:
* top-level (``from providers.base import ...``) — resolved by the repository
root already sitting on ``sys.path`` when pytest is launched from it;
* fully qualified (``from cowork_local.core.routing.service import ...``) —
which only resolves when a directory literally named ``cowork_local`` is
importable.
Simply appending the repository's PARENT directory to ``sys.path`` (the previous
behaviour) makes the second style resolve against *whatever* sibling folder
happens to be called ``cowork_local`` — on a developer machine that is often an
unrelated older checkout, so the whole suite silently exercises the wrong code
while still reporting green. Instead we bind the name ``cowork_local`` in
``sys.modules`` to the package rooted at THIS repository, so both import styles
always reach the working copy under test regardless of the checkout's directory
name.
"""
from __future__ import annotations
import importlib.util
import sys
from pathlib import Path
REPOSITORY_PARENT = Path(__file__).resolve().parents[2]
if str(REPOSITORY_PARENT) not in sys.path:
sys.path.insert(0, str(REPOSITORY_PARENT))
# .../<checkout>/tests/conftest.py -> .../<checkout>
PACKAGE_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "cowork_local"
# The repository root must stay importable so the top-level import style
# (``providers``/``domain``/``application``/``tests``) keeps working.
if str(PACKAGE_ROOT) not in sys.path:
sys.path.insert(0, str(PACKAGE_ROOT))
def _bind_checkout_as_package() -> None:
"""Register this checkout in ``sys.modules`` under the canonical package name.
Executed at import time of the conftest (i.e. before any test module is
imported) so that a stale same-named directory elsewhere on ``sys.path`` can
never win the lookup. A no-op when the package is already bound to this very
directory, which keeps repeated conftest loads (pytest-xdist, sub-sessions)
idempotent.
"""
existing = sys.modules.get(PACKAGE_NAME)
if existing is not None:
# Already bound. Only rebind when it points at a DIFFERENT checkout,
# otherwise re-executing the package __init__ would duplicate module
# state that tests may already hold references to.
origin = getattr(existing, "__file__", "") or ""
if Path(origin).resolve().parent == PACKAGE_ROOT:
return
spec = importlib.util.spec_from_file_location(
PACKAGE_NAME,
PACKAGE_ROOT / "__init__.py",
# Declaring the search locations is what turns the module into a real
# package, so ``cowork_local.core.routing`` and friends resolve as
# sub-modules of this directory.
submodule_search_locations=[str(PACKAGE_ROOT)],
)
if spec is None or spec.loader is None: # pragma: no cover — defensive
return
module = importlib.util.module_from_spec(spec)
# Insert BEFORE executing so that a circular ``import cowork_local`` from
# inside the package body resolves to the partially-initialised module
# instead of restarting the import (standard CPython import semantics).
sys.modules[PACKAGE_NAME] = module
spec.loader.exec_module(module)
_bind_checkout_as_package()
+7
View File
@@ -0,0 +1,7 @@
"""Contract tests: one shared specification every interchangeable adapter must satisfy.
Unlike unit tests (which pin ONE implementation's behaviour), a contract test is
parametrised over every implementation of an interface, so adding a new provider
means adding a row — not writing a new test file — and a provider that quietly
breaks the canonical shape fails here rather than in production.
"""
+178
View File
@@ -0,0 +1,178 @@
"""Offline transport doubles + per-protocol stream scripts for the provider contract tests.
Kept in its own module so ``test_providers.py`` stays a readable list of
assertions instead of a wall of SSE fixtures, and so the LOC ceiling (400 lines
per production file, applied here too) is comfortably met by both halves.
Nothing in here touches the network: :class:`FakeStreamResponse` mimics just
enough of ``requests.Response`` for the streaming loops in
``providers/openai_compat.py`` and ``providers/anthropic.py`` — status code,
mutable ``encoding``, ``iter_lines`` and ``close``.
"""
from __future__ import annotations
import json
from typing import Any, Dict, List, Optional
# Canonical turn every protocol script below must produce, so the contract test
# can assert one expected result no matter which provider produced it.
EXPECTED_TEXT = "Hello world"
EXPECTED_TOOL_CALL = {"id": "call-1", "name": "read_file", "arguments": {"path": "a.txt"}}
EXPECTED_INPUT_TOKENS = 11
EXPECTED_OUTPUT_TOKENS = 7
EXPECTED_CACHED_TOKENS = 3
class FakeStreamResponse:
"""A minimal stand-in for a streaming ``requests.Response``.
``iter_lines`` replays pre-baked SSE lines; ``closed`` records that the
provider released the connection, which the contract asserts because a
provider that leaks the response leaks a socket per turn.
"""
def __init__(
self,
lines: Optional[List[str]] = None,
status_code: int = 200,
body: str = "",
headers: Optional[Dict[str, str]] = None,
payload: Optional[Dict[str, Any]] = None,
) -> None:
self.status_code = status_code
self._lines = list(lines or ())
self.text = body
self.headers = dict(headers or {})
self._payload = payload
self.closed = False
# Providers force UTF-8 on the response before reading it; the attribute
# simply has to exist and be writable.
self.encoding = None
def iter_lines(self, decode_unicode: bool = False):
for line in self._lines:
yield line
def json(self) -> Any:
if self._payload is None:
raise ValueError("no JSON payload configured on this fake response")
return self._payload
def close(self) -> None:
self.closed = True
def _sse(payload: Dict[str, Any]) -> str:
"""One SSE ``data:`` line carrying a JSON event."""
return "data: " + json.dumps(payload, ensure_ascii=False)
def openai_stream_lines() -> List[str]:
"""A complete OpenAI Chat Completions stream: text, one tool call, usage.
Split across several deltas on purpose — chunk boundaries are where naive
stream parsers break, so the contract exercises them.
"""
return [
_sse({"choices": [{"delta": {"content": "Hello "}}]}),
_sse({"choices": [{"delta": {"content": "world"}}]}),
_sse({"choices": [{"delta": {"tool_calls": [{
"index": 0, "id": "call-1",
"function": {"name": "read_file", "arguments": '{"path":'},
}]}}]}),
# Arguments arrive fragmented; the provider must concatenate before parsing.
_sse({"choices": [{"delta": {"tool_calls": [{
"index": 0, "function": {"arguments": '"a.txt"}'},
}]}}]}),
_sse({
"choices": [{"delta": {}}],
"usage": {
"prompt_tokens": EXPECTED_INPUT_TOKENS,
"completion_tokens": EXPECTED_OUTPUT_TOKENS,
"prompt_tokens_details": {"cached_tokens": EXPECTED_CACHED_TOKENS},
},
}),
"data: [DONE]",
]
def anthropic_stream_lines() -> List[str]:
"""The same canonical turn expressed as an Anthropic Messages stream."""
return [
_sse({"type": "message_start", "message": {"usage": {
"input_tokens": EXPECTED_INPUT_TOKENS,
"cache_read_input_tokens": EXPECTED_CACHED_TOKENS,
}}}),
_sse({"type": "content_block_start", "index": 0,
"content_block": {"type": "text"}}),
_sse({"type": "content_block_delta", "index": 0,
"delta": {"type": "text_delta", "text": "Hello "}}),
_sse({"type": "content_block_delta", "index": 0,
"delta": {"type": "text_delta", "text": "world"}}),
_sse({"type": "content_block_start", "index": 1, "content_block": {
"type": "tool_use", "id": "call-1", "name": "read_file"}}),
_sse({"type": "content_block_delta", "index": 1,
"delta": {"type": "input_json_delta", "partial_json": '{"path":'}}),
_sse({"type": "content_block_delta", "index": 1,
"delta": {"type": "input_json_delta", "partial_json": '"a.txt"}'}}),
_sse({"type": "message_delta",
"usage": {"output_tokens": EXPECTED_OUTPUT_TOKENS}}),
_sse({"type": "message_stop"}),
]
# Per wire protocol: how to script a successful turn, and the model-list payload
# ``list_models()`` expects. Keyed by the descriptor's wire protocol value so a
# new provider that reuses an existing protocol needs no new entry here.
PROTOCOL_FIXTURES = {
"openai_compat": {
"stream_lines": openai_stream_lines,
"models_payload": {"data": [{"id": "gpt-4o-mini"}, {"id": "gpt-4o"}]},
"expected_models": ["gpt-4o-mini", "gpt-4o"],
},
"anthropic": {
"stream_lines": anthropic_stream_lines,
"models_payload": {"data": [{"id": "claude-sonnet-4-6"}]},
"expected_models": ["claude-sonnet-4-6"],
},
}
class ScriptedTransport:
"""Replaces ``Provider._request`` and hands back scripted responses.
Records every call so a test can assert *how* the provider talked to the
endpoint (method, url, JSON payload) without a socket ever being opened.
"""
def __init__(self, responses: List[FakeStreamResponse]) -> None:
self._responses = list(responses)
self.calls: List[Dict[str, Any]] = []
def __call__(self, method: str, url: str, **kwargs) -> FakeStreamResponse:
self.calls.append({"method": method, "url": url, **kwargs})
if not self._responses:
raise AssertionError(f"unexpected extra request: {method} {url}")
# Pop in order: a provider that retries gets the NEXT scripted response,
# which is how the retry/error paths are driven.
return self._responses.pop(0)
@property
def last_payload(self) -> Dict[str, Any]:
"""The JSON body of the most recent request."""
return self.calls[-1].get("json") or {}
__all__ = [
"EXPECTED_CACHED_TOKENS",
"EXPECTED_INPUT_TOKENS",
"EXPECTED_OUTPUT_TOKENS",
"EXPECTED_TEXT",
"EXPECTED_TOOL_CALL",
"FakeStreamResponse",
"PROTOCOL_FIXTURES",
"ScriptedTransport",
"anthropic_stream_lines",
"openai_stream_lines",
]
+279
View File
@@ -0,0 +1,279 @@
"""R03-T01 — the contract every LLM provider adapter must satisfy.
Parametrised over EVERY provider in the central registry
(``infrastructure/providers/provider_registry.py``), so registering a new
provider automatically subjects it to the same specification and a provider that
drifts from the canonical shapes fails here.
The contract, in one list:
* construction — the registry builds a real ``Provider`` for every id;
* ``chat()`` — canonical signature, canonical assistant message, streamed text
delivered through ``on_text``, tool calls normalised to
``{"id", "name", "arguments": dict}``, response always closed;
* tool schema translation matches the adapter's wire protocol;
* failures raise ``ProviderError`` — never a bare transport exception;
* ``list_models()`` / ``test_connection()`` report a reason instead of a silent
empty list;
* telemetry — exactly one ``UsageEvent`` per turn (R03-T06), with the real
counts when the stream reports them.
Everything runs offline: ``Provider._request`` is replaced by a scripted
transport, so the suite needs no network, no API key and no Qt event loop.
"""
from __future__ import annotations
import pytest
import requests
from cowork_local.infrastructure.providers.provider_registry import (
BUILTIN_DESCRIPTORS,
ProviderRegistry,
)
from cowork_local.infrastructure.telemetry import usage_sink
from cowork_local.providers.base import Provider, ProviderError, ToolSpec
from cowork_local.tests.contracts.provider_stubs import (
EXPECTED_CACHED_TOKENS,
EXPECTED_INPUT_TOKENS,
EXPECTED_OUTPUT_TOKENS,
EXPECTED_TEXT,
EXPECTED_TOOL_CALL,
PROTOCOL_FIXTURES,
FakeStreamResponse,
ScriptedTransport,
)
# Every provider id in the catalogue — the parametrisation that makes this a
# contract suite rather than a per-adapter unit test.
PROVIDER_IDS = [d.provider_id for d in BUILTIN_DESCRIPTORS]
# Minimal config: enough for any adapter to build a URL and headers offline.
BASE_CONF = {"base_url": "https://gateway.test/v1", "api_key": "test-key"}
SAMPLE_MESSAGES = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Say hello"},
]
SAMPLE_TOOL = ToolSpec(
name="read_file",
description="Read a file from disk",
parameters={"type": "object", "properties": {"path": {"type": "string"}}},
)
@pytest.fixture()
def registry() -> ProviderRegistry:
"""A private registry per test so registrations never leak between tests."""
return ProviderRegistry(BUILTIN_DESCRIPTORS)
@pytest.fixture()
def collected_usage(monkeypatch) -> usage_sink.InMemoryUsageSink:
"""Swap the process-wide telemetry sink for an in-memory one.
Restored by monkeypatch after each test, so a contract run never appends to
the developer's real ``~/.cowork_local/usage/`` files.
"""
sink = usage_sink.InMemoryUsageSink()
monkeypatch.setattr(usage_sink, "_sink", usage_sink.CompositeUsageSink([sink]))
return sink
def _fixtures_for(registry: ProviderRegistry, provider_id: str) -> dict:
"""The stream/model-list script matching this provider's wire protocol."""
protocol = registry.get(provider_id).wire_protocol.value
return PROTOCOL_FIXTURES[protocol]
def _build(registry: ProviderRegistry, provider_id: str, transport=None) -> Provider:
"""Build a provider and (optionally) replace its transport with a script."""
provider = registry.build(provider_id, dict(BASE_CONF))
if transport is not None:
# Patch the INSTANCE, not the class: parallel parametrised cases must
# not see each other's scripted transport.
provider._request = transport
return provider
# --------------------------------------------------------------------------- #
# Construction & interface shape
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_registry_builds_a_provider_for_every_registered_id(registry, provider_id) -> None:
"""Every catalogued provider must be constructible — a descriptor with no
working adapter is a broken entry, not a feature flag."""
provider = _build(registry, provider_id)
assert isinstance(provider, Provider)
# The registry fills in the descriptor's default model when config omits it,
# so a half-configured provider still names a concrete model.
assert provider.model, f"{provider_id} built without a model id"
assert provider.describe() == f"{provider.name}:{provider.model}"
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_chat_signature_is_uniform(registry, provider_id) -> None:
"""All adapters accept the same call, so the agent runtime can swap
providers without knowing which one it holds."""
import inspect
provider = _build(registry, provider_id)
params = list(inspect.signature(provider.chat).parameters)
assert params == ["messages", "tools", "on_text", "cancel", "on_reasoning"]
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_tool_schema_matches_the_wire_protocol(registry, provider_id) -> None:
"""A ToolSpec must translate into the exact shape the endpoint expects."""
descriptor = registry.get(provider_id)
if descriptor.wire_protocol.value == "anthropic":
translated = SAMPLE_TOOL.to_anthropic()
assert translated["input_schema"] == SAMPLE_TOOL.parameters
assert translated["name"] == "read_file"
else:
translated = SAMPLE_TOOL.to_openai()
assert translated["type"] == "function"
assert translated["function"]["parameters"] == SAMPLE_TOOL.parameters
# --------------------------------------------------------------------------- #
# The turn itself
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_chat_returns_the_canonical_assistant_message(registry, provider_id, collected_usage) -> None:
"""Whatever the wire format, one turn yields the same canonical result."""
fixtures = _fixtures_for(registry, provider_id)
response = FakeStreamResponse(lines=fixtures["stream_lines"]())
transport = ScriptedTransport([response])
provider = _build(registry, provider_id, transport)
streamed: list = []
result = provider.chat(
SAMPLE_MESSAGES, tools=[SAMPLE_TOOL], on_text=streamed.append,
)
assert result["role"] == "assistant"
assert result["content"] == EXPECTED_TEXT
# Text must arrive incrementally, not only in the final message — the chat
# UI streams from these callbacks.
assert "".join(streamed) == EXPECTED_TEXT
assert len(streamed) >= 2
# Tool calls are normalised: parsed arguments, never the raw JSON fragments.
assert result["tool_calls"] == [EXPECTED_TOOL_CALL]
assert response.closed, "provider left the streaming response open"
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_chat_publishes_exactly_one_usage_event(registry, provider_id, collected_usage) -> None:
"""R03-T06: a turn reports its token usage through the telemetry sink, with
the server's real counts when the stream carried them."""
fixtures = _fixtures_for(registry, provider_id)
transport = ScriptedTransport([FakeStreamResponse(lines=fixtures["stream_lines"]())])
provider = _build(registry, provider_id, transport)
provider.chat(SAMPLE_MESSAGES, tools=[SAMPLE_TOOL])
events = collected_usage.snapshot()
assert len(events) == 1, "a turn must publish exactly one usage event"
event = events[0]
assert event.provider == provider.name
assert event.model == provider.model
assert event.input_tokens == EXPECTED_INPUT_TOKENS
assert event.output_tokens == EXPECTED_OUTPUT_TOKENS
assert event.cached_tokens == EXPECTED_CACHED_TOKENS
# Real counts were available, so the event must NOT be flagged as a guess.
assert event.estimated is False
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_usage_is_estimated_when_the_stream_reports_none(registry, provider_id, collected_usage) -> None:
"""Gateways that never send usage still produce a dashboard row — clearly
flagged as an estimate rather than silently recorded as zero."""
# Only text; no usage block anywhere in the stream.
silent_stream = ['data: ' + '{"choices": [{"delta": {"content": "hi"}}]}', "data: [DONE]"]
if registry.get(provider_id).wire_protocol.value == "anthropic":
silent_stream = [
'data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text"}}',
'data: {"type": "content_block_delta", "index": 0,'
' "delta": {"type": "text_delta", "text": "hi"}}',
]
transport = ScriptedTransport([FakeStreamResponse(lines=silent_stream)])
provider = _build(registry, provider_id, transport)
provider.chat(SAMPLE_MESSAGES)
events = collected_usage.snapshot()
assert len(events) == 1
assert events[0].estimated is True
# An estimate still has to be a positive number to be worth showing.
assert events[0].total_tokens > 0
# --------------------------------------------------------------------------- #
# Failure behaviour
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_http_error_becomes_provider_error(registry, provider_id, collected_usage) -> None:
"""Callers handle exactly one exception type; adapters must not leak
transport- or JSON-level errors past their boundary."""
failing = FakeStreamResponse(status_code=401, body='{"error": {"message": "bad key"}}')
transport = ScriptedTransport([failing])
provider = _build(registry, provider_id, transport)
with pytest.raises(ProviderError):
provider.chat(SAMPLE_MESSAGES)
assert failing.closed, "provider left a failed response open"
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_list_models_and_test_connection_report_a_reason(registry, provider_id) -> None:
"""A failed model load must explain itself: ``last_error`` is what Settings
shows instead of an unexplained empty dropdown."""
def _boom(*_args, **_kwargs):
# A transport failure, i.e. what actually happens when the gateway is
# unreachable — adapters translate this class of error, not arbitrary
# programming errors, which must still surface as bugs.
raise requests.ConnectionError("network down")
provider = _build(registry, provider_id, _boom)
models = provider.list_models()
assert provider.last_error, f"{provider_id} swallowed a model-load failure"
ok, message = provider.test_connection()
assert ok is False
assert message
# Anthropic answers with a built-in fallback catalogue; a gateway answers
# with nothing. Both are acceptable — the contract is only that a failure is
# never reported as success.
assert isinstance(models, list)
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_list_models_returns_ids_on_success(registry, provider_id) -> None:
"""The happy path returns plain model-id strings, not raw API objects."""
fixtures = _fixtures_for(registry, provider_id)
transport = ScriptedTransport([
FakeStreamResponse(status_code=200, payload=fixtures["models_payload"]),
])
provider = _build(registry, provider_id, transport)
models = provider.list_models()
assert models == fixtures["expected_models"]
assert provider.last_error == ""
assert all(isinstance(m, str) for m in models)
@pytest.mark.parametrize("provider_id", PROVIDER_IDS)
def test_strip_think_removes_inline_reasoning(registry, provider_id) -> None:
"""Reasoning must never leak into a final answer, whichever adapter ran."""
provider = _build(registry, provider_id)
cleaned = provider.strip_think("<think>secret plan</think>Visible answer")
assert cleaned == "Visible answer"
+7
View File
@@ -0,0 +1,7 @@
"""Integration tests: several real layers wired together, still fully offline.
Where unit tests pin one class against fakes and contract tests pin an interface
across implementations, these exercise a real path end to end — e.g. the
application routing service on top of the real ``core/routing`` engine — so a
seam that only works against a mock is caught here.
"""
@@ -0,0 +1,249 @@
"""R03-T03/T04/T05 — the unified routing path over the REAL routing engine.
The unit tests drive ``RoutingApplicationService`` against fakes; this suite
proves the same service produces correct outcomes on top of the actual
``core/routing`` stack (classifier → assessment store → scorer → selector →
switch controller), which is what the three chat surfaces now call.
Offline by construction: a fake probe client answers benchmarks and judging, and
the assessment store is a temp file — no network, no Qt, no ``$HOME`` writes.
"""
from __future__ import annotations
import copy
import pytest
from cowork_local.application.model_routing import (
AppContextModeResolver,
CoreRoutingEngine,
RoutingApplicationService,
RoutingMode,
RoutingRequest,
)
from cowork_local.config import DEFAULT_CONFIG, AppConfig
from cowork_local.core import projects as projects_mod
from cowork_local.core.routing.clients import CompletionResult
from cowork_local.core.routing.service import RoutingService
from cowork_local.core.routing.store import AssessmentStore
from cowork_local.state import AppContext
STRONG_ANSWER = "STRONG-DETAILED-CORRECT-ANSWER"
WEAK_ANSWER = "weak"
class FakeProbeClient:
"""Deterministic stand-in for the provider layer used during assessment.
Mirrors ``tests/routing/test_service.py``'s client: benchmark prompts get a
per-model canned answer, and judge prompts are graded by looking up that
answer, so scores are stable and no model is ever really called.
"""
def __init__(self, answers, quality) -> None:
self.answers = answers
self.quality = quality
def complete(self, provider, model_id, messages) -> CompletionResult:
text = messages[0]["content"]
if "grading an AI assistant" in text: # the judge rubric prompt
score = 0.0
for answer, value in self.quality.items():
if answer and answer in text:
score = value
break
return CompletionResult(text='{"score": %s}' % score)
answer = self.answers.get((provider, model_id))
if answer is None:
return CompletionResult(error="unavailable")
return CompletionResult(text=answer, tokens_out=len(answer) // 4)
@pytest.fixture()
def ctx(tmp_path, monkeypatch):
"""An AppContext with two assessable models and temp-only persistence."""
# Keep workspace load/save off the developer's real ~/.cowork_local.
monkeypatch.setattr(projects_mod, "PROJECTS_DIR", tmp_path / "projects")
data = copy.deepcopy(DEFAULT_CONFIG)
data["providers"] = {
"anthropic": {"base_url": "x", "api_key": "x", "model": "strong-model"},
}
data["routing"]["candidates"] = [
{"provider": "anthropic", "model_id": "strong-model", "tier": "powerful"},
{"provider": "anthropic", "model_id": "weak-model", "tier": "fast"},
]
data["routing"]["judge_provider"] = "anthropic"
data["routing"]["judge_model"] = "judge-model"
data["routing"]["policy"] = "quality"
data["routing"]["min_score_gain"] = 0.05
return AppContext(AppConfig(data=data, path=tmp_path / "config.json"))
@pytest.fixture()
def routing_service(ctx, tmp_path) -> RoutingService:
"""A real RoutingService with a populated assessment store."""
client = FakeProbeClient(
answers={
("anthropic", "strong-model"): STRONG_ANSWER,
("anthropic", "weak-model"): WEAK_ANSWER,
},
quality={STRONG_ANSWER: 0.95, WEAK_ANSWER: 0.35},
)
store = AssessmentStore(store_path=tmp_path / "assess.json",
history_dir=tmp_path / "history")
service = RoutingService(ctx, store=store, client=client)
service.reassess() # populate real probe results + fit scores
return service
@pytest.fixture()
def app_service(ctx, routing_service) -> RoutingApplicationService:
"""The application service wired exactly the way the UI wires it."""
return RoutingApplicationService(
CoreRoutingEngine(routing_service),
AppContextModeResolver(ctx),
confirm_timeout_sec=lambda: float(ctx.config.routing["confirm_timeout_sec"]),
)
def coding_request(**overrides) -> RoutingRequest:
"""A coding turn currently pinned to the weaker model."""
fields = dict(
surface="cowork",
prompt="Write a Python function to reverse a linked list",
current_provider="anthropic",
current_model="weak-model",
)
fields.update(overrides)
return RoutingRequest(**fields)
# --------------------------------------------------------------------------- #
# Auto / Off / Manual over the real engine
# --------------------------------------------------------------------------- #
def test_auto_switches_to_the_better_assessed_model(app_service) -> None:
"""The real scorer must rank the strong model first and the service must
hand that model back as this turn's override."""
outcome = app_service.resolve(coding_request(mode=RoutingMode.AUTO))
assert outcome.switched is True
assert outcome.provider == "anthropic"
assert outcome.model == "strong-model"
assert outcome.task_type == "coding" # classified from the prompt
assert outcome.score_gain > 0
def test_off_keeps_the_pinned_model(app_service) -> None:
"""Off must not switch even when a clearly better model is assessed."""
outcome = app_service.resolve(coding_request(mode=RoutingMode.OFF))
assert outcome.switched is False
assert outcome.provider is None
def test_manual_asks_before_switching(app_service) -> None:
"""The confirm callback receives the engine's own decision object, which is
what ``ui/routing_toggle.py::confirm_switch`` renders."""
seen: list = []
outcome = app_service.resolve(
coding_request(mode=RoutingMode.MANUAL),
confirm=lambda decision, timeout: seen.append((decision, timeout)) or True,
)
assert outcome.switched is True
decision, timeout = seen[0]
assert decision.to_model == "anthropic/strong-model"
assert decision.reason # human-readable explanation
assert timeout == pytest.approx(60.0) # from DEFAULT_CONFIG
def test_manual_decline_keeps_the_pinned_model(app_service) -> None:
outcome = app_service.resolve(
coding_request(mode=RoutingMode.MANUAL),
confirm=lambda decision, timeout: False,
)
assert outcome.switched is False
assert outcome.declined is True
def test_already_best_model_is_left_alone(app_service) -> None:
"""No pointless churn: being on the best model is not a switch."""
outcome = app_service.resolve(
coding_request(mode=RoutingMode.AUTO, current_model="strong-model"))
assert outcome.switched is False
# --------------------------------------------------------------------------- #
# Fallback over the real engine
# --------------------------------------------------------------------------- #
def test_fallback_keeps_an_assessed_model_even_though_a_better_one_exists(app_service) -> None:
"""weak-model IS usable (it has a real probe score), so Fallback stays put
where Auto would switch — the behavioural difference between the modes."""
outcome = app_service.resolve(coding_request(mode=RoutingMode.FALLBACK))
assert outcome.switched is False
def test_fallback_rescues_a_model_the_engine_cannot_serve(app_service) -> None:
"""A model absent from the ranking (never assessed / unavailable) is exactly
the situation Fallback exists for."""
outcome = app_service.resolve(
coding_request(mode=RoutingMode.FALLBACK, current_model="ghost-model"))
assert outcome.switched is True
assert outcome.model == "strong-model"
# --------------------------------------------------------------------------- #
# Surface parity — the point of R03-T04/T05
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("surface", ["cowork", "co4e", "ai_edit"])
def test_every_surface_gets_the_same_decision(app_service, surface) -> None:
"""Chat, Co4E and AI-Edit used to hold three copies of this logic. Given the
same inputs they must now be indistinguishable."""
outcome = app_service.resolve(coding_request(surface=surface, mode=RoutingMode.AUTO))
assert outcome.switched is True
assert outcome.model == "strong-model"
def test_ai_edit_pinned_task_type_reaches_the_engine(app_service) -> None:
"""AI-Edit pins "coding" instead of classifying; the engine must honour it
even when the instruction text reads like something else entirely."""
outcome = app_service.resolve(coding_request(
surface="ai_edit",
prompt="Write a poem about the ocean", # classifier would say "creative"
task_type="coding",
mode=RoutingMode.AUTO,
))
assert outcome.task_type == "coding"
def test_mode_comes_from_the_workspace_when_not_pinned(ctx, app_service) -> None:
"""With no explicit mode, the service reads the per-workspace setting — the
lookup the widgets used to do themselves."""
ctx.config.data["routing"]["switch_mode"] = "auto"
outcome = app_service.resolve(coding_request())
assert outcome.mode is RoutingMode.AUTO
assert outcome.switched is True
def test_fallback_mode_survives_a_round_trip_through_config(ctx) -> None:
"""The new mode must be persistable, or the toggle could never select it."""
ctx.config.set_routing_mode_for("cowork", "fallback")
assert ctx.config.routing_mode_for("cowork") == "fallback"
assert ctx.project_routing_mode("cowork") == "fallback"
def test_unknown_persisted_mode_degrades_to_off(ctx) -> None:
"""A hand-edited config must not enable routing by accident."""
ctx.config.routing["surface_modes"]["cowork"] = "turbo"
assert ctx.config.routing_mode_for("cowork") == "off"
+5 -13
View File
@@ -1,17 +1,9 @@
"""Pytest fixtures/shared helpers for the routing test suite.
Ensures the ``cowork_local`` package is importable when pytest is invoked from
the package directory itself (so ``import cowork_local.core.routing...`` works
regardless of the working directory the suite is launched from).
Package importability is handled once and for all by ``tests/conftest.py``,
which binds THIS checkout to the ``cowork_local`` name in ``sys.modules``.
This file used to push the checkout's PARENT directory onto ``sys.path``, which
let an unrelated sibling folder named ``cowork_local`` shadow the working copy —
so that logic is intentionally gone; keep it that way.
"""
from __future__ import annotations
import sys
from pathlib import Path
# .../cowork_local/tests/routing/conftest.py → parent of the package dir
_PKG_DIR = Path(__file__).resolve().parents[2] # .../cowork_local
_REPO_ROOT = _PKG_DIR.parent # .../cowork_local_20260722
for p in (str(_REPO_ROOT), str(_PKG_DIR)):
if p not in sys.path:
sys.path.insert(0, p)
+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)