fix(worker): persist activity and aux usage off Relay's event loop

Relay drives Hermes' provider and tool callbacks synchronously on its
private asyncio loop. In a managed worker, _touch_activity inside those
callbacks reached RuntimeSessionStore -> WorkerRPC on that loop, and the
#139472 guard refused it (synchronous_rpc_on_event_loop). The refusal is
sticky on the store, so later append_message / finalize writes of the
turn failed the same way and the worker died as managed_worker_failed.

The reporter's caller diagnostic on #139029 names two stacks:
relay_llm.run_callback -> interruptible_api_call -> _emit_wait_notice ->
_touch_activity, and concurrent tool dispatch -> _begin_tool_execution ->
_touch_activity. Both end in _persist_session_activity_if_due.

call_off_event_loop() (runtime_session_store) runs a synchronous write on
a one-shot worker thread when the current thread is running a loop, with
the context copied and errors re-raised. The callback already blocks that
loop, so waiting costs it nothing. The guard is unchanged. Callers:
- _persist_session_activity_if_due (both reported stacks)
- _reset_activity_labels_after_turn (same mixin, same shape)
- record_aux_usage: a sync auxiliary call inside a managed tool callback
  (the async aux path already hops with asyncio.to_thread)

Tests drive the real Relay pipeline (relay_llm.execute,
relay_tools.execute), the real RuntimeSessionStore and the real guard.
Only the owner's socket is replaced.

Refs #139029
Refs #134889 (KB-2)
Refs #141321

agent/activity_tracking.py

@@ -112,8 +112,10 @@ class ActivityTrackingMixin:
         if last_mono and (now_mono - last_mono) < SESSION_ACTIVITY_HEARTBEAT_MIN_INTERVAL_SECONDS:
             return
         self._session_activity_last_persist_mono = now_mono
+        from agent.runtime_session_store import call_off_event_loop  # Relay callbacks run on a loop (#139029)
         try:
-            touch(
+            call_off_event_loop(
+                touch,
                 session_id,
                 getattr(self, "_last_activity_ts", None),
                 description=getattr(self, "_last_activity_desc", None),
@@ -142,5 +144,6 @@ class ActivityTrackingMixin:
         clear = getattr(session_db, "clear_session_activity_labels", None)
         if not callable(clear):
             return
+        from agent.runtime_session_store import call_off_event_loop
         with suppress(Exception):  # never let durable cleanup I/O break turn teardown
-            clear(session_id)
+            call_off_event_loop(clear, session_id)

agent/aux_accounting.py

@@ -81,7 +81,11 @@ def record_aux_usage(
                 estimated_cost = float(cost.amount_usd)
         except Exception:
             logger.debug("Aux usage cost estimation failed", exc_info=True)
-        session_db.record_auxiliary_usage(
+        # Inside a Relay-driven tool callback the thread runs Relay's loop; a worker store refuses
+        # synchronous RPC there (#139029).
+        from agent.runtime_session_store import call_off_event_loop
+        call_off_event_loop(
+            session_db.record_auxiliary_usage,
             session_id, task, model=model, billing_provider=provider, billing_base_url=base_url,
             input_tokens=usage.input_tokens, output_tokens=usage.output_tokens,
             cache_read_tokens=usage.cache_read_tokens, cache_write_tokens=usage.cache_write_tokens,

agent/runtime_session_store.py

@@ -7,6 +7,8 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
+import contextvars
+from concurrent.futures import ThreadPoolExecutor
 import json
 import logging
 import os
@@ -57,6 +59,24 @@ def _caller_summary(limit=8):
     return ' <- '.join(frames)
 
 
+def call_off_event_loop(fn, *args, **kwargs):
+    """Call ``fn`` here, or on a worker thread when this thread is running an event loop.
+
+    For synchronous best-effort writers that can run inside a callback Relay drives on its private
+    loop (provider/tool callbacks): there ``WorkerRPC`` refuses the call (``synchronous_rpc_on_event_loop``,
+    sticky for the rest of the turn; #139029). The callback already blocks that loop, so waiting for the
+    thread costs it nothing; the context is copied so session-scoped state survives the hop, and errors
+    re-raise here. Async code awaits ``asyncio.to_thread`` instead.
+    """
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return fn(*args, **kwargs)
+    context = contextvars.copy_context()
+    with ThreadPoolExecutor(max_workers=1, thread_name_prefix='worker-persist-off-loop') as pool:
+        return pool.submit(context.run, fn, *args, **kwargs).result()
+
+
 class _OwnerLink:
     """One owner WebSocket and the only thread that reads it. Events the owner broadcasts to every
     transport are discarded as they arrive, idle or not: left queued, websockets' bounded receive

tests/agent/test_activity_persist_off_relay_loop.py

@@ -0,0 +1,190 @@
+"""A managed worker's activity heartbeat must not reach the synchronous persistence RPC on
+Relay's event loop (#139029).
+
+Relay drives Hermes' provider and tool callbacks synchronously on its private ``asyncio.run``
+loop, so ``_touch_activity`` inside them runs with a loop on the thread. In a managed worker
+``_session_db`` is a ``RuntimeSessionStore`` whose ``WorkerRPC`` refuses on-loop calls
+(``synchronous_rpc_on_event_loop``) and the refusal is sticky: every later write of the turn
+failed. These drive the real Relay pipeline, the real store and the real guard; only the owner's
+socket is replaced by an in-process reply.
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("nemo_relay")
+
+import run_agent
+from agent import relay_llm, relay_runtime, relay_tools
+from agent.runtime_session_store import RuntimeSessionStore, WorkerRPC
+from agent.session_activity import ActivityProvenance
+
+SESSION = "session-1"
+
+
+class _OwnerReplies:
+    """Stands in for the owner's socket (``WorkerRPC._connection``); records what reached it."""
+
+    def __init__(self):
+        self.persisted = []
+        self.on_loop = []
+
+    def send(self, request_id, message):
+        import json
+        try:
+            asyncio.get_running_loop()
+            self.on_loop.append(True)
+        except RuntimeError:
+            self.on_loop.append(False)
+        frame = json.loads(message)
+        self.persisted.append(frame["params"]["operation"])
+        self._reply = {"jsonrpc": "2.0", "id": request_id, "result": {"ok": True}}
+
+    def receive(self, timeout):
+        return self._reply
+
+
+@pytest.fixture()
+def relay_turn(tmp_path, monkeypatch):
+    monkeypatch.setenv("HERMES_HOME", str(tmp_path / "profile"))
+    monkeypatch.delenv("HERMES_KANBAN_TASK", raising=False)
+    relay_runtime._reset_for_tests()
+    lease = relay_runtime.SESSION_COORDINATOR.acquire_conversation(
+        profile_key=relay_runtime.current_profile_key(), session_id=SESSION, platform="cli",
+    )
+    turn = relay_runtime.SESSION_COORDINATOR.begin_turn(lease, turn_id="turn-1", task_id="task-1")
+    lease.host.retain_managed_execution("test.activity_off_loop")
+    try:
+        yield lease.host
+    finally:
+        lease.host.release_managed_execution("test.activity_off_loop")
+        relay_runtime.SESSION_COORDINATOR.end_turn(turn, outcome="success")
+        relay_runtime.SESSION_COORDINATOR.release_conversation(lease)
+        relay_runtime._reset_for_tests()
+
+
+@pytest.fixture()
+def worker_agent(tmp_path, monkeypatch):
+    owner = _OwnerReplies()
+    monkeypatch.setattr(WorkerRPC, "_connection", lambda self: owner)
+    store = RuntimeSessionStore(WorkerRPC(tmp_path / "home"), {"session_id": SESSION}, tmp_path / "outbox")
+    agent = SimpleNamespace(
+        session_id=SESSION, _session_db=store, _last_activity_ts=0.0, _last_activity_desc="",
+        _last_activity_provenance=ActivityProvenance.UNKNOWN, _session_activity_last_persist_mono=0.0,
+    )
+    agent._touch_activity = run_agent.AIAgent._touch_activity.__get__(agent, SimpleNamespace)
+    agent._persist_session_activity_if_due = (
+        run_agent.AIAgent._persist_session_activity_if_due.__get__(agent, SimpleNamespace))
+    agent._reset_activity_labels_after_turn = (
+        run_agent.AIAgent._reset_activity_labels_after_turn.__get__(agent, SimpleNamespace))
+    try:
+        yield agent, store, owner
+    finally:
+        store._outbox_owner.close()
+
+
+def _guard_warnings(caplog):
+    return [r.getMessage() for r in caplog.records
+            if r.name == "agent.runtime_session_store" and "synchronous_rpc_on_event_loop" in r.getMessage()]
+
+
+def _assert_persisted_off_loop(store, owner, caplog, *operations):
+    assert _guard_warnings(caplog) == []
+    assert store.failure is None
+    assert owner.persisted == list(operations)
+    assert owner.on_loop == [False] * len(operations)
+
+
+def test_provider_callback_heartbeat_persists_off_relay_loop(relay_turn, worker_agent, caplog):
+    """Reporter stack 1: relay_llm.run_callback -> interruptible_api_call -> _emit_wait_notice
+    -> _touch_activity -> worker.persist, on the loop ``relay_llm.execute`` runs."""
+    agent, store, owner = worker_agent
+    caplog.set_level(logging.WARNING, logger="agent.runtime_session_store")
+
+    def provider(request):
+        assert asyncio.get_running_loop() is not None  # Relay's loop is live on this thread
+        agent._touch_activity("waiting for provider response")
+        return {"id": "r", "object": "chat.completion", "choices": [{
+            "index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}]}
+
+    relay_llm.execute(
+        {"messages": [{"role": "user", "content": "hi"}], "model": "m"}, provider,
+        name="prov", model_name="m", session_id=SESSION,
+        metadata={"api_mode": "chat_completions", "api_request_id": "req-1"},
+    )
+    _assert_persisted_off_loop(store, owner, caplog, "session.activity")
+
+
+def test_concurrent_tool_callback_heartbeat_persists_off_relay_loop(relay_turn, worker_agent, caplog):
+    """Reporter stack 3: concurrent tool dispatch on a ThreadPoolExecutor thread ->
+    _begin_tool_execution -> _touch_activity -> worker.persist, inside relay_tools.execute."""
+    agent, store, owner = worker_agent
+    caplog.set_level(logging.WARNING, logger="agent.runtime_session_store")
+
+    def tool(args):
+        assert asyncio.get_running_loop() is not None
+        agent._touch_activity("executing tool: image_generate")
+        return {"ok": True}
+
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        pool.submit(relay_tools.execute, "image_generate", {"prompt": "x"}, tool,
+                    session_id=SESSION, tool_call_id="call-1").result(timeout=30)
+    _assert_persisted_off_loop(store, owner, caplog, "session.activity")
+
+
+def test_later_turn_writes_still_land_after_an_on_loop_heartbeat(relay_turn, worker_agent, caplog):
+    """The sticky follow-on in the reporter's timeline: append_message / finalize failed with
+    the same error after the first refusal. The heartbeat must not poison the store."""
+    agent, store, owner = worker_agent
+    caplog.set_level(logging.WARNING, logger="agent.runtime_session_store")
+
+    def tool(args):
+        agent._touch_activity("executing tool: terminal")
+        return {"ok": True}
+
+    relay_tools.execute("terminal", {"command": "true"}, tool, session_id=SESSION)
+    store.clear_session_activity_labels(SESSION)  # an ordinary off-loop write after the tool
+    _assert_persisted_off_loop(store, owner, caplog, "session.activity", "session.activity_clear")
+
+
+def test_label_reset_on_a_running_loop_persists_off_loop(worker_agent, caplog):
+    """Sibling writer in the same mixin: turn-teardown label clear reached on a loop."""
+    agent, store, owner = worker_agent
+    caplog.set_level(logging.WARNING, logger="agent.runtime_session_store")
+
+    async def teardown():
+        agent._reset_activity_labels_after_turn()
+
+    asyncio.run(teardown())
+    _assert_persisted_off_loop(store, owner, caplog, "session.activity_clear")
+
+
+def test_sync_aux_usage_inside_managed_tool_persists_off_relay_loop(relay_turn, worker_agent, caplog):
+    """Sibling: a synchronous auxiliary call made by a tool (vision, web_extract, ...) records its
+    usage at ``_validate_llm_response`` -> ``record_aux_usage`` while the tool callback still holds
+    Relay's loop. The async aux path already hops with ``asyncio.to_thread``; the sync one did not."""
+    from agent.aux_accounting import reset_accounting_context, set_accounting_context
+    from agent.auxiliary_client import _validate_llm_response
+
+    _agent, store, owner = worker_agent
+    caplog.set_level(logging.WARNING, logger="agent.runtime_session_store")
+    response = SimpleNamespace(
+        model="m", usage={"prompt_tokens": 7, "completion_tokens": 3},
+        choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))])
+
+    def tool(args):
+        assert asyncio.get_running_loop() is not None
+        _validate_llm_response(response, "vision", provider="prov")
+        return {"ok": True}
+
+    token = set_accounting_context(store, SESSION)
+    try:
+        relay_tools.execute("vision_analyze", {"image_url": "x"}, tool, session_id=SESSION)
+    finally:
+        reset_accounting_context(token)
+    _assert_persisted_off_loop(store, owner, caplog, "usage.auxiliary")