wehub-resource-sync
127 lines
4.2 KiB
Python
127 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from opensquilla.gateway.config import SquillaRouterConfig, _router_tier_profile_defaults
|
|
from opensquilla.router_control import RouterControlHoldStore
|
|
from opensquilla.tool_boundary import ToolCall
|
|
from opensquilla.tools import get_default_registry
|
|
from opensquilla.tools.dispatch import build_tool_handler
|
|
from opensquilla.tools.types import CallerKind, ToolContext
|
|
|
|
|
|
def _ctx(*, replay_depth: int = 0, hold_applied: bool = False) -> ToolContext:
|
|
return ToolContext(
|
|
is_owner=True,
|
|
caller_kind=CallerKind.AGENT,
|
|
session_key="agent:main:test-router-control-tool",
|
|
router_control_config=SquillaRouterConfig(
|
|
enabled=True,
|
|
rollout_phase="full",
|
|
require_router_runtime=False,
|
|
tiers=_router_tier_profile_defaults("openrouter"),
|
|
),
|
|
router_control_hold_store=RouterControlHoldStore(),
|
|
router_control_replay_depth=replay_depth,
|
|
router_control_turn_hold_applied=hold_applied,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_control_set_hold_writes_store_and_requests_replay() -> None:
|
|
ctx = _ctx()
|
|
handler = build_tool_handler(get_default_registry(), ctx)
|
|
|
|
result = await handler(
|
|
ToolCall(
|
|
tool_use_id="call-1",
|
|
tool_name="router_control",
|
|
arguments={
|
|
"action": "set_hold",
|
|
"target_id": "tier:c3",
|
|
"evidence": "use c3",
|
|
},
|
|
)
|
|
)
|
|
|
|
payload = json.loads(result.content)
|
|
assert result.is_error is False
|
|
assert result.terminates_turn is True
|
|
assert payload["accepted"] is True
|
|
assert payload["target_tier"] == "c3"
|
|
assert payload["target_model"] == "anthropic/claude-opus-4.8"
|
|
assert payload["replay_required"] is True
|
|
hold = ctx.router_control_hold_store.get_valid(ctx.session_key or "")
|
|
assert hold is not None
|
|
assert hold.tier == "c3"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_control_rejects_alias_without_mutating_hold() -> None:
|
|
ctx = _ctx()
|
|
handler = build_tool_handler(get_default_registry(), ctx)
|
|
|
|
result = await handler(
|
|
ToolCall(
|
|
tool_use_id="call-1",
|
|
tool_name="router_control",
|
|
arguments={
|
|
"action": "set_hold",
|
|
"target_id": "Claude Opus 4.8",
|
|
"evidence": "use Claude Opus 4.8",
|
|
},
|
|
)
|
|
)
|
|
|
|
payload = json.loads(result.content)
|
|
assert result.is_error is False
|
|
assert result.terminates_turn is False
|
|
assert payload["accepted"] is False
|
|
assert payload["replay_required"] is False
|
|
assert ctx.router_control_hold_store.get_valid(ctx.session_key or "") is None
|
|
|
|
|
|
def test_router_control_tool_schema_includes_dynamic_target_enum() -> None:
|
|
ctx = _ctx()
|
|
definitions = get_default_registry().to_tool_definitions(ctx)
|
|
router_tool = next(tool for tool in definitions if tool.name == "router_control")
|
|
|
|
target_schema = router_tool.input_schema.properties["target_id"]
|
|
|
|
assert "enum" in target_schema
|
|
assert "tier:c3" in target_schema["enum"]
|
|
assert "tier:t3" not in target_schema["enum"]
|
|
assert "model:z-ai/glm-5.2" not in target_schema["enum"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_control_clear_hold_replays_when_hold_already_selected_turn() -> None:
|
|
ctx = _ctx(hold_applied=True)
|
|
target = next(
|
|
target
|
|
for target in ctx.router_control_hold_store.build_targets(ctx.router_control_config)
|
|
if target.target_id == "tier:c3"
|
|
)
|
|
ctx.router_control_hold_store.set_hold(ctx.session_key or "", target, evidence="use c3")
|
|
handler = build_tool_handler(get_default_registry(), ctx)
|
|
|
|
result = await handler(
|
|
ToolCall(
|
|
tool_use_id="call-1",
|
|
tool_name="router_control",
|
|
arguments={
|
|
"action": "clear_hold",
|
|
"evidence": "back to automatic routing",
|
|
},
|
|
)
|
|
)
|
|
|
|
payload = json.loads(result.content)
|
|
assert result.terminates_turn is True
|
|
assert payload["accepted"] is True
|
|
assert payload["action"] == "clear_hold"
|
|
assert payload["replay_required"] is True
|
|
assert ctx.router_control_hold_store.get_valid(ctx.session_key or "") is None
|