Files
2026-07-13 13:12:33 +08:00

292 lines
8.6 KiB
Python

from __future__ import annotations
from collections.abc import AsyncIterator
from typing import Any
import pytest
from opensquilla.cli.chat.turn_stream import (
default_turn_stream_dependencies,
stream_response_gateway,
stream_response_turnrunner,
)
from opensquilla.cli.tui.adapters.turn_stream_defaults import router_hud_event_sink_factory
from opensquilla.cli.tui.backend.domain_events import (
KIND_DONE,
KIND_ROUTER_DECISION,
KIND_TEXT_FLUSH,
TuiDomainEvent,
now_ms,
)
from opensquilla.cli.tui.backend.plugins import TuiPluginManager
from opensquilla.cli.tui.plugins.router_hud import (
ROUTER_HUD_SLOT,
RouterHudPlugin,
RouterHudSnapshot,
)
from opensquilla.engine.types import DoneEvent, RouterDecisionEvent
from opensquilla.tools.types import CallerKind, ToolContext
ROUTER_PAYLOAD: dict[str, object] = {
"tier": "t2",
"tier_index": 2,
"model": "anthropic/claude-sonnet-4.6",
"baseline_model": "anthropic/claude-opus-4.7",
"source": "router",
"confidence": 0.71,
"probs": [0.1, 0.2, 0.7],
"savings_pct": 64.0,
"fallback": False,
"thinking_mode": "balanced",
"prompt_policy": "default",
"routing_applied": True,
"rollout_phase": "full",
"context_window": 200_000,
}
class _GatewayClient:
def __init__(self, events: list[dict[str, Any]]) -> None:
self._events = events
async def send_message(self, *_args: Any, **_kwargs: Any):
for event in self._events:
yield event
async def resolve_approval(self, *_args: Any, **_kwargs: Any) -> None:
return None
async def abort_session(self, _key: str) -> None:
return None
class _TurnRunner:
def __init__(self, events: list[object]) -> None:
self._events = events
async def run(self, *_args: Any, **_kwargs: Any) -> AsyncIterator[object]:
for event in self._events:
yield event
def _router_event(payload: dict[str, object]) -> TuiDomainEvent:
return TuiDomainEvent(
kind=KIND_ROUTER_DECISION,
source="gateway",
payload=payload,
turn_id="agent:main:test",
timestamp_ms=now_ms(),
)
def _snapshot_for(payload: dict[str, object]) -> RouterHudSnapshot:
plugin = RouterHudPlugin()
plugin.on_event(_router_event(payload), context=object())
snapshot = plugin.snapshot(ROUTER_HUD_SLOT)
assert isinstance(snapshot, RouterHudSnapshot)
return snapshot
@pytest.mark.asyncio
async def test_gateway_stream_surfaces_normalized_router_decision_domain_event() -> None:
events: list[TuiDomainEvent] = []
deps = default_turn_stream_dependencies(tui_event_sink=events.append)
await stream_response_gateway(
_GatewayClient(
[
{"event": "session.event.router_decision", **ROUTER_PAYLOAD},
{"event": "session.event.done", "model": "anthropic/claude-sonnet-4.6"},
]
),
"agent:main:test",
"hello",
deps=deps,
)
assert [event.kind for event in events] == [KIND_ROUTER_DECISION, KIND_DONE]
assert events[0].source == "gateway"
assert events[0].payload == ROUTER_PAYLOAD
@pytest.mark.asyncio
async def test_turnrunner_stream_surfaces_matching_router_decision_domain_event(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router_event = RouterDecisionEvent(**ROUTER_PAYLOAD)
turn_runner = _TurnRunner([router_event, DoneEvent(model="anthropic/claude-sonnet-4.6")])
monkeypatch.setattr("opensquilla.engine.runtime.TurnRunner", _TurnRunner)
events: list[TuiDomainEvent] = []
deps = default_turn_stream_dependencies(
stream_wrapper=lambda stream, _svc: stream,
tui_event_sink=events.append,
)
tool_ctx = ToolContext(
caller_kind=CallerKind.CLI,
channel_kind="cli",
channel_id="cli:chat",
)
await stream_response_turnrunner(
turn_runner,
"agent:main:test",
tool_ctx,
"hello",
deps=deps,
)
assert [event.kind for event in events] == [KIND_ROUTER_DECISION, KIND_DONE]
assert events[0].source == "turn_runner"
assert events[0].payload == ROUTER_PAYLOAD
def test_router_hud_formats_full_route_label_and_snapshot_fields() -> None:
snapshot = _snapshot_for(ROUTER_PAYLOAD)
assert snapshot.tier == "t2"
assert snapshot.tier_index == 2
assert snapshot.model == "anthropic/claude-sonnet-4.6"
assert snapshot.baseline_model == "anthropic/claude-opus-4.7"
assert snapshot.confidence == 0.71
assert snapshot.savings_pct == 64.0
assert snapshot.label == "route t2 -> claude-sonnet-4.6 71% save 64%"
assert snapshot.style == "normal"
def test_router_hud_uses_natural_tier_indexes_and_observe_style() -> None:
payload = {
**ROUTER_PAYLOAD,
"tier": "t1",
"tier_index": None,
"routing_applied": False,
"rollout_phase": "observe",
}
snapshot = _snapshot_for(payload)
assert snapshot.tier_index == 1
assert snapshot.label == "observe t1 -> claude-sonnet-4.6 71%"
assert snapshot.style == "dim"
def test_router_hud_highlights_fallback_routes() -> None:
snapshot = _snapshot_for(
{
**ROUTER_PAYLOAD,
"source": "fallback",
"fallback": True,
"savings_pct": 0.0,
}
)
assert snapshot.fallback is True
assert snapshot.label == "fallback -> claude-sonnet-4.6"
assert snapshot.style == "warning"
def test_router_hud_formats_forced_and_no_baseline_without_savings() -> None:
snapshot = _snapshot_for(
{
**ROUTER_PAYLOAD,
"tier": "t0",
"tier_index": 0,
"source": "forced",
"baseline_model": "",
"savings_pct": 64.0,
}
)
assert snapshot.tier_index == 0
assert snapshot.label == "forced t0 -> claude-sonnet-4.6 71%"
assert "save" not in snapshot.label
def test_router_hud_omits_malformed_confidence() -> None:
snapshot = _snapshot_for(
{
**ROUTER_PAYLOAD,
"confidence": "not-a-number",
"savings_pct": "also-bad",
}
)
assert snapshot.confidence is None
assert snapshot.savings_pct is None
assert snapshot.label == "route t2 -> claude-sonnet-4.6"
def test_opentui_router_sink_updates_toolbar_only_for_router_events() -> None:
class _Output:
def __init__(self) -> None:
self.updates: list[tuple[str, object | None]] = []
self.invalidations = 0
def set_toolbar(self, key: str, value: object | None) -> None:
self.updates.append((key, value))
def invalidate(self) -> None:
self.invalidations += 1
output = _Output()
sink = router_hud_event_sink_factory(output)
sink(_router_event(ROUTER_PAYLOAD))
after_router = list(output.updates)
sink(
TuiDomainEvent(
kind=KIND_TEXT_FLUSH,
source="renderer",
payload={"text": "hello"},
turn_id="agent:main:test",
timestamp_ms=now_ms(),
)
)
assert ("router_hud", "route t2 -> claude-sonnet-4.6 71% save 64%") in after_router
assert ("router_hud_style", "normal") in after_router
assert output.updates == after_router
assert output.invalidations == 1
def test_opentui_router_sink_writes_context_window_from_payload() -> None:
class _Output:
def __init__(self) -> None:
self.updates: list[tuple[str, object | None]] = []
def set_toolbar(self, key: str, value: object | None) -> None:
self.updates.append((key, value))
def invalidate(self) -> None:
return None
output = _Output()
sink = router_hud_event_sink_factory(output)
sink(_router_event({**ROUTER_PAYLOAD, "context_window": 1_000_000}))
assert ("router_context_window", 1_000_000) in output.updates
def test_opentui_router_sink_reuses_launch_scoped_plugin_manager() -> None:
class _Output:
def __init__(self, plugin_manager: TuiPluginManager) -> None:
self.plugin_manager = plugin_manager
self.updates: list[tuple[str, object | None]] = []
def set_toolbar(self, key: str, value: object | None) -> None:
self.updates.append((key, value))
def invalidate(self) -> None:
return None
manager = TuiPluginManager([RouterHudPlugin()])
output = _Output(manager)
sink = router_hud_event_sink_factory(output)
sink(_router_event(ROUTER_PAYLOAD))
snapshot = manager.snapshot(ROUTER_HUD_SLOT)
assert isinstance(snapshot, RouterHudSnapshot)
assert snapshot.label == "route t2 -> claude-sonnet-4.6 71% save 64%"
assert ("router_hud", snapshot.label) in output.updates