From 6fc5794dc28b0624339a96080d7583a7c32da5fc Mon Sep 17 00:00:00 2001
From: Xubin Ren <52506698+Re-bin@users.noreply.github.com>
Date: Sun, 23 Aug 2026 21:59:11 +0800
Subject: [PATCH] fix(runtime): normalize recovery session routing

---
 nanobot/agent/loop.py                 | 15 ++++++------
 tests/agent/test_runner_injections.py | 33 +++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py
index 3782dae4b..233e74628 100644
--- a/nanobot/agent/loop.py
+++ b/nanobot/agent/loop.py
@@ -1303,6 +1303,12 @@ class AgentLoop:
                         break
                 if deferred:
                     continue
+                routed_msg = msg
+                if effective_key != msg.session_key:
+                    routed_msg = dataclasses.replace(
+                        msg,
+                        session_key_override=effective_key,
+                    )
                 # A newer WebUI message must supersede an explicit recovery
                 # before it is injected into that recovery's pending queue.
                 # Without this admission point, a recovered turn could finish
@@ -1311,7 +1317,7 @@ class AgentLoop:
                     effective_key in self._pending_queues
                     and msg.channel == "websocket"
                     and self._recovery_admission is not None
-                    and not await self._recovery_admission.admit(msg)
+                    and not await self._recovery_admission.admit(routed_msg)
                 ):
                     continue
                 # If this session already has an active pending queue (i.e. a task
@@ -1326,12 +1332,7 @@ class AgentLoop:
                             self.commands.dispatch,
                         )
                         continue
-                    pending_msg = msg
-                    if effective_key != msg.session_key:
-                        pending_msg = dataclasses.replace(
-                            msg,
-                            session_key_override=effective_key,
-                        )
+                    pending_msg = routed_msg
                     session = self.sessions.get_or_create(effective_key)
                     followup_id = record_pending_followup(session, pending_msg)
                     if followup_id is not None:
diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py
index f4fb7f757..38a4a977f 100644
--- a/tests/agent/test_runner_injections.py
+++ b/tests/agent/test_runner_injections.py
@@ -1018,6 +1018,39 @@ async def test_websocket_followup_is_admitted_before_recovery_queue(tmp_path):
     await asyncio.wait_for(run_task, timeout=2)
 
 
+@pytest.mark.asyncio
+async def test_unified_websocket_followup_admits_effective_session(tmp_path):
+    """Recovery admission and the pending queue must use the same session key."""
+    from nanobot.bus.events import InboundMessage
+    from nanobot.session.keys import UNIFIED_SESSION_KEY
+
+    admission = MagicMock()
+    admission.admit = AsyncMock(return_value=True)
+    loop = _make_loop(tmp_path, recovery_admission=admission)
+    loop._unified_session = True
+    loop._dispatch = AsyncMock()  # type: ignore[method-assign]
+
+    pending = asyncio.Queue(maxsize=20)
+    loop._pending_queues[UNIFIED_SESSION_KEY] = pending
+
+    run_task = asyncio.create_task(loop.run())
+    msg = InboundMessage(
+        channel="websocket",
+        sender_id="u",
+        chat_id="chat",
+        content="new request",
+    )
+    await loop.bus.publish_inbound(msg)
+
+    queued_msg = await asyncio.wait_for(pending.get(), timeout=2)
+    admitted_msg = admission.admit.await_args.args[0]
+    assert admitted_msg.session_key == UNIFIED_SESSION_KEY
+    assert queued_msg.session_key == UNIFIED_SESSION_KEY
+
+    loop.stop()
+    await asyncio.wait_for(run_task, timeout=2)
+
+
 @pytest.mark.asyncio
 async def test_mid_turn_subagent_result_does_not_resolve_a_new_turn_route(tmp_path):
     """Injected results stay inside the active turn instead of opening a side turn."""