fix(gateway): claim session slot before auto-resume task to prevent duplicate agents

When the gateway restarts and auto-resumes an interrupted session, an
inbound message arriving in the window between `asyncio.create_task()`
and the task's first await could spin up a second AIAgent for the same
session.  Both agents would then process messages concurrently,
producing interleaved duplicate responses (#45456).

Fix: set `_AGENT_PENDING_SENTINEL` in `_running_agents` immediately
after the "already running" check, before creating the task.  This
closes the race window — any inbound message sees the slot as occupied
and queues behind the auto-resume.

A `_guarded_handle_message` wrapper ensures the pre-claimed sentinel is
always released, even if `handle_message` raises before reaching
`_process_message_background` (whose `finally` block handles normal
cleanup).

(cherry picked from commit 85150c976b)
This commit is contained in:
liuhao1024
2026-06-13 23:36:51 +05:30
committed by kshitijk4poor
parent 78c11d99e3
commit 6e2fd955ca
2 changed files with 122 additions and 1 deletions
+31 -1
View File
@@ -4556,6 +4556,14 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
)
continue
# Claim the session slot *before* spawning the task so that an
# inbound message arriving between task creation and the task's
# first await (where _process_message_background sets the real
# sentinel) sees the slot as occupied and queues behind it
# instead of spinning up a duplicate AIAgent (#45456).
self._running_agents[entry.session_key] = _AGENT_PENDING_SENTINEL
self._running_agents_ts[entry.session_key] = time.time()
# Empty-text internal event — the _is_resume_pending branch in
# _handle_message_with_agent prepends the proper reason-aware
# system note before the turn runs.
@@ -4565,7 +4573,29 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
source=source,
internal=True,
)
task = asyncio.create_task(adapter.handle_message(event))
async def _guarded_handle_message(
_adapter: Any, _event: MessageEvent, _key: str = entry.session_key,
) -> None:
"""Ensure the pre-claimed sentinel is always released.
In the normal flow, ``_process_message_background`` sets
its own sentinel and releases it in its ``finally`` block.
If ``handle_message`` raises *before* reaching that
method (e.g. during topic recovery or session key
resolution), nobody clears our pre-claim so we do it
here unconditionally.
"""
try:
await _adapter.handle_message(_event)
finally:
# Only release if the sentinel we set is still there
# (i.e. _process_message_background hasn't replaced
# and cleaned it already).
if self._running_agents.get(_key) is _AGENT_PENDING_SENTINEL:
self._release_running_agent_state(_key)
task = asyncio.create_task(_guarded_handle_message(adapter, event))
self._background_tasks.add(task)
task.add_done_callback(self._background_tasks.discard)
scheduled += 1