Fix three runtime defects that silently corrupted metering and parsing

An unreachable max_seconds poisoned the deadline guard process-wide.
setitimer raises OverflowError past the platform's time_t, and it raises
*after* the SIGALRM handler is installed and the process-wide slot taken —
both were left that way, so every later guard found the slot held and fell
back to the async-exception mechanism, which cannot unwind a blocking
syscall. Measured: a 0.3s ceiling honoured at 0.30s, then at 5.00s after a
single 1e10 run. Arming is undone on failure now, and the armed delay is
clamped to what both mechanisms accept.

Every async def node double-charged its token re-reports. The re-report
ledger was keyed by threading.get_ident(), but on_llm_end is sync, so under
ainvoke LangChain dispatches it to a worker thread while the body stays on
the event loop. The automatic charge found no ledger, never recorded the
call, and the node's named re-report — the documented free path — was
charged again. Any node using charge_usage, AgentNode._charge_tokens or
planner.proposal._charge reported double its real spend and hit max_tokens
at half its declared allowance. The ledger is a contextvars scope now,
verified to propagate across that hop; nesting also stops discarding the
enclosing node's ledger.

A bracket in the model's prose hijacked JSON extraction, because only the
first { or [ was ever tried. "Based on the context [lines 3-5]: {...}" was
rejected as unparseable, and — worse — 'Analysis (note [1]): {"supported":
false}' returned a valid [1], substituting a fabricated value for the
verifier's actual answer. Every opener is tried now and the longest parse
wins, which also prefers a complete structure over a nested fragment. Junk
still returns None, so the caller's fail-closed path is unchanged.

A node's end event also reported the movement of the run's *shared* meter
rather than its own spend, so overlapping fan-out workers each absorbed
their siblings' concurrent charges: three workers costing 8 tokens each
traced as 24/16/8, and metrics and cost both reported 48 for 24 tokens of
real work, doubling the estimated bill purely because the work ran in
parallel. Attribution now comes from a per-node scope on the meter.

Each fix ships a test that fails without it, confirmed by stashing the
source and watching them go red.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: Shashankss1205

grapharc/runtime/budget.py

@@ -13,11 +13,12 @@
 
 from __future__ import annotations
 
+import contextvars
 import ctypes
 import signal
 import threading
 import time
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from typing import Any
 
@@ -71,7 +72,7 @@ def _call_key(source: Any) -> Any:
 
 
 class _MeteredCalls:
-    """The model calls the usage callback charged on one thread, by identity.
+    """The model calls the usage callback charged in one node scope, by identity.
 
     Not a count of tokens: a re-report is recognised because it names the same
     call, never because it happens to be the same number of tokens.
@@ -101,6 +102,30 @@ def claim(self, source: Any) -> bool:
         return True
 
 
+class _NodeScope:
+    """One node execution's re-report ledger and its own share of the spend.
+
+    `tokens` exists because the run's meter is shared: subtracting a
+    before/after reading off it attributed every *sibling's* concurrent spend to
+    whichever fan-out worker happened to be running at the time. This counts only
+    the charges made inside this scope, so the same work costs the same whether
+    it runs in parallel or one at a time.
+    """
+
+    __slots__ = ("ledger", "tokens")
+
+    def __init__(self) -> None:
+        self.ledger = _MeteredCalls()
+        self.tokens = 0
+
+
+# The node scope a re-report may claim a metered call from, as (meter, scope).
+# Context-scoped rather than thread-scoped: see `BudgetMeter.automatic_scope`.
+_METERED_SCOPE: contextvars.ContextVar[tuple[Any, _NodeScope] | None] = (
+    contextvars.ContextVar("grapharc_metered_scope", default=None)
+)
+
+
 class BudgetMeter:
     """Thread-safe per-run usage accountant.
 
@@ -116,7 +141,6 @@ def __init__(self, budget: Budget) -> None:
         self._iterations = 0
         self._tokens = 0
         self._started_at = time.monotonic()
-        self._metered: dict[int, _MeteredCalls] = {}
 
     def charge_iteration(self, n: int = 1) -> None:
         with self._lock:
@@ -142,39 +166,80 @@ def charge_tokens(self, n: int, *, automatic: bool = False, source: Any = None)
         run early and is visible in `snapshot()`; under-reporting is invisible
         and arrives on the bill. Pass `source=` to be exact.
 
-        Two shipped callers still charge an unnamed integer for a call the
-        callback already metered — `grapharc.testing.charge_usage` and
-        `AgentNode._charge_tokens` — so a node using either pays twice. Both
-        hold the message: `charge_tokens(total, source=message)` fixes them, and
-        keeps working when they run outside a graph, where nothing metered the
-        call and the charge must land.
+        All three shipped re-reporters now name their source — `testing.
+        charge_usage`, `AgentNode._charge_tokens` and `planner.proposal._charge`
+        each pass `source=message` — so none of them pays twice, and each still
+        works outside a graph, where nothing metered the call and the charge must
+        land.
+
+        What the source is matched against is a *node scope*, not a thread; see
+        `automatic_scope` for why that distinction was load-bearing.
         """
         with self._lock:
-            metered = self._metered.get(threading.get_ident())
+            scope = self._current_scope()
             if automatic:
                 self._tokens += n
-                if metered is not None and source is not None:
-                    metered.record(source)
+                self._attribute(scope, n)
+                if scope is not None and source is not None:
+                    scope.ledger.record(source)
                 return
-            if source is not None and metered is not None and metered.claim(source):
+            if source is not None and scope is not None and scope.ledger.claim(source):
                 return
             self._tokens += n
+            self._attribute(scope, n)
+
+    @staticmethod
+    def _attribute(scope: _NodeScope | None, n: int) -> None:
+        """Credit `n` to the node scope that spent it. Only charges that actually
+        landed on the run total get here, so a dropped re-report is not counted."""
+        if scope is not None:
+            scope.tokens += n
+
+    def _current_scope(self) -> _NodeScope | None:
+        """This node execution's scope, or None outside any. Caller holds the lock."""
+        entry = _METERED_SCOPE.get()
+        # The meter is part of the key so a scope opened by one meter — the
+        # planner runs a sub-meter inside the run's — cannot lend its scope to
+        # a charge made against another.
+        return entry[1] if entry is not None and entry[0] is self else None
+
+    def scope_tokens(self) -> int | None:
+        """Tokens charged inside the current node scope, or None outside one.
+
+        What a single node execution actually spent, as opposed to how the run's
+        shared total moved while it happened to be running.
+        """
+        with self._lock:
+            scope = self._current_scope()
+            return None if scope is None else scope.tokens
 
     @contextmanager
     def automatic_scope(self) -> Iterator[None]:
         """Bound the window in which a metered call may be re-reported for free.
 
-        Scoped per thread and opened once per node, so a re-report in one node
-        can never claim a call metered by another.
+        Scoped to the calling *context* and opened once per node, so a re-report
+        in one node can never claim a call metered by another.
+
+        A `contextvars` scope rather than a thread-keyed one, because the runtime's
+        usage callback does not always run on the thread the node body runs on.
+        `on_llm_end` is sync, so under `ainvoke`/`astream` LangChain dispatches it
+        to a worker thread while the body stays on the event loop. Keyed by thread
+        ident, the automatic charge then found no ledger, never recorded the call,
+        and the node's *named* re-report — the documented free path — was charged a
+        second time: every `async def` node using `charge_usage`, `AgentNode.
+        _charge_tokens` or `planner.proposal._charge` reported double its real
+        spend and hit `max_tokens` at half its declared allowance. LangChain copies
+        the context across that hop, so the callback and the body share one ledger.
+
+        Contexts also nest properly. `reset(token)` restores an enclosing scope's
+        ledger, where popping a thread-keyed entry discarded it — so an inner scope
+        used to make the outer node's remaining re-reports pay twice.
         """
-        ident = threading.get_ident()
-        with self._lock:
-            self._metered[ident] = _MeteredCalls()
+        token = _METERED_SCOPE.set((self, _NodeScope()))
         try:
             yield
         finally:
-            with self._lock:
-                self._metered.pop(ident, None)
+            _METERED_SCOPE.reset(token)
 
     @property
     def iterations(self) -> int:
@@ -255,6 +320,15 @@ def snapshot(self) -> dict[str, float | int]:
 # node alive; the cost while a node is being torn down is one timer per 50ms.
 _REARM_SECONDS = 0.05
 
+# The longest delay both mechanisms can actually be armed with. `setitimer`
+# raises `OverflowError` past the platform's `time_t` (~2**31 seconds), and
+# `threading.Timer` accepts a larger value but crashes its own thread once the
+# underlying `wait` exceeds `threading.TIMEOUT_MAX`. A `max_seconds` beyond this
+# is ~68 years, which no process reaches, so clamping the *armed delay* costs no
+# enforcement: the deadline is still computed from the meter, and the guard's
+# exit-time check still refuses a node that overran.
+_MAX_ARMABLE_SECONDS = min(2.0**31 - 1, threading.TIMEOUT_MAX)
+
 
 def _async_raise(thread_id: int, exc: type[BaseException] | None) -> None:
     """Queue `exc` in another thread, or clear a queued one when `exc` is None."""
@@ -329,18 +403,29 @@ def detail() -> str:
     state: dict[str, Any] = {"armed": True, "fired": False, "timer": None}
     lock = threading.Lock()
     thread_id = threading.get_ident()
-    use_signal = _signal_slot_available() and _SIGNAL_SLOT.acquire(blocking=False)
-
-    if use_signal:
+    # What the timers are armed with, as opposed to what the deadline *is*.
+    armable = min(remaining, _MAX_ARMABLE_SECONDS)
+
+    def arm_signal() -> Callable[[], None] | None:
+        """Arm SIGALRM and return its disarm, or return None to use mechanism 2.
+
+        Arming is undone on failure rather than left half-done. `setitimer`
+        rejects a `remaining` beyond the platform's `time_t` — `float("inf")`,
+        or a plausible "effectively unlimited" like `1e10` — and it raises
+        *after* the handler is installed and the slot is taken. Letting that
+        propagate leaked both for the life of the process: every later guard
+        found the slot held and silently degraded to mechanism 2, which cannot
+        unwind a blocking syscall, and a stray SIGALRM anywhere in the program
+        would raise `NodeDeadlineExceeded` citing a finished run's meter.
+        """
+        if not (_signal_slot_available() and _SIGNAL_SLOT.acquire(blocking=False)):
+            return None
 
         def on_alarm(signum: int, frame: object) -> None:
             state["fired"] = True
             raise NodeDeadlineExceeded(detail())
 
         previous_handler = signal.signal(signal.SIGALRM, on_alarm)
-        # The third argument is the repeat interval: the kernel re-arms the
-        # alarm for us, so a swallowed SIGALRM is followed by another one.
-        signal.setitimer(signal.ITIMER_REAL, remaining, _REARM_SECONDS)
 
         def disarm() -> None:
             # An alarm landing mid-disarm raises out of `setitimer`, so the
@@ -350,6 +435,20 @@ def disarm() -> None:
             finally:
                 signal.signal(signal.SIGALRM, previous_handler)
                 _SIGNAL_SLOT.release()
+
+        try:
+            # The third argument is the repeat interval: the kernel re-arms the
+            # alarm for us, so a swallowed SIGALRM is followed by another one.
+            signal.setitimer(signal.ITIMER_REAL, armable, _REARM_SECONDS)
+        except (OverflowError, OSError, ValueError):
+            disarm()
+            return None
+        return disarm
+
+    disarm_signal = arm_signal()
+
+    if disarm_signal is not None:
+        disarm = disarm_signal
     else:
 
         def fire() -> None:
@@ -371,7 +470,7 @@ def rearm(delay: float) -> None:
             timer.start()
 
         with lock:
-            rearm(remaining)
+            rearm(armable)
 
         def disarm() -> None:
             with lock:

grapharc/runtime/graph.py

@@ -657,11 +657,20 @@ def _leave(
         # one, so `observe.cost` can keep a measured figure apart from a
         # rate-card estimate instead of recording a guess as a fact.
         models = getattr(usage, "models", ())
+        # This node's own spend, from the meter's per-node scope, not the movement
+        # of the run's shared total while the node ran. Under fan-out those differ:
+        # the workers' windows overlap, so every worker used to be credited with
+        # each sibling's concurrent spend — three workers costing 8 tokens each
+        # traced as 24/16/8, and `metrics` and `cost` both reported 48 for 24
+        # tokens of real work, doubling the estimated bill purely because the work
+        # ran in parallel. Falls back to the difference when no scope was open, so
+        # a caller driving `_leave` outside `charging()` still gets a figure.
+        spent = getattr(usage, "tokens", None)
         emit(
             "end",
             state_delta=delta,
             duration_ms=duration_ms,
-            tokens=ctx.meter.tokens - tokens_before,
+            tokens=ctx.meter.tokens - tokens_before if spent is None else spent,
             cost_usd=getattr(usage, "cost_usd", None),
             model=models[0] if len(models) == 1 else None,
         )

grapharc/runtime/parsing.py

@@ -21,11 +21,8 @@
 _FENCE = re.compile(r"```(?:json|JSON)?\s*(.*?)```", re.DOTALL)
 
 
-def _balanced_span(text: str) -> str | None:
-    """The first balanced {...} or [...] region, ignoring braces inside strings."""
-    start = next((i for i, ch in enumerate(text) if ch in "{["), None)
-    if start is None:
-        return None
+def _span_from(text: str, start: int) -> str | None:
+    """The balanced {...} or [...] region opening at `start`, or None if unclosed."""
     opener = text[start]
     closer = "}" if opener == "{" else "]"
     depth = 0
@@ -52,20 +49,44 @@ def _balanced_span(text: str) -> str | None:
     return None
 
 
+def _balanced_spans(text: str) -> list[str]:
+    """Every balanced {...} or [...] region, longest first.
+
+    Every opener is tried, not just the first one in the text. Taking only the
+    first meant any bracket in the model's *prose* hijacked the span and the real
+    JSON was never reached: `Based on the context [lines 3-5]: {...}` yielded the
+    unparseable `[lines 3-5]` and the reply was rejected, and — worse —
+    `Analysis (note [1]): {"supported": false}` yielded a perfectly valid `[1]`,
+    substituting a fabricated value for the model's actual answer.
+
+    Longest first is what makes the ranking safe. It prefers a complete structure
+    over both a prose fragment that happens to parse and a nested piece of the
+    answer itself, so `{"claims": [{...}]}` returns the whole object rather than
+    the inner list.
+    """
+    spans = [
+        span
+        for i, ch in enumerate(text)
+        if ch in "{[" and (span := _span_from(text, i)) is not None
+    ]
+    return sorted(spans, key=len, reverse=True)
+
+
 def extract_json(content: Any) -> Any | None:
     """Best-effort JSON from a model reply. None when nothing valid is found."""
     text = content if isinstance(content, str) else str(content)
     text = text.strip()
     if not text:
         return None
 
+    # Whole reply first, then the fence, then balanced spans: the earlier a
+    # candidate is, the more of the model's reply it accounts for, so a fenced
+    # top-level array still wins over any span found inside it.
     candidates = [text]
     fenced = _FENCE.search(text)
     if fenced:
         candidates.append(fenced.group(1).strip())
-    span = _balanced_span(text)
-    if span:
-        candidates.append(span)
+    candidates.extend(_balanced_spans(text))
 
     for candidate in candidates:
         try:

grapharc/runtime/usage.py

@@ -76,6 +76,11 @@ def __init__(self, meter: BudgetMeter) -> None:
         # `observe.cost` keeps recorded and estimated figures apart.
         self.cost_usd: float | None = None
         self.models: list[str] = []
+        # This node execution's own token spend, stamped by `charging` when the
+        # scope closes. Read from the meter's node scope rather than differenced
+        # off the run total, which credited a fan-out worker with whatever its
+        # siblings spent while it was running.
+        self.tokens = 0
 
     def _record_price(self, response: LLMResult) -> None:
         """Accumulate the price a backend reported through `llm_output`.
@@ -129,12 +134,19 @@ def charging(meter: BudgetMeter) -> Iterator[MeterCallbackHandler]:
 
     Also opens the meter's automatic scope, so a call metered here can be
     re-reported by hand inside this block — and only inside it.
+
+    On the way out it stamps the scope's token tally onto the handler. Callers
+    read `handler.tokens` after the block has closed — the node wrapper reports
+    it on the node's `end` event — and by then the scope itself is gone.
     """
     handler = MeterCallbackHandler(meter)
     token = _ACTIVE_HANDLER.set(handler)
     try:
         with meter.automatic_scope():
-            yield handler
+            try:
+                yield handler
+            finally:
+                handler.tokens = meter.scope_tokens() or 0
     finally:
         _ACTIVE_HANDLER.reset(token)