Merge pull request #74864 from bbednarski9/fix/relay-concurrent-turn-scopes

fix(relay): avoid concurrent turn scope corruption
This commit is contained in:
Jeffrey Quesnelle 2026-08-04 12:04:42 -04:00 committed by GitHub
commit 42708f8bb3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 317 additions and 35 deletions

View File

@ -482,11 +482,9 @@ 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
@ -600,7 +598,28 @@ class RelaySessionCoordinator:
if lease.released:
raise RuntimeError("Hermes Relay conversation lease is released")
turn = RelayTurnContext(lease=lease, turn_id=turn_id, task_id=task_id)
if isinstance(lease.host, RelayRuntime) and lease.session is not None:
key = (lease.profile_key, lease.session_id)
with self._active_turns_lock:
active = self._active_turns.get(key)
if active:
# A Relay session owns one physical scope stack. Concurrent
# Hermes turns would create sibling scopes on that stack, but
# their completion order is not guaranteed to be LIFO.
turn.relay_enabled = False
logger.warning(
"Skipping Relay instrumentation for concurrent Hermes turn "
"%s in session %s",
turn_id,
lease.session_id,
)
else:
self._active_turns[key] = {id(turn)}
turn._active_registered = True
if (
turn.relay_enabled
and isinstance(lease.host, RelayRuntime)
and lease.session is not None
):
try:
turn.handle = lease.host.run_in_session(
lease.session,
@ -617,11 +636,8 @@ class RelaySessionCoordinator:
)
except Exception:
logger.warning("Hermes Relay turn initialization failed", exc_info=True)
turn._token = _CURRENT_TURN.set(turn)
key = (lease.profile_key, lease.session_id)
with self._active_turns_lock:
self._active_turns.setdefault(key, set()).add(id(turn))
turn._active_registered = True
turn._previous_turn = _CURRENT_TURN.get()
_CURRENT_TURN.set(turn)
return turn
def end_turn(
@ -755,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:
@ -793,10 +811,21 @@ def current_turn() -> RelayTurnContext | None:
return _CURRENT_TURN.get()
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 and not turn.closed)
def active_turn(session_id: str | None = None) -> RelayTurnContext | None:
"""Return a live turn only when it belongs to the active profile/session."""
turn = current_turn()
if turn is None or turn.closed or turn.lease.released:
if (
turn is None
or not turn.relay_enabled
or turn.closed
or turn.lease.released
):
return None
if turn.lease.profile_key != current_profile_key():
return None
@ -814,6 +843,11 @@ def resolve_execution_context(
session_id: str,
) -> 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 or inherited_turn.closed
):
return None, None, None
turn = active_turn(session_id)
if (
turn is not None

View File

@ -671,6 +671,8 @@ def observe_lifecycle(hook_name: str, **kwargs: Any) -> None:
"""Project one Hermes lifecycle event into the core Relay integration."""
if not handles_hook(hook_name):
return
if not relay_runtime.relay_instrumentation_enabled():
return
runtime = _get_runtime()
if runtime is None:
return

View File

@ -7639,11 +7639,14 @@ class AIAgent:
turn_id=relay_turn_id,
task_id=effective_task_id,
)
start_task_run(
**task_context,
parent_session_id=getattr(self, "_parent_session_id", None) or "",
)
task_started = True
# Keep existing tests and external relay-runtime shims that return
# a minimal turn object compatible with the new opt-out flag.
if getattr(relay_turn, "relay_enabled", True):
start_task_run(
**task_context,
parent_session_id=getattr(self, "_parent_session_id", None) or "",
)
task_started = True
# Publish the conversation id for ambient Nous Portal tagging. Every
# LLM call made inside this turn — main loop, compression, vision,
# web_extract, session_search, MoA slots, background-review forks
@ -7689,8 +7692,9 @@ class AIAgent:
relay_turn,
outcome=relay_outcome,
)
task_finished = True
finish_task_run(**task_context, result=result)
if task_started:
task_finished = True
finish_task_run(**task_context, result=result)
return result
except BaseException as exc:
if isinstance(exc, (KeyboardInterrupt, InterruptedError)) or (

View File

@ -32,6 +32,7 @@ class _Relay:
self._starts: dict[Any, dict[str, Any]] = {}
self._scope_starts: dict[Any, dict[str, Any]] = {}
self._scope = contextvars.ContextVar("relay_scope", default=None)
self._scope_stack = contextvars.ContextVar("relay_scope_stack", default=None)
self._scope_serial = 0
self.ScopeType = SimpleNamespace(Agent="agent", Function="function")
self.LLMRequest = _Request
@ -51,6 +52,11 @@ class _Relay:
def _scope_push(self, name: str, scope_type: Any, **kwargs: Any) -> Any:
self._scope_serial += 1
handle = ("scope", name, self._scope_serial)
stack = self._scope_stack.get()
if stack is None:
stack = []
self._scope_stack.set(stack)
stack.append(handle)
self._scope.set(handle)
self.events.append(("scope.push", name, scope_type, kwargs))
if scope_type == self.ScopeType.Function:
@ -69,6 +75,13 @@ class _Relay:
return handle
def _scope_pop(self, handle: Any, **kwargs: Any) -> None:
stack = self._scope_stack.get()
if not stack or stack[-1] != handle:
current = stack[-1] if stack else None
self.events.append(("scope.pop.rejected", handle, current))
raise RuntimeError("scope handle is not at the top of the stack")
stack.pop()
self._scope.set(stack[-1] if stack else None)
self.events.append(("scope.pop", handle, kwargs))
start = self._scope_starts.pop(handle, None)
if start is not None:
@ -91,7 +104,9 @@ class _Relay:
self.events.append(("scope.event", name, kwargs))
def _get_scope_stack(self) -> Any:
current = self._scope.get()
stack = self._scope_stack.get()
current = stack[-1] if stack else None
self._scope.set(current)
self.events.append(("scope.sync", current))
return current
@ -632,6 +647,75 @@ def test_core_task_instrumentation_preserves_prompt_history_and_tool_schema(
assert json.dumps(agent.tools, ensure_ascii=False, sort_keys=True) == tools_before
def test_skipped_turn_does_not_finish_another_sessions_matching_task(
direct_runtime,
monkeypatch,
):
"""A skipped turn must not use shared-metrics' task-id fallback on finish."""
from run_agent import AIAgent
owner_session = "instrumented-session"
shared_task_id = "caller-supplied-task-id"
relay_shared_metrics.start_task_run(
session_id=owner_session,
task_id=shared_task_id,
platform="cli",
)
runtime = relay_shared_metrics._get_runtime()
assert runtime is not None
assert (owner_session, shared_task_id) in runtime._task_sessions
agent = object.__new__(AIAgent)
agent.session_id = "skipped-session"
agent.platform = "cli"
agent._parent_session_id = None
agent._session_db = None
agent._cached_system_prompt = "stable"
agent.tools = []
skipped_turn = SimpleNamespace(relay_enabled=False)
monkeypatch.setattr(
relay_runtime.SESSION_COORDINATOR,
"begin_turn",
lambda *_args, **_kwargs: skipped_turn,
)
monkeypatch.setattr(
relay_runtime.SESSION_COORDINATOR,
"finish_logical_calls",
lambda *_args, **_kwargs: None,
)
monkeypatch.setattr(
relay_runtime.SESSION_COORDINATOR,
"end_turn",
lambda *_args, **_kwargs: None,
)
monkeypatch.setattr(
"agent.conversation_loop.run_conversation",
lambda *_args, **_kwargs: {"final_response": "ok", "completed": True},
)
result = AIAgent.run_conversation(
agent,
"hello",
conversation_history=[],
task_id=shared_task_id,
)
assert result["completed"] is True
assert (owner_session, shared_task_id) in runtime._task_sessions
assert not [
event
for event in direct_runtime.events
if event[0] == "scope.pop" and event[1][1] == relay_shared_metrics.TASK_SCOPE
]
relay_shared_metrics.finish_task_run(
session_id=owner_session,
task_id=shared_task_id,
platform="cli",
result={"completed": True},
)
@ -792,6 +876,171 @@ def test_sync_session_runner_releases_lock_before_callback(direct_runtime):
assert contender.is_alive() is False
def test_direct_runtime_fake_enforces_lifo_scope_contract(direct_runtime):
runtime = relay_runtime.get_runtime()
assert runtime is not None
session = runtime.ensure_session({"session_id": "lifo-contract"})
assert session is not None
first = runtime.run_in_session(
session,
direct_runtime.scope.push,
"first",
direct_runtime.ScopeType.Function,
)
second = runtime.run_in_session(
session,
direct_runtime.scope.push,
"second",
direct_runtime.ScopeType.Function,
)
with pytest.raises(RuntimeError, match="not at the top"):
runtime.run_in_session(session, direct_runtime.scope.pop, first)
runtime.run_in_session(session, direct_runtime.scope.pop, second)
runtime.run_in_session(session, direct_runtime.scope.pop, first)
def test_concurrent_turn_skips_relay_before_scope_stack_can_interleave(
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",
)
assert first.relay_enabled is True
assert first.handle is not None
assert second.relay_enabled is False
assert second.handle is None
assert relay_runtime.resolve_execution_context("shared-session") == (
None,
None,
None,
)
coordinator.end_turn(first, outcome="success")
coordinator.end_turn(second, outcome="success")
coordinator.release_conversation(lease)
coordinator.finalize_conversation(
profile_key=profile_key,
session_id="shared-session",
)
turn_closes = [
event
for event in direct_runtime.events
if event[0] == "scope.pop" and event[1] == first.handle
]
assert len(turn_closes) == 1
assert not [
event
for event in direct_runtime.events
if event[0] == "scope.pop.rejected"
]
def test_concurrent_turn_skips_shared_metrics_scope_creation(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")
relay_shared_metrics.observe_lifecycle(
"pre_llm_call",
session_id="shared-session",
task_id="second-task",
platform="cli",
)
relay_shared_metrics.observe_lifecycle(
"pre_api_request",
session_id="shared-session",
task_id="second-task",
api_request_id="second-request",
platform="cli",
)
assert second.relay_enabled is False
assert not [
event
for event in direct_runtime.events
if event[0] == "scope.push" and event[1] == relay_shared_metrics.TASK_SCOPE
]
coordinator.end_turn(first, outcome="success")
coordinator.end_turn(second, outcome="success")
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)
@ -968,10 +1217,3 @@ def test_failed_flush_keeps_daily_export_open_for_later_task(
assert flush_attempts == 2
assert "Hermes shared-metrics task flush failed" in caplog.text