fix: close MCP stdio transports from agent task
This commit is contained in:
Xubin Ren
+77
-73
@@ -873,89 +873,93 @@ class AgentLoop:
|
||||
async def run(self) -> None:
|
||||
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
|
||||
self._running = True
|
||||
await self._connect_mcp()
|
||||
logger.info("Agent loop started")
|
||||
try:
|
||||
await self._connect_mcp()
|
||||
logger.info("Agent loop started")
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
|
||||
except asyncio.TimeoutError:
|
||||
self.auto_compact.check_expired(
|
||||
self._schedule_background,
|
||||
active_session_keys=self._pending_queues.keys(),
|
||||
)
|
||||
continue
|
||||
except asyncio.CancelledError:
|
||||
# Preserve real task cancellation so shutdown can complete cleanly.
|
||||
# Only ignore non-task CancelledError signals that may leak from integrations.
|
||||
if not self._running or asyncio.current_task().cancelling():
|
||||
raise
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning("Error consuming inbound message: {}, continuing...", e)
|
||||
continue
|
||||
while self._running:
|
||||
try:
|
||||
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
|
||||
except asyncio.TimeoutError:
|
||||
self.auto_compact.check_expired(
|
||||
self._schedule_background,
|
||||
active_session_keys=self._pending_queues.keys(),
|
||||
)
|
||||
continue
|
||||
except asyncio.CancelledError:
|
||||
# Preserve real task cancellation so shutdown can complete cleanly.
|
||||
# Only ignore non-task CancelledError signals that may leak from integrations.
|
||||
if not self._running or asyncio.current_task().cancelling():
|
||||
raise
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning("Error consuming inbound message: {}, continuing...", e)
|
||||
continue
|
||||
|
||||
raw = msg.content.strip()
|
||||
effective_key = self._effective_session_key(msg)
|
||||
if await agent_context.handle_runtime_control(self, msg, self.tools):
|
||||
continue
|
||||
if self.commands.is_priority(raw):
|
||||
await self._dispatch_command_inline(
|
||||
msg, effective_key, raw,
|
||||
self.commands.dispatch_priority,
|
||||
)
|
||||
continue
|
||||
if self._cron_turns.defer_if_active(
|
||||
msg,
|
||||
session_key=effective_key,
|
||||
active_session_keys=self._pending_queues.keys(),
|
||||
):
|
||||
logger.info(
|
||||
"Deferred cron turn for active session {}",
|
||||
effective_key,
|
||||
)
|
||||
continue
|
||||
# If this session already has an active pending queue (i.e. a task
|
||||
# is processing this session), route the message there for mid-turn
|
||||
# injection instead of creating a competing task.
|
||||
if effective_key in self._pending_queues:
|
||||
# Non-priority commands must not be queued for injection;
|
||||
# dispatch them directly (same pattern as priority commands).
|
||||
if self.commands.is_dispatchable_command(raw):
|
||||
raw = msg.content.strip()
|
||||
effective_key = self._effective_session_key(msg)
|
||||
if await agent_context.handle_runtime_control(self, msg, self.tools):
|
||||
continue
|
||||
if self.commands.is_priority(raw):
|
||||
await self._dispatch_command_inline(
|
||||
msg, effective_key, raw,
|
||||
self.commands.dispatch,
|
||||
self.commands.dispatch_priority,
|
||||
)
|
||||
continue
|
||||
pending_msg = msg
|
||||
if effective_key != msg.session_key:
|
||||
pending_msg = dataclasses.replace(
|
||||
msg,
|
||||
session_key_override=effective_key,
|
||||
)
|
||||
try:
|
||||
self._pending_queues[effective_key].put_nowait(pending_msg)
|
||||
except asyncio.QueueFull:
|
||||
logger.warning(
|
||||
"Pending queue full for session {}, falling back to queued task",
|
||||
effective_key,
|
||||
)
|
||||
else:
|
||||
if self._cron_turns.defer_if_active(
|
||||
msg,
|
||||
session_key=effective_key,
|
||||
active_session_keys=self._pending_queues.keys(),
|
||||
):
|
||||
logger.info(
|
||||
"Routed follow-up message to pending queue for session {}",
|
||||
"Deferred cron turn for active session {}",
|
||||
effective_key,
|
||||
)
|
||||
continue
|
||||
# Compute the effective session key before dispatching
|
||||
# This ensures /stop command can find tasks correctly when unified session is enabled
|
||||
task = asyncio.create_task(self._dispatch(msg))
|
||||
self._active_tasks.setdefault(effective_key, []).append(task)
|
||||
task.add_done_callback(
|
||||
lambda t, k=effective_key: self._active_tasks.get(k, [])
|
||||
and self._active_tasks[k].remove(t)
|
||||
if t in self._active_tasks.get(k, [])
|
||||
else None
|
||||
)
|
||||
# If this session already has an active pending queue (i.e. a task
|
||||
# is processing this session), route the message there for mid-turn
|
||||
# injection instead of creating a competing task.
|
||||
if effective_key in self._pending_queues:
|
||||
# Non-priority commands must not be queued for injection;
|
||||
# dispatch them directly (same pattern as priority commands).
|
||||
if self.commands.is_dispatchable_command(raw):
|
||||
await self._dispatch_command_inline(
|
||||
msg, effective_key, raw,
|
||||
self.commands.dispatch,
|
||||
)
|
||||
continue
|
||||
pending_msg = msg
|
||||
if effective_key != msg.session_key:
|
||||
pending_msg = dataclasses.replace(
|
||||
msg,
|
||||
session_key_override=effective_key,
|
||||
)
|
||||
try:
|
||||
self._pending_queues[effective_key].put_nowait(pending_msg)
|
||||
except asyncio.QueueFull:
|
||||
logger.warning(
|
||||
"Pending queue full for session {}, falling back to queued task",
|
||||
effective_key,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Routed follow-up message to pending queue for session {}",
|
||||
effective_key,
|
||||
)
|
||||
continue
|
||||
# Compute the effective session key before dispatching
|
||||
# This ensures /stop command can find tasks correctly when unified session is enabled
|
||||
task = asyncio.create_task(self._dispatch(msg))
|
||||
self._active_tasks.setdefault(effective_key, []).append(task)
|
||||
task.add_done_callback(
|
||||
lambda t, k=effective_key: self._active_tasks.get(k, [])
|
||||
and self._active_tasks[k].remove(t)
|
||||
if t in self._active_tasks.get(k, [])
|
||||
else None
|
||||
)
|
||||
finally:
|
||||
# MCP stdio transports use AnyIO cancel scopes; close them from the task that opened them.
|
||||
await self.close_mcp()
|
||||
|
||||
async def _dispatch(self, msg: InboundMessage) -> None:
|
||||
"""Process a message: per-session serial, cross-session concurrent."""
|
||||
|
||||
+16
-5
@@ -1105,16 +1105,23 @@ def _run_gateway(
|
||||
console.print(f"[yellow]Could not open browser ({e}); visit {open_browser_url}[/yellow]")
|
||||
|
||||
async def run():
|
||||
tasks: list[asyncio.Task] = []
|
||||
try:
|
||||
await cron.start()
|
||||
tasks = [
|
||||
agent.run(),
|
||||
channels.start_all(),
|
||||
asyncio.create_task(agent.run(), name="nanobot-agent-loop"),
|
||||
asyncio.create_task(channels.start_all(), name="nanobot-channels"),
|
||||
]
|
||||
if health_server_enabled:
|
||||
tasks.append(_health_server(config.gateway.host, port))
|
||||
tasks.append(asyncio.create_task(
|
||||
_health_server(config.gateway.host, port),
|
||||
name="nanobot-health-server",
|
||||
))
|
||||
if open_browser_url:
|
||||
tasks.append(_open_browser_when_ready())
|
||||
tasks.append(asyncio.create_task(
|
||||
_open_browser_when_ready(),
|
||||
name="nanobot-open-browser",
|
||||
))
|
||||
await asyncio.gather(*tasks)
|
||||
except KeyboardInterrupt:
|
||||
console.print("\nShutting down...")
|
||||
@@ -1124,9 +1131,13 @@ def _run_gateway(
|
||||
console.print("\n[red]Error: Gateway crashed unexpectedly[/red]")
|
||||
console.print(traceback.format_exc())
|
||||
finally:
|
||||
await agent.close_mcp()
|
||||
cron.stop()
|
||||
agent.stop()
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
await channels.stop_all()
|
||||
# Flush all cached sessions to durable storage before exit.
|
||||
# This prevents data loss on filesystems with write-back
|
||||
|
||||
Reference in New Issue
Block a user