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:
committed by
kshitijk4poor
parent
78c11d99e3
commit
6e2fd955ca
+31
-1
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user