From 1d14260479a1b4dad7b2203467a3d30d9f0d1d55 Mon Sep 17 00:00:00 2001 From: romgenie <5861166+romgenie@users.noreply.github.com> Date: Sat, 9 May 2026 20:26:25 -0400 Subject: [PATCH 1/3] Add structured MCP run logging --- .gitignore | 2 + README.md | 4 + mcp/README.md | 42 +++- mcp/run_logger.py | 166 +++++++++++++++ mcp/server.py | 512 ++++++++++++++++++++++++++++++++++++++-------- 5 files changed, 636 insertions(+), 90 deletions(-) create mode 100644 mcp/run_logger.py diff --git a/.gitignore b/.gitignore index 825ceace..78427bca 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,8 @@ __pycache__/ *.pyc *.pyo mcp/.venv/ +logs/ +mcp/logs/ # Claude Code # .claude/ diff --git a/README.md b/README.md index 596839ee..5647d49f 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,10 @@ The MCP server accepts `--host` and `--port` options if you need non-default set Flag `--no-trust-env` can be used to disable `requests` from picking up proxy settings from the environment, which can cause connection issues if you are running the server in a container. +By default, the MCP bridge records a structured JSONL run log in `logs/run_-.jsonl` relative to the server working directory. The log captures MCP tool calls, STS2_MCP HTTP requests/responses, smart state-poll decisions, timings, SHA-256 hashes, bounded previews, and explicit `log_agent_decision` annotations. Use `--disable-run-log` to turn it off, `--log-dir` or `STS2_MCP_LOG_DIR` to choose a location, `--log-preview-chars` to adjust previews, and `--log-full-text` when you need complete response text for replay or evaluation. + +`get_game_state` and `mp_get_game_state` wait briefly for actionable states by default, so agents do not immediately receive enemy-turn combat states or transient event/reward screens. Pass `wait_for_actionable=false` to either tool for immediate state reads. + ## For Developers ### Build & Install diff --git a/mcp/README.md b/mcp/README.md index d36f9048..3cc5157c 100644 --- a/mcp/README.md +++ b/mcp/README.md @@ -4,7 +4,8 @@ | Tool | Scope | Description | |---|---|---| -| `get_game_state(format?)` | General | Get current game state (`markdown` or `json`) | +| `get_game_state(format?, wait_for_actionable?, wait_timeout?, poll_interval?)` | General | Get current game state (`markdown` or `json`), optionally waiting through transient non-actionable states | +| `log_agent_decision(summary, reasoning?, intended_action?, alternatives?, confidence?, tags?)` | General | Add a structured decision annotation to the run log | | `menu_select(option, seed?)` | General | Select a visible menu/game-over option | | `get_profile()` | Profiles | Get active profile progress | | `list_profiles()` | Profiles | List profile slots and active slot | @@ -44,7 +45,7 @@ All multiplayer tools are prefixed with `mp_`. They route through `/api/v1/multi | Tool | Scope | Description | |---|---|---| -| `mp_get_game_state(format?)` | General | Get multiplayer game state (all players, votes, bids) | +| `mp_get_game_state(format?, wait_for_actionable?, wait_timeout?, poll_interval?)` | General | Get multiplayer game state (all players, votes, bids), optionally waiting through transient non-actionable states | | `mp_combat_play_card(card_index, target?)` | Combat | Play a card from the local player's hand | | `mp_combat_end_turn()` | Combat | Submit end-turn vote (turn ends when all players submit) | | `mp_combat_undo_end_turn()` | Combat | Retract end-turn vote | @@ -73,3 +74,40 @@ All multiplayer tools are prefixed with `mp_`. They route through `/api/v1/multi | `mp_crystal_sphere_set_tool(tool)` | Crystal Sphere | Switch the active divination tool | | `mp_crystal_sphere_click_cell(x, y)` | Crystal Sphere | Click a hidden cell in the grid | | `mp_crystal_sphere_proceed()` | Crystal Sphere | Continue after the minigame finishes | + +## Run Logging + +The MCP bridge writes structured JSONL logs by default under `logs/run_-.jsonl` relative to the server working directory. Each line has a stable envelope with `schema_version`, `run_id`, `sequence`, UTC `timestamp`, `monotonic_ms`, `event_type`, and when applicable `tool_call_id` / `tool_name`. + +Logged events include: + +- `session_start` with bridge configuration and runtime metadata. +- `tool_call_start`, `tool_call_result`, and `tool_call_error` for every MCP tool. +- `http_request`, `http_response`, and `http_error` for calls to the STS2_MCP REST API. +- `state_poll` and `state_poll_final_format` for smart polling decisions. +- `agent_decision` entries from `log_agent_decision`. + +Tool and HTTP results include length, byte count, SHA-256, preview text, and truncation status. Keys containing `authorization`, `cookie`, `password`, `secret`, `token`, `api_key`, or `apikey` are redacted before writing. + +Logging options: + +```bash +python server.py --log-dir logs +python server.py --disable-run-log +python server.py --log-preview-chars 8000 +python server.py --log-full-text +``` + +`STS2_MCP_LOG_DIR` can also set the default log directory. Use `--log-full-text` when you need complete replayable tool/API text for evaluation; otherwise previews plus hashes keep the log smaller while preserving integrity checks. + +## Smart State Polling + +`get_game_state` and `mp_get_game_state` default to `wait_for_actionable=true`. The bridge polls JSON state until one of these conditions is met, then returns the requested format: + +- combat reaches the player's play phase; +- event dialogue/options are available; +- reward, rest, shop, treasure, or Crystal Sphere controls are actionable; +- another non-transient state is reached; +- `wait_timeout` expires. + +Set `wait_for_actionable=false` to get the immediate raw state. `wait_timeout` is capped at 60 seconds and `poll_interval` is capped between 0.1 and 5 seconds. diff --git a/mcp/run_logger.py b/mcp/run_logger.py new file mode 100644 index 00000000..0eeeb375 --- /dev/null +++ b/mcp/run_logger.py @@ -0,0 +1,166 @@ +"""Structured JSONL run logging for the STS2 MCP bridge.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import os +import time +import uuid +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + + +SCHEMA_VERSION = "2026-05-10" + +DEFAULT_REDACT_KEY_PARTS = ( + "authorization", + "cookie", + "password", + "secret", + "token", + "api_key", + "apikey", +) + + +class RunLogger: + """Append-only JSONL logger with stable event envelopes.""" + + def __init__( + self, + *, + enabled: bool = False, + log_dir: str | os.PathLike[str] = "logs", + preview_chars: int = 4000, + include_full_text: bool = False, + redact_key_parts: tuple[str, ...] = DEFAULT_REDACT_KEY_PARTS, + ) -> None: + self.enabled = enabled + self.log_dir = Path(log_dir) + self.preview_chars = max(0, preview_chars) + self.include_full_text = include_full_text + self.redact_key_parts = tuple(part.lower() for part in redact_key_parts) + self.run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + "-" + uuid.uuid4().hex[:8] + self.path = self.log_dir / f"run_{self.run_id}.jsonl" + self._sequence = 0 + self._started_at = time.monotonic() + self._lock = asyncio.Lock() + + def start(self, metadata: dict[str, Any] | None = None) -> None: + if not self.enabled: + return + + self.log_dir.mkdir(parents=True, exist_ok=True) + self._write_sync( + self._envelope( + "session_start", + { + "log_path": str(self.path), + "metadata": self.redact(metadata or {}), + }, + ) + ) + + async def log( + self, + event_type: str, + payload: dict[str, Any] | None = None, + *, + tool_call_id: str | None = None, + tool_name: str | None = None, + ) -> None: + if not self.enabled: + return + + async with self._lock: + record = self._envelope( + event_type, + self.redact(payload or {}), + tool_call_id=tool_call_id, + tool_name=tool_name, + ) + self._write_sync(record) + + def redact(self, value: Any) -> Any: + if isinstance(value, dict): + redacted: dict[str, Any] = {} + for key, item in value.items(): + key_text = str(key) + if self._is_sensitive_key(key_text): + redacted[key_text] = "[REDACTED]" + else: + redacted[key_text] = self.redact(item) + return redacted + if isinstance(value, list): + return [self.redact(item) for item in value] + if isinstance(value, tuple): + return [self.redact(item) for item in value] + return value + + def summarize_text(self, text: str | bytes | None) -> dict[str, Any]: + if text is None: + return {"length": 0, "sha256": None, "preview": ""} + if isinstance(text, bytes): + raw = text + display = text.decode("utf-8", errors="replace") + else: + display = text + raw = text.encode("utf-8", errors="replace") + + summary: dict[str, Any] = { + "length": len(display), + "bytes": len(raw), + "sha256": hashlib.sha256(raw).hexdigest(), + "preview": display[: self.preview_chars], + "truncated": len(display) > self.preview_chars, + } + if self.include_full_text: + summary["text"] = display + return summary + + def summarize_jsonable(self, value: Any) -> dict[str, Any]: + redacted = self.redact(value) + try: + serialized = json.dumps(redacted, ensure_ascii=False, sort_keys=True, default=str) + except TypeError: + serialized = json.dumps(str(redacted), ensure_ascii=False) + + summary = self.summarize_text(serialized) + summary["json_type"] = type(value).__name__ + return summary + + def _envelope( + self, + event_type: str, + payload: dict[str, Any], + *, + tool_call_id: str | None = None, + tool_name: str | None = None, + ) -> dict[str, Any]: + self._sequence += 1 + envelope: dict[str, Any] = { + "schema_version": SCHEMA_VERSION, + "run_id": self.run_id, + "sequence": self._sequence, + "timestamp": datetime.now(UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z"), + "monotonic_ms": round((time.monotonic() - self._started_at) * 1000, 3), + "event_type": event_type, + "payload": payload, + } + if tool_call_id is not None: + envelope["tool_call_id"] = tool_call_id + if tool_name is not None: + envelope["tool_name"] = tool_name + return envelope + + def _write_sync(self, record: dict[str, Any]) -> None: + with self.path.open("a", encoding="utf-8") as handle: + json.dump(record, handle, ensure_ascii=False, separators=(",", ":"), default=str) + handle.write("\n") + + def _is_sensitive_key(self, key: str) -> bool: + key_lower = key.lower() + return any(part in key_lower for part in self.redact_key_parts) diff --git a/mcp/server.py b/mcp/server.py index 5716310a..a949bc93 100644 --- a/mcp/server.py +++ b/mcp/server.py @@ -6,17 +6,31 @@ import argparse import asyncio +import contextvars +import functools +import inspect import json +import os +import platform import sys +import time +import uuid +from typing import Any, Awaitable, Callable import httpx from mcp.server.fastmcp import FastMCP +from run_logger import RunLogger mcp = FastMCP("sts2") _base_url: str = "http://localhost:15526" _trust_env: bool = True _http: httpx.AsyncClient | None = None +_run_logger = RunLogger(enabled=False) +_tool_call_id: contextvars.ContextVar[str | None] = contextvars.ContextVar("tool_call_id", default=None) +_tool_name: contextvars.ContextVar[str | None] = contextvars.ContextVar("tool_name", default=None) + +COMBAT_STATE_TYPES = {"monster", "elite", "boss"} def _sp_url() -> str: @@ -42,46 +56,95 @@ def _get_client() -> httpx.AsyncClient: return _http +async def _request_text( + method: str, + url: str, + *, + params: dict | None = None, + json_body: dict | None = None, +) -> str: + started = time.perf_counter() + current_tool_call_id = _tool_call_id.get() + current_tool_name = _tool_name.get() + await _run_logger.log( + "http_request", + { + "method": method, + "url": url, + "params": params or {}, + "json": json_body or {}, + }, + tool_call_id=current_tool_call_id, + tool_name=current_tool_name, + ) + try: + client = _get_client() + if method == "GET": + response = await client.get(url, params=params) + elif method == "POST": + response = await client.post(url, json=json_body) + else: + raise ValueError(f"Unsupported HTTP method: {method}") + + text = response.text + elapsed_ms = round((time.perf_counter() - started) * 1000, 3) + await _run_logger.log( + "http_response", + { + "method": method, + "url": url, + "status_code": response.status_code, + "elapsed_ms": elapsed_ms, + "response": _run_logger.summarize_text(text), + }, + tool_call_id=current_tool_call_id, + tool_name=current_tool_name, + ) + response.raise_for_status() + return text + except Exception as exc: + elapsed_ms = round((time.perf_counter() - started) * 1000, 3) + await _run_logger.log( + "http_error", + { + "method": method, + "url": url, + "elapsed_ms": elapsed_ms, + "error_type": type(exc).__name__, + "error": str(exc), + }, + tool_call_id=current_tool_call_id, + tool_name=current_tool_name, + ) + raise + + async def _get(params: dict | None = None) -> str: - r = await _get_client().get(_sp_url(), params=params) - r.raise_for_status() - return r.text + return await _request_text("GET", _sp_url(), params=params) async def _post(body: dict) -> str: - r = await _get_client().post(_sp_url(), json=body) - r.raise_for_status() - return r.text + return await _request_text("POST", _sp_url(), json_body=body) async def _mp_get(params: dict | None = None) -> str: - r = await _get_client().get(_mp_url(), params=params) - r.raise_for_status() - return r.text + return await _request_text("GET", _mp_url(), params=params) async def _mp_post(body: dict) -> str: - r = await _get_client().post(_mp_url(), json=body) - r.raise_for_status() - return r.text + return await _request_text("POST", _mp_url(), json_body=body) async def _profile_get() -> str: - r = await _get_client().get(_profile_url()) - r.raise_for_status() - return r.text + return await _request_text("GET", _profile_url()) async def _profiles_get() -> str: - r = await _get_client().get(_profiles_url()) - r.raise_for_status() - return r.text + return await _request_text("GET", _profiles_url()) async def _profiles_post(body: dict) -> str: - r = await _get_client().post(_profiles_url(), json=body) - r.raise_for_status() - return r.text + return await _request_text("POST", _profiles_url(), json_body=body) async def _wait_for_profile(profile_id: int, fallback: str) -> str: @@ -130,13 +193,188 @@ def _handle_error(e: Exception) -> str: return f"Error: {e}" +def logged_tool(*tool_args: Any, **tool_kwargs: Any) -> Callable[[Callable[..., Awaitable[str]]], Callable[..., Awaitable[str]]]: + def decorator(func: Callable[..., Awaitable[str]]) -> Callable[..., Awaitable[str]]: + @functools.wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> str: + call_id = uuid.uuid4().hex + call_id_token = _tool_call_id.set(call_id) + name_token = _tool_name.set(func.__name__) + started = time.perf_counter() + try: + bound = inspect.signature(func).bind_partial(*args, **kwargs) + tool_args_payload = dict(bound.arguments) + except Exception: + tool_args_payload = {"args": list(args), "kwargs": kwargs} + + await _run_logger.log( + "tool_call_start", + {"args": tool_args_payload}, + tool_call_id=call_id, + tool_name=func.__name__, + ) + try: + result = await func(*args, **kwargs) + except Exception as exc: + elapsed_ms = round((time.perf_counter() - started) * 1000, 3) + await _run_logger.log( + "tool_call_error", + { + "elapsed_ms": elapsed_ms, + "error_type": type(exc).__name__, + "error": str(exc), + }, + tool_call_id=call_id, + tool_name=func.__name__, + ) + raise + else: + elapsed_ms = round((time.perf_counter() - started) * 1000, 3) + await _run_logger.log( + "tool_call_result", + { + "elapsed_ms": elapsed_ms, + "result": _run_logger.summarize_text(result), + }, + tool_call_id=call_id, + tool_name=func.__name__, + ) + return result + finally: + _tool_call_id.reset(call_id_token) + _tool_name.reset(name_token) + + return mcp.tool(*tool_args, **tool_kwargs)(wrapper) + + return decorator + + +def _coerce_wait_timeout(wait_timeout: float) -> float: + return max(0.0, min(60.0, wait_timeout)) + + +def _coerce_poll_interval(poll_interval: float) -> float: + return max(0.1, min(5.0, poll_interval)) + + +def _state_actionability(state: dict[str, Any]) -> tuple[bool, str]: + state_type = str(state.get("state_type") or "unknown") + + if state_type in COMBAT_STATE_TYPES: + battle = state.get("battle") if isinstance(state.get("battle"), dict) else {} + is_play_phase = battle.get("is_play_phase") is True + turn = str(battle.get("turn") or "").lower() + if is_play_phase and turn == "player": + return True, "combat_player_turn" + return False, "combat_waiting_for_player_turn" + + if state_type == "event": + event = state.get("event") if isinstance(state.get("event"), dict) else {} + if event.get("in_dialogue") is True: + return True, "event_dialogue_actionable" + options = event.get("options") if isinstance(event.get("options"), list) else [] + if any(isinstance(opt, dict) and opt.get("is_locked") is not True for opt in options): + return True, "event_option_actionable" + return False, "event_waiting_for_options" + + if state_type in {"rewards", "rest_site", "shop", "fake_merchant", "crystal_sphere", "treasure"}: + payload = state.get(state_type) if isinstance(state.get(state_type), dict) else {} + if payload.get("can_proceed") is True: + return True, f"{state_type}_can_proceed" + for key in ("items", "options", "cards", "relics", "cells"): + value = payload.get(key) + if isinstance(value, list) and len(value) > 0: + return True, f"{state_type}_{key}_available" + return False, f"{state_type}_waiting" + + return True, f"{state_type}_actionable" + + +async def _get_smart_state( + getter: Callable[[dict | None], Awaitable[str]], + params: dict[str, Any], + *, + wait_for_actionable: bool, + wait_timeout: float, + poll_interval: float, +) -> str: + if not wait_for_actionable: + return await getter(params) + + timeout = _coerce_wait_timeout(wait_timeout) + interval = _coerce_poll_interval(poll_interval) + requested_format = str(params.get("format") or "json") + detection_params = dict(params) + detection_params["format"] = "json" + deadline = time.monotonic() + timeout + attempt = 0 + last_json_text = "" + last_reason = "not_polled" + + while True: + attempt += 1 + last_json_text = await getter(detection_params) + try: + state = json.loads(last_json_text) + except json.JSONDecodeError: + await _run_logger.log( + "state_poll", + { + "attempt": attempt, + "actionable": True, + "reason": "state_json_parse_failed", + }, + tool_call_id=_tool_call_id.get(), + tool_name=_tool_name.get(), + ) + return await getter(params) if requested_format != "json" else last_json_text + + actionable, reason = _state_actionability(state) + last_reason = reason + await _run_logger.log( + "state_poll", + { + "attempt": attempt, + "state_type": state.get("state_type"), + "actionable": actionable, + "reason": reason, + "timeout_seconds": timeout, + }, + tool_call_id=_tool_call_id.get(), + tool_name=_tool_name.get(), + ) + if actionable or time.monotonic() >= deadline: + break + await asyncio.sleep(min(interval, max(0.0, deadline - time.monotonic()))) + + if requested_format == "json": + return last_json_text + + await _run_logger.log( + "state_poll_final_format", + { + "format": requested_format, + "last_reason": last_reason, + "attempts": attempt, + }, + tool_call_id=_tool_call_id.get(), + tool_name=_tool_name.get(), + ) + return await getter(params) + + # --------------------------------------------------------------------------- # General # --------------------------------------------------------------------------- -@mcp.tool() -async def get_game_state(format: str = "markdown") -> str: +@logged_tool() +async def get_game_state( + format: str = "markdown", + wait_for_actionable: bool = True, + wait_timeout: float = 8.0, + poll_interval: float = 1.0, +) -> str: """Get the current Slay the Spire 2 game state. Returns the full game state including player stats, hand, enemies, potions, etc. @@ -145,14 +383,62 @@ async def get_game_state(format: str = "markdown") -> str: Args: format: "markdown" for human-readable output, "json" for structured data. + wait_for_actionable: Poll briefly through transient states such as enemy turns. + wait_timeout: Maximum seconds to wait for an actionable state. + poll_interval: Seconds between state polls while waiting. """ try: - return await _get({"format": format}) + return await _get_smart_state( + _get, + {"format": format}, + wait_for_actionable=wait_for_actionable, + wait_timeout=wait_timeout, + poll_interval=poll_interval, + ) except Exception as e: return _handle_error(e) -@mcp.tool() +@logged_tool() +async def log_agent_decision( + summary: str, + reasoning: str | None = None, + intended_action: str | None = None, + alternatives: list[str] | None = None, + confidence: float | None = None, + tags: list[str] | None = None, +) -> str: + """Record an explicit agent decision in the run log. + + Use before important choices so later run review can connect game state, + candidate actions, and the agent's stated decision. + + Args: + summary: Short decision summary. + reasoning: Optional rationale or policy note. + intended_action: Optional next tool/action the agent expects to call. + alternatives: Optional rejected alternatives. + confidence: Optional confidence score from 0.0 to 1.0. + tags: Optional labels for later filtering. + """ + event = { + "summary": summary, + "reasoning": reasoning, + "intended_action": intended_action, + "alternatives": alternatives or [], + "confidence": confidence, + "tags": tags or [], + } + await _run_logger.log( + "agent_decision", + event, + tool_call_id=_tool_call_id.get(), + tool_name=_tool_name.get(), + ) + return json.dumps({"status": "ok", "logged": event}, indent=2) + + +@logged_tool() async def menu_select(option: str, seed: str | None = None) -> str: """Select a visible menu option. @@ -186,7 +472,7 @@ async def menu_select(option: str, seed: str | None = None) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def get_profile() -> str: """Get the current profile's persistent progress summary. @@ -199,7 +485,7 @@ async def get_profile() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def list_profiles() -> str: """List the three profile slots and identify the active slot. @@ -211,7 +497,7 @@ async def list_profiles() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def switch_profile(profile_id: int) -> str: """Switch to a profile slot through the game's profile UI. @@ -249,7 +535,7 @@ async def switch_profile(profile_id: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def delete_profile(profile_id: int) -> str: """Delete an inactive profile slot. @@ -265,7 +551,7 @@ async def delete_profile(profile_id: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def use_potion(slot: int, target: str | None = None) -> str: """Use a potion from the player's potion slots. @@ -284,7 +570,7 @@ async def use_potion(slot: int, target: str | None = None) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def discard_potion(slot: int) -> str: """Discard a potion from the player's potion slots to free up space. @@ -300,7 +586,7 @@ async def discard_potion(slot: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def proceed_to_map() -> str: """Proceed from the current screen to the map. @@ -318,7 +604,7 @@ async def proceed_to_map() -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def combat_play_card(card_index: int, target: str | None = None) -> str: """[Combat] Play a card from the player's hand. @@ -338,7 +624,7 @@ async def combat_play_card(card_index: int, target: str | None = None) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def combat_end_turn() -> str: """[Combat] End the player's current turn.""" try: @@ -352,7 +638,7 @@ async def combat_end_turn() -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def combat_select_card(card_index: int) -> str: """[Combat Selection] Select a card from hand during an in-combat card selection prompt. @@ -368,7 +654,7 @@ async def combat_select_card(card_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def combat_confirm_selection() -> str: """[Combat Selection] Confirm the in-combat card selection. @@ -386,7 +672,7 @@ async def combat_confirm_selection() -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def rewards_claim(reward_index: int) -> str: """[Rewards] Claim a reward from the post-combat rewards screen. @@ -405,7 +691,7 @@ async def rewards_claim(reward_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def rewards_pick_card(card_index: int) -> str: """[Rewards] Select a card from the card reward selection screen. @@ -418,7 +704,7 @@ async def rewards_pick_card(card_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def rewards_skip_card() -> str: """[Rewards] Skip the card reward without selecting a card.""" try: @@ -432,7 +718,7 @@ async def rewards_skip_card() -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def map_choose_node(node_index: int) -> str: """[Map] Choose a map node to travel to. @@ -450,7 +736,7 @@ async def map_choose_node(node_index: int) -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def rest_choose_option(option_index: int) -> str: """[Rest Site] Choose a rest site option (rest, smith, etc.). @@ -468,7 +754,7 @@ async def rest_choose_option(option_index: int) -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def shop_purchase(item_index: int) -> str: """[Shop / Fake Merchant] Purchase an item from the shop. @@ -489,7 +775,7 @@ async def shop_purchase(item_index: int) -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def event_choose_option(option_index: int) -> str: """[Event] Choose an event option. @@ -505,7 +791,7 @@ async def event_choose_option(option_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def event_advance_dialogue() -> str: """[Event] Advance ancient event dialogue. @@ -522,7 +808,7 @@ async def event_advance_dialogue() -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def deck_select_card(card_index: int) -> str: """[Card Selection] Select or deselect a card in the card selection screen. @@ -540,7 +826,7 @@ async def deck_select_card(card_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def deck_confirm_selection() -> str: """[Card Selection] Confirm the current card selection. @@ -554,7 +840,7 @@ async def deck_confirm_selection() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def deck_cancel_selection() -> str: """[Card Selection] Cancel the current card selection. @@ -573,7 +859,7 @@ async def deck_cancel_selection() -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def bundle_select(bundle_index: int) -> str: """[Bundle Selection] Open a bundle preview. @@ -586,7 +872,7 @@ async def bundle_select(bundle_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def bundle_confirm_selection() -> str: """[Bundle Selection] Confirm the currently previewed bundle.""" try: @@ -595,7 +881,7 @@ async def bundle_confirm_selection() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def bundle_cancel_selection() -> str: """[Bundle Selection] Cancel the current bundle preview.""" try: @@ -609,7 +895,7 @@ async def bundle_cancel_selection() -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def relic_select(relic_index: int) -> str: """[Relic Selection] Select a relic from the relic selection screen. @@ -624,7 +910,7 @@ async def relic_select(relic_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def relic_skip() -> str: """[Relic Selection] Skip the relic selection without choosing a relic.""" try: @@ -638,7 +924,7 @@ async def relic_skip() -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def treasure_claim_relic(relic_index: int) -> str: """[Treasure] Claim a relic from the treasure chest. @@ -659,7 +945,7 @@ async def treasure_claim_relic(relic_index: int) -> str: # --------------------------------------------------------------------------- -@mcp.tool() +@logged_tool() async def crystal_sphere_set_tool(tool: str) -> str: """[Crystal Sphere] Switch the active divination tool. @@ -672,7 +958,7 @@ async def crystal_sphere_set_tool(tool: str) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def crystal_sphere_click_cell(x: int, y: int) -> str: """[Crystal Sphere] Click a hidden cell on the Crystal Sphere grid. @@ -686,7 +972,7 @@ async def crystal_sphere_click_cell(x: int, y: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def crystal_sphere_proceed() -> str: """[Crystal Sphere] Continue after the Crystal Sphere minigame finishes.""" try: @@ -700,8 +986,13 @@ async def crystal_sphere_proceed() -> str: # =========================================================================== -@mcp.tool() -async def mp_get_game_state(format: str = "markdown") -> str: +@logged_tool() +async def mp_get_game_state( + format: str = "markdown", + wait_for_actionable: bool = True, + wait_timeout: float = 8.0, + poll_interval: float = 1.0, +) -> str: """[Multiplayer] Get the current multiplayer game state. Returns a summary of all players (HP, gold, alive status) plus full @@ -711,14 +1002,23 @@ async def mp_get_game_state(format: str = "markdown") -> str: Args: format: "markdown" for human-readable output, "json" for structured data. + wait_for_actionable: Poll briefly through transient states such as enemy turns. + wait_timeout: Maximum seconds to wait for an actionable state. + poll_interval: Seconds between state polls while waiting. """ try: - return await _mp_get({"format": format}) + return await _get_smart_state( + _mp_get, + {"format": format}, + wait_for_actionable=wait_for_actionable, + wait_timeout=wait_timeout, + poll_interval=poll_interval, + ) except Exception as e: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_combat_play_card(card_index: int, target: str | None = None) -> str: """[Multiplayer Combat] Play a card from the local player's hand. @@ -738,7 +1038,7 @@ async def mp_combat_play_card(card_index: int, target: str | None = None) -> str return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_combat_end_turn() -> str: """[Multiplayer Combat] Submit end-turn vote. @@ -751,7 +1051,7 @@ async def mp_combat_end_turn() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_combat_undo_end_turn() -> str: """[Multiplayer Combat] Retract end-turn vote. @@ -764,7 +1064,7 @@ async def mp_combat_undo_end_turn() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_use_potion(slot: int, target: str | None = None) -> str: """[Multiplayer] Use a potion from the local player's potion slots. @@ -781,7 +1081,7 @@ async def mp_use_potion(slot: int, target: str | None = None) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_discard_potion(slot: int) -> str: """[Multiplayer] Discard a potion from the local player's potion slots to free up space. @@ -794,7 +1094,7 @@ async def mp_discard_potion(slot: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_map_vote(node_index: int) -> str: """[Multiplayer Map] Vote for a map node to travel to. @@ -810,7 +1110,7 @@ async def mp_map_vote(node_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_event_choose_option(option_index: int) -> str: """[Multiplayer Event] Choose or vote for an event option. @@ -826,7 +1126,7 @@ async def mp_event_choose_option(option_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_event_advance_dialogue() -> str: """[Multiplayer Event] Advance ancient event dialogue.""" try: @@ -835,7 +1135,7 @@ async def mp_event_advance_dialogue() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_rest_choose_option(option_index: int) -> str: """[Multiplayer Rest Site] Choose a rest site option (rest, smith, etc.). @@ -850,7 +1150,7 @@ async def mp_rest_choose_option(option_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_shop_purchase(item_index: int) -> str: """[Multiplayer Shop] Purchase an item from the shop. @@ -865,7 +1165,7 @@ async def mp_shop_purchase(item_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_rewards_claim(reward_index: int) -> str: """[Multiplayer Rewards] Claim a reward from the post-combat rewards screen. @@ -878,7 +1178,7 @@ async def mp_rewards_claim(reward_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_rewards_pick_card(card_index: int) -> str: """[Multiplayer Rewards] Select a card from the card reward screen. @@ -891,7 +1191,7 @@ async def mp_rewards_pick_card(card_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_rewards_skip_card() -> str: """[Multiplayer Rewards] Skip the card reward.""" try: @@ -900,7 +1200,7 @@ async def mp_rewards_skip_card() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_proceed_to_map() -> str: """[Multiplayer] Proceed from the current screen to the map. @@ -912,7 +1212,7 @@ async def mp_proceed_to_map() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_deck_select_card(card_index: int) -> str: """[Multiplayer Card Selection] Select or deselect a card in the card selection screen. @@ -925,7 +1225,7 @@ async def mp_deck_select_card(card_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_deck_confirm_selection() -> str: """[Multiplayer Card Selection] Confirm the current card selection.""" try: @@ -934,7 +1234,7 @@ async def mp_deck_confirm_selection() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_deck_cancel_selection() -> str: """[Multiplayer Card Selection] Cancel the current card selection.""" try: @@ -943,7 +1243,7 @@ async def mp_deck_cancel_selection() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_bundle_select(bundle_index: int) -> str: """[Multiplayer Bundle Selection] Open a bundle preview. @@ -956,7 +1256,7 @@ async def mp_bundle_select(bundle_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_bundle_confirm_selection() -> str: """[Multiplayer Bundle Selection] Confirm the currently previewed bundle.""" try: @@ -965,7 +1265,7 @@ async def mp_bundle_confirm_selection() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_bundle_cancel_selection() -> str: """[Multiplayer Bundle Selection] Cancel the current bundle preview.""" try: @@ -974,7 +1274,7 @@ async def mp_bundle_cancel_selection() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_combat_select_card(card_index: int) -> str: """[Multiplayer Combat Selection] Select a card from hand during in-combat card selection. @@ -987,7 +1287,7 @@ async def mp_combat_select_card(card_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_combat_confirm_selection() -> str: """[Multiplayer Combat Selection] Confirm the in-combat card selection.""" try: @@ -996,7 +1296,7 @@ async def mp_combat_confirm_selection() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_relic_select(relic_index: int) -> str: """[Multiplayer Relic Selection] Select a relic (boss relic rewards). @@ -1009,7 +1309,7 @@ async def mp_relic_select(relic_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_relic_skip() -> str: """[Multiplayer Relic Selection] Skip the relic selection.""" try: @@ -1018,7 +1318,7 @@ async def mp_relic_skip() -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_treasure_claim_relic(relic_index: int) -> str: """[Multiplayer Treasure] Bid on / claim a relic from the treasure chest. @@ -1034,7 +1334,7 @@ async def mp_treasure_claim_relic(relic_index: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_crystal_sphere_set_tool(tool: str) -> str: """[Multiplayer Crystal Sphere] Switch the active divination tool. @@ -1047,7 +1347,7 @@ async def mp_crystal_sphere_set_tool(tool: str) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_crystal_sphere_click_cell(x: int, y: int) -> str: """[Multiplayer Crystal Sphere] Click a hidden cell on the Crystal Sphere grid. @@ -1061,7 +1361,7 @@ async def mp_crystal_sphere_click_cell(x: int, y: int) -> str: return _handle_error(e) -@mcp.tool() +@logged_tool() async def mp_crystal_sphere_proceed() -> str: """[Multiplayer Crystal Sphere] Continue after the Crystal Sphere minigame finishes.""" try: @@ -1075,11 +1375,47 @@ def main(): parser.add_argument("--port", type=int, default=15526, help="Game HTTP server port") parser.add_argument("--host", type=str, default="localhost", help="Game HTTP server host") parser.add_argument("--no-trust-env", action="store_true", help="Ignore HTTP_PROXY/HTTPS_PROXY environment variables") + parser.add_argument( + "--log-dir", + default=os.environ.get("STS2_MCP_LOG_DIR", "logs"), + help="Directory for structured JSONL run logs (default: logs)", + ) + parser.add_argument("--disable-run-log", action="store_true", help="Disable MCP tool/http run logging") + parser.add_argument( + "--log-preview-chars", + type=int, + default=4000, + help="Maximum characters stored in each logged text preview", + ) + parser.add_argument( + "--log-full-text", + action="store_true", + help="Store complete tool and HTTP response text in logs instead of previews plus hashes only", + ) args = parser.parse_args() - global _base_url, _trust_env + global _base_url, _trust_env, _run_logger _base_url = f"http://{args.host}:{args.port}" _trust_env = not args.no_trust_env + _run_logger = RunLogger( + enabled=not args.disable_run_log, + log_dir=args.log_dir, + preview_chars=args.log_preview_chars, + include_full_text=args.log_full_text, + ) + _run_logger.start( + { + "base_url": _base_url, + "trust_env": _trust_env, + "argv": sys.argv, + "pid": os.getpid(), + "python_version": sys.version, + "platform": platform.platform(), + "cwd": os.getcwd(), + "log_full_text": args.log_full_text, + "log_preview_chars": args.log_preview_chars, + } + ) # Eagerly initialize the shared httpx client so the first request is fast _get_client() From 73945fd763fe8844d93340472c940a1ddd51b72f Mon Sep 17 00:00:00 2001 From: romgenie <5861166+romgenie@users.noreply.github.com> Date: Sat, 9 May 2026 22:31:12 -0400 Subject: [PATCH 2/3] Add token accounting to run logs --- README.md | 2 +- mcp/README.md | 52 ++++++- mcp/run_logger.py | 334 ++++++++++++++++++++++++++++++++++++++-- mcp/server.py | 73 ++++++++- mcp/token_usage.py | 173 +++++++++++++++++++++ mcp/validate_run_log.py | 137 ++++++++++++++++ 6 files changed, 755 insertions(+), 16 deletions(-) create mode 100644 mcp/token_usage.py create mode 100644 mcp/validate_run_log.py diff --git a/README.md b/README.md index 5647d49f..73f00cef 100644 --- a/README.md +++ b/README.md @@ -104,7 +104,7 @@ The MCP server accepts `--host` and `--port` options if you need non-default set Flag `--no-trust-env` can be used to disable `requests` from picking up proxy settings from the environment, which can cause connection issues if you are running the server in a container. -By default, the MCP bridge records a structured JSONL run log in `logs/run_-.jsonl` relative to the server working directory. The log captures MCP tool calls, STS2_MCP HTTP requests/responses, smart state-poll decisions, timings, SHA-256 hashes, bounded previews, and explicit `log_agent_decision` annotations. Use `--disable-run-log` to turn it off, `--log-dir` or `STS2_MCP_LOG_DIR` to choose a location, `--log-preview-chars` to adjust previews, and `--log-full-text` when you need complete response text for replay or evaluation. +By default, the MCP bridge records a structured JSONL run log in `logs/run_-.jsonl` and a token rollup summary in `logs/run_-.summary.json` relative to the server working directory. The log captures MCP tool calls, STS2_MCP HTTP requests/responses, smart state-poll decisions, timings, SHA-256 hashes, bounded previews, estimated token usage, and explicit `log_agent_decision` annotations. Use `--disable-run-log` to turn it off, `--log-dir` or `STS2_MCP_LOG_DIR` to choose a location, `--log-preview-chars` to adjust previews, `--log-full-text` when you need complete response text for replay or evaluation, and `--tokenizer-profile` / `--tokenizer-name` options to label the estimator profile. Clients with exact model API usage can call `log_external_token_usage` to reconcile provider token counts with local estimates. `get_game_state` and `mp_get_game_state` wait briefly for actionable states by default, so agents do not immediately receive enemy-turn combat states or transient event/reward screens. Pass `wait_for_actionable=false` to either tool for immediate state reads. diff --git a/mcp/README.md b/mcp/README.md index 3cc5157c..cea3a798 100644 --- a/mcp/README.md +++ b/mcp/README.md @@ -6,6 +6,7 @@ |---|---|---| | `get_game_state(format?, wait_for_actionable?, wait_timeout?, poll_interval?)` | General | Get current game state (`markdown` or `json`), optionally waiting through transient non-actionable states | | `log_agent_decision(summary, reasoning?, intended_action?, alternatives?, confidence?, tags?)` | General | Add a structured decision annotation to the run log | +| `log_external_token_usage(usage_json)` | General | Attach exact/external model token usage to a run log event or tool call | | `menu_select(option, seed?)` | General | Select a visible menu/game-over option | | `get_profile()` | Profiles | Get active profile progress | | `list_profiles()` | Profiles | List profile slots and active slot | @@ -77,7 +78,7 @@ All multiplayer tools are prefixed with `mp_`. They route through `/api/v1/multi ## Run Logging -The MCP bridge writes structured JSONL logs by default under `logs/run_-.jsonl` relative to the server working directory. Each line has a stable envelope with `schema_version`, `run_id`, `sequence`, UTC `timestamp`, `monotonic_ms`, `event_type`, and when applicable `tool_call_id` / `tool_name`. +The MCP bridge writes structured JSONL logs by default under `logs/run_-.jsonl` relative to the server working directory. It also writes `logs/run_-.summary.json` with run-level token rollups. Each JSONL line has a stable envelope with `schema_version`, `token_schema_version`, `run_id`, `event_id`, `sequence`, UTC `timestamp`, `monotonic_ms`, `event_type`, token fields, and when applicable `tool_call_id` / `tool_name`. Logged events include: @@ -86,8 +87,36 @@ Logged events include: - `http_request`, `http_response`, and `http_error` for calls to the STS2_MCP REST API. - `state_poll` and `state_poll_final_format` for smart polling decisions. - `agent_decision` entries from `log_agent_decision`. +- `external_token_usage` entries from `log_external_token_usage`. + +Tool and HTTP results include length, byte count, SHA-256, preview text, truncation status, and estimated token count. Keys containing `authorization`, `cookie`, `password`, `secret`, `token`, `api_key`, or `apikey` are redacted before writing. + +Every record includes: + +- `input_tokens` +- `output_tokens` +- `tool_response_tokens` +- `hidden_poll_tokens` +- `total_tokens` +- `token_source` +- `tokenizer_name` +- `tokenizer_version` +- `model_family` +- `estimation_method` + +By default these counts use a deterministic regex estimator (`generic_regex_v1`) so runs remain comparable without installing model-specific tokenizers. Treat these as estimates, not provider-billed tokens. Use `log_external_token_usage` to attach exact model API usage when a client has it: + +```json +{ + "related_tool_call_id": "9f...", + "model": "claude-sonnet-4.6", + "input_tokens": 1200, + "output_tokens": 400, + "token_source": "exact" +} +``` -Tool and HTTP results include length, byte count, SHA-256, preview text, and truncation status. Keys containing `authorization`, `cookie`, `password`, `secret`, `token`, `api_key`, or `apikey` are redacted before writing. +The summary artifact rolls up totals by tool name, event type, state type, game mode, action category, run phase, and floor where available. It also tracks hidden polling cost, largest payloads, repeated-state cost, invalid-action cost, replay artifact size/token metadata, and externally supplied usage records. Logging options: @@ -96,9 +125,26 @@ python server.py --log-dir logs python server.py --disable-run-log python server.py --log-preview-chars 8000 python server.py --log-full-text +python server.py --tokenizer-profile openai_cl100k_proxy +python server.py --token-model-family anthropic --tokenizer-name claude_proxy_regex +``` + +`STS2_MCP_LOG_DIR` can also set the default log directory. Use `--log-full-text` when you need complete replayable tool/API text for evaluation; otherwise previews plus hashes keep the log smaller while preserving integrity checks. Token estimates are computed from the full text before truncation, even when only previews are written. + +Validate a log and its summary: + +```bash +python validate_run_log.py --self-test +python validate_run_log.py logs/run_-.jsonl ``` -`STS2_MCP_LOG_DIR` can also set the default log directory. Use `--log-full-text` when you need complete replayable tool/API text for evaluation; otherwise previews plus hashes keep the log smaller while preserving integrity checks. +The validator checks JSONL parseability, strictly increasing sequence numbers, unique `event_id` values, non-decreasing monotonic time, token field consistency, summary rollup consistency, and a deterministic tokenizer fixture. + +Privacy notes: + +- Token counts are computed after recursive redaction for structured arguments and metadata. +- Response text is counted before preview truncation, but full response text is stored only with `--log-full-text`. +- External usage records may reveal provider/model information supplied by the client; redact those fields client-side if needed. ## Smart State Polling diff --git a/mcp/run_logger.py b/mcp/run_logger.py index 0eeeb375..9595ebb5 100644 --- a/mcp/run_logger.py +++ b/mcp/run_logger.py @@ -12,6 +12,8 @@ from pathlib import Path from typing import Any +from token_usage import TOKEN_FIELDS, TOKEN_SCHEMA_VERSION, TokenEstimator, TokenProfile + SCHEMA_VERSION = "2026-05-10" @@ -20,11 +22,19 @@ "cookie", "password", "secret", - "token", "api_key", "apikey", ) +DEFAULT_REDACT_TOKEN_KEYS = ( + "token", + "access_token", + "auth_token", + "bearer_token", + "refresh_token", + "session_token", +) + class RunLogger: """Append-only JSONL logger with stable event envelopes.""" @@ -37,32 +47,114 @@ def __init__( preview_chars: int = 4000, include_full_text: bool = False, redact_key_parts: tuple[str, ...] = DEFAULT_REDACT_KEY_PARTS, + token_profile: TokenProfile | None = None, ) -> None: self.enabled = enabled self.log_dir = Path(log_dir) self.preview_chars = max(0, preview_chars) self.include_full_text = include_full_text self.redact_key_parts = tuple(part.lower() for part in redact_key_parts) + self.token_estimator = TokenEstimator(token_profile) + self.token_profile = self.token_estimator.profile self.run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + "-" + uuid.uuid4().hex[:8] self.path = self.log_dir / f"run_{self.run_id}.jsonl" + self.summary_path = self.log_dir / f"run_{self.run_id}.summary.json" self._sequence = 0 self._started_at = time.monotonic() self._lock = asyncio.Lock() + self._summary = self._new_summary() def start(self, metadata: dict[str, Any] | None = None) -> None: if not self.enabled: return self.log_dir.mkdir(parents=True, exist_ok=True) - self._write_sync( - self._envelope( - "session_start", - { - "log_path": str(self.path), - "metadata": self.redact(metadata or {}), - }, - ) + payload = { + "log_path": str(self.path), + "summary_path": str(self.summary_path), + "metadata": self.redact(metadata or {}), + "token_profile": self.token_profile.to_dict(), + } + self._summary["session"] = payload + record = self._envelope( + "session_start", + payload, + self._infer_token_usage("session_start", payload, None), ) + self._write_sync(record) + self._update_summary(record) + self.write_summary() + + def _new_summary(self) -> dict[str, Any]: + totals = {field: 0 for field in TOKEN_FIELDS} + return { + "schema_version": SCHEMA_VERSION, + "token_schema_version": TOKEN_SCHEMA_VERSION, + "run_id": self.run_id, + "log_path": str(self.path), + "summary_path": str(self.summary_path), + "token_profile": self.token_profile.to_dict(), + "totals": totals.copy(), + "by_tool_name": {}, + "by_event_type": {}, + "by_state_type": {}, + "by_game_mode": {}, + "by_action_category": {}, + "by_run_phase": {}, + "by_floor": {}, + "polling": { + "poll_events": 0, + "hidden_poll_events": 0, + "hidden_poll_tokens": 0, + }, + "largest_payloads": [], + "external_usage": { + "records": 0, + "total_tokens": 0, + "reconciled_tool_call_ids": [], + "reconciled_event_ids": [], + }, + "invalid_actions": { + "events": 0, + "tokens": 0, + }, + "repeated_states": { + "events": 0, + "tokens": 0, + }, + "replay_artifacts": { + "log_path": str(self.path), + "summary_path": str(self.summary_path), + "log_bytes": 0, + "summary_bytes": 0, + "summary_estimated_tokens": 0, + }, + } + + def write_summary(self) -> None: + if not self.enabled: + return + self._summary["replay_artifacts"]["log_bytes"] = self.path.stat().st_size if self.path.exists() else 0 + summary_preview = dict(self._summary) + summary_text = json.dumps(summary_preview, ensure_ascii=False, sort_keys=True, default=str) + self._summary["replay_artifacts"]["summary_estimated_tokens"] = self.token_estimator.estimate_text(summary_text) + self._write_summary_once() + summary_bytes = self.summary_path.stat().st_size + if self._summary["replay_artifacts"]["summary_bytes"] != summary_bytes: + self._summary["replay_artifacts"]["summary_bytes"] = summary_bytes + self._write_summary_once() + + def _write_summary_once(self) -> None: + with self.summary_path.open("w", encoding="utf-8") as handle: + json.dump( + self._summary, + handle, + ensure_ascii=False, + indent=2, + sort_keys=True, + default=str, + ) + handle.write("\n") async def log( self, @@ -71,18 +163,23 @@ async def log( *, tool_call_id: str | None = None, tool_name: str | None = None, + token_usage: dict[str, Any] | None = None, ) -> None: if not self.enabled: return async with self._lock: + redacted_payload = self.redact(payload or {}) record = self._envelope( event_type, - self.redact(payload or {}), + redacted_payload, + self._infer_token_usage(event_type, redacted_payload, token_usage), tool_call_id=tool_call_id, tool_name=tool_name, ) self._write_sync(record) + self._update_summary(record) + self.write_summary() def redact(self, value: Any) -> Any: if isinstance(value, dict): @@ -116,6 +213,7 @@ def summarize_text(self, text: str | bytes | None) -> dict[str, Any]: "sha256": hashlib.sha256(raw).hexdigest(), "preview": display[: self.preview_chars], "truncated": len(display) > self.preview_chars, + "estimated_tokens": self.token_estimator.estimate_text(display), } if self.include_full_text: summary["text"] = display @@ -136,26 +234,240 @@ def _envelope( self, event_type: str, payload: dict[str, Any], + token_usage: dict[str, Any], *, tool_call_id: str | None = None, tool_name: str | None = None, ) -> dict[str, Any]: self._sequence += 1 + event_id = f"{self.run_id}:{self._sequence}" envelope: dict[str, Any] = { "schema_version": SCHEMA_VERSION, + "token_schema_version": TOKEN_SCHEMA_VERSION, "run_id": self.run_id, + "event_id": event_id, "sequence": self._sequence, "timestamp": datetime.now(UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z"), "monotonic_ms": round((time.monotonic() - self._started_at) * 1000, 3), "event_type": event_type, "payload": payload, } + envelope.update(token_usage) if tool_call_id is not None: envelope["tool_call_id"] = tool_call_id if tool_name is not None: envelope["tool_name"] = tool_name return envelope + def _infer_token_usage( + self, + event_type: str, + payload: dict[str, Any], + supplied: dict[str, Any] | None, + ) -> dict[str, Any]: + if supplied is not None: + usage = self.token_estimator.empty_usage(source=supplied.get("token_source")) + usage.update(supplied) + usage["total_tokens"] = int( + supplied.get("total_tokens") + if supplied.get("total_tokens") is not None + else sum(int(usage.get(field) or 0) for field in TOKEN_FIELDS if field != "total_tokens") + ) + return usage + + if event_type == "tool_call_start": + return self.token_estimator.usage( + input_tokens=self.token_estimator.estimate_jsonable(payload.get("args", {})) + ) + if event_type == "tool_call_result": + return self.token_estimator.usage( + tool_response_tokens=self._summary_estimated_tokens(payload.get("result")) + ) + if event_type == "tool_call_error": + return self.token_estimator.usage(output_tokens=self.token_estimator.estimate_jsonable(payload)) + if event_type == "http_request": + return self.token_estimator.usage(input_tokens=self.token_estimator.estimate_jsonable(payload)) + if event_type == "http_response": + return self.token_estimator.usage(output_tokens=self._summary_estimated_tokens(payload.get("response"))) + if event_type == "http_error": + return self.token_estimator.usage(output_tokens=self.token_estimator.estimate_jsonable(payload)) + if event_type == "agent_decision": + return self.token_estimator.usage(input_tokens=self.token_estimator.estimate_jsonable(payload)) + if event_type == "state_poll": + return self.token_estimator.usage(hidden_poll_tokens=int(payload.get("hidden_poll_tokens") or 0)) + if event_type == "external_token_usage": + return TokenEstimator.normalize_external_usage(payload, self.token_profile) + + return self.token_estimator.empty_usage() + + @staticmethod + def _summary_estimated_tokens(summary: Any) -> int: + return int(summary.get("estimated_tokens") or 0) if isinstance(summary, dict) else 0 + + def _update_summary(self, record: dict[str, Any]) -> None: + usage = {field: int(record.get(field) or 0) for field in TOKEN_FIELDS} + for field, value in usage.items(): + self._summary["totals"][field] += value + + context = self._record_context(record) + self._add_group_tokens(self._summary["by_event_type"], record["event_type"], usage) + if tool_name := record.get("tool_name"): + self._add_group_tokens(self._summary["by_tool_name"], str(tool_name), usage) + for key, group_name in ( + ("state_type", "by_state_type"), + ("game_mode", "by_game_mode"), + ("action_category", "by_action_category"), + ("run_phase", "by_run_phase"), + ("floor", "by_floor"), + ): + if value := context.get(key): + self._add_group_tokens(self._summary[group_name], str(value), usage) + + if record["event_type"] == "state_poll": + self._summary["polling"]["poll_events"] += 1 + if usage["hidden_poll_tokens"] > 0: + self._summary["polling"]["hidden_poll_events"] += 1 + self._summary["polling"]["hidden_poll_tokens"] += usage["hidden_poll_tokens"] + if context.get("state_repeated"): + self._summary["repeated_states"]["events"] += 1 + self._summary["repeated_states"]["tokens"] += usage["total_tokens"] + + if record["event_type"] == "external_token_usage": + self._summary["external_usage"]["records"] += 1 + self._summary["external_usage"]["total_tokens"] += usage["total_tokens"] + payload = record.get("payload", {}) + if isinstance(payload, dict): + if payload.get("related_tool_call_id"): + self._summary["external_usage"]["reconciled_tool_call_ids"].append(payload["related_tool_call_id"]) + if payload.get("related_event_id"): + self._summary["external_usage"]["reconciled_event_ids"].append(payload["related_event_id"]) + + if context.get("invalid_action"): + self._summary["invalid_actions"]["events"] += 1 + self._summary["invalid_actions"]["tokens"] += usage["total_tokens"] + + self._track_largest_payload(record, context) + + @staticmethod + def _add_group_tokens(group: dict[str, Any], key: str, usage: dict[str, int]) -> None: + bucket = group.setdefault(key, {field: 0 for field in TOKEN_FIELDS} | {"events": 0}) + bucket["events"] += 1 + for field, value in usage.items(): + bucket[field] += value + + def _record_context(self, record: dict[str, Any]) -> dict[str, Any]: + payload = record.get("payload", {}) + context: dict[str, Any] = { + "action_category": self._action_category(record.get("tool_name")), + } + if isinstance(payload, dict): + for key in ("state_type", "game_mode", "run_phase"): + if payload.get(key) is not None: + context[key] = payload[key] + if payload.get("state_repeated") is True or payload.get("reason") == "state_repeated": + context["state_repeated"] = True + if str(payload.get("status") or "").lower() == "error": + context["invalid_action"] = True + + for summary_key in ("result", "response"): + parsed = self._parse_summary_json(payload.get(summary_key)) + if isinstance(parsed, dict): + for key in ("state_type", "game_mode"): + if parsed.get(key) is not None: + context.setdefault(key, parsed[key]) + floor = self._extract_floor(parsed) + if floor is not None: + context.setdefault("floor", floor) + if parsed.get("status") == "error": + context["invalid_action"] = True + + if "state_type" in context: + context.setdefault("run_phase", self._run_phase_from_state(str(context["state_type"]))) + return context + + @staticmethod + def _extract_floor(state: dict[str, Any]) -> int | None: + run = state.get("run") + if not isinstance(run, dict): + return None + for key in ("floor", "current_floor", "floor_num"): + value = run.get(key) + if isinstance(value, int): + return value + return None + + @staticmethod + def _parse_summary_json(summary: Any) -> Any: + if not isinstance(summary, dict) or summary.get("truncated"): + return None + text = summary.get("text") or summary.get("preview") + if not isinstance(text, str) or not text: + return None + try: + return json.loads(text) + except json.JSONDecodeError: + return None + + @staticmethod + def _action_category(tool_name: Any) -> str | None: + if not isinstance(tool_name, str): + return None + normalized = tool_name.removeprefix("mp_") + for prefix, category in ( + ("combat_", "combat"), + ("rewards_", "reward"), + ("map_", "map"), + ("event_", "event"), + ("rest_", "rest_site"), + ("shop_", "shop"), + ("deck_", "card_select"), + ("bundle_", "bundle_select"), + ("relic_", "relic_select"), + ("treasure_", "treasure"), + ("crystal_sphere_", "crystal_sphere"), + ("menu_", "menu"), + ): + if normalized.startswith(prefix): + return category + if normalized in {"get_game_state", "log_agent_decision", "log_external_token_usage"}: + return "general" + return normalized + + @staticmethod + def _run_phase_from_state(state_type: str) -> str: + if state_type in {"monster", "elite", "boss", "hand_select"}: + return "combat" + if state_type in {"rewards", "card_reward"}: + return "reward" + if state_type in {"menu", "game_over"}: + return "menu" + return state_type + + def _track_largest_payload(self, record: dict[str, Any], context: dict[str, Any]) -> None: + candidates: list[tuple[str, dict[str, Any]]] = [] + payload = record.get("payload", {}) + if isinstance(payload, dict): + for key in ("result", "response"): + value = payload.get(key) + if isinstance(value, dict) and value.get("estimated_tokens"): + candidates.append((key, value)) + for payload_kind, summary in candidates: + entry = { + "event_id": record["event_id"], + "event_type": record["event_type"], + "tool_name": record.get("tool_name"), + "payload_kind": payload_kind, + "estimated_tokens": int(summary.get("estimated_tokens") or 0), + "length": int(summary.get("length") or 0), + "sha256": summary.get("sha256"), + "state_type": context.get("state_type"), + "game_mode": context.get("game_mode"), + } + largest = self._summary["largest_payloads"] + largest.append(entry) + largest.sort(key=lambda item: item["estimated_tokens"], reverse=True) + del largest[10:] + def _write_sync(self, record: dict[str, Any]) -> None: with self.path.open("a", encoding="utf-8") as handle: json.dump(record, handle, ensure_ascii=False, separators=(",", ":"), default=str) @@ -163,4 +475,4 @@ def _write_sync(self, record: dict[str, Any]) -> None: def _is_sensitive_key(self, key: str) -> bool: key_lower = key.lower() - return any(part in key_lower for part in self.redact_key_parts) + return key_lower in DEFAULT_REDACT_TOKEN_KEYS or any(part in key_lower for part in self.redact_key_parts) diff --git a/mcp/server.py b/mcp/server.py index a949bc93..47bdfc0b 100644 --- a/mcp/server.py +++ b/mcp/server.py @@ -8,6 +8,7 @@ import asyncio import contextvars import functools +import hashlib import inspect import json import os @@ -20,6 +21,7 @@ import httpx from mcp.server.fastmcp import FastMCP from run_logger import RunLogger +from token_usage import TokenEstimator, TokenProfile mcp = FastMCP("sts2") @@ -310,19 +312,27 @@ async def _get_smart_state( attempt = 0 last_json_text = "" last_reason = "not_polled" + previous_state_hash: str | None = None while True: attempt += 1 last_json_text = await getter(detection_params) + state_hash = hashlib.sha256(last_json_text.encode("utf-8", errors="replace")).hexdigest() + state_repeated = previous_state_hash == state_hash + previous_state_hash = state_hash try: state = json.loads(last_json_text) except json.JSONDecodeError: + hidden_tokens = 0 if requested_format == "json" else _run_logger.token_estimator.estimate_text(last_json_text) await _run_logger.log( "state_poll", { "attempt": attempt, "actionable": True, "reason": "state_json_parse_failed", + "hidden_poll_tokens": hidden_tokens, + "state_sha256": state_hash, + "state_repeated": state_repeated, }, tool_call_id=_tool_call_id.get(), tool_name=_tool_name.get(), @@ -331,19 +341,26 @@ async def _get_smart_state( actionable, reason = _state_actionability(state) last_reason = reason + timed_out = time.monotonic() >= deadline + returns_detection_payload = requested_format == "json" and (actionable or timed_out) + hidden_poll_tokens = 0 if returns_detection_payload else _run_logger.token_estimator.estimate_text(last_json_text) await _run_logger.log( "state_poll", { "attempt": attempt, "state_type": state.get("state_type"), + "game_mode": state.get("game_mode"), "actionable": actionable, "reason": reason, "timeout_seconds": timeout, + "hidden_poll_tokens": hidden_poll_tokens, + "state_sha256": state_hash, + "state_repeated": state_repeated, }, tool_call_id=_tool_call_id.get(), tool_name=_tool_name.get(), ) - if actionable or time.monotonic() >= deadline: + if actionable or timed_out: break await asyncio.sleep(min(interval, max(0.0, deadline - time.monotonic()))) @@ -438,6 +455,43 @@ async def log_agent_decision( return json.dumps({"status": "ok", "logged": event}, indent=2) +@logged_tool() +async def log_external_token_usage(usage_json: str) -> str: + """Record exact or externally supplied model token usage. + + Use this when an agent/client has model API usage data that should be + reconciled with the local estimates in the run log. The JSON object should + include a stable reference such as `related_tool_call_id`, `related_event_id`, + or `stable_id`, plus token fields from the provider. + + Example: + {"related_tool_call_id":"...", "model":"claude-sonnet-4.6", + "input_tokens":1200, "output_tokens":400, "token_source":"exact"} + + Args: + usage_json: JSON object with externally supplied usage metadata. + """ + try: + usage = json.loads(usage_json) + except json.JSONDecodeError as exc: + return json.dumps({"status": "error", "error": f"Invalid usage_json: {exc}"}, indent=2) + + if not isinstance(usage, dict): + return json.dumps({"status": "error", "error": "usage_json must be a JSON object"}, indent=2) + + token_usage = TokenEstimator.normalize_external_usage(usage, _run_logger.token_profile) + payload = dict(usage) + payload.update(token_usage) + await _run_logger.log( + "external_token_usage", + payload, + tool_call_id=_tool_call_id.get(), + tool_name=_tool_name.get(), + token_usage=token_usage, + ) + return json.dumps({"status": "ok", "logged": payload}, indent=2, sort_keys=True) + + @logged_tool() async def menu_select(option: str, seed: str | None = None) -> str: """Select a visible menu option. @@ -1392,16 +1446,32 @@ def main(): action="store_true", help="Store complete tool and HTTP response text in logs instead of previews plus hashes only", ) + parser.add_argument( + "--tokenizer-profile", + default="generic_regex_v1", + help="Token accounting profile: generic_regex_v1, openai_cl100k_proxy, anthropic_claude_proxy, or a custom name", + ) + parser.add_argument("--token-model-family", default=None, help="Override token accounting model family") + parser.add_argument("--tokenizer-name", default=None, help="Override tokenizer name recorded in logs") + parser.add_argument("--tokenizer-version", default=None, help="Override tokenizer version recorded in logs") + parser.add_argument("--token-estimation-method", default=None, help="Override token estimation method recorded in logs") args = parser.parse_args() global _base_url, _trust_env, _run_logger _base_url = f"http://{args.host}:{args.port}" _trust_env = not args.no_trust_env + token_profile = TokenProfile.from_profile_name(args.tokenizer_profile).with_overrides( + model_family=args.token_model_family, + tokenizer_name=args.tokenizer_name, + tokenizer_version=args.tokenizer_version, + estimation_method=args.token_estimation_method, + ) _run_logger = RunLogger( enabled=not args.disable_run_log, log_dir=args.log_dir, preview_chars=args.log_preview_chars, include_full_text=args.log_full_text, + token_profile=token_profile, ) _run_logger.start( { @@ -1414,6 +1484,7 @@ def main(): "cwd": os.getcwd(), "log_full_text": args.log_full_text, "log_preview_chars": args.log_preview_chars, + "token_profile": token_profile.to_dict(), } ) diff --git a/mcp/token_usage.py b/mcp/token_usage.py new file mode 100644 index 00000000..c86bef01 --- /dev/null +++ b/mcp/token_usage.py @@ -0,0 +1,173 @@ +"""Deterministic token accounting helpers for STS2 MCP run logs.""" + +from __future__ import annotations + +import json +import re +from dataclasses import asdict, dataclass +from typing import Any + + +TOKEN_SCHEMA_VERSION = "2026-05-10" + +TOKEN_FIELDS = ( + "input_tokens", + "output_tokens", + "tool_response_tokens", + "hidden_poll_tokens", + "total_tokens", +) + +_TOKEN_PATTERN = re.compile(r"\w+|[^\w\s]", re.UNICODE) + + +@dataclass(frozen=True) +class TokenProfile: + """Describes how local token counts were produced.""" + + profile_name: str = "generic_regex_v1" + model_family: str = "generic" + tokenizer_name: str = "sts2_regex" + tokenizer_version: str = TOKEN_SCHEMA_VERSION + estimation_method: str = "regex_words_and_punctuation_v1" + token_source: str = "estimated" + + @classmethod + def from_profile_name(cls, profile_name: str) -> "TokenProfile": + normalized = profile_name.strip().lower() + profiles = { + "generic": cls(profile_name="generic_regex_v1"), + "generic_regex_v1": cls(profile_name="generic_regex_v1"), + "openai": cls( + profile_name="openai_cl100k_proxy", + model_family="openai", + tokenizer_name="cl100k_proxy_regex", + ), + "openai_cl100k_proxy": cls( + profile_name="openai_cl100k_proxy", + model_family="openai", + tokenizer_name="cl100k_proxy_regex", + ), + "anthropic": cls( + profile_name="anthropic_claude_proxy", + model_family="anthropic", + tokenizer_name="claude_proxy_regex", + ), + "anthropic_claude_proxy": cls( + profile_name="anthropic_claude_proxy", + model_family="anthropic", + tokenizer_name="claude_proxy_regex", + ), + } + return profiles.get(normalized, cls(profile_name=profile_name)) + + def with_overrides( + self, + *, + model_family: str | None = None, + tokenizer_name: str | None = None, + tokenizer_version: str | None = None, + estimation_method: str | None = None, + ) -> "TokenProfile": + return TokenProfile( + profile_name=self.profile_name, + model_family=model_family or self.model_family, + tokenizer_name=tokenizer_name or self.tokenizer_name, + tokenizer_version=tokenizer_version or self.tokenizer_version, + estimation_method=estimation_method or self.estimation_method, + token_source=self.token_source, + ) + + def to_dict(self) -> dict[str, str]: + return asdict(self) + + +class TokenEstimator: + """Small deterministic estimator used when exact model usage is unavailable.""" + + def __init__(self, profile: TokenProfile | None = None) -> None: + self.profile = profile or TokenProfile() + + def estimate_text(self, text: str | bytes | None) -> int: + if text is None: + return 0 + if isinstance(text, bytes): + text = text.decode("utf-8", errors="replace") + if not text: + return 0 + return len(_TOKEN_PATTERN.findall(text)) + + def estimate_jsonable(self, value: Any) -> int: + try: + text = json.dumps(value, ensure_ascii=False, sort_keys=True, default=str) + except TypeError: + text = json.dumps(str(value), ensure_ascii=False) + return self.estimate_text(text) + + def empty_usage(self, *, source: str | None = None) -> dict[str, Any]: + usage: dict[str, Any] = {field: 0 for field in TOKEN_FIELDS} + usage.update( + { + "token_source": source or self.profile.token_source, + "tokenizer_name": self.profile.tokenizer_name, + "tokenizer_version": self.profile.tokenizer_version, + "model_family": self.profile.model_family, + "estimation_method": self.profile.estimation_method, + } + ) + return usage + + def usage( + self, + *, + input_tokens: int = 0, + output_tokens: int = 0, + tool_response_tokens: int = 0, + hidden_poll_tokens: int = 0, + source: str | None = None, + overrides: dict[str, Any] | None = None, + ) -> dict[str, Any]: + usage = self.empty_usage(source=source) + usage["input_tokens"] = max(0, int(input_tokens or 0)) + usage["output_tokens"] = max(0, int(output_tokens or 0)) + usage["tool_response_tokens"] = max(0, int(tool_response_tokens or 0)) + usage["hidden_poll_tokens"] = max(0, int(hidden_poll_tokens or 0)) + usage["total_tokens"] = ( + usage["input_tokens"] + + usage["output_tokens"] + + usage["tool_response_tokens"] + + usage["hidden_poll_tokens"] + ) + if overrides: + for key, value in overrides.items(): + if value is not None: + usage[key] = value + return usage + + @staticmethod + def normalize_external_usage( + payload: dict[str, Any], + fallback_profile: TokenProfile | None = None, + ) -> dict[str, Any]: + profile = fallback_profile or TokenProfile() + usage = { + "input_tokens": int(payload.get("input_tokens") or payload.get("prompt_tokens") or 0), + "output_tokens": int(payload.get("output_tokens") or payload.get("completion_tokens") or 0), + "tool_response_tokens": int(payload.get("tool_response_tokens") or 0), + "hidden_poll_tokens": int(payload.get("hidden_poll_tokens") or 0), + "token_source": payload.get("token_source") or "external", + "tokenizer_name": payload.get("tokenizer_name") or payload.get("model") or profile.tokenizer_name, + "tokenizer_version": payload.get("tokenizer_version") or profile.tokenizer_version, + "model_family": payload.get("model_family") or profile.model_family, + "estimation_method": payload.get("estimation_method") or "externally_supplied", + } + supplied_total = payload.get("total_tokens") + usage["total_tokens"] = ( + int(supplied_total) + if supplied_total is not None + else usage["input_tokens"] + + usage["output_tokens"] + + usage["tool_response_tokens"] + + usage["hidden_poll_tokens"] + ) + return usage diff --git a/mcp/validate_run_log.py b/mcp/validate_run_log.py new file mode 100644 index 00000000..0a8c5170 --- /dev/null +++ b/mcp/validate_run_log.py @@ -0,0 +1,137 @@ +"""Validate STS2 MCP JSONL run logs and token summary artifacts.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +from token_usage import TOKEN_FIELDS, TokenEstimator, TokenProfile + + +def _load_jsonl(path: Path) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + with path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError(f"{path}:{line_number}: invalid JSON: {exc}") from exc + if not isinstance(record, dict): + raise ValueError(f"{path}:{line_number}: record must be a JSON object") + records.append(record) + if not records: + raise ValueError(f"{path}: no JSONL records found") + return records + + +def _usage_from_records(records: list[dict[str, Any]]) -> dict[str, int]: + totals = {field: 0 for field in TOKEN_FIELDS} + seen_sequences: set[int] = set() + seen_event_ids: set[str] = set() + previous_sequence = 0 + previous_monotonic = -1.0 + + for index, record in enumerate(records, start=1): + sequence = record.get("sequence") + if not isinstance(sequence, int): + raise ValueError(f"record {index}: sequence must be an integer") + if sequence <= previous_sequence: + raise ValueError(f"record {index}: sequence is not strictly increasing") + if sequence in seen_sequences: + raise ValueError(f"record {index}: duplicate sequence {sequence}") + seen_sequences.add(sequence) + previous_sequence = sequence + + event_id = record.get("event_id") + if not isinstance(event_id, str) or not event_id: + raise ValueError(f"record {index}: missing event_id") + if event_id in seen_event_ids: + raise ValueError(f"record {index}: duplicate event_id {event_id}") + seen_event_ids.add(event_id) + + monotonic_ms = record.get("monotonic_ms") + if not isinstance(monotonic_ms, int | float): + raise ValueError(f"record {index}: monotonic_ms must be numeric") + if monotonic_ms < previous_monotonic: + raise ValueError(f"record {index}: monotonic_ms decreased") + previous_monotonic = float(monotonic_ms) + + for field in TOKEN_FIELDS: + value = record.get(field) + if not isinstance(value, int): + raise ValueError(f"record {index}: {field} must be an integer") + if value < 0: + raise ValueError(f"record {index}: {field} must be non-negative") + totals[field] += value + + expected_total = sum(int(record[field]) for field in TOKEN_FIELDS if field != "total_tokens") + if int(record["total_tokens"]) != expected_total: + raise ValueError( + f"record {index}: total_tokens {record['total_tokens']} != component sum {expected_total}" + ) + + if not record.get("token_source"): + raise ValueError(f"record {index}: missing token_source") + if not record.get("tokenizer_version"): + raise ValueError(f"record {index}: missing tokenizer_version") + + return totals + + +def _validate_summary(summary_path: Path, totals: dict[str, int]) -> None: + summary = json.loads(summary_path.read_text(encoding="utf-8")) + if not isinstance(summary, dict): + raise ValueError(f"{summary_path}: summary must be a JSON object") + summary_totals = summary.get("totals") + if not isinstance(summary_totals, dict): + raise ValueError(f"{summary_path}: missing totals object") + for field, expected in totals.items(): + actual = summary_totals.get(field) + if actual != expected: + raise ValueError(f"{summary_path}: totals.{field} {actual} != JSONL sum {expected}") + + +def _self_test() -> None: + estimator = TokenEstimator(TokenProfile.from_profile_name("generic_regex_v1")) + fixture = "Deal 6 damage." + expected = 4 + actual = estimator.estimate_text(fixture) + if actual != expected: + raise ValueError(f"token fixture mismatch: {fixture!r} expected {expected}, got {actual}") + + usage = estimator.usage(input_tokens=2, output_tokens=3, tool_response_tokens=5, hidden_poll_tokens=7) + if usage["total_tokens"] != 17: + raise ValueError(f"usage rollup mismatch: {usage}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Validate STS2 MCP run log token accounting") + parser.add_argument("log_path", nargs="?", help="Path to run_*.jsonl") + parser.add_argument("--summary", help="Path to run_*.summary.json") + parser.add_argument("--self-test", action="store_true", help="Run deterministic tokenizer fixture checks") + args = parser.parse_args() + + if args.self_test: + _self_test() + + if not args.log_path: + if args.self_test: + print("self-test ok") + return + raise SystemExit("log_path is required unless --self-test is used") + + log_path = Path(args.log_path) + records = _load_jsonl(log_path) + totals = _usage_from_records(records) + summary_path = Path(args.summary) if args.summary else log_path.with_name(log_path.name.replace(".jsonl", ".summary.json")) + if summary_path.exists(): + _validate_summary(summary_path, totals) + print(json.dumps({"status": "ok", "records": len(records), "totals": totals}, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() From e2b7a02f61bcc3c313ed4b3a946b211a2702c970 Mon Sep 17 00:00:00 2001 From: romgenie <5861166+romgenie@users.noreply.github.com> Date: Sun, 10 May 2026 09:23:08 -0400 Subject: [PATCH 3/3] Add prompt and turn accounting to run logs --- README.md | 2 +- mcp/README.md | 23 ++++++- mcp/run_logger.py | 141 ++++++++++++++++++++++++++++++++++++++++ mcp/server.py | 119 ++++++++++++++++++++++++++++++++- mcp/validate_run_log.py | 68 ++++++++++++++++++- 5 files changed, 347 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 73f00cef..c49fc4f3 100644 --- a/README.md +++ b/README.md @@ -104,7 +104,7 @@ The MCP server accepts `--host` and `--port` options if you need non-default set Flag `--no-trust-env` can be used to disable `requests` from picking up proxy settings from the environment, which can cause connection issues if you are running the server in a container. -By default, the MCP bridge records a structured JSONL run log in `logs/run_-.jsonl` and a token rollup summary in `logs/run_-.summary.json` relative to the server working directory. The log captures MCP tool calls, STS2_MCP HTTP requests/responses, smart state-poll decisions, timings, SHA-256 hashes, bounded previews, estimated token usage, and explicit `log_agent_decision` annotations. Use `--disable-run-log` to turn it off, `--log-dir` or `STS2_MCP_LOG_DIR` to choose a location, `--log-preview-chars` to adjust previews, `--log-full-text` when you need complete response text for replay or evaluation, and `--tokenizer-profile` / `--tokenizer-name` options to label the estimator profile. Clients with exact model API usage can call `log_external_token_usage` to reconcile provider token counts with local estimates. +By default, the MCP bridge records a structured JSONL run log in `logs/run_-.jsonl` and a token rollup summary in `logs/run_-.summary.json` relative to the server working directory. The log captures MCP tool calls, STS2_MCP HTTP requests/responses, smart state-poll decisions, timings, SHA-256 hashes, bounded previews, estimated token usage, prompt/message turns, and explicit `log_agent_decision` annotations. Use `--disable-run-log` to turn it off, `--log-dir` or `STS2_MCP_LOG_DIR` to choose a location, `--log-preview-chars` to adjust previews, `--log-full-text` when you need complete response text for replay or evaluation, and `--tokenizer-profile` / `--tokenizer-name` options to label the estimator profile. Clients can call `log_model_message` for first-class prompt/turn accounting and `log_external_token_usage` to reconcile provider token counts with local estimates. `get_game_state` and `mp_get_game_state` wait briefly for actionable states by default, so agents do not immediately receive enemy-turn combat states or transient event/reward screens. Pass `wait_for_actionable=false` to either tool for immediate state reads. diff --git a/mcp/README.md b/mcp/README.md index cea3a798..5f8a6c6a 100644 --- a/mcp/README.md +++ b/mcp/README.md @@ -7,6 +7,7 @@ | `get_game_state(format?, wait_for_actionable?, wait_timeout?, poll_interval?)` | General | Get current game state (`markdown` or `json`), optionally waiting through transient non-actionable states | | `log_agent_decision(summary, reasoning?, intended_action?, alternatives?, confidence?, tags?)` | General | Add a structured decision annotation to the run log | | `log_external_token_usage(usage_json)` | General | Attach exact/external model token usage to a run log event or tool call | +| `log_model_message(role, content, source?, turn_id?, turn_index?, message_id?, related_tool_call_id?, related_event_id?, state_sha256?, content_preview?, exact_*_tokens?, token_source?, model?)` | General | Attach prompt/completion/message records to a stable conversation turn | | `menu_select(option, seed?)` | General | Select a visible menu/game-over option | | `get_profile()` | Profiles | Get active profile progress | | `list_profiles()` | Profiles | List profile slots and active slot | @@ -88,6 +89,7 @@ Logged events include: - `state_poll` and `state_poll_final_format` for smart polling decisions. - `agent_decision` entries from `log_agent_decision`. - `external_token_usage` entries from `log_external_token_usage`. +- `model_message` entries from `log_model_message`. Tool and HTTP results include length, byte count, SHA-256, preview text, truncation status, and estimated token count. Keys containing `authorization`, `cookie`, `password`, `secret`, `token`, `api_key`, or `apikey` are redacted before writing. @@ -116,7 +118,23 @@ By default these counts use a deterministic regex estimator (`generic_regex_v1`) } ``` -The summary artifact rolls up totals by tool name, event type, state type, game mode, action category, run phase, and floor where available. It also tracks hidden polling cost, largest payloads, repeated-state cost, invalid-action cost, replay artifact size/token metadata, and externally supplied usage records. +Use `log_model_message` to attach first-class prompt and turn structure: + +```json +{ + "role": "user", + "content": "Choose the best card to play.", + "source": "agent-wrapper", + "turn_id": "run-1-turn-7", + "turn_index": 7, + "message_id": "msg-007-user", + "state_sha256": "..." +} +``` + +Message roles are `system`, `user`, `assistant`, `tool`, `developer`, and `external`. Each `model_message` record stores `message_id`, `turn_id`, `turn_index`, role/source, content hash, bounded preview metadata, estimated tokens, optional exact token usage, privacy flags, and related `tool_call_id` / `event_id` links when supplied. If `turn_id` or `message_id` is omitted, the bridge generates stable IDs for the current run. Set `content_preview=false` to keep hashes and counts while omitting prompt preview text. + +The summary artifact rolls up totals by tool name, event type, state type, game mode, action category, run phase, floor, token source, turn, role, and message source where available. It also tracks total prompts, total turns, average tokens per turn, tool-result token share, hidden polling token share, hidden polling cost, largest prompts, largest payloads, repeated-state cost, invalid-action cost, replay artifact size/token metadata, and externally supplied usage records. Logging options: @@ -138,12 +156,13 @@ python validate_run_log.py --self-test python validate_run_log.py logs/run_-.jsonl ``` -The validator checks JSONL parseability, strictly increasing sequence numbers, unique `event_id` values, non-decreasing monotonic time, token field consistency, summary rollup consistency, and a deterministic tokenizer fixture. +The validator checks JSONL parseability, strictly increasing sequence numbers, unique `event_id` values, non-decreasing monotonic time, token field consistency, summary rollup consistency, model-message roles, monotonic first-seen turn ordering, message-to-turn linkage, external-usage-to-message reconciliation, and a deterministic tokenizer fixture. Privacy notes: - Token counts are computed after recursive redaction for structured arguments and metadata. - Response text is counted before preview truncation, but full response text is stored only with `--log-full-text`. +- `log_model_message` stores prompt previews by default. Pass `content_preview=false` when clients need hash/token accounting without prompt text snippets. - External usage records may reveal provider/model information supplied by the client; redact those fields client-side if needed. ## Smart State Polling diff --git a/mcp/run_logger.py b/mcp/run_logger.py index 9595ebb5..7f75f469 100644 --- a/mcp/run_logger.py +++ b/mcp/run_logger.py @@ -62,6 +62,7 @@ def __init__( self._sequence = 0 self._started_at = time.monotonic() self._lock = asyncio.Lock() + self._turn_indices: dict[str, int] = {} self._summary = self._new_summary() def start(self, metadata: dict[str, Any] | None = None) -> None: @@ -102,6 +103,7 @@ def _new_summary(self) -> dict[str, Any]: "by_action_category": {}, "by_run_phase": {}, "by_floor": {}, + "by_token_source": {}, "polling": { "poll_events": 0, "hidden_poll_events": 0, @@ -129,6 +131,21 @@ def _new_summary(self) -> dict[str, Any]: "summary_bytes": 0, "summary_estimated_tokens": 0, }, + "conversation": { + "total_prompts": 0, + "total_turns": 0, + "total_messages": 0, + "model_messages": 0, + "tool_result_messages": 0, + "agent_decision_messages": 0, + "average_tokens_per_turn": 0.0, + "tool_result_token_share": 0.0, + "hidden_polling_token_share": 0.0, + "by_turn": {}, + "by_role": {}, + "by_source": {}, + "largest_prompts": [], + }, } def write_summary(self) -> None: @@ -181,6 +198,17 @@ async def log( self._update_summary(record) self.write_summary() + def resolve_turn(self, turn_id: str | None = None, turn_index: int | None = None) -> tuple[str, int]: + if turn_id is None or not str(turn_id).strip(): + if turn_index is not None: + turn_id = f"turn-{turn_index}" + else: + turn_id = f"turn-{len(self._turn_indices) + 1}" + turn_id = str(turn_id) + if turn_id not in self._turn_indices: + self._turn_indices[turn_id] = int(turn_index) if turn_index is not None else len(self._turn_indices) + 1 + return turn_id, self._turn_indices[turn_id] + def redact(self, value: Any) -> Any: if isinstance(value, dict): redacted: dict[str, Any] = {} @@ -297,9 +325,24 @@ def _infer_token_usage( return self.token_estimator.usage(hidden_poll_tokens=int(payload.get("hidden_poll_tokens") or 0)) if event_type == "external_token_usage": return TokenEstimator.normalize_external_usage(payload, self.token_profile) + if event_type == "model_message": + return self._model_message_usage(payload) return self.token_estimator.empty_usage() + def _model_message_usage(self, payload: dict[str, Any]) -> dict[str, Any]: + supplied = payload.get("exact_token_usage") + if isinstance(supplied, dict): + return TokenEstimator.normalize_external_usage(supplied, self.token_profile) + + content_tokens = self._summary_estimated_tokens(payload.get("content")) + role = str(payload.get("role") or "external").lower() + if role == "assistant": + return self.token_estimator.usage(output_tokens=content_tokens) + if role == "tool": + return self.token_estimator.usage(tool_response_tokens=content_tokens) + return self.token_estimator.usage(input_tokens=content_tokens) + @staticmethod def _summary_estimated_tokens(summary: Any) -> int: return int(summary.get("estimated_tokens") or 0) if isinstance(summary, dict) else 0 @@ -322,6 +365,7 @@ def _update_summary(self, record: dict[str, Any]) -> None: ): if value := context.get(key): self._add_group_tokens(self._summary[group_name], str(value), usage) + self._add_group_tokens(self._summary["by_token_source"], str(record.get("token_source") or "unknown"), usage) if record["event_type"] == "state_poll": self._summary["polling"]["poll_events"] += 1 @@ -341,11 +385,19 @@ def _update_summary(self, record: dict[str, Any]) -> None: self._summary["external_usage"]["reconciled_tool_call_ids"].append(payload["related_tool_call_id"]) if payload.get("related_event_id"): self._summary["external_usage"]["reconciled_event_ids"].append(payload["related_event_id"]) + if payload.get("related_message_id"): + self._summary["external_usage"].setdefault("reconciled_message_ids", []).append(payload["related_message_id"]) if context.get("invalid_action"): self._summary["invalid_actions"]["events"] += 1 self._summary["invalid_actions"]["tokens"] += usage["total_tokens"] + if record["event_type"] == "agent_decision": + self._summary["conversation"]["agent_decision_messages"] += 1 + if record["event_type"] == "model_message": + self._update_conversation_summary(record, usage) + + self._update_conversation_derived() self._track_largest_payload(record, context) @staticmethod @@ -468,6 +520,95 @@ def _track_largest_payload(self, record: dict[str, Any], context: dict[str, Any] largest.sort(key=lambda item: item["estimated_tokens"], reverse=True) del largest[10:] + def _update_conversation_summary(self, record: dict[str, Any], usage: dict[str, int]) -> None: + payload = record.get("payload", {}) + if not isinstance(payload, dict): + return + + conversation = self._summary["conversation"] + role = str(payload.get("role") or "external").lower() + source = str(payload.get("source") or "unknown") + turn_id = str(payload.get("turn_id") or "turn-unknown") + turn_index = int(payload.get("turn_index") or 0) + + conversation["total_messages"] += 1 + if role in {"system", "user", "developer", "external"}: + conversation["total_prompts"] += 1 + if role == "assistant": + conversation["model_messages"] += 1 + if role == "tool": + conversation["tool_result_messages"] += 1 + + self._add_group_tokens(conversation["by_role"], role, usage) + self._add_group_tokens(conversation["by_source"], source, usage) + + turn = conversation["by_turn"].setdefault( + turn_id, + { + "turn_id": turn_id, + "turn_index": turn_index, + "start_timestamp": record["timestamp"], + "end_timestamp": record["timestamp"], + "start_monotonic_ms": record["monotonic_ms"], + "end_monotonic_ms": record["monotonic_ms"], + "elapsed_ms": 0, + "event_ids": [], + "message_ids": [], + "related_tool_call_ids": [], + "state_hashes": [], + "roles": {}, + **{field: 0 for field in TOKEN_FIELDS}, + }, + ) + turn["start_timestamp"] = min(turn["start_timestamp"], record["timestamp"]) + turn["end_timestamp"] = max(turn["end_timestamp"], record["timestamp"]) + turn["start_monotonic_ms"] = min(float(turn["start_monotonic_ms"]), float(record["monotonic_ms"])) + turn["end_monotonic_ms"] = max(float(turn["end_monotonic_ms"]), float(record["monotonic_ms"])) + turn["elapsed_ms"] = round(float(turn["end_monotonic_ms"]) - float(turn["start_monotonic_ms"]), 3) + turn["event_ids"].append(record["event_id"]) + if payload.get("message_id"): + turn["message_ids"].append(payload["message_id"]) + if payload.get("related_tool_call_id"): + turn["related_tool_call_ids"].append(payload["related_tool_call_id"]) + if payload.get("state_sha256"): + turn["state_hashes"].append(payload["state_sha256"]) + turn["roles"][role] = turn["roles"].get(role, 0) + 1 + for field, value in usage.items(): + turn[field] += value + + content = payload.get("content") + content_tokens = self._summary_estimated_tokens(content) + if role in {"system", "user", "developer", "external"} and content_tokens: + prompt_entry = { + "event_id": record["event_id"], + "message_id": payload.get("message_id"), + "turn_id": turn_id, + "turn_index": turn_index, + "role": role, + "source": source, + "estimated_tokens": content_tokens, + "length": content.get("length") if isinstance(content, dict) else None, + "sha256": content.get("sha256") if isinstance(content, dict) else None, + "privacy": payload.get("privacy", {}), + } + largest_prompts = conversation["largest_prompts"] + largest_prompts.append(prompt_entry) + largest_prompts.sort(key=lambda item: item["estimated_tokens"], reverse=True) + del largest_prompts[10:] + + def _update_conversation_derived(self) -> None: + conversation = self._summary["conversation"] + conversation["total_turns"] = len(conversation["by_turn"]) + total_tokens = int(self._summary["totals"]["total_tokens"] or 0) + turn_count = int(conversation["total_turns"] or 0) + conversation["average_tokens_per_turn"] = round(total_tokens / turn_count, 3) if turn_count else 0.0 + conversation["tool_result_token_share"] = ( + round(self._summary["totals"]["tool_response_tokens"] / total_tokens, 6) if total_tokens else 0.0 + ) + conversation["hidden_polling_token_share"] = ( + round(self._summary["totals"]["hidden_poll_tokens"] / total_tokens, 6) if total_tokens else 0.0 + ) + def _write_sync(self, record: dict[str, Any]) -> None: with self.path.open("a", encoding="utf-8") as handle: json.dump(record, handle, ensure_ascii=False, separators=(",", ":"), default=str) diff --git a/mcp/server.py b/mcp/server.py index 47bdfc0b..02bcaa3e 100644 --- a/mcp/server.py +++ b/mcp/server.py @@ -195,6 +195,19 @@ def _handle_error(e: Exception) -> str: return f"Error: {e}" +def _sanitize_tool_args(tool_name: str, args: dict[str, Any]) -> dict[str, Any]: + sanitized = dict(args) + if tool_name == "log_model_message" and "content" in sanitized: + content = sanitized["content"] + summary = _run_logger.summarize_text(content) + if sanitized.get("content_preview") is False: + summary.pop("preview", None) + summary.pop("text", None) + summary["preview_redacted"] = True + sanitized["content"] = summary + return sanitized + + def logged_tool(*tool_args: Any, **tool_kwargs: Any) -> Callable[[Callable[..., Awaitable[str]]], Callable[..., Awaitable[str]]]: def decorator(func: Callable[..., Awaitable[str]]) -> Callable[..., Awaitable[str]]: @functools.wraps(func) @@ -205,7 +218,7 @@ async def wrapper(*args: Any, **kwargs: Any) -> str: started = time.perf_counter() try: bound = inspect.signature(func).bind_partial(*args, **kwargs) - tool_args_payload = dict(bound.arguments) + tool_args_payload = _sanitize_tool_args(func.__name__, dict(bound.arguments)) except Exception: tool_args_payload = {"args": list(args), "kwargs": kwargs} @@ -492,6 +505,110 @@ async def log_external_token_usage(usage_json: str) -> str: return json.dumps({"status": "ok", "logged": payload}, indent=2, sort_keys=True) +@logged_tool() +async def log_model_message( + role: str, + content: str, + source: str = "client", + turn_id: str | None = None, + turn_index: int | None = None, + message_id: str | None = None, + related_tool_call_id: str | None = None, + related_event_id: str | None = None, + state_sha256: str | None = None, + content_preview: bool = True, + exact_input_tokens: int | None = None, + exact_output_tokens: int | None = None, + exact_tool_response_tokens: int | None = None, + exact_total_tokens: int | None = None, + token_source: str | None = None, + model: str | None = None, +) -> str: + """Attach a prompt/model/tool message to the current run log. + + This records conversation structure that the MCP bridge cannot infer from + tool calls alone. Use it from an agent wrapper when prompts, completions, or + provider usage are available. + + Args: + role: Message role: system, user, assistant, tool, developer, or external. + content: Message text. Stored as hash/metadata and preview unless disabled. + source: Client, provider, or wrapper name that supplied the message. + turn_id: Stable turn identifier. Generated if omitted. + turn_index: Stable 1-based turn index. Assigned if omitted. + message_id: Stable message identifier. Generated if omitted. + related_tool_call_id: Existing tool_call_id to reconcile against. + related_event_id: Existing event_id to reconcile against. + state_sha256: Optional state hash visible to the model for this turn. + content_preview: Store bounded preview text when true. + exact_*_tokens: Optional exact provider token counts. + token_source: Token source label such as estimated, exact, or external. + model: Provider/model label for exact usage records. + """ + normalized_role = role.lower().strip() + allowed_roles = {"system", "user", "assistant", "tool", "developer", "external"} + if normalized_role not in allowed_roles: + return json.dumps( + {"status": "error", "error": f"role must be one of: {', '.join(sorted(allowed_roles))}"}, + indent=2, + ) + + resolved_turn_id, resolved_turn_index = _run_logger.resolve_turn(turn_id, turn_index) + resolved_message_id = message_id or f"{resolved_turn_id}:message-{uuid.uuid4().hex[:8]}" + content_summary = _run_logger.summarize_text(content) + if not content_preview: + content_summary.pop("preview", None) + content_summary.pop("text", None) + content_summary["preview_redacted"] = True + + exact_usage = None + if any( + value is not None + for value in ( + exact_input_tokens, + exact_output_tokens, + exact_tool_response_tokens, + exact_total_tokens, + ) + ): + exact_usage = { + "input_tokens": exact_input_tokens or 0, + "output_tokens": exact_output_tokens or 0, + "tool_response_tokens": exact_tool_response_tokens or 0, + "total_tokens": exact_total_tokens, + "token_source": token_source or "exact", + "model": model, + } + + payload = { + "message_id": resolved_message_id, + "turn_id": resolved_turn_id, + "turn_index": resolved_turn_index, + "role": normalized_role, + "source": source, + "content": content_summary, + "related_tool_call_id": related_tool_call_id, + "related_event_id": related_event_id, + "state_sha256": state_sha256, + "privacy": { + "content_preview": content_preview, + "content_hash": content_summary.get("sha256"), + "content_truncated": content_summary.get("truncated", False), + "content_preview_redacted": not content_preview, + }, + } + if exact_usage is not None: + payload["exact_token_usage"] = exact_usage + + await _run_logger.log( + "model_message", + payload, + tool_call_id=_tool_call_id.get(), + tool_name=_tool_name.get(), + ) + return json.dumps({"status": "ok", "logged": payload}, indent=2, sort_keys=True) + + @logged_tool() async def menu_select(option: str, seed: str | None = None) -> str: """Select a visible menu option. diff --git a/mcp/validate_run_log.py b/mcp/validate_run_log.py index 0a8c5170..f37b2a96 100644 --- a/mcp/validate_run_log.py +++ b/mcp/validate_run_log.py @@ -82,7 +82,7 @@ def _usage_from_records(records: list[dict[str, Any]]) -> dict[str, int]: return totals -def _validate_summary(summary_path: Path, totals: dict[str, int]) -> None: +def _validate_summary(summary_path: Path, records: list[dict[str, Any]], totals: dict[str, int]) -> None: summary = json.loads(summary_path.read_text(encoding="utf-8")) if not isinstance(summary, dict): raise ValueError(f"{summary_path}: summary must be a JSON object") @@ -93,6 +93,70 @@ def _validate_summary(summary_path: Path, totals: dict[str, int]) -> None: actual = summary_totals.get(field) if actual != expected: raise ValueError(f"{summary_path}: totals.{field} {actual} != JSONL sum {expected}") + _validate_conversation_summary(summary_path, records, summary) + + +def _validate_conversation_summary(summary_path: Path, records: list[dict[str, Any]], summary: dict[str, Any]) -> None: + conversation = summary.get("conversation") + if not isinstance(conversation, dict): + return + + model_messages = [ + record for record in records + if record.get("event_type") == "model_message" and isinstance(record.get("payload"), dict) + ] + turns: dict[str, int] = {} + message_ids: set[str] = set() + previous_first_seen_turn_index = 0 + + for record in model_messages: + payload = record["payload"] + message_id = payload.get("message_id") + turn_id = payload.get("turn_id") + turn_index = payload.get("turn_index") + role = payload.get("role") + content = payload.get("content") + + if not isinstance(message_id, str) or not message_id: + raise ValueError(f"record {record['sequence']}: model_message missing message_id") + if message_id in message_ids: + raise ValueError(f"record {record['sequence']}: duplicate message_id {message_id}") + message_ids.add(message_id) + + if not isinstance(turn_id, str) or not turn_id: + raise ValueError(f"record {record['sequence']}: model_message missing turn_id") + if not isinstance(turn_index, int) or turn_index < 1: + raise ValueError(f"record {record['sequence']}: model_message turn_index must be a positive integer") + if turn_id not in turns: + if turn_index < previous_first_seen_turn_index: + raise ValueError(f"record {record['sequence']}: turn_index decreased on first turn observation") + previous_first_seen_turn_index = turn_index + turns[turn_id] = turn_index + elif turns[turn_id] != turn_index: + raise ValueError(f"record {record['sequence']}: inconsistent turn_index for {turn_id}") + + if role not in {"system", "user", "assistant", "tool", "developer", "external"}: + raise ValueError(f"record {record['sequence']}: invalid model_message role {role!r}") + if not isinstance(content, dict) or not content.get("sha256"): + raise ValueError(f"record {record['sequence']}: model_message content summary missing sha256") + + if conversation.get("total_messages") != len(model_messages): + raise ValueError( + f"{summary_path}: conversation.total_messages {conversation.get('total_messages')} != {len(model_messages)}" + ) + if conversation.get("total_turns") != len(turns): + raise ValueError( + f"{summary_path}: conversation.total_turns {conversation.get('total_turns')} != {len(turns)}" + ) + + external_related_messages = [ + record.get("payload", {}).get("related_message_id") + for record in records + if record.get("event_type") == "external_token_usage" and isinstance(record.get("payload"), dict) + ] + for related_message_id in external_related_messages: + if related_message_id and related_message_id not in message_ids: + raise ValueError(f"{summary_path}: external usage references unknown message_id {related_message_id}") def _self_test() -> None: @@ -129,7 +193,7 @@ def main() -> None: totals = _usage_from_records(records) summary_path = Path(args.summary) if args.summary else log_path.with_name(log_path.name.replace(".jsonl", ".summary.json")) if summary_path.exists(): - _validate_summary(summary_path, totals) + _validate_summary(summary_path, records, totals) print(json.dumps({"status": "ok", "records": len(records), "totals": totals}, indent=2, sort_keys=True))