feat(agent): replace global lock with per-session locks for concurrent dispatch

Replace the single _processing_lock (asyncio.Lock) with per-session locks
so that different sessions can process LLM requests concurrently, while
messages within the same session remain serialised.

An optional global concurrency cap is available via the
NANOBOT_MAX_CONCURRENT_REQUESTS env var (default 3, <=0 for unlimited).

Also re-binds tool context before each tool execution round to prevent
concurrent sessions from clobbering each other's routing info.

Tested in production and manually reviewed.
This commit is contained in:
gem12
2026-03-21 23:03:15 +08:00
parent 939af8898b
commit c397bb4229
+25 -5
View File
@@ -7,7 +7,7 @@ import json
import os import os
import re import re
import sys import sys
from contextlib import AsyncExitStack from contextlib import AsyncExitStack, nullcontext
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, Awaitable, Callable from typing import TYPE_CHECKING, Any, Awaitable, Callable
@@ -103,7 +103,12 @@ class AgentLoop:
self._mcp_connecting = False self._mcp_connecting = False
self._active_tasks: dict[str, list[asyncio.Task]] = {} # session_key -> tasks self._active_tasks: dict[str, list[asyncio.Task]] = {} # session_key -> tasks
self._background_tasks: list[asyncio.Task] = [] self._background_tasks: list[asyncio.Task] = []
self._processing_lock = asyncio.Lock() self._session_locks: dict[str, asyncio.Lock] = {}
# NANOBOT_MAX_CONCURRENT_REQUESTS: <=0 means unlimited; default 3.
_max = int(os.environ.get("NANOBOT_MAX_CONCURRENT_REQUESTS", "3"))
self._concurrency_gate: asyncio.Semaphore | None = (
asyncio.Semaphore(_max) if _max > 0 else None
)
self.memory_consolidator = MemoryConsolidator( self.memory_consolidator = MemoryConsolidator(
workspace=workspace, workspace=workspace,
provider=provider, provider=provider,
@@ -206,6 +211,10 @@ class AgentLoop:
self, self,
initial_messages: list[dict], initial_messages: list[dict],
on_progress: Callable[..., Awaitable[None]] | None = None, on_progress: Callable[..., Awaitable[None]] | None = None,
*,
channel: str = "cli",
chat_id: str = "direct",
message_id: str | None = None,
) -> tuple[str | None, list[str], list[dict]]: ) -> tuple[str | None, list[str], list[dict]]:
"""Run the agent iteration loop.""" """Run the agent iteration loop."""
messages = initial_messages messages = initial_messages
@@ -248,6 +257,10 @@ class AgentLoop:
args_str = json.dumps(tc.arguments, ensure_ascii=False) args_str = json.dumps(tc.arguments, ensure_ascii=False)
logger.info("Tool call: {}({})", tc.name, args_str[:200]) logger.info("Tool call: {}({})", tc.name, args_str[:200])
# Re-bind tool context right before execution so that
# concurrent sessions don't clobber each other's routing.
self._set_tool_context(channel, chat_id, message_id)
# Execute all tool calls concurrently — the LLM batches # Execute all tool calls concurrently — the LLM batches
# independent calls in a single response on purpose. # independent calls in a single response on purpose.
# return_exceptions=True ensures all results are collected # return_exceptions=True ensures all results are collected
@@ -343,8 +356,10 @@ class AgentLoop:
asyncio.create_task(_do_restart()) asyncio.create_task(_do_restart())
async def _dispatch(self, msg: InboundMessage) -> None: async def _dispatch(self, msg: InboundMessage) -> None:
"""Process a message under the global lock.""" """Process a message: per-session serial, cross-session concurrent."""
async with self._processing_lock: lock = self._session_locks.setdefault(msg.session_key, asyncio.Lock())
gate = self._concurrency_gate or nullcontext()
async with lock, gate:
try: try:
response = await self._process_message(msg) response = await self._process_message(msg)
if response is not None: if response is not None:
@@ -411,7 +426,10 @@ class AgentLoop:
current_message=msg.content, channel=channel, chat_id=chat_id, current_message=msg.content, channel=channel, chat_id=chat_id,
current_role=current_role, current_role=current_role,
) )
final_content, _, all_msgs = await self._run_agent_loop(messages) final_content, _, all_msgs = await self._run_agent_loop(
messages, channel=channel, chat_id=chat_id,
message_id=msg.metadata.get("message_id"),
)
self._save_turn(session, all_msgs, 1 + len(history)) self._save_turn(session, all_msgs, 1 + len(history))
self.sessions.save(session) self.sessions.save(session)
self._schedule_background(self.memory_consolidator.maybe_consolidate_by_tokens(session)) self._schedule_background(self.memory_consolidator.maybe_consolidate_by_tokens(session))
@@ -473,6 +491,8 @@ class AgentLoop:
final_content, _, all_msgs = await self._run_agent_loop( final_content, _, all_msgs = await self._run_agent_loop(
initial_messages, on_progress=on_progress or _bus_progress, initial_messages, on_progress=on_progress or _bus_progress,
channel=msg.channel, chat_id=msg.chat_id,
message_id=msg.metadata.get("message_id"),
) )
if final_content is None: if final_content is None: