diff --git a/scripts/analyze_spread.py b/scripts/analyze_spread.py index 8ca14e1..67488b4 100644 --- a/scripts/analyze_spread.py +++ b/scripts/analyze_spread.py @@ -122,6 +122,10 @@ def main() -> None: delta = [round(a - b, 4) for a, b in zip(traj, base_slice)] ticks_above = sum(1 for x in delta if x > 0) c = m["summary"]["cost_snapshot"] + # Billed ground truth (run_scenario reconciles via OpenRouter's + # generation API) — same schema we hand-patched into the v0.3.0 + # leaderboard from the dashboard. Estimates stay for comparison. + billed = m["summary"].get("billed_cost") rows.append({ "model": m["summary"]["model"], "name": name, @@ -140,6 +144,10 @@ def main() -> None: "avg_confidence": m["deliberation"]["avg_confidence"], "wall_min": round(c["wall_time_s"] / 60, 1), "est_cost_usd": round(c["estimated_cost_usd"], 2), + "real_cost_usd": ( + round(billed["billed_cost_usd"], 3) if billed else None + ), + "cost_source": billed["source"] if billed else None, "trajectory": traj, "days": days, "baseline_slice": [round(x, 4) for x in base_slice], diff --git a/scripts/run_benchmark.py b/scripts/run_benchmark.py index ae07e08..dedd3d9 100644 --- a/scripts/run_benchmark.py +++ b/scripts/run_benchmark.py @@ -6,6 +6,7 @@ import asyncio import json import logging +import os import random import time from pathlib import Path @@ -35,7 +36,7 @@ from rle.scoring.composite import CompositeScorer from rle.scoring.delta import PairedResult from rle.scoring.recorder import TimeSeriesRecorder -from rle.tracking.cost_tracker import CostTracker, create_cost_tracker +from rle.tracking.cost_tracker import CostTracker, create_cost_tracker, fetch_billed_costs from rle.tracking.event_log import EventLog from rle.tracking.hf_logger import HFLogger from rle.tracking.history import append_history, get_run_dir, update_baseline @@ -692,6 +693,19 @@ async def main(args: argparse.Namespace) -> None: else: _print_leaderboard(results, model=args.model) + # Reconcile estimates against OpenRouter's billed ground truth + # (token-count estimates diverged up to 4x on the v0.3.0 spread). + billed_report = None + effective_base_url = args.base_url or config.provider_base_url or "" + openai_key = os.environ.get("OPENAI_API_KEY", "") + generation_ids = cost_tracker.generation_ids + if "openrouter.ai" in effective_base_url and openai_key and generation_ids: + print( + f"\nReconciling billed cost for {len(generation_ids)} " + "generations against OpenRouter...", + ) + billed_report = await fetch_billed_costs(generation_ids, openai_key) + # Build enriched summary with metadata metadata = collect_metadata(random_seed=args.seed) summary: dict[str, Any] = { @@ -708,6 +722,8 @@ async def main(args: argparse.Namespace) -> None: "scenarios": results, "cost_snapshot": cost_tracker.snapshot().model_dump(), } + if billed_report: + summary["billed_cost"] = billed_report.model_dump() if event_log: summary["event_summary"] = event_log.summary().model_dump() if is_paired and paired_results: diff --git a/scripts/run_scenario.py b/scripts/run_scenario.py index a82ec74..5d8b370 100644 --- a/scripts/run_scenario.py +++ b/scripts/run_scenario.py @@ -6,6 +6,7 @@ import asyncio import json import logging +import os import random import sys from pathlib import Path @@ -31,7 +32,7 @@ from rle.scenarios.loader import list_scenarios, load_scenario from rle.scoring.composite import CompositeScorer from rle.scoring.recorder import TimeSeriesRecorder -from rle.tracking.cost_tracker import create_cost_tracker +from rle.tracking.cost_tracker import create_cost_tracker, fetch_billed_costs from rle.tracking.event_log import EventLog from rle.tracking.history import append_history from rle.tracking.metadata import collect_metadata @@ -78,6 +79,7 @@ def _build_run_summary( # noqa: PLR0913 ticks_run: int, cost_snapshot_dict: dict[str, object], event_summary_dict: dict[str, object] | None, + billed_cost_dict: dict[str, object] | None = None, ) -> dict[str, object]: """Compose the per-scenario summary JSON (metadata + config + result).""" summary: dict[str, object] = { @@ -98,6 +100,8 @@ def _build_run_summary( # noqa: PLR0913 "ticks_run": ticks_run, "cost_snapshot": cost_snapshot_dict, } + if billed_cost_dict is not None: + summary["billed_cost"] = billed_cost_dict if event_summary_dict is not None: summary["event_summary"] = event_summary_dict return summary @@ -316,6 +320,20 @@ async def main(args: argparse.Namespace) -> None: # Output _print_results(loop, recorder) + # Reconcile estimates against OpenRouter's billed ground truth — the + # token-count estimator diverged up to 4x in both directions on the + # v0.3.0 spread (thinking-model usage shapes, caching discounts). + billed_report = None + effective_base_url = args.base_url or config.provider_base_url or "" + openai_key = os.environ.get("OPENAI_API_KEY", "") + generation_ids = cost_tracker.generation_ids + if "openrouter.ai" in effective_base_url and openai_key and generation_ids: + print( + f"\nReconciling billed cost for {len(generation_ids)} generations " + "against OpenRouter...", + ) + billed_report = await fetch_billed_costs(generation_ids, openai_key) + if args.output: output_dir = Path(args.output) output_dir.mkdir(parents=True, exist_ok=True) @@ -350,6 +368,9 @@ async def main(args: argparse.Namespace) -> None: event_summary_dict=( event_log.summary().model_dump() if event_log else None ), + billed_cost_dict=( + billed_report.model_dump() if billed_report else None + ), ) summary_path = output_dir / f"{scenario_path.stem}_summary.json" summary_path.write_text(json.dumps(summary, indent=2, default=str)) @@ -382,6 +403,16 @@ async def main(args: argparse.Namespace) -> None: f"| Est. cost: ${snap.estimated_cost_usd:.4f} " f"| Wall time: {snap.wall_time_s:.1f}s", ) + if billed_report: + unbilled = ( + f" ({billed_report.missing_generations} unbilled — lower bound)" + if billed_report.missing_generations else "" + ) + print( + f"Billed cost (OpenRouter ground truth): " + f"${billed_report.billed_cost_usd:.4f} over " + f"{billed_report.billed_generations} generations{unbilled}", + ) if event_log: event_log.close() diff --git a/src/rle/agents/base_role.py b/src/rle/agents/base_role.py index 4f4704d..bc8f678 100644 --- a/src/rle/agents/base_role.py +++ b/src/rle/agents/base_role.py @@ -180,6 +180,10 @@ def __init__( self._pending_events: list[RimAPIEvent] = [] self._last_usage: dict[str, int] | None = None self._last_raw_output: str | None = None + # Provider generation IDs since the last drain — one per completion + # call (parse retries included), consumed by the game loop for + # billed-cost reconciliation against OpenRouter. + self._generation_ids: list[str] = [] self._weave_op: Any = None # Set via enable_weave() for LLM call tracing def set_provider_kwargs(self, **kwargs: Any) -> None: @@ -259,8 +263,37 @@ def _record_completion(self, result: CompletionResult) -> CompletionResult: if reasoning: usage["reasoning_tokens"] = reasoning self._last_usage = usage + gen_id = self._extract_generation_id(getattr(result, "raw_response", None)) + if gen_id is not None: + self._generation_ids.append(gen_id) return result + def drain_generation_ids(self) -> list[str]: + """Return and clear generation IDs accumulated since the last drain. + + Captures EVERY provider call — including parse retries and + deliberations that later failed to parse — because those bill tokens + too. The game loop drains this each tick into the cost tracker. + """ + ids, self._generation_ids = self._generation_ids, [] + return ids + + @staticmethod + def _extract_generation_id(raw_response: Any) -> str | None: + """Pull the provider's generation/completion ID from the raw response. + + OpenRouter's chat completion ``id`` (``gen-...``) keys its + ``/api/v1/generation`` billing endpoint. Handles both the OpenAI SDK + object shape and the plain dict shape. None when absent. + """ + if raw_response is None: + return None + if isinstance(raw_response, dict): + gen_id = raw_response.get("id") + else: + gen_id = getattr(raw_response, "id", None) + return gen_id if isinstance(gen_id, str) and gen_id else None + @staticmethod def _extract_reasoning_tokens(raw_response: Any) -> int: """Pull reasoning-token count from a provider's raw usage payload. diff --git a/src/rle/orchestration/game_loop.py b/src/rle/orchestration/game_loop.py index 1e9bf5a..8b36458 100644 --- a/src/rle/orchestration/game_loop.py +++ b/src/rle/orchestration/game_loop.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import contextlib import json import logging import time as _time @@ -86,6 +87,7 @@ def __init__( screenshots_enabled: bool = False, auto_dismiss_dialogs: bool = True, camera_director: CameraDirector | None = None, + speed_keepalive_s: float = 10.0, ) -> None: self._config = config self._client = client @@ -129,6 +131,7 @@ def __init__( self._screenshots_enabled = screenshots_enabled self._auto_dismiss_dialogs = auto_dismiss_dialogs self._camera_director = camera_director + self._speed_keepalive_s = speed_keepalive_s # Hub-spoke communication — agents read messages from their spokes self._hub = CentralPost(max_agents=len(agents)) @@ -619,6 +622,15 @@ async def run_tick(self) -> TickResult: }, ) + # Capture generation IDs from every provider call this tick — + # parse retries and failed deliberations bill tokens too — for + # billed-cost reconciliation against OpenRouter (the token-count + # estimator diverged up to 4x on the v0.3.0 spread). + if self._cost_tracker: + for any_agent in self._agents: + for gen_id in any_agent.drain_generation_ids(): + self._cost_tracker.record_generation_id(gen_id) + # Resolve conflicts resolved, resolver_stats = self._resolver.resolve(plans, state) self._metric_context.conflicts_total += resolver_stats.conflicts_total @@ -745,6 +757,24 @@ async def run_tick(self) -> TickResult: return result + async def _speed_keepalive(self) -> None: + """Re-assert game speed periodically (no-pause mode only). + + RimWorld force-pauses on threat letters (mad animal, raid). With + --no-pause the loop previously re-asserted speed only at tick + boundaries, so a slow model left the game frozen for its entire + deliberation window (minutes on kimi/qwen in the v0.3.0 spread) — + frozen footage and wildly non-uniform game-time per tick. Setting + the speed is idempotent, so this just fires every few seconds for + the whole run. Cancelled by run(). + """ + while True: + await asyncio.sleep(self._speed_keepalive_s) + try: + await self._client.unpause_game() + except Exception: + logger.debug("Speed keepalive failed", exc_info=True) + async def run(self, max_ticks: int | None = None) -> list[TickResult]: """Run the game loop for N ticks or until stopped.""" self._running = True @@ -756,26 +786,41 @@ async def run(self, max_ticks: int | None = None) -> list[TickResult]: "the colony-name dialog and dev debug log to need manual " "dismissal.", ) + keepalive: asyncio.Task[None] | None = None if self._no_pause: await self._client.unpause_game() + if self._speed_keepalive_s > 0: + keepalive = asyncio.create_task(self._speed_keepalive()) tick_count = 0 - while self._running: - result = await self.run_tick() - tick_count += 1 - score_str = "" - if result.score: - score_str = f" | score={result.score.composite:.3f}" - logger.info( - "Tick %d (day %d): %d actions, %d executed%s", - tick_count, - result.day, - result.execution.total, - result.execution.executed, - score_str, - ) - if max_ticks and tick_count >= max_ticks: - break - await asyncio.sleep(self._config.tick_interval) + try: + while self._running: + result = await self.run_tick() + tick_count += 1 + score_str = "" + if result.score: + score_str = f" | score={result.score.composite:.3f}" + logger.info( + "Tick %d (day %d): %d actions, %d executed%s", + tick_count, + result.day, + result.execution.total, + result.execution.executed, + score_str, + ) + if max_ticks and tick_count >= max_ticks: + break + await asyncio.sleep(self._config.tick_interval) + finally: + if keepalive is not None: + keepalive.cancel() + with contextlib.suppress(asyncio.CancelledError): + await keepalive + # Final drain: a deliberation that timed out on the last tick may + # have completed in its worker thread after the per-tick drain. + if self._cost_tracker: + for agent in self._agents: + for gen_id in agent.drain_generation_ids(): + self._cost_tracker.record_generation_id(gen_id) return self._tick_results def stop(self) -> None: diff --git a/src/rle/tracking/cost_tracker.py b/src/rle/tracking/cost_tracker.py index 03ae790..8ffdc31 100644 --- a/src/rle/tracking/cost_tracker.py +++ b/src/rle/tracking/cost_tracker.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import logging import time @@ -11,6 +12,7 @@ logger = logging.getLogger(__name__) OPENROUTER_MODELS_URL = "https://openrouter.ai/api/v1/models" +OPENROUTER_GENERATION_URL = "https://openrouter.ai/api/v1/generation" class TokenUsage(BaseModel): @@ -63,6 +65,25 @@ class CostSnapshot(BaseModel): estimated_cost_usd will be $0 and should be ignored.""" +class BilledCostReport(BaseModel): + """Ground-truth billed cost reconciled from OpenRouter's generation API. + + Token-count × price estimates diverge from the OpenRouter dashboard by up + to 4x in both directions (reasoning-token shapes vary per provider; + prompt-caching discounts aren't modeled). The ``/api/v1/generation`` + endpoint returns the exact billed cost per call, so summing it over a + run's generation IDs gives the real spend — see the v0.3.0 spread + reconciliation in results/spread/real_costs_openrouter.json. + """ + + model_config = ConfigDict(frozen=True) + + billed_cost_usd: float + billed_generations: int + missing_generations: int + source: str = "openrouter_generation_api" + + class CostTracker: """Accumulates token usage and estimates cost across a benchmark run.""" @@ -82,6 +103,21 @@ def __init__( self._total_reasoning = 0 self._num_calls = 0 self._start_time = time.monotonic() + self._generation_ids: list[str] = [] + + def record_generation_id(self, generation_id: str | None) -> None: + """Remember a provider generation ID for billed-cost reconciliation. + + IDs accumulate for EVERY provider call — including parse retries and + deliberations that later failed to parse — because those bill tokens + too. No-op on None/empty (non-OpenRouter providers may not set one). + """ + if generation_id: + self._generation_ids.append(generation_id) + + @property + def generation_ids(self) -> list[str]: + return list(self._generation_ids) def record(self, usage: TokenUsage) -> None: """Record token usage from one LLM call.""" @@ -203,3 +239,74 @@ async def create_cost_tracker( source, ) return CostTracker(model, prompt_price, completion_price, pricing_source=source) + + +async def fetch_billed_costs( + generation_ids: list[str], + api_key: str, + *, + url: str = OPENROUTER_GENERATION_URL, + concurrency: int = 8, + timeout: float = 15.0, + retry_delay_s: float = 2.0, + transport: httpx.AsyncBaseTransport | None = None, +) -> BilledCostReport | None: + """Sum the exact billed cost for a run from OpenRouter's generation API. + + GET https://openrouter.ai/api/v1/generation?id= (authenticated) + returns ``data.total_cost`` — the ground-truth USD charge for that call. + Generation stats lag the completion by a moment, so a miss is retried + once after ``retry_delay_s``. Returns None when there is nothing to + reconcile (no IDs, or every lookup failed — e.g. wrong key), never a + misleading $0 report. ``transport`` is injectable for tests. + """ + if not generation_ids: + return None + semaphore = asyncio.Semaphore(concurrency) + + async def _fetch_one(client: httpx.AsyncClient, gen_id: str) -> float | None: + for attempt in range(2): + try: + resp = await client.get(url, params={"id": gen_id}) + if resp.status_code == 200: + cost = resp.json().get("data", {}).get("total_cost") + return float(cost) if cost is not None else 0.0 + except Exception: + logger.debug("Generation lookup failed for %s", gen_id, exc_info=True) + if attempt == 0: + await asyncio.sleep(retry_delay_s) + return None + + async def _bounded(client: httpx.AsyncClient, gen_id: str) -> float | None: + async with semaphore: + return await _fetch_one(client, gen_id) + + async with httpx.AsyncClient( + timeout=timeout, + headers={"Authorization": f"Bearer {api_key}"}, + transport=transport, + ) as client: + results = await asyncio.gather( + *(_bounded(client, g) for g in generation_ids), + ) + + costs = [r for r in results if r is not None] + missing = len(results) - len(costs) + if not costs: + logger.warning( + "Billed-cost reconciliation found none of %d generations — " + "check the API key / provider. Falling back to estimates.", + len(generation_ids), + ) + return None + if missing: + logger.warning( + "Billed-cost reconciliation missing %d/%d generations — " + "billed_cost_usd is a LOWER BOUND for this run.", + missing, len(generation_ids), + ) + return BilledCostReport( + billed_cost_usd=round(sum(costs), 6), + billed_generations=len(costs), + missing_generations=missing, + ) diff --git a/tests/unit/test_base_role.py b/tests/unit/test_base_role.py index fff9582..a1dc8de 100644 --- a/tests/unit/test_base_role.py +++ b/tests/unit/test_base_role.py @@ -503,3 +503,66 @@ def test_record_completion_merges_reasoning_into_usage( assert agent._last_usage["reasoning_tokens"] == 900 # original usage dict on the result is not mutated assert "reasoning_tokens" not in result.usage + + +# ------------------------------------------------------------------ +# generation-ID capture (billed-cost reconciliation against OpenRouter) +# ------------------------------------------------------------------ + + +class TestGenerationIdExtraction: + def test_none_raw_response(self) -> None: + assert RimWorldRoleAgent._extract_generation_id(None) is None + + def test_dict_shape(self) -> None: + assert RimWorldRoleAgent._extract_generation_id({"id": "gen-abc123"}) == "gen-abc123" + + def test_object_shape(self) -> None: + raw = SimpleNamespace(id="gen-xyz") + assert RimWorldRoleAgent._extract_generation_id(raw) == "gen-xyz" + + def test_non_string_id_returns_none(self) -> None: + assert RimWorldRoleAgent._extract_generation_id({"id": 12345}) is None + + def test_empty_string_returns_none(self) -> None: + assert RimWorldRoleAgent._extract_generation_id({"id": ""}) is None + + +class TestGenerationIdAccumulation: + def _result(self, gen_id: str) -> CompletionResult: + return CompletionResult( + content="{}", + model="mock", + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + raw_response={"id": gen_id, "usage": {}}, + ) + + def test_every_completion_accumulates( + self, mock_provider: MagicMock, helix: HelixGeometry, + ) -> None: + """Parse retries call _record_completion again — both IDs must + survive so failed-parse tokens are still billed (issue #04).""" + agent = _DummyRoleAgent("d-01", mock_provider, helix, spawn_time=0.0) + agent._record_completion(self._result("gen-1")) + agent._record_completion(self._result("gen-2")) + assert agent.drain_generation_ids() == ["gen-1", "gen-2"] + + def test_drain_clears( + self, mock_provider: MagicMock, helix: HelixGeometry, + ) -> None: + agent = _DummyRoleAgent("d-01", mock_provider, helix, spawn_time=0.0) + agent._record_completion(self._result("gen-1")) + agent.drain_generation_ids() + assert agent.drain_generation_ids() == [] + + def test_no_id_in_raw_response_is_skipped( + self, mock_provider: MagicMock, helix: HelixGeometry, + ) -> None: + agent = _DummyRoleAgent("d-01", mock_provider, helix, spawn_time=0.0) + result = CompletionResult( + content="{}", model="mock", + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + raw_response={"usage": {}}, + ) + agent._record_completion(result) + assert agent.drain_generation_ids() == [] diff --git a/tests/unit/test_cost_tracker.py b/tests/unit/test_cost_tracker.py index 349576a..8c830b3 100644 --- a/tests/unit/test_cost_tracker.py +++ b/tests/unit/test_cost_tracker.py @@ -10,10 +10,12 @@ import pytest from rle.tracking.cost_tracker import ( + BilledCostReport, CostSnapshot, CostTracker, TokenUsage, create_cost_tracker, + fetch_billed_costs, fetch_pricing, ) @@ -190,6 +192,127 @@ def test_wall_time_increases_over_time(self) -> None: assert snap2.wall_time_s >= snap1.wall_time_s +# ------------------------------------------------------------------ +# Generation-ID accumulation (billed-cost reconciliation) +# ------------------------------------------------------------------ + + +class TestGenerationIds: + def test_records_in_order(self) -> None: + tracker = CostTracker("model") + tracker.record_generation_id("gen-1") + tracker.record_generation_id("gen-2") + assert tracker.generation_ids == ["gen-1", "gen-2"] + + def test_ignores_none_and_empty(self) -> None: + tracker = CostTracker("model") + tracker.record_generation_id(None) + tracker.record_generation_id("") + tracker.record_generation_id("gen-1") + assert tracker.generation_ids == ["gen-1"] + + def test_property_returns_copy(self) -> None: + tracker = CostTracker("model") + tracker.record_generation_id("gen-1") + ids = tracker.generation_ids + ids.append("gen-tampered") + assert tracker.generation_ids == ["gen-1"] + + +# ------------------------------------------------------------------ +# fetch_billed_costs tests (OpenRouter /generation ground truth) +# ------------------------------------------------------------------ + + +def _generation_transport(costs: dict[str, float | None]) -> httpx.MockTransport: + """Transport serving /generation?id=... from a gen_id -> total_cost map. + + IDs absent from the map 404 (like OpenRouter does for unknown ids). + """ + + def handler(request: httpx.Request) -> httpx.Response: + gen_id = request.url.params.get("id") + if gen_id not in costs: + return httpx.Response(status_code=404, json={"error": "not found"}) + return httpx.Response( + status_code=200, json={"data": {"id": gen_id, "total_cost": costs[gen_id]}}, + ) + + return httpx.MockTransport(handler) + + +class TestFetchBilledCosts: + async def test_empty_ids_returns_none(self) -> None: + assert await fetch_billed_costs([], "key") is None + + async def test_sums_billed_costs(self) -> None: + transport = _generation_transport({"gen-1": 0.012, "gen-2": 0.03, "gen-3": 0.0005}) + report = await fetch_billed_costs( + ["gen-1", "gen-2", "gen-3"], "key", + transport=transport, retry_delay_s=0.0, + ) + assert report is not None + assert report.billed_cost_usd == pytest.approx(0.0425) + assert report.billed_generations == 3 + assert report.missing_generations == 0 + assert report.source == "openrouter_generation_api" + + async def test_missing_generation_counts_as_lower_bound(self) -> None: + transport = _generation_transport({"gen-1": 0.01}) + report = await fetch_billed_costs( + ["gen-1", "gen-unknown"], "key", + transport=transport, retry_delay_s=0.0, + ) + assert report is not None + assert report.billed_cost_usd == pytest.approx(0.01) + assert report.billed_generations == 1 + assert report.missing_generations == 1 + + async def test_all_missing_returns_none(self) -> None: + transport = _generation_transport({}) + report = await fetch_billed_costs( + ["gen-a", "gen-b"], "key", transport=transport, retry_delay_s=0.0, + ) + assert report is None + + async def test_null_total_cost_treated_as_free(self) -> None: + transport = _generation_transport({"gen-free": None}) + report = await fetch_billed_costs( + ["gen-free"], "key", transport=transport, retry_delay_s=0.0, + ) + assert report is not None + assert report.billed_cost_usd == 0.0 + assert report.billed_generations == 1 + + async def test_connection_errors_count_as_missing(self) -> None: + transport = _make_error_transport() + report = await fetch_billed_costs( + ["gen-1"], "key", transport=transport, retry_delay_s=0.0, + ) + assert report is None + + async def test_sends_bearer_auth(self) -> None: + seen_auth: list[str | None] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen_auth.append(request.headers.get("authorization")) + return httpx.Response(status_code=200, json={"data": {"total_cost": 0.01}}) + + report = await fetch_billed_costs( + ["gen-1"], "sk-or-test", transport=httpx.MockTransport(handler), + retry_delay_s=0.0, + ) + assert report is not None + assert seen_auth == ["Bearer sk-or-test"] + + def test_report_is_frozen(self) -> None: + report = BilledCostReport( + billed_cost_usd=1.0, billed_generations=2, missing_generations=0, + ) + with pytest.raises(Exception): + report.billed_cost_usd = 9.9 # type: ignore[misc] + + # ------------------------------------------------------------------ # fetch_pricing tests # ------------------------------------------------------------------ diff --git a/tests/unit/test_game_loop_keepalive.py b/tests/unit/test_game_loop_keepalive.py new file mode 100644 index 0000000..e92bde9 --- /dev/null +++ b/tests/unit/test_game_loop_keepalive.py @@ -0,0 +1,66 @@ +"""Tests for the no-pause speed keepalive (issue #05: pause-on-threat). + +RimWorld force-pauses on threat letters; in --no-pause mode the loop must +re-assert speed continuously, not just at tick boundaries, or slow models +leave the game frozen for their whole deliberation window. + +The keepalive only touches ``_speed_keepalive_s`` and ``_client``, so the +tests build a bare instance instead of wiring the full 7-agent loop. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from unittest.mock import AsyncMock, MagicMock + +from rle.orchestration.game_loop import RLEGameLoop + + +def _bare_loop(keepalive_s: float, client: MagicMock) -> RLEGameLoop: + loop = object.__new__(RLEGameLoop) + loop._speed_keepalive_s = keepalive_s + loop._client = client + return loop + + +class TestSpeedKeepalive: + async def test_reasserts_speed_repeatedly(self) -> None: + client = MagicMock() + client.unpause_game = AsyncMock() + loop = _bare_loop(0.01, client) + + task = asyncio.create_task(loop._speed_keepalive()) + await asyncio.sleep(0.1) + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + assert client.unpause_game.await_count >= 2 + + async def test_survives_client_errors(self) -> None: + """A flaky RIMAPI call must not kill the keepalive for the run.""" + client = MagicMock() + client.unpause_game = AsyncMock(side_effect=ConnectionError("rimapi down")) + loop = _bare_loop(0.01, client) + + task = asyncio.create_task(loop._speed_keepalive()) + await asyncio.sleep(0.1) + assert not task.done() # still looping despite every call failing + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + assert client.unpause_game.await_count >= 2 + + async def test_cancellation_stops_cleanly(self) -> None: + client = MagicMock() + client.unpause_game = AsyncMock() + loop = _bare_loop(10.0, client) + + task = asyncio.create_task(loop._speed_keepalive()) + await asyncio.sleep(0) + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + assert task.cancelled()