mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
Renn F
2089b9e765
Claude Code's per-agent token-budget hook fires against the SDK :9000 server; opencode exposes NO usage/budget hook to a plugin (confirmed against its plugin docs), so the budget kill-switch can't be a plugin/sidecar — the orchestrator enforces it instead. _enforce_grok_cost_budget runs each dispatch tick: for every ACTIVE GROK container it reads cumulative cost from the opencode store (the Phase-2 reader) and kills + evicts it past ROBOCO_GROK_MAX_COST_USD (0 = off), after which the reaper releases the freed task. This also catches a runaway loop that keeps firing verbs (so it evades the idle watchdog) but still burns cost. Covers the budget/runaway-burn slice of guardrail parity. The remaining Claude hooks (prompt-injection PRE-gate, stop-guard terminal-verb) have no blocking opencode equivalent — opencode's message/stop hooks are observe-only — and the interactive reasoning-variant has no opencode.json/serve knob (CLI-flag only); both are pinned for a live probe rather than shipped as a guess.
84 lines
3.1 KiB
Python
84 lines
3.1 KiB
Python
"""GROK cost budget kill-switch: kill a live container over the cost ceiling.
|
|
|
|
opencode exposes no usage hook to a plugin, so the budget kill-switch lives in
|
|
the orchestrator: it reads each live GROK container's cumulative opencode cost
|
|
and kills + evicts it past ROBOCO_GROK_MAX_COST_USD (also catching runaway-loop
|
|
token burn). The cost computation itself is covered in opencode_usage tests; here
|
|
cost_for_session is stubbed so the kill DECISION is exercised deterministically.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from roboco.models.runtime import AgentInstance
|
|
from roboco.runtime.orchestrator import AgentOrchestrator, AgentState
|
|
|
|
_COST_FN = "roboco.llm.providers.opencode_usage.cost_for_session"
|
|
|
|
|
|
def _grok_instance(provider_type: str = "grok") -> AgentInstance:
|
|
cfg = type("C", (), {"provider_type": provider_type, "model": "grok-build-0.1"})()
|
|
return AgentInstance(agent_id="be-dev-1", state=AgentState.ACTIVE, config=cfg)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_over_cap_kills_and_evicts(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
orch._grok_max_cost_usd = 5.0
|
|
orch._instances = {"be-dev-1": _grok_instance()}
|
|
remove_mock = AsyncMock()
|
|
monkeypatch.setattr(orch, "_remove_container", remove_mock)
|
|
monkeypatch.setattr(_COST_FN, lambda *_a, **_k: (None, 7.5))
|
|
|
|
await orch._enforce_grok_cost_budget()
|
|
|
|
remove_mock.assert_awaited_once_with("roboco-agent-be-dev-1")
|
|
assert "be-dev-1" not in orch._instances
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_under_cap_spares(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
orch._grok_max_cost_usd = 5.0
|
|
orch._instances = {"be-dev-1": _grok_instance()}
|
|
remove_mock = AsyncMock()
|
|
monkeypatch.setattr(orch, "_remove_container", remove_mock)
|
|
monkeypatch.setattr(_COST_FN, lambda *_a, **_k: (None, 1.0))
|
|
|
|
await orch._enforce_grok_cost_budget()
|
|
|
|
remove_mock.assert_not_awaited()
|
|
assert "be-dev-1" in orch._instances
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cap_zero_disables_the_sweep(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
orch._grok_max_cost_usd = 0.0
|
|
orch._instances = {"be-dev-1": _grok_instance()}
|
|
remove_mock = AsyncMock()
|
|
monkeypatch.setattr(orch, "_remove_container", remove_mock)
|
|
monkeypatch.setattr(_COST_FN, lambda *_a, **_k: (None, 999.0))
|
|
|
|
await orch._enforce_grok_cost_budget()
|
|
|
|
remove_mock.assert_not_awaited()
|
|
assert "be-dev-1" in orch._instances
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_grok_container_is_ignored(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
orch._grok_max_cost_usd = 5.0
|
|
orch._instances = {"be-dev-1": _grok_instance(provider_type="anthropic")}
|
|
remove_mock = AsyncMock()
|
|
monkeypatch.setattr(orch, "_remove_container", remove_mock)
|
|
monkeypatch.setattr(_COST_FN, lambda *_a, **_k: (None, 999.0))
|
|
|
|
await orch._enforce_grok_cost_budget()
|
|
|
|
remove_mock.assert_not_awaited()
|
|
assert "be-dev-1" in orch._instances
|