From e1caa611bf034edfe85f89a41628140e05e1b11c Mon Sep 17 00:00:00 2001
From: Bryan Bednarski <bbednarski@nvidia.com>
Date: Mon, 3 Aug 2026 12:28:47 -0600
Subject: [PATCH] fix(relay): preserve skipped turn context

Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>
---
 agent/relay_runtime.py                        | 34 +++++------
 .../test_relay_shared_metrics_runtime.py      | 56 ++++++++++++++++++-
 2 files changed, 72 insertions(+), 18 deletions(-)

diff --git a/agent/relay_runtime.py b/agent/relay_runtime.py
index 3e5082e0664bb..a0a7315796d4d 100644
--- a/agent/relay_runtime.py
+++ b/agent/relay_runtime.py
@@ -482,10 +482,7 @@ class RelayTurnContext:
         default_factory=threading.RLock,
         repr=False,
     )
-    _token: contextvars.Token[RelayTurnContext | None] | None = field(
-        default=None,
-        repr=False,
-    )
+    _previous_turn: RelayTurnContext | None = field(default=None, repr=False)
     _active_registered: bool = field(default=False, repr=False)
     relay_enabled: bool = True
     closed: bool = False
@@ -639,7 +636,8 @@ class RelaySessionCoordinator:
                 )
             except Exception:
                 logger.warning("Hermes Relay turn initialization failed", exc_info=True)
-        turn._token = _CURRENT_TURN.set(turn)
+        turn._previous_turn = _CURRENT_TURN.get()
+        _CURRENT_TURN.set(turn)
         return turn
 
     def end_turn(
@@ -773,16 +771,18 @@ class RelaySessionCoordinator:
 
     @staticmethod
     def _reset_turn_context(turn: RelayTurnContext) -> None:
-        """Reset the originating ContextVar token when called in that context."""
-        if turn._token is None:
+        """Unwind ``turn`` without disturbing a newer context-local turn."""
+        if _CURRENT_TURN.get() is not turn:
             return
-        try:
-            _CURRENT_TURN.reset(turn._token)
-        except ValueError:
-            # A copied async/thread context may own terminal cleanup. Keep the
-            # token so the originating context can clear its stale reference.
-            return
-        turn._token = None
+        previous = turn._previous_turn
+        seen = {id(turn)}
+        while previous is not None and previous.closed:
+            if id(previous) in seen:
+                previous = None
+                break
+            seen.add(id(previous))
+            previous = previous._previous_turn
+        _CURRENT_TURN.set(previous)
 
     @staticmethod
     def release_conversation(lease: ConversationLease) -> None:
@@ -814,7 +814,7 @@ def current_turn() -> RelayTurnContext | None:
 def relay_instrumentation_enabled() -> bool:
     """Return whether this inherited turn may create Relay instrumentation."""
     turn = current_turn()
-    return turn is None or turn.relay_enabled
+    return turn is None or (turn.relay_enabled and not turn.closed)
 
 
 def active_turn(session_id: str | None = None) -> RelayTurnContext | None:
@@ -844,7 +844,9 @@ def resolve_execution_context(
 ) -> tuple[RelayRuntime | None, RelaySession | None, Any]:
     """Resolve one active turn/session parent for managed Relay execution."""
     inherited_turn = current_turn()
-    if inherited_turn is not None and not inherited_turn.relay_enabled:
+    if inherited_turn is not None and (
+        not inherited_turn.relay_enabled or inherited_turn.closed
+    ):
         return None, None, None
     turn = active_turn(session_id)
     if (
diff --git a/tests/hermes_cli/test_relay_shared_metrics_runtime.py b/tests/hermes_cli/test_relay_shared_metrics_runtime.py
index 8f500f4ff5631..5fcb1db626c3b 100644
--- a/tests/hermes_cli/test_relay_shared_metrics_runtime.py
+++ b/tests/hermes_cli/test_relay_shared_metrics_runtime.py
@@ -918,6 +918,60 @@ def test_concurrent_turn_skips_shared_metrics_scope_creation(direct_runtime):
     coordinator.release_conversation(lease)
 
 
+def test_skipped_turn_stays_gated_after_instrumented_turn_ends(direct_runtime):
+    coordinator = relay_runtime.SESSION_COORDINATOR
+    profile_key = relay_runtime.current_profile_key()
+    lease = coordinator.acquire_conversation(
+        profile_key=profile_key,
+        session_id="shared-session",
+        platform="cli",
+    )
+    first = coordinator.begin_turn(lease, turn_id="first", task_id="first-task")
+    second = coordinator.begin_turn(lease, turn_id="second", task_id="second-task")
+    inherited = contextvars.copy_context()
+
+    coordinator.end_turn(first, outcome="success")
+
+    assert relay_runtime.current_turn() is second
+    assert inherited.run(relay_runtime.current_turn) is second
+    assert not relay_runtime.relay_instrumentation_enabled()
+    assert not inherited.run(relay_runtime.relay_instrumentation_enabled)
+    assert relay_runtime.resolve_execution_context("shared-session") == (
+        None,
+        None,
+        None,
+    )
+
+    relay_shared_metrics.observe_lifecycle(
+        "pre_llm_call",
+        session_id="shared-session",
+        task_id="second-task",
+        platform="cli",
+    )
+    inherited.run(
+        relay_shared_metrics.observe_lifecycle,
+        "pre_api_request",
+        session_id="shared-session",
+        task_id="second-task",
+        api_request_id="second-request",
+        platform="cli",
+    )
+
+    assert not [
+        event
+        for event in direct_runtime.events
+        if event[0] == "scope.push"
+        and event[1]
+        in {relay_shared_metrics.TASK_SCOPE, relay_shared_metrics.MODEL_CALL_SCOPE}
+    ]
+
+    coordinator.end_turn(second, outcome="success")
+    assert relay_runtime.current_turn() is None
+    assert inherited.run(relay_runtime.current_turn) is second
+    assert not inherited.run(relay_runtime.relay_instrumentation_enabled)
+    coordinator.release_conversation(lease)
+
+
 
 
 
@@ -1095,5 +1149,3 @@ def test_failed_flush_keeps_daily_export_open_for_later_task(
     assert "Hermes shared-metrics task flush failed" in caplog.text
 
 
-
-