# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/which-invariants-prevent-silent-token-accounting-corruption-in-opentelemetry-genai-agent-spans-when-tool-calls-re-enter-the-model-loop.md
# Remnant Warrant:    https://corpusfetch.com/c/which-invariants-prevent-silent-token-accounting-corruption-in-opentelemetry-genai-agent-spans-when-tool-calls-re-enter-the-model-loop.warrant.json
# Markdown SHA-256:   4914c462da9480378bcdd79f6201e0f47296fc123bc77ca4e839b818004aed15
# Code AST SHA-256:   90224f2b722ec4356b41bdff03d16e9d62b48075a60538947c52dcac8dbe8abb
# Verified Classes:   TokenAccountingInvariantError, RoundTrip, AgentSpan
# Defensive Guards:   5 fail-closed invariants (zero external deps)

from __future__ import annotations

from dataclasses import dataclass, field


class TokenAccountingInvariantError(RuntimeError):
    """Raised when a GenAI agent span violates the round-trip partition invariant."""


@dataclass
class RoundTrip:
    span_id: str
    input_tokens: int
    output_tokens: int
    tool_call_ids: list[str] = field(default_factory=list)


@dataclass
class AgentSpan:
    span_id: str
    round_trips: list[RoundTrip] = field(default_factory=list)
    emitted_usage_on_agent: bool = False

    def record_round_trip(self, rt: RoundTrip) -> None:
        if any(existing.span_id == rt.span_id for existing in self.round_trips):
            raise TokenAccountingInvariantError(
                f"span_id {rt.span_id!r} reused across round-trips; "
                "each provider call must get a fresh span_id"
            )
        if rt.input_tokens < 0 or rt.output_tokens < 0:
            raise TokenAccountingInvariantError(
                f"negative usage on span {rt.span_id!r}: "
                f"in={rt.input_tokens} out={rt.output_tokens}"
            )
        self.round_trips.append(rt)

    def assert_conformant(self) -> None:
        if self.emitted_usage_on_agent:
            raise TokenAccountingInvariantError(
                "agent span carries gen_ai.usage.* attributes; usage must live "
                "only on child chat spans to keep descendant-summing correct"
            )
        seen_tool_ids: set[str] = set()
        for rt in self.round_trips:
            for call_id in rt.tool_call_ids:
                if call_id in seen_tool_ids:
                    raise TokenAccountingInvariantError(
                        f"duplicate gen_ai.tool.call.id {call_id!r}; "
                        "tool-call ids must be unique within an agent span"
                    )
                seen_tool_ids.add(call_id)

    def total_tokens(self) -> tuple[int, int]:
        self.assert_conformant()
        return (
            sum(rt.input_tokens for rt in self.round_trips),
            sum(rt.output_tokens for rt in self.round_trips),
        )


def reconcile(agent: AgentSpan, provider_reported_output: int) -> None:
    agent.assert_conformant()
    _, traced_out = agent.total_tokens()
    if traced_out != provider_reported_output:
        raise TokenAccountingInvariantError(
            f"trace-derived output tokens ({traced_out}) != provider report "
            f"({provider_reported_output}); likely a collapsed or double-counted "
            "round-trip in the agent span"
        )