Merge pull request #74864 from bbednarski9/fix/relay-concurrent-turn-scopes
fix(relay): avoid concurrent turn scope corruption
This commit is contained in:
commit
42708f8bb3
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
18
run_agent.py
18
run_agent.py
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue