Compare commits
29
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
34584c3a2e | ||
|
|
53e09b924c | ||
|
|
1b302ab4bf | ||
|
|
3c681f1639 | ||
|
|
1ff3356d1b | ||
|
|
5193e34803 | ||
|
|
8f8fc81135 | ||
|
|
d4abb3d06f | ||
|
|
b2570f1a62 | ||
|
|
f19b5f5929 | ||
|
|
8e829396b2 | ||
|
|
e8e8ca6700 | ||
|
|
f1cbd4d730 | ||
|
|
f7cebfe7f3 | ||
|
|
b854d9a888 | ||
|
|
83d2acf07f | ||
|
|
eee9c38953 | ||
|
|
e782318338 | ||
|
|
dc94aa76cc | ||
|
|
5cf019c21e | ||
|
|
790bdd6b8a | ||
|
|
b25c09f5ed | ||
|
|
9e8c910ab1 | ||
|
|
cc10e20a47 | ||
|
|
34ed4345fc | ||
|
|
1a85333e4c | ||
|
|
3c587c788a | ||
|
|
303d123527 | ||
|
|
61c2cb4ac4 |
+150
-9
@@ -40,8 +40,8 @@ class AgentLoop:
|
|||||||
5. Sends responses back
|
5. Sends responses back
|
||||||
"""
|
"""
|
||||||
|
|
||||||
# Server-side context management: Anthropic trims old tool results and preserves all
|
# Server-side context management: Anthropic preserves all thinking blocks
|
||||||
# thinking blocks (keep="all" maximises cache hits). Client keeps full history.
|
# and clears old tool results only when approaching the 200k context limit.
|
||||||
CONTEXT_MANAGEMENT = {
|
CONTEXT_MANAGEMENT = {
|
||||||
"edits": [
|
"edits": [
|
||||||
{
|
{
|
||||||
@@ -50,7 +50,11 @@ class AgentLoop:
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
"type": "clear_tool_uses_20250919",
|
"type": "clear_tool_uses_20250919",
|
||||||
"trigger": {"type": "input_tokens", "value": 80000},
|
# Raised from 80k to 195k to avoid premature cache invalidation.
|
||||||
|
# For conversations with few tool uses (e.g., 18 uses over 182k tokens),
|
||||||
|
# cache stability (saves 169k/turn) >> clearing benefit (13-26k one-time).
|
||||||
|
# Leaves 5k headroom before hitting 200k standard context limit.
|
||||||
|
"trigger": {"type": "input_tokens", "value": 195000},
|
||||||
"keep": {"type": "tool_uses", "value": 5},
|
"keep": {"type": "tool_uses", "value": 5},
|
||||||
},
|
},
|
||||||
]
|
]
|
||||||
@@ -151,9 +155,28 @@ class AgentLoop:
|
|||||||
# Register native Anthropic tools
|
# Register native Anthropic tools
|
||||||
self.tools.register(BashTool20250124())
|
self.tools.register(BashTool20250124())
|
||||||
self.tools.register(EditTool20250728())
|
self.tools.register(EditTool20250728())
|
||||||
self.tools.register(ComputerTool20251124())
|
# self.tools.register(ComputerTool20251124()) # Disabled - VM unavailable
|
||||||
|
|
||||||
logger.info("Registered native Anthropic tools: bash, text_editor, computer")
|
logger.info("Registered native Anthropic tools: bash, text_editor")
|
||||||
|
|
||||||
|
# Register mem0 memory tools (if enabled)
|
||||||
|
from nanobot.agent.memory_mem0 import HAS_MEM0
|
||||||
|
if self.mem0_config and self.mem0_config.get("enabled") and HAS_MEM0:
|
||||||
|
from nanobot.agent.memory_mem0 import Mem0MemoryStore
|
||||||
|
from nanobot.agent.tools.memory_tools import (
|
||||||
|
Mem0ToolContext, MemorySearchTool, MemoryListTool,
|
||||||
|
MemoryAddTool, MemoryUpdateTool, MemoryDeleteTool,
|
||||||
|
MemoryConsolidateTool,
|
||||||
|
)
|
||||||
|
store = Mem0MemoryStore(self.workspace, config=self.mem0_config)
|
||||||
|
self._mem0_ctx = Mem0ToolContext(store, self._consolidate_memory)
|
||||||
|
self.tools.register(MemorySearchTool(self._mem0_ctx))
|
||||||
|
self.tools.register(MemoryListTool(self._mem0_ctx))
|
||||||
|
self.tools.register(MemoryAddTool(self._mem0_ctx))
|
||||||
|
self.tools.register(MemoryUpdateTool(self._mem0_ctx))
|
||||||
|
self.tools.register(MemoryDeleteTool(self._mem0_ctx))
|
||||||
|
self.tools.register(MemoryConsolidateTool(self._mem0_ctx))
|
||||||
|
logger.info("Registered mem0 memory tools")
|
||||||
|
|
||||||
async def run(self) -> None:
|
async def run(self) -> None:
|
||||||
"""Run the agent loop, processing messages from the bus."""
|
"""Run the agent loop, processing messages from the bus."""
|
||||||
@@ -323,6 +346,7 @@ class AgentLoop:
|
|||||||
message_tool = self.tools.get("message")
|
message_tool = self.tools.get("message")
|
||||||
if isinstance(message_tool, MessageTool):
|
if isinstance(message_tool, MessageTool):
|
||||||
message_tool.set_context(msg.channel, msg.chat_id)
|
message_tool.set_context(msg.channel, msg.chat_id)
|
||||||
|
message_tool.start_turn()
|
||||||
|
|
||||||
spawn_tool = self.tools.get("spawn")
|
spawn_tool = self.tools.get("spawn")
|
||||||
if isinstance(spawn_tool, SpawnTool):
|
if isinstance(spawn_tool, SpawnTool):
|
||||||
@@ -332,6 +356,9 @@ class AgentLoop:
|
|||||||
if isinstance(cron_tool, CronTool):
|
if isinstance(cron_tool, CronTool):
|
||||||
cron_tool.set_context(msg.channel, msg.chat_id)
|
cron_tool.set_context(msg.channel, msg.chat_id)
|
||||||
|
|
||||||
|
if hasattr(self, '_mem0_ctx'):
|
||||||
|
self._mem0_ctx.set_context(msg.channel, msg.chat_id, session)
|
||||||
|
|
||||||
# Track media for this turn (screenshots from computer tool)
|
# Track media for this turn (screenshots from computer tool)
|
||||||
media_paths_for_turn: list[str] = []
|
media_paths_for_turn: list[str] = []
|
||||||
|
|
||||||
@@ -518,6 +545,12 @@ class AgentLoop:
|
|||||||
else:
|
else:
|
||||||
final_content = "I've completed processing but have no response to give."
|
final_content = "I've completed processing but have no response to give."
|
||||||
|
|
||||||
|
# Check if message tool already sent to same target (suppress final reply)
|
||||||
|
message_tool = self.tools.get("message")
|
||||||
|
if isinstance(message_tool, MessageTool) and message_tool._sent_in_turn:
|
||||||
|
logger.info(f"Suppressing final reply to {msg.channel}:{msg.chat_id} (message tool already sent)")
|
||||||
|
return None
|
||||||
|
|
||||||
# Log response preview
|
# Log response preview
|
||||||
preview = final_content[:120] + "..." if len(final_content) > 120 else final_content
|
preview = final_content[:120] + "..." if len(final_content) > 120 else final_content
|
||||||
logger.info(f"Response to {msg.channel}:{msg.sender_id}: {preview}")
|
logger.info(f"Response to {msg.channel}:{msg.sender_id}: {preview}")
|
||||||
@@ -541,15 +574,39 @@ class AgentLoop:
|
|||||||
reasoning_content=final_reasoning,
|
reasoning_content=final_reasoning,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Save to session: user message + full tool chain (tool_use, tool_results, thinking, final reply)
|
# Save to session: mem0 context (if present) + user message + full tool chain
|
||||||
# Store current_message (not msg.content) so the time prefix is preserved
|
# Store current_message (not msg.content) so the time prefix is preserved
|
||||||
# and cache keys match on subsequent turns
|
|
||||||
# Include sender_id to distinguish real user messages from system-generated ones
|
# Include sender_id to distinguish real user messages from system-generated ones
|
||||||
|
|
||||||
|
# Find and save mem0 injection (appears just before current user message)
|
||||||
|
# build_messages returns: [...history, mem0_user, mem0_asst, current_user]
|
||||||
|
# turn_start = len(messages), so mem0 is at turn_start-3 and turn_start-2
|
||||||
|
# This makes mem0 part of immutable history, stabilizing cache across turns
|
||||||
|
if turn_start >= 3:
|
||||||
|
potential_mem0_user = messages[turn_start - 3]
|
||||||
|
potential_mem0_asst = messages[turn_start - 2]
|
||||||
|
if (potential_mem0_user.get("role") == "user" and
|
||||||
|
potential_mem0_user.get("content") == "[Memory context]" and
|
||||||
|
potential_mem0_asst.get("role") == "assistant"):
|
||||||
|
session.add_raw_message(potential_mem0_user)
|
||||||
|
session.add_raw_message(potential_mem0_asst)
|
||||||
|
|
||||||
session.add_message("user", current_message, sender_id=msg.sender_id)
|
session.add_message("user", current_message, sender_id=msg.sender_id)
|
||||||
for chain_msg in messages[turn_start:]:
|
for chain_msg in messages[turn_start:]:
|
||||||
session.add_raw_message(chain_msg)
|
session.add_raw_message(chain_msg)
|
||||||
self.sessions.save(session)
|
self.sessions.save(session)
|
||||||
|
|
||||||
|
# Deferred trim: if memory_consolidate ran mid-turn, it set a checkpoint
|
||||||
|
# marking where to trim. Now that the turn's tool chain is fully saved,
|
||||||
|
# we can safely trim to that checkpoint.
|
||||||
|
checkpoint = getattr(session, '_trim_checkpoint', None)
|
||||||
|
if checkpoint is not None:
|
||||||
|
old_size = len(session.messages)
|
||||||
|
session.messages = session.messages[checkpoint:]
|
||||||
|
session._trim_checkpoint = None
|
||||||
|
self.sessions.save(session)
|
||||||
|
logger.info(f"Deferred trim applied: {old_size} -> {len(session.messages)} messages (checkpoint={checkpoint})")
|
||||||
|
|
||||||
return OutboundMessage(
|
return OutboundMessage(
|
||||||
channel=msg.channel,
|
channel=msg.channel,
|
||||||
chat_id=msg.chat_id,
|
chat_id=msg.chat_id,
|
||||||
@@ -737,12 +794,32 @@ class AgentLoop:
|
|||||||
reasoning_content=final_reasoning,
|
reasoning_content=final_reasoning,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Save to session: user message + full tool chain
|
# Save to session: mem0 (if present) + user message + full tool chain
|
||||||
|
# Find and save mem0 injection for cache stability
|
||||||
|
if turn_start >= 3:
|
||||||
|
potential_mem0_user = messages[turn_start - 3]
|
||||||
|
potential_mem0_asst = messages[turn_start - 2]
|
||||||
|
if (potential_mem0_user.get("role") == "user" and
|
||||||
|
potential_mem0_user.get("content") == "[Memory context]" and
|
||||||
|
potential_mem0_asst.get("role") == "assistant"):
|
||||||
|
session.add_raw_message(potential_mem0_user)
|
||||||
|
session.add_raw_message(potential_mem0_asst)
|
||||||
|
|
||||||
session.add_message("user", f"[System: {msg.sender_id}] {msg.content}")
|
session.add_message("user", f"[System: {msg.sender_id}] {msg.content}")
|
||||||
for chain_msg in messages[turn_start:]:
|
for chain_msg in messages[turn_start:]:
|
||||||
session.add_raw_message(chain_msg)
|
session.add_raw_message(chain_msg)
|
||||||
self.sessions.save(session)
|
self.sessions.save(session)
|
||||||
|
|
||||||
|
# Deferred trim: same logic as _process_message
|
||||||
|
# System messages (including subagents) can trigger consolidation
|
||||||
|
checkpoint = getattr(session, '_trim_checkpoint', None)
|
||||||
|
if checkpoint is not None:
|
||||||
|
old_size = len(session.messages)
|
||||||
|
session.messages = session.messages[checkpoint:]
|
||||||
|
session._trim_checkpoint = None
|
||||||
|
self.sessions.save(session)
|
||||||
|
logger.info(f"Deferred trim applied: {old_size} -> {len(session.messages)} messages (checkpoint={checkpoint})")
|
||||||
|
|
||||||
# Return original content (not signed) for outbound, but with suppressed metadata
|
# Return original content (not signed) for outbound, but with suppressed metadata
|
||||||
return OutboundMessage(
|
return OutboundMessage(
|
||||||
channel=origin_channel,
|
channel=origin_channel,
|
||||||
@@ -751,6 +828,54 @@ class AgentLoop:
|
|||||||
metadata=outbound_metadata,
|
metadata=outbound_metadata,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _find_clean_boundary_before(messages: list[dict], target_pos: int) -> int:
|
||||||
|
"""Find a clean user message boundary at or before target position.
|
||||||
|
|
||||||
|
Returns the index of a user message at or before target_pos,
|
||||||
|
or target_pos if no user message is found.
|
||||||
|
"""
|
||||||
|
if not messages or target_pos <= 0:
|
||||||
|
return 0
|
||||||
|
if target_pos >= len(messages):
|
||||||
|
return len(messages)
|
||||||
|
|
||||||
|
# Walk backward from target to find a user message
|
||||||
|
for i in range(target_pos, -1, -1):
|
||||||
|
if messages[i].get("role") == "user":
|
||||||
|
return i
|
||||||
|
|
||||||
|
# No user message found, return target position
|
||||||
|
return target_pos
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _trim_to_clean_boundary(messages: list[dict], keep_count: int) -> list[dict]:
|
||||||
|
"""Trim messages to approximately keep_count, starting at a user message boundary.
|
||||||
|
|
||||||
|
Naive slicing (messages[-keep_count:]) can cut into a tool chain, leaving
|
||||||
|
orphaned tool_result messages at the start. This finds the nearest user
|
||||||
|
message (role="user") at or before the cut point and trims there.
|
||||||
|
"""
|
||||||
|
if not messages or keep_count <= 0:
|
||||||
|
return []
|
||||||
|
if keep_count >= len(messages):
|
||||||
|
return messages
|
||||||
|
|
||||||
|
cut = len(messages) - keep_count
|
||||||
|
# Walk forward from cut to find a "user" role message (start of a turn)
|
||||||
|
# that isn't a tool result. Tool results have role="tool", user messages
|
||||||
|
# have role="user" — but after conversion, tool results ARE user messages.
|
||||||
|
# In session storage, they're still role="tool", so we look for role="user".
|
||||||
|
for i in range(cut, len(messages)):
|
||||||
|
if messages[i].get("role") == "user":
|
||||||
|
return messages[i:]
|
||||||
|
# If no user message found after cut, try walking backward
|
||||||
|
for i in range(cut - 1, -1, -1):
|
||||||
|
if messages[i].get("role") == "user":
|
||||||
|
return messages[i:]
|
||||||
|
# Fallback: return everything (shouldn't happen in practice)
|
||||||
|
return messages
|
||||||
|
|
||||||
async def _consolidate_memory(self, session, archive_all: bool = False) -> None:
|
async def _consolidate_memory(self, session, archive_all: bool = False) -> None:
|
||||||
"""Consolidate session into MEMORY.md + HISTORY.md.
|
"""Consolidate session into MEMORY.md + HISTORY.md.
|
||||||
|
|
||||||
@@ -774,6 +899,22 @@ class AgentLoop:
|
|||||||
archive_all=archive_all,
|
archive_all=archive_all,
|
||||||
memory_window=self.memory_window,
|
memory_window=self.memory_window,
|
||||||
)
|
)
|
||||||
|
# archive_all (/new) runs at a turn boundary — safe to trim now.
|
||||||
|
# Mid-turn (memory_consolidate tool) — defer trim to end of turn
|
||||||
|
# to avoid orphaning tool_use IDs in the active tool chain.
|
||||||
|
if archive_all:
|
||||||
|
session.messages = []
|
||||||
|
self.sessions.save(session)
|
||||||
|
logger.info("Mem0 consolidation done, session cleared (archive_all)")
|
||||||
|
else:
|
||||||
|
keep_count = min(10, max(2, self.memory_window // 2))
|
||||||
|
# Set checkpoint at current session size minus keep_count
|
||||||
|
# This preserves the intended trim point regardless of messages added later
|
||||||
|
checkpoint = max(0, len(session.messages) - keep_count)
|
||||||
|
# Find clean boundary at or before checkpoint
|
||||||
|
checkpoint = self._find_clean_boundary_before(session.messages, checkpoint)
|
||||||
|
session._trim_checkpoint = checkpoint
|
||||||
|
logger.info(f"Mem0 consolidation done, trim deferred (checkpoint={checkpoint}, current_size={len(session.messages)})")
|
||||||
return
|
return
|
||||||
else:
|
else:
|
||||||
memory = MemoryStore(self.workspace)
|
memory = MemoryStore(self.workspace)
|
||||||
@@ -864,7 +1005,7 @@ Respond with ONLY valid JSON, no markdown fences."""
|
|||||||
if update != current_memory:
|
if update != current_memory:
|
||||||
memory.write_long_term(update)
|
memory.write_long_term(update)
|
||||||
|
|
||||||
session.messages = session.messages[-keep_count:] if keep_count else []
|
session.messages = self._trim_to_clean_boundary(session.messages, keep_count) if keep_count else []
|
||||||
self.sessions.save(session)
|
self.sessions.save(session)
|
||||||
logger.info(f"Memory consolidation done, session trimmed to {len(session.messages)} messages")
|
logger.info(f"Memory consolidation done, session trimmed to {len(session.messages)} messages")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -87,11 +87,7 @@ class MemoryStore:
|
|||||||
keep_count = memory_window // 2
|
keep_count = memory_window // 2
|
||||||
if len(session.messages) <= keep_count:
|
if len(session.messages) <= keep_count:
|
||||||
return True
|
return True
|
||||||
if len(session.messages) - session.last_consolidated <= 0:
|
old_messages = session.messages[:-keep_count]
|
||||||
return True
|
|
||||||
old_messages = session.messages[session.last_consolidated:-keep_count]
|
|
||||||
if not old_messages:
|
|
||||||
return True
|
|
||||||
logger.info("Memory consolidation: {} to consolidate, {} keep", len(old_messages), keep_count)
|
logger.info("Memory consolidation: {} to consolidate, {} keep", len(old_messages), keep_count)
|
||||||
|
|
||||||
lines = []
|
lines = []
|
||||||
@@ -142,8 +138,7 @@ class MemoryStore:
|
|||||||
if update != current_memory:
|
if update != current_memory:
|
||||||
self.write_long_term(update)
|
self.write_long_term(update)
|
||||||
|
|
||||||
session.last_consolidated = 0 if archive_all else len(session.messages) - keep_count
|
logger.info("Memory consolidation done: {} messages total", len(session.messages))
|
||||||
logger.info("Memory consolidation done: {} messages, last_consolidated={}", len(session.messages), session.last_consolidated)
|
|
||||||
return True
|
return True
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("Memory consolidation failed")
|
logger.exception("Memory consolidation failed")
|
||||||
|
|||||||
@@ -47,6 +47,7 @@ class Mem0MemoryStore:
|
|||||||
|
|
||||||
today = datetime.now().strftime("%Y-%m-%d")
|
today = datetime.now().strftime("%Y-%m-%d")
|
||||||
custom_prompt = f"Extract dated facts from this conversation as JSON: {{\"facts\": [...]}}. Today is {today}.\n\n"
|
custom_prompt = f"Extract dated facts from this conversation as JSON: {{\"facts\": [...]}}. Today is {today}.\n\n"
|
||||||
|
self.custom_prompt = custom_prompt
|
||||||
|
|
||||||
# Initialize mem0 with optional config + custom prompt
|
# Initialize mem0 with optional config + custom prompt
|
||||||
# Extract only MemoryConfig-relevant fields
|
# Extract only MemoryConfig-relevant fields
|
||||||
@@ -57,9 +58,6 @@ class Mem0MemoryStore:
|
|||||||
if key in raw_config:
|
if key in raw_config:
|
||||||
mem0_cfg_dict[key] = raw_config[key]
|
mem0_cfg_dict[key] = raw_config[key]
|
||||||
logger.debug(f"Extracted for MemoryConfig: {list(mem0_cfg_dict.keys())}")
|
logger.debug(f"Extracted for MemoryConfig: {list(mem0_cfg_dict.keys())}")
|
||||||
logger.debug(f"Custom prompt length: {len(custom_prompt)} chars")
|
|
||||||
self.custom_prompt = custom_prompt
|
|
||||||
mem0_cfg_dict["custom_fact_extraction_prompt"] = custom_prompt
|
|
||||||
mem0_config = MemoryConfig(**mem0_cfg_dict)
|
mem0_config = MemoryConfig(**mem0_cfg_dict)
|
||||||
logger.debug(f"MemoryConfig created: vector_store={mem0_config.vector_store.provider if mem0_config.vector_store else None}")
|
logger.debug(f"MemoryConfig created: vector_store={mem0_config.vector_store.provider if mem0_config.vector_store else None}")
|
||||||
self.memory = Memory(config=mem0_config)
|
self.memory = Memory(config=mem0_config)
|
||||||
@@ -156,21 +154,15 @@ class Mem0MemoryStore:
|
|||||||
provider: Any,
|
provider: Any,
|
||||||
model: str,
|
model: str,
|
||||||
) -> list[str]:
|
) -> list[str]:
|
||||||
"""
|
"""Extract facts from conversation using the main agent's LLM provider."""
|
||||||
Extract facts from conversation using the main agent's LLM provider.
|
|
||||||
|
|
||||||
Uses the same provider/model already running (e.g. Haiku via Claude Max),
|
|
||||||
avoiding a separate LLM call to mem0's default GPT-nano.
|
|
||||||
"""
|
|
||||||
import json as _json
|
import json as _json
|
||||||
|
|
||||||
# Build conversation text for extraction
|
|
||||||
conv_text = ""
|
conv_text = ""
|
||||||
for msg in messages:
|
for msg in messages:
|
||||||
role = msg.get("role", "unknown")
|
role = msg.get("role", "unknown")
|
||||||
content = msg.get("content", "")
|
content_val = msg.get("content", "")
|
||||||
if isinstance(content, str) and content.strip():
|
if isinstance(content_val, str) and content_val.strip():
|
||||||
conv_text += f"{role}: {content}\n\n"
|
conv_text += f"{role}: {content_val}\n\n"
|
||||||
|
|
||||||
if not conv_text.strip():
|
if not conv_text.strip():
|
||||||
return []
|
return []
|
||||||
@@ -186,22 +178,19 @@ class Mem0MemoryStore:
|
|||||||
max_tokens=2000,
|
max_tokens=2000,
|
||||||
temperature=0.3,
|
temperature=0.3,
|
||||||
)
|
)
|
||||||
|
text = (response.content or "").strip()
|
||||||
# Parse the JSON response — LLMResponse.content is a string
|
|
||||||
text = response.content or ""
|
|
||||||
# Strip markdown code fences if present
|
|
||||||
text = text.strip()
|
|
||||||
if text.startswith("```"):
|
if text.startswith("```"):
|
||||||
text = text.split("\n", 1)[1] if "\n" in text else text[3:]
|
text = text.split("```")[1]
|
||||||
if text.endswith("```"):
|
if text.startswith("json"):
|
||||||
text = text[:-3]
|
text = text[4:]
|
||||||
text = text.strip()
|
text = text.strip()
|
||||||
|
|
||||||
data = _json.loads(text)
|
data = _json.loads(text)
|
||||||
facts = data.get("facts", [])
|
facts = data.get("facts", [])
|
||||||
|
if not isinstance(facts, list):
|
||||||
|
logger.warning(f"LLM returned non-list facts: {type(facts)}")
|
||||||
|
return []
|
||||||
logger.debug(f"Extracted {len(facts)} facts using {model}")
|
logger.debug(f"Extracted {len(facts)} facts using {model}")
|
||||||
return facts
|
return facts
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Fact extraction failed: {e}")
|
logger.error(f"Fact extraction failed: {e}")
|
||||||
return []
|
return []
|
||||||
@@ -212,12 +201,7 @@ class Mem0MemoryStore:
|
|||||||
user_id: str,
|
user_id: str,
|
||||||
session_id: str | None = None,
|
session_id: str | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""Store pre-extracted facts in mem0 with infer=False."""
|
||||||
Store pre-extracted facts in mem0 with infer=False.
|
|
||||||
|
|
||||||
Bypasses mem0's built-in LLM extraction — facts are already
|
|
||||||
in final form from extract_facts().
|
|
||||||
"""
|
|
||||||
if not facts:
|
if not facts:
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -309,7 +293,8 @@ class Mem0MemoryStore:
|
|||||||
"""
|
"""
|
||||||
Consolidate session messages into mem0 memory.
|
Consolidate session messages into mem0 memory.
|
||||||
|
|
||||||
Facts are extracted using the main agent's LLM provider, then stored with infer=False.
|
Unlike the original MemoryStore, mem0 handles extraction automatically,
|
||||||
|
so this just needs to feed recent messages to mem0.
|
||||||
|
|
||||||
Returns True on success.
|
Returns True on success.
|
||||||
"""
|
"""
|
||||||
@@ -328,8 +313,8 @@ class Mem0MemoryStore:
|
|||||||
if len(session.messages) <= keep_count:
|
if len(session.messages) <= keep_count:
|
||||||
return True
|
return True
|
||||||
|
|
||||||
# Get unconsolidated messages
|
# Consolidate messages except the most recent (kept for context)
|
||||||
start_idx = session.last_consolidated
|
start_idx = 0
|
||||||
end_idx = len(session.messages) - keep_count
|
end_idx = len(session.messages) - keep_count
|
||||||
|
|
||||||
if end_idx <= start_idx:
|
if end_idx <= start_idx:
|
||||||
@@ -351,21 +336,9 @@ class Mem0MemoryStore:
|
|||||||
role = msg.get("role")
|
role = msg.get("role")
|
||||||
content = msg.get("content")
|
content = msg.get("content")
|
||||||
|
|
||||||
# Keep tool results but truncate long ones — they often contain
|
# Skip tool results — raw bash output, file contents, and JSON
|
||||||
# the actual substance (file reads, search results, web pages).
|
# get misinterpreted by the extraction LLM as user interests
|
||||||
# The extraction prompt handles ignoring code/JSON noise.
|
|
||||||
if role == "tool":
|
if role == "tool":
|
||||||
if isinstance(content, list):
|
|
||||||
text_parts = [
|
|
||||||
block.get("content", "") if isinstance(block, dict) else str(block)
|
|
||||||
for block in content
|
|
||||||
]
|
|
||||||
content = " ".join(text_parts).strip()
|
|
||||||
if isinstance(content, str) and len(content) > 2000:
|
|
||||||
content = content[:2000]
|
|
||||||
if not content or (isinstance(content, str) and len(content.strip()) < 10):
|
|
||||||
continue
|
|
||||||
mem0_messages.append({"role": "user", "content": content})
|
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Skip system messages — they're boilerplate instructions, not facts
|
# Skip system messages — they're boilerplate instructions, not facts
|
||||||
@@ -413,15 +386,8 @@ class Mem0MemoryStore:
|
|||||||
facts = await self.extract_facts(mem0_messages, provider, model)
|
facts = await self.extract_facts(mem0_messages, provider, model)
|
||||||
self.store_facts(facts, user_id=user_id, session_id=session.key)
|
self.store_facts(facts, user_id=user_id, session_id=session.key)
|
||||||
|
|
||||||
# Update consolidation marker
|
|
||||||
if archive_all:
|
|
||||||
session.last_consolidated = len(session.messages)
|
|
||||||
else:
|
|
||||||
session.last_consolidated = end_idx
|
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
f"Mem0 consolidation done: {len(session.messages)} messages, "
|
f"Mem0 consolidation done: {len(session.messages)} messages total"
|
||||||
f"last_consolidated={session.last_consolidated}"
|
|
||||||
)
|
)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
|||||||
@@ -73,7 +73,7 @@ class SubagentManager:
|
|||||||
origin_metadata: Optional metadata to propagate to announcement (e.g. suppress_output).
|
origin_metadata: Optional metadata to propagate to announcement (e.g. suppress_output).
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Status message indicating the subagent was started.
|
Task ID of the spawned subagent.
|
||||||
"""
|
"""
|
||||||
task_id = str(uuid.uuid4())[:8]
|
task_id = str(uuid.uuid4())[:8]
|
||||||
display_label = label or task[:30] + ("..." if len(task) > 30 else "")
|
display_label = label or task[:30] + ("..." if len(task) > 30 else "")
|
||||||
@@ -83,18 +83,18 @@ class SubagentManager:
|
|||||||
"chat_id": origin_chat_id,
|
"chat_id": origin_chat_id,
|
||||||
"metadata": origin_metadata or {},
|
"metadata": origin_metadata or {},
|
||||||
}
|
}
|
||||||
|
|
||||||
# Create background task
|
# Create background task
|
||||||
bg_task = asyncio.create_task(
|
bg_task = asyncio.create_task(
|
||||||
self._run_subagent(task_id, task, display_label, origin, model=model)
|
self._run_subagent(task_id, task, display_label, origin, model=model)
|
||||||
)
|
)
|
||||||
self._running_tasks[task_id] = bg_task
|
self._running_tasks[task_id] = bg_task
|
||||||
|
|
||||||
# Cleanup when done
|
# Cleanup when done
|
||||||
bg_task.add_done_callback(lambda _: self._running_tasks.pop(task_id, None))
|
bg_task.add_done_callback(lambda _: self._running_tasks.pop(task_id, None))
|
||||||
|
|
||||||
logger.info(f"Spawned subagent [{task_id}]: {display_label}")
|
logger.info(f"Spawned subagent [{task_id}]: {display_label}")
|
||||||
return f"Subagent [{display_label}] started. Task ID: {task_id}"
|
return task_id
|
||||||
|
|
||||||
async def _run_subagent(
|
async def _run_subagent(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -1,114 +1,117 @@
|
|||||||
"""BashTool20250124 - Persistent bash session with sentinel-based output.
|
"""BashTool20250124 - Persistent bash session with async buffer polling.
|
||||||
|
|
||||||
Anthropic's native bash_20250124 tool with a long-running session.
|
Based on Anthropic's reference implementation from anthropic-quickstarts.
|
||||||
|
Uses asyncio.create_subprocess_shell + direct buffer reads instead of
|
||||||
|
threaded readline, which avoids exhausting the default ThreadPoolExecutor.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import subprocess
|
import os
|
||||||
import uuid
|
|
||||||
from typing import Any, Literal
|
from typing import Any, Literal
|
||||||
|
|
||||||
from nanobot.agent.tools.anthropic.base import BaseAnthropicTool, ToolResult
|
from nanobot.agent.tools.anthropic.base import BaseAnthropicTool, ToolResult, ToolError
|
||||||
|
|
||||||
|
|
||||||
class _BashSession:
|
class _BashSession:
|
||||||
"""Manages a persistent bash subprocess with sentinel-based output reading."""
|
"""A session of a bash shell.
|
||||||
|
|
||||||
|
Uses asyncio subprocess with direct buffer polling — no threads.
|
||||||
|
Based on anthropics/anthropic-quickstarts computer-use-demo.
|
||||||
|
"""
|
||||||
|
|
||||||
|
command: str = "/bin/bash"
|
||||||
|
_output_delay: float = 0.2 # seconds between buffer polls
|
||||||
|
_timeout: float = 120.0 # seconds
|
||||||
|
_sentinel: str = "<<exit>>"
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.process: subprocess.Popen | None = None
|
self._started = False
|
||||||
self._start()
|
self._timed_out = False
|
||||||
|
self._process: asyncio.subprocess.Process | None = None
|
||||||
|
|
||||||
def _start(self):
|
async def start(self):
|
||||||
"""Start the bash process."""
|
if self._started:
|
||||||
self.process = subprocess.Popen(
|
return
|
||||||
["bash"],
|
|
||||||
stdin=subprocess.PIPE,
|
self._process = await asyncio.create_subprocess_shell(
|
||||||
stdout=subprocess.PIPE,
|
self.command,
|
||||||
stderr=subprocess.STDOUT,
|
preexec_fn=os.setsid,
|
||||||
text=True,
|
shell=True,
|
||||||
bufsize=1,
|
bufsize=0,
|
||||||
|
stdin=asyncio.subprocess.PIPE,
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE,
|
||||||
)
|
)
|
||||||
|
self._started = True
|
||||||
|
|
||||||
def restart(self):
|
def stop(self):
|
||||||
"""Restart the bash session."""
|
"""Terminate the bash shell."""
|
||||||
if self.process:
|
if not self._started:
|
||||||
self.process.terminate()
|
return
|
||||||
try:
|
if self._process and self._process.returncode is None:
|
||||||
self.process.wait(timeout=5)
|
self._process.terminate()
|
||||||
except subprocess.TimeoutExpired:
|
|
||||||
self.process.kill()
|
|
||||||
self.process.wait()
|
|
||||||
self._start()
|
|
||||||
|
|
||||||
async def run_command(self, command: str, timeout: float = 120.0) -> str:
|
async def run(self, command: str) -> ToolResult:
|
||||||
"""Run a command in the persistent bash session.
|
"""Execute a command in the bash shell."""
|
||||||
|
if not self._started:
|
||||||
|
raise ToolError("Session has not started.")
|
||||||
|
if self._process is None or self._process.returncode is not None:
|
||||||
|
return ToolResult(
|
||||||
|
system="tool must be restarted",
|
||||||
|
error=f"bash has exited with returncode "
|
||||||
|
f"{self._process.returncode if self._process else 'unknown'}",
|
||||||
|
)
|
||||||
|
if self._timed_out:
|
||||||
|
raise ToolError(
|
||||||
|
f"timed out: bash has not returned in {self._timeout} seconds "
|
||||||
|
"and must be restarted",
|
||||||
|
)
|
||||||
|
|
||||||
Uses a unique sentinel to detect command completion.
|
assert self._process.stdin
|
||||||
|
assert self._process.stdout
|
||||||
|
assert self._process.stderr
|
||||||
|
|
||||||
Args:
|
# Send command + sentinel on its own line so heredoc terminators
|
||||||
command: Bash command to execute
|
# aren't corrupted (EOF; echo '...' ≠ EOF)
|
||||||
timeout: Maximum time to wait for command completion (seconds)
|
self._process.stdin.write(
|
||||||
|
command.encode() + f"\necho '{self._sentinel}'\n".encode()
|
||||||
|
)
|
||||||
|
await self._process.stdin.drain()
|
||||||
|
|
||||||
Returns:
|
# Poll stdout buffer until sentinel appears — no threads involved
|
||||||
Command output (stdout + stderr combined)
|
try:
|
||||||
|
async with asyncio.timeout(self._timeout):
|
||||||
|
while True:
|
||||||
|
await asyncio.sleep(self._output_delay)
|
||||||
|
output = self._process.stdout._buffer.decode()
|
||||||
|
if self._sentinel in output:
|
||||||
|
output = output[: output.index(self._sentinel)]
|
||||||
|
break
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
self._timed_out = True
|
||||||
|
raise ToolError(
|
||||||
|
f"timed out: bash has not returned in {self._timeout} seconds "
|
||||||
|
"and must be restarted",
|
||||||
|
) from None
|
||||||
|
|
||||||
Raises:
|
if output.endswith("\n"):
|
||||||
asyncio.TimeoutError: If command doesn't complete within timeout
|
output = output[:-1]
|
||||||
RuntimeError: If bash process has died
|
|
||||||
"""
|
|
||||||
if not self.process or self.process.poll() is not None:
|
|
||||||
raise RuntimeError("Bash process has died")
|
|
||||||
|
|
||||||
# Generate unique sentinel
|
error = self._process.stderr._buffer.decode()
|
||||||
sentinel = f"<<BASH_COMMAND_DONE_{uuid.uuid4().hex}>>"
|
if error.endswith("\n"):
|
||||||
|
error = error[:-1]
|
||||||
|
|
||||||
# Send command + sentinel
|
# Clear buffers for next command
|
||||||
full_command = f"{command}\necho '{sentinel}'\n"
|
self._process.stdout._buffer.clear()
|
||||||
self.process.stdin.write(full_command)
|
self._process.stderr._buffer.clear()
|
||||||
self.process.stdin.flush()
|
|
||||||
|
|
||||||
# Read output until sentinel appears
|
# Return as ToolResult (our loop handles this type)
|
||||||
output_lines = []
|
if error and output:
|
||||||
start_time = asyncio.get_event_loop().time()
|
return ToolResult(output=f"{output}\n\nstderr: {error}")
|
||||||
|
elif error:
|
||||||
while True:
|
return ToolResult(output=error)
|
||||||
# Check timeout
|
else:
|
||||||
elapsed = asyncio.get_event_loop().time() - start_time
|
return ToolResult(output=output if output else "(no output)")
|
||||||
if elapsed > timeout:
|
|
||||||
raise asyncio.TimeoutError(
|
|
||||||
f"Command timed out after {timeout}s: {command[:50]}..."
|
|
||||||
)
|
|
||||||
|
|
||||||
# Read line (non-blocking via asyncio)
|
|
||||||
try:
|
|
||||||
line = await asyncio.wait_for(
|
|
||||||
asyncio.to_thread(self.process.stdout.readline),
|
|
||||||
timeout=1.0,
|
|
||||||
)
|
|
||||||
except asyncio.TimeoutError:
|
|
||||||
# No output yet, continue waiting
|
|
||||||
continue
|
|
||||||
|
|
||||||
if not line:
|
|
||||||
# EOF - process died
|
|
||||||
raise RuntimeError("Bash process terminated unexpectedly")
|
|
||||||
|
|
||||||
# Check for sentinel
|
|
||||||
if sentinel in line:
|
|
||||||
break
|
|
||||||
|
|
||||||
output_lines.append(line.rstrip("\n"))
|
|
||||||
|
|
||||||
return "\n".join(output_lines)
|
|
||||||
|
|
||||||
def __del__(self):
|
|
||||||
"""Clean up bash process on deletion."""
|
|
||||||
if self.process:
|
|
||||||
self.process.terminate()
|
|
||||||
try:
|
|
||||||
self.process.wait(timeout=2)
|
|
||||||
except subprocess.TimeoutExpired:
|
|
||||||
self.process.kill()
|
|
||||||
|
|
||||||
|
|
||||||
class BashTool20250124(BaseAnthropicTool):
|
class BashTool20250124(BaseAnthropicTool):
|
||||||
@@ -124,10 +127,10 @@ class BashTool20250124(BaseAnthropicTool):
|
|||||||
|
|
||||||
api_type: Literal["bash_20250124"] = "bash_20250124"
|
api_type: Literal["bash_20250124"] = "bash_20250124"
|
||||||
name: Literal["bash"] = "bash"
|
name: Literal["bash"] = "bash"
|
||||||
beta_flag: str = "computer-use-2025-11-24"
|
beta_flag: str | None = None
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self._session = _BashSession()
|
self._session: _BashSession | None = None
|
||||||
|
|
||||||
async def __call__(
|
async def __call__(
|
||||||
self,
|
self,
|
||||||
@@ -135,39 +138,26 @@ class BashTool20250124(BaseAnthropicTool):
|
|||||||
restart: bool = False,
|
restart: bool = False,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> ToolResult:
|
) -> ToolResult:
|
||||||
"""Execute bash command or restart session.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
command: Bash command to execute (optional)
|
|
||||||
restart: Restart the bash session (optional)
|
|
||||||
**kwargs: Additional arguments (ignored)
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
ToolResult with command output or error
|
|
||||||
"""
|
|
||||||
if restart:
|
if restart:
|
||||||
self._session.restart()
|
if self._session:
|
||||||
return ToolResult(output="Bash session restarted successfully.")
|
self._session.stop()
|
||||||
|
self._session = _BashSession()
|
||||||
|
await self._session.start()
|
||||||
|
return ToolResult(system="tool has been restarted.")
|
||||||
|
|
||||||
if not command:
|
if self._session is None:
|
||||||
return ToolResult(
|
self._session = _BashSession()
|
||||||
error="Either 'command' or 'restart=True' must be provided."
|
await self._session.start()
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
if command is not None:
|
||||||
output = await self._session.run_command(command)
|
try:
|
||||||
return ToolResult(output=output if output else "(no output)")
|
return await self._session.run(command)
|
||||||
except asyncio.TimeoutError as e:
|
except ToolError as e:
|
||||||
return ToolResult(error=f"Command timed out: {e}")
|
return ToolResult(error=str(e))
|
||||||
except Exception as e:
|
|
||||||
return ToolResult(error=f"{e}")
|
return ToolResult(error="Either 'command' or 'restart=True' must be provided.")
|
||||||
|
|
||||||
def to_params(self) -> dict[str, Any]:
|
def to_params(self) -> dict[str, Any]:
|
||||||
"""Convert to Anthropic API tool parameter format.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Tool definition for Anthropic API with bash_20250124 type
|
|
||||||
"""
|
|
||||||
return {
|
return {
|
||||||
"type": self.api_type,
|
"type": self.api_type,
|
||||||
"name": self.name,
|
"name": self.name,
|
||||||
|
|||||||
@@ -67,13 +67,14 @@ class ComputerTool20251124(BaseAnthropicTool):
|
|||||||
self.display_height_px = display_height_px
|
self.display_height_px = display_height_px
|
||||||
|
|
||||||
def to_params(self):
|
def to_params(self):
|
||||||
"""Return tool definition for API."""
|
"""Return tool definition for API.
|
||||||
|
|
||||||
|
NOTE: display_width_px, display_height_px, and enable_zoom are NOT
|
||||||
|
valid parameters for computer_20251124 and cause API hangs if sent.
|
||||||
|
"""
|
||||||
return {
|
return {
|
||||||
"type": self.api_type,
|
"type": self.api_type,
|
||||||
"name": self.name,
|
"name": self.name,
|
||||||
"display_width_px": self.display_width_px,
|
|
||||||
"display_height_px": self.display_height_px,
|
|
||||||
"enable_zoom": True,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async def __call__(
|
async def __call__(
|
||||||
|
|||||||
@@ -20,7 +20,7 @@ class EditTool20250728(BaseAnthropicTool):
|
|||||||
|
|
||||||
api_type: Literal["text_editor_20250728"] = "text_editor_20250728"
|
api_type: Literal["text_editor_20250728"] = "text_editor_20250728"
|
||||||
name: Literal["str_replace_based_edit_tool"] = "str_replace_based_edit_tool"
|
name: Literal["str_replace_based_edit_tool"] = "str_replace_based_edit_tool"
|
||||||
beta_flag: str = "computer-use-2025-11-24"
|
beta_flag: str | None = None
|
||||||
|
|
||||||
async def __call__(
|
async def __call__(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -0,0 +1,230 @@
|
|||||||
|
"""Mem0 memory tools — expose semantic memory to the agent."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from typing import Any, TYPE_CHECKING
|
||||||
|
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
from nanobot.agent.tools.base import Tool
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from nanobot.agent.memory_mem0 import Mem0MemoryStore
|
||||||
|
|
||||||
|
|
||||||
|
class Mem0ToolContext:
|
||||||
|
"""Shared mutable state injected into every mem0 tool."""
|
||||||
|
|
||||||
|
def __init__(self, store: Mem0MemoryStore, consolidate_fn):
|
||||||
|
self.store = store
|
||||||
|
self.consolidate_fn = consolidate_fn # async (session, archive_all) -> None
|
||||||
|
self.user_id: str = "unknown"
|
||||||
|
self.session = None
|
||||||
|
|
||||||
|
def set_context(self, channel: str, chat_id: str, session=None):
|
||||||
|
self.user_id = f"{channel}_{chat_id}"
|
||||||
|
self.session = session
|
||||||
|
|
||||||
|
|
||||||
|
class MemorySearchTool(Tool):
|
||||||
|
"""Search memories semantically."""
|
||||||
|
|
||||||
|
name = "memory_search"
|
||||||
|
description = (
|
||||||
|
"Search your long-term memory for facts relevant to a query. "
|
||||||
|
"Returns the most relevant memories ranked by similarity."
|
||||||
|
)
|
||||||
|
parameters = {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"query": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Natural-language search query",
|
||||||
|
},
|
||||||
|
"limit": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "Max results to return (default 5)",
|
||||||
|
"minimum": 1,
|
||||||
|
"maximum": 20,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["query"],
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(self, ctx: Mem0ToolContext):
|
||||||
|
self._ctx = ctx
|
||||||
|
|
||||||
|
async def execute(self, query: str, limit: int = 5, **kw: Any) -> str:
|
||||||
|
results = self._ctx.store.search_memories(
|
||||||
|
query=query,
|
||||||
|
user_id=self._ctx.user_id,
|
||||||
|
limit=limit,
|
||||||
|
)
|
||||||
|
if not results:
|
||||||
|
return "No memories found."
|
||||||
|
lines = []
|
||||||
|
for i, mem in enumerate(results, 1):
|
||||||
|
text = mem.get("memory", "")
|
||||||
|
score = mem.get("score")
|
||||||
|
mid = mem.get("id", "")
|
||||||
|
score_str = f" (score: {score:.2f})" if score else ""
|
||||||
|
lines.append(f"{i}. [{mid}] {text}{score_str}")
|
||||||
|
return "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryListTool(Tool):
|
||||||
|
"""List all memories for the current user."""
|
||||||
|
|
||||||
|
name = "memory_list"
|
||||||
|
description = (
|
||||||
|
"List ALL stored memories for the current user. "
|
||||||
|
"Use memory_search for targeted lookup; use this to browse everything."
|
||||||
|
)
|
||||||
|
parameters = {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {},
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(self, ctx: Mem0ToolContext):
|
||||||
|
self._ctx = ctx
|
||||||
|
|
||||||
|
async def execute(self, **kw: Any) -> str:
|
||||||
|
memories = self._ctx.store.get_all_memories(self._ctx.user_id)
|
||||||
|
if not memories:
|
||||||
|
return "No memories stored."
|
||||||
|
lines = []
|
||||||
|
for i, mem in enumerate(memories, 1):
|
||||||
|
text = mem.get("memory", "")
|
||||||
|
mid = mem.get("id", "")
|
||||||
|
lines.append(f"{i}. [{mid}] {text}")
|
||||||
|
return f"{len(memories)} memories:\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryAddTool(Tool):
|
||||||
|
"""Add a fact to long-term memory."""
|
||||||
|
|
||||||
|
name = "memory_add"
|
||||||
|
description = (
|
||||||
|
"Store a new fact or piece of information in long-term memory. "
|
||||||
|
"The content will be processed by the extraction LLM and stored as one or more facts."
|
||||||
|
)
|
||||||
|
parameters = {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"content": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "The fact or information to remember",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["content"],
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(self, ctx: Mem0ToolContext):
|
||||||
|
self._ctx = ctx
|
||||||
|
|
||||||
|
async def execute(self, content: str, **kw: Any) -> str:
|
||||||
|
try:
|
||||||
|
result = self._ctx.store.memory.add(
|
||||||
|
[{"role": "user", "content": content}],
|
||||||
|
user_id=self._ctx.user_id,
|
||||||
|
)
|
||||||
|
facts_count = len(result.get("results", [])) if result else 0
|
||||||
|
return f"Added to memory. {facts_count} fact(s) extracted."
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"memory_add failed: {e}")
|
||||||
|
return f"Error adding memory: {e}"
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryUpdateTool(Tool):
|
||||||
|
"""Update an existing memory by ID."""
|
||||||
|
|
||||||
|
name = "memory_update"
|
||||||
|
description = (
|
||||||
|
"Update the content of an existing memory. "
|
||||||
|
"Use memory_list or memory_search first to find the memory ID."
|
||||||
|
)
|
||||||
|
parameters = {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"memory_id": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "The memory ID to update",
|
||||||
|
},
|
||||||
|
"content": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "The new content for this memory",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["memory_id", "content"],
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(self, ctx: Mem0ToolContext):
|
||||||
|
self._ctx = ctx
|
||||||
|
|
||||||
|
async def execute(self, memory_id: str, content: str, **kw: Any) -> str:
|
||||||
|
try:
|
||||||
|
self._ctx.store.update_memory(memory_id, content)
|
||||||
|
return f"Memory {memory_id} updated."
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"memory_update failed: {e}")
|
||||||
|
return f"Error updating memory: {e}"
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryDeleteTool(Tool):
|
||||||
|
"""Delete a memory by ID."""
|
||||||
|
|
||||||
|
name = "memory_delete"
|
||||||
|
description = (
|
||||||
|
"Delete a specific memory by its ID. "
|
||||||
|
"Use memory_list or memory_search first to find the memory ID."
|
||||||
|
)
|
||||||
|
parameters = {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"memory_id": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "The memory ID to delete",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["memory_id"],
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(self, ctx: Mem0ToolContext):
|
||||||
|
self._ctx = ctx
|
||||||
|
|
||||||
|
async def execute(self, memory_id: str, **kw: Any) -> str:
|
||||||
|
try:
|
||||||
|
self._ctx.store.delete_memory(memory_id)
|
||||||
|
return f"Memory {memory_id} deleted."
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"memory_delete failed: {e}")
|
||||||
|
return f"Error deleting memory: {e}"
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryConsolidateTool(Tool):
|
||||||
|
"""Trigger memory consolidation for the current session."""
|
||||||
|
|
||||||
|
name = "memory_consolidate"
|
||||||
|
description = (
|
||||||
|
"Extract and store facts from the current conversation into long-term memory. "
|
||||||
|
"Normally this happens automatically on /new, but you can trigger it manually."
|
||||||
|
)
|
||||||
|
parameters = {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {},
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(self, ctx: Mem0ToolContext):
|
||||||
|
self._ctx = ctx
|
||||||
|
|
||||||
|
async def execute(self, **kw: Any) -> str:
|
||||||
|
session = self._ctx.session
|
||||||
|
if not session:
|
||||||
|
return "Error: no active session."
|
||||||
|
try:
|
||||||
|
await self._ctx.consolidate_fn(session, archive_all=False)
|
||||||
|
return "Memory consolidation complete."
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"memory_consolidate failed: {e}")
|
||||||
|
return f"Error during consolidation: {e}"
|
||||||
@@ -21,6 +21,7 @@ class MessageTool(Tool):
|
|||||||
self._sessions = sessions
|
self._sessions = sessions
|
||||||
self._default_channel = default_channel
|
self._default_channel = default_channel
|
||||||
self._default_chat_id = default_chat_id
|
self._default_chat_id = default_chat_id
|
||||||
|
self._sent_in_turn: bool = False
|
||||||
|
|
||||||
def set_context(self, channel: str, chat_id: str) -> None:
|
def set_context(self, channel: str, chat_id: str) -> None:
|
||||||
"""Set the current message context."""
|
"""Set the current message context."""
|
||||||
@@ -30,6 +31,10 @@ class MessageTool(Tool):
|
|||||||
def set_send_callback(self, callback: Callable[[OutboundMessage], Awaitable[None]]) -> None:
|
def set_send_callback(self, callback: Callable[[OutboundMessage], Awaitable[None]]) -> None:
|
||||||
"""Set the callback for sending messages."""
|
"""Set the callback for sending messages."""
|
||||||
self._send_callback = callback
|
self._send_callback = callback
|
||||||
|
|
||||||
|
def start_turn(self) -> None:
|
||||||
|
"""Reset per-turn send tracking."""
|
||||||
|
self._sent_in_turn = False
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def name(self) -> str:
|
def name(self) -> str:
|
||||||
@@ -92,6 +97,10 @@ class MessageTool(Tool):
|
|||||||
try:
|
try:
|
||||||
await self._send_callback(msg)
|
await self._send_callback(msg)
|
||||||
|
|
||||||
|
# Track if sent to same target as current context
|
||||||
|
if channel == self._default_channel and chat_id == self._default_chat_id:
|
||||||
|
self._sent_in_turn = True
|
||||||
|
|
||||||
if self._sessions:
|
if self._sessions:
|
||||||
session_key = f"{channel}:{chat_id}"
|
session_key = f"{channel}:{chat_id}"
|
||||||
session = self._sessions.get_or_create(session_key)
|
session = self._sessions.get_or_create(session_key)
|
||||||
|
|||||||
@@ -87,6 +87,10 @@ class HeartbeatService:
|
|||||||
logger.info("Heartbeat disabled")
|
logger.info("Heartbeat disabled")
|
||||||
return
|
return
|
||||||
|
|
||||||
|
# Idempotent: don't create a new task if already running
|
||||||
|
if self._task is not None and not self._task.done():
|
||||||
|
return
|
||||||
|
|
||||||
self._running = True
|
self._running = True
|
||||||
self._task = asyncio.create_task(self._run_loop())
|
self._task = asyncio.create_task(self._run_loop())
|
||||||
logger.info(f"Heartbeat started (every {self.interval_s}s)")
|
logger.info(f"Heartbeat started (every {self.interval_s}s)")
|
||||||
|
|||||||
@@ -59,9 +59,83 @@ class AnthropicOAuthProvider(LLMProvider):
|
|||||||
async def _get_client(self) -> httpx.AsyncClient:
|
async def _get_client(self) -> httpx.AsyncClient:
|
||||||
"""Get or create async HTTP client."""
|
"""Get or create async HTTP client."""
|
||||||
if self._client is None:
|
if self._client is None:
|
||||||
self._client = httpx.AsyncClient(timeout=300.0)
|
self._client = httpx.AsyncClient(
|
||||||
|
timeout=httpx.Timeout(300.0, pool=30.0),
|
||||||
|
)
|
||||||
return self._client
|
return self._client
|
||||||
|
|
||||||
|
async def _reset_client(self) -> None:
|
||||||
|
"""Destroy and recreate the HTTP client after connection errors."""
|
||||||
|
old = self._client
|
||||||
|
self._client = None
|
||||||
|
if old:
|
||||||
|
try:
|
||||||
|
await old.aclose()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
logger.warning("Reset httpx client (pool recycled)")
|
||||||
|
|
||||||
|
async def _diagnose_connectivity(self) -> None:
|
||||||
|
"""Run diagnostics when ConnectTimeout occurs to understand why."""
|
||||||
|
import socket
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
# 1. Raw socket test (bypasses httpx entirely)
|
||||||
|
try:
|
||||||
|
t0 = __import__('time').monotonic()
|
||||||
|
s = socket.create_connection(('api.anthropic.com', 443), timeout=10)
|
||||||
|
elapsed = __import__('time').monotonic() - t0
|
||||||
|
s.close()
|
||||||
|
logger.warning(f"DIAG: raw socket connect OK in {elapsed:.3f}s")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"DIAG: raw socket connect FAILED: {e}")
|
||||||
|
|
||||||
|
# 2. asyncio connect test (same event loop)
|
||||||
|
try:
|
||||||
|
t0 = __import__('time').monotonic()
|
||||||
|
reader, writer = await asyncio.wait_for(
|
||||||
|
asyncio.open_connection('api.anthropic.com', 443),
|
||||||
|
timeout=10.0,
|
||||||
|
)
|
||||||
|
elapsed = __import__('time').monotonic() - t0
|
||||||
|
writer.close()
|
||||||
|
await writer.wait_closed()
|
||||||
|
logger.warning(f"DIAG: asyncio connect OK in {elapsed:.3f}s")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"DIAG: asyncio connect FAILED: {e}")
|
||||||
|
|
||||||
|
# 3. Fresh httpx client test (new pool)
|
||||||
|
try:
|
||||||
|
t0 = __import__('time').monotonic()
|
||||||
|
async with httpx.AsyncClient(timeout=10.0) as fresh:
|
||||||
|
r = await fresh.get('https://api.anthropic.com/')
|
||||||
|
elapsed = __import__('time').monotonic() - t0
|
||||||
|
logger.warning(f"DIAG: fresh httpx OK in {elapsed:.3f}s (status={r.status_code})")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"DIAG: fresh httpx FAILED: {e}")
|
||||||
|
|
||||||
|
# 4. DNS resolution
|
||||||
|
try:
|
||||||
|
ips = socket.getaddrinfo('api.anthropic.com', 443)
|
||||||
|
logger.warning(f"DIAG: DNS resolved to {len(ips)} entries, first={ips[0][4][0]}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"DIAG: DNS FAILED: {e}")
|
||||||
|
|
||||||
|
# 5. Connection pool state of the broken client
|
||||||
|
if self._client:
|
||||||
|
transport = self._client._transport
|
||||||
|
if hasattr(transport, '_pool'):
|
||||||
|
pool = transport._pool
|
||||||
|
conns = getattr(pool, '_connections', [])
|
||||||
|
reqs = getattr(pool, '_requests', [])
|
||||||
|
logger.warning(
|
||||||
|
f"DIAG: pool state: {len(conns)} connections, "
|
||||||
|
f"{len(reqs)} pending requests"
|
||||||
|
)
|
||||||
|
for i, conn in enumerate(conns[:5]):
|
||||||
|
state = getattr(conn, '_state', 'unknown')
|
||||||
|
logger.warning(f"DIAG: conn[{i}] state={state}")
|
||||||
|
|
||||||
def _prepare_messages(
|
def _prepare_messages(
|
||||||
self,
|
self,
|
||||||
messages: list[dict[str, Any]]
|
messages: list[dict[str, Any]]
|
||||||
@@ -252,18 +326,35 @@ class AnthropicOAuthProvider(LLMProvider):
|
|||||||
"""Make request to Anthropic API."""
|
"""Make request to Anthropic API."""
|
||||||
client = await self._get_client()
|
client = await self._get_client()
|
||||||
|
|
||||||
# Cache the last user message so conversation history is cached across turns
|
# Add cache breakpoints on the last TWO user messages (4-breakpoint strategy):
|
||||||
if messages:
|
# BP3: Second-to-last user message (stable history from previous turn)
|
||||||
last = messages[-1]
|
# BP4: Last user message (current turn, will become BP3 next turn)
|
||||||
if last.get("role") == "user":
|
# This allows BP3 to reuse what BP4 cached last turn.
|
||||||
content = last["content"]
|
user_indices = [i for i, m in enumerate(messages) if m.get("role") == "user"]
|
||||||
if isinstance(content, str):
|
|
||||||
last = {**last, "content": [{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}]}
|
if len(user_indices) >= 2:
|
||||||
elif isinstance(content, list) and content:
|
# BP3: Second-to-last user message
|
||||||
new_content = list(content)
|
idx = user_indices[-2]
|
||||||
new_content[-1] = {**new_content[-1], "cache_control": {"type": "ephemeral"}}
|
msg = messages[idx]
|
||||||
last = {**last, "content": new_content}
|
content = msg["content"]
|
||||||
messages = messages[:-1] + [last]
|
if isinstance(content, str):
|
||||||
|
messages[idx] = {**msg, "content": [{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}]}
|
||||||
|
elif isinstance(content, list) and content:
|
||||||
|
new_content = list(content)
|
||||||
|
new_content[-1] = {**new_content[-1], "cache_control": {"type": "ephemeral"}}
|
||||||
|
messages[idx] = {**msg, "content": new_content}
|
||||||
|
|
||||||
|
if len(user_indices) >= 1:
|
||||||
|
# BP4: Last user message
|
||||||
|
idx = user_indices[-1]
|
||||||
|
msg = messages[idx]
|
||||||
|
content = msg["content"]
|
||||||
|
if isinstance(content, str):
|
||||||
|
messages[idx] = {**msg, "content": [{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}]}
|
||||||
|
elif isinstance(content, list) and content:
|
||||||
|
new_content = list(content)
|
||||||
|
new_content[-1] = {**new_content[-1], "cache_control": {"type": "ephemeral"}}
|
||||||
|
messages[idx] = {**msg, "content": new_content}
|
||||||
|
|
||||||
payload: dict[str, Any] = {
|
payload: dict[str, Any] = {
|
||||||
"model": model,
|
"model": model,
|
||||||
@@ -321,11 +412,43 @@ class AnthropicOAuthProvider(LLMProvider):
|
|||||||
tool_names = [t.get("name", "unnamed") for t in payload["tools"]]
|
tool_names = [t.get("name", "unnamed") for t in payload["tools"]]
|
||||||
logger.debug(f"Tool names in request: {tool_names}")
|
logger.debug(f"Tool names in request: {tool_names}")
|
||||||
|
|
||||||
response = await client.post(
|
# Debug: Log message structure to diagnose orphaned tool_result errors
|
||||||
self._get_api_url(),
|
for idx, m in enumerate(payload.get("messages", [])):
|
||||||
headers=headers,
|
role = m.get("role", "?")
|
||||||
json=payload,
|
content = m.get("content", "")
|
||||||
)
|
if isinstance(content, list):
|
||||||
|
block_types = [b.get("type", "?") for b in content]
|
||||||
|
logger.debug(f" msg[{idx}] role={role} blocks={block_types}")
|
||||||
|
else:
|
||||||
|
logger.debug(f" msg[{idx}] role={role} text={str(content)[:80]}")
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import time as _time
|
||||||
|
_t0 = _time.monotonic()
|
||||||
|
try:
|
||||||
|
response = await client.post(
|
||||||
|
self._get_api_url(),
|
||||||
|
headers=headers,
|
||||||
|
json=payload,
|
||||||
|
)
|
||||||
|
except httpx.ConnectTimeout:
|
||||||
|
elapsed = _time.monotonic() - _t0
|
||||||
|
logger.error(f"ConnectTimeout after {elapsed:.1f}s — running diagnostics")
|
||||||
|
await self._diagnose_connectivity()
|
||||||
|
await self._reset_client()
|
||||||
|
raise
|
||||||
|
except httpx.PoolTimeout:
|
||||||
|
elapsed = _time.monotonic() - _t0
|
||||||
|
logger.error(f"PoolTimeout after {elapsed:.1f}s — resetting client")
|
||||||
|
await self._reset_client()
|
||||||
|
raise
|
||||||
|
except (httpx.ConnectError, httpx.TimeoutException) as e:
|
||||||
|
elapsed = _time.monotonic() - _t0
|
||||||
|
logger.error(f"{type(e).__name__} after {elapsed:.1f}s")
|
||||||
|
raise
|
||||||
|
elapsed = _time.monotonic() - _t0
|
||||||
|
if elapsed > 30:
|
||||||
|
logger.warning(f"Anthropic API slow response: {elapsed:.1f}s")
|
||||||
|
|
||||||
# Dump rate limit headers for analysis
|
# Dump rate limit headers for analysis
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -19,9 +19,7 @@ class Session:
|
|||||||
|
|
||||||
Stores messages in JSONL format for easy reading and persistence.
|
Stores messages in JSONL format for easy reading and persistence.
|
||||||
|
|
||||||
Important: Messages are append-only for LLM cache efficiency.
|
Messages are trimmed after consolidation to keep session size manageable.
|
||||||
The consolidation process writes summaries to MEMORY.md/HISTORY.md
|
|
||||||
but does NOT modify the messages list or get_history() output.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
key: str # channel:chat_id
|
key: str # channel:chat_id
|
||||||
@@ -29,7 +27,6 @@ class Session:
|
|||||||
created_at: datetime = field(default_factory=datetime.now)
|
created_at: datetime = field(default_factory=datetime.now)
|
||||||
updated_at: datetime = field(default_factory=datetime.now)
|
updated_at: datetime = field(default_factory=datetime.now)
|
||||||
metadata: dict[str, Any] = field(default_factory=dict)
|
metadata: dict[str, Any] = field(default_factory=dict)
|
||||||
last_consolidated: int = 0 # Number of messages already consolidated to files
|
|
||||||
|
|
||||||
def add_message(self, role: str, content: str, **kwargs: Any) -> None:
|
def add_message(self, role: str, content: str, **kwargs: Any) -> None:
|
||||||
"""Add a message to the session."""
|
"""Add a message to the session."""
|
||||||
@@ -73,7 +70,6 @@ class Session:
|
|||||||
def clear(self) -> None:
|
def clear(self) -> None:
|
||||||
"""Clear all messages and reset session to initial state."""
|
"""Clear all messages and reset session to initial state."""
|
||||||
self.messages = []
|
self.messages = []
|
||||||
self.last_consolidated = 0
|
|
||||||
self.updated_at = datetime.now()
|
self.updated_at = datetime.now()
|
||||||
|
|
||||||
|
|
||||||
@@ -139,7 +135,6 @@ class SessionManager:
|
|||||||
messages = []
|
messages = []
|
||||||
metadata = {}
|
metadata = {}
|
||||||
created_at = None
|
created_at = None
|
||||||
last_consolidated = 0
|
|
||||||
|
|
||||||
with open(path, encoding="utf-8") as f:
|
with open(path, encoding="utf-8") as f:
|
||||||
for line in f:
|
for line in f:
|
||||||
@@ -152,7 +147,7 @@ class SessionManager:
|
|||||||
if data.get("_type") == "metadata":
|
if data.get("_type") == "metadata":
|
||||||
metadata = data.get("metadata", {})
|
metadata = data.get("metadata", {})
|
||||||
created_at = datetime.fromisoformat(data["created_at"]) if data.get("created_at") else None
|
created_at = datetime.fromisoformat(data["created_at"]) if data.get("created_at") else None
|
||||||
last_consolidated = data.get("last_consolidated", 0)
|
# Ignore legacy last_consolidated field
|
||||||
else:
|
else:
|
||||||
messages.append(data)
|
messages.append(data)
|
||||||
|
|
||||||
@@ -161,7 +156,6 @@ class SessionManager:
|
|||||||
messages=messages,
|
messages=messages,
|
||||||
created_at=created_at or datetime.now(),
|
created_at=created_at or datetime.now(),
|
||||||
metadata=metadata,
|
metadata=metadata,
|
||||||
last_consolidated=last_consolidated
|
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning("Failed to load session {}: {}", key, e)
|
logger.warning("Failed to load session {}: {}", key, e)
|
||||||
@@ -178,7 +172,6 @@ class SessionManager:
|
|||||||
"created_at": session.created_at.isoformat(),
|
"created_at": session.created_at.isoformat(),
|
||||||
"updated_at": session.updated_at.isoformat(),
|
"updated_at": session.updated_at.isoformat(),
|
||||||
"metadata": session.metadata,
|
"metadata": session.metadata,
|
||||||
"last_consolidated": session.last_consolidated
|
|
||||||
}
|
}
|
||||||
f.write(json.dumps(metadata_line, ensure_ascii=False) + "\n")
|
f.write(json.dumps(metadata_line, ensure_ascii=False) + "\n")
|
||||||
for msg in session.messages:
|
for msg in session.messages:
|
||||||
|
|||||||
@@ -50,6 +50,11 @@ dev = [
|
|||||||
mem0 = [
|
mem0 = [
|
||||||
"mem0ai>=0.1.0",
|
"mem0ai>=0.1.0",
|
||||||
]
|
]
|
||||||
|
matrix = [
|
||||||
|
"matrix-nio>=0.20.0",
|
||||||
|
"mistune>=3.0.0",
|
||||||
|
"nh3>=0.2.0",
|
||||||
|
]
|
||||||
|
|
||||||
[project.scripts]
|
[project.scripts]
|
||||||
nanobot = "nanobot.cli.commands:app"
|
nanobot = "nanobot.cli.commands:app"
|
||||||
|
|||||||
@@ -50,11 +50,12 @@ async def test_beta_flags_collected_from_tools():
|
|||||||
tools=tools_with_flags
|
tools=tools_with_flags
|
||||||
)
|
)
|
||||||
|
|
||||||
# Check that beta flag was added to headers
|
# Check that beta flag was added to headers (merged with hardcoded flags)
|
||||||
call_args = mock_client.post.call_args
|
call_args = mock_client.post.call_args
|
||||||
headers = call_args[1]["headers"]
|
headers = call_args[1]["headers"]
|
||||||
assert "anthropic-beta" in headers
|
assert "anthropic-beta" in headers
|
||||||
assert headers["anthropic-beta"] == "computer-use-2025-11-24"
|
# Should include hardcoded flags + tool flag, sorted alphabetically
|
||||||
|
assert headers["anthropic-beta"] == "claude-code-20250219,computer-use-2025-11-24,context-management-2025-06-27,oauth-2025-04-20"
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -99,5 +100,5 @@ async def test_multiple_beta_flags_joined():
|
|||||||
call_args = mock_client.post.call_args
|
call_args = mock_client.post.call_args
|
||||||
headers = call_args[1]["headers"]
|
headers = call_args[1]["headers"]
|
||||||
assert "anthropic-beta" in headers
|
assert "anthropic-beta" in headers
|
||||||
# Should be sorted alphabetically and joined with comma
|
# Should include hardcoded flags + tool flags, sorted alphabetically and joined with comma
|
||||||
assert headers["anthropic-beta"] == "flag-a,flag-b"
|
assert headers["anthropic-beta"] == "claude-code-20250219,context-management-2025-06-27,flag-a,flag-b,oauth-2025-04-20"
|
||||||
|
|||||||
@@ -1,828 +0,0 @@
|
|||||||
"""Test session management with cache-friendly message handling."""
|
|
||||||
|
|
||||||
import asyncio
|
|
||||||
from unittest.mock import AsyncMock, MagicMock
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
from pathlib import Path
|
|
||||||
from nanobot.session.manager import Session, SessionManager
|
|
||||||
|
|
||||||
# Test constants
|
|
||||||
MEMORY_WINDOW = 50
|
|
||||||
KEEP_COUNT = MEMORY_WINDOW // 2 # 25
|
|
||||||
|
|
||||||
|
|
||||||
def create_session_with_messages(key: str, count: int, role: str = "user") -> Session:
|
|
||||||
"""Create a session and add the specified number of messages.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
key: Session identifier
|
|
||||||
count: Number of messages to add
|
|
||||||
role: Message role (default: "user")
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Session with the specified messages
|
|
||||||
"""
|
|
||||||
session = Session(key=key)
|
|
||||||
for i in range(count):
|
|
||||||
session.add_message(role, f"msg{i}")
|
|
||||||
return session
|
|
||||||
|
|
||||||
|
|
||||||
def assert_messages_content(messages: list, start_index: int, end_index: int) -> None:
|
|
||||||
"""Assert that messages contain expected content from start to end index.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
messages: List of message dictionaries
|
|
||||||
start_index: Expected first message index
|
|
||||||
end_index: Expected last message index
|
|
||||||
"""
|
|
||||||
assert len(messages) > 0
|
|
||||||
assert messages[0]["content"] == f"msg{start_index}"
|
|
||||||
assert messages[-1]["content"] == f"msg{end_index}"
|
|
||||||
|
|
||||||
|
|
||||||
def get_old_messages(session: Session, last_consolidated: int, keep_count: int) -> list:
|
|
||||||
"""Extract messages that would be consolidated using the standard slice logic.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
session: The session containing messages
|
|
||||||
last_consolidated: Index of last consolidated message
|
|
||||||
keep_count: Number of recent messages to keep
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
List of messages that would be consolidated
|
|
||||||
"""
|
|
||||||
return session.messages[last_consolidated:-keep_count]
|
|
||||||
|
|
||||||
|
|
||||||
class TestSessionLastConsolidated:
|
|
||||||
"""Test last_consolidated tracking to avoid duplicate processing."""
|
|
||||||
|
|
||||||
def test_initial_last_consolidated_zero(self) -> None:
|
|
||||||
"""Test that new session starts with last_consolidated=0."""
|
|
||||||
session = Session(key="test:initial")
|
|
||||||
assert session.last_consolidated == 0
|
|
||||||
|
|
||||||
def test_last_consolidated_persistence(self, tmp_path) -> None:
|
|
||||||
"""Test that last_consolidated persists across save/load."""
|
|
||||||
manager = SessionManager(Path(tmp_path))
|
|
||||||
session1 = create_session_with_messages("test:persist", 20)
|
|
||||||
session1.last_consolidated = 15
|
|
||||||
manager.save(session1)
|
|
||||||
|
|
||||||
session2 = manager.get_or_create("test:persist")
|
|
||||||
assert session2.last_consolidated == 15
|
|
||||||
assert len(session2.messages) == 20
|
|
||||||
|
|
||||||
def test_clear_resets_last_consolidated(self) -> None:
|
|
||||||
"""Test that clear() resets last_consolidated to 0."""
|
|
||||||
session = create_session_with_messages("test:clear", 10)
|
|
||||||
session.last_consolidated = 5
|
|
||||||
|
|
||||||
session.clear()
|
|
||||||
assert len(session.messages) == 0
|
|
||||||
assert session.last_consolidated == 0
|
|
||||||
|
|
||||||
|
|
||||||
class TestSessionImmutableHistory:
|
|
||||||
"""Test Session message immutability for cache efficiency."""
|
|
||||||
|
|
||||||
def test_initial_state(self) -> None:
|
|
||||||
"""Test that new session has empty messages list."""
|
|
||||||
session = Session(key="test:initial")
|
|
||||||
assert len(session.messages) == 0
|
|
||||||
|
|
||||||
def test_add_messages_appends_only(self) -> None:
|
|
||||||
"""Test that adding messages only appends, never modifies."""
|
|
||||||
session = Session(key="test:preserve")
|
|
||||||
session.add_message("user", "msg1")
|
|
||||||
session.add_message("assistant", "resp1")
|
|
||||||
session.add_message("user", "msg2")
|
|
||||||
assert len(session.messages) == 3
|
|
||||||
assert session.messages[0]["content"] == "msg1"
|
|
||||||
|
|
||||||
def test_get_history_returns_most_recent(self) -> None:
|
|
||||||
"""Test get_history returns the most recent messages."""
|
|
||||||
session = Session(key="test:history")
|
|
||||||
for i in range(10):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
session.add_message("assistant", f"resp{i}")
|
|
||||||
|
|
||||||
history = session.get_history(max_messages=6)
|
|
||||||
assert len(history) == 6
|
|
||||||
assert history[0]["content"] == "msg7"
|
|
||||||
assert history[-1]["content"] == "resp9"
|
|
||||||
|
|
||||||
def test_get_history_with_all_messages(self) -> None:
|
|
||||||
"""Test get_history with max_messages larger than actual."""
|
|
||||||
session = create_session_with_messages("test:all", 5)
|
|
||||||
history = session.get_history(max_messages=100)
|
|
||||||
assert len(history) == 5
|
|
||||||
assert history[0]["content"] == "msg0"
|
|
||||||
|
|
||||||
def test_get_history_stable_for_same_session(self) -> None:
|
|
||||||
"""Test that get_history returns same content for same max_messages."""
|
|
||||||
session = create_session_with_messages("test:stable", 20)
|
|
||||||
history1 = session.get_history(max_messages=10)
|
|
||||||
history2 = session.get_history(max_messages=10)
|
|
||||||
assert history1 == history2
|
|
||||||
|
|
||||||
def test_messages_list_never_modified(self) -> None:
|
|
||||||
"""Test that messages list is never modified after creation."""
|
|
||||||
session = create_session_with_messages("test:immutable", 5)
|
|
||||||
original_len = len(session.messages)
|
|
||||||
|
|
||||||
session.get_history(max_messages=2)
|
|
||||||
assert len(session.messages) == original_len
|
|
||||||
|
|
||||||
for _ in range(10):
|
|
||||||
session.get_history(max_messages=3)
|
|
||||||
assert len(session.messages) == original_len
|
|
||||||
|
|
||||||
|
|
||||||
class TestSessionPersistence:
|
|
||||||
"""Test Session persistence and reload."""
|
|
||||||
|
|
||||||
@pytest.fixture
|
|
||||||
def temp_manager(self, tmp_path):
|
|
||||||
return SessionManager(Path(tmp_path))
|
|
||||||
|
|
||||||
def test_persistence_roundtrip(self, temp_manager):
|
|
||||||
"""Test that messages persist across save/load."""
|
|
||||||
session1 = create_session_with_messages("test:persistence", 20)
|
|
||||||
temp_manager.save(session1)
|
|
||||||
|
|
||||||
session2 = temp_manager.get_or_create("test:persistence")
|
|
||||||
assert len(session2.messages) == 20
|
|
||||||
assert session2.messages[0]["content"] == "msg0"
|
|
||||||
assert session2.messages[-1]["content"] == "msg19"
|
|
||||||
|
|
||||||
def test_get_history_after_reload(self, temp_manager):
|
|
||||||
"""Test that get_history works correctly after reload."""
|
|
||||||
session1 = create_session_with_messages("test:reload", 30)
|
|
||||||
temp_manager.save(session1)
|
|
||||||
|
|
||||||
session2 = temp_manager.get_or_create("test:reload")
|
|
||||||
history = session2.get_history(max_messages=10)
|
|
||||||
assert len(history) == 10
|
|
||||||
assert history[0]["content"] == "msg20"
|
|
||||||
assert history[-1]["content"] == "msg29"
|
|
||||||
|
|
||||||
def test_clear_resets_session(self, temp_manager):
|
|
||||||
"""Test that clear() properly resets session."""
|
|
||||||
session = create_session_with_messages("test:clear", 10)
|
|
||||||
assert len(session.messages) == 10
|
|
||||||
|
|
||||||
session.clear()
|
|
||||||
assert len(session.messages) == 0
|
|
||||||
|
|
||||||
|
|
||||||
class TestConsolidationTriggerConditions:
|
|
||||||
"""Test consolidation trigger conditions and logic."""
|
|
||||||
|
|
||||||
def test_consolidation_needed_when_messages_exceed_window(self):
|
|
||||||
"""Test consolidation logic: should trigger when messages > memory_window."""
|
|
||||||
session = create_session_with_messages("test:trigger", 60)
|
|
||||||
|
|
||||||
total_messages = len(session.messages)
|
|
||||||
messages_to_process = total_messages - session.last_consolidated
|
|
||||||
|
|
||||||
assert total_messages > MEMORY_WINDOW
|
|
||||||
assert messages_to_process > 0
|
|
||||||
|
|
||||||
expected_consolidate_count = total_messages - KEEP_COUNT
|
|
||||||
assert expected_consolidate_count == 35
|
|
||||||
|
|
||||||
def test_consolidation_skipped_when_within_keep_count(self):
|
|
||||||
"""Test consolidation skipped when total messages <= keep_count."""
|
|
||||||
session = create_session_with_messages("test:skip", 20)
|
|
||||||
|
|
||||||
total_messages = len(session.messages)
|
|
||||||
assert total_messages <= KEEP_COUNT
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
|
|
||||||
assert len(old_messages) == 0
|
|
||||||
|
|
||||||
def test_consolidation_skipped_when_no_new_messages(self):
|
|
||||||
"""Test consolidation skipped when messages_to_process <= 0."""
|
|
||||||
session = create_session_with_messages("test:already_consolidated", 40)
|
|
||||||
session.last_consolidated = len(session.messages) - KEEP_COUNT # 15
|
|
||||||
|
|
||||||
# Add a few more messages
|
|
||||||
for i in range(40, 42):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
|
|
||||||
total_messages = len(session.messages)
|
|
||||||
messages_to_process = total_messages - session.last_consolidated
|
|
||||||
assert messages_to_process > 0
|
|
||||||
|
|
||||||
# Simulate last_consolidated catching up
|
|
||||||
session.last_consolidated = total_messages - KEEP_COUNT
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
|
|
||||||
assert len(old_messages) == 0
|
|
||||||
|
|
||||||
|
|
||||||
class TestLastConsolidatedEdgeCases:
|
|
||||||
"""Test last_consolidated edge cases and data corruption scenarios."""
|
|
||||||
|
|
||||||
def test_last_consolidated_exceeds_message_count(self):
|
|
||||||
"""Test behavior when last_consolidated > len(messages) (data corruption)."""
|
|
||||||
session = create_session_with_messages("test:corruption", 10)
|
|
||||||
session.last_consolidated = 20
|
|
||||||
|
|
||||||
total_messages = len(session.messages)
|
|
||||||
messages_to_process = total_messages - session.last_consolidated
|
|
||||||
assert messages_to_process <= 0
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, 5)
|
|
||||||
assert len(old_messages) == 0
|
|
||||||
|
|
||||||
def test_last_consolidated_negative_value(self):
|
|
||||||
"""Test behavior with negative last_consolidated (invalid state)."""
|
|
||||||
session = create_session_with_messages("test:negative", 10)
|
|
||||||
session.last_consolidated = -5
|
|
||||||
|
|
||||||
keep_count = 3
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, keep_count)
|
|
||||||
|
|
||||||
# messages[-5:-3] with 10 messages gives indices 5,6
|
|
||||||
assert len(old_messages) == 2
|
|
||||||
assert old_messages[0]["content"] == "msg5"
|
|
||||||
assert old_messages[-1]["content"] == "msg6"
|
|
||||||
|
|
||||||
def test_messages_added_after_consolidation(self):
|
|
||||||
"""Test correct behavior when new messages arrive after consolidation."""
|
|
||||||
session = create_session_with_messages("test:new_messages", 40)
|
|
||||||
session.last_consolidated = len(session.messages) - KEEP_COUNT # 15
|
|
||||||
|
|
||||||
# Add new messages after consolidation
|
|
||||||
for i in range(40, 50):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
|
|
||||||
total_messages = len(session.messages)
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
|
|
||||||
expected_consolidate_count = total_messages - KEEP_COUNT - session.last_consolidated
|
|
||||||
|
|
||||||
assert len(old_messages) == expected_consolidate_count
|
|
||||||
assert_messages_content(old_messages, 15, 24)
|
|
||||||
|
|
||||||
def test_slice_behavior_when_indices_overlap(self):
|
|
||||||
"""Test slice behavior when last_consolidated >= total - keep_count."""
|
|
||||||
session = create_session_with_messages("test:overlap", 30)
|
|
||||||
session.last_consolidated = 12
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, 20)
|
|
||||||
assert len(old_messages) == 0
|
|
||||||
|
|
||||||
|
|
||||||
class TestArchiveAllMode:
|
|
||||||
"""Test archive_all mode (used by /new command)."""
|
|
||||||
|
|
||||||
def test_archive_all_consolidates_everything(self):
|
|
||||||
"""Test archive_all=True consolidates all messages."""
|
|
||||||
session = create_session_with_messages("test:archive_all", 50)
|
|
||||||
|
|
||||||
archive_all = True
|
|
||||||
if archive_all:
|
|
||||||
old_messages = session.messages
|
|
||||||
assert len(old_messages) == 50
|
|
||||||
|
|
||||||
assert session.last_consolidated == 0
|
|
||||||
|
|
||||||
def test_archive_all_resets_last_consolidated(self):
|
|
||||||
"""Test that archive_all mode resets last_consolidated to 0."""
|
|
||||||
session = create_session_with_messages("test:reset", 40)
|
|
||||||
session.last_consolidated = 15
|
|
||||||
|
|
||||||
archive_all = True
|
|
||||||
if archive_all:
|
|
||||||
session.last_consolidated = 0
|
|
||||||
|
|
||||||
assert session.last_consolidated == 0
|
|
||||||
assert len(session.messages) == 40
|
|
||||||
|
|
||||||
def test_archive_all_vs_normal_consolidation(self):
|
|
||||||
"""Test difference between archive_all and normal consolidation."""
|
|
||||||
# Normal consolidation
|
|
||||||
session1 = create_session_with_messages("test:normal", 60)
|
|
||||||
session1.last_consolidated = len(session1.messages) - KEEP_COUNT
|
|
||||||
|
|
||||||
# archive_all mode
|
|
||||||
session2 = create_session_with_messages("test:all", 60)
|
|
||||||
session2.last_consolidated = 0
|
|
||||||
|
|
||||||
assert session1.last_consolidated == 35
|
|
||||||
assert len(session1.messages) == 60
|
|
||||||
assert session2.last_consolidated == 0
|
|
||||||
assert len(session2.messages) == 60
|
|
||||||
|
|
||||||
|
|
||||||
class TestCacheImmutability:
|
|
||||||
"""Test that consolidation doesn't modify session.messages (cache safety)."""
|
|
||||||
|
|
||||||
def test_consolidation_does_not_modify_messages_list(self):
|
|
||||||
"""Test that consolidation leaves messages list unchanged."""
|
|
||||||
session = create_session_with_messages("test:immutable", 50)
|
|
||||||
|
|
||||||
original_messages = session.messages.copy()
|
|
||||||
original_len = len(session.messages)
|
|
||||||
session.last_consolidated = original_len - KEEP_COUNT
|
|
||||||
|
|
||||||
assert len(session.messages) == original_len
|
|
||||||
assert session.messages == original_messages
|
|
||||||
|
|
||||||
def test_get_history_does_not_modify_messages(self):
|
|
||||||
"""Test that get_history doesn't modify messages list."""
|
|
||||||
session = create_session_with_messages("test:history_immutable", 40)
|
|
||||||
original_messages = [m.copy() for m in session.messages]
|
|
||||||
|
|
||||||
for _ in range(5):
|
|
||||||
history = session.get_history(max_messages=10)
|
|
||||||
assert len(history) == 10
|
|
||||||
|
|
||||||
assert len(session.messages) == 40
|
|
||||||
for i, msg in enumerate(session.messages):
|
|
||||||
assert msg["content"] == original_messages[i]["content"]
|
|
||||||
|
|
||||||
def test_consolidation_only_updates_last_consolidated(self):
|
|
||||||
"""Test that consolidation only updates last_consolidated field."""
|
|
||||||
session = create_session_with_messages("test:field_only", 60)
|
|
||||||
|
|
||||||
original_messages = session.messages.copy()
|
|
||||||
original_key = session.key
|
|
||||||
original_metadata = session.metadata.copy()
|
|
||||||
|
|
||||||
session.last_consolidated = len(session.messages) - KEEP_COUNT
|
|
||||||
|
|
||||||
assert session.messages == original_messages
|
|
||||||
assert session.key == original_key
|
|
||||||
assert session.metadata == original_metadata
|
|
||||||
assert session.last_consolidated == 35
|
|
||||||
|
|
||||||
|
|
||||||
class TestSliceLogic:
|
|
||||||
"""Test the slice logic: messages[last_consolidated:-keep_count]."""
|
|
||||||
|
|
||||||
def test_slice_extracts_correct_range(self):
|
|
||||||
"""Test that slice extracts the correct message range."""
|
|
||||||
session = create_session_with_messages("test:slice", 60)
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, 0, KEEP_COUNT)
|
|
||||||
|
|
||||||
assert len(old_messages) == 35
|
|
||||||
assert_messages_content(old_messages, 0, 34)
|
|
||||||
|
|
||||||
remaining = session.messages[-KEEP_COUNT:]
|
|
||||||
assert len(remaining) == 25
|
|
||||||
assert_messages_content(remaining, 35, 59)
|
|
||||||
|
|
||||||
def test_slice_with_partial_consolidation(self):
|
|
||||||
"""Test slice when some messages already consolidated."""
|
|
||||||
session = create_session_with_messages("test:partial", 70)
|
|
||||||
|
|
||||||
last_consolidated = 30
|
|
||||||
old_messages = get_old_messages(session, last_consolidated, KEEP_COUNT)
|
|
||||||
|
|
||||||
assert len(old_messages) == 15
|
|
||||||
assert_messages_content(old_messages, 30, 44)
|
|
||||||
|
|
||||||
def test_slice_with_various_keep_counts(self):
|
|
||||||
"""Test slice behavior with different keep_count values."""
|
|
||||||
session = create_session_with_messages("test:keep_counts", 50)
|
|
||||||
|
|
||||||
test_cases = [(10, 40), (20, 30), (30, 20), (40, 10)]
|
|
||||||
|
|
||||||
for keep_count, expected_count in test_cases:
|
|
||||||
old_messages = session.messages[0:-keep_count]
|
|
||||||
assert len(old_messages) == expected_count
|
|
||||||
|
|
||||||
def test_slice_when_keep_count_exceeds_messages(self):
|
|
||||||
"""Test slice when keep_count > len(messages)."""
|
|
||||||
session = create_session_with_messages("test:exceed", 10)
|
|
||||||
|
|
||||||
old_messages = session.messages[0:-20]
|
|
||||||
assert len(old_messages) == 0
|
|
||||||
|
|
||||||
|
|
||||||
class TestEmptyAndBoundarySessions:
|
|
||||||
"""Test empty sessions and boundary conditions."""
|
|
||||||
|
|
||||||
def test_empty_session_consolidation(self):
|
|
||||||
"""Test consolidation behavior with empty session."""
|
|
||||||
session = Session(key="test:empty")
|
|
||||||
|
|
||||||
assert len(session.messages) == 0
|
|
||||||
assert session.last_consolidated == 0
|
|
||||||
|
|
||||||
messages_to_process = len(session.messages) - session.last_consolidated
|
|
||||||
assert messages_to_process == 0
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
|
|
||||||
assert len(old_messages) == 0
|
|
||||||
|
|
||||||
def test_single_message_session(self):
|
|
||||||
"""Test consolidation with single message."""
|
|
||||||
session = Session(key="test:single")
|
|
||||||
session.add_message("user", "only message")
|
|
||||||
|
|
||||||
assert len(session.messages) == 1
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
|
|
||||||
assert len(old_messages) == 0
|
|
||||||
|
|
||||||
def test_exactly_keep_count_messages(self):
|
|
||||||
"""Test session with exactly keep_count messages."""
|
|
||||||
session = create_session_with_messages("test:exact", KEEP_COUNT)
|
|
||||||
|
|
||||||
assert len(session.messages) == KEEP_COUNT
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
|
|
||||||
assert len(old_messages) == 0
|
|
||||||
|
|
||||||
def test_just_over_keep_count(self):
|
|
||||||
"""Test session with one message over keep_count."""
|
|
||||||
session = create_session_with_messages("test:over", KEEP_COUNT + 1)
|
|
||||||
|
|
||||||
assert len(session.messages) == 26
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
|
|
||||||
assert len(old_messages) == 1
|
|
||||||
assert old_messages[0]["content"] == "msg0"
|
|
||||||
|
|
||||||
def test_very_large_session(self):
|
|
||||||
"""Test consolidation with very large message count."""
|
|
||||||
session = create_session_with_messages("test:large", 1000)
|
|
||||||
|
|
||||||
assert len(session.messages) == 1000
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
|
|
||||||
assert len(old_messages) == 975
|
|
||||||
assert_messages_content(old_messages, 0, 974)
|
|
||||||
|
|
||||||
remaining = session.messages[-KEEP_COUNT:]
|
|
||||||
assert len(remaining) == 25
|
|
||||||
assert_messages_content(remaining, 975, 999)
|
|
||||||
|
|
||||||
def test_session_with_gaps_in_consolidation(self):
|
|
||||||
"""Test session with potential gaps in consolidation history."""
|
|
||||||
session = create_session_with_messages("test:gaps", 50)
|
|
||||||
session.last_consolidated = 10
|
|
||||||
|
|
||||||
# Add more messages
|
|
||||||
for i in range(50, 60):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
|
|
||||||
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
|
|
||||||
|
|
||||||
expected_count = 60 - KEEP_COUNT - 10
|
|
||||||
assert len(old_messages) == expected_count
|
|
||||||
assert_messages_content(old_messages, 10, 34)
|
|
||||||
|
|
||||||
|
|
||||||
class TestConsolidationDeduplicationGuard:
|
|
||||||
"""Test that consolidation tasks are deduplicated and serialized."""
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_consolidation_guard_prevents_duplicate_tasks(self, tmp_path: Path) -> None:
|
|
||||||
"""Concurrent messages above memory_window spawn only one consolidation task."""
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.providers.base import LLMResponse
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
loop = AgentLoop(
|
|
||||||
bus=bus, provider=provider, workspace=tmp_path, model="test-model", memory_window=10
|
|
||||||
)
|
|
||||||
|
|
||||||
loop.provider.chat = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
|
|
||||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
for i in range(15):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
session.add_message("assistant", f"resp{i}")
|
|
||||||
loop.sessions.save(session)
|
|
||||||
|
|
||||||
consolidation_calls = 0
|
|
||||||
|
|
||||||
async def _fake_consolidate(_session, archive_all: bool = False) -> None:
|
|
||||||
nonlocal consolidation_calls
|
|
||||||
consolidation_calls += 1
|
|
||||||
await asyncio.sleep(0.05)
|
|
||||||
|
|
||||||
loop._consolidate_memory = _fake_consolidate # type: ignore[method-assign]
|
|
||||||
|
|
||||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="hello")
|
|
||||||
await loop._process_message(msg)
|
|
||||||
await loop._process_message(msg)
|
|
||||||
await asyncio.sleep(0.1)
|
|
||||||
|
|
||||||
assert consolidation_calls == 1, (
|
|
||||||
f"Expected exactly 1 consolidation, got {consolidation_calls}"
|
|
||||||
)
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_new_command_guard_prevents_concurrent_consolidation(
|
|
||||||
self, tmp_path: Path
|
|
||||||
) -> None:
|
|
||||||
"""/new command does not run consolidation concurrently with in-flight consolidation."""
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.providers.base import LLMResponse
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
loop = AgentLoop(
|
|
||||||
bus=bus, provider=provider, workspace=tmp_path, model="test-model", memory_window=10
|
|
||||||
)
|
|
||||||
|
|
||||||
loop.provider.chat = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
|
|
||||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
for i in range(15):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
session.add_message("assistant", f"resp{i}")
|
|
||||||
loop.sessions.save(session)
|
|
||||||
|
|
||||||
consolidation_calls = 0
|
|
||||||
active = 0
|
|
||||||
max_active = 0
|
|
||||||
|
|
||||||
async def _fake_consolidate(_session, archive_all: bool = False) -> None:
|
|
||||||
nonlocal consolidation_calls, active, max_active
|
|
||||||
consolidation_calls += 1
|
|
||||||
active += 1
|
|
||||||
max_active = max(max_active, active)
|
|
||||||
await asyncio.sleep(0.05)
|
|
||||||
active -= 1
|
|
||||||
|
|
||||||
loop._consolidate_memory = _fake_consolidate # type: ignore[method-assign]
|
|
||||||
|
|
||||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="hello")
|
|
||||||
await loop._process_message(msg)
|
|
||||||
|
|
||||||
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
|
||||||
await loop._process_message(new_msg)
|
|
||||||
await asyncio.sleep(0.1)
|
|
||||||
|
|
||||||
assert consolidation_calls == 2, (
|
|
||||||
f"Expected normal + /new consolidations, got {consolidation_calls}"
|
|
||||||
)
|
|
||||||
assert max_active == 1, (
|
|
||||||
f"Expected serialized consolidation, observed concurrency={max_active}"
|
|
||||||
)
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_consolidation_tasks_are_referenced(self, tmp_path: Path) -> None:
|
|
||||||
"""create_task results are tracked in _consolidation_tasks while in flight."""
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.providers.base import LLMResponse
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
loop = AgentLoop(
|
|
||||||
bus=bus, provider=provider, workspace=tmp_path, model="test-model", memory_window=10
|
|
||||||
)
|
|
||||||
|
|
||||||
loop.provider.chat = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
|
|
||||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
for i in range(15):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
session.add_message("assistant", f"resp{i}")
|
|
||||||
loop.sessions.save(session)
|
|
||||||
|
|
||||||
started = asyncio.Event()
|
|
||||||
|
|
||||||
async def _slow_consolidate(_session, archive_all: bool = False) -> None:
|
|
||||||
started.set()
|
|
||||||
await asyncio.sleep(0.1)
|
|
||||||
|
|
||||||
loop._consolidate_memory = _slow_consolidate # type: ignore[method-assign]
|
|
||||||
|
|
||||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="hello")
|
|
||||||
await loop._process_message(msg)
|
|
||||||
|
|
||||||
await started.wait()
|
|
||||||
assert len(loop._consolidation_tasks) == 1, "Task must be referenced while in-flight"
|
|
||||||
|
|
||||||
await asyncio.sleep(0.15)
|
|
||||||
assert len(loop._consolidation_tasks) == 0, (
|
|
||||||
"Task reference must be removed after completion"
|
|
||||||
)
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_new_waits_for_inflight_consolidation_and_preserves_messages(
|
|
||||||
self, tmp_path: Path
|
|
||||||
) -> None:
|
|
||||||
"""/new waits for in-flight consolidation and archives before clear."""
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.providers.base import LLMResponse
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
loop = AgentLoop(
|
|
||||||
bus=bus, provider=provider, workspace=tmp_path, model="test-model", memory_window=10
|
|
||||||
)
|
|
||||||
|
|
||||||
loop.provider.chat = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
|
|
||||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
for i in range(15):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
session.add_message("assistant", f"resp{i}")
|
|
||||||
loop.sessions.save(session)
|
|
||||||
|
|
||||||
started = asyncio.Event()
|
|
||||||
release = asyncio.Event()
|
|
||||||
archived_count = 0
|
|
||||||
|
|
||||||
async def _fake_consolidate(sess, archive_all: bool = False) -> bool:
|
|
||||||
nonlocal archived_count
|
|
||||||
if archive_all:
|
|
||||||
archived_count = len(sess.messages)
|
|
||||||
return True
|
|
||||||
started.set()
|
|
||||||
await release.wait()
|
|
||||||
return True
|
|
||||||
|
|
||||||
loop._consolidate_memory = _fake_consolidate # type: ignore[method-assign]
|
|
||||||
|
|
||||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="hello")
|
|
||||||
await loop._process_message(msg)
|
|
||||||
await started.wait()
|
|
||||||
|
|
||||||
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
|
||||||
pending_new = asyncio.create_task(loop._process_message(new_msg))
|
|
||||||
|
|
||||||
await asyncio.sleep(0.02)
|
|
||||||
assert not pending_new.done(), "/new should wait while consolidation is in-flight"
|
|
||||||
|
|
||||||
release.set()
|
|
||||||
response = await pending_new
|
|
||||||
assert response is not None
|
|
||||||
assert "new session started" in response.content.lower()
|
|
||||||
assert archived_count > 0, "Expected /new archival to process a non-empty snapshot"
|
|
||||||
|
|
||||||
session_after = loop.sessions.get_or_create("cli:test")
|
|
||||||
assert session_after.messages == [], "Session should be cleared after successful archival"
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_new_does_not_clear_session_when_archive_fails(self, tmp_path: Path) -> None:
|
|
||||||
"""/new must keep session data if archive step reports failure."""
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.providers.base import LLMResponse
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
loop = AgentLoop(
|
|
||||||
bus=bus, provider=provider, workspace=tmp_path, model="test-model", memory_window=10
|
|
||||||
)
|
|
||||||
|
|
||||||
loop.provider.chat = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
|
|
||||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
for i in range(5):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
session.add_message("assistant", f"resp{i}")
|
|
||||||
loop.sessions.save(session)
|
|
||||||
before_count = len(session.messages)
|
|
||||||
|
|
||||||
async def _failing_consolidate(sess, archive_all: bool = False) -> bool:
|
|
||||||
if archive_all:
|
|
||||||
return False
|
|
||||||
return True
|
|
||||||
|
|
||||||
loop._consolidate_memory = _failing_consolidate # type: ignore[method-assign]
|
|
||||||
|
|
||||||
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
|
||||||
response = await loop._process_message(new_msg)
|
|
||||||
|
|
||||||
assert response is not None
|
|
||||||
assert "failed" in response.content.lower()
|
|
||||||
session_after = loop.sessions.get_or_create("cli:test")
|
|
||||||
assert len(session_after.messages) == before_count, (
|
|
||||||
"Session must remain intact when /new archival fails"
|
|
||||||
)
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_new_archives_only_unconsolidated_messages_after_inflight_task(
|
|
||||||
self, tmp_path: Path
|
|
||||||
) -> None:
|
|
||||||
"""/new should archive only messages not yet consolidated by prior task."""
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.providers.base import LLMResponse
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
loop = AgentLoop(
|
|
||||||
bus=bus, provider=provider, workspace=tmp_path, model="test-model", memory_window=10
|
|
||||||
)
|
|
||||||
|
|
||||||
loop.provider.chat = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
|
|
||||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
for i in range(15):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
session.add_message("assistant", f"resp{i}")
|
|
||||||
loop.sessions.save(session)
|
|
||||||
|
|
||||||
started = asyncio.Event()
|
|
||||||
release = asyncio.Event()
|
|
||||||
archived_count = -1
|
|
||||||
|
|
||||||
async def _fake_consolidate(sess, archive_all: bool = False) -> bool:
|
|
||||||
nonlocal archived_count
|
|
||||||
if archive_all:
|
|
||||||
archived_count = len(sess.messages)
|
|
||||||
return True
|
|
||||||
|
|
||||||
started.set()
|
|
||||||
await release.wait()
|
|
||||||
sess.last_consolidated = len(sess.messages) - 3
|
|
||||||
return True
|
|
||||||
|
|
||||||
loop._consolidate_memory = _fake_consolidate # type: ignore[method-assign]
|
|
||||||
|
|
||||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="hello")
|
|
||||||
await loop._process_message(msg)
|
|
||||||
await started.wait()
|
|
||||||
|
|
||||||
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
|
||||||
pending_new = asyncio.create_task(loop._process_message(new_msg))
|
|
||||||
await asyncio.sleep(0.02)
|
|
||||||
assert not pending_new.done()
|
|
||||||
|
|
||||||
release.set()
|
|
||||||
response = await pending_new
|
|
||||||
|
|
||||||
assert response is not None
|
|
||||||
assert "new session started" in response.content.lower()
|
|
||||||
assert archived_count == 3, (
|
|
||||||
f"Expected only unconsolidated tail to archive, got {archived_count}"
|
|
||||||
)
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_new_cleans_up_consolidation_lock_for_invalidated_session(
|
|
||||||
self, tmp_path: Path
|
|
||||||
) -> None:
|
|
||||||
"""/new should remove lock entry for fully invalidated session key."""
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.providers.base import LLMResponse
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
loop = AgentLoop(
|
|
||||||
bus=bus, provider=provider, workspace=tmp_path, model="test-model", memory_window=10
|
|
||||||
)
|
|
||||||
|
|
||||||
loop.provider.chat = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
|
|
||||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
for i in range(3):
|
|
||||||
session.add_message("user", f"msg{i}")
|
|
||||||
session.add_message("assistant", f"resp{i}")
|
|
||||||
loop.sessions.save(session)
|
|
||||||
|
|
||||||
# Ensure lock exists before /new.
|
|
||||||
loop._consolidation_locks.setdefault(session.key, asyncio.Lock())
|
|
||||||
assert session.key in loop._consolidation_locks
|
|
||||||
|
|
||||||
async def _ok_consolidate(sess, archive_all: bool = False) -> bool:
|
|
||||||
return True
|
|
||||||
|
|
||||||
loop._consolidate_memory = _ok_consolidate # type: ignore[method-assign]
|
|
||||||
|
|
||||||
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
|
||||||
response = await loop._process_message(new_msg)
|
|
||||||
|
|
||||||
assert response is not None
|
|
||||||
assert "new session started" in response.content.lower()
|
|
||||||
assert session.key not in loop._consolidation_locks
|
|
||||||
@@ -40,7 +40,7 @@ def test_system_prompt_stays_stable_when_clock_changes(tmp_path, monkeypatch) ->
|
|||||||
|
|
||||||
|
|
||||||
def test_runtime_context_is_separate_untrusted_user_message(tmp_path) -> None:
|
def test_runtime_context_is_separate_untrusted_user_message(tmp_path) -> None:
|
||||||
"""Runtime metadata should be a separate user message before the actual user message."""
|
"""Runtime metadata should be included in the system prompt."""
|
||||||
workspace = _make_workspace(tmp_path)
|
workspace = _make_workspace(tmp_path)
|
||||||
builder = ContextBuilder(workspace)
|
builder = ContextBuilder(workspace)
|
||||||
|
|
||||||
@@ -51,16 +51,12 @@ def test_runtime_context_is_separate_untrusted_user_message(tmp_path) -> None:
|
|||||||
chat_id="direct",
|
chat_id="direct",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Runtime context should be in the system prompt
|
||||||
assert messages[0]["role"] == "system"
|
assert messages[0]["role"] == "system"
|
||||||
assert "## Current Session" not in messages[0]["content"]
|
assert "## Current Session" in messages[0]["content"]
|
||||||
|
assert "Channel: cli" in messages[0]["content"]
|
||||||
assert messages[-2]["role"] == "user"
|
assert "Chat ID: direct" in messages[0]["content"]
|
||||||
runtime_content = messages[-2]["content"]
|
|
||||||
assert isinstance(runtime_content, str)
|
|
||||||
assert ContextBuilder._RUNTIME_CONTEXT_TAG in runtime_content
|
|
||||||
assert "Current Time:" in runtime_content
|
|
||||||
assert "Channel: cli" in runtime_content
|
|
||||||
assert "Chat ID: direct" in runtime_content
|
|
||||||
|
|
||||||
|
# The actual user message should be the last message
|
||||||
assert messages[-1]["role"] == "user"
|
assert messages[-1]["role"] == "user"
|
||||||
assert messages[-1]["content"] == "Return exactly: OK"
|
assert messages[-1]["content"] == "Return exactly: OK"
|
||||||
|
|||||||
@@ -113,4 +113,4 @@ def test_edit_tool_to_params():
|
|||||||
params = tool.to_params()
|
params = tool.to_params()
|
||||||
|
|
||||||
assert params["type"] == "text_editor_20250728"
|
assert params["type"] == "text_editor_20250728"
|
||||||
assert params["name"] == "str_replace_editor"
|
assert params["name"] == "str_replace_based_edit_tool"
|
||||||
|
|||||||
@@ -3,27 +3,12 @@ import asyncio
|
|||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from nanobot.heartbeat.service import HeartbeatService
|
from nanobot.heartbeat.service import HeartbeatService
|
||||||
from nanobot.providers.base import LLMResponse, ToolCallRequest
|
|
||||||
|
|
||||||
|
|
||||||
class DummyProvider:
|
|
||||||
def __init__(self, responses: list[LLMResponse]):
|
|
||||||
self._responses = list(responses)
|
|
||||||
|
|
||||||
async def chat(self, *args, **kwargs) -> LLMResponse:
|
|
||||||
if self._responses:
|
|
||||||
return self._responses.pop(0)
|
|
||||||
return LLMResponse(content="", tool_calls=[])
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_start_is_idempotent(tmp_path) -> None:
|
async def test_start_is_idempotent(tmp_path) -> None:
|
||||||
provider = DummyProvider([])
|
|
||||||
|
|
||||||
service = HeartbeatService(
|
service = HeartbeatService(
|
||||||
workspace=tmp_path,
|
workspace=tmp_path,
|
||||||
provider=provider,
|
|
||||||
model="openai/gpt-4o-mini",
|
|
||||||
interval_s=9999,
|
interval_s=9999,
|
||||||
enabled=True,
|
enabled=True,
|
||||||
)
|
)
|
||||||
@@ -38,80 +23,36 @@ async def test_start_is_idempotent(tmp_path) -> None:
|
|||||||
await asyncio.sleep(0)
|
await asyncio.sleep(0)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_decide_returns_skip_when_no_tool_call(tmp_path) -> None:
|
|
||||||
provider = DummyProvider([LLMResponse(content="no tool call", tool_calls=[])])
|
|
||||||
service = HeartbeatService(
|
|
||||||
workspace=tmp_path,
|
|
||||||
provider=provider,
|
|
||||||
model="openai/gpt-4o-mini",
|
|
||||||
)
|
|
||||||
|
|
||||||
action, tasks = await service._decide("heartbeat content")
|
|
||||||
assert action == "skip"
|
|
||||||
assert tasks == ""
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_trigger_now_executes_when_decision_is_run(tmp_path) -> None:
|
async def test_trigger_now_executes_when_decision_is_run(tmp_path) -> None:
|
||||||
(tmp_path / "HEARTBEAT.md").write_text("- [ ] do thing", encoding="utf-8")
|
(tmp_path / "HEARTBEAT.md").write_text("- [ ] do thing", encoding="utf-8")
|
||||||
|
|
||||||
provider = DummyProvider([
|
called_with: list[tuple[str, dict | None]] = []
|
||||||
LLMResponse(
|
|
||||||
content="",
|
|
||||||
tool_calls=[
|
|
||||||
ToolCallRequest(
|
|
||||||
id="hb_1",
|
|
||||||
name="heartbeat",
|
|
||||||
arguments={"action": "run", "tasks": "check open tasks"},
|
|
||||||
)
|
|
||||||
],
|
|
||||||
)
|
|
||||||
])
|
|
||||||
|
|
||||||
called_with: list[str] = []
|
async def _on_heartbeat(prompt: str, metadata: dict | None = None) -> str:
|
||||||
|
called_with.append((prompt, metadata))
|
||||||
async def _on_execute(tasks: str) -> str:
|
|
||||||
called_with.append(tasks)
|
|
||||||
return "done"
|
return "done"
|
||||||
|
|
||||||
service = HeartbeatService(
|
service = HeartbeatService(
|
||||||
workspace=tmp_path,
|
workspace=tmp_path,
|
||||||
provider=provider,
|
on_heartbeat=_on_heartbeat,
|
||||||
model="openai/gpt-4o-mini",
|
|
||||||
on_execute=_on_execute,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
result = await service.trigger_now()
|
result = await service.trigger_now()
|
||||||
assert result == "done"
|
assert result == "done"
|
||||||
assert called_with == ["check open tasks"]
|
assert len(called_with) == 1
|
||||||
|
prompt, metadata = called_with[0]
|
||||||
|
assert "HEARTBEAT.md" in prompt
|
||||||
|
assert metadata == {"suppress_output": True}
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_trigger_now_returns_none_when_decision_is_skip(tmp_path) -> None:
|
async def test_trigger_now_returns_none_when_no_callback(tmp_path) -> None:
|
||||||
(tmp_path / "HEARTBEAT.md").write_text("- [ ] do thing", encoding="utf-8")
|
(tmp_path / "HEARTBEAT.md").write_text("- [ ] do thing", encoding="utf-8")
|
||||||
|
|
||||||
provider = DummyProvider([
|
|
||||||
LLMResponse(
|
|
||||||
content="",
|
|
||||||
tool_calls=[
|
|
||||||
ToolCallRequest(
|
|
||||||
id="hb_1",
|
|
||||||
name="heartbeat",
|
|
||||||
arguments={"action": "skip"},
|
|
||||||
)
|
|
||||||
],
|
|
||||||
)
|
|
||||||
])
|
|
||||||
|
|
||||||
async def _on_execute(tasks: str) -> str:
|
|
||||||
return tasks
|
|
||||||
|
|
||||||
service = HeartbeatService(
|
service = HeartbeatService(
|
||||||
workspace=tmp_path,
|
workspace=tmp_path,
|
||||||
provider=provider,
|
on_heartbeat=None, # No callback
|
||||||
model="openai/gpt-4o-mini",
|
|
||||||
on_execute=_on_execute,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
assert await service.trigger_now() is None
|
assert await service.trigger_now() is None
|
||||||
|
|||||||
@@ -676,7 +676,7 @@ async def test_on_media_message_respects_declared_size_limit(
|
|||||||
assert client.download_calls == []
|
assert client.download_calls == []
|
||||||
assert len(handled) == 1
|
assert len(handled) == 1
|
||||||
assert handled[0]["media"] == []
|
assert handled[0]["media"] == []
|
||||||
assert handled[0]["metadata"]["attachments"] == []
|
assert handled[0]["metadata"].get("attachments", []) == []
|
||||||
assert "[attachment: large.bin - too large]" in handled[0]["content"]
|
assert "[attachment: large.bin - too large]" in handled[0]["content"]
|
||||||
|
|
||||||
|
|
||||||
@@ -712,7 +712,7 @@ async def test_on_media_message_uses_server_limit_when_smaller_than_local_limit(
|
|||||||
assert client.download_calls == []
|
assert client.download_calls == []
|
||||||
assert len(handled) == 1
|
assert len(handled) == 1
|
||||||
assert handled[0]["media"] == []
|
assert handled[0]["media"] == []
|
||||||
assert handled[0]["metadata"]["attachments"] == []
|
assert handled[0]["metadata"].get("attachments", []) == []
|
||||||
assert "[attachment: large.bin - too large]" in handled[0]["content"]
|
assert "[attachment: large.bin - too large]" in handled[0]["content"]
|
||||||
|
|
||||||
|
|
||||||
@@ -746,7 +746,7 @@ async def test_on_media_message_handles_download_error(monkeypatch, tmp_path) ->
|
|||||||
assert len(client.download_calls) == 1
|
assert len(client.download_calls) == 1
|
||||||
assert len(handled) == 1
|
assert len(handled) == 1
|
||||||
assert handled[0]["media"] == []
|
assert handled[0]["media"] == []
|
||||||
assert handled[0]["metadata"]["attachments"] == []
|
assert handled[0]["metadata"].get("attachments", []) == []
|
||||||
assert "[attachment: photo.png - download failed]" in handled[0]["content"]
|
assert "[attachment: photo.png - download failed]" in handled[0]["content"]
|
||||||
|
|
||||||
|
|
||||||
@@ -830,7 +830,7 @@ async def test_on_media_message_handles_decrypt_error(monkeypatch, tmp_path) ->
|
|||||||
|
|
||||||
assert len(handled) == 1
|
assert len(handled) == 1
|
||||||
assert handled[0]["media"] == []
|
assert handled[0]["media"] == []
|
||||||
assert handled[0]["metadata"]["attachments"] == []
|
assert handled[0]["metadata"].get("attachments", []) == []
|
||||||
assert "[attachment: secret.txt - download failed]" in handled[0]["content"]
|
assert "[attachment: secret.txt - download failed]" in handled[0]["content"]
|
||||||
|
|
||||||
|
|
||||||
@@ -972,7 +972,6 @@ async def test_send_passes_thread_relates_to_to_attachment_upload(monkeypatch) -
|
|||||||
captured: dict[str, object] = {}
|
captured: dict[str, object] = {}
|
||||||
|
|
||||||
async def _fake_upload_and_send_attachment(
|
async def _fake_upload_and_send_attachment(
|
||||||
*,
|
|
||||||
room_id: str,
|
room_id: str,
|
||||||
path: Path,
|
path: Path,
|
||||||
limit_bytes: int,
|
limit_bytes: int,
|
||||||
|
|||||||
@@ -43,15 +43,12 @@ def test_native_tools_registered(mock_provider, mock_bus, tmp_path):
|
|||||||
|
|
||||||
# Verify native tools are registered (using their internal names)
|
# Verify native tools are registered (using their internal names)
|
||||||
assert "bash" in tool_names, "bash tool should be registered"
|
assert "bash" in tool_names, "bash tool should be registered"
|
||||||
assert "str_replace_editor" in tool_names, "str_replace_editor tool should be registered"
|
assert "str_replace_based_edit_tool" in tool_names, "str_replace_based_edit_tool tool should be registered"
|
||||||
assert "computer" in tool_names, "computer tool should be registered"
|
# Note: computer tool is intentionally disabled by default (requires VNC setup)
|
||||||
|
|
||||||
# Verify we can get the tool instances
|
# Verify we can get the tool instances
|
||||||
bash_tool = loop.tools.get("bash")
|
bash_tool = loop.tools.get("bash")
|
||||||
assert isinstance(bash_tool, BashTool20250124)
|
assert isinstance(bash_tool, BashTool20250124)
|
||||||
|
|
||||||
editor_tool = loop.tools.get("str_replace_editor")
|
editor_tool = loop.tools.get("str_replace_based_edit_tool")
|
||||||
assert isinstance(editor_tool, EditTool20250728)
|
assert isinstance(editor_tool, EditTool20250728)
|
||||||
|
|
||||||
computer_tool = loop.tools.get("computer")
|
|
||||||
assert isinstance(computer_tool, ComputerTool20251124)
|
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ def test_get_auth_headers_oauth():
|
|||||||
assert "Authorization" in headers
|
assert "Authorization" in headers
|
||||||
assert headers["Authorization"] == "Bearer sk-ant-oat01-xxx"
|
assert headers["Authorization"] == "Bearer sk-ant-oat01-xxx"
|
||||||
assert "x-api-key" not in headers
|
assert "x-api-key" not in headers
|
||||||
assert headers["anthropic-beta"] == "claude-code-20250219,oauth-2025-04-20"
|
assert headers["anthropic-beta"] == "claude-code-20250219,oauth-2025-04-20,context-management-2025-06-27"
|
||||||
|
|
||||||
|
|
||||||
def test_get_auth_headers_api_key():
|
def test_get_auth_headers_api_key():
|
||||||
|
|||||||
@@ -31,7 +31,7 @@ async def test_registry_executes_edit_tool():
|
|||||||
with tempfile.TemporaryDirectory() as tmpdir:
|
with tempfile.TemporaryDirectory() as tmpdir:
|
||||||
test_file = str(Path(tmpdir) / "test.txt")
|
test_file = str(Path(tmpdir) / "test.txt")
|
||||||
|
|
||||||
result = await registry.execute("str_replace_editor", {
|
result = await registry.execute("str_replace_based_edit_tool", {
|
||||||
"command": "create",
|
"command": "create",
|
||||||
"path": test_file,
|
"path": test_file,
|
||||||
"file_text": "Hello, world!"
|
"file_text": "Hello, world!"
|
||||||
|
|||||||
@@ -0,0 +1,139 @@
|
|||||||
|
# tests/test_subagent_wait.py
|
||||||
|
"""Tests for wait_for_subagents with top-level and child subagents."""
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
from nanobot.agent.subagent import SubagentManager
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.providers.base import LLMProvider, LLMResponse
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_wait_for_top_level_subagent():
|
||||||
|
"""Test that wait_for works for top-level subagents spawned from telegram."""
|
||||||
|
bus = MessageBus()
|
||||||
|
provider = MagicMock(spec=LLMProvider)
|
||||||
|
provider.chat = AsyncMock(return_value=LLMResponse(
|
||||||
|
content="Task completed",
|
||||||
|
tool_calls=[]
|
||||||
|
))
|
||||||
|
provider.get_default_model = MagicMock(return_value="test-model")
|
||||||
|
provider.thinking_budget = 0
|
||||||
|
|
||||||
|
workspace = Path("/tmp/test-subagent")
|
||||||
|
workspace.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
manager = SubagentManager(
|
||||||
|
bus=bus,
|
||||||
|
provider=provider,
|
||||||
|
workspace=workspace
|
||||||
|
)
|
||||||
|
|
||||||
|
# Spawn a top-level subagent (origin channel = "telegram")
|
||||||
|
task_id = await manager.spawn(
|
||||||
|
task="Test task",
|
||||||
|
label="test",
|
||||||
|
model=None,
|
||||||
|
origin_channel="telegram",
|
||||||
|
origin_chat_id="12345"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Wait for it to complete
|
||||||
|
result = await manager.wait_for([task_id])
|
||||||
|
|
||||||
|
# Should find the result (not "No result found")
|
||||||
|
assert "No result found" not in result
|
||||||
|
assert task_id in result
|
||||||
|
assert "Task completed" in result or "completed" in result.lower()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_wait_for_child_subagent():
|
||||||
|
"""Test that wait_for works for child subagents (orchestrator pattern)."""
|
||||||
|
bus = MessageBus()
|
||||||
|
provider = MagicMock(spec=LLMProvider)
|
||||||
|
provider.chat = AsyncMock(return_value=LLMResponse(
|
||||||
|
content="Child task completed",
|
||||||
|
tool_calls=[]
|
||||||
|
))
|
||||||
|
provider.get_default_model = MagicMock(return_value="test-model")
|
||||||
|
provider.thinking_budget = 0
|
||||||
|
|
||||||
|
workspace = Path("/tmp/test-subagent")
|
||||||
|
workspace.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
manager = SubagentManager(
|
||||||
|
bus=bus,
|
||||||
|
provider=provider,
|
||||||
|
workspace=workspace
|
||||||
|
)
|
||||||
|
|
||||||
|
# Spawn a child subagent (origin channel = "subagent")
|
||||||
|
task_id = await manager.spawn(
|
||||||
|
task="Child test task",
|
||||||
|
label="test-child",
|
||||||
|
model=None,
|
||||||
|
origin_channel="subagent",
|
||||||
|
origin_chat_id="parent-id"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Wait for it to complete
|
||||||
|
result = await manager.wait_for([task_id])
|
||||||
|
|
||||||
|
# Should find the result (not "No result found")
|
||||||
|
assert "No result found" not in result
|
||||||
|
assert task_id in result
|
||||||
|
assert "Child task completed" in result or "completed" in result.lower()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_wait_for_multiple_subagents():
|
||||||
|
"""Test waiting for multiple subagents of different types."""
|
||||||
|
bus = MessageBus()
|
||||||
|
call_count = 0
|
||||||
|
|
||||||
|
async def chat_response(*args, **kwargs):
|
||||||
|
nonlocal call_count
|
||||||
|
call_count += 1
|
||||||
|
return LLMResponse(content=f"Task {call_count} completed", tool_calls=[])
|
||||||
|
|
||||||
|
provider = MagicMock(spec=LLMProvider)
|
||||||
|
provider.chat = AsyncMock(side_effect=chat_response)
|
||||||
|
provider.get_default_model = MagicMock(return_value="test-model")
|
||||||
|
provider.thinking_budget = 0
|
||||||
|
|
||||||
|
workspace = Path("/tmp/test-subagent")
|
||||||
|
workspace.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
manager = SubagentManager(
|
||||||
|
bus=bus,
|
||||||
|
provider=provider,
|
||||||
|
workspace=workspace
|
||||||
|
)
|
||||||
|
|
||||||
|
# Spawn one top-level and one child subagent
|
||||||
|
task_id_1 = await manager.spawn(
|
||||||
|
task="Top-level task",
|
||||||
|
label="test-top",
|
||||||
|
model=None,
|
||||||
|
origin_channel="telegram",
|
||||||
|
origin_chat_id="12345"
|
||||||
|
)
|
||||||
|
|
||||||
|
task_id_2 = await manager.spawn(
|
||||||
|
task="Child task",
|
||||||
|
label="test-child",
|
||||||
|
model=None,
|
||||||
|
origin_channel="subagent",
|
||||||
|
origin_chat_id="parent"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Wait for both
|
||||||
|
result = await manager.wait_for([task_id_1, task_id_2])
|
||||||
|
|
||||||
|
# Should find both results
|
||||||
|
assert "No result found" not in result
|
||||||
|
assert task_id_1 in result
|
||||||
|
assert task_id_2 in result
|
||||||
|
assert "Task 1 completed" in result or "completed" in result.lower()
|
||||||
|
assert "Task 2 completed" in result or "completed" in result.lower()
|
||||||
@@ -1,167 +0,0 @@
|
|||||||
"""Tests for /stop task cancellation."""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import asyncio
|
|
||||||
from unittest.mock import AsyncMock, MagicMock, patch
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
|
|
||||||
|
|
||||||
def _make_loop():
|
|
||||||
"""Create a minimal AgentLoop with mocked dependencies."""
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
workspace = MagicMock()
|
|
||||||
workspace.__truediv__ = MagicMock(return_value=MagicMock())
|
|
||||||
|
|
||||||
with patch("nanobot.agent.loop.ContextBuilder"), \
|
|
||||||
patch("nanobot.agent.loop.SessionManager"), \
|
|
||||||
patch("nanobot.agent.loop.SubagentManager") as MockSubMgr:
|
|
||||||
MockSubMgr.return_value.cancel_by_session = AsyncMock(return_value=0)
|
|
||||||
loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
|
|
||||||
return loop, bus
|
|
||||||
|
|
||||||
|
|
||||||
class TestHandleStop:
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_stop_no_active_task(self):
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
|
|
||||||
loop, bus = _make_loop()
|
|
||||||
msg = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="/stop")
|
|
||||||
await loop._handle_stop(msg)
|
|
||||||
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
|
||||||
assert "No active task" in out.content
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_stop_cancels_active_task(self):
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
|
|
||||||
loop, bus = _make_loop()
|
|
||||||
cancelled = asyncio.Event()
|
|
||||||
|
|
||||||
async def slow_task():
|
|
||||||
try:
|
|
||||||
await asyncio.sleep(60)
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
cancelled.set()
|
|
||||||
raise
|
|
||||||
|
|
||||||
task = asyncio.create_task(slow_task())
|
|
||||||
await asyncio.sleep(0)
|
|
||||||
loop._active_tasks["test:c1"] = [task]
|
|
||||||
|
|
||||||
msg = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="/stop")
|
|
||||||
await loop._handle_stop(msg)
|
|
||||||
|
|
||||||
assert cancelled.is_set()
|
|
||||||
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
|
||||||
assert "stopped" in out.content.lower()
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_stop_cancels_multiple_tasks(self):
|
|
||||||
from nanobot.bus.events import InboundMessage
|
|
||||||
|
|
||||||
loop, bus = _make_loop()
|
|
||||||
events = [asyncio.Event(), asyncio.Event()]
|
|
||||||
|
|
||||||
async def slow(idx):
|
|
||||||
try:
|
|
||||||
await asyncio.sleep(60)
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
events[idx].set()
|
|
||||||
raise
|
|
||||||
|
|
||||||
tasks = [asyncio.create_task(slow(i)) for i in range(2)]
|
|
||||||
await asyncio.sleep(0)
|
|
||||||
loop._active_tasks["test:c1"] = tasks
|
|
||||||
|
|
||||||
msg = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="/stop")
|
|
||||||
await loop._handle_stop(msg)
|
|
||||||
|
|
||||||
assert all(e.is_set() for e in events)
|
|
||||||
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
|
||||||
assert "2 task" in out.content
|
|
||||||
|
|
||||||
|
|
||||||
class TestDispatch:
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_dispatch_processes_and_publishes(self):
|
|
||||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
|
||||||
|
|
||||||
loop, bus = _make_loop()
|
|
||||||
msg = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="hello")
|
|
||||||
loop._process_message = AsyncMock(
|
|
||||||
return_value=OutboundMessage(channel="test", chat_id="c1", content="hi")
|
|
||||||
)
|
|
||||||
await loop._dispatch(msg)
|
|
||||||
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
|
||||||
assert out.content == "hi"
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_processing_lock_serializes(self):
|
|
||||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
|
||||||
|
|
||||||
loop, bus = _make_loop()
|
|
||||||
order = []
|
|
||||||
|
|
||||||
async def mock_process(m, **kwargs):
|
|
||||||
order.append(f"start-{m.content}")
|
|
||||||
await asyncio.sleep(0.05)
|
|
||||||
order.append(f"end-{m.content}")
|
|
||||||
return OutboundMessage(channel="test", chat_id="c1", content=m.content)
|
|
||||||
|
|
||||||
loop._process_message = mock_process
|
|
||||||
msg1 = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="a")
|
|
||||||
msg2 = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="b")
|
|
||||||
|
|
||||||
t1 = asyncio.create_task(loop._dispatch(msg1))
|
|
||||||
t2 = asyncio.create_task(loop._dispatch(msg2))
|
|
||||||
await asyncio.gather(t1, t2)
|
|
||||||
assert order == ["start-a", "end-a", "start-b", "end-b"]
|
|
||||||
|
|
||||||
|
|
||||||
class TestSubagentCancellation:
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_cancel_by_session(self):
|
|
||||||
from nanobot.agent.subagent import SubagentManager
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
mgr = SubagentManager(provider=provider, workspace=MagicMock(), bus=bus)
|
|
||||||
|
|
||||||
cancelled = asyncio.Event()
|
|
||||||
|
|
||||||
async def slow():
|
|
||||||
try:
|
|
||||||
await asyncio.sleep(60)
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
cancelled.set()
|
|
||||||
raise
|
|
||||||
|
|
||||||
task = asyncio.create_task(slow())
|
|
||||||
await asyncio.sleep(0)
|
|
||||||
mgr._running_tasks["sub-1"] = task
|
|
||||||
mgr._session_tasks["test:c1"] = {"sub-1"}
|
|
||||||
|
|
||||||
count = await mgr.cancel_by_session("test:c1")
|
|
||||||
assert count == 1
|
|
||||||
assert cancelled.is_set()
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_cancel_by_session_no_tasks(self):
|
|
||||||
from nanobot.agent.subagent import SubagentManager
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
|
|
||||||
bus = MessageBus()
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
mgr = SubagentManager(provider=provider, workspace=MagicMock(), bus=bus)
|
|
||||||
assert await mgr.cancel_by_session("nonexistent") == 0
|
|
||||||
Reference in New Issue
Block a user