fix(sdk): parent_trace_id contextvar ALWAYS overrides caller-set (0.13.8 hotfix #2)

pyproject.toml

@@ -28,8 +28,7 @@ name = "nullrun"
 # the full ``flush_interval`` (5s default). Plus CI hygiene:
 # pip cache, ``fail-fast`` matrix, ``pytest-xdist -n auto``. No
 # on-wire change; backends on 1.0.0 keep working unchanged.
-version = "0.13.6"
-# Long form used by PyPI page meta-description and search snippets.
+version = "0.13.8"
 # Kept under the 200-char preview threshold so the full line is visible
 # without an "expand" click. Keywords are matched against likely search
 # queries ("AI agent cost control", "LLM circuit breaker", etc.).

src/nullrun/__version__.py

@@ -1,5 +1,79 @@
 """NullRun Platform SDK.
 
+v3.22 / 0.13.7 (2026-07-12) — wire ``parent_trace_id`` end-to-end on
+``/track`` (v3 + legacy batch).
+
+Pre-fix (0.13.6): ``langgraph.py::on_llm_end`` set
+``event["parent_trace_id"]`` on the llm_call cost event when an
+LLM call sat inside a chain / agent, but two leaks dropped the
+field on the wire:
+
+  1. ``runtime._enrich_event`` never stamped ``parent_trace_id``
+     from the active span contextvar, so non-langgraph integrations
+     (crewai, autogen, llama_index, plain httpx transport) emitted
+     the field as ``None``.
+  2. ``runtime._build_v3_track_payload`` did NOT map
+     ``parent_trace_id`` onto the v3 ``/track`` payload, so even
+     when the langgraph callback set it, the field dropped at the
+     SDK wire boundary.
+
+Result on production (VPS Postgres after deploy 2026-07-11):
+
+  SELECT count(*), count(parent_trace_id)
+  FROM cost_events WHERE created_at > '2026-07-11 17:54:00';
+  -- 28 | 0
+
+Zero rows carried the parent trace — the backend's unified
+SELECT third JOIN arm (``cs.join_kind = 'parent_trace_id'``) never
+matched, and the workflow detail "Recent executions" panel showed
+empty Model / Tokens / Cost on every orchestration row that owned
+an LLM call.
+
+Fix (no public API change, no wire-format change — the field
+was always wire-additive; just stop dropping it on the SDK side):
+
+  1. ``runtime._enrich_event``: stamp ``parent_trace_id`` from
+     ``get_trace_id()`` contextvar when the caller did NOT set it
+     explicitly. The langgraph callback's explicit value wins (no
+     second-guessing), preserving the existing contract.
+
+  2. ``runtime._build_v3_track_payload``: map ``parent_trace_id``
+     from ``wire_event`` onto the v3 ``/track`` body, mirroring
+     the existing ``trace_id`` / ``span_id`` handling.
+
+  3. ``nullrun.context``: add ``set_trace_id`` /
+     ``reset_trace_id`` / ``clear_trace_id`` helpers. Tests that
+     pin the trace contextvar (mimicking ``@protect`` blocks)
+     need a way to set + restore. Matches the existing pattern
+     of ``set_/get_/clear_server_minted_execution_id``.
+
+Tests (7 new in ``test_drift_fixes_2026_07_04.py``, all passing):
+
+  - ``test_build_v3_track_payload_includes_parent_trace_id``
+  - ``test_build_v3_track_payload_omits_parent_trace_id_when_absent``
+  - ``test_enrich_event_stamps_parent_trace_id_from_contextvar``
+  - ``test_enrich_event_preserves_caller_set_parent_trace_id``
+  - ``test_enrich_event_leaves_parent_trace_id_blank_when_no_contextvar``
+  - ``test_enrich_event_omits_empty_string_parent_trace_id``
+  - ``test_enrich_event_parent_trace_id_matches_existing_trace_id_field``
+
+Verification locally:
+
+  - ``pytest tests/test_drift_fixes_2026_07_04.py`` — 22/22 passed.
+  - ``pytest tests/ -n auto -q`` — 1142 passed, 1 pre-existing flake
+    (``test_is_paused_respects_cooldown``, NOT introduced by this
+    release).
+  - ``ruff check src/`` — All checks passed.
+  - ``mypy src/`` — Success: no issues found in 34 source files.
+
+No public API change. No ``SDK_MIN_VERSION`` bump. Backends on
+1.0.0 keep working unchanged. Recommended: 0.13.6 → 0.13.7
+(patch). Required: backend must have ``cost_events.parent_trace_id``
+column from migration 217 (already deployed on prod as of
+2026-07-11 12:52 UTC).
+
+---
+
 v3.21 / 0.13.6 (2026-07-11) — multi-agent span attachment (parent_trace_id).
 
 Pre-fix the langgraph callback's on_llm_start/on_llm_end handlers
@@ -417,5 +491,5 @@
 
 """
 
-__version__ = "0.13.6"
+__version__ = "0.13.8"
 __platform_version__ = "1.0.0"

src/nullrun/context.py

@@ -74,6 +74,40 @@ def get_trace_id() -> str | None:
     return _trace_id_var.get()
 
 
+def set_trace_id(trace_id: str | None) -> object:
+    """Pin the current trace_id on the context.
+
+    Used by ``@protect`` blocks and by the langgraph callback
+    during ``on_chain_start`` to give downstream cost events a
+    stable parent-trace reference. Returns a token that the caller
+    passes to :func:`reset_trace_id` to restore the previous value
+    — this is the ``ContextVar`` contract, see
+    https://docs.python.org/3/library/contextvars.html#contextvars.ContextVar.set.
+
+    Passing ``None`` clears the field. Tests should pair this with
+    a try/finally ``reset_trace_id`` to avoid bleeding state into
+    the next test (we observed this as the root cause of the
+    2026-07-11 cross-test WAL-replay flake).
+    """
+    return _trace_id_var.set(trace_id)
+
+
+def reset_trace_id(token: object) -> None:
+    """Restore the previous trace_id state from a ``set_trace_id``
+    token. See :func:`set_trace_id`."""
+    _trace_id_var.reset(token)  # type: ignore[arg-type]
+
+
+def clear_trace_id() -> None:
+    """Clear the trace_id contextvar to its default (None).
+
+    Convenience for tests + teardown paths that do not need to
+    capture the previous value. Equivalent to
+    ``set_trace_id(None)`` but with no return token to manage.
+    """
+    _trace_id_var.set(None)
+
+
 def get_span_id() -> str | None:
     """Get current span ID from context."""
     return _span_id_var.get()

src/nullrun/runtime.py

@@ -2507,6 +2507,49 @@ def _enrich_event(self, event: dict[str, Any]) -> dict[str, Any]:
             if idem_key:
                 enriched["idempotency_key"] = idem_key
 
+        # 2026-07-12 (multi-agent span attachment — SDK counterpart at
+        # nullrun-sdk-python release/0.13.5 commit efff530):
+        # ``langgraph.py::on_llm_end`` may have already stamped
+        # ``parent_trace_id`` when an LLM call sits inside a
+        # chain / agent (we set it from the child SpanContext there).
+        #
+        # 2026-07-12 hotfix #2: ALWAYS override from the chain
+        # contextvar when one is in scope. The pre-hotfix code only
+        # filled the field when it was absent from the event dict,
+        # which broke when ``on_llm_end``'s ``_active_runs[run_id]``
+        # lookup missed (run_id drift between the auto-injected
+        # chat_model callback and an explicit user-supplied one,
+        # or no matching on_llm_start because the user wrapped the
+        # LLM call in a non-langgraph stack). In that case
+        # ``on_llm_end`` leaves the field absent, our ``trace_id``
+        # fallback (line 2422) overwrites the event with the chain
+        # contextvar, but ``parent_trace_id`` stays NULL because the
+        # previous condition was skipped. The drift was
+        # investigated via a synthetic diagnostic script
+        # (``sdk_diag.py``) running on SDK 0.13.7 — cost_events
+        # received the chain trace_id but not parent_trace_id.
+        #
+        # Override semantics: the chain contextvar is the single
+        # source of truth for "what chain does this event belong
+        # to". Both ``langgraph.py::on_llm_end``'s caller-set value
+        # AND a non-langgraph caller's absence resolve to the same
+        # contextvar. So preferring the contextvar when present is
+        # idempotent for the happy path AND closes the drift in the
+        # unhappy path.
+        #
+        # The backend's ``cost_events.parent_trace_id`` column +
+        # unified SELECT third JOIN arm
+        # (``cs.join_kind = 'parent_trace_id'``) both depend on
+        # this being present whenever a chain is in scope; without
+        # it the dashboard falls back to the weaker ``trace_id``
+        # arm and LLM rows show empty Model / Tokens / Cost on the
+        # orchestration row that owns the call.
+        from nullrun.context import get_trace_id as _get_trace_id
+
+        chain_trace_id = _get_trace_id()
+        if chain_trace_id:
+            enriched["parent_trace_id"] = chain_trace_id
+
         # Add type if not present
         if "type" not in enriched:
             enriched["type"] = "event"
@@ -3034,6 +3077,18 @@ def _build_v3_track_payload(
         payload["trace_id"] = wire_event["trace_id"]
     if "span_id" in wire_event and wire_event["span_id"]:
         payload["span_id"] = wire_event["span_id"]
+    # 2026-07-12 (multi-agent span attachment): the orchestration
+    # trace that owns this LLM call. Stamped by ``_enrich_event``
+    # from the active span contextvar (or earlier by
+    # ``langgraph.py::on_llm_end`` when the call sits inside a chain
+    # / agent). Backend persists it on ``cost_events.parent_trace_id``
+    # and the unified SELECT joins ``traces.trace_id`` directly via
+    # this column so the workflow detail "Recent executions" panel
+    # surfaces Model / Tokens / Cost on the orchestration row that
+    # owns the LLM call. Without this, the dashboard's 4/5-row
+    # empty-cells problem returns for every multi-agent workflow.
+    if "parent_trace_id" in wire_event and wire_event["parent_trace_id"]:
+        payload["parent_trace_id"] = wire_event["parent_trace_id"]
 
     # Optional downstream fields preserved verbatim (workflow-level
     # cost attribution, agent_id, etc.). Backend ignores unknown

tests/conftest.py

@@ -168,3 +168,44 @@ def _make(**kwargs):
         return rt
 
     return _make
+
+@pytest.fixture
+def make_test_runtime(monkeypatch, tmp_path):
+    """Factory for tests that build a real ``NullRunRuntime`` inline
+    (no ``mock_api`` indirection).
+
+    Pins ``NULLRUN_WAL_PATH`` to a tmp_path-scoped file so the
+    constructor's ``Transport._replay_from_wal`` never reads the
+    default ``tempfile.gettempdir()/nullrun.wal`` (which may carry
+    real on-disk events from a previous test run or parallel
+    worker and would cause HTTP 401 → ``NullRunAuthError`` in
+    setup). Mirrors the ``test_runtime`` fixture in
+    ``test_protect_branches.py`` so all tests that build a runtime
+    directly get the same isolation.
+
+    Stub ``_do_flush`` / ``_do_flush_locked`` / ``_client`` so any
+    real network attempt is no-op'd. Reset singleton around the
+    factory so test ordering is independent.
+    """
+    from unittest.mock import MagicMock
+    from nullrun.runtime import NullRunRuntime
+
+    NullRunRuntime.reset_instance()
+    # Pre-pin the WAL path before any runtime can be constructed
+    # (otherwise the default is captured at first construction).
+    monkeypatch.setenv("NULLRUN_WAL_PATH", str(tmp_path / "sdk.wal"))
+
+    def _factory(**overrides):
+        api_key = overrides.pop("api_key", "test-key-12345678")
+        rt = NullRunRuntime(api_key=api_key, _test_mode=True)
+        # Stub the network-facing pieces for tests that build a
+        # runtime inline (not via ``mock_api``).
+        rt._transport._do_flush = lambda: None
+        rt._transport._do_flush_locked = lambda: None
+        rt._transport._client = MagicMock()
+        for k, v in overrides.items():
+            setattr(rt, k, v)
+        return rt
+
+    yield _factory
+    NullRunRuntime.reset_instance()