fix(relay): preserve skipped turn context
Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>
This commit is contained in:
parent
2e65b0c604
commit
e1caa611bf
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue