Merge remote-tracking branch 'origin/main' into pr-986

This commit is contained in:
Re-bin
2026-02-22 18:13:37 +00:00
13 changed files with 535 additions and 56 deletions
+53 -16
View File
@@ -7,7 +7,7 @@ import json
import re
from contextlib import AsyncExitStack
from pathlib import Path
from typing import TYPE_CHECKING, Awaitable, Callable
from typing import TYPE_CHECKING, Any, Awaitable, Callable
from loguru import logger
@@ -95,6 +95,8 @@ class AgentLoop:
self._mcp_connected = False
self._mcp_connecting = False
self._consolidating: set[str] = set() # Session keys with consolidation in progress
self._consolidation_tasks: set[asyncio.Task] = set() # Strong refs to in-flight tasks
self._consolidation_locks: dict[str, asyncio.Lock] = {}
self._register_default_tools()
def _register_default_tools(self) -> None:
@@ -194,7 +196,8 @@ class AgentLoop:
clean = self._strip_think(response.content)
if clean:
await on_progress(clean)
await on_progress(self._tool_hint(response.tool_calls))
else:
await on_progress(self._tool_hint(response.tool_calls))
tool_call_dicts = [
{
@@ -270,6 +273,18 @@ class AgentLoop:
self._running = False
logger.info("Agent loop stopping")
def _get_consolidation_lock(self, session_key: str) -> asyncio.Lock:
lock = self._consolidation_locks.get(session_key)
if lock is None:
lock = asyncio.Lock()
self._consolidation_locks[session_key] = lock
return lock
def _prune_consolidation_lock(self, session_key: str, lock: asyncio.Lock) -> None:
"""Drop lock entry if no longer in use."""
if not lock.locked():
self._consolidation_locks.pop(session_key, None)
async def _process_message(
self,
msg: InboundMessage,
@@ -305,33 +320,55 @@ class AgentLoop:
# Slash commands
cmd = msg.content.strip().lower()
if cmd == "/new":
messages_to_archive = session.messages.copy()
lock = self._get_consolidation_lock(session.key)
self._consolidating.add(session.key)
try:
async with lock:
snapshot = session.messages[session.last_consolidated:]
if snapshot:
temp = Session(key=session.key)
temp.messages = list(snapshot)
if not await self._consolidate_memory(temp, archive_all=True):
return OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content="Memory archival failed, session not cleared. Please try again.",
)
except Exception:
logger.exception("/new archival failed for {}", session.key)
return OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content="Memory archival failed, session not cleared. Please try again.",
)
finally:
self._consolidating.discard(session.key)
self._prune_consolidation_lock(session.key, lock)
session.clear()
self.sessions.save(session)
self.sessions.invalidate(session.key)
async def _consolidate_and_cleanup():
temp = Session(key=session.key)
temp.messages = messages_to_archive
await self._consolidate_memory(temp, archive_all=True)
asyncio.create_task(_consolidate_and_cleanup())
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id,
content="New session started. Memory consolidation in progress.")
content="New session started.")
if cmd == "/help":
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id,
content="🐈 nanobot commands:\n/new — Start a new conversation\n/help — Show available commands")
if len(session.messages) > self.memory_window and session.key not in self._consolidating:
self._consolidating.add(session.key)
lock = self._get_consolidation_lock(session.key)
async def _consolidate_and_unlock():
try:
await self._consolidate_memory(session)
async with lock:
await self._consolidate_memory(session)
finally:
self._consolidating.discard(session.key)
self._prune_consolidation_lock(session.key, lock)
_task = asyncio.current_task()
if _task is not None:
self._consolidation_tasks.discard(_task)
asyncio.create_task(_consolidate_and_unlock())
_task = asyncio.create_task(_consolidate_and_unlock())
self._consolidation_tasks.add(_task)
self._set_tool_context(msg.channel, msg.chat_id, msg.metadata.get("message_id"))
if message_tool := self.tools.get("message"):
@@ -376,9 +413,9 @@ class AgentLoop:
metadata=msg.metadata or {},
)
async def _consolidate_memory(self, session, archive_all: bool = False) -> None:
"""Delegate to MemoryStore.consolidate()."""
await MemoryStore(self.workspace).consolidate(
async def _consolidate_memory(self, session, archive_all: bool = False) -> bool:
"""Delegate to MemoryStore.consolidate(). Returns True on success."""
return await MemoryStore(self.workspace).consolidate(
session, self.provider, self.model,
archive_all=archive_all, memory_window=self.memory_window,
)
+13 -8
View File
@@ -74,8 +74,11 @@ class MemoryStore:
*,
archive_all: bool = False,
memory_window: int = 50,
) -> None:
"""Consolidate old messages into MEMORY.md + HISTORY.md via LLM tool call."""
) -> bool:
"""Consolidate old messages into MEMORY.md + HISTORY.md via LLM tool call.
Returns True on success (including no-op), False on failure.
"""
if archive_all:
old_messages = session.messages
keep_count = 0
@@ -83,12 +86,12 @@ class MemoryStore:
else:
keep_count = memory_window // 2
if len(session.messages) <= keep_count:
return
return True
if len(session.messages) - session.last_consolidated <= 0:
return
return True
old_messages = session.messages[session.last_consolidated:-keep_count]
if not old_messages:
return
return True
logger.info("Memory consolidation: {} to consolidate, {} keep", len(old_messages), keep_count)
lines = []
@@ -119,7 +122,7 @@ class MemoryStore:
if not response.has_tool_calls:
logger.warning("Memory consolidation: LLM did not call save_memory, skipping")
return
return False
args = response.tool_calls[0].arguments
if entry := args.get("history_entry"):
@@ -134,5 +137,7 @@ class MemoryStore:
session.last_consolidated = 0 if archive_all else len(session.messages) - keep_count
logger.info("Memory consolidation done: {} messages, last_consolidated={}", len(session.messages), session.last_consolidated)
except Exception as e:
logger.error("Memory consolidation failed: {}", e)
return True
except Exception:
logger.exception("Memory consolidation failed")
return False
+5 -2
View File
@@ -13,8 +13,11 @@ def _resolve_path(path: str, workspace: Path | None = None, allowed_dir: Path |
if not p.is_absolute() and workspace:
p = workspace / p
resolved = p.resolve()
if allowed_dir and not str(resolved).startswith(str(allowed_dir.resolve())):
raise PermissionError(f"Path {path} is outside allowed directory {allowed_dir}")
if allowed_dir:
try:
resolved.relative_to(allowed_dir.resolve())
except ValueError:
raise PermissionError(f"Path {path} is outside allowed directory {allowed_dir}")
return resolved
+12 -3
View File
@@ -1,5 +1,6 @@
"""MCP client: connects to MCP servers and wraps their tools as native nanobot tools."""
import asyncio
from contextlib import AsyncExitStack
from typing import Any
@@ -13,12 +14,13 @@ from nanobot.agent.tools.registry import ToolRegistry
class MCPToolWrapper(Tool):
"""Wraps a single MCP server tool as a nanobot Tool."""
def __init__(self, session, server_name: str, tool_def):
def __init__(self, session, server_name: str, tool_def, tool_timeout: int = 30):
self._session = session
self._original_name = tool_def.name
self._name = f"mcp_{server_name}_{tool_def.name}"
self._description = tool_def.description or tool_def.name
self._parameters = tool_def.inputSchema or {"type": "object", "properties": {}}
self._tool_timeout = tool_timeout
@property
def name(self) -> str:
@@ -34,7 +36,14 @@ class MCPToolWrapper(Tool):
async def execute(self, **kwargs: Any) -> str:
from mcp import types
result = await self._session.call_tool(self._original_name, arguments=kwargs)
try:
result = await asyncio.wait_for(
self._session.call_tool(self._original_name, arguments=kwargs),
timeout=self._tool_timeout,
)
except asyncio.TimeoutError:
logger.warning("MCP tool '{}' timed out after {}s", self._name, self._tool_timeout)
return f"(MCP tool call timed out after {self._tool_timeout}s)"
parts = []
for block in result.content:
if isinstance(block, types.TextContent):
@@ -83,7 +92,7 @@ async def connect_mcp_servers(
tools = await session.list_tools()
for tool_def in tools.tools:
wrapper = MCPToolWrapper(session, name, tool_def)
wrapper = MCPToolWrapper(session, name, tool_def, tool_timeout=cfg.tool_timeout)
registry.register(wrapper)
logger.debug("MCP: registered tool '{}' from server '{}'", wrapper.name, name)
+2 -1
View File
@@ -304,7 +304,8 @@ class EmailChannel(BaseChannel):
self._processed_uids.add(uid)
# mark_seen is the primary dedup; this set is a safety net
if len(self._processed_uids) > self._MAX_PROCESSED_UIDS:
self._processed_uids.clear()
# Evict a random half to cap memory; mark_seen is the primary dedup
self._processed_uids = set(list(self._processed_uids)[len(self._processed_uids) // 2:])
if mark_seen:
client.store(imap_id, "+FLAGS", "\\Seen")
+6 -8
View File
@@ -55,7 +55,6 @@ class QQChannel(BaseChannel):
self.config: QQConfig = config
self._client: "botpy.Client | None" = None
self._processed_ids: deque = deque(maxlen=1000)
self._bot_task: asyncio.Task | None = None
async def start(self) -> None:
"""Start the QQ bot."""
@@ -71,8 +70,8 @@ class QQChannel(BaseChannel):
BotClass = _make_bot_class(self)
self._client = BotClass()
self._bot_task = asyncio.create_task(self._run_bot())
logger.info("QQ bot started (C2C private message)")
await self._run_bot()
async def _run_bot(self) -> None:
"""Run the bot connection with auto-reconnect."""
@@ -88,11 +87,10 @@ class QQChannel(BaseChannel):
async def stop(self) -> None:
"""Stop the QQ bot."""
self._running = False
if self._bot_task:
self._bot_task.cancel()
if self._client:
try:
await self._bot_task
except asyncio.CancelledError:
await self._client.close()
except Exception:
pass
logger.info("QQ bot stopped")
@@ -130,5 +128,5 @@ class QQChannel(BaseChannel):
content=content,
metadata={"message_id": data.id},
)
except Exception as e:
logger.error("Error handling QQ message: {}", e)
except Exception:
logger.exception("Error handling QQ message")
+15 -12
View File
@@ -179,18 +179,21 @@ class SlackChannel(BaseChannel):
except Exception as e:
logger.debug("Slack reactions_add failed: {}", e)
await self._handle_message(
sender_id=sender_id,
chat_id=chat_id,
content=text,
metadata={
"slack": {
"event": event,
"thread_ts": thread_ts,
"channel_type": channel_type,
}
},
)
try:
await self._handle_message(
sender_id=sender_id,
chat_id=chat_id,
content=text,
metadata={
"slack": {
"event": event,
"thread_ts": thread_ts,
"channel_type": channel_type,
}
},
)
except Exception:
logger.exception("Error handling Slack message from {}", sender_id)
def _is_allowed(self, sender_id: str, chat_id: str, channel_type: str) -> bool:
if channel_type == "im":
+1
View File
@@ -260,6 +260,7 @@ class MCPServerConfig(Base):
env: dict[str, str] = Field(default_factory=dict) # Stdio: extra env vars
url: str = "" # HTTP: streamable HTTP endpoint URL
headers: dict[str, str] = Field(default_factory=dict) # HTTP: Custom HTTP Headers
tool_timeout: int = 30 # Seconds before a tool call is cancelled
class ToolsConfig(Base):
+1 -1
View File
@@ -40,7 +40,7 @@ class CustomProvider(LLMProvider):
return LLMResponse(
content=msg.content, tool_calls=tool_calls, finish_reason=choice.finish_reason or "stop",
usage={"prompt_tokens": u.prompt_tokens, "completion_tokens": u.completion_tokens, "total_tokens": u.total_tokens} if u else {},
reasoning_content=getattr(msg, "reasoning_content", None),
reasoning_content=getattr(msg, "reasoning_content", None) or None,
)
def get_default_model(self) -> str:
+1 -1
View File
@@ -257,7 +257,7 @@ class LiteLLMProvider(LLMProvider):
"total_tokens": response.usage.total_tokens,
}
reasoning_content = getattr(message, "reasoning_content", None)
reasoning_content = getattr(message, "reasoning_content", None) or None
return LLMResponse(
content=message.content,
+6 -3
View File
@@ -1,6 +1,7 @@
"""Session management for conversation history."""
import json
import shutil
from pathlib import Path
from dataclasses import dataclass, field
from datetime import datetime
@@ -108,9 +109,11 @@ class SessionManager:
if not path.exists():
legacy_path = self._get_legacy_session_path(key)
if legacy_path.exists():
import shutil
shutil.move(str(legacy_path), str(path))
logger.info("Migrated session {} from legacy path", key)
try:
shutil.move(str(legacy_path), str(path))
logger.info("Migrated session {} from legacy path", key)
except Exception:
logger.exception("Failed to migrate session {}", key)
if not path.exists():
return None