Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(agent): address session and streaming concurrency bugs

Browse Source
This commit is contained in:
hamb1y
2026-05-28 22:54:46 +08:00
committed by Xubin Ren
parent 1a4ae8994d
commit 0df60416ba
10 changed files with 250 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/loop.py
+40 -29
View File
@@ -812,16 +812,16 @@ class AgentLoop:
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
raw = msg.content.strip()
if self.commands.is_priority(raw):
await self._dispatch_command_inline(
msg, msg.session_key, raw,
msg, effective_key, raw,
self.commands.dispatch_priority,
)
continue
effective_key = self._effective_session_key(msg)
# 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.
@@ -872,13 +872,13 @@ class AgentLoop:
lock = self._session_locks.setdefault(session_key, asyncio.Lock())
gate = self._concurrency_gate or nullcontext()
# Register a pending queue so follow-up messages for this session are
# routed here (mid-turn injection) instead of spawning a new task.
pending = asyncio.Queue(maxsize=20)
self._pending_queues[session_key] = pending
pending: asyncio.Queue | None = None
try:
async with lock, gate:
# Only the task that owns the session lock may publish the
# active mid-turn injection queue for this session.
pending = asyncio.Queue(maxsize=20)
self._pending_queues[session_key] = pending
try:
on_stream = on_stream_end = None
if msg.metadata.get("_wants_stream"):
@@ -962,28 +962,39 @@ class AgentLoop:
channel=msg.channel, chat_id=msg.chat_id,
content="Sorry, I encountered an error.",
))
finally:
# Drain any messages still in the pending queue and re-publish
# them to the bus so they are processed as fresh inbound messages
# rather than silently lost. Only remove our own queue; a
# later task waiting on the lock must not be able to steal
# cleanup ownership.
queue = None
if self._pending_queues.get(session_key) is pending:
queue = self._pending_queues.pop(session_key, None)
else:
queue = pending
if queue is not None:
leftover = 0
while True:
try:
item = queue.get_nowait()
except asyncio.QueueEmpty:
break
await self.bus.publish_inbound(item)
leftover += 1
if leftover:
logger.info(
"Re-published {} leftover message(s) to bus for session {}",
leftover, session_key,
)
await self._webui_turns.publish_run_status(msg, "idle")
self._pending_turn_latency_ms.pop(session_key, None)
self._webui_turns.discard(session_key)
finally:
# Drain any messages still in the pending queue and re-publish
# them to the bus so they are processed as fresh inbound messages
# rather than silently lost.
queue = self._pending_queues.pop(session_key, None)
if queue is not None:
leftover = 0
while True:
try:
item = queue.get_nowait()
except asyncio.QueueEmpty:
break
await self.bus.publish_inbound(item)
leftover += 1
if leftover:
logger.info(
"Re-published {} leftover message(s) to bus for session {}",
leftover, session_key,
)
await self._webui_turns.publish_run_status(msg, "idle")
self._pending_turn_latency_ms.pop(session_key, None)
self._webui_turns.discard(session_key)
if pending is None:
await self._webui_turns.publish_run_status(msg, "idle")
self._pending_turn_latency_ms.pop(session_key, None)
self._webui_turns.discard(session_key)
async def close_mcp(self) -> None:
"""Drain pending background archives, then close MCP connections."""
nanobot/agent/runner.py
+7 -1
View File
@@ -1273,7 +1273,13 @@ class AgentRunner:
return messages
system_tokens = sum(estimate_message_tokens(msg) for msg in system_messages)
remaining_budget = max(128, budget - system_tokens)
fixed_tokens, _ = estimate_prompt_tokens_chain(
self.provider,
spec.model,
system_messages,
spec.tools.get_definitions(),
)
remaining_budget = max(0, budget - max(system_tokens, fixed_tokens))
kept: list[dict[str, Any]] = []
kept_tokens = 0
for message in reversed(non_system):
nanobot/agent/tools/long_task.py
+10 -6
View File
@@ -16,6 +16,7 @@ There is **no** sub-agent orchestrator and **no** special WebSocket ``agent_ui``
from __future__ import annotations
from contextvars import ContextVar
from datetime import datetime
from typing import TYPE_CHECKING, Any
@@ -45,15 +46,19 @@ class _GoalToolsMixin(ContextAware):
def __init__(self, sessions: SessionManager, bus: Any | None = None) -> None:
self._sessions = sessions
self._bus = bus
self._request_ctx: RequestContext | None = None
self._request_ctx: ContextVar[RequestContext | None] = ContextVar(
f"{self.__class__.__name__}_request_ctx",
default=None,
)
def set_context(self, ctx: RequestContext) -> None:
self._request_ctx = ctx
self._request_ctx.set(ctx)
def _session(self):
if self._request_ctx is None:
request_ctx = self._request_ctx.get()
if request_ctx is None:
return None
key = self._request_ctx.session_key
key = request_ctx.session_key
if not key:
return None
return self._sessions.get_or_create(key)
@@ -61,7 +66,7 @@ class _GoalToolsMixin(ContextAware):
async def _publish_goal_state_ws(self, metadata: dict[str, Any]) -> None:
"""Fan-out authoritative goal snapshot for this WebSocket chat only."""
bus = self._bus
rc = self._request_ctx
rc = self._request_ctx.get()
if bus is None or rc is None or rc.channel != "websocket":
return
cid = (rc.chat_id or "").strip()
@@ -224,4 +229,3 @@ class CompleteGoalTool(Tool, _GoalToolsMixin):
if tail:
return f"Goal marked complete ({ended}). Recap:\n{tail}"
return f"Goal marked complete ({ended})."
nanobot/command/builtin.py
+1 -1
View File
@@ -123,7 +123,7 @@ async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
"""Cancel all active tasks and subagents for the session."""
loop = ctx.loop
msg = ctx.msg
total = await loop._cancel_active_tasks(msg.session_key)
total = await loop._cancel_active_tasks(ctx.key)
content = f"Stopped {total} task(s)." if total else "No active task to stop."
return OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id, content=content,
nanobot/providers/base.py
+17 -1
View File
@@ -557,11 +557,20 @@ class LLMProvider(ABC):
if reasoning_effort is self._SENTINEL:
reasoning_effort = self.generation.reasoning_effort
has_streamed_content = False
async def _tracking_delta(text: str) -> None:
nonlocal has_streamed_content
if text:
has_streamed_content = True
if on_content_delta:
await on_content_delta(text)
kw: dict[str, Any] = dict(
messages=messages, tools=tools, model=model,
max_tokens=max_tokens, temperature=temperature,
reasoning_effort=reasoning_effort, tool_choice=tool_choice,
on_content_delta=on_content_delta,
on_content_delta=_tracking_delta if on_content_delta is not None else None,
on_thinking_delta=on_thinking_delta,
on_tool_call_delta=on_tool_call_delta,
)
@@ -571,6 +580,7 @@ class LLMProvider(ABC):
messages,
retry_mode=retry_mode,
on_retry_wait=on_retry_wait,
should_retry_guard=lambda: not has_streamed_content,
)
async def chat_with_retry(
@@ -717,6 +727,7 @@ class LLMProvider(ABC):
*,
retry_mode: str,
on_retry_wait: Callable[[str], Awaitable[None]] | None,
should_retry_guard: Callable[[], bool] | None = None,
) -> LLMResponse:
attempt = 0
delays = list(self._CHAT_RETRY_DELAYS)
@@ -730,6 +741,11 @@ class LLMProvider(ABC):
if response.finish_reason != "error":
return response
last_response = response
if should_retry_guard is not None and not should_retry_guard():
logger.warning(
"LLM stream failed after content was emitted; skipping retry"
)
return response
error_key = ((response.content or "").strip().lower() or None)
if error_key and error_key == last_error_key:
identical_error_count += 1