fix(relay): preserve skipped turn context

Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>
This commit is contained in:
Bryan Bednarski 2026-08-03 12:28:47 -06:00
parent 2e65b0c604
commit e1caa611bf
No known key found for this signature in database
2 changed files with 72 additions and 18 deletions

View File

@ -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 (

View File

@ -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