atlas-iac/services/hermes/scripts/suite_cli_diagnostics.py

"""Allowlisted CLI metadata; never persist event bodies or provider error text."""
from __future__ import annotations

import json
import math

from suite_contract import CLAUDE_MAX_TURNS, EXECUTION_REVISION

SUBTYPES = {"success", "error_max_turns", "error_max_structured_output_retries",
            "error_max_budget_usd", "error_during_execution"}
STOP_REASONS = {"end_turn", "tool_use", "max_tokens", "stop_sequence", "refusal",
                "pause_turn", "model_context_window_exceeded"}


def number(value):
    """Retain finite nonnegative measurements, excluding booleans and strings."""
    return value if type(value) in (int, float) and 0 <= value <= 2**63 and math.isfinite(value) else None


def enum(value, allowed):
    """Unknown labels are not safe strings: they can contain provider content."""
    return value if isinstance(value, str) and value in allowed else (None if value is None else "other")


def usage_counts(value):
    """Project token counts without carrying arbitrary provider dictionary keys."""
    if not isinstance(value, dict):
        return None
    result = {key: number(value[key]) for key in (
        "input_tokens", "output_tokens", "cache_read_input_tokens", "cache_creation_input_tokens")
        if key in value}
    for key, fields in (("server_tool_use", ("web_search_requests", "web_fetch_requests")),
                        ("cache_creation", ("ephemeral_1h_input_tokens", "ephemeral_5m_input_tokens"))):
        if isinstance(value.get(key), dict):
            result[key] = {field: number(value[key][field]) for field in fields if field in value[key]}
    return result


def object_text(value):
    """Report a possible JSON envelope mismatch without retaining its contents."""
    if not isinstance(value, str):
        return False
    try:
        return isinstance(json.loads(value), dict)
    except ValueError:
        return False


def snapshot(raw, *, exit_code=None, subprocess_timeout_seconds=None, termination_reason=None):
    """Return bounded diagnostic fields from a CLI stream, including failed runs."""
    final, last_assistant = None, {}
    initialized = compacted = tool_candidate = text_candidate = False
    invalid_events = 0
    for line in raw.splitlines():
        try:
            event = json.loads(line)
        except ValueError:
            invalid_events += 1
            continue
        if not isinstance(event, dict):
            invalid_events += 1
            continue
        initialized |= event.get("type") == "system" and event.get("subtype") == "init"
        compacted |= "compact" in str(event.get("subtype", ""))
        if event.get("type") == "result":
            final = event
        if event.get("type") == "assistant" and isinstance(event.get("message"), dict):
            last_assistant = event["message"]
            content = last_assistant.get("content", [])
            for block in content if isinstance(content, list) else []:
                if isinstance(block, dict):
                    tool_candidate |= (block.get("type") == "tool_use" and
                                       block.get("name") == "StructuredOutput" and
                                       isinstance(block.get("input"), dict))
                    text_candidate |= block.get("type") == "text" and object_text(block.get("text"))
    result = final or {}
    subtype = enum(result.get("subtype"), SUBTYPES)
    stop_reason = enum(result.get("stop_reason"), STOP_REASONS)
    limits = result.get("modelUsage", {})
    limits = [{key: number(entry.get(key)) for key in ("contextWindow", "maxOutputTokens")}
              for entry in limits.values() if isinstance(entry, dict)] if isinstance(limits, dict) else []
    return {
        "execution_revision": EXECUTION_REVISION,
        "exit_code": exit_code if type(exit_code) is int and exit_code >= 0 else None,
        "termination_signal": -exit_code if type(exit_code) is int and exit_code < 0 else None,
        "termination_reason": termination_reason or (
            "signal" if type(exit_code) is int and exit_code < 0 else
            "nonzero_exit" if exit_code else "exited" if exit_code == 0 else None),
        "subprocess_timeout_seconds": number(subprocess_timeout_seconds),
        "provider_timeout_seconds": None,
        "max_turns": CLAUDE_MAX_TURNS,
        "turns": number(result.get("num_turns")),
        "turn_limit_reached": subtype == "error_max_turns" if subtype in SUBTYPES else None,
        "structured_retry_limit_reached": subtype == "error_max_structured_output_retries" if subtype in SUBTYPES else None,
        "init_event_seen": initialized,
        "final_event_seen": final is not None,
        "final_event_type": "result" if final is not None else None,
        "final_event_subtype": subtype,
        "final_is_error": result.get("is_error") if type(result.get("is_error")) is bool else None,
        "provider_stop_reason": stop_reason,
        "last_assistant_stop_reason": enum(last_assistant.get("stop_reason"), STOP_REASONS),
        "api_error_status": number(result.get("api_error_status")),
        "invalid_event_count": invalid_events,
        "structured_output_present": result.get("structured_output") is not None if final is not None else None,
        "structured_output_is_object": isinstance(result.get("structured_output"), dict) if final is not None else None,
        "structured_output_location": "result.structured_output" if result.get("structured_output") is not None else None,
        "assistant_structured_tool_input_present": tool_candidate,
        "assistant_json_text_present": text_candidate,
        "final_json_text_present": object_text(result.get("result")),
        "compaction_event_seen": compacted,
        "usage": usage_counts(result.get("usage")),
        "duration_api_ms": number(result.get("duration_api_ms")),
        "cost_usd_estimate": number(result.get("total_cost_usd")),
        "observed_model_limits": limits[:8],
        "configured_max_output_tokens": 64000,
        "reasoning_effort": "medium",
        "reasoning_token_limit": None,
    }