# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/why-does-an-instance-of-a-genai-agent-span-silently-lose-token-attribution-when-a-tool-call-event-lands-after-the-span-has-been-ended.md
# Remnant Warrant:    https://corpusfetch.com/c/why-does-an-instance-of-a-genai-agent-span-silently-lose-token-attribution-when-a-tool-call-event-lands-after-the-span-has-been-ended.warrant.json
# Markdown SHA-256:   a953d9ae38f265a2731e7aff6fa5c48c2b0f81ae1d0454801625044b77a68815
# Code AST SHA-256:   869018d2ec7f7f87e0af011865c5d4932fc824a956099dc2ae15baa8e439179d
# Verified Classes:   TokenAttributionLineageError, Span
# Defensive Guards:   4 fail-closed invariants (zero external deps)

import time
from dataclasses import dataclass, field
from typing import Optional


class TokenAttributionLineageError(RuntimeError):
    """Raised when a tool-call signal is emitted after its owning span ended,
    or when a token-attribution join cannot find a matching tool event."""
    pass


@dataclass
class Span:
    name: str
    span_id: str
    parent_span_id: Optional[str] = None
    events: list = field(default_factory=list)
    attributes: dict = field(default_factory=dict)
    ended: bool = False
    end_time_ns: Optional[int] = None

    def add_event(self, name: str, attributes: dict) -> None:
        if self.ended:
            # This is the silent-loss point in the real SDK: the event is dropped.
            # We surface it instead, because dropping it breaks token lineage.
            raise TokenAttributionLineageError(
                f"event {name!r} emitted after span {self.span_id!r} ended; "
                f"tool-call token attribution for {attributes.get('gen_ai.tool.call.id')!r} "
                f"will be lost"
            )
        self.events.append({"name": name, "attributes": attributes, "time_ns": time.time_ns()})

    def end(self) -> None:
        self.ended = True
        self.end_time_ns = time.time_ns()


def attribute_tokens(agent_span: Span, tool_call_id: str, tokens: int) -> None:
    """Join a tool-call id to the agent span's token accounting.

    Fails if the tool call has no corresponding event on the (already ended) span,
    which is exactly the lineage break observed in the field.
    """
    if agent_span.ended:
        raise TokenAttributionLineageError(
            f"cannot attribute {tokens} tokens to tool call {tool_call_id!r}: "
            f"agent span {agent_span.span_id!r} already ended"
        )
    matched = any(
        e["attributes"].get("gen_ai.tool.call.id") == tool_call_id
        for e in agent_span.events
    )
    if not matched:
        raise TokenAttributionLineageError(
            f"token attribution lineage broken: tool call {tool_call_id!r} "
            f"has no event on span {agent_span.span_id!r}"
        )
    agent_span.attributes["gen_ai.usage.input_tokens"] = (
        agent_span.attributes.get("gen_ai.usage.input_tokens", 0) + tokens
    )


def run_agent_turn(agent_span: Span, tool_call_id: str, tool_latency_s: float) -> None:
    # Correct pattern: tool call is emitted and attributed BEFORE end().
    agent_span.add_event(
        "gen_ai.tool.call",
        {"gen_ai.tool.name": "search", "gen_ai.tool.call.id": tool_call_id},
    )
    attribute_tokens(agent_span, tool_call_id, tokens=1841)
    agent_span.end()

    # Simulate late async callback arriving after end().
    time.sleep(tool_latency_s)
    try:
        agent_span.add_event(
            "gen_ai.tool.call",
            {"gen_ai.tool.name": "search", "gen_ai.tool.call.id": tool_call_id + "-late"},
        )
    except TokenAttributionLineageError:
        raise  # in production this surfaces as a dropped event + broken lineage


if __name__ == "__main__":
    span = Span(name="gen_ai.agent.invoke", span_id="a1b2c3")
    run_agent_turn(span, tool_call_id="call_001", tool_latency_s=0.0)